refactor chat_with_tools agent

Author: ks6088ts

template_langgraph/agents/chat_with_tools_agent/agent.py

@@ -46,7 +46,6 @@ def create_graph(self):
         workflow = StateGraph(AgentState)
 
         # Create nodes
-        workflow.add_node("initialize", self.initialize)
         workflow.add_node("chat_with_tools", self.chat_with_tools)
         workflow.add_node(
             "tools",
@@ -57,35 +56,24 @@ def create_graph(self):
                 ]
             ),
         )
-        workflow.add_node("finalize", self.finalize)
 
         # Create edges
-        workflow.set_entry_point("initialize")
-        workflow.add_edge("initialize", "chat_with_tools")
+        workflow.set_entry_point("chat_with_tools")
         workflow.add_conditional_edges(
-            "chat_with_tools",
-            self.route_tools,
-            # The following dictionary lets you tell the graph to interpret the condition's outputs as a specific node
-            # It defaults to the identity function, but if you
-            # want to use a node named something else apart from "tools",
-            # You can update the value of the dictionary to something else
-            # e.g., "tools": "my_tools"
-            {"tools": "tools", END: "finalize"},
+            source="chat_with_tools",
+            path=self.route_tools,
+            path_map={
+                "tools": "tools",
+                END: END,
+            },
         )
         workflow.add_edge("tools", "chat_with_tools")
-        workflow.add_edge("finalize", END)
 
         # Compile the graph
         return workflow.compile(
             name=ChatWithToolsAgent.__name__,
         )
 
-    def initialize(self, state: AgentState) -> AgentState:
-        """Initialize the agent with the given state."""
-        logger.info(f"Initializing ChatWithToolsAgent with state: {state}")
-        # Here you can add any initialization logic if needed
-        return state
-
     def chat_with_tools(self, state: AgentState) -> AgentState:
         """Chat with tools using the state."""
         logger.info(f"Chatting with tools using state: {state}")
@@ -119,15 +107,5 @@ def route_tools(
             return "tools"
         return END
 
-    def finalize(self, state: AgentState) -> AgentState:
-        """Finalize the agent's work and prepare the output."""
-        logger.info(f"Finalizing ChatWithToolsAgent with state: {state}")
-        # Here you can add any finalization logic if needed
-        return state
-
-    def draw_mermaid_png(self) -> bytes:
-        """Draw the graph in Mermaid format."""
-        return self.create_graph().get_graph().draw_mermaid_png()
-
 
 graph = ChatWithToolsAgent().create_graph()

template_langgraph/agents/chat_with_tools_agent/models.py

@@ -8,15 +8,6 @@
     BaseMessage,
 )
 from langgraph.graph.message import add_messages
-from pydantic import BaseModel, Field
-
-
-class AgentInput(BaseModel):
-    request: str = Field(..., description="ユーザーからのリクエスト")
-
-
-class AgentOutput(BaseModel):
-    response: str = Field(..., description="エージェントの応答")
 
 
 class AgentState(TypedDict):

template_langgraph/tasks/draw_chat_with_tools_agent_mermaid_png.py

@@ -1,11 +1,11 @@
 import sys
 
-from template_langgraph.agents.chat_with_tools_agent.agent import ChatWithToolsAgent
+from template_langgraph.agents.chat_with_tools_agent.agent import graph
 
 if __name__ == "__main__":
     png_path = "data/chat_with_tools_agent.png"
     if len(sys.argv) > 1:
         png_path = sys.argv[1]
 
     with open(png_path, "wb") as f:
-        f.write(ChatWithToolsAgent().draw_mermaid_png())
+        f.write(graph.get_graph().draw_mermaid_png())

template_langgraph/tasks/run_chat_with_tools_agent.py

@@ -19,15 +19,15 @@ def stream_graph_updates(
 
 if __name__ == "__main__":
     user_input = input("User: ")
-    state = AgentState(
-        messages=[
-            {
-                "role": "user",
-                "content": user_input,
-            }
-        ],
-        profile=None,
-    )
-    last_event = stream_graph_updates(state)
-    for value in last_event.values():
-        logger.info(f"Final state: {value['messages'][-1].content}")  # noqa: E501
+    for event in chat_with_tools_agent_graph.stream(
+        input=AgentState(
+            messages=[
+                {
+                    "role": "user",
+                    "content": user_input,
+                }
+            ],
+        )
+    ):
+        logger.info("-" * 20)
+        logger.info(f"Event: {event}")