Allow models.json to override built-in providers (#406)

* Allow models.json to override built-in providers

When a provider is defined in models.json with the same name as a
built-in provider (e.g., 'anthropic', 'google'), the built-in models
for that provider are completely replaced by the custom definition.

This enables users to:
- Use custom base URLs (proxies, self-hosted endpoints)
- Define a subset of models they want available
- Customize model configurations for built-in providers

Example usage in ~/.pi/agent/models.json:
{
  "providers": {
    "anthropic": {
      "baseUrl": "https://my-proxy.example.com/v1",
      "apiKey": "ANTHROPIC_API_KEY",
      "api": "anthropic-messages",
      "models": [...]
    }
  }
}

* Refactor model-registry for readability

- Extract CustomModelsResult type and emptyCustomModelsResult helper
- Extract loadBuiltInModels method with clear skip logic
- Simplify loadModels with destructuring and ternary
- Reduce repetition in error handling paths

* Refactor model-registry tests for readability

- Extract providerConfig() helper to hide irrelevant model fields
- Extract writeModelsJson() helper for file writing
- Extract getModelsForProvider() helper for filtering
- Move modelsJsonPath to beforeEach

Reduces test file from 262 to 130 lines while maintaining same coverage.
This commit is contained in:
Yevhen Bobrov 2026-01-03 01:59:59 +02:00 committed by GitHub
parent 6186e497c5
commit 243104fa18
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 173 additions and 33 deletions

View file

@ -0,0 +1,136 @@
import { existsSync, mkdirSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { afterEach, beforeEach, describe, expect, test } from "vitest";
import { AuthStorage } from "../src/core/auth-storage.js";
import { ModelRegistry } from "../src/core/model-registry.js";
describe("ModelRegistry", () => {
let tempDir: string;
let modelsJsonPath: string;
let authStorage: AuthStorage;
beforeEach(() => {
tempDir = join(tmpdir(), `pi-test-model-registry-${Date.now()}-${Math.random().toString(36).slice(2)}`);
mkdirSync(tempDir, { recursive: true });
modelsJsonPath = join(tempDir, "models.json");
authStorage = new AuthStorage(join(tempDir, "auth.json"));
});
afterEach(() => {
if (tempDir && existsSync(tempDir)) {
rmSync(tempDir, { recursive: true });
}
});
/** Create minimal provider config */
function providerConfig(
baseUrl: string,
models: Array<{ id: string; name?: string }>,
api: string = "anthropic-messages",
) {
return {
baseUrl,
apiKey: "TEST_KEY",
api,
models: models.map((m) => ({
id: m.id,
name: m.name ?? m.id,
reasoning: false,
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 100000,
maxTokens: 8000,
})),
};
}
function writeModelsJson(providers: Record<string, ReturnType<typeof providerConfig>>) {
writeFileSync(modelsJsonPath, JSON.stringify({ providers }));
}
function getModelsForProvider(registry: ModelRegistry, provider: string) {
return registry.getAll().filter((m) => m.provider === provider);
}
describe("provider override", () => {
test("custom provider with same name as built-in replaces built-in models", () => {
writeModelsJson({
anthropic: providerConfig("https://my-proxy.example.com/v1", [{ id: "claude-custom" }]),
});
const registry = new ModelRegistry(authStorage, modelsJsonPath);
const anthropicModels = getModelsForProvider(registry, "anthropic");
expect(anthropicModels).toHaveLength(1);
expect(anthropicModels[0].id).toBe("claude-custom");
expect(anthropicModels[0].baseUrl).toBe("https://my-proxy.example.com/v1");
});
test("custom provider with same name as built-in does not affect other built-in providers", () => {
writeModelsJson({
anthropic: providerConfig("https://my-proxy.example.com/v1", [{ id: "claude-custom" }]),
});
const registry = new ModelRegistry(authStorage, modelsJsonPath);
expect(getModelsForProvider(registry, "google").length).toBeGreaterThan(0);
expect(getModelsForProvider(registry, "openai").length).toBeGreaterThan(0);
});
test("multiple built-in providers can be overridden", () => {
writeModelsJson({
anthropic: providerConfig("https://anthropic-proxy.example.com/v1", [{ id: "claude-proxy" }]),
google: providerConfig("https://google-proxy.example.com/v1", [{ id: "gemini-proxy" }], "google-generative-ai"),
});
const registry = new ModelRegistry(authStorage, modelsJsonPath);
const anthropicModels = getModelsForProvider(registry, "anthropic");
const googleModels = getModelsForProvider(registry, "google");
expect(anthropicModels).toHaveLength(1);
expect(anthropicModels[0].id).toBe("claude-proxy");
expect(anthropicModels[0].baseUrl).toBe("https://anthropic-proxy.example.com/v1");
expect(googleModels).toHaveLength(1);
expect(googleModels[0].id).toBe("gemini-proxy");
expect(googleModels[0].baseUrl).toBe("https://google-proxy.example.com/v1");
});
test("refresh() reloads overrides from disk", () => {
writeModelsJson({
anthropic: providerConfig("https://first-proxy.example.com/v1", [{ id: "claude-first" }]),
});
const registry = new ModelRegistry(authStorage, modelsJsonPath);
expect(getModelsForProvider(registry, "anthropic")[0].id).toBe("claude-first");
// Update and refresh
writeModelsJson({
anthropic: providerConfig("https://second-proxy.example.com/v1", [{ id: "claude-second" }]),
});
registry.refresh();
const anthropicModels = getModelsForProvider(registry, "anthropic");
expect(anthropicModels[0].id).toBe("claude-second");
expect(anthropicModels[0].baseUrl).toBe("https://second-proxy.example.com/v1");
});
test("removing override from models.json restores built-in provider", () => {
writeModelsJson({
anthropic: providerConfig("https://proxy.example.com/v1", [{ id: "claude-custom" }]),
});
const registry = new ModelRegistry(authStorage, modelsJsonPath);
expect(getModelsForProvider(registry, "anthropic")).toHaveLength(1);
// Remove override and refresh
writeModelsJson({});
registry.refresh();
const anthropicModels = getModelsForProvider(registry, "anthropic");
expect(anthropicModels.length).toBeGreaterThan(1);
expect(anthropicModels.some((m) => m.id.includes("claude"))).toBe(true);
});
});
});