| import os |
| from datetime import datetime, timezone |
| from dotenv import load_dotenv |
| from colorama import Fore, Style |
| from langchain.agents import create_agent |
| from langchain_core.messages import HumanMessage |
| from agent.tools.math_solver import math_solver |
| from agent.agents.websearch import websearch_agent |
|
|
| load_dotenv() |
|
|
|
|
| def supervisor_agent(): |
| """Return a supervisor agent instance with math_solver and websearch_agent.""" |
| return create_agent( |
| model="google_genai:gemini-3-flash-preview", |
| tools=[math_solver, websearch_agent], |
| system_prompt=( |
| f"You are a supervisor agent. " |
| f"Current time is: {datetime.now(timezone.utc).isoformat()}. " |
| f"Your memory are out of date. " |
| f"For math or calculation questions, use the math_solver tool. " |
| f"For questions that need real-time, use the websearch_agent tool. " |
| f"Provide a concise and accurate final answer." |
| ), |
| ) |
|
|
|
|
| def run(query: str) -> str: |
| """Entry point: let the supervisor agent finish the work.""" |
| print(f"{Fore.CYAN}[Supervisor] Processing query...{Style.RESET_ALL}") |
| agent = supervisor_agent() |
| result = agent.invoke({"messages": [HumanMessage(content=query)]}) |
| content = result["messages"][-1].content |
| if isinstance(content, list): |
| return content[0].get("text", "") |
| return str(content) |
|
|
|
|
| if __name__ == "__main__": |
| agent = supervisor_agent() |
| chat_history: list = [] |
| while True: |
| query = input("\nYou: ") |
| if query.lower() in ("exit", "quit"): |
| break |
| chat_history.append(HumanMessage(content=query)) |
| result = agent.invoke({"messages": chat_history}) |
| chat_history = result["messages"] |
| content = chat_history[-1].content |
| if isinstance(content, list): |
| content = content[0].get("text", "") |
| print(f"Agent: {content}") |
|
|
