Category: Blog

# Create a virtual environment (recommended)
python -m venv .venv

# Activate the virtual environment
# On macOS/Linux:
source .venv/bin/activate
# On Windows (CMD):
.venv\Scripts\activate.bat
# On Windows (PowerShell):
.venv\Scripts\Activate.ps1

# Install ADK
pip install google-adk

parent_folder/
    weather_time_agent/
        __init__.py
        agent.py
        .env

mkdir -p weather_time_agent
echo "from . import agent" > weather_time_agent/__init__.py
touch weather_time_agent/agent.py
touch weather_time_agent/.env

import datetime
from zoneinfo import ZoneInfo
from google.adk.agents import Agent

def get_weather(city: str) -> dict:
    """Retrieves the current weather report for a specified city.

    Args:
        city (str): The name of the city for which to retrieve the weather report.

    Returns:
        dict: status and result or error msg.
    """
    if city.lower() == "new york":
        return {
            "status": "success",
            "report": (
                "The weather in New York is sunny with a temperature of 25 degrees"
                " Celsius (41 degrees Fahrenheit)."
            ),
        }
    else:
        return {
            "status": "error",
            "error_message": f"Weather information for '{city}' is not available.",
        }


def get_current_time(city: str) -> dict:
    """Returns the current time in a specified city.

    Args:
        city (str): The name of the city for which to retrieve the current time.

    Returns:
        dict: status and result or error msg.
    """

    if city.lower() == "new york":
        tz_identifier = "America/New_York"
    else:
        return {
            "status": "error",
            "error_message": (
                f"Sorry, I don't have timezone information for {city}."
            ),
        }

    tz = ZoneInfo(tz_identifier)
    now = datetime.datetime.now(tz)
    report = (
        f'The current time in {city} is {now.strftime("%Y-%m-%d %H:%M:%S %Z%z")}'
    )
    return {"status": "success", "report": report}


weather_time_agent = Agent(
    name="weather_time_agent",
    model="gemini-2.0-flash-exp",
    description=(
        "Agent to answer questions about the time and weather in a city."
    ),
    instruction=(
        "I can answer your questions about the time and weather in a city."
    ),
    tools=[get_weather, get_current_time],
)

from google.adk.agents import Agent  # This is actually an LlmAgent

my_agent = Agent(
    name="my_first_agent",
    model="gemini-2.0-flash-exp",
    description="A helpful assistant that answers general questions.",
    instruction="You are a friendly AI assistant. Be concise and helpful.",
    tools=[]  # Optional tools
)

from google.adk.agents import SequentialAgent

step1 = Agent(name="data_collector", model="gemini-2.0-flash-exp")
step2 = Agent(name="data_analyzer", model="gemini-2.0-flash-exp")

pipeline = SequentialAgent(
    name="analysis_pipeline",
    sub_agents=[step1, step2]  # Will execute in this order
)

from google.adk.agents import ParallelAgent

fetch_weather = Agent(name="weather_fetcher", model="gemini-2.0-flash-exp")
fetch_news = Agent(name="news_fetcher", model="gemini-2.0-flash-exp")

parallel_agent = ParallelAgent(
    name="information_gatherer",
    sub_agents=[fetch_weather, fetch_news]  # Will execute in parallel
)

from google.adk.agents import LoopAgent

process_step = Agent(name="process_item", model="gemini-2.0-flash-exp")
check_condition = Agent(name="check_complete", model="gemini-2.0-flash-exp")

loop_agent = LoopAgent(
    name="processing_loop",
    sub_agents=[process_step, check_condition],
    max_iterations=5  # Optional maximum iterations
)

from google.adk.agents import BaseAgent
from google.adk.agents.invocation_context import InvocationContext
from google.adk.events import Event
from typing import AsyncGenerator

class MyCustomAgent(BaseAgent):
    name: str = "custom_agent"
    description: str = "A specialized agent with custom behavior"
    
    async def _run_async_impl(self, context: InvocationContext) -> AsyncGenerator[Event, None]:
        # Custom implementation logic here
        # You must yield at least one Event
        yield Event(author=self.name, content=...)

# Using a Gemini model directly
agent = Agent(
    name="gemini_agent",
    model="gemini-2.0-flash-exp",  # Choose model variant based on needs
    # Other parameters...
)

agent = Agent(
    name="customer_support",
    model="gemini-2.0-flash-exp",
    instruction="""
    You are a customer support agent for TechGadgets Inc.
    
    When helping customers:
    1. Greet them politely and introduce yourself
    2. Ask clarifying questions if the issue isn't clear
    3. Provide step-by-step troubleshooting when appropriate
    4. For billing issues, use the check_account_status tool
    5. For technical problems, use the diagnostic_tool
    6. Always end by asking if there's anything else you can help with
    
    Never share internal company information or promise specific refund amounts.
    """
)

agent = Agent(
    name="billing_specialist",
    description="Handles customer billing inquiries and invoice issues.",
    # Other parameters...
)

recommendation_agent = Agent(
    name="product_recommender",
    # Other parameters...
    output_key="product_recommendation"
)

from google.adk.agents import Agent
from google.adk.models.lite_llm import LiteLlm

# Using OpenAI's GPT-4o
gpt_agent = Agent(
    name="gpt_agent",
    model=LiteLlm(model="openai/gpt-4o"),
    description="A GPT-powered agent",
    # Other parameters...
)

# Using Anthropic's Claude Sonnet
claude_agent = Agent(
    name="claude_agent",
    model=LiteLlm(model="anthropic/claude-3-sonnet-20240229"),
    description="A Claude-powered agent",
    # Other parameters...
)

# Using Mistral AI's model
mistral_agent = Agent(
    name="mistral_agent",
    model=LiteLlm(model="mistral/mistral-medium"),
    description="A Mistral-powered agent",
    # Other parameters...
)

