Agentic_Rag/app/tools.py

from __future__ import annotations

from time import perf_counter
from typing import Protocol

from app.schemas import Evidence, ToolResult
from app.sql_store import OrderRepository


class VectorStore(Protocol):
    def search(
        self,
        query: str,
        top_k: int = 5,
        policy_type: str = "",
    ) -> list[Evidence]: ...


class WebSearchProvider(Protocol):
    def search(self, query: str, max_results: int = 5) -> list[Evidence]: ...


class TavilyWebSearchProvider:
    """把 Tavily 搜索结果转换为项目统一的 Evidence。"""

    def __init__(self, api_key: str) -> None:
        self.api_key = api_key

    def search(self, query: str, max_results: int = 5) -> list[Evidence]:
        if not self.api_key:
            raise RuntimeError("网页搜索未配置 TAVILY_API_KEY")
        try:
            from tavily import TavilyClient
        except ImportError as exc:
            raise RuntimeError(
                "缺少网页搜索依赖，请执行 uv sync --extra test --extra web"
            ) from exc

        client = TavilyClient(api_key=self.api_key)
        response = client.search(
            query=query,
            search_depth="basic",
            topic="general",
            max_results=max_results,
            include_answer=False,
            include_raw_content=False,
        )

        # 只保留可引用片段和来源，不把搜索服务生成的 answer 当作事实答案。
        evidence: list[Evidence] = []
        for item in response.get("results", []):
            url = str(item.get("url", "")).strip()
            content = str(item.get("content", "")).strip()
            if not url or not content:
                continue
            title = str(item.get("title", "网页结果")).strip()
            evidence.append(
                Evidence(
                    source_type="web",
                    source=url,
                    content=f"{title}\n{content}",
                    score=float(item.get("score", 0.0)),
                    metadata={
                        "title": title,
                        "url": url,
                        "published_date": item.get("published_date"),
                    },
                )
            )
        return evidence


class VectorSearchTool:
    name = "milvus_search"

    def __init__(self, store: VectorStore) -> None:
        self.store = store

    def invoke(self, query: str, policy_type: str = "", top_k: int = 5) -> ToolResult:
        started = perf_counter()
        try:
            # Adapter 负责把数据源异常统一转换为 ToolResult，Graph 不捕获 SDK 异常。
            evidence = self.store.search(query, top_k=top_k, policy_type=policy_type)
            return ToolResult(
                status="success" if evidence else "empty",
                tool=self.name,
                evidence=evidence,
                latency_ms=int((perf_counter() - started) * 1000),
            )
        except Exception as exc:
            return ToolResult(
                status="error",
                tool=self.name,
                latency_ms=int((perf_counter() - started) * 1000),
                retryable=True,
                error_code="MILVUS_SEARCH_FAILED",
                error_message=str(exc),
            )


class SQLQueryTool:
    name = "sql_query"

    def __init__(self, repository: OrderRepository) -> None:
        self.repository = repository

    def invoke(
        self,
        user_id: str,
        days: int = 30,
        product_keyword: str = "",
    ) -> ToolResult:
        started = perf_counter()
        try:
            data, evidence = self.repository.summarize_recent(
                user_id=user_id,
                days=days,
                product_keyword=product_keyword,
            )
            return ToolResult(
                status="success",
                tool=self.name,
                data=[data],
                evidence=[evidence],
                latency_ms=int((perf_counter() - started) * 1000),
            )
        except (ValueError, RuntimeError) as exc:
            return ToolResult(
                status="error",
                tool=self.name,
                latency_ms=int((perf_counter() - started) * 1000),
                retryable=False,
                error_code="SQL_QUERY_REJECTED",
                error_message=str(exc),
            )


class WebSearchTool:
    name = "web_search"

    def __init__(self, provider: WebSearchProvider) -> None:
        self.provider = provider

    def invoke(self, query: str, max_results: int = 5) -> ToolResult:
        started = perf_counter()
        try:
            evidence = self.provider.search(query, max_results=max_results)
            return ToolResult(
                status="success" if evidence else "empty",
                tool=self.name,
                data=[item.metadata for item in evidence],
                evidence=evidence,
                latency_ms=int((perf_counter() - started) * 1000),
            )
        except RuntimeError as exc:
            return ToolResult(
                status="error",
                tool=self.name,
                latency_ms=int((perf_counter() - started) * 1000),
                retryable=False,
                error_code="WEB_SEARCH_NOT_CONFIGURED",
                error_message=str(exc),
            )
        except Exception as exc:
            return ToolResult(
                status="error",
                tool=self.name,
                latency_ms=int((perf_counter() - started) * 1000),
                retryable=True,
                error_code="WEB_SEARCH_FAILED",
                error_message=str(exc),
            )