02-agent-framework

# Agent Framework and ResponsesAgent Build your agent with the **OpenAI Agents SDK** and wrap it with **MLflow ResponsesAgent** for Databricks-compatible serving, tracing, and evaluation. ## When to Use - You have a running local agent from Track A Step 1 and need to customize its behavior. - You need to understand the two-layer architecture: **Agent Framework** (OpenAI Agents SDK) on top of **MLflow ResponsesAgent** (serving contract). - You want to add custom instructions, swap models, configure streaming, or set up handoffs. - You need to enable MLflow tracing for observability. **Foundation prerequisite:** This step consumes `mlflow_environment` from [F1](../../../foundation/01-mlflow-genai-foundation/SKILL.md) and uses the experiment/tracing infrastructure from [F2](../../../foundation/02-experiment-tracing-and-uc-storage/SKILL.md). Both must be complete. --- ## Two-Layer Architecture ``` ┌─────────────────────────────────────────────┐ │ OpenAI Agents SDK (Agent Framework) │ │ ├─ Agent class (instructions, model, tools)│ │ ├─ Runner.run / Runner.run_streamed │ │ ├─ @function_tool definitions │ │ └─ Handoffs for multi-agent routing │ ├─────────────────────────────────────────────┤ │ MLflow Host Runtime (Serving Contract) │ │ ├─ Option B (canonical): @invoke / @stream │ │ │ module-level async handlers, hosted │ │ │ on Databricks Apps │ │ └─ Option A (legacy): ResponsesAgent class │ │ hosted on Databricks Model Serving │ └─────────────────────────────────────────────┘ ``` The **OpenAI Agents SDK** handles conversation management, tool orchestration, and multi-agent routing. The **MLflow host runtime** wraps it for serving, tracing, and evaluation. Pick the host runtime by deployment target (next section). --- ## Deployment Target Decides the Wrapping Pattern The Track A agent runs under one of two MLflow wrapping patterns, chosen by **where you plan to deploy**. Pick the deployment target first; the authoring pattern and the downstream AppKit wiring skill follow automatically. | Option | Authoring pattern | Deployment target | AppKit wiring skill | |--------|-------------------|-------------------|---------------------| | **B (canonical)** | Module-level `@mlflow.genai.agent_server.invoke` / `@mlflow.genai.agent_server.stream` async handlers | Databricks **Apps** (as its own App, via `databricks apps deploy`) | [`06d-appkit-agent-app-proxy`](../../../../apps_lakebase/skills/06d-appkit-agent-app-proxy/SKILL.md) | | **A (legacy)** | Class-based `mlflow.pyfunc.ResponsesAgent` wrapping the Agents SDK `Runner` | Databricks **Model Serving** endpoint (via `databricks.agents.deploy()`) | [`06-appkit-serving-wiring`](../../../../apps_lakebase/skills/06-appkit-serving-wiring/SKILL.md) | ### When to pick which - **Pick Option B for any new use case.** Databricks recommends Apps as the default agent host runtime — per the 2026 [author-agent](https://docs.databricks.com/aws/en/generative-ai/agent-framework/author-agent) and [migrate-agent-to-apps](https://docs.databricks.com/aws/en/generative-ai/agent-framework/migrate-agent-to-apps) docs. This is the canonical Track A path used by the [`agent-openai-agents-sdk`](https://github.com/databricks/app-templates/tree/main/agent-openai-agents-sdk), [`agent-openai-agents-sdk-multiagent`](https://github.com/databricks/app-templates/tree/main/agent-openai-agents-sdk-multiagent), and [`agent-openai-advanced`](https://github.com/databricks/app-templates/tree/main/agent-openai-advanced) templates. The canonical SkyLoyalty walkthrough uses Option B. - **Pick Option A only when** you have a hard requirement for the legacy path — e.g. an existing Model Serving endpoint, a multi-agent Genie notebook packaged as `pyfunc`, a Knowledge Assistant wrapper that must register through UC Model Registry with `@champion`/`@production` aliases on the critical deploy path. - **Both options** still benefit from MLflow 3 tracing, `agent-evaluate`, and `mlflow.openai.autolog()`. The difference is the serving contract (`/invocations` on an App URL vs. `/serving-endpoints/:name/invocations`) and the auth model (explicit `x-forwarded-access-token` forwarding vs. the Serving plugin's `.asUser(req)`). --- ## Option B — Module-level `@invoke` / `@stream` (canonical) The Apps host runtime invokes module-level **async** functions decorated with `mlflow.genai.agent_server.invoke` and `mlflow.genai.agent_server.stream`. No class, no `mlflow.models.set_model()`, no `predict` / `predict_stream`. End-user identity flows via `http_request` → `get_user_workspace_client(http_request)`, so every tool call runs **on-behalf-of** the end user when the AppKit proxy forwards `x-forwarded-access-token`. ```python # agent.py — Option B canonical layout (mirrors agent-openai-agents-sdk template) import os import mlflow from agents import Agent, Runner, function_tool from mlflow.genai import agent_server from mlflow.models import ModelConfig from databricks_app.utils import get_user_workspace_client mlflow.openai.autolog() config = ModelConfig(development_config="config.yml") @function_tool def get_current_time() -> str: """Get the current ISO timestamp.""" from datetime import datetime return datetime.now().isoformat() def build_agent(ws) -> Agent: """Build the agent for this request, optionally per-user-scoped via ws.""" return Agent( name="loyalty-assistant", instructions="You are a helpful loyalty program assistant.", model=config.get("llm_endpoint"), tools=[get_current_time], ) @agent_server.invoke async def handle_invoke(request: dict, http_request) -> dict: ws = get_user_workspace_client(http_request) agent = build_agent(ws) result = await Runner.run(agent, request["input"]) return {"output": result.final_output} @agent_server.stream async def handle_stream(request: dict, http_request): ws = get_user_workspace_client(http_request) agent = build_agent(ws) async for event in Runner.run_streamed(agent, request["input"]): if event.type == "raw_response_event": yield event.data ``` ### Trace context: user, session, environment, request id The `@invoke` / `@stream` handler is the canonical call site for attaching trace context — user, session, environment override, and `client_request_id` for end-user feedback correlation. Set them all on the trace root **before** the `Runner.run` call so they're present even if the agent fails partway: ```python import os import mlflow @agent_server.invoke async def handle_invoke(request: dict, http_request) -> dict: ws = get_user_workspace_client(http_request) user_id = ws.current_user.me().user_name session_id = request.get("session_id") or http_request.headers.get("x-session-id", "anon") mlflow.update_current_trace( client_request_id=request.get("client_request_id"), # for /feedback correlation metadata={ "mlflow.trace.user": user_id, "mlflow.trace.session": session_id, "mlflow.source.type": os.getenv("APP_ENVIRONMENT", "development"), "agent_version": os.getenv("AGENT_VERSION", "unknown"), }, ) agent = build_agent(ws) result = await Runner.run(agent, request["input"]) return {"output": result.final_output, "trace_id": mlflow.get_current_active_span().trace_id} ``` The reserved metadata fields (`mlflow.trace.user` / `mlflow.trace.session`) are **immutable** post-log and light up the Trace UI's user / session facets. The `APP_ENVIRONMENT` env var is set in `app.yaml` per deployment. For the full pattern (auto-populated metadata, custom deployment metadata, search-by-metadata examples), see [F2c — Trace context and environments](../../../foundation/02c-trace-context-and-environments/SKILL.md). For end-user feedback correlation via `trace_id` / `client_request_id`, see [04c — End-user feedback](../../../sdlc/04c-end-user-feedback/SKILL.md). ### `app.yaml` for Option B The Apps host serves the module directly — no `start_server.py` is required. Point `command` at the MLflow agent server entrypoint and the agent module: ```yaml # app.yaml command: - mlflow - genai - agent-server - serve - --module - agent env: - name: LLM_ENDPOINT value: databricks-claude-sonnet-4-6 # F2c: app code reads APP_ENVIRONMENT and overrides mlflow.source.type via metadata. - name: APP_ENVIRONMENT value: production ``` For deployment, see [`07-deploy-and-query`](../07-deploy-and-query/SKILL.md) with `target=databricks_apps`. For wiring the AppKit frontend, see [`06d-appkit-agent-app-proxy`](../../../../apps_lakebase/skills/06d-appkit-agent-app-proxy/SKILL.md). ### Runtime model route from Tool Plan Track A agents consume `docs/agent_tool_plan.yaml.runtime_config.llm` through `ModelConfig`, never by hardcoding model names in Python. Minimum `config.yml` shape: ```yaml llm_endpoint: "databricks-claude-sonnet-4-6" llm_api_base_url: null llm_api_mode: "databricks_openai_compatible" ``` Agent construction must read: ```python config = ModelConfig(development_config="config.yml") agent = Agent( name="loyalty-assistant", instructions="...", model=config.get("llm_endpoint"), ) ``` If `llm_api_base_url` is non-null in a future Gateway route, the client factory may pass it to the OpenAI-compatible client. The core workshop does not require that path. ### Multi-agent variant For a triage / handoffs pattern, see the [`agent-openai-agents-sdk-multiagent`](https://github.com/databricks/app-templates/tree/main/agent-openai-agents-sdk-multiagent) template and the Handoffs section below — both apply identically inside `build_agent()` under Option B. ### Advanced variant (Lakebase memory + custom tools) For long-term memory backed by Lakebase, see the [`agent-openai-advanced`](https://github.com/databricks/app-templates/tree/main/agent-openai-advanced) template and [A5: Lakebase Memory](../05-lakebase-memory/SKILL.md). Memory state is loaded inside `handle_invoke` / `handle_stream` per request. ### Lakebase cold-start retry policy and lazy init The Track A agent owns the Lakebase **client policy** contract — the endpoint URI it connects to, the retry behavior on cold-start failures that Lakebase Autoscaling raises during the first request after idle, and the pool-close policy on graceful shutdown. Capture these fields in state and read them at runtime; **do not** hard-code retry counts or error classes in the agent module. Capture in state: ```yaml # state://Track_A.a2_agent_framework lakebase_endpoint_uri: "postgresql://<host>:<port>/<database>" # resolved from databricks_app.utils.get_user_workspace_client at runtime lakebase_cold_start_retry_policy: retry_on: - AdminShutdown # psycopg.errors.AdminShutdown — Lakebase autoscaler closed the prior connection - psycopg_pool.PoolClosed # the pool was closed mid-request and must be re-opened max_attempts: 3 initial_backoff_seconds: 5 # exponential: 5s, 10s, 20s lakebase_pool_close_policy: "close_on_app_shutdown" # allowed values: # - close_on_app_shutdown (default; close pool in `@app.on_event("shutdown")` or equivalent) # - close_on_request_end (only for short-lived dev runs — produces excessive cold starts in prod) # - never (only for diagnostic harnesses; leaks connections) ``` The retry-on classes (`psycopg.errors.AdminShutdown`, `psycopg_pool.PoolClosed`) are the canonical first-request-after-idle failure modes documented in the retrospective. The agent's connection helper MUST catch exactly these classes (not a bare `Exception`) and back off per `initial_backoff_seconds * 2 ** attempt` up to `max_attempts`. `preflight_check_registry.lakebase_cold_start_retry_policy_present` gates downstream prompt roles until `agent.retry_policy.lakebase_cold_start_max_retries >= 1` is recorded in state. #### Lazy init: Lakebase config MUST NOT execute at import time Connecting to Lakebase from module-level code (top of `agent.py`) is a recurring failure: imports happen during `mlflow models log` and during `databricks apps deploy` on machines that have no path to the Lakebase endpoint, blowing up the deploy. **All Lakebase-related setup runs inside `handle_invoke` / `handle_stream` (or a request-scoped factory they call), never at module import.** ```python # agent.py — Lakebase lazy init pattern (Option B) import os import asyncio import mlflow from agents import Agent, Runner from mlflow.genai import agent_server from databricks_app.utils import get_user_workspace_client mlflow.openai.autolog() # DO NOT instantiate the Lakebase pool here — that runs at import time # and breaks `mlflow models log` and `databricks apps deploy`. _LAKEBASE_POOL = None _LAKEBASE_LOCK = asyncio.Lock() async def _get_lakebase_pool(ws): """Build the pool the first time a request needs it; cache for the lifetime of the process.""" global _LAKEBASE_POOL if _LAKEBASE_POOL is not None: return _LAKEBASE_POOL async with _LAKEBASE_LOCK: if _LAKEBASE_POOL is None: from psycopg_pool import AsyncConnectionPool uri = os.environ["LAKEBASE_ENDPOINT_URI"] # from state://lakebase_endpoint_uri _LAKEBASE_POOL = AsyncConnectionPool(uri, open=False) await _LAKEBASE_POOL.open() return _LAKEBASE_POOL async def _with_cold_start_retry(coro_factory):