kg_agent/utils/text_utils.py

#!/usr/bin/env python3
"""
文本处理工具模块
提供文本相似度、编辑距离等功能
"""

from typing import Tuple


def edit_distance(str1: str, str2: str) -> int:
    """
    计算两个字符串的编辑距离（Levenshtein距离）

    编辑距离是指将一个字符串转换为另一个字符串所需的最少编辑操作次数。
    允许的操作包括：插入、删除、替换字符。

    Args:
        str1: 第一个字符串
        str2: 第二个字符串

    Returns:
        int: 编辑距离（最少操作次数）

    Examples:
        >>> edit_distance("kitten", "sitting")
        3
        >>> edit_distance("hello", "hello")
        0
        >>> edit_distance("", "abc")
        3
    """
    len1, len2 = len(str1), len(str2)

    # 创建 DP 表格，dp[i][j] 表示 str1[:i] 转换为 str2[:j] 的编辑距离
    dp = [[0] * (len2 + 1) for _ in range(len1 + 1)]

    # 初始化第一行和第一列
    for i in range(len1 + 1):
        dp[i][0] = i  # str1[:i] 转换为空字符串需要 i 次删除

    for j in range(len2 + 1):
        dp[0][j] = j  # 空字符串转换为 str2[:j] 需要 j 次插入

    # 动态规划填充表格
    for i in range(1, len1 + 1):
        for j in range(1, len2 + 1):
            if str1[i - 1] == str2[j - 1]:
                # 字符相同，不需要操作
                dp[i][j] = dp[i - 1][j - 1]
            else:
                # 取三种操作的最小值
                dp[i][j] = min(
                    dp[i - 1][j] + 1,      # 删除 str1[i-1]
                    dp[i][j - 1] + 1,      # 插入 str2[j-1]
                    dp[i - 1][j - 1] + 1   # 替换 str1[i-1] 为 str2[j-1]
                )

    return dp[len1][len2]


def similarity(str1: str, str2: str) -> float:
    """
    基于编辑距离计算两个字符串的相似度

    相似度 = 1 - (编辑距离 / 较长字符串的长度)
    返回值在 [0, 1] 区间，1 表示完全相同，0 表示完全不同

    Args:
        str1: 第一个字符串
        str2: 第二个字符串

    Returns:
        float: 相似度，范围 [0, 1]

    Examples:
        >>> similarity("hello", "hello")
        1.0
        >>> similarity("hello", "hallo")
        0.8
        >>> similarity("abc", "xyz")
        0.0
    """
    if not str1 and not str2:
        return 1.0

    max_len = max(len(str1), len(str2))
    if max_len == 0:
        return 1.0

    distance = edit_distance(str1, str2)
    return 1 - (distance / max_len)


def jaccard_similarity(str1: str, str2: str) -> float:
    """
    计算两个字符串的Jaccard相似度（基于字符集合）

    Jaccard相似度 = 交集大小 / 并集大小
    不考虑字符位置，只考虑字符是否出现

    Args:
        str1: 第一个字符串
        str2: 第二个字符串

    Returns:
        float: Jaccard相似度，范围 [0, 1]

    Examples:
        >>> jaccard_similarity("牛逼坏了", "我的牛逼")
        0.5  # 交集{'牛','逼'}, 并集{'牛','逼','坏','了','我','的'}
        >>> jaccard_similarity("hello", "hallo")
        0.8  # 交集{'h','a','l','o'}, 并集{'h','e','l','o','a'}
    """
    if not str1 and not str2:
        return 1.0

    set1 = set(str1)
    set2 = set(str2)

    intersection = set1 & set2
    union = set1 | set2

    if len(union) == 0:
        return 1.0

    return len(intersection) / len(union)


