diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index de9d38963c1..66beecbd2e6 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -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/`. 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/` 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 diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index dfc70a8df7c..f55cf9893e7 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -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", diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py new file mode 100644 index 00000000000..bd933874a10 --- /dev/null +++ b/litellm/proxy/client/cli/commands/pi.py @@ -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", +] diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index afd1696a89f..fad0d3842fa 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -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 = {} diff --git a/tests/test_litellm/proxy/client/cli/test_pi.py b/tests/test_litellm/proxy/client/cli/test_pi.py new file mode 100644 index 00000000000..99ed2734b76 --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_pi.py @@ -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