Update app.py

app.py

@@ -1,41 +1,39 @@
+# app.py
+
 import os
 import sys
 import json
 import re
 import logging
+import asyncio
 from datetime import datetime
 from typing import List, Dict, Optional
+
 from fastapi import FastAPI, HTTPException, UploadFile, File
 from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 import markdown
 import PyPDF2
-import asyncio
 
-# Setup logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
+# Logging setup
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("TxAgentAPI")
 
-# Adjust sys path
+# Path setup
 current_dir = os.path.dirname(os.path.abspath(__file__))
-src_path = os.path.abspath(os.path.join(current_dir, "src"))
-sys.path.insert(0, src_path)
+sys.path.insert(0, os.path.join(current_dir, "src"))
 
-# Import TxAgent
-try:
-    from txagent.txagent import TxAgent
-except ImportError as e:
-    logger.error(f"Failed to import TxAgent: {str(e)}")
-    raise
+# TxAgent import
+from txagent.txagent import TxAgent
 
-# Init FastAPI
+# MongoDB collections (shared URI via Hugging Face secrets)
+from db.mongo import patients_collection, results_collection
+
+# FastAPI app
 app = FastAPI(title="TxAgent API", version="2.1.0")
 
-# CORS
+# CORS config
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -44,7 +42,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Request schema
+# Pydantic schema
 class ChatRequest(BaseModel):
     message: str
     temperature: float = 0.7
@@ -52,149 +50,155 @@ class ChatRequest(BaseModel):
     history: Optional[List[Dict]] = None
     format: Optional[str] = "clean"
 
-# Response formatting
-def clean_text_response(text: str) -> str:
+# Utils
+def clean_text(text: str) -> str:
     text = re.sub(r'\n\s*\n', '\n\n', text)
     text = re.sub(r'[ ]+', ' ', text)
-    text = text.replace("**", "").replace("__", "")
-    return text.strip()
+    return text.strip().replace("**", "").replace("__", "")
+
+def extract_section(text: str, heading: str) -> str:
+    try:
+        pattern = rf"{heading}:\n(.*?)(?=\n[A-Z]|\Z)"
+        match = re.search(pattern, text, re.DOTALL)
+        return clean_text(match.group(1)) if match else ""
+    except:
+        return ""
 
 def structure_medical_response(text: str) -> Dict:
     return {
         "summary": extract_section(text, "Summary"),
         "risks": extract_section(text, "Risks or Red Flags"),
         "missed_issues": extract_section(text, "What the doctor might have missed"),
-        "recommendations": extract_section(text, "Suggested Clinical Actions")
+        "recommendations": extract_section(text, "Suggested Clinical Actions"),
     }
 
-def extract_section(text: str, heading: str) -> str:
+# Global agent
+agent = None
+
+# Background logic
+async def analyze_and_store_result(patient: dict):
     try:
-        pattern = rf"{heading}:\n(.*?)(?=\n\w|\Z)"
-        match = re.search(pattern, text, re.DOTALL)
-        return clean_text_response(match.group(1)) if match else ""
+        content = json.dumps(patient, indent=2)[:10000]
+        message = (
+            "You are a clinical AI assistant.\n\n"
+            "Analyze this patient's record and:\n"
+            "1. Summarize conditions and history.\n"
+            "2. Identify red flags.\n"
+            "3. Detect missed issues.\n"
+            "4. Suggest clinical actions.\n\n"
+            f"Patient Data:\n{content}"
+        )
+
+        raw = agent.chat(message=message, history=[], temperature=0.7, max_new_tokens=1024)
+        structured = structure_medical_response(raw)
+
+        await results_collection.update_one(
+            {"patient_id": patient.get("fhir_id")},
+            {
+                "$set": {
+                    "patient_id": patient.get("fhir_id"),
+                    "full_name": patient.get("full_name"),
+                    "raw": raw,
+                    "structured": structured,
+                    "analyzed_at": datetime.utcnow()
+                }
+            },
+            upsert=True
+        )
+        logger.info(f"Stored analysis for {patient.get('fhir_id')}")
     except Exception as e:
-        logger.error(f"Section extraction failed: {e}")
-        return ""
+        logger.error(f"Error analyzing patient: {e}")
 
-# Agent init
-agent = None
+async def analyze_existing_patients():
+    try:
+        patients = await patients_collection.find({}).to_list(length=None)
+        for patient in patients:
+            await analyze_and_store_result(patient)
+            await asyncio.sleep(0.3)
+    except Exception as e:
+        logger.error(f"Batch analysis failed: {e}")
+
+async def watch_new_patients():
+    try:
+        logger.info("Watching for new patient inserts...")
+        pipeline = [{'$match': {'operationType': 'insert'}}]
+        async with patients_collection.watch(pipeline) as stream:
+            async for change in stream:
+                patient = change["fullDocument"]
+                await analyze_and_store_result(patient)
+    except Exception as e:
+        logger.error(f"Change stream error: {e}")
 
 @app.on_event("startup")
 async def startup_event():
     global agent
