Skip to content

Support custom Model Provider Services in Pi #289

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
FantomKun wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Support custom Model Provider Services in Pi #289

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 26 additions & 11 deletions src/ucode/agents/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
from ucode.config_io import ToolSpec
from ucode.databricks import (
BEDROCK_PROVIDER_TYPES,
CUSTOM_PROVIDER_TYPES,
custom_openai_chat_targets,
get_databricks_token,
install_ai_tools,
install_databricks_cli,
Expand Down Expand Up @@ -282,16 +284,18 @@ def resolve_launch_model(

def resolve_provider_models(
tool: str, state: dict, provider: str | None
) -> tuple[dict | None, str | None, bool]:
) -> tuple[dict[str, str] | list[str] | None, str | None, bool]:
"""Validate ``provider`` for ``tool`` and return the model ids to pin.

Returns ``(provider_models, error, relayed)``. ``provider_models`` is a
``{family: model_id}`` dict for a Bedrock-backed claude service (whose
provider-side ids must be pinned explicitly), or None for an Anthropic/
canonical service or when ``provider`` is None. ``relayed`` is True for a
credential-less Anthropic subscription relay, which the launch path wires
with the relayed overlay + refresh proxy. A non-None ``error`` means the
provider is invalid for the tool and the caller should not launch.
provider-side ids must be pinned explicitly), a list of target ids for a
custom-backed service (routed by header, with the bare target name as the
request body's `model`), or None for an Anthropic/canonical service or when
``provider`` is None. ``relayed`` is True for a credential-less Anthropic
subscription relay, which the launch path wires with the relayed overlay +
refresh proxy. A non-None ``error`` means the provider is invalid for the
tool and the caller should not launch.
"""
if not provider:
return None, None, False
Expand All @@ -302,6 +306,8 @@ def resolve_provider_models(
relayed = bool(service.get("relayed"))
if service["provider_type"] in BEDROCK_PROVIDER_TYPES:
return map_bedrock_claude_models(service.get("targets") or []), None, relayed
if service["provider_type"] in CUSTOM_PROVIDER_TYPES:
return custom_openai_chat_targets(service), None, relayed
return None, None, relayed


Expand All @@ -310,7 +316,7 @@ def configure_tool(
state: dict,
model: str | None = None,
provider: str | None = None,
provider_models: dict[str, str] | None = None,
provider_models: dict[str, str] | list[str] | None = None,
relayed: bool = False,
route_root_model: str | None = None,
) -> dict:
Expand All @@ -326,20 +332,29 @@ def configure_tool(
state,
model,
provider=provider,
provider_models=provider_models,
provider_models=provider_models if isinstance(provider_models, dict) else None,
relayed=relayed,
route_root_model=route_root_model,
)
elif tool == "pi":
# As with claude, a Model Provider Service pins no Databricks model —
# the service's targets are the models.
if not model and not provider:
raise RuntimeError(f"A {tool} model must be selected before configuration.")
result = pi.write_tool_config(
state,
model,
provider=provider,
provider_models=provider_models if isinstance(provider_models, list) else None,
)
else:
# provider routing is claude/codex-only; every other tool needs a model.
# provider routing is claude/codex/pi-only; every other tool needs a model.
if not model:
raise RuntimeError(f"A {tool} model must be selected before configuration.")
if tool == "gemini":
result = gemini.write_tool_config(state, model)
elif tool == "copilot":
result = copilot.write_tool_config(state, model)
elif tool == "pi":
result = pi.write_tool_config(state, model)
else:
result = opencode.write_tool_config(state, model)
# gemini/opencode/copilot/pi return (state, token); codex/claude return state
Expand Down
140 changes: 131 additions & 9 deletions src/ucode/agents/pi.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,13 +8,22 @@
- `databricks-openai` (api: openai-responses) → /ai-gateway/codex/v1
- `databricks-gemini` (api: google-generative-ai) → /ai-gateway/gemini/v1beta

A fourth provider, `databricks-custom` (api: openai-completions →
/ai-gateway/openai/v1), is written only when launching through a `custom` Model
Provider Service — a self-hosted, OpenAI-compatible model registered in Unity
Catalog. The service is selected by the `Databricks-Model-Provider-Service`
header and its targets are the models, so no Databricks model is pinned.

Per-provider `compat` flags work around fields the gateway translators reject:

- claude: `supportsEagerToolInputStreaming: false` — the Anthropic translator
rejects `tools[].eager_input_streaming` on the streaming + tools path that
pi uses for every request. With this flag pi omits the per-tool field and
sends the legacy `anthropic-beta: fine-grained-tool-streaming-...` header
instead, which the gateway accepts.
- custom: three flags turn off assumptions pi makes for openai.com. Pi's own
auto-detection already gets the rest right for an unrecognized base URL, so
restating them would be noise.

OSS / Databricks-foundation models (Llama, Qwen, etc.) are not exposed via
pi today — they live behind /ai-gateway/mlflow/v1 with per-model
Expand Down Expand Up @@ -42,7 +51,9 @@
write_json_file,
)
from ucode.databricks import (
OPENAI_CHAT_NATIVE_API_TYPE,
TOKEN_REFRESH_INTERVAL_SECONDS,
build_native_api_base_url,
build_pi_base_urls,
get_databricks_token,
)
Expand All @@ -64,14 +75,33 @@
"backup_path": PI_BACKUP_PATH,
}

