From 917f98bff61097fb283eafb48363c3e3428ad12a Mon Sep 17 00:00:00 2001 From: Jenny Date: Tue, 3 Feb 2026 15:57:49 -0800 Subject: [PATCH 01/14] openai agents stateful example --- .../.claude/skills/add-tools/SKILL.md | 83 +++ .../add-tools/examples/custom-mcp-server.md | 60 ++ .../skills/add-tools/examples/experiment.yaml | 8 + .../add-tools/examples/genie-space.yaml | 9 + .../add-tools/examples/serving-endpoint.yaml | 7 + .../add-tools/examples/sql-warehouse.yaml | 7 + .../add-tools/examples/uc-connection.yaml | 9 + .../add-tools/examples/uc-function.yaml | 9 + .../add-tools/examples/vector-search.yaml | 9 + .../.claude/skills/deploy/SKILL.md | 222 +++++++ .../.claude/skills/discover-tools/SKILL.md | 47 ++ .../.claude/skills/modify-agent/SKILL.md | 146 +++++ .../.claude/skills/quickstart/SKILL.md | 83 +++ .../.claude/skills/run-locally/SKILL.md | 90 +++ .../.env.example | 17 + .../.gitignore | 219 +++++++ .../AGENTS.md | 109 ++++ .../CLAUDE.md | 3 + .../README.md | 241 ++++++++ .../agent_server/__init__.py | 0 .../agent_server/agent.py | 84 +++ .../agent_server/evaluate_agent.py | 53 ++ .../agent_server/start_server.py | 17 + .../agent_server/utils.py | 44 ++ .../app.yaml | 16 + .../databricks.yml | 37 ++ .../pyproject.toml | 44 ++ .../requirements.txt | 1 + .../scripts/__init__.py | 0 .../scripts/discover_tools.py | 432 +++++++++++++ .../scripts/quickstart.py | 575 ++++++++++++++++++ .../scripts/start_app.py | 253 ++++++++ 32 files changed, 2934 insertions(+) create mode 100644 agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/SKILL.md create mode 100644 agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/custom-mcp-server.md create mode 100644 agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/experiment.yaml create mode 100644 agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/genie-space.yaml create mode 100644 agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/serving-endpoint.yaml create mode 100644 agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/sql-warehouse.yaml create mode 100644 agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/uc-connection.yaml create mode 100644 agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/uc-function.yaml create mode 100644 agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/vector-search.yaml create mode 100644 agent-openai-agents-sdk-stateful-memory/.claude/skills/deploy/SKILL.md create mode 100644 agent-openai-agents-sdk-stateful-memory/.claude/skills/discover-tools/SKILL.md create mode 100644 agent-openai-agents-sdk-stateful-memory/.claude/skills/modify-agent/SKILL.md create mode 100644 agent-openai-agents-sdk-stateful-memory/.claude/skills/quickstart/SKILL.md create mode 100644 agent-openai-agents-sdk-stateful-memory/.claude/skills/run-locally/SKILL.md create mode 100644 agent-openai-agents-sdk-stateful-memory/.env.example create mode 100644 agent-openai-agents-sdk-stateful-memory/.gitignore create mode 100644 agent-openai-agents-sdk-stateful-memory/AGENTS.md create mode 100644 agent-openai-agents-sdk-stateful-memory/CLAUDE.md create mode 100644 agent-openai-agents-sdk-stateful-memory/README.md create mode 100644 agent-openai-agents-sdk-stateful-memory/agent_server/__init__.py create mode 100644 agent-openai-agents-sdk-stateful-memory/agent_server/agent.py create mode 100644 agent-openai-agents-sdk-stateful-memory/agent_server/evaluate_agent.py create mode 100644 agent-openai-agents-sdk-stateful-memory/agent_server/start_server.py create mode 100644 agent-openai-agents-sdk-stateful-memory/agent_server/utils.py create mode 100644 agent-openai-agents-sdk-stateful-memory/app.yaml create mode 100644 agent-openai-agents-sdk-stateful-memory/databricks.yml create mode 100644 agent-openai-agents-sdk-stateful-memory/pyproject.toml create mode 100644 agent-openai-agents-sdk-stateful-memory/requirements.txt create mode 100644 agent-openai-agents-sdk-stateful-memory/scripts/__init__.py create mode 100755 agent-openai-agents-sdk-stateful-memory/scripts/discover_tools.py create mode 100644 agent-openai-agents-sdk-stateful-memory/scripts/quickstart.py create mode 100644 agent-openai-agents-sdk-stateful-memory/scripts/start_app.py diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/SKILL.md b/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/SKILL.md new file mode 100644 index 00000000..6639e557 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/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/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/custom-mcp-server.md b/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/custom-mcp-server.md new file mode 100644 index 00000000..804bb679 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/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/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/experiment.yaml b/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/experiment.yaml new file mode 100644 index 00000000..ac5c626a --- /dev/null +++ b/agent-openai-agents-sdk-stateful-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-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/genie-space.yaml b/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/genie-space.yaml new file mode 100644 index 00000000..71589d52 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-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-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/serving-endpoint.yaml b/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/serving-endpoint.yaml new file mode 100644 index 00000000..b49ce9da --- /dev/null +++ b/agent-openai-agents-sdk-stateful-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-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/sql-warehouse.yaml b/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/sql-warehouse.yaml new file mode 100644 index 00000000..a6ce9446 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-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-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/uc-connection.yaml b/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/uc-connection.yaml new file mode 100644 index 00000000..316675fe --- /dev/null +++ b/agent-openai-agents-sdk-stateful-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-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/uc-function.yaml b/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/uc-function.yaml new file mode 100644 index 00000000..43f938a9 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-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-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/vector-search.yaml b/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/vector-search.yaml new file mode 100644 index 00000000..0ba39027 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-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-openai-agents-sdk-stateful-memory/.claude/skills/deploy/SKILL.md b/agent-openai-agents-sdk-stateful-memory/.claude/skills/deploy/SKILL.md new file mode 100644 index 00000000..d072e595 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-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_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 +# 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_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 + +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_openai_agents_sdk: + 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_openai_agents_sdk --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_openai_agents_sdk +``` + +### 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_openai_agents_sdk +``` + +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_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-stateful-memory/.claude/skills/discover-tools/SKILL.md b/agent-openai-agents-sdk-stateful-memory/.claude/skills/discover-tools/SKILL.md new file mode 100644 index 00000000..87c3f519 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-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-openai-agents-sdk-stateful-memory/.claude/skills/modify-agent/SKILL.md b/agent-openai-agents-sdk-stateful-memory/.claude/skills/modify-agent/SKILL.md new file mode 100644 index 00000000..e577b455 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/.claude/skills/modify-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/agent-openai-agents-sdk-stateful-memory/.claude/skills/quickstart/SKILL.md b/agent-openai-agents-sdk-stateful-memory/.claude/skills/quickstart/SKILL.md new file mode 100644 index 00000000..e550162c --- /dev/null +++ b/agent-openai-agents-sdk-stateful-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-openai-agents-sdk-stateful-memory/.claude/skills/run-locally/SKILL.md b/agent-openai-agents-sdk-stateful-memory/.claude/skills/run-locally/SKILL.md new file mode 100644 index 00000000..3eb83c82 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-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-openai-agents-sdk-stateful-memory/.env.example b/agent-openai-agents-sdk-stateful-memory/.env.example new file mode 100644 index 00000000..10a56922 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/.env.example @@ -0,0 +1,17 @@ +# Make a copy of this to set environment variables for local development +# cp .env.example .env + +# TODO: Fill in auth related env vars +DATABRICKS_CONFIG_PROFILE=DEFAULT +# DATABRICKS_HOST=https://.databricks.com +# DATABRICKS_TOKEN=dapi.... + +# TODO: Update with the MLflow experiment you want to log traces and models to +MLFLOW_EXPERIMENT_ID= + +CHAT_APP_PORT=3000 +CHAT_PROXY_TIMEOUT_SECONDS=300 +# IMPORTANT: For local development, use databricks (for default profile) or databricks:// to specify which Databricks CLI profile to use +# This is automatically configured by the quickstart script +MLFLOW_TRACKING_URI="databricks" +MLFLOW_REGISTRY_URI="databricks-uc" diff --git a/agent-openai-agents-sdk-stateful-memory/.gitignore b/agent-openai-agents-sdk-stateful-memory/.gitignore new file mode 100644 index 00000000..42eec123 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/.gitignore @@ -0,0 +1,219 @@ +# Created by https://www.toptal.com/developers/gitignore/api/python +# Edit at https://www.toptal.com/developers/gitignore?templates=python + +### Python ### +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + +# VS Code +.vscode/ + +### Python Patch ### +# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration +poetry.toml + +# ruff +.ruff_cache/ + +# LSP config files +pyrightconfig.json + +# End of https://www.toptal.com/developers/gitignore/api/python + +# Created by https://www.toptal.com/developers/gitignore/api/react +# Edit at https://www.toptal.com/developers/gitignore?templates=react + +### react ### +.DS_* +*.log +logs +**/*.backup.* +**/*.back.* + +node_modules +bower_components + +*.sublime* + +psd +thumb +sketch + +# End of https://www.toptal.com/developers/gitignore/api/react + +**/uv.lock +**/mlruns/ +**/.vite/ +**/.databricks + +# Ignore .claude directory but track template-provided skills +# User-created skills will be ignored by default (no conflicts) +.claude/* +!.claude/skills/ +.claude/skills/* +!.claude/skills/quickstart/ +!.claude/skills/discover-tools/ +!.claude/skills/deploy/ +!.claude/skills/add-tools/ +!.claude/skills/run-locally/ +!.claude/skills/modify-agent/ + +**/.env +**/.env.local diff --git a/agent-openai-agents-sdk-stateful-memory/AGENTS.md b/agent-openai-agents-sdk-stateful-memory/AGENTS.md new file mode 100644 index 00000000..ef319b9d --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/AGENTS.md @@ -0,0 +1,109 @@ +# Agent Development Guide + +## MANDATORY First Actions + +**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 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 + +**Ask the user questions to understand what they're building:** + +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?** + +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.** + +## Handling Deployment Errors + +**If `databricks bundle deploy` fails with "An app with the same name already exists":** + +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?" + +- **If they want to bind**: See the **deploy** skill for binding steps +- **If they want to delete**: Run `databricks apps delete ` then deploy again + +--- + +## Available Skills + +**Before executing any task, read the relevant skill file in `.claude/skills/`** - they contain tested commands, patterns, and troubleshooting steps. + +| 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` | + +**Note:** All agent skills are located in `.claude/skills/` directory. + +--- + +## Quick Commands + +| 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_openai_agents_sdk` | +| View logs | `databricks apps logs --follow` | + +--- + +## 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 | +| `databricks.yml` | Bundle config & resource permissions | +| `scripts/quickstart.py` | One-command setup script | +| `scripts/discover_tools.py` | Discovers available workspace resources | + +--- + +## Agent Framework Capabilities + +> **⚠️ 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:** +- **system.ai.python_exec** - Execute Python code dynamically within agent queries (code interpreter) + +**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 + +Reference: https://docs.databricks.com/aws/en/generative-ai/agent-framework/ diff --git a/agent-openai-agents-sdk-stateful-memory/CLAUDE.md b/agent-openai-agents-sdk-stateful-memory/CLAUDE.md new file mode 100644 index 00000000..4202b8f0 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/CLAUDE.md @@ -0,0 +1,3 @@ +@AGENTS.md + + diff --git a/agent-openai-agents-sdk-stateful-memory/README.md b/agent-openai-agents-sdk-stateful-memory/README.md new file mode 100644 index 00000000..6c171db9 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/README.md @@ -0,0 +1,241 @@ +# Responses API Agent + +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). + +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. + +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. + +This script will: + +1. Verify uv, nvm, and Databricks CLI installations +2. Configure Databricks authentication +3. Configure agent tracing, by creating and linking an MLflow experiment to your app +4. Start the agent server and chat app + +```bash +uv run quickstart +``` + +After the setup is complete, you can start the agent server and the chat app locally with: + +```bash +uv run start-app +``` + +This will start the agent server and the chat app at http://localhost:8000. + +**Next steps**: see [modifying your agent](#modifying-your-agent) to customize and iterate on the agent code. + +## Manual local development loop setup + +1. **Set up your local environment** + Install `uv` (python package manager), `nvm` (node version manager), and the Databricks CLI: + + - [`uv` installation docs](https://docs.astral.sh/uv/getting-started/installation/) + - [`nvm` installation](https://github.com/nvm-sh/nvm?tab=readme-ov-file#installing-and-updating) + - Run the following to use Node 20 LTS: + ```bash + nvm use 20 + ``` + - [`databricks CLI` installation](https://docs.databricks.com/aws/en/dev-tools/cli/install) + +2. **Set up local authentication to Databricks** + + In order to access Databricks resources from your local machine while developing your agent, you need to authenticate with Databricks. Choose one of the following options: + + **Option 1: OAuth via Databricks CLI (Recommended)** + + Authenticate with Databricks using the CLI. See the [CLI OAuth documentation](https://docs.databricks.com/aws/en/dev-tools/cli/authentication#oauth-user-to-machine-u2m-authentication). + + ```bash + databricks auth login + ``` + + Set the `DATABRICKS_CONFIG_PROFILE` environment variable in your .env file to the profile you used to authenticate: + + ```bash + DATABRICKS_CONFIG_PROFILE="DEFAULT" # change to the profile name you chose + ``` + + **Option 2: Personal Access Token (PAT)** + + See the [PAT documentation](https://docs.databricks.com/aws/en/dev-tools/auth/pat#databricks-personal-access-tokens-for-workspace-users). + + ```bash + # Add these to your .env file + DATABRICKS_HOST="https://host.databricks.com" + DATABRICKS_TOKEN="dapi_token" + ``` + + See the [Databricks SDK authentication docs](https://docs.databricks.com/aws/en/dev-tools/sdk-python#authenticate-the-databricks-sdk-for-python-with-your-databricks-account-or-workspace). + +3. **Create and link an MLflow experiment to your app** + + Create an MLflow experiment to enable tracing and version tracking. This is automatically done by the `uv run quickstart` script. + + Create the MLflow experiment via the CLI: + + ```bash + DATABRICKS_USERNAME=$(databricks current-user me | jq -r .userName) + databricks experiments create-experiment /Users/$DATABRICKS_USERNAME/agents-on-apps + ``` + + Make a copy of `.env.example` to `.env` and update the `MLFLOW_EXPERIMENT_ID` in your `.env` file with the experiment ID you created. The `.env` file will be automatically loaded when starting the server. + + ```bash + cp .env.example .env + # Edit .env and fill in your experiment ID + ``` + + See the [MLflow experiments documentation](https://docs.databricks.com/aws/en/mlflow/experiments#create-experiment-from-the-workspace). + +4. **Test your agent locally** + + Start up the agent server and chat UI locally: + + ```bash + uv run start-app + ``` + + Query your agent via the UI (http://localhost:8000) or REST API: + + **Advanced server options:** + + ```bash + uv run start-server --reload # hot-reload the server on code changes + uv run start-server --port 8001 # change the port the server listens on + uv run start-server --workers 4 # run the server with multiple workers + ``` + + - Example streaming request: + ```bash + curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' + ``` + - Example non-streaming request: + ```bash + curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }] }' + ``` + +## Modifying your agent + +See the [OpenAI Agents SDK documentation](https://platform.openai.com/docs/guides/agents-sdk) for more information on how to edit your own agent. + +Required files for hosting with MLflow `AgentServer`: + +- `agent.py`: Contains your agent logic. Modify this file to create your custom agent. For example, you can [add agent tools](https://docs.databricks.com/aws/en/generative-ai/agent-framework/agent-tool) to give your agent additional capabilities +- `start_server.py`: Initializes and runs the MLflow `AgentServer` with agent_type="ResponsesAgent". You don't have to modify this file for most common use cases, but can add additional server routes (e.g. a `/metrics` endpoint) here + +**Common customization questions:** + +**Q: Can I add additional files or folders to my agent?** +Yes. Add additional files or folders as needed. Ensure the script within `pyproject.toml` runs the correct script that starts the server and sets up MLflow tracing. + +**Q: How do I add dependencies to my agent?** +Run `uv add ` (e.g., `uv add "mlflow-skinny[databricks]"`). See the [python pyproject.toml guide](https://packaging.python.org/en/latest/guides/writing-pyproject-toml/#dependencies-and-requirements). + +**Q: Can I add custom tracing beyond the built-in tracing?** +Yes. This template uses MLflow's agent server, which comes with automatic tracing for agent logic decorated with `@invoke()` and `@stream()`. It also uses [MLflow autologging APIs](https://mlflow.org/docs/latest/genai/tracing/#one-line-auto-tracing-integrations) to capture traces from LLM invocations. However, you can add additional instrumentation to capture more granular trace information when your agent runs. See the [MLflow tracing documentation](https://docs.databricks.com/aws/en/mlflow3/genai/tracing/app-instrumentation/). + +**Q: How can I extend this example with additional tools and capabilities?** +This template can be extended by integrating additional MCP servers, Vector Search Indexes, UC Functions, and other Databricks tools. See the ["Agent Framework Tools Documentation"](https://docs.databricks.com/aws/en/generative-ai/agent-framework/agent-tool). + +## Evaluating your agent + +Evaluate your agent by calling the invoke function you defined for the agent locally. + +- Update your `evaluate_agent.py` file with the preferred evaluation dataset and scorers. + +Run the evaluation using the evaluation script: + +```bash +uv run agent-evaluate +``` + +After it completes, open the MLflow UI link for your experiment to inspect results. + +## Deploying to Databricks Apps + +0. **Create a Databricks App**: + Ensure you have the [Databricks CLI](https://docs.databricks.com/aws/en/dev-tools/cli/tutorial) installed and configured. + + ```bash + databricks apps create agent-openai-agents-sdk + ``` + +1. **Set up authentication to Databricks resources** + + For this example, you need to add an MLflow Experiment as a resource to your app. Grant the App's Service Principal (SP) permission to edit the experiment by clicking `edit` on your app home page. See the [Databricks Apps MLflow experiment documentation](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/mlflow) for more information. + + To grant access to other resources like serving endpoints, genie spaces, UC Functions, and Vector Search Indexes, click `edit` on your app home page to grant the App's SP permission. See the [Databricks Apps resources documentation](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/resources). + + For resources that are not supported yet, see the [Agent Framework authentication documentation](https://docs.databricks.com/aws/en/generative-ai/agent-framework/deploy-agent#automatic-authentication-passthrough) for the correct permission level to grant to your app SP. + + **On-behalf-of (OBO) User Authentication**: Use `get_user_workspace_client()` from `agent_server.utils` to authenticate as the requesting user instead of the app service principal. See the [OBO authentication documentation](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/auth?language=Streamlit#retrieve-user-authorization-credentials). + +2. **Sync local files to your workspace** + + See the [Databricks Apps deploy documentation](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/deploy?language=Databricks+CLI#deploy-the-app). + + ```bash + DATABRICKS_USERNAME=$(databricks current-user me | jq -r .userName) + databricks sync . "/Users/$DATABRICKS_USERNAME/agent-openai-agents-sdk" + ``` + +3. **Deploy your Databricks App** + + See the [Databricks Apps deploy documentation](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/deploy?language=Databricks+CLI#deploy-the-app). + + ```bash + databricks apps deploy agent-openai-agents-sdk --source-code-path /Workspace/Users/$DATABRICKS_USERNAME/agent-openai-agents-sdk + ``` + +4. **Query your agent hosted on Databricks Apps** + + Databricks Apps are _only_ queryable via OAuth token. You cannot use a PAT to query your agent. Generate an [OAuth token with your credentials using the Databricks CLI](https://docs.databricks.com/aws/en/dev-tools/cli/authentication#u2m-auth): + + ```bash + databricks auth login --host + databricks auth token + ``` + + Send a request to the `/invocations` endpoint: + + - Example streaming request: + + ```bash + curl -X POST /invocations \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' + ``` + + - Example non-streaming request: + + ```bash + curl -X POST /invocations \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }] }' + ``` + +For future updates to the agent, sync and redeploy your agent. + +### FAQ + +- For a streaming response, I see a 200 OK in the logs, but an error in the actual stream. What's going on? + - This is expected behavior. The initial 200 OK confirms stream setup; streaming errors don't affect this status. +- When querying my agent, I get a 302 error. What's going on? + - Use an OAuth token. PATs are not supported for querying agents. diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/__init__.py b/agent-openai-agents-sdk-stateful-memory/agent_server/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py new file mode 100644 index 00000000..2c48891a --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py @@ -0,0 +1,84 @@ +import os +from typing import AsyncGenerator + +import mlflow +from agents import Agent, Runner, set_default_openai_api, set_default_openai_client +from agents.tracing import set_trace_processors +from databricks_openai import AsyncDatabricksOpenAI +from databricks_openai.agents import McpServer +from databricks_openai.agents.session import MemorySession +from mlflow.genai.agent_server import invoke, stream +from mlflow.types.responses import ( + ResponsesAgentRequest, + ResponsesAgentResponse, + ResponsesAgentStreamEvent, +) + +from agent_server.utils import ( + get_databricks_host_from_env, + get_user_workspace_client, + process_agent_stream_events, +) + +# Lakebase instance name for persistent session storage +LAKEBASE_INSTANCE_NAME = os.environ.get("LAKEBASE_INSTANCE_NAME", "lakebase") + +# NOTE: this will work for all databricks models OTHER than GPT-OSS, which uses a slightly different API +set_default_openai_client(AsyncDatabricksOpenAI()) +set_default_openai_api("chat_completions") +set_trace_processors([]) # only use mlflow for trace processing +mlflow.openai.autolog() + + +async def init_mcp_server(): + return McpServer( + url=f"{get_databricks_host_from_env()}/api/2.0/mcp/functions/system/ai", + name="system.ai uc function mcp server", + ) + + +def create_coding_agent(mcp_server: McpServer) -> Agent: + return Agent( + name="code execution agent", + instructions="You are a code execution agent. You can execute code and return the results.", + model="databricks-gpt-5-2", + mcp_servers=[mcp_server], + ) + + +@invoke() +async def invoke(request: ResponsesAgentRequest) -> ResponsesAgentResponse: + # Optionally use the user's workspace client for on-behalf-of authentication + # user_workspace_client = get_user_workspace_client() + + # Create session for persistent conversation history with your Databricks Lakebase instance + session = MemorySession( + session_id=request.session_id, + instance_name=LAKEBASE_INSTANCE_NAME, + ) + + async with await init_mcp_server() as mcp_server: + agent = create_coding_agent(mcp_server) + messages = [i.model_dump() for i in request.input] + result = await Runner.run(agent, messages, session=session) + return ResponsesAgentResponse(output=[item.to_input_item() for item in result.new_items]) + + +@stream() +async def stream(request: ResponsesAgentRequest) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + # Optionally use the user's workspace client for on-behalf-of authentication + # user_workspace_client = get_user_workspace_client() + + # Create session for persistent conversation history with your Databricks Lakebase instance + session = MemorySession( + session_id=request.session_id, + instance_name=LAKEBASE_INSTANCE_NAME, + ) + + async with await init_mcp_server() as mcp_server: + agent = create_coding_agent(mcp_server) + messages = [i.model_dump() for i in request.input] + result = Runner.run_streamed(agent, input=messages, session=session) + + async for event in process_agent_stream_events(result.stream_events()): + yield event diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/evaluate_agent.py b/agent-openai-agents-sdk-stateful-memory/agent_server/evaluate_agent.py new file mode 100644 index 00000000..a1449535 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/evaluate_agent.py @@ -0,0 +1,53 @@ +import asyncio + +import mlflow +from dotenv import load_dotenv +from mlflow.genai.agent_server import get_invoke_function +from mlflow.genai.scorers import RelevanceToQuery, Safety +from mlflow.types.responses import ResponsesAgentRequest, ResponsesAgentResponse + +# Load environment variables from .env if it exists +load_dotenv(dotenv_path=".env", override=True) + +# Import agent for our @invoke function to be found +from agent_server import agent # noqa: F401 + +# Create your evaluation dataset +# Refer to documentation for evaluations: +# Scorers: https://docs.databricks.com/aws/en/mlflow3/genai/eval-monitor/concepts/scorers +# Predefined LLM scorers: https://mlflow.org/docs/latest/genai/eval-monitor/scorers/llm-judge/predefined +# Defining custom scorers: https://docs.databricks.com/aws/en/mlflow3/genai/eval-monitor/custom-scorers +eval_dataset = [ + { + "inputs": { + "request": { + "input": [{"role": "user", "content": "Calculate the 15th Fibonacci number"}] + } + }, + "expected_response": "The 15th Fibonacci number is 610.", + } +] + +# Get the invoke function that was registered via @invoke decorator in your agent +invoke_fn = get_invoke_function() +assert invoke_fn is not None, ( + "No function registered with the `@invoke` decorator found." + "Ensure you have a function decorated with `@invoke()`." +) + +# if invoke function is async, then we need to wrap it in a sync function +if asyncio.iscoroutinefunction(invoke_fn): + + def sync_invoke_fn(request: dict) -> ResponsesAgentResponse: + req = ResponsesAgentRequest(**request) + return asyncio.run(invoke_fn(req)) +else: + sync_invoke_fn = invoke_fn + + +def evaluate(): + mlflow.genai.evaluate( + data=eval_dataset, + predict_fn=sync_invoke_fn, + scorers=[RelevanceToQuery(), Safety()], + ) diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/start_server.py b/agent-openai-agents-sdk-stateful-memory/agent_server/start_server.py new file mode 100644 index 00000000..1d5ecd0c --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/start_server.py @@ -0,0 +1,17 @@ +from dotenv import load_dotenv +from mlflow.genai.agent_server import AgentServer, setup_mlflow_git_based_version_tracking + +# Load env vars from .env before importing the agent for proper auth +load_dotenv(dotenv_path=".env", override=True) + +# Need to import the agent to register the functions with the server +import agent_server.agent # noqa: E402 + +agent_server = AgentServer("ResponsesAgent", enable_chat_proxy=True) +# Define the app as a module level variable to enable multiple workers +app = agent_server.app # noqa: F841 +setup_mlflow_git_based_version_tracking() + + +def main(): + agent_server.run(app_import_string="agent_server.start_server:app") diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py b/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py new file mode 100644 index 00000000..03e52e9e --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py @@ -0,0 +1,44 @@ +import logging +from typing import AsyncGenerator, AsyncIterator, Optional +from uuid import uuid4 + +from agents.result import StreamEvent +from databricks.sdk import WorkspaceClient +from mlflow.genai.agent_server import get_request_headers +from mlflow.types.responses import ResponsesAgentStreamEvent + + +def get_databricks_host_from_env() -> Optional[str]: + try: + w = WorkspaceClient() + return w.config.host + except Exception as e: + logging.exception(f"Error getting databricks host from env: {e}") + return None + + +def get_user_workspace_client() -> WorkspaceClient: + token = get_request_headers().get("x-forwarded-access-token") + return WorkspaceClient(token=token, auth_type="pat") + + +async def process_agent_stream_events( + async_stream: AsyncIterator[StreamEvent], +) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + curr_item_id = str(uuid4()) + async for event in async_stream: + if event.type == "raw_response_event": + event_data = event.data.model_dump() + if event_data["type"] == "response.output_item.added": + curr_item_id = str(uuid4()) + event_data["item"]["id"] = curr_item_id + elif event_data.get("item") is not None and event_data["item"].get("id") is not None: + event_data["item"]["id"] = curr_item_id + elif event_data.get("item_id") is not None: + event_data["item_id"] = curr_item_id + yield event_data + elif event.type == "run_item_stream_event" and event.item.type == "tool_call_output_item": + yield ResponsesAgentStreamEvent( + type="response.output_item.done", + item=event.item.to_input_item(), + ) diff --git a/agent-openai-agents-sdk-stateful-memory/app.yaml b/agent-openai-agents-sdk-stateful-memory/app.yaml new file mode 100644 index 00000000..34465373 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/app.yaml @@ -0,0 +1,16 @@ +command: ["uv", "run", "start-app"] +# databricks apps listen by default on port 8000 + +env: + - name: MLFLOW_TRACKING_URI + value: "databricks" + - name: MLFLOW_REGISTRY_URI + value: "databricks-uc" + - name: API_PROXY + value: "http://localhost:8000/invocations" + - name: CHAT_APP_PORT + value: "3000" + - name: CHAT_PROXY_TIMEOUT_SECONDS + value: "300" + - name: MLFLOW_EXPERIMENT_ID + valueFrom: "experiment" diff --git a/agent-openai-agents-sdk-stateful-memory/databricks.yml b/agent-openai-agents-sdk-stateful-memory/databricks.yml new file mode 100644 index 00000000..48e2794f --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/databricks.yml @@ -0,0 +1,37 @@ +bundle: + name: agent_openai_agents_sdk + +resources: + # MLflow experiment for agent tracing - automatically created by bundle + experiments: + agent_openai_agents_sdk_experiment: + name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} + + apps: + agent_openai_agents_sdk: + name: "${bundle.target}-agent-openai-agents-sdk" + description: "OpenAI Agents SDK agent application" + source_code_path: ./ + + # Resources which this app has access to + resources: + - name: 'experiment' + experiment: + experiment_id: "${resources.experiments.agent_openai_agents_sdk_experiment.id}" + permission: 'CAN_MANAGE' + +targets: + dev: + mode: development + default: true + # workspace: + # host: https://... + + prod: + mode: production + # workspace: + # host: https://... + resources: + apps: + agent_openai_agents_sdk: + name: agent-openai-agents-sdk diff --git a/agent-openai-agents-sdk-stateful-memory/pyproject.toml b/agent-openai-agents-sdk-stateful-memory/pyproject.toml new file mode 100644 index 00000000..1fe2aa61 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/pyproject.toml @@ -0,0 +1,44 @@ +[project] +name = "agent-server-stateful-memory" +version = "0.1.0" +description = "MLflow-compatible agent server with FastAPI and persistent memory using Lakebase" +readme = "README.md" +authors = [ + { name = "Agent Developer", email = "developer@example.com" } +] +requires-python = ">=3.11" +dependencies = [ + "fastapi>=0.115.12", + "uvicorn>=0.34.2", + "mlflow>=3.8.0rc0", + "openai-agents>=0.4.1", + "python-dotenv", + "uuid-utils>=0.10.0", + # Note: Using git reference until official release + "databricks-openai[memory] @ git+https://github.com/databricks/databricks-ai-bridge@openai-memory-session#subdirectory=integrations/openai", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.metadata] +# Note: Required for direct git references. Delete when official databricks-openai[memory] release is available +allow-direct-references = true + +[tool.uv] +prerelease = "allow" + +[dependency-groups] +dev = [ + "hatchling>=1.27.0", + "pytest>=7.0.0", +] + + +[project.scripts] +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-openai-agents-sdk-stateful-memory/requirements.txt b/agent-openai-agents-sdk-stateful-memory/requirements.txt new file mode 100644 index 00000000..60cc5e6a --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/requirements.txt @@ -0,0 +1 @@ +uv diff --git a/agent-openai-agents-sdk-stateful-memory/scripts/__init__.py b/agent-openai-agents-sdk-stateful-memory/scripts/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/agent-openai-agents-sdk-stateful-memory/scripts/discover_tools.py b/agent-openai-agents-sdk-stateful-memory/scripts/discover_tools.py new file mode 100755 index 00000000..3eb37963 --- /dev/null +++ b/agent-openai-agents-sdk-stateful-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-openai-agents-sdk-stateful-memory/scripts/quickstart.py b/agent-openai-agents-sdk-stateful-memory/scripts/quickstart.py new file mode 100644 index 00000000..4147f1ea --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/scripts/quickstart.py @@ -0,0 +1,575 @@ +#!/usr/bin/env python3 +""" +Quickstart setup script for Databricks agent development. + +This script handles: +- Checking prerequisites (uv, nvm, Node 20, Databricks CLI) +- Databricks authentication (OAuth) +- MLflow experiment creation +- Environment variable configuration (.env) +- Lakebase instance setup (for memory-enabled templates) + +Usage: + uv run quickstart [OPTIONS] + +Options: + --profile NAME Use specified Databricks profile (non-interactive) + --host URL Databricks workspace URL (for initial setup) + --lakebase NAME Lakebase instance name (for memory features) + -h, --help Show this help message +""" + +import argparse +import json +import os +import platform +import re +import secrets +import shutil +import subprocess +import sys +from pathlib import Path + + +def print_header(text: str) -> None: + """Print a section header.""" + print(f"\n{'=' * 67}") + print(text) + print('=' * 67) + + +def print_step(text: str) -> None: + """Print a step indicator.""" + print(f"\n{text}") + + +def print_success(text: str) -> None: + """Print a success message.""" + print(f"✓ {text}") + + +def print_error(text: str) -> None: + """Print an error message.""" + print(f"✗ {text}", file=sys.stderr) + + +def print_troubleshooting_auth() -> None: + print("\nTroubleshooting tips:") + print(" • Ensure you have network connectivity to your Databricks workspace") + print(" • Try running 'databricks auth login' manually to see detailed errors") + print(" • Check that your workspace URL is correct") + print(" • If using a browser for OAuth, ensure popups are not blocked") + + +def print_troubleshooting_api() -> None: + print("\nTroubleshooting tips:") + print(" • Your authentication token may have expired - try 'databricks auth login' to refresh") + print(" • Verify your profile is valid with 'databricks auth profiles'") + print(" • Check network connectivity to your Databricks workspace") + + +def command_exists(cmd: str) -> bool: + """Check if a command exists in PATH.""" + return shutil.which(cmd) is not None + + +def run_command( + cmd: list[str], + capture_output: bool = True, + check: bool = True, + env: dict = None, + show_output: bool = False, +) -> subprocess.CompletedProcess: + """Run a command and return the result.""" + merged_env = {**os.environ, **(env or {})} + if show_output: + return subprocess.run(cmd, check=check, env=merged_env) + return subprocess.run(cmd, capture_output=capture_output, text=True, check=check, env=merged_env) + + +def get_command_output(cmd: list[str], env: dict = None) -> str: + """Run a command and return its stdout.""" + result = run_command(cmd, env=env) + return result.stdout.strip() + + +def check_prerequisites() -> dict[str, bool]: + """Check which prerequisites are installed.""" + print_step("Checking prerequisites...") + + prereqs = { + "uv": command_exists("uv"), + "node": command_exists("node"), + "npm": command_exists("npm"), + "databricks": command_exists("databricks"), + } + + for name, installed in prereqs.items(): + if installed: + try: + if name == "uv": + version = get_command_output(["uv", "--version"]) + elif name == "node": + version = get_command_output(["node", "--version"]) + elif name == "npm": + version = get_command_output(["npm", "--version"]) + elif name == "databricks": + version = get_command_output(["databricks", "--version"]) + print_success(f"{name} is installed: {version}") + except Exception: + print_success(f"{name} is installed") + else: + print(f" {name} is not installed") + + return prereqs + + +def check_missing_prerequisites(prereqs: dict[str, bool]) -> list[str]: + """Return list of missing prerequisites with install instructions.""" + missing = [] + + if not prereqs["uv"]: + missing.append("uv - Install with: curl -LsSf https://astral.sh/uv/install.sh | sh") + + if not prereqs["node"] or not prereqs["npm"]: + missing.append("Node.js 20 - Install with: nvm install 20 (or download from nodejs.org)") + + if not prereqs["databricks"]: + if platform.system() == "Darwin": + missing.append("Databricks CLI - Install with: brew install databricks/tap/databricks") + else: + missing.append("Databricks CLI - Install with: curl -fsSL https://raw.githubusercontent.com/databricks/setup-cli/main/install.sh | sh") + + if missing: + missing.append("Note: These install commands are for Unix/macOS. For Windows, please visit the official documentation for each tool.") + + return missing + + +def setup_env_file() -> None: + """Copy .env.example to .env if it doesn't exist.""" + print_step("Setting up configuration files...") + + env_local = Path(".env") + env_example = Path(".env.example") + + if env_local.exists(): + print(" .env already exists, skipping copy...") + elif env_example.exists(): + shutil.copy(env_example, env_local) + print_success("Copied .env.example to .env") + else: + # Create a minimal .env + env_local.write_text( + "# Databricks configuration\n" + "DATABRICKS_CONFIG_PROFILE=DEFAULT\n" + "MLFLOW_EXPERIMENT_ID=\n" + 'MLFLOW_TRACKING_URI="databricks"\n' + 'MLFLOW_REGISTRY_URI="databricks-uc"\n' + ) + print_success("Created .env") + + +def update_env_file(key: str, value: str) -> None: + """Update or add a key-value pair in .env.""" + env_file = Path(".env") + + if not env_file.exists(): + env_file.write_text(f"{key}={value}\n") + return + + content = env_file.read_text() + + # Check if key exists (with or without quotes, with any value) + pattern = rf'^{re.escape(key)}=.*$' + if re.search(pattern, content, re.MULTILINE): + # Replace existing key + content = re.sub(pattern, f"{key}={value}", content, flags=re.MULTILINE) + else: + # Add new key + if not content.endswith("\n"): + content += "\n" + content += f"{key}={value}\n" + + env_file.write_text(content) + + +def get_databricks_profiles() -> list[dict]: + """Get list of existing Databricks profiles.""" + try: + result = run_command(["databricks", "auth", "profiles"], check=False) + if result.returncode != 0 or not result.stdout.strip(): + return [] + + lines = result.stdout.strip().split("\n") + if len(lines) <= 1: # Only header or empty + return [] + + # Parse the output - first line is header + profiles = [] + for line in lines[1:]: + if line.strip(): + # Profile name is the first column + parts = line.split() + if parts: + profiles.append({ + "name": parts[0], + "line": line, + }) + + return profiles + except Exception: + return [] + + +def validate_profile(profile_name: str) -> bool: + """Test if a Databricks profile is authenticated.""" + try: + env = {"DATABRICKS_CONFIG_PROFILE": profile_name} + result = run_command( + ["databricks", "current-user", "me"], + check=False, + env=env, + ) + return result.returncode == 0 + except Exception: + return False + + +def authenticate_profile(profile_name: str, host: str = None) -> bool: + """Authenticate a Databricks profile.""" + print(f"\nAuthenticating profile '{profile_name}'...") + print("You will be prompted to log in to Databricks in your browser.\n") + + cmd = ["databricks", "auth", "login", "--profile", profile_name] + if host: + cmd.extend(["--host", host]) + + try: + # Run interactively so user can see browser prompt + result = subprocess.run(cmd) + return result.returncode == 0 + except Exception as e: + print_error(f"Authentication failed: {e}") + return False + + +def select_profile_interactive(profiles: list[dict]) -> str: + """Let user select a profile interactively.""" + print("\nFound existing Databricks profiles:\n") + + # Print header and profiles + for i, profile in enumerate(profiles, 1): + print(f" {i}) {profile['line']}") + + print() + + while True: + choice = input("Enter the number of the profile you want to use: ").strip() + if not choice: + print_error("Profile selection is required") + continue + + try: + index = int(choice) - 1 + if 0 <= index < len(profiles): + return profiles[index]["name"] + else: + print_error(f"Please choose a number between 1 and {len(profiles)}") + except ValueError: + print_error("Please enter a valid number") + + +def setup_databricks_auth(profile_arg: str = None, host_arg: str = None) -> str: + """Set up Databricks authentication and return the profile name.""" + print_step("Setting up Databricks authentication...") + + # If profile was specified via CLI, use it directly + if profile_arg: + profile_name = profile_arg + print(f"Using specified profile: {profile_name}") + else: + # Check for existing profiles + profiles = get_databricks_profiles() + + if profiles: + profile_name = select_profile_interactive(profiles) + print(f"\nSelected profile: {profile_name}") + else: + # No profiles exist - need to create one + profile_name = None + + # Validate or authenticate the profile + if profile_name: + if validate_profile(profile_name): + print_success(f"Successfully validated profile '{profile_name}'") + else: + print(f"Profile '{profile_name}' is not authenticated.") + if not authenticate_profile(profile_name): + print_error(f"Failed to authenticate profile '{profile_name}'") + print_troubleshooting_auth() + sys.exit(1) + print_success(f"Successfully authenticated profile '{profile_name}'") + else: + # Create new profile + print("No existing profiles found. Setting up Databricks authentication...") + + if host_arg: + host = host_arg + print(f"Using specified host: {host}") + else: + host = input("\nPlease enter your Databricks host URL\n(e.g., https://your-workspace.cloud.databricks.com): ").strip() + + if not host: + print_error("Databricks host is required") + sys.exit(1) + + profile_name = "DEFAULT" + if not authenticate_profile(profile_name, host): + print_error("Databricks authentication failed") + print_troubleshooting_auth() + sys.exit(1) + print_success(f"Successfully authenticated with Databricks") + + # Update .env with profile + update_env_file("DATABRICKS_CONFIG_PROFILE", profile_name) + update_env_file("MLFLOW_TRACKING_URI", f'"databricks://{profile_name}"') + print_success(f"Databricks profile '{profile_name}' saved to .env") + + return profile_name + + +def get_databricks_username(profile_name: str) -> str: + """Get the current Databricks username.""" + try: + result = run_command( + ["databricks", "-p", profile_name, "current-user", "me", "--output", "json"] + ) + user_data = json.loads(result.stdout) + return user_data.get("userName", "") + except Exception as e: + print_error(f"Failed to get Databricks username: {e}") + print_troubleshooting_api() + sys.exit(1) + + +def create_mlflow_experiment(profile_name: str, username: str) -> tuple[str, str]: + """Create an MLflow experiment and return (name, id).""" + print_step("Creating MLflow experiment...") + + experiment_name = f"/Users/{username}/agents-on-apps" + + try: + # Try to create with default name + result = run_command( + ["databricks", "-p", profile_name, "experiments", "create-experiment", + experiment_name, "--output", "json"], + check=False + ) + + if result.returncode == 0: + experiment_id = json.loads(result.stdout).get("experiment_id", "") + print_success(f"Created experiment '{experiment_name}' with ID: {experiment_id}") + return experiment_name, experiment_id + + # Name already exists, try with random suffix + print("Experiment name already exists, creating with random suffix...") + random_suffix = secrets.token_hex(4) + experiment_name = f"/Users/{username}/agents-on-apps-{random_suffix}" + + result = run_command( + ["databricks", "-p", profile_name, "experiments", "create-experiment", + experiment_name, "--output", "json"] + ) + experiment_id = json.loads(result.stdout).get("experiment_id", "") + print_success(f"Created experiment '{experiment_name}' with ID: {experiment_id}") + return experiment_name, experiment_id + + except Exception as e: + print_error(f"Failed to create MLflow experiment: {e}") + print_troubleshooting_api() + sys.exit(1) + + +def check_lakebase_required() -> bool: + """Check if app.yaml has LAKEBASE_INSTANCE_NAME configured.""" + app_yaml = Path("app.yaml") + if not app_yaml.exists(): + return False + + content = app_yaml.read_text() + return "LAKEBASE_INSTANCE_NAME" in content + + +def get_env_value(key: str) -> str: + """Get a value from .env file.""" + env_file = Path(".env") + if not env_file.exists(): + return "" + + content = env_file.read_text() + pattern = rf'^{re.escape(key)}=(.*)$' + match = re.search(pattern, content, re.MULTILINE) + if match: + return match.group(1).strip().strip('"').strip("'") + return "" + + +def validate_lakebase_instance(profile_name: str, lakebase_name: str) -> bool: + """Validate that the Lakebase instance exists and user has access.""" + print(f"Validating Lakebase instance '{lakebase_name}'...") + + result = run_command( + ["databricks", "-p", profile_name, "database", "get-database-instance", + lakebase_name, "--output", "json"], + check=False + ) + + if result.returncode == 0: + print_success(f"Lakebase instance '{lakebase_name}' validated") + return True + + # Check if database command is not recognized (old CLI version) + if 'unknown command "database" for "databricks"' in (result.stderr or ""): + print_error("The 'databricks database' command requires a newer version of the Databricks CLI.") + print(" Please upgrade: https://docs.databricks.com/dev-tools/cli/install.html") + return False + + error_msg = result.stderr.lower() if result.stderr else "" + if "not found" in error_msg: + print_error(f"Lakebase instance '{lakebase_name}' not found. Please check the instance name.") + elif "permission" in error_msg or "forbidden" in error_msg or "unauthorized" in error_msg: + print_error(f"No permission to access Lakebase instance '{lakebase_name}'") + else: + print_error(f"Failed to validate Lakebase instance: {result.stderr.strip() if result.stderr else 'Unknown error'}") + return False + + +def setup_lakebase(profile_name: str, lakebase_arg: str = None) -> str: + """Set up Lakebase instance for memory features.""" + print_step("Setting up Lakebase instance for memory...") + + lakebase_name = None + + # If --lakebase was provided, use it directly + if lakebase_arg: + lakebase_name = lakebase_arg + print(f"Using provided Lakebase instance: {lakebase_name}") + else: + # Check if already set in .env + existing = get_env_value("LAKEBASE_INSTANCE_NAME") + if existing: + print(f"Found existing Lakebase instance in .env: {existing}") + new_value = input("Press Enter to keep this value, or enter a new instance name: ").strip() + lakebase_name = new_value if new_value else existing + else: + # Interactive mode - prompt for instance name + lakebase_name = input("Please enter your Lakebase instance name: ").strip() + + if not lakebase_name: + print_error("Lakebase instance name is required for memory features") + sys.exit(1) + + # Validate that the Lakebase instance exists and user has access + if not validate_lakebase_instance(profile_name, lakebase_name): + sys.exit(1) + + # Update .env with the Lakebase instance name + update_env_file("LAKEBASE_INSTANCE_NAME", lakebase_name) + print_success(f"Lakebase instance name '{lakebase_name}' saved to .env") + + return lakebase_name + + +def main(): + parser = argparse.ArgumentParser( + description="Quickstart setup for Databricks agent development", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + uv run quickstart # Interactive setup + uv run quickstart --profile DEFAULT # Use existing profile (non-interactive) + uv run quickstart --host https://... # Set up new profile with host + uv run quickstart --lakebase my-db # Include Lakebase setup for memory + """ + ) + parser.add_argument( + "--profile", + help="Use specified Databricks profile (non-interactive)", + metavar="NAME", + ) + parser.add_argument( + "--host", + help="Databricks workspace URL (for initial setup)", + metavar="URL", + ) + parser.add_argument( + "--lakebase", + help="Lakebase instance name (for memory features)", + metavar="NAME", + ) + + args = parser.parse_args() + + try: + print_header("Agent on Apps - Quickstart Setup") + + # Step 1: Check prerequisites + prereqs = check_prerequisites() + missing = check_missing_prerequisites(prereqs) + + if missing: + print_step("Missing prerequisites:") + for item in missing: + print(f" • {item}") + print("\nPlease install the missing prerequisites and run this script again.") + sys.exit(1) + + # Step 2: Set up .env + setup_env_file() + + # Step 3: Databricks authentication + profile_name = setup_databricks_auth(args.profile, args.host) + + # Step 4: Get username and create MLflow experiment + print_step("Getting Databricks username...") + username = get_databricks_username(profile_name) + print(f"Username: {username}") + + experiment_name, experiment_id = create_mlflow_experiment(profile_name, username) + + # Step 5: Update .env with experiment ID + update_env_file("MLFLOW_EXPERIMENT_ID", experiment_id) + print_success("Updated .env with experiment ID") + + # Step 6: Lakebase setup (if needed for memory features) + lakebase_name = None + lakebase_required = args.lakebase or check_lakebase_required() + if lakebase_required: + lakebase_name = setup_lakebase(profile_name, args.lakebase) + + # Final summary + print_header("Setup Complete!") + summary = f""" +✓ Prerequisites verified (uv, Node.js, Databricks CLI) +✓ Databricks authenticated with profile: {profile_name} +✓ Configuration files created (.env) +✓ MLflow experiment created: {experiment_name} +✓ Experiment ID: {experiment_id}""" + + if lakebase_name: + summary += f"\n✓ Lakebase instance: {lakebase_name}" + + summary += """ + +Next step: Run 'uv run start-app' to start the agent locally +""" + print(summary) + + except KeyboardInterrupt: + print("\n\nSetup cancelled.") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/agent-openai-agents-sdk-stateful-memory/scripts/start_app.py b/agent-openai-agents-sdk-stateful-memory/scripts/start_app.py new file mode 100644 index 00000000..9fe60cde --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/scripts/start_app.py @@ -0,0 +1,253 @@ +#!/usr/bin/env python3 +""" +Start script for running frontend and backend processes concurrently. + +Requirements: +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 +import sys +import threading +import time +from pathlib import Path + +from dotenv import load_dotenv + +# Readiness patterns +BACKEND_READY = [r"Uvicorn running on", r"Application startup complete", r"Started server process"] +FRONTEND_READY = [r"Server is running on http://localhost"] + + +class ProcessManager: + def __init__(self, port=8000): + self.backend_process = None + self.frontend_process = None + self.backend_ready = False + self.frontend_ready = False + 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 + try: + for line in iter(process.stdout.readline, ""): + if not line: + break + + line = line.rstrip() + log_file.write(line + "\n") + print(f"[{name}] {line}") + + # Check readiness + if not is_ready and any(re.search(p, line, re.IGNORECASE) for p in patterns): + is_ready = True + if name == "backend": + self.backend_ready = True + else: + self.frontend_ready = True + print(f"✓ {name.capitalize()} is ready!") + + if self.backend_ready and self.frontend_ready: + print("\n" + "=" * 50) + print("✓ Both frontend and backend are ready!") + print(f"✓ Open the frontend at http://localhost:{self.port}") + print("=" * 50 + "\n") + + process.wait() + if process.returncode != 0: + self.failed.set() + + except Exception as e: + print(f"Error monitoring {name}: {e}") + self.failed.set() + + def clone_frontend_if_needed(self): + if Path("e2e-chatbot-app-next").exists(): + return True + + print("Cloning e2e-chatbot-app-next...") + for url in [ + "https://github.com/databricks/app-templates.git", + "git@github.com:databricks/app-templates.git", + ]: + try: + subprocess.run( + ["git", "clone", "--filter=blob:none", "--sparse", url, "temp-app-templates"], + check=True, + capture_output=True, + ) + break + except subprocess.CalledProcessError: + continue + else: + print("ERROR: Failed to clone repository.") + print( + "Manually download from: https://download-directory.github.io/?url=https://github.com/databricks/app-templates/tree/main/e2e-chatbot-app-next" + ) + return False + + subprocess.run( + ["git", "sparse-checkout", "set", "e2e-chatbot-app-next"], + cwd="temp-app-templates", + check=True, + ) + Path("temp-app-templates/e2e-chatbot-app-next").rename("e2e-chatbot-app-next") + shutil.rmtree("temp-app-templates", ignore_errors=True) + return True + + def start_process(self, cmd, name, log_file, patterns, cwd=None): + print(f"Starting {name}...") + process = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, cwd=cwd + ) + + thread = threading.Thread( + target=self.monitor_process, args=(process, name, log_file, patterns), daemon=True + ) + thread.start() + return process + + def print_logs(self, log_path): + print(f"\nLast 50 lines of {log_path}:") + print("-" * 40) + try: + lines = Path(log_path).read_text().splitlines() + print("\n".join(lines[-50:])) + except FileNotFoundError: + print(f"(no {log_path} found)") + print("-" * 40) + + def cleanup(self): + print("\n" + "=" * 42) + print("Shutting down both processes...") + print("=" * 42) + + for proc in [self.backend_process, self.frontend_process]: + if proc: + try: + proc.terminate() + proc.wait(timeout=5) + except (subprocess.TimeoutExpired, Exception): + proc.kill() + + if self.backend_log: + self.backend_log.close() + if self.frontend_log: + self.frontend_log.close() + + 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( + backend_cmd, "backend", self.backend_log, BACKEND_READY + ) + + # Setup and start frontend + frontend_dir = Path("e2e-chatbot-app-next") + for cmd, desc in [("npm install", "install"), ("npm run build", "build")]: + print(f"Running npm {desc}...") + result = subprocess.run( + cmd.split(), cwd=frontend_dir, capture_output=True, text=True + ) + if result.returncode != 0: + print(f"npm {desc} failed: {result.stderr}") + return 1 + + self.frontend_process = self.start_process( + ["npm", "run", "start"], + "frontend", + self.frontend_log, + FRONTEND_READY, + cwd=frontend_dir, + ) + + print( + f"\nMonitoring processes (Backend PID: {self.backend_process.pid}, Frontend PID: {self.frontend_process.pid})\n" + ) + + # Wait for failure + while not self.failed.is_set(): + time.sleep(0.1) + for proc in [self.backend_process, self.frontend_process]: + if proc.poll() is not None: + self.failed.set() + break + + # Determine which failed + failed_name = "backend" if self.backend_process.poll() is not None else "frontend" + failed_proc = ( + self.backend_process if failed_name == "backend" else self.frontend_process + ) + exit_code = failed_proc.returncode if failed_proc else 1 + + print( + f"\n{'=' * 42}\nERROR: {failed_name} process exited with code {exit_code}\n{'=' * 42}" + ) + self.print_logs("backend.log") + self.print_logs("frontend.log") + return exit_code + + except KeyboardInterrupt: + print("\nInterrupted") + return 0 + + finally: + self.cleanup() + + +def main(): + 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__": + main() From 48745b9ad53c1cc3d77df23aa4992678d4e6aa5a Mon Sep 17 00:00:00 2001 From: Jenny Date: Tue, 3 Feb 2026 16:08:37 -0800 Subject: [PATCH 02/14] add session id to outputs --- .../agent_server/agent.py | 22 ++++++++++++++++--- .../pyproject.toml | 3 +++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py index 2c48891a..1095fdf9 100644 --- a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py @@ -1,4 +1,5 @@ import os +import uuid from typing import AsyncGenerator import mlflow @@ -23,6 +24,16 @@ # Lakebase instance name for persistent session storage LAKEBASE_INSTANCE_NAME = os.environ.get("LAKEBASE_INSTANCE_NAME", "lakebase") + +def get_session_id(request: ResponsesAgentRequest) -> str: + """Extract session_id from request or generate a new one.""" + # Try to get session_id from custom_inputs if provided + if hasattr(request, "custom_inputs") and request.custom_inputs: + if "session_id" in request.custom_inputs: + return request.custom_inputs["session_id"] + # Fall back to generating a new session_id + return str(uuid.uuid4()) + # NOTE: this will work for all databricks models OTHER than GPT-OSS, which uses a slightly different API set_default_openai_client(AsyncDatabricksOpenAI()) set_default_openai_api("chat_completions") @@ -51,9 +62,11 @@ async def invoke(request: ResponsesAgentRequest) -> ResponsesAgentResponse: # Optionally use the user's workspace client for on-behalf-of authentication # user_workspace_client = get_user_workspace_client() + session_id = get_session_id(request) + # Create session for persistent conversation history with your Databricks Lakebase instance session = MemorySession( - session_id=request.session_id, + session_id=session_id, instance_name=LAKEBASE_INSTANCE_NAME, ) @@ -61,7 +74,10 @@ async def invoke(request: ResponsesAgentRequest) -> ResponsesAgentResponse: agent = create_coding_agent(mcp_server) messages = [i.model_dump() for i in request.input] result = await Runner.run(agent, messages, session=session) - return ResponsesAgentResponse(output=[item.to_input_item() for item in result.new_items]) + return ResponsesAgentResponse( + output=[item.to_input_item() for item in result.new_items], + custom_outputs={"session_id": session_id}, + ) @stream() @@ -71,7 +87,7 @@ async def stream(request: ResponsesAgentRequest) -> AsyncGenerator[ResponsesAgen # Create session for persistent conversation history with your Databricks Lakebase instance session = MemorySession( - session_id=request.session_id, + session_id=get_session_id(request), instance_name=LAKEBASE_INSTANCE_NAME, ) diff --git a/agent-openai-agents-sdk-stateful-memory/pyproject.toml b/agent-openai-agents-sdk-stateful-memory/pyproject.toml index 1fe2aa61..c0a18c89 100644 --- a/agent-openai-agents-sdk-stateful-memory/pyproject.toml +++ b/agent-openai-agents-sdk-stateful-memory/pyproject.toml @@ -26,6 +26,9 @@ build-backend = "hatchling.build" # Note: Required for direct git references. Delete when official databricks-openai[memory] release is available allow-direct-references = true +[tool.hatch.build.targets.wheel] +packages = ["agent_server", "scripts"] + [tool.uv] prerelease = "allow" From 62b2010e7a893eb18a694a6dd04000d0415297a1 Mon Sep 17 00:00:00 2001 From: Jenny Date: Wed, 4 Feb 2026 17:06:12 -0800 Subject: [PATCH 03/14] update example w/ asyncdatabrickssession --- .../agent_server/agent.py | 7 +++---- agent-openai-agents-sdk-stateful-memory/pyproject.toml | 2 +- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py index 1095fdf9..eef2f503 100644 --- a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py @@ -6,8 +6,7 @@ from agents import Agent, Runner, set_default_openai_api, set_default_openai_client from agents.tracing import set_trace_processors from databricks_openai import AsyncDatabricksOpenAI -from databricks_openai.agents import McpServer -from databricks_openai.agents.session import MemorySession +from databricks_openai.agents import AsyncDatabricksSession, McpServer from mlflow.genai.agent_server import invoke, stream from mlflow.types.responses import ( ResponsesAgentRequest, @@ -65,7 +64,7 @@ async def invoke(request: ResponsesAgentRequest) -> ResponsesAgentResponse: session_id = get_session_id(request) # Create session for persistent conversation history with your Databricks Lakebase instance - session = MemorySession( + session = AsyncDatabricksSession( session_id=session_id, instance_name=LAKEBASE_INSTANCE_NAME, ) @@ -86,7 +85,7 @@ async def stream(request: ResponsesAgentRequest) -> AsyncGenerator[ResponsesAgen # user_workspace_client = get_user_workspace_client() # Create session for persistent conversation history with your Databricks Lakebase instance - session = MemorySession( + session = AsyncDatabricksSession( session_id=get_session_id(request), instance_name=LAKEBASE_INSTANCE_NAME, ) diff --git a/agent-openai-agents-sdk-stateful-memory/pyproject.toml b/agent-openai-agents-sdk-stateful-memory/pyproject.toml index c0a18c89..4795e1be 100644 --- a/agent-openai-agents-sdk-stateful-memory/pyproject.toml +++ b/agent-openai-agents-sdk-stateful-memory/pyproject.toml @@ -15,7 +15,7 @@ dependencies = [ "python-dotenv", "uuid-utils>=0.10.0", # Note: Using git reference until official release - "databricks-openai[memory] @ git+https://github.com/databricks/databricks-ai-bridge@openai-memory-session#subdirectory=integrations/openai", + "databricks-openai[memory] @ git+https://github.com/databricks/databricks-ai-bridge@openai-lakebasesqlalchemysession#subdirectory=integrations/openai", ] [build-system] From d1619bfaa23913344bb1a89f4e1c49054fd79329 Mon Sep 17 00:00:00 2001 From: Jenny Date: Tue, 10 Feb 2026 18:03:20 -0800 Subject: [PATCH 04/14] package release agent updates --- .../agent_server/agent.py | 10 ++++------ agent-openai-agents-sdk-stateful-memory/pyproject.toml | 7 +------ 2 files changed, 5 insertions(+), 12 deletions(-) diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py index eef2f503..d914d844 100644 --- a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py @@ -57,15 +57,13 @@ def create_coding_agent(mcp_server: McpServer) -> Agent: @invoke() -async def invoke(request: ResponsesAgentRequest) -> ResponsesAgentResponse: +async def invoke_handler(request: ResponsesAgentRequest) -> ResponsesAgentResponse: # Optionally use the user's workspace client for on-behalf-of authentication # user_workspace_client = get_user_workspace_client() - session_id = get_session_id(request) - # Create session for persistent conversation history with your Databricks Lakebase instance session = AsyncDatabricksSession( - session_id=session_id, + session_id=get_session_id(request), instance_name=LAKEBASE_INSTANCE_NAME, ) @@ -75,12 +73,12 @@ async def invoke(request: ResponsesAgentRequest) -> ResponsesAgentResponse: result = await Runner.run(agent, messages, session=session) return ResponsesAgentResponse( output=[item.to_input_item() for item in result.new_items], - custom_outputs={"session_id": session_id}, + custom_outputs={"session_id": session.session_id}, ) @stream() -async def stream(request: ResponsesAgentRequest) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: +async def stream_handler(request: ResponsesAgentRequest) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: # Optionally use the user's workspace client for on-behalf-of authentication # user_workspace_client = get_user_workspace_client() diff --git a/agent-openai-agents-sdk-stateful-memory/pyproject.toml b/agent-openai-agents-sdk-stateful-memory/pyproject.toml index 4795e1be..53c93ccb 100644 --- a/agent-openai-agents-sdk-stateful-memory/pyproject.toml +++ b/agent-openai-agents-sdk-stateful-memory/pyproject.toml @@ -14,18 +14,13 @@ dependencies = [ "openai-agents>=0.4.1", "python-dotenv", "uuid-utils>=0.10.0", - # Note: Using git reference until official release - "databricks-openai[memory] @ git+https://github.com/databricks/databricks-ai-bridge@openai-lakebasesqlalchemysession#subdirectory=integrations/openai", + "databricks-openai[memory]>=0.11.1", ] [build-system] requires = ["hatchling"] build-backend = "hatchling.build" -[tool.hatch.metadata] -# Note: Required for direct git references. Delete when official databricks-openai[memory] release is available -allow-direct-references = true - [tool.hatch.build.targets.wheel] packages = ["agent_server", "scripts"] From d8da4837321bbddb78d031620cd16baf2b80f73f Mon Sep 17 00:00:00 2001 From: Jenny Date: Tue, 10 Feb 2026 18:09:43 -0800 Subject: [PATCH 05/14] use uuid7 for example --- .../agent_server/agent.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py index d914d844..c48f2849 100644 --- a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py @@ -1,7 +1,8 @@ import os -import uuid from typing import AsyncGenerator +from uuid_utils import uuid7 + import mlflow from agents import Agent, Runner, set_default_openai_api, set_default_openai_client from agents.tracing import set_trace_processors @@ -31,7 +32,7 @@ def get_session_id(request: ResponsesAgentRequest) -> str: if "session_id" in request.custom_inputs: return request.custom_inputs["session_id"] # Fall back to generating a new session_id - return str(uuid.uuid4()) + return str(uuid7()) # NOTE: this will work for all databricks models OTHER than GPT-OSS, which uses a slightly different API set_default_openai_client(AsyncDatabricksOpenAI()) From 013443c6a44cca92ae1285817ae171655f2fd15f Mon Sep 17 00:00:00 2001 From: Jenny Date: Wed, 11 Feb 2026 13:35:51 -0800 Subject: [PATCH 06/14] pr review updates --- .../README.md | 14 +++++- .../agent_server/agent.py | 11 ++++- .../agent_server/utils.py | 44 +++++++++++++++++++ .../app.yaml | 2 + .../databricks.yml | 16 +++---- 5 files changed, 76 insertions(+), 11 deletions(-) diff --git a/agent-openai-agents-sdk-stateful-memory/README.md b/agent-openai-agents-sdk-stateful-memory/README.md index 6c171db9..fd858761 100644 --- a/agent-openai-agents-sdk-stateful-memory/README.md +++ b/agent-openai-agents-sdk-stateful-memory/README.md @@ -1,9 +1,19 @@ -# Responses API Agent +# Responses API Agent with Stateful 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). +This template defines a **stateful** conversational agent app with persistent conversation history backed by [Databricks Lakebase](https://docs.databricks.com/aws/en/lakebase/). 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. +### Stateful Sessions + +This template uses OpenAI Agents SDK [Sessions](https://openai.github.io/openai-agents-python/sessions/) to automatically manage conversation history. Sessions store conversation history for a specific session, allowing agents to maintain context without requiring explicit manual memory management. This is particularly useful for building chat applications or multi-turn conversations where you want the agent to remember previous interactions. + +How it works: +- **Before each run**: The session retrieves prior conversation history and prepends it to the input +- **After each run**: New items (user messages, assistant responses, tool calls) are automatically stored in the session + +This template uses `AsyncDatabricksSession` from `databricks-openai`, which persists session data to a Databricks Lakebase instance. + 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 diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py index c48f2849..fdbd7b86 100644 --- a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py @@ -19,10 +19,19 @@ get_databricks_host_from_env, get_user_workspace_client, process_agent_stream_events, + resolve_lakebase_instance_name, ) # Lakebase instance name for persistent session storage -LAKEBASE_INSTANCE_NAME = os.environ.get("LAKEBASE_INSTANCE_NAME", "lakebase") +_LAKEBASE_INSTANCE_NAME_RAW = os.environ.get("LAKEBASE_INSTANCE_NAME") +if not _LAKEBASE_INSTANCE_NAME_RAW: + raise ValueError( + "LAKEBASE_INSTANCE_NAME environment variable is required but not set. " + "Please set it in your environment:\n" + " LAKEBASE_INSTANCE_NAME=\n" + ) +# Resolve hostname to instance name if needed (if given hostname of lakebase instead of name) +LAKEBASE_INSTANCE_NAME = resolve_lakebase_instance_name(_LAKEBASE_INSTANCE_NAME_RAW) def get_session_id(request: ResponsesAgentRequest) -> str: diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py b/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py index 03e52e9e..1877206d 100644 --- a/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py @@ -8,6 +8,50 @@ from mlflow.types.responses import ResponsesAgentStreamEvent +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), + resolve it to the actual instance name by looking up database instances. + """ + 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: + if hostname in (instance.read_write_dns, instance.read_only_dns): + if not instance.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 '{instance.name}'") + return instance.name + + raise ValueError( + f"Unable to find database instance matching hostname '{hostname}'. " + "Ensure the hostname is correct and the instance exists." + ) + + def get_databricks_host_from_env() -> Optional[str]: try: w = WorkspaceClient() diff --git a/agent-openai-agents-sdk-stateful-memory/app.yaml b/agent-openai-agents-sdk-stateful-memory/app.yaml index 34465373..f5bc6567 100644 --- a/agent-openai-agents-sdk-stateful-memory/app.yaml +++ b/agent-openai-agents-sdk-stateful-memory/app.yaml @@ -14,3 +14,5 @@ env: value: "300" - name: MLFLOW_EXPERIMENT_ID valueFrom: "experiment" + - name: LAKEBASE_INSTANCE_NAME + valueFrom: "database" diff --git a/agent-openai-agents-sdk-stateful-memory/databricks.yml b/agent-openai-agents-sdk-stateful-memory/databricks.yml index 48e2794f..00e3520f 100644 --- a/agent-openai-agents-sdk-stateful-memory/databricks.yml +++ b/agent-openai-agents-sdk-stateful-memory/databricks.yml @@ -1,23 +1,23 @@ bundle: - name: agent_openai_agents_sdk + name: agent_openai_agents_sdk_stateful_memory resources: # MLflow experiment for agent tracing - automatically created by bundle experiments: - agent_openai_agents_sdk_experiment: + agent_openai_agents_sdk_stateful_memory_experiment: name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} apps: - agent_openai_agents_sdk: - name: "${bundle.target}-agent-openai-agents-sdk" - description: "OpenAI Agents SDK agent application" + agent_openai_agents_sdk_stateful_memory: + name: "${bundle.target}-agent-openai-agents-sdk-stateful-memory" + description: "OpenAI Agents SDK agent application with stateful memory" source_code_path: ./ # Resources which this app has access to resources: - name: 'experiment' experiment: - experiment_id: "${resources.experiments.agent_openai_agents_sdk_experiment.id}" + experiment_id: "${resources.experiments.agent_openai_agents_sdk_stateful_memory_experiment.id}" permission: 'CAN_MANAGE' targets: @@ -33,5 +33,5 @@ targets: # host: https://... resources: apps: - agent_openai_agents_sdk: - name: agent-openai-agents-sdk + agent_openai_agents_sdk_stateful_memory: + name: agent-openai-agents-sdk-stateful-memory From ebbeb19e62c836fbf9d6777be204dcbe22166c32 Mon Sep 17 00:00:00 2001 From: Jenny Date: Wed, 11 Feb 2026 13:48:28 -0800 Subject: [PATCH 07/14] add openai agent memory skill --- .claude/skills/agent-openai-memory/SKILL.md | 175 ++++++++++++++++++ .gitignore | 17 ++ .../skills/agent-openai-memory/SKILL.md | 175 ++++++++++++++++++ .../.gitignore | 1 + 4 files changed, 368 insertions(+) create mode 100644 .claude/skills/agent-openai-memory/SKILL.md create mode 100644 agent-openai-agents-sdk-stateful-memory/.claude/skills/agent-openai-memory/SKILL.md diff --git a/.claude/skills/agent-openai-memory/SKILL.md b/.claude/skills/agent-openai-memory/SKILL.md new file mode 100644 index 00000000..b7ed11ad --- /dev/null +++ b/.claude/skills/agent-openai-memory/SKILL.md @@ -0,0 +1,175 @@ +--- +name: agent-openai-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." +--- + +# Stateful Memory with OpenAI Agents SDK Sessions + +This template uses OpenAI Agents SDK [Sessions](https://openai.github.io/openai-agents-python/sessions/) with `AsyncDatabricksSession` to persist conversation history to a Databricks Lakebase instance. + +## How Sessions Work + +Sessions automatically manage conversation history for multi-turn interactions: + +1. **Before each run**: The session retrieves prior conversation history and prepends it to input +2. **During the run**: New items (user messages, responses, tool calls) are generated +3. **After each run**: All new items are automatically stored in the session + +This eliminates the need to manually manage conversation state between runs. + +## Key Concepts + +| Concept | Description | +|---------|-------------| +| **Session** | Stores conversation history for a specific `session_id` | +| **`session_id`** | Unique identifier linking requests to the same conversation | +| **`AsyncDatabricksSession`** | Session implementation backed by Databricks Lakebase | +| **`LAKEBASE_INSTANCE_NAME`** | Environment variable specifying the Lakebase instance | + +## How This Template Uses Sessions + +### Session Creation (`agent_server/agent.py`) + +```python +from databricks_openai.agents import AsyncDatabricksSession + +session = AsyncDatabricksSession( + session_id=get_session_id(request), + instance_name=LAKEBASE_INSTANCE_NAME, +) + +result = await Runner.run(agent, messages, session=session) +``` + +### Session ID Extraction (`agent_server/agent.py`) + +The `session_id` is extracted from `custom_inputs` or auto-generated: + +```python +def get_session_id(request: ResponsesAgentRequest) -> str: + if hasattr(request, "custom_inputs") and request.custom_inputs: + if "session_id" in request.custom_inputs: + return request.custom_inputs["session_id"] + return str(uuid7()) +``` + +### Lakebase Instance Resolution (`agent_server/utils.py`) + +The `LAKEBASE_INSTANCE_NAME` env var can be either an instance name or a hostname. The `resolve_lakebase_instance_name()` function handles both cases: + +```python +_LAKEBASE_INSTANCE_NAME_RAW = os.environ.get("LAKEBASE_INSTANCE_NAME") +LAKEBASE_INSTANCE_NAME = resolve_lakebase_instance_name(_LAKEBASE_INSTANCE_NAME_RAW) +``` + +--- + +## Prerequisites + +1. **Dependency**: `databricks-openai[memory]` must be in `pyproject.toml` (already included) + +2. **Lakebase instance**: You need a Databricks Lakebase instance. See the **lakebase-setup** skill for creating and configuring one. + +3. **Environment variable**: Set `LAKEBASE_INSTANCE_NAME` in your `.env` file: + ```bash + LAKEBASE_INSTANCE_NAME= + ``` + +--- + +## Configuration Files + +### databricks.yml (Lakebase Resource) + +Add the Lakebase database resource to your app: + +```yaml +resources: + apps: + agent_openai_agents_sdk_stateful_memory: + name: "your-app-name" + source_code_path: ./ + + resources: + # ... other resources (experiment, etc.) ... + + # Lakebase instance for session storage + - name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' +``` + +### app.yaml (Environment Variables) + +The `LAKEBASE_INSTANCE_NAME` env var is resolved from the database resource at deploy time: + +```yaml +env: + - name: LAKEBASE_INSTANCE_NAME + valueFrom: "database" +``` + +### .env (Local Development) + +```bash +LAKEBASE_INSTANCE_NAME= +``` + +--- + +## Testing Sessions + +### Test Multi-Turn Conversation Locally + +```bash +# Start the server +uv run start-app + +# First message - starts a new session +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{"input": [{"role": "user", "content": "Hello, I live in SF!"}]}' + +# Note the session_id from custom_outputs in the response + +# Second message - continues the same session +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "What city did I say I live in?"}], + "custom_inputs": {"session_id": ""} + }' +``` + +### Test Streaming + +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Hello!"}], + "stream": true + }' +``` + +--- + +## Troubleshooting + +| Issue | Cause | Solution | +|-------|-------|----------| +| **"LAKEBASE_INSTANCE_NAME environment variable is required"** | Missing env var | Set `LAKEBASE_INSTANCE_NAME` in `.env` | +| **SSL connection closed unexpectedly** | Network/instance issue | Verify Lakebase instance is running: `databricks lakebase instances get ` | +| **Agent doesn't remember previous messages** | Different session_id | Pass the same `session_id` via `custom_inputs` across requests | +| **"Unable to resolve hostname"** | Hostname doesn't match any instance | Verify the hostname or use the instance name directly | +| **Permission denied** | Missing Lakebase access | Add `database` resource to `databricks.yml` with `CAN_CONNECT_AND_CREATE` | + +--- + +## Next Steps + +- Configure Lakebase: see **lakebase-setup** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/.gitignore b/.gitignore index 358665c9..60d6d671 100644 --- a/.gitignore +++ b/.gitignore @@ -167,3 +167,20 @@ 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/ +!.claude/skills/agent-openai-memory/ diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/agent-openai-memory/SKILL.md b/agent-openai-agents-sdk-stateful-memory/.claude/skills/agent-openai-memory/SKILL.md new file mode 100644 index 00000000..b7ed11ad --- /dev/null +++ b/agent-openai-agents-sdk-stateful-memory/.claude/skills/agent-openai-memory/SKILL.md @@ -0,0 +1,175 @@ +--- +name: agent-openai-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." +--- + +# Stateful Memory with OpenAI Agents SDK Sessions + +This template uses OpenAI Agents SDK [Sessions](https://openai.github.io/openai-agents-python/sessions/) with `AsyncDatabricksSession` to persist conversation history to a Databricks Lakebase instance. + +## How Sessions Work + +Sessions automatically manage conversation history for multi-turn interactions: + +1. **Before each run**: The session retrieves prior conversation history and prepends it to input +2. **During the run**: New items (user messages, responses, tool calls) are generated +3. **After each run**: All new items are automatically stored in the session + +This eliminates the need to manually manage conversation state between runs. + +## Key Concepts + +| Concept | Description | +|---------|-------------| +| **Session** | Stores conversation history for a specific `session_id` | +| **`session_id`** | Unique identifier linking requests to the same conversation | +| **`AsyncDatabricksSession`** | Session implementation backed by Databricks Lakebase | +| **`LAKEBASE_INSTANCE_NAME`** | Environment variable specifying the Lakebase instance | + +## How This Template Uses Sessions + +### Session Creation (`agent_server/agent.py`) + +```python +from databricks_openai.agents import AsyncDatabricksSession + +session = AsyncDatabricksSession( + session_id=get_session_id(request), + instance_name=LAKEBASE_INSTANCE_NAME, +) + +result = await Runner.run(agent, messages, session=session) +``` + +### Session ID Extraction (`agent_server/agent.py`) + +The `session_id` is extracted from `custom_inputs` or auto-generated: + +```python +def get_session_id(request: ResponsesAgentRequest) -> str: + if hasattr(request, "custom_inputs") and request.custom_inputs: + if "session_id" in request.custom_inputs: + return request.custom_inputs["session_id"] + return str(uuid7()) +``` + +### Lakebase Instance Resolution (`agent_server/utils.py`) + +The `LAKEBASE_INSTANCE_NAME` env var can be either an instance name or a hostname. The `resolve_lakebase_instance_name()` function handles both cases: + +```python +_LAKEBASE_INSTANCE_NAME_RAW = os.environ.get("LAKEBASE_INSTANCE_NAME") +LAKEBASE_INSTANCE_NAME = resolve_lakebase_instance_name(_LAKEBASE_INSTANCE_NAME_RAW) +``` + +--- + +## Prerequisites + +1. **Dependency**: `databricks-openai[memory]` must be in `pyproject.toml` (already included) + +2. **Lakebase instance**: You need a Databricks Lakebase instance. See the **lakebase-setup** skill for creating and configuring one. + +3. **Environment variable**: Set `LAKEBASE_INSTANCE_NAME` in your `.env` file: + ```bash + LAKEBASE_INSTANCE_NAME= + ``` + +--- + +## Configuration Files + +### databricks.yml (Lakebase Resource) + +Add the Lakebase database resource to your app: + +```yaml +resources: + apps: + agent_openai_agents_sdk_stateful_memory: + name: "your-app-name" + source_code_path: ./ + + resources: + # ... other resources (experiment, etc.) ... + + # Lakebase instance for session storage + - name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' +``` + +### app.yaml (Environment Variables) + +The `LAKEBASE_INSTANCE_NAME` env var is resolved from the database resource at deploy time: + +```yaml +env: + - name: LAKEBASE_INSTANCE_NAME + valueFrom: "database" +``` + +### .env (Local Development) + +```bash +LAKEBASE_INSTANCE_NAME= +``` + +--- + +## Testing Sessions + +### Test Multi-Turn Conversation Locally + +```bash +# Start the server +uv run start-app + +# First message - starts a new session +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{"input": [{"role": "user", "content": "Hello, I live in SF!"}]}' + +# Note the session_id from custom_outputs in the response + +# Second message - continues the same session +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "What city did I say I live in?"}], + "custom_inputs": {"session_id": ""} + }' +``` + +### Test Streaming + +```bash +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ + "input": [{"role": "user", "content": "Hello!"}], + "stream": true + }' +``` + +--- + +## Troubleshooting + +| Issue | Cause | Solution | +|-------|-------|----------| +| **"LAKEBASE_INSTANCE_NAME environment variable is required"** | Missing env var | Set `LAKEBASE_INSTANCE_NAME` in `.env` | +| **SSL connection closed unexpectedly** | Network/instance issue | Verify Lakebase instance is running: `databricks lakebase instances get ` | +| **Agent doesn't remember previous messages** | Different session_id | Pass the same `session_id` via `custom_inputs` across requests | +| **"Unable to resolve hostname"** | Hostname doesn't match any instance | Verify the hostname or use the instance name directly | +| **Permission denied** | Missing Lakebase access | Add `database` resource to `databricks.yml` with `CAN_CONNECT_AND_CREATE` | + +--- + +## Next Steps + +- Configure Lakebase: see **lakebase-setup** skill +- Test locally: see **run-locally** skill +- Deploy: see **deploy** skill diff --git a/agent-openai-agents-sdk-stateful-memory/.gitignore b/agent-openai-agents-sdk-stateful-memory/.gitignore index 42eec123..875a7549 100644 --- a/agent-openai-agents-sdk-stateful-memory/.gitignore +++ b/agent-openai-agents-sdk-stateful-memory/.gitignore @@ -214,6 +214,7 @@ sketch !.claude/skills/add-tools/ !.claude/skills/run-locally/ !.claude/skills/modify-agent/ +!.claude/skills/agent-openai-memory/ **/.env **/.env.local From 283e1a03ea97c7eb06e95a7fe3544df69e8ad5b7 Mon Sep 17 00:00:00 2001 From: Jenny Date: Wed, 11 Feb 2026 14:16:41 -0800 Subject: [PATCH 08/14] add to openai templates sync script --- .claude/sync-skills.py | 110 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 110 insertions(+) create mode 100755 .claude/sync-skills.py diff --git a/.claude/sync-skills.py b/.claude/sync-skills.py new file mode 100755 index 00000000..db8b85de --- /dev/null +++ b/.claude/sync-skills.py @@ -0,0 +1,110 @@ +#!/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-openai-agents-sdk-stateful-memory": { + "sdk": "openai", + "bundle_name": "agent_openai_agents_sdk_stateful_memory", + }, + "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 - SDK-specific, 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") + elif sdk == "openai": + copy_skill(SOURCE / "lakebase-setup", dest / "lakebase-setup") + copy_skill(SOURCE / "agent-openai-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() From ac6c7590ffd0464c56e47a843c236bae414e3456 Mon Sep 17 00:00:00 2001 From: Jenny Date: Wed, 11 Feb 2026 14:34:25 -0800 Subject: [PATCH 09/14] run python sync skills --- .../skills/agent-openai-memory/SKILL.md | 175 ------------------ .../.claude/skills/deploy/SKILL.md | 16 +- 2 files changed, 8 insertions(+), 183 deletions(-) delete mode 100644 agent-openai-agents-sdk-stateful-memory/.claude/skills/agent-openai-memory/SKILL.md diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/agent-openai-memory/SKILL.md b/agent-openai-agents-sdk-stateful-memory/.claude/skills/agent-openai-memory/SKILL.md deleted file mode 100644 index b7ed11ad..00000000 --- a/agent-openai-agents-sdk-stateful-memory/.claude/skills/agent-openai-memory/SKILL.md +++ /dev/null @@ -1,175 +0,0 @@ ---- -name: agent-openai-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." ---- - -# Stateful Memory with OpenAI Agents SDK Sessions - -This template uses OpenAI Agents SDK [Sessions](https://openai.github.io/openai-agents-python/sessions/) with `AsyncDatabricksSession` to persist conversation history to a Databricks Lakebase instance. - -## How Sessions Work - -Sessions automatically manage conversation history for multi-turn interactions: - -1. **Before each run**: The session retrieves prior conversation history and prepends it to input -2. **During the run**: New items (user messages, responses, tool calls) are generated -3. **After each run**: All new items are automatically stored in the session - -This eliminates the need to manually manage conversation state between runs. - -## Key Concepts - -| Concept | Description | -|---------|-------------| -| **Session** | Stores conversation history for a specific `session_id` | -| **`session_id`** | Unique identifier linking requests to the same conversation | -| **`AsyncDatabricksSession`** | Session implementation backed by Databricks Lakebase | -| **`LAKEBASE_INSTANCE_NAME`** | Environment variable specifying the Lakebase instance | - -## How This Template Uses Sessions - -### Session Creation (`agent_server/agent.py`) - -```python -from databricks_openai.agents import AsyncDatabricksSession - -session = AsyncDatabricksSession( - session_id=get_session_id(request), - instance_name=LAKEBASE_INSTANCE_NAME, -) - -result = await Runner.run(agent, messages, session=session) -``` - -### Session ID Extraction (`agent_server/agent.py`) - -The `session_id` is extracted from `custom_inputs` or auto-generated: - -```python -def get_session_id(request: ResponsesAgentRequest) -> str: - if hasattr(request, "custom_inputs") and request.custom_inputs: - if "session_id" in request.custom_inputs: - return request.custom_inputs["session_id"] - return str(uuid7()) -``` - -### Lakebase Instance Resolution (`agent_server/utils.py`) - -The `LAKEBASE_INSTANCE_NAME` env var can be either an instance name or a hostname. The `resolve_lakebase_instance_name()` function handles both cases: - -```python -_LAKEBASE_INSTANCE_NAME_RAW = os.environ.get("LAKEBASE_INSTANCE_NAME") -LAKEBASE_INSTANCE_NAME = resolve_lakebase_instance_name(_LAKEBASE_INSTANCE_NAME_RAW) -``` - ---- - -## Prerequisites - -1. **Dependency**: `databricks-openai[memory]` must be in `pyproject.toml` (already included) - -2. **Lakebase instance**: You need a Databricks Lakebase instance. See the **lakebase-setup** skill for creating and configuring one. - -3. **Environment variable**: Set `LAKEBASE_INSTANCE_NAME` in your `.env` file: - ```bash - LAKEBASE_INSTANCE_NAME= - ``` - ---- - -## Configuration Files - -### databricks.yml (Lakebase Resource) - -Add the Lakebase database resource to your app: - -```yaml -resources: - apps: - agent_openai_agents_sdk_stateful_memory: - name: "your-app-name" - source_code_path: ./ - - resources: - # ... other resources (experiment, etc.) ... - - # Lakebase instance for session storage - - name: 'database' - database: - instance_name: '' - database_name: 'postgres' - permission: 'CAN_CONNECT_AND_CREATE' -``` - -### app.yaml (Environment Variables) - -The `LAKEBASE_INSTANCE_NAME` env var is resolved from the database resource at deploy time: - -```yaml -env: - - name: LAKEBASE_INSTANCE_NAME - valueFrom: "database" -``` - -### .env (Local Development) - -```bash -LAKEBASE_INSTANCE_NAME= -``` - ---- - -## Testing Sessions - -### Test Multi-Turn Conversation Locally - -```bash -# Start the server -uv run start-app - -# First message - starts a new session -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{"input": [{"role": "user", "content": "Hello, I live in SF!"}]}' - -# Note the session_id from custom_outputs in the response - -# Second message - continues the same session -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ - "input": [{"role": "user", "content": "What city did I say I live in?"}], - "custom_inputs": {"session_id": ""} - }' -``` - -### Test Streaming - -```bash -curl -X POST http://localhost:8000/invocations \ - -H "Content-Type: application/json" \ - -d '{ - "input": [{"role": "user", "content": "Hello!"}], - "stream": true - }' -``` - ---- - -## Troubleshooting - -| Issue | Cause | Solution | -|-------|-------|----------| -| **"LAKEBASE_INSTANCE_NAME environment variable is required"** | Missing env var | Set `LAKEBASE_INSTANCE_NAME` in `.env` | -| **SSL connection closed unexpectedly** | Network/instance issue | Verify Lakebase instance is running: `databricks lakebase instances get ` | -| **Agent doesn't remember previous messages** | Different session_id | Pass the same `session_id` via `custom_inputs` across requests | -| **"Unable to resolve hostname"** | Hostname doesn't match any instance | Verify the hostname or use the instance name directly | -| **Permission denied** | Missing Lakebase access | Add `database` resource to `databricks.yml` with `CAN_CONNECT_AND_CREATE` | - ---- - -## Next Steps - -- Configure Lakebase: see **lakebase-setup** skill -- Test locally: see **run-locally** skill -- Deploy: see **deploy** skill diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/deploy/SKILL.md b/agent-openai-agents-sdk-stateful-memory/.claude/skills/deploy/SKILL.md index d072e595..99736e6d 100644 --- a/agent-openai-agents-sdk-stateful-memory/.claude/skills/deploy/SKILL.md +++ b/agent-openai-agents-sdk-stateful-memory/.claude/skills/deploy/SKILL.md @@ -16,7 +16,7 @@ Update the app name in `databricks.yml`: ```yaml resources: apps: - agent_openai_agents_sdk: + agent_openai_agents_sdk_stateful_memory: name: "agent-your-app-name" # Use agent-* prefix ``` @@ -32,12 +32,12 @@ databricks bundle validate databricks bundle deploy # 3. Run the app (starts/restarts with uploaded source code) - REQUIRED! -databricks bundle run agent_openai_agents_sdk +databricks bundle run agent_openai_agents_sdk_stateful_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_openai_agents_sdk` matches the app name in `databricks.yml` under `resources.apps`. +The resource key `agent_openai_agents_sdk_stateful_memory` matches the app name in `databricks.yml` under `resources.apps`. ## Handling "App Already Exists" Error @@ -61,7 +61,7 @@ databricks apps get --output json | jq '{name, budget_policy ```yaml resources: apps: - agent_openai_agents_sdk: + agent_openai_agents_sdk_stateful_memory: name: "existing-app-name" # Must match exactly budget_policy_id: "xxx-xxx-xxx" # Copy from step 1 if present ``` @@ -86,7 +86,7 @@ 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_openai_agents_sdk --auto-approve +databricks bundle deployment bind agent_openai_agents_sdk_stateful_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. @@ -94,7 +94,7 @@ databricks bundle deployment bind agent_openai_agents_sdk -- **Step 5:** Deploy: ```bash databricks bundle deploy -databricks bundle run agent_openai_agents_sdk +databricks bundle run agent_openai_agents_sdk_stateful_memory ``` ### Option 2: Delete and Recreate @@ -111,7 +111,7 @@ databricks bundle deploy To remove the link between bundle and deployed app: ```bash -databricks bundle deployment unbind agent_openai_agents_sdk +databricks bundle deployment unbind agent_openai_agents_sdk_stateful_memory ``` Use when: @@ -218,5 +218,5 @@ uv add | 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 | +| App running old code after deploy | Run `databricks bundle run agent_openai_agents_sdk_stateful_memory` after deploy | | Env var is None in deployed app | Check `valueFrom` in app.yaml matches resource `name` in databricks.yml | From dfcb42f4ca5aa252a4fa8b2f875cc693ca2a25e6 Mon Sep 17 00:00:00 2001 From: Jenny Date: Thu, 12 Feb 2026 15:39:09 -0800 Subject: [PATCH 10/14] databricks yml and use chatcontext convo id --- .../agent_server/agent.py | 20 +++++++++++++------ .../databricks.yml | 6 ++++++ 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py index fdbd7b86..07e994fb 100644 --- a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py @@ -36,11 +36,19 @@ def get_session_id(request: ResponsesAgentRequest) -> str: """Extract session_id from request or generate a new one.""" - # Try to get session_id from custom_inputs if provided - if hasattr(request, "custom_inputs") and request.custom_inputs: - if "session_id" in request.custom_inputs: - return request.custom_inputs["session_id"] - # Fall back to generating a new session_id + # Priority: + # 1. Use session_id from custom_inputs + # 2. Use conversation_id from ChatContext + # https://mlflow.org/docs/latest/api_reference/python_api/mlflow.types.html#mlflow.types.agent.ChatContext + # 3. Generate a new UUID + ci = dict(request.custom_inputs or {}) + + if "session_id" in ci and ci["session_id"]: + return str(ci["session_id"]) + + if request.context and getattr(request.context, "conversation_id", None): + return str(request.context.conversation_id) + return str(uuid7()) # NOTE: this will work for all databricks models OTHER than GPT-OSS, which uses a slightly different API @@ -71,7 +79,7 @@ async def invoke_handler(request: ResponsesAgentRequest) -> ResponsesAgentRespon # Optionally use the user's workspace client for on-behalf-of authentication # user_workspace_client = get_user_workspace_client() - # Create session for persistent conversation history with your Databricks Lakebase instance + # Create session for stateful, short-term conversation history with your Databricks Lakebase instance session = AsyncDatabricksSession( session_id=get_session_id(request), instance_name=LAKEBASE_INSTANCE_NAME, diff --git a/agent-openai-agents-sdk-stateful-memory/databricks.yml b/agent-openai-agents-sdk-stateful-memory/databricks.yml index 00e3520f..f76ec74c 100644 --- a/agent-openai-agents-sdk-stateful-memory/databricks.yml +++ b/agent-openai-agents-sdk-stateful-memory/databricks.yml @@ -19,6 +19,12 @@ resources: experiment: experiment_id: "${resources.experiments.agent_openai_agents_sdk_stateful_memory_experiment.id}" permission: 'CAN_MANAGE' + # Lakebase instance for persistent session storage + - name: 'database' + database: + instance_name: '' + database_name: 'postgres' + permission: 'CAN_CONNECT_AND_CREATE' targets: dev: From b33b62807ef3a8bb64876d0b15ff8027aa813c3a Mon Sep 17 00:00:00 2001 From: Jenny Date: Thu, 12 Feb 2026 16:01:51 -0800 Subject: [PATCH 11/14] sanitize mcp tool output items https://github.com/databricks/app-templates/pull/119/changes --- .../agent_server/agent.py | 3 ++- .../agent_server/utils.py | 23 ++++++++++++++++++- 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py index 07e994fb..9a045529 100644 --- a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py @@ -20,6 +20,7 @@ get_user_workspace_client, process_agent_stream_events, resolve_lakebase_instance_name, + sanitize_output_items, ) # Lakebase instance name for persistent session storage @@ -90,7 +91,7 @@ async def invoke_handler(request: ResponsesAgentRequest) -> ResponsesAgentRespon messages = [i.model_dump() for i in request.input] result = await Runner.run(agent, messages, session=session) return ResponsesAgentResponse( - output=[item.to_input_item() for item in result.new_items], + output=sanitize_output_items(result.new_items), custom_outputs={"session_id": session.session_id}, ) diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py b/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py index 1877206d..3ad82acf 100644 --- a/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py @@ -1,3 +1,4 @@ +import json import logging from typing import AsyncGenerator, AsyncIterator, Optional from uuid import uuid4 @@ -66,6 +67,26 @@ def get_user_workspace_client() -> WorkspaceClient: return WorkspaceClient(token=token, auth_type="pat") +def _sanitize_item(input_item: dict) -> dict: + """Sanitize a single output item dict for Pydantic validation. + + MCP tool calls (e.g. Genie) can return items where the `output` field is + a *list* of content objects instead of a plain string. MLflow's Pydantic + models expect `output` to be a string, so this serialises any non-string + values to JSON. + + TODO: Remove once https://github.com/mlflow/mlflow/pull/20777 is released. + """ + if isinstance(input_item.get("output"), list): + input_item["output"] = json.dumps(input_item["output"]) + return input_item + + +def sanitize_output_items(items) -> list[dict]: + """Convert agent output items to dicts safe for ResponsesAgentResponse.""" + return [_sanitize_item(item.to_input_item()) for item in items] + + async def process_agent_stream_events( async_stream: AsyncIterator[StreamEvent], ) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: @@ -84,5 +105,5 @@ async def process_agent_stream_events( elif event.type == "run_item_stream_event" and event.item.type == "tool_call_output_item": yield ResponsesAgentStreamEvent( type="response.output_item.done", - item=event.item.to_input_item(), + item=_sanitize_item(event.item.to_input_item()), ) From 4e3f78723814dc7e550670e6013f954e730708bc Mon Sep 17 00:00:00 2001 From: Jenny Date: Thu, 12 Feb 2026 17:28:42 -0800 Subject: [PATCH 12/14] deduplicate input logic --- .../agent_server/agent.py | 5 +++-- .../agent_server/utils.py | 18 +++++++++++++++++- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py index 9a045529..65325529 100644 --- a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py @@ -16,6 +16,7 @@ ) from agent_server.utils import ( + deduplicate_input, get_databricks_host_from_env, get_user_workspace_client, process_agent_stream_events, @@ -88,7 +89,7 @@ async def invoke_handler(request: ResponsesAgentRequest) -> ResponsesAgentRespon async with await init_mcp_server() as mcp_server: agent = create_coding_agent(mcp_server) - messages = [i.model_dump() for i in request.input] + messages = await deduplicate_input(request, session) result = await Runner.run(agent, messages, session=session) return ResponsesAgentResponse( output=sanitize_output_items(result.new_items), @@ -109,7 +110,7 @@ async def stream_handler(request: ResponsesAgentRequest) -> AsyncGenerator[Respo async with await init_mcp_server() as mcp_server: agent = create_coding_agent(mcp_server) - messages = [i.model_dump() for i in request.input] + messages = await deduplicate_input(request, session) result = Runner.run_streamed(agent, input=messages, session=session) async for event in process_agent_stream_events(result.stream_events()): diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py b/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py index 3ad82acf..211c9410 100644 --- a/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py @@ -5,8 +5,9 @@ from agents.result import StreamEvent from databricks.sdk import WorkspaceClient +from databricks_openai.agents import AsyncDatabricksSession from mlflow.genai.agent_server import get_request_headers -from mlflow.types.responses import ResponsesAgentStreamEvent +from mlflow.types.responses import ResponsesAgentRequest, ResponsesAgentStreamEvent def _is_lakebase_hostname(value: str) -> bool: @@ -67,6 +68,21 @@ def get_user_workspace_client() -> WorkspaceClient: return WorkspaceClient(token=token, auth_type="pat") +async def deduplicate_input(request: ResponsesAgentRequest, session: AsyncDatabricksSession) -> list[dict]: + """Return the input messages to pass to the Runner, avoiding duplication with session history. + + When a client sends the full conversation history AND the session already has + that history persisted, passing everything through would duplicate messages. + If the session already covers the prior turns, only the latest message is needed + since the session will prepend the full history automatically. + """ + messages = [i.model_dump() for i in request.input] + session_items = await session.get_items() + if len(session_items) >= len(messages) - 1: + return [messages[-1]] + return messages + + def _sanitize_item(input_item: dict) -> dict: """Sanitize a single output item dict for Pydantic validation. From 23be8cb5cf5d65b302f5c306745e8ba2be9abb3c Mon Sep 17 00:00:00 2001 From: Jenny Date: Thu, 12 Feb 2026 17:30:51 -0800 Subject: [PATCH 13/14] update sanitize mcp handler to be more defensive --- .../agent_server/agent.py | 2 +- .../agent_server/utils.py | 7 +++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py index 65325529..572be9bc 100644 --- a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py @@ -102,7 +102,7 @@ async def stream_handler(request: ResponsesAgentRequest) -> AsyncGenerator[Respo # Optionally use the user's workspace client for on-behalf-of authentication # user_workspace_client = get_user_workspace_client() - # Create session for persistent conversation history with your Databricks Lakebase instance + # Create session for stateful, short-term conversation history with your Databricks Lakebase instance session = AsyncDatabricksSession( session_id=get_session_id(request), instance_name=LAKEBASE_INSTANCE_NAME, diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py b/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py index 211c9410..d140d649 100644 --- a/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py +++ b/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py @@ -93,8 +93,11 @@ def _sanitize_item(input_item: dict) -> dict: TODO: Remove once https://github.com/mlflow/mlflow/pull/20777 is released. """ - if isinstance(input_item.get("output"), list): - input_item["output"] = json.dumps(input_item["output"]) + if not isinstance(input_item.get("output"), str): + try: + input_item["output"] = json.dumps(input_item.get("output")) + except (TypeError, ValueError): + input_item["output"] = str(input_item.get("output")) return input_item From 438a38e84cb7fb9e7287429e4843bd7a88f041b7 Mon Sep 17 00:00:00 2001 From: Jenny Date: Thu, 12 Feb 2026 17:36:00 -0800 Subject: [PATCH 14/14] rename from agent-openai-agents-sdk-stateful-memory to agent-openai-agents-sdk-short-term-memory --- .claude/skills/agent-openai-memory/SKILL.md | 2 +- .claude/sync-skills.py | 4 ++-- .../.claude/skills/add-tools/SKILL.md | 0 .../add-tools/examples/custom-mcp-server.md | 0 .../skills/add-tools/examples/experiment.yaml | 0 .../skills/add-tools/examples/genie-space.yaml | 0 .../add-tools/examples/serving-endpoint.yaml | 0 .../skills/add-tools/examples/sql-warehouse.yaml | 0 .../skills/add-tools/examples/uc-connection.yaml | 0 .../skills/add-tools/examples/uc-function.yaml | 0 .../skills/add-tools/examples/vector-search.yaml | 0 .../.claude/skills/deploy/SKILL.md | 16 ++++++++-------- .../.claude/skills/discover-tools/SKILL.md | 0 .../.claude/skills/modify-agent/SKILL.md | 0 .../.claude/skills/quickstart/SKILL.md | 0 .../.claude/skills/run-locally/SKILL.md | 0 .../.env.example | 0 .../.gitignore | 0 .../AGENTS.md | 0 .../CLAUDE.md | 0 .../README.md | 0 .../agent_server/__init__.py | 0 .../agent_server/agent.py | 0 .../agent_server/evaluate_agent.py | 0 .../agent_server/start_server.py | 0 .../agent_server/utils.py | 0 .../app.yaml | 0 .../databricks.yml | 16 ++++++++-------- .../pyproject.toml | 2 +- .../requirements.txt | 0 .../scripts/__init__.py | 0 .../scripts/discover_tools.py | 0 .../scripts/quickstart.py | 0 .../scripts/start_app.py | 0 34 files changed, 20 insertions(+), 20 deletions(-) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/.claude/skills/add-tools/SKILL.md (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/.claude/skills/add-tools/examples/custom-mcp-server.md (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/.claude/skills/add-tools/examples/experiment.yaml (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/.claude/skills/add-tools/examples/genie-space.yaml (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/.claude/skills/add-tools/examples/serving-endpoint.yaml (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/.claude/skills/add-tools/examples/sql-warehouse.yaml (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/.claude/skills/add-tools/examples/uc-connection.yaml (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/.claude/skills/add-tools/examples/uc-function.yaml (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/.claude/skills/add-tools/examples/vector-search.yaml (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/.claude/skills/deploy/SKILL.md (92%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/.claude/skills/discover-tools/SKILL.md (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/.claude/skills/modify-agent/SKILL.md (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/.claude/skills/quickstart/SKILL.md (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/.claude/skills/run-locally/SKILL.md (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/.env.example (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/.gitignore (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/AGENTS.md (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/CLAUDE.md (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/README.md (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/agent_server/__init__.py (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/agent_server/agent.py (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/agent_server/evaluate_agent.py (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/agent_server/start_server.py (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/agent_server/utils.py (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/app.yaml (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/databricks.yml (66%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/pyproject.toml (96%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/requirements.txt (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/scripts/__init__.py (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/scripts/discover_tools.py (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/scripts/quickstart.py (100%) rename {agent-openai-agents-sdk-stateful-memory => agent-openai-agents-sdk-short-term-memory}/scripts/start_app.py (100%) diff --git a/.claude/skills/agent-openai-memory/SKILL.md b/.claude/skills/agent-openai-memory/SKILL.md index b7ed11ad..e648879b 100644 --- a/.claude/skills/agent-openai-memory/SKILL.md +++ b/.claude/skills/agent-openai-memory/SKILL.md @@ -86,7 +86,7 @@ Add the Lakebase database resource to your app: ```yaml resources: apps: - agent_openai_agents_sdk_stateful_memory: + agent_openai_agents_sdk_short_term_memory: name: "your-app-name" source_code_path: ./ diff --git a/.claude/sync-skills.py b/.claude/sync-skills.py index 4db449a6..bea8ae67 100755 --- a/.claude/sync-skills.py +++ b/.claude/sync-skills.py @@ -34,9 +34,9 @@ "sdk": "openai", "bundle_name": "agent_openai_agents_sdk", }, - "agent-openai-agents-sdk-stateful-memory": { + "agent-openai-agents-sdk-short-term-memory": { "sdk": "openai", - "bundle_name": "agent_openai_agents_sdk_stateful_memory", + "bundle_name": "agent_openai_agents_sdk_short_term_memory", }, "agent-non-conversational": { "sdk": "langgraph", diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/SKILL.md b/agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/SKILL.md similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/SKILL.md rename to agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/SKILL.md diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/custom-mcp-server.md b/agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/custom-mcp-server.md rename to agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/examples/custom-mcp-server.md diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/experiment.yaml b/agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/examples/experiment.yaml similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/experiment.yaml rename to agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/examples/experiment.yaml diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/genie-space.yaml b/agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/examples/genie-space.yaml similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/genie-space.yaml rename to agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/examples/genie-space.yaml diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/serving-endpoint.yaml b/agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/examples/serving-endpoint.yaml similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/serving-endpoint.yaml rename to agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/examples/serving-endpoint.yaml diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/sql-warehouse.yaml b/agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/examples/sql-warehouse.yaml similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/sql-warehouse.yaml rename to agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/examples/sql-warehouse.yaml diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/uc-connection.yaml b/agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/examples/uc-connection.yaml similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/uc-connection.yaml rename to agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/examples/uc-connection.yaml diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/uc-function.yaml b/agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/examples/uc-function.yaml similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/uc-function.yaml rename to agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/examples/uc-function.yaml diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/vector-search.yaml b/agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/examples/vector-search.yaml similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/.claude/skills/add-tools/examples/vector-search.yaml rename to agent-openai-agents-sdk-short-term-memory/.claude/skills/add-tools/examples/vector-search.yaml diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/deploy/SKILL.md b/agent-openai-agents-sdk-short-term-memory/.claude/skills/deploy/SKILL.md similarity index 92% rename from agent-openai-agents-sdk-stateful-memory/.claude/skills/deploy/SKILL.md rename to agent-openai-agents-sdk-short-term-memory/.claude/skills/deploy/SKILL.md index 99736e6d..215205af 100644 --- a/agent-openai-agents-sdk-stateful-memory/.claude/skills/deploy/SKILL.md +++ b/agent-openai-agents-sdk-short-term-memory/.claude/skills/deploy/SKILL.md @@ -16,7 +16,7 @@ Update the app name in `databricks.yml`: ```yaml resources: apps: - agent_openai_agents_sdk_stateful_memory: + agent_openai_agents_sdk_short_term_memory: name: "agent-your-app-name" # Use agent-* prefix ``` @@ -32,12 +32,12 @@ databricks bundle validate databricks bundle deploy # 3. Run the app (starts/restarts with uploaded source code) - REQUIRED! -databricks bundle run agent_openai_agents_sdk_stateful_memory +databricks bundle run agent_openai_agents_sdk_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_openai_agents_sdk_stateful_memory` matches the app name in `databricks.yml` under `resources.apps`. +The resource key `agent_openai_agents_sdk_short_term_memory` matches the app name in `databricks.yml` under `resources.apps`. ## Handling "App Already Exists" Error @@ -61,7 +61,7 @@ databricks apps get --output json | jq '{name, budget_policy ```yaml resources: apps: - agent_openai_agents_sdk_stateful_memory: + agent_openai_agents_sdk_short_term_memory: name: "existing-app-name" # Must match exactly budget_policy_id: "xxx-xxx-xxx" # Copy from step 1 if present ``` @@ -86,7 +86,7 @@ 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_openai_agents_sdk_stateful_memory --auto-approve +databricks bundle deployment bind agent_openai_agents_sdk_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. @@ -94,7 +94,7 @@ databricks bundle deployment bind agent_openai_agents_sdk_stateful_memory | 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_stateful_memory` after deploy | +| App running old code after deploy | Run `databricks bundle run agent_openai_agents_sdk_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-openai-agents-sdk-stateful-memory/.claude/skills/discover-tools/SKILL.md b/agent-openai-agents-sdk-short-term-memory/.claude/skills/discover-tools/SKILL.md similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/.claude/skills/discover-tools/SKILL.md rename to agent-openai-agents-sdk-short-term-memory/.claude/skills/discover-tools/SKILL.md diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/modify-agent/SKILL.md b/agent-openai-agents-sdk-short-term-memory/.claude/skills/modify-agent/SKILL.md similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/.claude/skills/modify-agent/SKILL.md rename to agent-openai-agents-sdk-short-term-memory/.claude/skills/modify-agent/SKILL.md diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/quickstart/SKILL.md b/agent-openai-agents-sdk-short-term-memory/.claude/skills/quickstart/SKILL.md similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/.claude/skills/quickstart/SKILL.md rename to agent-openai-agents-sdk-short-term-memory/.claude/skills/quickstart/SKILL.md diff --git a/agent-openai-agents-sdk-stateful-memory/.claude/skills/run-locally/SKILL.md b/agent-openai-agents-sdk-short-term-memory/.claude/skills/run-locally/SKILL.md similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/.claude/skills/run-locally/SKILL.md rename to agent-openai-agents-sdk-short-term-memory/.claude/skills/run-locally/SKILL.md diff --git a/agent-openai-agents-sdk-stateful-memory/.env.example b/agent-openai-agents-sdk-short-term-memory/.env.example similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/.env.example rename to agent-openai-agents-sdk-short-term-memory/.env.example diff --git a/agent-openai-agents-sdk-stateful-memory/.gitignore b/agent-openai-agents-sdk-short-term-memory/.gitignore similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/.gitignore rename to agent-openai-agents-sdk-short-term-memory/.gitignore diff --git a/agent-openai-agents-sdk-stateful-memory/AGENTS.md b/agent-openai-agents-sdk-short-term-memory/AGENTS.md similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/AGENTS.md rename to agent-openai-agents-sdk-short-term-memory/AGENTS.md diff --git a/agent-openai-agents-sdk-stateful-memory/CLAUDE.md b/agent-openai-agents-sdk-short-term-memory/CLAUDE.md similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/CLAUDE.md rename to agent-openai-agents-sdk-short-term-memory/CLAUDE.md diff --git a/agent-openai-agents-sdk-stateful-memory/README.md b/agent-openai-agents-sdk-short-term-memory/README.md similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/README.md rename to agent-openai-agents-sdk-short-term-memory/README.md diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/__init__.py b/agent-openai-agents-sdk-short-term-memory/agent_server/__init__.py similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/agent_server/__init__.py rename to agent-openai-agents-sdk-short-term-memory/agent_server/__init__.py diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/agent.py b/agent-openai-agents-sdk-short-term-memory/agent_server/agent.py similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/agent_server/agent.py rename to agent-openai-agents-sdk-short-term-memory/agent_server/agent.py diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/evaluate_agent.py b/agent-openai-agents-sdk-short-term-memory/agent_server/evaluate_agent.py similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/agent_server/evaluate_agent.py rename to agent-openai-agents-sdk-short-term-memory/agent_server/evaluate_agent.py diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/start_server.py b/agent-openai-agents-sdk-short-term-memory/agent_server/start_server.py similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/agent_server/start_server.py rename to agent-openai-agents-sdk-short-term-memory/agent_server/start_server.py diff --git a/agent-openai-agents-sdk-stateful-memory/agent_server/utils.py b/agent-openai-agents-sdk-short-term-memory/agent_server/utils.py similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/agent_server/utils.py rename to agent-openai-agents-sdk-short-term-memory/agent_server/utils.py diff --git a/agent-openai-agents-sdk-stateful-memory/app.yaml b/agent-openai-agents-sdk-short-term-memory/app.yaml similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/app.yaml rename to agent-openai-agents-sdk-short-term-memory/app.yaml diff --git a/agent-openai-agents-sdk-stateful-memory/databricks.yml b/agent-openai-agents-sdk-short-term-memory/databricks.yml similarity index 66% rename from agent-openai-agents-sdk-stateful-memory/databricks.yml rename to agent-openai-agents-sdk-short-term-memory/databricks.yml index f76ec74c..fb17fa35 100644 --- a/agent-openai-agents-sdk-stateful-memory/databricks.yml +++ b/agent-openai-agents-sdk-short-term-memory/databricks.yml @@ -1,23 +1,23 @@ bundle: - name: agent_openai_agents_sdk_stateful_memory + name: agent_openai_agents_sdk_short_term_memory resources: # MLflow experiment for agent tracing - automatically created by bundle experiments: - agent_openai_agents_sdk_stateful_memory_experiment: + agent_openai_agents_sdk_short_term_memory_experiment: name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target} apps: - agent_openai_agents_sdk_stateful_memory: - name: "${bundle.target}-agent-openai-agents-sdk-stateful-memory" - description: "OpenAI Agents SDK agent application with stateful memory" + agent_openai_agents_sdk_short_term_memory: + name: "${bundle.target}-agent-openai-agents-sdk-short-term-memory" + description: "OpenAI Agents SDK agent application with short-term memory" source_code_path: ./ # Resources which this app has access to resources: - name: 'experiment' experiment: - experiment_id: "${resources.experiments.agent_openai_agents_sdk_stateful_memory_experiment.id}" + experiment_id: "${resources.experiments.agent_openai_agents_sdk_short_term_memory_experiment.id}" permission: 'CAN_MANAGE' # Lakebase instance for persistent session storage - name: 'database' @@ -39,5 +39,5 @@ targets: # host: https://... resources: apps: - agent_openai_agents_sdk_stateful_memory: - name: agent-openai-agents-sdk-stateful-memory + agent_openai_agents_sdk_short_term_memory: + name: agent-openai-agents-sdk-short-term-memory diff --git a/agent-openai-agents-sdk-stateful-memory/pyproject.toml b/agent-openai-agents-sdk-short-term-memory/pyproject.toml similarity index 96% rename from agent-openai-agents-sdk-stateful-memory/pyproject.toml rename to agent-openai-agents-sdk-short-term-memory/pyproject.toml index 53c93ccb..dbb118f0 100644 --- a/agent-openai-agents-sdk-stateful-memory/pyproject.toml +++ b/agent-openai-agents-sdk-short-term-memory/pyproject.toml @@ -1,5 +1,5 @@ [project] -name = "agent-server-stateful-memory" +name = "agent-server-short-term-memory" version = "0.1.0" description = "MLflow-compatible agent server with FastAPI and persistent memory using Lakebase" readme = "README.md" diff --git a/agent-openai-agents-sdk-stateful-memory/requirements.txt b/agent-openai-agents-sdk-short-term-memory/requirements.txt similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/requirements.txt rename to agent-openai-agents-sdk-short-term-memory/requirements.txt diff --git a/agent-openai-agents-sdk-stateful-memory/scripts/__init__.py b/agent-openai-agents-sdk-short-term-memory/scripts/__init__.py similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/scripts/__init__.py rename to agent-openai-agents-sdk-short-term-memory/scripts/__init__.py diff --git a/agent-openai-agents-sdk-stateful-memory/scripts/discover_tools.py b/agent-openai-agents-sdk-short-term-memory/scripts/discover_tools.py similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/scripts/discover_tools.py rename to agent-openai-agents-sdk-short-term-memory/scripts/discover_tools.py diff --git a/agent-openai-agents-sdk-stateful-memory/scripts/quickstart.py b/agent-openai-agents-sdk-short-term-memory/scripts/quickstart.py similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/scripts/quickstart.py rename to agent-openai-agents-sdk-short-term-memory/scripts/quickstart.py diff --git a/agent-openai-agents-sdk-stateful-memory/scripts/start_app.py b/agent-openai-agents-sdk-short-term-memory/scripts/start_app.py similarity index 100% rename from agent-openai-agents-sdk-stateful-memory/scripts/start_app.py rename to agent-openai-agents-sdk-short-term-memory/scripts/start_app.py