feat(examples): Databricks recipe on traceai-openai — model services and serving endpoints (TH-8316)

[nik13]

Summary

This PR adds a Databricks recipe at python/examples/databricks/ (Linear TH-8316, parent TH-8103). Customers who call their Databricks workspace with the official openai SDK get traces through the existing traceai-openai instrumentor. The PR adds no package and no instrumentor, and changes nothing outside python/examples/databricks/.

Databricks documents two OpenAI-SDK surfaces, kept separate here:

Surface	Base URL form	model
Unity Gateway model services	https://<workspace-host>/ai-gateway/mlflow/v1	model service name, e.g. system.ai.claude-sonnet-4-5
Model Serving endpoints	https://<workspace-host>/serving-endpoints	serving endpoint name

• src/app.py: calls register() and instruments with OpenAIInstrumentor before creating the client. The Databricks token (DATABRICKS_TOKEN) goes only to the OpenAI client; the Future AGI keys go only to the tracer.
  • DATABRICKS_BASE_URL is required, because there is no public default.
  • check_base_url() refuses placeholders (literal or percent-encoded) and Databricks' own sample host example.staging.cloud.databricks.com.
  • On Databricks hosts (*.cloud.databricks.com, *.azuredatabricks.net, *.gcp.databricks.com) it refuses anything but the two suffixes: the workspace root, /api/2.0/serving-endpoints, /serving-endpoints/<name>/invocations (a REST URL, not an SDK base URL), and the gateway's native /ai-gateway/gemini and /ai-gateway/anthropic paths. Upper-case and trailing-dot spellings are caught too. It never rewrites a URL.
  • Databricks hosts must use https (plain http would send the token in cleartext before any redirect), and a query string or fragment is refused. Host and path checks run on the URL as given; only the placeholder check decodes.
  • main() returns 2 before tracing on a refused URL or missing configuration.
• requirements.txt: the tested pins openai==3.24.0, traceAI-openai==0.1.10 and fi-instrumentation-otel==1.1.0.
• README.md: install, configure, run, code, what you see in Future AGI, and both surfaces. It also covers:
  • pricing: a price row keyed by a foundation-model id will not match an endpoint name;
  • embeddings, only when the endpoint task is embeddings;
  • DatabricksOpenAI, source-checked only;
  • privacy, limits, both test commands and the tested versions.
