Update agent.py

agent.py

@@ -1,100 +1,116 @@
-from typing import Dict, Any, Optional
-from langgraph.graph import StateGraph, START, END
-from langgraph.prebuilt import ToolNode
-from langgraph.checkpoint.sqlite import SqliteSaver
-from ai_tools import AITools
+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
 
-class CustomAgent:
-    def __init__(self):
-        self.tools = AITools()
-        self.workflow = self._build_workflow()
+# 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"]
+    )
     
-    def _build_workflow(self):
-        # 定义工具节点
-        tool_node = ToolNode([self._route_to_tool])
-        
-        # 构建工作流
-        workflow = StateGraph(State)
-        workflow.add_node("tools", tool_node)
-        workflow.add_node("process", self._process_result)
-        
-        # 设置入口点
-        workflow.set_entry_point("tools")
-        
-        # 添加条件边
-        workflow.add_conditional_edges(
-            "tools",
-            self._decide_next_step,
-            {
-                "continue": "process",
-                "end": END
-            }
-        )
-        
-        workflow.add_edge("process", END)
-        
-        # 添加持久化
-        workflow.checkpointer = SqliteSaver.from_conn_string(":memory:")
-        
-        return workflow.compile()
+    response = llm_pipeline(
+        prompt,
+        max_new_tokens=150,
+        do_sample=True,
+        temperature=0.2,
+        pad_token_id=tokenizer.eos_token_id
+    )[0]['generated_text']
     
-    def _route_to_tool(self, state: Dict[str, Any]):
-        """路由问题到适当的工具"""
-        question = state.get("question", "")
-        file_name = state.get("file_name", "")
-        
-        # 处理反转文本问题
-        if "rewsna" in question or "dnatsrednu" in question:
-            return {"result": self.tools.reverse_text(question.split('"')[1])}
-        
-        # 处理蔬菜分类问题
-        if "grocery list" in question.lower() or "vegetables" in question.lower():
-            items = re.findall(r"[a-zA-Z]+(?=\W|\Z)", question)
-            return {"result": ", ".join(self.tools.categorize_vegetables(items))}
-        
-        # 处理棋局问题
-        if "chess position" in question.lower() and file_name.endswith(".png"):
-            return {"result": self.tools.analyze_chess_position(file_name)}
-        
-        # 处理音频文件问题
-        if file_name.endswith(".mp3"):
-            transcript = self.tools.extract_audio_transcript(file_name)
-            if "page numbers" in question.lower():
-                return {"result": transcript}
-            else:
-                return {"result": ", ".join(sorted(transcript.split(", ")))}
-        
-        # 处理表格运算问题
-        if "* on the set S" in question:
-            table_data = {"operation": "*", "set": ["a", "b", "c", "d", "e"]}
-            return {"result": self.tools.process_table_operation(table_data)}
-        
-        # 处理Python代码问题
-        if file_name.endswith(".py"):
-            return {"result": self.tools.analyze_python_code(file_name)}
-        
-        # 处理Excel文件问题
-        if file_name.endswith(".xlsx"):
-            return {"result": self.tools.process_excel_file(file_name)}
-        
-        return {"result": "I don't have a tool to answer this question."}
+    # Extract tool call
+    action_match = re.search(r"Action: (\w+)", response)
+    action_input_match = re.search(r"Action Input: (.+?)\n", response)
     
-    def _process_result(self, state: Dict[str, Any]):
-        """处理工具返回的结果"""
-        result = state.get("result", "")
-        return {"answer": result}
+    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": ""}
     
-    def _decide_next_step(self, state: Dict[str, Any]):
-        """决定下一步"""
-        result = state.get("result", "")
-        if result == "I don't have a tool to answer this question.":
-            return "end"
-        return "continue"
+    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 __call__(self, question: str, file_name: str = "") -> str:
-        """执行Agent"""
-        state = {"question": question, "file_name": file_name}
-        for step in self.workflow.stream(state):
-            if "__end__" in step:
-                return step["__end__"]["answer"]
-        return "No answer generated."
+    return state["context"][-1] if state["context"] else "No output"