From 2d48ec0a826fb85b802d2c127de9ff12b014a174 Mon Sep 17 00:00:00 2001 From: Dhruv Gupta Date: Tue, 27 Jan 2026 11:00:17 -0800 Subject: [PATCH 01/14] Added Skills to Agents Langgraph Template --- .../.claude/skills/add-tools/SKILL.md | 80 ++++ .../add-tools/examples/custom-mcp-server.md | 57 +++ .../skills/add-tools/examples/experiment.yaml | 8 + .../add-tools/examples/genie-space.yaml | 9 + .../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/deploy/SKILL.md | 145 ++++++ .../.claude/skills/discover-tools/SKILL.md | 82 ++++ .../.claude/skills/modify-agent/SKILL.md | 270 +++++++++++ .../.claude/skills/quickstart/SKILL.md | 83 ++++ .../.claude/skills/run-locally/SKILL.md | 90 ++++ agent-langgraph/.gitignore | 16 +- agent-langgraph/AGENTS.md | 361 +++------------ agent-langgraph/README.md | 4 + agent-langgraph/databricks.yml | 37 ++ agent-langgraph/pyproject.toml | 1 + agent-langgraph/scripts/discover_tools.py | 432 ++++++++++++++++++ agent-langgraph/scripts/start_app.py | 45 +- 21 files changed, 1448 insertions(+), 313 deletions(-) create mode 100644 agent-langgraph/.claude/skills/add-tools/SKILL.md create mode 100644 agent-langgraph/.claude/skills/add-tools/examples/custom-mcp-server.md create mode 100644 agent-langgraph/.claude/skills/add-tools/examples/experiment.yaml create mode 100644 agent-langgraph/.claude/skills/add-tools/examples/genie-space.yaml create mode 100644 agent-langgraph/.claude/skills/add-tools/examples/serving-endpoint.yaml create mode 100644 agent-langgraph/.claude/skills/add-tools/examples/sql-warehouse.yaml create mode 100644 agent-langgraph/.claude/skills/add-tools/examples/uc-connection.yaml create mode 100644 agent-langgraph/.claude/skills/add-tools/examples/uc-function.yaml create mode 100644 agent-langgraph/.claude/skills/add-tools/examples/vector-search.yaml create mode 100644 agent-langgraph/.claude/skills/deploy/SKILL.md create mode 100644 agent-langgraph/.claude/skills/discover-tools/SKILL.md create mode 100644 agent-langgraph/.claude/skills/modify-agent/SKILL.md create mode 100644 agent-langgraph/.claude/skills/quickstart/SKILL.md create mode 100644 agent-langgraph/.claude/skills/run-locally/SKILL.md create mode 100644 agent-langgraph/databricks.yml create mode 100644 agent-langgraph/scripts/discover_tools.py diff --git a/agent-langgraph/.claude/skills/add-tools/SKILL.md b/agent-langgraph/.claude/skills/add-tools/SKILL.md new file mode 100644 index 00000000..7d00e104 --- /dev/null +++ b/agent-langgraph/.claude/skills/add-tools/SKILL.md @@ -0,0 +1,80 @@ +--- +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 + +**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. + +## Workflow + +**Step 1:** Add MCP server in `agent_server/agent.py`: +```python +from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient + +genie_server = DatabricksMCPServer( + url=f"{host}/api/2.0/mcp/genie/01234567-89ab-cdef", + name="my genie space", +) + +mcp_client = DatabricksMultiServerMCPClient([genie_server]) +tools = await mcp_client.get_tools() +``` + +**Step 2:** Grant access in `databricks.yml`: +```yaml +resources: + apps: + agent_langgraph: + resources: + - name: 'my_genie_space' + genie_space: + name: 'My Genie Space' + space_id: '01234567-89ab-cdef' + permission: 'CAN_RUN' +``` + +**Step 3:** Deploy with `databricks bundle deploy` (see **deploy** skill) + +## 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 | +| `experiment.yaml` | MLflow experiment | Tracing (already configured) | +| `custom-mcp-server.md` | Custom MCP apps | Apps starting with `mcp-*` | + +## Custom MCP Servers (Databricks Apps) + +Apps are **not yet supported** as resource dependencies in `databricks.yml`. Manual permission grant required: + +**Step 1:** Get your agent app's service principal: +```bash +databricks apps get --output json | jq -r '.service_principal_name' +``` + +**Step 2:** Grant permission on the MCP server app: +```bash +databricks apps update-permissions \ + --service-principal \ + --permission-level CAN_USE +``` + +See `examples/custom-mcp-server.md` for detailed steps. + +## 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 after changes**: Run `databricks bundle deploy` after modifying `databricks.yml` diff --git a/agent-langgraph/.claude/skills/add-tools/examples/custom-mcp-server.md b/agent-langgraph/.claude/skills/add-tools/examples/custom-mcp-server.md new file mode 100644 index 00000000..1324e6c5 --- /dev/null +++ b/agent-langgraph/.claude/skills/add-tools/examples/custom-mcp-server.md @@ -0,0 +1,57 @@ +# Custom MCP Server (Databricks App) + +Custom MCP servers are Databricks Apps with names starting with `mcp-*`. + +**Apps are not yet supported as resource dependencies in `databricks.yml`**, so manual permission grant is required. + +## 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. Deploy your agent app first + +```bash +databricks bundle deploy +databricks bundle run agent_langgraph +``` + +### 3. Get your agent app's service principal + +```bash +databricks apps get --output json | jq -r '.service_principal_name' +``` + +Example output: `sp-abc123-def456` + +### 4. Grant permission on the MCP server app + +```bash +databricks apps update-permissions \ + --service-principal \ + --permission-level CAN_USE +``` + +Example: +```bash +databricks apps update-permissions mcp-my-server \ + --service-principal sp-abc123-def456 \ + --permission-level CAN_USE +``` + +## Notes + +- This manual step is required each time you connect to a new custom MCP server +- The permission grant persists across deployments +- If you redeploy the agent app with a new service principal, you'll need to grant permissions again diff --git a/agent-langgraph/.claude/skills/add-tools/examples/experiment.yaml b/agent-langgraph/.claude/skills/add-tools/examples/experiment.yaml new file mode 100644 index 00000000..ac5c626a --- /dev/null +++ b/agent-langgraph/.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-langgraph/.claude/skills/add-tools/examples/genie-space.yaml b/agent-langgraph/.claude/skills/add-tools/examples/genie-space.yaml new file mode 100644 index 00000000..71589d52 --- /dev/null +++ b/agent-langgraph/.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-langgraph/.claude/skills/add-tools/examples/serving-endpoint.yaml b/agent-langgraph/.claude/skills/add-tools/examples/serving-endpoint.yaml new file mode 100644 index 00000000..b49ce9da --- /dev/null +++ b/agent-langgraph/.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-langgraph/.claude/skills/add-tools/examples/sql-warehouse.yaml b/agent-langgraph/.claude/skills/add-tools/examples/sql-warehouse.yaml new file mode 100644 index 00000000..a6ce9446 --- /dev/null +++ b/agent-langgraph/.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-langgraph/.claude/skills/add-tools/examples/uc-connection.yaml b/agent-langgraph/.claude/skills/add-tools/examples/uc-connection.yaml new file mode 100644 index 00000000..316675fe --- /dev/null +++ b/agent-langgraph/.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-langgraph/.claude/skills/add-tools/examples/uc-function.yaml b/agent-langgraph/.claude/skills/add-tools/examples/uc-function.yaml new file mode 100644 index 00000000..43f938a9 --- /dev/null +++ b/agent-langgraph/.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-langgraph/.claude/skills/add-tools/examples/vector-search.yaml b/agent-langgraph/.claude/skills/add-tools/examples/vector-search.yaml new file mode 100644 index 00000000..0ba39027 --- /dev/null +++ b/agent-langgraph/.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-langgraph/.claude/skills/deploy/SKILL.md b/agent-langgraph/.claude/skills/deploy/SKILL.md new file mode 100644 index 00000000..47c88f4b --- /dev/null +++ b/agent-langgraph/.claude/skills/deploy/SKILL.md @@ -0,0 +1,145 @@ +--- +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 + +## Deploy Commands + +```bash +# Deploy the bundle (creates/updates resources, uploads files) +databricks bundle deploy + +# Run the app (starts/restarts with uploaded source code) +databricks bundle run agent_langgraph +``` + +The resource key `agent_langgraph` 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_langgraph: + 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_langgraph --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_langgraph +``` + +### 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_langgraph +``` + +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 + +**Get OAuth token** (PATs not supported): +```bash +databricks auth token +``` + +**Send request:** +```bash +curl -X POST /invocations \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' +``` + +## 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' +``` + +## 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 + +## Troubleshooting + +| Issue | Solution | +|-------|----------| +| Permission errors at runtime | Grant resources in `databricks.yml` (see **add-tools** skill) | +| App not starting | Check `databricks apps logs ` | +| Auth token expired | Run `databricks auth token` again | +| "Provider produced inconsistent result" | Sync app config to `databricks.yml`| +| "should set workspace.root_path" | Add `root_path` to production target | diff --git a/agent-langgraph/.claude/skills/discover-tools/SKILL.md b/agent-langgraph/.claude/skills/discover-tools/SKILL.md new file mode 100644 index 00000000..eabac2a5 --- /dev/null +++ b/agent-langgraph/.claude/skills/discover-tools/SKILL.md @@ -0,0 +1,82 @@ +--- +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}` | + +## Using Discovered Tools in Code + +After discovering tools, add them to your agent in `agent_server/agent.py`: + +```python +from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient + +# Example: Add UC functions from a schema +uc_functions_server = DatabricksMCPServer( + url=f"{host}/api/2.0/mcp/functions/{catalog}/{schema}", + name="my uc functions", +) + +# Example: Add a Genie space +genie_server = DatabricksMCPServer( + url=f"{host}/api/2.0/mcp/genie/{space_id}", + name="my genie space", +) + +# Example: Add vector search +vector_server = DatabricksMCPServer( + url=f"{host}/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name}", + name="my vector index", +) + +# Create MCP client with all servers +mcp_client = DatabricksMultiServerMCPClient([ + uc_functions_server, + genie_server, + vector_server, +]) + +# Get tools for the agent +tools = await mcp_client.get_tools() +``` + +## Next Steps + +After adding MCP servers to your agent: +1. **Grant permissions** in `databricks.yml` (see **add-tools** skill) +2. Test locally with `uv run start-app` (see **run-locally** skill) diff --git a/agent-langgraph/.claude/skills/modify-agent/SKILL.md b/agent-langgraph/.claude/skills/modify-agent/SKILL.md new file mode 100644 index 00000000..b42a77d4 --- /dev/null +++ b/agent-langgraph/.claude/skills/modify-agent/SKILL.md @@ -0,0 +1,270 @@ +--- +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 + +## Main File + +**`agent_server/agent.py`** - Agent logic, model selection, instructions, MCP servers + +## Key Files + +| File | Purpose | +|------|---------| +| `agent_server/agent.py` | Agent logic, model, instructions, MCP servers | +| `agent_server/start_server.py` | FastAPI server + MLflow setup | +| `agent_server/evaluate_agent.py` | Agent evaluation with MLflow scorers | +| `agent_server/utils.py` | Databricks auth helpers, stream processing | +| `databricks.yml` | Bundle config & resource permissions | + +## SDK Setup + +```python +import mlflow +from databricks.sdk import WorkspaceClient +from databricks_langchain import ChatDatabricks, DatabricksMCPServer, DatabricksMultiServerMCPClient +from langchain.agents import create_agent + +# Enable autologging for tracing +mlflow.langchain.autolog() + +# Initialize workspace client +workspace_client = WorkspaceClient() +``` + +--- + +## databricks-langchain SDK Overview + +**SDK Location:** https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/langchain + +Before making any changes, ensure that the APIs actually exist in the SDK. If something is missing from the documentation here, look in the venv's `site-packages` directory for the `databricks_langchain` package. If it's not installed, run `uv sync` to create the .venv and install the package. + +--- + +### ChatDatabricks - LLM Chat Interface + +Connects to Databricks Model Serving endpoints for LLM inference. + +```python +from databricks_langchain import ChatDatabricks + +llm = ChatDatabricks( + endpoint="databricks-claude-3-7-sonnet", # or databricks-meta-llama-3-1-70b-instruct + temperature=0, + max_tokens=500, +) + +# For Responses API agents: +llm = ChatDatabricks(endpoint="my-agent-endpoint", use_responses_api=True) +``` + +Available models (check workspace for current list): +- `databricks-claude-3-7-sonnet` +- `databricks-claude-3-5-sonnet` +- `databricks-meta-llama-3-3-70b-instruct` + +**Note:** Some workspaces require granting the app access to the serving endpoint in `databricks.yml`. See the **add-tools** skill and `examples/serving-endpoint.yaml`. + +--- + +### DatabricksEmbeddings - Generate Embeddings + +Query Databricks embedding model endpoints. + +```python +from databricks_langchain import DatabricksEmbeddings + +embeddings = DatabricksEmbeddings(endpoint="databricks-bge-large-en") +vector = embeddings.embed_query("The meaning of life is 42") +vectors = embeddings.embed_documents(["doc1", "doc2"]) +``` + +--- + +### DatabricksVectorSearch - Vector Store + +Connect to Databricks Vector Search indexes for similarity search. + +```python +from databricks_langchain import DatabricksVectorSearch + +# Delta-sync index with Databricks-managed embeddings +vs = DatabricksVectorSearch(index_name="catalog.schema.index_name") + +# Direct-access or self-managed embeddings +vs = DatabricksVectorSearch( + index_name="catalog.schema.index_name", + embedding=embeddings, + text_column="content", +) + +docs = vs.similarity_search("query", k=5) +``` + +--- + +### MCP Client - Tool Integration + +Connect to MCP (Model Context Protocol) servers to get tools for your agent. + +**Basic MCP Server (manual URL):** + +```python +from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient + +client = DatabricksMultiServerMCPClient([ + DatabricksMCPServer( + name="system-ai", + url=f"{host}/api/2.0/mcp/functions/system/ai", + ) +]) +tools = await client.get_tools() +``` + +**From UC Function (convenience helper):** + +Creates MCP server for Unity Catalog functions. If `function_name` is omitted, exposes all functions in the schema. + +```python +server = DatabricksMCPServer.from_uc_function( + catalog="main", + schema="tools", + function_name="send_email", # Optional - omit for all functions in schema + name="email-server", + timeout=30.0, + handle_tool_error=True, +) +``` + +**From Vector Search (convenience helper):** + +Creates MCP server for Vector Search indexes. If `index_name` is omitted, exposes all indexes in the schema. + +```python +server = DatabricksMCPServer.from_vector_search( + catalog="main", + schema="embeddings", + index_name="product_docs", # Optional - omit for all indexes in schema + name="docs-search", + timeout=30.0, +) +``` + +**From Genie Space:** + +Create MCP server from Genie Space. Get the genie space ID from the URL. + +Example: `https://workspace.cloud.databricks.com/genie/rooms/01f0515f6739169283ef2c39b7329700?o=123` means the genie space ID is `01f0515f6739169283ef2c39b7329700` + +```python +DatabricksMCPServer( + name="genie", + url=f"{host_name}/api/2.0/mcp/genie/01f0515f6739169283ef2c39b7329700", +) +``` + +**Non-Databricks MCP Server:** + +```python +from databricks_langchain import MCPServer + +server = MCPServer( + name="external-server", + url="https://other-server.com/mcp", + headers={"X-API-Key": "secret"}, + timeout=15.0, +) +``` + +**After adding MCP servers:** Grant permissions in `databricks.yml` (see **add-tools** skill) + +--- + +## Running the Agent + +```python +from langchain.agents import create_agent + +# Create agent +agent = create_agent(tools=tools, model=llm) + +# Non-streaming +messages = {"messages": [{"role": "user", "content": "hi"}]} +result = await agent.ainvoke(messages) + +# Streaming +async for event in agent.astream(input=messages, stream_mode=["updates", "messages"]): + # Process stream events + pass +``` + +**Converting to Responses API format:** Use `process_agent_astream_events()` from `agent_server/utils.py`: + +```python +from agent_server.utils import process_agent_astream_events + +async for event in process_agent_astream_events( + agent.astream(input=messages, stream_mode=["updates", "messages"]) +): + yield event # Yields ResponsesAgentStreamEvent objects +``` + +--- + +## Customizing Agent Behavior + +In LangGraph, agent behavior is customized via system messages and agent configuration rather than an `instructions` parameter. + +**Add system message to conversation:** +```python +messages = { + "messages": [ + {"role": "system", "content": """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."""}, + {"role": "user", "content": "What were Q4 sales?"} + ] +} +result = await agent.ainvoke(messages) +``` + +For advanced customization (routing, state management, custom graphs), refer to the [LangGraph documentation](https://docs.langchain.com/oss/python/langgraph/overview). + +--- + +## External Connection Tools + +Connect to external services via Unity Catalog HTTP connections: + +- **Slack** - Post messages to channels +- **Google Calendar** - Calendar operations +- **Microsoft Graph API** - Office 365 services +- **Azure AI Search** - Search functionality +- **Any HTTP API** - Use `http_request` from databricks-sdk + +Example: Create UC function wrapping HTTP request for Slack, then expose via MCP. + +--- + +## External Resources + +1. [databricks-langchain SDK](https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/langchain) +2. [Agent examples](https://github.com/bbqiu/agent-on-app-prototype) +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. [LangGraph documentation](https://docs.langchain.com/oss/python/langgraph/overview) +6. [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/agent-langgraph/.claude/skills/quickstart/SKILL.md b/agent-langgraph/.claude/skills/quickstart/SKILL.md new file mode 100644 index 00000000..0ca105e6 --- /dev/null +++ b/agent-langgraph/.claude/skills/quickstart/SKILL.md @@ -0,0 +1,83 @@ +--- +name: quickstart +description: "Set up Databricks agent development environment. Use when: (1) First time setup, (2) Configuring Databricks authentication, (3) User says 'quickstart', 'set up', 'authenticate', or 'configure databricks', (4) No .env.local file exists." +--- + +# Quickstart & Authentication + +## Prerequisites + +- **uv** (Python package manager) +- **nvm** with Node 20 (for frontend) +- **Databricks CLI v0.283.0+** + +Check CLI version: +```bash +databricks -v # Must be v0.283.0 or above +brew upgrade databricks # If version is too old +``` + +## Run Quickstart + +```bash +uv run quickstart +``` + +**Options:** +- `--profile NAME`: Use specified profile (non-interactive) +- `--host URL`: Workspace URL for initial setup +- `-h, --help`: Show help + +**Examples:** +```bash +# Interactive (prompts for profile selection) +uv run quickstart + +# Non-interactive with existing profile +uv run quickstart --profile DEFAULT + +# New workspace setup +uv run quickstart --host https://your-workspace.cloud.databricks.com +``` + +## What Quickstart Configures + +Creates/updates `.env.local` with: +- `DATABRICKS_CONFIG_PROFILE` - Selected CLI profile +- `MLFLOW_TRACKING_URI` - Set to `databricks://` for local auth +- `MLFLOW_EXPERIMENT_ID` - Auto-created experiment ID + +## Manual Authentication (Fallback) + +If quickstart fails: + +```bash +# Create new profile +databricks auth login --host https://your-workspace.cloud.databricks.com + +# Verify +databricks auth profiles +``` + +Then manually create `.env.local` (copy from `.env.example`): +```bash +# Authentication (choose one method) +DATABRICKS_CONFIG_PROFILE=DEFAULT +# DATABRICKS_HOST=https://.databricks.com +# DATABRICKS_TOKEN=dapi.... + +# MLflow configuration +MLFLOW_EXPERIMENT_ID= +MLFLOW_TRACKING_URI="databricks://DEFAULT" +MLFLOW_REGISTRY_URI="databricks-uc" + +# Frontend proxy settings +CHAT_APP_PORT=3000 +CHAT_PROXY_TIMEOUT_SECONDS=300 +``` + +## Next Steps + +After quickstart completes: +1. Run `uv run discover-tools` to find available workspace resources (see **discover-tools** skill) +2. Run `uv run start-app` to test locally (see **run-locally** skill) diff --git a/agent-langgraph/.claude/skills/run-locally/SKILL.md b/agent-langgraph/.claude/skills/run-locally/SKILL.md new file mode 100644 index 00000000..294be8c8 --- /dev/null +++ b/agent-langgraph/.claude/skills/run-locally/SKILL.md @@ -0,0 +1,90 @@ +--- +name: run-locally +description: "Run and test the agent locally. Use when: (1) User says 'run locally', 'start server', 'test agent', or 'localhost', (2) Need curl commands to test API, (3) Troubleshooting local development issues, (4) Configuring server options like port or hot-reload." +--- + +# Run Agent Locally + +## Start the Server + +```bash +uv run start-app +``` + +This starts the agent at http://localhost:8000 + +## Server Options + +```bash +# Hot-reload on code changes (development) +uv run start-server --reload + +# Custom port +uv run start-server --port 8001 + +# Multiple workers (production-like) +uv run start-server --workers 4 + +# Combine options +uv run start-server --reload --port 8001 +``` + +## Test the API + +**Streaming request:** +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' +``` + +**Non-streaming request:** +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }] }' +``` + +## Run Evaluation + +```bash +uv run agent-evaluate +``` + +Uses MLflow scorers (RelevanceToQuery, Safety). + +## Run Unit Tests + +```bash +pytest [path] +``` + +## Troubleshooting + +| Issue | Solution | +|-------|----------| +| **Port already in use** | Use `--port 8001` or kill existing process | +| **Authentication errors** | Verify `.env.local` is correct; run **quickstart** skill | +| **Module not found** | Run `uv sync` to install dependencies | +| **MLflow experiment not found** | Ensure `MLFLOW_TRACKING_URI` in `.env.local` is `databricks://` | + +### MLflow Experiment Not Found + +If you see: "The provided MLFLOW_EXPERIMENT_ID environment variable value does not exist" + +**Verify the experiment exists:** +```bash +databricks -p experiments get-experiment +``` + +**Fix:** Ensure `.env.local` has the correct tracking URI format: +```bash +MLFLOW_TRACKING_URI="databricks://DEFAULT" # Include profile name +``` + +The quickstart script configures this automatically. If you manually edited `.env.local`, ensure the profile name is included. + +## Next Steps + +- Modify your agent: see **modify-agent** skill +- Deploy to Databricks: see **deploy** skill diff --git a/agent-langgraph/.gitignore b/agent-langgraph/.gitignore index a058c65c..d4a984d8 100644 --- a/agent-langgraph/.gitignore +++ b/agent-langgraph/.gitignore @@ -1,8 +1,6 @@ # Created by https://www.toptal.com/developers/gitignore/api/python # Edit at https://www.toptal.com/developers/gitignore?templates=python -databricks.yml - ### Python ### # Byte-compiled / optimized / DLL files __pycache__/ @@ -204,6 +202,18 @@ sketch **/mlruns/ **/.vite/ **/.databricks -**/.claude + +# Ignore .claude directory but track template-provided skills +# User-created skills will be ignored by default (no conflicts) +.claude/* +!.claude/skills/ +.claude/skills/* +!.claude/skills/quickstart/ +!.claude/skills/discover-tools/ +!.claude/skills/deploy/ +!.claude/skills/add-tools/ +!.claude/skills/run-locally/ +!.claude/skills/modify-agent/ + **/.env **/.env.local \ No newline at end of file diff --git a/agent-langgraph/AGENTS.md b/agent-langgraph/AGENTS.md index fcd86db5..a4168cd7 100644 --- a/agent-langgraph/AGENTS.md +++ b/agent-langgraph/AGENTS.md @@ -1,348 +1,99 @@ -# Agent LangGraph Development Guide +# Agent Development Guide -## Running the App +## MANDATORY First Action -**Prerequisites:** uv, nvm (Node 20), Databricks CLI +**BEFORE any other action, run `databricks auth profiles` to check authentication status.** -**Quick Start:** +This helps you understand: +- Which Databricks profiles are configured +- Whether authentication is already set up +- Which profile to use for subsequent commands -```bash -uv run quickstart # First-time setup (auth, MLflow experiment, env) -uv run start-app # Start app at http://localhost:8000 -``` +If no profiles exist, guide the user through running `uv run quickstart` to set up authentication. -**Advanced Server Options:** +## Understanding User Goals -```bash -uv run start-server --reload # Hot-reload on code changes during development -uv run start-server --port 8001 -uv run start-server --workers 4 -``` +**Ask the user questions to understand what they're building:** -**Test API:** +1. **What is the agent's purpose?** (e.g., data analyst assistant, customer support, code helper) +2. **What data or tools does it need access to?** + - Databases/tables (Unity Catalog) + - Documents for RAG (Vector Search) + - Natural language data queries (Genie Spaces) + - External APIs or services +3. **Any specific Databricks resources they want to connect?** -```bash -# Streaming request -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' +Use `uv run discover-tools` to show them available resources in their workspace, then help them select the right ones for their use case. **See the `add-tools` skill for how to connect tools and grant permissions.** -# Non-streaming request -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ "input": [{ "role": "user", "content": "hi" }] }' -``` +## Handling Deployment Errors ---- - -## Testing the Agent - -**Run evaluation:** +**If `databricks bundle deploy` fails with "An app with the same name already exists":** -```bash -uv run agent-evaluate # Uses MLflow scorers (RelevanceToQuery, Safety) -``` +Ask the user: "I see there's an existing app with the same name. Would you like me to bind it to this bundle so we can manage it, or delete it and create a new one?" -**Run unit tests:** - -```bash -pytest [path] # Standard pytest execution -``` +- **If they want to bind**: See the **deploy** skill for binding steps +- **If they want to delete**: Run `databricks apps delete ` then deploy again --- -## Modifying the Agent - -Anytime the user wants to modify the agent, look through each of the following resources to help them accomplish their goal: +## Available Skills -If the user wants to convert something into Responses API, refer to https://mlflow.org/docs/latest/genai/serving/responses-agent/ for more information. +**Before executing any task, read the relevant skill file in `.claude/skills/`** - they contain tested commands, patterns, and troubleshooting steps. -1. Look through existing databricks-langchain APIs to see if they can use one of these to accomplish their goal. -2. Look through the folders in https://github.com/bbqiu/agent-on-app-prototype to see if there's an existing example similar to what they're looking to do. -3. Reference the documentation available under https://docs.databricks.com/aws/en/generative-ai/agent-framework/ and its subpages. -4. For adding tools and capabilities, refer to: https://docs.databricks.com/aws/en/generative-ai/agent-framework/agent-tool -5. For stuff like LangGraph routing, configuration, and customization, refer to the LangGraph documentation: https://docs.langchain.com/oss/python/langgraph/overview. +| Task | Skill | Path | +|------|-------|------| +| Setup, auth, first-time | **quickstart** | `.claude/skills/quickstart/SKILL.md` | +| Find tools/resources | **discover-tools** | `.claude/skills/discover-tools/SKILL.md` | +| Deploy to Databricks | **deploy** | `.claude/skills/deploy/SKILL.md` | +| Add tools & permissions | **add-tools** | `.claude/skills/add-tools/SKILL.md` | +| Run/test locally | **run-locally** | `.claude/skills/run-locally/SKILL.md` | +| Modify agent code | **modify-agent** | `.claude/skills/modify-agent/SKILL.md` | -**Main file to modify:** `agent_server/agent.py` +**Note:** All agent skills are located in `.claude/skills/` directory. --- -## databricks-langchain SDK overview - -**SDK Location:** `https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/langchain` +## Quick Commands -**Development Workflow:** - -```bash -uv add databricks-langchain -``` - -Before making any changes, ensure that the APIs actually exist in the SDK. If something is missing from the documentation here, feel free to look in the venv's `site-packages` directory for the `databricks_langchain` package. If it's not installed, run `uv sync` in this folder to create the .venv and install the package. +| Task | Command | +|------|---------| +| Setup | `uv run quickstart` | +| Discover tools | `uv run discover-tools` | +| Run locally | `uv run start-app` | +| Deploy | `databricks bundle deploy && databricks bundle run agent_langgraph` | +| View logs | `databricks apps logs --follow` | --- -### ChatDatabricks - LLM Chat Interface - -Connects to Databricks Model Serving endpoints for LLM inference. - -```python -from databricks_langchain import ChatDatabricks - -llm = ChatDatabricks( - endpoint="databricks-claude-3-7-sonnet", # or databricks-meta-llama-3-1-70b-instruct - temperature=0, - max_tokens=500, -) - -# For Responses API agents: -llm = ChatDatabricks(endpoint="my-agent-endpoint", use_responses_api=True) -``` - ---- - -### DatabricksEmbeddings - Generate Embeddings - -Query Databricks embedding model endpoints. - -```python -from databricks_langchain import DatabricksEmbeddings - -embeddings = DatabricksEmbeddings(endpoint="databricks-bge-large-en") -vector = embeddings.embed_query("The meaning of life is 42") -vectors = embeddings.embed_documents(["doc1", "doc2"]) -``` - ---- - -### DatabricksVectorSearch - Vector Store - -Connect to Databricks Vector Search indexes for similarity search. - -```python -from databricks_langchain import DatabricksVectorSearch - -# Delta-sync index with Databricks-managed embeddings -vs = DatabricksVectorSearch(index_name="catalog.schema.index_name") - -# Direct-access or self-managed embeddings -vs = DatabricksVectorSearch( - index_name="catalog.schema.index_name", - embedding=embeddings, - text_column="content", -) +## Key Files -docs = vs.similarity_search("query", k=5) -``` +| File | Purpose | +|------|---------| +| `agent_server/agent.py` | Agent logic, model, instructions, MCP servers | +| `agent_server/start_server.py` | FastAPI server + MLflow setup | +| `agent_server/evaluate_agent.py` | Agent evaluation with MLflow scorers | +| `databricks.yml` | Bundle config & resource permissions | +| `scripts/quickstart.py` | One-command setup script | +| `scripts/discover_tools.py` | Discovers available workspace resources | --- -### MCP Client - Tool Integration - -Connect to MCP (Model Context Protocol) servers to get tools for your agent. - -**Basic MCP Server (manual URL):** - -```python -from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient - -client = DatabricksMultiServerMCPClient([ - DatabricksMCPServer( - name="system-ai", - url=f"{host}/api/2.0/mcp/functions/system/ai", - ) -]) -tools = await client.get_tools() -``` - -**From UC Function (convenience helper):** -Creates MCP server for Unity Catalog functions. If `function_name` is omitted, exposes all functions in the schema. - -```python -server = DatabricksMCPServer.from_uc_function( - catalog="main", - schema="tools", - function_name="send_email", # Optional - omit for all functions in schema - name="email-server", - timeout=30.0, - handle_tool_error=True, -) -``` - -**From Vector Search (convenience helper):** -Creates MCP server for Vector Search indexes. If `index_name` is omitted, exposes all indexes in the schema. - -```python -server = DatabricksMCPServer.from_vector_search( - catalog="main", - schema="embeddings", - index_name="product_docs", # Optional - omit for all indexes in schema - name="docs-search", - timeout=30.0, -) -``` - -**From Genie Space:** -Create MCP server from Genie Space. Need to get the genie space ID. Can prompt the user to retrieve this via the UI by getting the link to the genie space. - -Ex: https://db-ml-models-dev-us-west.cloud.databricks.com/genie/rooms/01f0515f6739169283ef2c39b7329700?o=3217006663075879 means the genie space ID is 01f0515f6739169283ef2c39b7329700 - -```python -DatabricksMCPServer( - name="genie", - url=f"{host_name}/api/2.0/mcp/genie/01f0515f6739169283ef2c39b7329700", -), -``` - -**Non-Databricks MCP Server:** - -```python -from databricks_langchain import MCPServer - -server = MCPServer( - name="external-server", - url="https://other-server.com/mcp", - headers={"X-API-Key": "secret"}, - timeout=15.0, -) -``` - -### Stateful LangGraph agent - -To enable statefulness in a LangGraph agent, we need to install `databricks-langchain[memory]`. - -Look through the package files for the latest on stateful langgraph agents. Can start by looking at the databricks_langchain/checkpoints.py and databricks_langchain/store.py files. - -## Lakebase instance setup for stateful agents - -Add the lakebase name to `.env`: - -```bash -LAKEBASE_INSTANCE_NAME= -``` - ## Agent Framework Capabilities -Reference: https://docs.databricks.com/aws/en/generative-ai/agent-framework/ - -### Tool Types +> **IMPORTANT:** When adding any tool to the agent, you MUST also grant permissions in `databricks.yml`. See the **add-tools** skill for required steps and examples. +**Tool Types:** 1. **Unity Catalog Function Tools** - SQL UDFs managed in UC with built-in governance 2. **Agent Code Tools** - Defined directly in agent code for REST APIs and low-latency operations 3. **MCP Tools** - Interoperable tools via Model Context Protocol (Databricks-managed, external, or self-hosted) -### Built-in Tools - +**Built-in Tools:** - **system.ai.python_exec** - Execute Python code dynamically within agent queries (code interpreter) -### External Connection Tools - -Connect to external services via Unity Catalog HTTP connections: - -- **Slack** - Post messages to channels -- **Google Calendar** - Calendar operations -- **Microsoft Graph API** - Office 365 services -- **Azure AI Search** - Search functionality -- **Any HTTP API** - Use `http_request` from databricks-sdk - -Example: Create UC function wrapping HTTP request for Slack, then expose via MCP. - -### Common Patterns - +**Common Patterns:** - **Structured data retrieval** - Query SQL tables/databases - **Unstructured data retrieval** - Document search and RAG via Vector Search - **Code interpreter** - Python execution for analysis via system.ai.python_exec - **External connections** - Integrate services like Slack via HTTP connections ---- - -## Authentication Setup - -**Option 1: OAuth (Recommended)** - -```bash -databricks auth login -``` - -Set in `.env`: - -```bash -DATABRICKS_CONFIG_PROFILE=DEFAULT -``` - -**Option 2: Personal Access Token** - -Set in `.env`: - -```bash -DATABRICKS_HOST="https://host.databricks.com" -DATABRICKS_TOKEN="dapi_token" -``` - ---- - -## MLflow Experiment Setup - -Create and link an MLflow experiment: - -```bash -DATABRICKS_USERNAME=$(databricks current-user me | jq -r .userName) -databricks experiments create-experiment /Users/$DATABRICKS_USERNAME/agents-on-apps -``` - -Add the experiment ID to `.env`: - -```bash -MLFLOW_EXPERIMENT_ID= -``` - ---- - -## Key Files - -| File | Purpose | -| -------------------------------- | --------------------------------------------- | -| `agent_server/agent.py` | Agent logic, model, instructions, MCP servers | -| `agent_server/start_server.py` | FastAPI server + MLflow setup | -| `agent_server/evaluate_agent.py` | Agent evaluation with MLflow scorers | -| `agent_server/utils.py` | Databricks auth helpers, stream processing | -| `scripts/start_app.py` | Manages backend+frontend startup | - ---- - -## Deploying to Databricks Apps - -**Create app:** - -```bash -databricks apps create agent-langgraph -``` - -**Sync files:** - -```bash -DATABRICKS_USERNAME=$(databricks current-user me | jq -r .userName) -databricks sync . "/Users/$DATABRICKS_USERNAME/agent-langgraph" -``` - -**Deploy:** - -```bash -databricks apps deploy agent-langgraph --source-code-path /Workspace/Users/$DATABRICKS_USERNAME/agent-langgraph -``` - -**Query deployed app:** - -Generate OAuth token (PATs are not supported): - -```bash -databricks auth token -``` - -Send request: - -```bash -curl -X POST /invocations \ - -H "Authorization: Bearer " \ - -H "Content-Type: application/json" \ - -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' -``` +Reference: https://docs.databricks.com/aws/en/generative-ai/agent-framework/ diff --git a/agent-langgraph/README.md b/agent-langgraph/README.md index 25f7f94b..0f006940 100644 --- a/agent-langgraph/README.md +++ b/agent-langgraph/README.md @@ -6,6 +6,10 @@ The agent in this template implements the [OpenAI Responses API](https://platfor The agent input and output format are defined by MLflow's ResponsesAgent interface, which closely follows the [OpenAI Responses API](https://platform.openai.com/docs/api-reference/responses) interface. See [the MLflow docs](https://mlflow.org/docs/latest/genai/flavors/responses-agent-intro/) for input and output formats for streaming and non-streaming requests, tracing requirements, and other agent authoring details. +## Build with AI Assistance + +We recommend using AI coding assistants (Claude Code, Cursor, GitHub Copilot) to customize and deploy this template. Agent Skills in `.claude/skills/` provide step-by-step guidance for common tasks like setup, adding tools, and deployment. These skills are automatically detected by Claude, Cursor, and GitHub Copilot. + ## Quick start Run the `uv run quickstart` script to quickly set up your local environment and start the agent server. At any step, if there are issues, refer to the manual local development loop setup below. diff --git a/agent-langgraph/databricks.yml b/agent-langgraph/databricks.yml new file mode 100644 index 00000000..0287fb61 --- /dev/null +++ b/agent-langgraph/databricks.yml @@ -0,0 +1,37 @@ +bundle: + name: agent_langgraph + +resources: + # MLflow experiment for agent tracing - automatically created by bundle + experiments: + agent_langgraph_experiment: + name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} + + apps: + agent_langgraph: + name: "${bundle.target}-agent-langgraph" + description: "LangGraph agent application" + source_code_path: ./ + + # Resources which this app has access to + resources: + - name: 'experiment' + experiment: + experiment_id: "${resources.experiments.agent_langgraph_experiment.id}" + permission: 'CAN_MANAGE' + +targets: + dev: + mode: development + default: true + # workspace: + # host: https://... + + prod: + mode: production + # workspace: + # host: https://... + resources: + apps: + agent_langgraph: + name: agent-langgraph diff --git a/agent-langgraph/pyproject.toml b/agent-langgraph/pyproject.toml index 9b62817d..16c2651e 100644 --- a/agent-langgraph/pyproject.toml +++ b/agent-langgraph/pyproject.toml @@ -36,3 +36,4 @@ quickstart = "scripts.quickstart:main" start-app = "scripts.start_app:main" start-server = "agent_server.start_server:main" agent-evaluate = "agent_server.evaluate_agent:evaluate" +discover-tools = "scripts.discover_tools:main" diff --git a/agent-langgraph/scripts/discover_tools.py b/agent-langgraph/scripts/discover_tools.py new file mode 100644 index 00000000..3eb37963 --- /dev/null +++ b/agent-langgraph/scripts/discover_tools.py @@ -0,0 +1,432 @@ +#!/usr/bin/env python3 +""" +Discover available tools and data sources for Databricks agents. + +This script scans for: +- Unity Catalog functions (data retrieval tools e.g. SQL UDFs) +- Unity Catalog tables (data sources) +- Vector search indexes (RAG data sources) +- Genie spaces (conversational interface over structured data) +- Custom MCP servers (Databricks apps with name mcp-*) +- External MCP servers (via Unity Catalog connections) +""" + +import json +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict, List + +from databricks.sdk import WorkspaceClient + +DEFAULT_MAX_RESULTS = 100 +DEFAULT_MAX_SCHEMAS = 25 + +def run_databricks_cli(args: List[str]) -> str: + """Run databricks CLI command and return output.""" + try: + result = subprocess.run( + ["databricks"] + args, + capture_output=True, + text=True, + check=True, + ) + return result.stdout + except subprocess.CalledProcessError as e: + print(f"Error running databricks CLI: {e.stderr}", file=sys.stderr) + return "" + + +def discover_uc_functions(w: WorkspaceClient, catalog: str = None, max_schemas: int = DEFAULT_MAX_SCHEMAS) -> List[Dict[str, Any]]: + """Discover Unity Catalog functions that could be used as tools. + + Args: + w: WorkspaceClient instance + catalog: Optional specific catalog to search + max_schemas: Total number of schemas to search across all catalogs + """ + functions = [] + schemas_searched = 0 + + try: + catalogs = [catalog] if catalog else [c.name for c in w.catalogs.list()] + + for cat in catalogs: + if schemas_searched >= max_schemas: + break + + try: + all_schemas = list(w.schemas.list(catalog_name=cat)) + # Take schemas from this catalog until we hit the global budget + schemas_to_search = all_schemas[:max_schemas - schemas_searched] + + for schema in schemas_to_search: + schema_name = f"{cat}.{schema.name}" + try: + funcs = list(w.functions.list(catalog_name=cat, schema_name=schema.name)) + for func in funcs: + functions.append({ + "type": "uc_function", + "name": func.full_name, + "catalog": cat, + "schema": schema.name, + "function_name": func.name, + "comment": func.comment, + "routine_definition": getattr(func, "routine_definition", None), + }) + except Exception as e: + # Skip schemas we can't access + continue + finally: + schemas_searched += 1 + except Exception as e: + # Skip catalogs we can't access + continue + + except Exception as e: + print(f"Error discovering UC functions: {e}", file=sys.stderr) + + return functions + + +def discover_uc_tables(w: WorkspaceClient, catalog: str = None, schema: str = None, max_schemas: int = DEFAULT_MAX_SCHEMAS) -> List[Dict[str, Any]]: + """Discover Unity Catalog tables that could be queried. + + Args: + w: WorkspaceClient instance + catalog: Optional specific catalog to search + schema: Optional specific schema to search (requires catalog) + max_schemas: Total number of schemas to search across all catalogs + """ + tables = [] + schemas_searched = 0 + + try: + catalogs = [catalog] if catalog else [c.name for c in w.catalogs.list()] + + for cat in catalogs: + if cat in ["__databricks_internal", "system"]: + continue + + if schemas_searched >= max_schemas: + break + + try: + if schema: + schemas_to_search = [schema] + else: + all_schemas = [s.name for s in w.schemas.list(catalog_name=cat)] + # Take schemas from this catalog until we hit the global budget + schemas_to_search = all_schemas[:max_schemas - schemas_searched] + + for sch in schemas_to_search: + if sch == "information_schema": + schemas_searched += 1 + continue + + try: + tbls = list(w.tables.list(catalog_name=cat, schema_name=sch)) + for tbl in tbls: + # Get column info + columns = [] + if hasattr(tbl, "columns") and tbl.columns: + columns = [ + {"name": col.name, "type": col.type_name.value if hasattr(col.type_name, "value") else str(col.type_name)} + for col in tbl.columns + ] + + tables.append({ + "type": "uc_table", + "name": tbl.full_name, + "catalog": cat, + "schema": sch, + "table_name": tbl.name, + "table_type": tbl.table_type.value if tbl.table_type else None, + "comment": tbl.comment, + "columns": columns, + }) + except Exception as e: + # Skip schemas we can't access + pass + finally: + schemas_searched += 1 + except Exception as e: + # Skip catalogs we can't access + continue + + except Exception as e: + print(f"Error discovering UC tables: {e}", file=sys.stderr) + + return tables + + +def discover_vector_search_indexes(w: WorkspaceClient) -> List[Dict[str, Any]]: + """Discover Vector Search indexes for RAG applications.""" + indexes = [] + + try: + # List all vector search endpoints + endpoints = list(w.vector_search_endpoints.list_endpoints()) + + for endpoint in endpoints: + try: + # List indexes for each endpoint + endpoint_indexes = list(w.vector_search_indexes.list_indexes(endpoint_name=endpoint.name)) + for idx in endpoint_indexes: + indexes.append({ + "type": "vector_search_index", + "name": idx.name, + "endpoint": endpoint.name, + "primary_key": idx.primary_key, + "index_type": idx.index_type.value if idx.index_type else None, + "status": idx.status.state.value if idx.status and idx.status.state else None, + }) + except Exception as e: + # Skip endpoints we can't access + continue + + except Exception as e: + print(f"Error discovering vector search indexes: {e}", file=sys.stderr) + + return indexes + + +def discover_genie_spaces(w: WorkspaceClient) -> List[Dict[str, Any]]: + """Discover Genie spaces for conversational data access.""" + spaces = [] + + try: + # Use SDK to list genie spaces + response = w.genie.list_spaces() + genie_spaces = response.spaces if hasattr(response, "spaces") else [] + for space in genie_spaces: + spaces.append({ + "type": "genie_space", + "id": space.space_id, + "name": space.title, + "description": space.description, + }) + except Exception as e: + print(f"Error discovering Genie spaces: {e}", file=sys.stderr) + + return spaces + + + +def discover_custom_mcp_servers(w: WorkspaceClient) -> List[Dict[str, Any]]: + """Discover custom MCP servers deployed as Databricks apps.""" + custom_servers = [] + + try: + # List all apps and filter for those starting with mcp- + apps = w.apps.list() + for app in apps: + if app.name and app.name.startswith("mcp-"): + custom_servers.append({ + "type": "custom_mcp_server", + "name": app.name, + "url": app.url, + "status": app.app_status.state.value if app.app_status and app.app_status.state else None, + "description": app.description, + }) + except Exception as e: + print(f"Error discovering custom MCP servers: {e}", file=sys.stderr) + + return custom_servers + + +def discover_external_mcp_servers(w: WorkspaceClient) -> List[Dict[str, Any]]: + """Discover external MCP servers configured via Unity Catalog connections.""" + external_servers = [] + + try: + # List all connections and filter for MCP connections + connections = w.connections.list() + for conn in connections: + # Check if this is an MCP connection + if conn.options and conn.options.get("is_mcp_connection") == "true": + external_servers.append({ + "type": "external_mcp_server", + "name": conn.name, + "connection_type": conn.connection_type.value if hasattr(conn.connection_type, "value") else str(conn.connection_type), + "comment": conn.comment, + "full_name": conn.full_name, + }) + except Exception as e: + print(f"Error discovering external MCP servers: {e}", file=sys.stderr) + + return external_servers + + +def format_output_markdown(results: Dict[str, List[Dict[str, Any]]]) -> str: + """Format discovery results as markdown.""" + lines = ["# Agent Tools and Data Sources Discovery\n"] + + # UC Functions + functions = results.get("uc_functions", []) + if functions: + lines.append(f"## Unity Catalog Functions ({len(functions)})\n") + lines.append("**What they are:** SQL UDFs that can be used as agent tools.\n") + lines.append("**How to use:** Access via UC functions MCP server:") + lines.append("- All functions in a schema: `{workspace_host}/api/2.0/mcp/functions/{catalog}/{schema}`") + lines.append("- Single function: `{workspace_host}/api/2.0/mcp/functions/{catalog}/{schema}/{function_name}`\n") + for func in functions[:10]: # Show first 10 + lines.append(f"- `{func['name']}`") + if func.get("comment"): + lines.append(f" - {func['comment']}") + if len(functions) > 10: + lines.append(f"\n*...and {len(functions) - 10} more*\n") + lines.append("") + + # UC Tables + tables = results.get("uc_tables", []) + if tables: + lines.append(f"## Unity Catalog Tables ({len(tables)})\n") + lines.append("Structured data that agents can query via UC SQL functions.\n") + for table in tables[:10]: # Show first 10 + lines.append(f"- `{table['name']}` ({table['table_type']})") + if table.get("comment"): + lines.append(f" - {table['comment']}") + if table.get("columns"): + col_names = [c["name"] for c in table["columns"][:5]] + lines.append(f" - Columns: {', '.join(col_names)}") + if len(tables) > 10: + lines.append(f"\n*...and {len(tables) - 10} more*\n") + lines.append("") + + # Vector Search Indexes + indexes = results.get("vector_search_indexes", []) + if indexes: + lines.append(f"## Vector Search Indexes ({len(indexes)})\n") + lines.append("These can be used for RAG applications with unstructured data.\n") + lines.append("**How to use:** Connect via MCP server at `{workspace_host}/api/2.0/mcp/vector-search/{catalog}/{schema}` or\n") + lines.append("`{workspace_host}/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name}`\n") + for idx in indexes: + lines.append(f"- `{idx['name']}`") + lines.append(f" - Endpoint: {idx['endpoint']}") + lines.append(f" - Status: {idx['status']}") + lines.append("") + + # Genie Spaces + spaces = results.get("genie_spaces", []) + if spaces: + lines.append(f"## Genie Spaces ({len(spaces)})\n") + lines.append("**What they are:** Natural language interface to your data\n") + lines.append("**How to use:** Connect via Genie MCP server at `{workspace_host}/api/2.0/mcp/genie/{space_id}`\n") + for space in spaces: + lines.append(f"- `{space['name']}` (ID: {space['id']})") + if space.get("description"): + lines.append(f" - {space['description']}") + lines.append("") + + # Custom MCP Servers (Databricks Apps) + custom_servers = results.get("custom_mcp_servers", []) + if custom_servers: + lines.append(f"## Custom MCP Servers ({len(custom_servers)})\n") + lines.append("**What:** Your own MCP servers deployed as Databricks Apps (names starting with mcp-)\n") + lines.append("**How to use:** Access via `{app_url}/mcp`\n") + lines.append("**⚠️ Important:** Custom MCP server apps require manual permission grants:") + lines.append("1. Get your agent app's service principal: `databricks apps get --output json | jq -r '.service_principal_name'`") + lines.append("2. Grant permission: `databricks apps update-permissions --service-principal --permission-level CAN_USE`") + lines.append("(Apps are not yet supported as resource dependencies in databricks.yml)\n") + for server in custom_servers: + lines.append(f"- `{server['name']}`") + if server.get("url"): + lines.append(f" - URL: {server['url']}") + if server.get("status"): + lines.append(f" - Status: {server['status']}") + if server.get("description"): + lines.append(f" - {server['description']}") + lines.append("") + + # External MCP Servers (UC Connections) + external_servers = results.get("external_mcp_servers", []) + if external_servers: + lines.append(f"## External MCP Servers ({len(external_servers)})\n") + lines.append("**What:** Third-party MCP servers via Unity Catalog connections\n") + lines.append("**How to use:** Connect via `{workspace_host}/api/2.0/mcp/external/{connection_name}`\n") + lines.append("**Benefits:** Secure access to external APIs through UC governance\n") + for server in external_servers: + lines.append(f"- `{server['name']}`") + if server.get("full_name"): + lines.append(f" - Full name: {server['full_name']}") + if server.get("comment"): + lines.append(f" - {server['comment']}") + lines.append("") + return "\n".join(lines) + + +def main(): + """Main discovery function.""" + import argparse + + parser = argparse.ArgumentParser(description="Discover available agent tools and data sources") + parser.add_argument("--catalog", help="Limit discovery to specific catalog") + parser.add_argument("--schema", help="Limit discovery to specific schema (requires --catalog)") + parser.add_argument("--format", choices=["json", "markdown"], default="markdown", help="Output format") + parser.add_argument("--output", help="Output file (default: stdout)") + parser.add_argument("--profile", help="Databricks CLI profile to use (default: uses default profile)") + parser.add_argument("--max-results", type=int, default=DEFAULT_MAX_RESULTS, help=f"Maximum results per resource type (default: {DEFAULT_MAX_RESULTS})") + parser.add_argument("--max-schemas", type=int, default=DEFAULT_MAX_SCHEMAS, help=f"Total schemas to search across all catalogs (default: {DEFAULT_MAX_SCHEMAS})") + + args = parser.parse_args() + + if args.schema and not args.catalog: + print("Error: --schema requires --catalog", file=sys.stderr) + sys.exit(1) + + print("Discovering available tools and data sources...", file=sys.stderr) + + # Initialize Databricks workspace client + # Only pass profile if specified, otherwise use default + if args.profile: + w = WorkspaceClient(profile=args.profile) + else: + w = WorkspaceClient() + + results = {} + + # Discover each type with configurable limits + print("- UC Functions...", file=sys.stderr) + results["uc_functions"] = discover_uc_functions(w, catalog=args.catalog, max_schemas=args.max_schemas)[:args.max_results] + + print("- UC Tables...", file=sys.stderr) + results["uc_tables"] = discover_uc_tables(w, catalog=args.catalog, schema=args.schema, max_schemas=args.max_schemas)[:args.max_results] + + print("- Vector Search Indexes...", file=sys.stderr) + results["vector_search_indexes"] = discover_vector_search_indexes(w)[:args.max_results] + + print("- Genie Spaces...", file=sys.stderr) + results["genie_spaces"] = discover_genie_spaces(w)[:args.max_results] + + print("- Custom MCP Servers (Apps)...", file=sys.stderr) + results["custom_mcp_servers"] = discover_custom_mcp_servers(w)[:args.max_results] + + print("- External MCP Servers (Connections)...", file=sys.stderr) + results["external_mcp_servers"] = discover_external_mcp_servers(w)[:args.max_results] + + # Format output + if args.format == "json": + output = json.dumps(results, indent=2) + else: + output = format_output_markdown(results) + + # Write output + if args.output: + Path(args.output).write_text(output) + print(f"\nResults written to {args.output}", file=sys.stderr) + else: + print("\n" + output) + + # Print summary + print("\n=== Discovery Summary ===", file=sys.stderr) + print(f"UC Functions: {len(results['uc_functions'])}", file=sys.stderr) + print(f"UC Tables: {len(results['uc_tables'])}", file=sys.stderr) + print(f"Vector Search Indexes: {len(results['vector_search_indexes'])}", file=sys.stderr) + print(f"Genie Spaces: {len(results['genie_spaces'])}", file=sys.stderr) + print(f"Custom MCP Servers: {len(results['custom_mcp_servers'])}", file=sys.stderr) + print(f"External MCP Servers: {len(results['external_mcp_servers'])}", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/agent-langgraph/scripts/start_app.py b/agent-langgraph/scripts/start_app.py index 46997762..9fe60cde 100644 --- a/agent-langgraph/scripts/start_app.py +++ b/agent-langgraph/scripts/start_app.py @@ -6,8 +6,16 @@ 1. Not reporting ready until BOTH frontend and backend processes are ready 2. Exiting as soon as EITHER process fails 3. Printing error logs if either process fails + +Usage: + start-app [OPTIONS] + +All options are passed through to the backend server (start-server). +See 'uv run start-server --help' for available options. """ +import argparse +import os import re import shutil import subprocess @@ -24,7 +32,7 @@ class ProcessManager: - def __init__(self): + def __init__(self, port=8000): self.backend_process = None self.frontend_process = None self.backend_ready = False @@ -32,6 +40,7 @@ def __init__(self): self.failed = threading.Event() self.backend_log = None self.frontend_log = None + self.port = port def monitor_process(self, process, name, log_file, patterns): is_ready = False @@ -56,7 +65,7 @@ def monitor_process(self, process, name, log_file, patterns): if self.backend_ready and self.frontend_ready: print("\n" + "=" * 50) print("✓ Both frontend and backend are ready!") - print("✓ Open the frontend at http://localhost:8000") + print(f"✓ Open the frontend at http://localhost:{self.port}") print("=" * 50 + "\n") process.wait() @@ -141,20 +150,28 @@ def cleanup(self): if self.frontend_log: self.frontend_log.close() - def run(self): + def run(self, backend_args=None): load_dotenv(dotenv_path=".env", override=True) if not self.clone_frontend_if_needed(): return 1 + # Set API_PROXY environment variable for frontend to connect to backend + os.environ["API_PROXY"] = f"http://localhost:{self.port}/invocations" + # Open log files self.backend_log = open("backend.log", "w", buffering=1) self.frontend_log = open("frontend.log", "w", buffering=1) try: + # Build backend command, passing through all arguments + backend_cmd = ["uv", "run", "start-server"] + if backend_args: + backend_cmd.extend(backend_args) + # Start backend self.backend_process = self.start_process( - ["uv", "run", "start-server"], "backend", self.backend_log, BACKEND_READY + backend_cmd, "backend", self.backend_log, BACKEND_READY ) # Setup and start frontend @@ -211,7 +228,25 @@ def run(self): def main(): - sys.exit(ProcessManager().run()) + parser = argparse.ArgumentParser( + description="Start agent frontend and backend", + usage="%(prog)s [OPTIONS]\n\nAll options are passed through to start-server. " + "Use 'uv run start-server --help' for available options." + ) + # Parse known args (none currently) and pass remaining to backend + _, backend_args = parser.parse_known_args() + + # Extract port from backend_args if specified + port = 8000 + for i, arg in enumerate(backend_args): + if arg == "--port" and i + 1 < len(backend_args): + try: + port = int(backend_args[i + 1]) + except ValueError: + pass + break + + sys.exit(ProcessManager(port=port).run(backend_args)) if __name__ == "__main__": From 671804932bd409c2a74f646128b8d88bde9edd9b Mon Sep 17 00:00:00 2001 From: Dhruv Gupta Date: Tue, 27 Jan 2026 11:53:52 -0800 Subject: [PATCH 02/14] Agent langgraph modify agent skill update --- .../.claude/skills/modify-agent/SKILL.md | 46 ++++++++++++++----- 1 file changed, 34 insertions(+), 12 deletions(-) diff --git a/agent-langgraph/.claude/skills/modify-agent/SKILL.md b/agent-langgraph/.claude/skills/modify-agent/SKILL.md index b42a77d4..eec80096 100644 --- a/agent-langgraph/.claude/skills/modify-agent/SKILL.md +++ b/agent-langgraph/.claude/skills/modify-agent/SKILL.md @@ -188,7 +188,7 @@ server = MCPServer( ```python from langchain.agents import create_agent -# Create agent +# Create agent - ONLY accepts tools and model, NO prompt/instructions parameter agent = create_agent(tools=tools, model=llm) # Non-streaming @@ -214,25 +214,47 @@ async for event in process_agent_astream_events( --- -## Customizing Agent Behavior +## Customizing Agent Behavior (System Instructions) -In LangGraph, agent behavior is customized via system messages and agent configuration rather than an `instructions` parameter. +> **IMPORTANT:** `create_agent()` does NOT accept `prompt`, `instructions`, or `system_message` parameters. Attempting to pass these will cause a runtime error. -**Add system message to conversation:** +In LangGraph, agent behavior is customized by prepending a system message to the conversation messages. + +**Correct pattern in `agent.py`:** + +1. Define instructions as a constant: ```python -messages = { - "messages": [ - {"role": "system", "content": """You are a helpful data analyst assistant. +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."""}, - {"role": "user", "content": "What were Q4 sales?"} - ] -} -result = await agent.ainvoke(messages) +Always cite your sources when answering questions.""" +``` + +2. Prepend to messages in the `streaming()` function: +```python +@stream() +async def streaming(request: ResponsesAgentRequest) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + agent = await init_agent() + # Prepend system instructions to user messages + user_messages = to_chat_completions_input([i.model_dump() for i in request.input]) + messages = {"messages": [{"role": "system", "content": AGENT_INSTRUCTIONS}] + user_messages} + + async for event in process_agent_astream_events( + agent.astream(input=messages, stream_mode=["updates", "messages"]) + ): + yield event +``` + +**Common mistake to avoid:** +```python +# WRONG - will cause "unexpected keyword argument" error +agent = create_agent(tools=tools, model=llm, prompt=AGENT_INSTRUCTIONS) + +# CORRECT - add instructions via messages +messages = {"messages": [{"role": "system", "content": AGENT_INSTRUCTIONS}] + user_messages} ``` For advanced customization (routing, state management, custom graphs), refer to the [LangGraph documentation](https://docs.langchain.com/oss/python/langgraph/overview). From c681775107574cb28f878dbc68a65af3e4d0298e Mon Sep 17 00:00:00 2001 From: Dhruv Gupta Date: Wed, 28 Jan 2026 10:28:22 -0800 Subject: [PATCH 03/14] Added lakebase template agent skills --- .../.claude/skills/add-tools/SKILL.md | 88 ++++ .../add-tools/examples/custom-mcp-server.md | 57 +++ .../skills/add-tools/examples/experiment.yaml | 8 + .../add-tools/examples/genie-space.yaml | 9 + .../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/agent-memory/SKILL.md | 220 +++++++++ .../.claude/skills/deploy/SKILL.md | 216 +++++++++ .../.claude/skills/discover-tools/SKILL.md | 82 ++++ .../.claude/skills/lakebase-setup/SKILL.md | 169 +++++++ .../.claude/skills/modify-agent/SKILL.md | 296 ++++++++++++ .../.claude/skills/quickstart/SKILL.md | 84 ++++ .../.claude/skills/run-locally/SKILL.md | 130 ++++++ agent-langgraph-long-term-memory/.gitignore | 15 +- agent-langgraph-long-term-memory/AGENTS.md | 365 +++------------ agent-langgraph-long-term-memory/README.md | 18 +- .../pyproject.toml | 1 + .../scripts/discover_tools.py | 432 ++++++++++++++++++ .../.claude/skills/add-tools/SKILL.md | 88 ++++ .../add-tools/examples/custom-mcp-server.md | 57 +++ .../skills/add-tools/examples/experiment.yaml | 8 + .../add-tools/examples/genie-space.yaml | 9 + .../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/agent-memory/SKILL.md | 196 ++++++++ .../.claude/skills/deploy/SKILL.md | 205 +++++++++ .../.claude/skills/discover-tools/SKILL.md | 82 ++++ .../.claude/skills/lakebase-setup/SKILL.md | 168 +++++++ .../.claude/skills/modify-agent/SKILL.md | 296 ++++++++++++ .../.claude/skills/quickstart/SKILL.md | 84 ++++ .../.claude/skills/run-locally/SKILL.md | 130 ++++++ agent-langgraph-short-term-memory/.gitignore | 16 +- agent-langgraph-short-term-memory/AGENTS.md | 365 +++------------ agent-langgraph-short-term-memory/README.md | 18 +- .../pyproject.toml | 1 + .../scripts/discover_tools.py | 432 ++++++++++++++++++ 42 files changed, 3816 insertions(+), 611 deletions(-) create mode 100644 agent-langgraph-long-term-memory/.claude/skills/add-tools/SKILL.md create mode 100644 agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md create mode 100644 agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/experiment.yaml create mode 100644 agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/genie-space.yaml create mode 100644 agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/serving-endpoint.yaml create mode 100644 agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/sql-warehouse.yaml create mode 100644 agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/uc-connection.yaml create mode 100644 agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/uc-function.yaml create mode 100644 agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/vector-search.yaml create mode 100644 agent-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md create mode 100644 agent-langgraph-long-term-memory/.claude/skills/deploy/SKILL.md create mode 100644 agent-langgraph-long-term-memory/.claude/skills/discover-tools/SKILL.md create mode 100644 agent-langgraph-long-term-memory/.claude/skills/lakebase-setup/SKILL.md create mode 100644 agent-langgraph-long-term-memory/.claude/skills/modify-agent/SKILL.md create mode 100644 agent-langgraph-long-term-memory/.claude/skills/quickstart/SKILL.md create mode 100644 agent-langgraph-long-term-memory/.claude/skills/run-locally/SKILL.md create mode 100644 agent-langgraph-long-term-memory/scripts/discover_tools.py create mode 100644 agent-langgraph-short-term-memory/.claude/skills/add-tools/SKILL.md create mode 100644 agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md create mode 100644 agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/experiment.yaml create mode 100644 agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/genie-space.yaml create mode 100644 agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/serving-endpoint.yaml create mode 100644 agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/sql-warehouse.yaml create mode 100644 agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/uc-connection.yaml create mode 100644 agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/uc-function.yaml create mode 100644 agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/vector-search.yaml create mode 100644 agent-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md create mode 100644 agent-langgraph-short-term-memory/.claude/skills/deploy/SKILL.md create mode 100644 agent-langgraph-short-term-memory/.claude/skills/discover-tools/SKILL.md create mode 100644 agent-langgraph-short-term-memory/.claude/skills/lakebase-setup/SKILL.md create mode 100644 agent-langgraph-short-term-memory/.claude/skills/modify-agent/SKILL.md create mode 100644 agent-langgraph-short-term-memory/.claude/skills/quickstart/SKILL.md create mode 100644 agent-langgraph-short-term-memory/.claude/skills/run-locally/SKILL.md create mode 100644 agent-langgraph-short-term-memory/scripts/discover_tools.py diff --git a/agent-langgraph-long-term-memory/.claude/skills/add-tools/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/add-tools/SKILL.md new file mode 100644 index 00000000..9c1c9101 --- /dev/null +++ b/agent-langgraph-long-term-memory/.claude/skills/add-tools/SKILL.md @@ -0,0 +1,88 @@ +--- +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 + +**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. + +## Workflow + +**Step 1:** Add MCP server in `agent_server/agent.py`: +```python +from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient + +genie_server = DatabricksMCPServer( + url=f"{host}/api/2.0/mcp/genie/01234567-89ab-cdef", + name="my genie space", +) + +mcp_client = DatabricksMultiServerMCPClient([genie_server]) +tools = await mcp_client.get_tools() +``` + +**Step 2:** Grant access in `databricks.yml`: +```yaml +resources: + apps: + agent_langgraph_long_term_memory: + resources: + - name: 'my_genie_space' + genie_space: + name: 'My Genie Space' + space_id: '01234567-89ab-cdef' + permission: 'CAN_RUN' +``` + +**Step 3:** Deploy with `databricks bundle deploy` (see **deploy** skill) + +## 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 | +| `experiment.yaml` | MLflow experiment | Tracing (already configured) | +| `custom-mcp-server.md` | Custom MCP apps | Apps starting with `mcp-*` | + +## Custom MCP Servers (Databricks Apps) + +Apps are **not yet supported** as resource dependencies in `databricks.yml`. Manual permission grant required: + +**Step 1:** Get your agent app's service principal: +```bash +databricks apps get --output json | jq -r '.service_principal_name' +``` + +**Step 2:** Grant permission on the MCP server app: +```bash +databricks apps update-permissions \ + --service-principal \ + --permission-level CAN_USE +``` + +See `examples/custom-mcp-server.md` for detailed steps. + +## Important Notes + +- **MLflow experiment**: Already configured in template, no action needed +- **Lakebase permissions**: Memory storage requires separate Lakebase setup (see **lakebase-setup** skill) +- **Multiple resources**: Add multiple entries under `resources:` list +- **Permission types vary**: Each resource type has specific permission values +- **Deploy after changes**: Run `databricks bundle deploy` after modifying `databricks.yml` + +## Next Steps + +- Configure memory storage: see **lakebase-setup** skill +- Understand memory patterns: see **agent-memory** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md b/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md new file mode 100644 index 00000000..4a056c01 --- /dev/null +++ b/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md @@ -0,0 +1,57 @@ +# Custom MCP Server (Databricks App) + +Custom MCP servers are Databricks Apps with names starting with `mcp-*`. + +**Apps are not yet supported as resource dependencies in `databricks.yml`**, so manual permission grant is required. + +## 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. Deploy your agent app first + +```bash +databricks bundle deploy +databricks bundle run agent_langgraph_long_term_memory +``` + +### 3. Get your agent app's service principal + +```bash +databricks apps get --output json | jq -r '.service_principal_name' +``` + +Example output: `sp-abc123-def456` + +### 4. Grant permission on the MCP server app + +```bash +databricks apps update-permissions \ + --service-principal \ + --permission-level CAN_USE +``` + +Example: +```bash +databricks apps update-permissions mcp-my-server \ + --service-principal sp-abc123-def456 \ + --permission-level CAN_USE +``` + +## Notes + +- This manual step is required each time you connect to a new custom MCP server +- The permission grant persists across deployments +- If you redeploy the agent app with a new service principal, you'll need to grant permissions again diff --git a/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/experiment.yaml b/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/experiment.yaml new file mode 100644 index 00000000..ac5c626a --- /dev/null +++ b/agent-langgraph-long-term-memory/.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-langgraph-long-term-memory/.claude/skills/add-tools/examples/genie-space.yaml b/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/genie-space.yaml new file mode 100644 index 00000000..71589d52 --- /dev/null +++ b/agent-langgraph-long-term-memory/.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-langgraph-long-term-memory/.claude/skills/add-tools/examples/serving-endpoint.yaml b/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/serving-endpoint.yaml new file mode 100644 index 00000000..b49ce9da --- /dev/null +++ b/agent-langgraph-long-term-memory/.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-langgraph-long-term-memory/.claude/skills/add-tools/examples/sql-warehouse.yaml b/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/sql-warehouse.yaml new file mode 100644 index 00000000..a6ce9446 --- /dev/null +++ b/agent-langgraph-long-term-memory/.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-langgraph-long-term-memory/.claude/skills/add-tools/examples/uc-connection.yaml b/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/uc-connection.yaml new file mode 100644 index 00000000..316675fe --- /dev/null +++ b/agent-langgraph-long-term-memory/.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-langgraph-long-term-memory/.claude/skills/add-tools/examples/uc-function.yaml b/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/uc-function.yaml new file mode 100644 index 00000000..43f938a9 --- /dev/null +++ b/agent-langgraph-long-term-memory/.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-langgraph-long-term-memory/.claude/skills/add-tools/examples/vector-search.yaml b/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/vector-search.yaml new file mode 100644 index 00000000..0ba39027 --- /dev/null +++ b/agent-langgraph-long-term-memory/.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-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md new file mode 100644 index 00000000..b8f15aa1 --- /dev/null +++ b/agent-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md @@ -0,0 +1,220 @@ +--- +name: agent-memory +description: "Understand and modify agent memory patterns. Use when: (1) User asks about 'memory', 'state', 'user preferences', (2) Working with user_id or memory tools, (3) Debugging memory issues, (4) Adding short-term memory capabilities." +--- + +# Agent Memory Patterns + +This skill covers both long-term memory (facts that persist across sessions) and short-term memory (conversation history within a session). + +## Long-Term Memory (This Template) + +Long-term memory stores facts about users that persist across conversation sessions. The agent can remember preferences, facts, and information across multiple interactions. + +### Key Components + +**AsyncDatabricksStore** - Persists user memories to Lakebase with semantic search: + +```python +from databricks_langchain import AsyncDatabricksStore + +async with AsyncDatabricksStore( + instance_name=LAKEBASE_INSTANCE_NAME, + embedding_endpoint=EMBEDDING_ENDPOINT, + embedding_dims=EMBEDDING_DIMS, +) as store: + agent = await init_agent(store=store) + # Agent now has access to persistent memory store +``` + +**User ID** - Identifies a user across sessions: + +```python +def _get_user_id(request: ResponsesAgentRequest) -> Optional[str]: + custom_inputs = dict(request.custom_inputs or {}) + if "user_id" in custom_inputs: + return custom_inputs["user_id"] + if request.context and getattr(request.context, "user_id", None): + return request.context.user_id + return None +``` + +### Memory Tools + +The template includes three memory tools that the agent can use: + +**get_user_memory** - Search for relevant information about the user: + +```python +@tool +async def get_user_memory(query: str, config: RunnableConfig) -> str: + """Search for relevant information about the user from long-term memory.""" + user_id = config.get("configurable", {}).get("user_id") + store = config.get("configurable", {}).get("store") + + namespace = ("user_memories", user_id.replace(".", "-")) + results = await store.asearch(namespace, query=query, limit=5) + # Returns formatted memory items +``` + +**save_user_memory** - Save information about the user: + +```python +@tool +async def save_user_memory(memory_key: str, memory_data_json: str, config: RunnableConfig) -> str: + """Save information about the user to long-term memory.""" + # memory_data_json must be a valid JSON object string + await store.aput(namespace, memory_key, memory_data) +``` + +**delete_user_memory** - Remove specific memories: + +```python +@tool +async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: + """Delete a specific memory from the user's long-term memory.""" + await store.adelete(namespace, memory_key) +``` + +### Creating a Long-Term Memory Agent + +Pass `store` to `create_agent()`: + +```python +from langchain.agents import create_agent + +agent = create_agent( + model=ChatDatabricks(endpoint=LLM_ENDPOINT_NAME), + tools=mcp_tools + init_memory_tools(), # Include memory tools + system_prompt=SYSTEM_PROMPT, + store=store, # Enables long-term memory +) +``` + +### System Prompt for Memory + +The template includes instructions for using memory tools: + +```python +SYSTEM_PROMPT = """You are a helpful assistant. + +You have access to memory tools that allow you to remember information about users: +- Use get_user_memory to search for previously saved information about the user +- Use save_user_memory to remember important facts, preferences, or details the user shares +- Use delete_user_memory to forget specific information when asked + +Always check for relevant memories at the start of a conversation to provide personalized responses.""" +``` + +--- + +## API Requests with Long-Term Memory + +### Request with user_id + +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "My favorite color is blue, remember that"}], + "custom_inputs": {"user_id": "user@example.com"} + }' +``` + +### Query stored memories + +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "What do you remember about me?"}], + "custom_inputs": {"user_id": "user@example.com"} + }' +``` + +### Using context instead of custom_inputs + +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "What is my favorite color?"}], + "context": {"user_id": "user@example.com"} + }' +``` + +### Deployed App with Memory + +```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"} + }' +``` + +--- + +## Short-Term Memory (Reference) + +Short-term memory stores conversation history within a thread. This is available in the `agent-langgraph-short-term-memory` template. + +### Key Components (Short-Term) + +**AsyncCheckpointSaver** - Persists LangGraph state to Lakebase: + +```python +from databricks_langchain import AsyncCheckpointSaver + +async with AsyncCheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME) as checkpointer: + agent = await init_agent(checkpointer=checkpointer) +``` + +**Thread ID** - Identifies a conversation thread: + +```python +thread_id = request.custom_inputs.get("thread_id") or str(uuid_utils.uuid7()) +config = {"configurable": {"thread_id": thread_id}} +``` + +### Short-Term Memory Request Example + +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "What did we discuss earlier?"}], + "custom_inputs": {"thread_id": ""} + }' +``` + +--- + +## Troubleshooting + +| Issue | Solution | +|-------|----------| +| **Memories not persisting** | Verify `user_id` is passed consistently | +| **"Cannot connect to Lakebase"** | See **lakebase-setup** skill | +| **Permission errors** | Grant store table permissions (see **lakebase-setup** skill) | +| **Memory not available warning** | Ensure `user_id` is provided in request | +| **No memories found** | User hasn't saved any memories yet | + +## Dependencies + +Long-term memory requires `databricks-langchain[memory]`: + +```toml +# pyproject.toml +dependencies = [ + "databricks-langchain[memory]>=0.13.0", +] +``` + +## Next Steps + +- Configure Lakebase: see **lakebase-setup** skill +- Modify agent behavior: see **modify-agent** skill +- Test locally: see **run-locally** skill diff --git a/agent-langgraph-long-term-memory/.claude/skills/deploy/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/deploy/SKILL.md new file mode 100644 index 00000000..118f2a59 --- /dev/null +++ b/agent-langgraph-long-term-memory/.claude/skills/deploy/SKILL.md @@ -0,0 +1,216 @@ +--- +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 + +> **Memory Template:** Before deploying, ensure Lakebase is configured. See the **lakebase-setup** skill for permissions setup required after deployment. + +## Deploy Commands + +```bash +# Deploy the bundle (creates/updates resources, uploads files) +databricks bundle deploy + +# Run the app (starts/restarts with uploaded source code) +databricks bundle run agent_langgraph_long_term_memory +``` + +The resource key `agent_langgraph_long_term_memory` 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_langgraph_long_term_memory: + 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_langgraph_long_term_memory --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_langgraph_long_term_memory +``` + +### 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. + +## Post-Deployment: Lakebase Permissions + +**After deploying a memory-enabled agent, you must grant Lakebase permissions.** + +See the **lakebase-setup** skill for: +- Adding Lakebase as an app resource +- SDK or SQL commands to grant store table permissions +- Troubleshooting access errors + +## Unbinding an App + +To remove the link between bundle and deployed app: + +```bash +databricks bundle deployment unbind agent_langgraph_long_term_memory +``` + +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 +``` + +**Send request:** +```bash +curl -X POST /invocations \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' +``` + +**Request with user_id (for long-term memory):** +```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"} + }' +``` + +For more memory-specific request examples, see the **agent-memory** skill. + +## 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' +``` + +## 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 | +|-------|----------| +| Permission errors at runtime | Grant resources in `databricks.yml` (see **add-tools** skill) | +| Lakebase access errors | See **lakebase-setup** skill for permissions | +| 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 | + +## Next Steps + +- Grant Lakebase permissions: see **lakebase-setup** skill +- Understand memory patterns: see **agent-memory** skill +- Add tools and permissions: see **add-tools** skill diff --git a/agent-langgraph-long-term-memory/.claude/skills/discover-tools/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/discover-tools/SKILL.md new file mode 100644 index 00000000..eabac2a5 --- /dev/null +++ b/agent-langgraph-long-term-memory/.claude/skills/discover-tools/SKILL.md @@ -0,0 +1,82 @@ +--- +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}` | + +## Using Discovered Tools in Code + +After discovering tools, add them to your agent in `agent_server/agent.py`: + +```python +from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient + +# Example: Add UC functions from a schema +uc_functions_server = DatabricksMCPServer( + url=f"{host}/api/2.0/mcp/functions/{catalog}/{schema}", + name="my uc functions", +) + +# Example: Add a Genie space +genie_server = DatabricksMCPServer( + url=f"{host}/api/2.0/mcp/genie/{space_id}", + name="my genie space", +) + +# Example: Add vector search +vector_server = DatabricksMCPServer( + url=f"{host}/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name}", + name="my vector index", +) + +# Create MCP client with all servers +mcp_client = DatabricksMultiServerMCPClient([ + uc_functions_server, + genie_server, + vector_server, +]) + +# Get tools for the agent +tools = await mcp_client.get_tools() +``` + +## Next Steps + +After adding MCP servers to your agent: +1. **Grant permissions** in `databricks.yml` (see **add-tools** skill) +2. Test locally with `uv run start-app` (see **run-locally** skill) diff --git a/agent-langgraph-long-term-memory/.claude/skills/lakebase-setup/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/lakebase-setup/SKILL.md new file mode 100644 index 00000000..30d7825f --- /dev/null +++ b/agent-langgraph-long-term-memory/.claude/skills/lakebase-setup/SKILL.md @@ -0,0 +1,169 @@ +--- +name: lakebase-setup +description: "Configure Lakebase for agent memory storage. Use when: (1) First-time memory setup, (2) 'Failed to connect to Lakebase' errors, (3) Permission errors on store tables, (4) User says 'lakebase', 'memory setup', or 'store'." +--- + +# Lakebase Setup for Memory + +This template uses Lakebase (Databricks-managed PostgreSQL) to store long-term user memories. You must configure Lakebase before the agent can persist memories. + +## Prerequisites + +- A Lakebase instance in your Databricks workspace +- The instance name (found in Databricks UI under SQL Warehouses > Lakebase) + +## Local Development Setup + +**Step 1:** Add the Lakebase instance name to `.env.local`: + +```bash +LAKEBASE_INSTANCE_NAME= +``` + +**Step 2:** Run `uv run start-app` to test locally. The agent will automatically create store tables on first run. + +## Deployed App Setup + +After deploying your agent with `databricks bundle deploy`, you must grant the app's service principal access to Lakebase. + +### Step 1: Add Lakebase as App Resource + +1. Go to the Databricks UI +2. Navigate to your app and click **Edit** +3. Go to **App resources** → **Add resource** +4. Add your Lakebase instance with **Connect + Create** permissions + +### Step 2: Get App Service Principal ID + +```bash +databricks apps get --output json | jq -r '.service_principal_id' +``` + +### Step 3: Grant Permissions + +Choose **Option A** (SDK - Recommended) or **Option B** (SQL). + +#### Option A: Using LakebaseClient SDK (Recommended) + +The `databricks-ai-bridge` package includes a `LakebaseClient` for programmatic permission management: + +```python +from databricks_ai_bridge.lakebase import LakebaseClient, SchemaPrivilege, TablePrivilege + +# Initialize client +client = LakebaseClient(instance_name="") + +app_sp = "" # From Step 2 + +# Create role for the service principal +client.create_role(app_sp, "SERVICE_PRINCIPAL") + +# Grant schema privileges +client.grant_schema( + grantee=app_sp, + privileges=[SchemaPrivilege.USAGE, SchemaPrivilege.CREATE], + schemas=["drizzle", "ai_chatbot", "public"], +) + +# Grant table privileges on all tables in schemas +client.grant_all_tables_in_schema( + grantee=app_sp, + privileges=[TablePrivilege.SELECT, TablePrivilege.INSERT, TablePrivilege.UPDATE], + schemas=["drizzle", "ai_chatbot"], +) + +# Grant privileges on specific checkpoint tables +client.grant_table( + grantee=app_sp, + privileges=[TablePrivilege.SELECT, TablePrivilege.INSERT, TablePrivilege.UPDATE], + tables=[ + "public.checkpoint_migrations", + "public.checkpoint_writes", + "public.checkpoints", + "public.checkpoint_blobs", + ], +) +``` + +**Benefits of SDK approach:** +- Type-safe privilege enums prevent typos +- Cleaner Python code vs raw SQL +- Easier to integrate into setup scripts + +#### Option B: Using SQL + +Run the following SQL on your Lakebase instance (replace `app-sp-id` with your app's service principal ID): + +```sql +DO $$ +DECLARE + app_sp text := 'app-sp-id'; -- TODO: Replace with your App's Service Principal ID +BEGIN + ------------------------------------------------------------------- + -- Drizzle schema: migration metadata tables + ------------------------------------------------------------------- + EXECUTE format('GRANT USAGE, CREATE ON SCHEMA drizzle TO %I;', app_sp); + EXECUTE format('GRANT SELECT, INSERT, UPDATE ON ALL TABLES IN SCHEMA drizzle TO %I;', app_sp); + ------------------------------------------------------------------- + -- App schema: business tables (Chat, Message, etc.) + ------------------------------------------------------------------- + EXECUTE format('GRANT USAGE, CREATE ON SCHEMA ai_chatbot TO %I;', app_sp); + EXECUTE format('GRANT SELECT, INSERT, UPDATE ON ALL TABLES IN SCHEMA ai_chatbot TO %I;', app_sp); + ------------------------------------------------------------------- + -- Public schema for checkpoint tables + ------------------------------------------------------------------- + EXECUTE format('GRANT USAGE, CREATE ON SCHEMA public TO %I;', app_sp); + EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.checkpoint_migrations TO %I;', app_sp); + EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.checkpoint_writes TO %I;', app_sp); + EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.checkpoints TO %I;', app_sp); + EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.checkpoint_blobs TO %I;', app_sp); +END $$; +``` + +**Schema Reference:** +| Schema | Purpose | +|--------|---------| +| `drizzle` | Migration metadata tables | +| `ai_chatbot` | Business tables (Chat, Message, etc.) | +| `public` | Store tables for user memories | + +## Troubleshooting + +| Issue | Solution | +|-------|----------| +| **"LAKEBASE_INSTANCE_NAME environment variable is required"** | Add `LAKEBASE_INSTANCE_NAME=` to `.env.local` | +| **"Failed to connect to Lakebase instance"** | Verify instance name is correct and your profile has access | +| **Permission errors on store tables** | Run the SDK script or SQL grant commands above | +| **Deployed app can't access Lakebase** | Add Lakebase as app resource in Databricks UI | +| **"role does not exist"** | Run `client.create_role()` first, or ensure SP ID is correct | + +### Common Error Messages + +**Local development:** +``` +Failed to connect to Lakebase instance ''. Please verify: +1. The instance name is correct +2. You have the necessary permissions to access the instance +3. Your Databricks authentication is configured correctly +``` + +**Deployed app:** +``` +Failed to connect to Lakebase instance ''. The App Service Principal for '' may not have access. +``` + +Both errors indicate Lakebase access issues. Follow the setup steps above for your environment. + +## How Memory Works + +See the **agent-memory** skill for: +- How `AsyncDatabricksStore` persists user memories +- Using `user_id` to scope memories per user +- Memory tools (get, save, delete) +- API request examples with user_id + +## Next Steps + +- Understand memory patterns: see **agent-memory** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/agent-langgraph-long-term-memory/.claude/skills/modify-agent/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/modify-agent/SKILL.md new file mode 100644 index 00000000..4932f002 --- /dev/null +++ b/agent-langgraph-long-term-memory/.claude/skills/modify-agent/SKILL.md @@ -0,0 +1,296 @@ +--- +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 + +> **Memory Template:** This template uses long-term memory (facts that persist across sessions). See the **agent-memory** skill for memory-specific patterns. + +## Main File + +**`agent_server/agent.py`** - Agent logic, model selection, instructions, MCP servers + +## Key Files + +| File | Purpose | +|------|---------| +| `agent_server/agent.py` | Agent logic, model, instructions, MCP servers | +| `agent_server/start_server.py` | FastAPI server + MLflow setup | +| `agent_server/evaluate_agent.py` | Agent evaluation with MLflow scorers | +| `agent_server/utils.py` | Databricks auth helpers, stream processing | +| `databricks.yml` | Bundle config & resource permissions | + +## SDK Setup + +```python +import mlflow +from databricks.sdk import WorkspaceClient +from databricks_langchain import ChatDatabricks, DatabricksMCPServer, DatabricksMultiServerMCPClient +from langchain.agents import create_agent + +# Enable autologging for tracing +mlflow.langchain.autolog() + +# Initialize workspace client +workspace_client = WorkspaceClient() +``` + +--- + +## databricks-langchain SDK Overview + +**SDK Location:** https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/langchain + +Before making any changes, ensure that the APIs actually exist in the SDK. If something is missing from the documentation here, look in the venv's `site-packages` directory for the `databricks_langchain` package. If it's not installed, run `uv sync` to create the .venv and install the package. + +--- + +### ChatDatabricks - LLM Chat Interface + +Connects to Databricks Model Serving endpoints for LLM inference. + +```python +from databricks_langchain import ChatDatabricks + +llm = ChatDatabricks( + endpoint="databricks-claude-3-7-sonnet", # or databricks-meta-llama-3-1-70b-instruct + temperature=0, + max_tokens=500, +) + +# For Responses API agents: +llm = ChatDatabricks(endpoint="my-agent-endpoint", use_responses_api=True) +``` + +Available models (check workspace for current list): +- `databricks-claude-3-7-sonnet` +- `databricks-claude-3-5-sonnet` +- `databricks-meta-llama-3-3-70b-instruct` + +**Note:** Some workspaces require granting the app access to the serving endpoint in `databricks.yml`. See the **add-tools** skill and `examples/serving-endpoint.yaml`. + +--- + +### DatabricksEmbeddings - Generate Embeddings + +Query Databricks embedding model endpoints. + +```python +from databricks_langchain import DatabricksEmbeddings + +embeddings = DatabricksEmbeddings(endpoint="databricks-bge-large-en") +vector = embeddings.embed_query("The meaning of life is 42") +vectors = embeddings.embed_documents(["doc1", "doc2"]) +``` + +--- + +### DatabricksVectorSearch - Vector Store + +Connect to Databricks Vector Search indexes for similarity search. + +```python +from databricks_langchain import DatabricksVectorSearch + +# Delta-sync index with Databricks-managed embeddings +vs = DatabricksVectorSearch(index_name="catalog.schema.index_name") + +# Direct-access or self-managed embeddings +vs = DatabricksVectorSearch( + index_name="catalog.schema.index_name", + embedding=embeddings, + text_column="content", +) + +docs = vs.similarity_search("query", k=5) +``` + +--- + +### MCP Client - Tool Integration + +Connect to MCP (Model Context Protocol) servers to get tools for your agent. + +**Basic MCP Server (manual URL):** + +```python +from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient + +client = DatabricksMultiServerMCPClient([ + DatabricksMCPServer( + name="system-ai", + url=f"{host}/api/2.0/mcp/functions/system/ai", + ) +]) +tools = await client.get_tools() +``` + +**From UC Function (convenience helper):** + +Creates MCP server for Unity Catalog functions. If `function_name` is omitted, exposes all functions in the schema. + +```python +server = DatabricksMCPServer.from_uc_function( + catalog="main", + schema="tools", + function_name="send_email", # Optional - omit for all functions in schema + name="email-server", + timeout=30.0, + handle_tool_error=True, +) +``` + +**From Vector Search (convenience helper):** + +Creates MCP server for Vector Search indexes. If `index_name` is omitted, exposes all indexes in the schema. + +```python +server = DatabricksMCPServer.from_vector_search( + catalog="main", + schema="embeddings", + index_name="product_docs", # Optional - omit for all indexes in schema + name="docs-search", + timeout=30.0, +) +``` + +**From Genie Space:** + +Create MCP server from Genie Space. Get the genie space ID from the URL. + +Example: `https://workspace.cloud.databricks.com/genie/rooms/01f0515f6739169283ef2c39b7329700?o=123` means the genie space ID is `01f0515f6739169283ef2c39b7329700` + +```python +DatabricksMCPServer( + name="genie", + url=f"{host_name}/api/2.0/mcp/genie/01f0515f6739169283ef2c39b7329700", +) +``` + +**Non-Databricks MCP Server:** + +```python +from databricks_langchain import MCPServer + +server = MCPServer( + name="external-server", + url="https://other-server.com/mcp", + headers={"X-API-Key": "secret"}, + timeout=15.0, +) +``` + +**After adding MCP servers:** Grant permissions in `databricks.yml` (see **add-tools** skill) + +--- + +## Running the Agent + +```python +from langchain.agents import create_agent + +# Create agent - ONLY accepts tools and model, NO prompt/instructions parameter +agent = create_agent(tools=tools, model=llm) + +# Non-streaming +messages = {"messages": [{"role": "user", "content": "hi"}]} +result = await agent.ainvoke(messages) + +# Streaming +async for event in agent.astream(input=messages, stream_mode=["updates", "messages"]): + # Process stream events + pass +``` + +**Converting to Responses API format:** Use `process_agent_astream_events()` from `agent_server/utils.py`: + +```python +from agent_server.utils import process_agent_astream_events + +async for event in process_agent_astream_events( + agent.astream(input=messages, stream_mode=["updates", "messages"]) +): + yield event # Yields ResponsesAgentStreamEvent objects +``` + +--- + +## Customizing Agent Behavior (System Instructions) + +> **IMPORTANT:** `create_agent()` does NOT accept `prompt`, `instructions`, or `system_message` parameters. Attempting to pass these will cause a runtime error. + +In LangGraph, agent behavior is customized by prepending a system message to the conversation messages. + +**Correct pattern in `agent.py`:** + +1. Define instructions as a constant: +```python +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.""" +``` + +2. Prepend to messages in the `streaming()` function: +```python +@stream() +async def streaming(request: ResponsesAgentRequest) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + agent = await init_agent() + # Prepend system instructions to user messages + user_messages = to_chat_completions_input([i.model_dump() for i in request.input]) + messages = {"messages": [{"role": "system", "content": AGENT_INSTRUCTIONS}] + user_messages} + + async for event in process_agent_astream_events( + agent.astream(input=messages, stream_mode=["updates", "messages"]) + ): + yield event +``` + +**Common mistake to avoid:** +```python +# WRONG - will cause "unexpected keyword argument" error +agent = create_agent(tools=tools, model=llm, prompt=AGENT_INSTRUCTIONS) + +# CORRECT - add instructions via messages +messages = {"messages": [{"role": "system", "content": AGENT_INSTRUCTIONS}] + user_messages} +``` + +For advanced customization (routing, state management, custom graphs), refer to the [LangGraph documentation](https://docs.langchain.com/oss/python/langgraph/overview). + +--- + +## External Connection Tools + +Connect to external services via Unity Catalog HTTP connections: + +- **Slack** - Post messages to channels +- **Google Calendar** - Calendar operations +- **Microsoft Graph API** - Office 365 services +- **Azure AI Search** - Search functionality +- **Any HTTP API** - Use `http_request` from databricks-sdk + +Example: Create UC function wrapping HTTP request for Slack, then expose via MCP. + +--- + +## External Resources + +1. [databricks-langchain SDK](https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/langchain) +2. [Agent examples](https://github.com/bbqiu/agent-on-app-prototype) +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. [LangGraph documentation](https://docs.langchain.com/oss/python/langgraph/overview) +6. [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 +- Memory patterns: see **agent-memory** skill +- Lakebase setup: see **lakebase-setup** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/agent-langgraph-long-term-memory/.claude/skills/quickstart/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/quickstart/SKILL.md new file mode 100644 index 00000000..edfbfd33 --- /dev/null +++ b/agent-langgraph-long-term-memory/.claude/skills/quickstart/SKILL.md @@ -0,0 +1,84 @@ +--- +name: quickstart +description: "Set up Databricks agent development environment. Use when: (1) First time setup, (2) Configuring Databricks authentication, (3) User says 'quickstart', 'set up', 'authenticate', or 'configure databricks', (4) No .env.local file exists." +--- + +# Quickstart & Authentication + +## Prerequisites + +- **uv** (Python package manager) +- **nvm** with Node 20 (for frontend) +- **Databricks CLI v0.283.0+** + +Check CLI version: +```bash +databricks -v # Must be v0.283.0 or above +brew upgrade databricks # If version is too old +``` + +## Run Quickstart + +```bash +uv run quickstart +``` + +**Options:** +- `--profile NAME`: Use specified profile (non-interactive) +- `--host URL`: Workspace URL for initial setup +- `-h, --help`: Show help + +**Examples:** +```bash +# Interactive (prompts for profile selection) +uv run quickstart + +# Non-interactive with existing profile +uv run quickstart --profile DEFAULT + +# New workspace setup +uv run quickstart --host https://your-workspace.cloud.databricks.com +``` + +## What Quickstart Configures + +Creates/updates `.env.local` with: +- `DATABRICKS_CONFIG_PROFILE` - Selected CLI profile +- `MLFLOW_TRACKING_URI` - Set to `databricks://` for local auth +- `MLFLOW_EXPERIMENT_ID` - Auto-created experiment ID + +## Manual Authentication (Fallback) + +If quickstart fails: + +```bash +# Create new profile +databricks auth login --host https://your-workspace.cloud.databricks.com + +# Verify +databricks auth profiles +``` + +Then manually create `.env.local` (copy from `.env.example`): +```bash +# Authentication (choose one method) +DATABRICKS_CONFIG_PROFILE=DEFAULT +# DATABRICKS_HOST=https://.databricks.com +# DATABRICKS_TOKEN=dapi.... + +# MLflow configuration +MLFLOW_EXPERIMENT_ID= +MLFLOW_TRACKING_URI="databricks://DEFAULT" +MLFLOW_REGISTRY_URI="databricks-uc" + +# Frontend proxy settings +CHAT_APP_PORT=3000 +CHAT_PROXY_TIMEOUT_SECONDS=300 +``` + +## Next Steps + +After quickstart completes: +1. **Configure Lakebase** for memory storage (see **lakebase-setup** skill) +2. Run `uv run discover-tools` to find available workspace resources (see **discover-tools** skill) +3. Run `uv run start-app` to test locally (see **run-locally** skill) diff --git a/agent-langgraph-long-term-memory/.claude/skills/run-locally/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/run-locally/SKILL.md new file mode 100644 index 00000000..7d959670 --- /dev/null +++ b/agent-langgraph-long-term-memory/.claude/skills/run-locally/SKILL.md @@ -0,0 +1,130 @@ +--- +name: run-locally +description: "Run and test the agent locally. Use when: (1) User says 'run locally', 'start server', 'test agent', or 'localhost', (2) Need curl commands to test API, (3) Troubleshooting local development issues, (4) Configuring server options like port or hot-reload." +--- + +# Run Agent Locally + +## Start the Server + +```bash +uv run start-app +``` + +This starts the agent at http://localhost:8000 with both the API server and chat UI. + +## Server Options + +| Option | Command | When to Use | +|--------|---------|-------------| +| **Hot-reload** | `uv run start-server --reload` | During development - auto-restarts on code changes | +| **Custom port** | `uv run start-server --port 8001` | When port 8000 is in use | +| **Multiple workers** | `uv run start-server --workers 4` | Load testing or production-like simulation | +| **Combined** | `uv run start-server --reload --port 8001` | Development on alternate port | + +## Test the API + +The agent implements the [OpenAI Responses API](https://platform.openai.com/docs/api-reference/responses) interface via MLflow's ResponsesAgent. + +**Streaming request:** +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' +``` + +**Non-streaming request:** +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }] }' +``` + +**Request with user_id (for long-term memory):** +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "What do you remember about me?"}], + "custom_inputs": {"user_id": "user@example.com"} + }' +``` + +See the **agent-memory** skill for more memory-related request examples. + +## Run Evaluation + +Evaluate your agent using MLflow scorers: + +```bash +uv run agent-evaluate +``` + +**What it does:** +- Runs the agent against a test dataset +- Applies MLflow scorers (RelevanceToQuery, Safety) +- Records results to your MLflow experiment + +**Customize evaluation:** +Edit `agent_server/evaluate_agent.py` to: +- Change the evaluation dataset +- Add or modify scorers +- Adjust evaluation parameters + +After evaluation completes, open the MLflow UI link for your experiment to inspect results. + +## Run Unit Tests + +```bash +pytest [path] +``` + +## Adding Dependencies + +```bash +uv add +# Example: uv add "langchain-community" +``` + +Dependencies are managed in `pyproject.toml`. + +## Troubleshooting + +| Issue | Solution | +|-------|----------| +| **Port already in use** | Use `--port 8001` or kill existing process | +| **Authentication errors** | Verify `.env.local` is correct; run **quickstart** skill | +| **Module not found** | Run `uv sync` to install dependencies | +| **MLflow experiment not found** | Ensure `MLFLOW_TRACKING_URI` in `.env.local` is `databricks://` | +| **Lakebase connection errors** | Verify `LAKEBASE_INSTANCE_NAME` in `.env.local`; see **lakebase-setup** skill | + +### MLflow Experiment Not Found + +If you see: "The provided MLFLOW_EXPERIMENT_ID environment variable value does not exist" + +**Verify the experiment exists:** +```bash +databricks -p experiments get-experiment +``` + +**Fix:** Ensure `.env.local` has the correct tracking URI format: +```bash +MLFLOW_TRACKING_URI="databricks://DEFAULT" # Include profile name +``` + +The quickstart script configures this automatically. If you manually edited `.env.local`, ensure the profile name is included. + +## MLflow Tracing + +This template uses MLflow for automatic tracing: +- Agent logic decorated with `@invoke()` and `@stream()` is automatically traced +- LLM invocations are captured via MLflow autologging + +To add custom trace instrumentation, see: [MLflow tracing documentation](https://docs.databricks.com/aws/en/mlflow3/genai/tracing/app-instrumentation/) + +## Next Steps + +- Configure Lakebase: see **lakebase-setup** skill +- Understand memory patterns: see **agent-memory** skill +- Modify your agent: see **modify-agent** skill +- Deploy to Databricks: see **deploy** skill diff --git a/agent-langgraph-long-term-memory/.gitignore b/agent-langgraph-long-term-memory/.gitignore index c0ee193b..0c3cb57a 100644 --- a/agent-langgraph-long-term-memory/.gitignore +++ b/agent-langgraph-long-term-memory/.gitignore @@ -204,5 +204,18 @@ sketch **/mlruns/ **/.vite/ **/.databricks -**/.claude **/.env +**/.env.local + +# Claude Code - track skills only +.claude/* +!.claude/skills/ +.claude/skills/* +!.claude/skills/quickstart/ +!.claude/skills/discover-tools/ +!.claude/skills/deploy/ +!.claude/skills/add-tools/ +!.claude/skills/run-locally/ +!.claude/skills/modify-agent/ +!.claude/skills/lakebase-setup/ +!.claude/skills/agent-memory/ diff --git a/agent-langgraph-long-term-memory/AGENTS.md b/agent-langgraph-long-term-memory/AGENTS.md index 617c2c28..41046b58 100644 --- a/agent-langgraph-long-term-memory/AGENTS.md +++ b/agent-langgraph-long-term-memory/AGENTS.md @@ -1,348 +1,109 @@ -# Agent LangGraph Development Guide +# Agent Development Guide -## Running the App +## MANDATORY First Action -**Prerequisites:** uv, nvm (Node 20), Databricks CLI +**BEFORE any other action, run `databricks auth profiles` to check authentication status.** -**Quick Start:** +This helps you understand: +- Which Databricks profiles are configured +- Whether authentication is already set up +- Which profile to use for subsequent commands -```bash -uv run quickstart # First-time setup (auth, MLflow experiment, env) -uv run start-app # Start app at http://localhost:8000 -``` +If no profiles exist, guide the user through running `uv run quickstart` to set up authentication. -**Advanced Server Options:** +## Understanding User Goals -```bash -uv run start-server --reload # Hot-reload on code changes during development -uv run start-server --port 8001 -uv run start-server --workers 4 -``` +**Ask the user questions to understand what they're building:** -**Test API:** +1. **What is the agent's purpose?** (e.g., data analyst assistant, customer support, code helper) +2. **What data or tools does it need access to?** + - Databases/tables (Unity Catalog) + - Documents for RAG (Vector Search) + - Natural language data queries (Genie Spaces) + - External APIs or services +3. **Any specific Databricks resources they want to connect?** -```bash -# Streaming request -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' +Use `uv run discover-tools` to show them available resources in their workspace, then help them select the right ones for their use case. **See the `add-tools` skill for how to connect tools and grant permissions.** -# Non-streaming request -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ "input": [{ "role": "user", "content": "hi" }] }' -``` +## Memory Template Note ---- - -## Testing the Agent - -**Run evaluation:** - -```bash -uv run agent-evaluate # Uses MLflow scorers (RelevanceToQuery, Safety) -``` - -**Run unit tests:** - -```bash -pytest [path] # Standard pytest execution -``` - ---- - -## Modifying the Agent - -Anytime the user wants to modify the agent, look through each of the following resources to help them accomplish their goal: - -If the user wants to convert something into Responses API, refer to https://mlflow.org/docs/latest/genai/serving/responses-agent/ for more information. +This template includes **long-term memory** (facts that persist across conversation sessions). The agent can remember user preferences and information across multiple interactions. -1. Look through existing databricks-langchain APIs to see if they can use one of these to accomplish their goal. -2. Look through the folders in https://github.com/bbqiu/agent-on-app-prototype to see if there's an existing example similar to what they're looking to do. -3. Reference the documentation available under https://docs.databricks.com/aws/en/generative-ai/agent-framework/ and its subpages. -4. For adding tools and capabilities, refer to: https://docs.databricks.com/aws/en/generative-ai/agent-framework/agent-tool -5. For stuff like LangGraph routing, configuration, and customization, refer to the LangGraph documentation: https://docs.langchain.com/oss/python/langgraph/overview. +**Required setup:** +1. Configure Lakebase instance (see **lakebase-setup** skill) +2. Use `user_id` in requests to scope memories per user (see **agent-memory** skill) -**Main file to modify:** `agent_server/agent.py` +## Handling Deployment Errors ---- - -## databricks-langchain SDK overview - -**SDK Location:** `https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/langchain` +**If `databricks bundle deploy` fails with "An app with the same name already exists":** -**Development Workflow:** +Ask the user: "I see there's an existing app with the same name. Would you like me to bind it to this bundle so we can manage it, or delete it and create a new one?" -```bash -uv add databricks-langchain -``` - -Before making any changes, ensure that the APIs actually exist in the SDK. If something is missing from the documentation here, feel free to look in the venv's `site-packages` directory for the `databricks_langchain` package. If it's not installed, run `uv sync` in this folder to create the .venv and install the package. +- **If they want to bind**: See the **deploy** skill for binding steps +- **If they want to delete**: Run `databricks apps delete ` then deploy again --- -### ChatDatabricks - LLM Chat Interface - -Connects to Databricks Model Serving endpoints for LLM inference. +## Available Skills -```python -from databricks_langchain import ChatDatabricks +**Before executing any task, read the relevant skill file in `.claude/skills/`** - they contain tested commands, patterns, and troubleshooting steps. -llm = ChatDatabricks( - endpoint="databricks-claude-3-7-sonnet", # or databricks-meta-llama-3-1-70b-instruct - temperature=0, - max_tokens=500, -) +| Task | Skill | Path | +|------|-------|------| +| Setup, auth, first-time | **quickstart** | `.claude/skills/quickstart/SKILL.md` | +| Lakebase configuration | **lakebase-setup** | `.claude/skills/lakebase-setup/SKILL.md` | +| Memory patterns | **agent-memory** | `.claude/skills/agent-memory/SKILL.md` | +| Find tools/resources | **discover-tools** | `.claude/skills/discover-tools/SKILL.md` | +| Deploy to Databricks | **deploy** | `.claude/skills/deploy/SKILL.md` | +| Add tools & permissions | **add-tools** | `.claude/skills/add-tools/SKILL.md` | +| Run/test locally | **run-locally** | `.claude/skills/run-locally/SKILL.md` | +| Modify agent code | **modify-agent** | `.claude/skills/modify-agent/SKILL.md` | -# For Responses API agents: -llm = ChatDatabricks(endpoint="my-agent-endpoint", use_responses_api=True) -``` +**Note:** All agent skills are located in `.claude/skills/` directory. --- -### DatabricksEmbeddings - Generate Embeddings - -Query Databricks embedding model endpoints. +## Quick Commands -```python -from databricks_langchain import DatabricksEmbeddings - -embeddings = DatabricksEmbeddings(endpoint="databricks-bge-large-en") -vector = embeddings.embed_query("The meaning of life is 42") -vectors = embeddings.embed_documents(["doc1", "doc2"]) -``` +| Task | Command | +|------|---------| +| Setup | `uv run quickstart` | +| Discover tools | `uv run discover-tools` | +| Run locally | `uv run start-app` | +| Deploy | `databricks bundle deploy && databricks bundle run agent_langgraph_long_term_memory` | +| View logs | `databricks apps logs --follow` | --- -### DatabricksVectorSearch - Vector Store - -Connect to Databricks Vector Search indexes for similarity search. - -```python -from databricks_langchain import DatabricksVectorSearch - -# Delta-sync index with Databricks-managed embeddings -vs = DatabricksVectorSearch(index_name="catalog.schema.index_name") - -# Direct-access or self-managed embeddings -vs = DatabricksVectorSearch( - index_name="catalog.schema.index_name", - embedding=embeddings, - text_column="content", -) +## Key Files -docs = vs.similarity_search("query", k=5) -``` +| File | Purpose | +|------|---------| +| `agent_server/agent.py` | Agent logic, model, instructions, MCP servers, memory tools | +| `agent_server/start_server.py` | FastAPI server + MLflow setup | +| `agent_server/evaluate_agent.py` | Agent evaluation with MLflow scorers | +| `databricks.yml` | Bundle config & resource permissions | +| `scripts/quickstart.py` | One-command setup script | +| `scripts/discover_tools.py` | Discovers available workspace resources | --- -### MCP Client - Tool Integration - -Connect to MCP (Model Context Protocol) servers to get tools for your agent. - -**Basic MCP Server (manual URL):** - -```python -from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient - -client = DatabricksMultiServerMCPClient([ - DatabricksMCPServer( - name="system-ai", - url=f"{host}/api/2.0/mcp/functions/system/ai", - ) -]) -tools = await client.get_tools() -``` - -**From UC Function (convenience helper):** -Creates MCP server for Unity Catalog functions. If `function_name` is omitted, exposes all functions in the schema. - -```python -server = DatabricksMCPServer.from_uc_function( - catalog="main", - schema="tools", - function_name="send_email", # Optional - omit for all functions in schema - name="email-server", - timeout=30.0, - handle_tool_error=True, -) -``` - -**From Vector Search (convenience helper):** -Creates MCP server for Vector Search indexes. If `index_name` is omitted, exposes all indexes in the schema. - -```python -server = DatabricksMCPServer.from_vector_search( - catalog="main", - schema="embeddings", - index_name="product_docs", # Optional - omit for all indexes in schema - name="docs-search", - timeout=30.0, -) -``` - -**From Genie Space:** -Create MCP server from Genie Space. Need to get the genie space ID. Can prompt the user to retrieve this via the UI by getting the link to the genie space. - -Ex: https://db-ml-models-dev-us-west.cloud.databricks.com/genie/rooms/01f0515f6739169283ef2c39b7329700?o=3217006663075879 means the genie space ID is 01f0515f6739169283ef2c39b7329700 - -```python -DatabricksMCPServer( - name="genie", - url=f"{host_name}/api/2.0/mcp/genie/01f0515f6739169283ef2c39b7329700", -), -``` - -**Non-Databricks MCP Server:** - -```python -from databricks_langchain import MCPServer - -server = MCPServer( - name="external-server", - url="https://other-server.com/mcp", - headers={"X-API-Key": "secret"}, - timeout=15.0, -) -``` - -### Stateful LangGraph agent - -To enable statefulness in a LangGraph agent, we need to install `databricks-langchain[memory]`. - -Look through the package files for the latest on stateful langgraph agents. Can start by looking at the databricks_langchain/checkpoints.py and databricks_langchain/store.py files. - -## Lakebase instance setup for stateful agents - -Add the lakebase name to `.env`: - -```bash -LAKEBASE_INSTANCE_NAME= -``` - ## Agent Framework Capabilities -Reference: https://docs.databricks.com/aws/en/generative-ai/agent-framework/ - -### Tool Types +> **IMPORTANT:** When adding any tool to the agent, you MUST also grant permissions in `databricks.yml`. See the **add-tools** skill for required steps and examples. +**Tool Types:** 1. **Unity Catalog Function Tools** - SQL UDFs managed in UC with built-in governance 2. **Agent Code Tools** - Defined directly in agent code for REST APIs and low-latency operations 3. **MCP Tools** - Interoperable tools via Model Context Protocol (Databricks-managed, external, or self-hosted) -### Built-in Tools - +**Built-in Tools:** - **system.ai.python_exec** - Execute Python code dynamically within agent queries (code interpreter) -### External Connection Tools - -Connect to external services via Unity Catalog HTTP connections: - -- **Slack** - Post messages to channels -- **Google Calendar** - Calendar operations -- **Microsoft Graph API** - Office 365 services -- **Azure AI Search** - Search functionality -- **Any HTTP API** - Use `http_request` from databricks-sdk - -Example: Create UC function wrapping HTTP request for Slack, then expose via MCP. - -### Common Patterns - +**Common Patterns:** - **Structured data retrieval** - Query SQL tables/databases - **Unstructured data retrieval** - Document search and RAG via Vector Search - **Code interpreter** - Python execution for analysis via system.ai.python_exec - **External connections** - Integrate services like Slack via HTTP connections ---- - -## Authentication Setup - -**Option 1: OAuth (Recommended)** - -```bash -databricks auth login -``` - -Set in `.env`: - -```bash -DATABRICKS_CONFIG_PROFILE=DEFAULT -``` - -**Option 2: Personal Access Token** - -Set in `.env`: - -```bash -DATABRICKS_HOST="https://host.databricks.com" -DATABRICKS_TOKEN="dapi_token" -``` - ---- - -## MLflow Experiment Setup - -Create and link an MLflow experiment: - -```bash -DATABRICKS_USERNAME=$(databricks current-user me | jq -r .userName) -databricks experiments create-experiment /Users/$DATABRICKS_USERNAME/agents-on-apps -``` - -Add the experiment ID to `.env`: - -```bash -MLFLOW_EXPERIMENT_ID= -``` - ---- - -## Key Files - -| File | Purpose | -| -------------------------------- | --------------------------------------------- | -| `agent_server/agent.py` | Agent logic, model, instructions, MCP servers | -| `agent_server/start_server.py` | FastAPI server + MLflow setup | -| `agent_server/evaluate_agent.py` | Agent evaluation with MLflow scorers | -| `agent_server/utils.py` | Databricks auth helpers, stream processing | -| `scripts/start_app.py` | Manages backend+frontend startup | - ---- - -## Deploying to Databricks Apps - -**Create app:** - -```bash -databricks apps create agent-langgraph -``` - -**Sync files:** - -```bash -DATABRICKS_USERNAME=$(databricks current-user me | jq -r .userName) -databricks sync . "/Users/$DATABRICKS_USERNAME/agent-langgraph" -``` - -**Deploy:** - -```bash -databricks apps deploy agent-langgraph --source-code-path /Workspace/Users/$DATABRICKS_USERNAME/agent-langgraph -``` - -**Query deployed app:** - -Generate OAuth token (PATs are not supported): - -```bash -databricks auth token -``` - -Send request: - -```bash -curl -X POST /invocations \ - -H "Authorization: Bearer " \ - -H "Content-Type: application/json" \ - -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' -``` +Reference: https://docs.databricks.com/aws/en/generative-ai/agent-framework/ diff --git a/agent-langgraph-long-term-memory/README.md b/agent-langgraph-long-term-memory/README.md index 5a4f6312..bb0aec45 100644 --- a/agent-langgraph-long-term-memory/README.md +++ b/agent-langgraph-long-term-memory/README.md @@ -1,6 +1,20 @@ -# Responses API Agent +# Responses API Agent (Long-Term Memory) -This template defines a conversational agent app. The app comes with a built-in chat UI, but also exposes an API endpoint for invoking the agent so that you can serve your UI elsewhere (e.g. on your website or in a mobile app). +## Build with AI Assistance + +This template includes Claude Code skills in `.claude/skills/` for AI-assisted development. Use [Claude Code](https://docs.anthropic.com/en/docs/claude-code) to: + +- **Set up your environment**: "Run quickstart to configure authentication" +- **Add tools**: "Connect my agent to a Genie space" +- **Configure memory**: "Set up Lakebase for user memories" +- **Deploy**: "Deploy my agent to Databricks Apps" +- **Debug**: "Why am I getting a permission error?" + +The skills contain tested commands, code patterns, and troubleshooting steps. + +--- + +This template defines a conversational agent app with long-term memory. The app comes with a built-in chat UI, but also exposes an API endpoint for invoking the agent so that you can serve your UI elsewhere (e.g. on your website or in a mobile app). The agent in this template implements the [OpenAI Responses API](https://platform.openai.com/docs/api-reference/responses) interface. It has access to a single tool; the [built-in code interpreter tool](https://docs.databricks.com/aws/en/generative-ai/agent-framework/code-interpreter-tools#built-in-python-executor-tool) (`system.ai.python_exec`) on Databricks. You can customize agent code and test it via the API or UI. diff --git a/agent-langgraph-long-term-memory/pyproject.toml b/agent-langgraph-long-term-memory/pyproject.toml index ed038488..ae46c70b 100644 --- a/agent-langgraph-long-term-memory/pyproject.toml +++ b/agent-langgraph-long-term-memory/pyproject.toml @@ -36,3 +36,4 @@ quickstart = "scripts.quickstart:main" start-app = "scripts.start_app:main" start-server = "agent_server.start_server:main" agent-evaluate = "agent_server.evaluate_agent:evaluate" +discover-tools = "scripts.discover_tools:main" diff --git a/agent-langgraph-long-term-memory/scripts/discover_tools.py b/agent-langgraph-long-term-memory/scripts/discover_tools.py new file mode 100644 index 00000000..3eb37963 --- /dev/null +++ b/agent-langgraph-long-term-memory/scripts/discover_tools.py @@ -0,0 +1,432 @@ +#!/usr/bin/env python3 +""" +Discover available tools and data sources for Databricks agents. + +This script scans for: +- Unity Catalog functions (data retrieval tools e.g. SQL UDFs) +- Unity Catalog tables (data sources) +- Vector search indexes (RAG data sources) +- Genie spaces (conversational interface over structured data) +- Custom MCP servers (Databricks apps with name mcp-*) +- External MCP servers (via Unity Catalog connections) +""" + +import json +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict, List + +from databricks.sdk import WorkspaceClient + +DEFAULT_MAX_RESULTS = 100 +DEFAULT_MAX_SCHEMAS = 25 + +def run_databricks_cli(args: List[str]) -> str: + """Run databricks CLI command and return output.""" + try: + result = subprocess.run( + ["databricks"] + args, + capture_output=True, + text=True, + check=True, + ) + return result.stdout + except subprocess.CalledProcessError as e: + print(f"Error running databricks CLI: {e.stderr}", file=sys.stderr) + return "" + + +def discover_uc_functions(w: WorkspaceClient, catalog: str = None, max_schemas: int = DEFAULT_MAX_SCHEMAS) -> List[Dict[str, Any]]: + """Discover Unity Catalog functions that could be used as tools. + + Args: + w: WorkspaceClient instance + catalog: Optional specific catalog to search + max_schemas: Total number of schemas to search across all catalogs + """ + functions = [] + schemas_searched = 0 + + try: + catalogs = [catalog] if catalog else [c.name for c in w.catalogs.list()] + + for cat in catalogs: + if schemas_searched >= max_schemas: + break + + try: + all_schemas = list(w.schemas.list(catalog_name=cat)) + # Take schemas from this catalog until we hit the global budget + schemas_to_search = all_schemas[:max_schemas - schemas_searched] + + for schema in schemas_to_search: + schema_name = f"{cat}.{schema.name}" + try: + funcs = list(w.functions.list(catalog_name=cat, schema_name=schema.name)) + for func in funcs: + functions.append({ + "type": "uc_function", + "name": func.full_name, + "catalog": cat, + "schema": schema.name, + "function_name": func.name, + "comment": func.comment, + "routine_definition": getattr(func, "routine_definition", None), + }) + except Exception as e: + # Skip schemas we can't access + continue + finally: + schemas_searched += 1 + except Exception as e: + # Skip catalogs we can't access + continue + + except Exception as e: + print(f"Error discovering UC functions: {e}", file=sys.stderr) + + return functions + + +def discover_uc_tables(w: WorkspaceClient, catalog: str = None, schema: str = None, max_schemas: int = DEFAULT_MAX_SCHEMAS) -> List[Dict[str, Any]]: + """Discover Unity Catalog tables that could be queried. + + Args: + w: WorkspaceClient instance + catalog: Optional specific catalog to search + schema: Optional specific schema to search (requires catalog) + max_schemas: Total number of schemas to search across all catalogs + """ + tables = [] + schemas_searched = 0 + + try: + catalogs = [catalog] if catalog else [c.name for c in w.catalogs.list()] + + for cat in catalogs: + if cat in ["__databricks_internal", "system"]: + continue + + if schemas_searched >= max_schemas: + break + + try: + if schema: + schemas_to_search = [schema] + else: + all_schemas = [s.name for s in w.schemas.list(catalog_name=cat)] + # Take schemas from this catalog until we hit the global budget + schemas_to_search = all_schemas[:max_schemas - schemas_searched] + + for sch in schemas_to_search: + if sch == "information_schema": + schemas_searched += 1 + continue + + try: + tbls = list(w.tables.list(catalog_name=cat, schema_name=sch)) + for tbl in tbls: + # Get column info + columns = [] + if hasattr(tbl, "columns") and tbl.columns: + columns = [ + {"name": col.name, "type": col.type_name.value if hasattr(col.type_name, "value") else str(col.type_name)} + for col in tbl.columns + ] + + tables.append({ + "type": "uc_table", + "name": tbl.full_name, + "catalog": cat, + "schema": sch, + "table_name": tbl.name, + "table_type": tbl.table_type.value if tbl.table_type else None, + "comment": tbl.comment, + "columns": columns, + }) + except Exception as e: + # Skip schemas we can't access + pass + finally: + schemas_searched += 1 + except Exception as e: + # Skip catalogs we can't access + continue + + except Exception as e: + print(f"Error discovering UC tables: {e}", file=sys.stderr) + + return tables + + +def discover_vector_search_indexes(w: WorkspaceClient) -> List[Dict[str, Any]]: + """Discover Vector Search indexes for RAG applications.""" + indexes = [] + + try: + # List all vector search endpoints + endpoints = list(w.vector_search_endpoints.list_endpoints()) + + for endpoint in endpoints: + try: + # List indexes for each endpoint + endpoint_indexes = list(w.vector_search_indexes.list_indexes(endpoint_name=endpoint.name)) + for idx in endpoint_indexes: + indexes.append({ + "type": "vector_search_index", + "name": idx.name, + "endpoint": endpoint.name, + "primary_key": idx.primary_key, + "index_type": idx.index_type.value if idx.index_type else None, + "status": idx.status.state.value if idx.status and idx.status.state else None, + }) + except Exception as e: + # Skip endpoints we can't access + continue + + except Exception as e: + print(f"Error discovering vector search indexes: {e}", file=sys.stderr) + + return indexes + + +def discover_genie_spaces(w: WorkspaceClient) -> List[Dict[str, Any]]: + """Discover Genie spaces for conversational data access.""" + spaces = [] + + try: + # Use SDK to list genie spaces + response = w.genie.list_spaces() + genie_spaces = response.spaces if hasattr(response, "spaces") else [] + for space in genie_spaces: + spaces.append({ + "type": "genie_space", + "id": space.space_id, + "name": space.title, + "description": space.description, + }) + except Exception as e: + print(f"Error discovering Genie spaces: {e}", file=sys.stderr) + + return spaces + + + +def discover_custom_mcp_servers(w: WorkspaceClient) -> List[Dict[str, Any]]: + """Discover custom MCP servers deployed as Databricks apps.""" + custom_servers = [] + + try: + # List all apps and filter for those starting with mcp- + apps = w.apps.list() + for app in apps: + if app.name and app.name.startswith("mcp-"): + custom_servers.append({ + "type": "custom_mcp_server", + "name": app.name, + "url": app.url, + "status": app.app_status.state.value if app.app_status and app.app_status.state else None, + "description": app.description, + }) + except Exception as e: + print(f"Error discovering custom MCP servers: {e}", file=sys.stderr) + + return custom_servers + + +def discover_external_mcp_servers(w: WorkspaceClient) -> List[Dict[str, Any]]: + """Discover external MCP servers configured via Unity Catalog connections.""" + external_servers = [] + + try: + # List all connections and filter for MCP connections + connections = w.connections.list() + for conn in connections: + # Check if this is an MCP connection + if conn.options and conn.options.get("is_mcp_connection") == "true": + external_servers.append({ + "type": "external_mcp_server", + "name": conn.name, + "connection_type": conn.connection_type.value if hasattr(conn.connection_type, "value") else str(conn.connection_type), + "comment": conn.comment, + "full_name": conn.full_name, + }) + except Exception as e: + print(f"Error discovering external MCP servers: {e}", file=sys.stderr) + + return external_servers + + +def format_output_markdown(results: Dict[str, List[Dict[str, Any]]]) -> str: + """Format discovery results as markdown.""" + lines = ["# Agent Tools and Data Sources Discovery\n"] + + # UC Functions + functions = results.get("uc_functions", []) + if functions: + lines.append(f"## Unity Catalog Functions ({len(functions)})\n") + lines.append("**What they are:** SQL UDFs that can be used as agent tools.\n") + lines.append("**How to use:** Access via UC functions MCP server:") + lines.append("- All functions in a schema: `{workspace_host}/api/2.0/mcp/functions/{catalog}/{schema}`") + lines.append("- Single function: `{workspace_host}/api/2.0/mcp/functions/{catalog}/{schema}/{function_name}`\n") + for func in functions[:10]: # Show first 10 + lines.append(f"- `{func['name']}`") + if func.get("comment"): + lines.append(f" - {func['comment']}") + if len(functions) > 10: + lines.append(f"\n*...and {len(functions) - 10} more*\n") + lines.append("") + + # UC Tables + tables = results.get("uc_tables", []) + if tables: + lines.append(f"## Unity Catalog Tables ({len(tables)})\n") + lines.append("Structured data that agents can query via UC SQL functions.\n") + for table in tables[:10]: # Show first 10 + lines.append(f"- `{table['name']}` ({table['table_type']})") + if table.get("comment"): + lines.append(f" - {table['comment']}") + if table.get("columns"): + col_names = [c["name"] for c in table["columns"][:5]] + lines.append(f" - Columns: {', '.join(col_names)}") + if len(tables) > 10: + lines.append(f"\n*...and {len(tables) - 10} more*\n") + lines.append("") + + # Vector Search Indexes + indexes = results.get("vector_search_indexes", []) + if indexes: + lines.append(f"## Vector Search Indexes ({len(indexes)})\n") + lines.append("These can be used for RAG applications with unstructured data.\n") + lines.append("**How to use:** Connect via MCP server at `{workspace_host}/api/2.0/mcp/vector-search/{catalog}/{schema}` or\n") + lines.append("`{workspace_host}/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name}`\n") + for idx in indexes: + lines.append(f"- `{idx['name']}`") + lines.append(f" - Endpoint: {idx['endpoint']}") + lines.append(f" - Status: {idx['status']}") + lines.append("") + + # Genie Spaces + spaces = results.get("genie_spaces", []) + if spaces: + lines.append(f"## Genie Spaces ({len(spaces)})\n") + lines.append("**What they are:** Natural language interface to your data\n") + lines.append("**How to use:** Connect via Genie MCP server at `{workspace_host}/api/2.0/mcp/genie/{space_id}`\n") + for space in spaces: + lines.append(f"- `{space['name']}` (ID: {space['id']})") + if space.get("description"): + lines.append(f" - {space['description']}") + lines.append("") + + # Custom MCP Servers (Databricks Apps) + custom_servers = results.get("custom_mcp_servers", []) + if custom_servers: + lines.append(f"## Custom MCP Servers ({len(custom_servers)})\n") + lines.append("**What:** Your own MCP servers deployed as Databricks Apps (names starting with mcp-)\n") + lines.append("**How to use:** Access via `{app_url}/mcp`\n") + lines.append("**⚠️ Important:** Custom MCP server apps require manual permission grants:") + lines.append("1. Get your agent app's service principal: `databricks apps get --output json | jq -r '.service_principal_name'`") + lines.append("2. Grant permission: `databricks apps update-permissions --service-principal --permission-level CAN_USE`") + lines.append("(Apps are not yet supported as resource dependencies in databricks.yml)\n") + for server in custom_servers: + lines.append(f"- `{server['name']}`") + if server.get("url"): + lines.append(f" - URL: {server['url']}") + if server.get("status"): + lines.append(f" - Status: {server['status']}") + if server.get("description"): + lines.append(f" - {server['description']}") + lines.append("") + + # External MCP Servers (UC Connections) + external_servers = results.get("external_mcp_servers", []) + if external_servers: + lines.append(f"## External MCP Servers ({len(external_servers)})\n") + lines.append("**What:** Third-party MCP servers via Unity Catalog connections\n") + lines.append("**How to use:** Connect via `{workspace_host}/api/2.0/mcp/external/{connection_name}`\n") + lines.append("**Benefits:** Secure access to external APIs through UC governance\n") + for server in external_servers: + lines.append(f"- `{server['name']}`") + if server.get("full_name"): + lines.append(f" - Full name: {server['full_name']}") + if server.get("comment"): + lines.append(f" - {server['comment']}") + lines.append("") + return "\n".join(lines) + + +def main(): + """Main discovery function.""" + import argparse + + parser = argparse.ArgumentParser(description="Discover available agent tools and data sources") + parser.add_argument("--catalog", help="Limit discovery to specific catalog") + parser.add_argument("--schema", help="Limit discovery to specific schema (requires --catalog)") + parser.add_argument("--format", choices=["json", "markdown"], default="markdown", help="Output format") + parser.add_argument("--output", help="Output file (default: stdout)") + parser.add_argument("--profile", help="Databricks CLI profile to use (default: uses default profile)") + parser.add_argument("--max-results", type=int, default=DEFAULT_MAX_RESULTS, help=f"Maximum results per resource type (default: {DEFAULT_MAX_RESULTS})") + parser.add_argument("--max-schemas", type=int, default=DEFAULT_MAX_SCHEMAS, help=f"Total schemas to search across all catalogs (default: {DEFAULT_MAX_SCHEMAS})") + + args = parser.parse_args() + + if args.schema and not args.catalog: + print("Error: --schema requires --catalog", file=sys.stderr) + sys.exit(1) + + print("Discovering available tools and data sources...", file=sys.stderr) + + # Initialize Databricks workspace client + # Only pass profile if specified, otherwise use default + if args.profile: + w = WorkspaceClient(profile=args.profile) + else: + w = WorkspaceClient() + + results = {} + + # Discover each type with configurable limits + print("- UC Functions...", file=sys.stderr) + results["uc_functions"] = discover_uc_functions(w, catalog=args.catalog, max_schemas=args.max_schemas)[:args.max_results] + + print("- UC Tables...", file=sys.stderr) + results["uc_tables"] = discover_uc_tables(w, catalog=args.catalog, schema=args.schema, max_schemas=args.max_schemas)[:args.max_results] + + print("- Vector Search Indexes...", file=sys.stderr) + results["vector_search_indexes"] = discover_vector_search_indexes(w)[:args.max_results] + + print("- Genie Spaces...", file=sys.stderr) + results["genie_spaces"] = discover_genie_spaces(w)[:args.max_results] + + print("- Custom MCP Servers (Apps)...", file=sys.stderr) + results["custom_mcp_servers"] = discover_custom_mcp_servers(w)[:args.max_results] + + print("- External MCP Servers (Connections)...", file=sys.stderr) + results["external_mcp_servers"] = discover_external_mcp_servers(w)[:args.max_results] + + # Format output + if args.format == "json": + output = json.dumps(results, indent=2) + else: + output = format_output_markdown(results) + + # Write output + if args.output: + Path(args.output).write_text(output) + print(f"\nResults written to {args.output}", file=sys.stderr) + else: + print("\n" + output) + + # Print summary + print("\n=== Discovery Summary ===", file=sys.stderr) + print(f"UC Functions: {len(results['uc_functions'])}", file=sys.stderr) + print(f"UC Tables: {len(results['uc_tables'])}", file=sys.stderr) + print(f"Vector Search Indexes: {len(results['vector_search_indexes'])}", file=sys.stderr) + print(f"Genie Spaces: {len(results['genie_spaces'])}", file=sys.stderr) + print(f"Custom MCP Servers: {len(results['custom_mcp_servers'])}", file=sys.stderr) + print(f"External MCP Servers: {len(results['external_mcp_servers'])}", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/agent-langgraph-short-term-memory/.claude/skills/add-tools/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/add-tools/SKILL.md new file mode 100644 index 00000000..45f83a8f --- /dev/null +++ b/agent-langgraph-short-term-memory/.claude/skills/add-tools/SKILL.md @@ -0,0 +1,88 @@ +--- +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 + +**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. + +## Workflow + +**Step 1:** Add MCP server in `agent_server/agent.py`: +```python +from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient + +genie_server = DatabricksMCPServer( + url=f"{host}/api/2.0/mcp/genie/01234567-89ab-cdef", + name="my genie space", +) + +mcp_client = DatabricksMultiServerMCPClient([genie_server]) +tools = await mcp_client.get_tools() +``` + +**Step 2:** Grant access in `databricks.yml`: +```yaml +resources: + apps: + agent_langgraph_short_term_memory: + resources: + - name: 'my_genie_space' + genie_space: + name: 'My Genie Space' + space_id: '01234567-89ab-cdef' + permission: 'CAN_RUN' +``` + +**Step 3:** Deploy with `databricks bundle deploy` (see **deploy** skill) + +## 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 | +| `experiment.yaml` | MLflow experiment | Tracing (already configured) | +| `custom-mcp-server.md` | Custom MCP apps | Apps starting with `mcp-*` | + +## Custom MCP Servers (Databricks Apps) + +Apps are **not yet supported** as resource dependencies in `databricks.yml`. Manual permission grant required: + +**Step 1:** Get your agent app's service principal: +```bash +databricks apps get --output json | jq -r '.service_principal_name' +``` + +**Step 2:** Grant permission on the MCP server app: +```bash +databricks apps update-permissions \ + --service-principal \ + --permission-level CAN_USE +``` + +See `examples/custom-mcp-server.md` for detailed steps. + +## Important Notes + +- **MLflow experiment**: Already configured in template, no action needed +- **Lakebase permissions**: Memory storage requires separate Lakebase setup (see **lakebase-setup** skill) +- **Multiple resources**: Add multiple entries under `resources:` list +- **Permission types vary**: Each resource type has specific permission values +- **Deploy after changes**: Run `databricks bundle deploy` after modifying `databricks.yml` + +## Next Steps + +- Configure memory storage: see **lakebase-setup** skill +- Understand memory patterns: see **agent-memory** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md b/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md new file mode 100644 index 00000000..9e919ad2 --- /dev/null +++ b/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md @@ -0,0 +1,57 @@ +# Custom MCP Server (Databricks App) + +Custom MCP servers are Databricks Apps with names starting with `mcp-*`. + +**Apps are not yet supported as resource dependencies in `databricks.yml`**, so manual permission grant is required. + +## 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. Deploy your agent app first + +```bash +databricks bundle deploy +databricks bundle run agent_langgraph_short_term_memory +``` + +### 3. Get your agent app's service principal + +```bash +databricks apps get --output json | jq -r '.service_principal_name' +``` + +Example output: `sp-abc123-def456` + +### 4. Grant permission on the MCP server app + +```bash +databricks apps update-permissions \ + --service-principal \ + --permission-level CAN_USE +``` + +Example: +```bash +databricks apps update-permissions mcp-my-server \ + --service-principal sp-abc123-def456 \ + --permission-level CAN_USE +``` + +## Notes + +- This manual step is required each time you connect to a new custom MCP server +- The permission grant persists across deployments +- If you redeploy the agent app with a new service principal, you'll need to grant permissions again diff --git a/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/experiment.yaml b/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/experiment.yaml new file mode 100644 index 00000000..ac5c626a --- /dev/null +++ b/agent-langgraph-short-term-memory/.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-langgraph-short-term-memory/.claude/skills/add-tools/examples/genie-space.yaml b/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/genie-space.yaml new file mode 100644 index 00000000..71589d52 --- /dev/null +++ b/agent-langgraph-short-term-memory/.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-langgraph-short-term-memory/.claude/skills/add-tools/examples/serving-endpoint.yaml b/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/serving-endpoint.yaml new file mode 100644 index 00000000..b49ce9da --- /dev/null +++ b/agent-langgraph-short-term-memory/.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-langgraph-short-term-memory/.claude/skills/add-tools/examples/sql-warehouse.yaml b/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/sql-warehouse.yaml new file mode 100644 index 00000000..a6ce9446 --- /dev/null +++ b/agent-langgraph-short-term-memory/.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-langgraph-short-term-memory/.claude/skills/add-tools/examples/uc-connection.yaml b/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/uc-connection.yaml new file mode 100644 index 00000000..316675fe --- /dev/null +++ b/agent-langgraph-short-term-memory/.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-langgraph-short-term-memory/.claude/skills/add-tools/examples/uc-function.yaml b/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/uc-function.yaml new file mode 100644 index 00000000..43f938a9 --- /dev/null +++ b/agent-langgraph-short-term-memory/.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-langgraph-short-term-memory/.claude/skills/add-tools/examples/vector-search.yaml b/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/vector-search.yaml new file mode 100644 index 00000000..0ba39027 --- /dev/null +++ b/agent-langgraph-short-term-memory/.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-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md new file mode 100644 index 00000000..64b8895c --- /dev/null +++ b/agent-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md @@ -0,0 +1,196 @@ +--- +name: agent-memory +description: "Understand and modify agent memory patterns. Use when: (1) User asks about 'memory', 'state', 'conversation history', (2) Working with thread_id or user_id, (3) Debugging memory issues, (4) Adding long-term memory capabilities." +--- + +# Agent Memory Patterns + +This skill covers both short-term memory (conversation history within a session) and long-term memory (facts that persist across sessions). + +## Short-Term Memory (This Template) + +Short-term memory stores conversation history within a thread. The agent remembers what was said earlier in the conversation. + +### Key Components + +**AsyncCheckpointSaver** - Persists LangGraph state to Lakebase: + +```python +from databricks_langchain import AsyncCheckpointSaver + +async with AsyncCheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME) as checkpointer: + agent = await init_agent(checkpointer=checkpointer) + # Agent now persists state between calls +``` + +**StatefulAgentState** - Custom state schema for memory: + +```python +from typing import Any, Sequence, TypedDict +from langchain_core.messages import AnyMessage +from langgraph.graph.message import add_messages +from typing_extensions import Annotated + +class StatefulAgentState(TypedDict, total=False): + messages: Annotated[Sequence[AnyMessage], add_messages] + custom_inputs: dict[str, Any] + custom_outputs: dict[str, Any] +``` + +**Thread ID** - Identifies a conversation thread: + +```python +def _get_or_create_thread_id(request: ResponsesAgentRequest) -> str: + # Priority: + # 1. Use thread_id from custom_inputs + # 2. Use conversation_id from ChatContext + # 3. Generate random UUID + ci = dict(request.custom_inputs or {}) + + if "thread_id" in ci and ci["thread_id"]: + return str(ci["thread_id"]) + + if request.context and getattr(request.context, "conversation_id", None): + return str(request.context.conversation_id) + + return str(uuid_utils.uuid7()) +``` + +### Creating a Stateful Agent + +Pass `checkpointer` and `state_schema` to `create_agent()`: + +```python +from langchain.agents import create_agent + +agent = create_agent( + model=model, + tools=tools, + system_prompt=SYSTEM_PROMPT, + checkpointer=checkpointer, # Enables persistence + state_schema=StatefulAgentState, # Custom state with messages +) +``` + +### Using Thread ID in Requests + +The agent uses `config` to track threads: + +```python +config = {"configurable": {"thread_id": thread_id}} +input_state = { + "messages": to_chat_completions_input([i.model_dump() for i in request.input]), + "custom_inputs": dict(request.custom_inputs or {}), +} + +async for event in agent.astream(input_state, config, stream_mode=["updates", "messages"]): + # Process events +``` + +--- + +## API Requests with Memory + +### Continuing a Conversation (with thread_id) + +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "What did we discuss?"}], + "custom_inputs": {"thread_id": ""} + }' +``` + +### Starting a New Conversation + +Omit `thread_id` to start fresh (a new UUID will be generated): + +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{"input": [{"role": "user", "content": "Hello!"}], "stream": true}' +``` + +The response includes `thread_id` in `custom_outputs` - save it to continue the conversation. + +### Deployed App with Memory + +```bash +curl -X POST /invocations \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Remember this: my favorite color is blue"}], + "custom_inputs": {"thread_id": "user-123-session-1"} + }' +``` + +--- + +## Long-Term Memory (Reference) + +Long-term memory stores facts that persist across conversation sessions (e.g., user preferences). This is available in the `agent-langgraph-long-term-memory` template. + +### Key Components (Long-Term) + +**AsyncDatabricksStore** - Persists memories to Lakebase: + +```python +from databricks_langchain import AsyncDatabricksStore + +async with AsyncDatabricksStore(instance_name=LAKEBASE_INSTANCE_NAME) as store: + # Store persists memories across sessions +``` + +**User ID** - Identifies a user across sessions: + +```python +user_id = request.custom_inputs.get("user_id", "default-user") +``` + +**Memory Tools** - LangGraph provides tools for memory management: +- `manage_memory` - Save/update memories about the user +- The agent can reference stored memories in conversations + +### Long-Term Memory Request Example + +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "What do you remember about me?"}], + "custom_inputs": { + "thread_id": "", + "user_id": "user-123" + } + }' +``` + +--- + +## Troubleshooting + +| Issue | Solution | +|-------|----------| +| **Conversation not persisting** | Verify `thread_id` is passed and consistent | +| **"Cannot connect to Lakebase"** | See **lakebase-setup** skill | +| **Permission errors** | Grant checkpoint table permissions (see **lakebase-setup** skill) | +| **State not loading** | Check that `checkpointer` is passed to `create_agent()` | + +## Dependencies + +Short-term memory requires `databricks-langchain[memory]`: + +```toml +# pyproject.toml +dependencies = [ + "databricks-langchain[memory]>=0.13.0", +] +``` + +## Next Steps + +- Configure Lakebase: see **lakebase-setup** skill +- Modify agent behavior: see **modify-agent** skill +- Test locally: see **run-locally** skill diff --git a/agent-langgraph-short-term-memory/.claude/skills/deploy/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/deploy/SKILL.md new file mode 100644 index 00000000..ebd80e08 --- /dev/null +++ b/agent-langgraph-short-term-memory/.claude/skills/deploy/SKILL.md @@ -0,0 +1,205 @@ +--- +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 + +> **Memory Template:** Before deploying, ensure Lakebase is configured. See the **lakebase-setup** skill for permissions setup required after deployment. + +## Deploy Commands + +```bash +# Deploy the bundle (creates/updates resources, uploads files) +databricks bundle deploy + +# Run the app (starts/restarts with uploaded source code) +databricks bundle run agent_langgraph_short_term_memory +``` + +The resource key `agent_langgraph_short_term_memory` 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_langgraph_short_term_memory: + 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_langgraph_short_term_memory --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_langgraph_short_term_memory +``` + +### 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. + +## Post-Deployment: Lakebase Permissions + +**After deploying a memory-enabled agent, you must grant Lakebase permissions.** + +See the **lakebase-setup** skill for: +- Adding Lakebase as an app resource +- SDK or SQL commands to grant checkpoint table permissions +- Troubleshooting access errors + +## Unbinding an App + +To remove the link between bundle and deployed app: + +```bash +databricks bundle deployment unbind agent_langgraph_short_term_memory +``` + +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 +``` + +**Send request:** +```bash +curl -X POST /invocations \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' +``` + +For memory-specific request examples (with thread_id), see the **agent-memory** skill. + +## 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' +``` + +## 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 | +|-------|----------| +| Permission errors at runtime | Grant resources in `databricks.yml` (see **add-tools** skill) | +| Lakebase access errors | See **lakebase-setup** skill for permissions | +| 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 | + +## Next Steps + +- Grant Lakebase permissions: see **lakebase-setup** skill +- Understand memory patterns: see **agent-memory** skill +- Add tools and permissions: see **add-tools** skill diff --git a/agent-langgraph-short-term-memory/.claude/skills/discover-tools/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/discover-tools/SKILL.md new file mode 100644 index 00000000..eabac2a5 --- /dev/null +++ b/agent-langgraph-short-term-memory/.claude/skills/discover-tools/SKILL.md @@ -0,0 +1,82 @@ +--- +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}` | + +## Using Discovered Tools in Code + +After discovering tools, add them to your agent in `agent_server/agent.py`: + +```python +from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient + +# Example: Add UC functions from a schema +uc_functions_server = DatabricksMCPServer( + url=f"{host}/api/2.0/mcp/functions/{catalog}/{schema}", + name="my uc functions", +) + +# Example: Add a Genie space +genie_server = DatabricksMCPServer( + url=f"{host}/api/2.0/mcp/genie/{space_id}", + name="my genie space", +) + +# Example: Add vector search +vector_server = DatabricksMCPServer( + url=f"{host}/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name}", + name="my vector index", +) + +# Create MCP client with all servers +mcp_client = DatabricksMultiServerMCPClient([ + uc_functions_server, + genie_server, + vector_server, +]) + +# Get tools for the agent +tools = await mcp_client.get_tools() +``` + +## Next Steps + +After adding MCP servers to your agent: +1. **Grant permissions** in `databricks.yml` (see **add-tools** skill) +2. Test locally with `uv run start-app` (see **run-locally** skill) diff --git a/agent-langgraph-short-term-memory/.claude/skills/lakebase-setup/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/lakebase-setup/SKILL.md new file mode 100644 index 00000000..e37edbd5 --- /dev/null +++ b/agent-langgraph-short-term-memory/.claude/skills/lakebase-setup/SKILL.md @@ -0,0 +1,168 @@ +--- +name: lakebase-setup +description: "Configure Lakebase for agent memory storage. Use when: (1) First-time memory setup, (2) 'Failed to connect to Lakebase' errors, (3) Permission errors on checkpoint tables, (4) User says 'lakebase', 'memory setup', or 'checkpoint'." +--- + +# Lakebase Setup for Memory + +This template uses Lakebase (Databricks-managed PostgreSQL) to store conversation memory. You must configure Lakebase before the agent can persist state. + +## Prerequisites + +- A Lakebase instance in your Databricks workspace +- The instance name (found in Databricks UI under SQL Warehouses > Lakebase) + +## Local Development Setup + +**Step 1:** Add the Lakebase instance name to `.env.local`: + +```bash +LAKEBASE_INSTANCE_NAME= +``` + +**Step 2:** Run `uv run start-app` to test locally. The agent will automatically create checkpoint tables on first run. + +## Deployed App Setup + +After deploying your agent with `databricks bundle deploy`, you must grant the app's service principal access to Lakebase. + +### Step 1: Add Lakebase as App Resource + +1. Go to the Databricks UI +2. Navigate to your app and click **Edit** +3. Go to **App resources** → **Add resource** +4. Add your Lakebase instance with **Connect + Create** permissions + +### Step 2: Get App Service Principal ID + +```bash +databricks apps get --output json | jq -r '.service_principal_id' +``` + +### Step 3: Grant Permissions + +Choose **Option A** (SDK - Recommended) or **Option B** (SQL). + +#### Option A: Using LakebaseClient SDK (Recommended) + +The `databricks-ai-bridge` package includes a `LakebaseClient` for programmatic permission management: + +```python +from databricks_ai_bridge.lakebase import LakebaseClient, SchemaPrivilege, TablePrivilege + +# Initialize client +client = LakebaseClient(instance_name="") + +app_sp = "" # From Step 2 + +# Create role for the service principal +client.create_role(app_sp, "SERVICE_PRINCIPAL") + +# Grant schema privileges +client.grant_schema( + grantee=app_sp, + privileges=[SchemaPrivilege.USAGE, SchemaPrivilege.CREATE], + schemas=["drizzle", "ai_chatbot", "public"], +) + +# Grant table privileges on all tables in schemas +client.grant_all_tables_in_schema( + grantee=app_sp, + privileges=[TablePrivilege.SELECT, TablePrivilege.INSERT, TablePrivilege.UPDATE], + schemas=["drizzle", "ai_chatbot"], +) + +# Grant privileges on specific checkpoint tables +client.grant_table( + grantee=app_sp, + privileges=[TablePrivilege.SELECT, TablePrivilege.INSERT, TablePrivilege.UPDATE], + tables=[ + "public.checkpoint_migrations", + "public.checkpoint_writes", + "public.checkpoints", + "public.checkpoint_blobs", + ], +) +``` + +**Benefits of SDK approach:** +- Type-safe privilege enums prevent typos +- Cleaner Python code vs raw SQL +- Easier to integrate into setup scripts + +#### Option B: Using SQL + +Run the following SQL on your Lakebase instance (replace `app-sp-id` with your app's service principal ID): + +```sql +DO $$ +DECLARE + app_sp text := 'app-sp-id'; -- TODO: Replace with your App's Service Principal ID +BEGIN + ------------------------------------------------------------------- + -- Drizzle schema: migration metadata tables + ------------------------------------------------------------------- + EXECUTE format('GRANT USAGE, CREATE ON SCHEMA drizzle TO %I;', app_sp); + EXECUTE format('GRANT SELECT, INSERT, UPDATE ON ALL TABLES IN SCHEMA drizzle TO %I;', app_sp); + ------------------------------------------------------------------- + -- App schema: business tables (Chat, Message, etc.) + ------------------------------------------------------------------- + EXECUTE format('GRANT USAGE, CREATE ON SCHEMA ai_chatbot TO %I;', app_sp); + EXECUTE format('GRANT SELECT, INSERT, UPDATE ON ALL TABLES IN SCHEMA ai_chatbot TO %I;', app_sp); + ------------------------------------------------------------------- + -- Public schema for checkpoint tables + ------------------------------------------------------------------- + EXECUTE format('GRANT USAGE, CREATE ON SCHEMA public TO %I;', app_sp); + EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.checkpoint_migrations TO %I;', app_sp); + EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.checkpoint_writes TO %I;', app_sp); + EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.checkpoints TO %I;', app_sp); + EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.checkpoint_blobs TO %I;', app_sp); +END $$; +``` + +**Schema Reference:** +| Schema | Purpose | +|--------|---------| +| `drizzle` | Migration metadata tables | +| `ai_chatbot` | Business tables (Chat, Message, etc.) | +| `public` | Checkpoint tables for conversation state | + +## Troubleshooting + +| Issue | Solution | +|-------|----------| +| **"LAKEBASE_INSTANCE_NAME environment variable is required"** | Add `LAKEBASE_INSTANCE_NAME=` to `.env.local` | +| **"Failed to connect to Lakebase instance"** | Verify instance name is correct and your profile has access | +| **Permission errors on checkpoint tables** | Run the SDK script or SQL grant commands above | +| **Deployed app can't access Lakebase** | Add Lakebase as app resource in Databricks UI | +| **"role does not exist"** | Run `client.create_role()` first, or ensure SP ID is correct | + +### Common Error Messages + +**Local development:** +``` +Failed to connect to Lakebase instance ''. Please verify: +1. The instance name is correct +2. You have the necessary permissions to access the instance +3. Your Databricks authentication is configured correctly +``` + +**Deployed app:** +``` +Failed to connect to Lakebase instance ''. The App Service Principal for '' may not have access. +``` + +Both errors indicate Lakebase access issues. Follow the setup steps above for your environment. + +## How Memory Works + +See the **agent-memory** skill for: +- How `AsyncCheckpointSaver` persists conversation state +- Using `thread_id` to maintain conversation context +- API request examples with thread_id + +## Next Steps + +- Understand memory patterns: see **agent-memory** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/agent-langgraph-short-term-memory/.claude/skills/modify-agent/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/modify-agent/SKILL.md new file mode 100644 index 00000000..575080d4 --- /dev/null +++ b/agent-langgraph-short-term-memory/.claude/skills/modify-agent/SKILL.md @@ -0,0 +1,296 @@ +--- +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 + +> **Memory Template:** This template uses short-term memory (conversation history). See the **agent-memory** skill for memory-specific patterns. + +## Main File + +**`agent_server/agent.py`** - Agent logic, model selection, instructions, MCP servers + +## Key Files + +| File | Purpose | +|------|---------| +| `agent_server/agent.py` | Agent logic, model, instructions, MCP servers | +| `agent_server/start_server.py` | FastAPI server + MLflow setup | +| `agent_server/evaluate_agent.py` | Agent evaluation with MLflow scorers | +| `agent_server/utils.py` | Databricks auth helpers, stream processing | +| `databricks.yml` | Bundle config & resource permissions | + +## SDK Setup + +```python +import mlflow +from databricks.sdk import WorkspaceClient +from databricks_langchain import ChatDatabricks, DatabricksMCPServer, DatabricksMultiServerMCPClient +from langchain.agents import create_agent + +# Enable autologging for tracing +mlflow.langchain.autolog() + +# Initialize workspace client +workspace_client = WorkspaceClient() +``` + +--- + +## databricks-langchain SDK Overview + +**SDK Location:** https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/langchain + +Before making any changes, ensure that the APIs actually exist in the SDK. If something is missing from the documentation here, look in the venv's `site-packages` directory for the `databricks_langchain` package. If it's not installed, run `uv sync` to create the .venv and install the package. + +--- + +### ChatDatabricks - LLM Chat Interface + +Connects to Databricks Model Serving endpoints for LLM inference. + +```python +from databricks_langchain import ChatDatabricks + +llm = ChatDatabricks( + endpoint="databricks-claude-3-7-sonnet", # or databricks-meta-llama-3-1-70b-instruct + temperature=0, + max_tokens=500, +) + +# For Responses API agents: +llm = ChatDatabricks(endpoint="my-agent-endpoint", use_responses_api=True) +``` + +Available models (check workspace for current list): +- `databricks-claude-3-7-sonnet` +- `databricks-claude-3-5-sonnet` +- `databricks-meta-llama-3-3-70b-instruct` + +**Note:** Some workspaces require granting the app access to the serving endpoint in `databricks.yml`. See the **add-tools** skill and `examples/serving-endpoint.yaml`. + +--- + +### DatabricksEmbeddings - Generate Embeddings + +Query Databricks embedding model endpoints. + +```python +from databricks_langchain import DatabricksEmbeddings + +embeddings = DatabricksEmbeddings(endpoint="databricks-bge-large-en") +vector = embeddings.embed_query("The meaning of life is 42") +vectors = embeddings.embed_documents(["doc1", "doc2"]) +``` + +--- + +### DatabricksVectorSearch - Vector Store + +Connect to Databricks Vector Search indexes for similarity search. + +```python +from databricks_langchain import DatabricksVectorSearch + +# Delta-sync index with Databricks-managed embeddings +vs = DatabricksVectorSearch(index_name="catalog.schema.index_name") + +# Direct-access or self-managed embeddings +vs = DatabricksVectorSearch( + index_name="catalog.schema.index_name", + embedding=embeddings, + text_column="content", +) + +docs = vs.similarity_search("query", k=5) +``` + +--- + +### MCP Client - Tool Integration + +Connect to MCP (Model Context Protocol) servers to get tools for your agent. + +**Basic MCP Server (manual URL):** + +```python +from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient + +client = DatabricksMultiServerMCPClient([ + DatabricksMCPServer( + name="system-ai", + url=f"{host}/api/2.0/mcp/functions/system/ai", + ) +]) +tools = await client.get_tools() +``` + +**From UC Function (convenience helper):** + +Creates MCP server for Unity Catalog functions. If `function_name` is omitted, exposes all functions in the schema. + +```python +server = DatabricksMCPServer.from_uc_function( + catalog="main", + schema="tools", + function_name="send_email", # Optional - omit for all functions in schema + name="email-server", + timeout=30.0, + handle_tool_error=True, +) +``` + +**From Vector Search (convenience helper):** + +Creates MCP server for Vector Search indexes. If `index_name` is omitted, exposes all indexes in the schema. + +```python +server = DatabricksMCPServer.from_vector_search( + catalog="main", + schema="embeddings", + index_name="product_docs", # Optional - omit for all indexes in schema + name="docs-search", + timeout=30.0, +) +``` + +**From Genie Space:** + +Create MCP server from Genie Space. Get the genie space ID from the URL. + +Example: `https://workspace.cloud.databricks.com/genie/rooms/01f0515f6739169283ef2c39b7329700?o=123` means the genie space ID is `01f0515f6739169283ef2c39b7329700` + +```python +DatabricksMCPServer( + name="genie", + url=f"{host_name}/api/2.0/mcp/genie/01f0515f6739169283ef2c39b7329700", +) +``` + +**Non-Databricks MCP Server:** + +```python +from databricks_langchain import MCPServer + +server = MCPServer( + name="external-server", + url="https://other-server.com/mcp", + headers={"X-API-Key": "secret"}, + timeout=15.0, +) +``` + +**After adding MCP servers:** Grant permissions in `databricks.yml` (see **add-tools** skill) + +--- + +## Running the Agent + +```python +from langchain.agents import create_agent + +# Create agent - ONLY accepts tools and model, NO prompt/instructions parameter +agent = create_agent(tools=tools, model=llm) + +# Non-streaming +messages = {"messages": [{"role": "user", "content": "hi"}]} +result = await agent.ainvoke(messages) + +# Streaming +async for event in agent.astream(input=messages, stream_mode=["updates", "messages"]): + # Process stream events + pass +``` + +**Converting to Responses API format:** Use `process_agent_astream_events()` from `agent_server/utils.py`: + +```python +from agent_server.utils import process_agent_astream_events + +async for event in process_agent_astream_events( + agent.astream(input=messages, stream_mode=["updates", "messages"]) +): + yield event # Yields ResponsesAgentStreamEvent objects +``` + +--- + +## Customizing Agent Behavior (System Instructions) + +> **IMPORTANT:** `create_agent()` does NOT accept `prompt`, `instructions`, or `system_message` parameters. Attempting to pass these will cause a runtime error. + +In LangGraph, agent behavior is customized by prepending a system message to the conversation messages. + +**Correct pattern in `agent.py`:** + +1. Define instructions as a constant: +```python +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.""" +``` + +2. Prepend to messages in the `streaming()` function: +```python +@stream() +async def streaming(request: ResponsesAgentRequest) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + agent = await init_agent() + # Prepend system instructions to user messages + user_messages = to_chat_completions_input([i.model_dump() for i in request.input]) + messages = {"messages": [{"role": "system", "content": AGENT_INSTRUCTIONS}] + user_messages} + + async for event in process_agent_astream_events( + agent.astream(input=messages, stream_mode=["updates", "messages"]) + ): + yield event +``` + +**Common mistake to avoid:** +```python +# WRONG - will cause "unexpected keyword argument" error +agent = create_agent(tools=tools, model=llm, prompt=AGENT_INSTRUCTIONS) + +# CORRECT - add instructions via messages +messages = {"messages": [{"role": "system", "content": AGENT_INSTRUCTIONS}] + user_messages} +``` + +For advanced customization (routing, state management, custom graphs), refer to the [LangGraph documentation](https://docs.langchain.com/oss/python/langgraph/overview). + +--- + +## External Connection Tools + +Connect to external services via Unity Catalog HTTP connections: + +- **Slack** - Post messages to channels +- **Google Calendar** - Calendar operations +- **Microsoft Graph API** - Office 365 services +- **Azure AI Search** - Search functionality +- **Any HTTP API** - Use `http_request` from databricks-sdk + +Example: Create UC function wrapping HTTP request for Slack, then expose via MCP. + +--- + +## External Resources + +1. [databricks-langchain SDK](https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/langchain) +2. [Agent examples](https://github.com/bbqiu/agent-on-app-prototype) +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. [LangGraph documentation](https://docs.langchain.com/oss/python/langgraph/overview) +6. [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 +- Memory patterns: see **agent-memory** skill +- Lakebase setup: see **lakebase-setup** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/agent-langgraph-short-term-memory/.claude/skills/quickstart/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/quickstart/SKILL.md new file mode 100644 index 00000000..edfbfd33 --- /dev/null +++ b/agent-langgraph-short-term-memory/.claude/skills/quickstart/SKILL.md @@ -0,0 +1,84 @@ +--- +name: quickstart +description: "Set up Databricks agent development environment. Use when: (1) First time setup, (2) Configuring Databricks authentication, (3) User says 'quickstart', 'set up', 'authenticate', or 'configure databricks', (4) No .env.local file exists." +--- + +# Quickstart & Authentication + +## Prerequisites + +- **uv** (Python package manager) +- **nvm** with Node 20 (for frontend) +- **Databricks CLI v0.283.0+** + +Check CLI version: +```bash +databricks -v # Must be v0.283.0 or above +brew upgrade databricks # If version is too old +``` + +## Run Quickstart + +```bash +uv run quickstart +``` + +**Options:** +- `--profile NAME`: Use specified profile (non-interactive) +- `--host URL`: Workspace URL for initial setup +- `-h, --help`: Show help + +**Examples:** +```bash +# Interactive (prompts for profile selection) +uv run quickstart + +# Non-interactive with existing profile +uv run quickstart --profile DEFAULT + +# New workspace setup +uv run quickstart --host https://your-workspace.cloud.databricks.com +``` + +## What Quickstart Configures + +Creates/updates `.env.local` with: +- `DATABRICKS_CONFIG_PROFILE` - Selected CLI profile +- `MLFLOW_TRACKING_URI` - Set to `databricks://` for local auth +- `MLFLOW_EXPERIMENT_ID` - Auto-created experiment ID + +## Manual Authentication (Fallback) + +If quickstart fails: + +```bash +# Create new profile +databricks auth login --host https://your-workspace.cloud.databricks.com + +# Verify +databricks auth profiles +``` + +Then manually create `.env.local` (copy from `.env.example`): +```bash +# Authentication (choose one method) +DATABRICKS_CONFIG_PROFILE=DEFAULT +# DATABRICKS_HOST=https://.databricks.com +# DATABRICKS_TOKEN=dapi.... + +# MLflow configuration +MLFLOW_EXPERIMENT_ID= +MLFLOW_TRACKING_URI="databricks://DEFAULT" +MLFLOW_REGISTRY_URI="databricks-uc" + +# Frontend proxy settings +CHAT_APP_PORT=3000 +CHAT_PROXY_TIMEOUT_SECONDS=300 +``` + +## Next Steps + +After quickstart completes: +1. **Configure Lakebase** for memory storage (see **lakebase-setup** skill) +2. Run `uv run discover-tools` to find available workspace resources (see **discover-tools** skill) +3. Run `uv run start-app` to test locally (see **run-locally** skill) diff --git a/agent-langgraph-short-term-memory/.claude/skills/run-locally/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/run-locally/SKILL.md new file mode 100644 index 00000000..f3b517de --- /dev/null +++ b/agent-langgraph-short-term-memory/.claude/skills/run-locally/SKILL.md @@ -0,0 +1,130 @@ +--- +name: run-locally +description: "Run and test the agent locally. Use when: (1) User says 'run locally', 'start server', 'test agent', or 'localhost', (2) Need curl commands to test API, (3) Troubleshooting local development issues, (4) Configuring server options like port or hot-reload." +--- + +# Run Agent Locally + +## Start the Server + +```bash +uv run start-app +``` + +This starts the agent at http://localhost:8000 with both the API server and chat UI. + +## Server Options + +| Option | Command | When to Use | +|--------|---------|-------------| +| **Hot-reload** | `uv run start-server --reload` | During development - auto-restarts on code changes | +| **Custom port** | `uv run start-server --port 8001` | When port 8000 is in use | +| **Multiple workers** | `uv run start-server --workers 4` | Load testing or production-like simulation | +| **Combined** | `uv run start-server --reload --port 8001` | Development on alternate port | + +## Test the API + +The agent implements the [OpenAI Responses API](https://platform.openai.com/docs/api-reference/responses) interface via MLflow's ResponsesAgent. + +**Streaming request:** +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' +``` + +**Non-streaming request:** +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }] }' +``` + +**Request with thread_id (for conversation memory):** +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "What did we discuss?"}], + "custom_inputs": {"thread_id": ""} + }' +``` + +See the **agent-memory** skill for more memory-related request examples. + +## Run Evaluation + +Evaluate your agent using MLflow scorers: + +```bash +uv run agent-evaluate +``` + +**What it does:** +- Runs the agent against a test dataset +- Applies MLflow scorers (RelevanceToQuery, Safety) +- Records results to your MLflow experiment + +**Customize evaluation:** +Edit `agent_server/evaluate_agent.py` to: +- Change the evaluation dataset +- Add or modify scorers +- Adjust evaluation parameters + +After evaluation completes, open the MLflow UI link for your experiment to inspect results. + +## Run Unit Tests + +```bash +pytest [path] +``` + +## Adding Dependencies + +```bash +uv add +# Example: uv add "langchain-community" +``` + +Dependencies are managed in `pyproject.toml`. + +## Troubleshooting + +| Issue | Solution | +|-------|----------| +| **Port already in use** | Use `--port 8001` or kill existing process | +| **Authentication errors** | Verify `.env.local` is correct; run **quickstart** skill | +| **Module not found** | Run `uv sync` to install dependencies | +| **MLflow experiment not found** | Ensure `MLFLOW_TRACKING_URI` in `.env.local` is `databricks://` | +| **Lakebase connection errors** | Verify `LAKEBASE_INSTANCE_NAME` in `.env.local`; see **lakebase-setup** skill | + +### MLflow Experiment Not Found + +If you see: "The provided MLFLOW_EXPERIMENT_ID environment variable value does not exist" + +**Verify the experiment exists:** +```bash +databricks -p experiments get-experiment +``` + +**Fix:** Ensure `.env.local` has the correct tracking URI format: +```bash +MLFLOW_TRACKING_URI="databricks://DEFAULT" # Include profile name +``` + +The quickstart script configures this automatically. If you manually edited `.env.local`, ensure the profile name is included. + +## MLflow Tracing + +This template uses MLflow for automatic tracing: +- Agent logic decorated with `@invoke()` and `@stream()` is automatically traced +- LLM invocations are captured via MLflow autologging + +To add custom trace instrumentation, see: [MLflow tracing documentation](https://docs.databricks.com/aws/en/mlflow3/genai/tracing/app-instrumentation/) + +## Next Steps + +- Configure Lakebase: see **lakebase-setup** skill +- Understand memory patterns: see **agent-memory** skill +- Modify your agent: see **modify-agent** skill +- Deploy to Databricks: see **deploy** skill diff --git a/agent-langgraph-short-term-memory/.gitignore b/agent-langgraph-short-term-memory/.gitignore index a058c65c..f12e1ea0 100644 --- a/agent-langgraph-short-term-memory/.gitignore +++ b/agent-langgraph-short-term-memory/.gitignore @@ -204,6 +204,18 @@ sketch **/mlruns/ **/.vite/ **/.databricks -**/.claude **/.env -**/.env.local \ No newline at end of file +**/.env.local + +# Claude Code - track skills only +.claude/* +!.claude/skills/ +.claude/skills/* +!.claude/skills/quickstart/ +!.claude/skills/discover-tools/ +!.claude/skills/deploy/ +!.claude/skills/add-tools/ +!.claude/skills/run-locally/ +!.claude/skills/modify-agent/ +!.claude/skills/lakebase-setup/ +!.claude/skills/agent-memory/ \ No newline at end of file diff --git a/agent-langgraph-short-term-memory/AGENTS.md b/agent-langgraph-short-term-memory/AGENTS.md index 617c2c28..3d5f45c6 100644 --- a/agent-langgraph-short-term-memory/AGENTS.md +++ b/agent-langgraph-short-term-memory/AGENTS.md @@ -1,348 +1,109 @@ -# Agent LangGraph Development Guide +# Agent Development Guide -## Running the App +## MANDATORY First Action -**Prerequisites:** uv, nvm (Node 20), Databricks CLI +**BEFORE any other action, run `databricks auth profiles` to check authentication status.** -**Quick Start:** +This helps you understand: +- Which Databricks profiles are configured +- Whether authentication is already set up +- Which profile to use for subsequent commands -```bash -uv run quickstart # First-time setup (auth, MLflow experiment, env) -uv run start-app # Start app at http://localhost:8000 -``` +If no profiles exist, guide the user through running `uv run quickstart` to set up authentication. -**Advanced Server Options:** +## Understanding User Goals -```bash -uv run start-server --reload # Hot-reload on code changes during development -uv run start-server --port 8001 -uv run start-server --workers 4 -``` +**Ask the user questions to understand what they're building:** -**Test API:** +1. **What is the agent's purpose?** (e.g., data analyst assistant, customer support, code helper) +2. **What data or tools does it need access to?** + - Databases/tables (Unity Catalog) + - Documents for RAG (Vector Search) + - Natural language data queries (Genie Spaces) + - External APIs or services +3. **Any specific Databricks resources they want to connect?** -```bash -# Streaming request -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' +Use `uv run discover-tools` to show them available resources in their workspace, then help them select the right ones for their use case. **See the `add-tools` skill for how to connect tools and grant permissions.** -# Non-streaming request -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ "input": [{ "role": "user", "content": "hi" }] }' -``` +## Memory Template Note ---- - -## Testing the Agent - -**Run evaluation:** - -```bash -uv run agent-evaluate # Uses MLflow scorers (RelevanceToQuery, Safety) -``` - -**Run unit tests:** - -```bash -pytest [path] # Standard pytest execution -``` - ---- - -## Modifying the Agent - -Anytime the user wants to modify the agent, look through each of the following resources to help them accomplish their goal: - -If the user wants to convert something into Responses API, refer to https://mlflow.org/docs/latest/genai/serving/responses-agent/ for more information. +This template includes **short-term memory** (conversation history within a session). The agent remembers what was said earlier in the same conversation thread. -1. Look through existing databricks-langchain APIs to see if they can use one of these to accomplish their goal. -2. Look through the folders in https://github.com/bbqiu/agent-on-app-prototype to see if there's an existing example similar to what they're looking to do. -3. Reference the documentation available under https://docs.databricks.com/aws/en/generative-ai/agent-framework/ and its subpages. -4. For adding tools and capabilities, refer to: https://docs.databricks.com/aws/en/generative-ai/agent-framework/agent-tool -5. For stuff like LangGraph routing, configuration, and customization, refer to the LangGraph documentation: https://docs.langchain.com/oss/python/langgraph/overview. +**Required setup:** +1. Configure Lakebase instance (see **lakebase-setup** skill) +2. Use `thread_id` in requests to maintain conversation context (see **agent-memory** skill) -**Main file to modify:** `agent_server/agent.py` +## Handling Deployment Errors ---- - -## databricks-langchain SDK overview - -**SDK Location:** `https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/langchain` +**If `databricks bundle deploy` fails with "An app with the same name already exists":** -**Development Workflow:** +Ask the user: "I see there's an existing app with the same name. Would you like me to bind it to this bundle so we can manage it, or delete it and create a new one?" -```bash -uv add databricks-langchain -``` - -Before making any changes, ensure that the APIs actually exist in the SDK. If something is missing from the documentation here, feel free to look in the venv's `site-packages` directory for the `databricks_langchain` package. If it's not installed, run `uv sync` in this folder to create the .venv and install the package. +- **If they want to bind**: See the **deploy** skill for binding steps +- **If they want to delete**: Run `databricks apps delete ` then deploy again --- -### ChatDatabricks - LLM Chat Interface - -Connects to Databricks Model Serving endpoints for LLM inference. +## Available Skills -```python -from databricks_langchain import ChatDatabricks +**Before executing any task, read the relevant skill file in `.claude/skills/`** - they contain tested commands, patterns, and troubleshooting steps. -llm = ChatDatabricks( - endpoint="databricks-claude-3-7-sonnet", # or databricks-meta-llama-3-1-70b-instruct - temperature=0, - max_tokens=500, -) +| Task | Skill | Path | +|------|-------|------| +| Setup, auth, first-time | **quickstart** | `.claude/skills/quickstart/SKILL.md` | +| Lakebase configuration | **lakebase-setup** | `.claude/skills/lakebase-setup/SKILL.md` | +| Memory patterns | **agent-memory** | `.claude/skills/agent-memory/SKILL.md` | +| Find tools/resources | **discover-tools** | `.claude/skills/discover-tools/SKILL.md` | +| Deploy to Databricks | **deploy** | `.claude/skills/deploy/SKILL.md` | +| Add tools & permissions | **add-tools** | `.claude/skills/add-tools/SKILL.md` | +| Run/test locally | **run-locally** | `.claude/skills/run-locally/SKILL.md` | +| Modify agent code | **modify-agent** | `.claude/skills/modify-agent/SKILL.md` | -# For Responses API agents: -llm = ChatDatabricks(endpoint="my-agent-endpoint", use_responses_api=True) -``` +**Note:** All agent skills are located in `.claude/skills/` directory. --- -### DatabricksEmbeddings - Generate Embeddings - -Query Databricks embedding model endpoints. +## Quick Commands -```python -from databricks_langchain import DatabricksEmbeddings - -embeddings = DatabricksEmbeddings(endpoint="databricks-bge-large-en") -vector = embeddings.embed_query("The meaning of life is 42") -vectors = embeddings.embed_documents(["doc1", "doc2"]) -``` +| Task | Command | +|------|---------| +| Setup | `uv run quickstart` | +| Discover tools | `uv run discover-tools` | +| Run locally | `uv run start-app` | +| Deploy | `databricks bundle deploy && databricks bundle run agent_langgraph_short_term_memory` | +| View logs | `databricks apps logs --follow` | --- -### DatabricksVectorSearch - Vector Store - -Connect to Databricks Vector Search indexes for similarity search. - -```python -from databricks_langchain import DatabricksVectorSearch - -# Delta-sync index with Databricks-managed embeddings -vs = DatabricksVectorSearch(index_name="catalog.schema.index_name") - -# Direct-access or self-managed embeddings -vs = DatabricksVectorSearch( - index_name="catalog.schema.index_name", - embedding=embeddings, - text_column="content", -) +## Key Files -docs = vs.similarity_search("query", k=5) -``` +| File | Purpose | +|------|---------| +| `agent_server/agent.py` | Agent logic, model, instructions, MCP servers, memory setup | +| `agent_server/start_server.py` | FastAPI server + MLflow setup | +| `agent_server/evaluate_agent.py` | Agent evaluation with MLflow scorers | +| `databricks.yml` | Bundle config & resource permissions | +| `scripts/quickstart.py` | One-command setup script | +| `scripts/discover_tools.py` | Discovers available workspace resources | --- -### MCP Client - Tool Integration - -Connect to MCP (Model Context Protocol) servers to get tools for your agent. - -**Basic MCP Server (manual URL):** - -```python -from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient - -client = DatabricksMultiServerMCPClient([ - DatabricksMCPServer( - name="system-ai", - url=f"{host}/api/2.0/mcp/functions/system/ai", - ) -]) -tools = await client.get_tools() -``` - -**From UC Function (convenience helper):** -Creates MCP server for Unity Catalog functions. If `function_name` is omitted, exposes all functions in the schema. - -```python -server = DatabricksMCPServer.from_uc_function( - catalog="main", - schema="tools", - function_name="send_email", # Optional - omit for all functions in schema - name="email-server", - timeout=30.0, - handle_tool_error=True, -) -``` - -**From Vector Search (convenience helper):** -Creates MCP server for Vector Search indexes. If `index_name` is omitted, exposes all indexes in the schema. - -```python -server = DatabricksMCPServer.from_vector_search( - catalog="main", - schema="embeddings", - index_name="product_docs", # Optional - omit for all indexes in schema - name="docs-search", - timeout=30.0, -) -``` - -**From Genie Space:** -Create MCP server from Genie Space. Need to get the genie space ID. Can prompt the user to retrieve this via the UI by getting the link to the genie space. - -Ex: https://db-ml-models-dev-us-west.cloud.databricks.com/genie/rooms/01f0515f6739169283ef2c39b7329700?o=3217006663075879 means the genie space ID is 01f0515f6739169283ef2c39b7329700 - -```python -DatabricksMCPServer( - name="genie", - url=f"{host_name}/api/2.0/mcp/genie/01f0515f6739169283ef2c39b7329700", -), -``` - -**Non-Databricks MCP Server:** - -```python -from databricks_langchain import MCPServer - -server = MCPServer( - name="external-server", - url="https://other-server.com/mcp", - headers={"X-API-Key": "secret"}, - timeout=15.0, -) -``` - -### Stateful LangGraph agent - -To enable statefulness in a LangGraph agent, we need to install `databricks-langchain[memory]`. - -Look through the package files for the latest on stateful langgraph agents. Can start by looking at the databricks_langchain/checkpoints.py and databricks_langchain/store.py files. - -## Lakebase instance setup for stateful agents - -Add the lakebase name to `.env`: - -```bash -LAKEBASE_INSTANCE_NAME= -``` - ## Agent Framework Capabilities -Reference: https://docs.databricks.com/aws/en/generative-ai/agent-framework/ - -### Tool Types +> **IMPORTANT:** When adding any tool to the agent, you MUST also grant permissions in `databricks.yml`. See the **add-tools** skill for required steps and examples. +**Tool Types:** 1. **Unity Catalog Function Tools** - SQL UDFs managed in UC with built-in governance 2. **Agent Code Tools** - Defined directly in agent code for REST APIs and low-latency operations 3. **MCP Tools** - Interoperable tools via Model Context Protocol (Databricks-managed, external, or self-hosted) -### Built-in Tools - +**Built-in Tools:** - **system.ai.python_exec** - Execute Python code dynamically within agent queries (code interpreter) -### External Connection Tools - -Connect to external services via Unity Catalog HTTP connections: - -- **Slack** - Post messages to channels -- **Google Calendar** - Calendar operations -- **Microsoft Graph API** - Office 365 services -- **Azure AI Search** - Search functionality -- **Any HTTP API** - Use `http_request` from databricks-sdk - -Example: Create UC function wrapping HTTP request for Slack, then expose via MCP. - -### Common Patterns - +**Common Patterns:** - **Structured data retrieval** - Query SQL tables/databases - **Unstructured data retrieval** - Document search and RAG via Vector Search - **Code interpreter** - Python execution for analysis via system.ai.python_exec - **External connections** - Integrate services like Slack via HTTP connections ---- - -## Authentication Setup - -**Option 1: OAuth (Recommended)** - -```bash -databricks auth login -``` - -Set in `.env`: - -```bash -DATABRICKS_CONFIG_PROFILE=DEFAULT -``` - -**Option 2: Personal Access Token** - -Set in `.env`: - -```bash -DATABRICKS_HOST="https://host.databricks.com" -DATABRICKS_TOKEN="dapi_token" -``` - ---- - -## MLflow Experiment Setup - -Create and link an MLflow experiment: - -```bash -DATABRICKS_USERNAME=$(databricks current-user me | jq -r .userName) -databricks experiments create-experiment /Users/$DATABRICKS_USERNAME/agents-on-apps -``` - -Add the experiment ID to `.env`: - -```bash -MLFLOW_EXPERIMENT_ID= -``` - ---- - -## Key Files - -| File | Purpose | -| -------------------------------- | --------------------------------------------- | -| `agent_server/agent.py` | Agent logic, model, instructions, MCP servers | -| `agent_server/start_server.py` | FastAPI server + MLflow setup | -| `agent_server/evaluate_agent.py` | Agent evaluation with MLflow scorers | -| `agent_server/utils.py` | Databricks auth helpers, stream processing | -| `scripts/start_app.py` | Manages backend+frontend startup | - ---- - -## Deploying to Databricks Apps - -**Create app:** - -```bash -databricks apps create agent-langgraph -``` - -**Sync files:** - -```bash -DATABRICKS_USERNAME=$(databricks current-user me | jq -r .userName) -databricks sync . "/Users/$DATABRICKS_USERNAME/agent-langgraph" -``` - -**Deploy:** - -```bash -databricks apps deploy agent-langgraph --source-code-path /Workspace/Users/$DATABRICKS_USERNAME/agent-langgraph -``` - -**Query deployed app:** - -Generate OAuth token (PATs are not supported): - -```bash -databricks auth token -``` - -Send request: - -```bash -curl -X POST /invocations \ - -H "Authorization: Bearer " \ - -H "Content-Type: application/json" \ - -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' -``` +Reference: https://docs.databricks.com/aws/en/generative-ai/agent-framework/ diff --git a/agent-langgraph-short-term-memory/README.md b/agent-langgraph-short-term-memory/README.md index f1a57615..307882bc 100644 --- a/agent-langgraph-short-term-memory/README.md +++ b/agent-langgraph-short-term-memory/README.md @@ -1,6 +1,20 @@ -# Responses API Agent +# Responses API Agent (Short-Term Memory) -This template defines a conversational agent app. The app comes with a built-in chat UI, but also exposes an API endpoint for invoking the agent so that you can serve your UI elsewhere (e.g. on your website or in a mobile app). +## Build with AI Assistance + +This template includes Claude Code skills in `.claude/skills/` for AI-assisted development. Use [Claude Code](https://docs.anthropic.com/en/docs/claude-code) to: + +- **Set up your environment**: "Run quickstart to configure authentication" +- **Add tools**: "Connect my agent to a Genie space" +- **Configure memory**: "Set up Lakebase for conversation history" +- **Deploy**: "Deploy my agent to Databricks Apps" +- **Debug**: "Why am I getting a permission error?" + +The skills contain tested commands, code patterns, and troubleshooting steps. + +--- + +This template defines a conversational agent app with short-term memory. The app comes with a built-in chat UI, but also exposes an API endpoint for invoking the agent so that you can serve your UI elsewhere (e.g. on your website or in a mobile app). The agent in this template implements the [OpenAI Responses API](https://platform.openai.com/docs/api-reference/responses) interface. It has access to a single tool; the [built-in code interpreter tool](https://docs.databricks.com/aws/en/generative-ai/agent-framework/code-interpreter-tools#built-in-python-executor-tool) (`system.ai.python_exec`) on Databricks. You can customize agent code and test it via the API or UI. diff --git a/agent-langgraph-short-term-memory/pyproject.toml b/agent-langgraph-short-term-memory/pyproject.toml index ed038488..ae46c70b 100644 --- a/agent-langgraph-short-term-memory/pyproject.toml +++ b/agent-langgraph-short-term-memory/pyproject.toml @@ -36,3 +36,4 @@ quickstart = "scripts.quickstart:main" start-app = "scripts.start_app:main" start-server = "agent_server.start_server:main" agent-evaluate = "agent_server.evaluate_agent:evaluate" +discover-tools = "scripts.discover_tools:main" diff --git a/agent-langgraph-short-term-memory/scripts/discover_tools.py b/agent-langgraph-short-term-memory/scripts/discover_tools.py new file mode 100644 index 00000000..3eb37963 --- /dev/null +++ b/agent-langgraph-short-term-memory/scripts/discover_tools.py @@ -0,0 +1,432 @@ +#!/usr/bin/env python3 +""" +Discover available tools and data sources for Databricks agents. + +This script scans for: +- Unity Catalog functions (data retrieval tools e.g. SQL UDFs) +- Unity Catalog tables (data sources) +- Vector search indexes (RAG data sources) +- Genie spaces (conversational interface over structured data) +- Custom MCP servers (Databricks apps with name mcp-*) +- External MCP servers (via Unity Catalog connections) +""" + +import json +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict, List + +from databricks.sdk import WorkspaceClient + +DEFAULT_MAX_RESULTS = 100 +DEFAULT_MAX_SCHEMAS = 25 + +def run_databricks_cli(args: List[str]) -> str: + """Run databricks CLI command and return output.""" + try: + result = subprocess.run( + ["databricks"] + args, + capture_output=True, + text=True, + check=True, + ) + return result.stdout + except subprocess.CalledProcessError as e: + print(f"Error running databricks CLI: {e.stderr}", file=sys.stderr) + return "" + + +def discover_uc_functions(w: WorkspaceClient, catalog: str = None, max_schemas: int = DEFAULT_MAX_SCHEMAS) -> List[Dict[str, Any]]: + """Discover Unity Catalog functions that could be used as tools. + + Args: + w: WorkspaceClient instance + catalog: Optional specific catalog to search + max_schemas: Total number of schemas to search across all catalogs + """ + functions = [] + schemas_searched = 0 + + try: + catalogs = [catalog] if catalog else [c.name for c in w.catalogs.list()] + + for cat in catalogs: + if schemas_searched >= max_schemas: + break + + try: + all_schemas = list(w.schemas.list(catalog_name=cat)) + # Take schemas from this catalog until we hit the global budget + schemas_to_search = all_schemas[:max_schemas - schemas_searched] + + for schema in schemas_to_search: + schema_name = f"{cat}.{schema.name}" + try: + funcs = list(w.functions.list(catalog_name=cat, schema_name=schema.name)) + for func in funcs: + functions.append({ + "type": "uc_function", + "name": func.full_name, + "catalog": cat, + "schema": schema.name, + "function_name": func.name, + "comment": func.comment, + "routine_definition": getattr(func, "routine_definition", None), + }) + except Exception as e: + # Skip schemas we can't access + continue + finally: + schemas_searched += 1 + except Exception as e: + # Skip catalogs we can't access + continue + + except Exception as e: + print(f"Error discovering UC functions: {e}", file=sys.stderr) + + return functions + + +def discover_uc_tables(w: WorkspaceClient, catalog: str = None, schema: str = None, max_schemas: int = DEFAULT_MAX_SCHEMAS) -> List[Dict[str, Any]]: + """Discover Unity Catalog tables that could be queried. + + Args: + w: WorkspaceClient instance + catalog: Optional specific catalog to search + schema: Optional specific schema to search (requires catalog) + max_schemas: Total number of schemas to search across all catalogs + """ + tables = [] + schemas_searched = 0 + + try: + catalogs = [catalog] if catalog else [c.name for c in w.catalogs.list()] + + for cat in catalogs: + if cat in ["__databricks_internal", "system"]: + continue + + if schemas_searched >= max_schemas: + break + + try: + if schema: + schemas_to_search = [schema] + else: + all_schemas = [s.name for s in w.schemas.list(catalog_name=cat)] + # Take schemas from this catalog until we hit the global budget + schemas_to_search = all_schemas[:max_schemas - schemas_searched] + + for sch in schemas_to_search: + if sch == "information_schema": + schemas_searched += 1 + continue + + try: + tbls = list(w.tables.list(catalog_name=cat, schema_name=sch)) + for tbl in tbls: + # Get column info + columns = [] + if hasattr(tbl, "columns") and tbl.columns: + columns = [ + {"name": col.name, "type": col.type_name.value if hasattr(col.type_name, "value") else str(col.type_name)} + for col in tbl.columns + ] + + tables.append({ + "type": "uc_table", + "name": tbl.full_name, + "catalog": cat, + "schema": sch, + "table_name": tbl.name, + "table_type": tbl.table_type.value if tbl.table_type else None, + "comment": tbl.comment, + "columns": columns, + }) + except Exception as e: + # Skip schemas we can't access + pass + finally: + schemas_searched += 1 + except Exception as e: + # Skip catalogs we can't access + continue + + except Exception as e: + print(f"Error discovering UC tables: {e}", file=sys.stderr) + + return tables + + +def discover_vector_search_indexes(w: WorkspaceClient) -> List[Dict[str, Any]]: + """Discover Vector Search indexes for RAG applications.""" + indexes = [] + + try: + # List all vector search endpoints + endpoints = list(w.vector_search_endpoints.list_endpoints()) + + for endpoint in endpoints: + try: + # List indexes for each endpoint + endpoint_indexes = list(w.vector_search_indexes.list_indexes(endpoint_name=endpoint.name)) + for idx in endpoint_indexes: + indexes.append({ + "type": "vector_search_index", + "name": idx.name, + "endpoint": endpoint.name, + "primary_key": idx.primary_key, + "index_type": idx.index_type.value if idx.index_type else None, + "status": idx.status.state.value if idx.status and idx.status.state else None, + }) + except Exception as e: + # Skip endpoints we can't access + continue + + except Exception as e: + print(f"Error discovering vector search indexes: {e}", file=sys.stderr) + + return indexes + + +def discover_genie_spaces(w: WorkspaceClient) -> List[Dict[str, Any]]: + """Discover Genie spaces for conversational data access.""" + spaces = [] + + try: + # Use SDK to list genie spaces + response = w.genie.list_spaces() + genie_spaces = response.spaces if hasattr(response, "spaces") else [] + for space in genie_spaces: + spaces.append({ + "type": "genie_space", + "id": space.space_id, + "name": space.title, + "description": space.description, + }) + except Exception as e: + print(f"Error discovering Genie spaces: {e}", file=sys.stderr) + + return spaces + + + +def discover_custom_mcp_servers(w: WorkspaceClient) -> List[Dict[str, Any]]: + """Discover custom MCP servers deployed as Databricks apps.""" + custom_servers = [] + + try: + # List all apps and filter for those starting with mcp- + apps = w.apps.list() + for app in apps: + if app.name and app.name.startswith("mcp-"): + custom_servers.append({ + "type": "custom_mcp_server", + "name": app.name, + "url": app.url, + "status": app.app_status.state.value if app.app_status and app.app_status.state else None, + "description": app.description, + }) + except Exception as e: + print(f"Error discovering custom MCP servers: {e}", file=sys.stderr) + + return custom_servers + + +def discover_external_mcp_servers(w: WorkspaceClient) -> List[Dict[str, Any]]: + """Discover external MCP servers configured via Unity Catalog connections.""" + external_servers = [] + + try: + # List all connections and filter for MCP connections + connections = w.connections.list() + for conn in connections: + # Check if this is an MCP connection + if conn.options and conn.options.get("is_mcp_connection") == "true": + external_servers.append({ + "type": "external_mcp_server", + "name": conn.name, + "connection_type": conn.connection_type.value if hasattr(conn.connection_type, "value") else str(conn.connection_type), + "comment": conn.comment, + "full_name": conn.full_name, + }) + except Exception as e: + print(f"Error discovering external MCP servers: {e}", file=sys.stderr) + + return external_servers + + +def format_output_markdown(results: Dict[str, List[Dict[str, Any]]]) -> str: + """Format discovery results as markdown.""" + lines = ["# Agent Tools and Data Sources Discovery\n"] + + # UC Functions + functions = results.get("uc_functions", []) + if functions: + lines.append(f"## Unity Catalog Functions ({len(functions)})\n") + lines.append("**What they are:** SQL UDFs that can be used as agent tools.\n") + lines.append("**How to use:** Access via UC functions MCP server:") + lines.append("- All functions in a schema: `{workspace_host}/api/2.0/mcp/functions/{catalog}/{schema}`") + lines.append("- Single function: `{workspace_host}/api/2.0/mcp/functions/{catalog}/{schema}/{function_name}`\n") + for func in functions[:10]: # Show first 10 + lines.append(f"- `{func['name']}`") + if func.get("comment"): + lines.append(f" - {func['comment']}") + if len(functions) > 10: + lines.append(f"\n*...and {len(functions) - 10} more*\n") + lines.append("") + + # UC Tables + tables = results.get("uc_tables", []) + if tables: + lines.append(f"## Unity Catalog Tables ({len(tables)})\n") + lines.append("Structured data that agents can query via UC SQL functions.\n") + for table in tables[:10]: # Show first 10 + lines.append(f"- `{table['name']}` ({table['table_type']})") + if table.get("comment"): + lines.append(f" - {table['comment']}") + if table.get("columns"): + col_names = [c["name"] for c in table["columns"][:5]] + lines.append(f" - Columns: {', '.join(col_names)}") + if len(tables) > 10: + lines.append(f"\n*...and {len(tables) - 10} more*\n") + lines.append("") + + # Vector Search Indexes + indexes = results.get("vector_search_indexes", []) + if indexes: + lines.append(f"## Vector Search Indexes ({len(indexes)})\n") + lines.append("These can be used for RAG applications with unstructured data.\n") + lines.append("**How to use:** Connect via MCP server at `{workspace_host}/api/2.0/mcp/vector-search/{catalog}/{schema}` or\n") + lines.append("`{workspace_host}/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name}`\n") + for idx in indexes: + lines.append(f"- `{idx['name']}`") + lines.append(f" - Endpoint: {idx['endpoint']}") + lines.append(f" - Status: {idx['status']}") + lines.append("") + + # Genie Spaces + spaces = results.get("genie_spaces", []) + if spaces: + lines.append(f"## Genie Spaces ({len(spaces)})\n") + lines.append("**What they are:** Natural language interface to your data\n") + lines.append("**How to use:** Connect via Genie MCP server at `{workspace_host}/api/2.0/mcp/genie/{space_id}`\n") + for space in spaces: + lines.append(f"- `{space['name']}` (ID: {space['id']})") + if space.get("description"): + lines.append(f" - {space['description']}") + lines.append("") + + # Custom MCP Servers (Databricks Apps) + custom_servers = results.get("custom_mcp_servers", []) + if custom_servers: + lines.append(f"## Custom MCP Servers ({len(custom_servers)})\n") + lines.append("**What:** Your own MCP servers deployed as Databricks Apps (names starting with mcp-)\n") + lines.append("**How to use:** Access via `{app_url}/mcp`\n") + lines.append("**⚠️ Important:** Custom MCP server apps require manual permission grants:") + lines.append("1. Get your agent app's service principal: `databricks apps get --output json | jq -r '.service_principal_name'`") + lines.append("2. Grant permission: `databricks apps update-permissions --service-principal --permission-level CAN_USE`") + lines.append("(Apps are not yet supported as resource dependencies in databricks.yml)\n") + for server in custom_servers: + lines.append(f"- `{server['name']}`") + if server.get("url"): + lines.append(f" - URL: {server['url']}") + if server.get("status"): + lines.append(f" - Status: {server['status']}") + if server.get("description"): + lines.append(f" - {server['description']}") + lines.append("") + + # External MCP Servers (UC Connections) + external_servers = results.get("external_mcp_servers", []) + if external_servers: + lines.append(f"## External MCP Servers ({len(external_servers)})\n") + lines.append("**What:** Third-party MCP servers via Unity Catalog connections\n") + lines.append("**How to use:** Connect via `{workspace_host}/api/2.0/mcp/external/{connection_name}`\n") + lines.append("**Benefits:** Secure access to external APIs through UC governance\n") + for server in external_servers: + lines.append(f"- `{server['name']}`") + if server.get("full_name"): + lines.append(f" - Full name: {server['full_name']}") + if server.get("comment"): + lines.append(f" - {server['comment']}") + lines.append("") + return "\n".join(lines) + + +def main(): + """Main discovery function.""" + import argparse + + parser = argparse.ArgumentParser(description="Discover available agent tools and data sources") + parser.add_argument("--catalog", help="Limit discovery to specific catalog") + parser.add_argument("--schema", help="Limit discovery to specific schema (requires --catalog)") + parser.add_argument("--format", choices=["json", "markdown"], default="markdown", help="Output format") + parser.add_argument("--output", help="Output file (default: stdout)") + parser.add_argument("--profile", help="Databricks CLI profile to use (default: uses default profile)") + parser.add_argument("--max-results", type=int, default=DEFAULT_MAX_RESULTS, help=f"Maximum results per resource type (default: {DEFAULT_MAX_RESULTS})") + parser.add_argument("--max-schemas", type=int, default=DEFAULT_MAX_SCHEMAS, help=f"Total schemas to search across all catalogs (default: {DEFAULT_MAX_SCHEMAS})") + + args = parser.parse_args() + + if args.schema and not args.catalog: + print("Error: --schema requires --catalog", file=sys.stderr) + sys.exit(1) + + print("Discovering available tools and data sources...", file=sys.stderr) + + # Initialize Databricks workspace client + # Only pass profile if specified, otherwise use default + if args.profile: + w = WorkspaceClient(profile=args.profile) + else: + w = WorkspaceClient() + + results = {} + + # Discover each type with configurable limits + print("- UC Functions...", file=sys.stderr) + results["uc_functions"] = discover_uc_functions(w, catalog=args.catalog, max_schemas=args.max_schemas)[:args.max_results] + + print("- UC Tables...", file=sys.stderr) + results["uc_tables"] = discover_uc_tables(w, catalog=args.catalog, schema=args.schema, max_schemas=args.max_schemas)[:args.max_results] + + print("- Vector Search Indexes...", file=sys.stderr) + results["vector_search_indexes"] = discover_vector_search_indexes(w)[:args.max_results] + + print("- Genie Spaces...", file=sys.stderr) + results["genie_spaces"] = discover_genie_spaces(w)[:args.max_results] + + print("- Custom MCP Servers (Apps)...", file=sys.stderr) + results["custom_mcp_servers"] = discover_custom_mcp_servers(w)[:args.max_results] + + print("- External MCP Servers (Connections)...", file=sys.stderr) + results["external_mcp_servers"] = discover_external_mcp_servers(w)[:args.max_results] + + # Format output + if args.format == "json": + output = json.dumps(results, indent=2) + else: + output = format_output_markdown(results) + + # Write output + if args.output: + Path(args.output).write_text(output) + print(f"\nResults written to {args.output}", file=sys.stderr) + else: + print("\n" + output) + + # Print summary + print("\n=== Discovery Summary ===", file=sys.stderr) + print(f"UC Functions: {len(results['uc_functions'])}", file=sys.stderr) + print(f"UC Tables: {len(results['uc_tables'])}", file=sys.stderr) + print(f"Vector Search Indexes: {len(results['vector_search_indexes'])}", file=sys.stderr) + print(f"Genie Spaces: {len(results['genie_spaces'])}", file=sys.stderr) + print(f"Custom MCP Servers: {len(results['custom_mcp_servers'])}", file=sys.stderr) + print(f"External MCP Servers: {len(results['external_mcp_servers'])}", file=sys.stderr) + + +if __name__ == "__main__": + main() From ef4c96a3b97448b2b5b1cce2949cef16ac19895f Mon Sep 17 00:00:00 2001 From: Dhruv Gupta Date: Wed, 28 Jan 2026 14:44:10 -0800 Subject: [PATCH 04/14] Add memory skills to base agent-langgraph template Added lakebase-setup and agent-memory skills for users who want to add memory capabilities to the base template. Skills include references to pre-configured memory templates for easier discovery. Co-Authored-By: Claude (databricks-claude-opus-4-5) --- .../.claude/skills/agent-memory/SKILL.md | 204 ++++++++++++++++++ .../.claude/skills/lakebase-setup/SKILL.md | 110 ++++++++++ .../.claude/skills/modify-agent/SKILL.md | 1 + agent-langgraph/.gitignore | 2 + agent-langgraph/AGENTS.md | 4 + 5 files changed, 321 insertions(+) create mode 100644 agent-langgraph/.claude/skills/agent-memory/SKILL.md create mode 100644 agent-langgraph/.claude/skills/lakebase-setup/SKILL.md diff --git a/agent-langgraph/.claude/skills/agent-memory/SKILL.md b/agent-langgraph/.claude/skills/agent-memory/SKILL.md new file mode 100644 index 00000000..9d4e535f --- /dev/null +++ b/agent-langgraph/.claude/skills/agent-memory/SKILL.md @@ -0,0 +1,204 @@ +--- +name: agent-memory +description: "Add memory capabilities to your agent. Use when: (1) User asks about 'memory', 'state', 'remember', 'conversation history', (2) Want to persist conversations or user preferences, (3) Adding checkpointing or long-term storage." +--- + +# Adding Memory to Your Agent + +> **Note:** This template does not include memory by default. Use this skill to **add memory capabilities**. For pre-configured memory templates, see: +> - `agent-langgraph-short-term-memory` - Conversation history within a session +> - `agent-langgraph-long-term-memory` - User facts that persist across sessions + +## Memory Types + +| Type | Use Case | Storage | Identifier | +|------|----------|---------|------------| +| **Short-term** | Conversation history within a session | `AsyncCheckpointSaver` | `thread_id` | +| **Long-term** | User facts that persist across sessions | `AsyncDatabricksStore` | `user_id` | + +## Prerequisites + +1. Add memory dependency to `pyproject.toml`: + ```toml + dependencies = [ + "databricks-langchain[memory]>=0.13.0", + ] + ``` + +2. Configure Lakebase (see **lakebase-setup** skill) + +--- + +## Adding Short-Term Memory + +Short-term memory stores conversation history within a thread. The agent remembers what was said earlier in the conversation. + +### Step 1: Import and Initialize Checkpointer + +```python +from databricks_langchain import AsyncCheckpointSaver + +LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") + +async with AsyncCheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME) as checkpointer: + agent = create_react_agent( + model=model, + tools=tools, + checkpointer=checkpointer, # Enables persistence + ) +``` + +### Step 2: Use thread_id in Requests + +```python +def _get_or_create_thread_id(request: ResponsesAgentRequest) -> str: + custom_inputs = dict(request.custom_inputs or {}) + if "thread_id" in custom_inputs and custom_inputs["thread_id"]: + return str(custom_inputs["thread_id"]) + if request.context and getattr(request.context, "conversation_id", None): + return str(request.context.conversation_id) + return str(uuid.uuid4()) + +# Use in agent invocation +config = {"configurable": {"thread_id": thread_id}} +async for event in agent.astream(input_state, config, stream_mode=["updates", "messages"]): + # Process events +``` + +### Test Short-Term Memory + +```bash +# First message (starts new conversation) +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{"input": [{"role": "user", "content": "My name is Alice"}]}' + +# Continue conversation (use thread_id from response) +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "What is my name?"}], + "custom_inputs": {"thread_id": ""} + }' +``` + +--- + +## Adding Long-Term Memory + +Long-term memory stores facts about users that persist across conversation sessions. + +### Step 1: Import and Initialize Store + +```python +from databricks_langchain import AsyncDatabricksStore + +LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") +EMBEDDING_ENDPOINT = os.getenv("EMBEDDING_ENDPOINT", "databricks-gte-large-en") + +async with AsyncDatabricksStore( + instance_name=LAKEBASE_INSTANCE_NAME, + embedding_endpoint=EMBEDDING_ENDPOINT, +) as store: + agent = create_react_agent( + model=model, + tools=tools + memory_tools, # Add memory tools + store=store, + ) +``` + +### Step 2: Create Memory Tools + +```python +from langchain_core.tools import tool +from langchain_core.runnables import RunnableConfig + +@tool +async def get_user_memory(query: str, config: RunnableConfig) -> str: + """Search for relevant information about the user from long-term memory.""" + user_id = config.get("configurable", {}).get("user_id") + store = config.get("configurable", {}).get("store") + if not user_id or not store: + return "Memory not available" + + namespace = ("user_memories", user_id.replace(".", "-")) + results = await store.asearch(namespace, query=query, limit=5) + if not results: + return "No memories found" + return "\n".join([f"- {r.value}" for r in results]) + +@tool +async def save_user_memory(memory_key: str, memory_data: str, config: RunnableConfig) -> str: + """Save information about the user to long-term memory.""" + user_id = config.get("configurable", {}).get("user_id") + store = config.get("configurable", {}).get("store") + if not user_id or not store: + return "Memory not available" + + namespace = ("user_memories", user_id.replace(".", "-")) + await store.aput(namespace, memory_key, {"value": memory_data}) + return f"Saved memory: {memory_key}" +``` + +### Step 3: Use user_id in Requests + +```python +def _get_user_id(request: ResponsesAgentRequest) -> str: + custom_inputs = dict(request.custom_inputs or {}) + if "user_id" in custom_inputs: + return custom_inputs["user_id"] + if request.context and getattr(request.context, "user_id", None): + return request.context.user_id + return "default-user" + +# Use in agent invocation +config = {"configurable": {"user_id": user_id, "store": store}} +``` + +### Test Long-Term Memory + +```bash +# Save a preference +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Remember that my favorite color is blue"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' + +# Recall later (even in new conversation) +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "What is my favorite color?"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' +``` + +--- + +## Using Pre-Built Memory Templates + +For fully configured memory implementations, use these templates instead: + +| Template | Memory Type | Key Features | +|----------|-------------|--------------| +| `agent-langgraph-short-term-memory` | Short-term | AsyncCheckpointSaver, thread_id, conversation persistence | +| `agent-langgraph-long-term-memory` | Long-term | AsyncDatabricksStore, user_id, memory tools | + +These templates have memory fully integrated and tested. + +## Troubleshooting + +| Issue | Solution | +|-------|----------| +| **"Cannot connect to Lakebase"** | See **lakebase-setup** skill | +| **Memory not persisting** | Verify thread_id/user_id is passed consistently | +| **Permission errors** | Grant Lakebase permissions (see **lakebase-setup** skill) | +| **"Memory not available"** | Ensure user_id is provided in request | + +## Next Steps + +- Configure Lakebase: see **lakebase-setup** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md b/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md new file mode 100644 index 00000000..4e0f4646 --- /dev/null +++ b/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md @@ -0,0 +1,110 @@ +--- +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 Memory + +> **Note:** This template does not include memory by default. Use this skill if you want to **add memory capabilities** to your agent. For pre-configured memory templates, see: +> - `agent-langgraph-short-term-memory` - Conversation history within a session +> - `agent-langgraph-long-term-memory` - User facts that persist across sessions + +## Overview + +Lakebase provides persistent storage for agent memory: +- **Short-term memory**: Conversation history within a thread (`AsyncCheckpointSaver`) +- **Long-term memory**: User facts across sessions (`AsyncDatabricksStore`) + +## Step 1: Add Memory Dependency + +Add the memory extra to your `pyproject.toml`: + +```toml +dependencies = [ + "databricks-langchain[memory]>=0.13.0", + # ... other dependencies +] +``` + +Then sync dependencies: +```bash +uv sync +``` + +## Step 2: Configure Lakebase Instance + +Add to your `.env.local`: +```bash +LAKEBASE_INSTANCE_NAME="" +``` + +The quickstart script can help you find or create a Lakebase instance. + +## Step 3: Grant Permissions (After Deployment) + +After deploying your app, grant the app's service principal access to Lakebase. + +### Option A: Using LakebaseClient SDK (Recommended) + +```python +from databricks_ai_bridge.lakebase import LakebaseClient, SchemaPrivilege, TablePrivilege + +# Initialize client +client = LakebaseClient(instance_name="") + +# Get app service principal from: databricks apps get +app_sp = "" + +# Create role for the service principal +client.create_role(app_sp, "SERVICE_PRINCIPAL") + +# Grant schema privileges +client.grant_schema( + role=app_sp, + schema="public", + privileges=[SchemaPrivilege.USAGE, SchemaPrivilege.CREATE], +) + +# Grant table privileges (for existing tables) +client.grant_all_tables_in_schema( + role=app_sp, + schema="public", + privileges=[TablePrivilege.SELECT, TablePrivilege.INSERT, TablePrivilege.UPDATE, TablePrivilege.DELETE], +) +``` + +### Option B: Using SQL + +Connect to your Lakebase instance via SQL and run: + +```sql +-- Replace with your app's service principal +-- Get it from: databricks apps get --output json | jq -r '.service_principal_id' + +-- Grant schema access +GRANT USAGE, CREATE ON SCHEMA public TO ""; + +-- Grant table access for memory tables +GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO ""; +``` + +## Adding Memory to Your Agent + +See the **agent-memory** skill for code patterns to add: +- Short-term memory (conversation history) +- Long-term memory (persistent user facts) + +## Troubleshooting + +| Issue | Solution | +|-------|----------| +| **"Failed to connect to Lakebase"** | Verify `LAKEBASE_INSTANCE_NAME` in `.env.local` | +| **Permission denied on tables** | Run the SQL grants above after deployment | +| **"Role does not exist"** | Use the service principal ID, not name | +| **Tables not created** | Tables are created automatically on first use | + +## Next Steps + +- Add memory patterns: see **agent-memory** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/agent-langgraph/.claude/skills/modify-agent/SKILL.md b/agent-langgraph/.claude/skills/modify-agent/SKILL.md index eec80096..d7218637 100644 --- a/agent-langgraph/.claude/skills/modify-agent/SKILL.md +++ b/agent-langgraph/.claude/skills/modify-agent/SKILL.md @@ -288,5 +288,6 @@ Example: Create UC function wrapping HTTP request for Slack, then expose via MCP - Discover available tools: see **discover-tools** skill - Grant resource permissions: see **add-tools** skill +- Add memory capabilities: see **agent-memory** skill - Test locally: see **run-locally** skill - Deploy: see **deploy** skill diff --git a/agent-langgraph/.gitignore b/agent-langgraph/.gitignore index d4a984d8..601fbb1a 100644 --- a/agent-langgraph/.gitignore +++ b/agent-langgraph/.gitignore @@ -214,6 +214,8 @@ sketch !.claude/skills/add-tools/ !.claude/skills/run-locally/ !.claude/skills/modify-agent/ +!.claude/skills/lakebase-setup/ +!.claude/skills/agent-memory/ **/.env **/.env.local \ No newline at end of file diff --git a/agent-langgraph/AGENTS.md b/agent-langgraph/AGENTS.md index a4168cd7..896b1c9d 100644 --- a/agent-langgraph/AGENTS.md +++ b/agent-langgraph/AGENTS.md @@ -48,9 +48,13 @@ Ask the user: "I see there's an existing app with the same name. Would you like | Add tools & permissions | **add-tools** | `.claude/skills/add-tools/SKILL.md` | | Run/test locally | **run-locally** | `.claude/skills/run-locally/SKILL.md` | | Modify agent code | **modify-agent** | `.claude/skills/modify-agent/SKILL.md` | +| Configure Lakebase storage | **lakebase-setup** | `.claude/skills/lakebase-setup/SKILL.md` | +| Add memory capabilities | **agent-memory** | `.claude/skills/agent-memory/SKILL.md` | **Note:** All agent skills are located in `.claude/skills/` directory. +> **Adding Memory?** The **lakebase-setup** and **agent-memory** skills help you add conversation history or persistent user memory to this agent. For pre-configured memory, see the `agent-langgraph-short-term-memory` or `agent-langgraph-long-term-memory` templates. + --- ## Quick Commands From 266ea1b61473d62355f45def7de0315d28db41da Mon Sep 17 00:00:00 2001 From: Dhruv Gupta Date: Wed, 28 Jan 2026 15:44:46 -0800 Subject: [PATCH 05/14] Update skills to use .env instead of .env.local Aligns with PR #99 changes that standardized on .env file naming. Co-Authored-By: Claude (databricks-claude-opus-4-5) --- .../.claude/skills/lakebase-setup/SKILL.md | 4 ++-- .../.claude/skills/quickstart/SKILL.md | 6 +++--- .../.claude/skills/run-locally/SKILL.md | 10 +++++----- .../.claude/skills/lakebase-setup/SKILL.md | 4 ++-- .../.claude/skills/quickstart/SKILL.md | 6 +++--- .../.claude/skills/run-locally/SKILL.md | 10 +++++----- agent-langgraph/.claude/skills/lakebase-setup/SKILL.md | 4 ++-- agent-langgraph/.claude/skills/quickstart/SKILL.md | 6 +++--- agent-langgraph/.claude/skills/run-locally/SKILL.md | 8 ++++---- 9 files changed, 29 insertions(+), 29 deletions(-) diff --git a/agent-langgraph-long-term-memory/.claude/skills/lakebase-setup/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/lakebase-setup/SKILL.md index 30d7825f..036b1a1b 100644 --- a/agent-langgraph-long-term-memory/.claude/skills/lakebase-setup/SKILL.md +++ b/agent-langgraph-long-term-memory/.claude/skills/lakebase-setup/SKILL.md @@ -14,7 +14,7 @@ This template uses Lakebase (Databricks-managed PostgreSQL) to store long-term u ## Local Development Setup -**Step 1:** Add the Lakebase instance name to `.env.local`: +**Step 1:** Add the Lakebase instance name to `.env`: ```bash LAKEBASE_INSTANCE_NAME= @@ -131,7 +131,7 @@ END $$; | Issue | Solution | |-------|----------| -| **"LAKEBASE_INSTANCE_NAME environment variable is required"** | Add `LAKEBASE_INSTANCE_NAME=` to `.env.local` | +| **"LAKEBASE_INSTANCE_NAME environment variable is required"** | Add `LAKEBASE_INSTANCE_NAME=` to `.env` | | **"Failed to connect to Lakebase instance"** | Verify instance name is correct and your profile has access | | **Permission errors on store tables** | Run the SDK script or SQL grant commands above | | **Deployed app can't access Lakebase** | Add Lakebase as app resource in Databricks UI | diff --git a/agent-langgraph-long-term-memory/.claude/skills/quickstart/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/quickstart/SKILL.md index edfbfd33..853b78d0 100644 --- a/agent-langgraph-long-term-memory/.claude/skills/quickstart/SKILL.md +++ b/agent-langgraph-long-term-memory/.claude/skills/quickstart/SKILL.md @@ -1,6 +1,6 @@ --- name: quickstart -description: "Set up Databricks agent development environment. Use when: (1) First time setup, (2) Configuring Databricks authentication, (3) User says 'quickstart', 'set up', 'authenticate', or 'configure databricks', (4) No .env.local file exists." +description: "Set up Databricks agent development environment. Use when: (1) First time setup, (2) Configuring Databricks authentication, (3) User says 'quickstart', 'set up', 'authenticate', or 'configure databricks', (4) No .env file exists." --- # Quickstart & Authentication @@ -42,7 +42,7 @@ uv run quickstart --host https://your-workspace.cloud.databricks.com ## What Quickstart Configures -Creates/updates `.env.local` with: +Creates/updates `.env` with: - `DATABRICKS_CONFIG_PROFILE` - Selected CLI profile - `MLFLOW_TRACKING_URI` - Set to `databricks://` for local auth - `MLFLOW_EXPERIMENT_ID` - Auto-created experiment ID @@ -59,7 +59,7 @@ databricks auth login --host https://your-workspace.cloud.databricks.com databricks auth profiles ``` -Then manually create `.env.local` (copy from `.env.example`): +Then manually create `.env` (copy from `.env.example`): ```bash # Authentication (choose one method) DATABRICKS_CONFIG_PROFILE=DEFAULT diff --git a/agent-langgraph-long-term-memory/.claude/skills/run-locally/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/run-locally/SKILL.md index 7d959670..398d22a3 100644 --- a/agent-langgraph-long-term-memory/.claude/skills/run-locally/SKILL.md +++ b/agent-langgraph-long-term-memory/.claude/skills/run-locally/SKILL.md @@ -93,10 +93,10 @@ Dependencies are managed in `pyproject.toml`. | Issue | Solution | |-------|----------| | **Port already in use** | Use `--port 8001` or kill existing process | -| **Authentication errors** | Verify `.env.local` is correct; run **quickstart** skill | +| **Authentication errors** | Verify `.env` is correct; run **quickstart** skill | | **Module not found** | Run `uv sync` to install dependencies | -| **MLflow experiment not found** | Ensure `MLFLOW_TRACKING_URI` in `.env.local` is `databricks://` | -| **Lakebase connection errors** | Verify `LAKEBASE_INSTANCE_NAME` in `.env.local`; see **lakebase-setup** skill | +| **MLflow experiment not found** | Ensure `MLFLOW_TRACKING_URI` in `.env` is `databricks://` | +| **Lakebase connection errors** | Verify `LAKEBASE_INSTANCE_NAME` in `.env`; see **lakebase-setup** skill | ### MLflow Experiment Not Found @@ -107,12 +107,12 @@ If you see: "The provided MLFLOW_EXPERIMENT_ID environment variable value does n databricks -p experiments get-experiment ``` -**Fix:** Ensure `.env.local` has the correct tracking URI format: +**Fix:** Ensure `.env` has the correct tracking URI format: ```bash MLFLOW_TRACKING_URI="databricks://DEFAULT" # Include profile name ``` -The quickstart script configures this automatically. If you manually edited `.env.local`, ensure the profile name is included. +The quickstart script configures this automatically. If you manually edited `.env`, ensure the profile name is included. ## MLflow Tracing diff --git a/agent-langgraph-short-term-memory/.claude/skills/lakebase-setup/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/lakebase-setup/SKILL.md index e37edbd5..96b2375a 100644 --- a/agent-langgraph-short-term-memory/.claude/skills/lakebase-setup/SKILL.md +++ b/agent-langgraph-short-term-memory/.claude/skills/lakebase-setup/SKILL.md @@ -14,7 +14,7 @@ This template uses Lakebase (Databricks-managed PostgreSQL) to store conversatio ## Local Development Setup -**Step 1:** Add the Lakebase instance name to `.env.local`: +**Step 1:** Add the Lakebase instance name to `.env`: ```bash LAKEBASE_INSTANCE_NAME= @@ -131,7 +131,7 @@ END $$; | Issue | Solution | |-------|----------| -| **"LAKEBASE_INSTANCE_NAME environment variable is required"** | Add `LAKEBASE_INSTANCE_NAME=` to `.env.local` | +| **"LAKEBASE_INSTANCE_NAME environment variable is required"** | Add `LAKEBASE_INSTANCE_NAME=` to `.env` | | **"Failed to connect to Lakebase instance"** | Verify instance name is correct and your profile has access | | **Permission errors on checkpoint tables** | Run the SDK script or SQL grant commands above | | **Deployed app can't access Lakebase** | Add Lakebase as app resource in Databricks UI | diff --git a/agent-langgraph-short-term-memory/.claude/skills/quickstart/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/quickstart/SKILL.md index edfbfd33..853b78d0 100644 --- a/agent-langgraph-short-term-memory/.claude/skills/quickstart/SKILL.md +++ b/agent-langgraph-short-term-memory/.claude/skills/quickstart/SKILL.md @@ -1,6 +1,6 @@ --- name: quickstart -description: "Set up Databricks agent development environment. Use when: (1) First time setup, (2) Configuring Databricks authentication, (3) User says 'quickstart', 'set up', 'authenticate', or 'configure databricks', (4) No .env.local file exists." +description: "Set up Databricks agent development environment. Use when: (1) First time setup, (2) Configuring Databricks authentication, (3) User says 'quickstart', 'set up', 'authenticate', or 'configure databricks', (4) No .env file exists." --- # Quickstart & Authentication @@ -42,7 +42,7 @@ uv run quickstart --host https://your-workspace.cloud.databricks.com ## What Quickstart Configures -Creates/updates `.env.local` with: +Creates/updates `.env` with: - `DATABRICKS_CONFIG_PROFILE` - Selected CLI profile - `MLFLOW_TRACKING_URI` - Set to `databricks://` for local auth - `MLFLOW_EXPERIMENT_ID` - Auto-created experiment ID @@ -59,7 +59,7 @@ databricks auth login --host https://your-workspace.cloud.databricks.com databricks auth profiles ``` -Then manually create `.env.local` (copy from `.env.example`): +Then manually create `.env` (copy from `.env.example`): ```bash # Authentication (choose one method) DATABRICKS_CONFIG_PROFILE=DEFAULT diff --git a/agent-langgraph-short-term-memory/.claude/skills/run-locally/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/run-locally/SKILL.md index f3b517de..f33918eb 100644 --- a/agent-langgraph-short-term-memory/.claude/skills/run-locally/SKILL.md +++ b/agent-langgraph-short-term-memory/.claude/skills/run-locally/SKILL.md @@ -93,10 +93,10 @@ Dependencies are managed in `pyproject.toml`. | Issue | Solution | |-------|----------| | **Port already in use** | Use `--port 8001` or kill existing process | -| **Authentication errors** | Verify `.env.local` is correct; run **quickstart** skill | +| **Authentication errors** | Verify `.env` is correct; run **quickstart** skill | | **Module not found** | Run `uv sync` to install dependencies | -| **MLflow experiment not found** | Ensure `MLFLOW_TRACKING_URI` in `.env.local` is `databricks://` | -| **Lakebase connection errors** | Verify `LAKEBASE_INSTANCE_NAME` in `.env.local`; see **lakebase-setup** skill | +| **MLflow experiment not found** | Ensure `MLFLOW_TRACKING_URI` in `.env` is `databricks://` | +| **Lakebase connection errors** | Verify `LAKEBASE_INSTANCE_NAME` in `.env`; see **lakebase-setup** skill | ### MLflow Experiment Not Found @@ -107,12 +107,12 @@ If you see: "The provided MLFLOW_EXPERIMENT_ID environment variable value does n databricks -p experiments get-experiment ``` -**Fix:** Ensure `.env.local` has the correct tracking URI format: +**Fix:** Ensure `.env` has the correct tracking URI format: ```bash MLFLOW_TRACKING_URI="databricks://DEFAULT" # Include profile name ``` -The quickstart script configures this automatically. If you manually edited `.env.local`, ensure the profile name is included. +The quickstart script configures this automatically. If you manually edited `.env`, ensure the profile name is included. ## MLflow Tracing diff --git a/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md b/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md index 4e0f4646..45ad68b4 100644 --- a/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md +++ b/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md @@ -33,7 +33,7 @@ uv sync ## Step 2: Configure Lakebase Instance -Add to your `.env.local`: +Add to your `.env`: ```bash LAKEBASE_INSTANCE_NAME="" ``` @@ -98,7 +98,7 @@ See the **agent-memory** skill for code patterns to add: | Issue | Solution | |-------|----------| -| **"Failed to connect to Lakebase"** | Verify `LAKEBASE_INSTANCE_NAME` in `.env.local` | +| **"Failed to connect to Lakebase"** | Verify `LAKEBASE_INSTANCE_NAME` in `.env` | | **Permission denied on tables** | Run the SQL grants above after deployment | | **"Role does not exist"** | Use the service principal ID, not name | | **Tables not created** | Tables are created automatically on first use | diff --git a/agent-langgraph/.claude/skills/quickstart/SKILL.md b/agent-langgraph/.claude/skills/quickstart/SKILL.md index 0ca105e6..e550162c 100644 --- a/agent-langgraph/.claude/skills/quickstart/SKILL.md +++ b/agent-langgraph/.claude/skills/quickstart/SKILL.md @@ -1,6 +1,6 @@ --- name: quickstart -description: "Set up Databricks agent development environment. Use when: (1) First time setup, (2) Configuring Databricks authentication, (3) User says 'quickstart', 'set up', 'authenticate', or 'configure databricks', (4) No .env.local file exists." +description: "Set up Databricks agent development environment. Use when: (1) First time setup, (2) Configuring Databricks authentication, (3) User says 'quickstart', 'set up', 'authenticate', or 'configure databricks', (4) No .env file exists." --- # Quickstart & Authentication @@ -42,7 +42,7 @@ uv run quickstart --host https://your-workspace.cloud.databricks.com ## What Quickstart Configures -Creates/updates `.env.local` with: +Creates/updates `.env` with: - `DATABRICKS_CONFIG_PROFILE` - Selected CLI profile - `MLFLOW_TRACKING_URI` - Set to `databricks://` for local auth - `MLFLOW_EXPERIMENT_ID` - Auto-created experiment ID @@ -59,7 +59,7 @@ databricks auth login --host https://your-workspace.cloud.databricks.com databricks auth profiles ``` -Then manually create `.env.local` (copy from `.env.example`): +Then manually create `.env` (copy from `.env.example`): ```bash # Authentication (choose one method) DATABRICKS_CONFIG_PROFILE=DEFAULT diff --git a/agent-langgraph/.claude/skills/run-locally/SKILL.md b/agent-langgraph/.claude/skills/run-locally/SKILL.md index 294be8c8..3eb83c82 100644 --- a/agent-langgraph/.claude/skills/run-locally/SKILL.md +++ b/agent-langgraph/.claude/skills/run-locally/SKILL.md @@ -64,9 +64,9 @@ pytest [path] | Issue | Solution | |-------|----------| | **Port already in use** | Use `--port 8001` or kill existing process | -| **Authentication errors** | Verify `.env.local` is correct; run **quickstart** skill | +| **Authentication errors** | Verify `.env` is correct; run **quickstart** skill | | **Module not found** | Run `uv sync` to install dependencies | -| **MLflow experiment not found** | Ensure `MLFLOW_TRACKING_URI` in `.env.local` is `databricks://` | +| **MLflow experiment not found** | Ensure `MLFLOW_TRACKING_URI` in `.env` is `databricks://` | ### MLflow Experiment Not Found @@ -77,12 +77,12 @@ If you see: "The provided MLFLOW_EXPERIMENT_ID environment variable value does n databricks -p experiments get-experiment ``` -**Fix:** Ensure `.env.local` has the correct tracking URI format: +**Fix:** Ensure `.env` has the correct tracking URI format: ```bash MLFLOW_TRACKING_URI="databricks://DEFAULT" # Include profile name ``` -The quickstart script configures this automatically. If you manually edited `.env.local`, ensure the profile name is included. +The quickstart script configures this automatically. If you manually edited `.env`, ensure the profile name is included. ## Next Steps From 43eb7265a695043a63c6f5251cabadf80895194e Mon Sep 17 00:00:00 2001 From: Dhruv Gupta Date: Thu, 29 Jan 2026 10:26:04 -0800 Subject: [PATCH 06/14] Skill updates for lakebase memory Co-Authored-By: Claude (databricks-claude-opus-4-5) --- .../.claude/skills/add-tools/SKILL.md | 28 +- .../skills/add-tools/examples/lakebase.yaml | 18 + .../.claude/skills/agent-memory/SKILL.md | 476 ++++++++++++++---- .../.claude/skills/deploy/SKILL.md | 10 +- .../.claude/skills/lakebase-setup/SKILL.md | 312 ++++++++++-- agent-langgraph/AGENTS.md | 7 +- agent-langgraph/pyproject.toml | 2 +- 7 files changed, 709 insertions(+), 144 deletions(-) create mode 100644 agent-langgraph/.claude/skills/add-tools/examples/lakebase.yaml diff --git a/agent-langgraph/.claude/skills/add-tools/SKILL.md b/agent-langgraph/.claude/skills/add-tools/SKILL.md index 7d00e104..7719f198 100644 --- a/agent-langgraph/.claude/skills/add-tools/SKILL.md +++ b/agent-langgraph/.claude/skills/add-tools/SKILL.md @@ -37,7 +37,13 @@ resources: permission: 'CAN_RUN' ``` -**Step 3:** Deploy with `databricks bundle deploy` (see **deploy** skill) +**Step 3:** Deploy and run: +```bash +databricks bundle deploy +databricks bundle run agent_langgraph # Required to start app with new code! +``` + +See **deploy** skill for more details. ## Resource Type Examples @@ -51,6 +57,7 @@ See the `examples/` directory for complete YAML snippets: | `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 | Agent memory storage | | `experiment.yaml` | MLflow experiment | Tracing (already configured) | | `custom-mcp-server.md` | Custom MCP apps | Apps starting with `mcp-*` | @@ -72,9 +79,26 @@ databricks apps update-permissions \ See `examples/custom-mcp-server.md` for detailed steps. +## valueFrom Pattern (for app.yaml) + +**IMPORTANT**: Make sure all `valueFrom` references in `app.yaml` reference an existing key in the `databricks.yml` file. +Some resources need environment variables in your app. Use `valueFrom` in `app.yaml` to reference resources defined in `databricks.yml`: + +```yaml +# app.yaml +env: + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" # References resources.apps..resources[name='experiment'] + - name: LAKEBASE_INSTANCE_NAME + valueFrom: "database" # References resources.apps..resources[name='database'] +``` + +**Critical:** Every `valueFrom` value must match a `name` field in `databricks.yml` resources. + ## 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 after changes**: Run `databricks bundle deploy` after modifying `databricks.yml` +- **Deploy + Run after changes**: Run both `databricks bundle deploy` AND `databricks bundle run agent_langgraph` +- **valueFrom matching**: Ensure `app.yaml` `valueFrom` values match `databricks.yml` resource `name` values diff --git a/agent-langgraph/.claude/skills/add-tools/examples/lakebase.yaml b/agent-langgraph/.claude/skills/add-tools/examples/lakebase.yaml new file mode 100644 index 00000000..78f0bc72 --- /dev/null +++ b/agent-langgraph/.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: valueFrom reference in app.yaml + +# In databricks.yml - add to resources.apps..resources: +- name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +# In app.yaml - add to env: +# - name: LAKEBASE_INSTANCE_NAME +# valueFrom: "database" +# - name: EMBEDDING_ENDPOINT +# value: "databricks-gte-large-en" +# - name: EMBEDDING_DIMS +# value: "1024" diff --git a/agent-langgraph/.claude/skills/agent-memory/SKILL.md b/agent-langgraph/.claude/skills/agent-memory/SKILL.md index 9d4e535f..f06c5d23 100644 --- a/agent-langgraph/.claude/skills/agent-memory/SKILL.md +++ b/agent-langgraph/.claude/skills/agent-memory/SKILL.md @@ -18,184 +18,478 @@ description: "Add memory capabilities to your agent. Use when: (1) User asks abo ## Prerequisites -1. Add memory dependency to `pyproject.toml`: +1. **Add memory dependency** to `pyproject.toml`: ```toml dependencies = [ - "databricks-langchain[memory]>=0.13.0", + "databricks-langchain[memory]", ] ``` -2. Configure Lakebase (see **lakebase-setup** skill) + Then run `uv sync` + +2. **Configure Lakebase** - See **lakebase-setup** skill for: + - Creating/configuring Lakebase instance + - Initializing tables (CRITICAL first-time step) --- -## Adding Short-Term Memory +## Quick Setup Summary -Short-term memory stores conversation history within a thread. The agent remembers what was said earlier in the conversation. +Adding memory requires changes to **4 files**: -### Step 1: Import and Initialize Checkpointer +| File | What to Add | +|------|-------------| +| `pyproject.toml` | Memory dependency + hatch config | +| `.env.local` | Lakebase env vars (for local dev) | +| `databricks.yml` | Lakebase database resource | +| `app.yaml` | `valueFrom` reference to lakebase resource | +| `agent_server/agent.py` | Memory tools and AsyncDatabricksStore | -```python -from databricks_langchain import AsyncCheckpointSaver +--- -LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") +## Step 1: Configure databricks.yml (Lakebase Resource) -async with AsyncCheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME) as checkpointer: - agent = create_react_agent( - model=model, - tools=tools, - checkpointer=checkpointer, # Enables persistence - ) -``` +Add the Lakebase database resource to your app in `databricks.yml`: -### Step 2: Use thread_id in Requests +```yaml +resources: + apps: + agent_langgraph: + name: "your-app-name" + source_code_path: ./ -```python -def _get_or_create_thread_id(request: ResponsesAgentRequest) -> str: - custom_inputs = dict(request.custom_inputs or {}) - if "thread_id" in custom_inputs and custom_inputs["thread_id"]: - return str(custom_inputs["thread_id"]) - if request.context and getattr(request.context, "conversation_id", None): - return str(request.context.conversation_id) - return str(uuid.uuid4()) + resources: + # ... other resources (experiment, UC functions, etc.) ... -# Use in agent invocation -config = {"configurable": {"thread_id": thread_id}} -async for event in agent.astream(input_state, config, stream_mode=["updates", "messages"]): - # Process events + # Lakebase instance for long-term memory + - name: 'lakebadatabasese_memory' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' ``` -### Test Short-Term Memory +**Important:** The `name: 'database'` must match the `valueFrom` reference in `app.yaml`. -```bash -# First message (starts new conversation) -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{"input": [{"role": "user", "content": "My name is Alice"}]}' +--- -# Continue conversation (use thread_id from response) -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ - "input": [{"role": "user", "content": "What is my name?"}], - "custom_inputs": {"thread_id": ""} - }' +## Step 2: Configure app.yaml (Environment Variables) + +Update `app.yaml` with the Lakebase environment variables: + +```yaml +command: ["uv", "run", "start-app"] + +env: + # ... other env vars ... + + # MLflow experiment (uses valueFrom to reference databricks.yml resource) + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" + + # Lakebase instance name - must match the instance_name in databricks.yml database resource + # Note: Use 'value' (not 'valueFrom') because AsyncDatabricksStore needs the instance name, + # not the full connection string that valueFrom would provide + - name: LAKEBASE_INSTANCE_NAME + value: "" + + # Embedding configuration (static values) + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" ``` +**Important:** The `LAKEBASE_INSTANCE_NAME` value must match the `instance_name` in your `databricks.yml` database resource. The `database` resource handles permissions, while `app.yaml` provides the instance name to your code. + +### Why not valueFrom for Lakebase? + +The `database` resource's `valueFrom` provides the full connection string (e.g., `instance-xxx.database.staging.cloud.databricks.com`), but `AsyncDatabricksStore` expects just the instance name. So we use a static `value` instead. + --- -## Adding Long-Term Memory +## Step 3: Configure .env.local (Local Development) + +For local development, add to `.env.local`: + +```bash +# Lakebase configuration for long-term memory +LAKEBASE_INSTANCE_NAME= +EMBEDDING_ENDPOINT=databricks-gte-large-en +EMBEDDING_DIMS=1024 +``` + +> **Note:** `.env.local` is only for local development. When deployed, the app gets `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. + +--- -Long-term memory stores facts about users that persist across conversation sessions. +## Step 4: Update agent.py -### Step 1: Import and Initialize Store +### Add Imports and Configuration ```python -from databricks_langchain import AsyncDatabricksStore +import os +import uuid +from typing import AsyncGenerator + +from databricks_langchain import AsyncDatabricksStore, ChatDatabricks +from langchain_core.runnables import RunnableConfig +from langchain_core.tools import tool +from langgraph.prebuilt import create_react_agent +# Environment configuration LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") EMBEDDING_ENDPOINT = os.getenv("EMBEDDING_ENDPOINT", "databricks-gte-large-en") +EMBEDDING_DIMS = int(os.getenv("EMBEDDING_DIMS", "1024")) # REQUIRED! +``` -async with AsyncDatabricksStore( - instance_name=LAKEBASE_INSTANCE_NAME, - embedding_endpoint=EMBEDDING_ENDPOINT, -) as store: - agent = create_react_agent( - model=model, - tools=tools + memory_tools, # Add memory tools - store=store, - ) +**CRITICAL:** You MUST specify `embedding_dims` when using `embedding_endpoint`. Without it, you'll get: +``` +"embedding_dims is required when embedding_endpoint is specified" ``` -### Step 2: Create Memory Tools +### Create Memory Tools ```python -from langchain_core.tools import tool -from langchain_core.runnables import RunnableConfig - @tool async def get_user_memory(query: str, config: RunnableConfig) -> str: - """Search for relevant information about the user from long-term memory.""" + """Search for relevant information about the user from long-term memory. + Use this to recall preferences, past interactions, or other saved information. + """ user_id = config.get("configurable", {}).get("user_id") store = config.get("configurable", {}).get("store") if not user_id or not store: - return "Memory not available" + return "Memory not available - no user context" - namespace = ("user_memories", user_id.replace(".", "-")) + # Sanitize user_id for namespace (replace special chars) + namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) results = await store.asearch(namespace, query=query, limit=5) if not results: - return "No memories found" + return "No memories found for this query" return "\n".join([f"- {r.value}" for r in results]) + @tool async def save_user_memory(memory_key: str, memory_data: str, config: RunnableConfig) -> str: - """Save information about the user to long-term memory.""" + """Save information about the user to long-term memory. + Use this to remember user preferences, important details, or other + information that should persist across conversations. + + Args: + memory_key: A short descriptive key (e.g., "preferred_name", "team", "region") + memory_data: The information to save + """ user_id = config.get("configurable", {}).get("user_id") store = config.get("configurable", {}).get("store") if not user_id or not store: - return "Memory not available" + return "Memory not available - no user context" - namespace = ("user_memories", user_id.replace(".", "-")) + namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) await store.aput(namespace, memory_key, {"value": memory_data}) return f"Saved memory: {memory_key}" ``` -### Step 3: Use user_id in Requests +### Add Helper Functions ```python def _get_user_id(request: ResponsesAgentRequest) -> str: + """Extract user_id from request context or custom inputs.""" custom_inputs = dict(request.custom_inputs or {}) - if "user_id" in custom_inputs: - return custom_inputs["user_id"] + if "user_id" in custom_inputs and custom_inputs["user_id"]: + return str(custom_inputs["user_id"]) if request.context and getattr(request.context, "user_id", None): - return request.context.user_id + return str(request.context.user_id) return "default-user" -# Use in agent invocation -config = {"configurable": {"user_id": user_id, "store": store}} + +def _get_or_create_thread_id(request: ResponsesAgentRequest) -> str: + """Extract or create thread_id for conversation tracking.""" + custom_inputs = dict(request.custom_inputs or {}) + if "thread_id" in custom_inputs and custom_inputs["thread_id"]: + return str(custom_inputs["thread_id"]) + if request.context and getattr(request.context, "conversation_id", None): + return str(request.context.conversation_id) + return str(uuid.uuid4()) +``` + +### Update Streaming Function + +```python +@stream() +async def streaming( + request: ResponsesAgentRequest, +) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + # Get user context + user_id = _get_user_id(request) + thread_id = _get_or_create_thread_id(request) + + # Get other tools (MCP, etc.) + mcp_client = init_mcp_client(sp_workspace_client) + mcp_tools = await mcp_client.get_tools() + + # Memory tools + memory_tools = [get_user_memory, save_user_memory] + + # Combine all tools + all_tools = mcp_tools + memory_tools + + # Initialize model + model = ChatDatabricks(endpoint="databricks-claude-sonnet-4") + + # Use AsyncDatabricksStore for long-term memory + async with AsyncDatabricksStore( + instance_name=LAKEBASE_INSTANCE_NAME, + embedding_endpoint=EMBEDDING_ENDPOINT, + embedding_dims=EMBEDDING_DIMS, # REQUIRED! + ) as store: + # Create agent with tools + agent = create_react_agent( + model=model, + tools=all_tools, + prompt=AGENT_INSTRUCTIONS, + ) + + # Prepare input + messages = {"messages": to_chat_completions_input([i.model_dump() for i in request.input])} + + # Configure with user context for memory tools + config = { + "configurable": { + "user_id": user_id, + "thread_id": thread_id, + "store": store, # Pass store to tools via config + } + } + + # Stream agent responses + async for event in process_agent_astream_events( + agent.astream(input=messages, config=config, stream_mode=["updates", "messages"]) + ): + yield event +``` + +--- + +## Step 5: Initialize Tables and Deploy + +### Initialize Lakebase Tables (First Time Only) + +Before deploying, initialize the tables locally: + +```bash +uv run python -c "$(cat <<'EOF' +import asyncio +from databricks_langchain import AsyncDatabricksStore + +async def setup(): + async with AsyncDatabricksStore( + instance_name="", + embedding_endpoint="databricks-gte-large-en", + embedding_dims=1024, + ) as store: + await store.setup() + print("Tables created!") + +asyncio.run(setup()) +EOF +)" +``` + +### Deploy and Run + +**IMPORTANT:** Always run `databricks bundle run` after `databricks bundle deploy` to start/restart the app with the new code: + +```bash +# Deploy resources and upload files +databricks bundle deploy + +# Start/restart the app with new code (REQUIRED!) +databricks bundle run agent_langgraph +``` + +> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is required to actually start the app with the new code. + +--- + +## Complete Example Files + +### databricks.yml (with Lakebase) + +```yaml +bundle: + name: agent_langgraph + +resources: + experiments: + agent_langgraph_experiment: + name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} + + apps: + agent_langgraph: + name: "my-agent-app" + description: "Agent with long-term memory" + source_code_path: ./ + + resources: + - name: 'experiment' + experiment: + experiment_id: "${resources.experiments.agent_langgraph_experiment.id}" + permission: 'CAN_MANAGE' + + # Lakebase instance for long-term memory + - name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +targets: + dev: + mode: development + default: true +``` + +### app.yaml + +```yaml +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" + # Reference experiment resource from databricks.yml + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" + # Lakebase instance name (must match instance_name in databricks.yml) + - name: LAKEBASE_INSTANCE_NAME + value: "" + # Embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` + +--- + +## Adding Short-Term Memory + +Short-term memory stores conversation history within a thread. + +### Import and Initialize Checkpointer + +```python +from databricks_langchain import AsyncCheckpointSaver + +LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") + +async with AsyncCheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME) as checkpointer: + agent = create_react_agent( + model=model, + tools=tools, + checkpointer=checkpointer, # Enables conversation persistence + ) +``` + +### Use thread_id in Agent Config + +```python +config = {"configurable": {"thread_id": thread_id}} +async for event in agent.astream(input_state, config, stream_mode=["updates", "messages"]): + yield event ``` -### Test Long-Term Memory +--- + +## Testing Memory + +### Test Locally ```bash -# Save a preference +# Start the server +uv run start-app + +# Save a memory curl -X POST http://localhost:8000/invocations \ -H "Content-Type: application/json" \ -d '{ - "input": [{"role": "user", "content": "Remember that my favorite color is blue"}], + "input": [{"role": "user", "content": "Remember that I am on the shipping team"}], "custom_inputs": {"user_id": "alice@example.com"} }' -# Recall later (even in new conversation) +# Recall the memory curl -X POST http://localhost:8000/invocations \ -H "Content-Type: application/json" \ -d '{ - "input": [{"role": "user", "content": "What is my favorite color?"}], + "input": [{"role": "user", "content": "What team am I on?"}], "custom_inputs": {"user_id": "alice@example.com"} }' ``` ---- +### Test Deployed App + +```bash +# Get OAuth token (PATs don't work for apps) +TOKEN=$(databricks auth token --host | jq -r '.access_token') + +# Test memory save +curl -X POST https:///invocations \ + -H "Authorization: Bearer $TOKEN" \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Remember I prefer detailed explanations"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' +``` -## Using Pre-Built Memory Templates +--- -For fully configured memory implementations, use these templates instead: +## First-Time Setup Checklist -| Template | Memory Type | Key Features | -|----------|-------------|--------------| -| `agent-langgraph-short-term-memory` | Short-term | AsyncCheckpointSaver, thread_id, conversation persistence | -| `agent-langgraph-long-term-memory` | Long-term | AsyncDatabricksStore, user_id, memory tools | +- [ ] Added `databricks-langchain[memory]` to `pyproject.toml` +- [ ] Run `uv sync` to install dependencies +- [ ] Created or identified Lakebase instance +- [ ] Added Lakebase env vars to `.env.local` (for local dev) +- [ ] Added `database` resource to `databricks.yml` (for permissions) +- [ ] Added `LAKEBASE_INSTANCE_NAME` value to `app.yaml` (matching instance_name in databricks.yml) +- [ ] **Initialized tables locally** by running `await store.setup()` +- [ ] Deployed with `databricks bundle deploy` +- [ ] **Started app with `databricks bundle run agent_langgraph`** -These templates have memory fully integrated and tested. +--- ## Troubleshooting -| Issue | Solution | -|-------|----------| -| **"Cannot connect to Lakebase"** | See **lakebase-setup** skill | -| **Memory not persisting** | Verify thread_id/user_id is passed consistently | -| **Permission errors** | Grant Lakebase permissions (see **lakebase-setup** skill) | -| **"Memory not available"** | Ensure user_id is provided in request | +| Issue | Cause | Solution | +|-------|-------|----------| +| **"embedding_dims is required"** | Missing parameter | Add `embedding_dims=1024` to AsyncDatabricksStore | +| **"relation 'store' does not exist"** | Tables not created | Run `await store.setup()` locally first | +| **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `valueFrom: "database"` to app.yaml | +| **"permission denied for table store"** | Missing grants | Lakebase `database` resource in DAB should handle this | +| **"Memory not available"** | No user_id in request | Ensure `custom_inputs.user_id` is passed | +| **Memory not persisting** | Different user_ids | Use consistent user_id across requests | +| **App not updated after deploy** | Forgot to run bundle | Run `databricks bundle run agent_langgraph` after deploy | + +--- + +## Pre-Built Memory Templates + +For fully configured implementations without manual setup: + +| Template | Memory Type | Key Features | +|----------|-------------|--------------| +| `agent-langgraph-short-term-memory` | Short-term | AsyncCheckpointSaver, thread_id | +| `agent-langgraph-long-term-memory` | Long-term | AsyncDatabricksStore, memory tools | + +--- ## Next Steps diff --git a/agent-langgraph/.claude/skills/deploy/SKILL.md b/agent-langgraph/.claude/skills/deploy/SKILL.md index 47c88f4b..073dafd9 100644 --- a/agent-langgraph/.claude/skills/deploy/SKILL.md +++ b/agent-langgraph/.claude/skills/deploy/SKILL.md @@ -7,14 +7,18 @@ description: "Deploy agent to Databricks Apps using DAB (Databricks Asset Bundle ## Deploy Commands +**IMPORTANT:** Always run BOTH commands to deploy and start your app: + ```bash -# Deploy the bundle (creates/updates resources, uploads files) +# 1. Deploy the bundle (creates/updates resources, uploads files) databricks bundle deploy -# Run the app (starts/restarts with uploaded source code) +# 2. Run the app (starts/restarts with uploaded source code) - REQUIRED! databricks bundle run agent_langgraph ``` +> **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_langgraph` matches the app name in `databricks.yml` under `resources.apps`. ## Handling "App Already Exists" Error @@ -143,3 +147,5 @@ databricks apps get --output json | jq -r '.url' | Auth token expired | Run `databricks auth token` again | | "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_langgraph` after deploy | +| Env var is None in deployed app | Check `valueFrom` in app.yaml matches resource `name` in databricks.yml | diff --git a/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md b/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md index 45ad68b4..3dc912da 100644 --- a/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md +++ b/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md @@ -15,13 +15,24 @@ Lakebase provides persistent storage for agent memory: - **Short-term memory**: Conversation history within a thread (`AsyncCheckpointSaver`) - **Long-term memory**: User facts across sessions (`AsyncDatabricksStore`) +## Complete Setup Workflow + +``` +┌─────────────────────────────────────────────────────────────────────────────┐ +│ 1. Add dependency → 2. Get instance → 3. Configure DAB + app.yaml │ +│ 4. Configure .env.local → 5. Initialize tables → 6. Deploy + Run │ +└─────────────────────────────────────────────────────────────────────────────┘ +``` + +--- + ## Step 1: Add Memory Dependency Add the memory extra to your `pyproject.toml`: ```toml dependencies = [ - "databricks-langchain[memory]>=0.13.0", + "databricks-langchain[memory]", # ... other dependencies ] ``` @@ -31,80 +42,287 @@ Then sync dependencies: uv sync ``` -## Step 2: Configure Lakebase Instance +--- + +## Step 2: Create or Get Lakebase Instance + +### Option A: Create New Instance (via Databricks UI) + +1. Go to your Databricks workspace +2. Navigate to **Compute** → **Lakebase** +3. Click **Create Instance** +4. Note the instance name + +### Option B: Use Existing Instance + +If you have an existing instance, note its name for the next step. + +--- + +## Step 3: Configure databricks.yml (Lakebase Resource) + +Add the Lakebase `database` resource to your app in `databricks.yml`: + +```yaml +resources: + apps: + agent_langgraph: + 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: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' +``` + +**Important:** +- The `name: 'database'` must match the `valueFrom` reference in `app.yaml` +- Using the `database` resource type automatically grants the app's service principal access to Lakebase + +### Update app.yaml (Environment Variables) + +Update `app.yaml` with the Lakebase instance name: + +```yaml +env: + # ... other env vars ... + + # Lakebase instance name - must match instance_name in databricks.yml database resource + # Note: Use 'value' (not 'valueFrom') because AsyncDatabricksStore needs the instance name, + # not the full connection string that valueFrom would provide + - name: LAKEBASE_INSTANCE_NAME + value: "" + + # Static values for embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` + +**Important:** +- The `LAKEBASE_INSTANCE_NAME` value must match the `instance_name` in your `databricks.yml` database resource +- The `database` resource handles permissions; `app.yaml` provides the instance name to your code +- Don't use `valueFrom` for Lakebase - it provides the connection string, not the instance name + +--- + +## Step 4: Configure .env.local (Local Development) + +For local development, add to `.env.local`: -Add to your `.env`: ```bash -LAKEBASE_INSTANCE_NAME="" +# Lakebase configuration for long-term memory +LAKEBASE_INSTANCE_NAME= +EMBEDDING_ENDPOINT=databricks-gte-large-en +EMBEDDING_DIMS=1024 ``` -The quickstart script can help you find or create a Lakebase instance. +**Important:** `embedding_dims` must match the embedding endpoint: + +| Endpoint | Dimensions | +|----------|------------| +| `databricks-gte-large-en` | 1024 | +| `databricks-bge-large-en` | 1024 | -## Step 3: Grant Permissions (After Deployment) +> **Note:** `.env.local` is only for local development. When deployed, the app gets `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. + +--- -After deploying your app, grant the app's service principal access to Lakebase. +## Step 5: Initialize Store Tables (CRITICAL - First Time Only) -### Option A: Using LakebaseClient SDK (Recommended) +**Before deploying**, you must initialize the Lakebase tables. The `AsyncDatabricksStore` creates tables on first use, but you need to do this locally first: ```python -from databricks_ai_bridge.lakebase import LakebaseClient, SchemaPrivilege, TablePrivilege +# Run this script locally BEFORE first deployment +import asyncio +from databricks_langchain import AsyncDatabricksStore -# Initialize client -client = LakebaseClient(instance_name="") +async def setup_store(): + async with AsyncDatabricksStore( + instance_name="", + embedding_endpoint="databricks-gte-large-en", + embedding_dims=1024, + ) as store: + print("Setting up store tables...") + await store.setup() # Creates required tables + print("Store tables created!") -# Get app service principal from: databricks apps get -app_sp = "" + # Verify with a test write/read + await store.aput(("test", "init"), "test_key", {"value": "test_value"}) + results = await store.asearch(("test", "init"), query="test", limit=1) + print(f"Test successful: {results}") -# Create role for the service principal -client.create_role(app_sp, "SERVICE_PRINCIPAL") +asyncio.run(setup_store()) +``` -# Grant schema privileges -client.grant_schema( - role=app_sp, - schema="public", - privileges=[SchemaPrivilege.USAGE, SchemaPrivilege.CREATE], -) +Run with: +```bash +uv run python -c "$(cat <<'EOF' +import asyncio +from databricks_langchain import AsyncDatabricksStore -# Grant table privileges (for existing tables) -client.grant_all_tables_in_schema( - role=app_sp, - schema="public", - privileges=[TablePrivilege.SELECT, TablePrivilege.INSERT, TablePrivilege.UPDATE, TablePrivilege.DELETE], -) +async def setup(): + async with AsyncDatabricksStore( + instance_name="", + embedding_endpoint="databricks-gte-large-en", + embedding_dims=1024, + ) as store: + await store.setup() + print("Tables created!") + +asyncio.run(setup()) +EOF +)" +``` + +This creates these tables in the `public` schema: +- `store` - Key-value storage for memories +- `store_vectors` - Vector embeddings for semantic search +- `store_migrations` - Schema migration tracking +- `vector_migrations` - Vector schema migration tracking + +--- + +## Step 6: Deploy and Run Your App + +**IMPORTANT:** Always run both `deploy` AND `run` commands: + +```bash +# Deploy resources and upload files +databricks bundle deploy + +# Start/restart the app with new code (REQUIRED!) +databricks bundle run agent_langgraph ``` -### Option B: Using SQL +> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is required to actually start the app with the new code. + +--- + +## Complete Example: databricks.yml with Lakebase -Connect to your Lakebase instance via SQL and run: +```yaml +bundle: + name: agent_langgraph -```sql --- Replace with your app's service principal --- Get it from: databricks apps get --output json | jq -r '.service_principal_id' +resources: + experiments: + agent_langgraph_experiment: + name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} --- Grant schema access -GRANT USAGE, CREATE ON SCHEMA public TO ""; + apps: + agent_langgraph: + name: "my-agent-app" + description: "Agent with long-term memory" + source_code_path: ./ --- Grant table access for memory tables -GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO ""; + resources: + - name: 'experiment' + experiment: + experiment_id: "${resources.experiments.agent_langgraph_experiment.id}" + permission: 'CAN_MANAGE' + + # Lakebase instance for long-term memory + - name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +targets: + dev: + mode: development + default: true ``` -## Adding Memory to Your Agent +## Complete Example: app.yaml + +```yaml +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" + # Reference experiment resource from databricks.yml + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" + # Lakebase instance name (must match instance_name in databricks.yml) + - name: LAKEBASE_INSTANCE_NAME + value: "" + # Embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` -See the **agent-memory** skill for code patterns to add: -- Short-term memory (conversation history) -- Long-term memory (persistent user facts) +--- ## Troubleshooting -| Issue | Solution | -|-------|----------| -| **"Failed to connect to Lakebase"** | Verify `LAKEBASE_INSTANCE_NAME` in `.env` | -| **Permission denied on tables** | Run the SQL grants above after deployment | -| **"Role does not exist"** | Use the service principal ID, not name | -| **Tables not created** | Tables are created automatically on first use | +| 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 | Run `await store.setup()` locally first (Step 5) | +| **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `LAKEBASE_INSTANCE_NAME` value to app.yaml | +| **"Unable to resolve Lakebase instance '...database.cloud.databricks.com'"** | Used valueFrom instead of value | Use `value: ""` not `valueFrom` for Lakebase | +| **"permission denied for table store"** | Missing grants | The `database` resource in DAB should handle this; verify the resource is configured | +| **"Failed to connect to Lakebase"** | Wrong instance name | Verify instance name in databricks.yml and .env.local | +| **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 agent_langgraph` after deploy | +| **valueFrom not resolving** | Resource name mismatch | Ensure `valueFrom` value matches `name` in databricks.yml resources | + +--- + +## Quick Reference: LakebaseClient API + +For manual permission management (usually not needed with DAB `database` resource): + +```python +from databricks_ai_bridge.lakebase import LakebaseClient, SchemaPrivilege, TablePrivilege + +client = LakebaseClient(instance_name="...") + +# 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'") +``` + +--- ## Next Steps -- Add memory patterns: see **agent-memory** skill +- Add memory to agent code: see **agent-memory** skill - Test locally: see **run-locally** skill - Deploy: see **deploy** skill diff --git a/agent-langgraph/AGENTS.md b/agent-langgraph/AGENTS.md index 896b1c9d..0cafbcb7 100644 --- a/agent-langgraph/AGENTS.md +++ b/agent-langgraph/AGENTS.md @@ -1,6 +1,11 @@ # Agent Development Guide -## MANDATORY First Action +## MANDATORY First Actions + +**BEFORE any actions, ask for crucial follow-ups about the app deployment specifics (do this interactively!)** + +1. ALWAYS make sure to ask the user if they have an existing app name that they want to deploy to and ask for them to provide it to you first interactively! +2. (MEMORY ONLY OPTIONAL) ALWAYS make sure to ask the user if they have an existing lakebase instance that they want to use for memory and provide it to you first interactively! **BEFORE any other action, run `databricks auth profiles` to check authentication status.** diff --git a/agent-langgraph/pyproject.toml b/agent-langgraph/pyproject.toml index 16c2651e..75c70949 100644 --- a/agent-langgraph/pyproject.toml +++ b/agent-langgraph/pyproject.toml @@ -10,7 +10,7 @@ requires-python = ">=3.11" dependencies = [ "fastapi>=0.115.12", "uvicorn>=0.34.2", - "databricks-langchain>=0.12.0", + "databricks-langchain>=0.14.0", "mlflow>=3.8.0rc0", "langgraph>=1.0.1", "langchain-mcp-adapters>=0.1.11", From e49d924327c687a66b67d4f56eadf1874939faba Mon Sep 17 00:00:00 2001 From: Dhruv Gupta Date: Thu, 29 Jan 2026 11:44:19 -0800 Subject: [PATCH 07/14] Fixed .env.local references --- agent-langgraph/.claude/skills/agent-memory/SKILL.md | 10 +++++----- agent-langgraph/.claude/skills/lakebase-setup/SKILL.md | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/agent-langgraph/.claude/skills/agent-memory/SKILL.md b/agent-langgraph/.claude/skills/agent-memory/SKILL.md index f06c5d23..ae840930 100644 --- a/agent-langgraph/.claude/skills/agent-memory/SKILL.md +++ b/agent-langgraph/.claude/skills/agent-memory/SKILL.md @@ -40,7 +40,7 @@ Adding memory requires changes to **4 files**: | File | What to Add | |------|-------------| | `pyproject.toml` | Memory dependency + hatch config | -| `.env.local` | Lakebase env vars (for local dev) | +| `.env` | Lakebase env vars (for local dev) | | `databricks.yml` | Lakebase database resource | | `app.yaml` | `valueFrom` reference to lakebase resource | | `agent_server/agent.py` | Memory tools and AsyncDatabricksStore | @@ -108,9 +108,9 @@ The `database` resource's `valueFrom` provides the full connection string (e.g., --- -## Step 3: Configure .env.local (Local Development) +## Step 3: Configure .env (Local Development) -For local development, add to `.env.local`: +For local development, add to `.env`: ```bash # Lakebase configuration for long-term memory @@ -119,7 +119,7 @@ EMBEDDING_ENDPOINT=databricks-gte-large-en EMBEDDING_DIMS=1024 ``` -> **Note:** `.env.local` is only for local development. When deployed, the app gets `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. +> **Note:** `.env` is only for local development. When deployed, the app gets `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. --- @@ -457,7 +457,7 @@ curl -X POST https:///invocations \ - [ ] Added `databricks-langchain[memory]` to `pyproject.toml` - [ ] Run `uv sync` to install dependencies - [ ] Created or identified Lakebase instance -- [ ] Added Lakebase env vars to `.env.local` (for local dev) +- [ ] Added Lakebase env vars to `.env` (for local dev) - [ ] Added `database` resource to `databricks.yml` (for permissions) - [ ] Added `LAKEBASE_INSTANCE_NAME` value to `app.yaml` (matching instance_name in databricks.yml) - [ ] **Initialized tables locally** by running `await store.setup()` diff --git a/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md b/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md index 3dc912da..1fe954a4 100644 --- a/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md +++ b/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md @@ -20,7 +20,7 @@ Lakebase provides persistent storage for agent memory: ``` ┌─────────────────────────────────────────────────────────────────────────────┐ │ 1. Add dependency → 2. Get instance → 3. Configure DAB + app.yaml │ -│ 4. Configure .env.local → 5. Initialize tables → 6. Deploy + Run │ +│ 4. Configure .env → 5. Initialize tables → 6. Deploy + Run │ └─────────────────────────────────────────────────────────────────────────────┘ ``` @@ -113,9 +113,9 @@ env: --- -## Step 4: Configure .env.local (Local Development) +## Step 4: Configure .env (Local Development) -For local development, add to `.env.local`: +For local development, add to `.env`: ```bash # Lakebase configuration for long-term memory @@ -131,7 +131,7 @@ EMBEDDING_DIMS=1024 | `databricks-gte-large-en` | 1024 | | `databricks-bge-large-en` | 1024 | -> **Note:** `.env.local` is only for local development. When deployed, the app gets `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. +> **Note:** `.env` is only for local development. When deployed, the app gets `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. --- @@ -282,7 +282,7 @@ env: | **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `LAKEBASE_INSTANCE_NAME` value to app.yaml | | **"Unable to resolve Lakebase instance '...database.cloud.databricks.com'"** | Used valueFrom instead of value | Use `value: ""` not `valueFrom` for Lakebase | | **"permission denied for table store"** | Missing grants | The `database` resource in DAB should handle this; verify the resource is configured | -| **"Failed to connect to Lakebase"** | Wrong instance name | Verify instance name in databricks.yml and .env.local | +| **"Failed to connect to Lakebase"** | Wrong instance name | Verify instance name 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 agent_langgraph` after deploy | | **valueFrom not resolving** | Resource name mismatch | Ensure `valueFrom` value matches `name` in databricks.yml resources | From 21d3289bc482b7ebf6ebae01d9d6aa97525955ca Mon Sep 17 00:00:00 2001 From: Dhruv Gupta Date: Thu, 29 Jan 2026 11:46:56 -0800 Subject: [PATCH 08/14] Added service principal info --- .../.claude/skills/lakebase-setup/SKILL.md | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md b/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md index 1fe954a4..6b8c802e 100644 --- a/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md +++ b/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md @@ -319,6 +319,29 @@ client.grant_table( 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 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 From aa3519ebedfc0d37957fb097f5eb86062383bad2 Mon Sep 17 00:00:00 2001 From: Dhruv Gupta Date: Fri, 30 Jan 2026 12:01:09 -0800 Subject: [PATCH 09/14] Consolidate skills with sync script - Add .claude/skills/ as source of truth for all skills - Add .claude/sync-skills.py to copy skills to templates - Consolidate deploy and deploy-memory into single deploy skill - Make discover-tools SDK-agnostic (no code examples) - Update AGENTS.md with interactive first actions across all templates - Fix .gitignore files (remove databricks.yml, fix .claude tracking) - Add databricks bundle validate step to deploy skill Skills are now maintained in one place and synced to templates, enabling `databricks workspace export-dir` to work correctly. Co-Authored-By: Claude (databricks-claude-opus-4-5) --- .claude/skills/add-tools-langgraph/SKILL.md | 104 ++++ .../examples/custom-mcp-server.md | 57 ++ .../examples/experiment.yaml | 8 + .../examples/genie-space.yaml | 9 + .../examples/lakebase.yaml | 18 + .../examples/serving-endpoint.yaml | 7 + .../examples/sql-warehouse.yaml | 7 + .../examples/uc-connection.yaml | 9 + .../examples/uc-function.yaml | 9 + .../examples/vector-search.yaml | 9 + .claude/skills/add-tools-openai/SKILL.md | 83 +++ .../examples/custom-mcp-server.md | 60 ++ .../add-tools-openai/examples/experiment.yaml | 8 + .../examples/genie-space.yaml | 9 + .../examples/serving-endpoint.yaml | 7 + .../examples/sql-warehouse.yaml | 7 + .../examples/uc-connection.yaml | 9 + .../examples/uc-function.yaml | 9 + .../examples/vector-search.yaml | 9 + .claude/skills/agent-memory/SKILL.md | 524 +++++++++++++++++ .claude/skills/deploy/SKILL.md | 222 +++++++ .claude/skills/discover-tools/SKILL.md | 47 ++ .claude/skills/lakebase-setup/SKILL.md | 351 +++++++++++ .../skills/modify-langgraph-agent/SKILL.md | 293 ++++++++++ .claude/skills/modify-openai-agent/SKILL.md | 146 +++++ .claude/skills/quickstart/SKILL.md | 83 +++ .claude/skills/run-locally/SKILL.md | 90 +++ .claude/sync-skills.py | 103 ++++ .gitignore | 16 + .../.claude/skills/add-tools/SKILL.md | 38 +- .../add-tools/examples/custom-mcp-server.md | 2 +- .../skills/add-tools/examples/lakebase.yaml | 18 + .../.claude/skills/agent-memory/SKILL.md | 550 ++++++++++++++---- .../.claude/skills/deploy/SKILL.md | 52 +- .../.claude/skills/discover-tools/SKILL.md | 43 +- .../.claude/skills/lakebase-setup/SKILL.md | 416 +++++++++---- .../.claude/skills/modify-agent/SKILL.md | 5 +- .../.claude/skills/quickstart/SKILL.md | 5 +- .../.claude/skills/run-locally/SKILL.md | 70 +-- agent-langgraph-long-term-memory/.gitignore | 2 - agent-langgraph-long-term-memory/AGENTS.md | 16 +- .../.claude/skills/add-tools/SKILL.md | 38 +- .../add-tools/examples/custom-mcp-server.md | 2 +- .../skills/add-tools/examples/lakebase.yaml | 18 + .../.claude/skills/agent-memory/SKILL.md | 550 ++++++++++++++---- .../.claude/skills/deploy/SKILL.md | 59 +- .../.claude/skills/discover-tools/SKILL.md | 43 +- .../.claude/skills/lakebase-setup/SKILL.md | 415 +++++++++---- .../.claude/skills/modify-agent/SKILL.md | 5 +- .../.claude/skills/quickstart/SKILL.md | 5 +- .../.claude/skills/run-locally/SKILL.md | 70 +-- agent-langgraph-short-term-memory/.gitignore | 2 - agent-langgraph-short-term-memory/AGENTS.md | 16 +- .../.claude/skills/agent-memory/SKILL.md | 28 +- .../.claude/skills/deploy/SKILL.md | 79 ++- .../.claude/skills/discover-tools/SKILL.md | 43 +- agent-langgraph/AGENTS.md | 15 +- .../.claude/skills/add-tools/SKILL.md | 104 ++++ .../add-tools/examples/custom-mcp-server.md | 57 ++ .../skills/add-tools/examples/experiment.yaml | 8 + .../add-tools/examples/genie-space.yaml | 9 + .../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/agent-memory/SKILL.md | 524 +++++++++++++++++ .../.claude/skills/deploy/SKILL.md | 222 +++++++ .../.claude/skills/discover-tools/SKILL.md | 47 ++ .../.claude/skills/lakebase-setup/SKILL.md | 351 +++++++++++ .../.claude/skills/modify-agent/SKILL.md | 293 ++++++++++ .../.claude/skills/quickstart/SKILL.md | 83 +++ .../.claude/skills/run-locally/SKILL.md | 90 +++ agent-non-conversational/.gitignore | 18 +- .../.claude/skills/deploy/SKILL.md | 87 ++- .../.claude/skills/discover-tools/SKILL.md | 41 +- agent-openai-agents-sdk/AGENTS.md | 16 +- 78 files changed, 6083 insertions(+), 844 deletions(-) create mode 100644 .claude/skills/add-tools-langgraph/SKILL.md create mode 100644 .claude/skills/add-tools-langgraph/examples/custom-mcp-server.md create mode 100644 .claude/skills/add-tools-langgraph/examples/experiment.yaml create mode 100644 .claude/skills/add-tools-langgraph/examples/genie-space.yaml create mode 100644 .claude/skills/add-tools-langgraph/examples/lakebase.yaml create mode 100644 .claude/skills/add-tools-langgraph/examples/serving-endpoint.yaml create mode 100644 .claude/skills/add-tools-langgraph/examples/sql-warehouse.yaml create mode 100644 .claude/skills/add-tools-langgraph/examples/uc-connection.yaml create mode 100644 .claude/skills/add-tools-langgraph/examples/uc-function.yaml create mode 100644 .claude/skills/add-tools-langgraph/examples/vector-search.yaml create mode 100644 .claude/skills/add-tools-openai/SKILL.md create mode 100644 .claude/skills/add-tools-openai/examples/custom-mcp-server.md create mode 100644 .claude/skills/add-tools-openai/examples/experiment.yaml create mode 100644 .claude/skills/add-tools-openai/examples/genie-space.yaml create mode 100644 .claude/skills/add-tools-openai/examples/serving-endpoint.yaml create mode 100644 .claude/skills/add-tools-openai/examples/sql-warehouse.yaml create mode 100644 .claude/skills/add-tools-openai/examples/uc-connection.yaml create mode 100644 .claude/skills/add-tools-openai/examples/uc-function.yaml create mode 100644 .claude/skills/add-tools-openai/examples/vector-search.yaml create mode 100644 .claude/skills/agent-memory/SKILL.md create mode 100644 .claude/skills/deploy/SKILL.md create mode 100644 .claude/skills/discover-tools/SKILL.md create mode 100644 .claude/skills/lakebase-setup/SKILL.md create mode 100644 .claude/skills/modify-langgraph-agent/SKILL.md create mode 100644 .claude/skills/modify-openai-agent/SKILL.md create mode 100644 .claude/skills/quickstart/SKILL.md create mode 100644 .claude/skills/run-locally/SKILL.md create mode 100755 .claude/sync-skills.py create mode 100644 agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/lakebase.yaml create mode 100644 agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/lakebase.yaml create mode 100644 agent-non-conversational/.claude/skills/add-tools/SKILL.md create mode 100644 agent-non-conversational/.claude/skills/add-tools/examples/custom-mcp-server.md create mode 100644 agent-non-conversational/.claude/skills/add-tools/examples/experiment.yaml create mode 100644 agent-non-conversational/.claude/skills/add-tools/examples/genie-space.yaml create mode 100644 agent-non-conversational/.claude/skills/add-tools/examples/lakebase.yaml create mode 100644 agent-non-conversational/.claude/skills/add-tools/examples/serving-endpoint.yaml create mode 100644 agent-non-conversational/.claude/skills/add-tools/examples/sql-warehouse.yaml create mode 100644 agent-non-conversational/.claude/skills/add-tools/examples/uc-connection.yaml create mode 100644 agent-non-conversational/.claude/skills/add-tools/examples/uc-function.yaml create mode 100644 agent-non-conversational/.claude/skills/add-tools/examples/vector-search.yaml create mode 100644 agent-non-conversational/.claude/skills/agent-memory/SKILL.md create mode 100644 agent-non-conversational/.claude/skills/deploy/SKILL.md create mode 100644 agent-non-conversational/.claude/skills/discover-tools/SKILL.md create mode 100644 agent-non-conversational/.claude/skills/lakebase-setup/SKILL.md create mode 100644 agent-non-conversational/.claude/skills/modify-agent/SKILL.md create mode 100644 agent-non-conversational/.claude/skills/quickstart/SKILL.md create mode 100644 agent-non-conversational/.claude/skills/run-locally/SKILL.md diff --git a/.claude/skills/add-tools-langgraph/SKILL.md b/.claude/skills/add-tools-langgraph/SKILL.md new file mode 100644 index 00000000..7719f198 --- /dev/null +++ b/.claude/skills/add-tools-langgraph/SKILL.md @@ -0,0 +1,104 @@ +--- +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 + +**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. + +## Workflow + +**Step 1:** Add MCP server in `agent_server/agent.py`: +```python +from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient + +genie_server = DatabricksMCPServer( + url=f"{host}/api/2.0/mcp/genie/01234567-89ab-cdef", + name="my genie space", +) + +mcp_client = DatabricksMultiServerMCPClient([genie_server]) +tools = await mcp_client.get_tools() +``` + +**Step 2:** Grant access in `databricks.yml`: +```yaml +resources: + apps: + agent_langgraph: + 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_langgraph # 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 | Agent memory storage | +| `experiment.yaml` | MLflow experiment | Tracing (already configured) | +| `custom-mcp-server.md` | Custom MCP apps | Apps starting with `mcp-*` | + +## Custom MCP Servers (Databricks Apps) + +Apps are **not yet supported** as resource dependencies in `databricks.yml`. Manual permission grant required: + +**Step 1:** Get your agent app's service principal: +```bash +databricks apps get --output json | jq -r '.service_principal_name' +``` + +**Step 2:** Grant permission on the MCP server app: +```bash +databricks apps update-permissions \ + --service-principal \ + --permission-level CAN_USE +``` + +See `examples/custom-mcp-server.md` for detailed steps. + +## valueFrom Pattern (for app.yaml) + +**IMPORTANT**: Make sure all `valueFrom` references in `app.yaml` reference an existing key in the `databricks.yml` file. +Some resources need environment variables in your app. Use `valueFrom` in `app.yaml` to reference resources defined in `databricks.yml`: + +```yaml +# app.yaml +env: + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" # References resources.apps..resources[name='experiment'] + - name: LAKEBASE_INSTANCE_NAME + valueFrom: "database" # References resources.apps..resources[name='database'] +``` + +**Critical:** Every `valueFrom` value must match a `name` field in `databricks.yml` resources. + +## 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_langgraph` +- **valueFrom matching**: Ensure `app.yaml` `valueFrom` values match `databricks.yml` resource `name` values diff --git a/.claude/skills/add-tools-langgraph/examples/custom-mcp-server.md b/.claude/skills/add-tools-langgraph/examples/custom-mcp-server.md new file mode 100644 index 00000000..1324e6c5 --- /dev/null +++ b/.claude/skills/add-tools-langgraph/examples/custom-mcp-server.md @@ -0,0 +1,57 @@ +# Custom MCP Server (Databricks App) + +Custom MCP servers are Databricks Apps with names starting with `mcp-*`. + +**Apps are not yet supported as resource dependencies in `databricks.yml`**, so manual permission grant is required. + +## 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. Deploy your agent app first + +```bash +databricks bundle deploy +databricks bundle run agent_langgraph +``` + +### 3. Get your agent app's service principal + +```bash +databricks apps get --output json | jq -r '.service_principal_name' +``` + +Example output: `sp-abc123-def456` + +### 4. Grant permission on the MCP server app + +```bash +databricks apps update-permissions \ + --service-principal \ + --permission-level CAN_USE +``` + +Example: +```bash +databricks apps update-permissions mcp-my-server \ + --service-principal sp-abc123-def456 \ + --permission-level CAN_USE +``` + +## Notes + +- This manual step is required each time you connect to a new custom MCP server +- The permission grant persists across deployments +- If you redeploy the agent app with a new service principal, you'll need to grant permissions again diff --git a/.claude/skills/add-tools-langgraph/examples/experiment.yaml b/.claude/skills/add-tools-langgraph/examples/experiment.yaml new file mode 100644 index 00000000..ac5c626a --- /dev/null +++ b/.claude/skills/add-tools-langgraph/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-langgraph/examples/genie-space.yaml b/.claude/skills/add-tools-langgraph/examples/genie-space.yaml new file mode 100644 index 00000000..71589d52 --- /dev/null +++ b/.claude/skills/add-tools-langgraph/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-langgraph/examples/lakebase.yaml b/.claude/skills/add-tools-langgraph/examples/lakebase.yaml new file mode 100644 index 00000000..78f0bc72 --- /dev/null +++ b/.claude/skills/add-tools-langgraph/examples/lakebase.yaml @@ -0,0 +1,18 @@ +# Lakebase Database (for agent memory) +# Use for: Long-term memory storage via AsyncDatabricksStore +# Requires: valueFrom reference in app.yaml + +# In databricks.yml - add to resources.apps..resources: +- name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +# In app.yaml - add to env: +# - name: LAKEBASE_INSTANCE_NAME +# valueFrom: "database" +# - name: EMBEDDING_ENDPOINT +# value: "databricks-gte-large-en" +# - name: EMBEDDING_DIMS +# value: "1024" diff --git a/.claude/skills/add-tools-langgraph/examples/serving-endpoint.yaml b/.claude/skills/add-tools-langgraph/examples/serving-endpoint.yaml new file mode 100644 index 00000000..b49ce9da --- /dev/null +++ b/.claude/skills/add-tools-langgraph/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-langgraph/examples/sql-warehouse.yaml b/.claude/skills/add-tools-langgraph/examples/sql-warehouse.yaml new file mode 100644 index 00000000..a6ce9446 --- /dev/null +++ b/.claude/skills/add-tools-langgraph/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-langgraph/examples/uc-connection.yaml b/.claude/skills/add-tools-langgraph/examples/uc-connection.yaml new file mode 100644 index 00000000..316675fe --- /dev/null +++ b/.claude/skills/add-tools-langgraph/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-langgraph/examples/uc-function.yaml b/.claude/skills/add-tools-langgraph/examples/uc-function.yaml new file mode 100644 index 00000000..43f938a9 --- /dev/null +++ b/.claude/skills/add-tools-langgraph/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-langgraph/examples/vector-search.yaml b/.claude/skills/add-tools-langgraph/examples/vector-search.yaml new file mode 100644 index 00000000..0ba39027 --- /dev/null +++ b/.claude/skills/add-tools-langgraph/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/add-tools-openai/SKILL.md b/.claude/skills/add-tools-openai/SKILL.md new file mode 100644 index 00000000..6639e557 --- /dev/null +++ b/.claude/skills/add-tools-openai/SKILL.md @@ -0,0 +1,83 @@ +--- +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 + +**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. + +## Workflow + +**Step 1:** Add MCP server in `agent_server/agent.py`: +```python +from databricks_openai.agents import McpServer + +genie_server = McpServer( + url=f"{host}/api/2.0/mcp/genie/01234567-89ab-cdef", + name="my genie space", +) + +agent = Agent( + name="my agent", + model="databricks-claude-3-7-sonnet", + mcp_servers=[genie_server], +) +``` + +**Step 2:** Grant access in `databricks.yml`: +```yaml +resources: + apps: + agent_openai_agents_sdk: + resources: + - name: 'my_genie_space' + genie_space: + name: 'My Genie Space' + space_id: '01234567-89ab-cdef' + permission: 'CAN_RUN' +``` + +**Step 3:** Deploy with `databricks bundle deploy` (see **deploy** skill) + +## 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 | +| `experiment.yaml` | MLflow experiment | Tracing (already configured) | +| `custom-mcp-server.md` | Custom MCP apps | Apps starting with `mcp-*` | + +## Custom MCP Servers (Databricks Apps) + +Apps are **not yet supported** as resource dependencies in `databricks.yml`. Manual permission grant required: + +**Step 1:** Get your agent app's service principal: +```bash +databricks apps get --output json | jq -r '.service_principal_name' +``` + +**Step 2:** Grant permission on the MCP server app: +```bash +databricks apps update-permissions \ + --service-principal \ + --permission-level CAN_USE +``` + +See `examples/custom-mcp-server.md` for detailed steps. + +## 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 after changes**: Run `databricks bundle deploy` after modifying `databricks.yml` diff --git a/.claude/skills/add-tools-openai/examples/custom-mcp-server.md b/.claude/skills/add-tools-openai/examples/custom-mcp-server.md new file mode 100644 index 00000000..804bb679 --- /dev/null +++ b/.claude/skills/add-tools-openai/examples/custom-mcp-server.md @@ -0,0 +1,60 @@ +# Custom MCP Server (Databricks App) + +Custom MCP servers are Databricks Apps with names starting with `mcp-*`. + +**Apps are not yet supported as resource dependencies in `databricks.yml`**, so manual permission grant is required. + +## Steps + +### 1. Add MCP server in `agent_server/agent.py` + +```python +from databricks_openai.agents import McpServer + +custom_mcp = McpServer( + url="https://mcp-my-server.cloud.databricks.com/mcp", + name="my custom mcp server", +) + +agent = Agent( + name="my agent", + model="databricks-claude-3-7-sonnet", + mcp_servers=[custom_mcp], +) +``` + +### 2. Deploy your agent app first + +```bash +databricks bundle deploy +databricks bundle run agent_openai_agents_sdk +``` + +### 3. Get your agent app's service principal + +```bash +databricks apps get --output json | jq -r '.service_principal_name' +``` + +Example output: `sp-abc123-def456` + +### 4. Grant permission on the MCP server app + +```bash +databricks apps update-permissions \ + --service-principal \ + --permission-level CAN_USE +``` + +Example: +```bash +databricks apps update-permissions mcp-my-server \ + --service-principal sp-abc123-def456 \ + --permission-level CAN_USE +``` + +## Notes + +- This manual step is required each time you connect to a new custom MCP server +- The permission grant persists across deployments +- If you redeploy the agent app with a new service principal, you'll need to grant permissions again diff --git a/.claude/skills/add-tools-openai/examples/experiment.yaml b/.claude/skills/add-tools-openai/examples/experiment.yaml new file mode 100644 index 00000000..ac5c626a --- /dev/null +++ b/.claude/skills/add-tools-openai/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-openai/examples/genie-space.yaml b/.claude/skills/add-tools-openai/examples/genie-space.yaml new file mode 100644 index 00000000..71589d52 --- /dev/null +++ b/.claude/skills/add-tools-openai/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-openai/examples/serving-endpoint.yaml b/.claude/skills/add-tools-openai/examples/serving-endpoint.yaml new file mode 100644 index 00000000..b49ce9da --- /dev/null +++ b/.claude/skills/add-tools-openai/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-openai/examples/sql-warehouse.yaml b/.claude/skills/add-tools-openai/examples/sql-warehouse.yaml new file mode 100644 index 00000000..a6ce9446 --- /dev/null +++ b/.claude/skills/add-tools-openai/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-openai/examples/uc-connection.yaml b/.claude/skills/add-tools-openai/examples/uc-connection.yaml new file mode 100644 index 00000000..316675fe --- /dev/null +++ b/.claude/skills/add-tools-openai/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-openai/examples/uc-function.yaml b/.claude/skills/add-tools-openai/examples/uc-function.yaml new file mode 100644 index 00000000..43f938a9 --- /dev/null +++ b/.claude/skills/add-tools-openai/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-openai/examples/vector-search.yaml b/.claude/skills/add-tools-openai/examples/vector-search.yaml new file mode 100644 index 00000000..0ba39027 --- /dev/null +++ b/.claude/skills/add-tools-openai/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/agent-memory/SKILL.md b/.claude/skills/agent-memory/SKILL.md new file mode 100644 index 00000000..83d59668 --- /dev/null +++ b/.claude/skills/agent-memory/SKILL.md @@ -0,0 +1,524 @@ +--- +name: agent-memory +description: "Add memory capabilities to your agent. Use when: (1) User asks about 'memory', 'state', 'remember', 'conversation history', (2) Want to persist conversations or user preferences, (3) Adding checkpointing or long-term storage." +--- + +# Adding Memory to Your Agent + +> **Note:** This template does not include memory by default. Use this skill to **add memory capabilities**. For pre-configured memory templates, see: +> - `agent-langgraph-short-term-memory` - Conversation history within a session +> - `agent-langgraph-long-term-memory` - User facts that persist across sessions + +## Memory Types + +| Type | Use Case | Storage | Identifier | +|------|----------|---------|------------| +| **Short-term** | Conversation history within a session | `AsyncCheckpointSaver` | `thread_id` | +| **Long-term** | User facts that persist across sessions | `AsyncDatabricksStore` | `user_id` | + +## Prerequisites + +1. **Add memory dependency** to `pyproject.toml`: + ```toml + dependencies = [ + "databricks-langchain[memory]", + ] + ``` + + Then run `uv sync` + +2. **Configure Lakebase** - See **lakebase-setup** skill for: + - Creating/configuring Lakebase instance + - Initializing tables (CRITICAL first-time step) + +--- + +## Quick Setup Summary + +Adding memory requires changes to **4 files**: + +| File | What to Add | +|------|-------------| +| `pyproject.toml` | Memory dependency + hatch config | +| `.env` | Lakebase env vars (for local dev) | +| `databricks.yml` | Lakebase database resource | +| `app.yaml` | `valueFrom` reference to lakebase resource | +| `agent_server/agent.py` | Memory tools and AsyncDatabricksStore | + +--- + +## Step 1: Configure databricks.yml (Lakebase Resource) + +Add the Lakebase database resource to your app in `databricks.yml`: + +```yaml +resources: + apps: + agent_langgraph: + name: "your-app-name" + source_code_path: ./ + + resources: + # ... other resources (experiment, UC functions, etc.) ... + + # Lakebase instance for long-term memory + - name: 'lakebadatabasese_memory' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' +``` + +**Important:** The `name: 'database'` must match the `valueFrom` reference in `app.yaml`. + +--- + +## Step 2: Configure app.yaml (Environment Variables) + +Update `app.yaml` with the Lakebase environment variables: + +```yaml +command: ["uv", "run", "start-app"] + +env: + # ... other env vars ... + + # MLflow experiment (uses valueFrom to reference databricks.yml resource) + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" + + # Lakebase instance name - must match the instance_name in databricks.yml database resource + # Note: Use 'value' (not 'valueFrom') because AsyncDatabricksStore needs the instance name, + # not the full connection string that valueFrom would provide + - name: LAKEBASE_INSTANCE_NAME + value: "" + + # Embedding configuration (static values) + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` + +**Important:** The `LAKEBASE_INSTANCE_NAME` value must match the `instance_name` in your `databricks.yml` database resource. The `database` resource handles permissions, while `app.yaml` provides the instance name to your code. + +### Why not valueFrom for Lakebase? + +The `database` resource's `valueFrom` provides the full connection string (e.g., `instance-xxx.database.staging.cloud.databricks.com`), but `AsyncDatabricksStore` expects just the instance name. So we use a static `value` instead. + +--- + +## Step 3: Configure .env (Local Development) + +For local development, add to `.env`: + +```bash +# Lakebase configuration for long-term memory +LAKEBASE_INSTANCE_NAME= +EMBEDDING_ENDPOINT=databricks-gte-large-en +EMBEDDING_DIMS=1024 +``` + +> **Note:** `.env` is only for local development. When deployed, the app gets `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. + +--- + +## Step 4: Update agent.py + +### Add Imports and Configuration + +```python +import os +import uuid +from typing import AsyncGenerator + +from databricks_langchain import AsyncDatabricksStore, ChatDatabricks +from langchain_core.runnables import RunnableConfig +from langchain_core.tools import tool +from langgraph.prebuilt import create_react_agent + +# Environment configuration +LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") +EMBEDDING_ENDPOINT = os.getenv("EMBEDDING_ENDPOINT", "databricks-gte-large-en") +EMBEDDING_DIMS = int(os.getenv("EMBEDDING_DIMS", "1024")) # REQUIRED! +``` + +**CRITICAL:** You MUST specify `embedding_dims` when using `embedding_endpoint`. Without it, you'll get: +``` +"embedding_dims is required when embedding_endpoint is specified" +``` + +### Create Memory Tools + +```python +@tool +async def get_user_memory(query: str, config: RunnableConfig) -> str: + """Search for relevant information about the user from long-term memory. + Use this to recall preferences, past interactions, or other saved information. + """ + user_id = config.get("configurable", {}).get("user_id") + store = config.get("configurable", {}).get("store") + if not user_id or not store: + return "Memory not available - no user context" + + # Sanitize user_id for namespace (replace special chars) + namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) + results = await store.asearch(namespace, query=query, limit=5) + if not results: + return "No memories found for this query" + return "\n".join([f"- {r.value}" for r in results]) + + +@tool +async def save_user_memory(memory_key: str, memory_data: str, config: RunnableConfig) -> str: + """Save information about the user to long-term memory. + Use this to remember user preferences, important details, or other + information that should persist across conversations. + + Args: + memory_key: A short descriptive key (e.g., "preferred_name", "team", "region") + memory_data: The information to save + """ + user_id = config.get("configurable", {}).get("user_id") + store = config.get("configurable", {}).get("store") + if not user_id or not store: + return "Memory not available - no user context" + + namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) + await store.aput(namespace, memory_key, {"value": memory_data}) + return f"Saved memory: {memory_key}" + + +@tool +async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: + """Delete a specific memory from the user's long-term memory. + Use this when the user asks to forget something or correct stored information. + + Args: + memory_key: The key of the memory to delete (e.g., "preferred_name", "team") + """ + user_id = config.get("configurable", {}).get("user_id") + store = config.get("configurable", {}).get("store") + if not user_id or not store: + return "Memory not available - no user context" + + namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) + await store.adelete(namespace, memory_key) + return f"Deleted memory: {memory_key}" +``` + +### Add Helper Functions + +```python +def _get_user_id(request: ResponsesAgentRequest) -> str: + """Extract user_id from request context or custom inputs.""" + custom_inputs = dict(request.custom_inputs or {}) + if "user_id" in custom_inputs and custom_inputs["user_id"]: + return str(custom_inputs["user_id"]) + if request.context and getattr(request.context, "user_id", None): + return str(request.context.user_id) + return "default-user" + + +def _get_or_create_thread_id(request: ResponsesAgentRequest) -> str: + """Extract or create thread_id for conversation tracking.""" + custom_inputs = dict(request.custom_inputs or {}) + if "thread_id" in custom_inputs and custom_inputs["thread_id"]: + return str(custom_inputs["thread_id"]) + if request.context and getattr(request.context, "conversation_id", None): + return str(request.context.conversation_id) + return str(uuid.uuid4()) +``` + +### Update Streaming Function + +```python +@stream() +async def streaming( + request: ResponsesAgentRequest, +) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + # Get user context + user_id = _get_user_id(request) + thread_id = _get_or_create_thread_id(request) + + # Get other tools (MCP, etc.) + mcp_client = init_mcp_client(sp_workspace_client) + mcp_tools = await mcp_client.get_tools() + + # Memory tools + memory_tools = [get_user_memory, save_user_memory, delete_user_memory] + + # Combine all tools + all_tools = mcp_tools + memory_tools + + # Initialize model + model = ChatDatabricks(endpoint="databricks-claude-sonnet-4") + + # Use AsyncDatabricksStore for long-term memory + async with AsyncDatabricksStore( + instance_name=LAKEBASE_INSTANCE_NAME, + embedding_endpoint=EMBEDDING_ENDPOINT, + embedding_dims=EMBEDDING_DIMS, # REQUIRED! + ) as store: + # Create agent with tools + agent = create_react_agent( + model=model, + tools=all_tools, + prompt=AGENT_INSTRUCTIONS, + ) + + # Prepare input + messages = {"messages": to_chat_completions_input([i.model_dump() for i in request.input])} + + # Configure with user context for memory tools + config = { + "configurable": { + "user_id": user_id, + "thread_id": thread_id, + "store": store, # Pass store to tools via config + } + } + + # Stream agent responses + async for event in process_agent_astream_events( + agent.astream(input=messages, config=config, stream_mode=["updates", "messages"]) + ): + yield event +``` + +--- + +## Step 5: Initialize Tables and Deploy + +### Initialize Lakebase Tables (First Time Only) + +Before deploying, initialize the tables locally: + +```bash +uv run python -c "$(cat <<'EOF' +import asyncio +from databricks_langchain import AsyncDatabricksStore + +async def setup(): + async with AsyncDatabricksStore( + instance_name="", + embedding_endpoint="databricks-gte-large-en", + embedding_dims=1024, + ) as store: + await store.setup() + print("Tables created!") + +asyncio.run(setup()) +EOF +)" +``` + +### Deploy and Run + +**IMPORTANT:** Always run `databricks bundle run` after `databricks bundle deploy` to start/restart the app with the new code: + +```bash +# Deploy resources and upload files +databricks bundle deploy + +# Start/restart the app with new code (REQUIRED!) +databricks bundle run agent_langgraph +``` + +> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is required to actually start the app with the new code. + +--- + +## Complete Example Files + +### databricks.yml (with Lakebase) + +```yaml +bundle: + name: agent_langgraph + +resources: + experiments: + agent_langgraph_experiment: + name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} + + apps: + agent_langgraph: + name: "my-agent-app" + description: "Agent with long-term memory" + source_code_path: ./ + + resources: + - name: 'experiment' + experiment: + experiment_id: "${resources.experiments.agent_langgraph_experiment.id}" + permission: 'CAN_MANAGE' + + # Lakebase instance for long-term memory + - name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +targets: + dev: + mode: development + default: true +``` + +### app.yaml + +```yaml +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" + # Reference experiment resource from databricks.yml + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" + # Lakebase instance name (must match instance_name in databricks.yml) + - name: LAKEBASE_INSTANCE_NAME + value: "" + # Embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` + +--- + +## Adding Short-Term Memory + +Short-term memory stores conversation history within a thread. + +### Import and Initialize Checkpointer + +```python +from databricks_langchain import AsyncCheckpointSaver + +LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") + +async with AsyncCheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME) as checkpointer: + agent = create_react_agent( + model=model, + tools=tools, + checkpointer=checkpointer, # Enables conversation persistence + ) +``` + +### Use thread_id in Agent Config + +```python +config = {"configurable": {"thread_id": thread_id}} +async for event in agent.astream(input_state, config, stream_mode=["updates", "messages"]): + yield event +``` + +--- + +## Testing Memory + +### Test Locally + +```bash +# Start the server +uv run start-app + +# Save a memory +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Remember that I am on the shipping team"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' + +# Recall the memory +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "What team am I on?"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' + +# Delete a memory +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Forget what team I am on"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' +``` + +### Test Deployed App + +```bash +# Get OAuth token (PATs don't work for apps) +TOKEN=$(databricks auth token --host | jq -r '.access_token') + +# Test memory save +curl -X POST https:///invocations \ + -H "Authorization: Bearer $TOKEN" \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Remember I prefer detailed explanations"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' +``` + +--- + +## First-Time Setup Checklist + +- [ ] Added `databricks-langchain[memory]` to `pyproject.toml` +- [ ] Run `uv sync` to install dependencies +- [ ] Created or identified Lakebase instance +- [ ] Added Lakebase env vars to `.env` (for local dev) +- [ ] Added `database` resource to `databricks.yml` (for permissions) +- [ ] Added `LAKEBASE_INSTANCE_NAME` value to `app.yaml` (matching instance_name in databricks.yml) +- [ ] **Initialized tables locally** by running `await store.setup()` +- [ ] Deployed with `databricks bundle deploy` +- [ ] **Started app with `databricks bundle run agent_langgraph`** + +--- + +## Troubleshooting + +| Issue | Cause | Solution | +|-------|-------|----------| +| **"embedding_dims is required"** | Missing parameter | Add `embedding_dims=1024` to AsyncDatabricksStore | +| **"relation 'store' does not exist"** | Tables not created | Run `await store.setup()` locally first | +| **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `valueFrom: "database"` to app.yaml | +| **"permission denied for table store"** | Missing grants | Lakebase `database` resource in DAB should handle this | +| **"Memory not available"** | No user_id in request | Ensure `custom_inputs.user_id` is passed | +| **Memory not persisting** | Different user_ids | Use consistent user_id across requests | +| **App not updated after deploy** | Forgot to run bundle | Run `databricks bundle run agent_langgraph` after deploy | + +--- + +## Pre-Built Memory Templates + +For fully configured implementations without manual setup: + +| Template | Memory Type | Key Features | +|----------|-------------|--------------| +| `agent-langgraph-short-term-memory` | Short-term | AsyncCheckpointSaver, thread_id | +| `agent-langgraph-long-term-memory` | Long-term | AsyncDatabricksStore, memory tools | + +--- + +## Next Steps + +- Configure Lakebase: see **lakebase-setup** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/.claude/skills/deploy/SKILL.md b/.claude/skills/deploy/SKILL.md new file mode 100644 index 00000000..ff3f8875 --- /dev/null +++ b/.claude/skills/deploy/SKILL.md @@ -0,0 +1,222 @@ +--- +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 + +## 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: + {{BUNDLE_NAME}}: + name: "agent-your-app-name" # Use agent-* prefix +``` + +## Deploy Commands + +**IMPORTANT:** Always run BOTH commands to deploy and start your app: + +```bash +# 1. Validate bundle configuration (catches errors before deploy) +databricks bundle validate + +# 2. Deploy the bundle (creates/updates resources, uploads files) +databricks bundle deploy + +# 3. Run the app (starts/restarts with uploaded source code) - REQUIRED! +databricks bundle run {{BUNDLE_NAME}} +``` + +> **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 `{{BUNDLE_NAME}}` 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: + {{BUNDLE_NAME}}: + 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 {{BUNDLE_NAME}} --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 {{BUNDLE_NAME}} +``` + +### 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 {{BUNDLE_NAME}} +``` + +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 +``` + +**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' +``` + +## 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 {{BUNDLE_NAME}}` after deploy | +| Env var is None in deployed app | Check `valueFrom` in app.yaml matches resource `name` in databricks.yml | diff --git a/.claude/skills/discover-tools/SKILL.md b/.claude/skills/discover-tools/SKILL.md new file mode 100644 index 00000000..87c3f519 --- /dev/null +++ b/.claude/skills/discover-tools/SKILL.md @@ -0,0 +1,47 @@ +--- +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 diff --git a/.claude/skills/lakebase-setup/SKILL.md b/.claude/skills/lakebase-setup/SKILL.md new file mode 100644 index 00000000..6b8c802e --- /dev/null +++ b/.claude/skills/lakebase-setup/SKILL.md @@ -0,0 +1,351 @@ +--- +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 Memory + +> **Note:** This template does not include memory by default. Use this skill if you want to **add memory capabilities** to your agent. For pre-configured memory templates, see: +> - `agent-langgraph-short-term-memory` - Conversation history within a session +> - `agent-langgraph-long-term-memory` - User facts that persist across sessions + +## Overview + +Lakebase provides persistent storage for agent memory: +- **Short-term memory**: Conversation history within a thread (`AsyncCheckpointSaver`) +- **Long-term memory**: User facts across sessions (`AsyncDatabricksStore`) + +## Complete Setup Workflow + +``` +┌─────────────────────────────────────────────────────────────────────────────┐ +│ 1. Add dependency → 2. Get instance → 3. Configure DAB + app.yaml │ +│ 4. Configure .env → 5. Initialize tables → 6. Deploy + Run │ +└─────────────────────────────────────────────────────────────────────────────┘ +``` + +--- + +## 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: Create New Instance (via Databricks UI) + +1. Go to your Databricks workspace +2. Navigate to **Compute** → **Lakebase** +3. Click **Create Instance** +4. Note the instance name + +### Option B: Use Existing Instance + +If you have an existing instance, note its name for the next step. + +--- + +## Step 3: Configure databricks.yml (Lakebase Resource) + +Add the Lakebase `database` resource to your app in `databricks.yml`: + +```yaml +resources: + apps: + agent_langgraph: + 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: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' +``` + +**Important:** +- The `name: 'database'` must match the `valueFrom` reference in `app.yaml` +- Using the `database` resource type automatically grants the app's service principal access to Lakebase + +### Update app.yaml (Environment Variables) + +Update `app.yaml` with the Lakebase instance name: + +```yaml +env: + # ... other env vars ... + + # Lakebase instance name - must match instance_name in databricks.yml database resource + # Note: Use 'value' (not 'valueFrom') because AsyncDatabricksStore needs the instance name, + # not the full connection string that valueFrom would provide + - name: LAKEBASE_INSTANCE_NAME + value: "" + + # Static values for embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` + +**Important:** +- The `LAKEBASE_INSTANCE_NAME` value must match the `instance_name` in your `databricks.yml` database resource +- The `database` resource handles permissions; `app.yaml` provides the instance name to your code +- Don't use `valueFrom` for Lakebase - it provides the connection string, not the instance name + +--- + +## Step 4: Configure .env (Local Development) + +For local development, add to `.env`: + +```bash +# Lakebase configuration for long-term memory +LAKEBASE_INSTANCE_NAME= +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 `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. + +--- + +## Step 5: Initialize Store Tables (CRITICAL - First Time Only) + +**Before deploying**, you must initialize the Lakebase tables. The `AsyncDatabricksStore` creates tables on first use, but you need to do this locally first: + +```python +# Run this script locally BEFORE first deployment +import asyncio +from databricks_langchain import AsyncDatabricksStore + +async def setup_store(): + async with AsyncDatabricksStore( + instance_name="", + embedding_endpoint="databricks-gte-large-en", + embedding_dims=1024, + ) as store: + print("Setting up store tables...") + await store.setup() # Creates required tables + print("Store tables created!") + + # Verify with a test write/read + await store.aput(("test", "init"), "test_key", {"value": "test_value"}) + results = await store.asearch(("test", "init"), query="test", limit=1) + print(f"Test successful: {results}") + +asyncio.run(setup_store()) +``` + +Run with: +```bash +uv run python -c "$(cat <<'EOF' +import asyncio +from databricks_langchain import AsyncDatabricksStore + +async def setup(): + async with AsyncDatabricksStore( + instance_name="", + embedding_endpoint="databricks-gte-large-en", + embedding_dims=1024, + ) as store: + await store.setup() + print("Tables created!") + +asyncio.run(setup()) +EOF +)" +``` + +This creates these tables in the `public` schema: +- `store` - Key-value storage for memories +- `store_vectors` - Vector embeddings for semantic search +- `store_migrations` - Schema migration tracking +- `vector_migrations` - Vector schema migration tracking + +--- + +## Step 6: Deploy and Run Your App + +**IMPORTANT:** Always run both `deploy` AND `run` commands: + +```bash +# Deploy resources and upload files +databricks bundle deploy + +# Start/restart the app with new code (REQUIRED!) +databricks bundle run agent_langgraph +``` + +> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is required to actually start the app with the new code. + +--- + +## Complete Example: databricks.yml with Lakebase + +```yaml +bundle: + name: agent_langgraph + +resources: + experiments: + agent_langgraph_experiment: + name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} + + apps: + agent_langgraph: + name: "my-agent-app" + description: "Agent with long-term memory" + source_code_path: ./ + + resources: + - name: 'experiment' + experiment: + experiment_id: "${resources.experiments.agent_langgraph_experiment.id}" + permission: 'CAN_MANAGE' + + # Lakebase instance for long-term memory + - name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +targets: + dev: + mode: development + default: true +``` + +## Complete Example: app.yaml + +```yaml +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" + # Reference experiment resource from databricks.yml + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" + # Lakebase instance name (must match instance_name in databricks.yml) + - name: LAKEBASE_INSTANCE_NAME + value: "" + # Embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` + +--- + +## 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 | Run `await store.setup()` locally first (Step 5) | +| **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `LAKEBASE_INSTANCE_NAME` value to app.yaml | +| **"Unable to resolve Lakebase instance '...database.cloud.databricks.com'"** | Used valueFrom instead of value | Use `value: ""` not `valueFrom` for Lakebase | +| **"permission denied for table store"** | Missing grants | The `database` resource in DAB should handle this; verify the resource is configured | +| **"Failed to connect to Lakebase"** | Wrong instance name | Verify instance name 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 agent_langgraph` after deploy | +| **valueFrom not resolving** | Resource name mismatch | Ensure `valueFrom` value matches `name` in databricks.yml resources | + +--- + +## Quick Reference: LakebaseClient API + +For manual permission management (usually not needed with DAB `database` resource): + +```python +from databricks_ai_bridge.lakebase import LakebaseClient, SchemaPrivilege, TablePrivilege + +client = LakebaseClient(instance_name="...") + +# 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 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/.claude/skills/modify-langgraph-agent/SKILL.md b/.claude/skills/modify-langgraph-agent/SKILL.md new file mode 100644 index 00000000..d7218637 --- /dev/null +++ b/.claude/skills/modify-langgraph-agent/SKILL.md @@ -0,0 +1,293 @@ +--- +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 + +## Main File + +**`agent_server/agent.py`** - Agent logic, model selection, instructions, MCP servers + +## Key Files + +| File | Purpose | +|------|---------| +| `agent_server/agent.py` | Agent logic, model, instructions, MCP servers | +| `agent_server/start_server.py` | FastAPI server + MLflow setup | +| `agent_server/evaluate_agent.py` | Agent evaluation with MLflow scorers | +| `agent_server/utils.py` | Databricks auth helpers, stream processing | +| `databricks.yml` | Bundle config & resource permissions | + +## SDK Setup + +```python +import mlflow +from databricks.sdk import WorkspaceClient +from databricks_langchain import ChatDatabricks, DatabricksMCPServer, DatabricksMultiServerMCPClient +from langchain.agents import create_agent + +# Enable autologging for tracing +mlflow.langchain.autolog() + +# Initialize workspace client +workspace_client = WorkspaceClient() +``` + +--- + +## databricks-langchain SDK Overview + +**SDK Location:** https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/langchain + +Before making any changes, ensure that the APIs actually exist in the SDK. If something is missing from the documentation here, look in the venv's `site-packages` directory for the `databricks_langchain` package. If it's not installed, run `uv sync` to create the .venv and install the package. + +--- + +### ChatDatabricks - LLM Chat Interface + +Connects to Databricks Model Serving endpoints for LLM inference. + +```python +from databricks_langchain import ChatDatabricks + +llm = ChatDatabricks( + endpoint="databricks-claude-3-7-sonnet", # or databricks-meta-llama-3-1-70b-instruct + temperature=0, + max_tokens=500, +) + +# For Responses API agents: +llm = ChatDatabricks(endpoint="my-agent-endpoint", use_responses_api=True) +``` + +Available models (check workspace for current list): +- `databricks-claude-3-7-sonnet` +- `databricks-claude-3-5-sonnet` +- `databricks-meta-llama-3-3-70b-instruct` + +**Note:** Some workspaces require granting the app access to the serving endpoint in `databricks.yml`. See the **add-tools** skill and `examples/serving-endpoint.yaml`. + +--- + +### DatabricksEmbeddings - Generate Embeddings + +Query Databricks embedding model endpoints. + +```python +from databricks_langchain import DatabricksEmbeddings + +embeddings = DatabricksEmbeddings(endpoint="databricks-bge-large-en") +vector = embeddings.embed_query("The meaning of life is 42") +vectors = embeddings.embed_documents(["doc1", "doc2"]) +``` + +--- + +### DatabricksVectorSearch - Vector Store + +Connect to Databricks Vector Search indexes for similarity search. + +```python +from databricks_langchain import DatabricksVectorSearch + +# Delta-sync index with Databricks-managed embeddings +vs = DatabricksVectorSearch(index_name="catalog.schema.index_name") + +# Direct-access or self-managed embeddings +vs = DatabricksVectorSearch( + index_name="catalog.schema.index_name", + embedding=embeddings, + text_column="content", +) + +docs = vs.similarity_search("query", k=5) +``` + +--- + +### MCP Client - Tool Integration + +Connect to MCP (Model Context Protocol) servers to get tools for your agent. + +**Basic MCP Server (manual URL):** + +```python +from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient + +client = DatabricksMultiServerMCPClient([ + DatabricksMCPServer( + name="system-ai", + url=f"{host}/api/2.0/mcp/functions/system/ai", + ) +]) +tools = await client.get_tools() +``` + +**From UC Function (convenience helper):** + +Creates MCP server for Unity Catalog functions. If `function_name` is omitted, exposes all functions in the schema. + +```python +server = DatabricksMCPServer.from_uc_function( + catalog="main", + schema="tools", + function_name="send_email", # Optional - omit for all functions in schema + name="email-server", + timeout=30.0, + handle_tool_error=True, +) +``` + +**From Vector Search (convenience helper):** + +Creates MCP server for Vector Search indexes. If `index_name` is omitted, exposes all indexes in the schema. + +```python +server = DatabricksMCPServer.from_vector_search( + catalog="main", + schema="embeddings", + index_name="product_docs", # Optional - omit for all indexes in schema + name="docs-search", + timeout=30.0, +) +``` + +**From Genie Space:** + +Create MCP server from Genie Space. Get the genie space ID from the URL. + +Example: `https://workspace.cloud.databricks.com/genie/rooms/01f0515f6739169283ef2c39b7329700?o=123` means the genie space ID is `01f0515f6739169283ef2c39b7329700` + +```python +DatabricksMCPServer( + name="genie", + url=f"{host_name}/api/2.0/mcp/genie/01f0515f6739169283ef2c39b7329700", +) +``` + +**Non-Databricks MCP Server:** + +```python +from databricks_langchain import MCPServer + +server = MCPServer( + name="external-server", + url="https://other-server.com/mcp", + headers={"X-API-Key": "secret"}, + timeout=15.0, +) +``` + +**After adding MCP servers:** Grant permissions in `databricks.yml` (see **add-tools** skill) + +--- + +## Running the Agent + +```python +from langchain.agents import create_agent + +# Create agent - ONLY accepts tools and model, NO prompt/instructions parameter +agent = create_agent(tools=tools, model=llm) + +# Non-streaming +messages = {"messages": [{"role": "user", "content": "hi"}]} +result = await agent.ainvoke(messages) + +# Streaming +async for event in agent.astream(input=messages, stream_mode=["updates", "messages"]): + # Process stream events + pass +``` + +**Converting to Responses API format:** Use `process_agent_astream_events()` from `agent_server/utils.py`: + +```python +from agent_server.utils import process_agent_astream_events + +async for event in process_agent_astream_events( + agent.astream(input=messages, stream_mode=["updates", "messages"]) +): + yield event # Yields ResponsesAgentStreamEvent objects +``` + +--- + +## Customizing Agent Behavior (System Instructions) + +> **IMPORTANT:** `create_agent()` does NOT accept `prompt`, `instructions`, or `system_message` parameters. Attempting to pass these will cause a runtime error. + +In LangGraph, agent behavior is customized by prepending a system message to the conversation messages. + +**Correct pattern in `agent.py`:** + +1. Define instructions as a constant: +```python +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.""" +``` + +2. Prepend to messages in the `streaming()` function: +```python +@stream() +async def streaming(request: ResponsesAgentRequest) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + agent = await init_agent() + # Prepend system instructions to user messages + user_messages = to_chat_completions_input([i.model_dump() for i in request.input]) + messages = {"messages": [{"role": "system", "content": AGENT_INSTRUCTIONS}] + user_messages} + + async for event in process_agent_astream_events( + agent.astream(input=messages, stream_mode=["updates", "messages"]) + ): + yield event +``` + +**Common mistake to avoid:** +```python +# WRONG - will cause "unexpected keyword argument" error +agent = create_agent(tools=tools, model=llm, prompt=AGENT_INSTRUCTIONS) + +# CORRECT - add instructions via messages +messages = {"messages": [{"role": "system", "content": AGENT_INSTRUCTIONS}] + user_messages} +``` + +For advanced customization (routing, state management, custom graphs), refer to the [LangGraph documentation](https://docs.langchain.com/oss/python/langgraph/overview). + +--- + +## External Connection Tools + +Connect to external services via Unity Catalog HTTP connections: + +- **Slack** - Post messages to channels +- **Google Calendar** - Calendar operations +- **Microsoft Graph API** - Office 365 services +- **Azure AI Search** - Search functionality +- **Any HTTP API** - Use `http_request` from databricks-sdk + +Example: Create UC function wrapping HTTP request for Slack, then expose via MCP. + +--- + +## External Resources + +1. [databricks-langchain SDK](https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/langchain) +2. [Agent examples](https://github.com/bbqiu/agent-on-app-prototype) +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. [LangGraph documentation](https://docs.langchain.com/oss/python/langgraph/overview) +6. [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 +- Add memory capabilities: see **agent-memory** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/.claude/skills/modify-openai-agent/SKILL.md b/.claude/skills/modify-openai-agent/SKILL.md new file mode 100644 index 00000000..e577b455 --- /dev/null +++ b/.claude/skills/modify-openai-agent/SKILL.md @@ -0,0 +1,146 @@ +--- +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 + +## Main File + +**`agent_server/agent.py`** - Agent logic, model selection, instructions, MCP servers + +## Key Files + +| File | Purpose | +|------|---------| +| `agent_server/agent.py` | Agent logic, model, instructions, MCP servers | +| `agent_server/start_server.py` | FastAPI server + MLflow setup | +| `agent_server/evaluate_agent.py` | Agent evaluation with MLflow scorers | +| `agent_server/utils.py` | Databricks auth helpers, stream processing | +| `databricks.yml` | Bundle config & resource permissions | + +## SDK Setup + +```python +import mlflow +from databricks_openai import AsyncDatabricksOpenAI +from agents import set_default_openai_api, set_default_openai_client, Agent +from agents.tracing import set_trace_processors + +# Set up async client (recommended for agent servers) +set_default_openai_client(AsyncDatabricksOpenAI()) +set_default_openai_api("chat_completions") + +# Use MLflow for tracing (disables SDK's built-in tracing) +set_trace_processors([]) +mlflow.openai.autolog() +``` + +## Adding MCP Servers + +```python +from databricks_openai.agents import McpServer + +# UC Functions +uc_server = McpServer( + url=f"{host}/api/2.0/mcp/functions/{catalog}/{schema}", + name="uc functions", +) + +# Genie Space +genie_server = McpServer( + url=f"{host}/api/2.0/mcp/genie/{space_id}", + name="genie space", +) + +# Vector Search +vector_server = McpServer( + url=f"{host}/api/2.0/mcp/vector-search/{catalog}/{schema}/{index}", + name="vector search", +) + +# Add to agent +agent = Agent( + name="my agent", + instructions="You are a helpful agent.", + model="databricks-claude-3-7-sonnet", + mcp_servers=[uc_server, genie_server, vector_server], +) +``` + +**After adding MCP servers:** Grant permissions in `databricks.yml` (see **add-tools** skill) + +## Changing the Model + +Available models (check workspace for current list): +- `databricks-claude-3-7-sonnet` +- `databricks-claude-3-5-sonnet` +- `databricks-meta-llama-3-3-70b-instruct` + +```python +agent = Agent( + name="my agent", + model="databricks-claude-3-7-sonnet", # Change here + ... +) +``` + +**Note:** Some workspaces require granting the app access to the serving endpoint in `databricks.yml`. See the **add-tools** skill and `examples/serving-endpoint.yaml`. + +## Changing Instructions + +```python +agent = Agent( + name="my 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.""", + ... +) +``` + +## Running the Agent + +```python +from agents import Runner + +# Non-streaming +messages = [{"role": "user", "content": "hi"}] +result = await Runner.run(agent, messages) + +# Streaming +result = Runner.run_streamed(agent, input=messages) +async for event in result.stream_events(): + # Process stream events + pass +``` + +**Converting to Responses API format:** Use `process_agent_stream_events()` from `agent_server/utils.py` to convert streaming output to Responses API compatible format: + +```python +from agent_server.utils import process_agent_stream_events + +result = Runner.run_streamed(agent, input=messages) +async for event in process_agent_stream_events(result.stream_events()): + yield event # Yields ResponsesAgentStreamEvent objects +``` + +## External Resources + +1. [databricks-openai SDK](https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/openai) +2. [Agent examples](https://github.com/bbqiu/agent-on-app-prototype) +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. [OpenAI Agents SDK](https://platform.openai.com/docs/guides/agents-sdk) +6. [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/.claude/skills/quickstart/SKILL.md b/.claude/skills/quickstart/SKILL.md new file mode 100644 index 00000000..e550162c --- /dev/null +++ b/.claude/skills/quickstart/SKILL.md @@ -0,0 +1,83 @@ +--- +name: quickstart +description: "Set up Databricks agent development environment. Use when: (1) First time setup, (2) Configuring Databricks authentication, (3) User says 'quickstart', 'set up', 'authenticate', or 'configure databricks', (4) No .env file exists." +--- + +# Quickstart & Authentication + +## Prerequisites + +- **uv** (Python package manager) +- **nvm** with Node 20 (for frontend) +- **Databricks CLI v0.283.0+** + +Check CLI version: +```bash +databricks -v # Must be v0.283.0 or above +brew upgrade databricks # If version is too old +``` + +## Run Quickstart + +```bash +uv run quickstart +``` + +**Options:** +- `--profile NAME`: Use specified profile (non-interactive) +- `--host URL`: Workspace URL for initial setup +- `-h, --help`: Show help + +**Examples:** +```bash +# Interactive (prompts for profile selection) +uv run quickstart + +# Non-interactive with existing profile +uv run quickstart --profile DEFAULT + +# New workspace setup +uv run quickstart --host https://your-workspace.cloud.databricks.com +``` + +## What Quickstart Configures + +Creates/updates `.env` with: +- `DATABRICKS_CONFIG_PROFILE` - Selected CLI profile +- `MLFLOW_TRACKING_URI` - Set to `databricks://` for local auth +- `MLFLOW_EXPERIMENT_ID` - Auto-created experiment ID + +## Manual Authentication (Fallback) + +If quickstart fails: + +```bash +# Create new profile +databricks auth login --host https://your-workspace.cloud.databricks.com + +# Verify +databricks auth profiles +``` + +Then manually create `.env` (copy from `.env.example`): +```bash +# Authentication (choose one method) +DATABRICKS_CONFIG_PROFILE=DEFAULT +# DATABRICKS_HOST=https://.databricks.com +# DATABRICKS_TOKEN=dapi.... + +# MLflow configuration +MLFLOW_EXPERIMENT_ID= +MLFLOW_TRACKING_URI="databricks://DEFAULT" +MLFLOW_REGISTRY_URI="databricks-uc" + +# Frontend proxy settings +CHAT_APP_PORT=3000 +CHAT_PROXY_TIMEOUT_SECONDS=300 +``` + +## Next Steps + +After quickstart completes: +1. Run `uv run discover-tools` to find available workspace resources (see **discover-tools** skill) +2. Run `uv run start-app` to test locally (see **run-locally** skill) diff --git a/.claude/skills/run-locally/SKILL.md b/.claude/skills/run-locally/SKILL.md new file mode 100644 index 00000000..3eb83c82 --- /dev/null +++ b/.claude/skills/run-locally/SKILL.md @@ -0,0 +1,90 @@ +--- +name: run-locally +description: "Run and test the agent locally. Use when: (1) User says 'run locally', 'start server', 'test agent', or 'localhost', (2) Need curl commands to test API, (3) Troubleshooting local development issues, (4) Configuring server options like port or hot-reload." +--- + +# Run Agent Locally + +## Start the Server + +```bash +uv run start-app +``` + +This starts the agent at http://localhost:8000 + +## Server Options + +```bash +# Hot-reload on code changes (development) +uv run start-server --reload + +# Custom port +uv run start-server --port 8001 + +# Multiple workers (production-like) +uv run start-server --workers 4 + +# Combine options +uv run start-server --reload --port 8001 +``` + +## Test the API + +**Streaming request:** +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' +``` + +**Non-streaming request:** +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }] }' +``` + +## Run Evaluation + +```bash +uv run agent-evaluate +``` + +Uses MLflow scorers (RelevanceToQuery, Safety). + +## Run Unit Tests + +```bash +pytest [path] +``` + +## Troubleshooting + +| Issue | Solution | +|-------|----------| +| **Port already in use** | Use `--port 8001` or kill existing process | +| **Authentication errors** | Verify `.env` is correct; run **quickstart** skill | +| **Module not found** | Run `uv sync` to install dependencies | +| **MLflow experiment not found** | Ensure `MLFLOW_TRACKING_URI` in `.env` is `databricks://` | + +### MLflow Experiment Not Found + +If you see: "The provided MLFLOW_EXPERIMENT_ID environment variable value does not exist" + +**Verify the experiment exists:** +```bash +databricks -p experiments get-experiment +``` + +**Fix:** Ensure `.env` has the correct tracking URI format: +```bash +MLFLOW_TRACKING_URI="databricks://DEFAULT" # Include profile name +``` + +The quickstart script configures this automatically. If you manually edited `.env`, ensure the profile name is included. + +## Next Steps + +- Modify your agent: see **modify-agent** skill +- Deploy to Databricks: see **deploy** skill diff --git a/.claude/sync-skills.py b/.claude/sync-skills.py new file mode 100755 index 00000000..469b5572 --- /dev/null +++ b/.claude/sync-skills.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python3 +"""Sync skills from .claude/skills/ to all agent templates. + +This script copies skills from the source of truth (.claude/skills/) to each +template directory. Each template gets a complete copy of its skills (no symlinks) +so that `databricks workspace export-dir` works correctly. + +Usage: + python .claude/sync-skills.py +""" + +import os +import shutil +from pathlib import Path + +# Get repo root (parent of .claude directory where this script lives) +SCRIPT_DIR = Path(__file__).parent.resolve() +REPO_ROOT = SCRIPT_DIR.parent + +TEMPLATES = { + "agent-langgraph": { + "sdk": "langgraph", + "bundle_name": "agent_langgraph", + }, + "agent-langgraph-short-term-memory": { + "sdk": "langgraph", + "bundle_name": "agent_langgraph_short_term_memory", + }, + "agent-langgraph-long-term-memory": { + "sdk": "langgraph", + "bundle_name": "agent_langgraph_long_term_memory", + }, + "agent-openai-agents-sdk": { + "sdk": "openai", + "bundle_name": "agent_openai_agents_sdk", + }, + "agent-non-conversational": { + "sdk": "langgraph", + "bundle_name": "agent_non_conversational", + }, +} + +SOURCE = SCRIPT_DIR / "skills" + + +def copy_skill(src: Path, dest: Path, substitutions: dict = None): + """Copy skill directory, applying substitutions to SKILL.md.""" + dest.mkdir(parents=True, exist_ok=True) + + for item in src.iterdir(): + if item.is_dir(): + shutil.copytree(item, dest / item.name, dirs_exist_ok=True) + elif item.suffix == ".md" and substitutions: + content = item.read_text() + for placeholder, value in substitutions.items(): + content = content.replace(placeholder, value) + (dest / item.name).write_text(content) + else: + shutil.copy2(item, dest / item.name) + + +def sync_template(template: str, config: dict): + """Sync all skills to a single template.""" + dest = REPO_ROOT / template / ".claude" / "skills" + sdk = config["sdk"] + subs = {"{{BUNDLE_NAME}}": config["bundle_name"]} + + # Clear existing skills + if dest.exists(): + shutil.rmtree(dest) + dest.mkdir(parents=True) + + # Shared skills (no substitution needed) + for skill in ["quickstart", "run-locally", "discover-tools"]: + copy_skill(SOURCE / skill, dest / skill) + + # Deploy skill (with substitution) + copy_skill(SOURCE / "deploy", dest / "deploy", subs) + + # SDK-specific skills (renamed on copy) + copy_skill(SOURCE / f"add-tools-{sdk}", dest / "add-tools") + copy_skill(SOURCE / f"modify-{sdk}-agent", dest / "modify-agent") + + # Memory skills (all LangGraph templates - enables adding memory to any agent) + if sdk == "langgraph": + copy_skill(SOURCE / "lakebase-setup", dest / "lakebase-setup") + copy_skill(SOURCE / "agent-memory", dest / "agent-memory") + + +def main(): + """Sync skills to all templates.""" + for template, config in TEMPLATES.items(): + template_path = REPO_ROOT / template + if not template_path.exists(): + print(f"Skipping {template} (directory not found)") + continue + print(f"Syncing {template}...") + sync_template(template, config) + print("Done!") + + +if __name__ == "__main__": + main() diff --git a/.gitignore b/.gitignore index 38e33701..5e89ca89 100644 --- a/.gitignore +++ b/.gitignore @@ -168,3 +168,19 @@ yarn-error.log* /blob-report/ /playwright/* tsconfig.tsbuildinfo + +# Claude Code - track source of truth skills and sync script +.claude/* +!.claude/sync-skills.py +!.claude/skills/ +.claude/skills/* +!.claude/skills/quickstart/ +!.claude/skills/run-locally/ +!.claude/skills/discover-tools/ +!.claude/skills/deploy/ +!.claude/skills/add-tools-langgraph/ +!.claude/skills/add-tools-openai/ +!.claude/skills/modify-langgraph-agent/ +!.claude/skills/modify-openai-agent/ +!.claude/skills/lakebase-setup/ +!.claude/skills/agent-memory/ diff --git a/agent-langgraph-long-term-memory/.claude/skills/add-tools/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/add-tools/SKILL.md index 9c1c9101..7719f198 100644 --- a/agent-langgraph-long-term-memory/.claude/skills/add-tools/SKILL.md +++ b/agent-langgraph-long-term-memory/.claude/skills/add-tools/SKILL.md @@ -28,7 +28,7 @@ tools = await mcp_client.get_tools() ```yaml resources: apps: - agent_langgraph_long_term_memory: + agent_langgraph: resources: - name: 'my_genie_space' genie_space: @@ -37,7 +37,13 @@ resources: permission: 'CAN_RUN' ``` -**Step 3:** Deploy with `databricks bundle deploy` (see **deploy** skill) +**Step 3:** Deploy and run: +```bash +databricks bundle deploy +databricks bundle run agent_langgraph # Required to start app with new code! +``` + +See **deploy** skill for more details. ## Resource Type Examples @@ -51,6 +57,7 @@ See the `examples/` directory for complete YAML snippets: | `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 | Agent memory storage | | `experiment.yaml` | MLflow experiment | Tracing (already configured) | | `custom-mcp-server.md` | Custom MCP apps | Apps starting with `mcp-*` | @@ -72,17 +79,26 @@ databricks apps update-permissions \ See `examples/custom-mcp-server.md` for detailed steps. +## valueFrom Pattern (for app.yaml) + +**IMPORTANT**: Make sure all `valueFrom` references in `app.yaml` reference an existing key in the `databricks.yml` file. +Some resources need environment variables in your app. Use `valueFrom` in `app.yaml` to reference resources defined in `databricks.yml`: + +```yaml +# app.yaml +env: + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" # References resources.apps..resources[name='experiment'] + - name: LAKEBASE_INSTANCE_NAME + valueFrom: "database" # References resources.apps..resources[name='database'] +``` + +**Critical:** Every `valueFrom` value must match a `name` field in `databricks.yml` resources. + ## Important Notes - **MLflow experiment**: Already configured in template, no action needed -- **Lakebase permissions**: Memory storage requires separate Lakebase setup (see **lakebase-setup** skill) - **Multiple resources**: Add multiple entries under `resources:` list - **Permission types vary**: Each resource type has specific permission values -- **Deploy after changes**: Run `databricks bundle deploy` after modifying `databricks.yml` - -## Next Steps - -- Configure memory storage: see **lakebase-setup** skill -- Understand memory patterns: see **agent-memory** skill -- Test locally: see **run-locally** skill -- Deploy: see **deploy** skill +- **Deploy + Run after changes**: Run both `databricks bundle deploy` AND `databricks bundle run agent_langgraph` +- **valueFrom matching**: Ensure `app.yaml` `valueFrom` values match `databricks.yml` resource `name` values diff --git a/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md b/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md index 4a056c01..1324e6c5 100644 --- a/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md +++ b/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md @@ -24,7 +24,7 @@ tools = await mcp_client.get_tools() ```bash databricks bundle deploy -databricks bundle run agent_langgraph_long_term_memory +databricks bundle run agent_langgraph ``` ### 3. Get your agent app's service principal diff --git a/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/lakebase.yaml b/agent-langgraph-long-term-memory/.claude/skills/add-tools/examples/lakebase.yaml new file mode 100644 index 00000000..78f0bc72 --- /dev/null +++ b/agent-langgraph-long-term-memory/.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: valueFrom reference in app.yaml + +# In databricks.yml - add to resources.apps..resources: +- name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +# In app.yaml - add to env: +# - name: LAKEBASE_INSTANCE_NAME +# valueFrom: "database" +# - name: EMBEDDING_ENDPOINT +# value: "databricks-gte-large-en" +# - name: EMBEDDING_DIMS +# value: "1024" diff --git a/agent-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md index b8f15aa1..83d59668 100644 --- a/agent-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md +++ b/agent-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md @@ -1,220 +1,524 @@ --- name: agent-memory -description: "Understand and modify agent memory patterns. Use when: (1) User asks about 'memory', 'state', 'user preferences', (2) Working with user_id or memory tools, (3) Debugging memory issues, (4) Adding short-term memory capabilities." +description: "Add memory capabilities to your agent. Use when: (1) User asks about 'memory', 'state', 'remember', 'conversation history', (2) Want to persist conversations or user preferences, (3) Adding checkpointing or long-term storage." --- -# Agent Memory Patterns +# Adding Memory to Your Agent -This skill covers both long-term memory (facts that persist across sessions) and short-term memory (conversation history within a session). +> **Note:** This template does not include memory by default. Use this skill to **add memory capabilities**. For pre-configured memory templates, see: +> - `agent-langgraph-short-term-memory` - Conversation history within a session +> - `agent-langgraph-long-term-memory` - User facts that persist across sessions -## Long-Term Memory (This Template) +## Memory Types -Long-term memory stores facts about users that persist across conversation sessions. The agent can remember preferences, facts, and information across multiple interactions. +| Type | Use Case | Storage | Identifier | +|------|----------|---------|------------| +| **Short-term** | Conversation history within a session | `AsyncCheckpointSaver` | `thread_id` | +| **Long-term** | User facts that persist across sessions | `AsyncDatabricksStore` | `user_id` | -### Key Components +## Prerequisites -**AsyncDatabricksStore** - Persists user memories to Lakebase with semantic search: +1. **Add memory dependency** to `pyproject.toml`: + ```toml + dependencies = [ + "databricks-langchain[memory]", + ] + ``` -```python -from databricks_langchain import AsyncDatabricksStore + Then run `uv sync` + +2. **Configure Lakebase** - See **lakebase-setup** skill for: + - Creating/configuring Lakebase instance + - Initializing tables (CRITICAL first-time step) + +--- + +## Quick Setup Summary + +Adding memory requires changes to **4 files**: + +| File | What to Add | +|------|-------------| +| `pyproject.toml` | Memory dependency + hatch config | +| `.env` | Lakebase env vars (for local dev) | +| `databricks.yml` | Lakebase database resource | +| `app.yaml` | `valueFrom` reference to lakebase resource | +| `agent_server/agent.py` | Memory tools and AsyncDatabricksStore | -async with AsyncDatabricksStore( - instance_name=LAKEBASE_INSTANCE_NAME, - embedding_endpoint=EMBEDDING_ENDPOINT, - embedding_dims=EMBEDDING_DIMS, -) as store: - agent = await init_agent(store=store) - # Agent now has access to persistent memory store +--- + +## Step 1: Configure databricks.yml (Lakebase Resource) + +Add the Lakebase database resource to your app in `databricks.yml`: + +```yaml +resources: + apps: + agent_langgraph: + name: "your-app-name" + source_code_path: ./ + + resources: + # ... other resources (experiment, UC functions, etc.) ... + + # Lakebase instance for long-term memory + - name: 'lakebadatabasese_memory' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' ``` -**User ID** - Identifies a user across sessions: +**Important:** The `name: 'database'` must match the `valueFrom` reference in `app.yaml`. -```python -def _get_user_id(request: ResponsesAgentRequest) -> Optional[str]: - custom_inputs = dict(request.custom_inputs or {}) - if "user_id" in custom_inputs: - return custom_inputs["user_id"] - if request.context and getattr(request.context, "user_id", None): - return request.context.user_id - return None +--- + +## Step 2: Configure app.yaml (Environment Variables) + +Update `app.yaml` with the Lakebase environment variables: + +```yaml +command: ["uv", "run", "start-app"] + +env: + # ... other env vars ... + + # MLflow experiment (uses valueFrom to reference databricks.yml resource) + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" + + # Lakebase instance name - must match the instance_name in databricks.yml database resource + # Note: Use 'value' (not 'valueFrom') because AsyncDatabricksStore needs the instance name, + # not the full connection string that valueFrom would provide + - name: LAKEBASE_INSTANCE_NAME + value: "" + + # Embedding configuration (static values) + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" ``` -### Memory Tools +**Important:** The `LAKEBASE_INSTANCE_NAME` value must match the `instance_name` in your `databricks.yml` database resource. The `database` resource handles permissions, while `app.yaml` provides the instance name to your code. + +### Why not valueFrom for Lakebase? -The template includes three memory tools that the agent can use: +The `database` resource's `valueFrom` provides the full connection string (e.g., `instance-xxx.database.staging.cloud.databricks.com`), but `AsyncDatabricksStore` expects just the instance name. So we use a static `value` instead. + +--- -**get_user_memory** - Search for relevant information about the user: +## Step 3: Configure .env (Local Development) + +For local development, add to `.env`: + +```bash +# Lakebase configuration for long-term memory +LAKEBASE_INSTANCE_NAME= +EMBEDDING_ENDPOINT=databricks-gte-large-en +EMBEDDING_DIMS=1024 +``` + +> **Note:** `.env` is only for local development. When deployed, the app gets `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. + +--- + +## Step 4: Update agent.py + +### Add Imports and Configuration + +```python +import os +import uuid +from typing import AsyncGenerator + +from databricks_langchain import AsyncDatabricksStore, ChatDatabricks +from langchain_core.runnables import RunnableConfig +from langchain_core.tools import tool +from langgraph.prebuilt import create_react_agent + +# Environment configuration +LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") +EMBEDDING_ENDPOINT = os.getenv("EMBEDDING_ENDPOINT", "databricks-gte-large-en") +EMBEDDING_DIMS = int(os.getenv("EMBEDDING_DIMS", "1024")) # REQUIRED! +``` + +**CRITICAL:** You MUST specify `embedding_dims` when using `embedding_endpoint`. Without it, you'll get: +``` +"embedding_dims is required when embedding_endpoint is specified" +``` + +### Create Memory Tools ```python @tool async def get_user_memory(query: str, config: RunnableConfig) -> str: - """Search for relevant information about the user from long-term memory.""" + """Search for relevant information about the user from long-term memory. + Use this to recall preferences, past interactions, or other saved information. + """ user_id = config.get("configurable", {}).get("user_id") store = config.get("configurable", {}).get("store") + if not user_id or not store: + return "Memory not available - no user context" - namespace = ("user_memories", user_id.replace(".", "-")) + # Sanitize user_id for namespace (replace special chars) + namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) results = await store.asearch(namespace, query=query, limit=5) - # Returns formatted memory items -``` + if not results: + return "No memories found for this query" + return "\n".join([f"- {r.value}" for r in results]) -**save_user_memory** - Save information about the user: -```python @tool -async def save_user_memory(memory_key: str, memory_data_json: str, config: RunnableConfig) -> str: - """Save information about the user to long-term memory.""" - # memory_data_json must be a valid JSON object string - await store.aput(namespace, memory_key, memory_data) -``` +async def save_user_memory(memory_key: str, memory_data: str, config: RunnableConfig) -> str: + """Save information about the user to long-term memory. + Use this to remember user preferences, important details, or other + information that should persist across conversations. + + Args: + memory_key: A short descriptive key (e.g., "preferred_name", "team", "region") + memory_data: The information to save + """ + user_id = config.get("configurable", {}).get("user_id") + store = config.get("configurable", {}).get("store") + if not user_id or not store: + return "Memory not available - no user context" + + namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) + await store.aput(namespace, memory_key, {"value": memory_data}) + return f"Saved memory: {memory_key}" -**delete_user_memory** - Remove specific memories: -```python @tool async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: - """Delete a specific memory from the user's long-term memory.""" + """Delete a specific memory from the user's long-term memory. + Use this when the user asks to forget something or correct stored information. + + Args: + memory_key: The key of the memory to delete (e.g., "preferred_name", "team") + """ + user_id = config.get("configurable", {}).get("user_id") + store = config.get("configurable", {}).get("store") + if not user_id or not store: + return "Memory not available - no user context" + + namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) await store.adelete(namespace, memory_key) + return f"Deleted memory: {memory_key}" ``` -### Creating a Long-Term Memory Agent - -Pass `store` to `create_agent()`: +### Add Helper Functions ```python -from langchain.agents import create_agent +def _get_user_id(request: ResponsesAgentRequest) -> str: + """Extract user_id from request context or custom inputs.""" + custom_inputs = dict(request.custom_inputs or {}) + if "user_id" in custom_inputs and custom_inputs["user_id"]: + return str(custom_inputs["user_id"]) + if request.context and getattr(request.context, "user_id", None): + return str(request.context.user_id) + return "default-user" -agent = create_agent( - model=ChatDatabricks(endpoint=LLM_ENDPOINT_NAME), - tools=mcp_tools + init_memory_tools(), # Include memory tools - system_prompt=SYSTEM_PROMPT, - store=store, # Enables long-term memory -) -``` -### System Prompt for Memory +def _get_or_create_thread_id(request: ResponsesAgentRequest) -> str: + """Extract or create thread_id for conversation tracking.""" + custom_inputs = dict(request.custom_inputs or {}) + if "thread_id" in custom_inputs and custom_inputs["thread_id"]: + return str(custom_inputs["thread_id"]) + if request.context and getattr(request.context, "conversation_id", None): + return str(request.context.conversation_id) + return str(uuid.uuid4()) +``` -The template includes instructions for using memory tools: +### Update Streaming Function ```python -SYSTEM_PROMPT = """You are a helpful assistant. - -You have access to memory tools that allow you to remember information about users: -- Use get_user_memory to search for previously saved information about the user -- Use save_user_memory to remember important facts, preferences, or details the user shares -- Use delete_user_memory to forget specific information when asked - -Always check for relevant memories at the start of a conversation to provide personalized responses.""" +@stream() +async def streaming( + request: ResponsesAgentRequest, +) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + # Get user context + user_id = _get_user_id(request) + thread_id = _get_or_create_thread_id(request) + + # Get other tools (MCP, etc.) + mcp_client = init_mcp_client(sp_workspace_client) + mcp_tools = await mcp_client.get_tools() + + # Memory tools + memory_tools = [get_user_memory, save_user_memory, delete_user_memory] + + # Combine all tools + all_tools = mcp_tools + memory_tools + + # Initialize model + model = ChatDatabricks(endpoint="databricks-claude-sonnet-4") + + # Use AsyncDatabricksStore for long-term memory + async with AsyncDatabricksStore( + instance_name=LAKEBASE_INSTANCE_NAME, + embedding_endpoint=EMBEDDING_ENDPOINT, + embedding_dims=EMBEDDING_DIMS, # REQUIRED! + ) as store: + # Create agent with tools + agent = create_react_agent( + model=model, + tools=all_tools, + prompt=AGENT_INSTRUCTIONS, + ) + + # Prepare input + messages = {"messages": to_chat_completions_input([i.model_dump() for i in request.input])} + + # Configure with user context for memory tools + config = { + "configurable": { + "user_id": user_id, + "thread_id": thread_id, + "store": store, # Pass store to tools via config + } + } + + # Stream agent responses + async for event in process_agent_astream_events( + agent.astream(input=messages, config=config, stream_mode=["updates", "messages"]) + ): + yield event ``` --- -## API Requests with Long-Term Memory +## Step 5: Initialize Tables and Deploy -### Request with user_id +### Initialize Lakebase Tables (First Time Only) + +Before deploying, initialize the tables locally: ```bash -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ - "input": [{"role": "user", "content": "My favorite color is blue, remember that"}], - "custom_inputs": {"user_id": "user@example.com"} - }' +uv run python -c "$(cat <<'EOF' +import asyncio +from databricks_langchain import AsyncDatabricksStore + +async def setup(): + async with AsyncDatabricksStore( + instance_name="", + embedding_endpoint="databricks-gte-large-en", + embedding_dims=1024, + ) as store: + await store.setup() + print("Tables created!") + +asyncio.run(setup()) +EOF +)" ``` -### Query stored memories +### Deploy and Run + +**IMPORTANT:** Always run `databricks bundle run` after `databricks bundle deploy` to start/restart the app with the new code: ```bash -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ - "input": [{"role": "user", "content": "What do you remember about me?"}], - "custom_inputs": {"user_id": "user@example.com"} - }' +# Deploy resources and upload files +databricks bundle deploy + +# Start/restart the app with new code (REQUIRED!) +databricks bundle run agent_langgraph ``` -### Using context instead of custom_inputs +> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is required to actually start the app with the new code. -```bash -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ - "input": [{"role": "user", "content": "What is my favorite color?"}], - "context": {"user_id": "user@example.com"} - }' -``` +--- -### Deployed App with Memory +## Complete Example Files + +### databricks.yml (with Lakebase) + +```yaml +bundle: + name: agent_langgraph + +resources: + experiments: + agent_langgraph_experiment: + name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} + + apps: + agent_langgraph: + name: "my-agent-app" + description: "Agent with long-term memory" + source_code_path: ./ + + resources: + - name: 'experiment' + experiment: + experiment_id: "${resources.experiments.agent_langgraph_experiment.id}" + permission: 'CAN_MANAGE' + + # Lakebase instance for long-term memory + - name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +targets: + dev: + mode: development + default: true +``` -```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"} - }' +### app.yaml + +```yaml +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" + # Reference experiment resource from databricks.yml + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" + # Lakebase instance name (must match instance_name in databricks.yml) + - name: LAKEBASE_INSTANCE_NAME + value: "" + # Embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" ``` --- -## Short-Term Memory (Reference) +## Adding Short-Term Memory -Short-term memory stores conversation history within a thread. This is available in the `agent-langgraph-short-term-memory` template. +Short-term memory stores conversation history within a thread. -### Key Components (Short-Term) - -**AsyncCheckpointSaver** - Persists LangGraph state to Lakebase: +### Import and Initialize Checkpointer ```python from databricks_langchain import AsyncCheckpointSaver +LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") + async with AsyncCheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME) as checkpointer: - agent = await init_agent(checkpointer=checkpointer) + agent = create_react_agent( + model=model, + tools=tools, + checkpointer=checkpointer, # Enables conversation persistence + ) ``` -**Thread ID** - Identifies a conversation thread: +### Use thread_id in Agent Config ```python -thread_id = request.custom_inputs.get("thread_id") or str(uuid_utils.uuid7()) config = {"configurable": {"thread_id": thread_id}} +async for event in agent.astream(input_state, config, stream_mode=["updates", "messages"]): + yield event ``` -### Short-Term Memory Request Example +--- + +## Testing Memory + +### Test Locally ```bash +# Start the server +uv run start-app + +# Save a memory +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Remember that I am on the shipping team"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' + +# Recall the memory curl -X POST http://localhost:8000/invocations \ -H "Content-Type: application/json" \ -d '{ - "input": [{"role": "user", "content": "What did we discuss earlier?"}], - "custom_inputs": {"thread_id": ""} + "input": [{"role": "user", "content": "What team am I on?"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' + +# Delete a memory +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Forget what team I am on"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' +``` + +### Test Deployed App + +```bash +# Get OAuth token (PATs don't work for apps) +TOKEN=$(databricks auth token --host | jq -r '.access_token') + +# Test memory save +curl -X POST https:///invocations \ + -H "Authorization: Bearer $TOKEN" \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Remember I prefer detailed explanations"}], + "custom_inputs": {"user_id": "alice@example.com"} }' ``` --- +## First-Time Setup Checklist + +- [ ] Added `databricks-langchain[memory]` to `pyproject.toml` +- [ ] Run `uv sync` to install dependencies +- [ ] Created or identified Lakebase instance +- [ ] Added Lakebase env vars to `.env` (for local dev) +- [ ] Added `database` resource to `databricks.yml` (for permissions) +- [ ] Added `LAKEBASE_INSTANCE_NAME` value to `app.yaml` (matching instance_name in databricks.yml) +- [ ] **Initialized tables locally** by running `await store.setup()` +- [ ] Deployed with `databricks bundle deploy` +- [ ] **Started app with `databricks bundle run agent_langgraph`** + +--- + ## Troubleshooting -| Issue | Solution | -|-------|----------| -| **Memories not persisting** | Verify `user_id` is passed consistently | -| **"Cannot connect to Lakebase"** | See **lakebase-setup** skill | -| **Permission errors** | Grant store table permissions (see **lakebase-setup** skill) | -| **Memory not available warning** | Ensure `user_id` is provided in request | -| **No memories found** | User hasn't saved any memories yet | +| Issue | Cause | Solution | +|-------|-------|----------| +| **"embedding_dims is required"** | Missing parameter | Add `embedding_dims=1024` to AsyncDatabricksStore | +| **"relation 'store' does not exist"** | Tables not created | Run `await store.setup()` locally first | +| **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `valueFrom: "database"` to app.yaml | +| **"permission denied for table store"** | Missing grants | Lakebase `database` resource in DAB should handle this | +| **"Memory not available"** | No user_id in request | Ensure `custom_inputs.user_id` is passed | +| **Memory not persisting** | Different user_ids | Use consistent user_id across requests | +| **App not updated after deploy** | Forgot to run bundle | Run `databricks bundle run agent_langgraph` after deploy | -## Dependencies +--- -Long-term memory requires `databricks-langchain[memory]`: +## Pre-Built Memory Templates -```toml -# pyproject.toml -dependencies = [ - "databricks-langchain[memory]>=0.13.0", -] -``` +For fully configured implementations without manual setup: + +| Template | Memory Type | Key Features | +|----------|-------------|--------------| +| `agent-langgraph-short-term-memory` | Short-term | AsyncCheckpointSaver, thread_id | +| `agent-langgraph-long-term-memory` | Long-term | AsyncDatabricksStore, memory tools | + +--- ## Next Steps - Configure Lakebase: see **lakebase-setup** skill -- Modify agent behavior: see **modify-agent** skill - Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/agent-langgraph-long-term-memory/.claude/skills/deploy/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/deploy/SKILL.md index 118f2a59..3a0f3090 100644 --- a/agent-langgraph-long-term-memory/.claude/skills/deploy/SKILL.md +++ b/agent-langgraph-long-term-memory/.claude/skills/deploy/SKILL.md @@ -5,18 +5,38 @@ description: "Deploy agent to Databricks Apps using DAB (Databricks Asset Bundle # Deploy to Databricks Apps -> **Memory Template:** Before deploying, ensure Lakebase is configured. See the **lakebase-setup** skill for permissions setup required after deployment. +## 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_langgraph_long_term_memory: + name: "agent-your-app-name" # Use agent-* prefix +``` ## Deploy Commands +**IMPORTANT:** Always run BOTH commands to deploy and start your app: + ```bash -# Deploy the bundle (creates/updates resources, uploads files) +# 1. Validate bundle configuration (catches errors before deploy) +databricks bundle validate + +# 2. Deploy the bundle (creates/updates resources, uploads files) databricks bundle deploy -# Run the app (starts/restarts with uploaded source code) +# 3. Run the app (starts/restarts with uploaded source code) - REQUIRED! databricks bundle run agent_langgraph_long_term_memory ``` +> **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_langgraph_long_term_memory` matches the app name in `databricks.yml` under `resources.apps`. ## Handling "App Already Exists" Error @@ -86,15 +106,6 @@ databricks bundle deploy **Warning:** This permanently deletes the app's URL, OAuth credentials, and service principal. -## Post-Deployment: Lakebase Permissions - -**After deploying a memory-enabled agent, you must grant Lakebase permissions.** - -See the **lakebase-setup** skill for: -- Adding Lakebase as an app resource -- SDK or SQL commands to grant store table permissions -- Troubleshooting access errors - ## Unbinding an App To remove the link between bundle and deployed app: @@ -127,7 +138,7 @@ curl -X POST /invocations \ -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' ``` -**Request with user_id (for long-term memory):** +**If using memory** - include `user_id` to scope memories per user: ```bash curl -X POST /invocations \ -H "Authorization: Bearer " \ @@ -138,8 +149,6 @@ curl -X POST /invocations \ }' ``` -For more memory-specific request examples, see the **agent-memory** skill. - ## On-Behalf-Of (OBO) User Authentication To authenticate as the requesting user instead of the app service principal: @@ -201,16 +210,13 @@ uv add | 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 | +| 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`| +| "Provider produced inconsistent result" | Sync app config to `databricks.yml` | | "should set workspace.root_path" | Add `root_path` to production target | - -## Next Steps - -- Grant Lakebase permissions: see **lakebase-setup** skill -- Understand memory patterns: see **agent-memory** skill -- Add tools and permissions: see **add-tools** skill +| App running old code after deploy | Run `databricks bundle run agent_langgraph_long_term_memory` after deploy | +| Env var is None in deployed app | Check `valueFrom` in app.yaml matches resource `name` in databricks.yml | diff --git a/agent-langgraph-long-term-memory/.claude/skills/discover-tools/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/discover-tools/SKILL.md index eabac2a5..87c3f519 100644 --- a/agent-langgraph-long-term-memory/.claude/skills/discover-tools/SKILL.md +++ b/agent-langgraph-long-term-memory/.claude/skills/discover-tools/SKILL.md @@ -39,44 +39,9 @@ uv run discover-tools --profile DEFAULT | **Custom MCP Servers** | Apps starting with `mcp-*` | `{app_url}/mcp` | | **External MCP Servers** | Via UC connections | `{host}/api/2.0/mcp/external/{connection_name}` | -## Using Discovered Tools in Code - -After discovering tools, add them to your agent in `agent_server/agent.py`: - -```python -from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient - -# Example: Add UC functions from a schema -uc_functions_server = DatabricksMCPServer( - url=f"{host}/api/2.0/mcp/functions/{catalog}/{schema}", - name="my uc functions", -) - -# Example: Add a Genie space -genie_server = DatabricksMCPServer( - url=f"{host}/api/2.0/mcp/genie/{space_id}", - name="my genie space", -) - -# Example: Add vector search -vector_server = DatabricksMCPServer( - url=f"{host}/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name}", - name="my vector index", -) - -# Create MCP client with all servers -mcp_client = DatabricksMultiServerMCPClient([ - uc_functions_server, - genie_server, - vector_server, -]) - -# Get tools for the agent -tools = await mcp_client.get_tools() -``` - ## Next Steps -After adding MCP servers to your agent: -1. **Grant permissions** in `databricks.yml` (see **add-tools** skill) -2. Test locally with `uv run start-app` (see **run-locally** skill) +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 diff --git a/agent-langgraph-long-term-memory/.claude/skills/lakebase-setup/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/lakebase-setup/SKILL.md index 036b1a1b..6b8c802e 100644 --- a/agent-langgraph-long-term-memory/.claude/skills/lakebase-setup/SKILL.md +++ b/agent-langgraph-long-term-memory/.claude/skills/lakebase-setup/SKILL.md @@ -1,169 +1,351 @@ --- name: lakebase-setup -description: "Configure Lakebase for agent memory storage. Use when: (1) First-time memory setup, (2) 'Failed to connect to Lakebase' errors, (3) Permission errors on store tables, (4) User says 'lakebase', 'memory setup', or 'store'." +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 Memory +# Lakebase Setup for Agent Memory -This template uses Lakebase (Databricks-managed PostgreSQL) to store long-term user memories. You must configure Lakebase before the agent can persist memories. +> **Note:** This template does not include memory by default. Use this skill if you want to **add memory capabilities** to your agent. For pre-configured memory templates, see: +> - `agent-langgraph-short-term-memory` - Conversation history within a session +> - `agent-langgraph-long-term-memory` - User facts that persist across sessions -## Prerequisites +## Overview -- A Lakebase instance in your Databricks workspace -- The instance name (found in Databricks UI under SQL Warehouses > Lakebase) +Lakebase provides persistent storage for agent memory: +- **Short-term memory**: Conversation history within a thread (`AsyncCheckpointSaver`) +- **Long-term memory**: User facts across sessions (`AsyncDatabricksStore`) -## Local Development Setup +## Complete Setup Workflow -**Step 1:** Add the Lakebase instance name to `.env`: +``` +┌─────────────────────────────────────────────────────────────────────────────┐ +│ 1. Add dependency → 2. Get instance → 3. Configure DAB + app.yaml │ +│ 4. Configure .env → 5. Initialize tables → 6. Deploy + Run │ +└─────────────────────────────────────────────────────────────────────────────┘ +``` + +--- + +## Step 1: Add Memory Dependency +Add the memory extra to your `pyproject.toml`: + +```toml +dependencies = [ + "databricks-langchain[memory]", + # ... other dependencies +] +``` + +Then sync dependencies: ```bash -LAKEBASE_INSTANCE_NAME= +uv sync ``` -**Step 2:** Run `uv run start-app` to test locally. The agent will automatically create store tables on first run. +--- -## Deployed App Setup +## Step 2: Create or Get Lakebase Instance -After deploying your agent with `databricks bundle deploy`, you must grant the app's service principal access to Lakebase. +### Option A: Create New Instance (via Databricks UI) -### Step 1: Add Lakebase as App Resource +1. Go to your Databricks workspace +2. Navigate to **Compute** → **Lakebase** +3. Click **Create Instance** +4. Note the instance name -1. Go to the Databricks UI -2. Navigate to your app and click **Edit** -3. Go to **App resources** → **Add resource** -4. Add your Lakebase instance with **Connect + Create** permissions +### Option B: Use Existing Instance -### Step 2: Get App Service Principal ID +If you have an existing instance, note its name for the next step. -```bash -databricks apps get --output json | jq -r '.service_principal_id' +--- + +## Step 3: Configure databricks.yml (Lakebase Resource) + +Add the Lakebase `database` resource to your app in `databricks.yml`: + +```yaml +resources: + apps: + agent_langgraph: + 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: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' ``` -### Step 3: Grant Permissions +**Important:** +- The `name: 'database'` must match the `valueFrom` reference in `app.yaml` +- Using the `database` resource type automatically grants the app's service principal access to Lakebase -Choose **Option A** (SDK - Recommended) or **Option B** (SQL). +### Update app.yaml (Environment Variables) -#### Option A: Using LakebaseClient SDK (Recommended) +Update `app.yaml` with the Lakebase instance name: -The `databricks-ai-bridge` package includes a `LakebaseClient` for programmatic permission management: +```yaml +env: + # ... other env vars ... -```python -from databricks_ai_bridge.lakebase import LakebaseClient, SchemaPrivilege, TablePrivilege + # Lakebase instance name - must match instance_name in databricks.yml database resource + # Note: Use 'value' (not 'valueFrom') because AsyncDatabricksStore needs the instance name, + # not the full connection string that valueFrom would provide + - name: LAKEBASE_INSTANCE_NAME + value: "" -# Initialize client -client = LakebaseClient(instance_name="") + # Static values for embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` -app_sp = "" # From Step 2 +**Important:** +- The `LAKEBASE_INSTANCE_NAME` value must match the `instance_name` in your `databricks.yml` database resource +- The `database` resource handles permissions; `app.yaml` provides the instance name to your code +- Don't use `valueFrom` for Lakebase - it provides the connection string, not the instance name -# Create role for the service principal -client.create_role(app_sp, "SERVICE_PRINCIPAL") +--- -# Grant schema privileges -client.grant_schema( - grantee=app_sp, - privileges=[SchemaPrivilege.USAGE, SchemaPrivilege.CREATE], - schemas=["drizzle", "ai_chatbot", "public"], -) +## Step 4: Configure .env (Local Development) -# Grant table privileges on all tables in schemas -client.grant_all_tables_in_schema( - grantee=app_sp, - privileges=[TablePrivilege.SELECT, TablePrivilege.INSERT, TablePrivilege.UPDATE], - schemas=["drizzle", "ai_chatbot"], -) +For local development, add to `.env`: -# Grant privileges on specific checkpoint tables -client.grant_table( - grantee=app_sp, - privileges=[TablePrivilege.SELECT, TablePrivilege.INSERT, TablePrivilege.UPDATE], - tables=[ - "public.checkpoint_migrations", - "public.checkpoint_writes", - "public.checkpoints", - "public.checkpoint_blobs", - ], -) +```bash +# Lakebase configuration for long-term memory +LAKEBASE_INSTANCE_NAME= +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 `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. + +--- + +## Step 5: Initialize Store Tables (CRITICAL - First Time Only) + +**Before deploying**, you must initialize the Lakebase tables. The `AsyncDatabricksStore` creates tables on first use, but you need to do this locally first: + +```python +# Run this script locally BEFORE first deployment +import asyncio +from databricks_langchain import AsyncDatabricksStore + +async def setup_store(): + async with AsyncDatabricksStore( + instance_name="", + embedding_endpoint="databricks-gte-large-en", + embedding_dims=1024, + ) as store: + print("Setting up store tables...") + await store.setup() # Creates required tables + print("Store tables created!") + + # Verify with a test write/read + await store.aput(("test", "init"), "test_key", {"value": "test_value"}) + results = await store.asearch(("test", "init"), query="test", limit=1) + print(f"Test successful: {results}") + +asyncio.run(setup_store()) ``` -**Benefits of SDK approach:** -- Type-safe privilege enums prevent typos -- Cleaner Python code vs raw SQL -- Easier to integrate into setup scripts - -#### Option B: Using SQL - -Run the following SQL on your Lakebase instance (replace `app-sp-id` with your app's service principal ID): - -```sql -DO $$ -DECLARE - app_sp text := 'app-sp-id'; -- TODO: Replace with your App's Service Principal ID -BEGIN - ------------------------------------------------------------------- - -- Drizzle schema: migration metadata tables - ------------------------------------------------------------------- - EXECUTE format('GRANT USAGE, CREATE ON SCHEMA drizzle TO %I;', app_sp); - EXECUTE format('GRANT SELECT, INSERT, UPDATE ON ALL TABLES IN SCHEMA drizzle TO %I;', app_sp); - ------------------------------------------------------------------- - -- App schema: business tables (Chat, Message, etc.) - ------------------------------------------------------------------- - EXECUTE format('GRANT USAGE, CREATE ON SCHEMA ai_chatbot TO %I;', app_sp); - EXECUTE format('GRANT SELECT, INSERT, UPDATE ON ALL TABLES IN SCHEMA ai_chatbot TO %I;', app_sp); - ------------------------------------------------------------------- - -- Public schema for checkpoint tables - ------------------------------------------------------------------- - EXECUTE format('GRANT USAGE, CREATE ON SCHEMA public TO %I;', app_sp); - EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.checkpoint_migrations TO %I;', app_sp); - EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.checkpoint_writes TO %I;', app_sp); - EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.checkpoints TO %I;', app_sp); - EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.checkpoint_blobs TO %I;', app_sp); -END $$; +Run with: +```bash +uv run python -c "$(cat <<'EOF' +import asyncio +from databricks_langchain import AsyncDatabricksStore + +async def setup(): + async with AsyncDatabricksStore( + instance_name="", + embedding_endpoint="databricks-gte-large-en", + embedding_dims=1024, + ) as store: + await store.setup() + print("Tables created!") + +asyncio.run(setup()) +EOF +)" ``` -**Schema Reference:** -| Schema | Purpose | -|--------|---------| -| `drizzle` | Migration metadata tables | -| `ai_chatbot` | Business tables (Chat, Message, etc.) | -| `public` | Store tables for user memories | +This creates these tables in the `public` schema: +- `store` - Key-value storage for memories +- `store_vectors` - Vector embeddings for semantic search +- `store_migrations` - Schema migration tracking +- `vector_migrations` - Vector schema migration tracking -## Troubleshooting +--- -| Issue | Solution | -|-------|----------| -| **"LAKEBASE_INSTANCE_NAME environment variable is required"** | Add `LAKEBASE_INSTANCE_NAME=` to `.env` | -| **"Failed to connect to Lakebase instance"** | Verify instance name is correct and your profile has access | -| **Permission errors on store tables** | Run the SDK script or SQL grant commands above | -| **Deployed app can't access Lakebase** | Add Lakebase as app resource in Databricks UI | -| **"role does not exist"** | Run `client.create_role()` first, or ensure SP ID is correct | +## Step 6: Deploy and Run Your App -### Common Error Messages +**IMPORTANT:** Always run both `deploy` AND `run` commands: + +```bash +# Deploy resources and upload files +databricks bundle deploy -**Local development:** +# Start/restart the app with new code (REQUIRED!) +databricks bundle run agent_langgraph ``` -Failed to connect to Lakebase instance ''. Please verify: -1. The instance name is correct -2. You have the necessary permissions to access the instance -3. Your Databricks authentication is configured correctly + +> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is required to actually start the app with the new code. + +--- + +## Complete Example: databricks.yml with Lakebase + +```yaml +bundle: + name: agent_langgraph + +resources: + experiments: + agent_langgraph_experiment: + name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} + + apps: + agent_langgraph: + name: "my-agent-app" + description: "Agent with long-term memory" + source_code_path: ./ + + resources: + - name: 'experiment' + experiment: + experiment_id: "${resources.experiments.agent_langgraph_experiment.id}" + permission: 'CAN_MANAGE' + + # Lakebase instance for long-term memory + - name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +targets: + dev: + mode: development + default: true ``` -**Deployed app:** +## Complete Example: app.yaml + +```yaml +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" + # Reference experiment resource from databricks.yml + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" + # Lakebase instance name (must match instance_name in databricks.yml) + - name: LAKEBASE_INSTANCE_NAME + value: "" + # Embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" ``` -Failed to connect to Lakebase instance ''. The App Service Principal for '' may not have access. + +--- + +## 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 | Run `await store.setup()` locally first (Step 5) | +| **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `LAKEBASE_INSTANCE_NAME` value to app.yaml | +| **"Unable to resolve Lakebase instance '...database.cloud.databricks.com'"** | Used valueFrom instead of value | Use `value: ""` not `valueFrom` for Lakebase | +| **"permission denied for table store"** | Missing grants | The `database` resource in DAB should handle this; verify the resource is configured | +| **"Failed to connect to Lakebase"** | Wrong instance name | Verify instance name 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 agent_langgraph` after deploy | +| **valueFrom not resolving** | Resource name mismatch | Ensure `valueFrom` value matches `name` in databricks.yml resources | + +--- + +## Quick Reference: LakebaseClient API + +For manual permission management (usually not needed with DAB `database` resource): + +```python +from databricks_ai_bridge.lakebase import LakebaseClient, SchemaPrivilege, TablePrivilege + +client = LakebaseClient(instance_name="...") + +# 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'") ``` -Both errors indicate Lakebase access issues. Follow the setup steps above for your environment. +### Service Principal Identifiers -## How Memory Works +When granting permissions manually, note that Databricks apps have multiple identifiers: -See the **agent-memory** skill for: -- How `AsyncDatabricksStore` persists user memories -- Using `user_id` to scope memories per user -- Memory tools (get, save, delete) -- API request examples with user_id +| 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 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 -- Understand memory patterns: see **agent-memory** skill +- Add memory to agent code: see **agent-memory** skill - Test locally: see **run-locally** skill - Deploy: see **deploy** skill diff --git a/agent-langgraph-long-term-memory/.claude/skills/modify-agent/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/modify-agent/SKILL.md index 4932f002..d7218637 100644 --- a/agent-langgraph-long-term-memory/.claude/skills/modify-agent/SKILL.md +++ b/agent-langgraph-long-term-memory/.claude/skills/modify-agent/SKILL.md @@ -5,8 +5,6 @@ description: "Modify agent code, add tools, or change configuration. Use when: ( # Modify the Agent -> **Memory Template:** This template uses long-term memory (facts that persist across sessions). See the **agent-memory** skill for memory-specific patterns. - ## Main File **`agent_server/agent.py`** - Agent logic, model selection, instructions, MCP servers @@ -290,7 +288,6 @@ Example: Create UC function wrapping HTTP request for Slack, then expose via MCP - Discover available tools: see **discover-tools** skill - Grant resource permissions: see **add-tools** skill -- Memory patterns: see **agent-memory** skill -- Lakebase setup: see **lakebase-setup** skill +- Add memory capabilities: see **agent-memory** skill - Test locally: see **run-locally** skill - Deploy: see **deploy** skill diff --git a/agent-langgraph-long-term-memory/.claude/skills/quickstart/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/quickstart/SKILL.md index 853b78d0..e550162c 100644 --- a/agent-langgraph-long-term-memory/.claude/skills/quickstart/SKILL.md +++ b/agent-langgraph-long-term-memory/.claude/skills/quickstart/SKILL.md @@ -79,6 +79,5 @@ CHAT_PROXY_TIMEOUT_SECONDS=300 ## Next Steps After quickstart completes: -1. **Configure Lakebase** for memory storage (see **lakebase-setup** skill) -2. Run `uv run discover-tools` to find available workspace resources (see **discover-tools** skill) -3. Run `uv run start-app` to test locally (see **run-locally** skill) +1. Run `uv run discover-tools` to find available workspace resources (see **discover-tools** skill) +2. Run `uv run start-app` to test locally (see **run-locally** skill) diff --git a/agent-langgraph-long-term-memory/.claude/skills/run-locally/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/run-locally/SKILL.md index 398d22a3..3eb83c82 100644 --- a/agent-langgraph-long-term-memory/.claude/skills/run-locally/SKILL.md +++ b/agent-langgraph-long-term-memory/.claude/skills/run-locally/SKILL.md @@ -11,20 +11,25 @@ description: "Run and test the agent locally. Use when: (1) User says 'run local uv run start-app ``` -This starts the agent at http://localhost:8000 with both the API server and chat UI. +This starts the agent at http://localhost:8000 ## Server Options -| Option | Command | When to Use | -|--------|---------|-------------| -| **Hot-reload** | `uv run start-server --reload` | During development - auto-restarts on code changes | -| **Custom port** | `uv run start-server --port 8001` | When port 8000 is in use | -| **Multiple workers** | `uv run start-server --workers 4` | Load testing or production-like simulation | -| **Combined** | `uv run start-server --reload --port 8001` | Development on alternate port | +```bash +# Hot-reload on code changes (development) +uv run start-server --reload -## Test the API +# Custom port +uv run start-server --port 8001 + +# Multiple workers (production-like) +uv run start-server --workers 4 -The agent implements the [OpenAI Responses API](https://platform.openai.com/docs/api-reference/responses) interface via MLflow's ResponsesAgent. +# Combine options +uv run start-server --reload --port 8001 +``` + +## Test the API **Streaming request:** ```bash @@ -40,38 +45,13 @@ curl -X POST http://localhost:8000/invocations \ -d '{ "input": [{ "role": "user", "content": "hi" }] }' ``` -**Request with user_id (for long-term memory):** -```bash -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ - "input": [{"role": "user", "content": "What do you remember about me?"}], - "custom_inputs": {"user_id": "user@example.com"} - }' -``` - -See the **agent-memory** skill for more memory-related request examples. - ## Run Evaluation -Evaluate your agent using MLflow scorers: - ```bash uv run agent-evaluate ``` -**What it does:** -- Runs the agent against a test dataset -- Applies MLflow scorers (RelevanceToQuery, Safety) -- Records results to your MLflow experiment - -**Customize evaluation:** -Edit `agent_server/evaluate_agent.py` to: -- Change the evaluation dataset -- Add or modify scorers -- Adjust evaluation parameters - -After evaluation completes, open the MLflow UI link for your experiment to inspect results. +Uses MLflow scorers (RelevanceToQuery, Safety). ## Run Unit Tests @@ -79,15 +59,6 @@ After evaluation completes, open the MLflow UI link for your experiment to inspe pytest [path] ``` -## Adding Dependencies - -```bash -uv add -# Example: uv add "langchain-community" -``` - -Dependencies are managed in `pyproject.toml`. - ## Troubleshooting | Issue | Solution | @@ -96,7 +67,6 @@ Dependencies are managed in `pyproject.toml`. | **Authentication errors** | Verify `.env` is correct; run **quickstart** skill | | **Module not found** | Run `uv sync` to install dependencies | | **MLflow experiment not found** | Ensure `MLFLOW_TRACKING_URI` in `.env` is `databricks://` | -| **Lakebase connection errors** | Verify `LAKEBASE_INSTANCE_NAME` in `.env`; see **lakebase-setup** skill | ### MLflow Experiment Not Found @@ -114,17 +84,7 @@ MLFLOW_TRACKING_URI="databricks://DEFAULT" # Include profile name The quickstart script configures this automatically. If you manually edited `.env`, ensure the profile name is included. -## MLflow Tracing - -This template uses MLflow for automatic tracing: -- Agent logic decorated with `@invoke()` and `@stream()` is automatically traced -- LLM invocations are captured via MLflow autologging - -To add custom trace instrumentation, see: [MLflow tracing documentation](https://docs.databricks.com/aws/en/mlflow3/genai/tracing/app-instrumentation/) - ## Next Steps -- Configure Lakebase: see **lakebase-setup** skill -- Understand memory patterns: see **agent-memory** skill - Modify your agent: see **modify-agent** skill - Deploy to Databricks: see **deploy** skill diff --git a/agent-langgraph-long-term-memory/.gitignore b/agent-langgraph-long-term-memory/.gitignore index 0c3cb57a..14c5b636 100644 --- a/agent-langgraph-long-term-memory/.gitignore +++ b/agent-langgraph-long-term-memory/.gitignore @@ -1,8 +1,6 @@ # Created by https://www.toptal.com/developers/gitignore/api/python # Edit at https://www.toptal.com/developers/gitignore?templates=python -databricks.yml - ### Python ### # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/agent-langgraph-long-term-memory/AGENTS.md b/agent-langgraph-long-term-memory/AGENTS.md index 41046b58..d8f25998 100644 --- a/agent-langgraph-long-term-memory/AGENTS.md +++ b/agent-langgraph-long-term-memory/AGENTS.md @@ -1,15 +1,25 @@ # Agent Development Guide -## MANDATORY First Action +## MANDATORY First Actions -**BEFORE any other action, run `databricks auth profiles` to check authentication status.** +**Ask the user interactively:** + +1. **App deployment target:** + > "Do you have an existing Databricks app you want to deploy to, or should we create a new one? If existing, what's the app name?" + + *Note: New apps should use the `agent-*` prefix (e.g., `agent-data-analyst`) unless the user specifies otherwise.* + +2. **Lakebase instance (required for memory):** + > "This template requires Lakebase for memory. Do you have an existing Lakebase instance? If so, what's the instance name?" + +**Then check authentication status by running `databricks auth profiles`.** This helps you understand: - Which Databricks profiles are configured - Whether authentication is already set up - Which profile to use for subsequent commands -If no profiles exist, guide the user through running `uv run quickstart` to set up authentication. +If no profiles exist or `.env` is missing, guide the user through running `uv run quickstart` to set up authentication and configuration. See the **quickstart** skill for details. ## Understanding User Goals diff --git a/agent-langgraph-short-term-memory/.claude/skills/add-tools/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/add-tools/SKILL.md index 45f83a8f..7719f198 100644 --- a/agent-langgraph-short-term-memory/.claude/skills/add-tools/SKILL.md +++ b/agent-langgraph-short-term-memory/.claude/skills/add-tools/SKILL.md @@ -28,7 +28,7 @@ tools = await mcp_client.get_tools() ```yaml resources: apps: - agent_langgraph_short_term_memory: + agent_langgraph: resources: - name: 'my_genie_space' genie_space: @@ -37,7 +37,13 @@ resources: permission: 'CAN_RUN' ``` -**Step 3:** Deploy with `databricks bundle deploy` (see **deploy** skill) +**Step 3:** Deploy and run: +```bash +databricks bundle deploy +databricks bundle run agent_langgraph # Required to start app with new code! +``` + +See **deploy** skill for more details. ## Resource Type Examples @@ -51,6 +57,7 @@ See the `examples/` directory for complete YAML snippets: | `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 | Agent memory storage | | `experiment.yaml` | MLflow experiment | Tracing (already configured) | | `custom-mcp-server.md` | Custom MCP apps | Apps starting with `mcp-*` | @@ -72,17 +79,26 @@ databricks apps update-permissions \ See `examples/custom-mcp-server.md` for detailed steps. +## valueFrom Pattern (for app.yaml) + +**IMPORTANT**: Make sure all `valueFrom` references in `app.yaml` reference an existing key in the `databricks.yml` file. +Some resources need environment variables in your app. Use `valueFrom` in `app.yaml` to reference resources defined in `databricks.yml`: + +```yaml +# app.yaml +env: + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" # References resources.apps..resources[name='experiment'] + - name: LAKEBASE_INSTANCE_NAME + valueFrom: "database" # References resources.apps..resources[name='database'] +``` + +**Critical:** Every `valueFrom` value must match a `name` field in `databricks.yml` resources. + ## Important Notes - **MLflow experiment**: Already configured in template, no action needed -- **Lakebase permissions**: Memory storage requires separate Lakebase setup (see **lakebase-setup** skill) - **Multiple resources**: Add multiple entries under `resources:` list - **Permission types vary**: Each resource type has specific permission values -- **Deploy after changes**: Run `databricks bundle deploy` after modifying `databricks.yml` - -## Next Steps - -- Configure memory storage: see **lakebase-setup** skill -- Understand memory patterns: see **agent-memory** skill -- Test locally: see **run-locally** skill -- Deploy: see **deploy** skill +- **Deploy + Run after changes**: Run both `databricks bundle deploy` AND `databricks bundle run agent_langgraph` +- **valueFrom matching**: Ensure `app.yaml` `valueFrom` values match `databricks.yml` resource `name` values diff --git a/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md b/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md index 9e919ad2..1324e6c5 100644 --- a/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md +++ b/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md @@ -24,7 +24,7 @@ tools = await mcp_client.get_tools() ```bash databricks bundle deploy -databricks bundle run agent_langgraph_short_term_memory +databricks bundle run agent_langgraph ``` ### 3. Get your agent app's service principal diff --git a/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/lakebase.yaml b/agent-langgraph-short-term-memory/.claude/skills/add-tools/examples/lakebase.yaml new file mode 100644 index 00000000..78f0bc72 --- /dev/null +++ b/agent-langgraph-short-term-memory/.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: valueFrom reference in app.yaml + +# In databricks.yml - add to resources.apps..resources: +- name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +# In app.yaml - add to env: +# - name: LAKEBASE_INSTANCE_NAME +# valueFrom: "database" +# - name: EMBEDDING_ENDPOINT +# value: "databricks-gte-large-en" +# - name: EMBEDDING_DIMS +# value: "1024" diff --git a/agent-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md index 64b8895c..83d59668 100644 --- a/agent-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md +++ b/agent-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md @@ -1,196 +1,524 @@ --- name: agent-memory -description: "Understand and modify agent memory patterns. Use when: (1) User asks about 'memory', 'state', 'conversation history', (2) Working with thread_id or user_id, (3) Debugging memory issues, (4) Adding long-term memory capabilities." +description: "Add memory capabilities to your agent. Use when: (1) User asks about 'memory', 'state', 'remember', 'conversation history', (2) Want to persist conversations or user preferences, (3) Adding checkpointing or long-term storage." --- -# Agent Memory Patterns +# Adding Memory to Your Agent -This skill covers both short-term memory (conversation history within a session) and long-term memory (facts that persist across sessions). +> **Note:** This template does not include memory by default. Use this skill to **add memory capabilities**. For pre-configured memory templates, see: +> - `agent-langgraph-short-term-memory` - Conversation history within a session +> - `agent-langgraph-long-term-memory` - User facts that persist across sessions -## Short-Term Memory (This Template) +## Memory Types -Short-term memory stores conversation history within a thread. The agent remembers what was said earlier in the conversation. +| Type | Use Case | Storage | Identifier | +|------|----------|---------|------------| +| **Short-term** | Conversation history within a session | `AsyncCheckpointSaver` | `thread_id` | +| **Long-term** | User facts that persist across sessions | `AsyncDatabricksStore` | `user_id` | -### Key Components +## Prerequisites -**AsyncCheckpointSaver** - Persists LangGraph state to Lakebase: +1. **Add memory dependency** to `pyproject.toml`: + ```toml + dependencies = [ + "databricks-langchain[memory]", + ] + ``` -```python -from databricks_langchain import AsyncCheckpointSaver + Then run `uv sync` -async with AsyncCheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME) as checkpointer: - agent = await init_agent(checkpointer=checkpointer) - # Agent now persists state between calls +2. **Configure Lakebase** - See **lakebase-setup** skill for: + - Creating/configuring Lakebase instance + - Initializing tables (CRITICAL first-time step) + +--- + +## Quick Setup Summary + +Adding memory requires changes to **4 files**: + +| File | What to Add | +|------|-------------| +| `pyproject.toml` | Memory dependency + hatch config | +| `.env` | Lakebase env vars (for local dev) | +| `databricks.yml` | Lakebase database resource | +| `app.yaml` | `valueFrom` reference to lakebase resource | +| `agent_server/agent.py` | Memory tools and AsyncDatabricksStore | + +--- + +## Step 1: Configure databricks.yml (Lakebase Resource) + +Add the Lakebase database resource to your app in `databricks.yml`: + +```yaml +resources: + apps: + agent_langgraph: + name: "your-app-name" + source_code_path: ./ + + resources: + # ... other resources (experiment, UC functions, etc.) ... + + # Lakebase instance for long-term memory + - name: 'lakebadatabasese_memory' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' ``` -**StatefulAgentState** - Custom state schema for memory: +**Important:** The `name: 'database'` must match the `valueFrom` reference in `app.yaml`. -```python -from typing import Any, Sequence, TypedDict -from langchain_core.messages import AnyMessage -from langgraph.graph.message import add_messages -from typing_extensions import Annotated +--- + +## Step 2: Configure app.yaml (Environment Variables) + +Update `app.yaml` with the Lakebase environment variables: + +```yaml +command: ["uv", "run", "start-app"] -class StatefulAgentState(TypedDict, total=False): - messages: Annotated[Sequence[AnyMessage], add_messages] - custom_inputs: dict[str, Any] - custom_outputs: dict[str, Any] +env: + # ... other env vars ... + + # MLflow experiment (uses valueFrom to reference databricks.yml resource) + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" + + # Lakebase instance name - must match the instance_name in databricks.yml database resource + # Note: Use 'value' (not 'valueFrom') because AsyncDatabricksStore needs the instance name, + # not the full connection string that valueFrom would provide + - name: LAKEBASE_INSTANCE_NAME + value: "" + + # Embedding configuration (static values) + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" ``` -**Thread ID** - Identifies a conversation thread: +**Important:** The `LAKEBASE_INSTANCE_NAME` value must match the `instance_name` in your `databricks.yml` database resource. The `database` resource handles permissions, while `app.yaml` provides the instance name to your code. -```python -def _get_or_create_thread_id(request: ResponsesAgentRequest) -> str: - # Priority: - # 1. Use thread_id from custom_inputs - # 2. Use conversation_id from ChatContext - # 3. Generate random UUID - ci = dict(request.custom_inputs or {}) +### Why not valueFrom for Lakebase? - if "thread_id" in ci and ci["thread_id"]: - return str(ci["thread_id"]) +The `database` resource's `valueFrom` provides the full connection string (e.g., `instance-xxx.database.staging.cloud.databricks.com`), but `AsyncDatabricksStore` expects just the instance name. So we use a static `value` instead. - if request.context and getattr(request.context, "conversation_id", None): - return str(request.context.conversation_id) +--- + +## Step 3: Configure .env (Local Development) + +For local development, add to `.env`: - return str(uuid_utils.uuid7()) +```bash +# Lakebase configuration for long-term memory +LAKEBASE_INSTANCE_NAME= +EMBEDDING_ENDPOINT=databricks-gte-large-en +EMBEDDING_DIMS=1024 ``` -### Creating a Stateful Agent +> **Note:** `.env` is only for local development. When deployed, the app gets `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. + +--- -Pass `checkpointer` and `state_schema` to `create_agent()`: +## Step 4: Update agent.py + +### Add Imports and Configuration ```python -from langchain.agents import create_agent +import os +import uuid +from typing import AsyncGenerator + +from databricks_langchain import AsyncDatabricksStore, ChatDatabricks +from langchain_core.runnables import RunnableConfig +from langchain_core.tools import tool +from langgraph.prebuilt import create_react_agent + +# Environment configuration +LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") +EMBEDDING_ENDPOINT = os.getenv("EMBEDDING_ENDPOINT", "databricks-gte-large-en") +EMBEDDING_DIMS = int(os.getenv("EMBEDDING_DIMS", "1024")) # REQUIRED! +``` -agent = create_agent( - model=model, - tools=tools, - system_prompt=SYSTEM_PROMPT, - checkpointer=checkpointer, # Enables persistence - state_schema=StatefulAgentState, # Custom state with messages -) +**CRITICAL:** You MUST specify `embedding_dims` when using `embedding_endpoint`. Without it, you'll get: +``` +"embedding_dims is required when embedding_endpoint is specified" ``` -### Using Thread ID in Requests +### Create Memory Tools + +```python +@tool +async def get_user_memory(query: str, config: RunnableConfig) -> str: + """Search for relevant information about the user from long-term memory. + Use this to recall preferences, past interactions, or other saved information. + """ + user_id = config.get("configurable", {}).get("user_id") + store = config.get("configurable", {}).get("store") + if not user_id or not store: + return "Memory not available - no user context" + + # Sanitize user_id for namespace (replace special chars) + namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) + results = await store.asearch(namespace, query=query, limit=5) + if not results: + return "No memories found for this query" + return "\n".join([f"- {r.value}" for r in results]) + + +@tool +async def save_user_memory(memory_key: str, memory_data: str, config: RunnableConfig) -> str: + """Save information about the user to long-term memory. + Use this to remember user preferences, important details, or other + information that should persist across conversations. + + Args: + memory_key: A short descriptive key (e.g., "preferred_name", "team", "region") + memory_data: The information to save + """ + user_id = config.get("configurable", {}).get("user_id") + store = config.get("configurable", {}).get("store") + if not user_id or not store: + return "Memory not available - no user context" + + namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) + await store.aput(namespace, memory_key, {"value": memory_data}) + return f"Saved memory: {memory_key}" + + +@tool +async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: + """Delete a specific memory from the user's long-term memory. + Use this when the user asks to forget something or correct stored information. + + Args: + memory_key: The key of the memory to delete (e.g., "preferred_name", "team") + """ + user_id = config.get("configurable", {}).get("user_id") + store = config.get("configurable", {}).get("store") + if not user_id or not store: + return "Memory not available - no user context" + + namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) + await store.adelete(namespace, memory_key) + return f"Deleted memory: {memory_key}" +``` -The agent uses `config` to track threads: +### Add Helper Functions ```python -config = {"configurable": {"thread_id": thread_id}} -input_state = { - "messages": to_chat_completions_input([i.model_dump() for i in request.input]), - "custom_inputs": dict(request.custom_inputs or {}), -} +def _get_user_id(request: ResponsesAgentRequest) -> str: + """Extract user_id from request context or custom inputs.""" + custom_inputs = dict(request.custom_inputs or {}) + if "user_id" in custom_inputs and custom_inputs["user_id"]: + return str(custom_inputs["user_id"]) + if request.context and getattr(request.context, "user_id", None): + return str(request.context.user_id) + return "default-user" -async for event in agent.astream(input_state, config, stream_mode=["updates", "messages"]): - # Process events + +def _get_or_create_thread_id(request: ResponsesAgentRequest) -> str: + """Extract or create thread_id for conversation tracking.""" + custom_inputs = dict(request.custom_inputs or {}) + if "thread_id" in custom_inputs and custom_inputs["thread_id"]: + return str(custom_inputs["thread_id"]) + if request.context and getattr(request.context, "conversation_id", None): + return str(request.context.conversation_id) + return str(uuid.uuid4()) +``` + +### Update Streaming Function + +```python +@stream() +async def streaming( + request: ResponsesAgentRequest, +) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + # Get user context + user_id = _get_user_id(request) + thread_id = _get_or_create_thread_id(request) + + # Get other tools (MCP, etc.) + mcp_client = init_mcp_client(sp_workspace_client) + mcp_tools = await mcp_client.get_tools() + + # Memory tools + memory_tools = [get_user_memory, save_user_memory, delete_user_memory] + + # Combine all tools + all_tools = mcp_tools + memory_tools + + # Initialize model + model = ChatDatabricks(endpoint="databricks-claude-sonnet-4") + + # Use AsyncDatabricksStore for long-term memory + async with AsyncDatabricksStore( + instance_name=LAKEBASE_INSTANCE_NAME, + embedding_endpoint=EMBEDDING_ENDPOINT, + embedding_dims=EMBEDDING_DIMS, # REQUIRED! + ) as store: + # Create agent with tools + agent = create_react_agent( + model=model, + tools=all_tools, + prompt=AGENT_INSTRUCTIONS, + ) + + # Prepare input + messages = {"messages": to_chat_completions_input([i.model_dump() for i in request.input])} + + # Configure with user context for memory tools + config = { + "configurable": { + "user_id": user_id, + "thread_id": thread_id, + "store": store, # Pass store to tools via config + } + } + + # Stream agent responses + async for event in process_agent_astream_events( + agent.astream(input=messages, config=config, stream_mode=["updates", "messages"]) + ): + yield event ``` --- -## API Requests with Memory +## Step 5: Initialize Tables and Deploy -### Continuing a Conversation (with thread_id) +### Initialize Lakebase Tables (First Time Only) + +Before deploying, initialize the tables locally: ```bash -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ - "input": [{"role": "user", "content": "What did we discuss?"}], - "custom_inputs": {"thread_id": ""} - }' +uv run python -c "$(cat <<'EOF' +import asyncio +from databricks_langchain import AsyncDatabricksStore + +async def setup(): + async with AsyncDatabricksStore( + instance_name="", + embedding_endpoint="databricks-gte-large-en", + embedding_dims=1024, + ) as store: + await store.setup() + print("Tables created!") + +asyncio.run(setup()) +EOF +)" ``` -### Starting a New Conversation +### Deploy and Run -Omit `thread_id` to start fresh (a new UUID will be generated): +**IMPORTANT:** Always run `databricks bundle run` after `databricks bundle deploy` to start/restart the app with the new code: ```bash -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{"input": [{"role": "user", "content": "Hello!"}], "stream": true}' +# Deploy resources and upload files +databricks bundle deploy + +# Start/restart the app with new code (REQUIRED!) +databricks bundle run agent_langgraph ``` -The response includes `thread_id` in `custom_outputs` - save it to continue the conversation. +> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is required to actually start the app with the new code. -### Deployed App with Memory +--- -```bash -curl -X POST /invocations \ - -H "Authorization: Bearer " \ - -H "Content-Type: application/json" \ - -d '{ - "input": [{"role": "user", "content": "Remember this: my favorite color is blue"}], - "custom_inputs": {"thread_id": "user-123-session-1"} - }' +## Complete Example Files + +### databricks.yml (with Lakebase) + +```yaml +bundle: + name: agent_langgraph + +resources: + experiments: + agent_langgraph_experiment: + name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} + + apps: + agent_langgraph: + name: "my-agent-app" + description: "Agent with long-term memory" + source_code_path: ./ + + resources: + - name: 'experiment' + experiment: + experiment_id: "${resources.experiments.agent_langgraph_experiment.id}" + permission: 'CAN_MANAGE' + + # Lakebase instance for long-term memory + - name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +targets: + dev: + mode: development + default: true ``` ---- +### app.yaml + +```yaml +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" + # Reference experiment resource from databricks.yml + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" + # Lakebase instance name (must match instance_name in databricks.yml) + - name: LAKEBASE_INSTANCE_NAME + value: "" + # Embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` -## Long-Term Memory (Reference) +--- -Long-term memory stores facts that persist across conversation sessions (e.g., user preferences). This is available in the `agent-langgraph-long-term-memory` template. +## Adding Short-Term Memory -### Key Components (Long-Term) +Short-term memory stores conversation history within a thread. -**AsyncDatabricksStore** - Persists memories to Lakebase: +### Import and Initialize Checkpointer ```python -from databricks_langchain import AsyncDatabricksStore +from databricks_langchain import AsyncCheckpointSaver -async with AsyncDatabricksStore(instance_name=LAKEBASE_INSTANCE_NAME) as store: - # Store persists memories across sessions +LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") + +async with AsyncCheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME) as checkpointer: + agent = create_react_agent( + model=model, + tools=tools, + checkpointer=checkpointer, # Enables conversation persistence + ) ``` -**User ID** - Identifies a user across sessions: +### Use thread_id in Agent Config ```python -user_id = request.custom_inputs.get("user_id", "default-user") +config = {"configurable": {"thread_id": thread_id}} +async for event in agent.astream(input_state, config, stream_mode=["updates", "messages"]): + yield event ``` -**Memory Tools** - LangGraph provides tools for memory management: -- `manage_memory` - Save/update memories about the user -- The agent can reference stored memories in conversations +--- + +## Testing Memory -### Long-Term Memory Request Example +### Test Locally ```bash +# Start the server +uv run start-app + +# Save a memory curl -X POST http://localhost:8000/invocations \ -H "Content-Type: application/json" \ -d '{ - "input": [{"role": "user", "content": "What do you remember about me?"}], - "custom_inputs": { - "thread_id": "", - "user_id": "user-123" - } + "input": [{"role": "user", "content": "Remember that I am on the shipping team"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' + +# Recall the memory +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "What team am I on?"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' + +# Delete a memory +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Forget what team I am on"}], + "custom_inputs": {"user_id": "alice@example.com"} }' ``` +### Test Deployed App + +```bash +# Get OAuth token (PATs don't work for apps) +TOKEN=$(databricks auth token --host | jq -r '.access_token') + +# Test memory save +curl -X POST https:///invocations \ + -H "Authorization: Bearer $TOKEN" \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Remember I prefer detailed explanations"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' +``` + +--- + +## First-Time Setup Checklist + +- [ ] Added `databricks-langchain[memory]` to `pyproject.toml` +- [ ] Run `uv sync` to install dependencies +- [ ] Created or identified Lakebase instance +- [ ] Added Lakebase env vars to `.env` (for local dev) +- [ ] Added `database` resource to `databricks.yml` (for permissions) +- [ ] Added `LAKEBASE_INSTANCE_NAME` value to `app.yaml` (matching instance_name in databricks.yml) +- [ ] **Initialized tables locally** by running `await store.setup()` +- [ ] Deployed with `databricks bundle deploy` +- [ ] **Started app with `databricks bundle run agent_langgraph`** + --- ## Troubleshooting -| Issue | Solution | -|-------|----------| -| **Conversation not persisting** | Verify `thread_id` is passed and consistent | -| **"Cannot connect to Lakebase"** | See **lakebase-setup** skill | -| **Permission errors** | Grant checkpoint table permissions (see **lakebase-setup** skill) | -| **State not loading** | Check that `checkpointer` is passed to `create_agent()` | +| Issue | Cause | Solution | +|-------|-------|----------| +| **"embedding_dims is required"** | Missing parameter | Add `embedding_dims=1024` to AsyncDatabricksStore | +| **"relation 'store' does not exist"** | Tables not created | Run `await store.setup()` locally first | +| **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `valueFrom: "database"` to app.yaml | +| **"permission denied for table store"** | Missing grants | Lakebase `database` resource in DAB should handle this | +| **"Memory not available"** | No user_id in request | Ensure `custom_inputs.user_id` is passed | +| **Memory not persisting** | Different user_ids | Use consistent user_id across requests | +| **App not updated after deploy** | Forgot to run bundle | Run `databricks bundle run agent_langgraph` after deploy | -## Dependencies +--- -Short-term memory requires `databricks-langchain[memory]`: +## Pre-Built Memory Templates -```toml -# pyproject.toml -dependencies = [ - "databricks-langchain[memory]>=0.13.0", -] -``` +For fully configured implementations without manual setup: + +| Template | Memory Type | Key Features | +|----------|-------------|--------------| +| `agent-langgraph-short-term-memory` | Short-term | AsyncCheckpointSaver, thread_id | +| `agent-langgraph-long-term-memory` | Long-term | AsyncDatabricksStore, memory tools | + +--- ## Next Steps - Configure Lakebase: see **lakebase-setup** skill -- Modify agent behavior: see **modify-agent** skill - Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/agent-langgraph-short-term-memory/.claude/skills/deploy/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/deploy/SKILL.md index ebd80e08..17f08ba0 100644 --- a/agent-langgraph-short-term-memory/.claude/skills/deploy/SKILL.md +++ b/agent-langgraph-short-term-memory/.claude/skills/deploy/SKILL.md @@ -5,18 +5,38 @@ description: "Deploy agent to Databricks Apps using DAB (Databricks Asset Bundle # Deploy to Databricks Apps -> **Memory Template:** Before deploying, ensure Lakebase is configured. See the **lakebase-setup** skill for permissions setup required after deployment. +## 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_langgraph_short_term_memory: + name: "agent-your-app-name" # Use agent-* prefix +``` ## Deploy Commands +**IMPORTANT:** Always run BOTH commands to deploy and start your app: + ```bash -# Deploy the bundle (creates/updates resources, uploads files) +# 1. Validate bundle configuration (catches errors before deploy) +databricks bundle validate + +# 2. Deploy the bundle (creates/updates resources, uploads files) databricks bundle deploy -# Run the app (starts/restarts with uploaded source code) +# 3. Run the app (starts/restarts with uploaded source code) - REQUIRED! databricks bundle run agent_langgraph_short_term_memory ``` +> **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_langgraph_short_term_memory` matches the app name in `databricks.yml` under `resources.apps`. ## Handling "App Already Exists" Error @@ -86,15 +106,6 @@ databricks bundle deploy **Warning:** This permanently deletes the app's URL, OAuth credentials, and service principal. -## Post-Deployment: Lakebase Permissions - -**After deploying a memory-enabled agent, you must grant Lakebase permissions.** - -See the **lakebase-setup** skill for: -- Adding Lakebase as an app resource -- SDK or SQL commands to grant checkpoint table permissions -- Troubleshooting access errors - ## Unbinding an App To remove the link between bundle and deployed app: @@ -127,7 +138,16 @@ curl -X POST /invocations \ -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' ``` -For memory-specific request examples (with thread_id), see the **agent-memory** skill. +**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 @@ -190,16 +210,13 @@ uv add | 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 | +| 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`| +| "Provider produced inconsistent result" | Sync app config to `databricks.yml` | | "should set workspace.root_path" | Add `root_path` to production target | - -## Next Steps - -- Grant Lakebase permissions: see **lakebase-setup** skill -- Understand memory patterns: see **agent-memory** skill -- Add tools and permissions: see **add-tools** skill +| App running old code after deploy | Run `databricks bundle run agent_langgraph_short_term_memory` after deploy | +| Env var is None in deployed app | Check `valueFrom` in app.yaml matches resource `name` in databricks.yml | diff --git a/agent-langgraph-short-term-memory/.claude/skills/discover-tools/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/discover-tools/SKILL.md index eabac2a5..87c3f519 100644 --- a/agent-langgraph-short-term-memory/.claude/skills/discover-tools/SKILL.md +++ b/agent-langgraph-short-term-memory/.claude/skills/discover-tools/SKILL.md @@ -39,44 +39,9 @@ uv run discover-tools --profile DEFAULT | **Custom MCP Servers** | Apps starting with `mcp-*` | `{app_url}/mcp` | | **External MCP Servers** | Via UC connections | `{host}/api/2.0/mcp/external/{connection_name}` | -## Using Discovered Tools in Code - -After discovering tools, add them to your agent in `agent_server/agent.py`: - -```python -from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient - -# Example: Add UC functions from a schema -uc_functions_server = DatabricksMCPServer( - url=f"{host}/api/2.0/mcp/functions/{catalog}/{schema}", - name="my uc functions", -) - -# Example: Add a Genie space -genie_server = DatabricksMCPServer( - url=f"{host}/api/2.0/mcp/genie/{space_id}", - name="my genie space", -) - -# Example: Add vector search -vector_server = DatabricksMCPServer( - url=f"{host}/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name}", - name="my vector index", -) - -# Create MCP client with all servers -mcp_client = DatabricksMultiServerMCPClient([ - uc_functions_server, - genie_server, - vector_server, -]) - -# Get tools for the agent -tools = await mcp_client.get_tools() -``` - ## Next Steps -After adding MCP servers to your agent: -1. **Grant permissions** in `databricks.yml` (see **add-tools** skill) -2. Test locally with `uv run start-app` (see **run-locally** skill) +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 diff --git a/agent-langgraph-short-term-memory/.claude/skills/lakebase-setup/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/lakebase-setup/SKILL.md index 96b2375a..6b8c802e 100644 --- a/agent-langgraph-short-term-memory/.claude/skills/lakebase-setup/SKILL.md +++ b/agent-langgraph-short-term-memory/.claude/skills/lakebase-setup/SKILL.md @@ -1,168 +1,351 @@ --- name: lakebase-setup -description: "Configure Lakebase for agent memory storage. Use when: (1) First-time memory setup, (2) 'Failed to connect to Lakebase' errors, (3) Permission errors on checkpoint tables, (4) User says 'lakebase', 'memory setup', or 'checkpoint'." +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 Memory +# Lakebase Setup for Agent Memory -This template uses Lakebase (Databricks-managed PostgreSQL) to store conversation memory. You must configure Lakebase before the agent can persist state. +> **Note:** This template does not include memory by default. Use this skill if you want to **add memory capabilities** to your agent. For pre-configured memory templates, see: +> - `agent-langgraph-short-term-memory` - Conversation history within a session +> - `agent-langgraph-long-term-memory` - User facts that persist across sessions -## Prerequisites +## Overview -- A Lakebase instance in your Databricks workspace -- The instance name (found in Databricks UI under SQL Warehouses > Lakebase) +Lakebase provides persistent storage for agent memory: +- **Short-term memory**: Conversation history within a thread (`AsyncCheckpointSaver`) +- **Long-term memory**: User facts across sessions (`AsyncDatabricksStore`) -## Local Development Setup +## Complete Setup Workflow -**Step 1:** Add the Lakebase instance name to `.env`: +``` +┌─────────────────────────────────────────────────────────────────────────────┐ +│ 1. Add dependency → 2. Get instance → 3. Configure DAB + app.yaml │ +│ 4. Configure .env → 5. Initialize tables → 6. Deploy + Run │ +└─────────────────────────────────────────────────────────────────────────────┘ +``` + +--- + +## Step 1: Add Memory Dependency +Add the memory extra to your `pyproject.toml`: + +```toml +dependencies = [ + "databricks-langchain[memory]", + # ... other dependencies +] +``` + +Then sync dependencies: ```bash -LAKEBASE_INSTANCE_NAME= +uv sync ``` -**Step 2:** Run `uv run start-app` to test locally. The agent will automatically create checkpoint tables on first run. +--- -## Deployed App Setup +## Step 2: Create or Get Lakebase Instance -After deploying your agent with `databricks bundle deploy`, you must grant the app's service principal access to Lakebase. +### Option A: Create New Instance (via Databricks UI) -### Step 1: Add Lakebase as App Resource +1. Go to your Databricks workspace +2. Navigate to **Compute** → **Lakebase** +3. Click **Create Instance** +4. Note the instance name -1. Go to the Databricks UI -2. Navigate to your app and click **Edit** -3. Go to **App resources** → **Add resource** -4. Add your Lakebase instance with **Connect + Create** permissions +### Option B: Use Existing Instance -### Step 2: Get App Service Principal ID +If you have an existing instance, note its name for the next step. -```bash -databricks apps get --output json | jq -r '.service_principal_id' +--- + +## Step 3: Configure databricks.yml (Lakebase Resource) + +Add the Lakebase `database` resource to your app in `databricks.yml`: + +```yaml +resources: + apps: + agent_langgraph: + 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: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' ``` -### Step 3: Grant Permissions +**Important:** +- The `name: 'database'` must match the `valueFrom` reference in `app.yaml` +- Using the `database` resource type automatically grants the app's service principal access to Lakebase -Choose **Option A** (SDK - Recommended) or **Option B** (SQL). +### Update app.yaml (Environment Variables) -#### Option A: Using LakebaseClient SDK (Recommended) +Update `app.yaml` with the Lakebase instance name: -The `databricks-ai-bridge` package includes a `LakebaseClient` for programmatic permission management: +```yaml +env: + # ... other env vars ... -```python -from databricks_ai_bridge.lakebase import LakebaseClient, SchemaPrivilege, TablePrivilege + # Lakebase instance name - must match instance_name in databricks.yml database resource + # Note: Use 'value' (not 'valueFrom') because AsyncDatabricksStore needs the instance name, + # not the full connection string that valueFrom would provide + - name: LAKEBASE_INSTANCE_NAME + value: "" -# Initialize client -client = LakebaseClient(instance_name="") + # Static values for embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` -app_sp = "" # From Step 2 +**Important:** +- The `LAKEBASE_INSTANCE_NAME` value must match the `instance_name` in your `databricks.yml` database resource +- The `database` resource handles permissions; `app.yaml` provides the instance name to your code +- Don't use `valueFrom` for Lakebase - it provides the connection string, not the instance name -# Create role for the service principal -client.create_role(app_sp, "SERVICE_PRINCIPAL") +--- -# Grant schema privileges -client.grant_schema( - grantee=app_sp, - privileges=[SchemaPrivilege.USAGE, SchemaPrivilege.CREATE], - schemas=["drizzle", "ai_chatbot", "public"], -) +## Step 4: Configure .env (Local Development) -# Grant table privileges on all tables in schemas -client.grant_all_tables_in_schema( - grantee=app_sp, - privileges=[TablePrivilege.SELECT, TablePrivilege.INSERT, TablePrivilege.UPDATE], - schemas=["drizzle", "ai_chatbot"], -) +For local development, add to `.env`: -# Grant privileges on specific checkpoint tables -client.grant_table( - grantee=app_sp, - privileges=[TablePrivilege.SELECT, TablePrivilege.INSERT, TablePrivilege.UPDATE], - tables=[ - "public.checkpoint_migrations", - "public.checkpoint_writes", - "public.checkpoints", - "public.checkpoint_blobs", - ], -) +```bash +# Lakebase configuration for long-term memory +LAKEBASE_INSTANCE_NAME= +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 `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. + +--- + +## Step 5: Initialize Store Tables (CRITICAL - First Time Only) + +**Before deploying**, you must initialize the Lakebase tables. The `AsyncDatabricksStore` creates tables on first use, but you need to do this locally first: + +```python +# Run this script locally BEFORE first deployment +import asyncio +from databricks_langchain import AsyncDatabricksStore + +async def setup_store(): + async with AsyncDatabricksStore( + instance_name="", + embedding_endpoint="databricks-gte-large-en", + embedding_dims=1024, + ) as store: + print("Setting up store tables...") + await store.setup() # Creates required tables + print("Store tables created!") + + # Verify with a test write/read + await store.aput(("test", "init"), "test_key", {"value": "test_value"}) + results = await store.asearch(("test", "init"), query="test", limit=1) + print(f"Test successful: {results}") + +asyncio.run(setup_store()) ``` -**Benefits of SDK approach:** -- Type-safe privilege enums prevent typos -- Cleaner Python code vs raw SQL -- Easier to integrate into setup scripts - -#### Option B: Using SQL - -Run the following SQL on your Lakebase instance (replace `app-sp-id` with your app's service principal ID): - -```sql -DO $$ -DECLARE - app_sp text := 'app-sp-id'; -- TODO: Replace with your App's Service Principal ID -BEGIN - ------------------------------------------------------------------- - -- Drizzle schema: migration metadata tables - ------------------------------------------------------------------- - EXECUTE format('GRANT USAGE, CREATE ON SCHEMA drizzle TO %I;', app_sp); - EXECUTE format('GRANT SELECT, INSERT, UPDATE ON ALL TABLES IN SCHEMA drizzle TO %I;', app_sp); - ------------------------------------------------------------------- - -- App schema: business tables (Chat, Message, etc.) - ------------------------------------------------------------------- - EXECUTE format('GRANT USAGE, CREATE ON SCHEMA ai_chatbot TO %I;', app_sp); - EXECUTE format('GRANT SELECT, INSERT, UPDATE ON ALL TABLES IN SCHEMA ai_chatbot TO %I;', app_sp); - ------------------------------------------------------------------- - -- Public schema for checkpoint tables - ------------------------------------------------------------------- - EXECUTE format('GRANT USAGE, CREATE ON SCHEMA public TO %I;', app_sp); - EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.checkpoint_migrations TO %I;', app_sp); - EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.checkpoint_writes TO %I;', app_sp); - EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.checkpoints TO %I;', app_sp); - EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.checkpoint_blobs TO %I;', app_sp); -END $$; +Run with: +```bash +uv run python -c "$(cat <<'EOF' +import asyncio +from databricks_langchain import AsyncDatabricksStore + +async def setup(): + async with AsyncDatabricksStore( + instance_name="", + embedding_endpoint="databricks-gte-large-en", + embedding_dims=1024, + ) as store: + await store.setup() + print("Tables created!") + +asyncio.run(setup()) +EOF +)" ``` -**Schema Reference:** -| Schema | Purpose | -|--------|---------| -| `drizzle` | Migration metadata tables | -| `ai_chatbot` | Business tables (Chat, Message, etc.) | -| `public` | Checkpoint tables for conversation state | +This creates these tables in the `public` schema: +- `store` - Key-value storage for memories +- `store_vectors` - Vector embeddings for semantic search +- `store_migrations` - Schema migration tracking +- `vector_migrations` - Vector schema migration tracking -## Troubleshooting +--- -| Issue | Solution | -|-------|----------| -| **"LAKEBASE_INSTANCE_NAME environment variable is required"** | Add `LAKEBASE_INSTANCE_NAME=` to `.env` | -| **"Failed to connect to Lakebase instance"** | Verify instance name is correct and your profile has access | -| **Permission errors on checkpoint tables** | Run the SDK script or SQL grant commands above | -| **Deployed app can't access Lakebase** | Add Lakebase as app resource in Databricks UI | -| **"role does not exist"** | Run `client.create_role()` first, or ensure SP ID is correct | +## Step 6: Deploy and Run Your App -### Common Error Messages +**IMPORTANT:** Always run both `deploy` AND `run` commands: + +```bash +# Deploy resources and upload files +databricks bundle deploy -**Local development:** +# Start/restart the app with new code (REQUIRED!) +databricks bundle run agent_langgraph ``` -Failed to connect to Lakebase instance ''. Please verify: -1. The instance name is correct -2. You have the necessary permissions to access the instance -3. Your Databricks authentication is configured correctly + +> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is required to actually start the app with the new code. + +--- + +## Complete Example: databricks.yml with Lakebase + +```yaml +bundle: + name: agent_langgraph + +resources: + experiments: + agent_langgraph_experiment: + name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} + + apps: + agent_langgraph: + name: "my-agent-app" + description: "Agent with long-term memory" + source_code_path: ./ + + resources: + - name: 'experiment' + experiment: + experiment_id: "${resources.experiments.agent_langgraph_experiment.id}" + permission: 'CAN_MANAGE' + + # Lakebase instance for long-term memory + - name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +targets: + dev: + mode: development + default: true ``` -**Deployed app:** +## Complete Example: app.yaml + +```yaml +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" + # Reference experiment resource from databricks.yml + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" + # Lakebase instance name (must match instance_name in databricks.yml) + - name: LAKEBASE_INSTANCE_NAME + value: "" + # Embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" ``` -Failed to connect to Lakebase instance ''. The App Service Principal for '' may not have access. + +--- + +## 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 | Run `await store.setup()` locally first (Step 5) | +| **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `LAKEBASE_INSTANCE_NAME` value to app.yaml | +| **"Unable to resolve Lakebase instance '...database.cloud.databricks.com'"** | Used valueFrom instead of value | Use `value: ""` not `valueFrom` for Lakebase | +| **"permission denied for table store"** | Missing grants | The `database` resource in DAB should handle this; verify the resource is configured | +| **"Failed to connect to Lakebase"** | Wrong instance name | Verify instance name 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 agent_langgraph` after deploy | +| **valueFrom not resolving** | Resource name mismatch | Ensure `valueFrom` value matches `name` in databricks.yml resources | + +--- + +## Quick Reference: LakebaseClient API + +For manual permission management (usually not needed with DAB `database` resource): + +```python +from databricks_ai_bridge.lakebase import LakebaseClient, SchemaPrivilege, TablePrivilege + +client = LakebaseClient(instance_name="...") + +# 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'") ``` -Both errors indicate Lakebase access issues. Follow the setup steps above for your environment. +### Service Principal Identifiers -## How Memory Works +When granting permissions manually, note that Databricks apps have multiple identifiers: -See the **agent-memory** skill for: -- How `AsyncCheckpointSaver` persists conversation state -- Using `thread_id` to maintain conversation context -- API request examples with thread_id +| 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 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 -- Understand memory patterns: see **agent-memory** skill +- Add memory to agent code: see **agent-memory** skill - Test locally: see **run-locally** skill - Deploy: see **deploy** skill diff --git a/agent-langgraph-short-term-memory/.claude/skills/modify-agent/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/modify-agent/SKILL.md index 575080d4..d7218637 100644 --- a/agent-langgraph-short-term-memory/.claude/skills/modify-agent/SKILL.md +++ b/agent-langgraph-short-term-memory/.claude/skills/modify-agent/SKILL.md @@ -5,8 +5,6 @@ description: "Modify agent code, add tools, or change configuration. Use when: ( # Modify the Agent -> **Memory Template:** This template uses short-term memory (conversation history). See the **agent-memory** skill for memory-specific patterns. - ## Main File **`agent_server/agent.py`** - Agent logic, model selection, instructions, MCP servers @@ -290,7 +288,6 @@ Example: Create UC function wrapping HTTP request for Slack, then expose via MCP - Discover available tools: see **discover-tools** skill - Grant resource permissions: see **add-tools** skill -- Memory patterns: see **agent-memory** skill -- Lakebase setup: see **lakebase-setup** skill +- Add memory capabilities: see **agent-memory** skill - Test locally: see **run-locally** skill - Deploy: see **deploy** skill diff --git a/agent-langgraph-short-term-memory/.claude/skills/quickstart/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/quickstart/SKILL.md index 853b78d0..e550162c 100644 --- a/agent-langgraph-short-term-memory/.claude/skills/quickstart/SKILL.md +++ b/agent-langgraph-short-term-memory/.claude/skills/quickstart/SKILL.md @@ -79,6 +79,5 @@ CHAT_PROXY_TIMEOUT_SECONDS=300 ## Next Steps After quickstart completes: -1. **Configure Lakebase** for memory storage (see **lakebase-setup** skill) -2. Run `uv run discover-tools` to find available workspace resources (see **discover-tools** skill) -3. Run `uv run start-app` to test locally (see **run-locally** skill) +1. Run `uv run discover-tools` to find available workspace resources (see **discover-tools** skill) +2. Run `uv run start-app` to test locally (see **run-locally** skill) diff --git a/agent-langgraph-short-term-memory/.claude/skills/run-locally/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/run-locally/SKILL.md index f33918eb..3eb83c82 100644 --- a/agent-langgraph-short-term-memory/.claude/skills/run-locally/SKILL.md +++ b/agent-langgraph-short-term-memory/.claude/skills/run-locally/SKILL.md @@ -11,20 +11,25 @@ description: "Run and test the agent locally. Use when: (1) User says 'run local uv run start-app ``` -This starts the agent at http://localhost:8000 with both the API server and chat UI. +This starts the agent at http://localhost:8000 ## Server Options -| Option | Command | When to Use | -|--------|---------|-------------| -| **Hot-reload** | `uv run start-server --reload` | During development - auto-restarts on code changes | -| **Custom port** | `uv run start-server --port 8001` | When port 8000 is in use | -| **Multiple workers** | `uv run start-server --workers 4` | Load testing or production-like simulation | -| **Combined** | `uv run start-server --reload --port 8001` | Development on alternate port | +```bash +# Hot-reload on code changes (development) +uv run start-server --reload -## Test the API +# Custom port +uv run start-server --port 8001 + +# Multiple workers (production-like) +uv run start-server --workers 4 -The agent implements the [OpenAI Responses API](https://platform.openai.com/docs/api-reference/responses) interface via MLflow's ResponsesAgent. +# Combine options +uv run start-server --reload --port 8001 +``` + +## Test the API **Streaming request:** ```bash @@ -40,38 +45,13 @@ curl -X POST http://localhost:8000/invocations \ -d '{ "input": [{ "role": "user", "content": "hi" }] }' ``` -**Request with thread_id (for conversation memory):** -```bash -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ - "input": [{"role": "user", "content": "What did we discuss?"}], - "custom_inputs": {"thread_id": ""} - }' -``` - -See the **agent-memory** skill for more memory-related request examples. - ## Run Evaluation -Evaluate your agent using MLflow scorers: - ```bash uv run agent-evaluate ``` -**What it does:** -- Runs the agent against a test dataset -- Applies MLflow scorers (RelevanceToQuery, Safety) -- Records results to your MLflow experiment - -**Customize evaluation:** -Edit `agent_server/evaluate_agent.py` to: -- Change the evaluation dataset -- Add or modify scorers -- Adjust evaluation parameters - -After evaluation completes, open the MLflow UI link for your experiment to inspect results. +Uses MLflow scorers (RelevanceToQuery, Safety). ## Run Unit Tests @@ -79,15 +59,6 @@ After evaluation completes, open the MLflow UI link for your experiment to inspe pytest [path] ``` -## Adding Dependencies - -```bash -uv add -# Example: uv add "langchain-community" -``` - -Dependencies are managed in `pyproject.toml`. - ## Troubleshooting | Issue | Solution | @@ -96,7 +67,6 @@ Dependencies are managed in `pyproject.toml`. | **Authentication errors** | Verify `.env` is correct; run **quickstart** skill | | **Module not found** | Run `uv sync` to install dependencies | | **MLflow experiment not found** | Ensure `MLFLOW_TRACKING_URI` in `.env` is `databricks://` | -| **Lakebase connection errors** | Verify `LAKEBASE_INSTANCE_NAME` in `.env`; see **lakebase-setup** skill | ### MLflow Experiment Not Found @@ -114,17 +84,7 @@ MLFLOW_TRACKING_URI="databricks://DEFAULT" # Include profile name The quickstart script configures this automatically. If you manually edited `.env`, ensure the profile name is included. -## MLflow Tracing - -This template uses MLflow for automatic tracing: -- Agent logic decorated with `@invoke()` and `@stream()` is automatically traced -- LLM invocations are captured via MLflow autologging - -To add custom trace instrumentation, see: [MLflow tracing documentation](https://docs.databricks.com/aws/en/mlflow3/genai/tracing/app-instrumentation/) - ## Next Steps -- Configure Lakebase: see **lakebase-setup** skill -- Understand memory patterns: see **agent-memory** skill - Modify your agent: see **modify-agent** skill - Deploy to Databricks: see **deploy** skill diff --git a/agent-langgraph-short-term-memory/.gitignore b/agent-langgraph-short-term-memory/.gitignore index f12e1ea0..2158c685 100644 --- a/agent-langgraph-short-term-memory/.gitignore +++ b/agent-langgraph-short-term-memory/.gitignore @@ -1,8 +1,6 @@ # Created by https://www.toptal.com/developers/gitignore/api/python # Edit at https://www.toptal.com/developers/gitignore?templates=python -databricks.yml - ### Python ### # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/agent-langgraph-short-term-memory/AGENTS.md b/agent-langgraph-short-term-memory/AGENTS.md index 3d5f45c6..b21b1ff1 100644 --- a/agent-langgraph-short-term-memory/AGENTS.md +++ b/agent-langgraph-short-term-memory/AGENTS.md @@ -1,15 +1,25 @@ # Agent Development Guide -## MANDATORY First Action +## MANDATORY First Actions -**BEFORE any other action, run `databricks auth profiles` to check authentication status.** +**Ask the user interactively:** + +1. **App deployment target:** + > "Do you have an existing Databricks app you want to deploy to, or should we create a new one? If existing, what's the app name?" + + *Note: New apps should use the `agent-*` prefix (e.g., `agent-data-analyst`) unless the user specifies otherwise.* + +2. **Lakebase instance (required for memory):** + > "This template requires Lakebase for memory. Do you have an existing Lakebase instance? If so, what's the instance name?" + +**Then check authentication status by running `databricks auth profiles`.** This helps you understand: - Which Databricks profiles are configured - Whether authentication is already set up - Which profile to use for subsequent commands -If no profiles exist, guide the user through running `uv run quickstart` to set up authentication. +If no profiles exist or `.env` is missing, guide the user through running `uv run quickstart` to set up authentication and configuration. See the **quickstart** skill for details. ## Understanding User Goals diff --git a/agent-langgraph/.claude/skills/agent-memory/SKILL.md b/agent-langgraph/.claude/skills/agent-memory/SKILL.md index ae840930..83d59668 100644 --- a/agent-langgraph/.claude/skills/agent-memory/SKILL.md +++ b/agent-langgraph/.claude/skills/agent-memory/SKILL.md @@ -187,6 +187,24 @@ async def save_user_memory(memory_key: str, memory_data: str, config: RunnableCo namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) await store.aput(namespace, memory_key, {"value": memory_data}) return f"Saved memory: {memory_key}" + + +@tool +async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: + """Delete a specific memory from the user's long-term memory. + Use this when the user asks to forget something or correct stored information. + + Args: + memory_key: The key of the memory to delete (e.g., "preferred_name", "team") + """ + user_id = config.get("configurable", {}).get("user_id") + store = config.get("configurable", {}).get("store") + if not user_id or not store: + return "Memory not available - no user context" + + namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) + await store.adelete(namespace, memory_key) + return f"Deleted memory: {memory_key}" ``` ### Add Helper Functions @@ -228,7 +246,7 @@ async def streaming( mcp_tools = await mcp_client.get_tools() # Memory tools - memory_tools = [get_user_memory, save_user_memory] + memory_tools = [get_user_memory, save_user_memory, delete_user_memory] # Combine all tools all_tools = mcp_tools + memory_tools @@ -432,6 +450,14 @@ curl -X POST http://localhost:8000/invocations \ "input": [{"role": "user", "content": "What team am I on?"}], "custom_inputs": {"user_id": "alice@example.com"} }' + +# Delete a memory +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Forget what team I am on"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' ``` ### Test Deployed App diff --git a/agent-langgraph/.claude/skills/deploy/SKILL.md b/agent-langgraph/.claude/skills/deploy/SKILL.md index 073dafd9..dfaa72af 100644 --- a/agent-langgraph/.claude/skills/deploy/SKILL.md +++ b/agent-langgraph/.claude/skills/deploy/SKILL.md @@ -5,15 +5,33 @@ description: "Deploy agent to Databricks Apps using DAB (Databricks Asset Bundle # Deploy to Databricks Apps +## 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_langgraph: + name: "agent-your-app-name" # Use agent-* prefix +``` + ## Deploy Commands **IMPORTANT:** Always run BOTH commands to deploy and start your app: ```bash -# 1. Deploy the bundle (creates/updates resources, uploads files) +# 1. Validate bundle configuration (catches errors before deploy) +databricks bundle validate + +# 2. Deploy the bundle (creates/updates resources, uploads files) databricks bundle deploy -# 2. Run the app (starts/restarts with uploaded source code) - REQUIRED! +# 3. Run the app (starts/restarts with uploaded source code) - REQUIRED! databricks bundle run agent_langgraph ``` @@ -105,7 +123,9 @@ Note: Unbinding doesn't delete the deployed app. ## Query Deployed App -**Get OAuth token** (PATs not supported): +> **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 ``` @@ -118,6 +138,33 @@ curl -X POST /invocations \ -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 @@ -138,14 +185,38 @@ databricks apps get --output json | jq -r '.url' - **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 | -| "Provider produced inconsistent result" | Sync app config to `databricks.yml`| +| 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_langgraph` after deploy | | Env var is None in deployed app | Check `valueFrom` in app.yaml matches resource `name` in databricks.yml | diff --git a/agent-langgraph/.claude/skills/discover-tools/SKILL.md b/agent-langgraph/.claude/skills/discover-tools/SKILL.md index eabac2a5..87c3f519 100644 --- a/agent-langgraph/.claude/skills/discover-tools/SKILL.md +++ b/agent-langgraph/.claude/skills/discover-tools/SKILL.md @@ -39,44 +39,9 @@ uv run discover-tools --profile DEFAULT | **Custom MCP Servers** | Apps starting with `mcp-*` | `{app_url}/mcp` | | **External MCP Servers** | Via UC connections | `{host}/api/2.0/mcp/external/{connection_name}` | -## Using Discovered Tools in Code - -After discovering tools, add them to your agent in `agent_server/agent.py`: - -```python -from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient - -# Example: Add UC functions from a schema -uc_functions_server = DatabricksMCPServer( - url=f"{host}/api/2.0/mcp/functions/{catalog}/{schema}", - name="my uc functions", -) - -# Example: Add a Genie space -genie_server = DatabricksMCPServer( - url=f"{host}/api/2.0/mcp/genie/{space_id}", - name="my genie space", -) - -# Example: Add vector search -vector_server = DatabricksMCPServer( - url=f"{host}/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name}", - name="my vector index", -) - -# Create MCP client with all servers -mcp_client = DatabricksMultiServerMCPClient([ - uc_functions_server, - genie_server, - vector_server, -]) - -# Get tools for the agent -tools = await mcp_client.get_tools() -``` - ## Next Steps -After adding MCP servers to your agent: -1. **Grant permissions** in `databricks.yml` (see **add-tools** skill) -2. Test locally with `uv run start-app` (see **run-locally** skill) +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 diff --git a/agent-langgraph/AGENTS.md b/agent-langgraph/AGENTS.md index 0cafbcb7..fcac7dcd 100644 --- a/agent-langgraph/AGENTS.md +++ b/agent-langgraph/AGENTS.md @@ -2,19 +2,24 @@ ## MANDATORY First Actions -**BEFORE any actions, ask for crucial follow-ups about the app deployment specifics (do this interactively!)** +**Ask the user interactively:** -1. ALWAYS make sure to ask the user if they have an existing app name that they want to deploy to and ask for them to provide it to you first interactively! -2. (MEMORY ONLY OPTIONAL) ALWAYS make sure to ask the user if they have an existing lakebase instance that they want to use for memory and provide it to you first interactively! +1. **App deployment target:** + > "Do you have an existing Databricks app you want to deploy to, or should we create a new one? If existing, what's the app name?" -**BEFORE any other action, run `databricks auth profiles` to check authentication status.** + *Note: New apps should use the `agent-*` prefix (e.g., `agent-data-analyst`) unless the user specifies otherwise.* + +2. **If the user mentions memory, conversation history, or persistence:** + > "For memory capabilities, do you have an existing Lakebase instance? If so, what's the instance name?" + +**Then check authentication status by running `databricks auth profiles`.** This helps you understand: - Which Databricks profiles are configured - Whether authentication is already set up - Which profile to use for subsequent commands -If no profiles exist, guide the user through running `uv run quickstart` to set up authentication. +If no profiles exist or `.env` is missing, guide the user through running `uv run quickstart` to set up authentication and configuration. See the **quickstart** skill for details. ## Understanding User Goals diff --git a/agent-non-conversational/.claude/skills/add-tools/SKILL.md b/agent-non-conversational/.claude/skills/add-tools/SKILL.md new file mode 100644 index 00000000..7719f198 --- /dev/null +++ b/agent-non-conversational/.claude/skills/add-tools/SKILL.md @@ -0,0 +1,104 @@ +--- +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 + +**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. + +## Workflow + +**Step 1:** Add MCP server in `agent_server/agent.py`: +```python +from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient + +genie_server = DatabricksMCPServer( + url=f"{host}/api/2.0/mcp/genie/01234567-89ab-cdef", + name="my genie space", +) + +mcp_client = DatabricksMultiServerMCPClient([genie_server]) +tools = await mcp_client.get_tools() +``` + +**Step 2:** Grant access in `databricks.yml`: +```yaml +resources: + apps: + agent_langgraph: + 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_langgraph # 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 | Agent memory storage | +| `experiment.yaml` | MLflow experiment | Tracing (already configured) | +| `custom-mcp-server.md` | Custom MCP apps | Apps starting with `mcp-*` | + +## Custom MCP Servers (Databricks Apps) + +Apps are **not yet supported** as resource dependencies in `databricks.yml`. Manual permission grant required: + +**Step 1:** Get your agent app's service principal: +```bash +databricks apps get --output json | jq -r '.service_principal_name' +``` + +**Step 2:** Grant permission on the MCP server app: +```bash +databricks apps update-permissions \ + --service-principal \ + --permission-level CAN_USE +``` + +See `examples/custom-mcp-server.md` for detailed steps. + +## valueFrom Pattern (for app.yaml) + +**IMPORTANT**: Make sure all `valueFrom` references in `app.yaml` reference an existing key in the `databricks.yml` file. +Some resources need environment variables in your app. Use `valueFrom` in `app.yaml` to reference resources defined in `databricks.yml`: + +```yaml +# app.yaml +env: + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" # References resources.apps..resources[name='experiment'] + - name: LAKEBASE_INSTANCE_NAME + valueFrom: "database" # References resources.apps..resources[name='database'] +``` + +**Critical:** Every `valueFrom` value must match a `name` field in `databricks.yml` resources. + +## 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_langgraph` +- **valueFrom matching**: Ensure `app.yaml` `valueFrom` values match `databricks.yml` resource `name` values diff --git a/agent-non-conversational/.claude/skills/add-tools/examples/custom-mcp-server.md b/agent-non-conversational/.claude/skills/add-tools/examples/custom-mcp-server.md new file mode 100644 index 00000000..1324e6c5 --- /dev/null +++ b/agent-non-conversational/.claude/skills/add-tools/examples/custom-mcp-server.md @@ -0,0 +1,57 @@ +# Custom MCP Server (Databricks App) + +Custom MCP servers are Databricks Apps with names starting with `mcp-*`. + +**Apps are not yet supported as resource dependencies in `databricks.yml`**, so manual permission grant is required. + +## 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. Deploy your agent app first + +```bash +databricks bundle deploy +databricks bundle run agent_langgraph +``` + +### 3. Get your agent app's service principal + +```bash +databricks apps get --output json | jq -r '.service_principal_name' +``` + +Example output: `sp-abc123-def456` + +### 4. Grant permission on the MCP server app + +```bash +databricks apps update-permissions \ + --service-principal \ + --permission-level CAN_USE +``` + +Example: +```bash +databricks apps update-permissions mcp-my-server \ + --service-principal sp-abc123-def456 \ + --permission-level CAN_USE +``` + +## Notes + +- This manual step is required each time you connect to a new custom MCP server +- The permission grant persists across deployments +- If you redeploy the agent app with a new service principal, you'll need to grant permissions again diff --git a/agent-non-conversational/.claude/skills/add-tools/examples/experiment.yaml b/agent-non-conversational/.claude/skills/add-tools/examples/experiment.yaml new file mode 100644 index 00000000..ac5c626a --- /dev/null +++ b/agent-non-conversational/.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-non-conversational/.claude/skills/add-tools/examples/genie-space.yaml b/agent-non-conversational/.claude/skills/add-tools/examples/genie-space.yaml new file mode 100644 index 00000000..71589d52 --- /dev/null +++ b/agent-non-conversational/.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-non-conversational/.claude/skills/add-tools/examples/lakebase.yaml b/agent-non-conversational/.claude/skills/add-tools/examples/lakebase.yaml new file mode 100644 index 00000000..78f0bc72 --- /dev/null +++ b/agent-non-conversational/.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: valueFrom reference in app.yaml + +# In databricks.yml - add to resources.apps..resources: +- name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +# In app.yaml - add to env: +# - name: LAKEBASE_INSTANCE_NAME +# valueFrom: "database" +# - name: EMBEDDING_ENDPOINT +# value: "databricks-gte-large-en" +# - name: EMBEDDING_DIMS +# value: "1024" diff --git a/agent-non-conversational/.claude/skills/add-tools/examples/serving-endpoint.yaml b/agent-non-conversational/.claude/skills/add-tools/examples/serving-endpoint.yaml new file mode 100644 index 00000000..b49ce9da --- /dev/null +++ b/agent-non-conversational/.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-non-conversational/.claude/skills/add-tools/examples/sql-warehouse.yaml b/agent-non-conversational/.claude/skills/add-tools/examples/sql-warehouse.yaml new file mode 100644 index 00000000..a6ce9446 --- /dev/null +++ b/agent-non-conversational/.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-non-conversational/.claude/skills/add-tools/examples/uc-connection.yaml b/agent-non-conversational/.claude/skills/add-tools/examples/uc-connection.yaml new file mode 100644 index 00000000..316675fe --- /dev/null +++ b/agent-non-conversational/.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-non-conversational/.claude/skills/add-tools/examples/uc-function.yaml b/agent-non-conversational/.claude/skills/add-tools/examples/uc-function.yaml new file mode 100644 index 00000000..43f938a9 --- /dev/null +++ b/agent-non-conversational/.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-non-conversational/.claude/skills/add-tools/examples/vector-search.yaml b/agent-non-conversational/.claude/skills/add-tools/examples/vector-search.yaml new file mode 100644 index 00000000..0ba39027 --- /dev/null +++ b/agent-non-conversational/.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-non-conversational/.claude/skills/agent-memory/SKILL.md b/agent-non-conversational/.claude/skills/agent-memory/SKILL.md new file mode 100644 index 00000000..83d59668 --- /dev/null +++ b/agent-non-conversational/.claude/skills/agent-memory/SKILL.md @@ -0,0 +1,524 @@ +--- +name: agent-memory +description: "Add memory capabilities to your agent. Use when: (1) User asks about 'memory', 'state', 'remember', 'conversation history', (2) Want to persist conversations or user preferences, (3) Adding checkpointing or long-term storage." +--- + +# Adding Memory to Your Agent + +> **Note:** This template does not include memory by default. Use this skill to **add memory capabilities**. For pre-configured memory templates, see: +> - `agent-langgraph-short-term-memory` - Conversation history within a session +> - `agent-langgraph-long-term-memory` - User facts that persist across sessions + +## Memory Types + +| Type | Use Case | Storage | Identifier | +|------|----------|---------|------------| +| **Short-term** | Conversation history within a session | `AsyncCheckpointSaver` | `thread_id` | +| **Long-term** | User facts that persist across sessions | `AsyncDatabricksStore` | `user_id` | + +## Prerequisites + +1. **Add memory dependency** to `pyproject.toml`: + ```toml + dependencies = [ + "databricks-langchain[memory]", + ] + ``` + + Then run `uv sync` + +2. **Configure Lakebase** - See **lakebase-setup** skill for: + - Creating/configuring Lakebase instance + - Initializing tables (CRITICAL first-time step) + +--- + +## Quick Setup Summary + +Adding memory requires changes to **4 files**: + +| File | What to Add | +|------|-------------| +| `pyproject.toml` | Memory dependency + hatch config | +| `.env` | Lakebase env vars (for local dev) | +| `databricks.yml` | Lakebase database resource | +| `app.yaml` | `valueFrom` reference to lakebase resource | +| `agent_server/agent.py` | Memory tools and AsyncDatabricksStore | + +--- + +## Step 1: Configure databricks.yml (Lakebase Resource) + +Add the Lakebase database resource to your app in `databricks.yml`: + +```yaml +resources: + apps: + agent_langgraph: + name: "your-app-name" + source_code_path: ./ + + resources: + # ... other resources (experiment, UC functions, etc.) ... + + # Lakebase instance for long-term memory + - name: 'lakebadatabasese_memory' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' +``` + +**Important:** The `name: 'database'` must match the `valueFrom` reference in `app.yaml`. + +--- + +## Step 2: Configure app.yaml (Environment Variables) + +Update `app.yaml` with the Lakebase environment variables: + +```yaml +command: ["uv", "run", "start-app"] + +env: + # ... other env vars ... + + # MLflow experiment (uses valueFrom to reference databricks.yml resource) + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" + + # Lakebase instance name - must match the instance_name in databricks.yml database resource + # Note: Use 'value' (not 'valueFrom') because AsyncDatabricksStore needs the instance name, + # not the full connection string that valueFrom would provide + - name: LAKEBASE_INSTANCE_NAME + value: "" + + # Embedding configuration (static values) + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` + +**Important:** The `LAKEBASE_INSTANCE_NAME` value must match the `instance_name` in your `databricks.yml` database resource. The `database` resource handles permissions, while `app.yaml` provides the instance name to your code. + +### Why not valueFrom for Lakebase? + +The `database` resource's `valueFrom` provides the full connection string (e.g., `instance-xxx.database.staging.cloud.databricks.com`), but `AsyncDatabricksStore` expects just the instance name. So we use a static `value` instead. + +--- + +## Step 3: Configure .env (Local Development) + +For local development, add to `.env`: + +```bash +# Lakebase configuration for long-term memory +LAKEBASE_INSTANCE_NAME= +EMBEDDING_ENDPOINT=databricks-gte-large-en +EMBEDDING_DIMS=1024 +``` + +> **Note:** `.env` is only for local development. When deployed, the app gets `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. + +--- + +## Step 4: Update agent.py + +### Add Imports and Configuration + +```python +import os +import uuid +from typing import AsyncGenerator + +from databricks_langchain import AsyncDatabricksStore, ChatDatabricks +from langchain_core.runnables import RunnableConfig +from langchain_core.tools import tool +from langgraph.prebuilt import create_react_agent + +# Environment configuration +LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") +EMBEDDING_ENDPOINT = os.getenv("EMBEDDING_ENDPOINT", "databricks-gte-large-en") +EMBEDDING_DIMS = int(os.getenv("EMBEDDING_DIMS", "1024")) # REQUIRED! +``` + +**CRITICAL:** You MUST specify `embedding_dims` when using `embedding_endpoint`. Without it, you'll get: +``` +"embedding_dims is required when embedding_endpoint is specified" +``` + +### Create Memory Tools + +```python +@tool +async def get_user_memory(query: str, config: RunnableConfig) -> str: + """Search for relevant information about the user from long-term memory. + Use this to recall preferences, past interactions, or other saved information. + """ + user_id = config.get("configurable", {}).get("user_id") + store = config.get("configurable", {}).get("store") + if not user_id or not store: + return "Memory not available - no user context" + + # Sanitize user_id for namespace (replace special chars) + namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) + results = await store.asearch(namespace, query=query, limit=5) + if not results: + return "No memories found for this query" + return "\n".join([f"- {r.value}" for r in results]) + + +@tool +async def save_user_memory(memory_key: str, memory_data: str, config: RunnableConfig) -> str: + """Save information about the user to long-term memory. + Use this to remember user preferences, important details, or other + information that should persist across conversations. + + Args: + memory_key: A short descriptive key (e.g., "preferred_name", "team", "region") + memory_data: The information to save + """ + user_id = config.get("configurable", {}).get("user_id") + store = config.get("configurable", {}).get("store") + if not user_id or not store: + return "Memory not available - no user context" + + namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) + await store.aput(namespace, memory_key, {"value": memory_data}) + return f"Saved memory: {memory_key}" + + +@tool +async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: + """Delete a specific memory from the user's long-term memory. + Use this when the user asks to forget something or correct stored information. + + Args: + memory_key: The key of the memory to delete (e.g., "preferred_name", "team") + """ + user_id = config.get("configurable", {}).get("user_id") + store = config.get("configurable", {}).get("store") + if not user_id or not store: + return "Memory not available - no user context" + + namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) + await store.adelete(namespace, memory_key) + return f"Deleted memory: {memory_key}" +``` + +### Add Helper Functions + +```python +def _get_user_id(request: ResponsesAgentRequest) -> str: + """Extract user_id from request context or custom inputs.""" + custom_inputs = dict(request.custom_inputs or {}) + if "user_id" in custom_inputs and custom_inputs["user_id"]: + return str(custom_inputs["user_id"]) + if request.context and getattr(request.context, "user_id", None): + return str(request.context.user_id) + return "default-user" + + +def _get_or_create_thread_id(request: ResponsesAgentRequest) -> str: + """Extract or create thread_id for conversation tracking.""" + custom_inputs = dict(request.custom_inputs or {}) + if "thread_id" in custom_inputs and custom_inputs["thread_id"]: + return str(custom_inputs["thread_id"]) + if request.context and getattr(request.context, "conversation_id", None): + return str(request.context.conversation_id) + return str(uuid.uuid4()) +``` + +### Update Streaming Function + +```python +@stream() +async def streaming( + request: ResponsesAgentRequest, +) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + # Get user context + user_id = _get_user_id(request) + thread_id = _get_or_create_thread_id(request) + + # Get other tools (MCP, etc.) + mcp_client = init_mcp_client(sp_workspace_client) + mcp_tools = await mcp_client.get_tools() + + # Memory tools + memory_tools = [get_user_memory, save_user_memory, delete_user_memory] + + # Combine all tools + all_tools = mcp_tools + memory_tools + + # Initialize model + model = ChatDatabricks(endpoint="databricks-claude-sonnet-4") + + # Use AsyncDatabricksStore for long-term memory + async with AsyncDatabricksStore( + instance_name=LAKEBASE_INSTANCE_NAME, + embedding_endpoint=EMBEDDING_ENDPOINT, + embedding_dims=EMBEDDING_DIMS, # REQUIRED! + ) as store: + # Create agent with tools + agent = create_react_agent( + model=model, + tools=all_tools, + prompt=AGENT_INSTRUCTIONS, + ) + + # Prepare input + messages = {"messages": to_chat_completions_input([i.model_dump() for i in request.input])} + + # Configure with user context for memory tools + config = { + "configurable": { + "user_id": user_id, + "thread_id": thread_id, + "store": store, # Pass store to tools via config + } + } + + # Stream agent responses + async for event in process_agent_astream_events( + agent.astream(input=messages, config=config, stream_mode=["updates", "messages"]) + ): + yield event +``` + +--- + +## Step 5: Initialize Tables and Deploy + +### Initialize Lakebase Tables (First Time Only) + +Before deploying, initialize the tables locally: + +```bash +uv run python -c "$(cat <<'EOF' +import asyncio +from databricks_langchain import AsyncDatabricksStore + +async def setup(): + async with AsyncDatabricksStore( + instance_name="", + embedding_endpoint="databricks-gte-large-en", + embedding_dims=1024, + ) as store: + await store.setup() + print("Tables created!") + +asyncio.run(setup()) +EOF +)" +``` + +### Deploy and Run + +**IMPORTANT:** Always run `databricks bundle run` after `databricks bundle deploy` to start/restart the app with the new code: + +```bash +# Deploy resources and upload files +databricks bundle deploy + +# Start/restart the app with new code (REQUIRED!) +databricks bundle run agent_langgraph +``` + +> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is required to actually start the app with the new code. + +--- + +## Complete Example Files + +### databricks.yml (with Lakebase) + +```yaml +bundle: + name: agent_langgraph + +resources: + experiments: + agent_langgraph_experiment: + name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} + + apps: + agent_langgraph: + name: "my-agent-app" + description: "Agent with long-term memory" + source_code_path: ./ + + resources: + - name: 'experiment' + experiment: + experiment_id: "${resources.experiments.agent_langgraph_experiment.id}" + permission: 'CAN_MANAGE' + + # Lakebase instance for long-term memory + - name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +targets: + dev: + mode: development + default: true +``` + +### app.yaml + +```yaml +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" + # Reference experiment resource from databricks.yml + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" + # Lakebase instance name (must match instance_name in databricks.yml) + - name: LAKEBASE_INSTANCE_NAME + value: "" + # Embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` + +--- + +## Adding Short-Term Memory + +Short-term memory stores conversation history within a thread. + +### Import and Initialize Checkpointer + +```python +from databricks_langchain import AsyncCheckpointSaver + +LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") + +async with AsyncCheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME) as checkpointer: + agent = create_react_agent( + model=model, + tools=tools, + checkpointer=checkpointer, # Enables conversation persistence + ) +``` + +### Use thread_id in Agent Config + +```python +config = {"configurable": {"thread_id": thread_id}} +async for event in agent.astream(input_state, config, stream_mode=["updates", "messages"]): + yield event +``` + +--- + +## Testing Memory + +### Test Locally + +```bash +# Start the server +uv run start-app + +# Save a memory +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Remember that I am on the shipping team"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' + +# Recall the memory +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "What team am I on?"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' + +# Delete a memory +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Forget what team I am on"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' +``` + +### Test Deployed App + +```bash +# Get OAuth token (PATs don't work for apps) +TOKEN=$(databricks auth token --host | jq -r '.access_token') + +# Test memory save +curl -X POST https:///invocations \ + -H "Authorization: Bearer $TOKEN" \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Remember I prefer detailed explanations"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' +``` + +--- + +## First-Time Setup Checklist + +- [ ] Added `databricks-langchain[memory]` to `pyproject.toml` +- [ ] Run `uv sync` to install dependencies +- [ ] Created or identified Lakebase instance +- [ ] Added Lakebase env vars to `.env` (for local dev) +- [ ] Added `database` resource to `databricks.yml` (for permissions) +- [ ] Added `LAKEBASE_INSTANCE_NAME` value to `app.yaml` (matching instance_name in databricks.yml) +- [ ] **Initialized tables locally** by running `await store.setup()` +- [ ] Deployed with `databricks bundle deploy` +- [ ] **Started app with `databricks bundle run agent_langgraph`** + +--- + +## Troubleshooting + +| Issue | Cause | Solution | +|-------|-------|----------| +| **"embedding_dims is required"** | Missing parameter | Add `embedding_dims=1024` to AsyncDatabricksStore | +| **"relation 'store' does not exist"** | Tables not created | Run `await store.setup()` locally first | +| **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `valueFrom: "database"` to app.yaml | +| **"permission denied for table store"** | Missing grants | Lakebase `database` resource in DAB should handle this | +| **"Memory not available"** | No user_id in request | Ensure `custom_inputs.user_id` is passed | +| **Memory not persisting** | Different user_ids | Use consistent user_id across requests | +| **App not updated after deploy** | Forgot to run bundle | Run `databricks bundle run agent_langgraph` after deploy | + +--- + +## Pre-Built Memory Templates + +For fully configured implementations without manual setup: + +| Template | Memory Type | Key Features | +|----------|-------------|--------------| +| `agent-langgraph-short-term-memory` | Short-term | AsyncCheckpointSaver, thread_id | +| `agent-langgraph-long-term-memory` | Long-term | AsyncDatabricksStore, memory tools | + +--- + +## Next Steps + +- Configure Lakebase: see **lakebase-setup** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/agent-non-conversational/.claude/skills/deploy/SKILL.md b/agent-non-conversational/.claude/skills/deploy/SKILL.md new file mode 100644 index 00000000..dd0be7e1 --- /dev/null +++ b/agent-non-conversational/.claude/skills/deploy/SKILL.md @@ -0,0 +1,222 @@ +--- +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 + +## 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_non_conversational: + name: "agent-your-app-name" # Use agent-* prefix +``` + +## Deploy Commands + +**IMPORTANT:** Always run BOTH commands to deploy and start your app: + +```bash +# 1. Validate bundle configuration (catches errors before deploy) +databricks bundle validate + +# 2. Deploy the bundle (creates/updates resources, uploads files) +databricks bundle deploy + +# 3. Run the app (starts/restarts with uploaded source code) - REQUIRED! +databricks bundle run agent_non_conversational +``` + +> **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_non_conversational` 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_non_conversational: + 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_non_conversational --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_non_conversational +``` + +### 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_non_conversational +``` + +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 +``` + +**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' +``` + +## 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_non_conversational` after deploy | +| Env var is None in deployed app | Check `valueFrom` in app.yaml matches resource `name` in databricks.yml | diff --git a/agent-non-conversational/.claude/skills/discover-tools/SKILL.md b/agent-non-conversational/.claude/skills/discover-tools/SKILL.md new file mode 100644 index 00000000..87c3f519 --- /dev/null +++ b/agent-non-conversational/.claude/skills/discover-tools/SKILL.md @@ -0,0 +1,47 @@ +--- +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 diff --git a/agent-non-conversational/.claude/skills/lakebase-setup/SKILL.md b/agent-non-conversational/.claude/skills/lakebase-setup/SKILL.md new file mode 100644 index 00000000..6b8c802e --- /dev/null +++ b/agent-non-conversational/.claude/skills/lakebase-setup/SKILL.md @@ -0,0 +1,351 @@ +--- +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 Memory + +> **Note:** This template does not include memory by default. Use this skill if you want to **add memory capabilities** to your agent. For pre-configured memory templates, see: +> - `agent-langgraph-short-term-memory` - Conversation history within a session +> - `agent-langgraph-long-term-memory` - User facts that persist across sessions + +## Overview + +Lakebase provides persistent storage for agent memory: +- **Short-term memory**: Conversation history within a thread (`AsyncCheckpointSaver`) +- **Long-term memory**: User facts across sessions (`AsyncDatabricksStore`) + +## Complete Setup Workflow + +``` +┌─────────────────────────────────────────────────────────────────────────────┐ +│ 1. Add dependency → 2. Get instance → 3. Configure DAB + app.yaml │ +│ 4. Configure .env → 5. Initialize tables → 6. Deploy + Run │ +└─────────────────────────────────────────────────────────────────────────────┘ +``` + +--- + +## 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: Create New Instance (via Databricks UI) + +1. Go to your Databricks workspace +2. Navigate to **Compute** → **Lakebase** +3. Click **Create Instance** +4. Note the instance name + +### Option B: Use Existing Instance + +If you have an existing instance, note its name for the next step. + +--- + +## Step 3: Configure databricks.yml (Lakebase Resource) + +Add the Lakebase `database` resource to your app in `databricks.yml`: + +```yaml +resources: + apps: + agent_langgraph: + 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: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' +``` + +**Important:** +- The `name: 'database'` must match the `valueFrom` reference in `app.yaml` +- Using the `database` resource type automatically grants the app's service principal access to Lakebase + +### Update app.yaml (Environment Variables) + +Update `app.yaml` with the Lakebase instance name: + +```yaml +env: + # ... other env vars ... + + # Lakebase instance name - must match instance_name in databricks.yml database resource + # Note: Use 'value' (not 'valueFrom') because AsyncDatabricksStore needs the instance name, + # not the full connection string that valueFrom would provide + - name: LAKEBASE_INSTANCE_NAME + value: "" + + # Static values for embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` + +**Important:** +- The `LAKEBASE_INSTANCE_NAME` value must match the `instance_name` in your `databricks.yml` database resource +- The `database` resource handles permissions; `app.yaml` provides the instance name to your code +- Don't use `valueFrom` for Lakebase - it provides the connection string, not the instance name + +--- + +## Step 4: Configure .env (Local Development) + +For local development, add to `.env`: + +```bash +# Lakebase configuration for long-term memory +LAKEBASE_INSTANCE_NAME= +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 `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. + +--- + +## Step 5: Initialize Store Tables (CRITICAL - First Time Only) + +**Before deploying**, you must initialize the Lakebase tables. The `AsyncDatabricksStore` creates tables on first use, but you need to do this locally first: + +```python +# Run this script locally BEFORE first deployment +import asyncio +from databricks_langchain import AsyncDatabricksStore + +async def setup_store(): + async with AsyncDatabricksStore( + instance_name="", + embedding_endpoint="databricks-gte-large-en", + embedding_dims=1024, + ) as store: + print("Setting up store tables...") + await store.setup() # Creates required tables + print("Store tables created!") + + # Verify with a test write/read + await store.aput(("test", "init"), "test_key", {"value": "test_value"}) + results = await store.asearch(("test", "init"), query="test", limit=1) + print(f"Test successful: {results}") + +asyncio.run(setup_store()) +``` + +Run with: +```bash +uv run python -c "$(cat <<'EOF' +import asyncio +from databricks_langchain import AsyncDatabricksStore + +async def setup(): + async with AsyncDatabricksStore( + instance_name="", + embedding_endpoint="databricks-gte-large-en", + embedding_dims=1024, + ) as store: + await store.setup() + print("Tables created!") + +asyncio.run(setup()) +EOF +)" +``` + +This creates these tables in the `public` schema: +- `store` - Key-value storage for memories +- `store_vectors` - Vector embeddings for semantic search +- `store_migrations` - Schema migration tracking +- `vector_migrations` - Vector schema migration tracking + +--- + +## Step 6: Deploy and Run Your App + +**IMPORTANT:** Always run both `deploy` AND `run` commands: + +```bash +# Deploy resources and upload files +databricks bundle deploy + +# Start/restart the app with new code (REQUIRED!) +databricks bundle run agent_langgraph +``` + +> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is required to actually start the app with the new code. + +--- + +## Complete Example: databricks.yml with Lakebase + +```yaml +bundle: + name: agent_langgraph + +resources: + experiments: + agent_langgraph_experiment: + name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} + + apps: + agent_langgraph: + name: "my-agent-app" + description: "Agent with long-term memory" + source_code_path: ./ + + resources: + - name: 'experiment' + experiment: + experiment_id: "${resources.experiments.agent_langgraph_experiment.id}" + permission: 'CAN_MANAGE' + + # Lakebase instance for long-term memory + - name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' + +targets: + dev: + mode: development + default: true +``` + +## Complete Example: app.yaml + +```yaml +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" + # Reference experiment resource from databricks.yml + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" + # Lakebase instance name (must match instance_name in databricks.yml) + - name: LAKEBASE_INSTANCE_NAME + value: "" + # Embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` + +--- + +## 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 | Run `await store.setup()` locally first (Step 5) | +| **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `LAKEBASE_INSTANCE_NAME` value to app.yaml | +| **"Unable to resolve Lakebase instance '...database.cloud.databricks.com'"** | Used valueFrom instead of value | Use `value: ""` not `valueFrom` for Lakebase | +| **"permission denied for table store"** | Missing grants | The `database` resource in DAB should handle this; verify the resource is configured | +| **"Failed to connect to Lakebase"** | Wrong instance name | Verify instance name 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 agent_langgraph` after deploy | +| **valueFrom not resolving** | Resource name mismatch | Ensure `valueFrom` value matches `name` in databricks.yml resources | + +--- + +## Quick Reference: LakebaseClient API + +For manual permission management (usually not needed with DAB `database` resource): + +```python +from databricks_ai_bridge.lakebase import LakebaseClient, SchemaPrivilege, TablePrivilege + +client = LakebaseClient(instance_name="...") + +# 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 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-non-conversational/.claude/skills/modify-agent/SKILL.md b/agent-non-conversational/.claude/skills/modify-agent/SKILL.md new file mode 100644 index 00000000..d7218637 --- /dev/null +++ b/agent-non-conversational/.claude/skills/modify-agent/SKILL.md @@ -0,0 +1,293 @@ +--- +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 + +## Main File + +**`agent_server/agent.py`** - Agent logic, model selection, instructions, MCP servers + +## Key Files + +| File | Purpose | +|------|---------| +| `agent_server/agent.py` | Agent logic, model, instructions, MCP servers | +| `agent_server/start_server.py` | FastAPI server + MLflow setup | +| `agent_server/evaluate_agent.py` | Agent evaluation with MLflow scorers | +| `agent_server/utils.py` | Databricks auth helpers, stream processing | +| `databricks.yml` | Bundle config & resource permissions | + +## SDK Setup + +```python +import mlflow +from databricks.sdk import WorkspaceClient +from databricks_langchain import ChatDatabricks, DatabricksMCPServer, DatabricksMultiServerMCPClient +from langchain.agents import create_agent + +# Enable autologging for tracing +mlflow.langchain.autolog() + +# Initialize workspace client +workspace_client = WorkspaceClient() +``` + +--- + +## databricks-langchain SDK Overview + +**SDK Location:** https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/langchain + +Before making any changes, ensure that the APIs actually exist in the SDK. If something is missing from the documentation here, look in the venv's `site-packages` directory for the `databricks_langchain` package. If it's not installed, run `uv sync` to create the .venv and install the package. + +--- + +### ChatDatabricks - LLM Chat Interface + +Connects to Databricks Model Serving endpoints for LLM inference. + +```python +from databricks_langchain import ChatDatabricks + +llm = ChatDatabricks( + endpoint="databricks-claude-3-7-sonnet", # or databricks-meta-llama-3-1-70b-instruct + temperature=0, + max_tokens=500, +) + +# For Responses API agents: +llm = ChatDatabricks(endpoint="my-agent-endpoint", use_responses_api=True) +``` + +Available models (check workspace for current list): +- `databricks-claude-3-7-sonnet` +- `databricks-claude-3-5-sonnet` +- `databricks-meta-llama-3-3-70b-instruct` + +**Note:** Some workspaces require granting the app access to the serving endpoint in `databricks.yml`. See the **add-tools** skill and `examples/serving-endpoint.yaml`. + +--- + +### DatabricksEmbeddings - Generate Embeddings + +Query Databricks embedding model endpoints. + +```python +from databricks_langchain import DatabricksEmbeddings + +embeddings = DatabricksEmbeddings(endpoint="databricks-bge-large-en") +vector = embeddings.embed_query("The meaning of life is 42") +vectors = embeddings.embed_documents(["doc1", "doc2"]) +``` + +--- + +### DatabricksVectorSearch - Vector Store + +Connect to Databricks Vector Search indexes for similarity search. + +```python +from databricks_langchain import DatabricksVectorSearch + +# Delta-sync index with Databricks-managed embeddings +vs = DatabricksVectorSearch(index_name="catalog.schema.index_name") + +# Direct-access or self-managed embeddings +vs = DatabricksVectorSearch( + index_name="catalog.schema.index_name", + embedding=embeddings, + text_column="content", +) + +docs = vs.similarity_search("query", k=5) +``` + +--- + +### MCP Client - Tool Integration + +Connect to MCP (Model Context Protocol) servers to get tools for your agent. + +**Basic MCP Server (manual URL):** + +```python +from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient + +client = DatabricksMultiServerMCPClient([ + DatabricksMCPServer( + name="system-ai", + url=f"{host}/api/2.0/mcp/functions/system/ai", + ) +]) +tools = await client.get_tools() +``` + +**From UC Function (convenience helper):** + +Creates MCP server for Unity Catalog functions. If `function_name` is omitted, exposes all functions in the schema. + +```python +server = DatabricksMCPServer.from_uc_function( + catalog="main", + schema="tools", + function_name="send_email", # Optional - omit for all functions in schema + name="email-server", + timeout=30.0, + handle_tool_error=True, +) +``` + +**From Vector Search (convenience helper):** + +Creates MCP server for Vector Search indexes. If `index_name` is omitted, exposes all indexes in the schema. + +```python +server = DatabricksMCPServer.from_vector_search( + catalog="main", + schema="embeddings", + index_name="product_docs", # Optional - omit for all indexes in schema + name="docs-search", + timeout=30.0, +) +``` + +**From Genie Space:** + +Create MCP server from Genie Space. Get the genie space ID from the URL. + +Example: `https://workspace.cloud.databricks.com/genie/rooms/01f0515f6739169283ef2c39b7329700?o=123` means the genie space ID is `01f0515f6739169283ef2c39b7329700` + +```python +DatabricksMCPServer( + name="genie", + url=f"{host_name}/api/2.0/mcp/genie/01f0515f6739169283ef2c39b7329700", +) +``` + +**Non-Databricks MCP Server:** + +```python +from databricks_langchain import MCPServer + +server = MCPServer( + name="external-server", + url="https://other-server.com/mcp", + headers={"X-API-Key": "secret"}, + timeout=15.0, +) +``` + +**After adding MCP servers:** Grant permissions in `databricks.yml` (see **add-tools** skill) + +--- + +## Running the Agent + +```python +from langchain.agents import create_agent + +# Create agent - ONLY accepts tools and model, NO prompt/instructions parameter +agent = create_agent(tools=tools, model=llm) + +# Non-streaming +messages = {"messages": [{"role": "user", "content": "hi"}]} +result = await agent.ainvoke(messages) + +# Streaming +async for event in agent.astream(input=messages, stream_mode=["updates", "messages"]): + # Process stream events + pass +``` + +**Converting to Responses API format:** Use `process_agent_astream_events()` from `agent_server/utils.py`: + +```python +from agent_server.utils import process_agent_astream_events + +async for event in process_agent_astream_events( + agent.astream(input=messages, stream_mode=["updates", "messages"]) +): + yield event # Yields ResponsesAgentStreamEvent objects +``` + +--- + +## Customizing Agent Behavior (System Instructions) + +> **IMPORTANT:** `create_agent()` does NOT accept `prompt`, `instructions`, or `system_message` parameters. Attempting to pass these will cause a runtime error. + +In LangGraph, agent behavior is customized by prepending a system message to the conversation messages. + +**Correct pattern in `agent.py`:** + +1. Define instructions as a constant: +```python +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.""" +``` + +2. Prepend to messages in the `streaming()` function: +```python +@stream() +async def streaming(request: ResponsesAgentRequest) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + agent = await init_agent() + # Prepend system instructions to user messages + user_messages = to_chat_completions_input([i.model_dump() for i in request.input]) + messages = {"messages": [{"role": "system", "content": AGENT_INSTRUCTIONS}] + user_messages} + + async for event in process_agent_astream_events( + agent.astream(input=messages, stream_mode=["updates", "messages"]) + ): + yield event +``` + +**Common mistake to avoid:** +```python +# WRONG - will cause "unexpected keyword argument" error +agent = create_agent(tools=tools, model=llm, prompt=AGENT_INSTRUCTIONS) + +# CORRECT - add instructions via messages +messages = {"messages": [{"role": "system", "content": AGENT_INSTRUCTIONS}] + user_messages} +``` + +For advanced customization (routing, state management, custom graphs), refer to the [LangGraph documentation](https://docs.langchain.com/oss/python/langgraph/overview). + +--- + +## External Connection Tools + +Connect to external services via Unity Catalog HTTP connections: + +- **Slack** - Post messages to channels +- **Google Calendar** - Calendar operations +- **Microsoft Graph API** - Office 365 services +- **Azure AI Search** - Search functionality +- **Any HTTP API** - Use `http_request` from databricks-sdk + +Example: Create UC function wrapping HTTP request for Slack, then expose via MCP. + +--- + +## External Resources + +1. [databricks-langchain SDK](https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/langchain) +2. [Agent examples](https://github.com/bbqiu/agent-on-app-prototype) +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. [LangGraph documentation](https://docs.langchain.com/oss/python/langgraph/overview) +6. [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 +- Add memory capabilities: see **agent-memory** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/agent-non-conversational/.claude/skills/quickstart/SKILL.md b/agent-non-conversational/.claude/skills/quickstart/SKILL.md new file mode 100644 index 00000000..e550162c --- /dev/null +++ b/agent-non-conversational/.claude/skills/quickstart/SKILL.md @@ -0,0 +1,83 @@ +--- +name: quickstart +description: "Set up Databricks agent development environment. Use when: (1) First time setup, (2) Configuring Databricks authentication, (3) User says 'quickstart', 'set up', 'authenticate', or 'configure databricks', (4) No .env file exists." +--- + +# Quickstart & Authentication + +## Prerequisites + +- **uv** (Python package manager) +- **nvm** with Node 20 (for frontend) +- **Databricks CLI v0.283.0+** + +Check CLI version: +```bash +databricks -v # Must be v0.283.0 or above +brew upgrade databricks # If version is too old +``` + +## Run Quickstart + +```bash +uv run quickstart +``` + +**Options:** +- `--profile NAME`: Use specified profile (non-interactive) +- `--host URL`: Workspace URL for initial setup +- `-h, --help`: Show help + +**Examples:** +```bash +# Interactive (prompts for profile selection) +uv run quickstart + +# Non-interactive with existing profile +uv run quickstart --profile DEFAULT + +# New workspace setup +uv run quickstart --host https://your-workspace.cloud.databricks.com +``` + +## What Quickstart Configures + +Creates/updates `.env` with: +- `DATABRICKS_CONFIG_PROFILE` - Selected CLI profile +- `MLFLOW_TRACKING_URI` - Set to `databricks://` for local auth +- `MLFLOW_EXPERIMENT_ID` - Auto-created experiment ID + +## Manual Authentication (Fallback) + +If quickstart fails: + +```bash +# Create new profile +databricks auth login --host https://your-workspace.cloud.databricks.com + +# Verify +databricks auth profiles +``` + +Then manually create `.env` (copy from `.env.example`): +```bash +# Authentication (choose one method) +DATABRICKS_CONFIG_PROFILE=DEFAULT +# DATABRICKS_HOST=https://.databricks.com +# DATABRICKS_TOKEN=dapi.... + +# MLflow configuration +MLFLOW_EXPERIMENT_ID= +MLFLOW_TRACKING_URI="databricks://DEFAULT" +MLFLOW_REGISTRY_URI="databricks-uc" + +# Frontend proxy settings +CHAT_APP_PORT=3000 +CHAT_PROXY_TIMEOUT_SECONDS=300 +``` + +## Next Steps + +After quickstart completes: +1. Run `uv run discover-tools` to find available workspace resources (see **discover-tools** skill) +2. Run `uv run start-app` to test locally (see **run-locally** skill) diff --git a/agent-non-conversational/.claude/skills/run-locally/SKILL.md b/agent-non-conversational/.claude/skills/run-locally/SKILL.md new file mode 100644 index 00000000..3eb83c82 --- /dev/null +++ b/agent-non-conversational/.claude/skills/run-locally/SKILL.md @@ -0,0 +1,90 @@ +--- +name: run-locally +description: "Run and test the agent locally. Use when: (1) User says 'run locally', 'start server', 'test agent', or 'localhost', (2) Need curl commands to test API, (3) Troubleshooting local development issues, (4) Configuring server options like port or hot-reload." +--- + +# Run Agent Locally + +## Start the Server + +```bash +uv run start-app +``` + +This starts the agent at http://localhost:8000 + +## Server Options + +```bash +# Hot-reload on code changes (development) +uv run start-server --reload + +# Custom port +uv run start-server --port 8001 + +# Multiple workers (production-like) +uv run start-server --workers 4 + +# Combine options +uv run start-server --reload --port 8001 +``` + +## Test the API + +**Streaming request:** +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' +``` + +**Non-streaming request:** +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }] }' +``` + +## Run Evaluation + +```bash +uv run agent-evaluate +``` + +Uses MLflow scorers (RelevanceToQuery, Safety). + +## Run Unit Tests + +```bash +pytest [path] +``` + +## Troubleshooting + +| Issue | Solution | +|-------|----------| +| **Port already in use** | Use `--port 8001` or kill existing process | +| **Authentication errors** | Verify `.env` is correct; run **quickstart** skill | +| **Module not found** | Run `uv sync` to install dependencies | +| **MLflow experiment not found** | Ensure `MLFLOW_TRACKING_URI` in `.env` is `databricks://` | + +### MLflow Experiment Not Found + +If you see: "The provided MLFLOW_EXPERIMENT_ID environment variable value does not exist" + +**Verify the experiment exists:** +```bash +databricks -p experiments get-experiment +``` + +**Fix:** Ensure `.env` has the correct tracking URI format: +```bash +MLFLOW_TRACKING_URI="databricks://DEFAULT" # Include profile name +``` + +The quickstart script configures this automatically. If you manually edited `.env`, ensure the profile name is included. + +## Next Steps + +- Modify your agent: see **modify-agent** skill +- Deploy to Databricks: see **deploy** skill diff --git a/agent-non-conversational/.gitignore b/agent-non-conversational/.gitignore index a058c65c..601fbb1a 100644 --- a/agent-non-conversational/.gitignore +++ b/agent-non-conversational/.gitignore @@ -1,8 +1,6 @@ # Created by https://www.toptal.com/developers/gitignore/api/python # Edit at https://www.toptal.com/developers/gitignore?templates=python -databricks.yml - ### Python ### # Byte-compiled / optimized / DLL files __pycache__/ @@ -204,6 +202,20 @@ sketch **/mlruns/ **/.vite/ **/.databricks -**/.claude + +# Ignore .claude directory but track template-provided skills +# User-created skills will be ignored by default (no conflicts) +.claude/* +!.claude/skills/ +.claude/skills/* +!.claude/skills/quickstart/ +!.claude/skills/discover-tools/ +!.claude/skills/deploy/ +!.claude/skills/add-tools/ +!.claude/skills/run-locally/ +!.claude/skills/modify-agent/ +!.claude/skills/lakebase-setup/ +!.claude/skills/agent-memory/ + **/.env **/.env.local \ No newline at end of file diff --git a/agent-openai-agents-sdk/.claude/skills/deploy/SKILL.md b/agent-openai-agents-sdk/.claude/skills/deploy/SKILL.md index e93000b2..d072e595 100644 --- a/agent-openai-agents-sdk/.claude/skills/deploy/SKILL.md +++ b/agent-openai-agents-sdk/.claude/skills/deploy/SKILL.md @@ -5,16 +5,38 @@ description: "Deploy agent to Databricks Apps using DAB (Databricks Asset Bundle # Deploy to Databricks Apps +## 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_openai_agents_sdk: + name: "agent-your-app-name" # Use agent-* prefix +``` + ## Deploy Commands +**IMPORTANT:** Always run BOTH commands to deploy and start your app: + ```bash -# Deploy the bundle (creates/updates resources, uploads files) +# 1. Validate bundle configuration (catches errors before deploy) +databricks bundle validate + +# 2. Deploy the bundle (creates/updates resources, uploads files) databricks bundle deploy -# Run the app (starts/restarts with uploaded source code) +# 3. Run the app (starts/restarts with uploaded source code) - REQUIRED! databricks bundle run agent_openai_agents_sdk ``` +> **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_openai_agents_sdk` matches the app name in `databricks.yml` under `resources.apps`. ## Handling "App Already Exists" Error @@ -101,7 +123,9 @@ Note: Unbinding doesn't delete the deployed app. ## Query Deployed App -**Get OAuth token** (PATs not supported): +> **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 ``` @@ -114,6 +138,33 @@ curl -X POST /invocations \ -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 @@ -134,12 +185,38 @@ databricks apps get --output json | jq -r '.url' - **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 | -| "Provider produced inconsistent result" | Sync app config to `databricks.yml`| -| "should set workspace.root_path" | Add `root_path` to production target | \ No newline at end of file +| 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_openai_agents_sdk` after deploy | +| Env var is None in deployed app | Check `valueFrom` in app.yaml matches resource `name` in databricks.yml | diff --git a/agent-openai-agents-sdk/.claude/skills/discover-tools/SKILL.md b/agent-openai-agents-sdk/.claude/skills/discover-tools/SKILL.md index fc79e8c4..87c3f519 100644 --- a/agent-openai-agents-sdk/.claude/skills/discover-tools/SKILL.md +++ b/agent-openai-agents-sdk/.claude/skills/discover-tools/SKILL.md @@ -39,42 +39,9 @@ uv run discover-tools --profile DEFAULT | **Custom MCP Servers** | Apps starting with `mcp-*` | `{app_url}/mcp` | | **External MCP Servers** | Via UC connections | `{host}/api/2.0/mcp/external/{connection_name}` | -## Using Discovered Tools in Code - -After discovering tools, add them to your agent in `agent_server/agent.py`: - -```python -from databricks_openai.agents import McpServer - -# Example: Add UC functions from a schema -uc_functions_server = McpServer( - url=f"{host}/api/2.0/mcp/functions/{catalog}/{schema}", - name="my uc functions", -) - -# Example: Add a Genie space -genie_server = McpServer( - url=f"{host}/api/2.0/mcp/genie/{space_id}", - name="my genie space", -) - -# Example: Add vector search -vector_server = McpServer( - url=f"{host}/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name}", - name="my vector index", -) - -# Add to agent -agent = Agent( - name="my agent", - instructions="You are a helpful agent.", - model="databricks-claude-3-7-sonnet", - mcp_servers=[uc_functions_server, genie_server, vector_server], -) -``` - ## Next Steps -After adding MCP servers to your agent: -1. **Grant permissions** in `databricks.yml` (see **add-tools** skill) -2. Test locally with `uv run start-app` (see **run-locally** skill) +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 diff --git a/agent-openai-agents-sdk/AGENTS.md b/agent-openai-agents-sdk/AGENTS.md index 10a6171f..ef319b9d 100644 --- a/agent-openai-agents-sdk/AGENTS.md +++ b/agent-openai-agents-sdk/AGENTS.md @@ -1,15 +1,25 @@ # Agent Development Guide -## MANDATORY First Action +## MANDATORY First Actions -**BEFORE any other action, run `databricks auth profiles` to check authentication status.** +**Ask the user interactively:** + +1. **App deployment target:** + > "Do you have an existing Databricks app you want to deploy to, or should we create a new one? If existing, what's the app name?" + + *Note: New apps should use the `agent-*` prefix (e.g., `agent-data-analyst`) unless the user specifies otherwise.* + +2. **If the user mentions memory, conversation history, or persistence:** + > "For memory capabilities, do you have an existing Lakebase instance? If so, what's the instance name?" + +**Then check authentication status by running `databricks auth profiles`.** This helps you understand: - Which Databricks profiles are configured - Whether authentication is already set up - Which profile to use for subsequent commands -If no profiles exist, guide the user through running `uv run quickstart` to set up authentication. +If no profiles exist or `.env` is missing, guide the user through running `uv run quickstart` to set up authentication and configuration. See the **quickstart** skill for details. ## Understanding User Goals From 1e77bf4678971a975dfbc6ad3eb1addcf6af4695 Mon Sep 17 00:00:00 2001 From: Dhruv Gupta Date: Mon, 2 Feb 2026 09:42:02 -0800 Subject: [PATCH 10/14] Update .claude/skills/lakebase-setup/SKILL.md Co-authored-by: Jenny --- .claude/skills/lakebase-setup/SKILL.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.claude/skills/lakebase-setup/SKILL.md b/.claude/skills/lakebase-setup/SKILL.md index 6b8c802e..65c79099 100644 --- a/.claude/skills/lakebase-setup/SKILL.md +++ b/.claude/skills/lakebase-setup/SKILL.md @@ -82,7 +82,7 @@ resources: ``` **Important:** -- The `name: 'database'` must match the `valueFrom` reference in `app.yaml` +- The `instance_name: ''` must match the `value` reference in `app.yaml` - Using the `database` resource type automatically grants the app's service principal access to Lakebase ### Update app.yaml (Environment Variables) From 867ac803f6094f026863a89a7a6b055493b352e0 Mon Sep 17 00:00:00 2001 From: Dhruv Gupta Date: Mon, 2 Feb 2026 09:52:12 -0800 Subject: [PATCH 11/14] Add .gitattributes to hide synced skills in diffs Co-Authored-By: Claude Opus 4.5 --- .gitattributes | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 .gitattributes diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 00000000..e884306c --- /dev/null +++ b/.gitattributes @@ -0,0 +1,3 @@ +# Mark synced skills in agent subfolders as generated +# These are copied from .claude/skills/ by the sync script +agent-*/.claude/skills/** linguist-generated=true From c1bf24044ca5821d563a44ad4c8f047f120733cd Mon Sep 17 00:00:00 2001 From: Dhruv Gupta Date: Mon, 2 Feb 2026 11:07:24 -0800 Subject: [PATCH 12/14] Refactor agent-memory skill: principles over static code - Rename agent-memory -> agent-langgraph-memory in parent .claude/ - Add "Key Principles" section with 5 focused patterns - Add "Production Reference" section linking to utils_memory.py - Remove full tool implementations (~150 lines of duplicated code) - Reduce skill from 515 to 370 lines (~28% reduction) - Update sync script to copy agent-langgraph-memory -> agent-memory Co-Authored-By: Claude Opus 4.5 --- .../skills/agent-langgraph-memory/SKILL.md | 375 +++++++++++++ .claude/skills/agent-memory/SKILL.md | 524 ------------------ .claude/sync-skills.py | 3 +- .../.claude/skills/agent-memory/SKILL.md | 425 +++++--------- .../.claude/skills/lakebase-setup/SKILL.md | 2 +- .../.claude/skills/agent-memory/SKILL.md | 425 +++++--------- .../.claude/skills/lakebase-setup/SKILL.md | 2 +- .../.claude/skills/agent-memory/SKILL.md | 425 +++++--------- .../.claude/skills/lakebase-setup/SKILL.md | 2 +- .../.claude/skills/agent-memory/SKILL.md | 425 +++++--------- .../.claude/skills/lakebase-setup/SKILL.md | 2 +- 11 files changed, 933 insertions(+), 1677 deletions(-) create mode 100644 .claude/skills/agent-langgraph-memory/SKILL.md delete mode 100644 .claude/skills/agent-memory/SKILL.md diff --git a/.claude/skills/agent-langgraph-memory/SKILL.md b/.claude/skills/agent-langgraph-memory/SKILL.md new file mode 100644 index 00000000..c87e21a3 --- /dev/null +++ b/.claude/skills/agent-langgraph-memory/SKILL.md @@ -0,0 +1,375 @@ +--- +name: agent-memory +description: "Add memory capabilities to your agent. Use when: (1) User asks about 'memory', 'state', 'remember', 'conversation history', (2) Want to persist conversations or user preferences, (3) Adding checkpointing or long-term storage." +--- + +# Adding Memory to Your Agent + +> **Note:** This template does not include memory by default. Use this skill to **add memory capabilities**. For pre-configured memory templates, see: +> - [agent-langgraph-short-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-short-term-memory) - Conversation history within a session +> - [agent-langgraph-long-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory) - User facts that persist across sessions + +## Memory Types + +| Type | Use Case | Storage | Identifier | +|------|----------|---------|------------| +| **Short-term** | Conversation history within a session | `AsyncCheckpointSaver` | `thread_id` | +| **Long-term** | User facts that persist across sessions | `AsyncDatabricksStore` | `user_id` | + +## Prerequisites + +1. **Add memory dependency** to `pyproject.toml`: + ```toml + dependencies = [ + "databricks-langchain[memory]", + ] + ``` + + Then run `uv sync` + +2. **Configure Lakebase** - See **lakebase-setup** skill for: + - Creating/configuring Lakebase instance + - Initializing tables (CRITICAL first-time step) + +--- + +## Quick Setup Summary + +Adding memory requires changes to **4 files**: + +| File | What to Add | +|------|-------------| +| `pyproject.toml` | Memory dependency | +| `.env` | Lakebase env vars (for local dev) | +| `databricks.yml` | Lakebase database resource | +| `app.yaml` | Environment variables for Lakebase | +| `agent_server/agent.py` | Memory tools and AsyncDatabricksStore | + +--- + +## Key Principles + +Before implementing memory, understand these patterns from the production implementation. + +### 1. Factory Function Pattern + +Memory tools should be returned from a factory function, not defined as standalone functions: + +```python +def memory_tools(): + @tool + async def get_user_memory(query: str, config: RunnableConfig) -> str: + ... + @tool + async def save_user_memory(memory_key: str, memory_data_json: str, config: RunnableConfig) -> str: + ... + @tool + async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: + ... + return [get_user_memory, save_user_memory, delete_user_memory] +``` + +### 2. User ID Extraction + +Extract `user_id` from the request, checking `custom_inputs` first. Return `None` (not a default) to let the caller decide: + +```python +def get_user_id(request: ResponsesAgentRequest) -> Optional[str]: + custom_inputs = dict(request.custom_inputs or {}) + if "user_id" in custom_inputs: + return custom_inputs["user_id"] + if request.context and getattr(request.context, "user_id", None): + return request.context.user_id + return None +``` + +### 3. Separate Error Handling + +Check `user_id` and `store` separately with distinct error messages: + +```python +user_id = config.get("configurable", {}).get("user_id") +if not user_id: + return "Memory not available - no user_id provided." + +store: Optional[BaseStore] = config.get("configurable", {}).get("store") +if not store: + return "Memory not available - store not configured." +``` + +### 4. JSON Validation for Save + +Validate JSON input before storing - the LLM may pass invalid JSON: + +```python +try: + memory_data = json.loads(memory_data_json) + if not isinstance(memory_data, dict): + return f"Failed: memory_data must be a JSON object, not {type(memory_data).__name__}" + await store.aput(namespace, memory_key, memory_data) +except json.JSONDecodeError as e: + return f"Failed to save memory: Invalid JSON - {e}" +``` + +### 5. Pass Store via RunnableConfig + +Pass the store through config, not as a function parameter: + +```python +config = {"configurable": {"user_id": user_id, "store": store}} +# Tools access via: config.get("configurable", {}).get("store") +``` + +--- + +## Production Reference + +For complete, tested implementations: + +| File | Description | +|------|-------------| +| [`agent-langgraph-long-term-memory/agent_server/utils_memory.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/utils_memory.py) | Memory tools factory, helpers, error handling | +| [`agent-langgraph-long-term-memory/agent_server/agent.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/agent.py) | Integration with agent, store initialization | + +Key functions in `utils_memory.py`: +- `memory_tools()` - Factory returning get/save/delete tools +- `get_user_id()` - Extract user_id from request +- `resolve_lakebase_instance_name()` - Handle hostname vs instance name +- `get_lakebase_access_error_message()` - Helpful error messages + +--- + +## Configuration Files + +### Step 1: databricks.yml (Lakebase Resource) + +Add the Lakebase database resource to your app: + +```yaml +resources: + apps: + agent_langgraph: + 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: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' +``` + +**Important:** The `name: 'database'` must match the `valueFrom` reference in `app.yaml`. + +### Step 2: app.yaml (Environment Variables) + +```yaml +command: ["uv", "run", "start-app"] + +env: + # ... other env vars ... + + # Lakebase instance name + - name: LAKEBASE_INSTANCE_NAME + value: "" + + # Embedding configuration + - name: EMBEDDING_ENDPOINT + value: "databricks-gte-large-en" + - name: EMBEDDING_DIMS + value: "1024" +``` + +**Important:** `LAKEBASE_INSTANCE_NAME` must match `instance_name` in databricks.yml. + +### Step 3: .env (Local Development) + +```bash +# Lakebase configuration for long-term memory +LAKEBASE_INSTANCE_NAME= +EMBEDDING_ENDPOINT=databricks-gte-large-en +EMBEDDING_DIMS=1024 +``` + +--- + +## Integration Example + +Minimal example showing how to integrate memory into your streaming function: + +```python +from agent_server.utils_memory import memory_tools, get_user_id + +@stream() +async def streaming(request: ResponsesAgentRequest): + user_id = get_user_id(request) + + async with AsyncDatabricksStore( + instance_name=LAKEBASE_INSTANCE_NAME, + embedding_endpoint=EMBEDDING_ENDPOINT, + embedding_dims=EMBEDDING_DIMS, + ) as store: + await store.setup() # Creates tables if needed + + tools = await mcp_client.get_tools() + memory_tools() + config = {"configurable": {"user_id": user_id, "store": store}} + + agent = create_react_agent(model=model, tools=tools) + async for event in agent.astream(messages, config): + yield event +``` + +--- + +## Initialize Tables and Deploy + +### Initialize Lakebase Tables (First Time Only) + +Before deploying, initialize the tables locally: + +```bash +uv run python -c "$(cat <<'EOF' +import asyncio +from databricks_langchain import AsyncDatabricksStore + +async def setup(): + async with AsyncDatabricksStore( + instance_name="", + embedding_endpoint="databricks-gte-large-en", + embedding_dims=1024, + ) as store: + await store.setup() + print("Tables created!") + +asyncio.run(setup()) +EOF +)" +``` + +### Deploy + +After initializing tables, deploy your agent. See **deploy** skill for full instructions. + +--- + +## Short-Term Memory + +For conversation history within a session, use `AsyncCheckpointSaver`: + +```python +from databricks_langchain import AsyncCheckpointSaver + +async with AsyncCheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME) as checkpointer: + agent = create_react_agent( + model=model, + tools=tools, + checkpointer=checkpointer, + ) + + config = {"configurable": {"thread_id": thread_id}} + async for event in agent.astream(messages, config): + yield event +``` + +See the [agent-langgraph-short-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-short-term-memory) template for a complete implementation. + +--- + +## Testing Memory + +### Test Locally + +```bash +# Start the server +uv run start-app + +# Save a memory +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Remember that I am on the shipping team"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' + +# Recall the memory +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "What team am I on?"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' + +# Delete a memory +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Forget what team I am on"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' +``` + +### Test Deployed App + +```bash +# Get OAuth token (PATs don't work for apps) +TOKEN=$(databricks auth token --host | jq -r '.access_token') + +# Test memory save +curl -X POST https:///invocations \ + -H "Authorization: Bearer $TOKEN" \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Remember I prefer detailed explanations"}], + "custom_inputs": {"user_id": "alice@example.com"} + }' +``` + +--- + +## First-Time Setup Checklist + +- [ ] Added `databricks-langchain[memory]` to `pyproject.toml` +- [ ] Run `uv sync` to install dependencies +- [ ] Created or identified Lakebase instance +- [ ] Added Lakebase env vars to `.env` (for local dev) +- [ ] Added `database` resource to `databricks.yml` +- [ ] Added `LAKEBASE_INSTANCE_NAME` to `app.yaml` +- [ ] **Initialized tables locally** by running `await store.setup()` +- [ ] Deployed with `databricks bundle deploy && databricks bundle run` + +--- + +## Troubleshooting + +| Issue | Cause | Solution | +|-------|-------|----------| +| **"embedding_dims is required"** | Missing parameter | Add `embedding_dims=1024` to AsyncDatabricksStore | +| **"relation 'store' does not exist"** | Tables not created | Run `await store.setup()` locally first | +| **"Unable to resolve Lakebase instance 'None'"** | Missing env var | Check `LAKEBASE_INSTANCE_NAME` in app.yaml | +| **"permission denied for table store"** | Missing grants | Add `database` resource to databricks.yml | +| **"Memory not available - no user_id"** | Missing user_id | Pass `custom_inputs.user_id` in request | +| **Memory not persisting** | Different user_ids | Use consistent user_id across requests | +| **App not updated after deploy** | Forgot to run bundle | Run `databricks bundle run agent_langgraph` after deploy | + +--- + +## Pre-Built Memory Templates + +For fully configured implementations without manual setup: + +| Template | Memory Type | Key Features | +|----------|-------------|--------------| +| [agent-langgraph-short-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-short-term-memory) | Short-term | AsyncCheckpointSaver, thread_id | +| [agent-langgraph-long-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory) | Long-term | AsyncDatabricksStore, memory tools | + +--- + +## Next Steps + +- Configure Lakebase: see **lakebase-setup** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/.claude/skills/agent-memory/SKILL.md b/.claude/skills/agent-memory/SKILL.md deleted file mode 100644 index 83d59668..00000000 --- a/.claude/skills/agent-memory/SKILL.md +++ /dev/null @@ -1,524 +0,0 @@ ---- -name: agent-memory -description: "Add memory capabilities to your agent. Use when: (1) User asks about 'memory', 'state', 'remember', 'conversation history', (2) Want to persist conversations or user preferences, (3) Adding checkpointing or long-term storage." ---- - -# Adding Memory to Your Agent - -> **Note:** This template does not include memory by default. Use this skill to **add memory capabilities**. For pre-configured memory templates, see: -> - `agent-langgraph-short-term-memory` - Conversation history within a session -> - `agent-langgraph-long-term-memory` - User facts that persist across sessions - -## Memory Types - -| Type | Use Case | Storage | Identifier | -|------|----------|---------|------------| -| **Short-term** | Conversation history within a session | `AsyncCheckpointSaver` | `thread_id` | -| **Long-term** | User facts that persist across sessions | `AsyncDatabricksStore` | `user_id` | - -## Prerequisites - -1. **Add memory dependency** to `pyproject.toml`: - ```toml - dependencies = [ - "databricks-langchain[memory]", - ] - ``` - - Then run `uv sync` - -2. **Configure Lakebase** - See **lakebase-setup** skill for: - - Creating/configuring Lakebase instance - - Initializing tables (CRITICAL first-time step) - ---- - -## Quick Setup Summary - -Adding memory requires changes to **4 files**: - -| File | What to Add | -|------|-------------| -| `pyproject.toml` | Memory dependency + hatch config | -| `.env` | Lakebase env vars (for local dev) | -| `databricks.yml` | Lakebase database resource | -| `app.yaml` | `valueFrom` reference to lakebase resource | -| `agent_server/agent.py` | Memory tools and AsyncDatabricksStore | - ---- - -## Step 1: Configure databricks.yml (Lakebase Resource) - -Add the Lakebase database resource to your app in `databricks.yml`: - -```yaml -resources: - apps: - agent_langgraph: - name: "your-app-name" - source_code_path: ./ - - resources: - # ... other resources (experiment, UC functions, etc.) ... - - # Lakebase instance for long-term memory - - name: 'lakebadatabasese_memory' - database: - instance_name: '' - database_name: 'postgres' - permission: 'CAN_CONNECT_AND_CREATE' -``` - -**Important:** The `name: 'database'` must match the `valueFrom` reference in `app.yaml`. - ---- - -## Step 2: Configure app.yaml (Environment Variables) - -Update `app.yaml` with the Lakebase environment variables: - -```yaml -command: ["uv", "run", "start-app"] - -env: - # ... other env vars ... - - # MLflow experiment (uses valueFrom to reference databricks.yml resource) - - name: MLFLOW_EXPERIMENT_ID - valueFrom: "experiment" - - # Lakebase instance name - must match the instance_name in databricks.yml database resource - # Note: Use 'value' (not 'valueFrom') because AsyncDatabricksStore needs the instance name, - # not the full connection string that valueFrom would provide - - name: LAKEBASE_INSTANCE_NAME - value: "" - - # Embedding configuration (static values) - - name: EMBEDDING_ENDPOINT - value: "databricks-gte-large-en" - - name: EMBEDDING_DIMS - value: "1024" -``` - -**Important:** The `LAKEBASE_INSTANCE_NAME` value must match the `instance_name` in your `databricks.yml` database resource. The `database` resource handles permissions, while `app.yaml` provides the instance name to your code. - -### Why not valueFrom for Lakebase? - -The `database` resource's `valueFrom` provides the full connection string (e.g., `instance-xxx.database.staging.cloud.databricks.com`), but `AsyncDatabricksStore` expects just the instance name. So we use a static `value` instead. - ---- - -## Step 3: Configure .env (Local Development) - -For local development, add to `.env`: - -```bash -# Lakebase configuration for long-term memory -LAKEBASE_INSTANCE_NAME= -EMBEDDING_ENDPOINT=databricks-gte-large-en -EMBEDDING_DIMS=1024 -``` - -> **Note:** `.env` is only for local development. When deployed, the app gets `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. - ---- - -## Step 4: Update agent.py - -### Add Imports and Configuration - -```python -import os -import uuid -from typing import AsyncGenerator - -from databricks_langchain import AsyncDatabricksStore, ChatDatabricks -from langchain_core.runnables import RunnableConfig -from langchain_core.tools import tool -from langgraph.prebuilt import create_react_agent - -# Environment configuration -LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") -EMBEDDING_ENDPOINT = os.getenv("EMBEDDING_ENDPOINT", "databricks-gte-large-en") -EMBEDDING_DIMS = int(os.getenv("EMBEDDING_DIMS", "1024")) # REQUIRED! -``` - -**CRITICAL:** You MUST specify `embedding_dims` when using `embedding_endpoint`. Without it, you'll get: -``` -"embedding_dims is required when embedding_endpoint is specified" -``` - -### Create Memory Tools - -```python -@tool -async def get_user_memory(query: str, config: RunnableConfig) -> str: - """Search for relevant information about the user from long-term memory. - Use this to recall preferences, past interactions, or other saved information. - """ - user_id = config.get("configurable", {}).get("user_id") - store = config.get("configurable", {}).get("store") - if not user_id or not store: - return "Memory not available - no user context" - - # Sanitize user_id for namespace (replace special chars) - namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) - results = await store.asearch(namespace, query=query, limit=5) - if not results: - return "No memories found for this query" - return "\n".join([f"- {r.value}" for r in results]) - - -@tool -async def save_user_memory(memory_key: str, memory_data: str, config: RunnableConfig) -> str: - """Save information about the user to long-term memory. - Use this to remember user preferences, important details, or other - information that should persist across conversations. - - Args: - memory_key: A short descriptive key (e.g., "preferred_name", "team", "region") - memory_data: The information to save - """ - user_id = config.get("configurable", {}).get("user_id") - store = config.get("configurable", {}).get("store") - if not user_id or not store: - return "Memory not available - no user context" - - namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) - await store.aput(namespace, memory_key, {"value": memory_data}) - return f"Saved memory: {memory_key}" - - -@tool -async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: - """Delete a specific memory from the user's long-term memory. - Use this when the user asks to forget something or correct stored information. - - Args: - memory_key: The key of the memory to delete (e.g., "preferred_name", "team") - """ - user_id = config.get("configurable", {}).get("user_id") - store = config.get("configurable", {}).get("store") - if not user_id or not store: - return "Memory not available - no user context" - - namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) - await store.adelete(namespace, memory_key) - return f"Deleted memory: {memory_key}" -``` - -### Add Helper Functions - -```python -def _get_user_id(request: ResponsesAgentRequest) -> str: - """Extract user_id from request context or custom inputs.""" - custom_inputs = dict(request.custom_inputs or {}) - if "user_id" in custom_inputs and custom_inputs["user_id"]: - return str(custom_inputs["user_id"]) - if request.context and getattr(request.context, "user_id", None): - return str(request.context.user_id) - return "default-user" - - -def _get_or_create_thread_id(request: ResponsesAgentRequest) -> str: - """Extract or create thread_id for conversation tracking.""" - custom_inputs = dict(request.custom_inputs or {}) - if "thread_id" in custom_inputs and custom_inputs["thread_id"]: - return str(custom_inputs["thread_id"]) - if request.context and getattr(request.context, "conversation_id", None): - return str(request.context.conversation_id) - return str(uuid.uuid4()) -``` - -### Update Streaming Function - -```python -@stream() -async def streaming( - request: ResponsesAgentRequest, -) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: - # Get user context - user_id = _get_user_id(request) - thread_id = _get_or_create_thread_id(request) - - # Get other tools (MCP, etc.) - mcp_client = init_mcp_client(sp_workspace_client) - mcp_tools = await mcp_client.get_tools() - - # Memory tools - memory_tools = [get_user_memory, save_user_memory, delete_user_memory] - - # Combine all tools - all_tools = mcp_tools + memory_tools - - # Initialize model - model = ChatDatabricks(endpoint="databricks-claude-sonnet-4") - - # Use AsyncDatabricksStore for long-term memory - async with AsyncDatabricksStore( - instance_name=LAKEBASE_INSTANCE_NAME, - embedding_endpoint=EMBEDDING_ENDPOINT, - embedding_dims=EMBEDDING_DIMS, # REQUIRED! - ) as store: - # Create agent with tools - agent = create_react_agent( - model=model, - tools=all_tools, - prompt=AGENT_INSTRUCTIONS, - ) - - # Prepare input - messages = {"messages": to_chat_completions_input([i.model_dump() for i in request.input])} - - # Configure with user context for memory tools - config = { - "configurable": { - "user_id": user_id, - "thread_id": thread_id, - "store": store, # Pass store to tools via config - } - } - - # Stream agent responses - async for event in process_agent_astream_events( - agent.astream(input=messages, config=config, stream_mode=["updates", "messages"]) - ): - yield event -``` - ---- - -## Step 5: Initialize Tables and Deploy - -### Initialize Lakebase Tables (First Time Only) - -Before deploying, initialize the tables locally: - -```bash -uv run python -c "$(cat <<'EOF' -import asyncio -from databricks_langchain import AsyncDatabricksStore - -async def setup(): - async with AsyncDatabricksStore( - instance_name="", - embedding_endpoint="databricks-gte-large-en", - embedding_dims=1024, - ) as store: - await store.setup() - print("Tables created!") - -asyncio.run(setup()) -EOF -)" -``` - -### Deploy and Run - -**IMPORTANT:** Always run `databricks bundle run` after `databricks bundle deploy` to start/restart the app with the new code: - -```bash -# Deploy resources and upload files -databricks bundle deploy - -# Start/restart the app with new code (REQUIRED!) -databricks bundle run agent_langgraph -``` - -> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is required to actually start the app with the new code. - ---- - -## Complete Example Files - -### databricks.yml (with Lakebase) - -```yaml -bundle: - name: agent_langgraph - -resources: - experiments: - agent_langgraph_experiment: - name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} - - apps: - agent_langgraph: - name: "my-agent-app" - description: "Agent with long-term memory" - source_code_path: ./ - - resources: - - name: 'experiment' - experiment: - experiment_id: "${resources.experiments.agent_langgraph_experiment.id}" - permission: 'CAN_MANAGE' - - # Lakebase instance for long-term memory - - name: 'database' - database: - instance_name: '' - database_name: 'postgres' - permission: 'CAN_CONNECT_AND_CREATE' - -targets: - dev: - mode: development - default: true -``` - -### app.yaml - -```yaml -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" - # Reference experiment resource from databricks.yml - - name: MLFLOW_EXPERIMENT_ID - valueFrom: "experiment" - # Lakebase instance name (must match instance_name in databricks.yml) - - name: LAKEBASE_INSTANCE_NAME - value: "" - # Embedding configuration - - name: EMBEDDING_ENDPOINT - value: "databricks-gte-large-en" - - name: EMBEDDING_DIMS - value: "1024" -``` - ---- - -## Adding Short-Term Memory - -Short-term memory stores conversation history within a thread. - -### Import and Initialize Checkpointer - -```python -from databricks_langchain import AsyncCheckpointSaver - -LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") - -async with AsyncCheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME) as checkpointer: - agent = create_react_agent( - model=model, - tools=tools, - checkpointer=checkpointer, # Enables conversation persistence - ) -``` - -### Use thread_id in Agent Config - -```python -config = {"configurable": {"thread_id": thread_id}} -async for event in agent.astream(input_state, config, stream_mode=["updates", "messages"]): - yield event -``` - ---- - -## Testing Memory - -### Test Locally - -```bash -# Start the server -uv run start-app - -# Save a memory -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ - "input": [{"role": "user", "content": "Remember that I am on the shipping team"}], - "custom_inputs": {"user_id": "alice@example.com"} - }' - -# Recall the memory -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ - "input": [{"role": "user", "content": "What team am I on?"}], - "custom_inputs": {"user_id": "alice@example.com"} - }' - -# Delete a memory -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ - "input": [{"role": "user", "content": "Forget what team I am on"}], - "custom_inputs": {"user_id": "alice@example.com"} - }' -``` - -### Test Deployed App - -```bash -# Get OAuth token (PATs don't work for apps) -TOKEN=$(databricks auth token --host | jq -r '.access_token') - -# Test memory save -curl -X POST https:///invocations \ - -H "Authorization: Bearer $TOKEN" \ - -H "Content-Type: application/json" \ - -d '{ - "input": [{"role": "user", "content": "Remember I prefer detailed explanations"}], - "custom_inputs": {"user_id": "alice@example.com"} - }' -``` - ---- - -## First-Time Setup Checklist - -- [ ] Added `databricks-langchain[memory]` to `pyproject.toml` -- [ ] Run `uv sync` to install dependencies -- [ ] Created or identified Lakebase instance -- [ ] Added Lakebase env vars to `.env` (for local dev) -- [ ] Added `database` resource to `databricks.yml` (for permissions) -- [ ] Added `LAKEBASE_INSTANCE_NAME` value to `app.yaml` (matching instance_name in databricks.yml) -- [ ] **Initialized tables locally** by running `await store.setup()` -- [ ] Deployed with `databricks bundle deploy` -- [ ] **Started app with `databricks bundle run agent_langgraph`** - ---- - -## Troubleshooting - -| Issue | Cause | Solution | -|-------|-------|----------| -| **"embedding_dims is required"** | Missing parameter | Add `embedding_dims=1024` to AsyncDatabricksStore | -| **"relation 'store' does not exist"** | Tables not created | Run `await store.setup()` locally first | -| **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `valueFrom: "database"` to app.yaml | -| **"permission denied for table store"** | Missing grants | Lakebase `database` resource in DAB should handle this | -| **"Memory not available"** | No user_id in request | Ensure `custom_inputs.user_id` is passed | -| **Memory not persisting** | Different user_ids | Use consistent user_id across requests | -| **App not updated after deploy** | Forgot to run bundle | Run `databricks bundle run agent_langgraph` after deploy | - ---- - -## Pre-Built Memory Templates - -For fully configured implementations without manual setup: - -| Template | Memory Type | Key Features | -|----------|-------------|--------------| -| `agent-langgraph-short-term-memory` | Short-term | AsyncCheckpointSaver, thread_id | -| `agent-langgraph-long-term-memory` | Long-term | AsyncDatabricksStore, memory tools | - ---- - -## Next Steps - -- Configure Lakebase: see **lakebase-setup** skill -- Test locally: see **run-locally** skill -- Deploy: see **deploy** skill diff --git a/.claude/sync-skills.py b/.claude/sync-skills.py index 469b5572..5b782185 100755 --- a/.claude/sync-skills.py +++ b/.claude/sync-skills.py @@ -82,9 +82,10 @@ def sync_template(template: str, config: dict): copy_skill(SOURCE / f"modify-{sdk}-agent", dest / "modify-agent") # Memory skills (all LangGraph templates - enables adding memory to any agent) + # SDK-specific memory skills are renamed on copy (e.g., agent-langgraph-memory -> agent-memory) if sdk == "langgraph": copy_skill(SOURCE / "lakebase-setup", dest / "lakebase-setup") - copy_skill(SOURCE / "agent-memory", dest / "agent-memory") + copy_skill(SOURCE / "agent-langgraph-memory", dest / "agent-memory") def main(): diff --git a/agent-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md index 83d59668..c87e21a3 100644 --- a/agent-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md +++ b/agent-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md @@ -6,8 +6,8 @@ description: "Add memory capabilities to your agent. Use when: (1) User asks abo # Adding Memory to Your Agent > **Note:** This template does not include memory by default. Use this skill to **add memory capabilities**. For pre-configured memory templates, see: -> - `agent-langgraph-short-term-memory` - Conversation history within a session -> - `agent-langgraph-long-term-memory` - User facts that persist across sessions +> - [agent-langgraph-short-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-short-term-memory) - Conversation history within a session +> - [agent-langgraph-long-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory) - User facts that persist across sessions ## Memory Types @@ -39,17 +39,111 @@ Adding memory requires changes to **4 files**: | File | What to Add | |------|-------------| -| `pyproject.toml` | Memory dependency + hatch config | +| `pyproject.toml` | Memory dependency | | `.env` | Lakebase env vars (for local dev) | | `databricks.yml` | Lakebase database resource | -| `app.yaml` | `valueFrom` reference to lakebase resource | +| `app.yaml` | Environment variables for Lakebase | | `agent_server/agent.py` | Memory tools and AsyncDatabricksStore | --- -## Step 1: Configure databricks.yml (Lakebase Resource) +## Key Principles -Add the Lakebase database resource to your app in `databricks.yml`: +Before implementing memory, understand these patterns from the production implementation. + +### 1. Factory Function Pattern + +Memory tools should be returned from a factory function, not defined as standalone functions: + +```python +def memory_tools(): + @tool + async def get_user_memory(query: str, config: RunnableConfig) -> str: + ... + @tool + async def save_user_memory(memory_key: str, memory_data_json: str, config: RunnableConfig) -> str: + ... + @tool + async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: + ... + return [get_user_memory, save_user_memory, delete_user_memory] +``` + +### 2. User ID Extraction + +Extract `user_id` from the request, checking `custom_inputs` first. Return `None` (not a default) to let the caller decide: + +```python +def get_user_id(request: ResponsesAgentRequest) -> Optional[str]: + custom_inputs = dict(request.custom_inputs or {}) + if "user_id" in custom_inputs: + return custom_inputs["user_id"] + if request.context and getattr(request.context, "user_id", None): + return request.context.user_id + return None +``` + +### 3. Separate Error Handling + +Check `user_id` and `store` separately with distinct error messages: + +```python +user_id = config.get("configurable", {}).get("user_id") +if not user_id: + return "Memory not available - no user_id provided." + +store: Optional[BaseStore] = config.get("configurable", {}).get("store") +if not store: + return "Memory not available - store not configured." +``` + +### 4. JSON Validation for Save + +Validate JSON input before storing - the LLM may pass invalid JSON: + +```python +try: + memory_data = json.loads(memory_data_json) + if not isinstance(memory_data, dict): + return f"Failed: memory_data must be a JSON object, not {type(memory_data).__name__}" + await store.aput(namespace, memory_key, memory_data) +except json.JSONDecodeError as e: + return f"Failed to save memory: Invalid JSON - {e}" +``` + +### 5. Pass Store via RunnableConfig + +Pass the store through config, not as a function parameter: + +```python +config = {"configurable": {"user_id": user_id, "store": store}} +# Tools access via: config.get("configurable", {}).get("store") +``` + +--- + +## Production Reference + +For complete, tested implementations: + +| File | Description | +|------|-------------| +| [`agent-langgraph-long-term-memory/agent_server/utils_memory.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/utils_memory.py) | Memory tools factory, helpers, error handling | +| [`agent-langgraph-long-term-memory/agent_server/agent.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/agent.py) | Integration with agent, store initialization | + +Key functions in `utils_memory.py`: +- `memory_tools()` - Factory returning get/save/delete tools +- `get_user_id()` - Extract user_id from request +- `resolve_lakebase_instance_name()` - Handle hostname vs instance name +- `get_lakebase_access_error_message()` - Helpful error messages + +--- + +## Configuration Files + +### Step 1: databricks.yml (Lakebase Resource) + +Add the Lakebase database resource to your app: ```yaml resources: @@ -62,7 +156,7 @@ resources: # ... other resources (experiment, UC functions, etc.) ... # Lakebase instance for long-term memory - - name: 'lakebadatabasese_memory' + - name: 'database' database: instance_name: '' database_name: 'postgres' @@ -71,11 +165,7 @@ resources: **Important:** The `name: 'database'` must match the `valueFrom` reference in `app.yaml`. ---- - -## Step 2: Configure app.yaml (Environment Variables) - -Update `app.yaml` with the Lakebase environment variables: +### Step 2: app.yaml (Environment Variables) ```yaml command: ["uv", "run", "start-app"] @@ -83,34 +173,20 @@ command: ["uv", "run", "start-app"] env: # ... other env vars ... - # MLflow experiment (uses valueFrom to reference databricks.yml resource) - - name: MLFLOW_EXPERIMENT_ID - valueFrom: "experiment" - - # Lakebase instance name - must match the instance_name in databricks.yml database resource - # Note: Use 'value' (not 'valueFrom') because AsyncDatabricksStore needs the instance name, - # not the full connection string that valueFrom would provide + # Lakebase instance name - name: LAKEBASE_INSTANCE_NAME value: "" - # Embedding configuration (static values) + # Embedding configuration - name: EMBEDDING_ENDPOINT value: "databricks-gte-large-en" - name: EMBEDDING_DIMS value: "1024" ``` -**Important:** The `LAKEBASE_INSTANCE_NAME` value must match the `instance_name` in your `databricks.yml` database resource. The `database` resource handles permissions, while `app.yaml` provides the instance name to your code. +**Important:** `LAKEBASE_INSTANCE_NAME` must match `instance_name` in databricks.yml. -### Why not valueFrom for Lakebase? - -The `database` resource's `valueFrom` provides the full connection string (e.g., `instance-xxx.database.staging.cloud.databricks.com`), but `AsyncDatabricksStore` expects just the instance name. So we use a static `value` instead. - ---- - -## Step 3: Configure .env (Local Development) - -For local development, add to `.env`: +### Step 3: .env (Local Development) ```bash # Lakebase configuration for long-term memory @@ -119,176 +195,37 @@ EMBEDDING_ENDPOINT=databricks-gte-large-en EMBEDDING_DIMS=1024 ``` -> **Note:** `.env` is only for local development. When deployed, the app gets `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. - --- -## Step 4: Update agent.py +## Integration Example -### Add Imports and Configuration +Minimal example showing how to integrate memory into your streaming function: ```python -import os -import uuid -from typing import AsyncGenerator - -from databricks_langchain import AsyncDatabricksStore, ChatDatabricks -from langchain_core.runnables import RunnableConfig -from langchain_core.tools import tool -from langgraph.prebuilt import create_react_agent - -# Environment configuration -LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") -EMBEDDING_ENDPOINT = os.getenv("EMBEDDING_ENDPOINT", "databricks-gte-large-en") -EMBEDDING_DIMS = int(os.getenv("EMBEDDING_DIMS", "1024")) # REQUIRED! -``` - -**CRITICAL:** You MUST specify `embedding_dims` when using `embedding_endpoint`. Without it, you'll get: -``` -"embedding_dims is required when embedding_endpoint is specified" -``` +from agent_server.utils_memory import memory_tools, get_user_id -### Create Memory Tools - -```python -@tool -async def get_user_memory(query: str, config: RunnableConfig) -> str: - """Search for relevant information about the user from long-term memory. - Use this to recall preferences, past interactions, or other saved information. - """ - user_id = config.get("configurable", {}).get("user_id") - store = config.get("configurable", {}).get("store") - if not user_id or not store: - return "Memory not available - no user context" - - # Sanitize user_id for namespace (replace special chars) - namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) - results = await store.asearch(namespace, query=query, limit=5) - if not results: - return "No memories found for this query" - return "\n".join([f"- {r.value}" for r in results]) - - -@tool -async def save_user_memory(memory_key: str, memory_data: str, config: RunnableConfig) -> str: - """Save information about the user to long-term memory. - Use this to remember user preferences, important details, or other - information that should persist across conversations. - - Args: - memory_key: A short descriptive key (e.g., "preferred_name", "team", "region") - memory_data: The information to save - """ - user_id = config.get("configurable", {}).get("user_id") - store = config.get("configurable", {}).get("store") - if not user_id or not store: - return "Memory not available - no user context" - - namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) - await store.aput(namespace, memory_key, {"value": memory_data}) - return f"Saved memory: {memory_key}" - - -@tool -async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: - """Delete a specific memory from the user's long-term memory. - Use this when the user asks to forget something or correct stored information. - - Args: - memory_key: The key of the memory to delete (e.g., "preferred_name", "team") - """ - user_id = config.get("configurable", {}).get("user_id") - store = config.get("configurable", {}).get("store") - if not user_id or not store: - return "Memory not available - no user context" - - namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) - await store.adelete(namespace, memory_key) - return f"Deleted memory: {memory_key}" -``` - -### Add Helper Functions - -```python -def _get_user_id(request: ResponsesAgentRequest) -> str: - """Extract user_id from request context or custom inputs.""" - custom_inputs = dict(request.custom_inputs or {}) - if "user_id" in custom_inputs and custom_inputs["user_id"]: - return str(custom_inputs["user_id"]) - if request.context and getattr(request.context, "user_id", None): - return str(request.context.user_id) - return "default-user" - - -def _get_or_create_thread_id(request: ResponsesAgentRequest) -> str: - """Extract or create thread_id for conversation tracking.""" - custom_inputs = dict(request.custom_inputs or {}) - if "thread_id" in custom_inputs and custom_inputs["thread_id"]: - return str(custom_inputs["thread_id"]) - if request.context and getattr(request.context, "conversation_id", None): - return str(request.context.conversation_id) - return str(uuid.uuid4()) -``` - -### Update Streaming Function - -```python @stream() -async def streaming( - request: ResponsesAgentRequest, -) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: - # Get user context - user_id = _get_user_id(request) - thread_id = _get_or_create_thread_id(request) +async def streaming(request: ResponsesAgentRequest): + user_id = get_user_id(request) - # Get other tools (MCP, etc.) - mcp_client = init_mcp_client(sp_workspace_client) - mcp_tools = await mcp_client.get_tools() - - # Memory tools - memory_tools = [get_user_memory, save_user_memory, delete_user_memory] - - # Combine all tools - all_tools = mcp_tools + memory_tools - - # Initialize model - model = ChatDatabricks(endpoint="databricks-claude-sonnet-4") - - # Use AsyncDatabricksStore for long-term memory async with AsyncDatabricksStore( instance_name=LAKEBASE_INSTANCE_NAME, embedding_endpoint=EMBEDDING_ENDPOINT, - embedding_dims=EMBEDDING_DIMS, # REQUIRED! + embedding_dims=EMBEDDING_DIMS, ) as store: - # Create agent with tools - agent = create_react_agent( - model=model, - tools=all_tools, - prompt=AGENT_INSTRUCTIONS, - ) - - # Prepare input - messages = {"messages": to_chat_completions_input([i.model_dump() for i in request.input])} - - # Configure with user context for memory tools - config = { - "configurable": { - "user_id": user_id, - "thread_id": thread_id, - "store": store, # Pass store to tools via config - } - } - - # Stream agent responses - async for event in process_agent_astream_events( - agent.astream(input=messages, config=config, stream_mode=["updates", "messages"]) - ): + await store.setup() # Creates tables if needed + + tools = await mcp_client.get_tools() + memory_tools() + config = {"configurable": {"user_id": user_id, "store": store}} + + agent = create_react_agent(model=model, tools=tools) + async for event in agent.astream(messages, config): yield event ``` --- -## Step 5: Initialize Tables and Deploy +## Initialize Tables and Deploy ### Initialize Lakebase Tables (First Time Only) @@ -313,118 +250,33 @@ EOF )" ``` -### Deploy and Run - -**IMPORTANT:** Always run `databricks bundle run` after `databricks bundle deploy` to start/restart the app with the new code: - -```bash -# Deploy resources and upload files -databricks bundle deploy - -# Start/restart the app with new code (REQUIRED!) -databricks bundle run agent_langgraph -``` - -> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is required to actually start the app with the new code. - ---- - -## Complete Example Files - -### databricks.yml (with Lakebase) - -```yaml -bundle: - name: agent_langgraph - -resources: - experiments: - agent_langgraph_experiment: - name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} - - apps: - agent_langgraph: - name: "my-agent-app" - description: "Agent with long-term memory" - source_code_path: ./ - - resources: - - name: 'experiment' - experiment: - experiment_id: "${resources.experiments.agent_langgraph_experiment.id}" - permission: 'CAN_MANAGE' - - # Lakebase instance for long-term memory - - name: 'database' - database: - instance_name: '' - database_name: 'postgres' - permission: 'CAN_CONNECT_AND_CREATE' - -targets: - dev: - mode: development - default: true -``` - -### app.yaml +### Deploy -```yaml -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" - # Reference experiment resource from databricks.yml - - name: MLFLOW_EXPERIMENT_ID - valueFrom: "experiment" - # Lakebase instance name (must match instance_name in databricks.yml) - - name: LAKEBASE_INSTANCE_NAME - value: "" - # Embedding configuration - - name: EMBEDDING_ENDPOINT - value: "databricks-gte-large-en" - - name: EMBEDDING_DIMS - value: "1024" -``` +After initializing tables, deploy your agent. See **deploy** skill for full instructions. --- -## Adding Short-Term Memory - -Short-term memory stores conversation history within a thread. +## Short-Term Memory -### Import and Initialize Checkpointer +For conversation history within a session, use `AsyncCheckpointSaver`: ```python from databricks_langchain import AsyncCheckpointSaver -LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") - async with AsyncCheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME) as checkpointer: agent = create_react_agent( model=model, tools=tools, - checkpointer=checkpointer, # Enables conversation persistence + checkpointer=checkpointer, ) -``` -### Use thread_id in Agent Config - -```python -config = {"configurable": {"thread_id": thread_id}} -async for event in agent.astream(input_state, config, stream_mode=["updates", "messages"]): - yield event + config = {"configurable": {"thread_id": thread_id}} + async for event in agent.astream(messages, config): + yield event ``` +See the [agent-langgraph-short-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-short-term-memory) template for a complete implementation. + --- ## Testing Memory @@ -484,11 +336,10 @@ curl -X POST https:///invocations \ - [ ] Run `uv sync` to install dependencies - [ ] Created or identified Lakebase instance - [ ] Added Lakebase env vars to `.env` (for local dev) -- [ ] Added `database` resource to `databricks.yml` (for permissions) -- [ ] Added `LAKEBASE_INSTANCE_NAME` value to `app.yaml` (matching instance_name in databricks.yml) +- [ ] Added `database` resource to `databricks.yml` +- [ ] Added `LAKEBASE_INSTANCE_NAME` to `app.yaml` - [ ] **Initialized tables locally** by running `await store.setup()` -- [ ] Deployed with `databricks bundle deploy` -- [ ] **Started app with `databricks bundle run agent_langgraph`** +- [ ] Deployed with `databricks bundle deploy && databricks bundle run` --- @@ -498,9 +349,9 @@ curl -X POST https:///invocations \ |-------|-------|----------| | **"embedding_dims is required"** | Missing parameter | Add `embedding_dims=1024` to AsyncDatabricksStore | | **"relation 'store' does not exist"** | Tables not created | Run `await store.setup()` locally first | -| **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `valueFrom: "database"` to app.yaml | -| **"permission denied for table store"** | Missing grants | Lakebase `database` resource in DAB should handle this | -| **"Memory not available"** | No user_id in request | Ensure `custom_inputs.user_id` is passed | +| **"Unable to resolve Lakebase instance 'None'"** | Missing env var | Check `LAKEBASE_INSTANCE_NAME` in app.yaml | +| **"permission denied for table store"** | Missing grants | Add `database` resource to databricks.yml | +| **"Memory not available - no user_id"** | Missing user_id | Pass `custom_inputs.user_id` in request | | **Memory not persisting** | Different user_ids | Use consistent user_id across requests | | **App not updated after deploy** | Forgot to run bundle | Run `databricks bundle run agent_langgraph` after deploy | @@ -512,8 +363,8 @@ For fully configured implementations without manual setup: | Template | Memory Type | Key Features | |----------|-------------|--------------| -| `agent-langgraph-short-term-memory` | Short-term | AsyncCheckpointSaver, thread_id | -| `agent-langgraph-long-term-memory` | Long-term | AsyncDatabricksStore, memory tools | +| [agent-langgraph-short-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-short-term-memory) | Short-term | AsyncCheckpointSaver, thread_id | +| [agent-langgraph-long-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory) | Long-term | AsyncDatabricksStore, memory tools | --- diff --git a/agent-langgraph-long-term-memory/.claude/skills/lakebase-setup/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/lakebase-setup/SKILL.md index 6b8c802e..65c79099 100644 --- a/agent-langgraph-long-term-memory/.claude/skills/lakebase-setup/SKILL.md +++ b/agent-langgraph-long-term-memory/.claude/skills/lakebase-setup/SKILL.md @@ -82,7 +82,7 @@ resources: ``` **Important:** -- The `name: 'database'` must match the `valueFrom` reference in `app.yaml` +- The `instance_name: ''` must match the `value` reference in `app.yaml` - Using the `database` resource type automatically grants the app's service principal access to Lakebase ### Update app.yaml (Environment Variables) diff --git a/agent-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md index 83d59668..c87e21a3 100644 --- a/agent-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md +++ b/agent-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md @@ -6,8 +6,8 @@ description: "Add memory capabilities to your agent. Use when: (1) User asks abo # Adding Memory to Your Agent > **Note:** This template does not include memory by default. Use this skill to **add memory capabilities**. For pre-configured memory templates, see: -> - `agent-langgraph-short-term-memory` - Conversation history within a session -> - `agent-langgraph-long-term-memory` - User facts that persist across sessions +> - [agent-langgraph-short-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-short-term-memory) - Conversation history within a session +> - [agent-langgraph-long-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory) - User facts that persist across sessions ## Memory Types @@ -39,17 +39,111 @@ Adding memory requires changes to **4 files**: | File | What to Add | |------|-------------| -| `pyproject.toml` | Memory dependency + hatch config | +| `pyproject.toml` | Memory dependency | | `.env` | Lakebase env vars (for local dev) | | `databricks.yml` | Lakebase database resource | -| `app.yaml` | `valueFrom` reference to lakebase resource | +| `app.yaml` | Environment variables for Lakebase | | `agent_server/agent.py` | Memory tools and AsyncDatabricksStore | --- -## Step 1: Configure databricks.yml (Lakebase Resource) +## Key Principles -Add the Lakebase database resource to your app in `databricks.yml`: +Before implementing memory, understand these patterns from the production implementation. + +### 1. Factory Function Pattern + +Memory tools should be returned from a factory function, not defined as standalone functions: + +```python +def memory_tools(): + @tool + async def get_user_memory(query: str, config: RunnableConfig) -> str: + ... + @tool + async def save_user_memory(memory_key: str, memory_data_json: str, config: RunnableConfig) -> str: + ... + @tool + async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: + ... + return [get_user_memory, save_user_memory, delete_user_memory] +``` + +### 2. User ID Extraction + +Extract `user_id` from the request, checking `custom_inputs` first. Return `None` (not a default) to let the caller decide: + +```python +def get_user_id(request: ResponsesAgentRequest) -> Optional[str]: + custom_inputs = dict(request.custom_inputs or {}) + if "user_id" in custom_inputs: + return custom_inputs["user_id"] + if request.context and getattr(request.context, "user_id", None): + return request.context.user_id + return None +``` + +### 3. Separate Error Handling + +Check `user_id` and `store` separately with distinct error messages: + +```python +user_id = config.get("configurable", {}).get("user_id") +if not user_id: + return "Memory not available - no user_id provided." + +store: Optional[BaseStore] = config.get("configurable", {}).get("store") +if not store: + return "Memory not available - store not configured." +``` + +### 4. JSON Validation for Save + +Validate JSON input before storing - the LLM may pass invalid JSON: + +```python +try: + memory_data = json.loads(memory_data_json) + if not isinstance(memory_data, dict): + return f"Failed: memory_data must be a JSON object, not {type(memory_data).__name__}" + await store.aput(namespace, memory_key, memory_data) +except json.JSONDecodeError as e: + return f"Failed to save memory: Invalid JSON - {e}" +``` + +### 5. Pass Store via RunnableConfig + +Pass the store through config, not as a function parameter: + +```python +config = {"configurable": {"user_id": user_id, "store": store}} +# Tools access via: config.get("configurable", {}).get("store") +``` + +--- + +## Production Reference + +For complete, tested implementations: + +| File | Description | +|------|-------------| +| [`agent-langgraph-long-term-memory/agent_server/utils_memory.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/utils_memory.py) | Memory tools factory, helpers, error handling | +| [`agent-langgraph-long-term-memory/agent_server/agent.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/agent.py) | Integration with agent, store initialization | + +Key functions in `utils_memory.py`: +- `memory_tools()` - Factory returning get/save/delete tools +- `get_user_id()` - Extract user_id from request +- `resolve_lakebase_instance_name()` - Handle hostname vs instance name +- `get_lakebase_access_error_message()` - Helpful error messages + +--- + +## Configuration Files + +### Step 1: databricks.yml (Lakebase Resource) + +Add the Lakebase database resource to your app: ```yaml resources: @@ -62,7 +156,7 @@ resources: # ... other resources (experiment, UC functions, etc.) ... # Lakebase instance for long-term memory - - name: 'lakebadatabasese_memory' + - name: 'database' database: instance_name: '' database_name: 'postgres' @@ -71,11 +165,7 @@ resources: **Important:** The `name: 'database'` must match the `valueFrom` reference in `app.yaml`. ---- - -## Step 2: Configure app.yaml (Environment Variables) - -Update `app.yaml` with the Lakebase environment variables: +### Step 2: app.yaml (Environment Variables) ```yaml command: ["uv", "run", "start-app"] @@ -83,34 +173,20 @@ command: ["uv", "run", "start-app"] env: # ... other env vars ... - # MLflow experiment (uses valueFrom to reference databricks.yml resource) - - name: MLFLOW_EXPERIMENT_ID - valueFrom: "experiment" - - # Lakebase instance name - must match the instance_name in databricks.yml database resource - # Note: Use 'value' (not 'valueFrom') because AsyncDatabricksStore needs the instance name, - # not the full connection string that valueFrom would provide + # Lakebase instance name - name: LAKEBASE_INSTANCE_NAME value: "" - # Embedding configuration (static values) + # Embedding configuration - name: EMBEDDING_ENDPOINT value: "databricks-gte-large-en" - name: EMBEDDING_DIMS value: "1024" ``` -**Important:** The `LAKEBASE_INSTANCE_NAME` value must match the `instance_name` in your `databricks.yml` database resource. The `database` resource handles permissions, while `app.yaml` provides the instance name to your code. +**Important:** `LAKEBASE_INSTANCE_NAME` must match `instance_name` in databricks.yml. -### Why not valueFrom for Lakebase? - -The `database` resource's `valueFrom` provides the full connection string (e.g., `instance-xxx.database.staging.cloud.databricks.com`), but `AsyncDatabricksStore` expects just the instance name. So we use a static `value` instead. - ---- - -## Step 3: Configure .env (Local Development) - -For local development, add to `.env`: +### Step 3: .env (Local Development) ```bash # Lakebase configuration for long-term memory @@ -119,176 +195,37 @@ EMBEDDING_ENDPOINT=databricks-gte-large-en EMBEDDING_DIMS=1024 ``` -> **Note:** `.env` is only for local development. When deployed, the app gets `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. - --- -## Step 4: Update agent.py +## Integration Example -### Add Imports and Configuration +Minimal example showing how to integrate memory into your streaming function: ```python -import os -import uuid -from typing import AsyncGenerator - -from databricks_langchain import AsyncDatabricksStore, ChatDatabricks -from langchain_core.runnables import RunnableConfig -from langchain_core.tools import tool -from langgraph.prebuilt import create_react_agent - -# Environment configuration -LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") -EMBEDDING_ENDPOINT = os.getenv("EMBEDDING_ENDPOINT", "databricks-gte-large-en") -EMBEDDING_DIMS = int(os.getenv("EMBEDDING_DIMS", "1024")) # REQUIRED! -``` - -**CRITICAL:** You MUST specify `embedding_dims` when using `embedding_endpoint`. Without it, you'll get: -``` -"embedding_dims is required when embedding_endpoint is specified" -``` +from agent_server.utils_memory import memory_tools, get_user_id -### Create Memory Tools - -```python -@tool -async def get_user_memory(query: str, config: RunnableConfig) -> str: - """Search for relevant information about the user from long-term memory. - Use this to recall preferences, past interactions, or other saved information. - """ - user_id = config.get("configurable", {}).get("user_id") - store = config.get("configurable", {}).get("store") - if not user_id or not store: - return "Memory not available - no user context" - - # Sanitize user_id for namespace (replace special chars) - namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) - results = await store.asearch(namespace, query=query, limit=5) - if not results: - return "No memories found for this query" - return "\n".join([f"- {r.value}" for r in results]) - - -@tool -async def save_user_memory(memory_key: str, memory_data: str, config: RunnableConfig) -> str: - """Save information about the user to long-term memory. - Use this to remember user preferences, important details, or other - information that should persist across conversations. - - Args: - memory_key: A short descriptive key (e.g., "preferred_name", "team", "region") - memory_data: The information to save - """ - user_id = config.get("configurable", {}).get("user_id") - store = config.get("configurable", {}).get("store") - if not user_id or not store: - return "Memory not available - no user context" - - namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) - await store.aput(namespace, memory_key, {"value": memory_data}) - return f"Saved memory: {memory_key}" - - -@tool -async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: - """Delete a specific memory from the user's long-term memory. - Use this when the user asks to forget something or correct stored information. - - Args: - memory_key: The key of the memory to delete (e.g., "preferred_name", "team") - """ - user_id = config.get("configurable", {}).get("user_id") - store = config.get("configurable", {}).get("store") - if not user_id or not store: - return "Memory not available - no user context" - - namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) - await store.adelete(namespace, memory_key) - return f"Deleted memory: {memory_key}" -``` - -### Add Helper Functions - -```python -def _get_user_id(request: ResponsesAgentRequest) -> str: - """Extract user_id from request context or custom inputs.""" - custom_inputs = dict(request.custom_inputs or {}) - if "user_id" in custom_inputs and custom_inputs["user_id"]: - return str(custom_inputs["user_id"]) - if request.context and getattr(request.context, "user_id", None): - return str(request.context.user_id) - return "default-user" - - -def _get_or_create_thread_id(request: ResponsesAgentRequest) -> str: - """Extract or create thread_id for conversation tracking.""" - custom_inputs = dict(request.custom_inputs or {}) - if "thread_id" in custom_inputs and custom_inputs["thread_id"]: - return str(custom_inputs["thread_id"]) - if request.context and getattr(request.context, "conversation_id", None): - return str(request.context.conversation_id) - return str(uuid.uuid4()) -``` - -### Update Streaming Function - -```python @stream() -async def streaming( - request: ResponsesAgentRequest, -) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: - # Get user context - user_id = _get_user_id(request) - thread_id = _get_or_create_thread_id(request) +async def streaming(request: ResponsesAgentRequest): + user_id = get_user_id(request) - # Get other tools (MCP, etc.) - mcp_client = init_mcp_client(sp_workspace_client) - mcp_tools = await mcp_client.get_tools() - - # Memory tools - memory_tools = [get_user_memory, save_user_memory, delete_user_memory] - - # Combine all tools - all_tools = mcp_tools + memory_tools - - # Initialize model - model = ChatDatabricks(endpoint="databricks-claude-sonnet-4") - - # Use AsyncDatabricksStore for long-term memory async with AsyncDatabricksStore( instance_name=LAKEBASE_INSTANCE_NAME, embedding_endpoint=EMBEDDING_ENDPOINT, - embedding_dims=EMBEDDING_DIMS, # REQUIRED! + embedding_dims=EMBEDDING_DIMS, ) as store: - # Create agent with tools - agent = create_react_agent( - model=model, - tools=all_tools, - prompt=AGENT_INSTRUCTIONS, - ) - - # Prepare input - messages = {"messages": to_chat_completions_input([i.model_dump() for i in request.input])} - - # Configure with user context for memory tools - config = { - "configurable": { - "user_id": user_id, - "thread_id": thread_id, - "store": store, # Pass store to tools via config - } - } - - # Stream agent responses - async for event in process_agent_astream_events( - agent.astream(input=messages, config=config, stream_mode=["updates", "messages"]) - ): + await store.setup() # Creates tables if needed + + tools = await mcp_client.get_tools() + memory_tools() + config = {"configurable": {"user_id": user_id, "store": store}} + + agent = create_react_agent(model=model, tools=tools) + async for event in agent.astream(messages, config): yield event ``` --- -## Step 5: Initialize Tables and Deploy +## Initialize Tables and Deploy ### Initialize Lakebase Tables (First Time Only) @@ -313,118 +250,33 @@ EOF )" ``` -### Deploy and Run - -**IMPORTANT:** Always run `databricks bundle run` after `databricks bundle deploy` to start/restart the app with the new code: - -```bash -# Deploy resources and upload files -databricks bundle deploy - -# Start/restart the app with new code (REQUIRED!) -databricks bundle run agent_langgraph -``` - -> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is required to actually start the app with the new code. - ---- - -## Complete Example Files - -### databricks.yml (with Lakebase) - -```yaml -bundle: - name: agent_langgraph - -resources: - experiments: - agent_langgraph_experiment: - name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} - - apps: - agent_langgraph: - name: "my-agent-app" - description: "Agent with long-term memory" - source_code_path: ./ - - resources: - - name: 'experiment' - experiment: - experiment_id: "${resources.experiments.agent_langgraph_experiment.id}" - permission: 'CAN_MANAGE' - - # Lakebase instance for long-term memory - - name: 'database' - database: - instance_name: '' - database_name: 'postgres' - permission: 'CAN_CONNECT_AND_CREATE' - -targets: - dev: - mode: development - default: true -``` - -### app.yaml +### Deploy -```yaml -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" - # Reference experiment resource from databricks.yml - - name: MLFLOW_EXPERIMENT_ID - valueFrom: "experiment" - # Lakebase instance name (must match instance_name in databricks.yml) - - name: LAKEBASE_INSTANCE_NAME - value: "" - # Embedding configuration - - name: EMBEDDING_ENDPOINT - value: "databricks-gte-large-en" - - name: EMBEDDING_DIMS - value: "1024" -``` +After initializing tables, deploy your agent. See **deploy** skill for full instructions. --- -## Adding Short-Term Memory - -Short-term memory stores conversation history within a thread. +## Short-Term Memory -### Import and Initialize Checkpointer +For conversation history within a session, use `AsyncCheckpointSaver`: ```python from databricks_langchain import AsyncCheckpointSaver -LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") - async with AsyncCheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME) as checkpointer: agent = create_react_agent( model=model, tools=tools, - checkpointer=checkpointer, # Enables conversation persistence + checkpointer=checkpointer, ) -``` -### Use thread_id in Agent Config - -```python -config = {"configurable": {"thread_id": thread_id}} -async for event in agent.astream(input_state, config, stream_mode=["updates", "messages"]): - yield event + config = {"configurable": {"thread_id": thread_id}} + async for event in agent.astream(messages, config): + yield event ``` +See the [agent-langgraph-short-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-short-term-memory) template for a complete implementation. + --- ## Testing Memory @@ -484,11 +336,10 @@ curl -X POST https:///invocations \ - [ ] Run `uv sync` to install dependencies - [ ] Created or identified Lakebase instance - [ ] Added Lakebase env vars to `.env` (for local dev) -- [ ] Added `database` resource to `databricks.yml` (for permissions) -- [ ] Added `LAKEBASE_INSTANCE_NAME` value to `app.yaml` (matching instance_name in databricks.yml) +- [ ] Added `database` resource to `databricks.yml` +- [ ] Added `LAKEBASE_INSTANCE_NAME` to `app.yaml` - [ ] **Initialized tables locally** by running `await store.setup()` -- [ ] Deployed with `databricks bundle deploy` -- [ ] **Started app with `databricks bundle run agent_langgraph`** +- [ ] Deployed with `databricks bundle deploy && databricks bundle run` --- @@ -498,9 +349,9 @@ curl -X POST https:///invocations \ |-------|-------|----------| | **"embedding_dims is required"** | Missing parameter | Add `embedding_dims=1024` to AsyncDatabricksStore | | **"relation 'store' does not exist"** | Tables not created | Run `await store.setup()` locally first | -| **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `valueFrom: "database"` to app.yaml | -| **"permission denied for table store"** | Missing grants | Lakebase `database` resource in DAB should handle this | -| **"Memory not available"** | No user_id in request | Ensure `custom_inputs.user_id` is passed | +| **"Unable to resolve Lakebase instance 'None'"** | Missing env var | Check `LAKEBASE_INSTANCE_NAME` in app.yaml | +| **"permission denied for table store"** | Missing grants | Add `database` resource to databricks.yml | +| **"Memory not available - no user_id"** | Missing user_id | Pass `custom_inputs.user_id` in request | | **Memory not persisting** | Different user_ids | Use consistent user_id across requests | | **App not updated after deploy** | Forgot to run bundle | Run `databricks bundle run agent_langgraph` after deploy | @@ -512,8 +363,8 @@ For fully configured implementations without manual setup: | Template | Memory Type | Key Features | |----------|-------------|--------------| -| `agent-langgraph-short-term-memory` | Short-term | AsyncCheckpointSaver, thread_id | -| `agent-langgraph-long-term-memory` | Long-term | AsyncDatabricksStore, memory tools | +| [agent-langgraph-short-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-short-term-memory) | Short-term | AsyncCheckpointSaver, thread_id | +| [agent-langgraph-long-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory) | Long-term | AsyncDatabricksStore, memory tools | --- diff --git a/agent-langgraph-short-term-memory/.claude/skills/lakebase-setup/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/lakebase-setup/SKILL.md index 6b8c802e..65c79099 100644 --- a/agent-langgraph-short-term-memory/.claude/skills/lakebase-setup/SKILL.md +++ b/agent-langgraph-short-term-memory/.claude/skills/lakebase-setup/SKILL.md @@ -82,7 +82,7 @@ resources: ``` **Important:** -- The `name: 'database'` must match the `valueFrom` reference in `app.yaml` +- The `instance_name: ''` must match the `value` reference in `app.yaml` - Using the `database` resource type automatically grants the app's service principal access to Lakebase ### Update app.yaml (Environment Variables) diff --git a/agent-langgraph/.claude/skills/agent-memory/SKILL.md b/agent-langgraph/.claude/skills/agent-memory/SKILL.md index 83d59668..c87e21a3 100644 --- a/agent-langgraph/.claude/skills/agent-memory/SKILL.md +++ b/agent-langgraph/.claude/skills/agent-memory/SKILL.md @@ -6,8 +6,8 @@ description: "Add memory capabilities to your agent. Use when: (1) User asks abo # Adding Memory to Your Agent > **Note:** This template does not include memory by default. Use this skill to **add memory capabilities**. For pre-configured memory templates, see: -> - `agent-langgraph-short-term-memory` - Conversation history within a session -> - `agent-langgraph-long-term-memory` - User facts that persist across sessions +> - [agent-langgraph-short-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-short-term-memory) - Conversation history within a session +> - [agent-langgraph-long-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory) - User facts that persist across sessions ## Memory Types @@ -39,17 +39,111 @@ Adding memory requires changes to **4 files**: | File | What to Add | |------|-------------| -| `pyproject.toml` | Memory dependency + hatch config | +| `pyproject.toml` | Memory dependency | | `.env` | Lakebase env vars (for local dev) | | `databricks.yml` | Lakebase database resource | -| `app.yaml` | `valueFrom` reference to lakebase resource | +| `app.yaml` | Environment variables for Lakebase | | `agent_server/agent.py` | Memory tools and AsyncDatabricksStore | --- -## Step 1: Configure databricks.yml (Lakebase Resource) +## Key Principles -Add the Lakebase database resource to your app in `databricks.yml`: +Before implementing memory, understand these patterns from the production implementation. + +### 1. Factory Function Pattern + +Memory tools should be returned from a factory function, not defined as standalone functions: + +```python +def memory_tools(): + @tool + async def get_user_memory(query: str, config: RunnableConfig) -> str: + ... + @tool + async def save_user_memory(memory_key: str, memory_data_json: str, config: RunnableConfig) -> str: + ... + @tool + async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: + ... + return [get_user_memory, save_user_memory, delete_user_memory] +``` + +### 2. User ID Extraction + +Extract `user_id` from the request, checking `custom_inputs` first. Return `None` (not a default) to let the caller decide: + +```python +def get_user_id(request: ResponsesAgentRequest) -> Optional[str]: + custom_inputs = dict(request.custom_inputs or {}) + if "user_id" in custom_inputs: + return custom_inputs["user_id"] + if request.context and getattr(request.context, "user_id", None): + return request.context.user_id + return None +``` + +### 3. Separate Error Handling + +Check `user_id` and `store` separately with distinct error messages: + +```python +user_id = config.get("configurable", {}).get("user_id") +if not user_id: + return "Memory not available - no user_id provided." + +store: Optional[BaseStore] = config.get("configurable", {}).get("store") +if not store: + return "Memory not available - store not configured." +``` + +### 4. JSON Validation for Save + +Validate JSON input before storing - the LLM may pass invalid JSON: + +```python +try: + memory_data = json.loads(memory_data_json) + if not isinstance(memory_data, dict): + return f"Failed: memory_data must be a JSON object, not {type(memory_data).__name__}" + await store.aput(namespace, memory_key, memory_data) +except json.JSONDecodeError as e: + return f"Failed to save memory: Invalid JSON - {e}" +``` + +### 5. Pass Store via RunnableConfig + +Pass the store through config, not as a function parameter: + +```python +config = {"configurable": {"user_id": user_id, "store": store}} +# Tools access via: config.get("configurable", {}).get("store") +``` + +--- + +## Production Reference + +For complete, tested implementations: + +| File | Description | +|------|-------------| +| [`agent-langgraph-long-term-memory/agent_server/utils_memory.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/utils_memory.py) | Memory tools factory, helpers, error handling | +| [`agent-langgraph-long-term-memory/agent_server/agent.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/agent.py) | Integration with agent, store initialization | + +Key functions in `utils_memory.py`: +- `memory_tools()` - Factory returning get/save/delete tools +- `get_user_id()` - Extract user_id from request +- `resolve_lakebase_instance_name()` - Handle hostname vs instance name +- `get_lakebase_access_error_message()` - Helpful error messages + +--- + +## Configuration Files + +### Step 1: databricks.yml (Lakebase Resource) + +Add the Lakebase database resource to your app: ```yaml resources: @@ -62,7 +156,7 @@ resources: # ... other resources (experiment, UC functions, etc.) ... # Lakebase instance for long-term memory - - name: 'lakebadatabasese_memory' + - name: 'database' database: instance_name: '' database_name: 'postgres' @@ -71,11 +165,7 @@ resources: **Important:** The `name: 'database'` must match the `valueFrom` reference in `app.yaml`. ---- - -## Step 2: Configure app.yaml (Environment Variables) - -Update `app.yaml` with the Lakebase environment variables: +### Step 2: app.yaml (Environment Variables) ```yaml command: ["uv", "run", "start-app"] @@ -83,34 +173,20 @@ command: ["uv", "run", "start-app"] env: # ... other env vars ... - # MLflow experiment (uses valueFrom to reference databricks.yml resource) - - name: MLFLOW_EXPERIMENT_ID - valueFrom: "experiment" - - # Lakebase instance name - must match the instance_name in databricks.yml database resource - # Note: Use 'value' (not 'valueFrom') because AsyncDatabricksStore needs the instance name, - # not the full connection string that valueFrom would provide + # Lakebase instance name - name: LAKEBASE_INSTANCE_NAME value: "" - # Embedding configuration (static values) + # Embedding configuration - name: EMBEDDING_ENDPOINT value: "databricks-gte-large-en" - name: EMBEDDING_DIMS value: "1024" ``` -**Important:** The `LAKEBASE_INSTANCE_NAME` value must match the `instance_name` in your `databricks.yml` database resource. The `database` resource handles permissions, while `app.yaml` provides the instance name to your code. +**Important:** `LAKEBASE_INSTANCE_NAME` must match `instance_name` in databricks.yml. -### Why not valueFrom for Lakebase? - -The `database` resource's `valueFrom` provides the full connection string (e.g., `instance-xxx.database.staging.cloud.databricks.com`), but `AsyncDatabricksStore` expects just the instance name. So we use a static `value` instead. - ---- - -## Step 3: Configure .env (Local Development) - -For local development, add to `.env`: +### Step 3: .env (Local Development) ```bash # Lakebase configuration for long-term memory @@ -119,176 +195,37 @@ EMBEDDING_ENDPOINT=databricks-gte-large-en EMBEDDING_DIMS=1024 ``` -> **Note:** `.env` is only for local development. When deployed, the app gets `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. - --- -## Step 4: Update agent.py +## Integration Example -### Add Imports and Configuration +Minimal example showing how to integrate memory into your streaming function: ```python -import os -import uuid -from typing import AsyncGenerator - -from databricks_langchain import AsyncDatabricksStore, ChatDatabricks -from langchain_core.runnables import RunnableConfig -from langchain_core.tools import tool -from langgraph.prebuilt import create_react_agent - -# Environment configuration -LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") -EMBEDDING_ENDPOINT = os.getenv("EMBEDDING_ENDPOINT", "databricks-gte-large-en") -EMBEDDING_DIMS = int(os.getenv("EMBEDDING_DIMS", "1024")) # REQUIRED! -``` - -**CRITICAL:** You MUST specify `embedding_dims` when using `embedding_endpoint`. Without it, you'll get: -``` -"embedding_dims is required when embedding_endpoint is specified" -``` +from agent_server.utils_memory import memory_tools, get_user_id -### Create Memory Tools - -```python -@tool -async def get_user_memory(query: str, config: RunnableConfig) -> str: - """Search for relevant information about the user from long-term memory. - Use this to recall preferences, past interactions, or other saved information. - """ - user_id = config.get("configurable", {}).get("user_id") - store = config.get("configurable", {}).get("store") - if not user_id or not store: - return "Memory not available - no user context" - - # Sanitize user_id for namespace (replace special chars) - namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) - results = await store.asearch(namespace, query=query, limit=5) - if not results: - return "No memories found for this query" - return "\n".join([f"- {r.value}" for r in results]) - - -@tool -async def save_user_memory(memory_key: str, memory_data: str, config: RunnableConfig) -> str: - """Save information about the user to long-term memory. - Use this to remember user preferences, important details, or other - information that should persist across conversations. - - Args: - memory_key: A short descriptive key (e.g., "preferred_name", "team", "region") - memory_data: The information to save - """ - user_id = config.get("configurable", {}).get("user_id") - store = config.get("configurable", {}).get("store") - if not user_id or not store: - return "Memory not available - no user context" - - namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) - await store.aput(namespace, memory_key, {"value": memory_data}) - return f"Saved memory: {memory_key}" - - -@tool -async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: - """Delete a specific memory from the user's long-term memory. - Use this when the user asks to forget something or correct stored information. - - Args: - memory_key: The key of the memory to delete (e.g., "preferred_name", "team") - """ - user_id = config.get("configurable", {}).get("user_id") - store = config.get("configurable", {}).get("store") - if not user_id or not store: - return "Memory not available - no user context" - - namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) - await store.adelete(namespace, memory_key) - return f"Deleted memory: {memory_key}" -``` - -### Add Helper Functions - -```python -def _get_user_id(request: ResponsesAgentRequest) -> str: - """Extract user_id from request context or custom inputs.""" - custom_inputs = dict(request.custom_inputs or {}) - if "user_id" in custom_inputs and custom_inputs["user_id"]: - return str(custom_inputs["user_id"]) - if request.context and getattr(request.context, "user_id", None): - return str(request.context.user_id) - return "default-user" - - -def _get_or_create_thread_id(request: ResponsesAgentRequest) -> str: - """Extract or create thread_id for conversation tracking.""" - custom_inputs = dict(request.custom_inputs or {}) - if "thread_id" in custom_inputs and custom_inputs["thread_id"]: - return str(custom_inputs["thread_id"]) - if request.context and getattr(request.context, "conversation_id", None): - return str(request.context.conversation_id) - return str(uuid.uuid4()) -``` - -### Update Streaming Function - -```python @stream() -async def streaming( - request: ResponsesAgentRequest, -) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: - # Get user context - user_id = _get_user_id(request) - thread_id = _get_or_create_thread_id(request) +async def streaming(request: ResponsesAgentRequest): + user_id = get_user_id(request) - # Get other tools (MCP, etc.) - mcp_client = init_mcp_client(sp_workspace_client) - mcp_tools = await mcp_client.get_tools() - - # Memory tools - memory_tools = [get_user_memory, save_user_memory, delete_user_memory] - - # Combine all tools - all_tools = mcp_tools + memory_tools - - # Initialize model - model = ChatDatabricks(endpoint="databricks-claude-sonnet-4") - - # Use AsyncDatabricksStore for long-term memory async with AsyncDatabricksStore( instance_name=LAKEBASE_INSTANCE_NAME, embedding_endpoint=EMBEDDING_ENDPOINT, - embedding_dims=EMBEDDING_DIMS, # REQUIRED! + embedding_dims=EMBEDDING_DIMS, ) as store: - # Create agent with tools - agent = create_react_agent( - model=model, - tools=all_tools, - prompt=AGENT_INSTRUCTIONS, - ) - - # Prepare input - messages = {"messages": to_chat_completions_input([i.model_dump() for i in request.input])} - - # Configure with user context for memory tools - config = { - "configurable": { - "user_id": user_id, - "thread_id": thread_id, - "store": store, # Pass store to tools via config - } - } - - # Stream agent responses - async for event in process_agent_astream_events( - agent.astream(input=messages, config=config, stream_mode=["updates", "messages"]) - ): + await store.setup() # Creates tables if needed + + tools = await mcp_client.get_tools() + memory_tools() + config = {"configurable": {"user_id": user_id, "store": store}} + + agent = create_react_agent(model=model, tools=tools) + async for event in agent.astream(messages, config): yield event ``` --- -## Step 5: Initialize Tables and Deploy +## Initialize Tables and Deploy ### Initialize Lakebase Tables (First Time Only) @@ -313,118 +250,33 @@ EOF )" ``` -### Deploy and Run - -**IMPORTANT:** Always run `databricks bundle run` after `databricks bundle deploy` to start/restart the app with the new code: - -```bash -# Deploy resources and upload files -databricks bundle deploy - -# Start/restart the app with new code (REQUIRED!) -databricks bundle run agent_langgraph -``` - -> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is required to actually start the app with the new code. - ---- - -## Complete Example Files - -### databricks.yml (with Lakebase) - -```yaml -bundle: - name: agent_langgraph - -resources: - experiments: - agent_langgraph_experiment: - name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} - - apps: - agent_langgraph: - name: "my-agent-app" - description: "Agent with long-term memory" - source_code_path: ./ - - resources: - - name: 'experiment' - experiment: - experiment_id: "${resources.experiments.agent_langgraph_experiment.id}" - permission: 'CAN_MANAGE' - - # Lakebase instance for long-term memory - - name: 'database' - database: - instance_name: '' - database_name: 'postgres' - permission: 'CAN_CONNECT_AND_CREATE' - -targets: - dev: - mode: development - default: true -``` - -### app.yaml +### Deploy -```yaml -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" - # Reference experiment resource from databricks.yml - - name: MLFLOW_EXPERIMENT_ID - valueFrom: "experiment" - # Lakebase instance name (must match instance_name in databricks.yml) - - name: LAKEBASE_INSTANCE_NAME - value: "" - # Embedding configuration - - name: EMBEDDING_ENDPOINT - value: "databricks-gte-large-en" - - name: EMBEDDING_DIMS - value: "1024" -``` +After initializing tables, deploy your agent. See **deploy** skill for full instructions. --- -## Adding Short-Term Memory - -Short-term memory stores conversation history within a thread. +## Short-Term Memory -### Import and Initialize Checkpointer +For conversation history within a session, use `AsyncCheckpointSaver`: ```python from databricks_langchain import AsyncCheckpointSaver -LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") - async with AsyncCheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME) as checkpointer: agent = create_react_agent( model=model, tools=tools, - checkpointer=checkpointer, # Enables conversation persistence + checkpointer=checkpointer, ) -``` -### Use thread_id in Agent Config - -```python -config = {"configurable": {"thread_id": thread_id}} -async for event in agent.astream(input_state, config, stream_mode=["updates", "messages"]): - yield event + config = {"configurable": {"thread_id": thread_id}} + async for event in agent.astream(messages, config): + yield event ``` +See the [agent-langgraph-short-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-short-term-memory) template for a complete implementation. + --- ## Testing Memory @@ -484,11 +336,10 @@ curl -X POST https:///invocations \ - [ ] Run `uv sync` to install dependencies - [ ] Created or identified Lakebase instance - [ ] Added Lakebase env vars to `.env` (for local dev) -- [ ] Added `database` resource to `databricks.yml` (for permissions) -- [ ] Added `LAKEBASE_INSTANCE_NAME` value to `app.yaml` (matching instance_name in databricks.yml) +- [ ] Added `database` resource to `databricks.yml` +- [ ] Added `LAKEBASE_INSTANCE_NAME` to `app.yaml` - [ ] **Initialized tables locally** by running `await store.setup()` -- [ ] Deployed with `databricks bundle deploy` -- [ ] **Started app with `databricks bundle run agent_langgraph`** +- [ ] Deployed with `databricks bundle deploy && databricks bundle run` --- @@ -498,9 +349,9 @@ curl -X POST https:///invocations \ |-------|-------|----------| | **"embedding_dims is required"** | Missing parameter | Add `embedding_dims=1024` to AsyncDatabricksStore | | **"relation 'store' does not exist"** | Tables not created | Run `await store.setup()` locally first | -| **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `valueFrom: "database"` to app.yaml | -| **"permission denied for table store"** | Missing grants | Lakebase `database` resource in DAB should handle this | -| **"Memory not available"** | No user_id in request | Ensure `custom_inputs.user_id` is passed | +| **"Unable to resolve Lakebase instance 'None'"** | Missing env var | Check `LAKEBASE_INSTANCE_NAME` in app.yaml | +| **"permission denied for table store"** | Missing grants | Add `database` resource to databricks.yml | +| **"Memory not available - no user_id"** | Missing user_id | Pass `custom_inputs.user_id` in request | | **Memory not persisting** | Different user_ids | Use consistent user_id across requests | | **App not updated after deploy** | Forgot to run bundle | Run `databricks bundle run agent_langgraph` after deploy | @@ -512,8 +363,8 @@ For fully configured implementations without manual setup: | Template | Memory Type | Key Features | |----------|-------------|--------------| -| `agent-langgraph-short-term-memory` | Short-term | AsyncCheckpointSaver, thread_id | -| `agent-langgraph-long-term-memory` | Long-term | AsyncDatabricksStore, memory tools | +| [agent-langgraph-short-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-short-term-memory) | Short-term | AsyncCheckpointSaver, thread_id | +| [agent-langgraph-long-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory) | Long-term | AsyncDatabricksStore, memory tools | --- diff --git a/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md b/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md index 6b8c802e..65c79099 100644 --- a/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md +++ b/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md @@ -82,7 +82,7 @@ resources: ``` **Important:** -- The `name: 'database'` must match the `valueFrom` reference in `app.yaml` +- The `instance_name: ''` must match the `value` reference in `app.yaml` - Using the `database` resource type automatically grants the app's service principal access to Lakebase ### Update app.yaml (Environment Variables) diff --git a/agent-non-conversational/.claude/skills/agent-memory/SKILL.md b/agent-non-conversational/.claude/skills/agent-memory/SKILL.md index 83d59668..c87e21a3 100644 --- a/agent-non-conversational/.claude/skills/agent-memory/SKILL.md +++ b/agent-non-conversational/.claude/skills/agent-memory/SKILL.md @@ -6,8 +6,8 @@ description: "Add memory capabilities to your agent. Use when: (1) User asks abo # Adding Memory to Your Agent > **Note:** This template does not include memory by default. Use this skill to **add memory capabilities**. For pre-configured memory templates, see: -> - `agent-langgraph-short-term-memory` - Conversation history within a session -> - `agent-langgraph-long-term-memory` - User facts that persist across sessions +> - [agent-langgraph-short-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-short-term-memory) - Conversation history within a session +> - [agent-langgraph-long-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory) - User facts that persist across sessions ## Memory Types @@ -39,17 +39,111 @@ Adding memory requires changes to **4 files**: | File | What to Add | |------|-------------| -| `pyproject.toml` | Memory dependency + hatch config | +| `pyproject.toml` | Memory dependency | | `.env` | Lakebase env vars (for local dev) | | `databricks.yml` | Lakebase database resource | -| `app.yaml` | `valueFrom` reference to lakebase resource | +| `app.yaml` | Environment variables for Lakebase | | `agent_server/agent.py` | Memory tools and AsyncDatabricksStore | --- -## Step 1: Configure databricks.yml (Lakebase Resource) +## Key Principles -Add the Lakebase database resource to your app in `databricks.yml`: +Before implementing memory, understand these patterns from the production implementation. + +### 1. Factory Function Pattern + +Memory tools should be returned from a factory function, not defined as standalone functions: + +```python +def memory_tools(): + @tool + async def get_user_memory(query: str, config: RunnableConfig) -> str: + ... + @tool + async def save_user_memory(memory_key: str, memory_data_json: str, config: RunnableConfig) -> str: + ... + @tool + async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: + ... + return [get_user_memory, save_user_memory, delete_user_memory] +``` + +### 2. User ID Extraction + +Extract `user_id` from the request, checking `custom_inputs` first. Return `None` (not a default) to let the caller decide: + +```python +def get_user_id(request: ResponsesAgentRequest) -> Optional[str]: + custom_inputs = dict(request.custom_inputs or {}) + if "user_id" in custom_inputs: + return custom_inputs["user_id"] + if request.context and getattr(request.context, "user_id", None): + return request.context.user_id + return None +``` + +### 3. Separate Error Handling + +Check `user_id` and `store` separately with distinct error messages: + +```python +user_id = config.get("configurable", {}).get("user_id") +if not user_id: + return "Memory not available - no user_id provided." + +store: Optional[BaseStore] = config.get("configurable", {}).get("store") +if not store: + return "Memory not available - store not configured." +``` + +### 4. JSON Validation for Save + +Validate JSON input before storing - the LLM may pass invalid JSON: + +```python +try: + memory_data = json.loads(memory_data_json) + if not isinstance(memory_data, dict): + return f"Failed: memory_data must be a JSON object, not {type(memory_data).__name__}" + await store.aput(namespace, memory_key, memory_data) +except json.JSONDecodeError as e: + return f"Failed to save memory: Invalid JSON - {e}" +``` + +### 5. Pass Store via RunnableConfig + +Pass the store through config, not as a function parameter: + +```python +config = {"configurable": {"user_id": user_id, "store": store}} +# Tools access via: config.get("configurable", {}).get("store") +``` + +--- + +## Production Reference + +For complete, tested implementations: + +| File | Description | +|------|-------------| +| [`agent-langgraph-long-term-memory/agent_server/utils_memory.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/utils_memory.py) | Memory tools factory, helpers, error handling | +| [`agent-langgraph-long-term-memory/agent_server/agent.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/agent.py) | Integration with agent, store initialization | + +Key functions in `utils_memory.py`: +- `memory_tools()` - Factory returning get/save/delete tools +- `get_user_id()` - Extract user_id from request +- `resolve_lakebase_instance_name()` - Handle hostname vs instance name +- `get_lakebase_access_error_message()` - Helpful error messages + +--- + +## Configuration Files + +### Step 1: databricks.yml (Lakebase Resource) + +Add the Lakebase database resource to your app: ```yaml resources: @@ -62,7 +156,7 @@ resources: # ... other resources (experiment, UC functions, etc.) ... # Lakebase instance for long-term memory - - name: 'lakebadatabasese_memory' + - name: 'database' database: instance_name: '' database_name: 'postgres' @@ -71,11 +165,7 @@ resources: **Important:** The `name: 'database'` must match the `valueFrom` reference in `app.yaml`. ---- - -## Step 2: Configure app.yaml (Environment Variables) - -Update `app.yaml` with the Lakebase environment variables: +### Step 2: app.yaml (Environment Variables) ```yaml command: ["uv", "run", "start-app"] @@ -83,34 +173,20 @@ command: ["uv", "run", "start-app"] env: # ... other env vars ... - # MLflow experiment (uses valueFrom to reference databricks.yml resource) - - name: MLFLOW_EXPERIMENT_ID - valueFrom: "experiment" - - # Lakebase instance name - must match the instance_name in databricks.yml database resource - # Note: Use 'value' (not 'valueFrom') because AsyncDatabricksStore needs the instance name, - # not the full connection string that valueFrom would provide + # Lakebase instance name - name: LAKEBASE_INSTANCE_NAME value: "" - # Embedding configuration (static values) + # Embedding configuration - name: EMBEDDING_ENDPOINT value: "databricks-gte-large-en" - name: EMBEDDING_DIMS value: "1024" ``` -**Important:** The `LAKEBASE_INSTANCE_NAME` value must match the `instance_name` in your `databricks.yml` database resource. The `database` resource handles permissions, while `app.yaml` provides the instance name to your code. +**Important:** `LAKEBASE_INSTANCE_NAME` must match `instance_name` in databricks.yml. -### Why not valueFrom for Lakebase? - -The `database` resource's `valueFrom` provides the full connection string (e.g., `instance-xxx.database.staging.cloud.databricks.com`), but `AsyncDatabricksStore` expects just the instance name. So we use a static `value` instead. - ---- - -## Step 3: Configure .env (Local Development) - -For local development, add to `.env`: +### Step 3: .env (Local Development) ```bash # Lakebase configuration for long-term memory @@ -119,176 +195,37 @@ EMBEDDING_ENDPOINT=databricks-gte-large-en EMBEDDING_DIMS=1024 ``` -> **Note:** `.env` is only for local development. When deployed, the app gets `LAKEBASE_INSTANCE_NAME` from the `valueFrom` reference in `app.yaml`. - --- -## Step 4: Update agent.py +## Integration Example -### Add Imports and Configuration +Minimal example showing how to integrate memory into your streaming function: ```python -import os -import uuid -from typing import AsyncGenerator - -from databricks_langchain import AsyncDatabricksStore, ChatDatabricks -from langchain_core.runnables import RunnableConfig -from langchain_core.tools import tool -from langgraph.prebuilt import create_react_agent - -# Environment configuration -LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") -EMBEDDING_ENDPOINT = os.getenv("EMBEDDING_ENDPOINT", "databricks-gte-large-en") -EMBEDDING_DIMS = int(os.getenv("EMBEDDING_DIMS", "1024")) # REQUIRED! -``` - -**CRITICAL:** You MUST specify `embedding_dims` when using `embedding_endpoint`. Without it, you'll get: -``` -"embedding_dims is required when embedding_endpoint is specified" -``` +from agent_server.utils_memory import memory_tools, get_user_id -### Create Memory Tools - -```python -@tool -async def get_user_memory(query: str, config: RunnableConfig) -> str: - """Search for relevant information about the user from long-term memory. - Use this to recall preferences, past interactions, or other saved information. - """ - user_id = config.get("configurable", {}).get("user_id") - store = config.get("configurable", {}).get("store") - if not user_id or not store: - return "Memory not available - no user context" - - # Sanitize user_id for namespace (replace special chars) - namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) - results = await store.asearch(namespace, query=query, limit=5) - if not results: - return "No memories found for this query" - return "\n".join([f"- {r.value}" for r in results]) - - -@tool -async def save_user_memory(memory_key: str, memory_data: str, config: RunnableConfig) -> str: - """Save information about the user to long-term memory. - Use this to remember user preferences, important details, or other - information that should persist across conversations. - - Args: - memory_key: A short descriptive key (e.g., "preferred_name", "team", "region") - memory_data: The information to save - """ - user_id = config.get("configurable", {}).get("user_id") - store = config.get("configurable", {}).get("store") - if not user_id or not store: - return "Memory not available - no user context" - - namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) - await store.aput(namespace, memory_key, {"value": memory_data}) - return f"Saved memory: {memory_key}" - - -@tool -async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: - """Delete a specific memory from the user's long-term memory. - Use this when the user asks to forget something or correct stored information. - - Args: - memory_key: The key of the memory to delete (e.g., "preferred_name", "team") - """ - user_id = config.get("configurable", {}).get("user_id") - store = config.get("configurable", {}).get("store") - if not user_id or not store: - return "Memory not available - no user context" - - namespace = ("user_memories", user_id.replace(".", "-").replace("@", "-")) - await store.adelete(namespace, memory_key) - return f"Deleted memory: {memory_key}" -``` - -### Add Helper Functions - -```python -def _get_user_id(request: ResponsesAgentRequest) -> str: - """Extract user_id from request context or custom inputs.""" - custom_inputs = dict(request.custom_inputs or {}) - if "user_id" in custom_inputs and custom_inputs["user_id"]: - return str(custom_inputs["user_id"]) - if request.context and getattr(request.context, "user_id", None): - return str(request.context.user_id) - return "default-user" - - -def _get_or_create_thread_id(request: ResponsesAgentRequest) -> str: - """Extract or create thread_id for conversation tracking.""" - custom_inputs = dict(request.custom_inputs or {}) - if "thread_id" in custom_inputs and custom_inputs["thread_id"]: - return str(custom_inputs["thread_id"]) - if request.context and getattr(request.context, "conversation_id", None): - return str(request.context.conversation_id) - return str(uuid.uuid4()) -``` - -### Update Streaming Function - -```python @stream() -async def streaming( - request: ResponsesAgentRequest, -) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: - # Get user context - user_id = _get_user_id(request) - thread_id = _get_or_create_thread_id(request) +async def streaming(request: ResponsesAgentRequest): + user_id = get_user_id(request) - # Get other tools (MCP, etc.) - mcp_client = init_mcp_client(sp_workspace_client) - mcp_tools = await mcp_client.get_tools() - - # Memory tools - memory_tools = [get_user_memory, save_user_memory, delete_user_memory] - - # Combine all tools - all_tools = mcp_tools + memory_tools - - # Initialize model - model = ChatDatabricks(endpoint="databricks-claude-sonnet-4") - - # Use AsyncDatabricksStore for long-term memory async with AsyncDatabricksStore( instance_name=LAKEBASE_INSTANCE_NAME, embedding_endpoint=EMBEDDING_ENDPOINT, - embedding_dims=EMBEDDING_DIMS, # REQUIRED! + embedding_dims=EMBEDDING_DIMS, ) as store: - # Create agent with tools - agent = create_react_agent( - model=model, - tools=all_tools, - prompt=AGENT_INSTRUCTIONS, - ) - - # Prepare input - messages = {"messages": to_chat_completions_input([i.model_dump() for i in request.input])} - - # Configure with user context for memory tools - config = { - "configurable": { - "user_id": user_id, - "thread_id": thread_id, - "store": store, # Pass store to tools via config - } - } - - # Stream agent responses - async for event in process_agent_astream_events( - agent.astream(input=messages, config=config, stream_mode=["updates", "messages"]) - ): + await store.setup() # Creates tables if needed + + tools = await mcp_client.get_tools() + memory_tools() + config = {"configurable": {"user_id": user_id, "store": store}} + + agent = create_react_agent(model=model, tools=tools) + async for event in agent.astream(messages, config): yield event ``` --- -## Step 5: Initialize Tables and Deploy +## Initialize Tables and Deploy ### Initialize Lakebase Tables (First Time Only) @@ -313,118 +250,33 @@ EOF )" ``` -### Deploy and Run - -**IMPORTANT:** Always run `databricks bundle run` after `databricks bundle deploy` to start/restart the app with the new code: - -```bash -# Deploy resources and upload files -databricks bundle deploy - -# Start/restart the app with new code (REQUIRED!) -databricks bundle run agent_langgraph -``` - -> **Note:** `bundle deploy` only uploads files and configures resources. `bundle run` is required to actually start the app with the new code. - ---- - -## Complete Example Files - -### databricks.yml (with Lakebase) - -```yaml -bundle: - name: agent_langgraph - -resources: - experiments: - agent_langgraph_experiment: - name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} - - apps: - agent_langgraph: - name: "my-agent-app" - description: "Agent with long-term memory" - source_code_path: ./ - - resources: - - name: 'experiment' - experiment: - experiment_id: "${resources.experiments.agent_langgraph_experiment.id}" - permission: 'CAN_MANAGE' - - # Lakebase instance for long-term memory - - name: 'database' - database: - instance_name: '' - database_name: 'postgres' - permission: 'CAN_CONNECT_AND_CREATE' - -targets: - dev: - mode: development - default: true -``` - -### app.yaml +### Deploy -```yaml -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" - # Reference experiment resource from databricks.yml - - name: MLFLOW_EXPERIMENT_ID - valueFrom: "experiment" - # Lakebase instance name (must match instance_name in databricks.yml) - - name: LAKEBASE_INSTANCE_NAME - value: "" - # Embedding configuration - - name: EMBEDDING_ENDPOINT - value: "databricks-gte-large-en" - - name: EMBEDDING_DIMS - value: "1024" -``` +After initializing tables, deploy your agent. See **deploy** skill for full instructions. --- -## Adding Short-Term Memory - -Short-term memory stores conversation history within a thread. +## Short-Term Memory -### Import and Initialize Checkpointer +For conversation history within a session, use `AsyncCheckpointSaver`: ```python from databricks_langchain import AsyncCheckpointSaver -LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME") - async with AsyncCheckpointSaver(instance_name=LAKEBASE_INSTANCE_NAME) as checkpointer: agent = create_react_agent( model=model, tools=tools, - checkpointer=checkpointer, # Enables conversation persistence + checkpointer=checkpointer, ) -``` -### Use thread_id in Agent Config - -```python -config = {"configurable": {"thread_id": thread_id}} -async for event in agent.astream(input_state, config, stream_mode=["updates", "messages"]): - yield event + config = {"configurable": {"thread_id": thread_id}} + async for event in agent.astream(messages, config): + yield event ``` +See the [agent-langgraph-short-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-short-term-memory) template for a complete implementation. + --- ## Testing Memory @@ -484,11 +336,10 @@ curl -X POST https:///invocations \ - [ ] Run `uv sync` to install dependencies - [ ] Created or identified Lakebase instance - [ ] Added Lakebase env vars to `.env` (for local dev) -- [ ] Added `database` resource to `databricks.yml` (for permissions) -- [ ] Added `LAKEBASE_INSTANCE_NAME` value to `app.yaml` (matching instance_name in databricks.yml) +- [ ] Added `database` resource to `databricks.yml` +- [ ] Added `LAKEBASE_INSTANCE_NAME` to `app.yaml` - [ ] **Initialized tables locally** by running `await store.setup()` -- [ ] Deployed with `databricks bundle deploy` -- [ ] **Started app with `databricks bundle run agent_langgraph`** +- [ ] Deployed with `databricks bundle deploy && databricks bundle run` --- @@ -498,9 +349,9 @@ curl -X POST https:///invocations \ |-------|-------|----------| | **"embedding_dims is required"** | Missing parameter | Add `embedding_dims=1024` to AsyncDatabricksStore | | **"relation 'store' does not exist"** | Tables not created | Run `await store.setup()` locally first | -| **"Unable to resolve Lakebase instance 'None'"** | Missing env var in deployed app | Add `valueFrom: "database"` to app.yaml | -| **"permission denied for table store"** | Missing grants | Lakebase `database` resource in DAB should handle this | -| **"Memory not available"** | No user_id in request | Ensure `custom_inputs.user_id` is passed | +| **"Unable to resolve Lakebase instance 'None'"** | Missing env var | Check `LAKEBASE_INSTANCE_NAME` in app.yaml | +| **"permission denied for table store"** | Missing grants | Add `database` resource to databricks.yml | +| **"Memory not available - no user_id"** | Missing user_id | Pass `custom_inputs.user_id` in request | | **Memory not persisting** | Different user_ids | Use consistent user_id across requests | | **App not updated after deploy** | Forgot to run bundle | Run `databricks bundle run agent_langgraph` after deploy | @@ -512,8 +363,8 @@ For fully configured implementations without manual setup: | Template | Memory Type | Key Features | |----------|-------------|--------------| -| `agent-langgraph-short-term-memory` | Short-term | AsyncCheckpointSaver, thread_id | -| `agent-langgraph-long-term-memory` | Long-term | AsyncDatabricksStore, memory tools | +| [agent-langgraph-short-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-short-term-memory) | Short-term | AsyncCheckpointSaver, thread_id | +| [agent-langgraph-long-term-memory](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory) | Long-term | AsyncDatabricksStore, memory tools | --- diff --git a/agent-non-conversational/.claude/skills/lakebase-setup/SKILL.md b/agent-non-conversational/.claude/skills/lakebase-setup/SKILL.md index 6b8c802e..65c79099 100644 --- a/agent-non-conversational/.claude/skills/lakebase-setup/SKILL.md +++ b/agent-non-conversational/.claude/skills/lakebase-setup/SKILL.md @@ -82,7 +82,7 @@ resources: ``` **Important:** -- The `name: 'database'` must match the `valueFrom` reference in `app.yaml` +- The `instance_name: ''` must match the `value` reference in `app.yaml` - Using the `database` resource type automatically grants the app's service principal access to Lakebase ### Update app.yaml (Environment Variables) From cd36cf2b7f5cefc07bc18e6f546894655c9b0a9c Mon Sep 17 00:00:00 2001 From: Dhruv Gupta Date: Mon, 2 Feb 2026 11:15:50 -0800 Subject: [PATCH 13/14] Add memory_tools.py example for standalone template use - Add examples/memory_tools.py with full implementation - Include all helper functions (get_user_id, resolve_lakebase_instance_name, etc.) - Update SKILL.md with "Complete Example" section and copy command - Synced to all LangGraph templates Co-Authored-By: Claude Opus 4.5 --- .../skills/agent-langgraph-memory/SKILL.md | 15 +- .../examples/memory_tools.py | 228 ++++++++++++++++++ .../.claude/skills/agent-memory/SKILL.md | 15 +- .../agent-memory/examples/memory_tools.py | 228 ++++++++++++++++++ .../.claude/skills/agent-memory/SKILL.md | 15 +- .../agent-memory/examples/memory_tools.py | 228 ++++++++++++++++++ .../.claude/skills/agent-memory/SKILL.md | 15 +- .../agent-memory/examples/memory_tools.py | 228 ++++++++++++++++++ .../.claude/skills/agent-memory/SKILL.md | 15 +- .../agent-memory/examples/memory_tools.py | 228 ++++++++++++++++++ 10 files changed, 1205 insertions(+), 10 deletions(-) create mode 100644 .claude/skills/agent-langgraph-memory/examples/memory_tools.py create mode 100644 agent-langgraph-long-term-memory/.claude/skills/agent-memory/examples/memory_tools.py create mode 100644 agent-langgraph-short-term-memory/.claude/skills/agent-memory/examples/memory_tools.py create mode 100644 agent-langgraph/.claude/skills/agent-memory/examples/memory_tools.py create mode 100644 agent-non-conversational/.claude/skills/agent-memory/examples/memory_tools.py diff --git a/.claude/skills/agent-langgraph-memory/SKILL.md b/.claude/skills/agent-langgraph-memory/SKILL.md index c87e21a3..49b0cb6a 100644 --- a/.claude/skills/agent-langgraph-memory/SKILL.md +++ b/.claude/skills/agent-langgraph-memory/SKILL.md @@ -122,16 +122,27 @@ config = {"configurable": {"user_id": user_id, "store": store}} --- +## Complete Example + +A full implementation is available in this skill's examples folder: + +```bash +# Copy to your project +cp .claude/skills/agent-memory/examples/memory_tools.py agent_server/ +``` + +See `examples/memory_tools.py` for production-ready code including all helper functions. + ## Production Reference -For complete, tested implementations: +For implementations in the pre-built templates: | File | Description | |------|-------------| | [`agent-langgraph-long-term-memory/agent_server/utils_memory.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/utils_memory.py) | Memory tools factory, helpers, error handling | | [`agent-langgraph-long-term-memory/agent_server/agent.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/agent.py) | Integration with agent, store initialization | -Key functions in `utils_memory.py`: +Key functions: - `memory_tools()` - Factory returning get/save/delete tools - `get_user_id()` - Extract user_id from request - `resolve_lakebase_instance_name()` - Handle hostname vs instance name diff --git a/.claude/skills/agent-langgraph-memory/examples/memory_tools.py b/.claude/skills/agent-langgraph-memory/examples/memory_tools.py new file mode 100644 index 00000000..65356287 --- /dev/null +++ b/.claude/skills/agent-langgraph-memory/examples/memory_tools.py @@ -0,0 +1,228 @@ +"""Memory tools for LangGraph agents. + +This module provides tools for managing user long-term memory using +Databricks Lakebase. Copy this file to your agent_server/ directory. + +Usage: + from agent_server.memory_tools import memory_tools, get_user_id + + # In your streaming function: + user_id = get_user_id(request) + tools = await mcp_client.get_tools() + memory_tools() + config = {"configurable": {"user_id": user_id, "store": store}} +""" + +import json +import logging +import os +from typing import Optional + +from databricks.sdk import WorkspaceClient +from langchain_core.runnables import RunnableConfig +from langchain_core.tools import tool +from langgraph.store.base import BaseStore +from mlflow.types.responses import ResponsesAgentRequest + + +# ----------------------------------------------------------------------------- +# Helper Functions +# ----------------------------------------------------------------------------- + + +def get_user_id(request: ResponsesAgentRequest) -> Optional[str]: + """Extract user_id from request context or custom inputs. + + Checks custom_inputs first (for API calls), then request.context + (for Databricks Apps with OBO authentication). + + Returns None if no user_id found - let the caller decide the fallback. + """ + custom_inputs = dict(request.custom_inputs or {}) + if "user_id" in custom_inputs: + return custom_inputs["user_id"] + if request.context and getattr(request.context, "user_id", None): + return request.context.user_id + return None + + +def _is_lakebase_hostname(value: str) -> bool: + """Check if the value looks like a Lakebase hostname rather than an instance name.""" + return ".database." in value and value.endswith(".com") + + +def resolve_lakebase_instance_name( + instance_name: str, workspace_client: Optional[WorkspaceClient] = None +) -> str: + """Resolve a Lakebase instance name from a hostname if needed. + + If the input is a hostname (e.g., from Databricks Apps valueFrom resolution), + this will resolve it to the actual instance name by listing database instances. + + Args: + instance_name: Either an instance name or a hostname + workspace_client: Optional WorkspaceClient to use for resolution + + Returns: + The resolved instance name + + Raises: + ValueError: If the hostname cannot be resolved to an instance name + """ + if not _is_lakebase_hostname(instance_name): + return instance_name + + client = workspace_client or WorkspaceClient() + hostname = instance_name + + try: + instances = list(client.database.list_database_instances()) + except Exception as exc: + raise ValueError( + f"Unable to list database instances to resolve hostname '{hostname}'. " + "Ensure you have access to database instances." + ) from exc + + for instance in instances: + rw_dns = getattr(instance, "read_write_dns", None) + ro_dns = getattr(instance, "read_only_dns", None) + + if hostname in (rw_dns, ro_dns): + resolved_name = getattr(instance, "name", None) + if not resolved_name: + raise ValueError( + f"Found matching instance for hostname '{hostname}' " + "but instance name is not available." + ) + logging.info(f"Resolved Lakebase hostname '{hostname}' to instance name '{resolved_name}'") + return resolved_name + + raise ValueError( + f"Unable to find database instance matching hostname '{hostname}'. " + "Ensure the hostname is correct and the instance exists." + ) + + +def _is_databricks_app_env() -> bool: + """Check if running in a Databricks App environment.""" + return bool(os.getenv("DATABRICKS_APP_NAME")) + + +def get_lakebase_access_error_message(lakebase_instance_name: str) -> str: + """Generate a helpful error message for Lakebase access issues.""" + if _is_databricks_app_env(): + app_name = os.getenv("DATABRICKS_APP_NAME") + return ( + f"Failed to connect to Lakebase instance '{lakebase_instance_name}'. " + f"The App Service Principal for '{app_name}' may not have access.\n\n" + "To fix this:\n" + "1. Go to the Databricks UI and navigate to your app\n" + "2. Click 'Edit' → 'App resources' → 'Add resource'\n" + "3. Add your Lakebase instance as a resource\n" + "4. Grant the necessary permissions on your Lakebase instance." + ) + else: + return ( + f"Failed to connect to Lakebase instance '{lakebase_instance_name}'. " + "Please verify:\n" + "1. The instance name is correct\n" + "2. You have the necessary permissions to access the instance\n" + "3. Your Databricks authentication is configured correctly" + ) + + +# ----------------------------------------------------------------------------- +# Memory Tools Factory +# ----------------------------------------------------------------------------- + + +def memory_tools(): + """Factory function returning memory tools for the agent. + + Returns a list of tools that can be added to your agent: + - get_user_memory: Search for relevant information from long-term memory + - save_user_memory: Save information to long-term memory + - delete_user_memory: Delete a specific memory + + Usage: + tools = await mcp_client.get_tools() + memory_tools() + config = {"configurable": {"user_id": user_id, "store": store}} + """ + + @tool + async def get_user_memory(query: str, config: RunnableConfig) -> str: + """Search for relevant information about the user from long-term memory. + + Use this to recall preferences, past interactions, or other saved information. + + Args: + query: What to search for in the user's memories + """ + user_id = config.get("configurable", {}).get("user_id") + if not user_id: + return "Memory not available - no user_id provided." + + store: Optional[BaseStore] = config.get("configurable", {}).get("store") + if not store: + return "Memory not available - store not configured." + + namespace = ("user_memories", user_id.replace(".", "-")) + results = await store.asearch(namespace, query=query, limit=5) + + if not results: + return "No memories found for this user." + + memory_items = [f"- [{item.key}]: {json.dumps(item.value)}" for item in results] + return f"Found {len(results)} relevant memories:\n" + "\n".join(memory_items) + + @tool + async def save_user_memory(memory_key: str, memory_data_json: str, config: RunnableConfig) -> str: + """Save information about the user to long-term memory. + + Use this to remember user preferences, important details, or other + information that should persist across conversations. + + Args: + memory_key: A short descriptive key (e.g., "preferred_name", "team", "interests") + memory_data_json: JSON object to save (e.g., '{"value": "engineering"}') + """ + user_id = config.get("configurable", {}).get("user_id") + if not user_id: + return "Cannot save memory - no user_id provided." + + store: Optional[BaseStore] = config.get("configurable", {}).get("store") + if not store: + return "Cannot save memory - store not configured." + + namespace = ("user_memories", user_id.replace(".", "-")) + + try: + memory_data = json.loads(memory_data_json) + if not isinstance(memory_data, dict): + return f"Failed: memory_data must be a JSON object, not {type(memory_data).__name__}" + await store.aput(namespace, memory_key, memory_data) + return f"Successfully saved memory '{memory_key}' for user." + except json.JSONDecodeError as e: + return f"Failed to save memory: Invalid JSON - {e}" + + @tool + async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: + """Delete a specific memory from the user's long-term memory. + + Use this when the user asks to forget something or correct stored information. + + Args: + memory_key: The key of the memory to delete (e.g., "preferred_name", "team") + """ + user_id = config.get("configurable", {}).get("user_id") + if not user_id: + return "Cannot delete memory - no user_id provided." + + store: Optional[BaseStore] = config.get("configurable", {}).get("store") + if not store: + return "Cannot delete memory - store not configured." + + namespace = ("user_memories", user_id.replace(".", "-")) + await store.adelete(namespace, memory_key) + return f"Successfully deleted memory '{memory_key}' for user." + + return [get_user_memory, save_user_memory, delete_user_memory] diff --git a/agent-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md b/agent-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md index c87e21a3..49b0cb6a 100644 --- a/agent-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md +++ b/agent-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md @@ -122,16 +122,27 @@ config = {"configurable": {"user_id": user_id, "store": store}} --- +## Complete Example + +A full implementation is available in this skill's examples folder: + +```bash +# Copy to your project +cp .claude/skills/agent-memory/examples/memory_tools.py agent_server/ +``` + +See `examples/memory_tools.py` for production-ready code including all helper functions. + ## Production Reference -For complete, tested implementations: +For implementations in the pre-built templates: | File | Description | |------|-------------| | [`agent-langgraph-long-term-memory/agent_server/utils_memory.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/utils_memory.py) | Memory tools factory, helpers, error handling | | [`agent-langgraph-long-term-memory/agent_server/agent.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/agent.py) | Integration with agent, store initialization | -Key functions in `utils_memory.py`: +Key functions: - `memory_tools()` - Factory returning get/save/delete tools - `get_user_id()` - Extract user_id from request - `resolve_lakebase_instance_name()` - Handle hostname vs instance name diff --git a/agent-langgraph-long-term-memory/.claude/skills/agent-memory/examples/memory_tools.py b/agent-langgraph-long-term-memory/.claude/skills/agent-memory/examples/memory_tools.py new file mode 100644 index 00000000..65356287 --- /dev/null +++ b/agent-langgraph-long-term-memory/.claude/skills/agent-memory/examples/memory_tools.py @@ -0,0 +1,228 @@ +"""Memory tools for LangGraph agents. + +This module provides tools for managing user long-term memory using +Databricks Lakebase. Copy this file to your agent_server/ directory. + +Usage: + from agent_server.memory_tools import memory_tools, get_user_id + + # In your streaming function: + user_id = get_user_id(request) + tools = await mcp_client.get_tools() + memory_tools() + config = {"configurable": {"user_id": user_id, "store": store}} +""" + +import json +import logging +import os +from typing import Optional + +from databricks.sdk import WorkspaceClient +from langchain_core.runnables import RunnableConfig +from langchain_core.tools import tool +from langgraph.store.base import BaseStore +from mlflow.types.responses import ResponsesAgentRequest + + +# ----------------------------------------------------------------------------- +# Helper Functions +# ----------------------------------------------------------------------------- + + +def get_user_id(request: ResponsesAgentRequest) -> Optional[str]: + """Extract user_id from request context or custom inputs. + + Checks custom_inputs first (for API calls), then request.context + (for Databricks Apps with OBO authentication). + + Returns None if no user_id found - let the caller decide the fallback. + """ + custom_inputs = dict(request.custom_inputs or {}) + if "user_id" in custom_inputs: + return custom_inputs["user_id"] + if request.context and getattr(request.context, "user_id", None): + return request.context.user_id + return None + + +def _is_lakebase_hostname(value: str) -> bool: + """Check if the value looks like a Lakebase hostname rather than an instance name.""" + return ".database." in value and value.endswith(".com") + + +def resolve_lakebase_instance_name( + instance_name: str, workspace_client: Optional[WorkspaceClient] = None +) -> str: + """Resolve a Lakebase instance name from a hostname if needed. + + If the input is a hostname (e.g., from Databricks Apps valueFrom resolution), + this will resolve it to the actual instance name by listing database instances. + + Args: + instance_name: Either an instance name or a hostname + workspace_client: Optional WorkspaceClient to use for resolution + + Returns: + The resolved instance name + + Raises: + ValueError: If the hostname cannot be resolved to an instance name + """ + if not _is_lakebase_hostname(instance_name): + return instance_name + + client = workspace_client or WorkspaceClient() + hostname = instance_name + + try: + instances = list(client.database.list_database_instances()) + except Exception as exc: + raise ValueError( + f"Unable to list database instances to resolve hostname '{hostname}'. " + "Ensure you have access to database instances." + ) from exc + + for instance in instances: + rw_dns = getattr(instance, "read_write_dns", None) + ro_dns = getattr(instance, "read_only_dns", None) + + if hostname in (rw_dns, ro_dns): + resolved_name = getattr(instance, "name", None) + if not resolved_name: + raise ValueError( + f"Found matching instance for hostname '{hostname}' " + "but instance name is not available." + ) + logging.info(f"Resolved Lakebase hostname '{hostname}' to instance name '{resolved_name}'") + return resolved_name + + raise ValueError( + f"Unable to find database instance matching hostname '{hostname}'. " + "Ensure the hostname is correct and the instance exists." + ) + + +def _is_databricks_app_env() -> bool: + """Check if running in a Databricks App environment.""" + return bool(os.getenv("DATABRICKS_APP_NAME")) + + +def get_lakebase_access_error_message(lakebase_instance_name: str) -> str: + """Generate a helpful error message for Lakebase access issues.""" + if _is_databricks_app_env(): + app_name = os.getenv("DATABRICKS_APP_NAME") + return ( + f"Failed to connect to Lakebase instance '{lakebase_instance_name}'. " + f"The App Service Principal for '{app_name}' may not have access.\n\n" + "To fix this:\n" + "1. Go to the Databricks UI and navigate to your app\n" + "2. Click 'Edit' → 'App resources' → 'Add resource'\n" + "3. Add your Lakebase instance as a resource\n" + "4. Grant the necessary permissions on your Lakebase instance." + ) + else: + return ( + f"Failed to connect to Lakebase instance '{lakebase_instance_name}'. " + "Please verify:\n" + "1. The instance name is correct\n" + "2. You have the necessary permissions to access the instance\n" + "3. Your Databricks authentication is configured correctly" + ) + + +# ----------------------------------------------------------------------------- +# Memory Tools Factory +# ----------------------------------------------------------------------------- + + +def memory_tools(): + """Factory function returning memory tools for the agent. + + Returns a list of tools that can be added to your agent: + - get_user_memory: Search for relevant information from long-term memory + - save_user_memory: Save information to long-term memory + - delete_user_memory: Delete a specific memory + + Usage: + tools = await mcp_client.get_tools() + memory_tools() + config = {"configurable": {"user_id": user_id, "store": store}} + """ + + @tool + async def get_user_memory(query: str, config: RunnableConfig) -> str: + """Search for relevant information about the user from long-term memory. + + Use this to recall preferences, past interactions, or other saved information. + + Args: + query: What to search for in the user's memories + """ + user_id = config.get("configurable", {}).get("user_id") + if not user_id: + return "Memory not available - no user_id provided." + + store: Optional[BaseStore] = config.get("configurable", {}).get("store") + if not store: + return "Memory not available - store not configured." + + namespace = ("user_memories", user_id.replace(".", "-")) + results = await store.asearch(namespace, query=query, limit=5) + + if not results: + return "No memories found for this user." + + memory_items = [f"- [{item.key}]: {json.dumps(item.value)}" for item in results] + return f"Found {len(results)} relevant memories:\n" + "\n".join(memory_items) + + @tool + async def save_user_memory(memory_key: str, memory_data_json: str, config: RunnableConfig) -> str: + """Save information about the user to long-term memory. + + Use this to remember user preferences, important details, or other + information that should persist across conversations. + + Args: + memory_key: A short descriptive key (e.g., "preferred_name", "team", "interests") + memory_data_json: JSON object to save (e.g., '{"value": "engineering"}') + """ + user_id = config.get("configurable", {}).get("user_id") + if not user_id: + return "Cannot save memory - no user_id provided." + + store: Optional[BaseStore] = config.get("configurable", {}).get("store") + if not store: + return "Cannot save memory - store not configured." + + namespace = ("user_memories", user_id.replace(".", "-")) + + try: + memory_data = json.loads(memory_data_json) + if not isinstance(memory_data, dict): + return f"Failed: memory_data must be a JSON object, not {type(memory_data).__name__}" + await store.aput(namespace, memory_key, memory_data) + return f"Successfully saved memory '{memory_key}' for user." + except json.JSONDecodeError as e: + return f"Failed to save memory: Invalid JSON - {e}" + + @tool + async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: + """Delete a specific memory from the user's long-term memory. + + Use this when the user asks to forget something or correct stored information. + + Args: + memory_key: The key of the memory to delete (e.g., "preferred_name", "team") + """ + user_id = config.get("configurable", {}).get("user_id") + if not user_id: + return "Cannot delete memory - no user_id provided." + + store: Optional[BaseStore] = config.get("configurable", {}).get("store") + if not store: + return "Cannot delete memory - store not configured." + + namespace = ("user_memories", user_id.replace(".", "-")) + await store.adelete(namespace, memory_key) + return f"Successfully deleted memory '{memory_key}' for user." + + return [get_user_memory, save_user_memory, delete_user_memory] diff --git a/agent-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md b/agent-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md index c87e21a3..49b0cb6a 100644 --- a/agent-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md +++ b/agent-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md @@ -122,16 +122,27 @@ config = {"configurable": {"user_id": user_id, "store": store}} --- +## Complete Example + +A full implementation is available in this skill's examples folder: + +```bash +# Copy to your project +cp .claude/skills/agent-memory/examples/memory_tools.py agent_server/ +``` + +See `examples/memory_tools.py` for production-ready code including all helper functions. + ## Production Reference -For complete, tested implementations: +For implementations in the pre-built templates: | File | Description | |------|-------------| | [`agent-langgraph-long-term-memory/agent_server/utils_memory.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/utils_memory.py) | Memory tools factory, helpers, error handling | | [`agent-langgraph-long-term-memory/agent_server/agent.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/agent.py) | Integration with agent, store initialization | -Key functions in `utils_memory.py`: +Key functions: - `memory_tools()` - Factory returning get/save/delete tools - `get_user_id()` - Extract user_id from request - `resolve_lakebase_instance_name()` - Handle hostname vs instance name diff --git a/agent-langgraph-short-term-memory/.claude/skills/agent-memory/examples/memory_tools.py b/agent-langgraph-short-term-memory/.claude/skills/agent-memory/examples/memory_tools.py new file mode 100644 index 00000000..65356287 --- /dev/null +++ b/agent-langgraph-short-term-memory/.claude/skills/agent-memory/examples/memory_tools.py @@ -0,0 +1,228 @@ +"""Memory tools for LangGraph agents. + +This module provides tools for managing user long-term memory using +Databricks Lakebase. Copy this file to your agent_server/ directory. + +Usage: + from agent_server.memory_tools import memory_tools, get_user_id + + # In your streaming function: + user_id = get_user_id(request) + tools = await mcp_client.get_tools() + memory_tools() + config = {"configurable": {"user_id": user_id, "store": store}} +""" + +import json +import logging +import os +from typing import Optional + +from databricks.sdk import WorkspaceClient +from langchain_core.runnables import RunnableConfig +from langchain_core.tools import tool +from langgraph.store.base import BaseStore +from mlflow.types.responses import ResponsesAgentRequest + + +# ----------------------------------------------------------------------------- +# Helper Functions +# ----------------------------------------------------------------------------- + + +def get_user_id(request: ResponsesAgentRequest) -> Optional[str]: + """Extract user_id from request context or custom inputs. + + Checks custom_inputs first (for API calls), then request.context + (for Databricks Apps with OBO authentication). + + Returns None if no user_id found - let the caller decide the fallback. + """ + custom_inputs = dict(request.custom_inputs or {}) + if "user_id" in custom_inputs: + return custom_inputs["user_id"] + if request.context and getattr(request.context, "user_id", None): + return request.context.user_id + return None + + +def _is_lakebase_hostname(value: str) -> bool: + """Check if the value looks like a Lakebase hostname rather than an instance name.""" + return ".database." in value and value.endswith(".com") + + +def resolve_lakebase_instance_name( + instance_name: str, workspace_client: Optional[WorkspaceClient] = None +) -> str: + """Resolve a Lakebase instance name from a hostname if needed. + + If the input is a hostname (e.g., from Databricks Apps valueFrom resolution), + this will resolve it to the actual instance name by listing database instances. + + Args: + instance_name: Either an instance name or a hostname + workspace_client: Optional WorkspaceClient to use for resolution + + Returns: + The resolved instance name + + Raises: + ValueError: If the hostname cannot be resolved to an instance name + """ + if not _is_lakebase_hostname(instance_name): + return instance_name + + client = workspace_client or WorkspaceClient() + hostname = instance_name + + try: + instances = list(client.database.list_database_instances()) + except Exception as exc: + raise ValueError( + f"Unable to list database instances to resolve hostname '{hostname}'. " + "Ensure you have access to database instances." + ) from exc + + for instance in instances: + rw_dns = getattr(instance, "read_write_dns", None) + ro_dns = getattr(instance, "read_only_dns", None) + + if hostname in (rw_dns, ro_dns): + resolved_name = getattr(instance, "name", None) + if not resolved_name: + raise ValueError( + f"Found matching instance for hostname '{hostname}' " + "but instance name is not available." + ) + logging.info(f"Resolved Lakebase hostname '{hostname}' to instance name '{resolved_name}'") + return resolved_name + + raise ValueError( + f"Unable to find database instance matching hostname '{hostname}'. " + "Ensure the hostname is correct and the instance exists." + ) + + +def _is_databricks_app_env() -> bool: + """Check if running in a Databricks App environment.""" + return bool(os.getenv("DATABRICKS_APP_NAME")) + + +def get_lakebase_access_error_message(lakebase_instance_name: str) -> str: + """Generate a helpful error message for Lakebase access issues.""" + if _is_databricks_app_env(): + app_name = os.getenv("DATABRICKS_APP_NAME") + return ( + f"Failed to connect to Lakebase instance '{lakebase_instance_name}'. " + f"The App Service Principal for '{app_name}' may not have access.\n\n" + "To fix this:\n" + "1. Go to the Databricks UI and navigate to your app\n" + "2. Click 'Edit' → 'App resources' → 'Add resource'\n" + "3. Add your Lakebase instance as a resource\n" + "4. Grant the necessary permissions on your Lakebase instance." + ) + else: + return ( + f"Failed to connect to Lakebase instance '{lakebase_instance_name}'. " + "Please verify:\n" + "1. The instance name is correct\n" + "2. You have the necessary permissions to access the instance\n" + "3. Your Databricks authentication is configured correctly" + ) + + +# ----------------------------------------------------------------------------- +# Memory Tools Factory +# ----------------------------------------------------------------------------- + + +def memory_tools(): + """Factory function returning memory tools for the agent. + + Returns a list of tools that can be added to your agent: + - get_user_memory: Search for relevant information from long-term memory + - save_user_memory: Save information to long-term memory + - delete_user_memory: Delete a specific memory + + Usage: + tools = await mcp_client.get_tools() + memory_tools() + config = {"configurable": {"user_id": user_id, "store": store}} + """ + + @tool + async def get_user_memory(query: str, config: RunnableConfig) -> str: + """Search for relevant information about the user from long-term memory. + + Use this to recall preferences, past interactions, or other saved information. + + Args: + query: What to search for in the user's memories + """ + user_id = config.get("configurable", {}).get("user_id") + if not user_id: + return "Memory not available - no user_id provided." + + store: Optional[BaseStore] = config.get("configurable", {}).get("store") + if not store: + return "Memory not available - store not configured." + + namespace = ("user_memories", user_id.replace(".", "-")) + results = await store.asearch(namespace, query=query, limit=5) + + if not results: + return "No memories found for this user." + + memory_items = [f"- [{item.key}]: {json.dumps(item.value)}" for item in results] + return f"Found {len(results)} relevant memories:\n" + "\n".join(memory_items) + + @tool + async def save_user_memory(memory_key: str, memory_data_json: str, config: RunnableConfig) -> str: + """Save information about the user to long-term memory. + + Use this to remember user preferences, important details, or other + information that should persist across conversations. + + Args: + memory_key: A short descriptive key (e.g., "preferred_name", "team", "interests") + memory_data_json: JSON object to save (e.g., '{"value": "engineering"}') + """ + user_id = config.get("configurable", {}).get("user_id") + if not user_id: + return "Cannot save memory - no user_id provided." + + store: Optional[BaseStore] = config.get("configurable", {}).get("store") + if not store: + return "Cannot save memory - store not configured." + + namespace = ("user_memories", user_id.replace(".", "-")) + + try: + memory_data = json.loads(memory_data_json) + if not isinstance(memory_data, dict): + return f"Failed: memory_data must be a JSON object, not {type(memory_data).__name__}" + await store.aput(namespace, memory_key, memory_data) + return f"Successfully saved memory '{memory_key}' for user." + except json.JSONDecodeError as e: + return f"Failed to save memory: Invalid JSON - {e}" + + @tool + async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: + """Delete a specific memory from the user's long-term memory. + + Use this when the user asks to forget something or correct stored information. + + Args: + memory_key: The key of the memory to delete (e.g., "preferred_name", "team") + """ + user_id = config.get("configurable", {}).get("user_id") + if not user_id: + return "Cannot delete memory - no user_id provided." + + store: Optional[BaseStore] = config.get("configurable", {}).get("store") + if not store: + return "Cannot delete memory - store not configured." + + namespace = ("user_memories", user_id.replace(".", "-")) + await store.adelete(namespace, memory_key) + return f"Successfully deleted memory '{memory_key}' for user." + + return [get_user_memory, save_user_memory, delete_user_memory] diff --git a/agent-langgraph/.claude/skills/agent-memory/SKILL.md b/agent-langgraph/.claude/skills/agent-memory/SKILL.md index c87e21a3..49b0cb6a 100644 --- a/agent-langgraph/.claude/skills/agent-memory/SKILL.md +++ b/agent-langgraph/.claude/skills/agent-memory/SKILL.md @@ -122,16 +122,27 @@ config = {"configurable": {"user_id": user_id, "store": store}} --- +## Complete Example + +A full implementation is available in this skill's examples folder: + +```bash +# Copy to your project +cp .claude/skills/agent-memory/examples/memory_tools.py agent_server/ +``` + +See `examples/memory_tools.py` for production-ready code including all helper functions. + ## Production Reference -For complete, tested implementations: +For implementations in the pre-built templates: | File | Description | |------|-------------| | [`agent-langgraph-long-term-memory/agent_server/utils_memory.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/utils_memory.py) | Memory tools factory, helpers, error handling | | [`agent-langgraph-long-term-memory/agent_server/agent.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/agent.py) | Integration with agent, store initialization | -Key functions in `utils_memory.py`: +Key functions: - `memory_tools()` - Factory returning get/save/delete tools - `get_user_id()` - Extract user_id from request - `resolve_lakebase_instance_name()` - Handle hostname vs instance name diff --git a/agent-langgraph/.claude/skills/agent-memory/examples/memory_tools.py b/agent-langgraph/.claude/skills/agent-memory/examples/memory_tools.py new file mode 100644 index 00000000..65356287 --- /dev/null +++ b/agent-langgraph/.claude/skills/agent-memory/examples/memory_tools.py @@ -0,0 +1,228 @@ +"""Memory tools for LangGraph agents. + +This module provides tools for managing user long-term memory using +Databricks Lakebase. Copy this file to your agent_server/ directory. + +Usage: + from agent_server.memory_tools import memory_tools, get_user_id + + # In your streaming function: + user_id = get_user_id(request) + tools = await mcp_client.get_tools() + memory_tools() + config = {"configurable": {"user_id": user_id, "store": store}} +""" + +import json +import logging +import os +from typing import Optional + +from databricks.sdk import WorkspaceClient +from langchain_core.runnables import RunnableConfig +from langchain_core.tools import tool +from langgraph.store.base import BaseStore +from mlflow.types.responses import ResponsesAgentRequest + + +# ----------------------------------------------------------------------------- +# Helper Functions +# ----------------------------------------------------------------------------- + + +def get_user_id(request: ResponsesAgentRequest) -> Optional[str]: + """Extract user_id from request context or custom inputs. + + Checks custom_inputs first (for API calls), then request.context + (for Databricks Apps with OBO authentication). + + Returns None if no user_id found - let the caller decide the fallback. + """ + custom_inputs = dict(request.custom_inputs or {}) + if "user_id" in custom_inputs: + return custom_inputs["user_id"] + if request.context and getattr(request.context, "user_id", None): + return request.context.user_id + return None + + +def _is_lakebase_hostname(value: str) -> bool: + """Check if the value looks like a Lakebase hostname rather than an instance name.""" + return ".database." in value and value.endswith(".com") + + +def resolve_lakebase_instance_name( + instance_name: str, workspace_client: Optional[WorkspaceClient] = None +) -> str: + """Resolve a Lakebase instance name from a hostname if needed. + + If the input is a hostname (e.g., from Databricks Apps valueFrom resolution), + this will resolve it to the actual instance name by listing database instances. + + Args: + instance_name: Either an instance name or a hostname + workspace_client: Optional WorkspaceClient to use for resolution + + Returns: + The resolved instance name + + Raises: + ValueError: If the hostname cannot be resolved to an instance name + """ + if not _is_lakebase_hostname(instance_name): + return instance_name + + client = workspace_client or WorkspaceClient() + hostname = instance_name + + try: + instances = list(client.database.list_database_instances()) + except Exception as exc: + raise ValueError( + f"Unable to list database instances to resolve hostname '{hostname}'. " + "Ensure you have access to database instances." + ) from exc + + for instance in instances: + rw_dns = getattr(instance, "read_write_dns", None) + ro_dns = getattr(instance, "read_only_dns", None) + + if hostname in (rw_dns, ro_dns): + resolved_name = getattr(instance, "name", None) + if not resolved_name: + raise ValueError( + f"Found matching instance for hostname '{hostname}' " + "but instance name is not available." + ) + logging.info(f"Resolved Lakebase hostname '{hostname}' to instance name '{resolved_name}'") + return resolved_name + + raise ValueError( + f"Unable to find database instance matching hostname '{hostname}'. " + "Ensure the hostname is correct and the instance exists." + ) + + +def _is_databricks_app_env() -> bool: + """Check if running in a Databricks App environment.""" + return bool(os.getenv("DATABRICKS_APP_NAME")) + + +def get_lakebase_access_error_message(lakebase_instance_name: str) -> str: + """Generate a helpful error message for Lakebase access issues.""" + if _is_databricks_app_env(): + app_name = os.getenv("DATABRICKS_APP_NAME") + return ( + f"Failed to connect to Lakebase instance '{lakebase_instance_name}'. " + f"The App Service Principal for '{app_name}' may not have access.\n\n" + "To fix this:\n" + "1. Go to the Databricks UI and navigate to your app\n" + "2. Click 'Edit' → 'App resources' → 'Add resource'\n" + "3. Add your Lakebase instance as a resource\n" + "4. Grant the necessary permissions on your Lakebase instance." + ) + else: + return ( + f"Failed to connect to Lakebase instance '{lakebase_instance_name}'. " + "Please verify:\n" + "1. The instance name is correct\n" + "2. You have the necessary permissions to access the instance\n" + "3. Your Databricks authentication is configured correctly" + ) + + +# ----------------------------------------------------------------------------- +# Memory Tools Factory +# ----------------------------------------------------------------------------- + + +def memory_tools(): + """Factory function returning memory tools for the agent. + + Returns a list of tools that can be added to your agent: + - get_user_memory: Search for relevant information from long-term memory + - save_user_memory: Save information to long-term memory + - delete_user_memory: Delete a specific memory + + Usage: + tools = await mcp_client.get_tools() + memory_tools() + config = {"configurable": {"user_id": user_id, "store": store}} + """ + + @tool + async def get_user_memory(query: str, config: RunnableConfig) -> str: + """Search for relevant information about the user from long-term memory. + + Use this to recall preferences, past interactions, or other saved information. + + Args: + query: What to search for in the user's memories + """ + user_id = config.get("configurable", {}).get("user_id") + if not user_id: + return "Memory not available - no user_id provided." + + store: Optional[BaseStore] = config.get("configurable", {}).get("store") + if not store: + return "Memory not available - store not configured." + + namespace = ("user_memories", user_id.replace(".", "-")) + results = await store.asearch(namespace, query=query, limit=5) + + if not results: + return "No memories found for this user." + + memory_items = [f"- [{item.key}]: {json.dumps(item.value)}" for item in results] + return f"Found {len(results)} relevant memories:\n" + "\n".join(memory_items) + + @tool + async def save_user_memory(memory_key: str, memory_data_json: str, config: RunnableConfig) -> str: + """Save information about the user to long-term memory. + + Use this to remember user preferences, important details, or other + information that should persist across conversations. + + Args: + memory_key: A short descriptive key (e.g., "preferred_name", "team", "interests") + memory_data_json: JSON object to save (e.g., '{"value": "engineering"}') + """ + user_id = config.get("configurable", {}).get("user_id") + if not user_id: + return "Cannot save memory - no user_id provided." + + store: Optional[BaseStore] = config.get("configurable", {}).get("store") + if not store: + return "Cannot save memory - store not configured." + + namespace = ("user_memories", user_id.replace(".", "-")) + + try: + memory_data = json.loads(memory_data_json) + if not isinstance(memory_data, dict): + return f"Failed: memory_data must be a JSON object, not {type(memory_data).__name__}" + await store.aput(namespace, memory_key, memory_data) + return f"Successfully saved memory '{memory_key}' for user." + except json.JSONDecodeError as e: + return f"Failed to save memory: Invalid JSON - {e}" + + @tool + async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: + """Delete a specific memory from the user's long-term memory. + + Use this when the user asks to forget something or correct stored information. + + Args: + memory_key: The key of the memory to delete (e.g., "preferred_name", "team") + """ + user_id = config.get("configurable", {}).get("user_id") + if not user_id: + return "Cannot delete memory - no user_id provided." + + store: Optional[BaseStore] = config.get("configurable", {}).get("store") + if not store: + return "Cannot delete memory - store not configured." + + namespace = ("user_memories", user_id.replace(".", "-")) + await store.adelete(namespace, memory_key) + return f"Successfully deleted memory '{memory_key}' for user." + + return [get_user_memory, save_user_memory, delete_user_memory] diff --git a/agent-non-conversational/.claude/skills/agent-memory/SKILL.md b/agent-non-conversational/.claude/skills/agent-memory/SKILL.md index c87e21a3..49b0cb6a 100644 --- a/agent-non-conversational/.claude/skills/agent-memory/SKILL.md +++ b/agent-non-conversational/.claude/skills/agent-memory/SKILL.md @@ -122,16 +122,27 @@ config = {"configurable": {"user_id": user_id, "store": store}} --- +## Complete Example + +A full implementation is available in this skill's examples folder: + +```bash +# Copy to your project +cp .claude/skills/agent-memory/examples/memory_tools.py agent_server/ +``` + +See `examples/memory_tools.py` for production-ready code including all helper functions. + ## Production Reference -For complete, tested implementations: +For implementations in the pre-built templates: | File | Description | |------|-------------| | [`agent-langgraph-long-term-memory/agent_server/utils_memory.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/utils_memory.py) | Memory tools factory, helpers, error handling | | [`agent-langgraph-long-term-memory/agent_server/agent.py`](https://github.com/databricks/app-templates/tree/main/agent-langgraph-long-term-memory/agent_server/agent.py) | Integration with agent, store initialization | -Key functions in `utils_memory.py`: +Key functions: - `memory_tools()` - Factory returning get/save/delete tools - `get_user_id()` - Extract user_id from request - `resolve_lakebase_instance_name()` - Handle hostname vs instance name diff --git a/agent-non-conversational/.claude/skills/agent-memory/examples/memory_tools.py b/agent-non-conversational/.claude/skills/agent-memory/examples/memory_tools.py new file mode 100644 index 00000000..65356287 --- /dev/null +++ b/agent-non-conversational/.claude/skills/agent-memory/examples/memory_tools.py @@ -0,0 +1,228 @@ +"""Memory tools for LangGraph agents. + +This module provides tools for managing user long-term memory using +Databricks Lakebase. Copy this file to your agent_server/ directory. + +Usage: + from agent_server.memory_tools import memory_tools, get_user_id + + # In your streaming function: + user_id = get_user_id(request) + tools = await mcp_client.get_tools() + memory_tools() + config = {"configurable": {"user_id": user_id, "store": store}} +""" + +import json +import logging +import os +from typing import Optional + +from databricks.sdk import WorkspaceClient +from langchain_core.runnables import RunnableConfig +from langchain_core.tools import tool +from langgraph.store.base import BaseStore +from mlflow.types.responses import ResponsesAgentRequest + + +# ----------------------------------------------------------------------------- +# Helper Functions +# ----------------------------------------------------------------------------- + + +def get_user_id(request: ResponsesAgentRequest) -> Optional[str]: + """Extract user_id from request context or custom inputs. + + Checks custom_inputs first (for API calls), then request.context + (for Databricks Apps with OBO authentication). + + Returns None if no user_id found - let the caller decide the fallback. + """ + custom_inputs = dict(request.custom_inputs or {}) + if "user_id" in custom_inputs: + return custom_inputs["user_id"] + if request.context and getattr(request.context, "user_id", None): + return request.context.user_id + return None + + +def _is_lakebase_hostname(value: str) -> bool: + """Check if the value looks like a Lakebase hostname rather than an instance name.""" + return ".database." in value and value.endswith(".com") + + +def resolve_lakebase_instance_name( + instance_name: str, workspace_client: Optional[WorkspaceClient] = None +) -> str: + """Resolve a Lakebase instance name from a hostname if needed. + + If the input is a hostname (e.g., from Databricks Apps valueFrom resolution), + this will resolve it to the actual instance name by listing database instances. + + Args: + instance_name: Either an instance name or a hostname + workspace_client: Optional WorkspaceClient to use for resolution + + Returns: + The resolved instance name + + Raises: + ValueError: If the hostname cannot be resolved to an instance name + """ + if not _is_lakebase_hostname(instance_name): + return instance_name + + client = workspace_client or WorkspaceClient() + hostname = instance_name + + try: + instances = list(client.database.list_database_instances()) + except Exception as exc: + raise ValueError( + f"Unable to list database instances to resolve hostname '{hostname}'. " + "Ensure you have access to database instances." + ) from exc + + for instance in instances: + rw_dns = getattr(instance, "read_write_dns", None) + ro_dns = getattr(instance, "read_only_dns", None) + + if hostname in (rw_dns, ro_dns): + resolved_name = getattr(instance, "name", None) + if not resolved_name: + raise ValueError( + f"Found matching instance for hostname '{hostname}' " + "but instance name is not available." + ) + logging.info(f"Resolved Lakebase hostname '{hostname}' to instance name '{resolved_name}'") + return resolved_name + + raise ValueError( + f"Unable to find database instance matching hostname '{hostname}'. " + "Ensure the hostname is correct and the instance exists." + ) + + +def _is_databricks_app_env() -> bool: + """Check if running in a Databricks App environment.""" + return bool(os.getenv("DATABRICKS_APP_NAME")) + + +def get_lakebase_access_error_message(lakebase_instance_name: str) -> str: + """Generate a helpful error message for Lakebase access issues.""" + if _is_databricks_app_env(): + app_name = os.getenv("DATABRICKS_APP_NAME") + return ( + f"Failed to connect to Lakebase instance '{lakebase_instance_name}'. " + f"The App Service Principal for '{app_name}' may not have access.\n\n" + "To fix this:\n" + "1. Go to the Databricks UI and navigate to your app\n" + "2. Click 'Edit' → 'App resources' → 'Add resource'\n" + "3. Add your Lakebase instance as a resource\n" + "4. Grant the necessary permissions on your Lakebase instance." + ) + else: + return ( + f"Failed to connect to Lakebase instance '{lakebase_instance_name}'. " + "Please verify:\n" + "1. The instance name is correct\n" + "2. You have the necessary permissions to access the instance\n" + "3. Your Databricks authentication is configured correctly" + ) + + +# ----------------------------------------------------------------------------- +# Memory Tools Factory +# ----------------------------------------------------------------------------- + + +def memory_tools(): + """Factory function returning memory tools for the agent. + + Returns a list of tools that can be added to your agent: + - get_user_memory: Search for relevant information from long-term memory + - save_user_memory: Save information to long-term memory + - delete_user_memory: Delete a specific memory + + Usage: + tools = await mcp_client.get_tools() + memory_tools() + config = {"configurable": {"user_id": user_id, "store": store}} + """ + + @tool + async def get_user_memory(query: str, config: RunnableConfig) -> str: + """Search for relevant information about the user from long-term memory. + + Use this to recall preferences, past interactions, or other saved information. + + Args: + query: What to search for in the user's memories + """ + user_id = config.get("configurable", {}).get("user_id") + if not user_id: + return "Memory not available - no user_id provided." + + store: Optional[BaseStore] = config.get("configurable", {}).get("store") + if not store: + return "Memory not available - store not configured." + + namespace = ("user_memories", user_id.replace(".", "-")) + results = await store.asearch(namespace, query=query, limit=5) + + if not results: + return "No memories found for this user." + + memory_items = [f"- [{item.key}]: {json.dumps(item.value)}" for item in results] + return f"Found {len(results)} relevant memories:\n" + "\n".join(memory_items) + + @tool + async def save_user_memory(memory_key: str, memory_data_json: str, config: RunnableConfig) -> str: + """Save information about the user to long-term memory. + + Use this to remember user preferences, important details, or other + information that should persist across conversations. + + Args: + memory_key: A short descriptive key (e.g., "preferred_name", "team", "interests") + memory_data_json: JSON object to save (e.g., '{"value": "engineering"}') + """ + user_id = config.get("configurable", {}).get("user_id") + if not user_id: + return "Cannot save memory - no user_id provided." + + store: Optional[BaseStore] = config.get("configurable", {}).get("store") + if not store: + return "Cannot save memory - store not configured." + + namespace = ("user_memories", user_id.replace(".", "-")) + + try: + memory_data = json.loads(memory_data_json) + if not isinstance(memory_data, dict): + return f"Failed: memory_data must be a JSON object, not {type(memory_data).__name__}" + await store.aput(namespace, memory_key, memory_data) + return f"Successfully saved memory '{memory_key}' for user." + except json.JSONDecodeError as e: + return f"Failed to save memory: Invalid JSON - {e}" + + @tool + async def delete_user_memory(memory_key: str, config: RunnableConfig) -> str: + """Delete a specific memory from the user's long-term memory. + + Use this when the user asks to forget something or correct stored information. + + Args: + memory_key: The key of the memory to delete (e.g., "preferred_name", "team") + """ + user_id = config.get("configurable", {}).get("user_id") + if not user_id: + return "Cannot delete memory - no user_id provided." + + store: Optional[BaseStore] = config.get("configurable", {}).get("store") + if not store: + return "Cannot delete memory - store not configured." + + namespace = ("user_memories", user_id.replace(".", "-")) + await store.adelete(namespace, memory_key) + return f"Successfully deleted memory '{memory_key}' for user." + + return [get_user_memory, save_user_memory, delete_user_memory] From 73b0d118f86e58a8d6c7d982a836318fee5e7c56 Mon Sep 17 00:00:00 2001 From: Dhruv Gupta Date: Mon, 2 Feb 2026 11:16:33 -0800 Subject: [PATCH 14/14] Fix .gitignore to track renamed agent-langgraph-memory skill Co-Authored-By: Claude Opus 4.5 --- .gitignore | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 5e89ca89..13108d9d 100644 --- a/.gitignore +++ b/.gitignore @@ -183,4 +183,4 @@ tsconfig.tsbuildinfo !.claude/skills/modify-langgraph-agent/ !.claude/skills/modify-openai-agent/ !.claude/skills/lakebase-setup/ -!.claude/skills/agent-memory/ +!.claude/skills/agent-langgraph-memory/