def calculate_mortgage_payment(principal: float, annual_interest_rate: float, years: int) -> dict:
    """Calculates the monthly payment for a mortgage loan.
    
    Use this tool to determine monthly payments for a home loan based on
    principal amount, interest rate, and loan term.
    
    Args:
        principal: The initial loan amount in dollars.
        annual_interest_rate: The annual interest rate as a percentage (e.g., 5.5 for 5.5%).
        years: The loan term in years.
        
    Returns:
        dict: A dictionary containing the status ("success" or "error") and
              either the monthly payment or an error message.
    """
    try:
        # Convert annual interest rate to monthly decimal rate
        monthly_rate = (annual_interest_rate / 100) / 12
        # Calculate number of monthly payments
        num_payments = years * 12
        
        # Guard against division by zero or negative values
        if monthly_rate <= 0 or principal <= 0 or num_payments <= 0:
            return {
                "status": "error",
                "error_message": "All inputs must be positive, and interest rate cannot be zero."
            }
            
        # Calculate monthly payment using the mortgage formula
        if monthly_rate == 0:
            monthly_payment = principal / num_payments
        else:
            monthly_payment = principal * (monthly_rate * (1 + monthly_rate) ** num_payments) / ((1 + monthly_rate) ** num_payments - 1)
            
        return {
            "status": "success",
            "monthly_payment": round(monthly_payment, 2),
            "total_payments": round(monthly_payment * num_payments, 2),
            "total_interest": round((monthly_payment * num_payments) - principal, 2)
        }
    except Exception as e:
        return {
            "status": "error",
            "error_message": f"Failed to calculate mortgage payment: {str(e)}"
        }

# Add this tool to an agent
from google.adk.agents import Agent

mortgage_advisor = Agent(
    name="mortgage_advisor",
    model="gemini-2.0-flash-exp",
    description="Helps calculate and explain mortgage payments.",
    instruction="You are a mortgage advisor that helps users understand their potential mortgage payments. When asked about payments, use the calculate_mortgage_payment tool.",
    tools=[calculate_mortgage_payment]  # Simply include the function in the tools list
)

from google.adk.tools.tool_context import ToolContext

def update_user_preference(category: str, preference: str, tool_context: ToolContext) -> dict:
    """Updates a user's preference for a specific category.
    
    Args:
        category: The category for which to set a preference (e.g., "theme", "notifications").
        preference: The preference value to set.
        tool_context: Automatically provided by ADK, do not specify when calling.
        
    Returns:
        dict: Status of the preference update operation.
    """
    # Access current preferences or initialize if none exist
    user_prefs_key = "user:preferences"  # Using user: prefix makes this persistent across sessions
    preferences = tool_context.state.get(user_prefs_key, {})
    
    # Update the preferences
    preferences[category] = preference
    
    # Save back to state
    tool_context.state[user_prefs_key] = preferences
    
    print(f"Tool: Updated user preference '{category}' to '{preference}'")
    return {
        "status": "success", 
        "message": f"Your {category} preference has been set to {preference}"
    }

def escalate_to_support(issue_type: str, severity: int, tool_context: ToolContext) -> dict:
    """Escalates an issue to a human support agent.
    
    Args:
        issue_type: The type of issue being escalated.
        severity: The severity level (1-5, where 5 is most severe).
        tool_context: Automatically provided by ADK.
        
    Returns:
        dict: Status of the escalation.
    """
    # Record the escalation details in state
    tool_context.state["escalation_details"] = {
        "issue_type": issue_type,
        "severity": severity,
        "timestamp": datetime.datetime.now().isoformat()
    }
    
    # For high severity issues, transfer to the support agent
    if severity >= 4:
        tool_context.actions.transfer_to_agent = "human_support_agent"
        return {
            "status": "success",
            "message": "This is a high-severity issue. Transferring you to a human support specialist."
        }
    
    # For medium severity, just note it but don't transfer
    return {
        "status": "success",
        "message": f"Your {issue_type} issue has been logged with severity {severity}."
    }

weather_agent = Agent(
    name="weather_assistant",
    model="gemini-2.0-flash-exp",
    instruction="""
    You help users get weather information.
    
    When using the get_weather tool:
    1. Check the "status" field of the result.
    2. If status is "success", present the "report" information in a friendly way.
    3. If status is "error", apologize and share the "error_message" with the user.
    4. Always thank the user for their query.
    """,
    tools=[get_weather]
)

from google.adk.tools import google_search

search_agent = Agent(
    name="research_assistant",
    model="gemini-2.0-flash-exp",
    instruction="You help users research topics. When asked, use the google_search tool to find up-to-date information.",
    tools=[google_search]
)

from google.adk.tools import code_interpreter

coding_assistant = Agent(
    name="coding_assistant",
    model="gemini-2.0-flash-exp",
    instruction="You help users with coding tasks. When appropriate, use the code_interpreter to execute Python code and demonstrate solutions.",
    tools=[code_interpreter]
)

from google.adk.tools import rag_tool

# Configure RAG with your documents
my_rag_tool = rag_tool.configure(
    document_store="your-document-source",
    embedding_model="your-embedding-model"
)

documentation_assistant = Agent(
    name="docs_assistant",
    model="gemini-2.0-flash-exp",
    instruction="You help users find information in the company documentation. Use the RAG tool to retrieve relevant information.",
    tools=[my_rag_tool]
)

# LangChain tool example
from langchain.tools import BaseTool as LangChainTool
from google.adk.tools.langchain_tool import LangChainTool

langchain_tool = LangChainTool(langchain_tool=your_langchain_tool_instance)

# LlamaIndex tool example
from google.adk.tools.llama_index_tool import LlamaIndexTool

llama_index_tool = LlamaIndexTool(llama_index_tool=your_llama_index_tool_instance)

