agent.py

from typing import Annotated, Literal
from typing_extensions import TypedDict

from langchain_groq import ChatGroq
from langchain_community.tools.tavily_search import TavilySearchResults
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_core.messages import ToolMessage
import config

# 1. Define State
class AgentState(TypedDict):
    messages: Annotated[list, add_messages]

# 2. Define Tools
# TavilySearchResults is a pre-built tool that searches the web
tavily_tool = TavilySearchResults(max_results=3)
tools = [tavily_tool]

# 3. Define Model
# Llama 3.3 70b is a strong model available on Groq
llm = ChatGroq(model_name="llama-3.3-70b-versatile", temperature=0)
llm_with_tools = llm.bind_tools(tools)

# 4. Define Nodes
def agent_node(state: AgentState):
    """
    The agent node that calls the LLM.
    """
    messages = state["messages"]
    response = llm_with_tools.invoke(messages)
    return {"messages": [response]}

# Implement ToolNode logic manually since import failed
def tools_node(state: AgentState):
    """
    Executes tool calls found in the last message.
    """
    last_message = state["messages"][-1]
    
    # Simple tool execution logic
    results = []
    
    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
        for tool_call in last_message.tool_calls:
            tool_name = tool_call["name"]
            tool_args = tool_call["args"]
            tool_id = tool_call["id"]
            
            # Find the tool
            selected_tool = next((t for t in tools if t.name == tool_name), None)
            
            if selected_tool:
                try:
                    # Execute tool
                    tool_output = selected_tool.invoke(tool_args)
                except Exception as e:
                    tool_output = str(e)
            else:
                tool_output = f"Error: Tool {tool_name} not found."
                
            # Create ToolMessage
            results.append(ToolMessage(
                content=str(tool_output),
                tool_call_id=tool_id
            ))
            
    return {"messages": results}

# Implement tools_condition logic manually
def custom_tools_condition(state: AgentState):
    """
    Determine if we should go to tools or end.
    """
    last_message = state["messages"][-1]
    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
        return "tools"
    return END

# 5. Build Graph
builder = StateGraph(AgentState)

builder.add_node("agent", agent_node)
builder.add_node("tools", tools_node) # Use our manual node

builder.add_edge(START, "agent")
# Use our manual condition
builder.add_conditional_edges("agent", custom_tools_condition, {"tools": "tools", END: END})
builder.add_edge("tools", "agent")

# Compile the graph
graph = builder.compile()