06_advanced_tool_calling.py

"""
Advanced Tool Calling Demo

This demo demonstrates advanced tool calling concepts from LangChain:
- Tool schemas and validation
- Parallel tool execution
- Tool calling with structured outputs
- Error handling and retries
- Custom tool calling patterns
- Tool calling with different models

Based on: https://python.langchain.com/docs/concepts/tool_calling/

Prerequisites:
1. LangChain installed: uv pip install langchain
2. Model providers: uv pip install langchain-google-genai langchain-ollama
3. API keys configured
"""

import os
import getpass
import asyncio
import json
from datetime import datetime
from typing import List, Dict, Any, Optional
from pydantic import BaseModel, Field

from langchain.tools import tool, StructuredTool
from langchain.chat_models import init_chat_model
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from langchain_core.tools import BaseTool
from langchain_core.pydantic_v1 import validator

# Set up API keys
if not os.environ.get("GOOGLE_API_KEY"):
    os.environ["GOOGLE_API_KEY"] = getpass.getpass("Enter API key for Google Gemini: ")

# Initialize models
print("=== Initializing Models ===")

# Google Gemini (best for tool calling)
gemini_model = init_chat_model("gemini-2.5-flash", model_provider="google_genai", temperature=0.1)
print("✓ Gemini model initialized")

# Try Ollama as fallback
try:
    from langchain_ollama import ChatOllama
    ollama_model = ChatOllama(model="gpt-oss:20b", temperature=0.1)
    print("✓ Ollama model available")
except Exception as e:
    ollama_model = None
    print(f"✗ Ollama not available: {e}")

# Example 1: Basic Tool Schema Definition
print("\n=== Example 1: Tool Schema Definition ===")

class WeatherInput(BaseModel):
    """Input schema for weather tool."""
    location: str = Field(..., description="The city and state/country")
    units: str = Field(default="celsius", description="Temperature units: celsius or fahrenheit")
    
    @validator('units')
    def validate_units(cls, v):
        if v not in ['celsius', 'fahrenheit']:
            raise ValueError('Units must be celsius or fahrenheit')
        return v

@tool
def get_weather(location: str, units: str = "celsius") -> dict:
    """Get current weather for a location.
    
    Args:
        location: The city and state/country
        units: Temperature units (celsius or fahrenheit)
    """
    # Mock weather data
    temp = 22 if units == "celsius" else 72
    return {
        "location": location,
        "temperature": temp,
        "units": units,
        "condition": "sunny",
        "humidity": 65,
        "wind_speed": "10 km/h" if units == "celsius" else "6 mph"
    }

# Test tool schema
weather_result = get_weather.invoke({"location": "New York, NY", "units": "fahrenheit"})
print(f"Weather tool result: {weather_result}")

# # Example 2: Structured Tool with Custom Input/Output
print("\n=== Example 2: Structured Tool Creation ===")

class CalculatorInput(BaseModel):
    """Input for calculator operations."""
    operation: str = Field(..., description="Mathematical operation: add, subtract, multiply, divide")
    a: float = Field(..., description="First number")
    b: float = Field(..., description="Second number")

class CalculatorOutput(BaseModel):
    """Output from calculator operations."""
    result: float = Field(..., description="The calculation result")
    operation_performed: str = Field(..., description="Description of the operation")
    success: bool = Field(..., description="Whether the operation was successful")

def calculator_func(operation: str, a: float, b: float) -> CalculatorOutput:
    """Perform mathematical operations."""
    try:
        if operation == "add":
            result = a + b
        elif operation == "subtract":
            result = a - b
        elif operation == "multiply":
            result = a * b
        elif operation == "divide":
            if b == 0:
                return CalculatorOutput(
                    result=0,
                    operation_performed=f"Failed to divide {a} by {b}",
                    success=False
                )
            result = a / b
        else:
            return CalculatorOutput(
                result=0,
                operation_performed=f"Unknown operation: {operation}",
                success=False
            )
        
        return CalculatorOutput(
            result=result,
            operation_performed=f"{operation}: {a} and {b} = {result}",
            success=True
        )
    except Exception as e:
        return CalculatorOutput(
            result=0,
            operation_performed=f"Error: {str(e)}",
            success=False
        )

