ai-chatbot/scripts/sync-atlascloud-models.ts

126 lines
3.4 KiB
TypeScript

import { and, eq, inArray } from "drizzle-orm";
import {
fetchAtlasCloudImagePresets,
resolveAtlasCloudDisplayName,
} from "../lib/ai/atlascloud-models";
import { db } from "../lib/db";
import { aiModel, aiProvider } from "../lib/db/schema";
async function main() {
const providers = await db
.select()
.from(aiProvider)
.where(and(eq(aiProvider.type, "atlascloud"), eq(aiProvider.isEnabled, true)));
if (providers.length === 0) {
throw new Error("没有找到已启用的 AtlasCloud 服务商");
}
const catalog = await fetchAtlasCloudImagePresets();
if (catalog.errors.length > 0) {
throw new Error(
`AtlasCloud 目录有 ${catalog.errors.length} 个 schema 加载失败:${catalog.errors
.map((item) => item.modelId)
.join(", ")}`
);
}
const summaries = [];
const catalogModelIds = new Set(
catalog.presets.map((preset) => preset.modelId)
);
for (const provider of providers) {
let added = 0;
let updated = 0;
for (const preset of catalog.presets) {
const existing = await db.query.aiModel.findFirst({
where: and(
eq(aiModel.providerId, provider.id),
eq(aiModel.modelId, preset.modelId)
),
});
const values = {
displayName: resolveAtlasCloudDisplayName(
existing?.displayName,
preset.displayName,
preset.modelId
),
avatarUrl: preset.avatarUrl || null,
description: preset.description,
supportsVision: preset.supportsVision,
supportsImageGen: preset.supportsImageGen,
isEnabled: true,
isHidden: false,
sortOrder: preset.sortOrder,
capabilities: preset.capabilities,
defaultParams: preset.defaultParams,
imageGenConfig: preset.imageGenConfig,
updatedAt: new Date(),
};
if (existing) {
await db
.update(aiModel)
.set(values)
.where(eq(aiModel.id, existing.id));
updated += 1;
} else {
await db.insert(aiModel).values({
providerId: provider.id,
modelId: preset.modelId,
...values,
isDefault: false,
isSubModel: false,
parentModelId: null,
});
added += 1;
}
}
const providerModels = await db
.select({ modelId: aiModel.modelId })
.from(aiModel)
.where(eq(aiModel.providerId, provider.id));
const staleModelIds = providerModels
.map((model) => model.modelId)
.filter((modelId) => !catalogModelIds.has(modelId));
if (staleModelIds.length > 0) {
await db
.update(aiModel)
.set({
isEnabled: false,
isHidden: true,
updatedAt: new Date(),
})
.where(
and(
eq(aiModel.providerId, provider.id),
inArray(aiModel.modelId, staleModelIds)
)
);
}
summaries.push({
providerId: provider.id,
providerName: provider.name,
added,
updated,
disabledStale: staleModelIds.length,
});
}
console.log(
JSON.stringify(
{
catalogModels: catalog.catalogCount,
availableModels: catalog.presets.length,
providers: summaries,
},
null,
2
)
);
}
void main().catch((error) => {
console.error(error instanceof Error ? error.message : String(error));
process.exitCode = 1;
});