wx-agent/utils/image_utils.py

"""
图片处理工具模块 - 最终版 v4
核心思路（最简单直接）：
1. 找到所有 .jpg/.png 文件名及其周围的 ![图片]( 上下文
2. 不管周围格式如何，也不管文件名是否被模型改过
3. 统一提取文件名 -> 与 original_paths 模匹 -> 替换成正确的 ![图片](完整URL)
"""
import re
from typing import List
from difflib import SequenceMatcher


def extract_image_paths(text: str) -> List[str]:
    """从文本中提取图片路径"""
    pattern = r'images/[a-zA-Z0-9_\-./]+\.(?:jpg|jpeg|png|gif|bmp|webp)'
    matches = re.findall(pattern, text, re.IGNORECASE)
    return list(set(matches))


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 best_score > 0.5:
        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
    
    # ========== 第一轮：检查和修改格式（包括固定URL前缀） ==========
    result = _normalize_image_format(text, base_prefix)
    
    # ========== 第二轮：只检查和替换最后文件名部分 ==========
    result = _normalize_image_paths(result, base_prefix, original_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_EXTENSIONS = r'(?:jpg|jpeg|png|gif|bmp|webp)'
    img_pattern = rf'([a-zA-Z0-9_\-./]+)\.({IMG_EXTENSIONS})'
    
    result = text
    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
        
        # 提取纯文件名（去掉可能的路径前缀）
        pure_filename = filename_with_ext.split('/')[-1]
        
        # 构建完整的标准格式（包括固定URL前缀）
        full_url = base_prefix.rstrip('/') + '/images/' + pure_filename
        standard_format = f'![图片]({full_url})'
        
        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 = 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 best_score > 0.5:
            # 匹配成功，替换为正确路径（保持固定URL前缀）
            # 只替换最后文件名部分，保持前缀不变
            pure_orig_filename = best_match.split('/')[-1]
            full_url = base_prefix.rstrip('/') + '/images/' + pure_orig_filename
            
            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 """【图片格式特别要求】
- 检索结果中的图片链接格式通常为：`![图片](/api/v1/knowledge/files/images/xxx.jpg)`
- 输出时必须严格保留此格式，不得修改、省略或截断
- 图片的 alt 文本必须保持为「图片」
- 图片 URL 必须完整包含 `/api/v1/knowledge/files/images` 前缀
- 不要将图片链接转换为纯文本描述
- 如果有多个图片，在对应位置分别引用，保持原始顺序
- 严禁修改图片文件名或路径结构
- 如果需要展示图片，直接使用原始的 `![图片](URL)` 格式即可"""