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/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("stable_id", "")
                path = r.get("path", "")
                desc = r.get("description", "")
                lines.append(f"[分类] stable_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_stable_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="获取指定分类节点的完整路径、祖先和子孙结构（通过 stable_id 精确查询）")
async def get_category_tree(
    stable_id: int,
    source_type: str,
    include_ancestors: bool = True,
    descendant_depth: int = -1,
) -> ToolResult:
    """
    获取指定分类的完整路径和子树结构。

    Args:
        stable_id: 分类的 stable_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/search/category/{stable_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', '')} (stable_id={stable_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', '')} (stable_id={a.get('stable_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', '')} (stable_id={d.get('stable_id', '')}){leaf_mark}")

        return ToolResult(
            title=f"分类树: {current.get('name', stable_id)} (stable_id={stable_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}")