from google import genai
from google.genai import types
import os
 
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
 
def save_to_database(topic: str, summary: str, source_url: str) -> dict:
    """Save research findings to your database (Firestore, Supabase, etc.)"""
    print(f"Saving to DB: topic={topic}")
    # Replace with your actual DB logic
    return {
        "success": True,
        "record_id": f"rec_{hash(topic) % 10000:04d}",
        "message": f"Saved '{topic}' successfully",
    }
 
# Define the custom tool
save_tool = types.FunctionDeclaration(
    name="save_to_database",
    description="Save research results to the database. Requires topic, summary, and source URL.",
    parameters=types.Schema(
        type="OBJECT",
        properties={
            "topic": types.Schema(type="STRING", description="Research topic"),
            "summary": types.Schema(type="STRING", description="Summary of findings (max 200 chars)"),
            "source_url": types.Schema(type="STRING", description="Primary source URL"),
        },
        required=["topic", "summary", "source_url"],
    ),
)
 
# Single request combining Built-in Search + custom function
response = client.models.generate_content(
    model="gemini-3-pro",
    contents="Research Gemini 3.1 Pro's latest features and save the results to the database.",
    config=types.GenerateContentConfig(
        tools=[
            types.Tool(google_search=types.GoogleSearch()),  # Built-in
            types.Tool(function_declarations=[save_tool]),   # Custom
        ],
        thinking_config=types.ThinkingConfig(thinking_budget=8192),
    ),
)
 
print(response.text)
# Expected output:
# I've searched for Gemini 3.1 Pro's latest features and saved the findings to the database.
# Key findings include: enhanced Function Calling with context circulation, improved...

def run_agent(user_query: str, max_iterations: int = 10) -> str:
    """
    Production agent loop with Context Circulation.
 
    The 'contents' list grows with each iteration, carrying forward
    all tool calls and their responses — this is Context Circulation.
    """
    contents = [types.Content(role="user", parts=[types.Part(text=user_query)])]
 
    config = types.GenerateContentConfig(
        tools=[
            types.Tool(google_search=types.GoogleSearch()),
            types.Tool(google_maps=types.GoogleMaps()),
            types.Tool(function_declarations=[save_tool, analyze_tool, notify_tool]),
        ],
        automatic_function_calling=types.AutomaticFunctionCallingConfig(disable=True),
    )
 
    for i in range(max_iterations):
        response = client.models.generate_content(
            model="gemini-3-pro",
            contents=contents,   # Entire accumulated context sent each time
            config=config,
        )
 
        candidate = response.candidates[0]
 
        # Done when no more tool calls
        if candidate.finish_reason == "STOP":
            return response.text
 
        tool_calls = [p for p in candidate.content.parts if p.function_call]
        if not tool_calls:
            return response.text
 
        # Add model's response to context (this is what "circulation" means)
        contents.append(candidate.content)
 
        # Execute each tool call and collect responses
        tool_results = []
        for part in tool_calls:
            fc = part.function_call
            result = dispatch_function(fc.name, dict(fc.args))
 
            tool_results.append(
                types.Part(
                    function_response=types.FunctionResponse(
                        id=fc.id,      # Tool Call Identifier — required for parallel calls
                        name=fc.name,
                        response={"result": result},
                    )
                )
            )
 
        # Add tool results to context — model sees these in the next iteration
        contents.append(types.Content(role="tool", parts=tool_results))
 
    return "Reached maximum iterations."
 
def dispatch_function(name: str, args: dict) -> dict:
    """Route function calls to their implementations."""
    handlers = {
        "save_to_database": save_to_database,
        "analyze_data": analyze_data,
        "send_notification": send_notification,
    }
    handler = handlers.get(name)
    if handler:
        return handler(**args)
    return {"error": f"Unknown function: {name}"}

import asyncio
from concurrent.futures import ThreadPoolExecutor
 
