Update agent.py

agent.py

@@ -1,9 +1,11 @@
 import os
 
 from langchain_openai import ChatOpenAI
-from langchain_core.messages import SystemMessage, HumanMessage, ToolMessage
+from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, ToolMessage
+from langgraph.graph.message import add_messages
 from langgraph.graph import MessagesState
 from langgraph.graph import StateGraph, START, END
+from typing import Literal
 
 from langchain_community.tools import BraveSearch  # web search
 from langchain.tools import Calculator  # for basic math
@@ -29,7 +31,41 @@ search_tool = BraveSearch.from_api_key(
 )
 
 community_tools = [calculator_tool, python_tool, search_tool]
-custom_tools = [datetime_tools, transcribe_audio_tool].
+custom_tools = [datetime_tools, transcribe_audio_tool]
 
 tools = community_tools + custom_tools
-llm_with_tools = llm.bind_tools(tools)
+llm_with_tools = llm.bind_tools(tools)
+
+tools_by_name = {tool.name: tool for tool in tools}
+
+class MessagesState(TypedDict):  # creates the state (is like the agent's memory at any moment)
+    messages: Annotated[list[AnyMessage], add_messages]
+
+# LLM node
+def llm_call(state: MessagesState):
+    return {
+        "messages": [
+            llm_with_tools.invoke(
+                [SystemMessage(content=system_prompt)] + state["messages"]
+            )
+        ]
+    }
+
+# Tool node
+def tool_node(state: MessagesState):
+    result = []
+    for tool_call in state["messages"][-1].tool_calls:
+        tool = tools_by_name[tool_call["name"]]
+        observation = tool.invoke(tool_call["args"])
+        result.append(ToolMessage(content=observation, tool_call_id=tool_call["id"]))
+    return {"messages": result}
+
+def should_continue(state: MessagesState) -> Literal["Action", END]:
+    """Decide if we should continue the loop or stop based upon whether the LLM made a tool call"""
+
+    last_message = state["messages"][-1]
+    # If the LLM makes a tool call, then perform an action
+    if last_message.tool_calls:
+        return "Action"
+    # Otherwise, we stop (reply to the user)
+    return END