Update app.py

app.py

@@ -39,12 +39,12 @@ class BasicAgent:
     FINAL ANSWER: [YOUR FINAL ANSWER]"
     """
     
-    def __init__(self, model_name: str = "gpt-4.1", api_key: str = None):
+    def __init__(self, model_name: str = "gpt-3.5-turbo", api_key: str = None):
         """
         Initialize the BasicAgent with the specified LLM.
         
         Args:
-            model_name: The name of the model to use (defaults to gpt-4.1)
+            model_name: The name of the model to use (defaults to gpt-3.5-turbo)
             api_key: OpenAI API key (defaults to environment variable)
         """
         # Use the provided API key or fall back to the environment variable
@@ -55,56 +55,20 @@ class BasicAgent:
             
         self.llm = ChatOpenAI(
             model=model_name,
-            temperature=0.2,
+            temperature=0,
             api_key=api_key
         )
         
-        # Create state graph for agent
-        self.workflow = self._create_workflow()
         print(f"BasicAgent initialized with model: {model_name}")
     
-    def _create_workflow(self) -> StateGraph:
-        """Create the agent workflow using LangGraph."""
-        
-        # Define our state
-        class AgentState(BaseModel):
-            question: str = Field(description="The original question")
-            thoughts: List[str] = Field(default_factory=list, description="Agent's thoughts")
-            final_answer: Optional[str] = Field(default=None, description="The final answer")
-            need_calculator: bool = Field(default=False, description="Whether we need to use a calculator")
-            math_expression: Optional[str] = Field(default=None, description="Math expression to evaluate")
-            
-        # Create the graph
-        workflow = StateGraph(AgentState)
-        
-        # Define nodes
-        workflow.add_node("analyze_question", self._analyze_question)
-        workflow.add_node("solve_math", self._solve_math)
-        workflow.add_node("reasoning", self._reasoning)
-        workflow.add_node("format_answer", self._format_answer)
-        
-        # Define simple edges - this approach avoids conditional routing
-        # All requests go through analyze_question -> solve_math -> reasoning -> format_answer
-        # The solve_math step will be essentially a no-op if no calculation is needed
-        workflow.add_edge("analyze_question", "solve_math")
-        workflow.add_edge("solve_math", "reasoning")
-        workflow.add_edge("reasoning", "format_answer")
-        workflow.add_edge("format_answer", END)
-        
-        # Set entry point
-        workflow.set_entry_point("analyze_question")
-        
-        # Compile the graph
-        return workflow.compile()
-    
-    def _analyze_question(self, state: Dict[str, Any]) -> Dict[str, Any]:
+    def analyze_question(self, question: str):
         """Analyze the question to determine its type and approach."""
         prompt = [
             HumanMessage(content=f"""
                 You are an expert problem analyzer. Examine the following question to determine its type 
                 and if it requires mathematical calculation.
                 
-                Question: {state.question}
+                Question: {question}
                 
                 Determine if this question requires mathematical calculation. If yes, extract the 
                 mathematical expression that needs to be calculated.
@@ -122,64 +86,36 @@ class BasicAgent:
             response = self.llm.invoke(prompt)
             try:
                 result = json.loads(response.content)
-                state.thoughts.append(f"Analysis: {result.get('thought', 'No analysis provided')}")
-                state.need_calculator = result.get('need_calculator', False)
-                state.math_expression = result.get('math_expression')
+                thought = f"Analysis: {result.get('thought', 'No analysis provided')}"
+                need_calculator = result.get('need_calculator', False)
+                math_expression = result.get('math_expression')
+                return thought, need_calculator, math_expression
             except json.JSONDecodeError as e:
                 # Fallback in case of JSON parsing error
-                state.thoughts.append(f"Couldn't parse the analysis response as JSON: {e}")
-                state.thoughts.append(f"Raw response: {response.content[:100]}...")
-                state.need_calculator = False
+                return f"Couldn't parse the analysis response as JSON: {e}", False, None
             except Exception as e:
                 # Fallback for any other parsing error
-                state.thoughts.append(f"Error processing response: {str(e)}")
-                state.need_calculator = False
+                return f"Error processing response: {str(e)}", False, None
         except Exception as e:
             # Fallback for LLM invocation error
-            state.thoughts.append(f"Error invoking language model: {str(e)}")
-            state.need_calculator = False
-            
-        return state
+            return f"Error invoking language model: {str(e)}", False, None
     
-    def _solve_math(self, state: Dict[str, Any]) -> Dict[str, Any]:
+    def solve_math(self, math_expression):
         """Evaluate the mathematical expression."""
-        # Skip calculation if not needed
-        if not state.need_calculator:
-            state.thoughts.append("No mathematical calculation required.")
-            return state
-            
-        # Skip if no expression provided
-        if not state.math_expression:
-            state.thoughts.append("Math calculation was indicated but no expression was provided.")
-            return state
-            
         try:
             # Simple math expression evaluator
-            # This is a basic implementation - in production, you'd want more robust validation
-            result = eval(state.math_expression)
-            state.thoughts.append(f"Calculated {state.math_expression} = {result}")
-            state.calc_result = result
+            result = eval(math_expression)
+            return f"Calculated {math_expression} = {result}", result
         except Exception as e:
