Update agent.py

agent.py

@@ -43,8 +43,8 @@ class ClinicalPrompts:
 def wrap_message(msg: Any) -> AIMessage:
     """
     Ensures the given message is an AIMessage.
-    If it is a dict, extracts the 'content' field (or serializes the dict).
-    Otherwise, converts the message to a string.
+    If it is a dict, it extracts the 'content' field (or serializes the dict).
+    Otherwise, it converts the message to a string.
     """
     if isinstance(msg, AIMessage):
         return msg
@@ -358,7 +358,6 @@ def tool_node(state: AgentState) -> Dict[str, Any]:
         new_state = {"messages": []}
         return propagate_state(new_state, state)
     last = wrap_message(messages_list[-1])
-    # Safely retrieve pending tool_calls
     tool_calls = last.__dict__.get("tool_calls")
     if not (isinstance(last, AIMessage) and tool_calls):
         logger.warning("tool_node invoked without pending tool_calls")
@@ -464,7 +463,7 @@ def should_continue(state: AgentState) -> str:
         state["done"] = True
         return "end_conversation_turn"
     state["done"] = False
-    return "agent"
+    return "start"
 
 def after_tools_router(state: AgentState) -> str:
     if state.get("interaction_warnings"):
@@ -476,11 +475,11 @@ class ClinicalAgent:
     def __init__(self):
         logger.info("Building ClinicalAgent workflow")
         wf = StateGraph(AgentState)
-        wf.add_node("agent", agent_node)
+        wf.add_node("start", agent_node)
         wf.add_node("tools", tool_node)
         wf.add_node("reflection", reflection_node)
-        wf.set_entry_point("agent")
-        wf.add_conditional_edges("agent", should_continue, {
+        wf.set_entry_point("start")
+        wf.add_conditional_edges("start", should_continue, {
             "continue_tools": "tools",
             "end_conversation_turn": END
         })
@@ -488,7 +487,7 @@ class ClinicalAgent:
             "reflection": "reflection",
             "end_conversation_turn": END
         })
-        # Removed edge from reflection back to agent.
+        # Removed edge from reflection back to start.
         self.graph_app = wf.compile()
         logger.info("ClinicalAgent ready")