calculatorWithPydanticAdvanced.py

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

# Configuration
API_KEY = "YOUR API KEY"  # 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(question: str, max_iterations: int = 10):
    """
    Run the agentic loop following the diagram:
    User Question → Agent Memory → Gemini Reasoning → Tool Call Decision
    → Execute Tool → Store Observation → Continue Thinking
    """
    
    tools = create_tools()
    
    # Initialize conversation history (Agent Memory)
    conversation_history = []
    iteration = 0
    
    # Step 1: Initial user question
    print(f"\n{'='*60}")
    print(f"User Question: {question}")
    #question = input("Ask a math question: ") #uncomment to get user input
    print(f"{'='*60}\n")
    
    conversation_history.append({
        "role": "user",
        "parts": [genai.protos.Part(text=question)]
    })
    
    # Agentic loop
    while iteration < max_iterations:
        iteration += 1
        print(f"\n[Iteration {iteration}]")
        
        # Step 2: Gemini Reasoning
        response = model.generate_content(
            conversation_history,
            tools=tools,
            stream=False
        )
        
        candidate = response.candidates[0]
        
        # Add model's response to conversation history
        conversation_history.append({
            "role": "model",
            "parts": candidate.content.parts
        })
        
        # Step 3: Check if there's a tool call
        tool_called = False
        tool_results = []
        
        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
                
                # Unpack protobuf Map to standard Python dictionary
                args = {key: value for key, value in function_call.args.items()}
                
                print(f"🔧 Tool Call: {tool_name}")
                print(f"   Parameters: {args}")
                
                # Step 4: Execute Tool
                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}")
                        
                        # Step 5: Store Observation
                        tool_results.append({
                            "tool_name": tool_name,
                            "result": result
                        })
                    else:
                        raise ValueError(f"Unknown tool: {tool_name}")
                        
                except Exception as e:
                    print(f"   ✗ Error: {e}")
                    tool_results.append({
                        "tool_name": tool_name,
                        "result": f"Error: {str(e)}"
                    })
        
        # If no tool was called, the model has provided the final answer
        if not tool_called:
            print(f"\n✅ Final Answer:")
            print(f"{'-'*60}")
            # Extract and print the text response
            for part in candidate.content.parts:
                if hasattr(part, "text"):
                    print(part.text)
            print(f"{'-'*60}")
            break
        
        # Step 6: Continue Thinking - Add tool results to conversation
        # This allows the model to use the results and decide if it needs more tool calls
        function_response_parts = []
        
        for tool_result in tool_results:
            function_response_parts.append(
                genai.protos.Part(
                    function_response=genai.protos.FunctionResponse(
                        name=tool_result["tool_name"],
                        response={"result": tool_result["result"]}
                    )
                )
            )
        
        # Add function responses to history
        conversation_history.append({
            "role": "user",
            "parts": function_response_parts
        })
        
        print(f"📝 Observation stored, continuing thinking...")
    
    if iteration >= max_iterations:
        print(f"\n⚠️  Max iterations ({max_iterations}) reached")

# ============= 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(complex_question)