Agent/examples_how/overall_derivation/tools/find_pattern.py

"""
查找 Pattern Tool - 从 pattern 库中获取符合条件概率阈值的 pattern

功能：读取账号的 pattern 库，合并去重后按条件概率筛选，返回 topN 条 pattern（含 pattern 名称、条件概率）。
"""

import importlib.util
import json
from pathlib import Path
from typing import Any, Optional

try:
    from agent.tools import tool, ToolResult, ToolContext
except ImportError:
    def tool(*args, **kwargs):
        return lambda f: f
    ToolResult = None  # 仅用 main() 测核心逻辑时可无 agent
    ToolContext = None

# 与 pattern_data_process 一致的 key 定义
TOP_KEYS = [
    "depth_max_with_name",
    "depth_mixed",
    "depth_max_concrete",
    "depth2_medium",
    "depth1_abstract",
]
SUB_KEYS = ["two_x", "one_x", "zero_x"]

# 加载 conditional_ratio_calc（与 find_tree_node 一致）
_utils_dir = Path(__file__).resolve().parent.parent / "utils"
_cond_spec = importlib.util.spec_from_file_location(
    "conditional_ratio_calc",
    _utils_dir / "conditional_ratio_calc.py",
)
_cond_mod = importlib.util.module_from_spec(_cond_spec)
_cond_spec.loader.exec_module(_cond_mod)
calc_pattern_conditional_ratio = _cond_mod.calc_pattern_conditional_ratio

_BASE_INPUT = Path(__file__).resolve().parent.parent / "input"

# 加载 point_match（用于检查 pattern 元素是否匹配帖子选题点）
_point_match_spec = importlib.util.spec_from_file_location(
    "point_match",
    Path(__file__).resolve().parent / "point_match.py",
)
_point_match_mod = importlib.util.module_from_spec(_point_match_spec)
_point_match_spec.loader.exec_module(_point_match_mod)
_match_derivation_to_post_points = _point_match_mod.match_derivation_to_post_points
_load_match_data = _point_match_mod._load_match_data


def _pattern_file(account_name: str) -> Path:
    """pattern 库文件：../input/{account_name}/原始数据/pattern/processed_edge_data.json"""
    return _BASE_INPUT / account_name / "原始数据" / "pattern" / "processed_edge_data.json"


def _slim_pattern(p: dict) -> tuple[float, int, list[str], int]:
    """提取 name 列表（去重保序）、support、length、post_count。"""
    names = [item["name"] for item in (p.get("items") or [])]
    seen = set()
    unique = []
    for n in names:
        if n not in seen:
            seen.add(n)
            unique.append(n)
    support = round(float(p.get("support", 0)), 4)
    length = int(p.get("length", 0))
    post_count = int(p.get("post_count", 0))
    return support, length, unique, post_count


def _merge_and_dedupe(patterns: list[dict]) -> list[dict]:
    """
    按 items 的 name 集合去重（不区分顺序），留 support 最大；
    输出格式保留 s、l、i（nameA+nameB+nameC）及 post_count，供条件概率计算使用。
    """
    key_to_best: dict[tuple, tuple[float, int, int]] = {}
    for p in patterns:
        support, length, unique, post_count = _slim_pattern(p)
        if not unique:
            continue
        key = tuple(sorted(unique))
        if key not in key_to_best or support > key_to_best[key][0]:
            key_to_best[key] = (support, length, post_count)
    out = []
    for k, (s, l, post_count) in key_to_best.items():
        if s < 0.1:
            continue
        out.append({
            "s": s,
            "l": l,
            "i": "+".join(k),
            "post_count": post_count,
        })
    out.sort(key=lambda x: x["s"] * x["l"], reverse=True)
    return out


def _load_and_merge_patterns(account_name: str) -> list[dict]:
    """读取 pattern 库 JSON，按 TOP_KEYS/SUB_KEYS 合并为列表并做合并、去重。"""
    path = _pattern_file(account_name)
    if not path.is_file():
        return []
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    all_patterns = []
    for top in TOP_KEYS:
        if top not in data:
            continue
        block = data[top]
        for sub in SUB_KEYS:
            all_patterns.extend(block.get(sub) or [])
    return _merge_and_dedupe(all_patterns)


def _parse_derived_list(derived_items: list[dict[str, str]]) -> list[tuple[str, str]]:
    """将 agent 传入的 [{"topic": "x", "source_node": "y"}, ...] 转为 DerivedItem 列表。"""
    out = []
    for item in derived_items:
        if isinstance(item, dict):
            topic = item.get("topic") or item.get("已推导的选题点")
            source = item.get("source_node") or item.get("推导来源人设树节点")
            if topic is not None and source is not None:
                out.append((str(topic).strip(), str(source).strip()))
        elif isinstance(item, (list, tuple)) and len(item) >= 2:
            out.append((str(item[0]).strip(), str(item[1]).strip()))
    return out


