-
Notifications
You must be signed in to change notification settings - Fork 154
Add agent-openai-basic template #262
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
elainemwang
wants to merge
12
commits into
databricks:main
Choose a base branch
from
elainemwang:add-agent-openai-local-template
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
Show all changes
12 commits
Select commit
Hold shift + click to select a range
bbf59ff
Add agent-openai-basic template
elainemwang 76a3505
Fix README tracing gate + drop duplicate durability line
elainemwang e146e08
Extract tracing into agent/tracing.py; drop git-version tracking
elainemwang b9a9925
Move set_trace_processors into tracing.configure()
elainemwang f7b6704
Use McpServer.from_uc_function in mcps.py example (Databricks OAuth)
elainemwang d54a04b
Group SDK-migration plumbing under agent/mason/; read session db path…
elainemwang 86ff625
Move memory into agent/mason/memory.py with call-time memory_tools()
elainemwang 898aa14
Snapshot tracing state in configure(); drop E402 from start_server
elainemwang cfdb40c
Satisfy ruff isort in agent.py (comments on own lines)
elainemwang cc67a33
Move MCP connect() lifecycle into agent/mason/mcp_runtime.py
elainemwang 3cce7cf
Fix manifest: postgres resource backs LRAS store, not session history
elainemwang dacd8ff
Accept any valid MLflow tracing config combo
elainemwang File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,31 @@ | ||
| # Copy to .env for local development: cp .env.example .env | ||
| # | ||
| # The only required setting for local dev is a Databricks auth profile (used to call the model). | ||
| # Everything else — Lakebase durability, MLflow tracing — is optional and off by default. | ||
|
|
||
| # Databricks auth profile (from `databricks auth profiles`). Used to call the model endpoint. | ||
| DATABRICKS_CONFIG_PROFILE=DEFAULT | ||
| # Or use explicit host/token instead of a profile: | ||
| # DATABRICKS_HOST=https://<your-workspace>.databricks.com | ||
| # DATABRICKS_TOKEN=dapi... | ||
|
|
||
| # --- Optional: MLflow tracing --- | ||
| # Leave UNSET to skip tracing (local dev). To enable, set a destination AND an experiment (either | ||
| # form of each works): destination = MLFLOW_TRACKING_URI or MLFLOW_TRACING_DESTINATION; | ||
| # experiment = MLFLOW_EXPERIMENT_ID or MLFLOW_EXPERIMENT_NAME. | ||
| # MLFLOW_TRACKING_URI="databricks" | ||
| # MLFLOW_EXPERIMENT_ID= | ||
| # MLFLOW_EXPERIMENT_NAME= | ||
|
|
||
| # --- Optional: durable conversation history (managed session store) --- | ||
| # Leave UNSET to keep the local SQLite session store. Set to a managed session store name to | ||
| # persist the transcript to its agents/v1 items API (durable, shared across replicas). | ||
| # AGENT_SESSION_STORE=my-agent-sessions | ||
|
|
||
| # --- Optional: long-running background mode + crash recovery (Lakebase) --- | ||
| # Leave UNSET to serve in-request. Set the Lakebase endpoint to enable durable background mode. | ||
| # LAKEBASE_AUTOSCALING_ENDPOINT= | ||
|
|
||
| # --- Optional: local SQLite session store path (when AGENT_SESSION_STORE is unset) --- | ||
| # Defaults to a file so history survives restarts. Set ":memory:" for ephemeral storage. | ||
| # LOCAL_SESSION_DB_PATH=local_agent_sessions.db |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,62 @@ | ||
| # Deploy this Databricks App from GitHub Actions. | ||
| # | ||
| # Prerequisites: | ||
| # 1. Create a Databricks service principal and a GitHub Actions workload | ||
| # identity federation policy for it. See: | ||
| # https://docs.databricks.com/aws/en/dev-tools/auth/provider-github | ||
| # 2. In your GitHub repo, create a "prod" Environment | ||
| # (Settings -> Environments -> New environment). On that environment, | ||
| # add two Variables: | ||
| # - DATABRICKS_HOST e.g. https://my-workspace.cloud.databricks.com | ||
| # - DATABRICKS_CLIENT_ID the service principal application ID (UUID) | ||
| # 3. Grant the service principal CAN_MANAGE on the Databricks App. | ||
| # 4. In databricks.yml, configure targets.prod.workspace.host (or | ||
| # workspace.root_path) so production deploys have an explicit path. | ||
| # 5. If your agent uses Lakebase memory (advanced templates), run the | ||
| # grant script once before the first CI deploy. See: | ||
| # https://docs.databricks.com/aws/en/generative-ai/agent-framework/cicd-agent-app#lakebase | ||
| # | ||
| # Trigger manually from the Actions tab, or uncomment the `push` trigger | ||
| # below to deploy on every push to main. | ||
| # | ||
| # This file is synced from .scripts/source/.github/workflows/deploy.yml — | ||
| # edit the source, then run `uv run python .scripts/sync-scripts.py`. | ||
|
|
||
| name: Deploy to Databricks Apps | ||
|
|
||
| on: | ||
| workflow_dispatch: | ||
| # push: | ||
| # branches: [main] | ||
|
|
||
| permissions: | ||
| id-token: write # required for OIDC federation | ||
| contents: read | ||
|
|
||
| jobs: | ||
| deploy: | ||
| name: Deploy | ||
| runs-on: ubuntu-latest | ||
| environment: prod | ||
| env: | ||
| DATABRICKS_AUTH_TYPE: github-oidc | ||
| DATABRICKS_HOST: ${{ vars.DATABRICKS_HOST }} | ||
| DATABRICKS_CLIENT_ID: ${{ vars.DATABRICKS_CLIENT_ID }} | ||
|
|
||
| steps: | ||
| - name: Checkout | ||
| uses: actions/checkout@v4 | ||
|
|
||
| - name: Install Databricks CLI | ||
| uses: databricks/setup-cli@main | ||
|
|
||
| - name: Validate bundle | ||
| run: databricks bundle validate --target prod | ||
|
|
||
| - name: Deploy bundle | ||
| run: databricks bundle deploy --target prod | ||
|
|
||
| # `bundle deploy` uploads files and configures resources, but does NOT | ||
| # restart the app. Without `bundle run`, the app keeps serving old code. | ||
| - name: Start / restart app | ||
| run: databricks bundle run agent_openai_local --target prod | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,195 @@ | ||
| # 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 | ||
|
|
||
| .DS_* | ||
|
|
||
| # Databricks / MLflow local artifacts | ||
| **/mlruns/ | ||
| mlflow.db | ||
| **/.databricks | ||
| .claude/ | ||
|
|
||
| # Environment files | ||
| **/.env | ||
| **/.env.local | ||
|
|
||
| # Local (Lakebase-free) SQLite session store | ||
| local_agent_sessions.db | ||
| local_agent_sessions.db-* | ||
| *.db-journal |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,72 @@ | ||
| # Agent Development Guide | ||
|
|
||
| A lean OpenAI Agents SDK agent backend for Databricks Apps. Local-first: runs with no database and | ||
| no setup beyond a Databricks auth profile. Lakebase durability and MLflow tracing are optional. | ||
|
|
||
| See `README.md` for the full run / deploy / client-contract docs. This file is the quick map for | ||
| making changes. | ||
|
|
||
| ## Run it | ||
|
|
||
| ```bash | ||
| cp .env.example .env # set DATABRICKS_CONFIG_PROFILE=<your-profile> | ||
| uv run start-server # http://localhost:8000 | ||
| ``` | ||
|
|
||
| No database needed — sessions use a local SQLite file by default. | ||
|
|
||
| ## Where things live | ||
|
|
||
| | You want to… | Edit | | ||
| | --- | --- | | ||
| | Change model / instructions | `agent/agent.py` (`create_agent`) | | ||
| | Add a function tool | new `*.py` in `agent/tools/` with a `@function_tool` function (auto-collected) | | ||
| | Add an MCP server | append one to `build_mcp_servers()` in `agent/mcps.py` (e.g. `McpServer.from_uc_function(...)` from `databricks_openai.agents`) | | ||
| | Change how a request maps to a run | `agent/agent.py` (`@invoke` / `@stream` handlers) | | ||
| | Change the session store | `agent/mason/session_store.py` | | ||
| | Server / durability wiring | `server/start_server.py` (rarely needed) | | ||
| | Add a test | `tests/` (hermetic; gate model calls on a workspace profile — see `test_agent.py`) | | ||
|
|
||
| `agent/mason/` holds plumbing (session store, tracing, MCP connection lifecycle, wire translation) | ||
| slated to move into Databricks SDKs — grouped so that migration is localized. You rarely edit it; | ||
| build the agent in `agent/agent.py`, `agent/tools/`, and `agent/mcps.py`. | ||
|
|
||
| ## How tools register | ||
|
|
||
| `agent/tools/all_tools()` auto-imports every module in the package and collects every | ||
| `@function_tool`-decorated `FunctionTool` it finds. So a tool registers just by existing in a file | ||
| there — `create_agent()` calls `all_tools()`. **Do not** edit `agent/agent.py` to add a tool — just | ||
| add a file to `agent/tools/`. | ||
|
|
||
| ## Sessions & durability | ||
|
|
||
| - Default: `agent/mason/session_store.py`'s `create_session()` returns a local `SQLiteSession` — no database. | ||
| - **Two independent durable stores, each with its own env var:** | ||
| - `AGENT_SESSION_STORE` (a managed session store name) → `mason/session_store.py` persists the | ||
| transcript to that store's `agents/v1` items API (durable conversation history). | ||
| - `LAKEBASE_AUTOSCALING_ENDPOINT` → `start_server.py` passes it into `LongRunningAgentServer` for | ||
| its durable server store (background mode + crash recovery). | ||
| - Enable either/both/neither. Fully durable = set both. | ||
|
|
||
| ## MLflow tracing | ||
|
|
||
| Optional. Enabled when MLflow has both a destination (`MLFLOW_TRACKING_URI` or | ||
| `MLFLOW_TRACING_DESTINATION`) and an experiment (`MLFLOW_EXPERIMENT_ID` or `MLFLOW_EXPERIMENT_NAME`); | ||
| `mason/tracing.py` gates on any valid combo. Leave either half unset to skip. See the README's | ||
| tracing section. | ||
|
|
||
| ## Quick commands | ||
|
|
||
| | Task | Command | | ||
| | --- | --- | | ||
| | Run locally | `uv run start-server` | | ||
| | Run via CLI local App runner | `databricks apps run-local --prepare-environment -p <profile>` | | ||
| | Test | `uv run pytest` (hermetic; live model test runs only with a profile) | | ||
| | Deploy | `databricks apps deploy agent-openai-basic --source-code-path <path>` | | ||
|
|
||
| ## Notes for maintainers | ||
|
|
||
| - `agent/mason/wire/` is OpenAI-Agents-SDK-specific (inbound request→SDK input; outbound SDK events→wire, | ||
| surfacing tool-call outputs the Responses raw event stream omits). | ||
| - `mcp<2` is pinned in `pyproject.toml` because `databricks-openai` imports a symbol removed in | ||
| `mcp` 2.0; remove the pin when `databricks-openai` supports `mcp>=2`. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1 @@ | ||
| @AGENTS.md |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.