feat: expose MLflow chat-completions-only foundation models in Pi and OpenCode#216
Closed
dgokeeffe wants to merge 9 commits into
Closed
feat: expose MLflow chat-completions-only foundation models in Pi and OpenCode#216dgokeeffe wants to merge 9 commits into
dgokeeffe wants to merge 9 commits into
Conversation
…ks-mlflow provider Pi only registered three providers (claude/openai/gemini), each pinned to a family-specific gateway route. Foundation models served solely through mlflow/v1/chat/completions (Llama, Qwen, GLM, inkling, ...) matched none of the discovery filters and were invisible in Pi's model picker. Add a fourth `databricks-mlflow` provider (api: openai-completions -> /ai-gateway/mlflow/v1) plus `discover_oss_models`, which lists endpoints exposing mlflow/v1/chat/completions but none of the native api_types already served by the other providers (dedup) and excludes embeddings. Verified end-to-end: glm-5-2 answers through the new provider. inkling is selectable and dispatched but blocked by an upstream Pi streaming bug (its SSE omits finish_reason on natural completion) tracked separately.
databricks-inkling streams a spec-incomplete response: on natural completion the terminal chunk omits finish_reason and the stream ends with only data: [DONE]. Pi's openai-completions parser rejects this with 'Stream ended without finish_reason', so inkling was selectable but unusable. Add a stdlib-only SSE-repair proxy (agents/_mlflow_proxy.py) started in Pi's launch path. It forwards requests verbatim and passes response bytes through unchanged, injecting a synthetic finish_reason:"stop" chunk before [DONE] only when upstream omitted it. Well-behaved models (glm, llama, qwen) never trigger the injection. The mlflow provider baseUrl is rewired to the proxy in state before the initial config write and every token-refresh rewrite. Verified end-to-end: 'ucode pi' -> databricks-inkling now answers through Pi (6*7 -> 42). Workaround for the gateway bug tracked in databricks#215; removable once the gateway emits finish_reason.
New databricks-mlflow Pi provider + inkling SSE-repair proxy (databricks#216).
Two defects in the SSE-repair proxy surfaced as recurring 'Stream ended without finish_reason' errors across sessions: 1. Stale proxy port persisted to config. Each launch picks a fresh ephemeral port; the old one was baked into state/models.json. _start_oss_proxy then derived its forward origin FROM that persisted value, so a later session would proxy a dead port. Now the gateway origin is always re-derived from the workspace via build_pi_base_urls, never the persisted oss URL. 2. Proxy only repaired the clean [DONE] path. A gateway 429 mid-stream (inkling is throttled as a reasoning model) or any upstream drop ended the stream with no finish_reason and no [DONE], which Pi rejects. The proxy now synthesizes the terminator on ANY stream end after data — clean or abnormal — and tolerates upstream OSError. Verified: inkling answers through Pi even with a stale port pre-seeded in state (proxy rebinds to a fresh port).
The 0.2.0 version bump ran `uv lock` on a machine configured with the internal Databricks PyPI proxy (~/.config/uv/uv.toml), which rewrote all 323 package URLs from files.pythonhosted.org to pypi-proxy.cloud.databricks.com — breaking resolution for CI and contributors without proxy access. Restore the lock to main's public-URL form and apply only the intended one-line ucode version bump (0.1.0 -> 0.2.0).
…icks-oss provider OpenCode registered only databricks-anthropic and databricks-google, so mlflow-only foundation models (inkling, glm, llama, qwen, gpt-oss, kimi, gemma) were unselectable — the same discovery gap fixed for Pi. Add a databricks-oss provider (@ai-sdk/openai in chat-completions mode -> /ai-gateway/mlflow/v1), fed by the existing oss_models discovery. Unlike the codex provider it does NOT set useResponsesApi; unlike Pi it needs no SSE-repair proxy — verified @ai-sdk/openai tolerates the finish_reason the gateway omits for inkling. default_model falls back to oss last. Verified end-to-end: opencode runs databricks-inkling via ucode-generated config (6*7 -> 42, exit 0), no proxy.
dgokeeffe
changed the title
Adds integration coverage for configure_shared_state: mlflow-only (oss) models are fetched and routed into state['oss_models'] (Pi) and opencode_models['oss'] (OpenCode) when those tools are requested, and NOT fetched for unrelated tools. This discovery->state seam was previously only exercised via mocked configure_shared_state.
…capabilities
OSS (mlflow-only) models were configured with just an id/UA header, so clients
fell back to defaults — Pi showed GLM-5.2 as 128K context with no thinking,
though the gateway reports 1M context + openai_reasoning.
Add discover_oss_model_specs: returns per-model {id, reasoning, context_window}.
reasoning comes from the structured capabilities.openai_reasoning; context_window
is parsed from the free-text description (1M for glm, 128K for llama/gemma) and
is None when unstated — we never fabricate a window we don't know.
Adapters enrich per-model config from the specs:
- Pi: reasoning=true (enables thinking; Pi renders the reasoning_content the
gateway already streams) and contextWindow when known.
- OpenCode: limit.context when known (reasoning is automatic via @ai-sdk/openai).
Verified live: GLM-5.2 -> {reasoning: true, contextWindow: 1000000} in Pi.
…way caps
The mlflow gateway enforces a per-model max_tokens ceiling (glm/inkling/kimi
65536, gpt-oss/qwen35 25000, qwen3-next 10000, llama/gemma 8192) but exposes it
nowhere in metadata — only by 400ing an oversized request. Requests above the
cap fail with 'max_tokens (N) cannot exceed <cap>'.
Bake the probed ceilings into _OSS_MAX_OUTPUT_TOKENS (dated, empirically
verified 2026-07-16) and carry them in the model spec as max_tokens. Adapters
set Pi's maxTokens and OpenCode's limit.output so clients don't exceed the cap.
Unprobed/new models get no override (client default) rather than a guess.
Verified live: GLM-5.2 -> {reasoning: true, contextWindow: 1000000,
maxTokens: 65536}.
Author
|
Superseded by #217. This branch was based on a pre-#178/#186 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Makes Databricks MLflow chat-completions-only foundation models (Llama, Qwen, GLM, gpt-oss, Kimi, Gemma, inkling) selectable and correctly configured in both Pi and OpenCode, and repairs inkling's streaming for Pi.
Closes #204. The gateway
finish_reasonbug is tracked separately in #215.Why
Both agents only registered their native-dialect providers, so models served solely via
mlflow/v1/chat/completionsmatched no discovery filter and never appeared:databricks-claude/databricks-openai/databricks-geminionly.databricks-anthropic/databricks-googleonly.Even once selectable, these models were configured with just an id — so clients fell back to defaults (Pi showed GLM-5.2 as 128K context, no thinking) though the gateway reports 1M context + reasoning.
Changes
1. Discovery + providers
discover_oss_models/discover_oss_model_specs(databricks.py): list endpoints exposingmlflow/v1/chat/completionsbut none of the native api_types already served by other providers (dedup), excluding embeddings.databricks-mlflow(api: openai-completions→/ai-gateway/mlflow/v1,compat: {supportsStore, supportsStrictMode}=false).databricks-ossprovider (@ai-sdk/openaichat-completions mode →/ai-gateway/mlflow/v1; nouseResponsesApi, distinct from the codex provider in fix(opencode): wire databricks-openai provider for GPT-5 / Codex #101; composes with it).2. inkling streaming repair (Pi only)
finish_reasonon natural completion; Pi'sopenai-completionsparser rejects that.agents/_mlflow_proxy.pyis a stdlib SSE-repair proxy started in Pi'slaunch()that injectsfinish_reason:"stop"on any stream end (clean or abnormal, e.g. a mid-stream 429) after data. No-op for well-behaved models; removable once AI Gateway MLflow chat-completions omits finish_reason on natural stop (breaks inkling in Pi) #215 lands.@ai-sdk/openaitolerates the missingfinish_reason.3. Per-model capability enrichment
discover_oss_model_specsreturns{id, reasoning, context_window, max_tokens}.reasoningfrom structuredcapabilities.openai_reasoning→ Pireasoning:true(enables thinking; OpenCode surfaces reasoning automatically).context_windowparsed from the free-text description (glm 1M, llama/gemma 128K);Nonewhen unstated (never fabricated).max_tokensfrom_OSS_MAX_OUTPUT_TOKENS— per-model output ceilings probed from the gateway (glm/inkling/kimi 65536, gpt-oss/qwen35 25000, qwen3-next 10000, llama/gemma 8192; not in any metadata API). PimaxTokens, OpenCodelimit.output. Unprobed models omit it.4. Housekeeping
uv.locknet diff is one line (publicfiles.pythonhosted.orgURLs preserved).Tests
test_databricks.py::TestDiscoverOssModelstest_databricks.py::TestDiscoverOssModelSpecstest_databricks.py::TestParseContextWindowtest_databricks.py::TestBuild{Pi,Opencode}BaseUrlsoss→/ai-gateway/mlflow/v1test_agent_pi.py::TestRenderOverlay*test_agent_pi.py::TestStartOssProxytest_mlflow_proxy.pytest_agent_opencode.py::TestRenderOverlayOssdatabricks-ossprovider (@ai-sdk/openai, mlflow URL, no useResponsesApi);limit.context/limit.outputset/omittedtest_cli.py::TestConfigureSharedStateOssWiringopencode_models['oss']; not fetched for unrelated toolsTest evidence
uv run pytest --ignore=tests/test_e2e.py): 579 passed, plus 4 pre-existingtest_e2e_user_agentfailures that spawn real agent binaries (sandbox-only; touch none of this code).ruff check .andty(on changed source) clean.End-to-end (live AI Gateway v2 workspace)
6 × 7→42, exit 0.6 × 7→42, exit 0.{reasoning: true, contextWindow: 1000000, maxTokens: 65536}.Reviewer note
_OSS_MAX_OUTPUT_TOKENSis hardcoded from probing (dated 2026-07-16) because the caps aren't in any metadata API. It's the most likely thing to drift; unknown/new models fail safe (no override → client default). Alternative is to drop it and let the gateway's clear 400 enforce the cap.