Update tools.py

tools.py

@@ -3,6 +3,8 @@ from langgraph.graph import Graph, StateGraph
 from langgraph.prebuilt import ToolNode
 from duckduckgo_search import DDGS
 
+# --- Calculator Tool Types ---
+
 class CalculatorInput(TypedDict):
     operation: str
     numbers: List[float]
@@ -15,6 +17,8 @@ class CalculatorState(TypedDict):
     input: CalculatorInput
     output: Optional[CalculatorOutput]
 
+# --- Search Tool Types ---
+
 class SearchInput(TypedDict):
     query: str
     max_results: int
@@ -29,21 +33,23 @@ class SearchOutput(TypedDict):
     query: str
 
 class SearchState(TypedDict):
-    input: SearchInput  # Now typed correctly
+    input: SearchInput
     output: Optional[SearchOutput]
 
+# --- Calculator Tool Implementation ---
+
 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: CalculatorState) -> Dict[str, Any]:
         print("Calculator function called")
         input_data = state["input"]
+        
         if len(input_data["numbers"]) < 2:
             raise ValueError("At least two numbers are required for calculation")
-        
+
         result = input_data["numbers"][0]
-        
         for num in input_data["numbers"][1:]:
             if input_data["operation"] == "add":
                 result += num
@@ -57,7 +63,7 @@ def create_calculator_tool() -> Graph:
                 result /= num
             else:
                 raise ValueError(f"Unsupported operation: {input_data['operation']}")
-        
+
         return {
             "output": {
                 "result": result,
@@ -65,27 +71,26 @@ def create_calculator_tool() -> Graph:
             }
         }
 
-    # Create the graph with state schema
     workflow = StateGraph(state_schema=CalculatorState)
-    print("Calculator graph for workflow created")
-    # Add the calculator tool node
     workflow.add_node("calculator", ToolNode(calculator_function))
-    print("Calculator tool node added")
-    # Set the entry and exit points
     workflow.set_entry_point("calculator")
     workflow.set_finish_point("calculator")
-    print("Calculator workflow set")
+    print("Calculator workflow created and compiled")
     return workflow.compile()
 
+# --- Search Tool Implementation ---
+
 def create_search_tool() -> Graph:
     """Creates a search tool using DuckDuckGo that can search for information online."""
-    
+    print("Creating search tool")
+
     def search_function(state: SearchState) -> Dict[str, Any]:
+        print("Search function called")
+        query = state["input"]["query"]
+        max_results = state["input"].get("max_results", 3)
+
         with DDGS() as ddgs:
-            raw_results = list(ddgs.text(
-                state["input"]["query"],
-                max_results=state["input"]["max_results"]
-            ))
+            raw_results = list(ddgs.text(query, max_results=max_results))
 
         results = []
         for r in raw_results:
@@ -101,35 +106,13 @@ def create_search_tool() -> Graph:
         return {
             "output": {
                 "results": results,
-                "query": state["input"]["query"]
+                "query": query
             }
         }
 
-    # Create the graph with state schema
     workflow = StateGraph(state_schema=SearchState)
-    
-    # Add the search tool node
     workflow.add_node("search", ToolNode(search_function))
-    
-    # Set the entry and exit points
     workflow.set_entry_point("search")
     workflow.set_finish_point("search")
-    
+    print("Search workflow created and compiled")
     return workflow.compile()
-
-# Example usage:
-# if __name__ == "__main__":
-#     # Create the calculator tool
-#     calculator = create_calculator_tool()
-    
-#     # Example calculation
-#     result = calculator.invoke({
-#         "input": {
-#             "operation": "add",
-#             "numbers": [1, 2, 3, 4]
-#         }
-#     })
-    
-#     print(f"Result: {result['output'].result}")
-
-