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), )