calculatorWithPydanticRobust.py

gistfile1.txt

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

# Configuration
API_KEY = "AIzaSyDfSpmkAB8zM7LAOIEzzYWthX7xO9r92TA"  # Replace with your actual 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: float = Field(description="First number")
    b: float = Field(description="Second number")

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

class MultiplyInput(BaseModel):
    """Multiply two numbers."""
    a: float = Field(description="First number")
    b: float = 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: float, b: float) -> float:
    """Adds two numbers."""
    return a + b

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

def multiply(a: float, b: float) -> float:
    """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
def create_tools():
    """Create Gemini tool declarations from registry."""
    function_declarations = []
    
    for tool_name, tool_info in tool_registry.items():
        schema = tool_info["input_schema"].model_json_schema()

        # Map pydantic schema types to Gemini API type strings
        properties_dict = {}
        for prop, details in schema.get("properties", {}).items():
            p_type = "STRING" if details.get("type") == "string" else "NUMBER"
            
            properties_dict[prop] = {
                "type": p_type,
                "description": details.get("description", "")
            }

        # Create FunctionDeclaration
        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)
    
    return [genai.types.Tool(function_declarations=function_declarations)]

# ============= AGENTIC LOOP =============
def run_agent_loop_robust(question: str, max_iterations: int = 10):
    conversation_history = []
    iteration = 0
    
    conversation_history.append({
        "role": "user",
        "parts": [genai.protos.Part(text=question)]
    })

    tools = create_tools()
    
    while iteration < max_iterations:
        iteration += 1
        
        try:
            response = model.generate_content(
                conversation_history,
                tools=tools,
                stream=False
            )
        except Exception as e:
            print(f"❌ API Error at iteration {iteration}: {e}")
            break
        
        candidate = response.candidates[0]
        conversation_history.append({
            "role": "model",
            "parts": candidate.content.parts
        })
        
        tool_called = False
        
        for part in candidate.content.parts:
            if hasattr(part, "function_call") and part.function_call:
                tool_called = True
                function_call = part.function_call
                tool_name = function_call.name
                args = {key: value for key, value in function_call.args.items()}
                
                try:
                    # Execute with error handling
                    if tool_name not in tool_registry:
                        raise ValueError(f"Unknown tool: {tool_name}")
                    
                    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())
                    
                    error_msg = None
                    
                except Exception as e:
                    # Tool execution failed
                    print(f"❌ Tool {tool_name} failed: {e}")
                    result = None
                    error_msg = str(e)
                
                # Store result (success or error)
                if error_msg:
                    function_response_parts = [
                        genai.protos.Part(
                            function_response=genai.protos.FunctionResponse(
                                name=tool_name,
                                response={"error": error_msg}
                            )
                        )
                    ]
                else:
                    function_response_parts = [
                        genai.protos.Part(
                            function_response=genai.protos.FunctionResponse(
                                name=tool_name,
                                response={"result": result}
                            )
                        )
                    ]
                
                conversation_history.append({
                    "role": "user",
                    "parts": function_response_parts
                })
        
        if not tool_called:
            # Final answer
            for part in candidate.content.parts:
                if hasattr(part, "text"):
                    print(f"\n✅ Final Answer:\n{part.text}")
            break
    
    if iteration >= max_iterations:
        print(f"\n⚠️  Reached max iterations ({max_iterations})")

# ============= MAIN =============
if __name__ == "__main__":
    # Test with complex expression
    complex_question = "Calculate (100 + (200 / 2)) * 5 - 10"
    
    # You can also test with simpler questions
    # simple_question = "What is 5 + 3?"
    
    run_agent_loop_robust(complex_question)