antonioschiro's picture
Update agent.py
2171e7d verified
import os
from typing import TypedDict, Annotated
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, ToolMessage, SystemMessage
from langchain_groq import ChatGroq
from langgraph.prebuilt import ToolNode
from langgraph.graph import START, StateGraph, END
from langgraph.prebuilt import tools_condition
from tools import (retriever, web_search, wiki_search, youtube_analysis,
add_numbers, subtract_numbers, multiply_numbers, divide_numbers, modulus_numbers,
detect_objects, run_python
)
from prompt import text_prompt
from dotenv import load_dotenv
#from PIL import Image
#from io import StringIO
#load_dotenv()
os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")
class State(TypedDict):
messages: Annotated[list[AnyMessage], add_messages]
tools = [retriever, web_search, wiki_search, youtube_analysis,
add_numbers, subtract_numbers, multiply_numbers, divide_numbers, modulus_numbers,
detect_objects, run_python]
#model = "qwen/qwen3-32b"
model = "deepseek-r1-distill-llama-70b"
llm = ChatGroq(
model= model,
temperature=0.0,
max_tokens= None,
reasoning_format="parsed",
timeout=None,
max_retries=2,
)
llm_with_tools = llm.bind_tools(tools)
def ask_agent(agent_state: State):
system_prompt = SystemMessage(
content = text_prompt
)
query = agent_state["messages"][-1] # HumanMessage
response = llm_with_tools.invoke(text_prompt + query.content)
return {"messages": [response]}
graph_builder = StateGraph(State)
graph_builder.add_node("agent", ask_agent)
graph_builder.add_node("tools", ToolNode(tools))
graph_builder.add_edge(START, "agent")
graph_builder.add_conditional_edges(
"agent",
tools_condition
)
graph_builder.add_edge("tools", "agent")
alfred = graph_builder.compile()