Skip to content

(WIP) Migrate DAIS agents to Apps AI Gateway#70

Open
djliden wants to merge 3 commits into
demo/dais-2026from
codex/apps-ai-gateway-agents
Open

(WIP) Migrate DAIS agents to Apps AI Gateway#70
djliden wants to merge 3 commits into
demo/dais-2026from
codex/apps-ai-gateway-agents

Conversation

@djliden

@djliden djliden commented Jun 8, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Migrates the DAIS refund and complaint custom agents from MLflow model registration + agents.deploy to Databricks Apps served with MLflow AgentServer /responses.
  • Treats LLM_MODEL as a Unity AI Gateway endpoint name; refund uses ChatDatabricks(model=LLM_MODEL, use_ai_gateway=True), complaint configures DSPy through ${DATABRICKS_HOST}/ai-gateway/mlflow/v1.
  • Updates stream jobs, eval notebooks, and the ops dashboard to call Apps via app URLs /responses, while leaving support_request_agent, MAS, and Knowledge Assistant serving endpoints unchanged.
  • Adds a shared app helper for deterministic app names, notebook-token exchange, /responses calls, Gateway probing, and response text extraction.

Validation

  • python3 -m py_compile utils/agent_app_client.py apps/refund-agent/agent.py apps/refund-agent/start_server.py apps/complaint-agent/agent.py apps/complaint-agent/start_server.py apps/caspers-ops-dashboard/app/main.py
  • Parsed edited notebook code cells with ast.parse
  • Static scan: no refund/complaint migrated paths use agents.deploy, DatabricksServingEndpoint, custom-agent endpoint params, or /serving-endpoints LLM invocation routes. Remaining /serving-endpoints references are MAS/KA/support/terminal paths intentionally left unchanged.
  • databricks bundle validate -t all --profile sandbox
  • databricks bundle deploy -t all --profile sandbox
  • Verified sandbox sync contains apps/refund-agent and apps/complaint-agent; verified .git remains excluded

Not run

  • Did not run the full initializer DAG or smoke-call the deployed Apps after stages. The bundle source is deployed; the agent Apps are created when the refunder_agent and complaint_agent stages run.

@djliden djliden changed the title Migrate DAIS agents to Apps AI Gateway (WIP) Migrate DAIS agents to Apps AI Gateway Jun 8, 2026
djliden added 2 commits June 10, 2026 07:59
Reconciles the app-based agent migration on this branch with the
gateway-routing intent from demo/dais-2026 (fe14776): the Refund and
Complaint agents stay deployed as Databricks Apps AND route every LLM
call through Unity AI Gateway. Verified end-to-end on the sandbox `all`
target (25/25 tasks; streams produced 75k+ rows; live + 8-way concurrent
calls all 200).

Gateway routing (always-on, no model-serving fallback)
- New AI_GATEWAY_ENDPOINT_NAME job param (all target), distinct from
  LLM_MODEL: LLM_MODEL stays the FM name for generators/support; the
  gateway endpoint name is sent verbatim as the request `model` to
  <host>/ai-gateway/mlflow/v1. Default databricks-claude-sonnet-4-5 so
  existing deploys keep working; override per governed endpoint.
- Both apps/*/agent.py read GATEWAY_ENDPOINT_NAME from the env (with a
  transitional LLM_MODEL fallback); static app.yaml + the deploy stages
  inject it.

Correctness
- Complaint agent: configure DSPy once at import, apply a fresh-token LM
  per request via dspy.context() instead of dspy.configure() — avoids the
  AgentServer worker-thread thread-affinity error under concurrency.
- Refund agent's ChatDatabricks(use_ai_gateway=True) already refreshes the
  bearer per request via DatabricksOpenAI's BearerAuth (no rebuild needed).

App-name + warehouse drift fixes (--var catalog vs --params CATALOG)
- New OPS_WAREHOUSE_NAME, REFUND_AGENT_APP_NAME, COMPLAINT_AGENT_APP_NAME
  params baked from ${var.catalog} at deploy time.
- utils/agent_app_client.resolve_agent_app_name() prefers the baked param,
  re-sanitises, falls back to deriving; threaded through the stream jobs,
  eval stages, agent stages, and ops dashboard.

Runbook + docs
- SETUP.ipynb step 5 (create gateway endpoint; CAN_QUERY to each agent
  App SP, not account users) and MLflow.ipynb gateway demo beats, rewritten
  for the App architecture.
- README + AGENTS: --var vs --params drift table and the gateway/LLM_MODEL
  distinction.

Also folds in the app-server simplification already in the working tree
(start_server.py uses AgentServer's native /responses; ops dashboard
output extraction tightened) — both exercised by the end-to-end run.

Co-authored-by: Isaac
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant