ai-chatbot/lib/ai/tools-registry.ts

import { getAllToolConfigs } from "@/lib/db/queries";

export type RuntimeToolName =
  | "getWeather"
  | "createDocument"
  | "updateDocument"
  | "requestSuggestions"
  | "generateVideo";

export type AgentProjectToolName =
  | RuntimeToolName
  | "respondDirect"
  | "resolveTargetImage"
  | "analyzeTargetImage"
  | "generateNewImage"
  | "editTargetImage"
  | "publishResultToCanvas"
  | "analyzeImage"
  | "generateImage"
  | "editImage"
  | "removeBackground"
  | "enhanceImage"
  | "extendImage"
  | "resizeImage"
  | "boxCutout"
  | "publishToCanvas"
  | "useSelectedCanvasImage"
  | "useLatestResultImage"
  | "useOriginalUploadImage";

export type ToolRegistry = Record<RuntimeToolName, boolean>;

export interface AgentProjectToolDefinition {
  name: AgentProjectToolName;
  title: string;
  description: string;
  category: "conversation" | "image" | "video" | "document" | "utility" | "canvas";
  runtimeTool?: RuntimeToolName;
  plannerHidden?: boolean;
  requiresImage?: boolean;
  requiresCanvasSelection?: boolean;
  requiresHistoryImage?: boolean;
  requiresCurrentUpload?: boolean;
  producesMedia?: boolean;
}

export interface AgentToolContext {
  hasImage?: boolean;
  hasCanvasSelection?: boolean;
  hasHistoryImage?: boolean;
  hasCurrentUpload?: boolean;
}

const RUNTIME_TOOL_DEFINITIONS: Array<AgentProjectToolDefinition & { runtimeTool: RuntimeToolName }> = [
  {
    name: "getWeather",
    title: "天气查询",
    description: "查询天气并在对话中直接回答用户。",
    category: "utility",
    runtimeTool: "getWeather",
  },
  {
    name: "createDocument",
    title: "创建文档",
    description: "按用户要求创建文档、草稿或结构化内容。",
    category: "document",
    runtimeTool: "createDocument",
  },
  {
    name: "updateDocument",
    title: "更新文档",
    description: "继续修改已有文档、方案或草稿内容。",
    category: "document",
    runtimeTool: "updateDocument",
  },
  {
    name: "requestSuggestions",
    title: "内容建议",
    description: "生成补充建议、优化意见或下一步选项。",
    category: "document",
    runtimeTool: "requestSuggestions",
  },
  {
    name: "generateVideo",
    title: "视频生成",
    description: "根据文本或参考图生成视频。",
    category: "video",
    runtimeTool: "generateVideo",
    producesMedia: true,
  },
];

