Update agent.py

agent.py

@@ -198,8 +198,9 @@ class FlagRiskInput(BaseModel):
     urgency: str          = Field("High")
 
 # ── Tool Implementations ──────────────────────────────────────────────────────
-@tool("order_lab_test", args_schema=LabOrderInput)
+@tool("order_lab_test", args_schema=LabOrderInput, description="Place an order for a laboratory test.")
 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",
@@ -207,7 +208,7 @@ def order_lab_test(test_name: str, reason: str, priority: str = "Routine") -> st
         "details": f"Reason: {reason}"
     })
 
-@tool("prescribe_medication", args_schema=PrescriptionInput)
+@tool("prescribe_medication", args_schema=PrescriptionInput, description="Prepare a medication prescription.")
 def prescribe_medication(
     medication_name: str,
     dosage: str,
@@ -216,6 +217,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",
@@ -223,12 +225,13 @@ def prescribe_medication(
         "details": f"Duration: {duration}. Reason: {reason}"
     })
 
-@tool("check_drug_interactions", args_schema=InteractionCheckInput)
+@tool("check_drug_interactions", args_schema=InteractionCheckInput, description="Check for drug–drug interactions and allergy risks.")
 def check_drug_interactions(
     potential_prescription: str,
     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]
@@ -263,8 +266,9 @@ def check_drug_interactions(
     )
     return json.dumps({"status": status, "message": message, "warnings": warnings})
 
-@tool("flag_risk", args_schema=FlagRiskInput)
+@tool("flag_risk", args_schema=FlagRiskInput, description="Flag a clinical risk with given urgency.")
 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",
@@ -378,12 +382,10 @@ 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"
 
@@ -408,7 +410,6 @@ 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")