def get_patterns_by_conditional_ratio(
    account_name: str,
    derived_list: list[tuple[str, str]],
    conditional_ratio_threshold: float,
    top_n: int,
    post_id: str = "",
) -> list[dict[str, Any]]:
    """
    从 pattern 库中获取条件概率 >= 阈值的 pattern，按以下优先级排序后返回 top_n 条：
      1. pattern 元素中直接包含已推导选题点（topic）的排最前；
      2. pattern 元素与任意已推导选题点的匹配分 >= 0.8 的次之（从 match_data 文件读取，
         key 为 (帖子选题点, 人设树节点)，pattern 元素视为人设树节点）；
      3. 按条件概率降序；
      4. 按 length 降序。
    derived_list 为空时，条件概率使用 pattern 自身的 support（s）。
    返回每项：pattern名称（nameA+nameB+nameC）、条件概率。
    """
    merged = _load_and_merge_patterns(account_name)
    if not merged:
        return []
    base_dir = _BASE_INPUT
    scored: list[tuple[dict, float]] = []

    if not derived_list:
        # derived_items 为空：条件概率取 pattern 本身的 support (s)
        for p in merged:
            ratio = float(p.get("s", 0))
            if ratio >= conditional_ratio_threshold:
                scored.append((p, ratio))
    else:
        for p in merged:
            ratio = calc_pattern_conditional_ratio(
                account_name, derived_list, p, base_dir=base_dir
            )
            if ratio >= conditional_ratio_threshold:
                scored.append((p, ratio))

    derived_topics = {topic for topic, _ in derived_list} if derived_list else set()

    # 次优先：从 match_data 文件加载 (帖子选题点, 人设树节点) -> 匹配分，
    # 用已推导选题点（topic）作为帖子选题点，pattern 元素作为人设树节点，
    # 检查是否存在匹配分 >= 0.8 的组合。
    match_lookup: dict[tuple[str, str], float] = {}
    if derived_topics and post_id:
        match_lookup = _load_match_data(account_name, post_id)

    def _sort_key(x: tuple[dict, float]) -> tuple:
        p, ratio = x
        elements = set(p["i"].split("+"))
        has_derived = bool(elements & derived_topics)
        has_high_match = False
        if not has_derived and match_lookup:
            for elem in elements:
                for dt in derived_topics:
                    if match_lookup.get((dt, elem), 0.0) >= 0.8:
                        has_high_match = True
                        break
                if has_high_match:
                    break
        return (not has_derived, not has_high_match, -ratio, -p["l"])

    scored.sort(key=_sort_key)
    result = []
    for p, ratio in scored[:top_n]:
        result.append({
            "pattern名称": p["i"],
            "条件概率": round(ratio, 6),
        })
    return result