# # Create structured tool
calculator_tool = StructuredTool.from_function(
    func=calculator_func,
    name="calculator",
    description="Perform mathematical calculations",
    args_schema=CalculatorInput,
    return_schema=CalculatorOutput
)

calc_result = calculator_tool.invoke({"operation": "multiply", "a": 15, "b": 7})
print(f"Calculator result: {calc_result}")

# Example 3: Custom Tool Class
print("\n=== Example 3: Custom Tool Class ===")

class DateTimeTool(BaseTool):
    """Custom tool for date and time operations."""
    name: str = "datetime_tool"
    description: str = "Get current date/time or perform date calculations"
    
    def _run(self, query: str) -> str:
        """Execute the tool."""
        current_time = datetime.now()
        
        if "current" in query.lower():
            return f"Current date and time: {current_time.strftime('%Y-%m-%d %H:%M:%S')}"
        elif "date" in query.lower():
            return f"Current date: {current_time.strftime('%Y-%m-%d')}"
        elif "time" in query.lower():
            return f"Current time: {current_time.strftime('%H:%M:%S')}"
        else:
            return f"DateTime info: {current_time.isoformat()}"
    
    async def _arun(self, query: str) -> str:
        """Async execution."""
        return self._run(query)

datetime_tool = DateTimeTool()
datetime_result = datetime_tool.invoke({"query": "current date and time"})
print(f"DateTime result: {datetime_result}")

# Example 4: Parallel Tool Execution
print("\n=== Example 4: Parallel Tool Execution ===")

@tool
def fetch_user_data(user_id: int) -> dict:
    """Fetch user data by ID."""
    import time
    time.sleep(1)  # Simulate API call
    return {
        "user_id": user_id,
        "name": f"User {user_id}",
        "email": f"user{user_id}@example.com",
        "status": "active"
    }

@tool
def fetch_user_orders(user_id: int) -> list:
    """Fetch user orders by user ID."""
    import time
    time.sleep(1)  # Simulate API call
    return [
        {"order_id": f"ORD-{user_id}-001", "total": 99.99, "status": "completed"},
        {"order_id": f"ORD-{user_id}-002", "total": 149.50, "status": "pending"}
    ]

# Test with tool calling model
tools = [get_weather, calculator_tool, datetime_tool, fetch_user_data, fetch_user_orders]
model_with_tools = gemini_model.bind_tools(tools)

# Example of parallel execution
response = model_with_tools.invoke(
    "Get the current weather in London and calculate 25 * 4, and tell me the current time"
)

print(f"AI Response: {response.content}")
print(f"Tool calls made: {len(response.tool_calls) if hasattr(response, 'tool_calls') else 0}")

# Execute tools if called
if hasattr(response, 'tool_calls') and response.tool_calls:
    print("\nExecuting tools in parallel...")
    results = []
    
    for tool_call in response.tool_calls:
        tool_name = tool_call['name']
        tool_args = tool_call['args']
        
        # Find the tool
        tool_to_call = next((t for t in tools if hasattr(t, 'name') and t.name == tool_name), None)
        if not tool_to_call:
            tool_to_call = next((t for t in tools if t.__name__ == tool_name), None)
        
        if tool_to_call:
            try:
                result = tool_to_call.invoke(tool_args)
                results.append(f"✓ {tool_name}: {result}")
            except Exception as e:
                results.append(f"✗ {tool_name}: Error - {e}")
    
    for result in results:
        print(result)

# Example 5: Tool Calling with Error Handling and Retries
print("\n=== Example 5: Error Handling and Retries ===")

@tool
def unreliable_api_call(endpoint: str) -> dict:
    """Simulate an unreliable API that sometimes fails."""
    import random
    
    if random.random() < 0.3:  # 30% failure rate
        raise Exception(f"API endpoint {endpoint} is temporarily unavailable")
    
    return {
        "endpoint": endpoint,
        "status": "success",
        "data": f"Sample data from {endpoint}",
        "timestamp": datetime.now().isoformat()
    }

def safe_tool_execution(tool, args, max_retries=3):
    """Execute tool with retry logic."""
    for attempt in range(max_retries):
        try:
            result = tool.invoke(args)
            return {"success": True, "result": result, "attempts": attempt + 1}
        except Exception as e:
            if attempt == max_retries - 1:
                return {"success": False, "error": str(e), "attempts": attempt + 1}
            print(f"Attempt {attempt + 1} failed: {e}. Retrying...")
    
    return {"success": False, "error": "Max retries exceeded", "attempts": max_retries}

