Update app.py

app.py

@@ -46,10 +46,58 @@ class BasicAgent:
             "Content-Type": "application/json"
         }
         
-        # Create the agent workflow
+        # Create the agent workflow with proper state schema
         self.workflow = self._create_workflow()
         print("BasicAgent initialization complete.")
 
+    def _create_workflow(self) -> Graph:
+        """Create the agent workflow using LangGraph."""
+        # Create the workflow with proper state schema
+        workflow = StateGraph(
+            input_type=AgentState,
+            output_type=AgentState,
+            state_schema=AgentState
+        )
+        
+        # Add nodes
+        workflow.add_node("analyze", self._analyze_question)
+        workflow.add_node("calculator", self._use_calculator)
+        workflow.add_node("search", self._use_search)
+        workflow.add_node("final_answer", self._generate_final_answer)
+        
+        # Define edges
+        workflow.add_edge("analyze", "calculator")
+        workflow.add_edge("analyze", "search")
+        workflow.add_edge("analyze", "final_answer")
+        workflow.add_edge("calculator", "final_answer")
+        workflow.add_edge("search", "final_answer")
+        
+        # Define conditional edges
+        def router(state: AgentState) -> str:
+            if state.current_step == 'calculator':
+                return 'calculator'
+            elif state.current_step == 'search':
+                return 'search'
+            elif state.current_step == 'final_answer':
+                return 'final_answer'
+            return 'analyze'
+        
+        workflow.add_conditional_edges(
+            "analyze",
+            router,
+            {
+                "calculator": "calculator",
+                "search": "search",
+                "final_answer": "final_answer"
+            }
+        )
+        
+        # Set entry and exit points
+        workflow.set_entry_point("analyze")
+        workflow.set_finish_point("final_answer")
+        
+        return workflow.compile()
+
     def _call_llm_api(self, prompt: str) -> str:
         """Call the Qwen model through the Hugging Face API."""
         try:
@@ -151,53 +199,6 @@ class BasicAgent:
         state.final_answer = self._call_llm_api(prompt)
         return state
 
-    def _create_workflow(self) -> Graph:
-        """Create the agent workflow using LangGraph."""
-        # Create the workflow with explicit input and output types
-        workflow = StateGraph(
-            input_type=AgentState,
-            output_type=AgentState
-        )
-        
-        # Add nodes with explicit type casting
-        workflow.add_node("analyze", cast(Any, self._analyze_question))
-        workflow.add_node("calculator", cast(Any, self._use_calculator))
-        workflow.add_node("search", cast(Any, self._use_search))
-        workflow.add_node("final_answer", cast(Any, self._generate_final_answer))
-        
-        # Define edges
-        workflow.add_edge("analyze", "calculator")
-        workflow.add_edge("analyze", "search")
-        workflow.add_edge("analyze", "final_answer")
-        workflow.add_edge("calculator", "final_answer")
-        workflow.add_edge("search", "final_answer")
-        
-        # Define conditional edges
-        def router(state: AgentState) -> str:
-            if state.current_step == 'calculator':
-                return 'calculator'
-            elif state.current_step == 'search':
-                return 'search'
-            elif state.current_step == 'final_answer':
-                return 'final_answer'
-            return 'analyze'
-        
-        workflow.add_conditional_edges(
-            "analyze",
-            router,
-            {
-                "calculator": "calculator",
-                "search": "search",
-                "final_answer": "final_answer"
-            }
-        )
-        
-        # Set entry and exit points
-        workflow.set_entry_point("analyze")
-        workflow.set_finish_point("final_answer")
-        
-        return workflow.compile()
-
     def __call__(self, question: str) -> str:
         """Process a question through the agent workflow."""
         print(f"Agent received question: {question[:50]}...")