Update agent.py

agent.py

@@ -0,0 +1,72 @@
+import os
+from smolagents import CodeAgent, DuckDuckGoSearchTool, tool, LiteLLMModel
+from dotenv import load_dotenv
+load_dotenv()
+
+api_key = os.getenv("OPENAI_API_KEY")
+
+@tool
+def add(a: int, b: int) -> int:
+    """
+    Adds two numbers together.
+    Args:
+        a (int): The first number.
+        b (int): The second number.
+    Returns:
+        int: The result of a + b.
+    """
+    return a + b
+
+@tool
+def subtract(a: int, b: int) -> int:
+    """
+    Subtracts the second number from the first.
+    Args:
+        a (int): The number to subtract from.
+        b (int): The number to subtract.
+    Returns:
+        int: The result of a - b.
+    """
+    return a - b
+
+@tool
+def multiply(a: int, b: int) -> int:
+    """
+    Multiplies two numbers.
+    Args:
+        a (int): The first number.
+        b (int): The second number.
+    Returns:
+        int: The product of a and b.
+    """
+    return a * b
+
+@tool
+def divide(a: int, b: int) -> float:
+    """
+    Divides the first number by the second.
+    Args:
+        a (int): The numerator.
+        b (int): The denominator.
+    Returns:
+        float: The result of a / b.
+    Raises:
+        ZeroDivisionError: If b is zero.
+    """
+    return a / b
+
+
+class BasicAgent:
+    def __init__(self):
+        print("BasicAgent initialized.")
+        llm = LiteLLMModel("openai/gpt-4o-mini", api_key = api_key)
+        
+        self.agent = CodeAgent(
+            tools=[add,subtract,multiply,divide,DuckDuckGoSearchTool()],
+            model=llm
+        )
+    def __call__(self, question: str) -> str:
+        print(f"Agent received question (first 50 chars): {question[:50]}...")
+        fixed_answer = self.agent.run(question)
+        print(f"Agent returning fixed answer: {fixed_answer}")
+        return fixed_answer