From 16d616f1801b5718e0a9a6e8db7aa049843b0052 Mon Sep 17 00:00:00 2001 From: Andrew Fleischer Date: Fri, 10 Jul 2026 13:20:05 -0500 Subject: [PATCH] Add agent-adk: Google ADK conversational agent template New Python agent template built on Google's Agent Development Kit (ADK v2.x), mirroring the agent-langgraph / agent-openai-agents-sdk conventions: MLflow Responses API via @invoke()/@stream() handlers, the shared chat UI, MLflow tracing, and DAB deployment to Databricks Apps. Design: - Model access via LiteLLM (openai/ against {host}/serving-endpoints), so any Databricks / Foundation Model API endpoint works by name; per-request bearer token from the SDK credential chain. Default endpoint: databricks-claude-sonnet-4-5. - Stateless base: client-carried history is replayed into a per-request InMemorySessionService (seed_session_history); ADK Events are adapted to Responses stream events (process_adk_events). - Tracing via mlflow.litellm.autolog() (ADK routes model calls through LiteLLM), nested under the @invoke()/@stream() request trace. - MCP tools (UC functions, Genie, Vector Search, code interpreter) via ADK MCPToolset, commented out by default (needs google-adk[extensions]). - Pinned google-adk>=2.0.0,<3 (targets the 2.x API surface). process_adk_events unwraps ADK's {"result": } tool-response wrapping so tool outputs are a bare string, matching the other templates and the shared e2e validate_time contract. Repo wiring: - Register agent-adk (sdk="adk") in .scripts/templates.py; add e2e entry in template_config.py; new source skills add-tools-adk and modify-adk-agent; .gitignore skill allowlist; README and .claude/AGENTS.md conventions (autolog, MCP, session/memory, app.yaml). - Synced shared scripts and skills via sync-scripts.py / sync-skills.py. Co-authored-by: Isaac --- .claude/AGENTS.md | 13 +- .claude/skills/add-tools-adk/SKILL.md | 138 ++ .../skills/add-tools-adk/examples/app.yaml | 9 + .../examples/custom-mcp-server.md | 51 + .../add-tools-adk/examples/experiment.yaml | 8 + .../add-tools-adk/examples/genie-space.yaml | 9 + .../examples/lakebase-autoscaling.yaml | 21 + .../add-tools-adk/examples/lakebase.yaml | 18 + .../examples/serving-endpoint.yaml | 7 + .../add-tools-adk/examples/sql-warehouse.yaml | 7 + .../add-tools-adk/examples/uc-connection.yaml | 9 + .../add-tools-adk/examples/uc-function.yaml | 9 + .../add-tools-adk/examples/vector-search.yaml | 9 + .claude/skills/modify-adk-agent/SKILL.md | 194 ++ .gitignore | 2 + .../template_config.py | 1 + .scripts/templates.py | 5 + README.md | 1 + agent-adk/.claude/skills/add-tools/SKILL.md | 138 ++ .../skills/add-tools/examples/app.yaml | 9 + .../add-tools/examples/custom-mcp-server.md | 51 + .../skills/add-tools/examples/experiment.yaml | 8 + .../add-tools/examples/genie-space.yaml | 9 + .../examples/lakebase-autoscaling.yaml | 21 + .../skills/add-tools/examples/lakebase.yaml | 18 + .../add-tools/examples/serving-endpoint.yaml | 7 + .../add-tools/examples/sql-warehouse.yaml | 7 + .../add-tools/examples/uc-connection.yaml | 9 + .../add-tools/examples/uc-function.yaml | 9 + .../add-tools/examples/vector-search.yaml | 9 + .../.claude/skills/create-tools/SKILL.md | 26 + .../create-tools/examples/genie-space.md | 35 + .../examples/local-python-tools.md | 103 + .../create-tools/examples/uc-connection.md | 65 + .../create-tools/examples/uc-function.md | 67 + .../examples/vector-search-index.md | 78 + agent-adk/.claude/skills/deploy/SKILL.md | 248 +++ .../.claude/skills/discover-tools/SKILL.md | 49 + .../.claude/skills/lakebase-setup/SKILL.md | 467 +++++ .../.claude/skills/load-testing/SKILL.md | 319 +++ .../examples/mock_openai_client.py | 272 +++ .../skills/long-running-server/SKILL.md | 343 ++++ .../.claude/skills/managed-memory/SKILL.md | 463 +++++ .../migrate-from-model-serving/SKILL.md | 967 +++++++++ .../.claude/skills/modify-agent/SKILL.md | 194 ++ agent-adk/.claude/skills/quickstart/SKILL.md | 261 +++ agent-adk/.claude/skills/run-locally/SKILL.md | 90 + agent-adk/.env.example | 15 + agent-adk/.github/workflows/deploy.yml | 62 + agent-adk/.gitignore | 228 +++ agent-adk/AGENTS.md | 154 ++ agent-adk/CLAUDE.md | 3 + agent-adk/README.md | 314 +++ agent-adk/agent_server/__init__.py | 0 agent-adk/agent_server/agent.py | 130 ++ agent-adk/agent_server/evaluate_agent.py | 100 + agent-adk/agent_server/start_server.py | 20 + agent-adk/agent_server/utils.py | 200 ++ agent-adk/app.yaml | 16 + agent-adk/databricks.yml | 54 + agent-adk/manifest.yaml | 9 + agent-adk/pyproject.toml | 47 + agent-adk/scripts/__init__.py | 0 agent-adk/scripts/discover_tools.py | 432 ++++ .../scripts/grant_lakebase_permissions.py | 245 +++ agent-adk/scripts/preflight.py | 171 ++ agent-adk/scripts/quickstart.py | 1822 +++++++++++++++++ agent-adk/scripts/start_app.py | 333 +++ 68 files changed, 9206 insertions(+), 2 deletions(-) create mode 100644 .claude/skills/add-tools-adk/SKILL.md create mode 100644 .claude/skills/add-tools-adk/examples/app.yaml create mode 100644 .claude/skills/add-tools-adk/examples/custom-mcp-server.md create mode 100644 .claude/skills/add-tools-adk/examples/experiment.yaml create mode 100644 .claude/skills/add-tools-adk/examples/genie-space.yaml create mode 100644 .claude/skills/add-tools-adk/examples/lakebase-autoscaling.yaml create mode 100644 .claude/skills/add-tools-adk/examples/lakebase.yaml create mode 100644 .claude/skills/add-tools-adk/examples/serving-endpoint.yaml create mode 100644 .claude/skills/add-tools-adk/examples/sql-warehouse.yaml create mode 100644 .claude/skills/add-tools-adk/examples/uc-connection.yaml create mode 100644 .claude/skills/add-tools-adk/examples/uc-function.yaml create mode 100644 .claude/skills/add-tools-adk/examples/vector-search.yaml create mode 100644 .claude/skills/modify-adk-agent/SKILL.md create mode 100644 agent-adk/.claude/skills/add-tools/SKILL.md create mode 100644 agent-adk/.claude/skills/add-tools/examples/app.yaml create mode 100644 agent-adk/.claude/skills/add-tools/examples/custom-mcp-server.md create mode 100644 agent-adk/.claude/skills/add-tools/examples/experiment.yaml create mode 100644 agent-adk/.claude/skills/add-tools/examples/genie-space.yaml create mode 100644 agent-adk/.claude/skills/add-tools/examples/lakebase-autoscaling.yaml create mode 100644 agent-adk/.claude/skills/add-tools/examples/lakebase.yaml create mode 100644 agent-adk/.claude/skills/add-tools/examples/serving-endpoint.yaml create mode 100644 agent-adk/.claude/skills/add-tools/examples/sql-warehouse.yaml create mode 100644 agent-adk/.claude/skills/add-tools/examples/uc-connection.yaml create mode 100644 agent-adk/.claude/skills/add-tools/examples/uc-function.yaml create mode 100644 agent-adk/.claude/skills/add-tools/examples/vector-search.yaml create mode 100644 agent-adk/.claude/skills/create-tools/SKILL.md create mode 100644 agent-adk/.claude/skills/create-tools/examples/genie-space.md create mode 100644 agent-adk/.claude/skills/create-tools/examples/local-python-tools.md create mode 100644 agent-adk/.claude/skills/create-tools/examples/uc-connection.md create mode 100644 agent-adk/.claude/skills/create-tools/examples/uc-function.md create mode 100644 agent-adk/.claude/skills/create-tools/examples/vector-search-index.md create mode 100644 agent-adk/.claude/skills/deploy/SKILL.md create mode 100644 agent-adk/.claude/skills/discover-tools/SKILL.md create mode 100644 agent-adk/.claude/skills/lakebase-setup/SKILL.md create mode 100644 agent-adk/.claude/skills/load-testing/SKILL.md create mode 100644 agent-adk/.claude/skills/load-testing/examples/mock_openai_client.py create mode 100644 agent-adk/.claude/skills/long-running-server/SKILL.md create mode 100644 agent-adk/.claude/skills/managed-memory/SKILL.md create mode 100644 agent-adk/.claude/skills/migrate-from-model-serving/SKILL.md create mode 100644 agent-adk/.claude/skills/modify-agent/SKILL.md create mode 100644 agent-adk/.claude/skills/quickstart/SKILL.md create mode 100644 agent-adk/.claude/skills/run-locally/SKILL.md create mode 100644 agent-adk/.env.example create mode 100644 agent-adk/.github/workflows/deploy.yml create mode 100644 agent-adk/.gitignore create mode 100644 agent-adk/AGENTS.md create mode 100644 agent-adk/CLAUDE.md create mode 100644 agent-adk/README.md create mode 100644 agent-adk/agent_server/__init__.py create mode 100644 agent-adk/agent_server/agent.py create mode 100644 agent-adk/agent_server/evaluate_agent.py create mode 100644 agent-adk/agent_server/start_server.py create mode 100644 agent-adk/agent_server/utils.py create mode 100644 agent-adk/app.yaml create mode 100644 agent-adk/databricks.yml create mode 100644 agent-adk/manifest.yaml create mode 100644 agent-adk/pyproject.toml create mode 100644 agent-adk/scripts/__init__.py create mode 100644 agent-adk/scripts/discover_tools.py create mode 100644 agent-adk/scripts/grant_lakebase_permissions.py create mode 100644 agent-adk/scripts/preflight.py create mode 100644 agent-adk/scripts/quickstart.py create mode 100644 agent-adk/scripts/start_app.py diff --git a/.claude/AGENTS.md b/.claude/AGENTS.md index 824b91ef..a78bad36 100644 --- a/.claude/AGENTS.md +++ b/.claude/AGENTS.md @@ -33,12 +33,13 @@ async def invoke_handler(request: ResponsesAgentRequest) -> ResponsesAgentRespon async def stream_handler(request: ResponsesAgentRequest) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: ... ``` -LangGraph `invoke_handler` delegates to `stream_handler`. OpenAI SDK `invoke_handler` calls `Runner.run()` independently. +LangGraph `invoke_handler` delegates to `stream_handler`. OpenAI SDK `invoke_handler` calls `Runner.run()` independently. ADK `invoke_handler` delegates to `stream_handler` (collecting `response.output_item.done` items). ### MLflow autologging - LangGraph templates: `mlflow.langchain.autolog()` - OpenAI SDK templates: `mlflow.openai.autolog()` + `set_trace_processors([])` +- ADK templates: `mlflow.litellm.autolog()` (ADK calls the model through LiteLLM, so LiteLLM autolog captures the model spans) All handlers tag traces with: `mlflow.update_current_trace(metadata={"mlflow.trace.session": session_id})` @@ -46,6 +47,7 @@ All handlers tag traces with: `mlflow.update_current_trace(metadata={"mlflow.tra - **LangGraph**: `DatabricksMultiServerMCPClient` wrapping `DatabricksMCPServer` objects. Tools fetched once at agent init via `.get_tools()`. - **OpenAI SDK**: `McpServer` used as async context manager per-request: `async with await init_mcp_server() as mcp_server:` +- **ADK**: `MCPToolset` (with `StreamableHTTPConnectionParams`, `Authorization: Bearer` header) added to the agent's `tools` list. Requires the `mcp` extra (`google-adk[extensions]`). ### Session/memory patterns @@ -54,9 +56,16 @@ All handlers tag traces with: `mlflow.update_current_trace(metadata={"mlflow.tra | Short-term memory | LangGraph | `AsyncCheckpointSaver` | `thread_id` via `config["configurable"]` | | Long-term memory | LangGraph | `AsyncDatabricksStore` | `user_id` via `config["configurable"]` | | Short-term memory | OpenAI | `AsyncDatabricksSession` | `session_id` passed to `Runner.run(..., session=)` | +| Short-term memory | ADK | `DatabaseSessionService` (Lakebase) — *not shipped; base `agent-adk` is stateless* | `session_id` passed to `runner.run_async(..., session_id=)` | All memory templates return the ID in `custom_outputs` so clients can reuse it. +The base `agent-adk` template is stateless: it replays client-carried history into a fresh +`InMemorySessionService` each request via `agent_server.utils.seed_session_history`, and adapts ADK +`Event`s to `ResponsesAgentStreamEvent`s via `process_adk_events`. The model is reached through +`LiteLlm(model="openai/", api_base="{host}/serving-endpoints", api_key=)`; the +per-request bearer token comes from `get_bearer_token()` (the SDK credential chain). + ### `databricks.yml` conventions - `bundle.name` uses underscores: `agent_langgraph` @@ -68,7 +77,7 @@ All memory templates return the ID in `custom_outputs` so clients can reuse it. ### `app.yaml` files -The 3 base templates (`agent-langgraph`, `agent-openai-agents-sdk`, `agent-non-conversational`) have `app.yaml` files for UI-based template creation in the Databricks UI. These are separate from `databricks.yml` and use `valueFrom` (camelCase) for resource references. Memory/multiagent variants do not need separate `app.yaml` files. +The base conversational templates (`agent-langgraph`, `agent-adk`, `agent-openai-agents-sdk`) and `agent-non-conversational` have `app.yaml` files for UI-based template creation in the Databricks UI. These are separate from `databricks.yml` and use `valueFrom` (camelCase) for resource references. Multiagent variants do not need separate `app.yaml` files. ### Per-template AGENTS.md diff --git a/.claude/skills/add-tools-adk/SKILL.md b/.claude/skills/add-tools-adk/SKILL.md new file mode 100644 index 00000000..388ba782 --- /dev/null +++ b/.claude/skills/add-tools-adk/SKILL.md @@ -0,0 +1,138 @@ +--- +name: add-tools +description: "Add tools to your agent and grant required permissions in databricks.yml. Use when: (1) Adding MCP servers, Genie spaces, vector search, or UC functions to agent, (2) Permission errors at runtime, (3) User says 'add tool', 'connect to', 'grant permission', (4) Configuring databricks.yml resources." +--- + +# Add Tools & Grant Permissions + +> **Profile reminder:** All `databricks` CLI commands must include the profile from `.env`: `databricks --profile ` + +> Don't have the resource yet? See **create-tools** skill first. + +**After adding any MCP server to your agent, you MUST grant the app access in `databricks.yml`.** + +Without this, you'll get permission errors when the agent tries to use the resource. + +This template uses **Google ADK**. Databricks-hosted tools are connected with an ADK +`MCPToolset`, which needs the `mcp` package — run `uv add "google-adk[extensions]"` once. + +## Workflow + +**Step 1:** Add an MCP toolset in `agent_server/agent.py`: +```python +from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset +from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams +from agent_server.utils import get_bearer_token, get_databricks_host + +def create_agent(workspace_client): + host = get_databricks_host(workspace_client) + token = get_bearer_token(workspace_client) + genie = MCPToolset( + connection_params=StreamableHTTPConnectionParams( + url=f"{host}/api/2.0/mcp/genie/01234567-89ab-cdef", + headers={"Authorization": f"Bearer {token}"}, + ), + ) + return LlmAgent(name="agent", model=build_model(workspace_client), + instruction="You are a helpful assistant.", tools=[get_current_time, genie]) +``` + +Common Databricks MCP URL paths: UC functions `/api/2.0/mcp/functions/{catalog}/{schema}`, +Genie `/api/2.0/mcp/genie/{space_id}`, Vector Search `/api/2.0/mcp/vector-search/{catalog}/{schema}`, +code interpreter `/api/2.0/mcp/functions/system/ai`. + +**Step 2:** Grant access in `databricks.yml`: +```yaml +resources: + apps: + agent_adk: + resources: + - name: 'my_genie_space' + genie_space: + name: 'My Genie Space' + space_id: '01234567-89ab-cdef' + permission: 'CAN_RUN' +``` + +**Step 3:** Deploy and run: +```bash +databricks bundle deploy +databricks bundle run agent_adk # Required to start app with new code! +``` + +See **deploy** skill for more details. + +## Resource Type Examples + +See the `examples/` directory for complete YAML snippets: + +| File | Resource Type | When to Use | +|------|--------------|-------------| +| `uc-function.yaml` | Unity Catalog function | UC functions via MCP | +| `uc-connection.yaml` | UC connection | External MCP servers | +| `vector-search.yaml` | Vector search index | RAG applications | +| `sql-warehouse.yaml` | SQL warehouse | SQL execution | +| `serving-endpoint.yaml` | Model serving endpoint | Model inference | +| `genie-space.yaml` | Genie space | Natural language data | +| `lakebase.yaml` | Lakebase database | Session/memory storage (provisioned) | +| `lakebase-autoscaling.yaml` | Lakebase autoscaling postgres | Session/memory storage (autoscaling) | +| `experiment.yaml` | MLflow experiment | Tracing (already configured) | +| `app.yaml` | Databricks App (app-to-app) | Custom MCP servers hosted as Apps | +| `custom-mcp-server.md` | Custom MCP apps | Apps starting with `mcp-*` | + +## Custom MCP Servers (Databricks Apps) + +Declare the target app as an `app` resource in `databricks.yml` — the bundle grants `CAN_USE` on deploy. Requires Databricks CLI **v0.298.0+**. + +```yaml +resources: + apps: + agent_adk: + resources: + - name: 'mcp_server' + app: + name: 'mcp-my-server' + permission: CAN_USE +``` + +See `examples/custom-mcp-server.md` for the full flow (agent code + YAML + deploy). + +## value_from Pattern + +**IMPORTANT**: Make sure all `value_from` references in `databricks.yml` `config.env` reference an existing key in the `databricks.yml` `resources` list. +Some resources need environment variables in your app. Use `value_from` in `databricks.yml` `config.env` to reference resources defined in `databricks.yml`: + +```yaml +# In databricks.yml, under apps..config.env: +env: + - name: MLFLOW_EXPERIMENT_ID + value_from: "experiment" # References resources.apps..resources[name='experiment'] + - name: LAKEBASE_INSTANCE_NAME + value_from: "database" # References resources.apps..resources[name='database'] +``` + +**Critical:** Every `value_from` value must match a `name` field in `databricks.yml` resources. + +## MCP Error Handling + +MCP tool calls can fail (network issues, permission errors, timeouts). Give slow servers like +Genie a longer timeout, and wrap `runner.run_async(...)` in a try/except so one unavailable +server can't crash the request: + +```python +MCPToolset( + connection_params=StreamableHTTPConnectionParams( + url=f"{host}/api/2.0/mcp/genie/{space_id}", + headers={"Authorization": f"Bearer {token}"}, + timeout=60.0, # increase for slow tools like Genie + ), +) +``` + +## Important Notes + +- **MLflow experiment**: Already configured in template, no action needed +- **Multiple resources**: Add multiple entries under `resources:` list +- **Permission types vary**: Each resource type has specific permission values +- **Deploy + Run after changes**: Run both `databricks bundle deploy` AND `databricks bundle run {{BUNDLE_NAME}}` +- **value_from matching**: Ensure `config.env` `value_from` values match `databricks.yml` resource `name` values diff --git a/.claude/skills/add-tools-adk/examples/app.yaml b/.claude/skills/add-tools-adk/examples/app.yaml new file mode 100644 index 00000000..8d748c7b --- /dev/null +++ b/.claude/skills/add-tools-adk/examples/app.yaml @@ -0,0 +1,9 @@ +# Databricks App (for custom MCP servers hosted as Apps) +# Use for: Granting CAN_USE on another Databricks App (e.g., an mcp-* server app) +# Requires: CLI v0.298.0+ + +# In databricks.yml - add to resources.apps..resources: +- name: 'mcp_server' + app: + name: '' + permission: CAN_USE diff --git a/.claude/skills/add-tools-adk/examples/custom-mcp-server.md b/.claude/skills/add-tools-adk/examples/custom-mcp-server.md new file mode 100644 index 00000000..6fc58cb5 --- /dev/null +++ b/.claude/skills/add-tools-adk/examples/custom-mcp-server.md @@ -0,0 +1,51 @@ +# Custom MCP Server (Databricks App) + +Custom MCP servers are Databricks Apps with names starting with `mcp-*`. + +Declare the target app as an `app` resource in `databricks.yml` and the bundle will grant `CAN_USE` on deploy. Requires Databricks CLI **v0.298.0+**. + +## Steps + +### 1. Add MCP server in `agent_server/agent.py` + +```python +from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient + +custom_mcp = DatabricksMCPServer( + url="https://mcp-my-server.cloud.databricks.com/mcp", + name="my custom mcp server", +) + +mcp_client = DatabricksMultiServerMCPClient([custom_mcp]) +tools = await mcp_client.get_tools() +``` + +### 2. Grant access in `databricks.yml` + +Add the target app as a resource: + +```yaml +resources: + apps: + agent_langgraph: + resources: + - name: 'mcp_server' + app: + name: 'mcp-my-server' + permission: CAN_USE +``` + +### 3. Deploy + +```bash +databricks bundle deploy +databricks bundle run agent_langgraph +``` + +The bundle grants `CAN_USE` on the target app automatically — no manual permission steps needed. + +## Notes + +- Requires CLI v0.298.0+ (earlier versions will warn `unknown field: name` on `app.name`) +- The only supported permission is `CAN_USE` +- Subsequent `databricks bundle deploy` commands preserve the `app` resource diff --git a/.claude/skills/add-tools-adk/examples/experiment.yaml b/.claude/skills/add-tools-adk/examples/experiment.yaml new file mode 100644 index 00000000..ac5c626a --- /dev/null +++ b/.claude/skills/add-tools-adk/examples/experiment.yaml @@ -0,0 +1,8 @@ +# MLflow Experiment +# Use for: Tracing and model logging +# Note: Already configured in template's databricks.yml + +- name: 'my_experiment' + experiment: + experiment_id: '12349876' + permission: 'CAN_MANAGE' diff --git a/.claude/skills/add-tools-adk/examples/genie-space.yaml b/.claude/skills/add-tools-adk/examples/genie-space.yaml new file mode 100644 index 00000000..71589d52 --- /dev/null +++ b/.claude/skills/add-tools-adk/examples/genie-space.yaml @@ -0,0 +1,9 @@ +# Genie Space +# Use for: Natural language interface to data +# MCP URL: {host}/api/2.0/mcp/genie/{space_id} + +- name: 'my_genie_space' + genie_space: + name: 'My Genie Space' + space_id: '01234567-89ab-cdef' + permission: 'CAN_RUN' diff --git a/.claude/skills/add-tools-adk/examples/lakebase-autoscaling.yaml b/.claude/skills/add-tools-adk/examples/lakebase-autoscaling.yaml new file mode 100644 index 00000000..5d0e4314 --- /dev/null +++ b/.claude/skills/add-tools-adk/examples/lakebase-autoscaling.yaml @@ -0,0 +1,21 @@ +# Lakebase Autoscaling Postgres (for agent memory) +# Use for: Short-term or long-term memory storage with autoscaling Lakebase + +# In databricks.yml - add to resources.apps..resources: +- name: 'postgres' + postgres: + branch: "projects//branches/" + database: "projects//branches//databases/" + permission: CAN_CONNECT_AND_CREATE + +# In databricks.yml config block - add to env: +# - name: LAKEBASE_AUTOSCALING_PROJECT +# value: "" +# - name: LAKEBASE_AUTOSCALING_BRANCH +# value: "" + +# How to find the values: +# databricks api get /api/2.0/postgres/projects +# databricks api get /api/2.0/postgres/projects//branches +# databricks api get /api/2.0/postgres/projects//branches//databases +# The database-id is the internal ID (e.g., db-xxxx-xxxxxxxxxx), NOT "databricks_postgres" diff --git a/.claude/skills/add-tools-adk/examples/lakebase.yaml b/.claude/skills/add-tools-adk/examples/lakebase.yaml new file mode 100644 index 00000000..ce3b5d08 --- /dev/null +++ b/.claude/skills/add-tools-adk/examples/lakebase.yaml @@ -0,0 +1,18 @@ +# Lakebase Database (for agent memory) +# Use for: Long-term memory storage via AsyncDatabricksStore +# Requires: value_from reference in databricks.yml config block + +# In databricks.yml - add to resources.apps..resources: +- name: 'database' + database: + instance_name: '' + database_name: 'databricks_postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +# In databricks.yml config block - add to env: +# - name: LAKEBASE_INSTANCE_NAME +# value_from: "database" +# - name: EMBEDDING_ENDPOINT +# value: "databricks-gte-large-en" +# - name: EMBEDDING_DIMS +# value: "1024" diff --git a/.claude/skills/add-tools-adk/examples/serving-endpoint.yaml b/.claude/skills/add-tools-adk/examples/serving-endpoint.yaml new file mode 100644 index 00000000..b49ce9da --- /dev/null +++ b/.claude/skills/add-tools-adk/examples/serving-endpoint.yaml @@ -0,0 +1,7 @@ +# Model Serving Endpoint +# Use for: Model inference endpoints + +- name: 'my_endpoint' + serving_endpoint: + name: 'my_endpoint' + permission: 'CAN_QUERY' diff --git a/.claude/skills/add-tools-adk/examples/sql-warehouse.yaml b/.claude/skills/add-tools-adk/examples/sql-warehouse.yaml new file mode 100644 index 00000000..a6ce9446 --- /dev/null +++ b/.claude/skills/add-tools-adk/examples/sql-warehouse.yaml @@ -0,0 +1,7 @@ +# SQL Warehouse +# Use for: SQL query execution + +- name: 'my_warehouse' + sql_warehouse: + sql_warehouse_id: 'abc123def456' + permission: 'CAN_USE' diff --git a/.claude/skills/add-tools-adk/examples/uc-connection.yaml b/.claude/skills/add-tools-adk/examples/uc-connection.yaml new file mode 100644 index 00000000..316675fe --- /dev/null +++ b/.claude/skills/add-tools-adk/examples/uc-connection.yaml @@ -0,0 +1,9 @@ +# Unity Catalog Connection +# Use for: External MCP servers via UC connections +# MCP URL: {host}/api/2.0/mcp/external/{connection_name} + +- name: 'my_connection' + uc_securable: + securable_full_name: 'my-connection-name' + securable_type: 'CONNECTION' + permission: 'USE_CONNECTION' diff --git a/.claude/skills/add-tools-adk/examples/uc-function.yaml b/.claude/skills/add-tools-adk/examples/uc-function.yaml new file mode 100644 index 00000000..43f938a9 --- /dev/null +++ b/.claude/skills/add-tools-adk/examples/uc-function.yaml @@ -0,0 +1,9 @@ +# Unity Catalog Function +# Use for: UC functions accessed via MCP server +# MCP URL: {host}/api/2.0/mcp/functions/{catalog}/{schema}/{function_name} + +- name: 'my_uc_function' + uc_securable: + securable_full_name: 'catalog.schema.function_name' + securable_type: 'FUNCTION' + permission: 'EXECUTE' diff --git a/.claude/skills/add-tools-adk/examples/vector-search.yaml b/.claude/skills/add-tools-adk/examples/vector-search.yaml new file mode 100644 index 00000000..0ba39027 --- /dev/null +++ b/.claude/skills/add-tools-adk/examples/vector-search.yaml @@ -0,0 +1,9 @@ +# Vector Search Index +# Use for: RAG applications with unstructured data +# MCP URL: {host}/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name} + +- name: 'my_vector_index' + uc_securable: + securable_full_name: 'catalog.schema.index_name' + securable_type: 'TABLE' + permission: 'SELECT' diff --git a/.claude/skills/modify-adk-agent/SKILL.md b/.claude/skills/modify-adk-agent/SKILL.md new file mode 100644 index 00000000..4c9de98a --- /dev/null +++ b/.claude/skills/modify-adk-agent/SKILL.md @@ -0,0 +1,194 @@ +--- +name: modify-agent +description: "Modify agent code, add tools, or change configuration. Use when: (1) User says 'modify agent', 'add tool', 'change model', or 'edit agent.py', (2) Adding MCP servers to agent, (3) Changing agent instructions, (4) Understanding SDK patterns." +--- + +# Modify the Agent + +This template's agent is built with **Google's Agent Development Kit (ADK)**. ADK talks to +Databricks Model Serving through **LiteLLM** (`google.adk.models.lite_llm.LiteLlm`), and the +MLflow `AgentServer` exposes it over the Responses API. + +## Main File + +**`agent_server/agent.py`** — Agent logic, model selection, instructions, tools + +## Key Files + +| File | Purpose | +| -------------------------------- | --------------------------------------------------- | +| `agent_server/agent.py` | Agent logic, model, instructions, tools | +| `agent_server/start_server.py` | FastAPI server + MLflow setup | +| `agent_server/evaluate_agent.py` | Agent evaluation with MLflow scorers | +| `agent_server/utils.py` | Auth helpers, session seeding, ADK→Responses adapter| +| `databricks.yml` | Bundle config & resource permissions | + +## SDK Setup + +```python +import mlflow +from databricks.sdk import WorkspaceClient +from google.adk.agents import LlmAgent +from google.adk.models.lite_llm import LiteLlm + +# ADK drives the model via LiteLLM, so LiteLLM autolog captures every model call as a trace span. +mlflow.litellm.autolog() + +workspace_client = WorkspaceClient() +``` + +Before making changes, confirm APIs exist in the installed `google-adk` package (look in the +venv's `site-packages/google/adk`, or run `uv sync` first). ADK's API changed between the 1.x +and 2.x lines; this template targets 2.x. + +--- + +## Connecting to a Databricks model serving endpoint + +`LiteLlm` with the `openai/` provider prefix points ADK at any OpenAI-compatible endpoint. +Databricks Model Serving exposes exactly such a surface at `{host}/serving-endpoints`. + +```python +from google.adk.models.lite_llm import LiteLlm +from agent_server.utils import get_bearer_token, get_databricks_host + +def build_model(workspace_client): + host = get_databricks_host(workspace_client) + token = get_bearer_token(workspace_client) # works for CLI, PAT, and App OAuth creds + return LiteLlm( + model=f"openai/databricks-claude-sonnet-4-5", # openai/ + api_base=f"{host}/serving-endpoints", + api_key=token, + temperature=0.1, # any extra kwargs pass through to LiteLLM + ) +``` + +Change `databricks-claude-sonnet-4-5` to any endpoint your app can query (e.g. +`databricks-gpt-5-2`, `databricks-meta-llama-3-3-70b-instruct`). Some workspaces require +granting the app access to the serving endpoint in `databricks.yml` — see the **add-tools** +skill and `examples/serving-endpoint.yaml`. + +--- + +## Defining the agent and its instructions + +Unlike LangGraph's `create_agent`, ADK's `LlmAgent` takes the system prompt directly via +`instruction`: + +```python +from google.adk.agents import LlmAgent + +AGENT_INSTRUCTIONS = """You are a helpful data analyst assistant. + +You have access to: +- Company sales data via Genie +- Product documentation via vector search + +Always cite your sources when answering questions.""" + +def create_agent(workspace_client): + return LlmAgent( + name="agent", # must be a valid identifier + model=build_model(workspace_client), + instruction=AGENT_INSTRUCTIONS, + description="A helpful assistant.", + tools=[get_current_time], + ) +``` + +--- + +## Function tools + +ADK auto-wraps a plain Python function (with type hints + a docstring) into a tool — no +decorator needed. The docstring is shown to the model, so make it descriptive. + +```python +def get_current_time() -> str: + """Get the current date and time.""" + from datetime import datetime + return datetime.now().isoformat() + +# then: tools=[get_current_time, your_other_tool] +``` + +--- + +## MCP tools (UC functions, Genie, Vector Search, code interpreter) + +Databricks-hosted tools are exposed over MCP. ADK connects to an MCP server with an +`MCPToolset`. This needs the `mcp` package — run `uv add "google-adk[extensions]"`. + +```python +from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset +from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams +from agent_server.utils import get_bearer_token, get_databricks_host + +def init_mcp_toolset(workspace_client): + host = get_databricks_host(workspace_client) + token = get_bearer_token(workspace_client) + return MCPToolset( + connection_params=StreamableHTTPConnectionParams( + url=f"{host}/api/2.0/mcp/functions/system/ai", # UC functions in system.ai + headers={"Authorization": f"Bearer {token}"}, + ), + ) +# then: tools.append(init_mcp_toolset(workspace_client)) +``` + +Common Databricks MCP URLs (swap the path): + +| Tool | URL path | +|------|----------| +| UC functions (schema) | `/api/2.0/mcp/functions/{catalog}/{schema}` | +| Genie space | `/api/2.0/mcp/genie/{space_id}` | +| Vector Search (schema) | `/api/2.0/mcp/vector-search/{catalog}/{schema}` | +| Built-in code interpreter | `/api/2.0/mcp/functions/system/ai` (`system.ai.python_exec`) | + +**After adding MCP servers:** grant permissions in `databricks.yml` (see **add-tools** skill). + +**Robustness note:** the base template runs the ADK agent per request. If you add an +`MCPToolset`, an unreachable/unauthorized MCP server can raise when ADK lists its tools inside +`runner.run_async`. Wrap the run in a `try/except` (log and continue) so one bad server doesn't +crash the request, and give slow servers (e.g. Genie) a longer `timeout`. Unlike the OpenAI +Agents SDK template, ADK has no built-in per-server health check — add one if you connect +multiple MCP servers. + +--- + +## How requests flow (Responses API contract) + +You normally don't edit this, but it helps to understand it: + +- `stream_handler` (decorated `@stream()`) replays the client-carried conversation into a fresh + ADK session via `seed_session_history()`, runs `runner.run_async(..., RunConfig(streaming_mode=SSE))`, + and adapts ADK `Event`s into `ResponsesAgentStreamEvent`s with `process_adk_events()`. +- `invoke_handler` (decorated `@invoke()`) collects the `response.output_item.done` items from + the same stream to build the non-streaming `ResponsesAgentResponse`. + +To change the streamed shape, edit `process_adk_events()` in `agent_server/utils.py`. + +--- + +## Session / memory + +This template is stateless — the client sends the full history each turn. For durable +cross-session memory, back an ADK `DatabaseSessionService` with Lakebase Postgres +(see **lakebase-setup**), or use governed **managed-memory** tools. + +--- + +## External Resources + +1. [Google ADK documentation](https://google.github.io/adk-docs/) +2. [ADK + LiteLLM models](https://google.github.io/adk-docs/agents/models/) +3. [Agent Framework docs](https://docs.databricks.com/aws/en/generative-ai/agent-framework/) +4. [Adding tools](https://docs.databricks.com/aws/en/generative-ai/agent-framework/agent-tool) +5. [Responses API](https://mlflow.org/docs/latest/genai/serving/responses-agent/) + +## Next Steps + +- Discover available tools: see **discover-tools** skill +- Grant resource permissions: see **add-tools** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/.gitignore b/.gitignore index 24a4a7f4..1b5928ac 100644 --- a/.gitignore +++ b/.gitignore @@ -183,8 +183,10 @@ mlflow.db !.claude/skills/deploy/ !.claude/skills/add-tools-langgraph/ !.claude/skills/add-tools-openai/ +!.claude/skills/add-tools-adk/ !.claude/skills/modify-langgraph-agent/ !.claude/skills/modify-openai-agent/ +!.claude/skills/modify-adk-agent/ !.claude/skills/lakebase-setup/ !.claude/skills/agent-langgraph-memory/ !.claude/skills/agent-openai-memory/ diff --git a/.scripts/agent-integration-tests/template_config.py b/.scripts/agent-integration-tests/template_config.py index 5fad9651..7cc2bed3 100644 --- a/.scripts/agent-integration-tests/template_config.py +++ b/.scripts/agent-integration-tests/template_config.py @@ -181,6 +181,7 @@ def build_templates( # (name, needs_lakebase, overrides) configs: list[tuple[str, bool, dict]] = [ ("agent-langgraph", False, {}), + ("agent-adk", False, {}), ("agent-langgraph-advanced", True, {"is_advanced": True}), ("agent-openai-agents-sdk", False, {}), ("agent-openai-advanced", True, {"is_advanced": True}), diff --git a/.scripts/templates.py b/.scripts/templates.py index 91bb4d32..44415fc5 100644 --- a/.scripts/templates.py +++ b/.scripts/templates.py @@ -12,6 +12,11 @@ "has_memory": True, "has_actions": True, }, + "agent-adk": { + "sdk": "adk", + "bundle_name": "agent_adk", + "has_actions": True, + }, "agent-openai-agents-sdk": { "sdk": "openai", "bundle_name": "agent_openai_agents_sdk", diff --git a/README.md b/README.md index 8edada06..ec32c0d6 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,7 @@ See [Create an App from a Template](https://docs.databricks.com/aws/en/dev-tools | Template | Description | Dependencies | |----------|-------------|--------------| | `agent-langgraph` | A conversational agent using LangGraph and MLflow AgentServer | MLflow experiment | +| `agent-adk` | A conversational agent using Google's Agent Development Kit (ADK) and MLflow AgentServer | MLflow experiment | | `agent-langgraph-advanced` | LangGraph agent with short-term memory, long-term memory, and long-running background tasks | MLflow experiment, Database | | `agent-openai-agents-sdk` | A conversational agent using OpenAI Agents SDK and MLflow AgentServer | MLflow experiment | | `agent-openai-advanced` | OpenAI Agents SDK agent with short-term memory and long-running background tasks | MLflow experiment, Database | diff --git a/agent-adk/.claude/skills/add-tools/SKILL.md b/agent-adk/.claude/skills/add-tools/SKILL.md new file mode 100644 index 00000000..d6a37859 --- /dev/null +++ b/agent-adk/.claude/skills/add-tools/SKILL.md @@ -0,0 +1,138 @@ +--- +name: add-tools +description: "Add tools to your agent and grant required permissions in databricks.yml. Use when: (1) Adding MCP servers, Genie spaces, vector search, or UC functions to agent, (2) Permission errors at runtime, (3) User says 'add tool', 'connect to', 'grant permission', (4) Configuring databricks.yml resources." +--- + +# Add Tools & Grant Permissions + +> **Profile reminder:** All `databricks` CLI commands must include the profile from `.env`: `databricks --profile ` + +> Don't have the resource yet? See **create-tools** skill first. + +**After adding any MCP server to your agent, you MUST grant the app access in `databricks.yml`.** + +Without this, you'll get permission errors when the agent tries to use the resource. + +This template uses **Google ADK**. Databricks-hosted tools are connected with an ADK +`MCPToolset`, which needs the `mcp` package — run `uv add "google-adk[extensions]"` once. + +## Workflow + +**Step 1:** Add an MCP toolset in `agent_server/agent.py`: +```python +from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset +from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams +from agent_server.utils import get_bearer_token, get_databricks_host + +def create_agent(workspace_client): + host = get_databricks_host(workspace_client) + token = get_bearer_token(workspace_client) + genie = MCPToolset( + connection_params=StreamableHTTPConnectionParams( + url=f"{host}/api/2.0/mcp/genie/01234567-89ab-cdef", + headers={"Authorization": f"Bearer {token}"}, + ), + ) + return LlmAgent(name="agent", model=build_model(workspace_client), + instruction="You are a helpful assistant.", tools=[get_current_time, genie]) +``` + +Common Databricks MCP URL paths: UC functions `/api/2.0/mcp/functions/{catalog}/{schema}`, +Genie `/api/2.0/mcp/genie/{space_id}`, Vector Search `/api/2.0/mcp/vector-search/{catalog}/{schema}`, +code interpreter `/api/2.0/mcp/functions/system/ai`. + +**Step 2:** Grant access in `databricks.yml`: +```yaml +resources: + apps: + agent_adk: + resources: + - name: 'my_genie_space' + genie_space: + name: 'My Genie Space' + space_id: '01234567-89ab-cdef' + permission: 'CAN_RUN' +``` + +**Step 3:** Deploy and run: +```bash +databricks bundle deploy +databricks bundle run agent_adk # Required to start app with new code! +``` + +See **deploy** skill for more details. + +## Resource Type Examples + +See the `examples/` directory for complete YAML snippets: + +| File | Resource Type | When to Use | +|------|--------------|-------------| +| `uc-function.yaml` | Unity Catalog function | UC functions via MCP | +| `uc-connection.yaml` | UC connection | External MCP servers | +| `vector-search.yaml` | Vector search index | RAG applications | +| `sql-warehouse.yaml` | SQL warehouse | SQL execution | +| `serving-endpoint.yaml` | Model serving endpoint | Model inference | +| `genie-space.yaml` | Genie space | Natural language data | +| `lakebase.yaml` | Lakebase database | Session/memory storage (provisioned) | +| `lakebase-autoscaling.yaml` | Lakebase autoscaling postgres | Session/memory storage (autoscaling) | +| `experiment.yaml` | MLflow experiment | Tracing (already configured) | +| `app.yaml` | Databricks App (app-to-app) | Custom MCP servers hosted as Apps | +| `custom-mcp-server.md` | Custom MCP apps | Apps starting with `mcp-*` | + +## Custom MCP Servers (Databricks Apps) + +Declare the target app as an `app` resource in `databricks.yml` — the bundle grants `CAN_USE` on deploy. Requires Databricks CLI **v0.298.0+**. + +```yaml +resources: + apps: + agent_adk: + resources: + - name: 'mcp_server' + app: + name: 'mcp-my-server' + permission: CAN_USE +``` + +See `examples/custom-mcp-server.md` for the full flow (agent code + YAML + deploy). + +## value_from Pattern + +**IMPORTANT**: Make sure all `value_from` references in `databricks.yml` `config.env` reference an existing key in the `databricks.yml` `resources` list. +Some resources need environment variables in your app. Use `value_from` in `databricks.yml` `config.env` to reference resources defined in `databricks.yml`: + +```yaml +# In databricks.yml, under apps..config.env: +env: + - name: MLFLOW_EXPERIMENT_ID + value_from: "experiment" # References resources.apps..resources[name='experiment'] + - name: LAKEBASE_INSTANCE_NAME + value_from: "database" # References resources.apps..resources[name='database'] +``` + +**Critical:** Every `value_from` value must match a `name` field in `databricks.yml` resources. + +## MCP Error Handling + +MCP tool calls can fail (network issues, permission errors, timeouts). Give slow servers like +Genie a longer timeout, and wrap `runner.run_async(...)` in a try/except so one unavailable +server can't crash the request: + +```python +MCPToolset( + connection_params=StreamableHTTPConnectionParams( + url=f"{host}/api/2.0/mcp/genie/{space_id}", + headers={"Authorization": f"Bearer {token}"}, + timeout=60.0, # increase for slow tools like Genie + ), +) +``` + +## Important Notes + +- **MLflow experiment**: Already configured in template, no action needed +- **Multiple resources**: Add multiple entries under `resources:` list +- **Permission types vary**: Each resource type has specific permission values +- **Deploy + Run after changes**: Run both `databricks bundle deploy` AND `databricks bundle run agent_adk` +- **value_from matching**: Ensure `config.env` `value_from` values match `databricks.yml` resource `name` values diff --git a/agent-adk/.claude/skills/add-tools/examples/app.yaml b/agent-adk/.claude/skills/add-tools/examples/app.yaml new file mode 100644 index 00000000..8d748c7b --- /dev/null +++ b/agent-adk/.claude/skills/add-tools/examples/app.yaml @@ -0,0 +1,9 @@ +# Databricks App (for custom MCP servers hosted as Apps) +# Use for: Granting CAN_USE on another Databricks App (e.g., an mcp-* server app) +# Requires: CLI v0.298.0+ + +# In databricks.yml - add to resources.apps..resources: +- name: 'mcp_server' + app: + name: '' + permission: CAN_USE diff --git a/agent-adk/.claude/skills/add-tools/examples/custom-mcp-server.md b/agent-adk/.claude/skills/add-tools/examples/custom-mcp-server.md new file mode 100644 index 00000000..6fc58cb5 --- /dev/null +++ b/agent-adk/.claude/skills/add-tools/examples/custom-mcp-server.md @@ -0,0 +1,51 @@ +# Custom MCP Server (Databricks App) + +Custom MCP servers are Databricks Apps with names starting with `mcp-*`. + +Declare the target app as an `app` resource in `databricks.yml` and the bundle will grant `CAN_USE` on deploy. Requires Databricks CLI **v0.298.0+**. + +## Steps + +### 1. Add MCP server in `agent_server/agent.py` + +```python +from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient + +custom_mcp = DatabricksMCPServer( + url="https://mcp-my-server.cloud.databricks.com/mcp", + name="my custom mcp server", +) + +mcp_client = DatabricksMultiServerMCPClient([custom_mcp]) +tools = await mcp_client.get_tools() +``` + +### 2. Grant access in `databricks.yml` + +Add the target app as a resource: + +```yaml +resources: + apps: + agent_langgraph: + resources: + - name: 'mcp_server' + app: + name: 'mcp-my-server' + permission: CAN_USE +``` + +### 3. Deploy + +```bash +databricks bundle deploy +databricks bundle run agent_langgraph +``` + +The bundle grants `CAN_USE` on the target app automatically — no manual permission steps needed. + +## Notes + +- Requires CLI v0.298.0+ (earlier versions will warn `unknown field: name` on `app.name`) +- The only supported permission is `CAN_USE` +- Subsequent `databricks bundle deploy` commands preserve the `app` resource diff --git a/agent-adk/.claude/skills/add-tools/examples/experiment.yaml b/agent-adk/.claude/skills/add-tools/examples/experiment.yaml new file mode 100644 index 00000000..ac5c626a --- /dev/null +++ b/agent-adk/.claude/skills/add-tools/examples/experiment.yaml @@ -0,0 +1,8 @@ +# MLflow Experiment +# Use for: Tracing and model logging +# Note: Already configured in template's databricks.yml + +- name: 'my_experiment' + experiment: + experiment_id: '12349876' + permission: 'CAN_MANAGE' diff --git a/agent-adk/.claude/skills/add-tools/examples/genie-space.yaml b/agent-adk/.claude/skills/add-tools/examples/genie-space.yaml new file mode 100644 index 00000000..71589d52 --- /dev/null +++ b/agent-adk/.claude/skills/add-tools/examples/genie-space.yaml @@ -0,0 +1,9 @@ +# Genie Space +# Use for: Natural language interface to data +# MCP URL: {host}/api/2.0/mcp/genie/{space_id} + +- name: 'my_genie_space' + genie_space: + name: 'My Genie Space' + space_id: '01234567-89ab-cdef' + permission: 'CAN_RUN' diff --git a/agent-adk/.claude/skills/add-tools/examples/lakebase-autoscaling.yaml b/agent-adk/.claude/skills/add-tools/examples/lakebase-autoscaling.yaml new file mode 100644 index 00000000..5d0e4314 --- /dev/null +++ b/agent-adk/.claude/skills/add-tools/examples/lakebase-autoscaling.yaml @@ -0,0 +1,21 @@ +# Lakebase Autoscaling Postgres (for agent memory) +# Use for: Short-term or long-term memory storage with autoscaling Lakebase + +# In databricks.yml - add to resources.apps..resources: +- name: 'postgres' + postgres: + branch: "projects//branches/" + database: "projects//branches//databases/" + permission: CAN_CONNECT_AND_CREATE + +# In databricks.yml config block - add to env: +# - name: LAKEBASE_AUTOSCALING_PROJECT +# value: "" +# - name: LAKEBASE_AUTOSCALING_BRANCH +# value: "" + +# How to find the values: +# databricks api get /api/2.0/postgres/projects +# databricks api get /api/2.0/postgres/projects//branches +# databricks api get /api/2.0/postgres/projects//branches//databases +# The database-id is the internal ID (e.g., db-xxxx-xxxxxxxxxx), NOT "databricks_postgres" diff --git a/agent-adk/.claude/skills/add-tools/examples/lakebase.yaml b/agent-adk/.claude/skills/add-tools/examples/lakebase.yaml new file mode 100644 index 00000000..ce3b5d08 --- /dev/null +++ b/agent-adk/.claude/skills/add-tools/examples/lakebase.yaml @@ -0,0 +1,18 @@ +# Lakebase Database (for agent memory) +# Use for: Long-term memory storage via AsyncDatabricksStore +# Requires: value_from reference in databricks.yml config block + +# In databricks.yml - add to resources.apps..resources: +- name: 'database' + database: + instance_name: '' + database_name: 'databricks_postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +# In databricks.yml config block - add to env: +# - name: LAKEBASE_INSTANCE_NAME +# value_from: "database" +# - name: EMBEDDING_ENDPOINT +# value: "databricks-gte-large-en" +# - name: EMBEDDING_DIMS +# value: "1024" diff --git a/agent-adk/.claude/skills/add-tools/examples/serving-endpoint.yaml b/agent-adk/.claude/skills/add-tools/examples/serving-endpoint.yaml new file mode 100644 index 00000000..b49ce9da --- /dev/null +++ b/agent-adk/.claude/skills/add-tools/examples/serving-endpoint.yaml @@ -0,0 +1,7 @@ +# Model Serving Endpoint +# Use for: Model inference endpoints + +- name: 'my_endpoint' + serving_endpoint: + name: 'my_endpoint' + permission: 'CAN_QUERY' diff --git a/agent-adk/.claude/skills/add-tools/examples/sql-warehouse.yaml b/agent-adk/.claude/skills/add-tools/examples/sql-warehouse.yaml new file mode 100644 index 00000000..a6ce9446 --- /dev/null +++ b/agent-adk/.claude/skills/add-tools/examples/sql-warehouse.yaml @@ -0,0 +1,7 @@ +# SQL Warehouse +# Use for: SQL query execution + +- name: 'my_warehouse' + sql_warehouse: + sql_warehouse_id: 'abc123def456' + permission: 'CAN_USE' diff --git a/agent-adk/.claude/skills/add-tools/examples/uc-connection.yaml b/agent-adk/.claude/skills/add-tools/examples/uc-connection.yaml new file mode 100644 index 00000000..316675fe --- /dev/null +++ b/agent-adk/.claude/skills/add-tools/examples/uc-connection.yaml @@ -0,0 +1,9 @@ +# Unity Catalog Connection +# Use for: External MCP servers via UC connections +# MCP URL: {host}/api/2.0/mcp/external/{connection_name} + +- name: 'my_connection' + uc_securable: + securable_full_name: 'my-connection-name' + securable_type: 'CONNECTION' + permission: 'USE_CONNECTION' diff --git a/agent-adk/.claude/skills/add-tools/examples/uc-function.yaml b/agent-adk/.claude/skills/add-tools/examples/uc-function.yaml new file mode 100644 index 00000000..43f938a9 --- /dev/null +++ b/agent-adk/.claude/skills/add-tools/examples/uc-function.yaml @@ -0,0 +1,9 @@ +# Unity Catalog Function +# Use for: UC functions accessed via MCP server +# MCP URL: {host}/api/2.0/mcp/functions/{catalog}/{schema}/{function_name} + +- name: 'my_uc_function' + uc_securable: + securable_full_name: 'catalog.schema.function_name' + securable_type: 'FUNCTION' + permission: 'EXECUTE' diff --git a/agent-adk/.claude/skills/add-tools/examples/vector-search.yaml b/agent-adk/.claude/skills/add-tools/examples/vector-search.yaml new file mode 100644 index 00000000..0ba39027 --- /dev/null +++ b/agent-adk/.claude/skills/add-tools/examples/vector-search.yaml @@ -0,0 +1,9 @@ +# Vector Search Index +# Use for: RAG applications with unstructured data +# MCP URL: {host}/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name} + +- name: 'my_vector_index' + uc_securable: + securable_full_name: 'catalog.schema.index_name' + securable_type: 'TABLE' + permission: 'SELECT' diff --git a/agent-adk/.claude/skills/create-tools/SKILL.md b/agent-adk/.claude/skills/create-tools/SKILL.md new file mode 100644 index 00000000..4f2606ed --- /dev/null +++ b/agent-adk/.claude/skills/create-tools/SKILL.md @@ -0,0 +1,26 @@ +--- +name: create-tools +description: "Create Databricks resources that agents connect to as tools. Use when: (1) User needs to create a Genie space, vector search index, UC function, or UC connection, (2) User says 'create tool', 'set up genie', 'create vector search', 'register MCP server', (3) Before add-tools when the resource doesn't exist yet, (4) User asks 'what do I need to create before adding this tool'." +--- + +# Create Tool Resources + +> This skill covers creating the Databricks resources your agent connects to. +> After creating a resource, use the **add-tools** skill to wire it into your agent and grant permissions. + +## Which resource do you need? + +| I want my agent to... | Resource to create | Guide | +|---|---|---| +| Answer questions about structured data | Genie space | `examples/genie-space.md` | +| Search documents / RAG | Vector Search index | `examples/vector-search-index.md` | +| Call custom SQL/Python logic | UC function | `examples/uc-function.md` | +| Connect to an external MCP server | UC connection | `examples/uc-connection.md` | +| Add inline Python tools | Local function tools | `examples/local-python-tools.md` | + +## Workflow + +1. **Discover** existing resources: `uv run discover-tools` (see **discover-tools** skill) +2. **Create** the resource if it doesn't exist (this skill) +3. **Add** the MCP server to your agent code + grant permissions (see **add-tools** skill) +4. **Deploy** (see **deploy** skill) diff --git a/agent-adk/.claude/skills/create-tools/examples/genie-space.md b/agent-adk/.claude/skills/create-tools/examples/genie-space.md new file mode 100644 index 00000000..c29cf121 --- /dev/null +++ b/agent-adk/.claude/skills/create-tools/examples/genie-space.md @@ -0,0 +1,35 @@ +# Create a Genie Space + +Genie spaces let agents query structured data in Unity Catalog tables using natural language. A Genie space can include up to 30 tables or views. + +## Create via Databricks UI + +1. In your workspace, go to **Genie** in the left sidebar. +2. Click **New** to create a new Genie space. +3. Add the Unity Catalog tables or views your agent needs to query. +4. Configure instructions to guide how Genie interprets queries (optional but recommended). +5. Configure a default SQL warehouse: go to **Configure** > **Settings** > **Default warehouse**. +6. Share the space with the app's service principal: + - Click **Share** in the top right + - Enter the service principal name, click **Add**, and set the permission level to **CAN RUN** + - To find your app's service principal: `databricks apps get --output json --profile | jq -r '.service_principal_name'` + +## Find the space ID + +The space ID is in the URL when viewing the Genie space: + +``` +https://.databricks.com/genie/rooms/?o=... +``` + +To list all Genie spaces via CLI: + +```bash +databricks genie list-spaces --profile +``` + +## Next step + +Wire the Genie space into your agent and grant permissions. See the **add-tools** skill and use `examples/genie-space.yaml` for the `databricks.yml` resource grant. + +MCP URL: `{host}/api/2.0/mcp/genie/{space_id}` (OAuth scope for on-behalf-of-user auth: `genie`) diff --git a/agent-adk/.claude/skills/create-tools/examples/local-python-tools.md b/agent-adk/.claude/skills/create-tools/examples/local-python-tools.md new file mode 100644 index 00000000..c04598de --- /dev/null +++ b/agent-adk/.claude/skills/create-tools/examples/local-python-tools.md @@ -0,0 +1,103 @@ +# Local Python Function Tools + +For operations that don't need external data sources or MCP servers, define tools directly in your agent code. These run in the same process as your agent — no resource creation or `databricks.yml` permissions needed. + +## When to use local tools vs. MCP + +- **Local tools**: Simple logic, API calls with custom auth, data transformations, utility functions +- **MCP tools**: When you need Databricks-managed auth, UC governance, or access to Databricks resources (tables, indexes, Genie) + +## OpenAI Agents SDK + +```python +from agents import Agent, function_tool + +@function_tool +def get_current_time() -> str: + """Get the current date and time in ISO format.""" + from datetime import datetime + return datetime.now().isoformat() + +@function_tool +def calculate_discount(price: float, percent: float) -> str: + """Calculate a discounted price. Returns the new price after applying the discount.""" + discounted = price * (1 - percent / 100) + return f"${discounted:.2f}" + +agent = Agent( + name="My agent", + instructions="You are a helpful assistant.", + model="databricks-claude-sonnet-4-5", + tools=[get_current_time, calculate_discount], +) +``` + +## LangGraph + +```python +from langchain_core.tools import tool + +@tool +def get_current_time() -> str: + """Get the current date and time in ISO format.""" + from datetime import datetime + return datetime.now().isoformat() + +@tool +def calculate_discount(price: float, percent: float) -> str: + """Calculate a discounted price. Returns the new price after applying the discount.""" + discounted = price * (1 - percent / 100) + return f"${discounted:.2f}" + +# Pass to create_react_agent or add to your tools list +tools = [get_current_time, calculate_discount] +``` + +## Error handling + +Both SDKs handle tool errors gracefully by default — the error message is returned to the LLM so it can retry or respond to the user. For custom error messages, use the patterns below. + +### OpenAI Agents SDK + +`@function_tool` includes a built-in `default_tool_error_function` that catches exceptions and returns `"An error occurred while running the tool. Error: {error}"` to the LLM. To customize: + +```python +from agents import RunContextWrapper, function_tool + +def handle_api_error(ctx: RunContextWrapper, error: Exception) -> str: + """Return a helpful error message the LLM can act on.""" + return f"Tool failed: {error}. Try a different query or ask the user for clarification." + +@function_tool(failure_error_function=handle_api_error) +def call_external_api(query: str) -> str: + """Call an external API.""" + # If this raises, handle_api_error returns a message to the LLM + ... +``` + +### LangGraph + +LangGraph tools raise by default. To return errors to the LLM instead of crashing, raise `ToolException` and set `handle_tool_error`: + +```python +from langchain_core.tools import tool, ToolException + +@tool +def call_external_api(query: str) -> str: + """Call an external API.""" + try: + ... + except Exception as e: + raise ToolException(f"API call failed: {e}. Try a different query.") + +# Enable error handling on the tool +call_external_api.handle_tool_error = True +``` + +Set `handle_tool_error=True` for a generic message, or assign a string/callable for custom messages. Only `ToolException` is caught — other exceptions still raise. + +## Tips + +- The docstring becomes the tool description the LLM sees — make it clear and specific +- Type annotations on parameters help the LLM provide correct arguments +- Local tools can call the Databricks SDK, external APIs, or any Python library diff --git a/agent-adk/.claude/skills/create-tools/examples/uc-connection.md b/agent-adk/.claude/skills/create-tools/examples/uc-connection.md new file mode 100644 index 00000000..011baa0a --- /dev/null +++ b/agent-adk/.claude/skills/create-tools/examples/uc-connection.md @@ -0,0 +1,65 @@ +# Create a UC Connection for External MCP Servers + +Unity Catalog HTTP connections let you register external MCP servers so Databricks can securely proxy requests and manage credentials. After creating the connection, your agent accesses the external MCP server through a managed Databricks endpoint. + +## Create the connection + +### Option 1: Managed OAuth (Glean, GitHub, Atlassian, Google Drive, SharePoint) + +For supported providers, Databricks manages the OAuth credentials. Create the connection in the Databricks UI: + +1. Go to **Catalog** > **External Data** > **Connections** +2. Click **Create connection** +3. Select **HTTP** connection type +4. Choose **OAuth User to Machine Per User** auth type +5. Select the provider from the **OAuth Provider** drop-down +6. Configure scopes as needed + +You can also install pre-built integrations from the **Databricks Marketplace**. + +See [external MCP docs](https://docs.databricks.com/aws/en/generative-ai/mcp/external-mcp) for the full list of supported providers, scopes, and setup methods. + +### Option 2: CLI with bearer token + +```bash +databricks connections create --json '{ + "name": "my-external-mcp", + "connection_type": "HTTP", + "options": { + "host": "https://mcp.example.com", + "base_path": "/api", + "bearer_token": "" + } +}' --profile +``` + +### Option 3: CLI with OAuth M2M + +```bash +databricks connections create --json '{ + "name": "my-external-mcp", + "connection_type": "HTTP", + "options": { + "host": "https://mcp.example.com", + "base_path": "/mcp", + "client_id": "", + "client_secret": "", + "token_endpoint": "https://auth.example.com/oauth/token", + "oauth_scope": "read write" + } +}' --profile +``` + +## Verify + +```bash +databricks connections get my-external-mcp --profile +``` + +## Next step + +Wire the external MCP server into your agent. See the **add-tools** skill and use `examples/uc-connection.yaml` for the `databricks.yml` resource grant. + +MCP URL: `{host}/api/2.0/mcp/external/{connection_name}` + +> You can also access the external server through the [UC connections proxy](https://docs.databricks.com/aws/en/query-federation/http#proxy), which works with any HTTP or MCP client and supports arbitrary sub-paths and all HTTP methods: `{host}/api/2.0/unity-catalog/connections/{connection_name}/proxy[/]` diff --git a/agent-adk/.claude/skills/create-tools/examples/uc-function.md b/agent-adk/.claude/skills/create-tools/examples/uc-function.md new file mode 100644 index 00000000..d24e2637 --- /dev/null +++ b/agent-adk/.claude/skills/create-tools/examples/uc-function.md @@ -0,0 +1,67 @@ +# Create a Unity Catalog Function + +UC functions let agents run custom SQL or Python logic. Expose them as tools via the managed MCP server for UC functions. + +## Option 1: SQL function (recommended for data lookups) + +Run this in a SQL warehouse or notebook: + +```sql +CREATE OR REPLACE FUNCTION catalog.schema.lookup_customer( + customer_name STRING COMMENT 'Name of the customer to look up' +) +RETURNS STRING +COMMENT 'Returns customer metadata including email and account ID. Use this when the user asks about a specific customer.' +RETURN SELECT CONCAT( + 'Customer ID: ', customer_id, ', ', + 'Email: ', email + ) + FROM catalog.schema.customers + WHERE name = customer_name + LIMIT 1; +``` + +Via CLI: + +```bash +databricks api post /api/2.0/sql/statements --json '{ + "warehouse_id": "", + "statement": "CREATE OR REPLACE FUNCTION catalog.schema.my_func(...) ..." +}' --profile +``` + +## Option 2: Python function + +```sql +CREATE OR REPLACE FUNCTION catalog.schema.analyze_text( + text STRING COMMENT 'Text to analyze' +) +RETURNS STRING +LANGUAGE PYTHON +COMMENT 'Analyzes text and returns a summary of key entities found.' +AS $$ + # Python code runs in serverless compute + entities = [word for word in text.split() if word[0].isupper()] + return f"Found {len(entities)} potential entities: {', '.join(entities[:5])}" +$$; +``` + +## Writing effective tool descriptions + +The `COMMENT` clause is critical — the LLM uses it to decide when to call the tool. + +- **Function COMMENT**: Describe what the function does and when to use it +- **Parameter COMMENT**: Describe what values the parameter accepts +- Be specific: "Returns customer email and ID given a customer name" is better than "Looks up customer info" + +## Verify + +```bash +databricks functions get catalog.schema.my_func --profile +``` + +## Next step + +Wire the UC function into your agent. See the **add-tools** skill and use `examples/uc-function.yaml` for the `databricks.yml` resource grant. + +MCP URL: `{host}/api/2.0/mcp/functions/{catalog}/{schema}` (exposes all functions in the schema) or `{host}/api/2.0/mcp/functions/{catalog}/{schema}/{function_name}` (single function) (OAuth scope for on-behalf-of-user auth: `unity-catalog`) diff --git a/agent-adk/.claude/skills/create-tools/examples/vector-search-index.md b/agent-adk/.claude/skills/create-tools/examples/vector-search-index.md new file mode 100644 index 00000000..467ef594 --- /dev/null +++ b/agent-adk/.claude/skills/create-tools/examples/vector-search-index.md @@ -0,0 +1,78 @@ +# Create a Vector Search Index + +Vector Search indexes let agents search unstructured data (documents, knowledge bases) using semantic similarity. The managed MCP server handles embedding and retrieval automatically. + +## Prerequisites + +- Unity Catalog enabled in your workspace +- Serverless compute enabled +- A Delta table in Unity Catalog with a text column containing the content to search +- Change Data Feed enabled on the source table (for standard endpoints) +- The index must use **Databricks-managed embeddings** for the managed MCP server + +## Step 1: Create a Vector Search endpoint (if needed) + +```bash +databricks vector-search-endpoints create-endpoint my-vs-endpoint STANDARD --profile +``` + +Verify it exists: + +```bash +databricks vector-search-endpoints list-endpoints --profile +``` + +## Step 2: Create the index with managed embeddings + +When using `--json`, pass all required fields in the JSON body (name, endpoint_name, primary_key, index_type). Do not combine positional arguments with `--json`. + +```bash +databricks vector-search-indexes create-index --json '{ + "name": "..", + "endpoint_name": "my-vs-endpoint", + "primary_key": "id", + "index_type": "DELTA_SYNC", + "delta_sync_index_spec": { + "source_table": "..", + "pipeline_type": "TRIGGERED", + "embedding_source_columns": [ + { + "name": "content", + "embedding_model_endpoint_name": "databricks-gte-large-en" + } + ] + } +}' --profile +``` + +Key parameters: +- `name`: Full 3-part index name (catalog.schema.index) +- `endpoint_name`: The Vector Search endpoint that will serve the index +- `primary_key`: Unique row identifier in the source table +- `index_type`: `DELTA_SYNC` or `DIRECT_ACCESS` +- `source_table`: The Delta table to index +- `embedding_source_columns.name`: The text column to embed and search +- `embedding_model_endpoint_name`: Use `databricks-gte-large-en` (recommended) or another embedding endpoint +- `pipeline_type`: `TRIGGERED` (manual sync) or `CONTINUOUS` (auto-sync on table changes) + +## Step 3: Sync the index + +For `TRIGGERED` pipelines, start the initial sync: + +```bash +databricks vector-search-indexes sync-index .. --profile +``` + +## Verify + +```bash +databricks vector-search-indexes get-index .. --profile +``` + +Check that `status.ready` is `true` before connecting your agent. + +## Next step + +Wire the Vector Search index into your agent. See the **add-tools** skill and use `examples/vector-search.yaml` for the `databricks.yml` resource grant. + +MCP URL: `{host}/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name}` (OAuth scope for on-behalf-of-user auth: `vector-search`) diff --git a/agent-adk/.claude/skills/deploy/SKILL.md b/agent-adk/.claude/skills/deploy/SKILL.md new file mode 100644 index 00000000..7eec6958 --- /dev/null +++ b/agent-adk/.claude/skills/deploy/SKILL.md @@ -0,0 +1,248 @@ +--- +name: deploy +description: "Deploy agent to Databricks Apps using DAB (Databricks Asset Bundles). Use when: (1) User says 'deploy', 'push to databricks', or 'bundle deploy', (2) 'App already exists' error occurs, (3) Need to bind/unbind existing apps, (4) Debugging deployed apps, (5) Querying deployed app endpoints." +--- + +# Deploy to Databricks Apps + +## Profile Configuration + +**IMPORTANT:** Before running any `databricks` CLI command, read the `.env` file to get the `DATABRICKS_CONFIG_PROFILE` value. All commands must include the profile: + +```bash +databricks --profile +``` + +For example, if `.env` has `DATABRICKS_CONFIG_PROFILE=dev`, run `databricks bundle deploy --profile dev`. Without this, the CLI may target the wrong workspace. + +## App Naming Convention + +Unless the user specifies a different name, apps should use the prefix `agent-*`: +- `agent-data-analyst` +- `agent-customer-support` +- `agent-code-helper` + +Update the app name in `databricks.yml`: +```yaml +resources: + apps: + agent_adk: + name: "agent-your-app-name" # Use agent-* prefix +``` + +## Deploy Commands + +**IMPORTANT:** Run the pre-flight check before deploying to catch errors early, then run commands to deploy and start your app: + +```bash +# 1. Pre-flight check (starts server locally, sends test request, verifies response) +uv run preflight + +# 2. Validate bundle configuration (catches config errors before deploy) +databricks bundle validate + +# 3. Deploy the bundle (creates/updates resources, uploads files) +databricks bundle deploy + +# 4. Run the app (starts/restarts with uploaded source code) - REQUIRED! +databricks bundle run agent_adk +``` + +> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is **required** to actually start/restart the app with the new code. If you only run `deploy`, the app will continue running old code! + +The resource key `agent_adk` matches the app name in `databricks.yml` under `resources.apps`. + +## Handling "App Already Exists" Error + +If `databricks bundle deploy` fails with: +``` +Error: failed to create app +Failed to create app . An app with the same name already exists. +``` + +**Ask the user:** "Would you like to bind the existing app to this bundle, or delete it and create a new one?" + +### Option 1: Bind Existing App (Recommended) + +**Step 1:** Get the existing app's full configuration: +```bash +# Get app config including budget_policy_id and other server-side settings +databricks apps get --output json | jq '{name, budget_policy_id, description}' +``` + +**Step 2:** Update `databricks.yml` to match the existing app's configuration exactly: +```yaml +resources: + apps: + agent_adk: + name: "existing-app-name" # Must match exactly + budget_policy_id: "xxx-xxx-xxx" # Copy from step 1 if present +``` + +> **Why this matters:** Existing apps may have server-side configuration (like `budget_policy_id`) that isn't in your bundle. If these don't match, Terraform will fail with "Provider produced inconsistent result after apply". Always sync the app's current config to `databricks.yml` before binding. + +**Step 3:** If deploying to a `mode: production` target, set `workspace.root_path`: +```yaml +targets: + prod: + mode: production + workspace: + root_path: /Workspace/Users/${workspace.current_user.userName}/.bundle/${bundle.name}/${bundle.target} +``` + +> **Why this matters:** Production mode requires an explicit root path to ensure only one copy of the bundle is deployed. Without this, the deploy will fail with a recommendation to set `workspace.root_path`. + +**Step 4:** Check if already bound, then bind if needed: +```bash +# Check if resource is already managed by this bundle +databricks bundle summary --output json | jq '.resources.apps' + +# If the app appears in the summary, skip binding and go to Step 5 +# If NOT in summary, bind the resource: +databricks bundle deployment bind agent_adk --auto-approve +``` + +> **Note:** If bind fails with "Resource already managed by Terraform", the app is already bound to this bundle. Skip to Step 5 and deploy directly. + +**Step 5:** Deploy: +```bash +databricks bundle deploy +databricks bundle run agent_adk +``` + +### Option 2: Delete and Recreate + +```bash +databricks apps delete +databricks bundle deploy +``` + +**Warning:** This permanently deletes the app's URL, OAuth credentials, and service principal. + +## Unbinding an App + +To remove the link between bundle and deployed app: + +```bash +databricks bundle deployment unbind agent_adk +``` + +Use when: +- Switching to a different app +- Letting bundle create a new app +- Switching between deployed instances + +Note: Unbinding doesn't delete the deployed app. + +## Query Deployed App + +> **IMPORTANT:** Databricks Apps are **only** queryable via OAuth token. You **cannot** use a Personal Access Token (PAT) to query your agent. Attempting to use a PAT will result in a 302 redirect error. + +**Get OAuth token:** +```bash +databricks auth token | jq -r '.access_token' +``` + +**Send request:** +```bash +curl -X POST /invocations \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' +``` + +**If using memory** - include `user_id` to scope memories per user: +```bash +curl -X POST /invocations \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "What do you remember about me?"}], + "custom_inputs": {"user_id": "user@example.com"} + }' +``` + +## On-Behalf-Of (OBO) User Authentication + +To authenticate as the requesting user instead of the app service principal: + +```python +from agent_server.utils import get_user_workspace_client + +# In your agent code +user_client = get_user_workspace_client() +# Use user_client for operations that should run as the user +``` + +This is useful when you want the agent to access resources with the user's permissions rather than the app's service principal permissions. + +See: [OBO authentication documentation](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/auth#retrieve-user-authorization-credentials) + +## Debug Deployed Apps + +```bash +# View logs (follow mode) +databricks apps logs --follow + +# Check app status +databricks apps get --output json | jq '{app_status, compute_status}' + +# Get app URL +databricks apps get --output json | jq -r '.url' +``` + +## Post-Deploy: Autoscaling Lakebase Resources + +If the agent uses **autoscaling Lakebase** (user mentions "autoscaling", "project", or "branch" in the context of Lakebase), the postgres resource is declared natively in `databricks.yml` — `databricks bundle deploy` creates the app with it. You only need to grant table permissions to the app's service principal after deploy: + +```bash +# Find the SP client ID +databricks apps get --output json | jq -r '.service_principal_client_id' + +# Grant table permissions (see scripts/grant_lakebase_permissions.py) +``` + +**See `.claude/skills/add-tools/examples/lakebase-autoscaling.yaml` for the full resource snippet.** Requires CLI v0.295.0+ for native `postgres` resource support. + +## Important Notes + +- **App naming convention**: App names must be prefixed with `agent-` (e.g., `agent-my-assistant`, `agent-data-analyst`) +- **Name is immutable**: Changing the `name` field in `databricks.yml` forces app replacement (destroy + create) +- **Remote Terraform state**: Databricks stores state remotely; same app detected across directories +- **Review the plan**: Look for `# forces replacement` in Terraform output before confirming + +## FAQ + +**Q: I see a 200 OK in the logs, but get an error in the actual stream. What's going on?** + +This is expected behavior. The initial 200 OK confirms stream setup was successful. Errors that occur during streaming don't affect the initial HTTP status code. Check the stream content for the actual error message. + +**Q: When querying my agent, I get a 302 redirect error. What's wrong?** + +You're likely using a Personal Access Token (PAT). Databricks Apps only support OAuth tokens. Generate one with: +```bash +databricks auth token +``` + +**Q: How do I add dependencies to my agent?** + +Use `uv add`: +```bash +uv add +# Example: uv add "mlflow-skinny[databricks]" +``` + +## Troubleshooting + +| Issue | Solution | +|-------|----------| +| Validation errors | Run `databricks bundle validate` to see detailed errors before deploying | +| Permission errors at runtime | Grant resources in `databricks.yml` (see **add-tools** skill) | +| Lakebase access errors | See **lakebase-setup** skill for permissions (if using memory) | +| App not starting | Check `databricks apps logs ` | +| Auth token expired | Run `databricks auth token` again | +| 302 redirect error | Use OAuth token, not PAT | +| "Provider produced inconsistent result" | Sync app config to `databricks.yml` | +| "should set workspace.root_path" | Add `root_path` to production target | +| App running old code after deploy | Run `databricks bundle run agent_adk` after deploy | +| Env var is None in deployed app | Check `value_from` in databricks.yml `config.env` matches resource `name` | diff --git a/agent-adk/.claude/skills/discover-tools/SKILL.md b/agent-adk/.claude/skills/discover-tools/SKILL.md new file mode 100644 index 00000000..ae159f95 --- /dev/null +++ b/agent-adk/.claude/skills/discover-tools/SKILL.md @@ -0,0 +1,49 @@ +--- +name: discover-tools +description: "Discover available tools and resources in Databricks workspace. Use when: (1) User asks 'what tools are available', (2) Before writing agent code, (3) Looking for MCP servers, Genie spaces, UC functions, or vector search indexes, (4) User says 'discover', 'find resources', or 'what can I connect to'." +--- + +# Discover Available Tools + +**Run tool discovery BEFORE writing agent code** to understand what resources are available in the workspace. + +## Run Discovery + +```bash +uv run discover-tools +``` + +**Options:** +```bash +# Limit to specific catalog/schema +uv run discover-tools --catalog my_catalog --schema my_schema + +# Output as JSON +uv run discover-tools --format json --output tools.json + +# Save markdown report +uv run discover-tools --output tools.md + +# Use specific Databricks profile +uv run discover-tools --profile DEFAULT +``` + +## What Gets Discovered + +| Resource Type | Description | MCP URL Pattern | +|--------------|-------------|-----------------| +| **UC Functions** | SQL UDFs as agent tools | `{host}/api/2.0/mcp/functions/{catalog}/{schema}` | +| **UC Tables** | Structured data for querying | (via UC functions) | +| **Vector Search Indexes** | RAG applications | `{host}/api/2.0/mcp/vector-search/{catalog}/{schema}` | +| **Genie Spaces** | Natural language data interface | `{host}/api/2.0/mcp/genie/{space_id}` | +| **Custom MCP Servers** | Apps starting with `mcp-*` | `{app_url}/mcp` | +| **External MCP Servers** | Via UC connections | `{host}/api/2.0/mcp/external/{connection_name}` | + +## Next Steps + +After discovering tools: +1. **Add MCP servers to your agent** - See **modify-agent** skill for SDK-specific code examples +2. **Grant permissions** in `databricks.yml` - See **add-tools** skill for YAML snippets +3. **Test locally** with `uv run start-app` - See **run-locally** skill + +Need a resource that doesn't exist yet? See the **create-tools** skill. diff --git a/agent-adk/.claude/skills/lakebase-setup/SKILL.md b/agent-adk/.claude/skills/lakebase-setup/SKILL.md new file mode 100644 index 00000000..d2821ccc --- /dev/null +++ b/agent-adk/.claude/skills/lakebase-setup/SKILL.md @@ -0,0 +1,467 @@ +--- +name: lakebase-setup +description: "Configure Lakebase for agent memory storage. Use when: (1) Adding memory capabilities to the agent, (2) 'Failed to connect to Lakebase' errors, (3) Permission errors on checkpoint/store tables, (4) User says 'lakebase', 'memory setup', or 'add memory'." +--- + +# Lakebase Setup for Agent Persistence + +> **Profile reminder:** All `databricks` CLI commands must include the profile from `.env`: `databricks --profile ` or `DATABRICKS_CONFIG_PROFILE= databricks ` + +> **Two types of Lakebase:** Databricks supports **provisioned** instances (with instance name) and **autoscaling** instances (project/branch model). This skill covers both. Make sure you know which Lakebase instance the user is using, ask the user which type they are using if unclear. + +## Use Cases + +Lakebase is used for three distinct purposes across the agent templates: + +| Use case | Templates | Description | +|----------|-----------|-------------| +| **Chat UI conversation history** | All templates | The built-in chat UI (`e2e-chatbot-app-next`) can persist conversations across page refreshes and browser sessions. This is purely UI-side persistence — the agent itself is stateless. | +| **Agent short-term memory** | `agent-langgraph-advanced`, `agent-openai-advanced` | Conversation threads within a session via `AsyncCheckpointSaver` (LangGraph) or `AsyncDatabricksSession` (OpenAI SDK). The agent remembers what was said earlier in the same conversation. | +| **Agent long-term memory** | `agent-langgraph-advanced` | User facts across sessions via `AsyncDatabricksStore`. The agent remembers things about a user from previous conversations. | + +> **Note:** When the quickstart prompts for Lakebase on a non-memory template, it's for **chat UI history** only — not for the agent. Memory templates always require Lakebase. + +## Overview + +Lakebase provides persistent PostgreSQL storage for agents: +- **Short-term memory** (LangGraph): Conversation history within a thread (`AsyncCheckpointSaver`) +- **Long-term memory** (LangGraph): User facts across sessions (`AsyncDatabricksStore`) +- **Short-term memory** (OpenAI SDK): Conversation history via `AsyncDatabricksSession` +- **Long-running agent persistence** (OpenAI SDK): Background task state via custom SQLAlchemy tables (`agent_server` schema) + +> **Note:** For pre-configured memory templates, see: +> - `agent-langgraph-advanced` - Short-term memory, long-term memory, and long-running background tasks (LangGraph) +> - `agent-openai-advanced` - Short-term memory and long-running background tasks (OpenAI SDK) + +## Complete Setup Workflow + +``` +┌───────────────────────────────────────────────────────────────────────────┐ +│ 1. Add dependency → 2. Get instance → 3. Configure DAB │ +│ 4. Configure .env → 5. Deploy → 6. Grant SP permissions → 7. Run │ +└───────────────────────────────────────────────────────────────────────────┘ +``` + +> **Shortcut:** If using a pre-configured memory template, `uv run quickstart` with Lakebase flags handles steps 2-4 automatically. You still need to do steps 5-7 manually. + +--- + +## Step 1: Add Memory Dependency + +Add the memory extra to your `pyproject.toml`: + +```toml +dependencies = [ + "databricks-langchain[memory]", + # ... other dependencies +] +``` + +Then sync dependencies: +```bash +uv sync +``` + +--- + +## Step 2: Create or Get Lakebase Instance + +### Option A: Provisioned Instance + +1. Go to your Databricks workspace +2. Navigate to **Compute** → **Lakebase** +3. Click **Create Instance** (or use an existing one) +4. Note the **instance name** + +### Option B: Autoscaling Instance + +Autoscaling uses a **project/branch** model. You need three values: +- **Project name** (e.g., `my-project`) +- **Branch name** (e.g., `my-branch`) +- **Database ID** (e.g., `db-xxxx-xxxxxxxxxx`) + +Find these via the postgres API: + +```bash +# List projects +databricks api get /api/2.0/postgres/projects --profile + +# List branches for a project +databricks api get /api/2.0/postgres/projects//branches --profile + +# List databases for a branch +databricks api get /api/2.0/postgres/projects//branches//databases --profile +``` + +**Important:** The database ID is the internal ID (e.g., `db-xxxx-xxxxxxxxxx`), NOT `databricks_postgres`. + +--- + +## Step 3: Configure databricks.yml (Lakebase Resource) + +> **Note:** If you ran `uv run quickstart` with Lakebase flags (`--lakebase-provisioned-name` or `--lakebase-autoscaling-endpoint`), the quickstart already configured `databricks.yml` for you — including fetching the database ID for autoscaling. Manual configuration is only needed if you didn't use quickstart or need to change values. + +### Option A: Provisioned + +Add the `database` resource to your app in `databricks.yml`: + +```yaml +resources: + apps: + your_app: + name: "your-app-name" + source_code_path: ./ + resources: + # ... other resources (experiment, UC functions, etc.) ... + + # Lakebase instance for long-term memory + - name: 'database' + database: + instance_name: '' + database_name: 'databricks_postgres' + permission: 'CAN_CONNECT_AND_CREATE' +``` + +**Important:** +- The `instance_name: ''` must match the actual Lakebase instance name +- Using the `database` resource type automatically grants the app's service principal access to Lakebase +See `.claude/skills/add-tools/examples/lakebase.yaml` for the YAML snippet. + +### Option B: Autoscaling + +Add the `postgres` resource to your app in `databricks.yml`: + +```yaml +resources: + apps: + your_app: + name: "your-app-name" + source_code_path: ./ + resources: + # ... other resources (experiment, UC functions, etc.) ... + + # Autoscaling Lakebase instance for long-term memory + - name: 'postgres' + postgres: + branch: "projects//branches/" + database: "projects//branches//databases/" + permission: 'CAN_CONNECT_AND_CREATE' +``` + +**Important:** The `branch` and `database` fields use full resource path format. + +See `.claude/skills/add-tools/examples/lakebase-autoscaling.yaml` for the YAML snippet. + +### Add Environment Variables to databricks.yml config block + +**Provisioned:** +```yaml + config: + env: + # Lakebase instance name - resolved from database resource at deploy time + - name: LAKEBASE_INSTANCE_NAME + value_from: "database" + # Static values for embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` + +**Autoscaling:** +```yaml + config: + env: + # Autoscaling Lakebase config + - name: LAKEBASE_AUTOSCALING_PROJECT + value: "" + - name: LAKEBASE_AUTOSCALING_BRANCH + value: "" + # Static values for embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` + +--- + +## Step 4: Configure .env (Local Development) + +For local development, add to `.env`: + +**Provisioned:** +```bash +LAKEBASE_INSTANCE_NAME= +EMBEDDING_ENDPOINT=databricks-gte-large-en +EMBEDDING_DIMS=1024 +``` + +**Autoscaling:** +```bash +LAKEBASE_AUTOSCALING_PROJECT= +LAKEBASE_AUTOSCALING_BRANCH= +EMBEDDING_ENDPOINT=databricks-gte-large-en +EMBEDDING_DIMS=1024 +``` + +**Important:** `embedding_dims` must match the embedding endpoint: + +| Endpoint | Dimensions | +|----------|------------| +| `databricks-gte-large-en` | 1024 | +| `databricks-bge-large-en` | 1024 | + +> **Note:** `.env` is only for local development. When deployed, the app gets values from `databricks.yml` config env. + +--- + +## Step 5: Initialize Tables +## Step 5: Deploy + +Deploy the app so the service principal and resources are created: + +```bash +DATABRICKS_CONFIG_PROFILE= databricks bundle deploy +``` + +--- + +## Step 6: Grant SP Permissions (CRITICAL) + +> **WARNING:** You MUST complete this step before running the app. Without it, the app will fail with database migration errors like `CREATE TABLE IF NOT EXISTS "drizzle"."__drizzle_migrations"` — permission denied. + +After deploying, the app's service principal needs Postgres roles to access Lakebase tables. The DAB resource grants basic connectivity, but you must also grant Postgres-level schema and table permissions. + +**Step 1:** Get the app's service principal client ID: +```bash +DATABRICKS_CONFIG_PROFILE= databricks apps get --output json | jq -r '.service_principal_client_id' +``` + +**Step 2:** Grant permissions using the grant script: + +```bash +# Provisioned: +DATABRICKS_CONFIG_PROFILE= uv run python scripts/grant_lakebase_permissions.py \ + --memory-type --instance-name + +# Autoscaling (endpoint — reads LAKEBASE_AUTOSCALING_ENDPOINT from .env by default): +DATABRICKS_CONFIG_PROFILE= uv run python scripts/grant_lakebase_permissions.py \ + --memory-type --autoscaling-endpoint + +# Autoscaling (project + branch): +DATABRICKS_CONFIG_PROFILE= uv run python scripts/grant_lakebase_permissions.py \ + --memory-type --project --branch +``` + +**Memory type by template:** + +| Template | `--memory-type` value | +|----------|-----------------------| +| `agent-langgraph-advanced` | `langgraph` | +| `agent-openai-advanced` | `openai` | + +The script handles fresh branches gracefully (warns but doesn't fail if tables don't exist yet — they'll be created on first app startup). + +--- + +## Step 7: Run Your App + +```bash +DATABRICKS_CONFIG_PROFILE= databricks bundle run {{BUNDLE_NAME}} +``` + +> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is required to actually start the app with the new code. + +--- + +## Complete Examples: databricks.yml with Lakebase + +### Provisioned Lakebase + +```yaml +bundle: + name: agent_langgraph + +resources: + apps: + agent_langgraph: + name: "my-agent-app" + description: "Agent with long-term memory" + source_code_path: ./ + config: + command: ["uv", "run", "start-app"] + env: + - name: MLFLOW_TRACKING_URI + value: "databricks" + - name: MLFLOW_REGISTRY_URI + value: "databricks-uc" + - name: API_PROXY + value: "http://localhost:8000/invocations" + - name: CHAT_APP_PORT + value: "3000" + - name: CHAT_PROXY_TIMEOUT_SECONDS + value: "300" + - name: MLFLOW_EXPERIMENT_ID + value_from: "experiment" + # Lakebase instance name (resolved from database resource) + - name: LAKEBASE_INSTANCE_NAME + value_from: "database" + # Static values for embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" + + resources: + - name: 'experiment' + experiment: + experiment_id: "" + permission: 'CAN_MANAGE' + - name: 'database' + database: + instance_name: '' + database_name: 'databricks_postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +targets: + dev: + mode: development + default: true +``` + +### Autoscaling Lakebase + +```yaml +bundle: + name: agent_langgraph + +resources: + apps: + agent_langgraph: + name: "my-agent-app" + description: "Agent with long-term memory" + source_code_path: ./ + config: + command: ["uv", "run", "start-app"] + env: + - name: MLFLOW_TRACKING_URI + value: "databricks" + - name: MLFLOW_REGISTRY_URI + value: "databricks-uc" + - name: API_PROXY + value: "http://localhost:8000/invocations" + - name: CHAT_APP_PORT + value: "3000" + - name: CHAT_PROXY_TIMEOUT_SECONDS + value: "300" + - name: MLFLOW_EXPERIMENT_ID + value_from: "experiment" + # Autoscaling Lakebase config + - name: LAKEBASE_AUTOSCALING_PROJECT + value: "" + - name: LAKEBASE_AUTOSCALING_BRANCH + value: "" + # Static values for embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" + + resources: + - name: 'experiment' + experiment: + experiment_id: "" + permission: 'CAN_MANAGE' + - name: 'postgres' + postgres: + branch: "projects//branches/" + database: "projects//branches//databases/" + permission: 'CAN_CONNECT_AND_CREATE' + +targets: + dev: + mode: development + default: true +``` + +--- + +## Troubleshooting + +| Issue | Cause | Solution | +|-------|-------|----------| +| **"embedding_dims is required when embedding_endpoint is specified"** | Missing `embedding_dims` parameter | Add `embedding_dims=1024` to AsyncDatabricksStore | +| **"relation 'store' does not exist"** | Tables not initialized | The app creates tables on first use; ensure SP has CREATE permission | +| **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `LAKEBASE_INSTANCE_NAME` to databricks.yml `config.env` | +| **"permission denied for table store"** | Missing grants | Run `uv run python scripts/grant_lakebase_permissions.py ` to grant permissions | +| **"Failed to connect to Lakebase"** | Wrong instance name or project/branch | Verify values in databricks.yml and .env | +| **Connection pool errors on exit** | Python cleanup race | Ignore `PythonFinalizationError` - it's harmless | +| **App not updated after deploy** | Forgot to run bundle | Run `databricks bundle run ` after deploy | +| **value_from not resolving** | Resource name mismatch | Ensure `value_from` value matches `name` in databricks.yml resources | +| **"Invalid postgres resource parameters"** | Missing `database` field in postgres resource | Add full `database` path: `projects//branches//databases/` | +| **`CREATE TABLE IF NOT EXISTS "drizzle"."__drizzle_migrations"` fails** | Grant step was skipped — SP lacks Postgres permissions | Run `grant_lakebase_permissions.py` with `--memory-type`, then restart the app | + +--- + +## LakebaseClient API (for reference) + +```python +from databricks_ai_bridge.lakebase import LakebaseClient, SchemaPrivilege, TablePrivilege + +# Provisioned: +client = LakebaseClient(instance_name="...") +# Autoscaling: +client = LakebaseClient(project="...", branch="...") + +# Create role (must do first) +client.create_role(identity_name, "SERVICE_PRINCIPAL") + +# Grant schema (note: schemas is a list, grantee not role) +client.grant_schema( + grantee="...", + schemas=["public"], + privileges=[SchemaPrivilege.USAGE, SchemaPrivilege.CREATE], +) + +# Grant tables (note: tables includes schema prefix) +client.grant_table( + grantee="...", + tables=["public.store"], + privileges=[TablePrivilege.SELECT, TablePrivilege.INSERT, ...], +) + +# Execute raw SQL +client.execute("SELECT * FROM pg_tables WHERE schemaname = 'public'") +``` + +### Service Principal Identifiers + +When granting permissions manually, note that Databricks apps have multiple identifiers: + +| Field | Format | Example | +|-------|--------|---------| +| `service_principal_id` | Numeric ID | `1234567890123456` | +| `service_principal_client_id` | UUID | `a1b2c3d4-e5f6-7890-abcd-ef1234567890` | +| `service_principal_name` | String name | `my-app-service-principal` | + +**Get all identifiers:** +```bash +DATABRICKS_CONFIG_PROFILE= databricks apps get --output json | jq '{ + id: .service_principal_id, + client_id: .service_principal_client_id, + name: .service_principal_name +}' +``` + +**Which to use:** +- `LakebaseClient.create_role()` - Use `service_principal_client_id` (UUID) or `service_principal_name` +- Raw SQL grants - Use `service_principal_client_id` (UUID) + +--- + +## Next Steps + +- Add memory to agent code: see **agent-memory** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/agent-adk/.claude/skills/load-testing/SKILL.md b/agent-adk/.claude/skills/load-testing/SKILL.md new file mode 100644 index 00000000..d051aa27 --- /dev/null +++ b/agent-adk/.claude/skills/load-testing/SKILL.md @@ -0,0 +1,319 @@ +--- +name: load-testing +description: "Load test a Databricks App to find its maximum QPS. Use when: (1) User says 'load test', 'benchmark', 'QPS', 'throughput', or 'performance test', (2) User wants to find how many queries per second their app can handle, (3) User wants to set up load testing scripts for their agent, (4) User wants to view load test results/dashboard." +--- + +# Load Testing Your Databricks App + +**Goal:** Find the maximum QPS (queries per second) your Databricks App can support. + +## Before You Start — Gather Parameters + +Before beginning, use the `AskUserQuestion` tool to collect the following from the user: + +1. **Do they already have deployed apps to test, or do they need to set up new apps?** +2. **Do they want to mock LLM calls?** Mocking isolates infrastructure throughput from LLM latency — useful for capacity planning. Testing without mocks measures end-to-end performance. +3. **What compute sizes do they want to test?** (Medium, Large, or both) +4. **How many worker configurations do they want to test?** (e.g., 2, 4, 6, 8 workers) +5. **Do they have M2M OAuth credentials (service principal client_id/client_secret)?** — Recommended for tests longer than ~30 minutes. If not, guide them to create one. +6. **What is their `DATABRICKS_HOST`?** (workspace URL) + +--- + +## Step 1: Set Up Load Testing Scripts + +Create a `load-test-scripts/` directory in the project with the following files. These scripts are framework-agnostic and work with any Databricks App. + +### Directory Structure + +``` +/ + agent_server/ # Existing agent code + load-test-scripts/ # Load testing scripts (create this) + run_load_test.py # Main CLI — orchestrates Locust tests + locustfile.py # Locust test definition (SSE streaming, TTFT tracking) + dashboard_template.py # Generates interactive HTML dashboard from results + .env.example # Template for env vars + load-test-runs/ # Test results (auto-created per run) + / + dashboard.html # Interactive dashboard + test_config.json # Test parameters +