def search_product_catalog(
    query: str, 
    category: str = None, 
    price_max: float = None,
    sort_by: str = None,
    tool_context: ToolContext = None
) -> dict:
    """Searches the product catalog for items matching the query and filters.
    
    Use this tool to find products in our inventory based on customer requests.
    
    Args:
        query: The search term entered by the customer (required).
        category: Optional category to filter results (e.g., "electronics", "clothing").
        price_max: Optional maximum price filter.
        sort_by: Optional sorting method ("price_low", "price_high", "popularity", "rating").
        tool_context: Automatically provided by ADK.
        
    Returns:
        dict: A dictionary containing:
            - "status": "success" or "error"
            - If success: "products" list of matching products (up to 5 items)
            - If error: "error_message" explaining what went wrong
            
    Example success:
        {"status": "success", "products": [{"name": "42-inch TV", "price": 299.99, ...}, ...]}
    Example error:
        {"status": "error", "error_message": "No products found matching 'flying car'"}
    """
    try:
        # Log the tool execution for debugging
        print(f"Tool: search_product_catalog called with query='{query}', category='{category}', price_max={price_max}")
        
        # Track the search in user history if tool_context is available
        if tool_context:
            search_history = tool_context.state.get("user:search_history", [])
            search_history.append({
                "query": query,
                "timestamp": datetime.datetime.now().isoformat()
            })
            # Keep only last 10 searches
            if len(search_history) > 10:
                search_history = search_history[-10:]
            tool_context.state["user:search_history"] = search_history
        
        # ... actual catalog search implementation ...
        # (For demo, we'll return mock data)
        
        mock_products = [
            {"name": "42-inch Smart TV", "price": 299.99, "category": "electronics", "rating": 4.5},
            {"name": "Wireless Headphones", "price": 89.99, "category": "electronics", "rating": 4.2},
        ]
        
        # Apply filters if provided
        filtered_products = mock_products
        if category:
            filtered_products = [p for p in filtered_products if p["category"].lower() == category.lower()]
        if price_max:
            filtered_products = [p for p in filtered_products if p["price"] <= price_max]
            
        # Apply sorting if requested
        if sort_by == "price_low":
            filtered_products = sorted(filtered_products, key=lambda p: p["price"])
        elif sort_by == "price_high":
            filtered_products = sorted(filtered_products, key=lambda p: p["price"], reverse=True)
        elif sort_by == "rating":
            filtered_products = sorted(filtered_products, key=lambda p: p["rating"], reverse=True)
            
        # Return formatted response
        if filtered_products:
            return {
                "status": "success",
                "products": filtered_products[:5],  # Limit to 5 results
                "total_matches": len(filtered_products)
            }
        else:
            return {
                "status": "error",
                "error_message": f"No products found matching '{query}' with the specified filters."
            }
            
    except Exception as e:
        print(f"Tool Error: search_product_catalog failed: {str(e)}")
        return {
            "status": "error",
            "error_message": f"Failed to search catalog: {str(e)}"
        }

session.state = {
    # Session-specific state (default scope)
    "last_query": "What's the weather in London?",
    "current_step": 3,
    
    # User-specific state (persists across sessions)
    "user:preferred_temperature_unit": "Celsius",
    "user:name": "Alex",
    
    # Application-wide state (shared across all users)
    "app:version": "1.2.3",
    "app:maintenance_mode": False,
    
    # Temporary state (not persisted beyond current execution)
    "temp:calculation_result": 42
}

# Getting a session
from google.adk.sessions import InMemorySessionService

session_service = InMemorySessionService()
APP_NAME = "my_application"
USER_ID = "user_123"
SESSION_ID = "session_456"

# Create or retrieve a session
session = session_service.create_session(
    app_name=APP_NAME,
    user_id=USER_ID,
    session_id=SESSION_ID
)

# Reading from state
last_city = session.state.get("last_city", "New York")  # Default if key doesn't exist

# Writing to state
session.state["last_city"] = "London"

from google.adk.tools.tool_context import ToolContext

def remember_favorite_city(city: str, tool_context: ToolContext) -> dict:
    """Remembers the user's favorite city.
    
    Args:
        city: The city to remember as favorite.
        tool_context: Automatically provided by ADK.
        
    Returns:
        dict: Status of the operation.
    """
    # Store at user scope so it persists across sessions
    tool_context.state["user:favorite_city"] = city
    
    # Also store when this preference was set
    tool_context.state["user:favorite_city_set_at"] = datetime.datetime.now().isoformat()
    
    return {
        "status": "success",
        "message": f"I've remembered that your favorite city is {city}."
    }

weather_reporter = Agent(
    name="weather_reporter",
    model="gemini-2.0-flash-exp",
    instruction="You provide weather reports for cities. Be concise but informative.",
    tools=[get_weather],
    output_key="last_weather_report"  # Automatically saves response to this state key
)

personalized_agent = Agent(
    name="personalized_assistant",
    model="gemini-2.0-flash-exp",
    instruction="""
    You are a personalized assistant.
    
    CHECK THESE STATE VALUES AT THE START OF EACH INTERACTION:
    - If state["user:name"] exists, greet the user by name.
    - If state["user:favorite_city"] exists, personalize weather or travel recommendations.
    - If state["current_workflow"] exists, continue that workflow where you left off.
    
    MAINTAIN THESE STATE VALUES:
    - When the user mentions their name, use the remember_name tool to store it.
    - When discussing a city positively, use the remember_favorite_city tool.
    - When starting a multi-step workflow, set state["current_workflow"] and state["current_step"].
    """
)

def set_preference(category: str, value: str, tool_context: ToolContext) -> dict:
    """Stores a user preference.
    
    Args:
        category: The preference category (e.g., "language", "theme").
        value: The preference value.
        tool_context: Automatically provided by ADK.
        
    Returns:
        dict: Status of the operation.
    """
    preferences = tool_context.state.get("user:preferences", {})
    preferences[category] = value
    tool_context.state["user:preferences"] = preferences
    return {"status": "success", "message": f"Preference set: {category} = {value}"}

def get_preferences(tool_context: ToolContext) -> dict:
    """Retrieves all user preferences.
    
    Args:
        tool_context: Automatically provided by ADK.
        
    Returns:
        dict: The user's stored preferences.
    """
    preferences = tool_context.state.get("user:preferences", {})
    return {"status": "success", "preferences": preferences}

preference_agent = Agent(
    name="preference_aware_agent",
    model="gemini-2.0-flash-exp",
    instruction="""
    You help users and remember their preferences.
    
    At the start of each conversation:
    1. Use the get_preferences tool to check stored preferences.
    2. Adapt your responses based on these preferences.
    
    During conversations:
    1. When a user expresses a preference, use set_preference to store it.
    2. Acknowledge when you've saved a preference.
    
    Examples of preferences to track:
    - Language preferences
    - Communication style (brief/detailed)
    - Topic interests
    """,
    tools=[set_preference, get_preferences]
)

def start_workflow(workflow_name: str, tool_context: ToolContext) -> dict:
    """Starts a new workflow and tracks it in state.
    
    Args:
        workflow_name: The name of the workflow to start.
        tool_context: Automatically provided by ADK.
        
    Returns:
        dict: Status and the initial workflow state.
    """
    workflow = {
        "name": workflow_name,
        "current_step": 1,
        "started_at": datetime.datetime.now().isoformat(),
        "data": {}
    }
    tool_context.state["current_workflow"] = workflow
    return {"status": "success", "workflow": workflow}

