lovart3/lib/ai/prompt-optimizer.ts

/**
 * 提示词优化器
 * 将用户简单输入转换为专业的图片编辑提示词
 * 
 * 主要应用场景：图片编辑（抠图、换背景、增强等）
 */

// 优化结果类型
export interface OptimizeResult {
  originalPrompt: string;
  optimizedPrompt: string;
  negativePrompt: string;
  friendlyMessage?: string; // 友好的对话消息
}

// 子意图的友好名称和回复模板
const FRIENDLY_MESSAGES: Record<string, { name: string; template: string }> = {
  remove_background: {
    name: "抠图",
    template: "好的，我来帮你抠图，去除背景保留主体~",
  },
  replace_background: {
    name: "换背景",
    template: "收到，我来帮你把背景换成{background}~",
  },
  enhance: {
    name: "图片增强",
    template: "明白，我来帮你提升图片清晰度~",
  },
  extend: {
    name: "图像拓展",
    template: "好的，我来帮你扩展图片边界~",
  },
  inpaint: {
    name: "局部修改",
    template: "收到，我来帮你{action}{target}~",
  },
  style_transfer: {
    name: "风格迁移",
    template: "好的，我来帮你转换成{style}风格~",
  },
  resize: {
    name: "尺寸调整",
    template: "明白，我来帮你调整图片尺寸~",
  },
  color_adjust: {
    name: "调色",
    template: "收到，我来帮你调整图片颜色~",
  },
};

// 图片编辑子意图的提示词模板
const EDIT_TEMPLATES: Record<
  string,
  {
    template: string;
    negative: string;
    extractors: Record<string, RegExp>;
    defaults: Record<string, string>;
  }
> = {
  remove_background: {
    template:
      "精准抠图：去除图片背景，保留主体完整清晰，输出透明背景PNG格式。边缘处理要求：平滑自然无锯齿，保留精细细节（如发丝、毛发、边缘纹理）。主体保持原有清晰度和色彩。",
    negative: "残留背景色块，边缘毛糙，锯齿明显，主体不完整，边缘模糊，细节丢失，颜色溢出",
    extractors: {},
    defaults: {},
  },

  replace_background: {
    template:
      "背景替换：将图片背景替换为{background}。要求：保持主体完整清晰不变形，边缘过渡自然平滑，光影效果与新背景协调一致。{extra}",
    negative: "残留原背景，边缘不自然，透视错误，主体变形，光影不协调，明显P图痕迹，边缘光晕",
    extractors: {
      background: /(?:换成|改成|替换为|变成|换为|改为)(.+?)(?:背景|底|$)/,
    },
    defaults: {
      background: "纯白色(#FFFFFF)",
      extra: "保留主体自然阴影效果",
    },
  },

  enhance: {
    template:
      "图片增强：提升图片整体清晰度和画质。要求：保持自然色彩不失真，智能去除噪点，增强细节层次，适度锐化边缘。{extra}",
    negative: "过度锐化，色彩失真，噪点增加，伪影，画质下降，不自然的对比度，颜色过饱和",
    extractors: {},
    defaults: {
      extra: "整体效果自然真实",
    },
  },

  extend: {
    template:
      "图像拓展(Outpainting)：向{direction}方向扩展图片约{amount}。要求：扩展区域的内容、风格、色调与原图保持高度一致，光影过渡自然流畅，无明显接缝痕迹，透视关系正确。",
    negative: "风格不一致，明显接缝，透视错误，内容突变，色调不匹配，光影断裂",
    extractors: {
      direction: /向?(上|下|左|右|四周|周围)/,
      amount: /(\d+%?|一倍|两倍|50%|100%)/,
    },
    defaults: {
      direction: "四周",
      amount: "50%",
    },
  },

  inpaint: {
    template:
      "局部修改(Inpainting)：{action}图片中的{target}。要求：使用周围环境自然填充修改区域，保持整体风格、色调、光影一致，无任何修复痕迹，过渡自然。",
    negative: "明显修复痕迹，风格不一致，内容错误，边缘不自然，色差明显，纹理断裂",
    extractors: {
      target: /(?:去除|消除|擦掉|删除|移除|P掉|p掉)(.+?)(?:，|。|$| )/,
      action: /(去除|消除|擦掉|删除|移除|修改)/,
    },
    defaults: {
      target: "指定区域内的元素",
      action: "去除",
    },
  },

  style_transfer: {
    template:
      "风格迁移：将图片转换为{style}风格。要求：保持主体可识别性和基本特征，色彩风格协调统一，整体效果自然美观。",
    negative: "主体变形严重，特征丢失，风格不统一，色彩异常，人物不可识别",
    extractors: {
      style: /(?:转换为|变成|改成|转成)(.+?)(?:风格|$)/,
    },
    defaults: {
      style: "卡通动漫",
    },
  },

  resize: {
    template:
      "尺寸调整：将图片{action}至{size}。要求：保持画质清晰，比例协调，细节不丢失。",
    negative: "画质损失严重，比例失调，模糊，变形",
    extractors: {
      size: /(\d+[x×]\d+|\d+%)/,
      action: /(放大|缩小|调整|裁剪)/,
    },
    defaults: {
      size: "目标尺寸",
      action: "调整",
    },
  },

  color_adjust: {
    template:
      "颜色调整：对图片进行{adjustment}处理。要求：保持自然效果，色彩协调，不过度调整。",
    negative: "色彩失真，过度调整，不自然，偏色严重",
    extractors: {
      adjustment: /(?:调整|增加|减少|提高|降低)(.+?)(?:，|。|$| )/,
    },
    defaults: {
      adjustment: "色彩优化",
    },
  },
};

