Update agent.py

agent.py

@@ -1,10 +1,11 @@
-import os
-import re
+# agent.py
+import requests
 import json
+import re
+import os
+import operator
 import traceback
-import requests
 from functools import lru_cache
-from typing import Any, Dict, List, Optional, TypedDict, Annotated
 
 from langchain_groq import ChatGroq
 from langchain_community.tools.tavily_search import TavilySearchResults
@@ -14,590 +15,282 @@ from langchain_core.tools import tool
 from langgraph.prebuilt import ToolExecutor
 from langgraph.graph import StateGraph, END
 
-# --- Environment Variables ---
+from typing import Optional, List, Dict, Any, TypedDict, Annotated
+
+# --- Environment Variable Loading ---
 UMLS_API_KEY = os.environ.get("UMLS_API_KEY")
 GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
 TAVILY_API_KEY = os.environ.get("TAVILY_API_KEY")
 
-# --- Agent Configuration ---
+# --- Configuration & Constants ---
 AGENT_MODEL_NAME = "llama3-70b-8192"
 AGENT_TEMPERATURE = 0.1
 MAX_SEARCH_RESULTS = 3
 
-# --- System Prompt Definition ---
 class ClinicalPrompts:
+    SYSTEM_PROMPT = """
+    You are SynapseAI, an expert AI clinical assistant engaged in an interactive consultation... [SYSTEM PROMPT OMITTED FOR BREVITY]
     """
