Update agent.py

agent.py

@@ -43,10 +43,8 @@ from io import StringIO
 from transformers import BertTokenizer, BertModel
 import torch
 import torch.nn.functional as F
-#from langchain.tools import Tool
 from langchain.agents import initialize_agent, AgentType
-from langchain_community.chat_models import ChatOpenAI
-#from langchain_openai import ChatOpenAI
+from langchain_openai import ChatOpenAI
 from langchain_community.tools import Tool
 
 import time
@@ -104,27 +102,31 @@ def modulus(a: int, b: int) -> int:
     """
     return a % b
 
+
 @tool
-def calculator(a: int, b: int, operation: str) -> float:
+def calculator(data: dict) -> float:
     """
     Perform a calculation between two numbers.
 
     Args:
-        a: First number.
-        b: Second number.
-        operation: One of 'add', 'subtract', 'multiply', 'divide', 'modulus'.
+        data: A dictionary with keys 'a', 'b', and 'operation'
     """
-    operation = operation.lower()
+    a = data.get("a")
+    b = data.get("b")
+    operation = data.get("operation", "").lower()
+
     if operation == "add":
-        return add(a, b)
+        return a + b
     elif operation == "subtract":
-        return subtract(a, b)
+        return a - b
     elif operation == "multiply":
-        return multiply(a, b)
+        return a * b
     elif operation == "divide":
-        return divide(a, b)
+        if b == 0:
+            raise ValueError("Cannot divide by zero.")
+        return a / b
     elif operation == "modulus":
-        return modulus(a, b)
+        return a % b
     else:
         raise ValueError(f"Unsupported operation: {operation}")