-    try:
-        agent = TxAgent(
-            model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
-            rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
-            enable_finish=True,
-            enable_rag=False,
-            force_finish=True,
-            enable_checker=True,
-            step_rag_num=4,
-            seed=100
-        )
-        agent.chat_prompt = (
-            "You are a clinical decision support assistant for doctors. "
-            "You analyze patient documents, detect medical issues, identify missed diagnoses, "
-            "and provide treatment suggestions with rationale in concise, readable language."
-        )
-        agent.init_model()
-        logger.info("TxAgent initialized successfully")
-    except Exception as e:
-        logger.error(f"Startup error: {str(e)}")
+    agent = TxAgent(
+        model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
+        rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
+        enable_finish=True,
+        enable_checker=True,
+        force_finish=True,
+    )
+    agent.chat_prompt = (
+        "You are a clinical decision support AI helping doctors review patient records and suggest care plans."
+    )
+    agent.init_model()
+    logger.info("TxAgent loaded")
+    asyncio.create_task(analyze_existing_patients())
+    asyncio.create_task(watch_new_patients())
 
 @app.post("/chat-stream")
-async def chat_stream_endpoint(request: ChatRequest):
-    async def token_stream():
+async def chat_stream(request: ChatRequest):
+    async def stream():
         try:
-            conversation = []
-            conversation.append({"role": "system", "content": agent.chat_prompt})
+            msgs = [{"role": "system", "content": agent.chat_prompt}]
             if request.history:
-                for msg in request.history:
-                    conversation.append({"role": msg["role"], "content": msg["content"]})
-            conversation.append({"role": "user", "content": request.message})
-
-            input_ids = agent.tokenizer.apply_chat_template(
-                conversation,
-                add_generation_prompt=True,
-                return_tensors="pt"
-            ).to(agent.device)
-
-            streamer = agent.model.generate(
-                input_ids,
-                do_sample=True,
-                temperature=request.temperature,
-                max_new_tokens=request.max_new_tokens,
-                pad_token_id=agent.tokenizer.eos_token_id,
-                return_dict_in_generate=True,
-                output_scores=False
-            )
-
-            output = agent.tokenizer.decode(streamer["sequences"][0][input_ids.shape[1]:], skip_special_tokens=True)
-
-            for chunk in output.split():
+                msgs += request.history
+            msgs.append({"role": "user", "content": request.message})
+
+            input_ids = agent.tokenizer.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt").to(agent.device)
+            output = agent.model.generate(input_ids, do_sample=True, temperature=request.temperature, max_new_tokens=request.max_new_tokens)
+            text = agent.tokenizer.decode(output[0][input_ids.shape[1]:], skip_special_tokens=True)
+
+            for chunk in text.split():
                 yield chunk + " "
                 await asyncio.sleep(0.05)
-
         except Exception as e:
-            logger.error(f"Streaming chat error: {str(e)}")
-            yield f"\n⚠️ Error: {str(e)}"
+            yield f"\n⚠️ Error: {e}"
 
-    return StreamingResponse(token_stream(), media_type="text/plain")
+    return StreamingResponse(stream(), media_type="text/plain")
 
 @app.post("/upload")
 async def upload_file(file: UploadFile = File(...)):
     try:
-        logger.info(f"File upload received: {file.filename}")
-        content = ""
+        logger.info(f"Uploaded file: {file.filename}")
+        text = ""
         if file.filename.endswith(".pdf"):
-            pdf_reader = PyPDF2.PdfReader(file.file)
-            for page in pdf_reader.pages:
-                content += page.extract_text() or ""
+            pdf = PyPDF2.PdfReader(file.file)
+            text = "\n".join(p.extract_text() for p in pdf.pages if p.extract_text())
         else:
             content = await file.read()
-            content = content.decode("utf-8", errors="ignore")
-
-        message = f"""
-        You are a clinical decision support AI assisting physicians.
-
-        Given the following patient report, do the following:
-        1. Summarize the patient's main conditions and history.
-        2. Identify any potential clinical risks or red flags.
-        3. Highlight any important diagnoses or treatments the doctor might have missed.
-        4. Suggest next clinical steps, treatments, or referrals (if applicable).
-        5. Flag anything that could pose an urgent risk (e.g., suicide risk, untreated critical conditions).
-
-        Patient Document:
-        -----------------
-        {content[:10000]}
-        """
-
-        raw_response = agent.chat(message=message, history=[], temperature=0.7, max_new_tokens=1024)
-        formatted_response = {
-            "raw": raw_response,
-            "clean": clean_text_response(raw_response),
-            "structured": structure_medical_response(raw_response),
-            "html": markdown.markdown(raw_response)
-        }
-        return JSONResponse({
+            text = content.decode("utf-8", errors="ignore")
+
+        prompt = (
+            "You are a clinical support AI. Analyze the following:\n"
+            f"{text[:10000]}"
+        )
+        raw = agent.chat(message=prompt, history=[], temperature=0.7)
+        return {
             "status": "success",
-            "format": "clean",
-            "response": formatted_response["clean"],
-            "timestamp": datetime.now().isoformat(),
-            "available_formats": list(formatted_response.keys())
-        })
+            "response": clean_text(raw),
+            "structured": structure_medical_response(raw),
+            "timestamp": datetime.now().isoformat()
+        }
     except Exception as e:
-        logger.error(f"File upload error: {str(e)}")
+        logger.error(f"Upload error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
-    finally:
-        file.file.close()
 
 @app.get("/status")
 async def status():
     return {
         "status": "running",
-        "version": "2.1.0",
-        "model": agent.model_name if agent else "not loaded",
+        "model": agent.model_name,
         "timestamp": datetime.now().isoformat()
     }