• tests/: a loopback contract on the shared harness Receiver (test(harness): shared OTLP harness for TH-8103 contract tests #203), a socket guard with a negative control, and a fake OpenAI server.

What the tests pin

Every test runs offline. In-process tests use httpx.MockTransport at a syntactic test host that is not real (dbc-00000000-0000.cloud.databricks.com). Subprocess tests use a fake on 127.0.0.1 under the guard.

Both surfaces:

• the request goes to exactly <base>/chat/completions with Authorization: Bearer <placeholder> and no Future AGI header;
• exactly one ChatCompletion LLM span, with the model the fixture returns;
• gen_ai.provider.name is openai;
• usage is copied when present and omitted, never set to 0, when absent;
• the token appears nowhere in the export.

Embeddings (gateway form): one CreateEmbeddingResponse span with kind EMBEDDING and provider openai.

• the model is in embedding.model_name, not gen_ai.request.model;
• input and total tokens are recorded, with no output tokens;
• the input text is exported by default. The README says the chat masking tests do not cover embedding text.

Behaviour:

• FI_HIDE_INPUTS/FI_HIDE_OUTPUTS remove a marker, with visible controls;
• a 401 produces openai.AuthenticationError, a span with ERROR status and an exception event, with no token recorded;
• the CLI subprocess runs chat and stream, honours a model override, and keeps the guard log empty.

Current traceai-openai behaviour, pinned and stated in the README (D-F5):

• streamed and failed spans have no gen_ai.request.model; the model is still inside gen_ai.request.parameters;
• a default stream exports no usage; with include_usage and a final usage chunk, all three counts are exported.

The shared fix is tracked separately (Linear TH-8402).

Tests (exact head b46c12e80232022f482ec4512b0ce0725104c6c5, clean tree, actual exit codes + JUnit)

Cell	Python	openai	traceAI packages	Result
1–4	3.10 / 3.11 / 3.12 / 3.13	3.24.0	repository source	140 passed each, rc 0
5	3.11	1.69.0 (floor)	repository source	140 passed, rc 0
6	3.11	3.24.0	published traceAI-openai==0.1.10 + fi-instrumentation-otel==1.1.0 (imported from site-packages, no repo source on the path)	140 passed, rc 0

No files outside python/examples/databricks/ changed, and the head did not move during the run. The credential scan of the recipe and the evidence found 0 hits. The first commit, 684708b, had the same 6/6 matrix with 122 tests.

Decisions (approved: Nikhil 2026-10-03 blanket)

• D1: a recipe on traceai-openai, not a traceai-databricks package.
• D-8316a: two labelled surfaces, never merged. There is no default host, and placeholders are refused.
• D-8316b: refuse Databricks paths that are not SDK base URLs, and never rewrite a URL.
  • Rejected alternative: building the URL from DATABRICKS_HOST, because it hides which surface the customer chose.
• D-8316c: DatabricksOpenAI is source-checked only. In databricks-openai 0.17.1 it subclasses openai.OpenAI without overriding request. It is not a dependency and is not tested.
• D-8316d: one embeddings case, pinning what the instrumentor exports.
• The vendor docs were re-checked on 2026-10-06:
  • https://docs.databricks.com/aws/en/machine-learning/model-serving/query-chat-models
  • https://docs.databricks.com/aws/en/generative-ai/external-models/external-models-tutorial

Not covered

• SQL ai_query, the MLflow Deployments SDK, Agent / Genie code, and the gateway's Gemini/Anthropic native APIs.
• DatabricksOpenAI execution, and direct REST invocation.
• No live workspace call was made (no token, no spend). Live model behaviour, fi-collector auth and storage, and the trace view are not exercised here.
• CI: none configured on this repository for this path.

Review status

• Review r1 (pr-reviewer t_fdf90ae4, claude-opus-5-5, at 684708b): CHANGES_REQUESTED on one P2.

  • R1 (P2): plain http:// was accepted for Databricks workspace hosts, so the token could be sent in cleartext.

  Fixed in b46c12e, test-first (20 failed before, 140 passed after):

  • R1: Databricks hosts must use https;
  • F1: query and fragment refusal, and checks on the URL as given;
  • F2: one actionable message for a missing or empty token;
  • a README sentence on https and query refusal.

  Kept as P3:

  • F3: the product name. The live Databricks page (2026-10-06) says "Unity Gateway", so the README keeps that;
  • F4: the README snippet imports the recipe helper;
  • F5: the embeddings fixture returns floats while the SDK requests base64;
  • F6: GovCloud and sovereign domains are treated as proxies.

• Verification r1 (pr-verifier t_048d2b3a, claude-opus-5-5, at b46c12e): VERIFIED, no blocking findings.

  • R1 and F2 fixed. F1 partly fixed: an empty trailing ? still passes. F3–F6 accepted as P3.
  • Remaining P3 follow-ups, not fixed in this PR (no third round for P3-only findings):
    • V1: the host is classified on the Unicode hostname. A workspace URL typed with U+3002/U+FF0E/U+FF61 full stops is treated as a proxy, and httpx IDNA-maps it back to the real host. Over http that would send the token in cleartext. Fix: refuse non-ASCII hosts, or classify on httpx.URL(url).host.
    • V2: refuse an empty trailing ?/# by checking the delimiters, not truthiness.
    • V3: refuse whitespace and control characters up front. A trailing newline passes urlsplit and then fails in httpx, with exit 1 after tracing has started.
    • V4: test strength. Add a raw≠decoded allowed URL, an upper-case/trailing-dot http case, and a response model distinct from the request.
  • Pre-existing shared-code note: embedding spans also export the whole response JSON, vectors included, under embedding.embeddings. FI_HIDE_EMBEDDING_VECTORS masks only embedding.vector keys. The README makes no vector-masking claim. This is tracked with the shared traceai-openai work.

Video demo

A narrated CLI walkthrough, 5:42, recorded at the verified head b46c12e: 1080p H.264/AAC, burned captions and 10 embedded chapters. The video, captions, transcript, chapters, preview and both media-verification notes are private attachments on the Linear issue TH-8316 (Future AGI workspace access required). Every command chapter runs ./run_demo.sh <chapter> live in a real terminal and ends at exit 0.

Time	Chapter
00:00	Problem: two OpenAI-compatible surfaces at your own workspace host, the Unity Gateway (/ai-gateway/mlflow/v1, model service name) and Model Serving (/serving-endpoints, endpoint name) (card)
00:44	head: PR #237 at b46c12e, the recipe files, the env variable names and the required URL's placeholder forms
01:13	run: chat, then a stream; what Future AGI and the provider each received; the stream has no gen_ai.request.model
02:01	host: both surfaces, with and without the trailing slash, through MockTransport (the dbc-…, adb-… and workspace-… hosts are syntactic test hosts), plus a gateway embeddings call (an EMBEDDING span with embedding.model_name); 0 external DNS
02:39	refusals: 53 URLs refused (placeholders, the sample host, wrong paths on all three clouds, then the review hardening: https only, no query or fragment), each exit 2 with 0 requests and 0 spans; a missing URL also exits 2
03:44	error: a 401 ERROR span with no gen_ai.request.model, the model in parameters, and no token
04:01	privacy: FI_HIDE_INPUTS / FI_HIDE_OUTPUTS versus a control; Databricks still receives the prompt
04:14	tests: the LIVE Python 3.11 suite (140 passed) and the RECORDED six-cell exact-head matrix
04:44	limits
05:03	Review status (card)

Two waits are shortened in the edit, and both are labelled on screen: about 74 s of the refusals run and about 62 s of the live pytest run. Command output itself is not edited.

Not shown: a live Databricks call, workspace authentication, a running fi-collector, storage or the trace view. All runs use loopback fakes, and the narration says so.