From 6eb5bfee7f7855599801387d36235b79069c8132 Mon Sep 17 00:00:00 2001 From: Jenny Date: Tue, 20 Jan 2026 17:48:30 -0800 Subject: [PATCH 1/2] copy agent-openai-agents-sdk --- agent-openai-agents-sdk-stateful/.env.example | 17 + agent-openai-agents-sdk-stateful/.gitignore | 206 +++++++ agent-openai-agents-sdk-stateful/AGENTS.md | 530 ++++++++++++++++++ agent-openai-agents-sdk-stateful/CLAUDE.md | 3 + agent-openai-agents-sdk-stateful/README.md | 237 ++++++++ .../agent_server/__init__.py | 0 .../agent_server/agent.py | 65 +++ .../agent_server/evaluate_agent.py | 53 ++ .../agent_server/start_server.py | 17 + .../agent_server/utils.py | 44 ++ agent-openai-agents-sdk-stateful/app.yaml | 16 + .../databricks.yml | 37 ++ .../e2e-chatbot-app-next/databricks.yml | 63 +++ .../pyproject.toml | 37 ++ .../requirements.txt | 1 + .../scripts/__init__.py | 0 .../scripts/discover_tools.py | 432 ++++++++++++++ .../scripts/quickstart.sh | 392 +++++++++++++ .../scripts/start_app.py | 253 +++++++++ .../e2e-chatbot-app-next/databricks.yml | 63 +++ 20 files changed, 2466 insertions(+) create mode 100644 agent-openai-agents-sdk-stateful/.env.example create mode 100644 agent-openai-agents-sdk-stateful/.gitignore create mode 100644 agent-openai-agents-sdk-stateful/AGENTS.md create mode 100644 agent-openai-agents-sdk-stateful/CLAUDE.md create mode 100644 agent-openai-agents-sdk-stateful/README.md create mode 100644 agent-openai-agents-sdk-stateful/agent_server/__init__.py create mode 100644 agent-openai-agents-sdk-stateful/agent_server/agent.py create mode 100644 agent-openai-agents-sdk-stateful/agent_server/evaluate_agent.py create mode 100644 agent-openai-agents-sdk-stateful/agent_server/start_server.py create mode 100644 agent-openai-agents-sdk-stateful/agent_server/utils.py create mode 100644 agent-openai-agents-sdk-stateful/app.yaml create mode 100644 agent-openai-agents-sdk-stateful/databricks.yml create mode 100644 agent-openai-agents-sdk-stateful/e2e-chatbot-app-next/databricks.yml create mode 100644 agent-openai-agents-sdk-stateful/pyproject.toml create mode 100644 agent-openai-agents-sdk-stateful/requirements.txt create mode 100644 agent-openai-agents-sdk-stateful/scripts/__init__.py create mode 100755 agent-openai-agents-sdk-stateful/scripts/discover_tools.py create mode 100755 agent-openai-agents-sdk-stateful/scripts/quickstart.sh create mode 100644 agent-openai-agents-sdk-stateful/scripts/start_app.py create mode 100644 agent-openai-agents-sdk/e2e-chatbot-app-next/databricks.yml diff --git a/agent-openai-agents-sdk-stateful/.env.example b/agent-openai-agents-sdk-stateful/.env.example new file mode 100644 index 00000000..7c4ebf97 --- /dev/null +++ b/agent-openai-agents-sdk-stateful/.env.example @@ -0,0 +1,17 @@ +# Make a copy of this to set environment variables for local development +# cp .env.example .env.local + +# 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/.gitignore b/agent-openai-agents-sdk-stateful/.gitignore new file mode 100644 index 00000000..077d791b --- /dev/null +++ b/agent-openai-agents-sdk-stateful/.gitignore @@ -0,0 +1,206 @@ +# 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 +**/.claude +**/.env.local diff --git a/agent-openai-agents-sdk-stateful/AGENTS.md b/agent-openai-agents-sdk-stateful/AGENTS.md new file mode 100644 index 00000000..6e03e94f --- /dev/null +++ b/agent-openai-agents-sdk-stateful/AGENTS.md @@ -0,0 +1,530 @@ +# Agent Development Guide + +## For AI Agents: MANDATORY First Action + +**BEFORE any other action, run `databricks auth profiles` to check authentication status.** + +This helps you understand: +- Which Databricks profiles are configured +- Whether authentication is already set up +- Which profile to use for subsequent commands + +If no profiles exist, guide the user through running `./scripts/quickstart.sh` to set up authentication. + +## For AI Agents: 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**: Follow the steps in the "Binding an existing app to your bundle" section +- **If they want to delete**: Run `databricks apps delete ` then deploy again + +--- + +## Getting Started with Your Agent + +This guide walks you through the initial setup of your agent project: installing prerequisites, discovering available tools in your workspace, and testing the baseline template locally. + +**After completing these steps**, see the README.md for information on modifying your agent and deploying to Databricks. + +--- + +## Quick Setup + +**Prerequisites:** uv, nvm (Node 20), Databricks CLI (v0.283.0+) + +**What you need:** +- uv (Python package manager) +- nvm (Node.js version manager) with Node 20 +- Databricks CLI v0.283.0 or above + +**Quickest path to running:** + +```bash +# 0. Ensure you have the latest Databricks CLI +databricks -v # Should be v0.283.0 or above +brew upgrade databricks # Offer to run if the Databricks CLI version is too old + +# 1. Initialize git (recommended for version control) +git init + +# 2. Run quickstart for setup (auth, MLflow experiment) +./scripts/quickstart.sh + +# Or run non-interactively with a profile +./scripts/quickstart.sh --profile DEFAULT + +# Or with a host URL for initial setup +./scripts/quickstart.sh --host https://your-workspace.cloud.databricks.com + +# 3. Discover available tools (IMPORTANT - do this before coding!) +uv run discover-tools + +# 4. Start the agent server +uv run start-app +``` + +**Quickstart script handles:** +- Databricks authentication (OAuth) +- MLflow experiment creation +- Environment variable configuration (`.env.local`) + - Sets `DATABRICKS_CONFIG_PROFILE` to your selected profile + - Configures `MLFLOW_TRACKING_URI` as `databricks://` for proper local authentication + - Sets `MLFLOW_EXPERIMENT_ID` to the created experiment + +**Quickstart options:** +- `--profile NAME`: Use specified Databricks profile (non-interactive) +- `--host URL`: Databricks workspace URL (for initial setup) +- `-h, --help`: Show help message + +--- + +## Discovering Available Tools + +**⚠️ CRITICAL:** Always run tool discovery BEFORE writing agent code! + +This step helps you understand what resources are already available in your workspace, preventing duplicate work and showing you the best practices for connecting to each resource. + +```bash +# Discover all available resources (recommended) +uv run discover-tools + +# Limit to specific catalog/schema +uv run discover-tools --catalog my_catalog --schema my_schema + +# Output as JSON for programmatic use +uv run discover-tools --format json --output tools.json + +# Save markdown report +uv run discover-tools --output tools.md +``` + +**What gets discovered:** +1. **Unity Catalog Functions** - SQL UDFs usable as agent tools +2. **Unity Catalog Tables** - Structured data for querying +3. **Vector Search Indexes** - For RAG applications +4. **Genie Spaces** - Natural language interface to data +5. **Custom MCP Servers** - Your MCP servers deployed as Databricks Apps +6. **External MCP Servers** - Third-party MCP servers via UC connections + +**Using discovered tools in your agent:** + +After discovering tools, configure your agent to use them: + +```python +from databricks_openai.agents import McpServer + +async def init_mcp_server(): + return McpServer( + url=f"{host}/api/2.0/mcp/functions/{catalog}/{schema}", + name="my custom tools", + ) + +# Use in agent +agent = Agent( + name="my agent", + instructions="You are a helpful agent.", + model="databricks-claude-3-7-sonnet", + mcp_servers=[mcp_server], +) +``` + +See the [MCP documentation](https://docs.databricks.com/aws/en/generative-ai/mcp/) for more details. + +--- + +## Granting Access to App Resources + +### ⚠️ CRITICAL: Resource Permissions + +**After adding any MCP server to your agent, you MUST grant the app access to the server's dependent resource(s) in `databricks.yml`.** + +Without this, you'll get permission errors when the agent tries to use the resource. + +### Example Workflow + +**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", + instructions="You are a helpful agent.", + model="databricks-claude-3-7-sonnet", + mcp_servers=[genie_server], +) +``` + +**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' +``` + +### Resource Type Examples + +```yaml +# Unity Catalog function (for UC functions accessed via MCP) +- name: 'my_uc_function' + uc_securable: + securable_full_name: 'catalog.schema.function_name' + securable_type: 'FUNCTION' + permission: 'EXECUTE' + +# Unity Catalog connection (for external MCP servers via UC connections) +- name: 'my_connection' + uc_securable: + securable_full_name: 'my-connection-name' + securable_type: 'CONNECTION' + permission: 'USE_CONNECTION' + +# Vector search index +- name: 'my_vector_index' + uc_securable: + securable_full_name: 'catalog.schema.index_name' + securable_type: 'TABLE' + permission: 'SELECT' + +# SQL warehouse +- name: 'my_warehouse' + sql_warehouse: + sql_warehouse_id: 'abc123def456' + permission: 'CAN_USE' + +# Model serving endpoint +- name: 'my_endpoint' + serving_endpoint: + name: 'my_endpoint' + permission: 'CAN_QUERY' + +# Genie space +- name: 'my_genie_space' + genie_space: + name: 'My Genie Space' + space_id: '01234567-89ab-cdef' + permission: 'CAN_RUN' + +# MLflow experiment +- name: 'my_experiment' + experiment: + experiment_id: "12349876" + permission: 'CAN_MANAGE' +``` + +### Custom MCP Servers (Databricks Apps) + +If you're using custom MCP servers deployed as Databricks Apps (names starting with `mcp-`), you need to manually grant your agent app's service principal permission to access them: + +1. Find your agent app's service principal name: +```bash +databricks apps get --output json | jq -r '.service_principal_name' +``` + +2. Grant the service principal `CAN_USE` permission on the MCP server app: +```bash +databricks apps update-permissions --service-principal --permission-level CAN_USE +``` + +**Note:** Apps are not yet supported as resource dependencies for other apps in `databricks.yml`, so this manual permission grant is required for now. + +### Important Notes + +- The app automatically has access to the MLflow experiment (already configured in template) +- For all other resources (UC functions, Genie spaces, vector indexes, warehouses, etc.), you MUST add them +- Without proper resource grants, you'll see permission errors at runtime + +--- + +## Running the App Locally + +**Start the server:** + +```bash +uv run start-app +``` + +This starts the agent at http://localhost:8000 + +**Advanced server options:** + +```bash +uv run start-server --reload # Hot-reload on code changes during development +uv run start-server --port 8001 +uv run start-server --workers 4 +``` + +**Test the API:** + +```bash +# Streaming request +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' + +# Non-streaming request +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }] }' +``` + +**Common issues:** +- Port already in use: Use `--port` to specify a different port +- Authentication errors: Verify `.env.local` is correct +- Module not found: Run `uv sync` to install dependencies +- **MLflow experiment not found**: If you see an error like "The provided MLFLOW_EXPERIMENT_ID environment variable value does not exist", ensure your `MLFLOW_TRACKING_URI` in `.env.local` is set to `databricks://` (e.g., `databricks://DEFAULT-testing`). The quickstart script should configure this automatically, but if you manually edit `.env.local`, make sure to include the profile name in the tracking URI. + +--- + +## Modifying the Agent + +**Main file to modify:** `agent_server/agent.py` + +**Key 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/) + +**databricks-openai SDK basics:** + +```python +from databricks_openai import AsyncDatabricksOpenAI +from agents import set_default_openai_api, set_default_openai_client + +# Set up async client (recommended for agent servers) +set_default_openai_client(AsyncDatabricksOpenAI()) +set_default_openai_api("chat_completions") +``` + +--- + +## Testing the Agent + +```bash +# Run evaluation +uv run agent-evaluate + +# Run unit tests +pytest [path] +``` + +**Test API locally:** + +```bash +# Streaming request +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' + +# Non-streaming request +curl -X POST http://localhost:8000/invocations \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }] }' +``` + +--- + +## Deploying to Databricks Apps + +**Deploy using Databricks bundles:** + +```bash +# Deploy the bundle (creates/updates resources and uploads files) +databricks bundle deploy + +# Run the app (starts/restarts the app with uploaded source code) +databricks bundle run agent_openai_agents_sdk +``` + +The resource key `agent_openai_agents_sdk` matches the app name defined in `databricks.yml` under `resources.apps.agent_openai_agents_sdk`. + +**Error: "An app with the same name already exists"** + +If you see this error when running `databricks bundle deploy`: + +``` +Error: failed to create app + +Failed to create app . An app with the same name already exists. +``` + +This means you have an existing app that needs to be linked to your bundle. You have two options: + +1. **Bind the existing app to your bundle** (recommended if you want to manage the existing app): + - Follow the steps in [Binding an existing app to your bundle](#binding-an-existing-app-to-your-bundle) below + - This will link the existing app to your bundle so future deploys update it + +2. **Delete the existing app and let the bundle create a new one**: + ```bash + databricks apps delete + databricks bundle deploy + ``` + - ⚠️ This will permanently delete the existing app including its URL, OAuth credentials, and service principal + +**Binding an existing app to your bundle:** + +If you've already deployed an app from a different directory or through the UI and want to link it to this bundle, follow these steps: + +**Step 1: Update `databricks.yml` to match the existing app name** + +⚠️ **CRITICAL**: The app name in your `databricks.yml` **must match** the existing app name exactly, or Terraform will **destroy and recreate** the app (not update it in-place). + +First, find your existing app name: +```bash +# List existing apps to find the app name +databricks apps list --output json | jq '.[].name' +``` + +Then update `databricks.yml` to use that exact name: +```yaml +resources: + apps: + agent_openai_agents_sdk: + name: "openai-agents-sdk-agent" # Match your existing app name exactly + description: "OpenAI Agents SDK agent application" + source_code_path: ./ +``` + +The default configuration uses: +```yaml +name: "${bundle.target}-agent-openai-agents-sdk" # Evaluates to "dev-agent-openai-agents-sdk" +``` + +Make sure to replace this with your actual app name. + +**Step 2: Bind the resource to the existing app** + +```bash +# Bind the resource to the existing app +databricks bundle deployment bind agent_openai_agents_sdk + +# Example: +databricks bundle deployment bind agent_openai_agents_sdk openai-agents-sdk-agent + +# If the operation requires confirmation and you want to skip prompts: +databricks bundle deployment bind agent_openai_agents_sdk openai-agents-sdk-agent --auto-approve +``` + +This links your bundle configuration to the existing deployed app. Future `databricks bundle deploy` commands will update the existing app instead of creating a new one. + +**Important notes about binding:** +- **Remote Terraform state**: Databricks stores Terraform state remotely, so the same app can be detected across different local directories +- **Name is immutable**: The `name` field cannot be changed in-place; changing it forces replacement (destroy + create) +- **Review the plan**: When binding, carefully review the Terraform plan output. Look for `# forces replacement` which indicates the app will be destroyed and recreated +- **Existing binding**: If a resource is already bound to another app, you must unbind it first before binding to a different app + +**Unbinding a resource:** + +To remove the link between your bundle and the deployed app: + +```bash +databricks bundle deployment unbind agent_openai_agents_sdk +``` + +This is useful when: +- You want to bind to a different app +- You want to let the bundle create a new app on the next deploy +- You're switching between different deployed instances + +Note: Unbinding only removes the link in your bundle state - it does not delete the deployed app. + +**Query deployed app:** + +Generate OAuth token (PATs are not supported): + +```bash +databricks auth token +``` + +Send request: + +```bash +curl -X POST /invocations \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{ "input": [{ "role": "user", "content": "hi" }], "stream": true }' +``` + +**Debug deployed apps:** + +```bash +# View logs (use the deployed app name from databricks.yml) +databricks apps logs dev-agent-openai-agents-sdk --follow + +# Check status +databricks apps get dev-agent-openai-agents-sdk --output json | jq '{app_status, compute_status}' +``` + +--- + +## Key Files + +| File | Purpose | +| -------------------------------- | --------------------------------------------- | +| `agent_server/agent.py` | Agent logic, model, instructions, MCP servers | +| `agent_server/start_server.py` | FastAPI server + MLflow setup | +| `agent_server/evaluate_agent.py` | Agent evaluation with MLflow scorers | +| `agent_server/utils.py` | Databricks auth helpers, stream processing | +| `scripts/start_app.py` | Manages backend+frontend startup | +| `scripts/discover_tools.py` | Discovers available workspace resources | +| `scripts/quickstart.sh` | One-command setup script | + +--- + +## Agent Framework Capabilities + +**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/ + +--- + +## Next Steps + +✅ **You've completed the initial setup!** + +After running the quickstart script, you have: +- ✅ Installed prerequisites +- ✅ Authenticated with Databricks +- ✅ Created MLflow experiment +- ✅ Discovered available tools in your workspace + +Now you're ready to: +- Start the agent locally: `uv run start-app` +- Modify your agent to use the tools you discovered +- Deploy your agent to Databricks + +**See the README.md** for more information on: +- Modifying the agent and adding tools +- Evaluating your agent +- Deploying to Databricks Apps +- Debugging and monitoring deployed apps diff --git a/agent-openai-agents-sdk-stateful/CLAUDE.md b/agent-openai-agents-sdk-stateful/CLAUDE.md new file mode 100644 index 00000000..4202b8f0 --- /dev/null +++ b/agent-openai-agents-sdk-stateful/CLAUDE.md @@ -0,0 +1,3 @@ +@AGENTS.md + + diff --git a/agent-openai-agents-sdk-stateful/README.md b/agent-openai-agents-sdk-stateful/README.md new file mode 100644 index 00000000..fbff3176 --- /dev/null +++ b/agent-openai-agents-sdk-stateful/README.md @@ -0,0 +1,237 @@ +# 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. + +## Quick start + +Run the `./scripts/quickstart.sh` 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 +./scripts/quickstart.sh +``` + +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.local 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.local 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 `./scripts/quickstart.sh` 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.local` and update the `MLFLOW_EXPERIMENT_ID` in your `.env.local` file with the experiment ID you created. The `.env.local` file will be automatically loaded when starting the server. + + ```bash + cp .env.example .env.local + # Edit .env.local 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/agent_server/__init__.py b/agent-openai-agents-sdk-stateful/agent_server/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/agent-openai-agents-sdk-stateful/agent_server/agent.py b/agent-openai-agents-sdk-stateful/agent_server/agent.py new file mode 100644 index 00000000..13786126 --- /dev/null +++ b/agent-openai-agents-sdk-stateful/agent_server/agent.py @@ -0,0 +1,65 @@ +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 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, +) + +# 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() + 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) + return ResponsesAgentResponse(output=[item.to_input_item() for item in result.new_items]) + + +@stream() +async def stream(request: dict) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + # Optionally use the user's workspace client for on-behalf-of authentication + # user_workspace_client = get_user_workspace_client() + 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) + + async for event in process_agent_stream_events(result.stream_events()): + yield event diff --git a/agent-openai-agents-sdk-stateful/agent_server/evaluate_agent.py b/agent-openai-agents-sdk-stateful/agent_server/evaluate_agent.py new file mode 100644 index 00000000..1035b293 --- /dev/null +++ b/agent-openai-agents-sdk-stateful/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.local if it exists +load_dotenv(dotenv_path=".env.local", 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/agent_server/start_server.py b/agent-openai-agents-sdk-stateful/agent_server/start_server.py new file mode 100644 index 00000000..51388842 --- /dev/null +++ b/agent-openai-agents-sdk-stateful/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.local before importing the agent for proper auth +load_dotenv(dotenv_path=".env.local", 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/agent_server/utils.py b/agent-openai-agents-sdk-stateful/agent_server/utils.py new file mode 100644 index 00000000..03e52e9e --- /dev/null +++ b/agent-openai-agents-sdk-stateful/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/app.yaml b/agent-openai-agents-sdk-stateful/app.yaml new file mode 100644 index 00000000..34465373 --- /dev/null +++ b/agent-openai-agents-sdk-stateful/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/databricks.yml b/agent-openai-agents-sdk-stateful/databricks.yml new file mode 100644 index 00000000..48e2794f --- /dev/null +++ b/agent-openai-agents-sdk-stateful/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/e2e-chatbot-app-next/databricks.yml b/agent-openai-agents-sdk-stateful/e2e-chatbot-app-next/databricks.yml new file mode 100644 index 00000000..28801e62 --- /dev/null +++ b/agent-openai-agents-sdk-stateful/e2e-chatbot-app-next/databricks.yml @@ -0,0 +1,63 @@ +bundle: + name: databricks-chatbot + +variables: + serving_endpoint_name: + description: "Name of the model serving endpoint to be used by the app" + # TODO: uncomment the line below and specify a default value to avoid needing to specify it on each deployment + # default: "your-serving-endpoint-name-goes-here" + database_instance_name: + description: "Base name of the Lakebase database instance" + default: "chatbot-lakebase" + resource_name_suffix: + description: "Suffix for resource names (app/lakebase instance names etc). Helpful for avoiding name collisions if multiple users deploy this app to the same workspace" + +resources: + database_instances: + # TODO (optional): Uncomment the database resource below to deploy the app with a database + # chatbot_lakebase: + # # NOTE: customize the Lakebase instance name here, if desired + # name: ${var.database_instance_name}-${var.resource_name_suffix} + # capacity: CU_1 + + apps: + databricks_chatbot: + # NOTE: customize the app name here, if desired + name: db-chatbot-${var.resource_name_suffix} + description: "Agentic Chat application for Foundation Models, Agent Bricks, and custom code agents" + source_code_path: . + resources: + - name: serving-endpoint + description: "Databricks serving endpoint name for the AI agent" + # NOTE: If chatting with an Agent Bricks Multi-Agent Supervisor (https://docs.databricks.com/aws/en/generative-ai/agent-bricks/multi-agent-supervisor), + # you need to additionally grant the app service principal the `CAN_QUERY` permission on the + # underlying agent(s) that the MAS orchestrates. You can do this by adding those + # agent serving endpoints as additional resources here. + serving_endpoint: + name: ${var.serving_endpoint_name} + permission: CAN_QUERY + # TODO (optional): Uncomment the database resource below to bind the app to the deployed database instance + # and to enable persistent chat history + # - name: database + # description: "Lakebase database instance for the chat app" + # database: + # database_name: databricks_postgres + # instance_name: ${resources.database_instances.chatbot_lakebase.name} + # permission: CAN_CONNECT_AND_CREATE + +targets: + dev: + mode: development + default: true + variables: + resource_name_suffix: dev-${workspace.current_user.domain_friendly_name} + + staging: + mode: production + variables: + resource_name_suffix: "staging" + + prod: + mode: production + variables: + resource_name_suffix: "prod" diff --git a/agent-openai-agents-sdk-stateful/pyproject.toml b/agent-openai-agents-sdk-stateful/pyproject.toml new file mode 100644 index 00000000..1bb46f31 --- /dev/null +++ b/agent-openai-agents-sdk-stateful/pyproject.toml @@ -0,0 +1,37 @@ +[project] +name = "agent-server" +version = "0.1.0" +description = "MLflow-compatible agent server with FastAPI" +readme = "README.md" +authors = [ + { name = "Agent Developer", email = "developer@example.com" } +] +requires-python = ">=3.11" +dependencies = [ + "fastapi>=0.115.12", + "uvicorn>=0.34.2", + "databricks-openai>=0.9.0", + "mlflow>=3.8.0rc0", + "openai-agents>=0.4.1", + "python-dotenv", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.uv] +prerelease = "allow" + +[dependency-groups] +dev = [ + "hatchling>=1.27.0", + "pytest>=7.0.0", +] + + +[project.scripts] +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/requirements.txt b/agent-openai-agents-sdk-stateful/requirements.txt new file mode 100644 index 00000000..60cc5e6a --- /dev/null +++ b/agent-openai-agents-sdk-stateful/requirements.txt @@ -0,0 +1 @@ +uv diff --git a/agent-openai-agents-sdk-stateful/scripts/__init__.py b/agent-openai-agents-sdk-stateful/scripts/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/agent-openai-agents-sdk-stateful/scripts/discover_tools.py b/agent-openai-agents-sdk-stateful/scripts/discover_tools.py new file mode 100755 index 00000000..3eb37963 --- /dev/null +++ b/agent-openai-agents-sdk-stateful/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/scripts/quickstart.sh b/agent-openai-agents-sdk-stateful/scripts/quickstart.sh new file mode 100755 index 00000000..870351ca --- /dev/null +++ b/agent-openai-agents-sdk-stateful/scripts/quickstart.sh @@ -0,0 +1,392 @@ +#!/bin/bash +set -e + +# Parse command line arguments +PROFILE_ARG="" +HOST_ARG="" + +while [[ $# -gt 0 ]]; do + case $1 in + --profile) + PROFILE_ARG="$2" + shift 2 + ;; + --host) + HOST_ARG="$2" + shift 2 + ;; + -h|--help) + echo "Usage: $0 [OPTIONS]" + echo "" + echo "Options:" + echo " --profile NAME Use specified Databricks profile (non-interactive)" + echo " --host URL Databricks workspace URL (for initial setup)" + echo " -h, --help Show this help message" + exit 0 + ;; + *) + echo "Error: Unknown option $1" + echo "Use --help for usage information" + exit 1 + ;; + esac +done + +# Helper function to check if a command exists +command_exists() { + command -v "$1" >/dev/null 2>&1 +} + +# Helper function to check if Homebrew is available +has_brew() { + command_exists brew +} + +echo "===================================================================" +echo "Agent on Apps - Quickstart Setup" +echo "===================================================================" +echo + +# =================================================================== +# Section 1: Prerequisites Installation +# =================================================================== + +echo "Checking and installing prerequisites..." +echo + +# Check and install UV +if command_exists uv; then + echo "✓ UV is already installed" + uv --version +else + echo "Installing UV..." + if has_brew; then + echo "Using Homebrew to install UV..." + brew install uv + else + echo "Using curl to install UV..." + curl -LsSf https://astral.sh/uv/install.sh | sh + # Add UV to PATH for current session + export PATH="$HOME/.cargo/bin:$PATH" + fi + echo "✓ UV installed successfully" +fi + +# Check and install nvm +if [ -s "$HOME/.nvm/nvm.sh" ]; then + echo "✓ nvm is already installed" + # Load nvm for current session + export NVM_DIR="$HOME/.nvm" + [ -s "$NVM_DIR/nvm.sh" ] && \. "$NVM_DIR/nvm.sh" +else + echo "Installing nvm..." + if has_brew; then + echo "Using Homebrew to install nvm..." + brew install nvm + # Create nvm directory + mkdir -p ~/.nvm + # Add nvm to current session + export NVM_DIR="$HOME/.nvm" + [ -s "/opt/homebrew/opt/nvm/nvm.sh" ] && \. "/opt/homebrew/opt/nvm/nvm.sh" + [ -s "/usr/local/opt/nvm/nvm.sh" ] && \. "/usr/local/opt/nvm/nvm.sh" + else + echo "Using curl to install nvm..." + curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.39.0/install.sh | bash + # Load nvm for current session + export NVM_DIR="$HOME/.nvm" + [ -s "$NVM_DIR/nvm.sh" ] && \. "$NVM_DIR/nvm.sh" + fi + echo "✓ nvm installed successfully" +fi + +# Use Node 20 +echo "Setting up Node.js 20..." +nvm install 20 +nvm use 20 +echo "✓ Node.js 20 is now active" +node --version +npm --version +echo + +# Check and install Databricks CLI +if command_exists databricks; then + echo "✓ Databricks CLI is already installed" + databricks --version +else + echo "Installing Databricks CLI..." + if has_brew; then + echo "Using Homebrew to install Databricks CLI..." + brew tap databricks/tap + brew install databricks + else + echo "Using curl to install Databricks CLI..." + if curl -fsSL https://raw.githubusercontent.com/databricks/setup-cli/main/install.sh | sh; then + echo "✓ Databricks CLI installed successfully" + else + echo "Installation failed, trying with sudo..." + curl -fsSL https://raw.githubusercontent.com/databricks/setup-cli/main/install.sh | sudo sh + fi + fi + echo "✓ Databricks CLI installed successfully" +fi +echo + +# =================================================================== +# Section 2: Configuration Files Setup +# =================================================================== +echo "Setting up configuration files..." + +# Copy .env.example to .env.local if it doesn't exist +if [ ! -f ".env.local" ]; then + echo "Copying .env.example to .env.local..." + cp .env.example .env.local + echo +else + echo ".env.local already exists, skipping copy..." +fi +echo + +# =================================================================== +# Section 3: Databricks Authentication +# =================================================================== + +echo "Setting up Databricks authentication..." + +# If --profile flag was provided, use that directly +if [ -n "$PROFILE_ARG" ]; then + PROFILE_NAME="$PROFILE_ARG" + echo "Using specified profile: $PROFILE_NAME" +else + # Check if there are existing profiles + set +e + EXISTING_PROFILES=$(databricks auth profiles 2>/dev/null) + PROFILES_EXIT_CODE=$? + set -e + + if [ $PROFILES_EXIT_CODE -eq 0 ] && [ -n "$EXISTING_PROFILES" ]; then + # Profiles exist - let user select one + echo "Found existing Databricks profiles:" + echo + + # Parse profiles into an array (compatible with older bash) + # Skip the first line (header row) + PROFILE_ARRAY=() + PROFILE_NAMES=() + LINE_NUM=0 + while IFS= read -r line; do + if [ -n "$line" ]; then + if [ $LINE_NUM -eq 0 ]; then + # Print header without number + echo "$line" + else + # Add full line to display array + PROFILE_ARRAY+=("$line") + # Extract just the profile name (first column) for selection + PROFILE_NAME_ONLY=$(echo "$line" | awk '{print $1}') + PROFILE_NAMES+=("$PROFILE_NAME_ONLY") + fi + LINE_NUM=$((LINE_NUM + 1)) + fi + done <<< "$EXISTING_PROFILES" + echo + + # Display numbered list + for i in "${!PROFILE_ARRAY[@]}"; do + echo "$((i+1))) ${PROFILE_ARRAY[$i]}" + done + echo + + echo "Enter the number of the profile you want to use:" + read -r PROFILE_CHOICE + + if [ -z "$PROFILE_CHOICE" ]; then + echo "Error: Profile selection is required" + exit 1 + fi + + # Validate the choice is a number + if ! [[ "$PROFILE_CHOICE" =~ ^[0-9]+$ ]]; then + echo "Error: Please enter a valid number" + exit 1 + fi + + # Convert to array index (subtract 1) + PROFILE_INDEX=$((PROFILE_CHOICE - 1)) + + # Check if the index is valid + if [ $PROFILE_INDEX -lt 0 ] || [ $PROFILE_INDEX -ge ${#PROFILE_NAMES[@]} ]; then + echo "Error: Invalid selection. Please choose a number between 1 and ${#PROFILE_NAMES[@]}" + exit 1 + fi + + # Get the selected profile name (just the name, not the full line) + PROFILE_NAME="${PROFILE_NAMES[$PROFILE_INDEX]}" + echo "Selected profile: $PROFILE_NAME" + fi +fi + +# Validate the profile if it was specified +if [ -n "$PROFILE_NAME" ]; then + + # Test if the profile works + set +e + DATABRICKS_CONFIG_PROFILE="$PROFILE_NAME" databricks current-user me >/dev/null 2>&1 + PROFILE_TEST=$? + set -e + + if [ $PROFILE_TEST -eq 0 ]; then + echo "✓ Successfully validated profile '$PROFILE_NAME'" + else + # Profile exists but isn't authenticated - prompt to authenticate + echo "Profile '$PROFILE_NAME' is not authenticated." + echo "Authenticating profile '$PROFILE_NAME'..." + echo "You will be prompted to log in to Databricks in your browser." + echo + + # Temporarily disable exit on error for the auth command + set +e + + # Run auth login with the profile name and capture output while still showing it to the user + AUTH_LOG=$(mktemp) + databricks auth login --profile "$PROFILE_NAME" 2>&1 | tee "$AUTH_LOG" + AUTH_EXIT_CODE=$? + + set -e + + if [ $AUTH_EXIT_CODE -eq 0 ]; then + echo "✓ Successfully authenticated profile '$PROFILE_NAME'" + # Clean up temp file + rm -f "$AUTH_LOG" + else + # Clean up temp file + rm -f "$AUTH_LOG" + echo "Error: Profile '$PROFILE_NAME' authentication failed" + exit 1 + fi + fi + + # Update .env.local with the profile name + if grep -q "DATABRICKS_CONFIG_PROFILE=" .env.local; then + sed -i '' "s/DATABRICKS_CONFIG_PROFILE=.*/DATABRICKS_CONFIG_PROFILE=$PROFILE_NAME/" .env.local + else + echo "DATABRICKS_CONFIG_PROFILE=$PROFILE_NAME" >> .env.local + fi + + # Update MLFLOW_TRACKING_URI to use the profile for local development + if grep -q "MLFLOW_TRACKING_URI=" .env.local; then + sed -i '' "s|MLFLOW_TRACKING_URI=.*|MLFLOW_TRACKING_URI=\"databricks://$PROFILE_NAME\"|" .env.local + else + echo "MLFLOW_TRACKING_URI=\"databricks://$PROFILE_NAME\"" >> .env.local + fi + echo "✓ Databricks profile '$PROFILE_NAME' saved to .env.local" +else + # No profiles exist - create default one + echo "No existing profiles found. Setting up Databricks authentication..." + + # Use --host flag if provided, otherwise prompt + if [ -n "$HOST_ARG" ]; then + DATABRICKS_HOST="$HOST_ARG" + echo "Using specified host: $DATABRICKS_HOST" + else + echo "Please enter your Databricks host URL (e.g., https://your-workspace.cloud.databricks.com):" + read -r DATABRICKS_HOST + + if [ -z "$DATABRICKS_HOST" ]; then + echo "Error: Databricks host is required" + exit 1 + fi + fi + + echo "Authenticating with Databricks..." + echo "You will be prompted to log in to Databricks in your browser." + echo + + # Temporarily disable exit on error for the auth command + set +e + + # Run auth login with profile DEFAULT and host parameter + AUTH_LOG=$(mktemp) + databricks auth login --profile DEFAULT --host "$DATABRICKS_HOST" 2>&1 | tee "$AUTH_LOG" + AUTH_EXIT_CODE=$? + + set -e + + if [ $AUTH_EXIT_CODE -eq 0 ]; then + echo "✓ Successfully authenticated with Databricks" + + # Use DEFAULT as the profile name + PROFILE_NAME="DEFAULT" + + # Clean up temp file + rm -f "$AUTH_LOG" + + # Update .env.local with the profile name + if grep -q "DATABRICKS_CONFIG_PROFILE=" .env.local; then + sed -i '' "s/DATABRICKS_CONFIG_PROFILE=.*/DATABRICKS_CONFIG_PROFILE=$PROFILE_NAME/" .env.local + else + echo "DATABRICKS_CONFIG_PROFILE=$PROFILE_NAME" >> .env.local + fi + + # Update MLFLOW_TRACKING_URI to use the profile for local development + if grep -q "MLFLOW_TRACKING_URI=" .env.local; then + sed -i '' "s|MLFLOW_TRACKING_URI=.*|MLFLOW_TRACKING_URI=\"databricks://$PROFILE_NAME\"|" .env.local + else + echo "MLFLOW_TRACKING_URI=\"databricks://$PROFILE_NAME\"" >> .env.local + fi + + echo "✓ Databricks profile '$PROFILE_NAME' saved to .env.local" + else + # Clean up temp file + rm -f "$AUTH_LOG" + echo "Databricks authentication was cancelled or failed." + echo "Please run this script again when you're ready to authenticate." + exit 1 + fi +fi +echo + +# =================================================================== +# Section 4: MLflow Experiment Setup +# =================================================================== + + +# Get current Databricks username +echo "Getting Databricks username..." +DATABRICKS_USERNAME=$(databricks -p $PROFILE_NAME current-user me | jq -r .userName) +echo "Username: $DATABRICKS_USERNAME" +echo + +# Create MLflow experiment and capture the experiment ID +echo "Creating MLflow experiment..." +EXPERIMENT_NAME="/Users/$DATABRICKS_USERNAME/agents-on-apps" + +# Try to create the experiment with the default name first +if EXPERIMENT_RESPONSE=$(databricks -p $PROFILE_NAME experiments create-experiment $EXPERIMENT_NAME 2>/dev/null); then + EXPERIMENT_ID=$(echo $EXPERIMENT_RESPONSE | jq -r .experiment_id) + echo "Created experiment '$EXPERIMENT_NAME' with ID: $EXPERIMENT_ID" +else + echo "Experiment name already exists, creating with random suffix..." + RANDOM_SUFFIX=$(openssl rand -hex 4) + EXPERIMENT_NAME="/Users/$DATABRICKS_USERNAME/agents-on-apps-$RANDOM_SUFFIX" + EXPERIMENT_RESPONSE=$(databricks -p $PROFILE_NAME experiments create-experiment $EXPERIMENT_NAME) + EXPERIMENT_ID=$(echo $EXPERIMENT_RESPONSE | jq -r .experiment_id) + echo "Created experiment '$EXPERIMENT_NAME' with ID: $EXPERIMENT_ID" +fi +echo + +# Update .env.local with the experiment ID +echo "Updating .env.local with experiment ID..." +sed -i '' "s/MLFLOW_EXPERIMENT_ID=.*/MLFLOW_EXPERIMENT_ID=$EXPERIMENT_ID/" .env.local +echo + +echo "===================================================================" +echo "Setup Complete!" +echo "===================================================================" +echo "✓ Prerequisites installed (UV, nvm, Databricks CLI)" +echo "✓ Databricks authenticated with profile: $PROFILE_NAME" +echo "✓ Configuration files created (.env.local)" +echo "✓ MLflow experiment created: $EXPERIMENT_NAME" +echo "✓ Experiment ID: $EXPERIMENT_ID" +echo "✓ Configuration updated in .env.local" +echo "===================================================================" +echo diff --git a/agent-openai-agents-sdk-stateful/scripts/start_app.py b/agent-openai-agents-sdk-stateful/scripts/start_app.py new file mode 100644 index 00000000..a812880f --- /dev/null +++ b/agent-openai-agents-sdk-stateful/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.local", 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() diff --git a/agent-openai-agents-sdk/e2e-chatbot-app-next/databricks.yml b/agent-openai-agents-sdk/e2e-chatbot-app-next/databricks.yml new file mode 100644 index 00000000..28801e62 --- /dev/null +++ b/agent-openai-agents-sdk/e2e-chatbot-app-next/databricks.yml @@ -0,0 +1,63 @@ +bundle: + name: databricks-chatbot + +variables: + serving_endpoint_name: + description: "Name of the model serving endpoint to be used by the app" + # TODO: uncomment the line below and specify a default value to avoid needing to specify it on each deployment + # default: "your-serving-endpoint-name-goes-here" + database_instance_name: + description: "Base name of the Lakebase database instance" + default: "chatbot-lakebase" + resource_name_suffix: + description: "Suffix for resource names (app/lakebase instance names etc). Helpful for avoiding name collisions if multiple users deploy this app to the same workspace" + +resources: + database_instances: + # TODO (optional): Uncomment the database resource below to deploy the app with a database + # chatbot_lakebase: + # # NOTE: customize the Lakebase instance name here, if desired + # name: ${var.database_instance_name}-${var.resource_name_suffix} + # capacity: CU_1 + + apps: + databricks_chatbot: + # NOTE: customize the app name here, if desired + name: db-chatbot-${var.resource_name_suffix} + description: "Agentic Chat application for Foundation Models, Agent Bricks, and custom code agents" + source_code_path: . + resources: + - name: serving-endpoint + description: "Databricks serving endpoint name for the AI agent" + # NOTE: If chatting with an Agent Bricks Multi-Agent Supervisor (https://docs.databricks.com/aws/en/generative-ai/agent-bricks/multi-agent-supervisor), + # you need to additionally grant the app service principal the `CAN_QUERY` permission on the + # underlying agent(s) that the MAS orchestrates. You can do this by adding those + # agent serving endpoints as additional resources here. + serving_endpoint: + name: ${var.serving_endpoint_name} + permission: CAN_QUERY + # TODO (optional): Uncomment the database resource below to bind the app to the deployed database instance + # and to enable persistent chat history + # - name: database + # description: "Lakebase database instance for the chat app" + # database: + # database_name: databricks_postgres + # instance_name: ${resources.database_instances.chatbot_lakebase.name} + # permission: CAN_CONNECT_AND_CREATE + +targets: + dev: + mode: development + default: true + variables: + resource_name_suffix: dev-${workspace.current_user.domain_friendly_name} + + staging: + mode: production + variables: + resource_name_suffix: "staging" + + prod: + mode: production + variables: + resource_name_suffix: "prod" From 75ec2c7a0a1db3ea4ef4f01f21059d2388fda6fd Mon Sep 17 00:00:00 2001 From: Jenny Date: Wed, 21 Jan 2026 18:07:36 -0800 Subject: [PATCH 2/2] openai stateful agent updates --- agent-openai-agents-sdk-stateful/.env.example | 3 + agent-openai-agents-sdk-stateful/README.md | 48 +++++- .../agent_server/agent.py | 137 ++++++++++++++++-- agent-openai-agents-sdk-stateful/app.yaml | 2 + .../pyproject.toml | 8 +- 5 files changed, 184 insertions(+), 14 deletions(-) diff --git a/agent-openai-agents-sdk-stateful/.env.example b/agent-openai-agents-sdk-stateful/.env.example index 7c4ebf97..6748df29 100644 --- a/agent-openai-agents-sdk-stateful/.env.example +++ b/agent-openai-agents-sdk-stateful/.env.example @@ -9,6 +9,9 @@ DATABRICKS_CONFIG_PROFILE=DEFAULT # TODO: Update with the MLflow experiment you want to log traces and models to MLFLOW_EXPERIMENT_ID= +# TODO: Update with Lakebase instance name for stateful session memory +LAKEBASE_INSTANCE_NAME=lakebase + 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 diff --git a/agent-openai-agents-sdk-stateful/README.md b/agent-openai-agents-sdk-stateful/README.md index fbff3176..363fb349 100644 --- a/agent-openai-agents-sdk-stateful/README.md +++ b/agent-openai-agents-sdk-stateful/README.md @@ -173,7 +173,7 @@ After it completes, open the MLflow UI link for your experiment to inspect resul 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. + For this example, you need to add an MLflow Experiment and Lakebase instance 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) and [Databricks Apps Lakebase documentation](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/lakebase) 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). @@ -190,7 +190,49 @@ After it completes, open the MLflow UI link for your experiment to inspect resul databricks sync . "/Users/$DATABRICKS_USERNAME/agent-openai-agents-sdk" ``` -3. **Deploy your Databricks App** +3. **Grant Lakebase permissions to your App's Service Principal** + + Before deploying/querying your agent, you need to ensure your app has access to the necessary Lakebase tables for memory. + + First, add your Lakebase instance as a resource to your app: + - Go to the Databricks UI + - Navigate to your app and click **Edit** + - Go to **App resources** → **Add resource** + - Add your Lakebase instance that you are using for short-term memory store + + After adding your Lakebase as a resource to your app (with the Connect + Create permissions), you'll need to ensure access to certain schemas and tables that have already been created during local testing. To grant the necessary permissions on your Lakebase instance for your app's service principal, run the following SQL commands on your Lakebase instance (replace `app-sp-id` with your app's service principal UUID): + + ```sql + DO $$ + DECLARE + app_sp text := 'app-sp-id'; -- TODO: Replace with your App's Service Principal ID here + BEGIN + EXECUTE format('GRANT CREATE ON DATABASE databricks_postgres TO %I;', app_sp); + ------------------------------------------------------------------- + -- Drizzle schema: migration metadata tables + ------------------------------------------------------------------- + EXECUTE format('GRANT USAGE, CREATE ON SCHEMA drizzle TO %I;', app_sp); + EXECUTE format('GRANT SELECT, INSERT, UPDATE ON ALL TABLES IN SCHEMA drizzle TO %I;', app_sp); + ------------------------------------------------------------------- + -- App schema: business tables (Chat, Message, etc.) + ------------------------------------------------------------------- + EXECUTE format('GRANT USAGE, CREATE ON SCHEMA ai_chatbot TO %I;', app_sp); + EXECUTE format('GRANT SELECT, INSERT, UPDATE ON ALL TABLES IN SCHEMA ai_chatbot TO %I;', app_sp); + ------------------------------------------------------------------- + -- For agent memory/backend: Public schema for short-term memory tables + ------------------------------------------------------------------- + EXECUTE format('GRANT USAGE, CREATE ON SCHEMA public TO %I;', app_sp); + EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.agent_sessions TO %I;', app_sp); + EXECUTE format('GRANT SELECT, INSERT, UPDATE ON TABLE public.agent_messages TO %I;', app_sp); + -- For all sequences in public (short-term memory tables) + EXECUTE format( + 'GRANT USAGE, SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA public TO %I;', + app_sp + ); + END $$; + ``` + +4. **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). @@ -198,7 +240,7 @@ After it completes, open the MLflow UI link for your experiment to inspect resul 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** +5. **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): diff --git a/agent-openai-agents-sdk-stateful/agent_server/agent.py b/agent-openai-agents-sdk-stateful/agent_server/agent.py index 13786126..24e52b09 100644 --- a/agent-openai-agents-sdk-stateful/agent_server/agent.py +++ b/agent-openai-agents-sdk-stateful/agent_server/agent.py @@ -1,10 +1,12 @@ -from typing import AsyncGenerator - +from typing import AsyncGenerator, Optional +import os +import uuid_utils 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 AsyncMemorySession from mlflow.genai.agent_server import invoke, stream from mlflow.types.responses import ( ResponsesAgentRequest, @@ -24,6 +26,17 @@ set_trace_processors([]) # only use mlflow for trace processing mlflow.openai.autolog() +LAKEBASE_INSTANCE_NAME = os.getenv("LAKEBASE_INSTANCE_NAME", "") +print("lakebase instance name is") +print(LAKEBASE_INSTANCE_NAME) + +if not LAKEBASE_INSTANCE_NAME: + raise ValueError( + "LAKEBASE_INSTANCE_NAME environment variable is required but not set. " + "Please set it in your environment:\n" + " LAKEBASE_INSTANCE_NAME=\n" + ) + async def init_mcp_server(): return McpServer( @@ -34,32 +47,136 @@ async def init_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.", + name="assistant", + instructions="""You are a helpful assistant. You have access to a Python code execution tool + that you can use when needed for mathematical calculations, data analysis + and running code the user asks you to execute. For general conversation, questions, + and discussions, respond directly without using tools.""", model="databricks-gpt-5-2", mcp_servers=[mcp_server], ) +def extract_session_id(request: ResponsesAgentRequest) -> Optional[str]: + """ + Extract session_id from custom_inputs if present. + """ + if hasattr(request, "custom_inputs") and request.custom_inputs: + return request.custom_inputs.get("session_id") + return None + + +def extract_latest_user_message(request: ResponsesAgentRequest) -> str: + """ + Extract the latest user message content as a string. + + When using session memory, the OpenAI Agents SDK expects a string input + (the new user message). The session handles conversation history. + """ + # Find the last user message in the input + for item in reversed(request.input): + item_dict = item.model_dump() + if item_dict.get("role") == "user": + content = item_dict.get("content") + # Handle both string content and structured content + if isinstance(content, str): + return content + elif isinstance(content, list): + # Extract text from structured content (e.g., [{"type": "input_text", "text": "..."}]) + for part in content: + if isinstance(part, dict): + if part.get("type") in ("input_text", "text"): + return part.get("text", "") + elif "text" in part: + return part["text"] + return "" + + +def get_session(session_id: Optional[str] = None) -> tuple[AsyncMemorySession, str]: + """ + Get or create a AsyncMemorySession for the given session_id. + + Args: + session_id: Optional session ID. Can be any meaningful identifier for user, thread, or context. + If not provided, a new UUID is generated. + + Returns: + A tuple of (session, session_id) where session_id is the resolved ID. + """ + if not LAKEBASE_INSTANCE_NAME: + raise ValueError( + "LAKEBASE_INSTANCE_NAME environment variable is not set. " + "Please set it to your Lakebase instance name for session memory." + ) + + # Generate new session_id if not provided + resolved_session_id = session_id or str(uuid_utils.uuid7()) + + session = AsyncMemorySession( + session_id=resolved_session_id, + instance_name=LAKEBASE_INSTANCE_NAME, + ) + + return session, resolved_session_id @invoke() async def invoke(request: ResponsesAgentRequest) -> ResponsesAgentResponse: + """ + Handle non-streaming agent invocation with stateful session memory. + + Supports session_id in custom_inputs for conversation continuity. + The session_id can be any meaningful identifier: + - User-based: "user_12345" + - Thread-based: "thread_abc123" + - Context-based: "support_ticket_456" + + If session_id is not provided, a new UUID is generated and returned. + """ # Optionally use the user's workspace client for on-behalf-of authentication # user_workspace_client = get_user_workspace_client() + + session_id = extract_session_id(request) + 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) - return ResponsesAgentResponse(output=[item.to_input_item() for item in result.new_items]) + user_message = extract_latest_user_message(request) + session, resolved_session_id = get_session(session_id) + result = await Runner.run(agent, user_message, session=session) + return ResponsesAgentResponse( + output=[item.to_input_item() for item in result.new_items], + custom_outputs={"session_id": resolved_session_id}, + ) @stream() -async def stream(request: dict) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: +async def stream(request: ResponsesAgentRequest) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: + """ + Handle streaming agent invocation with stateful session memory. + + Supports session_id in custom_inputs for conversation continuity. + The session_id can be any meaningful identifier: + - User-based: "user_12345" + - Thread-based: "thread_abc123" + - Context-based: "support_ticket_456" + + If session_id is not provided, a new UUID is generated and returned. + """ # Optionally use the user's workspace client for on-behalf-of authentication # user_workspace_client = get_user_workspace_client() + + session_id = extract_session_id(request) + 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) + user_message = extract_latest_user_message(request) + session, resolved_session_id = get_session(session_id) + result = Runner.run_streamed(agent, input=user_message, session=session) async for event in process_agent_stream_events(result.stream_events()): yield event + + yield ResponsesAgentStreamEvent( + type="response.done", + response={ + "custom_outputs": {"session_id": resolved_session_id} + } + ) diff --git a/agent-openai-agents-sdk-stateful/app.yaml b/agent-openai-agents-sdk-stateful/app.yaml index 34465373..f5bc6567 100644 --- a/agent-openai-agents-sdk-stateful/app.yaml +++ b/agent-openai-agents-sdk-stateful/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/pyproject.toml b/agent-openai-agents-sdk-stateful/pyproject.toml index 1bb46f31..95a28a8b 100644 --- a/agent-openai-agents-sdk-stateful/pyproject.toml +++ b/agent-openai-agents-sdk-stateful/pyproject.toml @@ -10,10 +10,12 @@ requires-python = ">=3.11" dependencies = [ "fastapi>=0.115.12", "uvicorn>=0.34.2", - "databricks-openai>=0.9.0", "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] @@ -23,6 +25,10 @@ build-backend = "hatchling.build" [tool.uv] prerelease = "allow" +[tool.hatch.metadata] +# Note: Required for direct git references. Delete when official databricks-openai[memory] release is available +allow-direct-references = true + [dependency-groups] dev = [ "hatchling>=1.27.0",