req_pydantic

app.py

@@ -6,6 +6,7 @@ import pandas as pd
 from typing import Dict, Any, List, TypedDict
 from langgraph.graph import Graph, StateGraph
 from langgraph.prebuilt import ToolNode
+from pydantic import BaseModel, Field
 from tools import create_calculator_tool, create_search_tool
 print("trial")
 # (Keep Constants as is)
@@ -14,6 +15,16 @@ 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(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")
+
 class BasicAgent:
     def __init__(self):
         print("Initializing BasicAgent with Qwen2.5-Coder-32B-Instruct API...")
@@ -36,8 +47,8 @@ class BasicAgent:
 
     def _create_workflow(self) -> Graph:
         """Create the agent workflow using LangGraph."""
-        # Create the workflow with simple state
-        workflow = StateGraph(dict)
+        # Create the workflow with state schema
+        workflow = StateGraph(AgentState)
         
         # Add nodes
         workflow.add_node("analyze", self._analyze_question)
@@ -53,12 +64,12 @@ class BasicAgent:
         workflow.add_edge("search", "final_answer")
         
         # Define conditional edges
-        def router(state: Dict[str, Any]) -> str:
-            if state.get('current_step') == 'calculator':
+        def router(state: AgentState) -> str:
+            if state.current_step == 'calculator':
                 return 'calculator'
-            elif state.get('current_step') == 'search':
+            elif state.current_step == 'search':
                 return 'search'
-            elif state.get('current_step') == 'final_answer':
+            elif state.current_step == 'final_answer':
                 return 'final_answer'
             return 'analyze'
         
@@ -92,9 +103,9 @@ class BasicAgent:
             print(f"Error calling LLM API: {e}")
             return f"Error getting response from LLM: {str(e)}"
 
-    def _analyze_question(self, state: Dict[str, Any]) -> Dict[str, Any]:
+    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,
@@ -108,67 +119,67 @@ 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: Dict[str, Any]) -> Dict[str, Any]:
+    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: Dict[str, Any]) -> Dict[str, Any]:
+    def _use_search(self, state: AgentState) -> AgentState:
         """Use the search tool."""
         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: Dict[str, Any]) -> Dict[str, Any]:
+    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}
@@ -176,7 +187,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 __call__(self, question: str) -> str:
@@ -185,19 +196,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}")
@@ -269,15 +280,15 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
 
         try:
             # Initialize the state for this question
-            initial_state = {
-                'question': question_text,
-                'current_step': 'analyze',
-                'tool_output': '',
-                'final_answer': '',
-                'history': [],
-                'needs_more_info': False,
-                'search_query': ''
-            }
+            initial_state = AgentState(
+                question=question_text,
+                current_step="analyze",
+                tool_output="",
+                final_answer="",
+                history=[],
+                needs_more_info=False,
+                search_query=""
+            )
             
             # Run the workflow for this question
             print(f"\nProcessing question {task_id}: {question_text[:50]}...")
@@ -288,11 +299,11 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
                 f"Step: {h['step']}\n" +
                 f"Input: {h.get('input', h.get('query', ''))}\n" +
                 f"Output: {h.get('output', h.get('results', h.get('error', '')))}"
-                for h in final_state['history']
+                for h in final_state.history
             ])
             
             # Add to results
-            submitted_answer = final_state['final_answer']
+            submitted_answer = final_state.final_answer
             answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
             results_log.append({
                 "Task ID": task_id,
@@ -301,7 +312,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
                 "Workflow History": workflow_history
             })
             
-            print(f"Completed question {task_id} with {len(final_state['history'])} workflow steps")
+            print(f"Completed question {task_id} with {len(final_state.history)} workflow steps")
             
         except Exception as e:
             print(f"Error running agent workflow on task {task_id}: {e}")

requirements.txt

@@ -1,5 +1,5 @@
 gradio
 requests
-langgraph>=0.0.10
+langgraph==0.2.62
 pydantic>=2.0.0
 duckduckgo-search>=4.1.1