Update token handling as proposed by CodeRabbit

app/llm/external/PipelineEnum.py

@@ -11,4 +11,5 @@ class PipelineEnum(str, Enum):
     IRIS_CITATION_PIPELINE = "IRIS_CITATION_PIPELINE"
     IRIS_RERANKER_PIPELINE = "IRIS_RERANKER_PIPELINE"
     IRIS_SUMMARY_PIPELINE = "IRIS_SUMMARY_PIPELINE"
+    IRIS_LECTURE_RETRIEVAL_PIPELINE = "IRIS_LECTURE_RETRIEVAL_PIPELINE"
     NOT_SET = "NOT_SET"

app/llm/external/ollama.py

@@ -114,10 +114,10 @@ def chat(
             options=self.options,
         )
         return convert_to_iris_message(
-            response["message"],
-            response["prompt_eval_count"],
-            response["eval_count"],
-            response["model"],
+            response.get("message"),
+            response.get("prompt_eval_count", 0),
+            response.get("eval_count", 0),
+            response.get("model", self.model),
         )
 
     def embed(self, text: str) -> list[float]:

app/pipeline/chat/code_feedback_pipeline.py

@@ -144,7 +144,7 @@ def __call__(
                 }
             )
         )
-        num_tokens = self.llm.tokens
-        num_tokens.pipeline = PipelineEnum.IRIS_CODE_FEEDBACK
-        self.tokens = num_tokens
+        token_usage = self.llm.tokens
+        token_usage.pipeline = PipelineEnum.IRIS_CODE_FEEDBACK
+        self.tokens = token_usage
         return response.replace("{", "{{").replace("}", "}}")

app/pipeline/chat/exercise_chat_pipeline.py

@@ -116,10 +116,14 @@ def __call__(self, dto: ExerciseChatPipelineExecutionDTO):
                     suggestion_dto.last_message = self.exercise_chat_response
                     suggestion_dto.problem_statement = dto.exercise.problem_statement
                     suggestions = self.suggestion_pipeline(suggestion_dto)
+                    if self.suggestion_pipeline.tokens is not None:
+                        tokens = [self.suggestion_pipeline.tokens]
+                    else:
+                        tokens = []
                     self.callback.done(
                         final_result=None,
                         suggestions=suggestions,
-                        tokens=[self.suggestion_pipeline.tokens],
+                        tokens=tokens,
                     )
                 else:
                     # This should never happen but whatever
@@ -264,9 +268,7 @@ def _run_exercise_chat_pipeline(
                 .with_config({"run_name": "Response Drafting"})
                 .invoke({})
             )