const INTERNAL_AGENT_TOOLS: AgentProjectToolDefinition[] = [
  {
    name: "respondDirect",
    title: "直接回复",
    description: "仅对话、解释、建议或策略讨论，不执行生图生视频。",
    category: "conversation",
  },
  {
    name: "resolveTargetImage",
    title: "解析目标图",
    description:
      "先确定本轮要操作的是哪张图，可通过 args.source 指定 selected_canvas/latest_result/original_upload/auto。",
    category: "canvas",
    requiresImage: true,
  },
  {
    name: "analyzeTargetImage",
    title: "分析目标图",
    description: "对已解析的目标图做识别、风格分析、质量评估与改造建议。",
    category: "image",
    requiresImage: true,
  },
  {
    name: "generateNewImage",
    title: "生成新图",
    description: "根据文本或参考风格生成全新图片，需要新建结果图时使用。",
    category: "image",
    producesMedia: true,
  },
  {
    name: "editTargetImage",
    title: "编辑目标图",
    description: "基于已解析的目标图进行修改、换色、换背景、增强、扩图或局部重绘。",
    category: "image",
    requiresImage: true,
    producesMedia: true,
  },
  {
    name: "publishResultToCanvas",
    title: "发布结果到画布",
    description: "把当前结果图发布回无限画布，作为新的结果节点或继续编辑的输入。",
    category: "canvas",
  },
  {
    name: "analyzeImage",
    title: "图片分析",
    description: "对当前图或历史图做识别、风格分析、质量评估、品牌/营销建议。",
    category: "image",
    requiresImage: true,
    plannerHidden: true,
  },
  {
    name: "generateImage",
    title: "图片生成",
    description: "根据文本或参考风格生成新图片。",
    category: "image",
    producesMedia: true,
    plannerHidden: true,
  },
  {
    name: "editImage",
    title: "图片编辑",
    description: "基于当前图进行抠图、换背景、增强、调色、尺寸调整、局部重绘等。",
    category: "image",
    requiresImage: true,
    producesMedia: true,
    plannerHidden: true,
  },
  {
    name: "removeBackground",
    title: "抠图去背",
    description: "提取主体、去除背景，适合商品图、人像、元素裁切。",
    category: "image",
    requiresImage: true,
    producesMedia: true,
  },
  {
    name: "enhanceImage",
    title: "清晰增强",
    description: "提升清晰度、细节和质感，适合低清图修复与放大前预处理。",
    category: "image",
    requiresImage: true,
    producesMedia: true,
  },
  {
    name: "extendImage",
    title: "扩图补边",
    description: "向四周延展画面，适合改比例、补全背景或生成更大构图。",
    category: "image",
    requiresImage: true,
    producesMedia: true,
  },
  {
    name: "resizeImage",
    title: "尺寸调整",
    description: "调整画布尺寸、背景色与构图位置，适合电商主图、社媒比例适配。",
    category: "image",
    requiresImage: true,
    producesMedia: true,
  },
  {
    name: "boxCutout",
    title: "框选抠图",
    description: "对目标图中的局部区域做框选裁切，再提取被框选主体。",
    category: "image",
    requiresImage: true,
    producesMedia: true,
  },
  {
    name: "publishToCanvas",
    title: "发布到画布",
    description: "把当前结果图继续发送到无限画布或作为新的画布结果节点。",
    category: "canvas",
    producesMedia: false,
    plannerHidden: true,
  },
  {
    name: "useSelectedCanvasImage",
    title: "使用选中图",
    description: "将当前画布里选中的图片作为本轮分析或编辑目标。",
    category: "canvas",
    requiresImage: true,
    requiresCanvasSelection: true,
    plannerHidden: true,
  },
  {
    name: "useLatestResultImage",
    title: "使用最新结果图",
    description: "将上一轮生成结果作为连续对话和继续编辑的目标图。",
    category: "canvas",
    requiresImage: true,
    requiresHistoryImage: true,
    plannerHidden: true,
  },
  {
    name: "useOriginalUploadImage",
    title: "使用原上传图",
    description: "将最近上传的原图作为分析、修图或对比基准。",
    category: "canvas",
    requiresImage: true,
    requiresCurrentUpload: true,
    plannerHidden: true,
  },
];

const DEFAULT_RUNTIME_TOOLS = RUNTIME_TOOL_DEFINITIONS.map((tool) => tool.runtimeTool);

function getAllAgentToolDefinitions(): AgentProjectToolDefinition[] {
  return [...INTERNAL_AGENT_TOOLS, ...RUNTIME_TOOL_DEFINITIONS];
}

function matchesAgentToolContext(
  tool: AgentProjectToolDefinition,
  context?: AgentToolContext,
) {
  if (!context) return true;
  if (tool.requiresImage && !context.hasImage) return false;
  if (tool.requiresCanvasSelection && !context.hasCanvasSelection) return false;
  if (tool.requiresHistoryImage && !context.hasHistoryImage) return false;
  if (tool.requiresCurrentUpload && !context.hasCurrentUpload) return false;
  return true;
}

/**
 * 获取启用的运行时工具列表
 * 从数据库读取工具配置，如果工具未配置则默认启用
 */
export async function getEnabledTools(): Promise<RuntimeToolName[]> {
  try {
    const configs = await getAllToolConfigs();
    const enabledTools: RuntimeToolName[] = [];

    for (const toolName of DEFAULT_RUNTIME_TOOLS) {
      const config = configs.find((c) => c.toolName === toolName);
      if (!config || config.isEnabled) {
        enabledTools.push(toolName);
      }
    }

    return enabledTools;
  } catch (error) {
    console.error("Failed to get enabled tools, using defaults:", error);
    return [...DEFAULT_RUNTIME_TOOLS];
  }
}

export async function getEnabledAgentToolDefinitions(): Promise<AgentProjectToolDefinition[]> {
  return getEnabledAgentToolDefinitionsForContext();
}

export async function getEnabledAgentToolDefinitionsForContext(
  context?: AgentToolContext,
): Promise<AgentProjectToolDefinition[]> {
  const runtimeTools = await getEnabledTools();
  return getAllAgentToolDefinitions().filter((tool) => {
    if (!tool.runtimeTool) return true;
    return runtimeTools.includes(tool.runtimeTool);
  }).filter((tool) => matchesAgentToolContext(tool, context));
}