335 lines
15 KiB
TypeScript
335 lines
15 KiB
TypeScript
import type { AgentVisionExecutor } from "../runtime/agent-vision-failover";
|
||
import { resolveUserImageModel, applyAgentImagePreferences, agentProviderImageParameters, type AgentImagePreferences } from "./image-preferences";
|
||
import { tool, type LanguageModel, type ToolSet } from "ai";
|
||
import { z } from "zod";
|
||
import { createCanvasCapabilityTools } from "@/lib/canvas/capabilities/ai-sdk-tools";
|
||
import { CANVAS_CAPABILITY_REGISTRY } from "@/lib/canvas/capabilities/registry";
|
||
import { applyCanvasOperationPlan } from "@/lib/canvas/agent-mutations";
|
||
import { generateCanvasImage } from "@/lib/canvas/generate-image";
|
||
import { saveGeneratedImages } from "@/lib/ai/runtime/generated-image-persistence";
|
||
import { routeToModel } from "@/lib/ai/model-router";
|
||
import {
|
||
getModelGenerationCreditCost,
|
||
type DiscountableImageGenConfigLike,
|
||
} from "@/lib/ai/image-model-pricing";
|
||
import { checkUserImageQuota, getPointsSystemConfig, getUserById } from "@/lib/db/queries";
|
||
import {
|
||
getStyleProImageModelEntries,
|
||
sortStyleProImageModelEntries,
|
||
} from "@/lib/style-pro/model-options";
|
||
import type { CanvasAgentLoopPorts, CanvasAgentTool } from "./loop";
|
||
import { readCanvasImages, readOwnedCanvas } from "./canvas";
|
||
import { analyzeRuntimeImages } from "./review";
|
||
import { isExplicitDestructiveIntent } from "./capability-access";
|
||
import { registerImageCanvasAccessTools } from "./image-access-tools";
|
||
|
||
export async function buildCanvasRuntimeTools(params: {
|
||
userId: string;
|
||
chatId: string;
|
||
projectId: string;
|
||
runId: string;
|
||
assistantMessageId: string;
|
||
selectedNodeIds: string[];
|
||
request: Request;
|
||
model: LanguageModel;
|
||
supportsVision: boolean;
|
||
withModel?: AgentVisionExecutor;
|
||
imageSize?: string;
|
||
preferredModelId?: string;
|
||
imagePreferences?: AgentImagePreferences;
|
||
emit: (type: string, data: unknown) => void;
|
||
objective?: string;
|
||
operationSelectedNodeIds?: string[];
|
||
}) {
|
||
const { nodes } = await readOwnedCanvas(params.userId, params.projectId);
|
||
const imageNodeIds = nodes.filter((node) => node.type === "image").map((node) => node.id);
|
||
const defaults = params.selectedNodeIds.filter((id) => imageNodeIds.includes(id));
|
||
const sdk = await createCanvasCapabilityTools({
|
||
projectId: params.projectId,
|
||
userId: params.userId,
|
||
selectedNodeIds: defaults,
|
||
allowExplicitTargets: true,
|
||
defaultModelId: params.preferredModelId,
|
||
imagePreferences: params.imagePreferences,
|
||
authorizeDestructive: (ids) =>
|
||
isExplicitDestructiveIntent(params.objective || "") &&
|
||
ids.length > 0 &&
|
||
ids.every((id) => (params.operationSelectedNodeIds || defaults).includes(id)),
|
||
requestHeaders: params.request.headers,
|
||
requestUrl: params.request.url,
|
||
requestIdPrefix: params.runId,
|
||
callbacks: {
|
||
onClientCommand: (command) => params.emit("data-canvasCapabilityCommand", command),
|
||
onResult: (result, capability) => {
|
||
if (result.operation)
|
||
params.emit("data-canvasOperation", {
|
||
projectId: params.projectId,
|
||
operationId: result.task.requestId,
|
||
operation: capability.id,
|
||
status: "applied",
|
||
...result.operation,
|
||
});
|
||
},
|
||
},
|
||
});
|
||
const executable: Record<string, CanvasAgentTool> = {};
|
||
const definitions: ToolSet = {};
|
||
for (const [name, entry] of Object.entries(sdk)) {
|
||
const capability = CANVAS_CAPABILITY_REGISTRY.find((item) => item.agent.toolName === name)!;
|
||
definitions[name] = { ...entry, execute: undefined };
|
||
executable[name] = {
|
||
title: capability.title,
|
||
readOnly: capability.risk === "read",
|
||
concurrentSafe: false,
|
||
execute: async (input, callId) => {
|
||
if (!defaults.length && !Array.isArray(input.nodeIds))
|
||
throw new Error("请明确提供目标图片节点，不能猜测第一张图。");
|
||
// Generic redraw is the same constrained generator, including batch count,
|
||
// provider parameters and the no-silent-model-failover policy.
|
||
if (capability.id === "image.prompt_regenerate") {
|
||
const args = (input.args || {}) as Record<string, unknown>;
|
||
return executable.generateCanvasImages.execute({
|
||
...args, referenceNodeIds: input.nodeIds || defaults,
|
||
}, callId);
|
||
}
|
||
const result = await entry.execute!(input as never, { toolCallId: callId, messages: [] });
|
||
return result as Record<string, unknown>;
|
||
},
|
||
};
|
||
}
|
||
function register<T extends z.ZodTypeAny>(
|
||
name: string,
|
||
title: string,
|
||
description: string,
|
||
schema: T,
|
||
execute: (input: z.infer<T>, callId: string) => Promise<Record<string, unknown>>,
|
||
readOnly = true,
|
||
) {
|
||
definitions[name] = tool({ description, inputSchema: schema });
|
||
executable[name] = {
|
||
title,
|
||
readOnly,
|
||
concurrentSafe: readOnly,
|
||
execute: (input, callId) => execute(schema.parse(input), callId),
|
||
};
|
||
}
|
||
register("listAvailableImageModels", "可用图片模型", "列出已启用图片模型 ID、名称及默认状态，不包含密钥。生成时可用 modelId 选择其中一项。",
|
||
z.object({}), async () => ({ models: (await getStyleProImageModelEntries())
|
||
.filter(({ model }) => model.healthStatus !== "auto_disabled")
|
||
.map(({ model }) => ({ id: model.id, name: model.displayName, modelId: model.modelId, isDefault: model.isDefault })) }));
|
||
register(
|
||
"planCanvasTask",
|
||
"处理计划",
|
||
"先明确用户目标、步骤及最终出图数量。数量只计算最终交付图，不包含抠图等中间产物；用户未指定时可以不填。",
|
||
z.object({
|
||
steps: z.array(z.string().min(1).max(180)).min(1).max(12),
|
||
expectedImageCount: z.number().int().min(1).max(24).optional(),
|
||
aspectRatio: z.string().max(24).optional(),
|
||
visualCriteria: z.array(z.string().max(200)).max(10),
|
||
}),
|
||
async (input) => ({ ...input,
|
||
...(input.expectedImageCount !== undefined && params.imagePreferences?.parameters?.imageCount !== undefined
|
||
? { expectedImageCount: params.imagePreferences.parameters.imageCount } : {}),
|
||
...(params.imagePreferences?.parameters?.aspectRatio && params.imagePreferences.parameters.aspectRatio !== "auto"
|
||
? { aspectRatio: params.imagePreferences.parameters.aspectRatio } : {}),
|
||
}),
|
||
);
|
||
register(
|
||
"readCanvas",
|
||
"读取画布",
|
||
"读取真实画布节点及位置。先读取后操作；nodeIds 可选择具体对象，省略返回当前选中图及画布图片摘要。",
|
||
z.object({ nodeIds: z.array(z.string().min(1).max(180)).max(40).optional() }),
|
||
async (input) => {
|
||
const current = await readOwnedCanvas(params.userId, params.projectId);
|
||
const chosen = input.nodeIds?.length
|
||
? input.nodeIds
|
||
: defaults.length
|
||
? defaults
|
||
: current.nodes.map((node) => node.id);
|
||
return {
|
||
selectedNodeIds: defaults,
|
||
nodes: current.nodes
|
||
.filter((node) => chosen.includes(node.id))
|
||
.slice(0, 40)
|
||
.map((node) => ({
|
||
id: node.id,
|
||
type: node.type,
|
||
title: node.data.title,
|
||
prompt: node.data.prompt?.slice(0, 500),
|
||
position: node.position,
|
||
size: node.size,
|
||
status: node.data.status,
|
||
})),
|
||
};
|
||
},
|
||
);
|
||
register(
|
||
"analyzeCanvasImages",
|
||
"检查图片",
|
||
"查看真实图片，检查主体、背景及用户要求。不会自动修改图片，返回 passed 与具体问题。",
|
||
z.object({
|
||
nodeIds: z.array(z.string().min(1).max(180)).min(1).max(4),
|
||
question: z.string().min(1).max(2000),
|
||
}),
|
||
async (input) => analyzeRuntimeImages({ ...params, ...input }),
|
||
);
|
||
|
||
let generatedCount = 0;
|
||
register(
|
||
"generateCanvasImages",
|
||
"生成图片",
|
||
"使用本站配置的生图模型生成或修改图片，包括转正面、换角度、旋转主体视角。通过 referenceNodeIds 提供参考图；结果作为新节点放在参考图旁边，保留原图。",
|
||
z.object({
|
||
prompt: z.string().min(1).max(8000),
|
||
referenceNodeIds: z.array(z.string().min(1).max(180)).max(8).default([]),
|
||
imageCount: z.number().int().min(1).max(10).default(1),
|
||
aspectRatio: z.string().max(24).default("auto"),
|
||
imageSize: z.enum(["0.5K", "1K", "2K", "4K", "8K", "10K", "12K"]).optional(),
|
||
modelId: z.string().max(180).optional().describe("指定已配置的图片模型 ID；省略使用预设或系统模型"),
|
||
}),
|
||
async (toolInput, callId) => {
|
||
const input = applyAgentImagePreferences(toolInput, params.imagePreferences);
|
||
// Bound per-task spend without introducing any global/model concurrency queue.
|
||
if (generatedCount + input.imageCount > 24)
|
||
throw new Error("本次任务已达到 24 张生成预算，请检查结果后继续。");
|
||
const sources = await readCanvasImages(
|
||
params.userId,
|
||
params.projectId,
|
||
input.referenceNodeIds,
|
||
);
|
||
const models = sortStyleProImageModelEntries(await getStyleProImageModelEntries());
|
||
const route = await routeToModel({
|
||
intent: sources.length ? "image_edit" : "image_gen",
|
||
hasImage: sources.length > 0,
|
||
}).catch(() => null);
|
||
const pinnedModel = resolveUserImageModel(params.objective || "", models, params.imagePreferences?.modelId);
|
||
const requestedModel = pinnedModel || input.modelId || params.preferredModelId;
|
||
const requestedEntry = requestedModel ? models.find(({ model }) => model.healthStatus !== "auto_disabled" && (model.id === requestedModel || model.modelId === requestedModel)) : null;
|
||
if (pinnedModel && !requestedEntry) throw new Error("指定的图片模型不可用，未切换到其他模型收费生成");
|
||
const chosen = requestedEntry || models.find(
|
||
({ model }) => model.id === route?.id && model.healthStatus !== "auto_disabled",
|
||
) ||
|
||
models.find(({ model }) => model.isDefault && model.healthStatus !== "auto_disabled") ||
|
||
models.find(({ model }) => model.healthStatus !== "auto_disabled");
|
||
if (!chosen) throw new Error("系统未配置可用的图片生成模型");
|
||
const imageSize = input.imageSize || params.imageSize || "2K";
|
||
const quota = await checkUserImageQuota(params.userId);
|
||
if (!quota.allowed || (quota.remaining !== null && input.imageCount > quota.remaining))
|
||
throw new Error("本月图片生成配额不足。");
|
||
const points = await getPointsSystemConfig();
|
||
const user = await getUserById({ id: params.userId });
|
||
const cost =
|
||
getModelGenerationCreditCost(
|
||
{ imageGenConfig: chosen.model.imageGenConfig as DiscountableImageGenConfigLike | null },
|
||
"image",
|
||
{ imageSize },
|
||
) * input.imageCount;
|
||
if (
|
||
pinnedModel && points.enabled &&
|
||
points.blockOnInsufficient !== false &&
|
||
!user?.isHighVip &&
|
||
(user?.pointsBalance ?? 0) < cost
|
||
)
|
||
throw new Error("积分不足，请充值或联系管理员。");
|
||
generatedCount += input.imageCount;
|
||
const result = await generateCanvasImage({
|
||
prompt: input.prompt,
|
||
model: chosen.model.id,
|
||
imageCount: input.imageCount,
|
||
imageSize,
|
||
aspectRatio: input.aspectRatio,
|
||
negativePrompt: input.negativePrompt,
|
||
promptExtend: input.promptExtend,
|
||
providerImageParams: agentProviderImageParameters(input),
|
||
allowFailover: !pinnedModel,
|
||
agentImageMode: pinnedModel ? "pinned" : "auto",
|
||
agentImageMaxCreditCost: points.enabled && points.blockOnInsufficient !== false && !user?.isHighVip
|
||
? Math.max(0, user?.pointsBalance ?? 0) / input.imageCount : undefined,
|
||
ownerKey: params.userId,
|
||
referenceImages: sources.map((node) => node.data.src || node.data.url!),
|
||
debugLabel: "canvas-agent-runtime",
|
||
debugMeta: { runId: params.runId, callId },
|
||
});
|
||
if (!result.images.length) throw new Error("模型未返回有效图片。");
|
||
const ratio = input.aspectRatio.match(/^(\d+):(\d+)$/);
|
||
const operation = await applyCanvasOperationPlan({
|
||
userId: params.userId,
|
||
projectId: params.projectId,
|
||
plan: {
|
||
version: 1,
|
||
projectId: params.projectId,
|
||
operationId: `${params.runId}:${callId}`,
|
||
operation: "insert_beside",
|
||
sourceNodeIds: input.referenceNodeIds,
|
||
targetNodeIds: input.referenceNodeIds,
|
||
preserveOriginal: true,
|
||
placement: sources.length ? "right" : "viewport-center",
|
||
layoutStrategy: "row",
|
||
mediaItems: result.images.map((image, index) => ({
|
||
mediaItemId: `${callId}:${index}`,
|
||
type: "image",
|
||
url: image.originalUrl,
|
||
width: 1024,
|
||
height: ratio ? Math.round((1024 * Number(ratio[2])) / Number(ratio[1])) : 1024,
|
||
prompt: input.prompt,
|
||
sourceNodeIds: input.referenceNodeIds,
|
||
metadata: { canvasAgentRunId: params.runId, model: image.modelRef || result.modelRef },
|
||
})),
|
||
},
|
||
});
|
||
params.emit("data-canvasOperation", {
|
||
projectId: params.projectId,
|
||
operationId: `${params.runId}:${callId}`,
|
||
operation: "insert_beside",
|
||
status: "applied",
|
||
createdNodeIds: operation.createdNodeIds,
|
||
createdNodeIdsByMediaItemId: operation.createdNodeIdsByMediaItemId,
|
||
updatedNodeIds: operation.updatedNodeIds,
|
||
canvasRevision: operation.canvasRevision,
|
||
});
|
||
const billingGroups = new Map<string, typeof result.images>();
|
||
for (const image of result.images) {
|
||
const ref = image.modelRef || result.modelRef;
|
||
billingGroups.set(ref, [...(billingGroups.get(ref) || []), image]);
|
||
}
|
||
for (const [modelRef, group] of billingGroups) {
|
||
await saveGeneratedImages(
|
||
group
|
||
.map((image) => ``)
|
||
.join("\n"),
|
||
params.userId,
|
||
params.chatId,
|
||
params.assistantMessageId,
|
||
input.prompt,
|
||
modelRef,
|
||
{
|
||
projectId: params.projectId,
|
||
sourceImageUrls: sources.map((node) => node.data.src || node.data.url!),
|
||
relationType: sources.length ? "edited_from" : "generated_from",
|
||
imageSize,
|
||
aspectRatio: input.aspectRatio,
|
||
},
|
||
);
|
||
}
|
||
return {
|
||
status: result.images.length < input.imageCount ? "partial" : "succeeded",
|
||
appliedImagePreferences: { modelId: result.modelRef, imageSize, aspectRatio: input.aspectRatio, imageCount: input.imageCount },
|
||
createdNodeIds: operation.createdNodeIds,
|
||
outputs: result.images.map((image) => ({ url: image.originalUrl })),
|
||
model: result.modelRef,
|
||
};
|
||
},
|
||
false,
|
||
);
|
||
registerImageCanvasAccessTools({
|
||
...params,
|
||
selectedNodeIds: params.operationSelectedNodeIds || params.selectedNodeIds,
|
||
definitions,
|
||
executable,
|
||
});
|
||
return { definitions, executable } satisfies {
|
||
definitions: ToolSet;
|
||
executable: CanvasAgentLoopPorts["tools"];
|
||
};
|
||
}
|