@mastra/core
Version:
197 lines (191 loc) • 8.24 kB
JavaScript
import { i as getCapabilityFileName, n as getGatewayId, r as shouldEnableGateway } from "./gateway-helpers-DusR3xFY.js";
import path from "path";
import fs from "fs/promises";
//#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.writeFile(tempPath, content, encoding);
await fs.rename(tempPath, filePath);
} catch (error) {
try {
await fs.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 (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(getGatewayId(gateway));
continue;
}
const gatewayId = 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.dirname(jsonPath);
const typesDir = path.dirname(typesPath);
await fs.mkdir(jsonDir, { recursive: true });
await fs.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.join(jsonDir, "capabilities");
await fs.rm(capDir, {
recursive: true,
force: true
});
await fs.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.join(capDir, getCapabilityFileName(provider)), JSON.stringify(capData, null, 2), "utf-8");
}
}
}
//#endregion
export { fetchProvidersFromGateways, writeRegistryFiles };
//# sourceMappingURL=registry-generator-BTyp-Kqp.js.map