Agent/examples/piaoquan_needs/topic_build_pattern_tools.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Pattern 数据查询工具

提供给 TopicBuildClaudeCodeAgent 使用的 pattern 挖掘结果查询函数。
所有工具均为只读操作，不修改数据库。

工具列表:
  - get_category_tree: 分类树快照（紧凑文本格式）
  - get_frequent_itemsets: 搜索频繁项集（按分类ID、维度模式筛选，分组返回）
  - get_itemset_detail: 单个项集详情
  - get_post_elements: 帖子结构化元素
  - search_elements: 按关键词搜索元素
  - get_element_category_chain: 元素反查分类链
  - get_category_detail: 分类节点详情
  - search_categories: 按关键词搜索分类
  - get_category_elements: 分类下的元素列表

【重要】所有作为 Agent tool 注册的函数，必须包含完整的 docstring 签名。
"""
import json
from typing import Any

from agent import tool
from examples.piaoquan_needs.topic_build_agent_context import TopicBuildAgentContext
from log_capture import log
import pattern_service


def _log_tool_input(tool_name: str, params: dict):
    """工具调用前立即打印参数"""
    log(f"\n[FOLD:🔧 {tool_name}]")
    log(f"[FOLD:📥 调用参数]")
    log(json.dumps(params, ensure_ascii=False, indent=2))
    log(f"[/FOLD]")


def _log_tool_output(tool_name: str, result: str) -> str:
    """工具执行完成后打印返回内容"""
    log(f"[FOLD:📤 返回内容]")
    log(result)
    log(f"[/FOLD]")
    log(f"[/FOLD]\n")
    return result


# ============================================================================
# 执行 & 配置 & 分类树
# ============================================================================

@tool(
    "获取分类树结构快照。分类树是所有数据的骨架——帖子中的元素按'实质/形式/意图'三个维度归类到树形分类节点中。"
    "\n\n返回紧凑文本格式，包含每个分类节点的名称、层级和元素数量。这是理解数据整体结构的起点。"
    "\n\n使用场景："
    "\n- 启动时调用，了解数据涵盖哪些主题领域和维度"
    "\n- 按 source_type 筛选单个维度，聚焦分析"
    "\n- 从树中发现感兴趣的分类节点，获取其 ID 后用于 get_frequent_itemsets、get_category_detail 等"
)
def get_category_tree(source_type: str = None) -> str:
    """获取当前执行的分类树快照，返回紧凑文本格式（节省token）。

    Args:
        source_type: 按元素类型筛选：实质/形式/意图。不传则返回所有类型。

    Returns:
        分类树的紧凑文本字符串。
    """
    execution_id = TopicBuildAgentContext.get_execution_id()
    params = {"execution_id": execution_id, "source_type": source_type}
    _log_tool_input("get_category_tree", params)

    result = pattern_service.get_category_tree_compact(execution_id, source_type=source_type)

    return _log_tool_output("get_category_tree", result)


# ============================================================================
# 项集查询
# ============================================================================
@tool("获取频繁项集——即经常在同一帖子中共同出现的分类组合。结果按 dimension_mode/depth 分组返回，每组各返回 top_n 条。"
      "\n\n核心概念：每个频繁项集 = 一组高概率共现的分类节点。absolute_support = 同时出现在多少个帖子中。"
      "\n\n返回结构：groups 字典，key 为 'dimension_mode/target_depth'，每组含 dimension_mode、target_depth、total（该组总数）、itemsets 列表。"
      "\n\n使用场景："
      "\n- 全局探索：不传 category_ids，浏览支持度最高的共现模式"
      "\n- 定向查询：传入 category_ids=[A]，获取所有包含分类A的项集，即A和哪些分类共现"
      "\n- 交叉共现：传入 category_ids=[A,B]，获取同时包含A和B的项集，发现A+B还经常和什么一起出现"
      "\n- 渐进探索：先查 category_ids=[A] 发现B共现多，再查 category_ids=[A,B] 缩小范围"
      "\n- 维度聚焦：用 dimension_mode 筛选特定挖掘模式的结果"
      "\n- 帖子范围筛选：用 account_name/merge_leve2 限定帖子来源，筛选后重算 support"
      "\n\n提示：category_ids 需要分类节点ID，可先用 search_categories 按名称查ID。项集详情（匹配帖子等）通过 get_itemset_detail 获取。")
def get_frequent_itemsets(
        top_n: int = 20,
        category_ids: list = None,
        dimension_mode: str = None,
        min_support: int = None,
        min_item_count: int = None,
        max_item_count: int = None,
        sort_by: str = "absolute_support",
        account_name=None,
        merge_leve2=None,
) -> str:
    """获取频繁项集——即经常在同一帖子中共同出现的分类组合，按 dimension_mode/depth 分组返回。

    使用场景:
    - 全局探索: 不传 category_ids，浏览支持度最高的共现模式
    - 定向查询: category_ids=[A]，获取包含分类A的所有项集，即A和哪些分类共现
    - 交叉共现: category_ids=[A,B]，获取同时包含A和B的项集，发现A+B还经常和什么一起出现
    - 渐进探索: 先查 [A] 发现B共现多，再查 [A,B] 缩小范围，逐步聚焦
    - 维度聚焦: dimension_mode 筛选特定挖掘模式

    category_ids 需要分类节点ID，可先用 search_categories 按名称查找。
    返回精简信息，详情（匹配帖子等）通过 get_itemset_detail 获取。

    Args:
        top_n: 返回前N个项集，默认20。
        category_ids: 分类节点ID列表（AND逻辑）。传入后返回同时包含所有这些分类的项集；不传则返回全局Top。
        dimension_mode: 挖掘维度模式筛选。full=点类型×元素类型（混合），substance_form_only=仅元素类型，point_type_only=仅点类型。
        min_support: 最低绝对支持度（共现帖子数）阈值。
        min_item_count: 项集最少包含的分类数量。
        max_item_count: 项集最多包含的分类数量。
        sort_by: 排序方式：absolute_support=共现帖子数（默认），support=相对支持度，item_count=分类数量。
        account_name: 按账号名筛选，支持单个字符串或列表（多个取OR）。
        merge_leve2: 按二级品类筛选，支持单个字符串或列表（多个取OR）。

    Returns:
        按 dimension_mode/depth 分组的项集JSON，每组含 itemsets 列表。
    """
    execution_id = TopicBuildAgentContext.get_execution_id()
    params = {
        "execution_id": execution_id, "top_n": top_n,
        "category_ids": category_ids, "dimension_mode": dimension_mode,
        "min_support": min_support,
        "min_item_count": min_item_count, "max_item_count": max_item_count,
        "sort_by": sort_by,
        "account_name": account_name, "merge_leve2": merge_leve2,
    }
    _log_tool_input("get_frequent_itemsets", params)

    data = pattern_service.search_top_itemsets(
        execution_id=execution_id, top_n=top_n,
        category_ids=category_ids, dimension_mode=dimension_mode,
        min_support=min_support,
        min_item_count=min_item_count, max_item_count=max_item_count,
        sort_by=sort_by,
        account_name=account_name, merge_leve2=merge_leve2,
    )
    result = json.dumps(data, ensure_ascii=False, indent=2)
    return _log_tool_output("get_frequent_itemsets", result)

@tool("获取一个或多个频繁项集的完整详情。相比 get_frequent_itemsets 的精简列表，这里返回每个项集的所有信息。"
    "\n\n返回内容：每个项集的 dimension_mode（维度模式）、target_depth（挖掘深度）、items 完整结构（含分类路径、元素名称、维度、点类型）、post_ids（匹配的帖子ID列表）、支持度等。"
    "\n\n使用场景："
    "\n- 从 get_frequent_itemsets 中发现有价值的项集后，批量查看其匹配了哪些帖子"
    "\n- 获取 post_ids 后可传给 get_post_elements 查看帖子的具体内容"
    "\n- 支持传入多个 itemset_id，一次获取多个项集的详情，减少调用次数")
def get_itemset_detail(itemset_ids) -> str:
    """获取一个或多个频繁项集的详情，包括每个项集的维度模式、depth、items 结构（含分类路径、元素名称、维度、点类型）和匹配的帖子ID列表。

    Args:
        itemset_ids: 项集ID，单个整数或整数列表。

    Returns:
        项集详情列表的JSON字符串，每项含 id, dimension_mode, target_depth, items, post_ids, absolute_support。
    """
    if isinstance(itemset_ids, int):
        itemset_ids = [itemset_ids]
    params = {"itemset_ids": itemset_ids}
    _log_tool_input("get_itemset_detail", params)

    data = pattern_service.get_itemset_posts(itemset_ids)
    if not data:
        return _log_tool_output("get_itemset_detail", f"未找到 itemset_ids={itemset_ids} 的项集")

    result = json.dumps(data, ensure_ascii=False, indent=2)
    return _log_tool_output("get_itemset_detail", result)


# ============================================================================
# 帖子 & 元素
# ============================================================================
@tool("获取指定帖子的结构化元素数据。每个帖子的内容被拆解为多个'选题点'，每个点下有实质/形式/意图三个维度的元素。"
    "\n\n返回结构：按帖子分组 → 按点类型分组 → 每个点含 point_text（点的原文）和 elements（三维度元素列表）。"
    "\n\n使用场景："
    "\n- 从 get_itemset_detail 获取 post_ids 后，深入查看这些帖子的具体内容结构"
    "\n- 验证某个共现模式在帖子中的实际表现"
    "\n- 发现帖子中未被分类捕捉到的内容细节")
def get_post_elements(post_ids: list) -> str:
    """获取指定帖子的结构化元素数据。按帖子分组，每个帖子按 点类型→元素类型 组织。用于深入了解某个项集匹配的帖子具体内容。

    Args:
        post_ids: 帖子ID列表（建议每次不超过10个）。

    Returns:
        帖子元素数据的JSON字符串，结构: {post_id: {point_type: [{point_text, elements: {实质, 形式, 意图}}]}}。
    """
    execution_id = TopicBuildAgentContext.get_execution_id()
    params = {"execution_id": execution_id, "post_ids": post_ids}
    _log_tool_input("get_post_elements", params)

    if len(post_ids) > 20:
        return _log_tool_output("get_post_elements",
                                f"错误: post_ids 数量过多({len(post_ids)})，请每次不超过20个")

    data = pattern_service.get_post_elements(execution_id, post_ids)
    result = json.dumps(data, ensure_ascii=False, indent=2)
    return _log_tool_output("get_post_elements", result)

@tool("按名称关键词搜索元素。元素是帖子中的具体内容实体（如'水煮鱼'、'短视频'），归属于分类树的叶子节点下。"
    "\n\n返回去重聚合结果：每个元素附带 point_types（该元素出现在哪些点类型中，如['灵感点','关键点']）、"
    "所属分类（category_id、category_path）、出现次数和帖子数。"
    "\n\n使用场景："
    "\n- 从某个关键词出发，找到相关元素及其分类归属"
    "\n- 通过 point_types 了解元素在灵感点/目的点/关键点中的分布"
    "\n- 获取元素名称后，传给 get_element_co_occurrences 查共现关系"
    "\n- 通过元素的 category_id 桥接到 get_frequent_itemsets 做分类级分析")
def search_elements(keyword: str, element_type: str = None, limit: int = 50,
                    account_name=None, merge_leve2=None) -> str:
    """按名称关键词搜索元素。返回去重聚合后的元素列表，每个元素附带其所属分类信息（category_id、category_path）、出现次数和帖子数。

    使用场景:
    - 从某个关键词出发，找到相关元素及其分类归属
    - 了解某个元素在数据中出现的频率

    Args:
        keyword: 搜索关键词（模糊匹配元素名称）。
        element_type: 按维度筛选：实质/形式/意图。不传则搜索所有维度。
        limit: 最多返回数量，默认50。
        account_name: 按账号名筛选，支持单个字符串或列表（多个取OR）。
        merge_leve2: 按二级品类筛选，支持单个字符串或列表（多个取OR）。

    Returns:
        元素列表的JSON字符串，每个元素含 name、element_type、category_id、category_path、occurrence_count、post_count。
    """
    execution_id = TopicBuildAgentContext.get_execution_id()
    params = {"execution_id": execution_id, "keyword": keyword,
              "element_type": element_type, "limit": limit,
              "account_name": account_name, "merge_leve2": merge_leve2}
    _log_tool_input("search_elements", params)

    data = pattern_service.search_elements(execution_id, keyword, element_type=element_type, limit=limit,
                                           account_name=account_name, merge_leve2=merge_leve2)
    result = json.dumps({
        "keyword": keyword,
        "count": len(data),
        "elements": data,
    }, ensure_ascii=False, indent=2)
    return _log_tool_output("search_elements", result)

@tool("从元素名称反查其所属的完整分类链。一个元素可能归属于多个分类节点，此工具返回每个归属分类从根到叶的完整祖先路径。"
    "\n\n使用场景："
    "\n- 知道一个元素名称，想了解它在分类树中的位置"
    "\n- 从元素出发向上回溯分类层级，获取 category_id 用于 get_frequent_itemsets"
    "\n- 理解同一元素在不同维度下的分类归属差异")
def get_element_category_chain(element_name: str, element_type: str = None) -> str:
    """从元素名称反查其所属分类链。返回该元素出现在哪些分类下，以及每个分类的完整祖先路径（从根到叶）。

    使用场景:
    - 知道一个元素名称，想了解它在分类树中的位置
    - 从元素出发，向上回溯分类层级，为泛化推理提供路径

    Args:
        element_name: 元素名称（精确匹配）。
        element_type: 按维度筛选：实质/形式/意图。不传则查所有维度。

    Returns:
        分类链列表的JSON字符串，每条含 category_id、category_path、ancestors 链（从根到叶）。
    """
    execution_id = TopicBuildAgentContext.get_execution_id()
    params = {"execution_id": execution_id, "element_name": element_name,
              "element_type": element_type}
    _log_tool_input("get_element_category_chain", params)

    data = pattern_service.get_element_category_chain(execution_id, element_name, element_type=element_type)
    result = json.dumps({
        "element_name": element_name,
        "category_chains": data,
    }, ensure_ascii=False, indent=2)
    return _log_tool_output("get_element_category_chain", result)


# ============================================================================
# 分类导航
# ============================================================================
@tool("获取分类节点的完整上下文信息，用于在分类树中导航和理解某个分类的位置。"
    "\n\n返回内容：自身信息（名称、层级、元素数）、祖先链（从根到当前的路径）、直接子节点、同级兄弟节点、"
    "该分类下的元素列表（Top100，每个元素含 point_types 列表表示出现在哪些点类型中）。"
    "\n\n使用场景："
    "\n- 从 get_frequent_itemsets 中发现某个分类后，了解它的层级上下文"
    "\n- 向上泛化：查看祖先节点，理解更宏观的领域"
    "\n- 向下细化：查看子节点，找到更具体的方向"
    "\n- 平行探索：查看兄弟节点，发现同级别的其他内容领域"
    "\n- 获取子节点的 category_id 后可传给 get_frequent_itemsets 做进一步分析")
def get_category_detail(category_id: int) -> str:
    """获取分类节点的完整上下文。包括: 自身信息、祖先链（从根到当前节点的路径）、直接子节点列表、同级兄弟节点列表、该分类下的元素列表（去重聚合Top100）。

    使用场景:
    - 查看某个分类节点的全貌
    - 从分类出发向上回溯（祖先）或向下展开（子节点）
    - 查看同级兄弟节点，发现平行的内容领域

    Args:
        category_id: 分类节点ID（TopicPatternCategory.id）。

    Returns:
        分类详情的JSON字符串，含 category、ancestors、children、siblings、elements。
    """
    execution_id = TopicBuildAgentContext.get_execution_id()
    params = {"execution_id": execution_id, "category_id": category_id}
    _log_tool_input("get_category_detail", params)

    data = pattern_service.get_category_detail_with_context(execution_id, category_id)
    if not data:
        return _log_tool_output("get_category_detail", f"未找到 category_id={category_id}")

    result = json.dumps(data, ensure_ascii=False, indent=2)
    return _log_tool_output("get_category_detail", result)

@tool("按名称关键词搜索分类节点。分类节点是分类树上的一个层级（如'中餐'），区别于具体元素（如'水煮鱼'）。"
    "\n\n返回匹配的分类列表，含 id、name、path、level、element_count、point_types（该分类下元素涉及的点类型列表，如['灵感点','关键点']）。"
    "\n\n使用场景："
    "\n- 获取分类的 category_id，用于 get_frequent_itemsets(category_ids=[...]) 查共现"
    "\n- 通过 point_types 了解分类在灵感点/目的点/关键点中的分布"
    "\n- 获取 category_id 后传给 get_category_detail 查看层级上下文"
    "\n- 作为 get_frequent_itemsets 的前置步骤：先按名称找到 ID，再查频繁项集")
def search_categories(keyword: str, source_type: str = None) -> str:
    """按名称关键词搜索分类节点。返回匹配的分类列表，含 id、name、path、level、element_count 等。

    使用场景:
    - 用关键词定位分类节点，然后用 get_category_detail 或 get_frequent_itemsets_by_category 进一步探索

    Args:
        keyword: 搜索关键词（模糊匹配分类名称）。
        source_type: 按维度筛选：实质/形式/意图。不传则搜索所有维度。

    Returns:
        分类列表的JSON字符串。
    """
    execution_id = TopicBuildAgentContext.get_execution_id()
    params = {"execution_id": execution_id, "keyword": keyword, "source_type": source_type}
    _log_tool_input("search_categories", params)

    data = pattern_service.search_categories(execution_id, keyword, source_type=source_type)
    result = json.dumps({
        "keyword": keyword,
        "count": len(data),
        "categories": data,
    }, ensure_ascii=False, indent=2)
    return _log_tool_output("search_categories", result)

@tool("获取某个分类节点下的具体元素列表。分类是抽象方向（如'中餐'），元素是具体实例（如'水煮鱼'、'麻婆豆腐'）。"
    "\n\n返回按名称去重聚合的元素列表，按出现次数降序，含 name、element_type、point_types（该元素出现在哪些点类型中）、occurrence_count、post_count。"
    "\n\n使用场景："
    "\n- 从分类节点下钻到具体元素，了解该分类包含哪些内容"
    "\n- 通过 point_types 了解元素在灵感点/目的点/关键点中的分布"
    "\n- 获取元素名称后，传给 get_element_co_occurrences 查元素级共现"
    "\n- 从频繁项集的分类出发，落地到可用于选题的具体元素")
def get_category_elements(category_id: int, account_name=None, merge_leve2=None) -> str:
    """获取某个分类节点下的元素列表（按名称去重聚合），按出现次数降序。

    Args:
        category_id: 分类节点ID。
        account_name: 按账号名筛选，支持单个字符串或列表（多个取OR）。
        merge_leve2: 按二级品类筛选，支持单个字符串或列表（多个取OR）。

    Returns:
        元素列表的JSON字符串，每个元素含 name、element_type、occurrence_count、post_count。
    """
    execution_id = TopicBuildAgentContext.get_execution_id()
    params = {"category_id": category_id, "account_name": account_name, "merge_leve2": merge_leve2}
    _log_tool_input("get_category_elements", params)

    data = pattern_service.get_category_elements(category_id, execution_id=execution_id,
                                                 account_name=account_name, merge_leve2=merge_leve2)
    result = json.dumps({
        "category_id": category_id,
        "element_count": len(data),
        "elements": data,
    }, ensure_ascii=False, indent=2)
    return _log_tool_output("get_category_elements", result)


# ============================================================================
# 共现查询
# ============================================================================
@tool("查询分类级共现关系——找到同时包含指定分类下元素的帖子，统计这些帖子中其他分类的出现频率。"
    "\n\n核心概念：与 get_frequent_itemsets 不同，此工具是实时从帖子数据中计算共现，"
    "不依赖预计算的频繁项集，因此可以灵活组合任意分类进行探索。"
    "\n\n使用场景："
    "\n- 单分类探索：category_ids=[123]，发现经常和该分类一起出现的其他分类"
    "\n- 多分类交叉：category_ids=[123,456]，发现同时涉及这两个分类的帖子中还包含什么分类"
    "\n- 渐进聚焦：先查单个分类，发现高频共现后叠加查询缩小范围"
    "\n- 验证频繁项集：将 get_frequent_itemsets 中发现的模式用此工具做更细粒度的验证"
    "\n\n提示：需要分类节点ID，可先用 search_categories 按名称查找。")
def get_category_co_occurrences(category_ids: list, top_n: int = 30,
                                account_name=None, merge_leve2=None) -> str:
    """查询多个分类的共现关系。找到同时包含所有指定分类下元素的帖子，返回这些帖子中其他分类的出现频率。

    支持叠加多分类，传入越多分类，结果越精确（交集缩小）。

    Args:
        category_ids: 分类节点ID列表（AND逻辑）。传入多个时取帖子交集。
        top_n: 返回共现频率最高的前N个分类，默认30。
        account_name: 按账号名筛选帖子范围，支持单个字符串或列表（多个取OR）。
        merge_leve2: 按二级品类筛选帖子范围，支持单个字符串或列表（多个取OR）。

    Returns:
        共现分类列表的JSON字符串，含 matched_post_count（交集帖子数）和 co_categories（共现分类排名）。
    """
    execution_id = TopicBuildAgentContext.get_execution_id()
    params = {"execution_id": execution_id, "category_ids": category_ids, "top_n": top_n,
              "account_name": account_name, "merge_leve2": merge_leve2}
    _log_tool_input("get_category_co_occurrences", params)

    if not category_ids:
        return _log_tool_output("get_category_co_occurrences", "错误: category_ids 不能为空")

    data = pattern_service.get_category_co_occurrences(
        execution_id=execution_id, category_ids=category_ids, top_n=top_n,
        account_name=account_name, merge_leve2=merge_leve2,
    )
    result = json.dumps(data, ensure_ascii=False, indent=2)
    return _log_tool_output("get_category_co_occurrences", result)

@tool("查询元素级共现关系——找到同时包含指定元素的帖子，统计这些帖子中其他元素的出现频率。"
    "\n\n与 get_category_co_occurrences（分类级共现）互补：此工具在具体元素粒度上分析共现，更适合落地到选题细节。"
    "\n\n返回的每个共现元素含 point_types（该元素出现在哪些点类型中，如['灵感点','关键点']）。"
    "\n\n使用场景："
    "\n- 单元素探索：element_names=['猫咪']，发现经常和猫咪一起出现的其他元素"
    "\n- 多元素交叉：element_names=['猫咪','拟人化']，发现同时涉及这两个元素的帖子还包含什么"
    "\n- 渐进聚焦：先查单个元素，发现高频共现后叠加查询缩小范围"
    "\n\n提示：element_names 需要精确匹配，可先用 search_elements 按关键词查找确切名称。")
def get_element_co_occurrences(element_names: list, top_n: int = 30,
                               account_name=None, merge_leve2=None) -> str:
    """查询多个元素的共现关系。找到同时包含所有指定元素的帖子，返回这些帖子中其他元素的出现频率。

    支持叠加多元素，传入越多元素，结果越精确（交集缩小）。

    Args:
        element_names: 元素名称列表（精确匹配）。传入多个时取帖子交集。
        top_n: 返回共现频率最高的前N个元素，默认30。
        account_name: 按账号名筛选帖子范围，支持单个字符串或列表（多个取OR）。
        merge_leve2: 按二级品类筛选帖子范围，支持单个字符串或列表（多个取OR）。

    Returns:
        共现元素列表的JSON字符串，含 matched_post_count（交集帖子数）和 co_elements（共现元素排名）。
    """
    execution_id = TopicBuildAgentContext.get_execution_id()
    params = {"execution_id": execution_id, "element_names": element_names, "top_n": top_n,
              "account_name": account_name, "merge_leve2": merge_leve2}
    _log_tool_input("get_element_co_occurrences", params)

    if not element_names:
        return _log_tool_output("get_element_co_occurrences", "错误: element_names 不能为空")

    data = pattern_service.get_element_co_occurrences(
        execution_id=execution_id, element_names=element_names, top_n=top_n,
        account_name=account_name, merge_leve2=merge_leve2,
    )
    result = json.dumps(data, ensure_ascii=False, indent=2)
    return _log_tool_output("get_element_co_occurrences", result)