simple-agent-scaffold

# Simple Agent Scaffold Shortest reliable path from "I have Genie Spaces" to "I have a deployed agent endpoint." Follows the [OpenAI MCP Tool Calling Agent](https://docs.databricks.com/aws/en/notebooks/source/generative-ai/openai-mcp-tool-calling-agent.html) notebook pattern verbatim. ``` Step 1 Step 2 Step 3 Step 4 Step 5 Write ──► Test locally ──► Log MLflow ──► Register UC ──► Deploy Serving agent.py │ ▼ AI Playground (default) │ ┌──────────┴──────────┐ ▼ ▼ 06-appkit-serving Go Further -wiring (optional) (optional) ``` ## When to Use - Creating a simple tool-calling agent with Genie Spaces - Workshop quick-start: zero to deployed endpoint - Prototyping an agent before adding evaluation, memory, or monitoring **Not for production-grade multi-agent systems.** Use the full `00-course-orchestrator` for evaluation pipelines, Lakebase memory, prompt registries, and multi-domain orchestration. --- ## Prerequisites | Requirement | How to verify | |---|---| | Databricks workspace with Model Serving | `databricks serving-endpoints list` returns without error | | At least one Genie Space **answering questions** | Verify end-to-end: open the space in the UI and ask a benchmark question, or via CLI — `databricks genie start-conversation <SPACE_ID> --json '{"content":"<your question>"}' --profile $PROFILE` then `databricks genie create-message`. If it returns a permission error on the underlying tables, fix Genie Space permissions **before proceeding**. | | Foundation Model API endpoint | `databricks serving-endpoints get databricks-claude-sonnet-4-6` (or your chosen model) | | Python packages | `pip install databricks-agents databricks-openai "mlflow[databricks]" mcp nest_asyncio uv` (the `[databricks]` extra is required on Azure for `azure-core`) | | Unity Catalog schema for the registered model | `databricks schemas get <catalog>.<schema>` | | MLflow experiment (optional but recommended) | Create one in the workspace UI or `mlflow.set_experiment()` | > **Critical:** do not skip the Genie Space test — a space that exists but can't answer questions produces an agent that only greets and never exercises the tool-calling path. > **Genie Space fails to answer questions? Three remediation options:** > > 1. **Use a different space.** Run `databricks genie list-spaces --profile $PROFILE` and test each one. Pick the first that returns data. > 2. **Create a new space.** If all existing spaces reference tables you can't access, create a new Genie Space pointing to tables in YOUR gold schema (Workspace → Genie → New Space). > 3. **Fix permissions on the existing space.** Ask a catalog admin to grant you `SELECT` on the tables the Genie Space references. Check the space's "Tables" tab to see which tables it uses. > > **Do not proceed to Step 1 until the Genie Space answers a data question successfully.** An agent wired to a broken Genie Space will deploy but fail every data query, wasting the entire build-test-deploy cycle. --- ## Decision Defaults | Decision | Default | Go Further | |---|---|---| | Agent framework | `MCPToolCallingAgent(ResponsesAgent)` per MCP notebook | — | | LLM client | `DatabricksOpenAI` (OpenAI SDK compatible) | — | | Genie access | `McpServerToolkit` with MCP server URLs, built per-request | `05-multi-agent-genie-orchestration` for Conversation API | | Authentication | OBO-first (`auth_policy`: `mcp.genie`+`sql` scopes) with best-effort system-SP fallback | system-SP only (`resources=`) if no per-user access is needed | | Streaming | Yes (`predict_stream` + `output_to_responses_items_stream`) | — | | Memory | None (stateless) | `03-lakebase-memory-patterns` | | Evaluation | Skip | `02-mlflow-genai-evaluation` | | Prompt management | Inline system prompt via `ModelConfig` | `04-prompt-registry-patterns` | | Deployment | Bundle job `agent_deploy_job` (`bundle deploy` -> `bundle run`) whose notebook calls `agents.deploy()` | `06-deployment-automation` for CI/CD | | Frontend | AI Playground (default) | `06-appkit-serving-wiring` for AppKit UI | --- ## Step 1: Write `agent.py` Copy the template and its config file to your project directory: ```bash cp references/agent-template.py agent.py cp references/agent-config.yaml agent-config.yaml ``` Open `agent-config.yaml` and resolve the three TODO blocks: 1. **`llm_endpoint`** — Verify the Foundation Model API endpoint name exists in your workspace. 2. **`system_prompt`** — Write domain-specific instructions for your agent. 3. **`genie_spaces`** — Replace each `TODO_REPLACE_WITH_SPACE_ID` with a real Genie Space ID. Add or remove entries as needed. Finding Genie Space IDs: ``` Workspace → Genie → open a space → the ID is in the URL: https://<workspace>.databricks.com/spaces/<SPACE_ID>/... ``` The MCP server URL format for Genie Spaces is: ``` {host}/api/2.0/mcp/genie/{space_id} ``` ### What the template contains The template is the notebook's `MCPToolCallingAgent` class with one addition: `ModelConfig` for parameterization via `agent-config.yaml`. The class structure is identical to the canonical notebook: | Method | Purpose | |---|---| | `_obo_client()` | Module-level helper: returns an OBO `WorkspaceClient` in Model Serving (via `ModelServingUserCredentials`), else falls back to the default client (system SP). Called per request. | | `__init__` | Stores `llm_endpoint` + `genie_spaces`; creates the `DatabricksOpenAI` LLM client (system-SP, identity-stable) | | `_build_tools` | Builds the `McpServerToolkit`(s) **per request** with the OBO `WorkspaceClient` and assembles `tools_dict` | | `execute_tool` | Traced with `@mlflow.trace(span_type=SpanType.TOOL)` | | `call_llm` | Traced with `@mlflow.trace(span_type=SpanType.LLM)`, streams via `chat.completions.create` | | `handle_tool_call` | Parses arguments, executes tool, returns `ResponsesAgentStreamEvent` | | `call_and_run_tools` | Iterative tool-calling loop with `max_iter=10` | | `predict` | Non-streaming entry point, delegates to `predict_stream` | | `predict_stream` | Streaming entry point: builds the per-request OBO client + tools, then converts `request.input` to messages | At the bottom: `mlflow.openai.autolog()` enables automatic tracing and `mlflow.models.set_model(AGENT)` binds the model for logging. ### Critical rules (from `01-responses-agent-patterns`) - **ResponsesAgent is mandatory** — not ChatAgent, not PythonModel. - **Never pass a `signature` parameter** to `log_model()` — MLflow auto-infers it. - **Use `input` key, not `messages`** — `{"input": [{"role": "user", "content": "..."}]}`. - **`nest_asyncio` is required** — MCP servers use async internally; `nest_asyncio.apply()` avoids event loop conflicts in notebook environments. - **Build the `McpServerToolkit` per request, NOT at module load.** A toolkit built at import time hard-binds whatever identity existed then (the system SP) and defeats OBO. `predict_stream()` constructs `_obo_client()` and the toolkit on every call so the Genie MCP call runs as the invoking user. See `references/obo-authentication.md`. **Gate:** `agent.py` exists with all TODOs in `agent-config.yaml` resolved. No `TODO_REPLACE` strings remain. --- ## Step 2: Test locally ### Deploy Steps 2–5 as a bundle job (canonical — same spine as Bronze/Silver/Gold) Workshop workspaces usually have no interactive cluster, and — more importantly — the agent endpoint should be created the same versioned way as every other artifact: **by running a bundle job**, not by an ad-hoc `agents.deploy()` in a loose notebook or a one-off `jobs submit`. Combine Steps 2–5 into one notebook (`deploy_agent`) and run it as a serverless **bundle job**: 1. Copy the templates into your bundle (keep them beside `agent.py`/`agent-config.yaml` so `bundle deploy` syncs them): - `references/agent-deploy-notebook.py` → `<bundle>/agents/deploy_agent.py` — the notebook-task body: Step 0 schema creation (direct SQL), Steps 2–5, Step 5b auto-grant, Step 5c checkpoint. - `references/agent_deploy_job.yml` → `<bundle>/resources/agent_deploy_job.yml`. 2. Wire the `variables:` block in `databricks.yml` (`catalog`, `agent_schema`, `agent_model_name`, `gold_schema`, `semantic_warehouse_id`, `genie_space_id`, `agents_folder_ws_path`) — see the YAML header. 3. Validate → deploy → run: ```bash databricks bundle validate -t dev databricks bundle deploy -t dev databricks bundle run -t dev agent_deploy_job ``` The job's notebook creates the UC agent schema with direct SQL (`CREATE SCHEMA IF NOT EXISTS` — the schema exception, **not** a bundle resource), then logs, registers, deploys, auto-grants the endpoint system SP (Step 5b), and writes `DEPLOY_CHECKPOINT.md` (Step 5c). Always use serverless (`environment_key`), never classic clusters (`new_cluster`) — workshop workspaces block classic clusters with `NETWORK_CONFIGURATION_FAILURE`. When run as a job the working directory is NOT the notebook's directory, which is why `model_config="agent-config.yaml"` is required in Step 3. > **Genie Code:** run `bundle deploy`/`bundle run` through `runDatabricksCli` **from the bundle-editor page** (open `<bundle>/databricks.yml` → "Open in bundle editor"); omit `--profile` (pre-authenticated) and do NOT `databricks sync` (deploy syncs the source). If a `bundle` command is blocked you are not on the bundle page — navigate there; do not fall back to `jobs submit` or a direct `agents.deploy()`. Keep the agent source under `{REPO_ROOT}`, never a bare relative path. See `skills/genie-code-environment` §2, §8. > **Fallback (one-off, IDE only, no bundle):** if you explicitly need a non-bundle run, [references/job-submission.md](references/job-submission.md) documents a standalone `databricks jobs submit` + polling loop. It is a convenience escape hatch — the bundle job above is the canonical, version-controlled path. ### Running Steps 2–5 interactively In a notebook or Python REPL, import and test both prediction paths: ```python # Cell 1: Restart Python if you edited agent.py # dbutils.library.restartPython() # (uncomment in Databricks notebook) # Cell 2: Test non-streaming from agent import AGENT result = AGENT.predict( {"input": [{"role": "user", "content": "What were total sales last month?"}]} ) print(result.model_dump(exclude_none=True)) ``` ```python # Cell 3: Test streaming for chunk in AGENT.predict_stream( {"input": [{"role": "user", "content": "What were total sales last month?"}]} ): print(chunk.model_dump(exclude_none=True)) ``` Replace the test question with something your Genie Space can actually answer. **Gate:** Both `predict` and `predict_stream` return valid, non-empty responses. MLflow traces are visible in the experiment UI (check the Traces tab). --- ## Step 3: Log with MLflow (dual `auth_policy` — OBO-first) Log the agent as code. This captures the `agent.py` file, its dependencies, and a dual `auth_policy` so the deployed endpoint supports BOTH the system SP (for the LLM and evaluation) and On-Behalf-Of the calling user (for the Genie MCP call). ```python import mlflow from agent import LLM_ENDPOINT_NAME from mlflow.models.auth_policy import AuthPolicy, SystemAuthPolicy, UserAuthPolicy from mlflow.models.resources import ( DatabricksGenieSpace, DatabricksServingEndpoint, DatabricksSQLWarehouse, ) from pkg_resources import get_distribution GENIE_SPACE_ID = "<GENIE_SPACE_ID>" WAREHOUSE_ID = "<WAREHOUSE_ID>" # the warehouse the Genie Space runs its SQL on # Dual policy: # SystemAuthPolicy.resources → system SP gets CAN_QUERY (LLM), Can Run (Genie), # CAN USE (warehouse) automatically — used by the LLM call and by evaluation. # UserAuthPolicy.api_scopes → forwards the caller's token for OBO. The Managed # MCP path needs "mcp.genie" (NOT "dashboards.genie", which is the # Conversation API) plus "sql". auth_policy = AuthPolicy( system_auth_policy=SystemAuthPolicy( resources=[ DatabricksServingEndpoint(endpoint_name=LLM_ENDPOINT_NAME), DatabricksGenieSpace(genie_space_id=GENIE_SPACE_ID), # one per space DatabricksSQLWarehouse(warehouse_id=WAREHOUSE_ID), # MANDATORY for Genie ] ), user_auth_policy=UserAuthPolicy(api_scopes=["mcp.genie", "sql"]), ) with mlflow.start_run(): logged_agent_info = mlflow.pyfunc.log_model( name="agent", python_model="agent.py", model_config="agent-config.yaml", # REQUIRED — see note below auth_policy=auth_policy, pip_requirements=[ f"mlflow[databricks]=={get_distribution('mlflow').version}", f"mcp=={get_distribution('mcp').version}", f"databricks-openai=={get_distribution('databricks-openai').version}", "databricks-ai-bridge", # REQUIRED for OBO (ModelServingUserCredentials) "databricks-sdk", ], ) ``` Key points: - **NO `signature` parameter.** ResponsesAgent auto-infers it. - **`python_model="agent.py"`** logs as "models from code" — MLflow loads the file, not a pickled object. - **`model_config="agent-config.yaml"`** is required — MLflow copies `agent.py` to a temp dir for validation where the yaml isn't present. Without this parameter you get `FileNotFoundError: Config file is not provided`. Fixing `__file__` / path tricks won't help; the parameter bypasses the file lookup. - **Use `mlflow[databricks]`, not bare `mlflow`.** On Azure the `[databricks]` extra ships `azure-core` and related storage SDKs required by `register_model()`. On AWS/GCP it adds harmless extras. Matches the `pip install` line in Prerequisites. - **`auth_policy` and `resources=` are mutually exclusive.** Use `auth_policy`; put every resource inside `SystemAuthPolicy.resources`. Passing both raises a parameter conflict. - **`DatabricksSQLWarehouse` is mandatory** — the Genie Space executes its SQL on that warehouse; omitting it is a common silent failure (Genie fails while the LLM works). - **`databricks-ai-bridge` MUST be in `pip_requirements`** — without it `ModelServingUserCredentials` cannot be imported in the serving container and OBO silently degrades to the system SP. - Add one `DatabricksGenieSpace(genie_space_id="...")` per Genie Space. > **Why OBO instead of granting the system SP?** A dual `auth_policy` deploys the endpoint as `EMBEDDED_AND_USER_CREDENTIALS`: the Genie MCP call runs as the **calling user**, so it respects their existing UC grants, row filters, and column masks with **zero** post-deploy grants. The system-SP fallback (Step 5b) is best-effort and only matters for true machine-to-machine callers. See `references/post-deploy-permissions.md` and `genai-agents/.../references/obo-authentication.md`. ### Pre-deployment validation Before registering, run the pre-deployment check: ```python mlflow.models.predict( model_uri=f"runs:/{logged_agent_info.run_id}/agent",