Agent/examples/create/tool/topic_search.py

"""
选题检索工具 - 根据关键词在数据库中匹配已有帖子的选题

用于 Agent 执行时自主调取参考数据。
"""

import json
import os
from typing import Any, Dict, List, Optional

import httpx

from agent.tools import tool, ToolResult

# 选题检索 API 配置
TOPIC_SEARCH_BASE_URL = os.getenv("TOPIC_SEARCH_BASE_URL", "http://192.168.81.89:8000")
DEFAULT_TIMEOUT = 30.0


async def _call_search_api(keywords: List[str]) -> Optional[List[Dict[str, Any]]]:
    """调用选题检索 API，返回结果列表。"""
    url = f"{TOPIC_SEARCH_BASE_URL.rstrip('/')}/search"
    payload = {"keywords": keywords}

    try:
        async with httpx.AsyncClient(timeout=DEFAULT_TIMEOUT) as client:
            resp = await client.post(url, json=payload)
            resp.raise_for_status()
            data = resp.json()
    except httpx.HTTPStatusError as e:
        raise RuntimeError(f"API 请求失败: {e.response.status_code} - {e.response.text[:200]}")
    except Exception as e:
        raise RuntimeError(f"请求异常: {str(e)}")

    # 兼容多种响应格式
    if isinstance(data, list):
        return data[:5]
    if isinstance(data, dict):
        items = data.get("data") or data.get("results") or data.get("items") or []
        return list(items)[:5] if isinstance(items, (list, tuple)) else []
    return []


def _pick_first(results: List[Dict[str, Any]]) -> Dict[str, Any]:
    """从结果中取第一条。"""
    if not results:
        raise ValueError("无可用结果")
    return results[0]


@tool(
    description="根据关键词在数据库中检索已有帖子的选题，用于创作参考。最多返回5条，取第一条输出。",
    display={
        "zh": {
            "name": "爆款选题检索",
            "params": {
                "keywords": "关键词列表",
            },
        },
    },
)
async def topic_search(keywords: List[str]) -> ToolResult:
    """
    根据关键词检索数据库中已有帖子的选题，取第一条作为参考。

    Args:
        keywords: 关键词列表，如 ["中老年健康养生", "爆款", "知识科普"]

    Returns:
        ToolResult: 选题参考内容
    """
    if not keywords:
        return ToolResult(
            title="选题检索失败",
            output="",
            error="请提供至少一个关键词",
        )

    try:
        results = await _call_search_api(keywords)
    except RuntimeError as e:
        return ToolResult(
            title="选题检索失败",
            output="",
            error=str(e),
        )

    if not results:
        return ToolResult(
            title="选题检索",
            output=json.dumps({"message": "未找到匹配的选题", "keywords": keywords}, ensure_ascii=False, indent=2),
        )

    try:
        best = _pick_first(results)
    except ValueError:
        return ToolResult(
            title="选题检索",
            output=json.dumps({"message": "无可用结果", "keywords": keywords}, ensure_ascii=False, indent=2),
        )

    output = json.dumps(best, ensure_ascii=False, indent=2)
    return ToolResult(
        title="选题检索 - 参考数据",
        output=output,
        long_term_memory=f"检索到选题参考，关键词: {', '.join(keywords)}",
    )