Update agent.py

agent.py

@@ -55,18 +55,54 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndB
 from langchain_huggingface import HuggingFaceEndpoint
 from langchain.agents import initialize_agent
 from langchain.agents import AgentType
+from typing import Union, List
+from functools import reduce
+import operator
 
 
 load_dotenv()
 
 
+
+
 @tool
-def calculator(inputs: dict):
-    """Perform mathematical operations based on the operation provided."""
+def calculator(inputs: Union[str, dict]):
+    """
+    Perform mathematical operations based on the operation provided.
+    Supports both binary (a, b) operations and list operations.
+    """
+
+    # If input is a JSON string, parse it
+    if isinstance(inputs, str):
+        try:
+            import json
+            inputs = json.loads(inputs)
+        except Exception as e:
+            return f"Invalid input format: {e}"
+
+    # Handle list-based operations like SUM
+    if "list" in inputs:
+        nums = inputs.get("list", [])
+        op = inputs.get("operation", "").lower()
+
+        if not isinstance(nums, list) or not all(isinstance(n, (int, float)) for n in nums):
+            return "Invalid list input. Must be a list of numbers."
+
+        if op == "sum":
+            return sum(nums)
+        elif op == "multiply":
+            return reduce(operator.mul, nums, 1)
+        else:
+            return f"Unsupported list operation: {op}"
+
+    # Handle basic two-number operations
     a = inputs.get("a")
     b = inputs.get("b")
-    operation = inputs.get("operation")
-    
+    operation = inputs.get("operation", "").lower()
+
+    if a is None or b is None or not isinstance(a, (int, float)) or not isinstance(b, (int, float)):
+        return "Both 'a' and 'b' must be numbers."
+
     if operation == "add":
         return a + b
     elif operation == "subtract":
@@ -80,7 +116,8 @@ def calculator(inputs: dict):
     elif operation == "modulus":
         return a % b
     else:
-        return "Unknown operation"
+        return f"Unknown operation: {operation}"
+
 
 
 
@@ -638,8 +675,14 @@ def select_tool_and_run(question: str, tools: list) -> any:
         print(f"No matching tool found for {intent}")
         return None
 
-    # Step 4: Call the tool function and return the result
-    return tool_func.run(question)  # or tool_func(question), depending on how your tools are structured
+    # If question was transformed into JSON for tool input
+    try:
+        parsed_input = json.loads(question)
+    except json.JSONDecodeError:
+        parsed_input = question  # fallback if question is not JSON
+
+    return tool_func.run(parsed_input)
+
 
 # Run the function to select and execute the tool
 result = select_tool_and_run(question, tools)