Update app.py

app.py

@@ -7,27 +7,19 @@ from datetime import datetime
 from typing import List, Dict, Optional
 
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from pymongo import MongoClient
-from bson import ObjectId
 import asyncio
 
-# Adjust sys path
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "src")))
-
-# TxAgent
 from txagent.txagent import TxAgent
-
-# MongoDB
 from db.mongo import get_mongo_client
 
-# Setup logging
+# Logging
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger("TxAgentAPI")
 
-# FastAPI app
+# App
 app = FastAPI(title="TxAgent API", version="2.1.0")
 
 app.add_middleware(
@@ -36,7 +28,7 @@ app.add_middleware(
     allow_methods=["*"], allow_headers=["*"]
 )
 
-# Models
+# Pydantic
 class ChatRequest(BaseModel):
     message: str
     temperature: float = 0.7
@@ -46,7 +38,6 @@ class ChatRequest(BaseModel):
 
 # Globals
 agent = None
-mongo_client = None
 patients_collection = None
 analysis_collection = None
 
@@ -58,15 +49,13 @@ def clean_text_response(text: str) -> str:
 
 def extract_section(text: str, heading: str) -> str:
     try:
-        # Accept formats like "Heading:" or "**Heading:**"
-        pattern = rf"(\*\*)?{re.escape(heading)}(\*\*)?:\s*\n(.*?)(?=\n\s*\*?\*?\w|$)"
+        pattern = rf"{re.escape(heading)}:\s*\n(.*?)(?=\n[A-Z][^\n]*:|\Z)"
         match = re.search(pattern, text, re.DOTALL | re.IGNORECASE)
-        return clean_text_response(match.group(3)) if match else ""
+        return match.group(1).strip() if match else ""
     except Exception as e:
         logger.error(f"Section extraction failed for heading '{heading}': {e}")
         return ""
 
-
 def structure_medical_response(text: str) -> Dict:
     return {
         "summary": extract_section(text, "Summarize the patient's medical history"),
@@ -75,7 +64,6 @@ def structure_medical_response(text: str) -> Dict:
         "recommendations": extract_section(text, "Suggested Clinical Actions")
     }
 
-
 def serialize_patient(patient: dict) -> dict:
     patient_copy = patient.copy()
     if "_id" in patient_copy:
@@ -84,11 +72,15 @@ def serialize_patient(patient: dict) -> dict:
 
 async def analyze_patient(patient: dict):
     try:
-        doc = json.dumps(serialize_patient(patient), indent=2)
+        serialized = serialize_patient(patient)
+        doc = json.dumps(serialized, indent=2)
+        logger.info(f"🧾 Analyzing patient: {serialized.get('fhir_id')}")
+        logger.debug(f"🧠 Data passed to TxAgent:\n{doc[:1000]}")
+
         message = (
             "You are a clinical decision support AI.\n\n"
             "Given the patient document below:\n"
-            "1. Summarize their medical history.\n"
+            "1. Summarize the patient's medical history.\n"
             "2. Identify risks or red flags.\n"
             "3. Highlight missed diagnoses or treatments.\n"
             "4. Suggest next clinical steps.\n"
@@ -99,17 +91,17 @@ async def analyze_patient(patient: dict):
         structured = structure_medical_response(raw)
 
         analysis_doc = {
-            "patient_id": patient.get("fhir_id"),
+            "patient_id": serialized.get("fhir_id"),
             "timestamp": datetime.utcnow(),
             "summary": structured,
             "raw": raw
         }
         await analysis_collection.update_one(
-            {"patient_id": patient.get("fhir_id")},
+            {"patient_id": serialized.get("fhir_id")},
             {"$set": analysis_doc},
             upsert=True
         )
-        logger.info(f"✔️ Analysis stored for patient {patient.get('fhir_id')}")
+        logger.info(f"✅ Stored analysis for patient {serialized.get('fhir_id')}")
     except Exception as e:
         logger.error(f"Error analyzing patient: {e}")
 
@@ -119,12 +111,10 @@ async def analyze_all_patients():
         await analyze_patient(patient)
         await asyncio.sleep(0.1)
 
-# Startup logic
 @app.on_event("startup")
 async def startup_event():
-    global agent, mongo_client, patients_collection, analysis_collection
+    global agent, patients_collection, analysis_collection
 
-    # Init agent
     agent = TxAgent(
         model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
         rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
@@ -141,21 +131,17 @@ async def startup_event():
     agent.init_model()
     logger.info("✅ TxAgent initialized")
 
-    # MongoDB
-    mongo_client = get_mongo_client()
-    db = mongo_client["cps_db"]
-    patients_collection = db.get_collection("patients")
-    analysis_collection = db.get_collection("patient_analysis_results")
-
+    db = get_mongo_client()["cps_db"]
+    patients_collection = db["patients"]
+    analysis_collection = db["patient_analysis_results"]
     logger.info("📡 Connected to MongoDB")
+
     asyncio.create_task(analyze_all_patients())
 
-# Endpoints
 @app.get("/status")
 async def status():
     return {
         "status": "running",
-        "version": "2.1.0",
         "timestamp": datetime.utcnow().isoformat()
     }