Agent/examples/content_tree_analyst/tools/content_tree.py

"""
内容树 API 工具

封装内容树搜索接口：
1. search_content_tree - 关键词搜索分类和元素
2. get_category_tree - 获取指定分类的完整路径和子树
"""

import logging
from typing import Optional

import httpx

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

logger = logging.getLogger(__name__)

BASE_URL = "http://8.147.104.190:8001"


@tool(description="在内容树中搜索分类（category）和元素（element），支持获取祖先路径和子孙节点")
async def search_content_tree(
    q: str,
    source_type: str,
    entity_type: str = "all",
    top_k: int = 20,
    use_description: bool = False,
    include_ancestors: bool = False,
    descendant_depth: int = 0,
) -> ToolResult:
    """
    关键词搜索内容树中的分类和元素。

    Args:
        q: 搜索关键词
        source_type: 维度，必须是 "实质" / "形式" / "意图" 之一
        entity_type: 搜索对象类型，"category" / "element" / "all"（默认）
        top_k: 返回结果数量，1-100（默认20）
        use_description: 是否同时搜索描述字段（默认仅搜索名称）
        include_ancestors: 是否返回祖先路径
        descendant_depth: 返回子孙节点深度，0=不返回，1=直接子节点，2=子+孙...
    """
    params = {
        "q": q,
        "source_type": source_type,
        "entity_type": entity_type,
        "top_k": top_k,
        "use_description": str(use_description).lower(),
        "include_ancestors": str(include_ancestors).lower(),
        "descendant_depth": descendant_depth,
    }

    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            resp = await client.get(f"{BASE_URL}/api/agent/search", params=params)
            resp.raise_for_status()
            data = resp.json()

        count = data.get("count", 0)
        results = data.get("results", [])

        # 格式化输出
        lines = [f"搜索「{q}」（{source_type}维度）共找到 {count} 条结果：\n"]
        for r in results:
            etype = r.get("entity_type", "")
            name = r.get("name", "")
            score = r.get("score", 0)
            if etype == "category":
                sid = r.get("entity_id", "")
                path = r.get("path", "")
                desc = r.get("description", "")
                lines.append(f"[分类] entity_id={sid} | {path} | score={score:.2f}")
                if desc:
                    lines.append(f"  描述: {desc}")
                ancestors = r.get("ancestors", [])
                if ancestors:
                    anc_names = " > ".join(a["name"] for a in ancestors)
                    lines.append(f"  祖先: {anc_names}")
                descendants = r.get("descendants", [])
                if descendants:
                    desc_names = ", ".join(d["name"] for d in descendants[:10])
                    lines.append(f"  子孙({len(descendants)}): {desc_names}")
            else:
                eid = r.get("entity_id", "")
                belong = r.get("belong_category_entity_id", "")
                occ = r.get("occurrence_count", 0)
                lines.append(f"[元素] entity_id={eid} | {name} | belong_category={belong} | 出现次数={occ} | score={score:.2f}")
                edesc = r.get("description", "")
                if edesc:
                    lines.append(f"  描述: {edesc}")
            lines.append("")

        return ToolResult(
            title=f"内容树搜索: {q} ({source_type}) → {count} 条",
            output="\n".join(lines),
        )

    except httpx.HTTPError as e:
        return ToolResult(title="内容树搜索失败", output=f"HTTP 错误: {e}")
    except Exception as e:
        logger.exception("search_content_tree error")
        return ToolResult(title="内容树搜索失败", output=f"错误: {e}")


