245 lines
8.5 KiB
TypeScript
245 lines
8.5 KiB
TypeScript
import fs from 'node:fs';
|
|
import path from 'node:path';
|
|
import type { VaultStore } from '@waggle/core';
|
|
import type { FastifyInstance } from 'fastify';
|
|
import {
|
|
discoverProviderModels,
|
|
PROVIDER_MODEL_CATALOGS,
|
|
type DiscoveryOptions,
|
|
type DiscoveredProviderModel,
|
|
} from './provider-model-catalog.js';
|
|
import { applyProviderKeyToEnv, getProviderApiKeys } from './provider-env.js';
|
|
import { startLiteLLM, stopLiteLLM } from './lifecycle.js';
|
|
|
|
interface LiteLLMProviderRoute {
|
|
envName: string;
|
|
modelPrefix: string;
|
|
apiBase?: string;
|
|
}
|
|
|
|
interface LiteLLMModelEntry {
|
|
model_name: string;
|
|
litellm_params: {
|
|
model: string;
|
|
api_key: string;
|
|
api_base?: string;
|
|
};
|
|
}
|
|
|
|
export interface LiteLLMRuntimeConfigResult {
|
|
configPath: string | null;
|
|
modelIds: string[];
|
|
unavailableProviders: string[];
|
|
}
|
|
|
|
export interface LiteLLMRefreshResult {
|
|
managed: boolean;
|
|
ready: boolean;
|
|
port: number;
|
|
models: string[];
|
|
unavailableProviders: string[];
|
|
error?: string;
|
|
}
|
|
|
|
/** Routing metadata only. Model identities always come from provider APIs. */
|
|
const LITELLM_PROVIDER_ROUTES: Record<string, LiteLLMProviderRoute> = {
|
|
anthropic: { envName: 'ANTHROPIC_API_KEY', modelPrefix: 'anthropic' },
|
|
openai: { envName: 'OPENAI_API_KEY', modelPrefix: 'openai' },
|
|
google: { envName: 'GEMINI_API_KEY', modelPrefix: 'gemini' },
|
|
deepseek: { envName: 'DEEPSEEK_API_KEY', modelPrefix: 'deepseek' },
|
|
xai: { envName: 'XAI_API_KEY', modelPrefix: 'xai' },
|
|
mistral: { envName: 'MISTRAL_API_KEY', modelPrefix: 'mistral' },
|
|
alibaba: {
|
|
envName: 'DASHSCOPE_API_KEY',
|
|
modelPrefix: 'openai',
|
|
apiBase: 'https://dashscope-intl.aliyuncs.com/compatible-mode/v1',
|
|
},
|
|
minimax: {
|
|
envName: 'MINIMAX_API_KEY',
|
|
modelPrefix: 'openai',
|
|
apiBase: 'https://api.minimax.io/v1',
|
|
},
|
|
zhipu: {
|
|
envName: 'ZHIPU_API_KEY',
|
|
modelPrefix: 'openai',
|
|
apiBase: 'https://open.bigmodel.cn/api/paas/v4',
|
|
},
|
|
moonshot: {
|
|
envName: 'MOONSHOT_API_KEY',
|
|
modelPrefix: 'openai',
|
|
apiBase: 'https://api.moonshot.ai/v1',
|
|
},
|
|
perplexity: { envName: 'PERPLEXITY_API_KEY', modelPrefix: 'perplexity' },
|
|
openrouter: { envName: 'OPENROUTER_API_KEY', modelPrefix: 'openrouter' },
|
|
};
|
|
|
|
function routeModel(providerId: string, model: DiscoveredProviderModel): string {
|
|
const route = LITELLM_PROVIDER_ROUTES[providerId];
|
|
const providerPrefix = `${providerId}/`;
|
|
const providerModelId = model.id.startsWith(providerPrefix)
|
|
? model.id.slice(providerPrefix.length)
|
|
: model.id;
|
|
return `${route.modelPrefix}/${providerModelId}`;
|
|
}
|
|
|
|
export function buildLiteLLMRuntimeConfig(
|
|
catalogs: ReadonlyMap<string, DiscoveredProviderModel[]>,
|
|
customBaseUrls: ReadonlyMap<string, string> = new Map(),
|
|
): { model_list: LiteLLMModelEntry[] } {
|
|
const modelList: LiteLLMModelEntry[] = [];
|
|
for (const [providerId, models] of catalogs) {
|
|
const route = LITELLM_PROVIDER_ROUTES[providerId];
|
|
if (!route) continue;
|
|
for (const model of models) {
|
|
const apiBase = customBaseUrls.get(providerId) ?? route.apiBase;
|
|
modelList.push({
|
|
model_name: model.id,
|
|
litellm_params: {
|
|
model: routeModel(providerId, model),
|
|
api_key: `os.environ/${route.envName}`,
|
|
...(apiBase ? { api_base: apiBase } : {}),
|
|
},
|
|
});
|
|
}
|
|
}
|
|
return { model_list: modelList };
|
|
}
|
|
|
|
/** Build the executable router catalog from the same live APIs used by the UI. */
|
|
export async function prepareLiteLLMRuntimeConfig(
|
|
dataDir: string,
|
|
vault: VaultStore,
|
|
discoveryOptions: DiscoveryOptions = {},
|
|
): Promise<LiteLLMRuntimeConfigResult> {
|
|
const catalogs = new Map<string, DiscoveredProviderModel[]>();
|
|
const customBaseUrls = new Map<string, string>();
|
|
const unavailableProviders: string[] = [];
|
|
|
|
await Promise.all(Object.keys(PROVIDER_MODEL_CATALOGS).map(async (providerId) => {
|
|
const entry = vault.get(providerId);
|
|
const apiKeys = getProviderApiKeys(providerId, vault);
|
|
if (apiKeys.length === 0) return;
|
|
const baseUrl = typeof entry?.metadata?.baseUrl === 'string' ? entry.metadata.baseUrl : undefined;
|
|
if (baseUrl) customBaseUrls.set(providerId, baseUrl);
|
|
let result = await discoverProviderModels(providerId, apiKeys[0], baseUrl, discoveryOptions);
|
|
let workingKey = result.status === 'unavailable' ? null : apiKeys[0];
|
|
for (let index = 1; !workingKey && index < apiKeys.length; index += 1) {
|
|
result = await discoverProviderModels(providerId, apiKeys[index], baseUrl, discoveryOptions);
|
|
if (result.status !== 'unavailable') workingKey = apiKeys[index];
|
|
}
|
|
// Keep aliases such as GEMINI_API_KEY / GOOGLE_API_KEY aligned to the key
|
|
// that actually passed provider discovery. Otherwise a stale machine-level
|
|
// alias can poison LiteLLM even though another configured alias is valid.
|
|
if (workingKey) applyProviderKeyToEnv(providerId, workingKey, true);
|
|
if (result.models.length > 0) catalogs.set(providerId, result.models);
|
|
if (result.status === 'unavailable') unavailableProviders.push(providerId);
|
|
}));
|
|
|
|
const config = buildLiteLLMRuntimeConfig(catalogs, customBaseUrls);
|
|
const modelIds = config.model_list.map((entry) => entry.model_name);
|
|
if (modelIds.length === 0) {
|
|
return { configPath: null, modelIds, unavailableProviders };
|
|
}
|
|
|
|
const configPath = path.join(dataDir, 'litellm.runtime.json');
|
|
const tempPath = `${configPath}.tmp`;
|
|
fs.writeFileSync(tempPath, JSON.stringify(config, null, 2), { mode: 0o600 });
|
|
try {
|
|
fs.renameSync(tempPath, configPath);
|
|
} catch {
|
|
try { fs.unlinkSync(configPath); } catch { /* first write */ }
|
|
fs.renameSync(tempPath, configPath);
|
|
}
|
|
return { configPath, modelIds, unavailableProviders };
|
|
}
|
|
|
|
/** Rebuild and restart the managed router after credentials or catalogs change. */
|
|
export async function refreshManagedLiteLLM(server: FastifyInstance): Promise<LiteLLMRefreshResult> {
|
|
const port = server.localConfig.managedLiteLLMPort ?? 4000;
|
|
if (!server.localConfig.manageLiteLLM) {
|
|
return {
|
|
managed: false,
|
|
ready: false,
|
|
port,
|
|
models: [],
|
|
unavailableProviders: [],
|
|
error: 'LiteLLM lifecycle is managed outside this server.',
|
|
};
|
|
}
|
|
|
|
const runtime = await prepareLiteLLMRuntimeConfig(server.localConfig.dataDir, server.vault);
|
|
if (!runtime.configPath) {
|
|
return {
|
|
managed: true,
|
|
ready: false,
|
|
port,
|
|
models: runtime.modelIds,
|
|
unavailableProviders: runtime.unavailableProviders,
|
|
error: 'No provider model catalog is currently available.',
|
|
};
|
|
}
|
|
|
|
try {
|
|
await stopLiteLLM();
|
|
} catch {
|
|
// The prior child may already have exited. startLiteLLM checks the port.
|
|
}
|
|
const status = await startLiteLLM(port, runtime.configPath);
|
|
const ready = status.status === 'running' || status.status === 'started';
|
|
if (ready) {
|
|
server.localConfig.litellmUrl = `http://localhost:${port}`;
|
|
server.agentState.litellmApiKey = process.env.LITELLM_API_KEY
|
|
?? process.env.LITELLM_MASTER_KEY
|
|
?? 'sk-waggle-dev';
|
|
server.agentState.llmProvider = {
|
|
provider: 'litellm',
|
|
health: 'healthy',
|
|
detail: `LiteLLM on port ${port} (${runtime.modelIds.length} provider models)`,
|
|
checkedAt: new Date().toISOString(),
|
|
};
|
|
}
|
|
|
|
return {
|
|
managed: true,
|
|
ready,
|
|
port,
|
|
models: runtime.modelIds,
|
|
unavailableProviders: runtime.unavailableProviders,
|
|
...(!ready ? {
|
|
error: status.error ?? (status.status === 'timeout'
|
|
? 'LiteLLM did not start in time'
|
|
: 'LiteLLM did not become ready.'),
|
|
} : {}),
|
|
};
|
|
}
|
|
|
|
function readManagedModelIds(dataDir: string): string[] {
|
|
const configPath = path.join(dataDir, 'litellm.runtime.json');
|
|
try {
|
|
const parsed = JSON.parse(fs.readFileSync(configPath, 'utf8')) as {
|
|
model_list?: Array<{ model_name?: unknown }>;
|
|
};
|
|
return (parsed.model_list ?? [])
|
|
.map((entry) => entry.model_name)
|
|
.filter((modelName): modelName is string => typeof modelName === 'string' && modelName.length > 0);
|
|
} catch {
|
|
return [];
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Make a provider-discovered model executable in the already-running managed
|
|
* router. Existing models are a zero-cost read; only a genuinely new model
|
|
* causes catalog regeneration and one router restart.
|
|
*/
|
|
export async function ensureManagedLiteLLMModel(
|
|
server: FastifyInstance,
|
|
model: string,
|
|
): Promise<boolean> {
|
|
if (!server.localConfig.manageLiteLLM || model.startsWith('ollama/')) return true;
|
|
if (readManagedModelIds(server.localConfig.dataDir).includes(model)) return true;
|
|
|
|
const refresh = await refreshManagedLiteLLM(server);
|
|
return refresh.ready && refresh.models.includes(model);
|
|
}
|