Update agent.py

agent.py

@@ -1,5 +1,3 @@
-# agent.py
-
 import os
 import re
 import json
@@ -54,7 +52,6 @@ def get_rxcui(drug_name: str) -> Optional[str]:
         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()
@@ -62,7 +59,6 @@ def get_rxcui(drug_name: str) -> Optional[str]:
         if ids:
             logger.info(f"Found RxCUI {ids[0]} for '{drug_name}'")
             return ids[0]
-        # Fallback search
         r = requests.get(f"{RXNORM_API_BASE}/drugs.json", params={"name": drug_name}, timeout=10)
         r.raise_for_status()
         for grp in r.json().get("drugGroup", {}).get("conceptGroup", []):
@@ -126,8 +122,6 @@ def check_red_flags(patient_data: Dict[str, Any]) -> List[str]:
     hpi    = patient_data.get("hpi", {})
     vitals = patient_data.get("vitals", {})
     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",
@@ -138,14 +132,11 @@ def check_red_flags(patient_data: Dict[str, Any]) -> List[str]:
     for term, desc in mapping.items():
         if term in syms:
             flags.append(f"Red Flag: {desc}.")
-
-    # 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:
@@ -163,7 +154,6 @@ 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))  # dedupe, preserve order
 
 def format_patient_data_for_prompt(data: Dict[str, Any]) -> str:
@@ -210,7 +200,6 @@ class FlagRiskInput(BaseModel):
 # ── 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",
@@ -227,7 +216,6 @@ 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",
@@ -241,31 +229,22 @@ 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 exact match
     if potential_prescription.lower().strip() in al:
         warnings.append(f"CRITICAL ALLERGY: Patient allergic to '{potential_prescription}'.")
-
-    # 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):
         warnings.append(f"INFO: Could not identify '{potential_prescription}'. Checks may be incomplete.")
-
-    # Contraindications & warnings sections
     for section in ("contraindications", "warnings_and_cautions", "warnings"):
         items = label.get(section) if label else None
         if isinstance(items, list):
             snippets = search_text_list(items, al)
             if snippets:
                 warnings.append(f"ALLERGY RISK ({section}): {'; '.join(snippets)}")
-
-    # Drug–drug interactions
     for med in pm:
         mrxcui = get_rxcui(med)
         mlabel = get_openfda_label(rxcui=mrxcui, drug_name=med)
@@ -276,7 +255,6 @@ def check_drug_interactions(
                     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"
     message = (
         f"{len(warnings)} issue(s) found for '{potential_prescription}'."
@@ -287,7 +265,6 @@ 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",
@@ -309,6 +286,7 @@ class AgentState(TypedDict):
     patient_data: Optional[Dict[str, Any]]
     summary: Optional[str]
     interaction_warnings: Optional[List[str]]
+    done: Optional[bool]
 
 # ── Graph Nodes ───────────────────────────────────────────────────────────────
 def agent_node(state: AgentState) -> Dict[str, Any]:
@@ -328,7 +306,6 @@ def tool_node(state: AgentState) -> Dict[str, Any]:
     if not isinstance(last, AIMessage) or not getattr(last, "tool_calls", None):
         logger.warning("tool_node invoked without pending tool_calls")
         return {"messages": [], "interaction_warnings": None}
-
     calls = last.tool_calls
     blocked_ids = set()
     for call in calls:
@@ -341,14 +318,12 @@ def tool_node(state: AgentState) -> Dict[str, Any]:
             ):
                 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]
     pd = state.get("patient_data", {})
     for call in to_execute:
         if call["name"] == "check_drug_interactions":
             call["args"].setdefault("current_medications", pd.get("medications", {}).get("current", []))
             call["args"].setdefault("allergies", pd.get("allergies", []))
-
     messages: List[ToolMessage] = []
     warnings: List[str]     = []
     try:
@@ -379,16 +354,13 @@ def reflection_node(state: AgentState) -> Dict[str, Any]:
     if not warns:
         logger.warning("reflection_node called without warnings")
         return {"messages": [], "interaction_warnings": None}
-
     triggering = None
     for msg in reversed(state["messages"]):
         if isinstance(msg, AIMessage) and getattr(msg, "tool_calls", None):
             triggering = msg
             break
-
     if not triggering:
         return {"messages": [AIMessage(content="Internal Error: reflection context missing.")], "interaction_warnings": None}
-
     prompt = (
         "You are SynapseAI, performing a focused safety review of the following plan:\n\n"
         f"{triggering.content}\n\n"
@@ -406,9 +378,13 @@ def reflection_node(state: AgentState) -> Dict[str, Any]:
 def should_continue(state: AgentState) -> str:
     last = state["messages"][-1] if state["messages"] else None
     if not isinstance(last, AIMessage):
+        # Mark conversation as done
+        state["done"] = True
         return "end_conversation_turn"
     if getattr(last, "tool_calls", None):
         return "continue_tools"
+    # No further tool calls – conversation is finished.
+    state["done"] = True
     return "end_conversation_turn"
 
 def after_tools_router(state: AgentState) -> str:
@@ -432,6 +408,8 @@ class ClinicalAgent:
             "agent": "agent"
         })
         wf.add_edge("reflection", "agent")
+        # Set termination condition: stop when state["done"] is True.
+        wf.set_termination_condition(lambda state: state.get("done", False))
         self.graph_app = wf.compile()
         logger.info("ClinicalAgent ready")