Update app.py

app.py

@@ -15,15 +15,15 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 MODEL_API_URL = "https://api-inference.huggingface.co/models/Qwen/Qwen2.5-Coder-32B-Instruct"
 HF_TOKEN = os.getenv("HF_TOKEN")  # Make sure to set this environment variable
 
-class AgentState(TypedDict):
-    """Type definition for the agent's state."""
-    question: str
-    current_step: str
-    tool_output: str
-    final_answer: str
-    history: List[Dict[str, str]]
-    needs_more_info: bool
-    search_query: str
+class AgentState(BaseModel):
+    """Schema for the agent's state."""
+    question: str = Field(..., description="The original question")
+    current_step: str = Field(default="analyze", description="Current step in the workflow")
+    tool_output: str = Field(default="", description="Output from the last tool used")
+    final_answer: str = Field(default="", description="The final answer to be returned")
+    history: List[Dict[str, str]] = Field(default_factory=list, description="History of operations performed")
+    needs_more_info: bool = Field(default=False, description="Whether more information is needed")
+    search_query: str = Field(default="", description="Current search query if any")
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
@@ -63,7 +63,7 @@ class BasicAgent:
 
     def _analyze_question(self, state: AgentState) -> AgentState:
         """Analyze the question and determine the next step."""
-        prompt = f"""Analyze this question and determine what needs to be done: {state['question']}
+        prompt = f"""Analyze this question and determine what needs to be done: {state.question}
         Return your analysis in this format:
         {{
             "needs_calculation": true/false,
@@ -77,35 +77,35 @@ class BasicAgent:
         """
         
         analysis = eval(self._call_llm_api(prompt))
-        state['needs_more_info'] = analysis.get('needs_search', False)
-        state['search_query'] = analysis.get('search_query', '')
+        state.needs_more_info = analysis.get('needs_search', False)
+        state.search_query = analysis.get('search_query', '')
         
         if analysis.get('needs_calculation', False):
-            state['current_step'] = 'calculator'
-            state['tool_output'] = str(analysis['calculation'])
+            state.current_step = 'calculator'
+            state.tool_output = str(analysis['calculation'])
         elif analysis.get('needs_search', False):
-            state['current_step'] = 'search'
+            state.current_step = 'search'
         else:
-            state['current_step'] = 'final_answer'
+            state.current_step = 'final_answer'
         
         return state
 
     def _use_calculator(self, state: AgentState) -> AgentState:
         """Use the calculator tool."""
         try:
-            result = self.calculator.invoke({"input": eval(state['tool_output'])})
-            state['history'].append({
+            result = self.calculator.invoke({"input": eval(state.tool_output)})
+            state.history.append({
                 'step': 'calculator',
-                'input': state['tool_output'],
+                'input': state.tool_output,
                 'output': str(result['output'].result)
             })
-            state['current_step'] = 'final_answer'
+            state.current_step = 'final_answer'
         except Exception as e:
-            state['history'].append({
+            state.history.append({
                 'step': 'calculator_error',
                 'error': str(e)
             })
-            state['current_step'] = 'final_answer'
+            state.current_step = 'final_answer'
         return state
 
     def _use_search(self, state: AgentState) -> AgentState:
@@ -113,31 +113,31 @@ class BasicAgent:
         try:
             result = self.search_tool.invoke({
                 "input": {
-                    "query": state['search_query'],
+                    "query": state.search_query,
                     "max_results": 3
                 }
             })
-            state['history'].append({
+            state.history.append({
                 'step': 'search',
-                'query': state['search_query'],
+                'query': state.search_query,
                 'results': [str(r) for r in result['output'].results]
             })
-            state['needs_more_info'] = False
-            state['current_step'] = 'final_answer'
+            state.needs_more_info = False
+            state.current_step = 'final_answer'
         except Exception as e:
-            state['history'].append({
+            state.history.append({
                 'step': 'search_error',
                 'error': str(e)
             })
-            state['current_step'] = 'final_answer'
+            state.current_step = 'final_answer'
         return state
 
     def _generate_final_answer(self, state: AgentState) -> AgentState:
         """Generate the final answer based on all gathered information."""
         history_str = "\n".join([f"{h['step']}: {h.get('output', h.get('results', h.get('error', '')))}" 
-                               for h in state['history']])
+                               for h in state.history])
         
-        prompt = f"""Based on the following information and history, provide a final answer to the question: {state['question']}
+        prompt = f"""Based on the following information and history, provide a final answer to the question: {state.question}
         
         History of steps taken:
         {history_str}
@@ -145,7 +145,7 @@ class BasicAgent:
         Provide a clear, concise answer that addresses the original question.
         """
         
-        state['final_answer'] = self._call_llm_api(prompt)
+        state.final_answer = self._call_llm_api(prompt)
         return state
 
     def _create_workflow(self) -> Graph:
@@ -167,11 +167,11 @@ class BasicAgent:
         
         # Define conditional edges
         def router(state: AgentState) -> str:
-            if state['current_step'] == 'calculator':
+            if state.current_step == 'calculator':
                 return 'calculator'
-            elif state['current_step'] == 'search':
+            elif state.current_step == 'search':
                 return 'search'
-            elif state['current_step'] == 'final_answer':
+            elif state.current_step == 'final_answer':
                 return 'final_answer'
             return 'analyze'
         
@@ -197,19 +197,19 @@ class BasicAgent:
         
         try:
             # Initialize the state
-            initial_state = {
-                "question": question,
-                "current_step": "analyze",
-                "tool_output": "",
-                "final_answer": "",
-                "history": [],
-                "needs_more_info": False,
-                "search_query": ""
-            }
+            initial_state = AgentState(
+                question=question,
+                current_step="analyze",
+                tool_output="",
+                final_answer="",
+                history=[],
+                needs_more_info=False,
+                search_query=""
+            )
             
             # Run the workflow
             final_state = self.workflow.invoke(initial_state)
-            return final_state['final_answer']
+            return final_state.final_answer
             
         except Exception as e:
             print(f"Error in agent processing: {e}")