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-langgraph-memory/SKILL.md b/.claude/skills/agent-langgraph-memory/SKILL.md new file mode 100644 index 00000000..49b0cb6a --- /dev/null +++ b/.claude/skills/agent-langgraph-memory/SKILL.md @@ -0,0 +1,386 @@ +--- +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") +``` + +--- + +## 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 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: +- `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-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/.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..65c79099 --- /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 `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) + +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..5b782185 --- /dev/null +++ b/.claude/sync-skills.py @@ -0,0 +1,104 @@ +#!/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) + # 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-langgraph-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/.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 diff --git a/.gitignore b/.gitignore index 38e33701..13108d9d 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-langgraph-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 new file mode 100644 index 00000000..7719f198 --- /dev/null +++ b/agent-langgraph-long-term-memory/.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-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..1324e6c5 --- /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 +``` + +### 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/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/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..49b0cb6a --- /dev/null +++ b/agent-langgraph-long-term-memory/.claude/skills/agent-memory/SKILL.md @@ -0,0 +1,386 @@ +--- +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") +``` + +--- + +## 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 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: +- `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/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-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..3a0f3090 --- /dev/null +++ b/agent-langgraph-long-term-memory/.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_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 +# 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_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 + +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. + +## 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 }' +``` + +**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_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 new file mode 100644 index 00000000..87c3f519 --- /dev/null +++ b/agent-langgraph-long-term-memory/.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-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..65c79099 --- /dev/null +++ b/agent-langgraph-long-term-memory/.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 `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) + +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-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..d7218637 --- /dev/null +++ b/agent-langgraph-long-term-memory/.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-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..e550162c --- /dev/null +++ b/agent-langgraph-long-term-memory/.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-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..3eb83c82 --- /dev/null +++ b/agent-langgraph-long-term-memory/.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-langgraph-long-term-memory/.gitignore b/agent-langgraph-long-term-memory/.gitignore index c0ee193b..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__/ @@ -204,5 +202,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..d8f25998 100644 --- a/agent-langgraph-long-term-memory/AGENTS.md +++ b/agent-langgraph-long-term-memory/AGENTS.md @@ -1,348 +1,119 @@ -# Agent LangGraph Development Guide +# Agent Development Guide -## Running the App +## MANDATORY First Actions -**Prerequisites:** uv, nvm (Node 20), Databricks CLI +**Ask the user interactively:** -**Quick Start:** +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?" -```bash -uv run quickstart # First-time setup (auth, MLflow experiment, env) -uv run start-app # Start app at http://localhost:8000 -``` + *Note: New apps should use the `agent-*` prefix (e.g., `agent-data-analyst`) unless the user specifies otherwise.* -**Advanced Server Options:** +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?" -```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 -``` +**Then check authentication status by running `databricks auth profiles`.** -**Test API:** +This helps you understand: +- Which Databricks profiles are configured +- Whether authentication is already set up +- Which profile to use for subsequent commands -```bash -# Streaming request -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' +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. -# Non-streaming request -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ "input": [{ "role": "user", "content": "hi" }] }' -``` +## Understanding User Goals ---- - -## Testing the Agent - -**Run evaluation:** +**Ask the user questions to understand what they're building:** -```bash -uv run agent-evaluate # Uses MLflow scorers (RelevanceToQuery, Safety) -``` +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?** -**Run unit tests:** +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.** -```bash -pytest [path] # Standard pytest execution -``` +## Memory Template Note ---- +This template includes **long-term memory** (facts that persist across conversation sessions). The agent can remember user preferences and information across multiple interactions. -## Modifying the Agent +**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) -Anytime the user wants to modify the agent, look through each of the following resources to help them accomplish their goal: +## Handling Deployment Errors -If the user wants to convert something into Responses API, refer to https://mlflow.org/docs/latest/genai/serving/responses-agent/ for more information. +**If `databricks bundle deploy` fails with "An app with the same name already exists":** -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. +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?" -**Main file to modify:** `agent_server/agent.py` +- **If they want to bind**: See the **deploy** skill for binding steps +- **If they want to delete**: Run `databricks apps delete ` then deploy again --- -## databricks-langchain SDK overview - -**SDK Location:** `https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/langchain` +## Available Skills -**Development Workflow:** +**Before executing any task, read the relevant skill file in `.claude/skills/`** - they contain tested commands, patterns, and troubleshooting steps. -```bash -uv add databricks-langchain -``` +| 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` | -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. +**Note:** All agent skills are located in `.claude/skills/` directory. --- -### ChatDatabricks - LLM Chat Interface - -Connects to Databricks Model Serving endpoints for LLM inference. +## Quick Commands -```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) -``` +| 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` | --- -### 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, 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..7719f198 --- /dev/null +++ b/agent-langgraph-short-term-memory/.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-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..1324e6c5 --- /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 +``` + +### 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/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/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..49b0cb6a --- /dev/null +++ b/agent-langgraph-short-term-memory/.claude/skills/agent-memory/SKILL.md @@ -0,0 +1,386 @@ +--- +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") +``` + +--- + +## 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 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: +- `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/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-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..17f08ba0 --- /dev/null +++ b/agent-langgraph-short-term-memory/.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_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 +# 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_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 + +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. + +## 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 }' +``` + +**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_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 new file mode 100644 index 00000000..87c3f519 --- /dev/null +++ b/agent-langgraph-short-term-memory/.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-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..65c79099 --- /dev/null +++ b/agent-langgraph-short-term-memory/.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 `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) + +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-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..d7218637 --- /dev/null +++ b/agent-langgraph-short-term-memory/.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-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..e550162c --- /dev/null +++ b/agent-langgraph-short-term-memory/.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-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..3eb83c82 --- /dev/null +++ b/agent-langgraph-short-term-memory/.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-langgraph-short-term-memory/.gitignore b/agent-langgraph-short-term-memory/.gitignore index a058c65c..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__/ @@ -204,6 +202,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..b21b1ff1 100644 --- a/agent-langgraph-short-term-memory/AGENTS.md +++ b/agent-langgraph-short-term-memory/AGENTS.md @@ -1,348 +1,119 @@ -# Agent LangGraph Development Guide +# Agent Development Guide -## Running the App +## MANDATORY First Actions -**Prerequisites:** uv, nvm (Node 20), Databricks CLI +**Ask the user interactively:** -**Quick Start:** +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?" -```bash -uv run quickstart # First-time setup (auth, MLflow experiment, env) -uv run start-app # Start app at http://localhost:8000 -``` + *Note: New apps should use the `agent-*` prefix (e.g., `agent-data-analyst`) unless the user specifies otherwise.* -**Advanced Server Options:** +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?" -```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 -``` +**Then check authentication status by running `databricks auth profiles`.** -**Test API:** +This helps you understand: +- Which Databricks profiles are configured +- Whether authentication is already set up +- Which profile to use for subsequent commands -```bash -# Streaming request -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' +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. -# Non-streaming request -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ "input": [{ "role": "user", "content": "hi" }] }' -``` +## Understanding User Goals ---- - -## Testing the Agent - -**Run evaluation:** +**Ask the user questions to understand what they're building:** -```bash -uv run agent-evaluate # Uses MLflow scorers (RelevanceToQuery, Safety) -``` +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?** -**Run unit tests:** +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.** -```bash -pytest [path] # Standard pytest execution -``` +## Memory Template Note ---- +This template includes **short-term memory** (conversation history within a session). The agent remembers what was said earlier in the same conversation thread. -## Modifying the Agent +**Required setup:** +1. Configure Lakebase instance (see **lakebase-setup** skill) +2. Use `thread_id` in requests to maintain conversation context (see **agent-memory** skill) -Anytime the user wants to modify the agent, look through each of the following resources to help them accomplish their goal: +## Handling Deployment Errors -If the user wants to convert something into Responses API, refer to https://mlflow.org/docs/latest/genai/serving/responses-agent/ for more information. +**If `databricks bundle deploy` fails with "An app with the same name already exists":** -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. +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?" -**Main file to modify:** `agent_server/agent.py` +- **If they want to bind**: See the **deploy** skill for binding steps +- **If they want to delete**: Run `databricks apps delete ` then deploy again --- -## databricks-langchain SDK overview - -**SDK Location:** `https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/langchain` +## Available Skills -**Development Workflow:** +**Before executing any task, read the relevant skill file in `.claude/skills/`** - they contain tested commands, patterns, and troubleshooting steps. -```bash -uv add databricks-langchain -``` +| 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` | -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. +**Note:** All agent skills are located in `.claude/skills/` directory. --- -### ChatDatabricks - LLM Chat Interface - -Connects to Databricks Model Serving endpoints for LLM inference. +## Quick Commands -```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) -``` +| 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` | --- -### 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, 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() 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..7719f198 --- /dev/null +++ b/agent-langgraph/.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-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/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/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/agent-memory/SKILL.md b/agent-langgraph/.claude/skills/agent-memory/SKILL.md new file mode 100644 index 00000000..49b0cb6a --- /dev/null +++ b/agent-langgraph/.claude/skills/agent-memory/SKILL.md @@ -0,0 +1,386 @@ +--- +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") +``` + +--- + +## 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 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: +- `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/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-langgraph/.claude/skills/deploy/SKILL.md b/agent-langgraph/.claude/skills/deploy/SKILL.md new file mode 100644 index 00000000..dfaa72af --- /dev/null +++ b/agent-langgraph/.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_langgraph: + 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_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 + +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 + +> **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_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 new file mode 100644 index 00000000..87c3f519 --- /dev/null +++ b/agent-langgraph/.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-langgraph/.claude/skills/lakebase-setup/SKILL.md b/agent-langgraph/.claude/skills/lakebase-setup/SKILL.md new file mode 100644 index 00000000..65c79099 --- /dev/null +++ b/agent-langgraph/.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 `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) + +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-langgraph/.claude/skills/modify-agent/SKILL.md b/agent-langgraph/.claude/skills/modify-agent/SKILL.md new file mode 100644 index 00000000..d7218637 --- /dev/null +++ b/agent-langgraph/.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-langgraph/.claude/skills/quickstart/SKILL.md b/agent-langgraph/.claude/skills/quickstart/SKILL.md new file mode 100644 index 00000000..e550162c --- /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 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-langgraph/.claude/skills/run-locally/SKILL.md b/agent-langgraph/.claude/skills/run-locally/SKILL.md new file mode 100644 index 00000000..3eb83c82 --- /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` 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-langgraph/.gitignore b/agent-langgraph/.gitignore index a058c65c..601fbb1a 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,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-langgraph/AGENTS.md b/agent-langgraph/AGENTS.md index fcd86db5..fcac7dcd 100644 --- a/agent-langgraph/AGENTS.md +++ b/agent-langgraph/AGENTS.md @@ -1,348 +1,113 @@ -# Agent LangGraph Development Guide +# Agent Development Guide -## Running the App +## MANDATORY First Actions -**Prerequisites:** uv, nvm (Node 20), Databricks CLI +**Ask the user interactively:** -**Quick Start:** +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?" -```bash -uv run quickstart # First-time setup (auth, MLflow experiment, env) -uv run start-app # Start app at http://localhost:8000 -``` + *Note: New apps should use the `agent-*` prefix (e.g., `agent-data-analyst`) unless the user specifies otherwise.* -**Advanced Server Options:** +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?" -```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 -``` +**Then check authentication status by running `databricks auth profiles`.** -**Test API:** +This helps you understand: +- Which Databricks profiles are configured +- Whether authentication is already set up +- Which profile to use for subsequent commands -```bash -# Streaming request -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' +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. -# Non-streaming request -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ "input": [{ "role": "user", "content": "hi" }] }' -``` +## Understanding User Goals ---- - -## 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. +**Ask the user questions to understand what they're building:** -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. +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?** -**Main file to modify:** `agent_server/agent.py` +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.** ---- - -## databricks-langchain SDK overview - -**SDK Location:** `https://github.com/databricks/databricks-ai-bridge/tree/main/integrations/langchain` +## Handling Deployment Errors -**Development Workflow:** +**If `databricks bundle deploy` fails with "An app with the same name already exists":** -```bash -uv add databricks-langchain -``` +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?" -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 +## Available Skills -Connects to Databricks Model Serving endpoints for LLM inference. +**Before executing any task, read the relevant skill file in `.claude/skills/`** - they contain tested commands, patterns, and troubleshooting steps. -```python -from databricks_langchain import ChatDatabricks +| 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` | +| Configure Lakebase storage | **lakebase-setup** | `.claude/skills/lakebase-setup/SKILL.md` | +| Add memory capabilities | **agent-memory** | `.claude/skills/agent-memory/SKILL.md` | -llm = ChatDatabricks( - endpoint="databricks-claude-3-7-sonnet", # or databricks-meta-llama-3-1-70b-instruct - temperature=0, - max_tokens=500, -) +**Note:** All agent skills are located in `.claude/skills/` directory. -# For Responses API agents: -llm = ChatDatabricks(endpoint="my-agent-endpoint", use_responses_api=True) -``` +> **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. --- -### DatabricksEmbeddings - Generate Embeddings +## Quick Commands -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"]) -``` +| 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` | --- -### 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..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", @@ -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__": 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..49b0cb6a --- /dev/null +++ b/agent-non-conversational/.claude/skills/agent-memory/SKILL.md @@ -0,0 +1,386 @@ +--- +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") +``` + +--- + +## 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 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: +- `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/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] 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..65c79099 --- /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 `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) + +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