def update_workflow_step(step: int, data: dict, tool_context: ToolContext) -> dict:
    """Updates the current workflow step and associated data.
    
    Args:
        step: The new step number.
        data: Data to associate with this step.
        tool_context: Automatically provided by ADK.
        
    Returns:
        dict: Status and the updated workflow state.
    """
    workflow = tool_context.state.get("current_workflow", {})
    if not workflow:
        return {"status": "error", "message": "No active workflow found."}
    
    workflow["current_step"] = step
    workflow["last_updated"] = datetime.datetime.now().isoformat()
    workflow["data"].update(data)
    tool_context.state["current_workflow"] = workflow
    return {"status": "success", "workflow": workflow}

workflow_agent = Agent(
    name="workflow_agent",
    model="gemini-2.0-flash-exp",
    instruction="""
    You guide users through structured workflows.
    
    At the start of each interaction:
    1. Check if state["current_workflow"] exists.
    2. If it exists, continue from the current_step.
    3. If not, determine if the user wants to start a workflow.
    
    Available workflows:
    - "account_setup": A 3-step process to set up a new account
    - "support_request": A 4-step process to file a support ticket
    
    Use start_workflow and update_workflow_step to track progress.
    """,
    tools=[start_workflow, update_workflow_step]
)

def update_conversation_summary(new_insight: str, tool_context: ToolContext) -> dict:
    """Updates the running summary of the conversation with a new insight.
    
    Args:
        new_insight: New information to add to the summary.
        tool_context: Automatically provided by ADK.
        
    Returns:
        dict: Status and the updated summary.
    """
    summary = tool_context.state.get("conversation_summary", "")
    if summary:
        summary += "\n- " + new_insight
    else:
        summary = "Conversation Summary:\n- " + new_insight
    
    tool_context.state["conversation_summary"] = summary
    return {"status": "success", "summary": summary}

summarizing_agent = Agent(
    name="summarizing_agent",
    model="gemini-2.0-flash-exp",
    instruction="""
    You help users while maintaining a summary of key points.
    
    At the start of each interaction:
    1. Check state["conversation_summary"] to recall context.
    
    During conversations:
    1. When you learn important information (preferences, goals, constraints),
       use update_conversation_summary to store it.
    2. Focus on facts and insights, not general chat.
    
    Keep your internal summary up-to-date to provide consistent, contextual help.
    """,
    tools=[update_conversation_summary]
)

from google.adk.agents import Agent, SequentialAgent
from google.adk.tools.tool_context import ToolContext

# --- Tools for personalization ---

def get_user_profile(tool_context: ToolContext) -> dict:
    """Retrieves the user's stored profile information.
    
    Args:
        tool_context: Automatically provided by ADK.
        
    Returns:
        dict: The user's profile data.
    """
    profile = tool_context.state.get("user:profile", {})
    return {
        "status": "success",
        "profile": profile,
        "is_returning_user": bool(profile)
    }

def update_user_profile(field: str, value: str, tool_context: ToolContext) -> dict:
    """Updates a specific field in the user's profile.
    
    Args:
        field: The profile field to update (e.g., "name", "occupation").
        value: The value to store.
        tool_context: Automatically provided by ADK.
        
    Returns:
        dict: Status of the operation.
    """
    profile = tool_context.state.get("user:profile", {})
    profile[field] = value
    tool_context.state["user:profile"] = profile
    return {"status": "success", "field": field, "value": value}

def log_user_interest(topic: str, score: float, tool_context: ToolContext) -> dict:
    """Records a user's interest in a topic with a relevance score.
    
    Args:
        topic: The topic of interest.
        score: Relevance score (0.0-1.0, higher means more interested).
        tool_context: Automatically provided by ADK.
        
    Returns:
        dict: Status of the operation.
    """
    interests = tool_context.state.get("user:interests", {})
    interests[topic] = max(interests.get(topic, 0), score) # Take highest score
    tool_context.state["user:interests"] = interests
    return {"status": "success", "topic": topic, "score": score}

def get_personalization_strategy(tool_context: ToolContext) -> dict:
    """Analyzes user data and returns a personalization strategy.
    
    Args:
        tool_context: Automatically provided by ADK.
        
    Returns:
        dict: Personalization recommendations based on user data.
    """
    profile = tool_context.state.get("user:profile", {})
    interests = tool_context.state.get("user:interests", {})
    interaction_count = tool_context.state.get("user:interaction_count", 0)
    
    # Increment interaction count
    tool_context.state["user:interaction_count"] = interaction_count + 1
    
    # Determine name usage style
    name_style = "formal"
    if interaction_count > 5 and "name" in profile:
        name_style = "casual"
    
    # Identify top interests
    top_interests = sorted(
        [(topic, score) for topic, score in interests.items()], 
        key=lambda x: x[1], 
        reverse=True
    )[:3]
    
    return {
        "status": "success",
        "strategy": {
            "name_usage": {
                "style": name_style,
                "name": profile.get("name", ""),
                "use_name": "name" in profile
            },
            "experience_level": "new" if interaction_count < 3 else "returning",
            "top_interests": top_interests,
            "verbosity": profile.get("preferred_verbosity", "balanced")
        }
    }

# --- Creating a personalized agent ---

personalization_agent = Agent(
    name="profile_manager",
    model="gemini-2.0-flash-exp",
    instruction="""
    You manage user profile information and personalization strategy.
    Your job is to extract and store relevant user information, then provide
    personalization guidance to other agents.
    
    YOU MUST:
    1. Use get_user_profile at the start of conversation to check existing data.
    2. During conversation, identify personal details and preferences.
    3. Use update_user_profile to store name, age, occupation, etc.
    4. Use log_user_interest when the user shows interest in topics.
    5. Use get_personalization_strategy to generate guidance for personalization.
    
    Do not explicitly tell the user you are storing this information.
    """,
    tools=[get_user_profile, update_user_profile, log_user_interest, get_personalization_strategy],
    output_key="personalization_strategy"
)

response_agent = Agent(
    name="personalized_responder",
    model="gemini-2.0-flash-exp",
    instruction="""
    You provide personalized responses based on the personalization strategy.
    
    At the beginning of each interaction:
    1. Check state["personalization_strategy"] for guidance on personalization.
    2. Adapt your tone, detail level, and content based on this strategy.
    
    Personalization Elements:
    1. If strategy says to use name, address the user by name per the specified style.
    2. Adapt verbosity based on preference.
    3. Reference top interests when relevant.
    4. Provide more explanation for new users, be more direct with returning users.
    
    Always keep your personalization subtle and natural, never explicit.
    """,
)

# Combine as a sequential workflow
personalized_assistant = SequentialAgent(
    name="personalized_assistant",
    sub_agents=[personalization_agent, response_agent]
)

from google.adk.sessions import SessionService, Session
from typing import Optional, Dict, Any
import firebase_admin
from firebase_admin import firestore

class FirestoreSessionService(SessionService):
    """A session service that persists state in Firestore."""
    
    def __init__(self, collection_name: str = "adk_sessions"):
        """Initialize with a Firestore collection name."""
        self.collection_name = collection_name
        if not firebase_admin._apps:
            firebase_admin.initialize_app()
        self.db = firestore.client()
        
    def create_session(
        self, app_name: str, user_id: str, session_id: str, state: Optional[Dict[str, Any]] = None
    ) -> Session:
        """Create a new session or get existing session."""
        session_ref = self._get_session_ref(app_name, user_id, session_id)
        doc = session_ref.get()
        
        if doc.exists:
            # Session exists, retrieve it
            session_data = doc.to_dict()
            return Session(
                app_name=app_name,
                user_id=user_id,
                session_id=session_id,
                state=session_data.get("state", {}),
                last_update_time=session_data.get("last_update_time", 0)
            )
        else:
            # Create new session
            session = Session(
                app_name=app_name,
                user_id=user_id,
                session_id=session_id,
                state=state or {}
            )
            self._save_session(session)
            return session
            
    def get_session(
        self, app_name: str, user_id: str, session_id: str
    ) -> Optional[Session]:
        """Get an existing session."""
        session_ref = self._get_session_ref(app_name, user_id, session_id)
        doc = session_ref.get()
        
        if not doc.exists:
            return None
            
        session_data = doc.to_dict()
        return Session(
            app_name=app_name,
            user_id=user_id,
            session_id=session_id,
            state=session_data.get("state", {}),
            last_update_time=session_data.get("last_update_time", 0)
        )
        
    def update_session(self, session: Session) -> None:
        """Update a session in the database."""
        self._save_session(session)
        
    def _get_session_ref(self, app_name: str, user_id: str, session_id: str):
        """Get a reference to the session document."""
        return self.db.collection(self.collection_name).document(
            f"{app_name}_{user_id}_{session_id}"
        )
        
    def _save_session(self, session: Session) -> None:
        """Save a session to Firestore."""
        session_ref = self._get_session_ref(
            session.app_name, session.user_id, session.session_id
        )
        session_ref.set({
            "state": session.state,
            "last_update_time": session.last_update_time
        })

from google.adk.agents import Agent

# Create specialized sub-agents
weather_specialist = Agent(
    name="weather_specialist",
    model="gemini-2.0-flash-exp",
    description="Provides detailed weather information for any location.",
    instruction="You are a weather specialist. Provide accurate, detailed weather information when asked.",
    tools=[get_weather]  # Assume get_weather is defined
)

restaurant_specialist = Agent(
    name="restaurant_specialist",
    model="gemini-2.0-flash-exp",
    description="Recommends restaurants based on location, cuisine, and preferences.",
    instruction="You are a restaurant specialist. Recommend restaurants based on user preferences.",
    tools=[find_restaurants]  # Assume find_restaurants is defined
)

# Create a parent agent with sub-agents
coordinator = Agent(
    name="travel_assistant",
    model="gemini-2.0-flash-exp",
    description="Helps plan trips and activities.",
    instruction="""
    You are a travel assistant that helps users plan trips and activities.
    
    You have two specialized sub-agents:
    - weather_specialist: For weather-related questions
    - restaurant_specialist: For restaurant recommendations
    
    When a user asks about weather, delegate to the weather_specialist.
    When a user asks about restaurants or food, delegate to the restaurant_specialist.
    For general travel questions, handle them yourself.
    """,
    sub_agents=[weather_specialist, restaurant_specialist]
)

# LLM-driven delegation relies on clear agent descriptions
customer_service = Agent(
    name="customer_service",
    model="gemini-2.0-flash-exp",
    description="Handles general customer inquiries and routes to specialists.",
    instruction="""
    You are the main customer service agent. 
    
    Analyze each customer query and determine the best way to handle it:
    - For billing questions, transfer to the billing_specialist
    - For technical issues, transfer to the tech_support
    - For product questions, handle yourself
    
    Make your delegation decisions based on the query content.
    """,
    sub_agents=[
        Agent(
            name="billing_specialist",
            model="gemini-2.0-flash-exp",
            description="Handles all billing, payment, and invoice inquiries."
        ),
        Agent(
            name="tech_support",
            model="gemini-2.0-flash-exp",
            description="Resolves technical issues and troubleshooting problems."
        )
    ]
)

from google.adk.agents import Agent
from google.adk.tools import AgentTool

# Create a specialized agent
calculator_agent = Agent(
    name="calculator",
    model="gemini-2.0-flash-exp",
    description="Performs complex mathematical calculations.",
    instruction="You perform mathematical calculations with precision."
)

# Wrap it as a tool
calculator_tool = AgentTool(
    agent=calculator_agent,
    description="Use this tool to perform complex calculations."
)

# Create a parent agent that uses the agent tool
math_tutor = Agent(
    name="math_tutor",
    model="gemini-2.0-flash-exp",
    description="Helps students learn mathematics.",
    instruction="""
    You are a math tutor helping students learn.
    
    When a student asks a question requiring complex calculations:
    1. Explain the mathematical concept
    2. Use the calculator tool to compute the result
    3. Explain the significance of the result
    """,
    tools=[calculator_tool]
)

from google.adk.agents import Agent, SequentialAgent

# First agent gathers information and stores it in state
information_gatherer = Agent(
    name="information_gatherer",
    model="gemini-2.0-flash-exp",
    instruction="Gather travel information from the user and store it in state.",
    tools=[
        # Tool to save travel details to state
        save_travel_details  # Assume this is defined and writes to state
    ],
    output_key="information_gathering_complete"  # Saves final response to state
)

# Second agent uses information from state
recommendation_generator = Agent(
    name="recommendation_generator",
    model="gemini-2.0-flash-exp",
    instruction="""
    Generate travel recommendations based on information in state.
    
    Look for:
    - destination in state["travel_destination"]
    - dates in state["travel_dates"]
    - preferences in state["travel_preferences"]
    """,
    tools=[
        # Tool to retrieve recommendations based on state information
        get_recommendations  # Assume this is defined and reads from state
    ]
)

# Sequential agent ensures these run in order
travel_planner = SequentialAgent(
    name="travel_planner",
    sub_agents=[information_gatherer, recommendation_generator]
)

from google.adk.agents import SequentialAgent, Agent

data_processor = SequentialAgent(
    name="data_processor",
    sub_agents=[
        Agent(name="data_validator", output_key="validation_result"),
        Agent(name="data_transformer", output_key="transformed_data"),
        Agent(name="data_analyzer", output_key="analysis_result"),
        Agent(name="report_generator")
    ]
)

from google.adk.agents import ParallelAgent, Agent

data_gatherer = ParallelAgent(
    name="data_gatherer",
    sub_agents=[
        Agent(name="weather_fetcher", output_key="weather_data"),
        Agent(name="traffic_fetcher", output_key="traffic_data"),
        Agent(name="news_fetcher", output_key="news_data")
    ]
)

from google.adk.agents import LoopAgent, Agent, BaseAgent
from google.adk.agents.invocation_context import InvocationContext
from google.adk.events import Event, EventActions
from typing import AsyncGenerator

# Custom agent that checks if the loop should continue
class ConditionChecker(BaseAgent):
    name: str = "condition_checker"
    
    async def _run_async_impl(self, context: InvocationContext) -> AsyncGenerator[Event, None]:
        # Check if the condition for stopping the loop is met
        completed = context.session.state.get("task_completed", False)
        max_iterations = context.session.state.get("max_iterations", 5)
        current_iteration = context.session.state.get("current_iteration", 0)
        
        # Increment iteration counter
        context.session.state["current_iteration"] = current_iteration + 1
        
        # If task is completed or max iterations reached, escalate to stop the loop
        if completed or current_iteration >= max_iterations:
            yield Event(
                author=self.name,
                actions=EventActions(escalate=True)  # This signals loop termination
            )
        else:
            yield Event(
                author=self.name,
                content=None  # No content needed, just continuing the loop
            )

# Create task processor agent
task_processor = Agent(
    name="task_processor",
    model="gemini-2.0-flash-exp",
    instruction="""
    Process the current task step.
    
    Check state["current_iteration"] to see which step you're on.
    When the task is complete, set state["task_completed"] = True.
    """,
    tools=[
        # Tool to process the current step
        process_step,  # Assume this is defined
        # Tool to mark the task as completed
        mark_completed  # Assume this is defined
    ]
)

# Create loop agent that combines processing and condition checking
iterative_processor = LoopAgent(
    name="iterative_processor",
    sub_agents=[
        task_processor,
        ConditionChecker()
    ],
    max_iterations=10  # Optional backup limit
)

# Financial advisory team with clear specializations
mortgage_specialist = Agent(
    name="mortgage_specialist",
    description="Expert on mortgage products, rates, and qualification requirements.",
    # Other parameters...
)

investment_specialist = Agent(
    name="investment_specialist",
    description="Expert on investment strategies, market trends, and portfolio management.",
    # Other parameters...
)

tax_specialist = Agent(
    name="tax_specialist",
    description="Expert on tax planning, deductions, and regulatory compliance.",
    # Other parameters...
)

# Hub agent coordinates specialists
financial_advisor = Agent(
    name="financial_advisor",
    description="Coordinates financial advice across multiple domains.",
    instruction="""
    You are the main financial advisor.
    
    For mortgage questions, delegate to mortgage_specialist.
    For investment questions, delegate to investment_specialist.
    For tax questions, delegate to tax_specialist.
    
    Only handle general financial questions yourself.
    """,
    sub_agents=[mortgage_specialist, investment_specialist, tax_specialist]
)

# Sequential workflow for loan processing
loan_processor = SequentialAgent(
    name="loan_processor",
    sub_agents=[
        Agent(name="application_validator"),
        Agent(name="credit_checker"),
        Agent(name="risk_assessor"),
        Agent(name="decision_maker"),
        Agent(name="notification_sender")
    ]
)

# Multi-level hierarchy for complex tasks
project_manager = Agent(
    name="project_manager",
    sub_agents=[
        Agent(
            name="design_lead",
            sub_agents=[
                Agent(name="ui_designer"),
                Agent(name="ux_researcher")
            ]
        ),
        Agent(
            name="development_lead",
            sub_agents=[
                Agent(name="frontend_developer"),
                Agent(name="backend_developer")
            ]
        ),
        Agent(name="qa_lead")
    ]
)

# First agent gathers information
data_collector = Agent(
    name="data_collector",
    instruction="""
    Collect information from the user. Store each piece in the appropriate state key:
    - Personal details in state["user_details"]
    - Goals in state["financial_goals"]
    - Current situation in state["current_situation"]
    """,
    tools=[save_to_state],  # Assume this tool saves data to specific state keys
    output_key="collection_complete"
)

# Specialist agents use collected information
retirement_planner = Agent(
    name="retirement_planner",
    instruction="""
    Create a retirement plan based on information in state.
    Use state["user_details"] for age and income information.
    Use state["financial_goals"] for retirement targets.
    Store your plan in state["retirement_plan"].
    """,
    tools=[create_retirement_plan],  # Assume this tool creates and saves a plan
    output_key="retirement_planning_complete"
)

from google.adk.agents import Agent, SequentialAgent
from google.adk.tools.tool_context import ToolContext

# Tool to check if the previous agent encountered an error
def check_previous_result(tool_context: ToolContext) -> dict:
    """Checks if the previous agent step was successful.
    
    Returns:
        dict: Status and whether a fallback is needed.
    """
    error_detected = tool_context.state.get("error_detected", False)
    return {
        "status": "success",
        "fallback_needed": error_detected,
        "error_details": tool_context.state.get("error_details", "Unknown error")
    }

# Tool to handle error recovery
def recover_from_error(error_details: str, tool_context: ToolContext) -> dict:
    """Attempts to recover from an error.
    
    Args:
        error_details: Details about the error that occurred.
        
    Returns:
        dict: Status of recovery attempt.
    """
    # Record the recovery attempt
    tool_context.state["recovery_attempted"] = True
    
    # Clear the error flag
    tool_context.state["error_detected"] = False
    
    return {
        "status": "success",
        "message": f"Recovered from error: {error_details}"
    }

# Primary agent that might encounter errors
primary_handler = Agent(
    name="primary_handler",
    model="gemini-2.0-flash-exp",
    instruction="""
    You handle the primary task.
    If you encounter an error, set state["error_detected"] = True and
    state["error_details"] = "description of error".
    """,
    tools=[process_task, set_error_state]  # Assume these are defined
)