-            if self.llm.tokens is not None:
-                self.llm.tokens.pipeline = PipelineEnum.IRIS_CHAT_EXERCISE_MESSAGE
-                self.tokens.append(self.llm.tokens)
+            self._collect_llm_tokens()
             self.callback.done()
             self.prompt = ChatPromptTemplate.from_messages(
                 [
@@ -281,9 +283,7 @@ def _run_exercise_chat_pipeline(
                 .with_config({"run_name": "Response Refining"})
                 .invoke({})
             )
-            if self.llm.tokens is not None:
-                self.llm.tokens.pipeline = PipelineEnum.IRIS_CHAT_EXERCISE_MESSAGE
-                self.tokens.append(self.llm.tokens)
+            self._collect_llm_tokens()
 
             if "!ok!" in guide_response:
                 print("Response is ok and not rewritten!!!")
@@ -385,3 +385,8 @@ def should_execute_lecture_pipeline(self, course_id: int) -> bool:
             )
             return len(result.objects) > 0
         return False
+
+    def _collect_llm_tokens(self):
+        if self.llm.tokens is not None:
+            self.llm.tokens.pipeline = PipelineEnum.IRIS_CHAT_EXERCISE_MESSAGE
+            self.tokens.append(self.llm.tokens)

app/pipeline/chat/lecture_chat_pipeline.py

@@ -116,9 +116,9 @@ def __call__(self, dto: LectureChatPipelineExecutionDTO):
         self.prompt = ChatPromptTemplate.from_messages(prompt_val)
         try:
             response = (self.prompt | self.pipeline).invoke({})
-            num_tokens = self.llm.tokens
-            num_tokens.pipeline = PipelineEnum.IRIS_CHAT_LECTURE_MESSAGE
-            self.tokens.append(num_tokens)
+            token_usage = self.llm.tokens
+            token_usage.pipeline = PipelineEnum.IRIS_CHAT_LECTURE_MESSAGE
+            self.tokens.append(token_usage)
             response_with_citation = self.citation_pipeline(
                 retrieved_lecture_chunks, response
             )

app/pipeline/competency_extraction_pipeline.py

@@ -79,13 +79,13 @@ def __call__(
         response = self.request_handler.chat(
             [prompt], CompletionArguments(temperature=0.4)
         )
-        num_tokens = LLMTokenCount(
+        token_usage = LLMTokenCount(
             model_info=response.model_info,
             num_input_tokens=response.num_input_tokens,
             num_output_tokens=response.num_output_tokens,
             pipeline=PipelineEnum.IRIS_COMPETENCY_GENERATION,
         )
-        self.tokens.append(num_tokens)
+        self.tokens.append(token_usage)
         response = response.contents[0].text_content
 
         generated_competencies: list[Competency] = []

app/pipeline/lecture_ingestion_pipeline.py

@@ -6,7 +6,6 @@
 import fitz
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
-from sipbuild.generator.parser.tokens import tokens
 from unstructured.cleaners.core import clean
 from weaviate import WeaviateClient
 from weaviate.classes.query import Filter
@@ -280,9 +279,9 @@ def merge_page_content_and_image_interpretation(
             (prompt | self.pipeline).invoke({}), bullets=True, extra_whitespace=True
         )
         # TODO: send to artemis
-        num_tokens = self.llm.tokens
-        num_tokens.pipeline = PipelineEnum.IRIS_LECTURE_INGESTION
-        tokens.append(num_tokens)
+        token_usage = self.llm.tokens
+        token_usage.pipeline = PipelineEnum.IRIS_LECTURE_INGESTION
+        self.tokens.append(token_usage)
         return clean_output
 
     def get_course_language(self, page_content: str) -> str:

app/pipeline/pipeline.py

@@ -1,7 +1,7 @@
 from abc import ABCMeta
 from typing import List
 
-from app.llm.external import LLMTokenCount
+from app.llm.external.LLMTokenCount import LLMTokenCount
 
 
 class Pipeline(metaclass=ABCMeta):

app/pipeline/shared/reranker_pipeline.py

@@ -110,7 +110,7 @@ def __call__(
             prompt = self.default_prompt
 
         response = (prompt | self.pipeline).invoke(data)
-        num_tokens = self.llm.tokens
-        num_tokens.pipeline = PipelineEnum.IRIS_RERANKER_PIPELINE
-        self.tokens.append(num_tokens)
+        token_usage = self.llm.tokens
+        token_usage.pipeline = PipelineEnum.IRIS_RERANKER_PIPELINE
+        self.tokens.append(token_usage)
         return response.selected_paragraphs

app/retrieval/lecture_retrieval.py

@@ -7,6 +7,7 @@
 
 from ..common import convert_iris_message_to_langchain_message
 from ..llm.external.LLMTokenCount import LLMTokenCount
+from ..llm.external.PipelineEnum import PipelineEnum
 from ..llm.langchain import IrisLangchainChatModel
 from ..pipeline import Pipeline
 
@@ -82,7 +83,7 @@ class LectureRetrieval(Pipeline):
     Class for retrieving lecture data from the database.
     """
 
-    tokens: LLMTokenCount
+    tokens: [LLMTokenCount]
 
     def __init__(self, client: WeaviateClient, **kwargs):
         super().__init__(implementation_id="lecture_retrieval_pipeline")
@@ -101,6 +102,7 @@ def __init__(self, client: WeaviateClient, **kwargs):
         self.pipeline = self.llm | StrOutputParser()
         self.collection = init_lecture_schema(client)
         self.reranker_pipeline = RerankerPipeline()
+        self.tokens = []
 
     @traceable(name="Full Lecture Retrieval")
     def __call__(
@@ -239,7 +241,9 @@ def rewrite_student_query(
         prompt = ChatPromptTemplate.from_messages(prompt_val)
         try:
             response = (prompt | self.pipeline).invoke({})
-            self.tokens = self.llm.tokens
+            token_usage = self.llm.tokens
+            token_usage.pipeline = PipelineEnum.IRIS_LECTURE_RETRIEVAL_PIPELINE
+            self.tokens.append(self.llm.tokens)
             logger.info(f"Response from exercise chat pipeline: {response}")
             return response
         except Exception as e:
@@ -277,6 +281,9 @@ def rewrite_student_query_with_exercise_context(
         prompt = ChatPromptTemplate.from_messages(prompt_val)
         try:
             response = (prompt | self.pipeline).invoke({})
+            token_usage = self.llm.tokens
+            token_usage.pipeline = PipelineEnum.IRIS_LECTURE_RETRIEVAL_PIPELINE
+            self.tokens.append(self.llm.tokens)
             logger.info(f"Response from exercise chat pipeline: {response}")
             return response
         except Exception as e:
@@ -312,6 +319,9 @@ def rewrite_elaborated_query(
         prompt = ChatPromptTemplate.from_messages(prompt_val)
         try:
             response = (prompt | self.pipeline).invoke({})
+            token_usage = self.llm.tokens
+            token_usage.pipeline = PipelineEnum.IRIS_LECTURE_RETRIEVAL_PIPELINE
+            self.tokens.append(self.llm.tokens)
             logger.info(f"Response from retirval pipeline: {response}")
             return response
         except Exception as e:
@@ -351,6 +361,9 @@ def rewrite_elaborated_query_with_exercise_context(
         )
         try:
             response = (prompt | self.pipeline).invoke({})
+            token_usage = self.llm.tokens
+            token_usage.pipeline = PipelineEnum.IRIS_LECTURE_RETRIEVAL_PIPELINE
+            self.tokens.append(self.llm.tokens)
             logger.info(f"Response from exercise chat pipeline: {response}")
             return response
         except Exception as e: