Demo Integrating Google ADK Agent With Snowflake MCP

# churn_agent_snowflake.py
import asyncio
import os
import warnings
from dotenv import load_dotenv

# Load environment variables from .env (if present).
load_dotenv()

# Suppress ADK experimental feature warnings (MCP support is labeled experimental in some ADK versions).
# Filter before ADK imports so the first warning fired at import time is ignored.
warnings.filterwarnings("ignore", message=r".*\[EXPERIMENTAL\].*")

# ADK imports
from google.adk.agents import LlmAgent
from google.adk.models.lite_llm import LiteLlm
from google.adk.runners import InMemoryRunner
from google.genai import types

# MCP-specific imports
from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
from mcp import StdioServerParameters

# -------------------------
# Model configuration
# -------------------------
# Example: Claude via KodeKey. Replace model name and keys with your provider values.
MODEL = LiteLlm(
    model="openai/claude-haiku-4-5-20251001",
    api_base=os.environ["KODEKEY_BASE_URL"],
    api_key=os.environ["KODEKEY_API_KEY"],
)

# Path to the MCP Toolbox binary (TOOLBOX_BIN env overrides the default here).
TOOLBOX = os.environ.get("TOOLBOX_BIN", "./toolbox")

# Fully qualified Snowflake table name (database.schema.table) used by the agent instruction.
TABLE = (
    f"{os.environ['SNOWFLAKE_DATABASE']}"
    f".{os.environ['SNOWFLAKE_SCHEMA']}.customer_features"
)

# -------------------------
# MCP Toolbox setup
# -------------------------
# The only substantive change vs other warehouse demonstrations: choose the "snowflake"
# prebuilt config/adapter in your MCP server/toolbox args. The toolbox process reads
# Create stdio connection params for the MCP server/toolbox.
conn_params = StdioConnectionParams()

# Server parameters: include the args to select the Snowflake adapter/config.
# Replace the ... with adapter-specific flags required by your toolbox version.
server_params = StdioServerParameters(...)  # e.g., server_params = StdioServerParameters(adapter="snowflake", ...)

# Construct the McpToolset that exposes a SQL tool backed by Snowflake.
# The constructor args depend on toolbox version and prebuilt configs; the key is to
# reference the Snowflake adapter in the server params above.
snowflake_tools = McpToolset(...)

# -------------------------
# Agent configuration
# -------------------------
agent = LlmAgent(
    name="churn_agent",
    model=MODEL,
    instruction=(
        "You help a customer-success team spot churn risk. The data lives in "
        f"the Snowflake table {TABLE} - one row per customer with columns: "
        "customer_id, region, plan_tier, tenure_months, usage_last_30d, "
        "usage_delta_30d, support_tickets_90d, upgrade_eligible. Use the SQL "
        "tool to fetch the customer the user refers to. Do NOT query the `churned` "
        "column. Interpret the signals: low usage_last_30d or a sharply "
        "negative usage_delta_30d means disengagement; high support_tickets_90d "
        "means friction; short tenure_months is riskier. State the risk level "
        "(low/medium/high), cite the signals, and suggest one retention action."
    ),
    tools=[snowflake_tools],
)

# -------------------------
# Runner / session / example query
# -------------------------
async def main():
    # In-memory runner for quick demos. Create a session for user "me".
    runner = InMemoryRunner(agent=agent, app_name="churn")
    session = await runner.session_service.create_session(app_name="churn", user_id="me")

    # Example user query referencing a specific customer ID.
    msg = types.Content(
        role="user",
        parts=[types.Part(text="Is customer 23e950bd-4270-4f47-b49b-bfda0135ef3a at risk?")]
    )

    # Send the message; ADK versions may expose slightly different service APIs.
    response = await runner.session_service.create_message(session.session_id, content=msg)

    # Print the response (field access may vary by ADK response object structure).
    print(response)

if __name__ == "__main__":
    asyncio.run(main())

python churn_agent_snowflake.py