agent.py

import os
from dotenv import load_dotenv
from langchain.tools import Tool
from langchain.utilities import SerpAPIWrapper
from langgraph.graph.message import add_messages
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_core.messages import AnyMessage, HumanMessage
from langchain_groq import ChatGroq
from typing import TypedDict, Annotated

# Load environment variables
load_dotenv()
groq_api_key = os.getenv("GROQ_API_KEY")
serpapi_api_key = os.getenv("SERPAPI_API_KEY")

# --- Tool Definitions ---
def calculator_tool_func(query: str):
    try:
        result = str(eval(query, {"__builtins__": {}}))
        return result
    except Exception:
        return "Could not calculate the input expression."

calculator_tool = Tool(
    name="Calculator",
    func=calculator_tool_func,
    description="Performs basic arithmetic calculations."
)

def weather_tool_func(location: str):
    return f"The weather in {location}: Sunny, 25°C."

weather_tool = Tool(
    name="Weather",
    func=weather_tool_func,
    description="Provides mock weather information for a given location."
)

serpapi = SerpAPIWrapper(serpapi_api_key=serpapi_api_key)
web_search_tool = Tool(
    name="WebSearch",
    func=serpapi.run,
    description="Searches the web for recent information using SerpAPI."
)

# --- Model and Tools Binding ---
tools = [web_search_tool, calculator_tool, weather_tool]
llm = ChatGroq(model="deepseek-r1-distill-llama-70b", groq_api_key=groq_api_key)
llm_with_tools = llm.bind_tools(tools)

# --- Agent State ---
class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]

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

# --- Agent Graph ---
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")

# 👇 The compiled agent: use with .invoke({...})
ninu = builder.compile()

# --- Run Function ---
def run_ninu(query: str):
    conversation = []

    system_prompt = """
You are a general AI assistant. 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.
"""
    conversation.append(HumanMessage(content=system_prompt.strip()))
    conversation.append(HumanMessage(content=query))

    response = ninu.invoke({"messages": conversation})
    for message in reversed(response["messages"]):
        if isinstance(message.content, str) and "FINAL ANSWER:" in message.content:
            return message.content.split("FINAL ANSWER:")[-1].strip()
    return "FINAL ANSWER: No valid answer found."

# --- Example test ---
if __name__ == "__main__":
    result = run_ninu("What is the capital of Japan?")
    print("🧠 NINU replied:", result)