# Fallback agent for error recovery
fallback_handler = Agent(
    name="fallback_handler",
    model="gemini-2.0-flash-exp",
    instruction="""
    You handle error recovery when the primary agent fails.
    
    First, use check_previous_result to see if you need to act.
    If fallback is needed, use recover_from_error to attempt recovery.
    Provide a simplified but functional response to the user.
    """,
    tools=[check_previous_result, recover_from_error]
)

# Combine with sequential flow
robust_handler = SequentialAgent(
    name="robust_handler",
    sub_agents=[primary_handler, fallback_handler]
)

from google.adk.tools.tool_context import ToolContext

def log_agent_action(action: str, details: str, tool_context: ToolContext) -> dict:
    """Logs an agent action to the trace log in state.
    
    Args:
        action: The type of action being logged.
        details: Details about the action.
        
    Returns:
        dict: Status of the logging operation.
    """
    # Get existing log or initialize new one
    trace_log = tool_context.state.get("agent_trace_log", [])
    
    # Add new entry with timestamp
    import time
    trace_log.append({
        "timestamp": time.time(),
        "agent": tool_context.agent_name,
        "action": action,
        "details": details
    })
    
    # Update state with new log
    tool_context.state["agent_trace_log"] = trace_log
    
    return {
        "status": "success"
    }

# Add this tool to all agents in your system for comprehensive tracing

from google.adk.agents.callback_context import CallbackContext
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.genai import types
from typing import Optional
import re

def profanity_filter(
    callback_context: CallbackContext, llm_request: LlmRequest
) -> Optional[LlmResponse]:
    """
    Checks user input for profanity and blocks requests containing prohibited language.
    
    Args:
        callback_context: Provides context about the agent and session
        llm_request: The request about to be sent to the LLM
        
    Returns:
        LlmResponse if the request should be blocked, None if it should proceed
    """
    # Simple profanity detection (in a real system, use a more sophisticated approach)
    prohibited_terms = ["badword1", "badword2", "badword3"]
    
    # Extract the last user message
    last_user_message = ""
    if llm_request.contents:
        for content in reversed(llm_request.contents):
            if content.role == 'user' and content.parts:
                if content.parts[0].text:
                    last_user_message = content.parts[0].text
                    break
    
    # Check for prohibited terms
    contains_profanity = any(term in last_user_message.lower() for term in prohibited_terms)
    
    if contains_profanity:
        # Log the blocking action
        print(f"Profanity filter blocked message: '{last_user_message[:20]}...'")
        
        # Record the event in state
        callback_context.state["profanity_filter_triggered"] = True
        
        # Return a response that will be sent instead of calling the LLM
        return LlmResponse(
            content=types.Content(
                role="model",
                parts=[types.Part(text="I'm sorry, but I cannot respond to messages containing inappropriate language. Please rephrase your request without using prohibited terms.")]
            )
        )
    
    # If no profanity detected, return None to allow the request to proceed
    return None

# Add the callback to an agent
safe_agent = Agent(
    name="safe_agent",
    model="gemini-2.0-flash-exp",
    instruction="You are a helpful assistant.",
    before_model_callback=profanity_filter
)

from google.adk.tools.base_tool import BaseTool
from google.adk.tools.tool_context import ToolContext
from typing import Optional, Dict, Any

def restricted_city_guardrail(
    tool: BaseTool, args: Dict[str, Any], tool_context: ToolContext
) -> Optional[Dict]:
    """
    Prevents the get_weather tool from being called for restricted cities.
    
    Args:
        tool: Information about the tool being called
        args: The arguments passed to the tool
        tool_context: Access to session state and other context
        
    Returns:
        Dict if the tool call should be blocked, None if it should proceed
    """
    # Check if this is the get_weather tool
    if tool.name == "get_weather" and "city" in args:
        city = args["city"].lower()
        
        # List of restricted cities (example - could be loaded dynamically)
        restricted_cities = ["restricted_city_1", "restricted_city_2"]
        
        if city in restricted_cities:
            # Log the blocking action
            print(f"Blocked get_weather call for restricted city: {city}")
            
            # Record the event in state
            tool_context.state["restricted_city_blocked"] = city
            
            # Return a response that will be used instead of calling the tool
            return {
                "status": "error",
                "error_message": f"Sorry, weather information for {city} is not available due to policy restrictions."
            }
    
    # For other tools or non-restricted cities, allow the call to proceed
    return None

# Add the callback to an agent
restricted_agent = Agent(
    name="restricted_agent",
    model="gemini-2.0-flash-exp",
    instruction="You provide weather information using the get_weather tool.",
    tools=[get_weather],  # Assume get_weather is defined
    before_tool_callback=restricted_city_guardrail
)

# Agent with multiple safety measures
comprehensive_agent = Agent(
    name="comprehensive_agent",
    model="gemini-2.0-flash-exp",
    instruction="You help users with various tasks safely and responsibly.",
    tools=[get_weather, search_web, send_email],  # Assume these are defined
    before_model_callback=content_safety_filter,  # Filter unsafe user input
    after_model_callback=output_sanitizer,        # Clean up model responses
    before_tool_callback=tool_usage_validator,    # Validate tool usage
    after_tool_callback=tool_result_logger        # Log tool results
)

# Define test cases in a structured format
test_cases = [
    {
        "name": "Basic weather query",
        "input": "What's the weather in New York?",
        "expected_tool_calls": ["get_weather"],
        "expected_tool_args": {"city": "New York"},
        "expected_response_contains": ["weather", "New York"]
    },
    {
        "name": "Ambiguous city query",
        "input": "How's the weather in Springfield?",
        "expected_tool_calls": ["clarify_city"],
        "expected_response_contains": ["multiple cities", "which Springfield"]
    },
    {
        "name": "City not supported",
        "input": "What's the weather in Atlantis?",
        "expected_tool_calls": ["get_weather"],
        "expected_tool_args": {"city": "Atlantis"},
        "expected_response_contains": ["don't have information", "Atlantis"]
    }
]

from google.adk.evaluation import AgentEvaluator

# Create the evaluator
evaluator = AgentEvaluator(agent=weather_agent)

# Run evaluation
evaluation_results = evaluator.evaluate(test_cases=test_cases)

# Print results
for result in evaluation_results:
    print(f"Test: {result.test_case['name']}")
    print(f"  Status: {'PASS' if result.success else 'FAIL'}")
    print(f"  Feedback: {result.feedback}")
    if not result.success:
        print(f"  Expected: {result.expected}")
        print(f"  Actual: {result.actual}")
    print()

