Agent/examples/requirement_extract/requirement.py

"""
需求管理工具 - KnowHub Requirement API 封装

提供需求的向量检索和列表查询能力，
让 Agent 在提取新需求前可以先了解已有的需求库。
"""

import os
import json
import logging
import httpx
from typing import Optional
from agent.tools import tool, ToolResult, ToolContext

logger = logging.getLogger(__name__)

KNOWHUB_API = os.getenv("KNOWHUB_API", "http://localhost:8000").rstrip("/")


@tool(
    hidden_params=["context"],
    display={
        "zh": {
            "name": "需求检索",
            "params": {
                "query": "搜索关键词",
                "top_k": "返回数量",
            }
        },
        "en": {
            "name": "Search Requirements",
            "params": {
                "query": "Search query",
                "top_k": "Number of results",
            }
        }
    }
)
async def requirement_search(
    query: str,
    top_k: int = 20,
    context: Optional[ToolContext] = None,
) -> ToolResult:
    """向量检索需求库，根据语义相似度返回最相关的需求。用于在提取新需求前了解已有需求，避免重复。"""
    try:
        params = {"q": query, "top_k": top_k}
        async with httpx.AsyncClient(timeout=30.0) as client:
            res = await client.get(f"{KNOWHUB_API}/api/requirement/search", params=params)
            res.raise_for_status()
            data = res.json()

        results = data.get("results", data) if isinstance(data, dict) else data
        items = results if isinstance(results, list) else []
        count = len(items)

        if count == 0:
            return ToolResult(
                title="需求检索完成",
                output="未找到相关需求。",
                long_term_memory=f"需求检索: 未找到与 '{query[:50]}' 相关的需求",
            )

        output = json.dumps(data, ensure_ascii=False, indent=2)
        return ToolResult(
            title=f"需求检索完成 -- 找到 {count} 条",
            output=output,
            long_term_memory=f"需求检索: 找到 {count} 条与 '{query[:50]}' 相关的需求",
            metadata={"count": count, "items": items},
        )

    except httpx.ConnectError:
        msg = f"无法连接到 KnowHub 服务 ({KNOWHUB_API})，请确认服务是否启动。"
        return ToolResult(title="需求检索失败", output=msg, error=msg)
    except Exception as e:
        return ToolResult(title="需求检索失败", output=str(e), error=str(e))


@tool(
    hidden_params=["context"],
    display={
        "zh": {
            "name": "需求列表",
            "params": {
                "limit": "返回数量",
                "offset": "偏移量",
            }
        },
        "en": {
            "name": "List Requirements",
            "params": {
                "limit": "Number of results",
                "offset": "Offset",
            }
        }
    }
)
async def requirement_list(
    limit: int = 20,
    offset: int = 0,
    context: Optional[ToolContext] = None,
) -> ToolResult:
    """列出需求库中的已有需求，用于浏览和了解需求全貌。"""
    try:
        params = {"limit": limit, "offset": offset}
        async with httpx.AsyncClient(timeout=30.0) as client:
            res = await client.get(f"{KNOWHUB_API}/api/requirement/list", params=params)
            res.raise_for_status()
            data = res.json()

        results = data.get("results", data) if isinstance(data, dict) else data
        items = results if isinstance(results, list) else []
        count = len(items)

        if count == 0:
            return ToolResult(
                title="需求列表为空",
                output="需求库中暂无需求。",
                long_term_memory="需求列表: 需求库为空",
            )

        output = json.dumps(data, ensure_ascii=False, indent=2)
        return ToolResult(
            title=f"需求列表 -- 共 {count} 条",
            output=output,
            long_term_memory=f"需求列表: 获取了 {count} 条需求 (offset={offset})",
            metadata={"count": count, "items": items},
        )

    except httpx.ConnectError:
        msg = f"无法连接到 KnowHub 服务 ({KNOWHUB_API})，请确认服务是否启动。"
        return ToolResult(title="需求列表获取失败", output=msg, error=msg)
    except Exception as e:
        return ToolResult(title="需求列表获取失败", output=str(e), error=str(e))