Update agent.py

agent.py

@@ -1,116 +1,46 @@
-from typing import TypedDict, Annotated, Sequence
-import operator
-from langgraph.graph import StateGraph, END
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from ai_tools import Calculator, DocRetriever, WebSearcher
-
-# Configuration
-MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
-llm_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-# Define tools
-tools = [Calculator(), WebSearcher()]
-doc_retriever = DocRetriever()
-tool_map = {tool.name: tool for tool in tools}
-tool_map["DocRetriever"] = doc_retriever
-
-# Agent State
-class AgentState(TypedDict):
-    input: str
-    context: Annotated[Sequence[str], operator.add]
-    last_tool: str
-
-# Tool calling prompt template
-TOOL_PROMPT = """<|system|>
-You're an expert problem solver. Use these tools:
-{tool_descriptions}
-
-Respond ONLY in this format:
-Thought: <strategy>
-Action: <tool_name>
-Action Input: <input>
-</s>
-<|user|>
-{input}
-Context: {context}
-</s>
-<|assistant|>
-"""
-
-# Initialize graph
-graph = StateGraph(AgentState)
-
-# Node: Generate tool calls
-def agent_node(state):
-    tool_list = "\n".join([f"- {t.name}: {t.description}" for t in tools])
-    prompt = TOOL_PROMPT.format(
-        tool_descriptions=tool_list,
-        input=state["input"],
-        context=state["context"]
-    )
-    
-    response = llm_pipeline(
-        prompt,
-        max_new_tokens=150,
-        do_sample=True,
-        temperature=0.2,
-        pad_token_id=tokenizer.eos_token_id
-    )[0]['generated_text']
-    
-    # Extract tool call
-    action_match = re.search(r"Action: (\w+)", response)
-    action_input_match = re.search(r"Action Input: (.+?)\n", response)
-    
-    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}
-
-# Node: Execute tools
-def tool_node(state):
-    tool = 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}"}
-
-# Define graph structure
-graph.add_node("agent", agent_node)
-graph.add_node("tool", tool_node)
-graph.set_entry_point("agent")
-
-# Conditional edges
-def route_action(state):
-    if state["last_tool"] == "FINISH":
-        return END
-    return "tool"
-
-graph.add_edge("agent", "tool")
-graph.add_conditional_edges("tool", route_action, {"agent": "agent", END: END})
-graph.add_edge("tool", "agent")  # Loop back after tool use
-
-# Compile the agent
-agent = graph.compile()
-
-# Interface function
-def run_agent(query: str, document: str = ""):
-    doc_retriever.document = document  # Load document
-    state = {"input": query, "context": [], "last_tool": ""}
+from .ai_tools import Calculator, DocRetriever, WebSearcher
+from .graph import GaiaGraph
+
+class GaiaAgent:
+    def __init__(self, model_name="HuggingFaceH4/zephyr-7b-beta"):
+        self.model_name = model_name
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(model_name)
+        self.llm_pipeline = pipeline(
+            "text-generation", 
+            model=self.model, 
+            tokenizer=self.tokenizer
+        )
+        
+        # Initialize tools
+        self.calculator = Calculator()
+        self.doc_retriever = DocRetriever()
+        self.web_searcher = WebSearcher()
+        
+        # Create tool list
+        self.tools = [
+            self.calculator,
+            self.web_searcher,
+            self.doc_retriever
+        ]
+        
+        # Build LangGraph workflow
+        self.graph = GaiaGraph(
+            model=self.llm_pipeline,
+            tokenizer=self.tokenizer,
+            tools=self.tools
+        )
+        
+        print(f"GaiaAgent initialized with model: {model_name}")
     
-    for step in agent.stream(state):
-        for node, value in step.items():
-            if node == "agent":
-                print(f"THOUGHT: {value['thought']}")
-            if node == "tool":
-                print(f"TOOL RESULT: {value['context']}")
+    def load_document(self, document_text: str):
+        """Load document content for retrieval"""
+        self.doc_retriever.load_document(document_text)
+        print(f"Document loaded ({len(document_text)} characters)")
     
-    return state["context"][-1] if state["context"] else "No output"
+    def __call__(self, question: str) -> str:
+        print(f"Agent received question: {question[:50]}{'...' if len(question) > 50 else ''}")
+        result = self.graph.run(question)
+        print(f"Agent returning answer: {result[:50]}{'...' if len(result) > 50 else ''}")
+        return result