02-agent-framework

name	02-agent-framework
description	Use when building a custom agent with the OpenAI Agents SDK and MLflow ResponsesAgent. Covers Agent class, Runner, @function_tool, handoffs, streaming, ModelConfig, mlflow.models.set_model(), autolog, and manual tracing. Track A Step 2. Produces a customized, MLflow-compatible agent class ready for tools and deployment.
license	Apache-2.0
clients	["ide_cli","genie_code"]
bundle_resource	none
deploy_verb	none
deploy_note	Agent code (OpenAI Agents SDK + MLflow ResponsesAgent) — no deployed resource of its own; deployed later by Track A 07. Authored identically on both clients. On Genie Code write the agent source under the cloned repo root (`{REPO_ROOT}` = `state_file_root` from `skills/vibecoding-state`); see `skills/genie-code-environment`.
coverage	full
metadata	{"last_verified":"2026-06-05","volatility":"high","upstream_sources":[],"author":"prashanth-subrahmanyam","version":"3.1.0","domain":"genai-agents","pipeline_position":"A2","consumes":"mlflow_environment, running_local_agent","produces":"agent_class, responses_agent, customized_agent, lakebase_endpoint_uri, lakebase_cold_start_retry_policy, lakebase_pool_close_policy","grounded_in":"docs.databricks.com/aws/en/generative-ai/agent-framework/author-agent, openai.github.io/openai-agents-python, mlflow.org/docs/latest/genai/serving/responses-agent, github.com/databricks/app-templates/tree/main/agent-openai-agents-sdk"}
fields_read	["agent.system_prompt","agent.capabilities","agent.model"]

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 and uses the experiment/tracing infrastructure from F2. 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
A (legacy)	Class-based mlflow.pyfunc.ResponsesAgent wrapping the Agents SDK Runner	Databricks Model Serving endpoint (via databricks.agents.deploy())	06-appkit-serving-wiring

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 and migrate-agent-to-apps docs. This is the canonical Track A path used by the agent-openai-agents-sdk, agent-openai-agents-sdk-multiagent, and 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.

# 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:

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. For end-user feedback correlation via trace_id / client_request_id, see 04c — End-user feedback.

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:

# 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 with target=databricks_apps. For wiring the AppKit frontend, see 06d-appkit-agent-app-proxy.

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:

llm_endpoint: "databricks-claude-sonnet-4-6"
llm_api_base_url: null
llm_api_mode: "databricks_openai_compatible"

Agent construction must read:

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 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 template and A5: Lakebase Memory. 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:

# 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.

# 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):