Track token usage of iris requests

app/domain/data/token_usage_dto.py

@@ -0,0 +1,10 @@
+from pydantic import BaseModel
+
+from app.llm.external.PipelineEnum import PipelineEnum
+
+
+class TokenUsageDTO(BaseModel):
+    model_info: str
+    num_input_tokens: int
+    num_output_tokens: int
+    pipeline: PipelineEnum

app/domain/pyris_message.py

@@ -16,6 +16,10 @@ class IrisMessageRole(str, Enum):
 class PyrisMessage(BaseModel):
     model_config = ConfigDict(populate_by_name=True)
 
+    num_input_tokens: int = Field(alias="numInputTokens", default=0)
+    num_output_tokens: int = Field(alias="numOutputTokens", default=0)
+    model_info: str = Field(alias="modelInfo", default="")
+
     sent_at: datetime | None = Field(alias="sentAt", default=None)
     sender: IrisMessageRole
     contents: List[MessageContentDTO] = []

app/domain/status/status_update_dto.py

@@ -2,8 +2,10 @@
 
 from pydantic import BaseModel
 
+from ..data.token_usage_dto import TokenUsageDTO
 from ...domain.status.stage_dto import StageDTO
 
 
 class StatusUpdateDTO(BaseModel):
     stages: List[StageDTO]
+    tokens: List[TokenUsageDTO] = []

app/llm/external/LLMTokenCount.py

@@ -0,0 +1,24 @@
+from app.llm.external.PipelineEnum import PipelineEnum
+
+
+class LLMTokenCount:
+
+    model_info: str
+    num_input_tokens: int
+    num_output_tokens: int
+    pipeline: PipelineEnum
+
+    def __init__(
+        self,
+        model_info: str,
+        num_input_tokens: int,
+        num_output_tokens: int,
+        pipeline: PipelineEnum,
+    ):
+        self.model_info = model_info
+        self.num_input_tokens = num_input_tokens
+        self.num_output_tokens = num_output_tokens
+        self.pipeline = pipeline
+
+    def __str__(self):
+        return f"{self.model_info}: {self.num_input_tokens} in, {self.num_output_tokens} out, {self.pipeline} pipeline"

app/llm/external/PipelineEnum.py

@@ -0,0 +1,14 @@
+from enum import Enum
+
+
+class PipelineEnum(str, Enum):
+    IRIS_CODE_FEEDBACK = "IRIS_CODE_FEEDBACK"
+    IRIS_CHAT_COURSE_MESSAGE = "IRIS_CHAT_COURSE_MESSAGE"
+    IRIS_CHAT_EXERCISE_MESSAGE = "IRIS_CHAT_EXERCISE_MESSAGE"
+    IRIS_INTERACTION_SUGGESTION = "IRIS_INTERACTION_SUGGESTION"
+    IRIS_CHAT_LECTURE_MESSAGE = "IRIS_CHAT_LECTURE_MESSAGE"
+    IRIS_COMPETENCY_GENERATION = "IRIS_COMPETENCY_GENERATION"
+    IRIS_CITATION_PIPELINE = "IRIS_CITATION_PIPELINE"
+    IRIS_RERANKER_PIPELINE = "IRIS_RERANKER_PIPELINE"
+    IRIS_SUMMARY_PIPELINE = "IRIS_SUMMARY_PIPELINE"
+    NOT_SET = "NOT_SET"

app/llm/external/ollama.py

@@ -57,7 +57,9 @@ def convert_to_ollama_messages(messages: list[PyrisMessage]) -> list[Message]:
     return messages_to_return
 
 
-def convert_to_iris_message(message: Message) -> PyrisMessage:
+def convert_to_iris_message(
+    message: Message, num_input_tokens: int, num_output_tokens: int, model: str
+) -> PyrisMessage:
     """
     Convert a Message to a PyrisMessage
     """
@@ -66,6 +68,9 @@ def convert_to_iris_message(message: Message) -> PyrisMessage:
         sender=map_str_to_role(message["role"]),
         contents=contents,
         send_at=datetime.now(),