# Test retry logic
api_result = safe_tool_execution(unreliable_api_call, {"endpoint": "/users"})
print(f"API call result: {api_result}")

# # Example 6: Advanced Tool Calling Patterns
print("\n=== Example 6: Advanced Tool Calling Patterns ===")

class MultiStepWorkflow:
    """Manage multi-step tool calling workflows."""
    
    def __init__(self, model, tools):
        self.model = model.bind_tools(tools)
        self.tools = {t.name if hasattr(t, 'name') else t.__name__: t for t in tools}
        self.conversation = []
    
    def execute_workflow(self, initial_query: str, max_steps: int = 5):
        """Execute a multi-step workflow."""
        self.conversation = [HumanMessage(content=initial_query)]
        step = 0
        
        while step < max_steps:
            step += 1
            print(f"\n--- Step {step} ---")
            
            response = self.model.invoke(self.conversation)
            print(f"AI: {response.content}")
            
            # Add AI response to conversation
            self.conversation.append(response)
            
            # Execute any tool calls
            if hasattr(response, 'tool_calls') and response.tool_calls:
                tool_results = []
                
                for tool_call in response.tool_calls:
                    tool_name = tool_call['name']
                    tool_args = tool_call['args']
                    call_id = tool_call.get('id', f"call_{step}")
                    
                    print(f"Executing {tool_name} with args: {tool_args}")
                    
                    if tool_name in self.tools:
                        try:
                            result = self.tools[tool_name].invoke(tool_args)
                            print(f"Tool result: {result}")
                            
                            # Add tool result to conversation
                            tool_message = ToolMessage(
                                content=str(result),
                                tool_call_id=call_id
                            )
                            self.conversation.append(tool_message)
                            tool_results.append(result)
                            
                        except Exception as e:
                            error_message = f"Error executing {tool_name}: {str(e)}"
                            print(f"Tool error: {error_message}")
                            
                            tool_message = ToolMessage(
                                content=error_message,
                                tool_call_id=call_id
                            )
                            self.conversation.append(tool_message)
                
                # If tools were executed, continue the conversation
                if tool_results:
                    continue
            
            # No more tool calls, workflow complete
            break
        
        print(f"\nWorkflow completed in {step} steps")
        return self.conversation

# # Test multi-step workflow
workflow = MultiStepWorkflow(gemini_model, tools)
conversation = workflow.execute_workflow(
    "I need to check the weather in Tokyo, then calculate the tip for a $85 bill at 18%, and get the current time"
)

# Example 7: Tool Calling with Structured Outputs
print("\n=== Example 7: Structured Output Analysis ===")

class AnalysisResult(BaseModel):
    """Structured output for analysis results."""
    summary: str = Field(..., description="Brief summary of the analysis")
    key_points: List[str] = Field(..., description="List of key findings")
    confidence_score: float = Field(..., description="Confidence in analysis (0-1)")
    recommendations: List[str] = Field(..., description="Recommended actions")

@tool
def analyze_data(data_description: str) -> AnalysisResult:
    """Analyze data and return structured results."""
    # Mock analysis
    return AnalysisResult(
        summary=f"Analysis of {data_description} completed successfully",
        key_points=[
            "Data shows positive trends",
            "Some anomalies detected in recent data",
            "Overall performance is above average"
        ],
        confidence_score=0.85,
        recommendations=[
            "Continue monitoring trends",
            "Investigate anomalies further",
            "Maintain current strategy"
        ]
    )

analysis_result = analyze_data.invoke({"data_description": "quarterly sales performance"})
print(f"Structured analysis: {json.dumps(analysis_result.dict(), indent=2)}")

print("\n=== Demo Complete ===")
print("This demo covered advanced tool calling concepts:")
print("- Tool schema definition and validation")
print("- Structured tools with custom input/output")
print("- Custom tool classes")
print("- Parallel tool execution")
print("- Error handling and retries")
print("- Multi-step workflows")
print("- Structured output patterns")
print("\nThese patterns enable building sophisticated AI agents with robust tool integration.")