/**
 * 优化提示词
 */
export function optimizePrompt(params: {
  originalPrompt: string;
  intent: string;
  subIntent?: string;
}): OptimizeResult {
  const { originalPrompt, intent, subIntent } = params;

  console.log("[Prompt Optimizer] Input:", { originalPrompt, intent, subIntent });

  // 只对图片编辑类意图进行优化
  if (intent !== "image_edit") {
    return {
      originalPrompt,
      optimizedPrompt: originalPrompt,
      negativePrompt: "模糊，低质量，变形，水印，失真",
    };
  }

  // 如果没有子意图，返回通用优化
  if (!subIntent) {
    return {
      originalPrompt,
      optimizedPrompt: `图片编辑：${originalPrompt}。要求：效果自然，画质清晰，细节完整。`,
      negativePrompt: "模糊，低质量，变形，水印，失真，不自然",
    };
  }

  const template = EDIT_TEMPLATES[subIntent];
  if (!template) {
    console.log("[Prompt Optimizer] No template for subIntent:", subIntent);
    return {
      originalPrompt,
      optimizedPrompt: originalPrompt,
      negativePrompt: "模糊，低质量，变形，水印，失真",
    };
  }

  // 使用模板生成优化后的提示词
  let optimizedPrompt = template.template;

  // 从用户输入中提取参数并填充模板
  for (const [key, regex] of Object.entries(template.extractors)) {
    const match = originalPrompt.match(regex);
    let value = match ? match[1].trim() : template.defaults[key] || "";
    
    // 清理提取的值
    value = value.replace(/[，。！？,.!?]/g, "").trim();
    
    if (!value) {
      value = template.defaults[key] || "";
    }
    
    optimizedPrompt = optimizedPrompt.replace(`{${key}}`, value);
  }

  // 填充剩余的默认值
  for (const [key, defaultValue] of Object.entries(template.defaults)) {
    optimizedPrompt = optimizedPrompt.replace(`{${key}}`, defaultValue);
  }

  // 附加用户原始意图（如果优化后的提示词没有完全覆盖用户需求）
  if (originalPrompt.length > 10 && !optimizedPrompt.includes(originalPrompt)) {
    optimizedPrompt = `${optimizedPrompt}\n\n用户补充要求：${originalPrompt}`;
  }

  // 生成友好的对话消息
  let friendlyMessage = "";
  const friendlyTemplate = FRIENDLY_MESSAGES[subIntent];
  if (friendlyTemplate) {
    friendlyMessage = friendlyTemplate.template;
    
    // 填充友好消息中的变量
    for (const [key, regex] of Object.entries(template.extractors)) {
      const match = originalPrompt.match(regex);
      let value = match ? match[1].trim() : template.defaults[key] || "";
      value = value.replace(/[，。！？,.!?]/g, "").trim();
      if (!value) {
        value = template.defaults[key] || "";
      }
      friendlyMessage = friendlyMessage.replace(`{${key}}`, value);
    }
    
    // 填充剩余的默认值
    for (const [key, defaultValue] of Object.entries(template.defaults)) {
      friendlyMessage = friendlyMessage.replace(`{${key}}`, defaultValue);
    }
  }

  const result: OptimizeResult = {
    originalPrompt,
    optimizedPrompt,
    negativePrompt: template.negative,
    friendlyMessage: friendlyMessage || `好的，我来帮你处理这张图片~`,
  };

  console.log("[Prompt Optimizer] Output:", result);
  return result;
}

/**
 * 提取背景描述（辅助函数）
 */
export function extractBackgroundDescription(text: string): string {
  // 尝试提取用户描述的背景
  const patterns = [
    /(?:换成|改成|替换为|变成|换为|改为)(.+?)(?:背景|底)/,
    /(.+?)(?:背景|底)$/,
    /背景(?:换成|改成|替换为|变成)(.+)/,
  ];

  for (const pattern of patterns) {
    const match = text.match(pattern);
    if (match && match[1]) {
      return match[1].trim();
    }
  }

  return "纯白色";
}

/**
 * 检测是否需要保留阴影
 */
export function shouldPreserveShadow(text: string): boolean {
  // 产品图、商品图通常需要保留阴影
  if (/产品|商品|电商|白底|纯白/.test(text)) {
    return true;
  }
  return false;
}