please_work

agent.py

@@ -1,116 +1,54 @@
-from __future__ import annotations
-import os
 from langchain_openai import ChatOpenAI
-from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
-from typing import Any, Dict, List, Optional
-import json
-
-# ─────────────────────────── External tools ──────────────────────────────
+from langchain.schema import SystemMessage, HumanMessage
+from langgraph.prebuilt import create_react_agent
 from tools import (
-    wikipedia_search_tool,
-    arxiv_search_tool,
-    audio_transcriber_tool,
-    excel_tool,
-    analyze_code_tool,
-    image_tool,
-    add_tool,
-    subtract_tool,
-    multiply_tool,
-    divide_tool
+    wikipedia_search_tool, arxiv_search_tool,
+    audio_transcriber_tool, excel_tool, analyze_code_tool, image_tool,
+    add_tool, subtract_tool, multiply_tool, divide_tool
 )
 
-# ─────────────────────────── Configuration ───────────────────────────────
-MAX_TOOL_CALLS = 5
+# ──────────────────────────────── Config ────────────────────────────────
+SYSTEM_PROMPT = """
+You are a smart AI assistant using a tool-augmented reasoning strategy. Follow this loop:
+Thought: ...
+Action: ...
+Action Input: ...
 
-class SimpleToolExecutor:
-    """Simple tool executor that maps tool names to functions."""
-    
-    def __init__(self, tools: List):
-        self.tools_map = {tool.name: tool for tool in tools}
-    
-    def run(self, tool_name: str, tool_input: str) -> str:
-        """Execute a tool with the given input."""
-        if tool_name not in self.tools_map:
-            return f"Tool '{tool_name}' not found. Available tools: {list(self.tools_map.keys())}"
-        
-        tool = self.tools_map[tool_name]
-        try:
-            # Handle different input formats
-            if tool_name in ['add_tool', 'subtract_tool', 'multiply_tool', 'divide_tool']:
-                # Math tools expect two numbers
-                # Try to parse the input as two numbers
-                import re
-                numbers = re.findall(r'-?\d+(?:\.\d+)?', tool_input)
-                if len(numbers) >= 2:
-                    a, b = float(numbers[0]), float(numbers[1])
-                    return tool.run(a=a, b=b)
-                else:
-                    return f"Math tool requires two numbers. Got: {tool_input}"
-            else:
-                # Other tools expect a string input
-                # Clean the input - remove quotes if present
-                clean_input = tool_input.strip().strip('"').strip("'")
-                if "search" in tool_name:
-                    return tool.run(clean_input)
-                elif tool_name in ['audio_transcriber_tool', 'excel_tool', 'analyze_code_tool', 'image_tool']:
-                    return tool.run(task_id=clean_input)
-                else:
-                    return tool.run(clean_input)
-        except Exception as e:
-            return f"Error executing {tool_name}: {str(e)}"
+You will get an Observation, and then continue reasoning.
 
-class CustomReActAgent:
-    def __init__(self, tools: List, model_name="gpt-4o-mini"):
-        self.llm = ChatOpenAI(model_name=model_name, temperature=0.3)
-        self.tool_executor = SimpleToolExecutor(tools)
+Only end the loop with:
+FINAL ANSWER: [your short final answer here].
 
-    def run(self, question: str, task_id: Optional[str] = None, max_turns: int = 15, system_prompt: str = "") -> str:
-        messages = [SystemMessage(content=system_prompt)]
-        if task_id:
-            messages[0].content += f"\n\nIMPORTANT: Your current task_id is: {task_id}."
-        messages.append(HumanMessage(content=question))
+DO NOT answer directly without using this loop at least once, unless the answer is trivially obvious.
 
-        for _ in range(max_turns):
-            ai_response = self.llm.invoke(messages)
-            messages.append(ai_response)
-            print("AI said:", ai_response.content)
+When using task-based tools (audio_transcriber_tool, excel_tool, analyze_code_tool, image_tool), ONLY use the 'task_id' value.
+"""
 
