Agent/examples/content_tree_analyst/tools/frequent_itemsets.py

"""
频繁项集 API 工具

封装 pattern.aiddit.com 的频繁项集接口，用于查询与指定分类节点
在优质内容中共同出现的关联要素。
"""

import logging

import httpx

from agent.tools import tool
from agent.tools.models import ToolResult

logger = logging.getLogger(__name__)

ITEMSETS_URL = "https://pattern.aiddit.com/api/pattern/tools/get_frequent_itemsets/execute"

HEADERS = {
    "Accept": "*/*",
    "Accept-Language": "zh-CN,zh;q=0.9",
    "Connection": "keep-alive",
    "Content-Type": "application/json",
    "Origin": "https://pattern.aiddit.com",
    "Referer": "https://pattern.aiddit.com/execution/33",
    "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/145.0.0.0 Safari/537.36",
}


@tool(description="查询与指定分类节点在优质内容中共同出现的频繁项集（关联要素），用于扩展制作需求的关联维度")
async def get_frequent_itemsets(
    entity_ids: list,
    top_n: int = 20,
    execution_id: int = 33,
    sort_by: str = "absolute_support",
) -> ToolResult:
    """
    获取与指定分类节点关联的频繁项集。

    Args:
        entity_ids: 分类节点的 entity_id 列表（即搜索接口返回的 entity_id 字段，非 stable_id）
        top_n: 返回前 N 个项集，默认 20
        execution_id: 执行 ID，默认 33
        sort_by: 排序字段，默认 "absolute_support"
    """
    payload = {
        "execution_id": execution_id,
        "args": {
            "top_n": top_n,
            "category_ids": entity_ids,
            "sort_by": sort_by,
        },
    }

    try:
        import json as _json
        async with httpx.AsyncClient(timeout=30.0) as client:
            resp = await client.post(ITEMSETS_URL, json=payload, headers=HEADERS)
            resp.raise_for_status()
            outer = resp.json()

        # result 字段是 JSON 字符串，需要二次解析
        data = _json.loads(outer["result"])
        total = data.get("total", 0)
        groups = data.get("groups", {})

        # 收集所有 group 下的 itemsets
        all_itemsets = []
        for group_key, group in groups.items():
            for itemset in group.get("itemsets", []):
                itemset["_group"] = group_key
                all_itemsets.append(itemset)

        lines = [f"频繁项集查询 entity_ids={entity_ids}，共 {total} 条，返回 {len(all_itemsets)} 条：\n"]
        for i, itemset in enumerate(all_itemsets, 1):
            itemset_id = itemset.get("id", "")
            item_count = itemset.get("item_count", "")
            support = itemset.get("support", 0)
            abs_support = itemset.get("absolute_support", "")
            lines.append(f"{i}. 项集ID={itemset_id} | 项数={item_count} | support={support:.4f} | abs={abs_support}")
            for elem in itemset.get("items", []):
                dim = elem.get("dimension", "")
                path = elem.get("category_path", "")
                ename = elem.get("element_name") or ""
                label = f"{path}（{ename}）" if ename else path
                lines.append(f"   [{dim}] {label}")
            lines.append("")

        return ToolResult(
            title=f"频繁项集: entity_ids={entity_ids} → {total} 条",
            output="\n".join(lines),
        )
        return ToolResult(
            title="频繁项集查询失败",
            output=f"HTTP {e.response.status_code}: {e.response.text[:200]}",
        )
    except Exception as e:
        logger.exception("get_frequent_itemsets error")
        return ToolResult(title="频繁项集查询失败", output=f"错误: {e}")