File size: 3,541 Bytes
1033417
 
 
 
 
 
 
 
 
 
8987ddb
409dd63
1033417
 
 
6306cd3
409dd63
1033417
 
 
 
 
409dd63
6c496de
1033417
 
 
409dd63
1033417
6c496de
1033417
409dd63
1033417
 
 
 
409dd63
1033417
 
 
409dd63
1033417
 
409dd63
1033417
 
409dd63
 
1033417
 
 
409dd63
1033417
 
 
 
 
 
409dd63
1033417
 
 
 
 
 
 
409dd63
1033417
 
409dd63
1033417
 
 
409dd63
 
1033417
409dd63
1033417
 
 
5ffb4a3
 
 
409dd63
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
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)