Add mlflow query parameter logging

.gitignore

@@ -210,6 +210,7 @@ __marimo__/
 *.DS_Store
 **.DS_Store
 .idea/
+*.db
 
 # Chroma DB
 data/chroma_db/

README.md

@@ -16,7 +16,7 @@ Submit a clinical question. The system retrieves PubMed literature, generates a
 
 ## Motivation
 
-LLMs are increasingly being deployed in clinical settings, but they hallucinate — and in healthcare, hallucinations are dangerous. A model confidently stating an incorrect drug dosage or contraindication can directly harm patients.
+LLMs are increasingly being deployed in clinical settings, but they hallucinate. In healthcare, hallucinations are dangerous. A model confidently stating an incorrect drug dosage or contraindication can directly harm patients.
 
 SentinelMD addresses this by functioning as a **safety layer** that sits on top of any LLM, verifying its claims against authoritative medical literature in real time. Drawing on 8+ years of clinical experience in cardiac telemetry, this system was designed with a real understanding of how bad clinical information propagates through care workflows and what the consequences look like.
 
@@ -74,8 +74,6 @@ assembly                  Returns annotated response with claims, evidence, and
 
 ## Evaluation
 
-*RAG pipeline evaluation via RAGAS — coming in v1.1*
-
 | Metric | Score |
 |---|---|
 | Faithfulness | TBD |
@@ -193,13 +191,13 @@ Pinecone is a production-grade managed vector database used in real health tech
 General-purpose sentence transformers produce weak embeddings for clinical text because they weren't trained on biomedical language. BioBERT was pretrained on PubMed abstracts and fine-tuned on MedNLI, making it significantly better at capturing semantic similarity in clinical contexts.
 
 **Why NLI over cosine similarity for claim verification?**
-Cosine similarity tells you whether two pieces of text are topically related. NLI tells you whether one piece of text entails, contradicts, or is neutral toward another — which is the correct operation for hallucination detection.
+Cosine similarity tells you whether two pieces of text are topically related. NLI tells you whether one piece of text entails, contradicts, or is neutral toward another.
 
 ---
 
 ## Background
 
-Developed as a portfolio project demonstrating full-stack ML engineering in clinical AI safety. Informed by 8+ years of clinical experience in cardiac telemetry monitoring, with real-world awareness of how dangerous unverified clinical information is at the point of care — and what the consequences look like when it goes wrong.
+Developed as a portfolio project demonstrating full-stack ML engineering in clinical AI safety. Informed by 8+ years of clinical experience in cardiac telemetry monitoring, with real-world awareness of how dangerous unverified clinical information is at the point of care.
 
 ---
 

requirements.txt

@@ -4,6 +4,7 @@ python-dotenv>=1.0.0
 
 # Logging
 python-json-logger>=2.0.0
+mlflow>=3.11.0
 
 # Ruff Linting
 ruff>=0.4.0

src/agent/nodes.py

@@ -4,6 +4,7 @@
 from src.retrieval.cache import get_cache, set_cache
 from src.retrieval.vector_store import add_abstracts, query_abstracts
 from src.retrieval.pubmed import search_pubmed
+from src.monitoring.mlflow_logger import log_query_run
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_core.output_parsers import JsonOutputParser
 from sentence_transformers import CrossEncoder
@@ -117,6 +118,16 @@ def confidence_scoring(state: AgentState):
     return {"confidence_score": score}
 
 def assembly(state: AgentState):
+    final_response = {
+        "query": state["query"],
+        "response": state["llm_response"],
+        "confidence_score": state["confidence_score"],
+        "scored_claims": state["scored_claims"],
+        "abstracts": state["abstracts"]
+    }
+
+    log_query_run(final_response)
+
     return {"final_response": {
         "query": state["query"],
         "response": state["llm_response"],

src/monitoring/mlflow_logger.py

@@ -0,0 +1,18 @@
+import mlflow
+from src.core.config import settings
+
+def log_query_run(final_response: dict) -> None:
+    mlflow.set_tracking_uri(settings.MLFLOW_TRACKING_URI)
+    mlflow.set_experiment("SentinelMD")
+    supported_count = len([c for c in final_response["scored_claims"] if c["label"] == "Supported"])
+    unverifiable_count = len([c for c in final_response["scored_claims"] if c["label"] == "Unverifiable"])
+    contradicted_count = len([c for c in final_response["scored_claims"] if c["label"] == "Contradicted"])
+
+    with mlflow.start_run():
+        mlflow.log_param("query", final_response["query"])
+        mlflow.log_metric("abstracts_retrieved_count", len(final_response["abstracts"]))
+        mlflow.log_metric("confidence_score", final_response['confidence_score'])
+        mlflow.log_metric("supported_claims", supported_count)
+        mlflow.log_metric("unverifiable_claims", unverifiable_count)
+        mlflow.log_metric("contradicted_claims", contradicted_count)
+        mlflow.log_metric("total_claims", len(final_response["scored_claims"]))