def longest_common_subsequence(str1: str, str2: str) -> str:
    """
    计算两个字符串的最长公共子序列（LCS）

    子序列不要求连续，但要求保持相对顺序

    Args:
        str1: 第一个字符串
        str2: 第二个字符串

    Returns:
        str: 最长公共子序列字符串

    Examples:
        >>> longest_common_subsequence("牛逼坏了", "我的牛逼")
        "牛逼"
        >>> longest_common_subsequence("ABCDGH", "AEDFHR")
        "ADH"
    """
    len1, len2 = len(str1), len(str2)

    # 创建 DP 表格，dp[i][j] 表示 str1[:i] 和 str2[:j] 的 LCS 长度
    dp = [[0] * (len2 + 1) for _ in range(len1 + 1)]

    # 填充 DP 表格
    for i in range(1, len1 + 1):
        for j in range(1, len2 + 1):
            if str1[i - 1] == str2[j - 1]:
                dp[i][j] = dp[i - 1][j - 1] + 1
            else:
                dp[i][j] = max(dp[i - 1][j], dp[i][j - 1])

    # 回溯构建 LCS 字符串
    lcs = []
    i, j = len1, len2

    while i > 0 and j > 0:
        if str1[i - 1] == str2[j - 1]:
            lcs.append(str1[i - 1])
            i -= 1
            j -= 1
        elif dp[i - 1][j] > dp[i][j - 1]:
            i -= 1
        else:
            j -= 1

    return ''.join(reversed(lcs))


def lcs_similarity(str1: str, str2: str) -> float:
    """
    基于最长公共子序列（LCS）计算相似度

    相似度 = 2 * LCS长度 / (str1长度 + str2长度)

    Args:
        str1: 第一个字符串
        str2: 第二个字符串

    Returns:
        float: LCS相似度，范围 [0, 1]

    Examples:
        >>> lcs_similarity("牛逼坏了", "我的牛逼")
        0.5  # LCS="牛逼" (长度2), 2*2/(4+4)=0.5
        >>> lcs_similarity("hello", "hallo")
        0.8  # LCS="hllo" (长度4), 2*4/(5+5)=0.8
    """
    if not str1 and not str2:
        return 1.0

    lcs = longest_common_subsequence(str1, str2)
    lcs_len = len(lcs)

    total_len = len(str1) + len(str2)
    if total_len == 0:
        return 1.0

    return 2 * lcs_len / total_len


def text_similarity(str1: str, str2: str, method: str = "levenshtein") -> float:
    """
    计算两个字符串的相似度（统一接口）

    Args:
        str1: 第一个字符串
        str2: 第二个字符串
        method: 算法类型，可选值：
            - "levenshtein": 编辑距离相似度（考虑位置）
            - "jaccard": Jaccard相似度（字符集合）
            - "lcs": LCS相似度（保持顺序）

    Returns:
        float: 相似度，范围 [0, 1]

    Examples:
        >>> text_similarity("hello", "hallo", method="levenshtein")
        0.8
        >>> text_similarity("牛逼坏了", "我的牛逼", method="jaccard")
        0.33
        >>> text_similarity("牛逼坏了", "我的牛逼", method="lcs")
        0.5
    """
    method = method.lower()

    if method == "levenshtein":
        return similarity(str1, str2)
    elif method == "jaccard":
        return jaccard_similarity(str1, str2)
    elif method == "lcs":
        return lcs_similarity(str1, str2)
    else:
        raise ValueError(f"Unknown method: {method}. Choose from: levenshtein, jaccard, lcs")