# Calculate overall metrics
success_rate = sum(1 for r in evaluation_results if r.success) / len(evaluation_results)
print(f"Overall success rate: {success_rate:.2%}")

def evaluate_response_correctness(test_case, agent_response, tool_calls):
    """Evaluates the correctness of the agent's response for weather queries."""
    # Exact city match checker
    if "expected_tool_args" in test_case and "city" in test_case["expected_tool_args"]:
        expected_city = test_case["expected_tool_args"]["city"]
        
        # Find the actual city used in tool calls
        actual_city = None
        for call in tool_calls:
            if call["name"] == "get_weather" and "city" in call["args"]:
                actual_city = call["args"]["city"]
                break
        
        # Check city match
        city_match = (actual_city == expected_city)
        
        # Temperature format checker (should include °C or °F)
        temp_format_correct = False
        if "°C" in agent_response or "°F" in agent_response:
            temp_format_correct = True
        
        return {
            "city_match": city_match,
            "temp_format_correct": temp_format_correct,
            "overall_correct": city_match and temp_format_correct
        }
    
    return {"overall_correct": None}  # Not applicable for this test case

# Apply custom evaluation to results
for result in evaluation_results:
    correctness = evaluate_response_correctness(
        result.test_case, 
        result.actual_response, 
        result.actual_tool_calls
    )
    print(f"Test: {result.test_case['name']}")
    print(f"  Overall correct: {correctness['overall_correct']}")
    if "city_match" in correctness:
        print(f"  City match: {correctness['city_match']}")
    if "temp_format_correct" in correctness:
        print(f"  Temperature format: {correctness['temp_format_correct']}")
    print()

import unittest
from google.adk.evaluation import AgentEvaluator

class WeatherAgentTests(unittest.TestCase):
    def setUp(self):
        self.agent = create_weather_agent()  # Assume this function creates your agent
        self.evaluator = AgentEvaluator(agent=self.agent)
        
    def test_basic_weather_queries(self):
        results = self.evaluator.evaluate(test_cases=[
            {
                "name": "New York weather",
                "input": "What's the weather in New York?",
                "expected_tool_calls": ["get_weather"]
            }
        ])
        self.assertTrue(results[0].success, results[0].feedback)
        
    def test_ambiguous_cities(self):
        results = self.evaluator.evaluate(test_cases=[
            {
                "name": "Springfield ambiguity",
                "input": "How's the weather in Springfield?",
                "expected_response_contains": ["which Springfield", "multiple"]
            }
        ])
        self.assertTrue(results[0].success, results[0].feedback)
        
    def test_error_handling(self):
        results = self.evaluator.evaluate(test_cases=[
            {
                "name": "Nonexistent city",
                "input": "What's the weather in Narnia?",
                "expected_response_contains": ["don't have information", "Narnia"]
            }
        ])
        self.assertTrue(results[0].success, results[0].feedback)

if __name__ == "__main__":
    unittest.main()

import asyncio
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.genai import types

# Set up session and runner
session_service = InMemorySessionService()
APP_NAME = "streaming_app"
USER_ID = "user_123"
SESSION_ID = "session_456"

session = session_service.create_session(
    app_name=APP_NAME, 
    user_id=USER_ID,
    session_id=SESSION_ID
)

runner = Runner(
    agent=streaming_agent,  # Assume this is defined
    app_name=APP_NAME,
    session_service=session_service
)

async def stream_response(query: str):
    """Streams the agent's response token by token."""
    content = types.Content(role='user', parts=[types.Part(text=query)])
    
    print(f"User: {query}")
    print("Agent: ", end="", flush=True)
    
    # Process events as they arrive
    async for event in runner.run_async(
        user_id=USER_ID,
        session_id=SESSION_ID,
        new_message=content
    ):
        # For token-by-token streaming, look for ContentPartDelta events
        if hasattr(event, 'content_part_delta') and event.content_part_delta:
            delta = event.content_part_delta
            if delta.text:
                print(delta.text, end="", flush=True)
        
        # For final response
        if event.is_final_response():
            print()  # End line after response
            
    print("\n")  # Add space after complete response

# Run streaming interaction
async def main():
    queries = [
        "What's the weather in New York?",
        "How about London?",
        "Thanks for your help!"
    ]
    
    for query in queries:
        await stream_response(query)

# Run the async main function
asyncio.run(main())

import asyncio
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.genai import types
import sounddevice as sd
import numpy as np
import wave
import io

# Assume setup of session_service and runner as in previous example

async def audio_conversation():
    """Conducts a voice conversation with the agent."""
    # Audio recording parameters
    sample_rate = 16000
    recording_duration = 5  # seconds
    
    print("Press Enter to start recording your question...")
    input()
    
    # Record audio
    print("Recording... (5 seconds)")
    audio_data = sd.rec(
        int(recording_duration * sample_rate),
        samplerate=sample_rate,
        channels=1,
        dtype='int16'
    )
    sd.wait()  # Wait for recording to complete
    print("Recording complete.")
    
    # Convert audio to WAV format in memory
    audio_bytes = io.BytesIO()
    with wave.open(audio_bytes, 'wb') as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)  # 16-bit
        wf.setframerate(sample_rate)
        wf.writeframes(audio_data.tobytes())
    
    # Create audio content for the agent
    audio_part = types.Part.from_bytes(
        audio_bytes.getvalue(),
        mime_type="audio/wav"
    )
    content = types.Content(role='user', parts=[audio_part])
    
    print("Processing your question...")
    
    # Stream the response
    print("Agent response:")
    text_response = ""
    
    async for event in runner.run_async(
        user_id=USER_ID,
        session_id=SESSION_ID,
        new_message=content
    ):
        # Handle text streaming
        if hasattr(event, 'content_part_delta') and event.content_part_delta:
            delta = event.content_part_delta
            if delta.text:
                print(delta.text, end="", flush=True)
                text_response += delta.text
        
        # Handle final audio response
        if event.is_final_response() and event.content and event.content.parts:
            for part in event.content.parts:
                if part.mime_type and part.mime_type.startswith('audio/'):
                    # Play the audio response
                    audio_bytes = io.BytesIO(part.bytes_value)
                    with wave.open(audio_bytes, 'rb') as wf:
                        audio_data = np.frombuffer(
                            wf.readframes(wf.getnframes()),
                            dtype=np.int16
                        )
                        sd.play(audio_data, wf.getframerate())
                        sd.wait()
    
    print("\nConversation turn complete.")

# Run the audio conversation
asyncio.run(audio_conversation())