Demo Building Agent With Databricks

# --- Model (all scripts) - Claude via KodeKey ---
KODEKEY_BASE_URL=https://your-kodekey-endpoint/v1
KODEKEY_API_KEY=your-kodekey-api-key

# --- BigQuery settings ---
# Auth comes from ADC: run `gcloud auth application-default login` first.
GOOGLE_GENAI_USE_VERTEXAI=TRUE
GOOGLE_CLOUD_PROJECT=your-gcp-project-id
GOOGLE_CLOUD_LOCATION=us-xxxxx
# Optional: path to the toolbox binary (default ./toolbox)
# --- Databricks settings ---
DATABRICKS_HOST=https://xx-xxxxxx-xxxxxx.cloud.databricks.com
DATABRICKS_TOKEN=dapi-your-personal-access-token
# Genie-space-only:
GENIE_SPACE_ID=xxxxxxxxxxxxxxxxx
# SQL-endpoint-only - full Unity Catalog path to the churn table:
DATABRICKS_TABLE=main.analytics.customer_features

# --- Snowflake settings ---
SNOWFLAKE_ACCOUNT=MYORG-MYACCOUNT
SNOWFLAKE_USER=your-user
# A programmatic access token (PAT) works here as a drop-in password.
SNOWFLAKE_PASSWORD=your-password-or-pat
SNOWFLAKE_DATABASE=demo_db
SNOWFLAKE_SCHEMA=analytics
SNOWFLAKE_WAREHOUSE=DEMO_WH
# Use a read-only role (SELECT on customer_features only) - not ACCOUNTADMIN.
SNOWFLAKE_ROLE=ANALYST_READONLY

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

# Load environment variables from .env
load_dotenv()

# ADK warns once per experimental feature it touches (MCP support is one).
# Filter before ADK imports — the first warning fires at import time.
warnings.filterwarnings("ignore", message=r".*\[EXPERIMENTAL\].*")

from google.adk.agents import LlmAgent
from google.adk.models.lite_llm import LiteLlm
from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams

# Model (Claude via KodeKey)
MODEL = LiteLlm(
    "claude-4-5-20251001",
    api_base=os.environ["KODEKEY_BASE_URL"],
    api_key=os.environ["KODEKEY_API_KEY"],
)

# Databricks / Genie configuration
HOST = os.environ["DATABRICKS_HOST"]    # e.g. https://dbc-a1b2c3d4-e5f6.cloud.databricks.com
TOKEN = os.environ["DATABRICKS_TOKEN"]  # personal access token (dapi-...)
SPACE = os.environ["GENIE_SPACE_ID"]    # Genie space curated over customer_features

# No local MCP binary required: Databricks hosts the MCP server.
# We connect to the hosted MCP endpoint using streamable HTTP and a bearer token.
databricks_tools = McpToolset(
    connection_params=StreamableHTTPConnectionParams(
        url=f"{HOST}/api/2.0/mcp/genie/{SPACE}",
        headers={"Authorization": f"Bearer {TOKEN}"},
        timeout=60,
    )
)

# Initialize the LLM agent with a domain-specific system instruction.
agent = LlmAgent(
    name="churn_agent",
    model=MODEL,
    instruction=(
        "You help a customer-success team spot churn risk. The churn data lives "
        "in a Databricks Genie space — ask it plain-English questions to pull a "
        "customer's signals. Low usage or a sharp usage drop indicates "
        "disengagement; many support tickets indicates friction; short tenure is "
        "riskier. State the risk level (low/medium/high), cite the signals, and "
        "suggest one retention action."
    ),
    tools=[databricks_tools],
)

async def main():
    # Example placeholder for where you would run or interact with the agent.
    # The exact call depends on your ADK utilities (e.g., agent.run, agent.invoke, or a runner).
    # For this lesson, we keep the main function as the entry point that triggers the agent.
    # Replace with your project's preferred invocation pattern.
    ...
    
if __name__ == "__main__":
    asyncio.run(main())

$ python churn_agent_databricks.py

## Churn Risk Assessment: CUST-0003115

Risk Level: MEDIUM

Key Signals:
1. Very Short Tenure (36 days)
   - Customer onboarded ~2026-04-28
   - New customers within their first 60 days are statistically higher-risk for early churn

2. High Friction — 3 Support Tickets in 5 Weeks
   - 3 tickets in 36 days = high-touch customer experiencing issues
   - One billing ticket (high priority) resolved in 55 hours — slower than expected
   - One complaint ticket — suggests dissatisfaction
   - Pattern indicates friction in onboarding or setup

3. Recent Engagement / Resolution
   - All tickets resolved
   - Sentiment neutral on latest ticket

Suggested action:
- Send a personalized outreach email offering a dedicated onboarding session, expedited billing review, and a follow-up in two weeks to check progress.