agent.py

"""
Agent built with LangGraph.
"""

from typing import TypedDict, Annotated
from langgraph.graph import START, StateGraph, END
from langgraph.graph.message import add_messages
from langchain_huggingface import ChatHuggingFace
from langchain_groq import ChatGroq
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_core.tools import tool
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_core.messages import AnyMessage
import os

with open("system_prompt.txt") as f:
    SYSTEM_PROMPT = f.read()

@tool 
def add_tool(a: int, b: int) -> int:
    """Add two integers."""
    return a + b

@tool 
def subtract_tool(a: int, b: int) -> int:
    """Subtract two integers."""
    return a - b

@tool 
def multiply_tool(a: int, b: int) -> int:
    """Multiply two integers."""
    return a * b

@tool 
def divide_tool(a: int, b: int) -> int:
    """Divide two integers."""
    return a / b

@tool 
def modulus_tool(a: int, b: int) -> int:
    """Calculate the modulus of two integers."""
    return a % b

search_tool = DuckDuckGoSearchRun()

llm = ChatGroq(
    model="llama3-8b-8192",
    api_key=os.environ.get("GROQ_API_KEY")
)

chat = ChatHuggingFace(llm=llm, verbose=True)
tools = [add_tool, subtract_tool, multiply_tool, divide_tool, modulus_tool, search_tool]
chat_with_tools = chat.bind_tools(tools)

class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]

def assistant(state: AgentState):
    return {
        "messages": [chat_with_tools.invoke(state["messages"])],
    }

builder = StateGraph(AgentState)
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))
builder.add_edge(START, "assistant")
builder.add_conditional_edges(
    "assistant",
    tools_condition,
)
builder.add_edge("tools", "assistant")
graph = builder.compile()