Skip to content

feat(pi): OSS mlflow provider + broaden OSS cohort (reasoning, limits)#217

Open
dgokeeffe wants to merge 4 commits into
databricks:mainfrom
dgokeeffe:feat/pi-oss-provider-v2
Open

feat(pi): OSS mlflow provider + broaden OSS cohort (reasoning, limits)#217
dgokeeffe wants to merge 4 commits into
databricks:mainfrom
dgokeeffe:feat/pi-oss-provider-v2

Conversation

@dgokeeffe

@dgokeeffe dgokeeffe commented Jul 16, 2026

Copy link
Copy Markdown

What

Adds Pi support for OSS (mlflow chat-completions-only) foundation models and broadens the OSS cohort across both agents, building on the OpenCode OSS support already merged (#178 Kimi, #186 GLM).

Supersedes #216 (which was based on a pre-#178/#186 main and conflicted with the merged OSS work). Closes #204.

Gaps this closes on current main

main discovers OSS models via UC model-services and buckets them by name-substring allowlist _OSS_MODEL_FAMILIES = ("kimi-", "glm-"), wired into OpenCode only. That leaves:

  1. Pi has no OSS provider at all — GLM/inkling/etc. are unselectable in Pi.
  2. Only kimi/glm are bucketed — inkling, llama, qwen, gpt-oss, gemma are discovered from model-services but dropped.
  3. GLM limits are offmain pins 200k context / 25k output; the gateway actually accepts 1M / 65536 (probed 2026-07-16).
  4. No reasoning flag — reasoning models (glm, inkling, ...) don't surface thinking in Pi.

Changes

  • Pi databricks-mlflow provider (api: openai-completions/ai-gateway/mlflow/v1, compat {supportsStore, supportsStrictMode}=false). Selector/default-fallback/stale-cleanup wired; default_model falls back to OSS last.
  • _mlflow_proxy.py — a stdlib SSE-repair proxy started in Pi's launch(). Some OSS models (inkling) omit the spec-required finish_reason on the streaming terminal chunk; Pi's openai-completions parser rejects that. The proxy injects it on any stream end (clean or an abnormal/429 mid-stream drop) and re-derives its origin from the workspace each launch (never a stale persisted port). No-op for well-behaved models. OpenCode needs no proxy — verified @ai-sdk/openai tolerates the missing finish_reason.
  • Broaden _OSS_MODEL_FAMILIES to the full cohort (inkling, llama, qwen, gpt-oss, gemma + existing kimi/glm), with _is_oss_chat_model excluding non-chat services (embeddings/rerankers) that share a family substring — e.g. qwen3-embedding is dropped though it matches qwen.
  • Correct + extend _MODEL_TOKEN_LIMITS from probed gateway caps (glm 1M/65536, kimi/inkling 65536, gpt-oss/qwen35 25000, qwen3-next 10000, llama/gemma 8192). Add model_is_reasoning() from the probed openai_reasoning capability. Both agents read these shared tables; Pi sets reasoning/contextWindow/maxTokens, OpenCode sets limit.{context,output}.
  • Version 0.1.0 → 0.2.0. uv.lock net diff is the one ucode version line.

Verification

  • Full suite: 890 passed, 7 skipped (excl. e2e, which spawn real agents); ruff + ty clean.
  • Live (AI Gateway v2): discovery returns 11 chat OSS models (no embeddings); glm-5-2 answers 42 through Pi via the new provider.

Notes for reviewers

  • Kept main's model-services + per-family-table architecture rather than a capability-API approach (model-services doesn't expose api_types/capabilities). Reasoning/limits are therefore hardcoded per-family, consistent with the existing _MODEL_TOKEN_LIMITS.
  • _MODEL_TOKEN_LIMITS output caps are probed (dated 2026-07-16); they're the thing most likely to drift if the gateway changes. Unlisted/new models fail safe (no override → client default).
  • The GLM cap change (25k → 65536) reverses part of support glm-5-2 on opencode  #186. 65536 is the live gateway ceiling (30k is accepted, 70k is rejected), so 25k was leaving output on the table. Confirmed working end-to-end at these caps. Flagging only in case support glm-5-2 on opencode  #186 pinned 25k for a reason beyond the hard cap.
  • inkling and glm run end-to-end through the Pi provider (glm 6×7 → 42; inkling confirmed by the author). Note the selector uses the full UC id, e.g. databricks-mlflow/system.ai.inkling — a stale databricks-<name> selector from an earlier build won't resolve (Pi warns and falls back). The gateway finish_reason omission is tracked in AI Gateway MLflow chat-completions omits finish_reason on natural stop (breaks inkling in Pi) #215.
  • Reasoning models stream a correct delta.content after their reasoning_content deltas, and the proxy passes it through verbatim (verified in the emitted bytes at low/medium/high reasoning_effort). If a reasoning model appears to print nothing in Pi's non-interactive --print mode, the content is present in the stream — this is a Pi client rendering behavior, not a gap in what ucode produces.

