added mlflow tracing

Author: floschne

src/fundus_murag/agent/chat_assistant.py

@@ -4,10 +4,12 @@
 from typing import Any, Iterable
 
 import google.auth.transport.requests
+import mlflow
 import openai
 import pandas as pd
 from google.auth import default
 from loguru import logger
+from mlflow.entities import SpanType
 from openai.types.chat import ChatCompletionMessageParam
 from openai.types.chat.chat_completion import ChatCompletion
 from openai.types.chat.chat_completion_assistant_message_param import (
@@ -268,6 +270,7 @@ def _create_chat_completion_from_history(self) -> ChatCompletion:
             logger.error(f"[{self.assistant_name}] Unexpected Error of Type {type(e)}: {e}")
             raise e
 
+    @mlflow.trace(span_type=SpanType.AGENT)
     def _run_agentic_loop(self) -> str:
         # 1. Send the messages in the chat history to the model
         response = self._create_chat_completion_from_history()

src/fundus_murag/agent/tools/image_analyzer.py

@@ -1,7 +1,9 @@
 import json
 from enum import Enum
 
+import mlflow
 from loguru import logger
+from mlflow.entities import SpanType
 
 from fundus_murag.agent.chat_assistant import ChatAssistant
 from fundus_murag.agent.chat_assistant_factory import ChatAssistantFactory
@@ -49,6 +51,9 @@ def _get_image_analysis_assistant(self, task: ImageAnalysisTask) -> ChatAssistan
         )
         return assistant
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def answer_question_about_fundus_record_image(self, question: str, murag_id: str) -> str:
         """
         Generates an answer to the given question about the image of a `FundusRecord` specified by the given `murag_id`.
@@ -75,6 +80,9 @@ def answer_question_about_fundus_record_image(self, question: str, murag_id: str
             logger.error(msg)
             return msg
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def generate_caption_for_fundus_record_image(self, murag_id: str, detailed: bool = False) -> str:
         """
         Generates a caption for the image of a `FundusRecord` specified by the given `murag_id`.