async def handle_parallel_tool_calls(tool_calls: list) -> list:
    """
    Execute multiple tool calls in parallel and return results
    correctly mapped via Tool Call Identifiers.
 
    Expected behavior:
    - 3 tool calls execute concurrently (total time ≈ slowest tool)
    - Each result is uniquely identified by its tool_call.id
    """
    async def run_one(tool_call):
        loop = asyncio.get_event_loop()
        with ThreadPoolExecutor() as pool:
            result = await loop.run_in_executor(
                pool,
                lambda: dispatch_function(tool_call.name, dict(tool_call.args))
            )
        return types.Part(
            function_response=types.FunctionResponse(
                id=tool_call.id,    # Critical: maps result back to correct request
                name=tool_call.name,
                response={"result": result},
            )
        )
 
    results = await asyncio.gather(*[run_one(tc) for tc in tool_calls])
    return list(results)

query = """
Please do the following in sequence:
1. Use Google Search to research 'Gemini 3.1 Pro new features 2026'
2. Run analyze_data on what you found to extract competitive insights
3. Call send_notification to post the analysis to our Slack channel
"""
 
result = run_agent(query)
# Context Circulation ensures that Step 2's analysis
# has full access to Step 1's search results, and
# Step 3's notification includes Step 2's analysis.

response = client.models.generate_content(
    model="gemini-3-flash",   # Flash is cost-effective for simpler location tasks
    contents="Find AI-themed cafes in Shibuya, Tokyo with ratings above 4.0 and save the top 5 to the database.",
    config=types.GenerateContentConfig(
        tools=[
            types.Tool(google_maps=types.GoogleMaps()),
            types.Tool(function_declarations=[save_tool]),
        ],
    ),
)

complex_query = """
Compare Gemini API vs Claude API for a use case processing 1M tokens/day
with strict latency requirements. Consider cost, throughput, and reliability.
Use Google Search for the latest pricing and benchmark data.
"""
 
response = client.models.generate_content(
    model="gemini-3-pro",
    contents=complex_query,
    config=types.GenerateContentConfig(
        tools=[
            types.Tool(google_search=types.GoogleSearch()),
            types.Tool(function_declarations=[calculate_cost_tool, compare_latency_tool]),
        ],
        thinking_config=types.ThinkingConfig(thinking_budget=16384),
    ),
)
 
# Optionally inspect the thinking process
for part in response.candidates[0].content.parts:
    if part.thought:
        print(f"[Thinking]: {part.text[:300]}...")

# BAD: Omitting the id field causes mapping errors in parallel scenarios
function_response=types.FunctionResponse(
    name=fc.name,
    response={"result": result},
)
 
# GOOD: Always include the id
function_response=types.FunctionResponse(
    id=fc.id,   # Never skip this
    name=fc.name,
    response={"result": result},
)

# Cache invariant system context to avoid re-sending it every iteration
cached_content = client.caches.create(
    model="gemini-3-pro",
    contents=[system_context],
    ttl="3600s",
)
 
response = client.models.generate_content(
    model="gemini-3-pro",
    contents=conversation_history,
    config=types.GenerateContentConfig(
        cached_content=cached_content.name,
        tools=[...],
    ),
)

config = types.GenerateContentConfig(
    system_instruction="""
    Tool usage priority:
    1. Use google_search FIRST to gather current information
    2. Pass gathered data to analyze_data for processing
    3. Use save_to_database to persist the final output
    Always follow this sequence in order.
    """,
    tools=[...],
)

def choose_model(complexity: str, needs_thinking: bool) -> str:
    if complexity == "high" or needs_thinking:
        return "gemini-3-pro"
    elif complexity == "medium":
        return "gemini-3-flash"
    return "gemini-3-flash-lite"

import asyncio, time, statistics
 
async def mock_tool(name: str, delay: float) -> dict:
    await asyncio.sleep(delay)
    return {"tool": name, "ok": True}
 
async def parallel_run(specs):
    return await asyncio.gather(*[mock_tool(n, d) for n, d in specs])
 
def sequential_run(specs):
    for _, d in specs:
        time.sleep(d)
 
def timed(fn) -> float:
    start = time.perf_counter()
    fn()
    return time.perf_counter() - start
 
specs = [("search", 1.4), ("maps", 0.9), ("save", 0.8)]
seq = statistics.median(timed(lambda: sequential_run(specs)) for _ in range(50))
par = statistics.median(timed(lambda: asyncio.run(parallel_run(specs))) for _ in range(50))
print(f"sequential={seq:.2f}s parallel={par:.2f}s")