Skip to content

fix: support OSS model discovery in Pi - #686

Open
dairyqiu wants to merge 2 commits into
databricks:mainfrom
dairyqiu:fix/pi-oss-model-discovery
Open

dairyqiu wants to merge 2 commits into
databricks:mainfrom
dairyqiu:fix/pi-oss-model-discovery

Conversation

@dairyqiu

@dairyqiu dairyqiu commented Sep 16, 2026

Copy link
Copy Markdown

Problem

Pi only configured Claude, OpenAI Responses, and Gemini providers. Workspace model discovery already identifies Kimi, GLM, and DeepSeek services, but Pi did not consume that OSS family or have a provider that could route it.

OpenCode works because it already maps OSS models to the AI Gateway MLflow chat-completions route. This is a Pi integration gap, not a Databricks AI Gateway limitation.

Fix

  • Discover and persist OSS models when configuring Pi.
  • Add a managed databricks-oss provider using Pi openai-completions against /ai-gateway/mlflow/v1.
  • Keep per-request credentials through !ug auth-token.
  • Disable unsupported store and strict-mode fields for the MLflow route.
  • Apply known per-model context and output limits, including GLM output caps.
  • Support managed allowlists, default selection, availability checks, and explicit model selectors for OSS models.
  • Remove the stale databricks-kimi provider written by the earlier implementation.
  • Extend the live Pi e2e harness to discover, require, configure, and invoke every OSS model exposed by the test workspace.

This revisits #227 on current main. The earlier PR stopped at config-shape validation; this version validates live discovery and inference with the installed Pi client.

Verification

  • uv run ruff check .
  • uv run ruff format --check src/ tests/
  • Final focused suite: 387 passed.
  • Final full suite: 2309 passed, 42 skipped. The only failure is the existing installed Claude Code user-agent canary rejecting its synthetic model before reaching the capture server.
  • Isolated Pi 0.84.3 model listing: live workspace discovery returned seven OSS models, all shown under databricks-oss while the developer state and real Pi config remained unchanged.
  • Isolated live OSS inference: all seven discovered models passed, covering three DeepSeek, three GLM, and Kimi K3 models through /ai-gateway/mlflow/v1/chat/completions.
  • The aggregate local Pi e2e run also exercised its existing non-OSS families; its only failures were two deprecated Gemini 2.5 endpoints on the local workspace.

Residual risk

The live checks cover one-turn inference for every discovered OSS model. A multi-turn OSS tool-call loop is not yet covered.

@dairyqiu
dairyqiu force-pushed the fix/pi-oss-model-discovery branch from c8f45cb to 655d075 Compare September 16, 2026 20:02
@dairyqiu

Copy link
Copy Markdown
Author

@rohita5l, could you review this PR as the approver? It adds Pi OSS discovery and MLflow chat-completions routing, with live one-turn inference passing for all seven discovered DeepSeek, GLM, and Kimi models. The branch is rebased on current main, and the focused suite and static checks pass.

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