Update app.py

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
-from typing import Dict, Any, List, TypedDict, Annotated
+from typing import Dict, Any, List, TypedDict, Annotated, TypeVar, cast
 from langgraph.graph import Graph, StateGraph
 from langgraph.prebuilt import ToolNode
 from pydantic import BaseModel, Field
@@ -15,6 +15,9 @@ 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
 
+# Define the state type
+StateType = TypeVar("StateType", bound=BaseModel)
+
 class AgentState(BaseModel):
     """Schema for the agent's state."""
     question: str = Field(..., description="The original question")
@@ -150,13 +153,17 @@ class BasicAgent:
 
     def _create_workflow(self) -> Graph:
         """Create the agent workflow using LangGraph."""
-        workflow = StateGraph(AgentState)
+        # Create the workflow with explicit input and output types
+        workflow = StateGraph(
+            input_type=AgentState,
+            output_type=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)
+        # 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")