Skip to content

pi: support OSS models (kimi, glm) via managed databricks-oss provider#227

Closed
lilly-luo wants to merge 1 commit into
mainfrom
lilly/pi-oss-models
Closed

pi: support OSS models (kimi, glm) via managed databricks-oss provider#227
lilly-luo wants to merge 1 commit into
mainfrom
lilly/pi-oss-models

Conversation

@lilly-luo

@lilly-luo lilly-luo commented Jul 21, 2026

Copy link
Copy Markdown
Collaborator

Problem

Configuring Pi and calling an OSS model (e.g. kimi-k2-7-code) returned 401 against …/ai-gateway/mlflow/v1/responses, while Claude on the same endpoint worked.

Root cause: Pi only ever managed claude/openai/gemini providers. An early ucode build had written a databricks-kimi provider using the wrong openai-responses dialect against the MLflow chat-completions route (hence the /responses 401), and because that provider wasn't in any list ucode manages:

  • its baked bearer token was never refreshed (frozen from weeks ago → 401), and
  • the stale-provider cleanup couldn't remove it, so it survived every ucode configure.

Fix

Re-introduce OSS support to Pi as a first-class managed provider:

  • databricks.pybuild_pi_base_urls now returns the oss MLflow base URL.
  • agents/pi.py — add a managed databricks-oss provider:
    • api: "openai-completions"/ai-gateway/mlflow/v1 (correct dialect for the MLflow route).
    • compat.supportsStore=false + supportsStrictMode=false — the two fields the MLflow route rejects.
    • per-model contextWindow/maxTokens from model_token_limits, so GLM's 25k output cap is honored (Pi's 16k default would 400 GLM); kimi has no entry and uses Pi defaults.
    • threaded oss_models through render_overlay, write_tool_config, _resolve_model_selector, default_model.
    • added databricks-oss to PROVIDER_NAMES and moved databricks-kimi into LEGACY_PROVIDER_NAMES so the stale block is stripped on every write.
  • cli.pywant_oss now includes pi, so ucode configure pi discovers and refreshes OSS models; updated the discovery diagnostic hint.

Verified the models.json schema (api dialects, per-model contextWindow/maxTokens, supportsStore/supportsStrictMode compat) against the installed @earendil-works/pi-coding-agent package and its docs.

Tests

pi and oss are not ready yet

Screenshot 2026-07-21 at 10 44 44 AM

Added OSS coverage in tests/test_agent_pi.py (dialect, MLflow URL, compat flags, model listing, GLM limits pinned / kimi limits omitted, selector prefix, default-model fallback) and updated the legacy-removal test to assert databricks-kimi is stripped. Full suite passes (895 passed; the one failing test is a pre-existing e2e environment issue unrelated to this change).

Note

Not yet validated with a live end-to-end request through the gateway — the config shape is verified against Pi's schema/docs, but a real ucode pi launch against a workspace is needed to confirm kimi/glm return 200.

This pull request and its description were written by Isaac.

Pi previously exposed only claude/openai/gemini. OSS families (kimi, glm)
were left unmanaged, so an early build's `databricks-kimi` provider — wired
with the wrong `openai-responses` dialect against the MLflow chat-completions
route — lingered with a frozen, never-refreshed token and 401'd on every call.

Add a managed `databricks-oss` provider (api: openai-completions →
/ai-gateway/mlflow/v1) with per-model contextWindow/maxTokens from
model_token_limits (so GLM's 25k output cap is honored) and
supportsStore/supportsStrictMode compat flags the MLflow route requires.
`ucode configure pi` now discovers and refreshes OSS models, and the stale
`databricks-kimi` block is stripped on every write via LEGACY_PROVIDER_NAMES.

Co-authored-by: Isaac
@lilly-luo lilly-luo closed this Jul 21, 2026
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