wx-agent/utils/image_utils.py

"""图片处理工具模块。"""
import re
from typing import List, Tuple
from difflib import SequenceMatcher


IMG_EXTENSIONS = r'(?:jpg|jpeg|png|gif|bmp|webp)'


def _dedupe_keep_order(values: List[str]) -> List[str]:
    seen = set()
    result = []
    for value in values:
        cleaned = (value or "").strip().strip("`'\"，,。；;")
        if not cleaned:
            continue
        key = cleaned.lower()
        if key in seen:
            continue
        seen.add(key)
        result.append(cleaned)
    return result


def extract_image_paths(text: str) -> List[str]:
    """从文本中提取图片路径或 URL，并保持原始出现顺序。"""
    if not text:
        return []

    patterns = [
        # Markdown 图片：![alt](url)
        rf'!\[[^\]\n]*\]\(\s*([^)\s]+?\.{IMG_EXTENSIONS}(?:[?#][^)\s]*)?)\s*\)',
        # 完整 URL。
        rf'(https?://[^\s<>"\'\]\)]+?\.{IMG_EXTENSIONS}(?:[?#][^\s<>"\'\]\)]*)?)',
        # API 绝对路径。
        rf'(?<![\w:/.%-])(/api/v1/knowledge/files/images/[^\s<>"\'\]\)]+?\.{IMG_EXTENSIONS}(?:[?#][^\s<>"\'\]\)]*)?)',
        # 相对 images 路径。
        rf'(?<![\w/.-])(images/[A-Za-z0-9_\-./%]+?\.{IMG_EXTENSIONS}(?:[?#][^\s<>"\'\]\)]*)?)',
        # 裸文件名，兼容检索结果只给文件名的情况。
        rf'(?<![\w/.-])([A-Za-z0-9_-]{{8,}}\.{IMG_EXTENSIONS})(?![\w/.-])',
    ]

    matches: List[Tuple[int, str]] = []
    for pattern in patterns:
        matches.extend((m.start(1), m.group(1)) for m in re.finditer(pattern, text, re.IGNORECASE))

    matches.sort(key=lambda item: item[0])
    return _dedupe_keep_order([value for _, value in matches])


def _strip_url_suffix(filename: str) -> str:
    return re.sub(r'[?#].*$', '', filename or '').strip()


def _split_filename(filename: str) -> Tuple[str, str]:
    filename = _strip_url_suffix(filename).lower()
    if "." not in filename:
        return filename, ""
    stem, ext = filename.rsplit(".", 1)
    return stem, ext


def _path_to_image_url(path: str, base_prefix: str) -> str:
    cleaned = (path or "").strip()
    if cleaned.startswith(("http://", "https://", "/")):
        return cleaned
    if cleaned.lower().startswith("images/"):
        return base_prefix.rstrip("/") + "/" + cleaned
    return base_prefix.rstrip("/") + "/images/" + cleaned.split("/")[-1]


def _is_acceptable_image_match(generated_filename: str, original_filename: str, score: float) -> bool:
    """只接受完全一致、明确截断或高度相似的文件名，避免把幻觉图片错配到真实图片。"""
    generated_stem, generated_ext = _split_filename(generated_filename)
    original_stem, original_ext = _split_filename(original_filename)

    if not generated_stem or not original_stem or generated_ext != original_ext:
        return False
    if generated_stem == original_stem:
        return True
    if len(generated_stem) >= 12 and original_stem.startswith(generated_stem):
        return True
    if len(original_stem) >= 12 and generated_stem.startswith(original_stem):
        return True
    return score >= 0.85


def _remove_untrusted_image_references(text: str) -> str:
    """没有可信原图列表时，删除模型输出的图片标签和裸图片路径。"""
    if not text:
        return text

    result = re.sub(
        rf'!\[[^\]\n]*\]\([^\)\n]*\.{IMG_EXTENSIONS}(?:[?#][^\)\n]*)?\)',
        '',
        text,
        flags=re.IGNORECASE,
    )
    result = re.sub(
        rf'(?<![\w/.-])(?:https?://[^\s<>)]+/images/[^\s<>)]+|/api/v1/knowledge/files/images/[^\s<>)]+|images/[^\s<>)]+|[A-Za-z0-9_-]+\.{IMG_EXTENSIONS})(?:[?#][^\s<>)]+)?',
        '',
        result,
        flags=re.IGNORECASE,
    )
    result = re.sub(r'[ \t]+\n', '\n', result)
    result = re.sub(r'\n{3,}', '\n\n', result)
    return result.strip()