@tool(
    description="按条件概率从 pattern 库中筛选 pattern，优先返回包含已推导选题点的 pattern，并检查每个 pattern 的元素是否与帖子选题点匹配。"
    "功能：根据账号与已推导选题点（可选），筛选条件概率不低于阈值的 pattern；当 derived_items 非空时，优先返回 pattern 元素中包含已推导选题点的 pattern；同时对每个 pattern 的所有元素做帖子选题点匹配，匹配结果直接包含在返回数据中。"
    "参数：account_name 为账号名；post_id 为帖子ID，用于加载帖子选题点并做匹配判断；derived_items 为已推导选题点列表，每项含 topic（或已推导的选题点）与 source_node（或推导来源人设树节点），可为空，为空时条件概率使用 pattern 自身的 support；conditional_ratio_threshold 为条件概率阈值；top_n 为返回条数上限，默认 100。"
    "返回：ToolResult，output 为可读的 pattern 列表文本，metadata.items 为列表，每项含「pattern名称」（nameA+nameB+nameC 形式）、「条件概率」、「帖子选题点匹配」（匹配到帖子选题点的元素列表，每项含 pattern元素、帖子选题点与匹配分数；若无匹配则为字符串'无匹配帖子选题点'）。"
)
async def find_pattern(
    account_name: str,
    post_id: str,
    derived_items: list[dict[str, str]],
    conditional_ratio_threshold: float,
    top_n: int = 100,
    context: Optional[ToolContext] = None,
) -> ToolResult:
    """
    按条件概率阈值从 pattern 库筛选 pattern，返回最多 top_n 条（按条件概率降序）。
    当 derived_items 非空时，优先返回元素中包含已推导选题点的 pattern。
    返回前对每个 pattern 的所有元素做帖子选题点匹配，匹配结果直接包含在返回数据中。

    参数
    -------
    account_name : 账号名，用于定位该账号的 pattern 库。
    post_id : 帖子ID，用于加载帖子选题点并与 pattern 元素做匹配判断。
    derived_items : 已推导选题点列表，可为空。非空时每项为字典，需含 topic（或「已推导的选题点」）与 source_node（或「推导来源人设树节点」）；为空时各 pattern 的条件概率取其自身 support。
    conditional_ratio_threshold : 条件概率阈值，仅返回条件概率 >= 该值的 pattern。
    top_n : 返回条数上限，默认 100。
    context : 可选，Agent 工具上下文。

    返回
    -------
    ToolResult：
        - title: 结果标题。
        - output: 可读的 pattern 列表文本（每行：pattern名称、条件概率、帖子匹配情况）。
        - metadata: 含 account_name、conditional_ratio_threshold、top_n、count、items；
          items 为列表，每项为 {"pattern名称": str, "条件概率": float,
          "帖子选题点匹配": list[{"pattern元素": str, "帖子选题点": str, "匹配分数": float}] 或 "无匹配帖子选题点"}。
        - 出错时 error 为错误信息。
    """
    pattern_path = _pattern_file(account_name)
    if not pattern_path.is_file():
        return ToolResult(
            title="Pattern 库不存在",
            output=f"pattern 文件不存在: {pattern_path}",
            error="Pattern file not found",
        )
    try:
        derived_list = _parse_derived_list(derived_items or [])
        items = get_patterns_by_conditional_ratio(
            account_name, derived_list, conditional_ratio_threshold, top_n, post_id
        )
        # 批量收集所有 pattern 元素，统一做一次帖子选题点匹配
        if items and post_id:
            all_elements: list[str] = []
            seen_elements: set[str] = set()
            for item in items:
                for elem in item["pattern名称"].split("+"):
                    elem = elem.strip()
                    if elem and elem not in seen_elements:
                        all_elements.append(elem)
                        seen_elements.add(elem)
            matched_results = await _match_derivation_to_post_points(all_elements, account_name, post_id)
            elem_match_map: dict[str, list] = {}
            for m in matched_results:
                elem_match_map.setdefault(m["推导选题点"], []).append({
                    "帖子选题点": m["帖子选题点"],
                    "匹配分数": m["匹配分数"],
                })
            for item in items:
                pattern_matches = []
                for elem in item["pattern名称"].split("+"):
                    elem = elem.strip()
                    for post_match in elem_match_map.get(elem, []):
                        pattern_matches.append({
                            "pattern元素": elem,
                            "帖子选题点": post_match["帖子选题点"],
                            "匹配分数": post_match["匹配分数"],
                        })
                item["帖子选题点匹配"] = pattern_matches if pattern_matches else "无匹配帖子选题点"
        if not items:
            output = f"未找到条件概率 >= {conditional_ratio_threshold} 的 pattern"
        else:
            lines = []
            for x in items:
                match_info = x.get("帖子选题点匹配", "未查询")
                if isinstance(match_info, list):
                    match_str = "、".join(
                        f"{m['pattern元素']}→{m['帖子选题点']}({m['匹配分数']})" for m in match_info
                    )
                else:
                    match_str = str(match_info)
                lines.append(f"- {x['pattern名称']}\t条件概率={x['条件概率']}\t帖子匹配={match_str}")
            output = "\n".join(lines)
        return ToolResult(
            title=f"符合条件概率的 Pattern ({account_name}, 阈值={conditional_ratio_threshold})",
            output=output,
            metadata={
                "account_name": account_name,
                "conditional_ratio_threshold": conditional_ratio_threshold,
                "top_n": top_n,
                "count": len(items),
                "items": items,
            },
        )
    except Exception as e:
        return ToolResult(
            title="查找 Pattern 失败",
            output=str(e),
            error=str(e),
        )


def main() -> None:
    """本地测试：用家有大志账号、已推导选题点，查询符合条件概率阈值的 pattern（含帖子匹配）。"""
    import asyncio

    account_name = "家有大志"
    post_id = "68fb6a5c000000000302e5de"
    # 已推导选题点，每项：已推导的选题点 + 推导来源人设树节点
    derived_items = [
        # {"topic": "分享", "source_node": "分享"},
        {"topic": "柴犬", "source_node": "动物角色"},
        {"topic": "叙事结构", "source_node": "叙事逻辑"},
    ]
    conditional_ratio_threshold = 0.01
    top_n = 100

    # 1）直接调用核心函数（不含帖子匹配，仅验证排序逻辑）
    derived_list = _parse_derived_list(derived_items)
    items = get_patterns_by_conditional_ratio(
        account_name, derived_list, conditional_ratio_threshold, top_n, post_id
    )
    print(f"账号: {account_name}, 阈值: {conditional_ratio_threshold}, top_n: {top_n}")
    print(f"共 {len(items)} 条 pattern:\n")
    for x in items:
        print(f"  - {x['pattern名称']}\t条件概率={x['条件概率']}")

    # 2）有 agent 时通过 tool 接口再跑一遍（含帖子选题点匹配）
    if ToolResult is not None:
        async def run_tool():
            result = await find_pattern(
                account_name=account_name,
                post_id=post_id,
                derived_items=derived_items,
                conditional_ratio_threshold=conditional_ratio_threshold,
                top_n=top_n,
            )
            print("\n--- Tool 返回 ---")
            print(result.output)
        asyncio.run(run_tool())


if __name__ == "__main__":
    main()