Calculator Agent using Gemini AI

SimpleCalculatorAgent.py

import google.generativeai as genai
from pydantic import BaseModel, Field
from typing import Dict, Any

# Configuration
API_KEY = "YOUR API KEY"
genai.configure(api_key=API_KEY)
model = genai.GenerativeModel("gemini-2.5-flash")

# Define tool input schemas using Pydantic
class AddInput(BaseModel):
    """Add two numbers."""
    a: int = Field(description="First number")
    b: int = Field(description="Second number")

class SubtractInput(BaseModel):
    """Subtract second number from first."""
    a: int = Field(description="Number to subtract from")
    b: int = Field(description="Number to subtract")

class MultiplyInput(BaseModel):
    """Multiply two numbers."""
    a: int = Field(description="First number")
    b: int = Field(description="Second number")
    
class DivideInput(BaseModel):
    """Divide first number by second."""
    a: float = Field(description="Dividend")
    b: float = Field(description="Divisor")

# Define functions
def add(a: int, b: int) -> int:
    """Adds two numbers."""
    return a + b

def subtract(a: int, b: int) -> int:
    """Subtracts b from a."""
    return a - b

def multiply(a: int, b: int) -> int:
    """Multiplies two numbers."""
    return a * b

def divide(a: float, b: float) -> float:
    """Divides a by b."""
    if b == 0:
        raise ValueError("Cannot divide by zero")
    return a / b

# Register tools
tool_registry = {
    "add": {
        "function": add,
        "input_schema": AddInput,
        "description": "Adds two numbers"
    },
    "subtract": {
        "function": subtract,
        "input_schema": SubtractInput,
        "description": "Subtracts the second number from the first"
    },
    "multiply": {
        "function": multiply,
        "input_schema": MultiplyInput,
        "description": "Multiplies two numbers"
    },
    "divide": {
        "function": divide,
        "input_schema": DivideInput,
        "description": "Divides the first number by the second"
    }
}

# Convert to Gemini tool declarations
function_declarations = []
for tool_name, tool_info in tool_registry.items():
    schema = tool_info["input_schema"].model_json_schema()

    # Create properties using a standard dictionary matching Gemini's Schema format
    properties_dict = {}
    for prop, details in schema.get("properties", {}).items():
        # Map pydantic schema type fields to Gemini API type strings
        p_type = "STRING" if details.get("type") == "string" else "NUMBER"
        
        properties_dict[prop] = {
            "type": p_type,
            "description": details.get("description", "")
        }

    # Pass parameters as a standard Python dictionary
    declaration = genai.types.FunctionDeclaration(
        name=tool_name,
        description=tool_info["description"],
        parameters={
            "type": "OBJECT",
            "properties": properties_dict,
            "required": schema.get("required", [])
        }
    )
    function_declarations.append(declaration)
    
tools = [genai.types.Tool(function_declarations=function_declarations)]

# Get user input
question = input("Ask a math question: ")

# Send to model
response = model.generate_content(
    question,
    tools=tools
)

# Extract the first candidate and its primary content part
candidate = response.candidates[0]
first_part = candidate.content.parts[0]

# Check if the part contains a function call
if hasattr(first_part, "function_call") and first_part.function_call:
    function_call = first_part.function_call
    tool_name = function_call.name
    
    # Safely unpack the protobuf Map to a standard Python dictionary
    args = {key: value for key, value in function_call.args.items()}
    
    print(f"Selected tool: {tool_name}")
    print(f"Parameters: {args}")
    
    # Validate and execute
    try:
        if tool_name in tool_registry:
            input_schema = tool_registry[tool_name]["input_schema"]
            validated_input = input_schema(**args)
            
            tool_function = tool_registry[tool_name]["function"]
            result = tool_function(**validated_input.model_dump())
            
            print(f"Result: {result}")
        else:
            raise ValueError(f"Unknown tool: {tool_name}")
            
    except Exception as e:
        print(f"Error: {e}")
        result = None

    # Send the execution loop result back to Gemini
    if result is not None:
        # Re-construct user prompt using standard content role objects
        user_content = genai.protos.Content(role="user", parts=[genai.protos.Part(text=question)])
        model_content = candidate.content # Use the existing content object containing the function call
        
        function_content = genai.protos.Content(
            role="function",
            parts=[
                genai.protos.Part(
                    function_response=genai.protos.FunctionResponse(
                        name=tool_name,
                        response={"result": result}
                    )
                )
            ]
        )
        
        final_response = model.generate_content([user_content, model_content, function_content])
        print("\nFinal Answer:")
        print(final_response.text)
else:
    # Handle plain text answers gracefully (e.g. if you say "Hello")
    print("\nFinal Answer:")
    print(response.text)