+        num_input_tokens=num_input_tokens,
+        num_output_tokens=num_output_tokens,
+        model_info=model,
     )
 
 
@@ -108,7 +113,12 @@ def chat(
             format="json" if arguments.response_format == "JSON" else "",
             options=self.options,
         )
-        return convert_to_iris_message(response["message"])
+        return convert_to_iris_message(
+            response["message"],
+            response["prompt_eval_count"],
+            response["eval_count"],
+            response["model"],
+        )
 
     def embed(self, text: str) -> list[float]:
         response = self._client.embeddings(

app/llm/external/openai_chat.py

@@ -2,10 +2,11 @@
 import time
 import traceback
 from datetime import datetime
-from typing import Literal, Any
+from typing import Literal, Any, Optional
 
 from openai import OpenAI
 from openai.lib.azure import AzureOpenAI
+from openai.types import CompletionUsage
 from openai.types.chat import ChatCompletionMessage, ChatCompletionMessageParam
 from openai.types.shared_params import ResponseFormatJSONObject
 
@@ -61,15 +62,24 @@ def convert_to_open_ai_messages(
     return openai_messages
 
 
-def convert_to_iris_message(message: ChatCompletionMessage) -> PyrisMessage:
+def convert_to_iris_message(
+    message: ChatCompletionMessage, usage: Optional[CompletionUsage], model: str
+) -> PyrisMessage:
     """
     Convert a ChatCompletionMessage to a PyrisMessage
     """
-    return PyrisMessage(
+    num_input_tokens = getattr(usage, "prompt_tokens", -1)
+    num_output_tokens = getattr(usage, "completion_tokens", -1)
+
+    message = PyrisMessage(
         sender=map_str_to_role(message.role),
         contents=[TextMessageContentDTO(textContent=message.content)],
         send_at=datetime.now(),
+        num_input_tokens=num_input_tokens,
+        num_output_tokens=num_output_tokens,
+        model_info=model,
     )
+    return message
 
 
 class OpenAIChatModel(ChatModel):
@@ -103,7 +113,9 @@ def chat(
                         temperature=arguments.temperature,
                         max_tokens=arguments.max_tokens,
                     )
-                return convert_to_iris_message(response.choices[0].message)
+                return convert_to_iris_message(
+                    response.choices[0].message, response.usage, response.model
+                )
             except Exception as e:
                 wait_time = initial_delay * (backoff_factor**attempt)
                 logging.warning(f"Exception on attempt {attempt + 1}: {e}")

app/llm/langchain/iris_langchain_chat_model.py

@@ -5,9 +5,10 @@
     BaseChatModel,
 )
 from langchain_core.messages import BaseMessage
-from langchain_core.outputs import ChatResult
-from langchain_core.outputs.chat_generation import ChatGeneration
+from langchain_core.outputs import ChatResult, ChatGeneration
 
+from ..external.LLMTokenCount import LLMTokenCount
+from ..external.PipelineEnum import PipelineEnum
 from ...common import (
     convert_iris_message_to_langchain_message,
     convert_langchain_message_to_iris_message,
@@ -20,6 +21,7 @@ class IrisLangchainChatModel(BaseChatModel):
 
     request_handler: RequestHandler
     completion_args: CompletionArguments
+    tokens: LLMTokenCount = None
 
     def __init__(
         self,
@@ -43,6 +45,12 @@ def _generate(
         iris_message = self.request_handler.chat(iris_messages, self.completion_args)
         base_message = convert_iris_message_to_langchain_message(iris_message)
         chat_generation = ChatGeneration(message=base_message)
+        self.tokens = LLMTokenCount(
+            model_info=iris_message.model_info,
+            num_input_tokens=iris_message.num_input_tokens,
+            num_output_tokens=iris_message.num_output_tokens,
+            pipeline=PipelineEnum.NOT_SET,
+        )
         return ChatResult(generations=[chat_generation])
 
     @property

app/pipeline/chat/code_feedback_pipeline.py

@@ -13,6 +13,8 @@
 from ...domain.data.feedback_dto import FeedbackDTO
 from ...llm import CapabilityRequestHandler, RequirementList
 from ...llm import CompletionArguments
+from ...llm.external.LLMTokenCount import LLMTokenCount
+from ...llm.external.PipelineEnum import PipelineEnum
 from ...llm.langchain import IrisLangchainChatModel
 from ...pipeline import Pipeline
 from ...web.status.status_update import StatusCallback
@@ -40,6 +42,7 @@ class CodeFeedbackPipeline(Pipeline):
     callback: StatusCallback
     default_prompt: PromptTemplate
     output_parser: StrOutputParser
+    tokens: LLMTokenCount
 
     def __init__(self, callback: Optional[StatusCallback] = None):
         super().__init__(implementation_id="code_feedback_pipeline_reference_impl")
@@ -141,4 +144,7 @@ def __call__(
                 }
             )
         )
+        num_tokens = self.llm.tokens
+        num_tokens.pipeline = PipelineEnum.IRIS_CODE_FEEDBACK
+        self.tokens = num_tokens
         return response.replace("{", "{{").replace("}", "}}")

app/pipeline/chat/course_chat_pipeline.py

@@ -41,6 +41,7 @@
     elicit_begin_agent_jol_prompt,
 )
 from ...domain import CourseChatPipelineExecutionDTO
+from ...llm.external.PipelineEnum import PipelineEnum
 from ...retrieval.lecture_retrieval import LectureRetrieval
 from ...vector_database.database import VectorDatabase
 from ...vector_database.lecture_schema import LectureSchema
@@ -107,6 +108,7 @@ def __init__(self, callback: CourseChatStatusCallback, variant: str = "default")
 
         # Create the pipeline
         self.pipeline = self.llm | StrOutputParser()
+        self.tokens = []
 
     def __repr__(self):
         return f"{self.__class__.__name__}(llm={self.llm})"
@@ -406,14 +408,18 @@ def lecture_content_retrieval() -> str:
             self.callback.in_progress()
             for step in agent_executor.iter(params):
                 print("STEP:", step)
+                token_count = self.llm.tokens
+                token_count.pipeline = PipelineEnum.IRIS_CHAT_COURSE_MESSAGE
+                self.tokens.append(token_count)
                 if step.get("output", None):
                     out = step["output"]
 
             if self.retrieved_paragraphs:
                 self.callback.in_progress("Augmenting response ...")
                 out = self.citation_pipeline(self.retrieved_paragraphs, out)
+            self.tokens.extend(self.citation_pipeline.tokens)
 
-            self.callback.done("Response created", final_result=out)
+            self.callback.done("Response created", final_result=out, tokens=self.tokens)
 
             # try:
             #     # if out:

app/pipeline/chat/exercise_chat_pipeline.py

@@ -34,6 +34,7 @@
 from ...domain.data.programming_submission_dto import ProgrammingSubmissionDTO
 from ...llm import CapabilityRequestHandler, RequirementList
 from ...llm import CompletionArguments
+from ...llm.external.PipelineEnum import PipelineEnum
 from ...llm.langchain import IrisLangchainChatModel
 from ...retrieval.lecture_retrieval import LectureRetrieval
 from ...vector_database.database import VectorDatabase
@@ -78,6 +79,7 @@ def __init__(self, callback: ExerciseChatStatusCallback):
         self.code_feedback_pipeline = CodeFeedbackPipeline()
         self.pipeline = self.llm | StrOutputParser()
         self.citation_pipeline = CitationPipeline()
+        self.tokens = []
 
     def __repr__(self):
         return f"{self.__class__.__name__}(llm={self.llm})"
@@ -98,7 +100,9 @@ def __call__(self, dto: ExerciseChatPipelineExecutionDTO):
             )
             self._run_exercise_chat_pipeline(dto, should_execute_lecture_pipeline),
             self.callback.done(
-                "Generated response", final_result=self.exercise_chat_response
+                "Generated response",
+                final_result=self.exercise_chat_response,
+                tokens=self.tokens,
             )
 
             try:
@@ -112,7 +116,11 @@ 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)
-                    self.callback.done(final_result=None, suggestions=suggestions)
+                    self.callback.done(
+                        final_result=None,
+                        suggestions=suggestions,
+                        tokens=[self.suggestion_pipeline.tokens],
+                    )
                 else:
                     # This should never happen but whatever
                     self.callback.skip(
@@ -200,6 +208,8 @@ def _run_exercise_chat_pipeline(
             if submission:
                 try:
                     feedback = future_feedback.result()
+                    if self.code_feedback_pipeline.tokens is not None:
+                        self.tokens.append(self.code_feedback_pipeline.tokens)
                     self.prompt += SystemMessagePromptTemplate.from_template(
                         "Another AI has checked the code of the student and has found the following issues. "
                         "Use this information to help the student. "
@@ -220,6 +230,8 @@ def _run_exercise_chat_pipeline(
             if should_execute_lecture_pipeline:
                 try:
                     self.retrieved_lecture_chunks = future_lecture.result()
+                    if self.retriever.tokens is not None:
+                        self.tokens.append(self.retriever.tokens)
                     if len(self.retrieved_lecture_chunks) > 0:
                         self._add_relevant_chunks_to_prompt(
                             self.retrieved_lecture_chunks
@@ -252,6 +264,9 @@ 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.callback.done()
             self.prompt = ChatPromptTemplate.from_messages(
                 [
@@ -266,6 +281,9 @@ 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)
 
             if "!ok!" in guide_response:
                 print("Response is ok and not rewritten!!!")

app/pipeline/chat/interaction_suggestion_pipeline.py

@@ -34,6 +34,8 @@
 )
 
 from ...llm import CompletionArguments
+from ...llm.external.LLMTokenCount import LLMTokenCount
+from ...llm.external.PipelineEnum import PipelineEnum
 from ...llm.langchain import IrisLangchainChatModel
 
 from ..pipeline import Pipeline
@@ -52,6 +54,7 @@ class InteractionSuggestionPipeline(Pipeline):
     pipeline: Runnable
     prompt: ChatPromptTemplate
     variant: str
+    tokens: LLMTokenCount
 
     def __init__(self, variant: str = "default"):
         super().__init__(implementation_id="interaction_suggestion_pipeline")
@@ -164,6 +167,8 @@ def __call__(
             self.prompt = ChatPromptTemplate.from_messages(prompt_val)
 
             response: dict = (self.prompt | self.pipeline).invoke({})
+            self.tokens = self.llm.tokens
+            self.tokens.pipeline = PipelineEnum.IRIS_INTERACTION_SUGGESTION
             return response["questions"]
         except Exception as e:
             logger.error(

app/pipeline/chat/lecture_chat_pipeline.py

@@ -16,6 +16,7 @@
     LectureChatPipelineExecutionDTO,
 )
 from ...llm import CapabilityRequestHandler, RequirementList
+from ...llm.external.PipelineEnum import PipelineEnum
 from ...retrieval.lecture_retrieval import LectureRetrieval
 from ...vector_database.database import VectorDatabase
 from ...vector_database.lecture_schema import LectureSchema
@@ -74,6 +75,7 @@ def __init__(self):
         self.retriever = LectureRetrieval(self.db.client)
         self.pipeline = self.llm | StrOutputParser()
         self.citation_pipeline = CitationPipeline()
+        self.tokens = []
 
     def __repr__(self):
         return f"{self.__class__.__name__}(llm={self.llm})"
@@ -114,9 +116,13 @@ 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)
             response_with_citation = self.citation_pipeline(
                 retrieved_lecture_chunks, response
             )
+            self.tokens.extend(self.citation_pipeline.tokens)
             logger.info(f"Response from lecture chat pipeline: {response}")
             return response_with_citation
         except Exception as e: