119 lines
3.2 KiB
TypeScript
119 lines
3.2 KiB
TypeScript
import { and, eq, inArray } from "drizzle-orm";
|
|
import { fetchAtlasCloudImagePresets } 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: preset.displayName,
|
|
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;
|
|
});
|