-            # Final answer?
-            if "FINAL ANSWER:" in ai_response.content.upper():
-                return ai_response.content.split("FINAL ANSWER:")[-1].strip()
+TOOLS = [
+    wikipedia_search_tool, arxiv_search_tool,
+    audio_transcriber_tool, excel_tool, analyze_code_tool, image_tool,
+    add_tool, subtract_tool, multiply_tool, divide_tool
+]
 
-            # Tool call expected
-            if "Action:" in ai_response.content and "Action Input:" in ai_response.content:
-                try:
-                    tool_name = ai_response.content.split("Action:")[1].split("Action Input:")[0].strip()
-                    tool_input = ai_response.content.split("Action Input:")[1].strip()
-                    tool_result = self.tool_executor.run(tool_name, tool_input)
-                    messages.append(HumanMessage(content=f"Observation: {tool_result}"))
-                except Exception as e:
-                    messages.append(HumanMessage(content=f"Observation: Tool execution failed: {str(e)}"))
-            else:
-                # Instead of adding a confusing "observation"
-                messages.append(HumanMessage(content="Observation: No action detected. If you are done, write FINAL ANSWER."))
+# ───────────────────────────── Agent Wrapper ─────────────────────────────
+class SimpleLangGraphAgent:
+    def __init__(self, model_name="gpt-4o-mini"):
+        self.agent = create_react_agent(
+            model=ChatOpenAI(model_name=model_name, temperature=0),
+            tools=TOOLS,
+        )
+
+    def run(self, question: str, task_id: str = None, max_steps: int = 10) -> str:
+        messages = [SystemMessage(content=SYSTEM_PROMPT)]
+        if task_id:
+            messages[0].content += f"\n\nYour task_id is: {task_id}. Use it only with the tools that require it."
+        messages.append(HumanMessage(content=question))
 
-        return "No FINAL ANSWER found within allowed steps."
+        state = {"messages": messages}
+        final_state = self.agent.invoke(state, config={"recursion_limit": max_steps})
 
-def build_graph():
-    """Build and return a CustomReActAgent."""
-    
-    llm_tools = [
-        wikipedia_search_tool,
-        arxiv_search_tool,
-        audio_transcriber_tool,
-        excel_tool,
-        analyze_code_tool,
-        image_tool,
-        add_tool,
-        subtract_tool,
-        multiply_tool,
-        divide_tool
-    ]
-    
-    # Create the custom react agent
-    agent = CustomReActAgent(tools=llm_tools)
+        for msg in final_state["messages"][::-1]:
+            if "FINAL ANSWER:" in msg.content.upper():
+                return msg.content.split("FINAL ANSWER:")[-1].strip()
 
-    return agent 
+        return "No FINAL ANSWER found."

app.py

@@ -7,7 +7,7 @@ from langchain.schema import HumanMessage, SystemMessage
 from typing import Optional
 import re
 
-from agent import build_graph
+from agent import SimpleLangGraphAgent
 from state import AgentState
 
 # --- Constants ---
@@ -51,24 +51,16 @@ IMPORTANT:
 class BasicAgent:
     def __init__(self):
         print("BasicAgent initialized.")
-        self.graph = build_graph()
+        
 
     def __call__(self, question: str, task_id: Optional[str] = None) -> str:
         """Run the agent and return whatever FINAL_ANSWER the agent produces."""
         print(f"Agent received question: {question}")
-        
-        try:
-            # Run the custom react agent
-            result = self.graph.run(question=question, task_id=task_id, max_turns=15, system_prompt=SYSTEM_PROMPT)
-            print("Final result: ", result)
-            print("\n\n\n\n")
-            return result
-            
-        except Exception as e:
-            print(f"Agent execution error: {e}")
-            return f"Unable to process the question due to an error. Please try again."
-
+        print("\n\n\n")
 
+        agent = SimpleLangGraphAgent()
+        return agent.run(question, task_id)
+        
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,