Update tools.py

tools.py

@@ -12,6 +12,10 @@ class CalculatorOutput(BaseModel):
     result: float = Field(..., description="The result of the calculation")
     operation: str = Field(..., description="The operation that was performed")
 
+class CalculatorState(BaseModel):
+    input: CalculatorInput
+    output: CalculatorOutput = None
+
 class SearchInput(BaseModel):
     query: str = Field(..., description="The search query to look up")
     max_results: int = Field(default=3, description="Maximum number of results to return")
@@ -25,40 +29,42 @@ class SearchOutput(BaseModel):
     results: List[SearchResult] = Field(..., description="List of search results")
     query: str = Field(..., description="The original search query")
 
+class SearchState(BaseModel):
+    input: SearchInput
+    output: SearchOutput = None
+
 def create_calculator_tool() -> Graph:
     """Creates a calculator tool using LangGraph that can perform basic arithmetic operations."""
     
-    def calculator_function(state: Dict[str, Any]) -> Dict[str, Any]:
-        input_data = CalculatorInput(**state["input"])
-        
-        if len(input_data.numbers) < 2:
+    def calculator_function(state: CalculatorState) -> CalculatorState:
+        if len(state.input.numbers) < 2:
             raise ValueError("At least two numbers are required for calculation")
         
-        result = input_data.numbers[0]
+        result = state.input.numbers[0]
         
-        for num in input_data.numbers[1:]:
-            if input_data.operation == "add":
+        for num in state.input.numbers[1:]:
+            if state.input.operation == "add":
                 result += num
-            elif input_data.operation == "subtract":
+            elif state.input.operation == "subtract":
                 result -= num
-            elif input_data.operation == "multiply":
+            elif state.input.operation == "multiply":
                 result *= num
-            elif input_data.operation == "divide":
+            elif state.input.operation == "divide":
                 if num == 0:
                     raise ValueError("Cannot divide by zero")
                 result /= num
             else:
-                raise ValueError(f"Unsupported operation: {input_data.operation}")
+                raise ValueError(f"Unsupported operation: {state.input.operation}")
         
-        output = CalculatorOutput(
+        state.output = CalculatorOutput(
             result=result,
-            operation=input_data.operation
+            operation=state.input.operation
         )
         
-        return {"output": output}
+        return state
 
     # Create the graph with state schema
-    workflow = StateGraph(dict)
+    workflow = StateGraph(state_schema=CalculatorState)
     
     # Add the calculator tool node
     workflow.add_node("calculator", ToolNode(calculator_function))
@@ -72,15 +78,13 @@ 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[str, Any]:
-        input_data = SearchInput(**state["input"])
-        
+    def search_function(state: SearchState) -> SearchState:
         # Initialize DuckDuckGo search
         with DDGS() as ddgs:
             # Perform the search
             search_results = list(ddgs.text(
-                input_data.query,
-                max_results=input_data.max_results
+                state.input.query,
+                max_results=state.input.max_results
             ))
             
             # Convert results to our model
@@ -93,15 +97,15 @@ def create_search_tool() -> Graph:
                 for result in search_results
             ]
             
-            output = SearchOutput(
+            state.output = SearchOutput(
                 results=results,
-                query=input_data.query
+                query=state.input.query
             )
             
-            return {"output": output}
+            return state
 
     # Create the graph with state schema
-    workflow = StateGraph(dict)
+    workflow = StateGraph(state_schema=SearchState)
     
     # Add the search tool node
     workflow.add_node("search", ToolNode(search_function))