CUSTOM_PROVIDER_NAME = "databricks-custom"

PROVIDER_NAMES = (
"databricks-claude",
"databricks-openai",
"databricks-gemini",
# Written only under a custom Model Provider Service, but listed here
# unconditionally so it's stripped on every write: a later launch without
# `--provider` (or against a different service) must not leave a stale
# provider pointing at the old service's header.
CUSTOM_PROVIDER_NAME,
)

PROVIDER_KEYS: list[list[str]] = [["providers", name] for name in PROVIDER_NAMES]

# The Model Provider Service API exposes no context-window metadata, so ucode has
# to assume one for a custom target. The two directions are not symmetric:
# understating costs earlier compaction and shorter replies, while overstating is
# unrecoverable — pi compacts to `contextWindow - reserveTokens`, so a window
# above the server's real limit makes the compact-and-retry overflow again and
# the turn ends with "Context overflow recovery failed after one
# compact-and-retry attempt." Assume the conservative floor common to
# self-hosted OpenAI-compatible servers and let users raise it with
# `ucode configure --provider-context-window`.
PROVIDER_CONTEXT_WINDOW = 32768
PROVIDER_MAX_OUTPUT_TOKENS = 8192

# Old provider names earlier ucode versions wrote; cleaned up on each write so
# users don't end up with stale entries pointing at routes that 400.
LEGACY_PROVIDER_NAMES = ("databricks-anthropic", "databricks-codex", "databricks-oss")
Expand All @@ -82,15 +112,28 @@ def is_update_available() -> tuple[str, str] | None:


