Commit ·
40deb66
1
Parent(s): a877f54
feat(agent): implement supervisor agent with tool routing and math solver
Browse files- Replace single websearch agent with supervisor agent architecture
- Add math_solver tool using sympy for calculation queries
- Implement tool routing logic (math vs web search) in supervisor prompt
- Add interactive chat loop with conversation history
- Enhance terminal output with colored logging (colorama)
- Update dependencies: sympy, colorama
- agent/__pycache__/agent.cpython-314.pyc +0 -0
- agent/agent.py +38 -25
- agent/agents/__init__.py +3 -0
- agent/agents/__pycache__/__init__.cpython-314.pyc +0 -0
- agent/agents/__pycache__/websearch.cpython-314.pyc +0 -0
- agent/agents/websearch.py +47 -0
- agent/tools/__init__.py +2 -1
- agent/tools/__pycache__/__init__.cpython-314.pyc +0 -0
- agent/tools/__pycache__/math_solver.cpython-314.pyc +0 -0
- agent/tools/__pycache__/search.cpython-314.pyc +0 -0
- agent/tools/math_solver.py +47 -0
- agent/tools/search.py +3 -2
- requirements.txt +2 -0
agent/__pycache__/agent.cpython-314.pyc
CHANGED
|
Binary files a/agent/__pycache__/agent.cpython-314.pyc and b/agent/__pycache__/agent.cpython-314.pyc differ
|
|
|
agent/agent.py
CHANGED
|
@@ -1,40 +1,53 @@
|
|
| 1 |
import os
|
|
|
|
| 2 |
from dotenv import load_dotenv
|
|
|
|
| 3 |
from langchain.agents import create_agent
|
| 4 |
-
from
|
|
|
|
|
|
|
| 5 |
|
| 6 |
load_dotenv()
|
| 7 |
|
| 8 |
|
| 9 |
-
def
|
| 10 |
-
"""
|
| 11 |
-
|
| 12 |
model="google_genai:gemini-3-flash-preview",
|
| 13 |
-
tools=[
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 14 |
)
|
| 15 |
-
result = base_agent.invoke(
|
| 16 |
-
{
|
| 17 |
-
"messages": [
|
| 18 |
-
{
|
| 19 |
-
"role": "system",
|
| 20 |
-
"content": "Try to search web sites and get the answer.",
|
| 21 |
-
},
|
| 22 |
-
{"role": "user", "content": query},
|
| 23 |
-
]
|
| 24 |
-
}
|
| 25 |
-
)
|
| 26 |
-
content = result["messages"][-1].content[0]
|
| 27 |
-
return content["text"]
|
| 28 |
|
| 29 |
|
| 30 |
def run(query: str) -> str:
|
| 31 |
-
"""
|
| 32 |
-
|
| 33 |
-
|
| 34 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 35 |
|
| 36 |
|
| 37 |
if __name__ == "__main__":
|
| 38 |
-
|
| 39 |
-
|
| 40 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
import os
|
| 2 |
+
from datetime import datetime, timezone
|
| 3 |
from dotenv import load_dotenv
|
| 4 |
+
from colorama import Fore, Style
|
| 5 |
from langchain.agents import create_agent
|
| 6 |
+
from langchain_core.messages import HumanMessage
|
| 7 |
+
from agent.tools.math_solver import math_solver
|
| 8 |
+
from agent.agents.websearch import websearch_agent
|
| 9 |
|
| 10 |
load_dotenv()
|
| 11 |
|
| 12 |
|
| 13 |
+
def supervisor_agent():
|
| 14 |
+
"""Return a supervisor agent instance with math_solver and websearch_agent."""
|
| 15 |
+
return create_agent(
|
| 16 |
model="google_genai:gemini-3-flash-preview",
|
| 17 |
+
tools=[math_solver, websearch_agent],
|
| 18 |
+
system_prompt=(
|
| 19 |
+
f"You are a supervisor agent. "
|
| 20 |
+
f"Current time is: {datetime.now(timezone.utc).isoformat()}. "
|
| 21 |
+
f"Your memory are out of date. "
|
| 22 |
+
f"For math or calculation questions, use the math_solver tool. "
|
| 23 |
+
f"For questions that need real-time, use the websearch_agent tool. "
|
| 24 |
+
f"Provide a concise and accurate final answer."
|
| 25 |
+
),
|
| 26 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 27 |
|
| 28 |
|
| 29 |
def run(query: str) -> str:
|
| 30 |
+
"""Entry point: let the supervisor agent finish the work."""
|
| 31 |
+
print(f"{Fore.CYAN}[Supervisor] Processing query...{Style.RESET_ALL}")
|
| 32 |
+
agent = supervisor_agent()
|
| 33 |
+
result = agent.invoke({"messages": [HumanMessage(content=query)]})
|
| 34 |
+
content = result["messages"][-1].content
|
| 35 |
+
if isinstance(content, list):
|
| 36 |
+
return content[0].get("text", "")
|
| 37 |
+
return str(content)
|
| 38 |
|
| 39 |
|
| 40 |
if __name__ == "__main__":
|
| 41 |
+
agent = supervisor_agent()
|
| 42 |
+
chat_history: list = []
|
| 43 |
+
while True:
|
| 44 |
+
query = input("\nYou: ")
|
| 45 |
+
if query.lower() in ("exit", "quit"):
|
| 46 |
+
break
|
| 47 |
+
chat_history.append(HumanMessage(content=query))
|
| 48 |
+
result = agent.invoke({"messages": chat_history})
|
| 49 |
+
chat_history = result["messages"]
|
| 50 |
+
content = chat_history[-1].content
|
| 51 |
+
if isinstance(content, list):
|
| 52 |
+
content = content[0].get("text", "")
|
| 53 |
+
print(f"Agent: {content}")
|
agent/agents/__init__.py
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from agent.agents.websearch import websearch_agent
|
| 2 |
+
|
| 3 |
+
__all__ = ["websearch_agent"]
|
agent/agents/__pycache__/__init__.cpython-314.pyc
ADDED
|
Binary file (249 Bytes). View file
|
|
|
agent/agents/__pycache__/websearch.cpython-314.pyc
ADDED
|
Binary file (2.58 kB). View file
|
|
|
agent/agents/websearch.py
ADDED
|
@@ -0,0 +1,47 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from datetime import datetime, timezone
|
| 2 |
+
from colorama import Fore, Style # type: ignore[import]
|
| 3 |
+
from langchain_core.tools import tool
|
| 4 |
+
from langchain.agents import create_agent
|
| 5 |
+
from agent.tools.search import web_search
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
@tool
|
| 9 |
+
def websearch_agent(query: str) -> str:
|
| 10 |
+
"""
|
| 11 |
+
A web search agent that searches the internet and returns an answer.
|
| 12 |
+
Use this tool when you need to find real-time or factual information
|
| 13 |
+
from the web, such as current events, specific facts, or any
|
| 14 |
+
knowledge that may require up-to-date sources.
|
| 15 |
+
|
| 16 |
+
Args:
|
| 17 |
+
query: The question or search query to look up on the web.
|
| 18 |
+
"""
|
| 19 |
+
print(f"{Fore.YELLOW}[SupervisorAgent -> WebSearchAgent] {query}{Style.RESET_ALL}")
|
| 20 |
+
base_agent = create_agent(
|
| 21 |
+
model="google_genai:gemini-3-flash-preview",
|
| 22 |
+
tools=[web_search],
|
| 23 |
+
system_prompt=(
|
| 24 |
+
f"Current time is: {datetime.now(timezone.utc).isoformat()}. "
|
| 25 |
+
f"Your memory are out of date. "
|
| 26 |
+
f"All of truth that you believe without search are wrong. "
|
| 27 |
+
f"You must search the web and find the lastest answer."
|
| 28 |
+
),
|
| 29 |
+
)
|
| 30 |
+
result = base_agent.invoke({"messages": [{"role": "user", "content": query}]})
|
| 31 |
+
content = result["messages"][-1].content
|
| 32 |
+
if isinstance(content, list):
|
| 33 |
+
content = content[0].get("text", "")
|
| 34 |
+
else:
|
| 35 |
+
content = str(content)
|
| 36 |
+
print(
|
| 37 |
+
f"{Fore.YELLOW}[WebSearchAgent -> SupervisorAgent] {content}{Style.RESET_ALL}"
|
| 38 |
+
)
|
| 39 |
+
return content
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
if __name__ == "__main__":
|
| 43 |
+
from dotenv import load_dotenv
|
| 44 |
+
|
| 45 |
+
load_dotenv()
|
| 46 |
+
answer = websearch_agent.invoke({"query": "What is LangGraph?"})
|
| 47 |
+
print(answer)
|
agent/tools/__init__.py
CHANGED
|
@@ -1,3 +1,4 @@
|
|
| 1 |
from agent.tools.search import web_search
|
|
|
|
| 2 |
|
| 3 |
-
__all__ = ["web_search"]
|
|
|
|
| 1 |
from agent.tools.search import web_search
|
| 2 |
+
from agent.tools.math_solver import math_solver
|
| 3 |
|
| 4 |
+
__all__ = ["web_search", "math_solver"]
|
agent/tools/__pycache__/__init__.cpython-314.pyc
CHANGED
|
Binary files a/agent/tools/__pycache__/__init__.cpython-314.pyc and b/agent/tools/__pycache__/__init__.cpython-314.pyc differ
|
|
|
agent/tools/__pycache__/math_solver.cpython-314.pyc
ADDED
|
Binary file (2.42 kB). View file
|
|
|
agent/tools/__pycache__/search.cpython-314.pyc
CHANGED
|
Binary files a/agent/tools/__pycache__/search.cpython-314.pyc and b/agent/tools/__pycache__/search.cpython-314.pyc differ
|
|
|
agent/tools/math_solver.py
ADDED
|
@@ -0,0 +1,47 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from colorama import Fore, Style # type: ignore[import]
|
| 2 |
+
from langchain_core.tools import tool
|
| 3 |
+
from sympy import sympify, N # type: ignore[import]
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
@tool
|
| 7 |
+
def math_solver(expression: str) -> str:
|
| 8 |
+
"""
|
| 9 |
+
Evaluate a mathematical expression using SymPy symbolic computation.
|
| 10 |
+
|
| 11 |
+
This tool converts a string expression into a SymPy symbolic object and
|
| 12 |
+
evaluates it with arbitrary precision. Unlike standard floating-point
|
| 13 |
+
arithmetic, SymPy uses exact rational and symbolic representations
|
| 14 |
+
internally, so results are free from floating-point rounding errors.
|
| 15 |
+
|
| 16 |
+
Supported syntax includes:
|
| 17 |
+
- Basic arithmetic: +, -, *, /, ** (power)
|
| 18 |
+
- Fractions: "1/3 + 2/7" (exact rational arithmetic)
|
| 19 |
+
- Square roots: "sqrt(2)", "sqrt(3)/2"
|
| 20 |
+
- Trigonometric functions: "sin(pi/4)", "cos(pi/3)", "tan(pi/6)"
|
| 21 |
+
- Logarithms: "log(100)", "ln(e**2)"
|
| 22 |
+
- Constants: "pi", "E" (Euler's number)
|
| 23 |
+
- Factorials: "factorial(10)"
|
| 24 |
+
- Combinations: "binomial(10, 3)"
|
| 25 |
+
- Complex expressions: "1/3 + 2/7 * sqrt(2)"
|
| 26 |
+
|
| 27 |
+
Args:
|
| 28 |
+
expression: A mathematical expression string to evaluate.
|
| 29 |
+
Example: "1/3 + 2/7 * sqrt(2)"
|
| 30 |
+
|
| 31 |
+
Returns:
|
| 32 |
+
A string containing both the exact symbolic result and a numerical
|
| 33 |
+
approximation to 10 decimal places.
|
| 34 |
+
"""
|
| 35 |
+
try:
|
| 36 |
+
expr = sympify(expression)
|
| 37 |
+
exact = str(expr)
|
| 38 |
+
numeric = str(N(expr, 10))
|
| 39 |
+
print(f"{Fore.MAGENTA}[MathSolver] {expression}{Style.RESET_ALL} = {numeric}")
|
| 40 |
+
return f"Exact: {exact}\nNumeric (10 digits): {numeric}"
|
| 41 |
+
except Exception as e:
|
| 42 |
+
return f"Error evaluating expression: {e}"
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
if __name__ == "__main__":
|
| 46 |
+
result = math_solver.invoke({"expression": "1/3 + 2/7 * sqrt(2)"})
|
| 47 |
+
print(result)
|
agent/tools/search.py
CHANGED
|
@@ -1,4 +1,5 @@
|
|
| 1 |
import os
|
|
|
|
| 2 |
from langchain_core.tools import tool
|
| 3 |
from tavily import TavilyClient # type: ignore[import]
|
| 4 |
|
|
@@ -14,7 +15,7 @@ def web_search(query: str, depth: str = "advanced", max_results: int = 5) -> str
|
|
| 14 |
depth: Search depth, either "basic" or "advanced".
|
| 15 |
max_results: Number of results to return (recommended 3-5 to save tokens).
|
| 16 |
"""
|
| 17 |
-
print(f"[Search & Extract] {query}")
|
| 18 |
client = TavilyClient(api_key=os.environ["TAVILY_API_KEY"])
|
| 19 |
|
| 20 |
# 1. Search and get URLs
|
|
@@ -40,7 +41,7 @@ def web_search(query: str, depth: str = "advanced", max_results: int = 5) -> str
|
|
| 40 |
item["url"]: item["raw_content"] for item in extraction.get("results", [])
|
| 41 |
}
|
| 42 |
except Exception as e:
|
| 43 |
-
print(f"Extraction failed: {e}")
|
| 44 |
extracted_results = {}
|
| 45 |
|
| 46 |
# 3. Format output
|
|
|
|
| 1 |
import os
|
| 2 |
+
from colorama import Fore, Style
|
| 3 |
from langchain_core.tools import tool
|
| 4 |
from tavily import TavilyClient # type: ignore[import]
|
| 5 |
|
|
|
|
| 15 |
depth: Search depth, either "basic" or "advanced".
|
| 16 |
max_results: Number of results to return (recommended 3-5 to save tokens).
|
| 17 |
"""
|
| 18 |
+
print(f"{Fore.GREEN}[Search & Extract] {query}{Style.RESET_ALL}")
|
| 19 |
client = TavilyClient(api_key=os.environ["TAVILY_API_KEY"])
|
| 20 |
|
| 21 |
# 1. Search and get URLs
|
|
|
|
| 41 |
item["url"]: item["raw_content"] for item in extraction.get("results", [])
|
| 42 |
}
|
| 43 |
except Exception as e:
|
| 44 |
+
print(f"{Fore.RED}Extraction failed: {e}{Style.RESET_ALL}")
|
| 45 |
extracted_results = {}
|
| 46 |
|
| 47 |
# 3. Format output
|
requirements.txt
CHANGED
|
@@ -7,3 +7,5 @@ langchain-core
|
|
| 7 |
langchain-google-genai
|
| 8 |
langgraph
|
| 9 |
tavily-python
|
|
|
|
|
|
|
|
|
| 7 |
langchain-google-genai
|
| 8 |
langgraph
|
| 9 |
tavily-python
|
| 10 |
+
sympy
|
| 11 |
+
colorama
|