@@ -100,6 +108,9 @@ def generate_caption_for_fundus_record_image(self, murag_id: str, detailed: bool
             logger.error(msg)
             return msg
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def extract_text_from_fundus_record_image(self, murag_id: str) -> str:
         """
         Perform OCR on the image of a `FundusRecord` specified by the given `murag_id`.
@@ -124,6 +135,9 @@ def extract_text_from_fundus_record_image(self, murag_id: str) -> str:
             logger.error(msg)
             return msg
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def detect_objects_in_fundus_record_image(self, murag_id: str) -> str:
         """
         Perform object detection on the image of a `FundusRecord` specified by the given `murag_id`.

src/fundus_murag/api/lifespan.py

@@ -1,5 +1,6 @@
 from contextlib import asynccontextmanager
 
+import mlflow
 from fastapi import FastAPI
 from loguru import logger
 
@@ -12,6 +13,9 @@
 @asynccontextmanager
 async def api_lifespan(app: FastAPI):
     # Startup
+    mlflow.set_tracking_uri("http://localhost:58058")
+    mlflow.set_experiment("/fundus")
+    mlflow.openai.autolog()
     logger.info("Starting FUNDus! Data API")
     chat_assistant_factory = ChatAssistantFactory()
     fundus_agent_factory = FundusMultiAgentSystemFactory()

src/fundus_murag/data/vector_db.py

@@ -3,9 +3,11 @@
 from functools import reduce
 from typing import Any, Literal
 
+import mlflow
 import pandas as pd
 import weaviate
 from loguru import logger
+from mlflow.entities import SpanType
 from tqdm import tqdm
 from weaviate.classes.query import Filter, MetadataQuery, QueryNested, QueryReference
 
@@ -359,6 +361,9 @@ def _create_fundus_record_from_query_results(self, res: Any) -> list[FundusRecor
 
         return records
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def get_total_number_of_fundus_records(self) -> int:
         """
         Get the total number of FUNDus! records in the database.
@@ -375,6 +380,9 @@ def get_total_number_of_fundus_records(self) -> int:
         except Exception:
             return 0
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def get_total_number_of_fundus_collections(self) -> int:
         """
         Get the total number of FUNDus! collections in the database.
@@ -391,6 +399,9 @@ def get_total_number_of_fundus_collections(self) -> int:
         except Exception:
             return 0
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def list_all_fundus_collections(self) -> list[FundusCollection]:
         """
         List all `FundusCollection`s in the FUNDus! database.
@@ -403,6 +414,9 @@ def list_all_fundus_collections(self) -> list[FundusCollection]:
         results = self._create_fundus_collection_from_query_results(res)
         return results
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def get_fundus_collection_by_murag_id(
         self,
         murag_id: str,
@@ -478,6 +492,9 @@ def _resolve_collection_name(self, collection_name: str) -> str:
         results = matches[0][0]
         return results
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def get_fundus_collection_by_name(
         self,
         collection_name: str,
@@ -504,6 +521,9 @@ def get_fundus_collection_by_name(
         results = self._create_fundus_collection_from_query_results(res)
         return results[0]
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def get_random_fundus_collection(
         self,
         n: int = 1,
@@ -525,6 +545,9 @@ def get_random_fundus_collection(
         results = [self.get_fundus_collection_by_name(collection_name=cn) for cn in collection_names]
         return results
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def get_random_fundus_records(
         self,
         n: int = 1,
@@ -557,6 +580,9 @@ def get_random_fundus_records(
         results = [self.get_fundus_record_by_murag_id(murag_id=murag_id) for murag_id in murag_ids]
         return results
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def get_fundus_record_by_murag_id(
         self,
         murag_id: str,
@@ -612,6 +638,9 @@ def get_fundus_record_image_by_murag_id(
         )
         return result
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def get_fundus_records_by_fundus_id(
         self,
         fundus_id: int,
@@ -636,6 +665,9 @@ def get_fundus_records_by_fundus_id(
         results = self._create_fundus_record_from_query_results(res)
         return results
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def find_fundus_records_with_similar_image(
         self,
         murag_id: str,
@@ -664,6 +696,9 @@ def find_fundus_records_with_similar_image(
         )
         return results
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def find_fundus_records_with_images_similar_to_the_text_query(
         self,
         query: str,
@@ -692,6 +727,9 @@ def find_fundus_records_with_images_similar_to_the_text_query(
         )
         return results
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def find_fundus_records_with_images_similar_to_the_user_provided_image(
         self,
         user_image_id: str,
@@ -723,6 +761,9 @@ def find_fundus_records_with_images_similar_to_the_user_provided_image(
         )
         return results
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def find_fundus_records_with_titles_similar_to_the_text_query(
         self,
         query: str,
@@ -920,6 +961,9 @@ def get_fundus_record_internal_by_murag_id(
 
         return rec[0]
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def fundus_record_title_lexical_search(
         self,
         query: str,
@@ -992,6 +1036,9 @@ def _fundus_record_lexical_search(
         results = self._create_fundus_record_from_query_results(results)
         return results
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def fundus_collection_lexical_search(
         self,
         query: str,
@@ -1073,6 +1120,9 @@ def _fundus_collection_lexical_search(
         results = self._create_fundus_collection_from_query_results(res)
         return results
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def fundus_collection_description_similarity_search(
         self,
         query_embedding: list[float],
@@ -1095,6 +1145,9 @@ def fundus_collection_description_similarity_search(
         )
         return results
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def fundus_collection_title_similarity_search(
         self,
         query_embedding: list[float],
@@ -1152,6 +1205,9 @@ def _fundus_collection_similarity_search(
         results = simsearch_results
         return results
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def get_number_of_records_per_collection(self) -> dict[str, int]:
         """
         Get the number of records in each collection.
@@ -1162,6 +1218,9 @@ def get_number_of_records_per_collection(self) -> dict[str, int]:
         results = self._records_df.groupby("collection_name").size().to_dict()  # type: ignore
         return results
 
+    @mlflow.trace(
+        span_type=SpanType.TOOL,
+    )
     def get_number_of_records_in_collection(self, collection_name: str) -> int:
         """
         Get the number of records in the specified collection.