update
This commit is contained in:
ZacharyZhang-NY committed
1 parent
9779c4d1ca
commit
1b578c111a
719 files changed
+191593
-98
No files matched your search
@@ -0,0 +1,844 @@
|
||||
/**
|
||||
* Model registry - manages built-in and custom models, provides API key resolution.
|
||||
*/
|
||||
|
||||
import {
|
||||
type Api,
|
||||
type AssistantMessageEventStream,
|
||||
type Context,
|
||||
getModels,
|
||||
getProviders,
|
||||
type KnownProvider,
|
||||
type Model,
|
||||
type OAuthProviderInterface,
|
||||
type OpenAICompletionsCompat,
|
||||
type OpenAIResponsesCompat,
|
||||
registerApiProvider,
|
||||
resetApiProviders,
|
||||
type SimpleStreamOptions,
|
||||
} from "@mariozechner/pi-ai";
|
||||
import { registerOAuthProvider, resetOAuthProviders } from "@mariozechner/pi-ai/oauth";
|
||||
import { type Static, Type } from "@sinclair/typebox";
|
||||
import AjvModule from "ajv";
|
||||
import { existsSync, readFileSync } from "fs";
|
||||
import { join } from "path";
|
||||
import { getAgentDir } from "../config.js";
|
||||
import type { AuthStorage } from "./auth-storage.js";
|
||||
import {
|
||||
clearConfigValueCache,
|
||||
resolveConfigValueOrThrow,
|
||||
resolveConfigValueUncached,
|
||||
resolveHeadersOrThrow,
|
||||
} from "./resolve-config-value.js";
|
||||
|
||||
const Ajv = (AjvModule as any).default || AjvModule;
|
||||
const ajv = new Ajv();
|
||||
|
||||
// Schema for OpenRouter routing preferences
|
||||
const PercentileCutoffsSchema = Type.Object({
|
||||
p50: Type.Optional(Type.Number()),
|
||||
p75: Type.Optional(Type.Number()),
|
||||
p90: Type.Optional(Type.Number()),
|
||||
p99: Type.Optional(Type.Number()),
|
||||
});
|
||||
|
||||
const OpenRouterRoutingSchema = Type.Object({
|
||||
allow_fallbacks: Type.Optional(Type.Boolean()),
|
||||
require_parameters: Type.Optional(Type.Boolean()),
|
||||
data_collection: Type.Optional(Type.Union([Type.Literal("deny"), Type.Literal("allow")])),
|
||||
zdr: Type.Optional(Type.Boolean()),
|
||||
enforce_distillable_text: Type.Optional(Type.Boolean()),
|
||||
order: Type.Optional(Type.Array(Type.String())),
|
||||
only: Type.Optional(Type.Array(Type.String())),
|
||||
ignore: Type.Optional(Type.Array(Type.String())),
|
||||
quantizations: Type.Optional(Type.Array(Type.String())),
|
||||
sort: Type.Optional(
|
||||
Type.Union([
|
||||
Type.String(),
|
||||
Type.Object({
|
||||
by: Type.Optional(Type.String()),
|
||||
partition: Type.Optional(Type.Union([Type.String(), Type.Null()])),
|
||||
}),
|
||||
]),
|
||||
),
|
||||
max_price: Type.Optional(
|
||||
Type.Object({
|
||||
prompt: Type.Optional(Type.Union([Type.Number(), Type.String()])),
|
||||
completion: Type.Optional(Type.Union([Type.Number(), Type.String()])),
|
||||
image: Type.Optional(Type.Union([Type.Number(), Type.String()])),
|
||||
audio: Type.Optional(Type.Union([Type.Number(), Type.String()])),
|
||||
request: Type.Optional(Type.Union([Type.Number(), Type.String()])),
|
||||
}),
|
||||
),
|
||||
preferred_min_throughput: Type.Optional(Type.Union([Type.Number(), PercentileCutoffsSchema])),
|
||||
preferred_max_latency: Type.Optional(Type.Union([Type.Number(), PercentileCutoffsSchema])),
|
||||
});
|
||||
|
||||
// Schema for Vercel AI Gateway routing preferences
|
||||
const VercelGatewayRoutingSchema = Type.Object({
|
||||
only: Type.Optional(Type.Array(Type.String())),
|
||||
order: Type.Optional(Type.Array(Type.String())),
|
||||
});
|
||||
|
||||
// Schema for OpenAI compatibility settings
|
||||
const ReasoningEffortMapSchema = Type.Object({
|
||||
minimal: Type.Optional(Type.String()),
|
||||
low: Type.Optional(Type.String()),
|
||||
medium: Type.Optional(Type.String()),
|
||||
high: Type.Optional(Type.String()),
|
||||
xhigh: Type.Optional(Type.String()),
|
||||
});
|
||||
|
||||
const OpenAICompletionsCompatSchema = Type.Object({
|
||||
supportsStore: Type.Optional(Type.Boolean()),
|
||||
supportsDeveloperRole: Type.Optional(Type.Boolean()),
|
||||
supportsReasoningEffort: Type.Optional(Type.Boolean()),
|
||||
reasoningEffortMap: Type.Optional(ReasoningEffortMapSchema),
|
||||
supportsUsageInStreaming: Type.Optional(Type.Boolean()),
|
||||
maxTokensField: Type.Optional(Type.Union([Type.Literal("max_completion_tokens"), Type.Literal("max_tokens")])),
|
||||
requiresToolResultName: Type.Optional(Type.Boolean()),
|
||||
requiresAssistantAfterToolResult: Type.Optional(Type.Boolean()),
|
||||
requiresThinkingAsText: Type.Optional(Type.Boolean()),
|
||||
thinkingFormat: Type.Optional(
|
||||
Type.Union([
|
||||
Type.Literal("openai"),
|
||||
Type.Literal("openrouter"),
|
||||
Type.Literal("zai"),
|
||||
Type.Literal("qwen"),
|
||||
Type.Literal("qwen-chat-template"),
|
||||
]),
|
||||
),
|
||||
openRouterRouting: Type.Optional(OpenRouterRoutingSchema),
|
||||
vercelGatewayRouting: Type.Optional(VercelGatewayRoutingSchema),
|
||||
supportsStrictMode: Type.Optional(Type.Boolean()),
|
||||
});
|
||||
|
||||
const OpenAIResponsesCompatSchema = Type.Object({
|
||||
// Reserved for future use
|
||||
});
|
||||
|
||||
const OpenAICompatSchema = Type.Union([OpenAICompletionsCompatSchema, OpenAIResponsesCompatSchema]);
|
||||
|
||||
// Schema for custom model definition
|
||||
// Most fields are optional with sensible defaults for local models (Ollama, LM Studio, etc.)
|
||||
const ModelDefinitionSchema = Type.Object({
|
||||
id: Type.String({ minLength: 1 }),
|
||||
name: Type.Optional(Type.String({ minLength: 1 })),
|
||||
api: Type.Optional(Type.String({ minLength: 1 })),
|
||||
baseUrl: Type.Optional(Type.String({ minLength: 1 })),
|
||||
reasoning: Type.Optional(Type.Boolean()),
|
||||
input: Type.Optional(Type.Array(Type.Union([Type.Literal("text"), Type.Literal("image")]))),
|
||||
cost: Type.Optional(
|
||||
Type.Object({
|
||||
input: Type.Number(),
|
||||
output: Type.Number(),
|
||||
cacheRead: Type.Number(),
|
||||
cacheWrite: Type.Number(),
|
||||
}),
|
||||
),
|
||||
contextWindow: Type.Optional(Type.Number()),
|
||||
maxTokens: Type.Optional(Type.Number()),
|
||||
headers: Type.Optional(Type.Record(Type.String(), Type.String())),
|
||||
compat: Type.Optional(OpenAICompatSchema),
|
||||
});
|
||||
|
||||
// Schema for per-model overrides (all fields optional, merged with built-in model)
|
||||
const ModelOverrideSchema = Type.Object({
|
||||
name: Type.Optional(Type.String({ minLength: 1 })),
|
||||
reasoning: Type.Optional(Type.Boolean()),
|
||||
input: Type.Optional(Type.Array(Type.Union([Type.Literal("text"), Type.Literal("image")]))),
|
||||
cost: Type.Optional(
|
||||
Type.Object({
|
||||
input: Type.Optional(Type.Number()),
|
||||
output: Type.Optional(Type.Number()),
|
||||
cacheRead: Type.Optional(Type.Number()),
|
||||
cacheWrite: Type.Optional(Type.Number()),
|
||||
}),
|
||||
),
|
||||
contextWindow: Type.Optional(Type.Number()),
|
||||
maxTokens: Type.Optional(Type.Number()),
|
||||
headers: Type.Optional(Type.Record(Type.String(), Type.String())),
|
||||
compat: Type.Optional(OpenAICompatSchema),
|
||||
});
|
||||
|
||||
type ModelOverride = Static<typeof ModelOverrideSchema>;
|
||||
|
||||
const ProviderConfigSchema = Type.Object({
|
||||
baseUrl: Type.Optional(Type.String({ minLength: 1 })),
|
||||
apiKey: Type.Optional(Type.String({ minLength: 1 })),
|
||||
api: Type.Optional(Type.String({ minLength: 1 })),
|
||||
headers: Type.Optional(Type.Record(Type.String(), Type.String())),
|
||||
compat: Type.Optional(OpenAICompatSchema),
|
||||
authHeader: Type.Optional(Type.Boolean()),
|
||||
models: Type.Optional(Type.Array(ModelDefinitionSchema)),
|
||||
modelOverrides: Type.Optional(Type.Record(Type.String(), ModelOverrideSchema)),
|
||||
});
|
||||
|
||||
const ModelsConfigSchema = Type.Object({
|
||||
providers: Type.Record(Type.String(), ProviderConfigSchema),
|
||||
});
|
||||
|
||||
ajv.addSchema(ModelsConfigSchema, "ModelsConfig");
|
||||
|
||||
type ModelsConfig = Static<typeof ModelsConfigSchema>;
|
||||
|
||||
/** Provider override config (baseUrl, compat) without request auth/headers */
|
||||
interface ProviderOverride {
|
||||
baseUrl?: string;
|
||||
compat?: Model<Api>["compat"];
|
||||
}
|
||||
|
||||
interface ProviderRequestConfig {
|
||||
apiKey?: string;
|
||||
headers?: Record<string, string>;
|
||||
authHeader?: boolean;
|
||||
}
|
||||
|
||||
export type ResolvedRequestAuth =
|
||||
| {
|
||||
ok: true;
|
||||
apiKey?: string;
|
||||
headers?: Record<string, string>;
|
||||
}
|
||||
| {
|
||||
ok: false;
|
||||
error: string;
|
||||
};
|
||||
|
||||
/** Result of loading custom models from models.json */
|
||||
interface CustomModelsResult {
|
||||
models: Model<Api>[];
|
||||
/** Providers with baseUrl/headers/apiKey overrides for built-in models */
|
||||
overrides: Map<string, ProviderOverride>;
|
||||
/** Per-model overrides: provider -> modelId -> override */
|
||||
modelOverrides: Map<string, Map<string, ModelOverride>>;
|
||||
error: string | undefined;
|
||||
}
|
||||
|
||||
function emptyCustomModelsResult(error?: string): CustomModelsResult {
|
||||
return { models: [], overrides: new Map(), modelOverrides: new Map(), error };
|
||||
}
|
||||
|
||||
function mergeCompat(
|
||||
baseCompat: Model<Api>["compat"],
|
||||
overrideCompat: ModelOverride["compat"],
|
||||
): Model<Api>["compat"] | undefined {
|
||||
if (!overrideCompat) return baseCompat;
|
||||
|
||||
const base = baseCompat as OpenAICompletionsCompat | OpenAIResponsesCompat | undefined;
|
||||
const override = overrideCompat as OpenAICompletionsCompat | OpenAIResponsesCompat;
|
||||
const merged = { ...base, ...override } as OpenAICompletionsCompat | OpenAIResponsesCompat;
|
||||
|
||||
const baseCompletions = base as OpenAICompletionsCompat | undefined;
|
||||
const overrideCompletions = override as OpenAICompletionsCompat;
|
||||
const mergedCompletions = merged as OpenAICompletionsCompat;
|
||||
|
||||
if (baseCompletions?.openRouterRouting || overrideCompletions.openRouterRouting) {
|
||||
mergedCompletions.openRouterRouting = {
|
||||
...baseCompletions?.openRouterRouting,
|
||||
...overrideCompletions.openRouterRouting,
|
||||
};
|
||||
}
|
||||
|
||||
if (baseCompletions?.vercelGatewayRouting || overrideCompletions.vercelGatewayRouting) {
|
||||
mergedCompletions.vercelGatewayRouting = {
|
||||
...baseCompletions?.vercelGatewayRouting,
|
||||
...overrideCompletions.vercelGatewayRouting,
|
||||
};
|
||||
}
|
||||
|
||||
return merged as Model<Api>["compat"];
|
||||
}
|
||||
|
||||
/**
|
||||
* Deep merge a model override into a model.
|
||||
* Handles nested objects (cost, compat) by merging rather than replacing.
|
||||
*/
|
||||
function applyModelOverride(model: Model<Api>, override: ModelOverride): Model<Api> {
|
||||
const result = { ...model };
|
||||
|
||||
// Simple field overrides
|
||||
if (override.name !== undefined) result.name = override.name;
|
||||
if (override.reasoning !== undefined) result.reasoning = override.reasoning;
|
||||
if (override.input !== undefined) result.input = override.input as ("text" | "image")[];
|
||||
if (override.contextWindow !== undefined) result.contextWindow = override.contextWindow;
|
||||
if (override.maxTokens !== undefined) result.maxTokens = override.maxTokens;
|
||||
|
||||
// Merge cost (partial override)
|
||||
if (override.cost) {
|
||||
result.cost = {
|
||||
input: override.cost.input ?? model.cost.input,
|
||||
output: override.cost.output ?? model.cost.output,
|
||||
cacheRead: override.cost.cacheRead ?? model.cost.cacheRead,
|
||||
cacheWrite: override.cost.cacheWrite ?? model.cost.cacheWrite,
|
||||
};
|
||||
}
|
||||
|
||||
// Deep merge compat
|
||||
result.compat = mergeCompat(model.compat, override.compat);
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
/** Clear the config value command cache. Exported for testing. */
|
||||
export const clearApiKeyCache = clearConfigValueCache;
|
||||
|
||||
/**
|
||||
* Model registry - loads and manages models, resolves API keys via AuthStorage.
|
||||
*/
|
||||
export class ModelRegistry {
|
||||
private models: Model<Api>[] = [];
|
||||
private providerRequestConfigs: Map<string, ProviderRequestConfig> = new Map();
|
||||
private modelRequestHeaders: Map<string, Record<string, string>> = new Map();
|
||||
private registeredProviders: Map<string, ProviderConfigInput> = new Map();
|
||||
private loadError: string | undefined = undefined;
|
||||
|
||||
private constructor(
|
||||
readonly authStorage: AuthStorage,
|
||||
private modelsJsonPath: string | undefined,
|
||||
) {
|
||||
this.loadModels();
|
||||
}
|
||||
|
||||
static create(authStorage: AuthStorage, modelsJsonPath: string = join(getAgentDir(), "models.json")): ModelRegistry {
|
||||
return new ModelRegistry(authStorage, modelsJsonPath);
|
||||
}
|
||||
|
||||
static inMemory(authStorage: AuthStorage): ModelRegistry {
|
||||
return new ModelRegistry(authStorage, undefined);
|
||||
}
|
||||
|
||||
/**
|
||||
* Reload models from disk (built-in + custom from models.json).
|
||||
*/
|
||||
refresh(): void {
|
||||
this.providerRequestConfigs.clear();
|
||||
this.modelRequestHeaders.clear();
|
||||
this.loadError = undefined;
|
||||
|
||||
// Ensure dynamic API/OAuth registrations are rebuilt from current provider state.
|
||||
resetApiProviders();
|
||||
resetOAuthProviders();
|
||||
|
||||
this.loadModels();
|
||||
|
||||
for (const [providerName, config] of this.registeredProviders.entries()) {
|
||||
this.applyProviderConfig(providerName, config);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get any error from loading models.json (undefined if no error).
|
||||
*/
|
||||
getError(): string | undefined {
|
||||
return this.loadError;
|
||||
}
|
||||
|
||||
private loadModels(): void {
|
||||
// Load custom models and overrides from models.json
|
||||
const {
|
||||
models: customModels,
|
||||
overrides,
|
||||
modelOverrides,
|
||||
error,
|
||||
} = this.modelsJsonPath ? this.loadCustomModels(this.modelsJsonPath) : emptyCustomModelsResult();
|
||||
|
||||
if (error) {
|
||||
this.loadError = error;
|
||||
// Keep built-in models even if custom models failed to load
|
||||
}
|
||||
|
||||
const builtInModels = this.loadBuiltInModels(overrides, modelOverrides);
|
||||
let combined = this.mergeCustomModels(builtInModels, customModels);
|
||||
|
||||
// Let OAuth providers modify their models (e.g., update baseUrl)
|
||||
for (const oauthProvider of this.authStorage.getOAuthProviders()) {
|
||||
const cred = this.authStorage.get(oauthProvider.id);
|
||||
if (cred?.type === "oauth" && oauthProvider.modifyModels) {
|
||||
combined = oauthProvider.modifyModels(combined, cred);
|
||||
}
|
||||
}
|
||||
|
||||
this.models = combined;
|
||||
}
|
||||
|
||||
/** Load built-in models and apply provider/model overrides */
|
||||
private loadBuiltInModels(
|
||||
overrides: Map<string, ProviderOverride>,
|
||||
modelOverrides: Map<string, Map<string, ModelOverride>>,
|
||||
): Model<Api>[] {
|
||||
return getProviders().flatMap((provider) => {
|
||||
const models = getModels(provider as KnownProvider) as Model<Api>[];
|
||||
const providerOverride = overrides.get(provider);
|
||||
const perModelOverrides = modelOverrides.get(provider);
|
||||
|
||||
return models.map((m) => {
|
||||
let model = m;
|
||||
|
||||
// Apply provider-level baseUrl/headers/compat override
|
||||
if (providerOverride) {
|
||||
model = {
|
||||
...model,
|
||||
baseUrl: providerOverride.baseUrl ?? model.baseUrl,
|
||||
compat: mergeCompat(model.compat, providerOverride.compat),
|
||||
};
|
||||
}
|
||||
|
||||
// Apply per-model override
|
||||
const modelOverride = perModelOverrides?.get(m.id);
|
||||
if (modelOverride) {
|
||||
model = applyModelOverride(model, modelOverride);
|
||||
}
|
||||
|
||||
return model;
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** Merge custom models into built-in list by provider+id (custom wins on conflicts). */
|
||||
private mergeCustomModels(builtInModels: Model<Api>[], customModels: Model<Api>[]): Model<Api>[] {
|
||||
const merged = [...builtInModels];
|
||||
for (const customModel of customModels) {
|
||||
const existingIndex = merged.findIndex((m) => m.provider === customModel.provider && m.id === customModel.id);
|
||||
if (existingIndex >= 0) {
|
||||
merged[existingIndex] = customModel;
|
||||
} else {
|
||||
merged.push(customModel);
|
||||
}
|
||||
}
|
||||
return merged;
|
||||
}
|
||||
|
||||
private loadCustomModels(modelsJsonPath: string): CustomModelsResult {
|
||||
if (!existsSync(modelsJsonPath)) {
|
||||
return emptyCustomModelsResult();
|
||||
}
|
||||
|
||||
try {
|
||||
const content = readFileSync(modelsJsonPath, "utf-8");
|
||||
const config: ModelsConfig = JSON.parse(content);
|
||||
|
||||
// Validate schema
|
||||
const validate = ajv.getSchema("ModelsConfig")!;
|
||||
if (!validate(config)) {
|
||||
const errors =
|
||||
validate.errors?.map((e: any) => ` - ${e.instancePath || "root"}: ${e.message}`).join("\n") ||
|
||||
"Unknown schema error";
|
||||
return emptyCustomModelsResult(`Invalid models.json schema:\n${errors}\n\nFile: ${modelsJsonPath}`);
|
||||
}
|
||||
|
||||
// Additional validation
|
||||
this.validateConfig(config);
|
||||
|
||||
const overrides = new Map<string, ProviderOverride>();
|
||||
const modelOverrides = new Map<string, Map<string, ModelOverride>>();
|
||||
|
||||
for (const [providerName, providerConfig] of Object.entries(config.providers)) {
|
||||
if (providerConfig.baseUrl || providerConfig.compat) {
|
||||
overrides.set(providerName, {
|
||||
baseUrl: providerConfig.baseUrl,
|
||||
compat: providerConfig.compat,
|
||||
});
|
||||
}
|
||||
|
||||
this.storeProviderRequestConfig(providerName, providerConfig);
|
||||
|
||||
if (providerConfig.modelOverrides) {
|
||||
modelOverrides.set(providerName, new Map(Object.entries(providerConfig.modelOverrides)));
|
||||
for (const [modelId, modelOverride] of Object.entries(providerConfig.modelOverrides)) {
|
||||
this.storeModelHeaders(providerName, modelId, modelOverride.headers);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return { models: this.parseModels(config), overrides, modelOverrides, error: undefined };
|
||||
} catch (error) {
|
||||
if (error instanceof SyntaxError) {
|
||||
return emptyCustomModelsResult(`Failed to parse models.json: ${error.message}\n\nFile: ${modelsJsonPath}`);
|
||||
}
|
||||
return emptyCustomModelsResult(
|
||||
`Failed to load models.json: ${error instanceof Error ? error.message : error}\n\nFile: ${modelsJsonPath}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
private validateConfig(config: ModelsConfig): void {
|
||||
const builtInProviders = new Set<string>(getProviders());
|
||||
|
||||
for (const [providerName, providerConfig] of Object.entries(config.providers)) {
|
||||
const isBuiltIn = builtInProviders.has(providerName);
|
||||
const hasProviderApi = !!providerConfig.api;
|
||||
const models = providerConfig.models ?? [];
|
||||
const hasModelOverrides =
|
||||
providerConfig.modelOverrides && Object.keys(providerConfig.modelOverrides).length > 0;
|
||||
|
||||
if (models.length === 0) {
|
||||
// Override-only config: needs baseUrl, compat, modelOverrides, or some combination.
|
||||
if (!providerConfig.baseUrl && !providerConfig.compat && !hasModelOverrides) {
|
||||
throw new Error(
|
||||
`Provider ${providerName}: must specify "baseUrl", "compat", "modelOverrides", or "models".`,
|
||||
);
|
||||
}
|
||||
} else if (!isBuiltIn) {
|
||||
// Non-built-in providers with custom models require endpoint + auth.
|
||||
if (!providerConfig.baseUrl) {
|
||||
throw new Error(`Provider ${providerName}: "baseUrl" is required when defining custom models.`);
|
||||
}
|
||||
if (!providerConfig.apiKey) {
|
||||
throw new Error(`Provider ${providerName}: "apiKey" is required when defining custom models.`);
|
||||
}
|
||||
}
|
||||
// Built-in providers with custom models: baseUrl/apiKey/api are optional,
|
||||
// inherited from built-in models. Auth comes from env vars / auth storage.
|
||||
|
||||
for (const modelDef of models) {
|
||||
const hasModelApi = !!modelDef.api;
|
||||
|
||||
if (!hasProviderApi && !hasModelApi && !isBuiltIn) {
|
||||
throw new Error(
|
||||
`Provider ${providerName}, model ${modelDef.id}: no "api" specified. Set at provider or model level.`,
|
||||
);
|
||||
}
|
||||
// For built-in providers, api is optional — inherited from built-in models.
|
||||
|
||||
if (!modelDef.id) throw new Error(`Provider ${providerName}: model missing "id"`);
|
||||
// Validate contextWindow/maxTokens only if provided (they have defaults)
|
||||
if (modelDef.contextWindow !== undefined && modelDef.contextWindow <= 0)
|
||||
throw new Error(`Provider ${providerName}, model ${modelDef.id}: invalid contextWindow`);
|
||||
if (modelDef.maxTokens !== undefined && modelDef.maxTokens <= 0)
|
||||
throw new Error(`Provider ${providerName}, model ${modelDef.id}: invalid maxTokens`);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private parseModels(config: ModelsConfig): Model<Api>[] {
|
||||
const models: Model<Api>[] = [];
|
||||
const builtInProviders = new Set<string>(getProviders());
|
||||
|
||||
// Cache built-in defaults (api, baseUrl) per provider, extracted from first model.
|
||||
const builtInDefaultsCache = new Map<string, { api: string; baseUrl: string }>();
|
||||
const getBuiltInDefaults = (providerName: string): { api: string; baseUrl: string } | undefined => {
|
||||
if (!builtInProviders.has(providerName)) return undefined;
|
||||
if (builtInDefaultsCache.has(providerName)) return builtInDefaultsCache.get(providerName);
|
||||
const builtIn = getModels(providerName as KnownProvider) as Model<Api>[];
|
||||
if (builtIn.length === 0) return undefined;
|
||||
const defaults = { api: builtIn[0].api, baseUrl: builtIn[0].baseUrl };
|
||||
builtInDefaultsCache.set(providerName, defaults);
|
||||
return defaults;
|
||||
};
|
||||
|
||||
for (const [providerName, providerConfig] of Object.entries(config.providers)) {
|
||||
const modelDefs = providerConfig.models ?? [];
|
||||
if (modelDefs.length === 0) continue; // Override-only, no custom models
|
||||
|
||||
const builtInDefaults = getBuiltInDefaults(providerName);
|
||||
|
||||
for (const modelDef of modelDefs) {
|
||||
const api = modelDef.api ?? providerConfig.api ?? builtInDefaults?.api;
|
||||
if (!api) continue;
|
||||
|
||||
const baseUrl = modelDef.baseUrl ?? providerConfig.baseUrl ?? builtInDefaults?.baseUrl;
|
||||
if (!baseUrl) continue;
|
||||
|
||||
const compat = mergeCompat(providerConfig.compat, modelDef.compat);
|
||||
this.storeModelHeaders(providerName, modelDef.id, modelDef.headers);
|
||||
|
||||
const defaultCost = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 };
|
||||
models.push({
|
||||
id: modelDef.id,
|
||||
name: modelDef.name ?? modelDef.id,
|
||||
api: api as Api,
|
||||
provider: providerName,
|
||||
baseUrl,
|
||||
reasoning: modelDef.reasoning ?? false,
|
||||
input: (modelDef.input ?? ["text"]) as ("text" | "image")[],
|
||||
cost: modelDef.cost ?? defaultCost,
|
||||
contextWindow: modelDef.contextWindow ?? 128000,
|
||||
maxTokens: modelDef.maxTokens ?? 16384,
|
||||
headers: undefined,
|
||||
compat,
|
||||
} as Model<Api>);
|
||||
}
|
||||
}
|
||||
|
||||
return models;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get all models (built-in + custom).
|
||||
* If models.json had errors, returns only built-in models.
|
||||
*/
|
||||
getAll(): Model<Api>[] {
|
||||
return this.models;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get only models that have auth configured.
|
||||
* This is a fast check that doesn't refresh OAuth tokens.
|
||||
*/
|
||||
getAvailable(): Model<Api>[] {
|
||||
return this.models.filter((m) => this.hasConfiguredAuth(m));
|
||||
}
|
||||
|
||||
/**
|
||||
* Find a model by provider and ID.
|
||||
*/
|
||||
find(provider: string, modelId: string): Model<Api> | undefined {
|
||||
return this.models.find((m) => m.provider === provider && m.id === modelId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get API key for a model.
|
||||
*/
|
||||
hasConfiguredAuth(model: Model<Api>): boolean {
|
||||
return (
|
||||
this.authStorage.hasAuth(model.provider) ||
|
||||
this.providerRequestConfigs.get(model.provider)?.apiKey !== undefined
|
||||
);
|
||||
}
|
||||
|
||||
private getModelRequestKey(provider: string, modelId: string): string {
|
||||
return `${provider}:${modelId}`;
|
||||
}
|
||||
|
||||
private storeProviderRequestConfig(
|
||||
providerName: string,
|
||||
config: {
|
||||
apiKey?: string;
|
||||
headers?: Record<string, string>;
|
||||
authHeader?: boolean;
|
||||
},
|
||||
): void {
|
||||
if (!config.apiKey && !config.headers && !config.authHeader) {
|
||||
return;
|
||||
}
|
||||
|
||||
this.providerRequestConfigs.set(providerName, {
|
||||
apiKey: config.apiKey,
|
||||
headers: config.headers,
|
||||
authHeader: config.authHeader,
|
||||
});
|
||||
}
|
||||
|
||||
private storeModelHeaders(providerName: string, modelId: string, headers?: Record<string, string>): void {
|
||||
const key = this.getModelRequestKey(providerName, modelId);
|
||||
if (!headers || Object.keys(headers).length === 0) {
|
||||
this.modelRequestHeaders.delete(key);
|
||||
return;
|
||||
}
|
||||
this.modelRequestHeaders.set(key, headers);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get API key and request headers for a model.
|
||||
*/
|
||||
async getApiKeyAndHeaders(model: Model<Api>): Promise<ResolvedRequestAuth> {
|
||||
try {
|
||||
const providerConfig = this.providerRequestConfigs.get(model.provider);
|
||||
const apiKeyFromAuthStorage = await this.authStorage.getApiKey(model.provider, { includeFallback: false });
|
||||
const apiKey =
|
||||
apiKeyFromAuthStorage ??
|
||||
(providerConfig?.apiKey
|
||||
? resolveConfigValueOrThrow(providerConfig.apiKey, `API key for provider "${model.provider}"`)
|
||||
: undefined);
|
||||
|
||||
const providerHeaders = resolveHeadersOrThrow(providerConfig?.headers, `provider "${model.provider}"`);
|
||||
const modelHeaders = resolveHeadersOrThrow(
|
||||
this.modelRequestHeaders.get(this.getModelRequestKey(model.provider, model.id)),
|
||||
`model "${model.provider}/${model.id}"`,
|
||||
);
|
||||
|
||||
let headers =
|
||||
model.headers || providerHeaders || modelHeaders
|
||||
? { ...model.headers, ...providerHeaders, ...modelHeaders }
|
||||
: undefined;
|
||||
|
||||
if (providerConfig?.authHeader) {
|
||||
if (!apiKey) {
|
||||
return { ok: false, error: `No API key found for "${model.provider}"` };
|
||||
}
|
||||
headers = { ...headers, Authorization: `Bearer ${apiKey}` };
|
||||
}
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
apiKey,
|
||||
headers: headers && Object.keys(headers).length > 0 ? headers : undefined,
|
||||
};
|
||||
} catch (error) {
|
||||
return {
|
||||
ok: false,
|
||||
error: error instanceof Error ? error.message : String(error),
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get API key for a provider.
|
||||
*/
|
||||
async getApiKeyForProvider(provider: string): Promise<string | undefined> {
|
||||
const apiKey = await this.authStorage.getApiKey(provider, { includeFallback: false });
|
||||
if (apiKey !== undefined) {
|
||||
return apiKey;
|
||||
}
|
||||
|
||||
const providerApiKey = this.providerRequestConfigs.get(provider)?.apiKey;
|
||||
return providerApiKey ? resolveConfigValueUncached(providerApiKey) : undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a model is using OAuth credentials (subscription).
|
||||
*/
|
||||
isUsingOAuth(model: Model<Api>): boolean {
|
||||
const cred = this.authStorage.get(model.provider);
|
||||
return cred?.type === "oauth";
|
||||
}
|
||||
|
||||
/**
|
||||
* Register a provider dynamically (from extensions).
|
||||
*
|
||||
* If provider has models: replaces all existing models for this provider.
|
||||
* If provider has only baseUrl/headers: overrides existing models' URLs.
|
||||
* If provider has oauth: registers OAuth provider for /login support.
|
||||
*/
|
||||
registerProvider(providerName: string, config: ProviderConfigInput): void {
|
||||
this.validateProviderConfig(providerName, config);
|
||||
this.applyProviderConfig(providerName, config);
|
||||
this.registeredProviders.set(providerName, config);
|
||||
}
|
||||
|
||||
/**
|
||||
* Unregister a previously registered provider.
|
||||
*
|
||||
* Removes the provider from the registry and reloads models from disk so that
|
||||
* built-in models overridden by this provider are restored to their original state.
|
||||
* Also resets dynamic OAuth and API stream registrations before reapplying
|
||||
* remaining dynamic providers.
|
||||
* Has no effect if the provider was never registered.
|
||||
*/
|
||||
unregisterProvider(providerName: string): void {
|
||||
if (!this.registeredProviders.has(providerName)) return;
|
||||
this.registeredProviders.delete(providerName);
|
||||
this.refresh();
|
||||
}
|
||||
|
||||
private validateProviderConfig(providerName: string, config: ProviderConfigInput): void {
|
||||
if (config.streamSimple && !config.api) {
|
||||
throw new Error(`Provider ${providerName}: "api" is required when registering streamSimple.`);
|
||||
}
|
||||
|
||||
if (!config.models || config.models.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (!config.baseUrl) {
|
||||
throw new Error(`Provider ${providerName}: "baseUrl" is required when defining models.`);
|
||||
}
|
||||
if (!config.apiKey && !config.oauth) {
|
||||
throw new Error(`Provider ${providerName}: "apiKey" or "oauth" is required when defining models.`);
|
||||
}
|
||||
|
||||
for (const modelDef of config.models) {
|
||||
const api = modelDef.api || config.api;
|
||||
if (!api) {
|
||||
throw new Error(`Provider ${providerName}, model ${modelDef.id}: no "api" specified.`);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private applyProviderConfig(providerName: string, config: ProviderConfigInput): void {
|
||||
// Register OAuth provider if provided
|
||||
if (config.oauth) {
|
||||
// Ensure the OAuth provider ID matches the provider name
|
||||
const oauthProvider: OAuthProviderInterface = {
|
||||
...config.oauth,
|
||||
id: providerName,
|
||||
};
|
||||
registerOAuthProvider(oauthProvider);
|
||||
}
|
||||
|
||||
if (config.streamSimple) {
|
||||
const streamSimple = config.streamSimple;
|
||||
registerApiProvider(
|
||||
{
|
||||
api: config.api!,
|
||||
stream: (model, context, options) => streamSimple(model, context, options as SimpleStreamOptions),
|
||||
streamSimple,
|
||||
},
|
||||
`provider:${providerName}`,
|
||||
);
|
||||
}
|
||||
|
||||
this.storeProviderRequestConfig(providerName, config);
|
||||
|
||||
if (config.models && config.models.length > 0) {
|
||||
// Full replacement: remove existing models for this provider
|
||||
this.models = this.models.filter((m) => m.provider !== providerName);
|
||||
|
||||
// Parse and add new models
|
||||
for (const modelDef of config.models) {
|
||||
const api = modelDef.api || config.api;
|
||||
this.storeModelHeaders(providerName, modelDef.id, modelDef.headers);
|
||||
|
||||
this.models.push({
|
||||
id: modelDef.id,
|
||||
name: modelDef.name,
|
||||
api: api as Api,
|
||||
provider: providerName,
|
||||
baseUrl: config.baseUrl!,
|
||||
reasoning: modelDef.reasoning,
|
||||
input: modelDef.input as ("text" | "image")[],
|
||||
cost: modelDef.cost,
|
||||
contextWindow: modelDef.contextWindow,
|
||||
maxTokens: modelDef.maxTokens,
|
||||
headers: undefined,
|
||||
compat: modelDef.compat,
|
||||
} as Model<Api>);
|
||||
}
|
||||
|
||||
// Apply OAuth modifyModels if credentials exist (e.g., to update baseUrl)
|
||||
if (config.oauth?.modifyModels) {
|
||||
const cred = this.authStorage.get(providerName);
|
||||
if (cred?.type === "oauth") {
|
||||
this.models = config.oauth.modifyModels(this.models, cred);
|
||||
}
|
||||
}
|
||||
} else if (config.baseUrl || config.headers) {
|
||||
// Override-only: update baseUrl for existing models. Request headers are resolved per request.
|
||||
this.models = this.models.map((m) => {
|
||||
if (m.provider !== providerName) return m;
|
||||
return {
|
||||
...m,
|
||||
baseUrl: config.baseUrl ?? m.baseUrl,
|
||||
};
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Input type for registerProvider API.
|
||||
*/
|
||||
export interface ProviderConfigInput {
|
||||
baseUrl?: string;
|
||||
apiKey?: string;
|
||||
api?: Api;
|
||||
streamSimple?: (model: Model<Api>, context: Context, options?: SimpleStreamOptions) => AssistantMessageEventStream;
|
||||
headers?: Record<string, string>;
|
||||
authHeader?: boolean;
|
||||
/** OAuth provider for /login support */
|
||||
oauth?: Omit<OAuthProviderInterface, "id">;
|
||||
models?: Array<{
|
||||
id: string;
|
||||
name: string;
|
||||
api?: Api;
|
||||
baseUrl?: string;
|
||||
reasoning: boolean;
|
||||
input: ("text" | "image")[];
|
||||
cost: { input: number; output: number; cacheRead: number; cacheWrite: number };
|
||||
contextWindow: number;
|
||||
maxTokens: number;
|
||||
headers?: Record<string, string>;
|
||||
compat?: Model<Api>["compat"];
|
||||
}>;
|
||||
}
|
||||
Reference in new issue
Block a user