|
|
from typing import TypedDict, Annotated, Sequence |
|
|
import operator |
|
|
import re |
|
|
from langgraph.graph import StateGraph, END |
|
|
from .ai_tools import Calculator, DocRetriever, WebSearcher |
|
|
|
|
|
class AgentState(TypedDict): |
|
|
input: str |
|
|
context: Annotated[Sequence[str], operator.add] |
|
|
last_tool: str |
|
|
output: str |
|
|
|
|
|
class GaiaGraph: |
|
|
def __init__(self, model, tokenizer, tools): |
|
|
self.model = model |
|
|
self.tokenizer = tokenizer |
|
|
self.tools = tools |
|
|
self.tool_map = {tool.name: tool for tool in tools} |
|
|
self.graph = self._build_graph() |
|
|
|
|
|
def _build_graph(self): |
|
|
graph = StateGraph(AgentState) |
|
|
|
|
|
graph.add_node("agent", self._agent_node) |
|
|
graph.add_node("tool", self._tool_node) |
|
|
graph.set_entry_point("agent") |
|
|
|
|
|
graph.add_edge("agent", "tool") |
|
|
graph.add_conditional_edges( |
|
|
"tool", |
|
|
self._route_action, |
|
|
{"continue": "agent", "end": END} |
|
|
) |
|
|
|
|
|
return graph.compile() |
|
|
|
|
|
def _agent_node(self, state: AgentState) -> dict: |
|
|
tool_list = "\n".join([f"- {t.name}: {t.description}" for t in self.tools]) |
|
|
prompt = f"""<|system|> |
|
|
You're an expert problem solver. Use these tools when needed: |
|
|
{tool_list} |
|
|
|
|
|
Respond ONLY in this format: |
|
|
Thought: <your reasoning> |
|
|
Action: <tool_name or 'FINISH'> |
|
|
Action Input: <input for tool> |
|
|
</s> |
|
|
<|user|> |
|
|
{state['input']} |
|
|
Context: {state['context']} |
|
|
</s> |
|
|
<|assistant|> |
|
|
""" |
|
|
|
|
|
response = self.model( |
|
|
prompt, |
|
|
max_new_tokens=200, |
|
|
do_sample=True, |
|
|
temperature=0.2, |
|
|
pad_token_id=self.tokenizer.eos_token_id |
|
|
)[0]['generated_text'] |
|
|
|
|
|
|
|
|
action_match = re.search(r"Action: (\w+)", response) |
|
|
action_input_match = re.search(r"Action Input: (.+?)\n", response, re.DOTALL) |
|
|
|
|
|
if action_match and action_input_match: |
|
|
tool_name = action_match.group(1) |
|
|
tool_input = action_input_match.group(1).strip() |
|
|
return { |
|
|
"last_tool": tool_name, |
|
|
"tool_input": tool_input, |
|
|
"thought": response |
|
|
} |
|
|
else: |
|
|
return {"last_tool": "FINISH", "output": response} |
|
|
|
|
|
def _tool_node(self, state: AgentState) -> dict: |
|
|
if state["last_tool"] == "FINISH": |
|
|
return {"output": state.get("output", "No output generated")} |
|
|
|
|
|
tool = self.tool_map.get(state["last_tool"]) |
|
|
if not tool: |
|
|
return {"context": f"Error: Unknown tool {state['last_tool']}"} |
|
|
|
|
|
result = tool.run(state["tool_input"]) |
|
|
return {"context": f"Tool {tool.name} returned: {result}"} |
|
|
|
|
|
def _route_action(self, state: AgentState) -> str: |
|
|
return "end" if state["last_tool"] == "FINISH" else "continue" |
|
|
|
|
|
def run(self, input: str) -> str: |
|
|
state = {"input": input, "context": [], "last_tool": "", "output": ""} |
|
|
for step in self.graph.stream(state): |
|
|
for node, value in step.items(): |
|
|
if node == "__end__": |
|
|
return value["output"] |
|
|
return "Execution completed without output" |
