Agent 007¶

import os
import json
from typing import Annotated, TypedDict

from dotenv import load_dotenv
from langchain_core.messages import ToolMessage
from langchain_core.tools import Tool
from langchain_google_community import GoogleSearchAPIWrapper
from langchain_google_genai import ChatGoogleGenerativeAI
from langgraph.graph import StateGraph, END, START
from langgraph.graph.message import add_messages

# Load environment variables from .env
load_dotenv()

# We will set the GOOGLE_API_KEY and GOOGLE_CSE_ID
# in the .env file

# ... (imports and setup)

class State(TypedDict):
    """
    Represents the state of our agent.

    Attributes:
        messages: A list of messages that is dynamically updated.
                  The `add_messages` function ensures that new messages are
                  appended to the list, rather than replacing it.
    """
    messages: Annotated[list, add_messages]

# ... (after State class)

# 1. Define the primary LLM
# This is the "brain" of our agent.
llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash")

# 2. Define the tool(s)
# We're giving our agent a Google Search tool.
search = GoogleSearchAPIWrapper()
google_search_tool = Tool(
    name="google_search",
    description="Search Google for recent results.",
    func=search.run,
)
tools = [google_search_tool]

# 3. Bind the tools to the LLM
# This lets the LLM know what tools it can call.
llm_with_tools = llm.bind_tools(tools)

# 4. Define the chatbot node
# This function will be the first step in our graph.
def chatbot(state: State):
    """
    The chatbot node. This is the "brain" of the agent.
    It takes the current state and decides the next action.
    """
    print("---CHATBOT---")
    return {"messages": [llm_with_tools.invoke(state["messages"])]}

# ... (after chatbot function)

# Create a mapping of tool names to their actual functions
tool_map = {tool.name: tool for tool in tools}

def tool_node(state: State):
    """
    The tool node. This node executes the tools requested by the LLM.
    """
    print("---TOOL NODE---")
    # The last message will be the AI message with tool calls
    tool_calls = state["messages"][-1].tool_calls

    # We'll call the tools and aggregate the results
    tool_messages = []
    for tool_call in tool_calls:
        tool_name = tool_call["name"]
        print(f"Executing tool: {tool_name}")

        # Look up the tool from our map
        tool_to_call = tool_map[tool_name]

        # Execute the tool with the provided arguments
        tool_output = tool_to_call.invoke(tool_call["args"])

        # Format the output as a ToolMessage
        tool_messages.append(
            ToolMessage(
                content=json.dumps(tool_output),
                name=tool_name,
                tool_call_id=tool_call["id"],
            )
        )

    return {"messages": tool_messages}

# ... (after tool_node function)

def route_tools(state: State):
    """
    A router. This function decides which node to go to next.
    If the last message has tool calls, we go to the tool_node.
    Otherwise, we finish.
    """
    print("---ROUTING---")
    if state["messages"][-1].tool_calls:
        return "tools"
    return END

# Define the graph
graph_builder = StateGraph(State)

# Add the nodes
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_node("tools", tool_node)

# Set the entry point
graph_builder.add_edge(START, "chatbot")

# Add the conditional edge. This is the agent loop.
graph_builder.add_conditional_edges(
    "chatbot",
    route_tools,
    {"tools": "tools", END: END},
)

# After the tools are called, we always go back to the chatbot
graph_builder.add_edge("tools", "chatbot")

# Compile the graph into a runnable object
graph = graph_builder.compile()

# ... (after graph compilation)

# Let's run it!
while True:
    user_input = input("User: ")
    if user_input.lower() in ["quit", "exit", "q"]:
        print("Agent 007, signing off.")
        break

    # We stream the events from the graph
    for event in graph.stream({"messages": [("user", user_input)]}):
        # Each event is a dictionary with the node name as the key
        for key, value in event.items():
            # We'll print the final response from the chatbot
            if key == "chatbot" and value["messages"][-1].tool_calls == []:
                print(f"Assistant: {value['messages'][-1].content}")