added first versions for tracking for smaller pipelines

app/llm/external/PipelineEnum.py

@@ -9,4 +9,6 @@ class PipelineEnum(str, Enum):
     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/pipeline/lecture_ingestion_pipeline.py

@@ -6,6 +6,7 @@
 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
@@ -18,6 +19,7 @@
     IngestionPipelineExecutionDto,
 )
 from ..domain.data.text_message_content_dto import TextMessageContentDTO
+from ..llm.external.PipelineEnum import PipelineEnum
 from ..llm.langchain import IrisLangchainChatModel
 from ..vector_database.lecture_schema import init_lecture_schema, LectureSchema
 from ..ingestion.abstract_ingestion import AbstractIngestion
@@ -112,6 +114,7 @@ def __init__(
             request_handler=request_handler, completion_args=completion_args
         )
         self.pipeline = self.llm | StrOutputParser()
+        self.tokens = []
 
     def __call__(self) -> bool:
         try:
@@ -273,9 +276,14 @@ def merge_page_content_and_image_interpretation(
             image_interpretation=image_interpretation,
         )
         prompt = ChatPromptTemplate.from_messages(prompt_val)
-        return clean(
+        clean_output = clean(
             (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)
+        return clean_output
 
     def get_course_language(self, page_content: str) -> str:
         """

app/pipeline/shared/reranker_pipeline.py

@@ -6,9 +6,11 @@
 from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
 from langchain_core.runnables import Runnable
 from langsmith import traceable
+from sipbuild.generator.parser.tokens import tokens
 
 from app.domain import PyrisMessage
 from app.llm import CapabilityRequestHandler, RequirementList, CompletionArguments
+from app.llm.external.PipelineEnum import PipelineEnum
 from app.llm.langchain import IrisLangchainChatModel
 from app.pipeline import Pipeline
 from app.pipeline.chat.output_models.output_models.selected_paragraphs import (
@@ -56,6 +58,7 @@ def __init__(self):
         )
         logger.debug(self.output_parser.get_format_instructions())
         self.pipeline = self.llm | self.output_parser
+        self.tokens = []
 
     def __repr__(self):
         return f"{self.__class__.__name__}(llm={self.llm})"
@@ -108,4 +111,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)
         return response.selected_paragraphs

app/pipeline/shared/summary_pipeline.py

@@ -45,6 +45,7 @@ def __init__(self):
         )
         # Create the pipeline
         self.pipeline = self.prompt | self.llm | StrOutputParser()
+        self.tokens = []
 
     def __repr__(self):
         return f"{self.__class__.__name__}(llm={self.llm})"