> ## Documentation Index
> Fetch the complete documentation index at: https://notes.kodekloud.com/llms.txt
> Use this file to discover all available pages before exploring further.

# Demo Building Agent With Databricks

> Guide to building a Databricks-based churn detection LLM agent using MCP streamable HTTP, including environment setup, Python example, transport options, and sample assessment output

In this lesson we build a churn-detection LLM agent that queries customer signals stored in a Databricks Genie space. The implementation reuses the same LLM/agent logic as a BigQuery-based version — the only differences are the MCP transport (Databricks hosts an MCP endpoint) and the connection details (host URL + personal access token). The agent issues plain-English queries to the Genie space and reasons over streamed results to produce a churn-risk assessment.

Below you’ll find:

* a recommended `.env` layout (never commit `.env`),
* a consolidated, corrected Python example showing how to configure the Databricks MCP connection,
* a concise explanation of the two MCP transports (STDIO vs streamable HTTP),
* and an example of the script output for a sample customer assessment.

Never commit `.env` files. Keep credentials and tokens out of source control.

## .env example

Copy these entries into a local `.env` file and fill in the values appropriate for your environment. These variables are used by the sample script to configure the model provider (Claude via KodeKey in this demo), Google/BigQuery hints, Databricks Genie access, and an example Snowflake fallback.

```text theme={null}
# --- 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
```

Protect your Databricks personal access token (PAT) and any API keys. Rotate credentials regularly and grant the minimal permissions required (read-only where possible).

## Consolidated Python example

The following consolidated example highlights the essential pieces:

* environment loading,
* model initialization (Claude via KodeKey for this demo),
* and the Databricks MCP connection using streamable HTTP (no local MCP binary required).

Keep the code intact as your runnable starting point; tweak the agent invocation to match your project's runner or ADK utilities.

```python theme={null}
# 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())
```

Notes:

* The only Databricks-specific configuration is the `StreamableHTTPConnectionParams` URL (targets the hosted MCP endpoint for your Genie space) and the Authorization header with the Databricks PAT.
* Everything else — model initialization, agent instruction, and toolset usage — remains identical to a BigQuery-based agent implementation.
* For Databricks specifics and PAT guidance, see Databricks authentication docs: [https://docs.databricks.com/dev-tools/api/latest/authentication.html](https://docs.databricks.com/dev-tools/api/latest/authentication.html)

## How Databricks Genie space data maps to the agent

Within the Databricks Catalog, the Genie space exposes curated tables (or views) that the agent queries via MCP. In the churn demo the agent expects columns typically found in analytics and support tables, for example:

* customer identifiers (e.g., `customer_id`)
* usage metrics (e.g., `daily_active`, `30d_usage`, `usage_delta`)
* support ticket counts and ticket-level details (ticket id, category, priority, sentiment, resolved\_at)
* tenancy or signup timestamp (for tenure calculations)

The agent uses the customer ID as the primary lookup key and issues queries such as:

* "Show me the last 90 days of support tickets for CUST-0003115"
* "Return usage and delta for the last 30 days for CUST-0003115"

Because the Genie space exposes these tables through MCP, the MCP toolset streams row-level results back to the agent. The agent can then aggregate counts (e.g., ticket count, usage decline) or cite specific ticket rows (ID, priority, sentiment) when producing a churn assessment.

## Example script run & sample output

A typical run of the Databricks churn agent produces a concise assessment with risk level, supporting signals, and a recommended action. Below is a cleaned example of the script output for a sample customer:

```text theme={null}
$ 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.
```

The agent interprets aggregated signals (counts, deltas) and, when available, cites granular ticket-level details (ticket IDs, priorities, sentiments) streamed from Databricks tables.

## Key takeaway: MCP transports

MCP supports two primary transport modes. Choose the transport that matches how your MCP server is hosted or how your environment runs the MCP compatible tooling.

| Transport | Typical usage | What you provide in code |
| - | - | - |
| STDIO | When running an MCP-compatible binary locally (e.g., the MCP Toolbox). Local subprocess communicates via STDIO. | Start the local binary/subprocess and connect via an STDIO-based session manager (usually provided by the toolbox). |
| Streamable HTTP | When an MCP server is hosted remotely (Databricks Genie endpoint or another hosted MCP API). | Provide `StreamableHTTPConnectionParams` with `url`, `headers` (e.g., `Authorization: Bearer <token>`), and `timeout`. No local binary required. |

When the MCP server is hosted (Databricks or another provider exposing an MCP endpoint), prefer `StreamableHTTPConnectionParams` and pass the MCP URL and authentication headers. This approach simplifies deployment and removes the need to manage a local MCP binary.

## Links and references

* Databricks documentation: [https://docs.databricks.com/](https://docs.databricks.com/)
* Databricks API auth (PAT) docs: [https://docs.databricks.com/dev-tools/api/latest/authentication.html](https://docs.databricks.com/dev-tools/api/latest/authentication.html)
* Google Cloud SDK - Application Default Credentials: [https://cloud.google.com/sdk/gcloud/reference/auth/application-default/login](https://cloud.google.com/sdk/gcloud/reference/auth/application-default/login)

That’s it — by swapping only the MCP transport configuration and adding Databricks connection environment variables, you can reuse the same LLM agent logic used in BigQuery-based implementations against a Databricks-hosted Genie space.

> **[Watch Video](https://learn.kodekloud.com/user/courses/ai-agents-for-data-analytics/module/e7d9a0ef-37ce-482f-98c9-8e5332348342/lesson/198dad37-bbb4-4248-b5be-d23f95fa9f4a)**


This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.