Code-review fixes (post-initial-review)

Addressed a review pass on _mlflow_proxy.py: close the connection at stream end so HTTP/1.1 clients detect the boundary (was a potential hang); only repair text/event-stream responses and relay non-streaming/error bodies verbatim; return a clean 502 on connection-level upstream failures; restrict forwarded paths to /ai-gateway/mlflow/ (SSRF guard); tolerate client-disconnect in the terminator-repair path; parse data: lines with or without the optional space. Discovery: OSS chat models matching a broad family but lacking a specific limits entry now fall back to a conservative 128k/8192 cap so they're never offered uncapped (which would 400). Full suite: 898 passed.

…h reasoning/limits

Builds on the merged OpenCode OSS support (kimi/glm) to close the remaining gaps:

- Pi had NO OSS provider. Add databricks-mlflow (api openai-completions ->
  /ai-gateway/mlflow/v1) so Pi can select OSS chat models, with a local
  SSE-repair proxy (_mlflow_proxy) for models like inkling whose gateway stream
  omits finish_reason (Pi's parser rejects that; OpenCode's @ai-sdk/openai
  tolerates it, so it needs no proxy).
- Broaden _OSS_MODEL_FAMILIES from kimi/glm to the full chat-completions cohort
  (inkling, llama, qwen, gpt-oss, gemma), excluding non-chat services
  (embeddings/rerankers) that share a family substring (e.g. qwen3-embedding).
- Correct + extend _MODEL_TOKEN_LIMITS from probed gateway caps: glm is
  1M context / 65536 output (was 200k/25k); per-family output ceilings for the
  cohort. Add model_is_reasoning() from probed openai_reasoning capability;
  Pi sets reasoning:true so it renders streamed reasoning_content as thinking.

Both agents read the shared model_token_limits/model_is_reasoning tables.
Verified live: glm-5-2 answers through Pi (6*7 -> 42).
Pi databricks-mlflow OSS provider + broadened OSS cohort with reasoning/limits.
gpt-oss-* contains 'gpt-' so it matched the codex bucket, but it's a
chat-completions-only OSS model (served via /ai-gateway/mlflow/v1), not an
openai-responses model. Broadening _OSS_MODEL_FAMILIES to include gpt-oss put
it in BOTH codex_models and oss_models, so it was offered under Pi's
databricks-openai provider too — where the openai-responses route 400s for it.
Exclude gpt-oss from the codex bucket so it's OSS-only.
Proxy (_mlflow_proxy.py):
- Close the connection at stream end (self.close_connection = True) so HTTP/1.1
  clients detect the message boundary instead of hanging after [DONE] (the
  proxy sends no Content-Length/chunking). [critical]
- Only repair text/event-stream responses; relay non-streaming and error bodies
  verbatim with their real Content-Type instead of forcing SSE. [high]
- Catch URLError/OSError around urlopen and return a clean 502 instead of a dead
  socket on connection-level failures. [high]
- Restrict forwarded paths to /ai-gateway/mlflow/ (404 otherwise) so the
  localhost relay can't be used as an authenticated arbitrary-path workspace
  client. [security]
- Move terminator-repair writes inside the OSError guard so a client
  disconnect doesn't traceback on the daemon thread. [medium]
- Parse SSE data lines with or without the optional space after the colon. [low]

Discovery (databricks.py):
- OSS chat models matching a broad family (qwen, llama-) but lacking a specific
  _MODEL_TOKEN_LIMITS entry now fall back to a conservative floor
  (128k/8192) so they're never offered uncapped (which 400s). Non-OSS and
  embedding ids still return None.

Verified live: glm-5-2 streams content (42) + finish_reason + [DONE] through the
proxy; non-stream/error/404/502 paths covered by tests. 898 passed.
@dgokeeffe

Copy link
Copy Markdown
Author

@rohita5l @AarushiShah-db — would appreciate a review when you have a moment. This builds directly on your OSS work (#178 Kimi, #186 GLM): adds a Pi databricks-mlflow provider (Pi had none), broadens the OSS family bucket to the rest of the chat-completions cohort (inkling/llama/qwen/gpt-oss/gemma, excluding embeddings), corrects the GLM limits to the probed gateway caps (1M context / 65536 output), and adds a reasoning flag. One heads-up: it reverses the GLM output cap from #186's 25k to a verified 65536 — details in the PR description. A code-review pass was already applied (proxy framing/error-handling/SSRF-guard hardening). CI hasn't run — could use a maintainer approve-and-run on this fork PR.

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.

ucode pi automatic model detection

1 participant