@tool(description="获取指定分类节点的完整路径、祖先和子孙结构（通过 entity_id 精确查询）")
async def get_category_tree(
    entity_id: int,
    source_type: str,
    include_ancestors: bool = True,
    descendant_depth: int = -1,
) -> ToolResult:
    """
    获取指定分类的完整路径和子树结构。

    Args:
        entity_id: 分类的 entity_id
        source_type: 维度，"实质" / "形式" / "意图"
        include_ancestors: 是否返回祖先路径（默认 True）
        descendant_depth: 子孙深度，-1=全部，0=仅当前，1=子节点，2=子+孙...
    """
    params = {
        "source_type": source_type,
        "include_ancestors": str(include_ancestors).lower(),
        "descendant_depth": descendant_depth,
    }

    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            resp = await client.get(f"{BASE_URL}/api/agent/search/category/{entity_id}", params=params)
            resp.raise_for_status()
            data = resp.json()

        current = data.get("current", {})
        ancestors = data.get("ancestors", [])
        descendants = data.get("descendants", [])

        lines = []
        lines.append(f"分类节点: {current.get('name', '')} (entity_id={entity_id})")
        lines.append(f"路径: {current.get('path', '')}")
        if current.get("description"):
            lines.append(f"描述: {current['description']}")
        lines.append("")

        if ancestors:
            lines.append("祖先路径:")
            for a in ancestors:
                lines.append(f"  L{a.get('level', '?')} {a.get('name', '')} (entity_id={a.get('entity_id', '')})")
            lines.append("")

        if descendants:
            lines.append(f"子孙节点 ({len(descendants)} 个):")
            for d in descendants:
                indent = "  " * d.get("depth_from_parent", 1)
                leaf_mark = " [叶]" if d.get("is_leaf") else ""
                lines.append(f"{indent}L{d.get('level', '?')} {d.get('name', '')} (entity_id={d.get('entity_id', '')}){leaf_mark}")

        return ToolResult(
            title=f"分类树: {current.get('name', entity_id)} (entity_id={entity_id})",
            output="\n".join(lines),
        )

    except httpx.HTTPError as e:
        return ToolResult(title="获取分类树失败", output=f"HTTP 错误: {e}")
    except Exception as e:
        logger.exception("get_category_tree error")
        return ToolResult(title="获取分类树失败", output=f"错误: {e}")


@tool(description="获取指定分类下的所有元素，支持分页、排序和筛选")
async def get_category_elements(
    category_id: int,
    source_type: str,
    page_size: int = 50,
    sort_by: str = "occurrence_count",
    order: str = "desc",
    min_occurrence: Optional[int] = None,
) -> ToolResult:
    """
    获取指定分类下的所有元素。

    Args:
        category_id: 分类的 entity_id
        source_type: 维度，"实质" / "形式" / "意图"
        page_size: 每页数量，1-200（默认50）
        sort_by: 排序字段，"name" / "id" / "occurrence_count"（默认）
        order: 排序方向，"asc" / "desc"（默认）
        min_occurrence: 最小出现次数，用于过滤低频元素（可选）
    """
    params = {
        "source_type": source_type,
        "category_entity_id": category_id,
        "page_size": page_size,
        "sort_by": sort_by,
        "order": order,
    }
    if min_occurrence is not None:
        params["min_occurrence"] = min_occurrence

    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            resp = await client.get(f"{BASE_URL}/api/agent/search/elements", params=params)
            resp.raise_for_status()
            data = resp.json()

        total = data.get("total", 0)
        results = data.get("results", [])

        lines = [f"分类 {category_id} 下的元素（{source_type}维度）共 {total} 个，返回 {len(results)} 个：\n"]
        for r in results:
            eid = r.get("id", "")
            name = r.get("name", "")
            occ = r.get("occurrence_count", 0)
            desc = r.get("description", "")
            category = r.get("category", {})
            cat_path = category.get("path", "")

            lines.append(f"[元素] entity_id={eid} | {name} | 出现次数={occ}")
            if desc:
                lines.append(f"  描述: {desc}")
            if cat_path:
                lines.append(f"  所属分类: {cat_path}")
            lines.append("")

        return ToolResult(
            title=f"分类元素: category_id={category_id} → {total} 个",
            output="\n".join(lines),
        )

    except httpx.HTTPError as e:
        return ToolResult(title="获取分类元素失败", output=f"HTTP 错误: {e}")
    except Exception as e:
        logger.exception("get_category_elements error")
        return ToolResult(title="获取分类元素失败", output=f"错误: {e}")