scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

feat(cli): add lite pi to run the pi coding agent through the proxy

Browse source
This commit is contained in:
ryan-crabbe-berri 2026-08-13 15:33:36 -07:00
parent 9d069f21dc
commit 4e280ecc34
5 changed files with 591 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/proxy/client/cli/README.md
View file

@ -467,6 +467,7 @@ Launch a coding agent with all of its LLM traffic routed through your LiteLLM pr
lite claude
lite codex
lite opencode
lite pi
```
Anything you type after the agent name is forwarded to it untouched, so the usual flags keep working:
@ -480,17 +481,19 @@ Each command resolves your LiteLLM key (logging in via SSO when none is stored a
The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol).
pi ignores base-URL environment variables entirely, so `lite pi` wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/<id>`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. Your own flags come after the injected pin, so `lite pi --model litellm/<other-id>` wins, and inside the TUI the `/model` picker lists every synced litellm model.
Options (these belong to the wrapper, so put them before the agent's own flags):
- `--skip-verify`: Skip the pre-launch key check (useful offline or with non-standard auth).
To pin the model, pass the agent's own model flag (for example `lite claude --model my-proxy-model` or `lite codex -m my-proxy-model`), or export the variable the agent reads (`ANTHROPIC_MODEL` / `ANTHROPIC_SMALL_FAST_MODEL` for Claude Code); the wrapper preserves anything you already have set. Whatever model the agent ends up requesting must exist on the proxy, since requests land on the proxy's `/v1/messages` (Anthropic) or `/v1/chat/completions` and `/v1/responses` (OpenAI) endpoints.
To pin the model, pass the agent's own model flag (for example `lite claude --model my-proxy-model`, `lite codex -m my-proxy-model`, or `lite pi --model my-proxy-model`), or export the variable the agent reads (`ANTHROPIC_MODEL` / `ANTHROPIC_SMALL_FAST_MODEL` for Claude Code); the wrapper preserves anything you already have set. Whatever model the agent ends up requesting must exist on the proxy, since requests land on the proxy's `/v1/messages` (Anthropic) or `/v1/chat/completions` and `/v1/responses` (OpenAI) endpoints.
#### About the `lite login` credential
The token minted by `lite login` is a short-lived, per-session agent credential, not a managed virtual key. It is scoped to the user and team you authenticated as, inherits that user's and team's models and budgets, and is enforced on the proxy exactly like a virtual key on the same team (guardrails, routing, logging, spend). Spend is tracked against the shared team and user budgets, so running several agents (or logging in more than once) does not hand each session its own separate budget; they all draw down the same team/user allowance. There is no separate per-session cap, so sustained agent use is not capped at a small chat-session limit.
The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while it's still fresh and fails once it expires -- there is no silent renewal, so a long-running session needs a fresh `lite login` once a day. `lite claude`, `lite codex`, and `lite opencode` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead.
The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while it's still fresh and fails once it expires -- there is no silent renewal, so a long-running session needs a fresh `lite login` once a day. `lite claude`, `lite codex`, `lite opencode`, and `lite pi` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead.
### Route Every Claude Code Session Through the Proxy

71
litellm/proxy/client/cli/commands/agents.py
View file

@ -2,12 +2,22 @@ import os
import shutil
import sys
from collections.abc import Callable, Mapping, Sequence
from types import MappingProxyType
from typing import Final
import click
import requests
from .auth import get_stored_api_key, login
from .pi import (
LITELLM_PROXY_API_KEY_ENV,
PI_PROVIDER_NAME,
PiSyncError,
fetch_model_ids,
fetch_model_limits,
models_json_path,
sync_models_json,
)
ANTHROPIC_BASE_URL_ENV: Final = "ANTHROPIC_BASE_URL"
ANTHROPIC_AUTH_TOKEN_ENV: Final = "ANTHROPIC_AUTH_TOKEN"
@ -17,17 +27,20 @@ OPENAI_API_KEY_ENV: Final = "OPENAI_API_KEY"
PROFILE_ANTHROPIC: Final = "anthropic"
PROFILE_OPENAI: Final = "openai"
PROFILE_LITELLM: Final = "litellm"
_KNOWN_AGENTS: Final[dict[str, tuple[str, frozenset[str]]]] = {
"claude": ("Claude Code", frozenset({PROFILE_ANTHROPIC})),
"codex": ("Codex", frozenset({PROFILE_OPENAI})),
"opencode": ("OpenCode", frozenset({PROFILE_OPENAI})),
"pi": ("pi", frozenset({PROFILE_LITELLM})),
}
_INSTALL_DOCS: Final[dict[str, str]] = {
"claude": "https://docs.claude.com/en/docs/claude-code/setup",
"codex": "https://developers.openai.com/codex/cli",
"opencode": "https://opencode.ai/docs",
"pi": "https://pi.dev",
}
CODEX_PROXY_PROVIDER: Final = "litellm"
@ -60,7 +73,9 @@ def build_agent_env(
Anthropic clients (Claude Code) append /v1/messages to ANTHROPIC_BASE_URL,
so it stays the bare proxy root; OpenAI clients (Codex, OpenCode) expect the
/v1 suffix on OPENAI_BASE_URL. ANTHROPIC_API_KEY is dropped so a stray
Anthropic key cannot win over the bearer token we set.
Anthropic key cannot win over the bearer token we set. pi ignores both base
URL variables and instead resolves $LITELLM_PROXY_API_KEY from its synced
models.json provider entry.
"""
env: Final = dict(base_env)
root: Final = base_url.rstrip("/")
@ -71,6 +86,8 @@ def build_agent_env(
if PROFILE_OPENAI in profiles:
env[OPENAI_BASE_URL_ENV] = root + "/v1"
env[OPENAI_API_KEY_ENV] = api_key
if PROFILE_LITELLM in profiles:
env[LITELLM_PROXY_API_KEY_ENV] = api_key
return env
@ -106,6 +123,38 @@ _PROXY_ARGS: Final[dict[str, Callable[[str], list[str]]]] = {
}
def prepare_pi(
base_url: str,
api_key: str,
base_env: Mapping[str, str],
*,
get: Callable[..., requests.Response] = requests.get,
) -> list[str]:
"""Sync the proxy's model list into pi's models.json before handoff.
pi has no base-URL env vars, so this file is the only way to point it at the
proxy. Only the litellm provider entry is touched; the synced entry references
the key as $LITELLM_PROXY_API_KEY, which build_agent_env exports. The returned
--model pin is needed because pi ignores a bare --provider when picking the
interactive startup model; a user-supplied --model comes later in argv and wins.
"""
ids: Final = fetch_model_ids(base_url, api_key, get=get)
if isinstance(ids, PiSyncError):
raise AgentRunError(ids.message)
limits: Final = fetch_model_limits(base_url, api_key, get=get)
path: Final = models_json_path(base_env)
error: Final = sync_models_json(path, base_url, ids, limits)
if error is not None:
raise AgentRunError(error.message)
click.echo(f"litellm: synced {len(ids)} proxy models into {path}")
return ["--model", f"{PI_PROVIDER_NAME}/{ids[0]}"]
_PREPARERS: Final[dict[str, Callable[[str, str, Mapping[str, str]], Sequence[str]]]] = {
"pi": prepare_pi,
}
def agent_launch_args(command: str, base_url: str) -> list[str]:
"""Extra CLI args an agent needs to actually honor the proxy.
@ -177,12 +226,14 @@ def run_agent(
verify: Callable[[str, str], None] = verify_proxy_key,
launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _exec,
reattach_terminal: Callable[[], None] | None = None,
preparers: Mapping[str, Callable[[str, str, Mapping[str, str]], Sequence[str]]] = MappingProxyType(_PREPARERS),
) -> None:
"""Validate, wire the environment, and hand off to the agent.
On success this replaces the current process and never returns. Raises
AgentRunError for missing binaries, an unreachable proxy, or a rejected key.
reattach_terminal, when given, runs just before handoff to restore stdin.
AgentRunError for missing binaries, an unreachable proxy, a rejected key, or
a failed pre-launch config sync (pi). reattach_terminal, when given, runs
just before handoff to restore stdin.
"""
if not command:
raise AgentRunError("Nothing to run.")
@ -197,13 +248,12 @@ def run_agent(
if not skip_verify:
verify(base_url, api_key)
env: Final = build_agent_env(
base_env if base_env is not None else os.environ,
base_url,
api_key,
profiles,
)
extra_args: Final = agent_launch_args(command[0], base_url)
source_env: Final = base_env if base_env is not None else os.environ
prepare: Final = preparers.get(os.path.basename(command[0]))
prepared_args: Final = list(prepare(base_url, api_key, source_env)) if prepare is not None else []
env: Final = build_agent_env(source_env, base_url, api_key, profiles)
extra_args: Final = [*agent_launch_args(command[0], base_url), *prepared_args]
if reattach_terminal is not None:
reattach_terminal()
launcher(binary, [command[0], *extra_args, *command[1:]], env)
@ -288,6 +338,7 @@ __all__ = [
"agent_launch_args",
"agent_profile",
"build_agent_env",
"prepare_pi",
"resolve_api_key",
"run_agent",
"verify_proxy_key",

178
litellm/proxy/client/cli/commands/pi.py Normal file
View file

@ -0,0 +1,178 @@
"""Sync a LiteLLM provider into pi's models.json.
pi ignores ANTHROPIC_BASE_URL/OPENAI_BASE_URL, so `lite pi` routes it through the
proxy by writing a provider entry instead. The key is stored as a $-reference so
the short-lived login token never lands on disk.
"""
import json
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final
import requests
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
PI_CONFIG_DIR_ENV: Final = "PI_CODING_AGENT_DIR"
PI_PROVIDER_NAME: Final = "litellm"
LITELLM_PROXY_API_KEY_ENV: Final = "LITELLM_PROXY_API_KEY"
@dataclass(frozen=True, slots=True)
class PiSyncError:
message: str
@dataclass(frozen=True, slots=True)
class ModelLimits:
context_window: int | None
max_tokens: int | None
class _Model(BaseModel):
id: str
class _ModelList(BaseModel):
data: list[_Model]
class _ModelGroup(BaseModel):
model_group: str
max_input_tokens: float | None = None
max_output_tokens: float | None = None
class _ModelGroupList(BaseModel):
data: list[_ModelGroup]
def fetch_model_ids(
base_url: str,
api_key: str,
*,
get: Callable[..., requests.Response] = requests.get,
) -> tuple[str, ...] | PiSyncError:
url: Final = base_url.rstrip("/") + "/v1/models"
try:
resp: Final = get(url, headers={"Authorization": f"Bearer {api_key}"}, timeout=10)
except requests.RequestException as e:
return PiSyncError(f"Could not list models from the proxy: {e}")
if resp.status_code != 200:
return PiSyncError(f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot build pi's model list.")
try:
listing: Final = _ModelList.model_validate(resp.json())
except (ValueError, ValidationError) as e:
return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}")
ids: Final = tuple(dict.fromkeys(model.id for model in listing.data))
if not ids:
return PiSyncError("The proxy returned no models for your key, so pi would have nothing to run.")
return ids
_NO_LIMITS: Final[Mapping[str, ModelLimits]] = MappingProxyType({})
def fetch_model_limits(
base_url: str,
api_key: str,
*,
get: Callable[..., requests.Response] = requests.get,
) -> Mapping[str, ModelLimits]:
"""Best effort: pi falls back to its own defaults for models without limits,
so an unavailable /model_group/info must not block the launch."""
url: Final = base_url.rstrip("/") + "/model_group/info"
try:
resp: Final = get(url, headers={"Authorization": f"Bearer {api_key}"}, timeout=10)
if resp.status_code != 200:
return _NO_LIMITS
listing: Final = _ModelGroupList.model_validate(resp.json())
except (requests.RequestException, ValueError, ValidationError):
return _NO_LIMITS
return MappingProxyType(
{
group.model_group: ModelLimits(
context_window=int(group.max_input_tokens) if group.max_input_tokens else None,
max_tokens=int(group.max_output_tokens) if group.max_output_tokens else None,
)
for group in listing.data
}
)
def models_json_path(env: Mapping[str, str]) -> Path:
override: Final = env.get(PI_CONFIG_DIR_ENV)
root: Final = Path(override) if override else Path.home() / ".pi" / "agent"
return root / "models.json"
def _model_entry(model_id: str, limits: Mapping[str, ModelLimits]) -> dict[str, JsonValue]:
limit: Final = limits.get(model_id)
context: Final[dict[str, JsonValue]] = (
{"contextWindow": limit.context_window} if limit and limit.context_window else {}
)
output: Final[dict[str, JsonValue]] = {"maxTokens": limit.max_tokens} if limit and limit.max_tokens else {}
return {"id": model_id, **context, **output}
def provider_block(
base_url: str,
model_ids: tuple[str, ...],
limits: Mapping[str, ModelLimits] = _NO_LIMITS,
) -> dict[str, JsonValue]:
"""openai-completions is the one API shape every LiteLLM model serves.
Real contextWindow/maxTokens matter: pi otherwise assumes 128k/16384, which
breaks compaction thresholds and over-asks models with smaller output caps.
"""
return {
"baseUrl": base_url.rstrip("/") + "/v1",
"api": "openai-completions",
"apiKey": f"${LITELLM_PROXY_API_KEY_ENV}",
"models": [_model_entry(model_id, limits) for model_id in model_ids],
}
_MODELS_FILE_ADAPTER: Final = TypeAdapter(dict[str, JsonValue])
def sync_models_json(
path: Path,
base_url: str,
model_ids: tuple[str, ...],
limits: Mapping[str, ModelLimits] = _NO_LIMITS,
) -> PiSyncError | None:
"""Replace only the litellm provider entry, leaving the rest of the file intact."""
try:
current: Final = _MODELS_FILE_ADAPTER.validate_json(path.read_text()) if path.exists() else {}
except (OSError, ValidationError) as e:
return PiSyncError(f"Could not read {path} as a JSON object: {e}. Fix or move the file, then retry.")
existing_providers: Final = current.get("providers", {})
if not isinstance(existing_providers, dict):
return PiSyncError(f'"providers" in {path} is not an object; fix or move the file, then retry.')
updated: Final = {
**current,
"providers": {**existing_providers, PI_PROVIDER_NAME: provider_block(base_url, model_ids, limits)},
}
try:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(updated, indent=2) + "\n")
except OSError as e:
return PiSyncError(f"Could not write {path}: {e}")
return None
__all__ = [
"LITELLM_PROXY_API_KEY_ENV",
"PI_CONFIG_DIR_ENV",
"PI_PROVIDER_NAME",
"ModelLimits",
"PiSyncError",
"fetch_model_ids",
"fetch_model_limits",
"models_json_path",
"provider_block",
"sync_models_json",
]

147
tests/test_litellm/proxy/client/cli/test_agents.py
View file

@ -34,6 +34,15 @@ class _FakeResponse:
self.status_code = status_code
class _FakeJsonResponse:
def __init__(self, status_code, payload=None):
self.status_code = status_code
self._payload = payload
def json(self):
return self._payload
class TestAgentProfile:
def test_claude_is_anthropic(self):
name, profiles = agent_profile("claude")
@ -49,6 +58,9 @@ class TestAgentProfile:
assert agent_profile("codex") == ("Codex", frozenset({"openai"}))
assert agent_profile("opencode") == ("OpenCode", frozenset({"openai"}))
def test_pi_is_litellm(self):
assert agent_profile("pi") == ("pi", frozenset({"litellm"}))
def test_unknown_command_gets_both_profiles(self):
name, profiles = agent_profile("mytool")
assert name == "mytool"
@ -91,6 +103,15 @@ class TestBuildAgentEnv:
assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key"
assert env["OPENAI_API_KEY"] == "sk-key"
def test_litellm_profile_exports_only_the_proxy_key(self):
env = build_agent_env(
{}, "http://localhost:4000/", "sk-key", frozenset({"litellm"})
)
assert env["LITELLM_PROXY_API_KEY"] == "sk-key"
assert "ANTHROPIC_BASE_URL" not in env
assert "OPENAI_BASE_URL" not in env
assert "OPENAI_API_KEY" not in env
def test_preserves_unrelated_env_and_does_not_mutate_input(self):
base = {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"}
env = build_agent_env(
@ -123,6 +144,9 @@ class TestAgentLaunchArgs:
agent_launch_args("codex", "http://localhost:4000")
)
def test_pi_gets_no_static_args(self):
assert agent_launch_args("pi", "http://localhost:4000") == []
class TestVerifyProxyKey:
def test_ok_status_passes_and_uses_models_endpoint(self):
@ -223,6 +247,127 @@ class TestRunAgent:
# overrides must precede the codex subcommand so codex parses them
assert args.index('model_provider="litellm"') < args.index("exec")
def test_pi_preparer_runs_after_verify_and_before_launch(self):
order = []
captured = {}
def fake_prepare(base_url, api_key, base_env):
order.append("prepare")
captured["args"] = (base_url, api_key, dict(base_env))
return []
run_agent(
"http://localhost:4000",
"sk-key",
["pi"],
base_env={"HOME": "/home/u"},
which=lambda name: "/usr/local/bin/pi",
verify=lambda *a: order.append("verify"),
launcher=lambda *a: order.append("launch"),
preparers={"pi": fake_prepare},
)
assert order == ["verify", "prepare", "launch"]
assert captured["args"] == (
"http://localhost:4000",
"sk-key",
{"HOME": "/home/u"},
)
def test_pi_prepared_args_precede_user_args_and_env_has_proxy_key(self):
calls = {}
run_agent(
"http://localhost:4000",
"sk-key",
["pi", "-p", "hello"],
base_env={},
which=lambda name: "/usr/local/bin/pi",
verify=lambda *a: None,
launcher=lambda p, a, e: calls.update(args=tuple(a), env=dict(e)),
preparers={"pi": lambda *a: ["--model", "litellm/m-1"]},
)
# user args come last so a user-supplied --model wins in pi's parser
assert calls["args"] == ("pi", "--model", "litellm/m-1", "-p", "hello")
assert calls["env"]["LITELLM_PROXY_API_KEY"] == "sk-key"
assert "OPENAI_API_KEY" not in calls["env"]
assert "ANTHROPIC_BASE_URL" not in calls["env"]
def test_failed_preparer_aborts_before_launch(self):
launched = []
def boom(*a):
raise AgentRunError("sync failed")
with pytest.raises(AgentRunError, match="sync failed"):
run_agent(
"http://localhost:4000",
"sk-key",
["pi"],
base_env={},
which=lambda name: "/usr/local/bin/pi",
verify=lambda *a: None,
launcher=lambda *a: launched.append(a),
preparers={"pi": boom},
)
assert launched == []
def test_prepare_pi_syncs_models_json_and_pins_first_model(self, tmp_path):
from litellm.proxy.client.cli.commands.agents import prepare_pi
def fake_get(url, headers, timeout):
if url.endswith("/model_group/info"):
return _FakeJsonResponse(
200,
{"data": [{"model_group": "m-first", "max_input_tokens": 131072, "max_output_tokens": 8192}]},
)
return _FakeJsonResponse(200, {"data": [{"id": "m-first"}, {"id": "m-second"}]})
pin = prepare_pi(
"http://localhost:4000",
"sk-key",
{"PI_CODING_AGENT_DIR": str(tmp_path)},
get=fake_get,
)
assert pin == ["--model", "litellm/m-first"]
import json
written = json.loads((tmp_path / "models.json").read_text())
assert written["providers"]["litellm"]["apiKey"] == "$LITELLM_PROXY_API_KEY"
assert written["providers"]["litellm"]["models"] == [
{"id": "m-first", "contextWindow": 131072, "maxTokens": 8192},
{"id": "m-second"},
]
def test_prepare_pi_surfaces_fetch_failure_as_agent_error(self, tmp_path):
from litellm.proxy.client.cli.commands.agents import prepare_pi
with pytest.raises(AgentRunError, match="HTTP 500"):
prepare_pi(
"http://localhost:4000",
"sk-key",
{"PI_CODING_AGENT_DIR": str(tmp_path)},
get=lambda *a, **k: _FakeJsonResponse(500),
)
def test_claude_has_no_preparer(self):
prepared = []
def fake_prepare(*a):
prepared.append(a)
return []
run_agent(
"http://localhost:4000",
"sk-key",
["claude"],
base_env={},
which=lambda name: "/usr/local/bin/claude",
verify=lambda *a: None,
launcher=lambda *a: None,
preparers={"pi": fake_prepare},
)
assert prepared == []
def test_claude_launches_without_injected_args(self):
calls = {}
run_agent(
@ -319,7 +464,7 @@ class TestAgentCommands:
self.runner = CliRunner()
def test_one_command_per_known_agent(self):
assert {c.name for c in agent_commands()} == {"claude", "codex", "opencode"}
assert {c.name for c in agent_commands()} == {"claude", "codex", "opencode", "pi"}
def test_claude_launches_with_stored_key_and_forwards_args(self):
captured = {}

201
tests/test_litellm/proxy/client/cli/test_pi.py Normal file
View file

@ -0,0 +1,201 @@
import json
from pathlib import Path
import requests
from litellm.proxy.client.cli.commands.pi import (
ModelLimits,
PiSyncError,
fetch_model_ids,
fetch_model_limits,
models_json_path,
provider_block,
sync_models_json,
)
class _FakeResponse:
def __init__(self, status_code, payload=None):
self.status_code = status_code
self._payload = payload
def json(self):
if self._payload is None:
raise ValueError("not json")
return self._payload
class TestFetchModelIds:
def test_returns_ids_in_proxy_order_deduped(self):
captured = {}
def fake_get(url, headers, timeout):
captured["url"] = url
captured["headers"] = headers
return _FakeResponse(
200,
{"data": [{"id": "m-b"}, {"id": "m-a"}, {"id": "m-b"}]},
)
assert fetch_model_ids("http://localhost:4000/", "sk-key", get=fake_get) == ("m-b", "m-a")
assert captured["url"] == "http://localhost:4000/v1/models"
assert captured["headers"] == {"Authorization": "Bearer sk-key"}
def test_network_error_is_a_value(self):
def boom(*a, **k):
raise requests.ConnectionError("refused")
result = fetch_model_ids("http://localhost:4000", "sk-key", get=boom)
assert isinstance(result, PiSyncError)
assert "Could not list models" in result.message
def test_non_200_is_a_value(self):
result = fetch_model_ids(
"http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500)
)
assert isinstance(result, PiSyncError)
assert "HTTP 500" in result.message
def test_malformed_body_is_a_value(self):
result = fetch_model_ids(
"http://localhost:4000",
"sk-key",
get=lambda *a, **k: _FakeResponse(200, {"data": "nope"}),
)
assert isinstance(result, PiSyncError)
def test_empty_model_list_is_a_value(self):
result = fetch_model_ids(
"http://localhost:4000",
"sk-key",
get=lambda *a, **k: _FakeResponse(200, {"data": []}),
)
assert isinstance(result, PiSyncError)
assert "no models" in result.message
class TestFetchModelLimits:
def test_maps_group_limits_and_hits_model_group_info(self):
captured = {}
def fake_get(url, headers, timeout):
captured["url"] = url
return _FakeResponse(
200,
{
"data": [
{"model_group": "m-a", "max_input_tokens": 131072, "max_output_tokens": 8192},
{"model_group": "m-b", "max_input_tokens": None, "max_output_tokens": None},
]
},
)
limits = fetch_model_limits("http://localhost:4000/", "sk-key", get=fake_get)
assert captured["url"] == "http://localhost:4000/model_group/info"
assert limits["m-a"] == ModelLimits(context_window=131072, max_tokens=8192)
assert limits["m-b"] == ModelLimits(context_window=None, max_tokens=None)
def test_non_200_degrades_to_no_limits(self):
assert fetch_model_limits("http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(403)) == {}
def test_network_error_degrades_to_no_limits(self):
def boom(*a, **k):
raise requests.ConnectionError("refused")
assert fetch_model_limits("http://localhost:4000", "sk-key", get=boom) == {}
def test_malformed_body_degrades_to_no_limits(self):
assert (
fetch_model_limits(
"http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(200, {"data": "nope"})
)
== {}
)
class TestModelsJsonPath:
def test_env_override_wins(self):
assert models_json_path({"PI_CODING_AGENT_DIR": "/custom/dir"}) == Path("/custom/dir/models.json")
def test_defaults_to_home_pi_agent(self):
assert models_json_path({}) == Path.home() / ".pi" / "agent" / "models.json"
class TestProviderBlock:
def test_points_pi_at_proxy_with_env_interpolated_key(self):
block = provider_block("http://localhost:4000/", ("m-1", "m-2"))
assert block == {
"baseUrl": "http://localhost:4000/v1",
"api": "openai-completions",
"apiKey": "$LITELLM_PROXY_API_KEY",
"models": [{"id": "m-1"}, {"id": "m-2"}],
}
def test_known_limits_become_context_window_and_max_tokens(self):
block = provider_block(
"http://localhost:4000",
("m-1", "m-2"),
{
"m-1": ModelLimits(context_window=131072, max_tokens=8192),
"m-2": ModelLimits(context_window=None, max_tokens=None),
},
)
assert block["models"] == [
{"id": "m-1", "contextWindow": 131072, "maxTokens": 8192},
{"id": "m-2"},
]
class TestSyncModelsJson:
def test_creates_file_and_parent_dirs(self, tmp_path):
path = tmp_path / "agent" / "models.json"
assert sync_models_json(path, "http://localhost:4000", ("m-1",)) is None
written = json.loads(path.read_text())
assert written["providers"]["litellm"]["baseUrl"] == "http://localhost:4000/v1"
assert written["providers"]["litellm"]["models"] == [{"id": "m-1"}]
def test_preserves_other_providers_and_top_level_keys(self, tmp_path):
path = tmp_path / "models.json"
path.write_text(
json.dumps(
{
"somethingElse": True,
"providers": {
"ollama": {"baseUrl": "http://localhost:11434/v1"},
"litellm": {"baseUrl": "http://stale:1234/v1", "models": []},
},
}
)
)
assert sync_models_json(path, "http://localhost:4000", ("m-1",)) is None
written = json.loads(path.read_text())
assert written["somethingElse"] is True
assert written["providers"]["ollama"] == {"baseUrl": "http://localhost:11434/v1"}
assert written["providers"]["litellm"]["baseUrl"] == "http://localhost:4000/v1"
assert written["providers"]["litellm"]["models"] == [{"id": "m-1"}]
def test_invalid_json_is_a_value_and_file_untouched(self, tmp_path):
path = tmp_path / "models.json"
path.write_text("{not json")
result = sync_models_json(path, "http://localhost:4000", ("m-1",))
assert isinstance(result, PiSyncError)
assert path.read_text() == "{not json"
def test_non_object_providers_is_a_value(self, tmp_path):
path = tmp_path / "models.json"
path.write_text(json.dumps({"providers": ["nope"]}))
result = sync_models_json(path, "http://localhost:4000", ("m-1",))
assert isinstance(result, PiSyncError)
def test_top_level_non_object_is_a_value(self, tmp_path):
path = tmp_path / "models.json"
path.write_text(json.dumps(["nope"]))
result = sync_models_json(path, "http://localhost:4000", ("m-1",))
assert isinstance(result, PiSyncError)
def test_unwritable_path_is_a_value(self, tmp_path):
blocker = tmp_path / "agent"
blocker.write_text("i am a file, not a directory")
result = sync_models_json(blocker / "models.json", "http://localhost:4000", ("m-1",))
assert isinstance(result, PiSyncError)
assert "Could not" in result.message