some

app.py

@@ -50,7 +50,7 @@ class BasicAgent:
         """Create the agent workflow using LangGraph."""
         # Create the workflow with state schema
         print("Creating Stategraph : error happens here?")
-        workflow = StateGraph(state_schema=Dict[str, Any])
+        workflow = StateGraph(state_schema=AgentState)
         print("Stategraph created")
         # Add nodes
         workflow.add_node("analyze", self._analyze_question)
@@ -66,7 +66,7 @@ class BasicAgent:
         workflow.add_edge("search", "final_answer")
         
         # Define conditional edges
-        def router(state: Dict[str, Any]) -> str:
+        def router(state: AgentState) -> str:
             if state["current_step"] == 'calculator':
                 return 'calculator'
             elif state["current_step"] == 'search':
@@ -105,7 +105,7 @@ 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']}
         Return your analysis in this format:
@@ -134,7 +134,7 @@ class BasicAgent:
         
         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:
             # Create calculator state with input from tool_output
@@ -155,7 +155,7 @@ class BasicAgent:
             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:
             # Create search state with input from search_query
@@ -180,7 +180,7 @@ class BasicAgent:
             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"]])
@@ -202,7 +202,7 @@ class BasicAgent:
         
         try:
             # Initialize the state
-            initial_state = {
+            initial_state: AgentState = {
                 "question": question,
                 "current_step": "analyze",
                 "tool_output": "",

tools.py

@@ -11,6 +11,10 @@ class CalculatorOutput(TypedDict):
     result: float
     operation: str
 
+class CalculatorState(TypedDict):
+    input: CalculatorInput
+    output: Optional[CalculatorOutput]
+
 class SearchResult(TypedDict):
     title: str
     link: str
@@ -20,11 +24,15 @@ class SearchOutput(TypedDict):
     results: List[SearchResult]
     query: str
 
+class SearchState(TypedDict):
+    input: Dict[str, Any]  # Contains query and max_results
+    output: Optional[SearchOutput]
+
 def create_calculator_tool() -> Graph:
     """Creates a calculator tool using LangGraph that can perform basic arithmetic operations."""
     print("Creating calculator tool")
     
-    def calculator_function(state: Dict[str, Any]) -> dict:
+    def calculator_function(state: CalculatorState) -> Dict[str, Any]:
         print("Calculator function called")
         input_data = state["input"]
         if len(input_data["numbers"]) < 2:
@@ -54,7 +62,7 @@ def create_calculator_tool() -> Graph:
         }
 
     # Create the graph with state schema
-    workflow = StateGraph(state_schema=Dict[str, Any])
+    workflow = StateGraph(state_schema=CalculatorState)
     print("Calculator graph for workflow created")
     # Add the calculator tool node
     workflow.add_node("calculator", ToolNode(calculator_function))
@@ -68,7 +76,7 @@ def create_calculator_tool() -> Graph:
 def create_search_tool() -> Graph:
     """Creates a search tool using DuckDuckGo that can search for information online."""
     
-    def search_function(state: Dict[str, Any]) -> dict:
+    def search_function(state: SearchState) -> Dict[str, Any]:
         with DDGS() as ddgs:
             raw_results = list(ddgs.text(
                 state["input"]["query"],
@@ -94,7 +102,7 @@ def create_search_tool() -> Graph:
         }
 
     # Create the graph with state schema
-    workflow = StateGraph(state_schema=Dict[str, Any])
+    workflow = StateGraph(state_schema=SearchState)
     
     # Add the search tool node
     workflow.add_node("search", ToolNode(search_function))