def edit_distance_with_operations(str1: str, str2: str) -> Tuple[int, list]:
    """
    计算编辑距离并返回操作序列

    Args:
        str1: 第一个字符串
        str2: 第二个字符串

    Returns:
        Tuple[int, list]: (编辑距离, 操作列表)
        操作列表格式：[("operation", char, position), ...]
        operation 可以是: "insert", "delete", "replace"
    """
    len1, len2 = len(str1), len(str2)

    # 创建 DP 表格
    dp = [[0] * (len2 + 1) for _ in range(len1 + 1)]

    # 初始化
    for i in range(len1 + 1):
        dp[i][0] = i
    for j in range(len2 + 1):
        dp[0][j] = j

    # 填充 DP 表格
    for i in range(1, len1 + 1):
        for j in range(1, len2 + 1):
            if str1[i - 1] == str2[j - 1]:
                dp[i][j] = dp[i - 1][j - 1]
            else:
                dp[i][j] = min(
                    dp[i - 1][j] + 1,
                    dp[i][j - 1] + 1,
                    dp[i - 1][j - 1] + 1
                )

    # 回溯获取操作序列
    operations = []
    i, j = len1, len2

    while i > 0 or j > 0:
        if i == 0:
            operations.append(("insert", str2[j - 1], j - 1))
            j -= 1
        elif j == 0:
            operations.append(("delete", str1[i - 1], i - 1))
            i -= 1
        elif str1[i - 1] == str2[j - 1]:
            i -= 1
            j -= 1
        else:
            min_val = min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1])
            if dp[i - 1][j - 1] == min_val:
                operations.append(("replace", f"{str1[i-1]}->{str2[j-1]}", i - 1))
                i -= 1
                j -= 1
            elif dp[i - 1][j] == min_val:
                operations.append(("delete", str1[i - 1], i - 1))
                i -= 1
            else:
                operations.append(("insert", str2[j - 1], j - 1))
                j -= 1

    operations.reverse()
    return dp[len1][len2], operations


def main():
    """命令行接口"""
    import argparse

    parser = argparse.ArgumentParser(
        description="计算两个文本的编辑距离和相似度（支持多种算法）",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  python utils/text_utils.py --str1 "hello" --str2 "hallo"
  python utils/text_utils.py --str1 "牛逼坏了" --str2 "我的牛逼" --method jaccard
  python utils/text_utils.py --str1 "hello" --str2 "hallo" --verbose
        """
    )

    parser.add_argument("--str1", required=True, help="第一个字符串")
    parser.add_argument("--str2", required=True, help="第二个字符串")
    parser.add_argument("--method", "-m", default="all",
                       choices=["all", "levenshtein", "jaccard", "lcs"],
                       help="相似度算法（默认显示所有）")
    parser.add_argument("--verbose", "-v", action="store_true",
                       help="显示详细的操作步骤")

    args = parser.parse_args()

    print(f"\n文本对比结果:")
    print(f"{'='*50}")
    print(f"字符串1: {args.str1}")
    print(f"字符串2: {args.str2}")
    print(f"{'='*50}")

    # 根据method参数显示相应的结果
    if args.method == "all":
        distance = edit_distance(args.str1, args.str2)
        levenshtein_sim = similarity(args.str1, args.str2)
        jaccard_sim = jaccard_similarity(args.str1, args.str2)
        lcs_sim = lcs_similarity(args.str1, args.str2)
        lcs_str = longest_common_subsequence(args.str1, args.str2)

        print(f"编辑距离 (Levenshtein): {distance}")
        print(f"编辑距离相似度:         {levenshtein_sim:.2%}")
        print(f"Jaccard相似度:          {jaccard_sim:.2%} (基于字符集合)")
        print(f"LCS相似度:              {lcs_sim:.2%} (基于公共子序列)")
        print(f"最长公共子序列:         '{lcs_str}'")
    elif args.method == "levenshtein":
        distance = edit_distance(args.str1, args.str2)
        sim = similarity(args.str1, args.str2)
        print(f"编辑距离: {distance}")
        print(f"相似度:   {sim:.2%}")
    elif args.method == "jaccard":
        sim = jaccard_similarity(args.str1, args.str2)
        print(f"Jaccard相似度: {sim:.2%}")
    elif args.method == "lcs":
        sim = lcs_similarity(args.str1, args.str2)
        lcs_str = longest_common_subsequence(args.str1, args.str2)
        print(f"LCS相似度:      {sim:.2%}")
        print(f"最长公共子序列: '{lcs_str}'")

    if args.verbose:
        distance, operations = edit_distance_with_operations(args.str1, args.str2)
        print(f"\n编辑操作步骤 (共{len(operations)}步):")
        for idx, (op, char, pos) in enumerate(operations, 1):
            if op == "insert":
                print(f"  {idx}. 插入 '{char}' 到位置 {pos}")
            elif op == "delete":
                print(f"  {idx}. 删除 '{char}' 在位置 {pos}")
            elif op == "replace":
                print(f"  {idx}. 替换位置 {pos} 的字符: {char}")

    print()


if __name__ == "__main__":
    main()