-            state.thoughts.append(f"Error calculating expression: {e}")
-            
-        return state
+            return f"Error calculating expression: {e}", None
     
-    def _reasoning(self, state: Dict[str, Any]) -> Dict[str, Any]:
+    def reasoning(self, question, calc_context=""):
         """Perform step-by-step reasoning to solve the problem."""
-        # Build calculation context if available
-        calc_context = ""
-        if state.need_calculator and hasattr(state, "calc_result") and state.calc_result is not None:
-            calc_context = f"""
-            Mathematical calculation was needed for this problem.
-            Expression: {state.math_expression}
-            Calculation result: {state.calc_result}
-            """
-            
         prompt = [
             HumanMessage(content=f"""
                 You are a general AI assistant. Solve the following question using careful step-by-step reasoning.
                 
-                Question: {state.question}
+                Question: {question}
                 
                 {calc_context}
                 
@@ -190,22 +126,20 @@ class BasicAgent:
         
         try:
             response = self.llm.invoke(prompt)
-            state.thoughts.append(response.content)
+            return response.content
         except Exception as e:
-            state.thoughts.append(f"Error during reasoning step: {str(e)}")
-            
-        return state
+            return f"Error during reasoning step: {str(e)}"
     
-    def _format_answer(self, state: Dict[str, Any]) -> Dict[str, Any]:
+    def format_answer(self, question, thoughts):
         """Format the final answer according to Gaia Benchmark requirements."""
-        all_thoughts = "\n".join(state.thoughts)
+        all_thoughts = "\n".join(thoughts)
         
         prompt = [
             HumanMessage(content=f"""
                 You are a general AI assistant. Based on the following reasoning, provide a concise 
                 final answer to the question. The final answer format is extremely important.
                 
-                Question: {state.question}
+                Question: {question}
                 
                 Reasoning:
                 {all_thoughts}
@@ -229,23 +163,18 @@ class BasicAgent:
             # Extract the final answer using regex
             final_answer_match = re.search(r'FINAL ANSWER:\s*(.*?)(?:\n|$)', response.content)
             if final_answer_match:
-                state.final_answer = final_answer_match.group(1).strip()
-            else:
-                # If regex fails, try to extract anything after "FINAL ANSWER:"
-                if "FINAL ANSWER:" in response.content:
-                    parts = response.content.split("FINAL ANSWER:")
-                    if len(parts) > 1:
-                        state.final_answer = parts[1].strip()
-                    else:
-                        state.final_answer = "Error: Malformed answer format."
+                return final_answer_match.group(1).strip()
+            elif "FINAL ANSWER:" in response.content:
+                parts = response.content.split("FINAL ANSWER:")
+                if len(parts) > 1:
+                    return parts[1].strip()
                 else:
-                    # Fallback in case the model didn't follow the format at all
-                    state.final_answer = response.content.strip()
+                    return "Error: Malformed answer format."
+            else:
+                return response.content.strip()
         except Exception as e:
-            state.final_answer = f"Error formatting answer: {str(e)}"
-            
-        return state
-    
+            return f"Error formatting answer: {str(e)}"
+        
     def __call__(self, question: str) -> str:
         """
         Process the question and return a formatted response according to Gaia requirements.
@@ -258,19 +187,40 @@ class BasicAgent:
         """
         print(f"Agent received question (first 50 chars): {question[:50]}...")
         
-        # Initialize state
-        initial_state = {"question": question}
+        # Initialize an empty list to collect thoughts
+        thoughts = []
         
-        # Execute the workflow
         try:
-            result = self.workflow.invoke(initial_state)
+            # Step 1: Analyze the question
+            analysis_thought, need_calculator, math_expression = self.analyze_question(question)
+            thoughts.append(analysis_thought)
+            
+            # Step 2: Calculate if needed
+            calc_result = None
+            calc_context = ""
+            if need_calculator and math_expression:
+                calc_thought, calc_result = self.solve_math(math_expression)
+                thoughts.append(calc_thought)
+                if calc_result is not None:
+                    calc_context = f"""
+                    Mathematical calculation was needed for this problem.
+                    Expression: {math_expression}
+                    Calculation result: {calc_result}
+                    """
+            
+            # Step 3: Perform reasoning
+            reasoning_output = self.reasoning(question, calc_context)
+            thoughts.append(reasoning_output)
+            
+            # Step 4: Format final answer
+            final_answer = self.format_answer(question, thoughts)
             
             # Format the response according to Gaia requirements
             formatted_response = "You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template:\nFINAL ANSWER: [YOUR FINAL ANSWER]\n\n"
-            formatted_response += "\n\n".join(result.thoughts)
-            formatted_response += f"\n\nFINAL ANSWER: {result.final_answer}"
+            formatted_response += "\n\n".join(thoughts)
+            formatted_response += f"\n\nFINAL ANSWER: {final_answer}"
             
-            print(f"Agent returning answer with final part: {result.final_answer}")
+            print(f"Agent returning answer with final part: {final_answer}")
             return formatted_response
         except Exception as e:
             print(f"Error running agent workflow: {e}")