FinalAgentProject / agent.py
WeByT3's picture
Update agent.py
c8b3fa3 verified
raw
history blame
3.62 kB
import os
from langchain_groq import ChatGroq
from langgraph.graph import START, StateGraph, MessagesState
from langgraph.prebuilt import tools_condition
from langgraph.prebuilt import ToolNode
from langchain_core.messages import SystemMessage, HumanMessage
from tools import *
from typing import TypedDict, Optional
tools = [add, multiply, divide, subtract, search_wikipedia]
def build_agent():
class AgentState(TypedDict):
messages: list # chat history
wiki_table: Optional[str]
llm = ChatGroq(model="qwen-qwq-32b", temperature=0)
chat_with_tools = llm.bind_tools(tools)
def assistant(state: AgentState):
table = state.get("wiki_table")
context = []
if table:
context.append({"role": "system", "content": f"Here is a Wikipedia table you can use:\n\n{table}"})
response = chat_with_tools.invoke(context + state["messages"])
return {
"messages": [response],
"wiki_table": table,
}
def enhancer(state: AgentState):
sys_msg = """
You are a helpful assistant tasked with answering questions using a set of tools. You can reason through problems and take actions using available tools to solve complex tasks.
Follow this format strictly:
Message: {user question}
Thought: Describe your reasoning about the question.
Action: Call one of the available tools with a specific input.
Observation: Note what was returned from the tool.
Repeat the Thought → Action → Observation steps as many times as necessary until you have enough information to answer the original question.
Now, I will ask you a question. Report your thoughts, and finish your answer with the following template:
FINAL ANSWER: [YOUR FINAL ANSWER].
YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.
If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
Your answer should only start with "FINAL ANSWER: ", then follows with the answer.
"""
context = []
context.append({"role": "system", "content": sys_msg})
return context + state["messages"]
## The graph
builder = StateGraph(AgentState)
# Define nodes: these do the work
builder.add_node("enhancer", enhancer)
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))
# Define edges: these determine how the control flow moves
builder.add_edge(START, "enhancer")
builder.add_edge("enhancer", "assistant")
builder.add_conditional_edges(
"assistant",
# If the latest message requires a tool, route to tools
# Otherwise, provide a direct response
tools_condition,
)
builder.add_edge("tools", "assistant")
return builder.compile()