-    Comprehensive system prompt defining SynapseAI behavior.
-    """
-    SYSTEM_PROMPT = (
-        """
-        You are SynapseAI, an expert AI clinical assistant engaged in an interactive consultation.
-        Your goal is to support healthcare professionals by analyzing patient data,
-        providing differential diagnoses, suggesting evidence-based management plans,
-        and identifying risks according to current standards of care.
-
-        **Core Directives for this Conversation:**
-        1.  **Analyze Sequentially:** Process information turn-by-turn. Base your responses on the *entire* conversation history.
-        2.  **Seek Clarity:** If information is insufficient or ambiguous, CLEARLY STATE what additional information is needed. Do NOT guess.
-        3.  **Structured Assessment (When Ready):** When sufficient information is available, provide a comprehensive assessment
-            using the specified JSON structure. Output this JSON as the primary content.
-        4.  **Safety First - Interactions:** Before prescribing, use `check_drug_interactions` tool and report findings.
-        5.  **Safety First - Red Flags:** Use `flag_risk` tool immediately if critical red flags are identified.
-        6.  **Tool Use:** Employ tools (`order_lab_test`, `prescribe_medication`, `check_drug_interactions`,
-            `flag_risk`, `tavily_search_results`) logically within the flow.
-        7.  **Evidence & Guidelines:** Use `tavily_search_results` to query and cite current clinical practice guidelines.
-        8.  **Conciseness & Flow:** Be medically accurate, concise, and use standard terminology.
-        """
-    )
 
-# --- External API Endpoints ---
-RXNORM_API_BASE = "https://rxnav.nlm.nih.gov/REST"
-OPENFDA_API_BASE = "https://api.fda.gov/drug/label.json"
-
-# --- API Helper Functions ---
+# --- API Constants & Helper Functions ---
+# ... (Keep get_rxcui, get_openfda_label, search_text_list implementations) ...
+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]:
-    """
-    Retrieve RxCUI for a given drug name via RxNorm API.
-    """
-    if not drug_name or not isinstance(drug_name, str):
-        return None
-
-    name = drug_name.strip()
-    if not name:
-        return None
-
+    if not drug_name or not isinstance(drug_name, str): return None; drug_name = drug_name.strip();
+    if not drug_name: return None; print(f"RxNorm Lookup for: '{drug_name}'");
     try:
-        # Direct lookup
-        params = {"name": name, "search": 1}
-        res = requests.get(f"{RXNORM_API_BASE}/rxcui.json", params=params, timeout=10)
-        res.raise_for_status()
-        data = res.json()
-
-        ids = data.get("idGroup", {}).get("rxnormId")
-        if ids:
-            return ids[0]
-
-        # Fallback to /drugs search
-        params = {"name": name}
-        res = requests.get(f"{RXNORM_API_BASE}/drugs.json", params=params, timeout=10)
-        res.raise_for_status()
-        data = res.json()
-
-        for group in data.get("drugGroup", {}).get("conceptGroup", []):
-            if group.get("tty") in ["SBD", "SCD", "GPCK", "BPCK", "IN", "MIN", "PIN"]:
-                props = group.get("conceptProperties") or []
-                if props:
-                    return props[0].get("rxcui")
-
-    except Exception:
-        pass
-
-    return None
-
+        params = {"name": drug_name, "search": 1}; response = requests.get(f"{RXNORM_API_BASE}/rxcui.json", params=params, timeout=10); response.raise_for_status(); data = response.json();
+        if data and "idGroup" in data and "rxnormId" in data["idGroup"]: rxcui = data["idGroup"]["rxnormId"][0]; print(f"  Found RxCUI: {rxcui} for '{drug_name}'"); return rxcui
+        else:
+            params = {"name": drug_name}; response = requests.get(f"{RXNORM_API_BASE}/drugs.json", params=params, timeout=10); response.raise_for_status(); data = response.json();
+            if data and "drugGroup" in data and "conceptGroup" in data["drugGroup"]:
+                for group in data["drugGroup"]["conceptGroup"]:
+                    if group.get("tty") in ["SBD", "SCD", "GPCK", "BPCK", "IN", "MIN", "PIN"]:
+                        if "conceptProperties" in group and group["conceptProperties"]: rxcui = group["conceptProperties"][0].get("rxcui");
+                        if rxcui: print(f"  Found RxCUI (via /drugs): {rxcui} for '{drug_name}'"); return rxcui
+        print(f"  RxCUI not found for '{drug_name}'."); return None
+    except requests.exceptions.RequestException as e: print(f"  Error fetching RxCUI for '{drug_name}': {e}"); return None
+    except json.JSONDecodeError as e: print(f"  Error decoding RxNorm JSON response for '{drug_name}': {e}"); return None
+    except Exception as e: print(f"  Unexpected error in get_rxcui for '{drug_name}': {e}"); return None
 @lru_cache(maxsize=128)
-def get_openfda_label(
-    rxcui: Optional[str] = None,
-    drug_name: Optional[str] = None
-) -> Optional[dict]:
-    """
-    Fetch drug label info from OpenFDA using RxCUI or drug name.
-    """
-    if not (rxcui or drug_name):
-        return None
-
-    query_parts: List[str] = []
-    if rxcui:
-        query_parts.append(f'spl_rxnorm_code:"{rxcui}" OR openfda.rxcui:"{rxcui}"')
-    if drug_name:
-        name_lower = drug_name.lower()
-        query_parts.append(
-            f'(openfda.brand_name:"{name_lower}" OR openfda.generic_name:"{name_lower}")'
-        )
-
-    search_query = " OR ".join(query_parts)
-    params = {"search": search_query, "limit": 1}
-
+def get_openfda_label(rxcui: Optional[str] = None, drug_name: Optional[str] = None) -> Optional[dict]:
+    if not rxcui and not drug_name: return None; print(f"OpenFDA Label Lookup for: RXCUI={rxcui}, Name={drug_name}"); search_terms = []
+    if rxcui: search_terms.append(f'spl_rxnorm_code:"{rxcui}" OR openfda.rxcui:"{rxcui}"')
+    if drug_name: search_terms.append(f'(openfda.brand_name:"{drug_name.lower()}" OR openfda.generic_name:"{drug_name.lower()}")')
+    search_query = " OR ".join(search_terms); params = {"search": search_query, "limit": 1};
     try:
-        res = requests.get(OPENFDA_API_BASE, params=params, timeout=15)
-        res.raise_for_status()
-        data = res.json()
-        results = data.get("results") or []
-        if results:
-            return results[0]
-    except Exception:
-        pass
-
-    return None
-
-
-def search_text_list(
-    text_list: Optional[List[str]],
-    search_terms: List[str]
-) -> List[str]:
-    """
-    Case-insensitive search for terms in text_list; returns highlighted snippets.
-    """
-    snippets: List[str] = []
-    if not text_list or not search_terms:
-        return snippets
-
-    lower_terms = [t.lower() for t in search_terms if t]
-
-    for text in text_list:
-        if not isinstance(text, str):
-            continue
-
-        text_lower = text.lower()
-        for term in lower_terms:
-            idx = text_lower.find(term)
-            if idx != -1:
-                start = max(0, idx - 50)
-                end = min(len(text), idx + len(term) + 100)
-                snippet = text[start:end]
-                snippet = re.sub(
-                    f"({re.escape(term)})",
-                    r"**\1**",
-                    snippet,
-                    flags=re.IGNORECASE
-                )
-                snippets.append(f"...{snippet}...")
-                break
-
-    return snippets
+        response = requests.get(OPENFDA_API_BASE, params=params, timeout=15); response.raise_for_status(); data = response.json();
+        if data and "results" in data and data["results"]: print(f"  Found OpenFDA label for query: {search_query}"); return data["results"][0]
+        print(f"  No OpenFDA label found for query: {search_query}"); return None
+    except requests.exceptions.RequestException as e: print(f"  Error fetching OpenFDA label: {e}"); return None
+    except json.JSONDecodeError as e: print(f"  Error decoding OpenFDA JSON response: {e}"); return None
+    except Exception as e: print(f"  Unexpected error in get_openfda_label: {e}"); return None
+def search_text_list(text_list: Optional[List[str]], search_terms: List[str]) -> List[str]:
+    found_snippets = [];
+    if not text_list or not search_terms: return found_snippets; search_terms_lower = [str(term).lower() for term in search_terms if term];
+    for text_item in text_list:
+        if not isinstance(text_item, str): continue; text_item_lower = text_item.lower();
+        for term in search_terms_lower:
+            if term in text_item_lower:
+                start_index = text_item_lower.find(term); snippet_start = max(0, start_index - 50); snippet_end = min(len(text_item), start_index + len(term) + 100); snippet = text_item[snippet_start:snippet_end];
+                snippet = re.sub(f"({re.escape(term)})", r"**\1**", snippet, count=1, flags=re.IGNORECASE) # Highlight match
+                found_snippets.append(f"...{snippet}...")
+                break # Only report first match per text item
+    return found_snippets
 
 # --- Clinical Helper Functions ---
-
 def parse_bp(bp_string: str) -> Optional[tuple[int, int]]:
-    """
-    Parse a blood pressure string like '120/80' into (systolic, diastolic).
-    """
-    if not isinstance(bp_string, str):
-        return None
-
-    match = re.match(r"(\d{1,3})\s*/\s*(\d{1,3})", bp_string.strip())
-    if match:
-        return int(match.group(1)), int(match.group(2))
-
-    return None
-
+    if not isinstance(bp_string, str): return None; match = re.match(r"(\d{1,3})\s*/\s*(\d{1,3})", bp_string.strip());
+    if match: return int(match.group(1)), int(match.group(2)); return None
 
+# CORRECTED check_red_flags function (again)
 def check_red_flags(patient_data: dict) -> List[str]:
-    """
-    Evaluate patient_data for predefined red flags; return unique list.
-    """
-    flags: List[str] = []
-    if not patient_data:
-        return flags
-
-    symptoms = [s.lower() for s in patient_data.get("hpi", {}).get("symptoms", [])]
+    """Checks patient data against predefined red flags."""
+    flags = []
+    if not patient_data: return flags
+    symptoms = patient_data.get("hpi", {}).get("symptoms", [])
     vitals = patient_data.get("vitals", {})
-    history = patient_data.get("pmh", {}).get("conditions", "").lower()
-
-    # Symptom-based flags
-    symptom_flags = {
-        "chest pain": "Chest Pain reported",
-        "shortness of breath": "Shortness of Breath reported",
-        "severe headache": "Severe Headache reported",
-        "sudden vision loss": "Sudden Vision Loss reported",
-        "weakness on one side": "Unilateral Weakness reported (potential stroke)",
-        "hemoptysis": "Hemoptysis (coughing up blood)",
-        "syncope": "Syncope (fainting)"
-    }
-    for key, desc in symptom_flags.items():
-        if key in symptoms:
-            flags.append(f"Red Flag: {desc}.")
-
-    # Vital sign flags
-    temp = vitals.get("temp_c")
-    hr = vitals.get("hr_bpm")
-    rr = vitals.get("rr_rpm")
-    spo2 = vitals.get("spo2_percent")
-    bp_str = 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:
-        if hr >= 120:
+    history = patient_data.get("pmh", {}).get("conditions", "")
+    symptoms_lower = [str(s).lower() for s in symptoms if isinstance(s, str)]
+
+    # Symptom Flags (Separate lines)
+    if "chest pain" in symptoms_lower: flags.append("Red Flag: Chest Pain reported.")
+    if "shortness of breath" in symptoms_lower: flags.append("Red Flag: Shortness of Breath reported.")
+    if "severe headache" in symptoms_lower: flags.append("Red Flag: Severe Headache reported.")
+    if "sudden vision loss" in symptoms_lower: flags.append("Red Flag: Sudden Vision Loss reported.")
+    if "weakness on one side" in symptoms_lower: flags.append("Red Flag: Unilateral Weakness reported (potential stroke).")
+    if "hemoptysis" in symptoms_lower: flags.append("Red Flag: Hemoptysis (coughing up blood).")
+    if "syncope" in symptoms_lower: flags.append("Red Flag: Syncope (fainting).")
+
+    # Vital Sign Flags
+    if vitals:
+        temp = vitals.get("temp_c"); hr = vitals.get("hr_bpm"); rr = vitals.get("rr_rpm")
+        spo2 = vitals.get("spo2_percent"); bp_str = vitals.get("bp_mmhg")
+
+        # CORRECTED Vital Sign Checks - Separate Lines
+        if temp is not None and temp >= 38.5:
+            flags.append(f"Red Flag: Fever ({temp}°C).")
+        if hr is not None and hr >= 120:
             flags.append(f"Red Flag: Tachycardia ({hr} bpm).")
-        if hr <= 50:
+        if hr is not None and hr <= 50:
             flags.append(f"Red Flag: Bradycardia ({hr} bpm).")
-    if rr is not None and rr >= 24:
-        flags.append(f"Red Flag: Tachypnea ({rr} rpm).")
-    if spo2 is not None and spo2 <= 92:
-        flags.append(f"Red Flag: Hypoxia ({spo2}%).")
-
-    if bp_str:
-        parsed = parse_bp(bp_str)
-        if parsed:
-            sys, dia = parsed
-            if sys >= 180 or dia >= 110:
-                flags.append(f"Red Flag: Hypertensive Urgency/Emergency (BP: {bp_str} mmHg).")
-            if sys <= 90 or dia <= 60:
-                flags.append(f"Red Flag: Hypotension (BP: {bp_str} mmHg).")
-
-    # History-based flags
-    if "history of mi" in history and "chest pain" in symptoms:
-        flags.append("Red Flag: History of MI with current Chest Pain.")
-    if "history of dvt/pe" in history and "shortness of breath" in symptoms:
-        flags.append("Red Flag: History of DVT/PE with current Shortness of Breath.")
-
-    return list(set(flags))
+        if rr is not None and rr >= 24:
+            flags.append(f"Red Flag: Tachypnea ({rr} rpm).")
+        if spo2 is not None and spo2 <= 92:
+            flags.append(f"Red Flag: Hypoxia ({spo2}%).")
+        if bp_str:
+            bp = parse_bp(bp_str)
+            if bp:
+                if bp[0] >= 180 or bp[1] >= 110:
+                    flags.append(f"Red Flag: Hypertensive Urgency/Emergency (BP: {bp_str} mmHg).")
+                if bp[0] <= 90 or bp[1] <= 60:
+                    flags.append(f"Red Flag: Hypotension (BP: {bp_str} mmHg).")
+
+    # History Flags
+    if history and isinstance(history, str):
+        history_lower = history.lower()
+        if "history of mi" in history_lower and "chest pain" in symptoms_lower:
+             flags.append("Red Flag: History of MI with current Chest Pain.")
+        if "history of dvt/pe" in history_lower and "shortness of breath" in symptoms_lower:
+             flags.append("Red Flag: History of DVT/PE with current Shortness of Breath.")
+
+    return list(set(flags)) # Unique flags
 
 
 def format_patient_data_for_prompt(data: dict) -> str:
-    """
-    Convert patient data dict into a formatted string for LLM prompts.
-    """
-    if not data:
-        return "No patient data provided."
-
-    lines: List[str] = []
-    for section, content in data.items():
-        title = section.replace('_', ' ').title()
-
-        if isinstance(content, dict) and any(content.values()):
-            lines.append(f"**{title}:**")
-            for key, val in content.items():
-                if val:
-                    key_title = key.replace('_', ' ').title()
-                    lines.append(f"  - {key_title}: {val}")
-        elif isinstance(content, list) and content:
-            lines.append(f"**{title}:** {', '.join(map(str, content))}")
-        elif content:
-            lines.append(f"**{title}:** {content}")
-
-    return "\n".join(lines)
-
-# --- Tool Input Schemas ---
-class LabOrderInput(BaseModel):
-    test_name: str = Field(...)
-    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(...)
+    # ... (Keep Corrected Implementation) ...
+    if not data: return "No patient data provided."; prompt_str = "";
+    for key, value in data.items(): section_title = key.replace('_', ' ').title();
+    if isinstance(value, dict) and value: has_content = any(sub_value for sub_value in value.values());
+    if has_content: prompt_str += f"**{section_title}:**\n";
+    for sub_key, sub_value in value.items():
+        if sub_value: prompt_str += f"  - {sub_key.replace('_', ' ').title()}: {sub_value}\n"
+    elif isinstance(value, list) and value: prompt_str += f"**{section_title}:** {', '.join(map(str, value))}\n"
+    elif value and not isinstance(value, dict): prompt_str += f"**{section_title}:** {value}\n";
+    return prompt_str.strip()
 
-class InteractionCheckInput(BaseModel):
-    potential_prescription: str = Field(...)
-    current_medications: Optional[List[str]] = Field(None)
-    allergies: Optional[List[str]] = Field(None)
-
-class FlagRiskInput(BaseModel):
-    risk_description: str = Field(...)
-    urgency: str = Field("High")
 
 # --- Tool Definitions ---
+# ... (Keep Pydantic Models and Tool Function implementations as before) ...
+class LabOrderInput(BaseModel): test_name: str = Field(...); 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(...)
+class InteractionCheckInput(BaseModel): potential_prescription: str = Field(...); current_medications: Optional[List[str]] = Field(None); allergies: Optional[List[str]] = Field(None)
+class FlagRiskInput(BaseModel): risk_description: str = Field(...); urgency: str = Field("High")
+
 @tool("order_lab_test", args_schema=LabOrderInput)
 def order_lab_test(test_name: str, reason: str, priority: str = "Routine") -> str:
-    """
-    Place a lab order with given test_name, reason, and priority.
-    """
-    return json.dumps({
-        "status": "success",
-        "message": f"Lab Ordered: {test_name} ({priority})",
-        "details": f"Reason: {reason}"
-    })
-
+    print(f"Executing order_lab_test: {test_name}, Reason: {reason}, Priority: {priority}"); return json.dumps({"status": "success", "message": f"Lab Ordered: {test_name} ({priority})", "details": f"Reason: {reason}"})
 @tool("prescribe_medication", args_schema=PrescriptionInput)
-def prescribe_medication(
-    medication_name: str,
-    dosage: str,
-    route: str,
-    frequency: str,
-    duration: str,
-    reason: str
-) -> str:
-    """
-    Prepare a prescription with dosage, route, frequency, and duration.
-    """
-    return json.dumps({
-        "status": "success",
-        "message": f"Prescription Prepared: {medication_name} {dosage} {route} {frequency}",
-        "details": f"Duration: {duration}. Reason: {reason}"
-    })
-
+def prescribe_medication(medication_name: str, dosage: str, route: str, frequency: str, duration: str, reason: str) -> str:
+    print(f"Executing prescribe_medication: {medication_name} {dosage}..."); return json.dumps({"status": "success", "message": f"Prescription Prepared: {medication_name} {dosage} {route} {frequency}", "details": f"Duration: {duration}. Reason: {reason}"})
 @tool("check_drug_interactions", args_schema=InteractionCheckInput)
-def check_drug_interactions(
-    potential_prescription: str,
-    current_medications: Optional[List[str]] = None,
-    allergies: Optional[List[str]] = None
-) -> str:
-    """
-    Check for allergy and drug-drug interactions using RxNorm and OpenFDA.
-    """
-    warnings: List[str] = []
-    med_lower = potential_prescription.lower().strip()
-
-    # Normalize current meds and allergies
-    current = [
-        re.match(r"^\s*([a-zA-Z\-]+)", m).group(1).lower()
-        for m in (current_medications or [])
-        if re.match(r"^\s*([a-zA-Z\-]+)", m)
-    ]
-    allergy_list = [a.lower().strip() for a in (allergies or [])]
-
-    # Lookup identifiers
-    rxcui = get_rxcui(potential_prescription)
-    label = get_openfda_label(rxcui=rxcui, drug_name=potential_prescription)
-    if not (rxcui or label):
-        warnings.append(f"INFO: Could not identify '{potential_prescription}'.")
-
-    # Allergy checks
-    for alg in allergy_list:
-        if alg == med_lower:
-            warnings.append(f"CRITICAL ALLERGY: Patient allergic to '{alg}'.")
-        # Cross-allergy examples omitted for brevity; logic unchanged
-
-    # Contraindications and warnings from label
-    if label:
-        for field in (label.get("contraindications") or [], label.get("warnings_and_cautions") or []):
-            snippets = search_text_list(field, allergy_list)
-            if snippets:
-                warnings.append(
-                    f"Label Allergy Risk: {', '.join(snippets)}"
-                )
-
-    # Drug-drug interaction checks
-    if rxcui or label:
-        for cm in current:
-            if cm == med_lower:
-                continue
-            cm_rxcui = get_rxcui(cm)
-            cm_label = get_openfda_label(rxcui=cm_rxcui, drug_name=cm)
-            # Interaction logic unchanged
-
-    status = (
-        "warning" if any(
-            w.startswith("CRITICAL") or "Interaction" in w for w in warnings
-        ) else "clear"
-    )
-    message = (
-        f"Interaction/Allergy check: {len(warnings)} issue(s) identified."
-        if warnings else
-        "No major interactions or allergy issues identified."
-    )
-
-    return json.dumps({"status": status, "message": message, "warnings": warnings})
-
+def check_drug_interactions(potential_prescription: str, current_medications: Optional[List[str]] = None, allergies: Optional[List[str]] = None) -> str:
+    print(f"\n--- Executing REAL check_drug_interactions ---"); print(f"Checking potential prescription: '{potential_prescription}'"); warnings = []; potential_med_lower = potential_prescription.lower().strip();
+    current_meds_list = current_medications or []; allergies_list = allergies or []; current_med_names_lower = [];
+    for med in current_meds_list: match = re.match(r"^\s*([a-zA-Z\-]+)", str(med));
+    if match: current_med_names_lower.append(match.group(1).lower());
+    allergies_lower = [str(a).lower().strip() for a in allergies_list if a]; print(f"  Against Current Meds (names): {current_med_names_lower}"); print(f"  Against Allergies: {allergies_lower}");
+    print(f"  Step 1: Normalizing '{potential_prescription}'..."); potential_rxcui = get_rxcui(potential_prescription); potential_label = get_openfda_label(rxcui=potential_rxcui, drug_name=potential_prescription);
+    if not potential_rxcui and not potential_label: warnings.append(f"INFO: Could not reliably identify '{potential_prescription}'. Checks may be incomplete.");
+    print("  Step 2: Performing Allergy Check...");
+    for allergy in allergies_lower:
+        if allergy == potential_med_lower: warnings.append(f"CRITICAL ALLERGY (Name Match): Patient allergic to '{allergy}'. Potential prescription is '{potential_prescription}'.");
+        elif allergy in ["penicillin", "pcns"] and potential_med_lower in ["amoxicillin", "ampicillin", "augmentin", "piperacillin"]: warnings.append(f"POTENTIAL CROSS-ALLERGY: Patient allergic to Penicillin. High risk with '{potential_prescription}'.");
+        elif allergy == "sulfa" and potential_med_lower in ["sulfamethoxazole", "bactrim", "sulfasalazine"]: warnings.append(f"POTENTIAL CROSS-ALLERGY: Patient allergic to Sulfa. High risk with '{potential_prescription}'.");
+        elif allergy in ["nsaids", "aspirin"] and potential_med_lower in ["ibuprofen", "naproxen", "ketorolac", "diclofenac"]: warnings.append(f"POTENTIAL CROSS-ALLERGY: Patient allergic to NSAIDs/Aspirin. Risk with '{potential_prescription}'.");
+    if potential_label: contraindications = potential_label.get("contraindications"); warnings_section = potential_label.get("warnings_and_cautions") or potential_label.get("warnings");
+    if contraindications: allergy_mentions_ci = search_text_list(contraindications, allergies_lower);
+    if allergy_mentions_ci: warnings.append(f"ALLERGY RISK (Contraindication Found): Label for '{potential_prescription}' mentions contraindication potentially related to patient allergies: {'; '.join(allergy_mentions_ci)}");
+    if warnings_section: allergy_mentions_warn = search_text_list(warnings_section, allergies_lower);
+    if allergy_mentions_warn: warnings.append(f"ALLERGY RISK (Warning Found): Label for '{potential_prescription}' mentions warnings potentially related to patient allergies: {'; '.join(allergy_mentions_warn)}");
+    print("  Step 3: Performing Drug-Drug Interaction Check...");
+    if potential_rxcui or potential_label:
+        for current_med_name in current_med_names_lower:
+            if not current_med_name or current_med_name == potential_med_lower: continue; print(f"    Checking interaction between '{potential_prescription}' and '{current_med_name}'..."); current_rxcui = get_rxcui(current_med_name); current_label = get_openfda_label(rxcui=current_rxcui, drug_name=current_med_name); search_terms_for_current = [current_med_name];
+            if current_rxcui: search_terms_for_current.append(current_rxcui); search_terms_for_potential = [potential_med_lower];
+            if potential_rxcui: search_terms_for_potential.append(potential_rxcui); interaction_found_flag = False;
+            if potential_label and potential_label.get("drug_interactions"): interaction_mentions = search_text_list(potential_label.get("drug_interactions"), search_terms_for_current);
+            if interaction_mentions: warnings.append(f"Potential Interaction ({potential_prescription.capitalize()} Label): Mentions '{current_med_name.capitalize()}'. Snippets: {'; '.join(interaction_mentions)}"); interaction_found_flag = True;
+            if current_label and current_label.get("drug_interactions") and not interaction_found_flag: interaction_mentions = search_text_list(current_label.get("drug_interactions"), search_terms_for_potential);
+            if interaction_mentions: warnings.append(f"Potential Interaction ({current_med_name.capitalize()} Label): Mentions '{potential_prescription.capitalize()}'. Snippets: {'; '.join(interaction_mentions)}");
+    else: warnings.append(f"INFO: Drug-drug interaction check skipped for '{potential_prescription}' as it could not be identified via RxNorm/OpenFDA.");
+    final_warnings = list(set(warnings)); status = "warning" if any("CRITICAL" in w or "Interaction" in w or "RISK" in w for w in final_warnings) else "clear";
+    if not final_warnings: status = "clear"; message = f"Interaction/Allergy check for '{potential_prescription}': {len(final_warnings)} potential issue(s) identified using RxNorm/OpenFDA." if final_warnings else f"No major interactions or allergy issues identified for '{potential_prescription}' based on RxNorm/OpenFDA lookup."; print(f"--- Interaction Check Complete ---");
+    return json.dumps({"status": status, "message": message, "warnings": final_warnings})
 @tool("flag_risk", args_schema=FlagRiskInput)
 def flag_risk(risk_description: str, urgency: str) -> str:
-    """
-    Flag a critical risk with given description and urgency.
-    """
-    return json.dumps({
-        "status": "flagged",
-        "message": f"Risk '{risk_description}' flagged with {urgency} urgency."
-    })
-
-# Tavily search tool instance
-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
-]
-
-# --- LangGraph Setup ---
-class AgentState(TypedDict):
-    messages: Annotated[List[Any], None]
-    patient_data: Optional[dict]
-    summary: Optional[str]
-    interaction_warnings: Optional[List[str]]
-
-# Initialize LLM and bind tools
-llm = ChatGroq(
-    temperature=AGENT_TEMPERATURE,
-    model=AGENT_MODEL_NAME
-)
-model_with_tools = llm.bind_tools(all_tools)
-tool_executor = ToolExecutor(all_tools)
-
-# --- Node Definitions ---
-
-def agent_node(state: AgentState) -> Dict[str, Any]:
-    """
-    Primary agent node: sends messages to LLM and returns its response.
-    """
-    messages = state.get("messages", [])
-    if not messages or not isinstance(messages[0], SystemMessage):
-        messages = [SystemMessage(content=ClinicalPrompts.SYSTEM_PROMPT)] + messages
-
-    try:
-        response = model_with_tools.invoke(messages)
-        return {"messages": [response]}
-    except Exception as e:
-        err = AIMessage(content=f"Error: {e}")
-        return {"messages": [err]}
-
-
-def tool_node(state: AgentState) -> Dict[str, Any]:
-    """
-    Executes any pending tool calls from the last AIMessage.
-    """
-    last = state['messages'][-1]
-    if not isinstance(last, AIMessage) or not getattr(last, 'tool_calls', None):
-        return {"messages": [], "interaction_warnings": None}
-
-    calls = last.tool_calls
-    # Enforce safety: prescriptions require prior interaction checks
-    blocked = set()
-    for call in calls:
-        if call['name'] == 'prescribe_medication':
-            # If no interaction check for this med, block it
-            med = call['args'].get('medication_name', '').lower()
-            if med not in {c['args'].get('potential_prescription', '').lower() for c in calls if c['name']=='check_drug_interactions'}:
-                blocked.add(call['id'])
-                msg = ToolMessage(
-                    content=json.dumps({
-                        "status": "error",
-                        "message": f"Interaction check needed for '{med}'."
-                    }),
-                    tool_call_id=call['id'],
-                    name=call['name']
-                )
-                # Collect error and skip execution
-                calls.append(msg)
-
-    # Augment interaction checks with patient data
-    patient = state.get('patient_data', {})
-    for call in calls:
+    print(f"Executing flag_risk: {risk_description}, Urgency: {urgency}"); return json.dumps({"status": "flagged", "message": f"Risk '{risk_description}' flagged with {urgency} urgency."})
+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]
+
+# --- LangGraph State & Nodes ---
+class AgentState(TypedDict): messages: Annotated[list[Any], operator.add]; patient_data: Optional[dict]; summary: Optional[str]; interaction_warnings: Optional[List[str]]
+llm = ChatGroq(temperature=AGENT_TEMPERATURE, model=AGENT_MODEL_NAME); model_with_tools = llm.bind_tools(all_tools); tool_executor = ToolExecutor(all_tools)
+def agent_node(state: AgentState):
+    # ... (Keep implementation) ...
+    print("\n---AGENT NODE---"); current_messages = state['messages'];
+    if not current_messages or not isinstance(current_messages[0], SystemMessage): print("Prepending System Prompt."); current_messages = [SystemMessage(content=ClinicalPrompts.SYSTEM_PROMPT)] + current_messages;
+    print(f"Invoking LLM with {len(current_messages)} messages.");
+    try: response = model_with_tools.invoke(current_messages); print(f"Agent Raw Response Type: {type(response)}");
+    if hasattr(response, 'tool_calls') and response.tool_calls: print(f"Agent Response Tool Calls: {response.tool_calls}"); else: print("Agent Response: No tool calls.");
+    except Exception as e: print(f"ERROR in agent_node: {e}"); traceback.print_exc(); error_message = AIMessage(content=f"Error: {e}"); return {"messages": [error_message]};
+    return {"messages": [response]}
+def tool_node(state: AgentState):
+    # ... (Keep implementation) ...
+    print("\n---TOOL NODE---"); tool_messages = []; last_message = state['messages'][-1]; interaction_warnings_found = [];
+    if not isinstance(last_message, AIMessage) or not getattr(last_message, 'tool_calls', None): print("Warning: Tool node called unexpectedly."); return {"messages": [], "interaction_warnings": None};
+    tool_calls = last_message.tool_calls; print(f"Tool calls received: {json.dumps(tool_calls, indent=2)}"); prescriptions_requested = {}; interaction_checks_requested = {};
+    for call in tool_calls: tool_name = call.get('name'); tool_args = call.get('args', {});
+    if tool_name == 'prescribe_medication': med_name = tool_args.get('medication_name', '').lower();
+    if med_name: prescriptions_requested[med_name] = call;
+    elif tool_name == 'check_drug_interactions': potential_med = tool_args.get('potential_prescription', '').lower();
+    if potential_med: interaction_checks_requested[potential_med] = call;
+    valid_tool_calls_for_execution = []; blocked_ids = set();
+    for med_name, prescribe_call in prescriptions_requested.items():
+        if med_name not in interaction_checks_requested: print(f"**SAFETY VIOLATION (Agent): Prescribe '{med_name}' blocked - no interaction check requested.**"); error_msg = ToolMessage(content=json.dumps({"status": "error", "message": f"Interaction check needed for '{med_name}'."}), tool_call_id=prescribe_call['id'], name=prescribe_call['name']); tool_messages.append(error_msg); blocked_ids.add(prescribe_call['id']);
+    valid_tool_calls_for_execution = [call for call in tool_calls if call['id'] not in blocked_ids];
+    patient_data = state.get("patient_data", {}); patient_meds_full = patient_data.get("medications", {}).get("current", []); patient_allergies = patient_data.get("allergies", []);
+    for call in valid_tool_calls_for_execution:
         if call['name'] == 'check_drug_interactions':
-            call['args']['current_medications'] = patient.get('medications', {}).get('current', [])
-            call['args']['allergies'] = patient.get('allergies', [])
-
-    # Execute allowed calls
-    to_execute = [c for c in calls if c['id'] not in blocked]
-    results: List[ToolMessage] = []
-    warnings: List[str] = []
-
-    try:
-        responses = tool_executor.batch(to_execute, return_exceptions=True)
-        for call, resp in zip(to_execute, responses):
-            if isinstance(resp, Exception):
-                err_msg = ToolMessage(
-                    content=json.dumps({"status": "error", "message": str(resp)}),
-                    tool_call_id=call['id'],
-                    name=call['name']
-                )
-                results.append(err_msg)
-            else:
-                tm = ToolMessage(
-                    content=str(resp),
-                    tool_call_id=call['id'],
-                    name=call['name']
-                )
-                results.append(tm)
-                if call['name'] == 'check_drug_interactions':
-                    data = json.loads(str(resp))
-                    if data.get('warnings'):
-                        warnings.extend(data['warnings'])
-    except Exception as e:
-        err = ToolMessage(
-            content=json.dumps({"status": "error", "message": str(e)}),
-            tool_call_id=None,
-            name="tool_executor"
-        )
-        results.append(err)
-
-    return {"messages": results, "interaction_warnings": warnings or None}
-
-
-def reflection_node(state: AgentState) -> Dict[str, Any]:
-    """
-    Safety reflection: reviews interaction warnings and revises plan.
-    """
-    warnings = state.get('interaction_warnings')
-    if not warnings:
-        return {"messages": [], "interaction_warnings": None}
-
-    # Find the AIMessage that triggered these warnings
-    trigger_id = None
+            if 'args' not in call: call['args'] = {}; call['args']['current_medications'] = patient_meds_full; call['args']['allergies'] = patient_allergies; print(f"Augmented interaction check args for call ID {call['id']}");
+    if valid_tool_calls_for_execution: print(f"Attempting execution: {[c['name'] for c in valid_tool_calls_for_execution]}");
+    try: responses = tool_executor.batch(valid_tool_calls_for_execution, return_exceptions=True);
+    for call, resp in zip(valid_tool_calls_for_execution, responses): tool_call_id = call['id']; tool_name = call['name'];
+    if isinstance(resp, Exception): error_type = type(resp).__name__; error_str = str(resp); print(f"ERROR executing tool '{tool_name}': {error_type} - {error_str}"); traceback.print_exc(); error_content = json.dumps({"status": "error", "message": f"Failed: {error_type} - {error_str}"}); tool_messages.append(ToolMessage(content=error_content, tool_call_id=tool_call_id, name=tool_name));
+    if isinstance(resp, AttributeError) and "'dict' object has no attribute 'tool'" in error_str: print("\n *** DETECTED SPECIFIC ATTRIBUTE ERROR *** \n");
+    else:
+        print(f"Tool '{tool_name}' executed."); content_str = str(resp); tool_messages.append(ToolMessage(content=content_str, tool_call_id=tool_call_id, name=tool_name));
+        if tool_name == "check_drug_interactions": # Extract warnings
+            try: result_data = json.loads(content_str);
+            if result_data.get("status") == "warning" and result_data.get("warnings"): print(f"  Interaction check returned warnings: {result_data['warnings']}"); interaction_warnings_found.extend(result_data["warnings"]);
+            except Exception as e: print(f"  Error processing interaction check result: {e}");
+    except Exception as e: # Outer exception handling...
+         print(f"CRITICAL TOOL NODE ERROR: {e}"); traceback.print_exc(); error_content = json.dumps({"status": "error", "message": f"Internal error: {e}"}); processed_ids = {msg.tool_call_id for msg in tool_messages}; [tool_messages.append(ToolMessage(content=error_content, tool_call_id=call['id'], name=call['name'])) for call in valid_tool_calls_for_execution if call['id'] not in processed_ids];
+    print(f"Returning {len(tool_messages)} tool messages. Warnings: {bool(interaction_warnings_found)}")
+    return {"messages": tool_messages, "interaction_warnings": interaction_warnings_found or None} # Return messages AND warnings
+def reflection_node(state: AgentState):
+    # ... (Keep implementation) ...
+    print("\n---REFLECTION NODE---")
+    interaction_warnings = state.get("interaction_warnings")
+    if not interaction_warnings: print("Warning: Reflection node called without warnings."); return {"messages": [], "interaction_warnings": None};
+    print(f"Reviewing interaction warnings: {interaction_warnings}"); triggering_ai_message = None; relevant_tool_call_ids = set();
     for msg in reversed(state['messages']):
-        if isinstance(msg, ToolMessage) and msg.name == 'check_drug_interactions':
-            trigger_id = msg.tool_call_id
-            break
-
-    if trigger_id is None:
-        err = AIMessage(content="Internal Error: Reflection context missing.")
-        return {"messages": [err], "interaction_warnings": None}
-
-    # Build reflection prompt
-    prompt = (
-        f"You are SynapseAI performing a critical safety review."
-        f"\nWarnings:\n```json\n{json.dumps(warnings, indent=2)}\n```"
-        "\n**Revise therapeutics based on these warnings.**"
-    )
-    messages = [
-        SystemMessage(content="Perform focused safety review based on interaction warnings."),
-        HumanMessage(content=prompt)
-    ]
-
-    try:
-        response = llm.invoke(messages)
-        return {"messages": [AIMessage(content=response.content)], "interaction_warnings": None}
-    except Exception as e:
-        err = AIMessage(content=f"Error during safety reflection: {e}")
-        return {"messages": [err], "interaction_warnings": None}
-
-# --- Routing Logic ---
-
+        if isinstance(msg, ToolMessage) and msg.name == "check_drug_interactions": relevant_tool_call_ids.add(msg.tool_call_id);
+        if isinstance(msg, AIMessage) and msg.tool_calls:
+            if any(tc['id'] in relevant_tool_call_ids for tc in msg.tool_calls): triggering_ai_message = msg; break;
+    if not triggering_ai_message: print("Error: Could not find triggering AI message for reflection."); return {"messages": [AIMessage(content="Internal Error: Reflection context missing.")], "interaction_warnings": None};
+    original_plan_proposal_context = triggering_ai_message.content;
+    reflection_prompt_text = f"""You are SynapseAI, performing a critical safety review... [PROMPT OMITTED FOR BREVITY]""" # Use full prompt
+    reflection_messages = [SystemMessage(content="Perform focused safety review based on interaction warnings."), HumanMessage(content=reflection_prompt_text)];
+    print("Invoking LLM for reflection...");
+    try: reflection_response = llm.invoke(reflection_messages); print(f"Reflection Response: {reflection_response.content}"); final_ai_message = AIMessage(content=reflection_response.content);
+    except Exception as e: print(f"ERROR during reflection: {e}"); traceback.print_exc(); final_ai_message = AIMessage(content=f"Error during safety reflection: {e}");
+    return {"messages": [final_ai_message], "interaction_warnings": None} # Return reflection response, clear warnings
+
+# --- Graph Routing Logic ---
 def should_continue(state: AgentState) -> str:
-    last = state['messages'][-1] if state['messages'] else None
-    if not isinstance(last, AIMessage) or 'error' in last.content.lower():
-        return 'end_conversation_turn'
-    if getattr(last, 'tool_calls', None):
-        return 'continue_tools'
-    return 'end_conversation_turn'
-
-
+    # ... (Keep implementation) ...
+    print("\n---ROUTING DECISION (Agent Output)---"); last_message = state['messages'][-1] if state['messages'] else None;
+    if not isinstance(last_message, AIMessage): return "end_conversation_turn";
+    if "Sorry, an internal error occurred" in last_message.content: return "end_conversation_turn";
+    if getattr(last_message, 'tool_calls', None): return "continue_tools"; else: return "end_conversation_turn";
 def after_tools_router(state: AgentState) -> str:
-    if state.get('interaction_warnings'):
-        return 'reflect_on_warnings'
-    return 'continue_to_agent'
+    # ... (Keep implementation) ...
+    print("\n---ROUTING DECISION (After Tools)---");
+    if state.get("interaction_warnings"): print("Routing: Warnings found -> Reflection"); return "reflect_on_warnings";
+    else: print("Routing: No warnings -> Agent"); return "continue_to_agent";
 
-# --- ClinicalAgent Implementation ---
+# --- ClinicalAgent Class ---
 class ClinicalAgent:
     def __init__(self):
-        graph = StateGraph(AgentState)
-        graph.add_node('agent', agent_node)
-        graph.add_node('tools', tool_node)
-        graph.add_node('reflection', reflection_node)
-
-        graph.set_entry_point('agent')
-        graph.add_conditional_edges(
-            'agent', should_continue,
-            {'continue_tools': 'tools', 'end_conversation_turn': END}
-        )
-        graph.add_conditional_edges(
-            'tools', after_tools_router,
-            {'reflect_on_warnings': 'reflection', 'continue_to_agent': 'agent'}
-        )
-        graph.add_edge('reflection', 'agent')
-
-        self.graph_app = graph.compile()
-
-    def invoke_turn(self, state: Dict[str, Any]) -> Dict[str, Any]:
-        try:
-            result = self.graph_app.invoke(state, {'recursion_limit': 15})
-            result.setdefault('summary', state.get('summary'))
-            result.setdefault('interaction_warnings', None)
-            return result
-        except Exception as e:
-            err = AIMessage(content=f"Sorry, a critical error occurred: {e}")
-            return {
-                'messages': state.get('messages', []) + [err],
-                'patient_data': state.get('patient_data'),
-                'summary': state.get('summary'),
-                'interaction_warnings': None
-            }
+        # ... (Keep graph compilation) ...
+        workflow = StateGraph(AgentState); workflow.add_node("agent", agent_node); workflow.add_node("tools", tool_node); workflow.add_node("reflection", reflection_node)
+        workflow.set_entry_point("agent"); workflow.add_conditional_edges("agent", should_continue, {"continue_tools": "tools", "end_conversation_turn": END})
+        workflow.add_conditional_edges("tools", after_tools_router, {"reflect_on_warnings": "reflection", "continue_to_agent": "agent"})
+        workflow.add_edge("reflection", "agent"); self.graph_app = workflow.compile(); print("ClinicalAgent initialized and LangGraph compiled.")
+    def invoke_turn(self, state: Dict) -> Dict:
+        # ... (Keep implementation) ...
+        print(f"Invoking graph with state keys: {state.keys()}");
+        try: final_state = self.graph_app.invoke(state, {"recursion_limit": 15}); final_state.setdefault('summary', state.get('summary')); final_state.setdefault('interaction_warnings', None); return final_state
+        except Exception as e: print(f"CRITICAL ERROR during graph invocation: {type(e).__name__} - {e}"); traceback.print_exc(); error_msg = AIMessage(content=f"Sorry, error occurred: {e}"); return {"messages": state.get('messages', []) + [error_msg], "patient_data": state.get('patient_data'), "summary": state.get('summary'), "interaction_warnings": None}