| 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" |