def find_closest_image_path(generated_path: str, original_paths: List[str]) -> str:
    """
    找到与生成路径最接近的原始图片路径
    
    即使模型输出的是错误/近似/截断的文件名，也能匹配到正确的原始路径
    """
    if not original_paths:
        return generated_path

    generated_filename = generated_path.split('/')[-1].lower()
    best_match = None
    best_score = 0

    for orig_path in original_paths:
        orig_filename = orig_path.split('/')[-1].lower()
        score = SequenceMatcher(None, generated_filename, orig_filename).ratio()
        if score > best_score:
            best_score = score
            best_match = orig_path

    if best_match and _is_acceptable_image_match(generated_filename, best_match.split('/')[-1], best_score):
        return best_match
    return generated_path


def normalize_markdown_images(text: str, base_prefix: str = "/api/v1/knowledge/files/", original_paths: List[str] = None) -> str:
    """
    规范化 Markdown 图片链接 - 两轮处理版
    
    算法：
    第一轮：检查和修改格式，确保是 ![图片](/api/v1/knowledge/files/images/xxx.jpg) 格式
    第二轮：只检查和替换最后文件名部分（xxx.jpg），与原始路径匹配，匹配度低的删除
    """
    if not text or original_paths is None:
        return text
    trusted_paths = [path for path in original_paths if path]
    if not trusted_paths:
        return _remove_untrusted_image_references(text)
    
    # ========== 第一轮：检查和修改格式（包括固定URL前缀） ==========
    result = _normalize_image_format(text, base_prefix)
    
    # ========== 第二轮：只检查和替换最后文件名部分 ==========
    result = _normalize_image_paths(result, base_prefix, trusted_paths)
    
    return result


def _normalize_image_format(text: str, base_prefix: str) -> str:
    """
    第一轮：规范化图片格式，确保所有图片都是正确的 ![图片](/api/v1/knowledge/files/images/xxx.jpg) 格式
    
    处理各种格式破损的情况，并确保有正确的URL前缀
    """
    if not text:
        return text
    
    img_pattern = rf'([a-zA-Z0-9_\-./]+)\.({IMG_EXTENSIONS})'
    
    result = text

    def build_standard_tag(path: str) -> str:
        full_url = _path_to_image_url(path, base_prefix)
        return f'![图片]({full_url})'

    # 先处理已经是 Markdown 图片语法的内容，包含任意 alt 文本。
    # 例如：![图2-3 消磁电源柜功能框图](images/xxx.jpg)
    markdown_img_pattern = rf'!\[[^\]\n]*\]\(([^)\n]*?\.{IMG_EXTENSIONS})\)'
    result = re.sub(
        markdown_img_pattern,
        lambda m: build_standard_tag(m.group(1)),
        result,
        flags=re.IGNORECASE
    )

    # 再兜底清理模型或前序处理产生的嵌套图片语法。
    # 例如：![图2-3](![图片](/api/v1/knowledge/files/images/xxx.jpg)
    nested_any_alt_pattern = rf'!\[[^\]\n]*\]\([^)\n]*!\[[^\]\n]*\]\(([^)\n]+\.{IMG_EXTENSIONS})\)\)?'
    while re.search(nested_any_alt_pattern, result, re.IGNORECASE):
        result = re.sub(
            nested_any_alt_pattern,
            lambda m: build_standard_tag(m.group(1)),
            result,
            flags=re.IGNORECASE
        )

    matches = list(re.finditer(img_pattern, result, re.IGNORECASE))
    
    if not matches:
        return result
    
    # 从后往前处理，避免位置偏移
    for match in reversed(matches):
        filename_with_ext = match.group(0)
        file_start = match.start()
        file_end = match.end()
        
        mark_start = file_start
        mark_end = file_end
        
        # 向前查找是否已经有 ![图片]( 标记
        lookback_limit = max(0, file_start - 100)
        before = result[lookback_limit:file_start]
        open_marker_pos = before.rfind('![图片](')
        
        if open_marker_pos != -1:
            actual_open_pos = lookback_limit + open_marker_pos
            check_segment = result[actual_open_pos:file_start]
            if check_segment.startswith('![图片]('):
                mark_start = actual_open_pos
        
        # 向后查找是否有 )
        after = result[file_end:]
        close_paren_pos = after.find(')')
        if close_paren_pos != -1 and close_paren_pos < 20:
            middle = after[:close_paren_pos]
            if '\n' not in middle and '![图片](' not in middle:
                mark_end = file_end + close_paren_pos + 1
        
        # 构建完整的标准格式（包括固定URL前缀）
        standard_format = build_standard_tag(filename_with_ext)
        
        result = result[:mark_start] + standard_format + result[mark_end:]
    
    # 清理嵌套格式
    nested_pattern = r'!\[图片\]\([^)]*!\[图片\]\(([^)]+\.(?:jpg|jpeg|png|gif|bmp|webp))\)[^)]*\)'
    while re.search(nested_pattern, result, re.IGNORECASE):
        result = re.sub(nested_pattern, r'![图片](\1)', result, flags=re.IGNORECASE)
    
    return result