def _resolve_model_selector(
model: str,
model: str | None,
claude_models: dict[str, str],
codex_models: list[str],
gemini_models: list[str],
provider_models: list[str] | None = None,
) -> str:
"""Return a Pi model selector in `<provider>/<model>` form when possible."""
"""Return a Pi model selector in `<provider>/<model>` form when possible.

The provider-qualified form matters: a bare target id would match Pi's own
built-in provider of the same name (e.g. `deepseek`) and fail with "No API
key found for deepseek" rather than routing through the gateway.
"""
provider_models = provider_models or []
# Under a custom Model Provider Service no Databricks model is resolved, so
# the service's first target is the default.
if not model:
return f"{CUSTOM_PROVIDER_NAME}/{provider_models[0]}" if provider_models else ""
for name in PROVIDER_NAMES:
if model.startswith(f"{name}/"):
return model
if model in provider_models:
return f"{CUSTOM_PROVIDER_NAME}/{model}"
if model in claude_models.values():
return f"databricks-claude/{model}"
if model in codex_models:
Expand All @@ -101,14 +144,25 @@ def _resolve_model_selector(


def render_overlay(
model: str,
model: str | None,
token: str,
pi_base_urls: dict[str, str],
claude_models: dict[str, str],
codex_models: list[str],
gemini_models: list[str],
*,
provider: str | None = None,
provider_models: list[str] | None = None,
provider_base_url: str | None = None,
context_window: int | None = None,
) -> tuple[dict, list[list[str]]]:
"""Return (overlay, managed_key_paths) for ~/.pi/agent/models.json."""
"""Return (overlay, managed_key_paths) for ~/.pi/agent/models.json.

``provider`` is a `custom` Model Provider Service name; ``provider_models``
its routable targets. When both are set (plus a base URL for the dialect) a
`databricks-custom` provider is emitted alongside whichever Databricks
providers the workspace exposes.
"""
providers: dict = {}
keys: list[list[str]] = [["model"]]
# Pi expands header values that match an env var name. Our UA contains
Expand Down Expand Up @@ -150,8 +204,47 @@ def render_overlay(
"models": [{"id": m} for m in gemini_models],
}
keys.append(["providers", "databricks-gemini"])
if provider and provider_models and provider_base_url:
window = context_window or PROVIDER_CONTEXT_WINDOW
providers[CUSTOM_PROVIDER_NAME] = {
"baseUrl": provider_base_url,
"api": "openai-completions",
"apiKey": token,
"authHeader": True,
# The header selects the service; the body's `model` carries the bare
# target name (the fully-qualified service name returns NOT_FOUND).
"headers": {**ua_headers, "Databricks-Model-Provider-Service": provider},
# Only the flags that change pi's behavior against an unknown
# OpenAI-compatible backend. Pi already auto-detects
# supportsReasoningEffort / supportsUsageInStreaming /
# supportsStrictMode correctly for an unrecognized base URL.
"compat": {
# Older OSS servers accept `max_tokens`, not the newer
# `max_completion_tokens` pi would otherwise send.
"maxTokensField": "max_tokens",
# The `developer` role is an OpenAI-ism; `system` is universal.
"supportsDeveloperRole": False,
# Many OSS servers reject unknown body fields such as `store`.
"supportsStore": False,
},
# No `reasoning`/`thinkingLevelMap`: the service exposes no capability
# metadata, and claiming reasoning support would make pi send
# `reasoning_effort` on every request, which a non-reasoning server
# can reject outright.
"models": [
{
"id": target,
"contextWindow": window,
"maxTokens": min(PROVIDER_MAX_OUTPUT_TOKENS, window // 4),
}
for target in provider_models
],
}
keys.append(["providers", CUSTOM_PROVIDER_NAME])
overlay: dict = {
"model": _resolve_model_selector(model, claude_models, codex_models, gemini_models),
"model": _resolve_model_selector(
model, claude_models, codex_models, gemini_models, provider_models
),
}
if providers:
overlay["providers"] = providers
Expand All @@ -160,9 +253,11 @@ def render_overlay(

def write_tool_config(
state: dict,
model: str,
model: str | None,
token: str | None = None,
*,
provider: str | None = None,
provider_models: list[str] | None = None,
force_refresh: bool = False,
) -> tuple[dict, str]:
backup_existing_file(PI_CONFIG_PATH, PI_BACKUP_PATH)
Expand All @@ -178,7 +273,22 @@ def write_tool_config(
state.get("claude_models") or {},
state.get("codex_models") or [],
state.get("gemini_models") or [],
provider=provider,
provider_models=provider_models,
provider_base_url=build_native_api_base_url(
state["workspace"], OPENAI_CHAT_NATIVE_API_TYPE
),
context_window=state.get("provider_context_window"),
)
# Persist the resolved pair so the background refresh thread can rewrite the
# same provider config without re-resolving it, and clear it on a launch
# without a provider so the next refresh doesn't resurrect a stale one.
if provider and provider_models:
state["pi_provider"] = provider
state["pi_provider_models"] = list(provider_models)
else:
state.pop("pi_provider", None)
state.pop("pi_provider_models", None)
existing = read_json_safe(PI_CONFIG_PATH)
providers = existing.get("providers")
if isinstance(providers, dict):
Expand Down Expand Up @@ -219,10 +329,22 @@ def default_model(state: dict) -> str | None:


def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str:
model = default_model(state)
if not model:
# Under a Model Provider Service the workspace may expose no Databricks model
# at all — the service's targets are the models. Requiring one here would
# raise, and `_refresh_forever` swallows that, so the token would silently
# stop refreshing and the session would die when it expired.
provider = state.get("pi_provider")
provider_models = state.get("pi_provider_models") or []
model = None if provider else default_model(state)
if not model and not (provider and provider_models):
raise RuntimeError("No Pi model is available on this workspace.")
_, token = write_tool_config(state, model, force_refresh=force_refresh)
_, token = write_tool_config(
state,
model,
provider=provider,
provider_models=provider_models or None,
force_refresh=force_refresh,
)
return token


Expand Down
Loading