agent.py

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

tools = [add, multiply, divide, subtract, wikipedia_search_tool]

def build_agent():
    llm = ChatGroq(model="qwen-qwq-32b", temperature=0)
    chat_with_tools = llm.bind_tools(tools)

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

    def enhancer(state: MessagesState):
        sys_msg = """
                            You are a helpful assistant tasked with answering questions using a set of tools.
                            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.
                            Use the tool that is most appropiate on the task you are required to do. Whenever necessary, look for updated information through searches in the web, Wikipedia, Arxiv, etc.
                        """
        return {
            "messages": [sys_msg] + state["messages"]
        }

    ## The graph
    builder = StateGraph(MessagesState)
    
    # 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()