def _normalize_image_paths(text: str, base_prefix: str, original_paths: List[str]) -> str:
    """
    第二轮：只规范化最后文件名部分（xxx.jpg），与原始路径匹配，匹配度低的删除
    
    步骤：
    1. 找到所有 ![图片](URL) 格式的图片
    2. 对每个图片：
       a. 提取最后文件名部分（xxx.jpg）
       b. 只与 original_paths 中的最后文件名部分匹配
       c. 如果匹配度高（>0.5），替换为正确路径
       d. 如果匹配度低（<=0.5），删除整个图片
    """
    if not text or not original_paths:
        return text
    
    # 匹配所有 ![图片](...) 格式
    img_tag_pattern = r'!\[图片\]\(([^)]+)\)'
    matches = list(re.finditer(img_tag_pattern, text, re.IGNORECASE))
    
    if not matches:
        return text
    
    result = text
    
    # 从后往前处理，避免位置偏移
    for match in reversed(matches):
        full_tag = match.group(0)
        url = match.group(1)
        tag_start = match.start()
        tag_end = match.end()
        
        # 从 URL 中只提取最后的文件名部分（xxx.jpg）
        filename = url.split('/')[-1]
        
        # 尝试匹配 - 只匹配最后文件名部分
        best_match = None
        best_score = 0
        
        generated_filename = _strip_url_suffix(filename).lower()
        
        for orig_path in original_paths:
            orig_filename = orig_path.split('/')[-1].lower()
            score = SequenceMatcher(None, generated_filename, orig_filename).ratio()
            if score > best_score:
                best_score = score
                best_match = orig_path
        
        if best_match and _is_acceptable_image_match(generated_filename, best_match.split('/')[-1], best_score):
            # 匹配成功，替换为正确路径（保持固定URL前缀）
            # 只替换最后文件名部分，保持前缀不变
            full_url = _path_to_image_url(best_match, base_prefix)
            
            standard_format = f'![图片]({full_url})'
            result = result[:tag_start] + standard_format + result[tag_end:]
        else:
            # 匹配度低，删除整个图片
            result = result[:tag_start] + result[tag_end:]
    
    return result


def get_image_prompt_guidance() -> str:
    """获取图片格式处理的提示词指导"""
    return """【图片格式特别要求】
- 只允许引用本次参考资料中真实出现的图片链接
- 输出图片时使用 Markdown 图片标签：alt 文本固定为「图片」，括号内必须是参考资料里的原始完整 URL
- 图片文件名、路径和 URL 前缀必须逐字复制，不得补全、改写、截断、转义或根据记忆/示例构造
- 不要输出提示词中的占位符或示例文件名，不要把普通文字描述改写成图片链接
- 图片规则只用于生成图片标签，不要围绕图片资源状态输出任何文字说明
- 多个图片可在对应位置分别引用并保持原始顺序；同一图片只引用一次"""