UNPKG

@mastra/core

Version:
201 lines (195 loc) • 8.7 kB
const require_rolldown_runtime = require("./rolldown-runtime-uwYp4b74.cjs"); const require_gateway_helpers = require("./gateway-helpers-CdK-tcyA.cjs"); let path = require("path"); path = require_rolldown_runtime.__toESM(path, 1); let fs_promises = require("fs/promises"); fs_promises = require_rolldown_runtime.__toESM(fs_promises, 1); //#region src/llm/model/registry-generator.ts /** * Shared provider registry generation logic * Used by both the CLI generation script and runtime refresh */ function hasAttachmentCapabilities(gateway) { return "getAttachmentCapabilities" in gateway && typeof gateway.getAttachmentCapabilities === "function"; } function hasTemperatureCapabilities(gateway) { return "getTemperatureCapabilities" in gateway && typeof gateway.getTemperatureCapabilities === "function"; } function hasStructuredOutputCapabilities(gateway) { return "getStructuredOutputCapabilities" in gateway && typeof gateway.getStructuredOutputCapabilities === "function"; } /** * Write a file atomically using the write-to-temp-then-rename pattern. * This prevents file corruption when multiple processes write to the same file concurrently. * * The rename operation is atomic on POSIX systems when source and destination * are on the same filesystem. * * @param filePath - The target file path * @param content - The content to write * @param encoding - The encoding to use (default: 'utf-8') */ async function atomicWriteFile(filePath, content, encoding = "utf-8") { const randomSuffix = Math.random().toString(36).substring(2, 15); const tempPath = `${filePath}.${process.pid}.${Date.now()}.${randomSuffix}.tmp`; try { await fs_promises.default.writeFile(tempPath, content, encoding); await fs_promises.default.rename(tempPath, filePath); } catch (error) { try { await fs_promises.default.unlink(tempPath); } catch {} throw error; } } /** * Fetch providers from all enabled gateways with silent retry logic. * Retries up to 3 times per gateway with exponential backoff. If all * retries are exhausted the gateway is silently skipped (no error logging) * since the bundled registry already contains all model data. * @param gateways - Array of gateway instances to fetch from * @returns Object containing providers and models records */ async function fetchProvidersFromGateways(gateways) { const enabledGateways = []; for (const gateway of gateways) if (require_gateway_helpers.shouldEnableGateway(gateway)) enabledGateways.push(gateway); const allProviders = {}; const allModels = {}; const allAttachmentCapabilities = {}; const allTemperatureCapabilities = {}; const allStructuredOutputCapabilities = {}; const failedGateways = []; const maxRetries = 3; for (const gateway of enabledGateways) { let providers = null; for (let attempt = 1; attempt <= maxRetries; attempt++) try { providers = await gateway.fetchProviders(); break; } catch { if (attempt < maxRetries) { const delayMs = Math.min(1e3 * Math.pow(2, attempt - 1), 5e3); await new Promise((resolve) => setTimeout(resolve, delayMs)); } } if (!providers) { failedGateways.push(require_gateway_helpers.getGatewayId(gateway)); continue; } const gatewayId = require_gateway_helpers.getGatewayId(gateway); const isProviderRegistry = gatewayId === "models.dev"; const gatewayAttachmentCaps = hasAttachmentCapabilities(gateway) ? gateway.getAttachmentCapabilities() : void 0; const gatewayTemperatureCaps = hasTemperatureCapabilities(gateway) ? gateway.getTemperatureCapabilities() : void 0; const gatewayStructuredOutputCaps = hasStructuredOutputCapabilities(gateway) ? gateway.getStructuredOutputCapabilities() : void 0; for (const [providerId, config] of Object.entries(providers)) { const typeProviderId = isProviderRegistry ? providerId : providerId === gatewayId ? gatewayId : `${gatewayId}/${providerId}`; allProviders[typeProviderId] = config; allModels[typeProviderId] = config.models.sort(); if (gatewayAttachmentCaps?.[providerId]) allAttachmentCapabilities[typeProviderId] = gatewayAttachmentCaps[providerId]; if (gatewayTemperatureCaps?.[providerId]) allTemperatureCapabilities[typeProviderId] = gatewayTemperatureCaps[providerId]; if (gatewayStructuredOutputCaps?.[providerId]) allStructuredOutputCapabilities[typeProviderId] = gatewayStructuredOutputCaps[providerId]; } } return { providers: allProviders, models: allModels, attachmentCapabilities: allAttachmentCapabilities, temperatureCapabilities: allTemperatureCapabilities, structuredOutputCapabilities: allStructuredOutputCapabilities, failedGateways }; } /** * Generate TypeScript type definitions content * @param models - Record of provider IDs to model arrays * @returns Generated TypeScript type definitions as a string */ function generateTypesContent(models) { return `/** * THIS FILE IS AUTO-GENERATED - DO NOT EDIT * Generated from model gateway providers */ /** * Provider models mapping type * This is derived from the JSON data and provides type-safe access */ export type ProviderModelsMap = { ${Object.entries(models).map(([provider, modelList]) => { const modelsList = modelList.map((m) => `'${m}'`); const providerKey = !/^[a-zA-Z_$][a-zA-Z0-9_$]*$/.test(provider) ? `'${provider}'` : provider; const singleLine = ` readonly ${providerKey}: readonly [${modelsList.join(", ")}];`; if (singleLine.length > 120) return ` readonly ${providerKey}: readonly [\n${modelList.map((m) => ` '${m}',`).join("\n")}\n ];`; return singleLine; }).join("\n")} }; /** * Union type of all registered provider IDs */ export type Provider = keyof ProviderModelsMap; /** * Provider models mapping interface */ export interface ProviderModels { [key: string]: string[]; } /** * OpenAI-compatible model ID type * Dynamically derived from ProviderModelsMap * Full provider/model paths (e.g., "openai/gpt-4o", "anthropic/claude-3-5-sonnet-20241022") */ export type ModelRouterModelId = | { [P in Provider]: \`\${P}/\${ProviderModelsMap[P][number]}\`; }[Provider] | \`mastra/\${ProviderModelsMap['openrouter'][number]}\` | (string & {}); /** * Extract the model part from a ModelRouterModelId for a specific provider * Dynamically derived from ProviderModelsMap * Example: ModelForProvider<'openai'> = 'gpt-4o' | 'gpt-4-turbo' | ... */ export type ModelForProvider<P extends Provider> = ProviderModelsMap[P][number]; `; } /** * Write registry files to disk (JSON and .d.ts) * @param jsonPath - Path to write the JSON file * @param typesPath - Path to write the .d.ts file * @param providers - Provider configurations * @param models - Model lists by provider */ async function writeRegistryFiles(jsonPath, typesPath, providers, models, attachmentCapabilities, temperatureCapabilities, structuredOutputCapabilities) { const jsonDir = path.default.dirname(jsonPath); const typesDir = path.default.dirname(typesPath); await fs_promises.default.mkdir(jsonDir, { recursive: true }); await fs_promises.default.mkdir(typesDir, { recursive: true }); await atomicWriteFile(jsonPath, JSON.stringify({ providers, models, version: "1.0.0" }, null, 2), "utf-8"); await atomicWriteFile(typesPath, generateTypesContent(models), "utf-8"); if (attachmentCapabilities && Object.keys(attachmentCapabilities).length > 0 || temperatureCapabilities && Object.keys(temperatureCapabilities).length > 0 || structuredOutputCapabilities && Object.keys(structuredOutputCapabilities).length > 0) { const capDir = path.default.join(jsonDir, "capabilities"); await fs_promises.default.rm(capDir, { recursive: true, force: true }); await fs_promises.default.mkdir(capDir, { recursive: true }); const allProviderIds = /* @__PURE__ */ new Set([ ...attachmentCapabilities ? Object.keys(attachmentCapabilities) : [], ...temperatureCapabilities ? Object.keys(temperatureCapabilities) : [], ...structuredOutputCapabilities ? Object.keys(structuredOutputCapabilities) : [] ]); for (const provider of allProviderIds) { const capData = {}; if (attachmentCapabilities?.[provider]) capData.attachment = attachmentCapabilities[provider]; if (temperatureCapabilities?.[provider]) capData.temperature = temperatureCapabilities[provider]; if (structuredOutputCapabilities?.[provider]) capData.structuredOutput = structuredOutputCapabilities[provider]; await atomicWriteFile(path.default.join(capDir, require_gateway_helpers.getCapabilityFileName(provider)), JSON.stringify(capData, null, 2), "utf-8"); } } } //#endregion exports.fetchProvidersFromGateways = fetchProvidersFromGateways; exports.writeRegistryFiles = writeRegistryFiles; //# sourceMappingURL=registry-generator-DZsB-A6S.cjs.map