Update agent.py

agent.py

@@ -22,8 +22,8 @@ logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
 
 # ── Environment Variables ─────────────────────────────────────────────────────
-UMLS_API_KEY = os.getenv("UMLS_API_KEY")
-GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+UMLS_API_KEY   = os.getenv("UMLS_API_KEY")
+GROQ_API_KEY   = os.getenv("GROQ_API_KEY")
 TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
 
 if not all([UMLS_API_KEY, GROQ_API_KEY, TAVILY_API_KEY]):
@@ -31,8 +31,8 @@ if not all([UMLS_API_KEY, GROQ_API_KEY, TAVILY_API_KEY]):
     raise RuntimeError("Missing required API keys")
 
 # ── Agent Configuration ───────────────────────────────────────────────────────
-AGENT_MODEL_NAME = "llama3-70b-8192"
-AGENT_TEMPERATURE = 0.1
+AGENT_MODEL_NAME   = "llama3-70b-8192"
+AGENT_TEMPERATURE  = 0.1
 MAX_SEARCH_RESULTS = 3
 
 class ClinicalPrompts:
@@ -42,30 +42,32 @@ class ClinicalPrompts:
     """
 
 # ── Helper Functions ──────────────────────────────────────────────────────────
-UMLS_AUTH_ENDPOINT = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
-RXNORM_API_BASE = "https://rxnav.nlm.nih.gov/REST"
-OPENFDA_API_BASE = "https://api.fda.gov/drug/label.json"
+UMLS_AUTH_ENDPOINT  = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
+RXNORM_API_BASE     = "https://rxnav.nlm.nih.gov/REST"
+OPENFDA_API_BASE    = "https://api.fda.gov/drug/label.json"
 
 @lru_cache(maxsize=256)
 def get_rxcui(drug_name: str) -> Optional[str]:
-    """Lookup RxNorm CUI for a drug name."""
+    """Lookup RxNorm CUI for a given drug name."""
     drug_name = (drug_name or "").strip()
     if not drug_name:
         return None
     logger.info(f"Looking up RxCUI for '{drug_name}'")
     try:
+        # First attempt
         params = {"name": drug_name, "search": 1}
         r = requests.get(f"{RXNORM_API_BASE}/rxcui.json", params=params, timeout=10)
         r.raise_for_status()
-        data = r.json().get("idGroup", {})
-        if ids := data.get("rxnormId"):
+        ids = r.json().get("idGroup", {}).get("rxnormId")
+        if ids:
             logger.info(f"Found RxCUI {ids[0]} for '{drug_name}'")
             return ids[0]
-        # fallback to broader search
+        # Fallback search
         r = requests.get(f"{RXNORM_API_BASE}/drugs.json", params={"name": drug_name}, timeout=10)
         r.raise_for_status()
-        for group in r.json().get("drugGroup", {}).get("conceptGroup", []):
-            if props := group.get("conceptProperties"):
+        for grp in r.json().get("drugGroup", {}).get("conceptGroup", []):
+            props = grp.get("conceptProperties")
+            if props:
                 logger.info(f"Found RxCUI {props[0]['rxcui']} via /drugs for '{drug_name}'")
                 return props[0]["rxcui"]
     except Exception:
@@ -74,7 +76,7 @@ def get_rxcui(drug_name: str) -> Optional[str]:
 
 @lru_cache(maxsize=128)
 def get_openfda_label(rxcui: Optional[str] = None, drug_name: Optional[str] = None) -> Optional[Dict[str, Any]]:
-    """Fetch label data from OpenFDA by RxCUI or drug name."""
+    """Fetch the OpenFDA label for a drug by RxCUI or name."""
     if not (rxcui or drug_name):
         return None
     terms = []
@@ -96,7 +98,7 @@ def get_openfda_label(rxcui: Optional[str] = None, drug_name: Optional[str] = No
     return None
 
 def search_text_list(texts: List[str], terms: List[str]) -> List[str]:
-    """Return snippets from texts containing any of the search terms."""
+    """Return highlighted snippets from a list of texts containing any of the search terms."""
     snippets = []
     lowers = [t.lower() for t in terms if t]
     for text in texts or []:
@@ -105,7 +107,7 @@ def search_text_list(texts: List[str], terms: List[str]) -> List[str]:
             if term in tl:
                 i = tl.find(term)
                 start = max(0, i - 50)
-                end = min(len(text), i + len(term) + 100)
+                end   = min(len(text), i + len(term) + 100)
                 snippet = text[start:end]
                 snippet = re.sub(f"({re.escape(term)})", r"**\1**", snippet, flags=re.IGNORECASE)
                 snippets.append(f"...{snippet}...")
@@ -113,34 +115,37 @@ def search_text_list(texts: List[str], terms: List[str]) -> List[str]:
     return snippets
 
 def parse_bp(bp: str) -> Optional[tuple[int, int]]:
-    """Parse a 'SYS/DIA' blood pressure string into a tuple."""
+    """Parse 'SYS/DIA' blood pressure string into a (sys, dia) tuple."""
     if m := re.match(r"(\d{1,3})\s*/\s*(\d{1,3})", (bp or "").strip()):
         return int(m.group(1)), int(m.group(2))
     return None
 
 def check_red_flags(patient_data: Dict[str, Any]) -> List[str]:
-    """Identify immediate red flags in patient data."""
+    """Identify immediate red flags from patient_data."""
     flags: List[str] = []
-    hpi = patient_data.get("hpi", {})
+    hpi    = patient_data.get("hpi", {})
     vitals = patient_data.get("vitals", {})
-    symptoms = [s.lower() for s in hpi.get("symptoms", []) if isinstance(s, str)]
-    # Symptom-based
+    syms   = [s.lower() for s in hpi.get("symptoms", []) if isinstance(s, str)]
+
+    # Symptom-based flags
     mapping = {
         "chest pain": "Chest pain reported",
         "shortness of breath": "Shortness of breath reported",
         "severe headache": "Severe headache reported",
-        "syncope": "Syncope (fainting) reported",
-        "hemoptysis": "Hemoptysis (coughing blood) reported"
+        "syncope": "Syncope reported",
+        "hemoptysis": "Hemoptysis reported"
     }
     for term, desc in mapping.items():
-        if term in symptoms:
+        if term in syms:
             flags.append(f"Red Flag: {desc}.")
-    # Vitals-based
+
+    # Vitals-based flags
     temp = vitals.get("temp_c")
     hr   = vitals.get("hr_bpm")
     rr   = vitals.get("rr_rpm")
     spo2 = vitals.get("spo2_percent")
     bp   = parse_bp(vitals.get("bp_mmhg", ""))
+
     if temp is not None and temp >= 38.5:
         flags.append(f"Red Flag: Fever ({temp}°C).")
     if hr is not None:
@@ -158,10 +163,11 @@ def check_red_flags(patient_data: Dict[str, Any]) -> List[str]:
             flags.append(f"Red Flag: Hypertensive urgency/emergency ({sys}/{dia} mmHg).")
         if sys <= 90 or dia <= 60:
             flags.append(f"Red Flag: Hypotension ({sys}/{dia} mmHg).")
-    return list(dict.fromkeys(flags))  # preserve order, dedupe
+
+    return list(dict.fromkeys(flags))  # dedupe, preserve order
 
 def format_patient_data_for_prompt(data: Dict[str, Any]) -> str:
-    """Convert patient_data dict into a markdown-like prompt string."""
+    """Format patient_data dict into a markdown-like prompt section."""
     if not data:
         return "No patient data provided."
     lines: List[str] = []
@@ -181,29 +187,30 @@ def format_patient_data_for_prompt(data: Dict[str, Any]) -> str:
 # ── Tool Input Schemas ────────────────────────────────────────────────────────
 class LabOrderInput(BaseModel):
     test_name: str = Field(...)
-    reason: str = Field(...)
-    priority: str = Field("Routine")
+    reason: str    = Field(...)
+    priority: str  = Field("Routine")
 
 class PrescriptionInput(BaseModel):
-    medication_name: str = Field(...)
-    dosage: str = Field(...)
-    route: str = Field(...)
-    frequency: str = Field(...)
-    duration: str = Field("As directed")
-    reason: str = Field(...)
+    medication_name: str          = Field(...)
+    dosage: str                   = Field(...)
+    route: str                    = Field(...)
+    frequency: str                = Field(...)
+    duration: str                 = Field("As directed")
+    reason: str                   = Field(...)
 
 class InteractionCheckInput(BaseModel):
-    potential_prescription: str = Field(...)
+    potential_prescription: str
     current_medications: Optional[List[str]] = Field(None)
-    allergies: Optional[List[str]] = Field(None)
+    allergies: Optional[List[str]]             = Field(None)
 
 class FlagRiskInput(BaseModel):
     risk_description: str = Field(...)
-    urgency: str = Field("High")
+    urgency: str          = Field("High")
 
 # ── Tool Implementations ──────────────────────────────────────────────────────
 @tool("order_lab_test", args_schema=LabOrderInput)
 def order_lab_test(test_name: str, reason: str, priority: str = "Routine") -> str:
+    """Place an order for a laboratory test."""
     logger.info(f"Ordering lab test: {test_name}, reason: {reason}, priority: {priority}")
     return json.dumps({
         "status": "success",
@@ -220,6 +227,7 @@ def prescribe_medication(
     duration: str,
     reason: str
 ) -> str:
+    """Prepare a medication prescription."""
     logger.info(f"Preparing prescription: {medication_name} {dosage}, route: {route}, freq: {frequency}")
     return json.dumps({
         "status": "success",
@@ -233,15 +241,17 @@ def check_drug_interactions(
     current_medications: Optional[List[str]] = None,
     allergies: Optional[List[str]] = None
 ) -> str:
+    """Check for drug–drug interactions and allergy risks."""
     logger.info(f"Checking interactions for: {potential_prescription}")
     warnings: List[str] = []
     pm = [m.lower().strip() for m in (current_medications or []) if m]
     al = [a.lower().strip() for a in (allergies or []) if a]
 
-    # Allergy checks
+    # Allergy exact match
     if potential_prescription.lower().strip() in al:
         warnings.append(f"CRITICAL ALLERGY: Patient allergic to '{potential_prescription}'.")
-    # RxNorm/OpenFDA lookups
+
+    # Identify drug via RxNorm/OpenFDA
     rxcui = get_rxcui(potential_prescription)
     label = get_openfda_label(rxcui=rxcui, drug_name=potential_prescription)
     if not (rxcui or label):
@@ -259,16 +269,15 @@ def check_drug_interactions(
     for med in pm:
         mrxcui = get_rxcui(med)
         mlabel = get_openfda_label(rxcui=mrxcui, drug_name=med)
-        # check in both labels
         for sec in ("drug_interactions",):
             for src_label, src_name in ((label, potential_prescription), (mlabel, med)):
                 items = src_label.get(sec) if src_label else None
                 if isinstance(items, list):
-                    snippets = search_text_list(items, [med if src_name==potential_prescription else potential_prescription])
+                    snippets = search_text_list(items, [med if src_name == potential_prescription else potential_prescription])
                     if snippets:
                         warnings.append(f"Interaction ({src_name} label): {'; '.join(snippets)}")
 
-    status = "warning" if warnings else "clear"
+    status  = "warning" if warnings else "clear"
     message = (
         f"{len(warnings)} issue(s) found for '{potential_prescription}'."
         if warnings else
@@ -278,19 +287,21 @@ def check_drug_interactions(
 
 @tool("flag_risk", args_schema=FlagRiskInput)
 def flag_risk(risk_description: str, urgency: str = "High") -> str:
+    """Flag a clinical risk with given urgency."""
     logger.info(f"Flagging risk: {risk_description} (urgency={urgency})")
     return json.dumps({
         "status": "flagged",
         "message": f"Risk '{risk_description}' flagged with {urgency} urgency."
     })
 
+# Include the Tavily search tool
 search_tool = TavilySearchResults(max_results=MAX_SEARCH_RESULTS, name="tavily_search_results")
-all_tools = [order_lab_test, prescribe_medication, check_drug_interactions, flag_risk, search_tool]
+all_tools   = [order_lab_test, prescribe_medication, check_drug_interactions, flag_risk, search_tool]
 
 # ── LLM & Tool Executor ──────────────────────────────────────────────────────
-llm = ChatGroq(temperature=AGENT_TEMPERATURE, model=AGENT_MODEL_NAME)
+llm              = ChatGroq(temperature=AGENT_TEMPERATURE, model=AGENT_MODEL_NAME)
 model_with_tools = llm.bind_tools(all_tools)
-tool_executor = ToolExecutor(all_tools)
+tool_executor    = ToolExecutor(all_tools)
 
 # ── State Definition ──────────────────────────────────────────────────────────
 class AgentState(TypedDict):
@@ -319,17 +330,19 @@ def tool_node(state: AgentState) -> Dict[str, Any]:
         return {"messages": [], "interaction_warnings": None}
 
     calls = last.tool_calls
-    # Safety: require interaction check before prescribing
     blocked_ids = set()
     for call in calls:
         if call["name"] == "prescribe_medication":
             med = call["args"].get("medication_name", "").lower()
-            if not any(c["name"] == "check_drug_interactions" and c["args"].get("potential_prescription","").lower() == med for c in calls):
+            if not any(
+                c["name"] == "check_drug_interactions" and
+                c["args"].get("potential_prescription","").lower() == med
+                for c in calls
+            ):
                 logger.warning(f"Blocking prescribe_medication for '{med}' without interaction check")
                 blocked_ids.add(call["id"])
 
     to_execute = [c for c in calls if c["id"] not in blocked_ids]
-    # Augment interaction checks with patient data
     pd = state.get("patient_data", {})
     for call in to_execute:
         if call["name"] == "check_drug_interactions":
@@ -337,7 +350,7 @@ def tool_node(state: AgentState) -> Dict[str, Any]:
             call["args"].setdefault("allergies", pd.get("allergies", []))
 
     messages: List[ToolMessage] = []
-    warnings: List[str] = []
+    warnings: List[str]     = []
     try:
         responses = tool_executor.batch(to_execute, return_exceptions=True)
         for call, resp in zip(to_execute, responses):
@@ -353,7 +366,6 @@ def tool_node(state: AgentState) -> Dict[str, Any]:
             messages.append(ToolMessage(content=content, tool_call_id=call["id"], name=call["name"]))
     except Exception as e:
         logger.exception("Critical error in tool_node")
-        # return an error message for each pending call
         for call in to_execute:
             messages.append(ToolMessage(
                 content=json.dumps({"status": "error", "message": str(e)}),
@@ -368,8 +380,7 @@ def reflection_node(state: AgentState) -> Dict[str, Any]:
         logger.warning("reflection_node called without warnings")
         return {"messages": [], "interaction_warnings": None}
 
-    # Find the AIMessage that triggered the warnings
-    triggering: Optional[AIMessage] = None
+    triggering = None
     for msg in reversed(state["messages"]):
         if isinstance(msg, AIMessage) and getattr(msg, "tool_calls", None):
             triggering = msg
@@ -381,7 +392,8 @@ def reflection_node(state: AgentState) -> Dict[str, Any]:
     prompt = (
         "You are SynapseAI, performing a focused safety review of the following plan:\n\n"
         f"{triggering.content}\n\n"
-        "Highlight any issues based on these warnings:\n" + "\n".join(f"- {w}" for w in warns)
+        "Highlight any issues based on these warnings:\n" +
+        "\n".join(f"- {w}" for w in warns)
     )
     try:
         resp = llm.invoke([SystemMessage(content="Safety reflection"), HumanMessage(content=prompt)])