"""HTTP RAG Service - 通过 HTTP 调用业务系统知识库 API 配置驱动，不直接依赖业务系统代码，通过 base_url + api_key 连接。 """ from __future__ import annotations import logging from typing import TYPE_CHECKING, TypeAlias import httpx from agentkit.memory.base import MetadataDict if TYPE_CHECKING: from agentkit.llm.gateway import LLMGateway logger = logging.getLogger(__name__) # 标准化检索结果：id/content/score/source/document_id/document_title/ # knowledge_base_id/metadata — 值为原始标量或嵌套 dict。 RAGSearchResult: TypeAlias = dict[str, object] # ingest() 写入的文档负载：title/content/source_type/metadata。 RAGIngestPayload: TypeAlias = dict[str, object] class HttpRAGService: """HTTP 客户端，调用业务系统的知识库检索 API 适配任意提供以下接口的知识库服务： - POST {base_url}/search → 语义检索 - POST {base_url}/ingest → 文档写入（可选） 典型配置（agentkit.yaml）:: memory: semantic: enabled: true base_url: "http://localhost:8000/api/knowledge" api_key: "${GEO_API_KEY}" knowledge_base_ids: - "industry-kb-id" - "enterprise-kb-id" timeout: 30 contextual_chunking: false """ def __init__( self, base_url: str, api_key: str | None = None, knowledge_base_ids: list[str] | None = None, timeout: int = 30, contextual_chunking: bool = False, llm_gateway: LLMGateway | None = None, ): """ Args: base_url: 知识库 API 基础地址，如 http://localhost:8000/api/knowledge api_key: 认证 API Key（放在 Authorization: Bearer 头） knowledge_base_ids: 默认检索的知识库 ID 列表 timeout: HTTP 请求超时秒数 """ self._base_url = base_url.rstrip("/") self._api_key = api_key self._knowledge_base_ids = knowledge_base_ids or [] self._timeout = timeout self._client: httpx.AsyncClient | None = None self._contextual_chunking = contextual_chunking self._llm_gateway = llm_gateway def _get_client(self) -> httpx.AsyncClient: """懒初始化 httpx 客户端""" if self._client is None or self._client.is_closed: headers: dict[str, str] = {"Content-Type": "application/json"} if self._api_key: headers["Authorization"] = f"Bearer {self._api_key}" self._client = httpx.AsyncClient( base_url=self._base_url, headers=headers, timeout=self._timeout, ) return self._client async def search( self, query: str, knowledge_base_ids: list[str] | None = None, top_k: int = 5, ) -> list[RAGSearchResult]: """语义检索知识库 Args: query: 检索查询 knowledge_base_ids: 知识库 ID 列表（默认使用配置值） top_k: 返回结果数量 Returns: 检索结果列表，每项包含 content/score/document_id 等字段 """ kb_ids = knowledge_base_ids or self._knowledge_base_ids payload = { "query": query, "knowledge_base_ids": kb_ids, "top_k": top_k, } client = self._get_client() try: resp = await client.post("/search", json=payload) resp.raise_for_status() data = resp.json() # 兼容两种响应格式： # 1. {"results": [...]} — GEO 标准 SearchResponse # 2. [...] — 直接返回列表 if isinstance(data, dict) and "results" in data: results = data["results"] elif isinstance(data, list): results = data else: logger.warning(f"Unexpected search response format: {type(data)}") return [] # 标准化为 SemanticMemory 期望的格式 normalized = [] for r in results: if isinstance(r, dict): normalized.append( { "id": r.get("chunk_id", r.get("id", "")), "content": r.get("content", ""), "score": float(r.get("score", 0.0)), "source": r.get("source", "rag"), "document_id": r.get("document_id", ""), "document_title": r.get("document_title", ""), "metadata": r.get("metadata", {}), } ) return normalized except httpx.HTTPStatusError as e: logger.error( f"RAG search HTTP error: {e.response.status_code} — {e.response.text[:200]}" ) return [] except httpx.RequestError as e: logger.error(f"RAG search request error: {e}") return [] except Exception as e: logger.error(f"RAG search unexpected error: {e}") return [] async def enhanced_search( self, query: str, knowledge_base_ids: list[str] | None = None, top_k: int = 5, use_rerank: bool = True, use_compression: bool = False, ) -> list[RAGSearchResult]: """增强语义检索知识库（支持 rerank 和 compression） 对每个知识库分别调用 /bases/{kb_id}/retrieve 接口， 合并结果后按 score 降序返回 top_k 条。 Args: query: 检索查询 knowledge_base_ids: 知识库 ID 列表（默认使用配置值） top_k: 返回结果数量 use_rerank: 是否启用 rerank 重排序 use_compression: 是否启用上下文压缩 Returns: 检索结果列表，每项包含 content/score/document_id 等字段 """ kb_ids = knowledge_base_ids or self._knowledge_base_ids if not kb_ids: return [] payload = { "query": query, "top_k": top_k, "use_rerank": use_rerank, "use_compression": use_compression, } client = self._get_client() all_results: list[RAGSearchResult] = [] for kb_id in kb_ids: try: resp = await client.post(f"/bases/{kb_id}/retrieve", json=payload) resp.raise_for_status() data = resp.json() # 兼容两种响应格式 if isinstance(data, dict) and "results" in data: results = data["results"] elif isinstance(data, list): results = data else: logger.warning(f"Unexpected enhanced_search response format: {type(data)}") continue # 标准化 for r in results: if isinstance(r, dict): all_results.append( { "id": r.get("chunk_id", r.get("id", "")), "content": r.get("content", ""), "score": float(r.get("score", 0.0)), "source": r.get("source", "rag"), "document_id": r.get("document_id", ""), "document_title": r.get("document_title", ""), "knowledge_base_id": kb_id, "metadata": r.get("metadata", {}), } ) except httpx.HTTPStatusError as e: if e.response.status_code == 404: # This KB doesn't support enhanced search — fall back to # standard search for THIS KB only, not all KBs. logger.info( f"Enhanced search not available for KB {kb_id}, using standard search" ) std_result = await self.search(query, knowledge_base_ids=[kb_id], top_k=top_k) all_results.extend(std_result) else: logger.error( f"RAG enhanced_search HTTP error for KB {kb_id}: " f"{e.response.status_code} — {e.response.text[:200]}" ) raise except httpx.RequestError as e: logger.error(f"RAG enhanced_search request error for KB {kb_id}: {e}") raise except Exception as e: logger.error(f"RAG enhanced_search unexpected error for KB {kb_id}: {e}") raise # 按 score 降序排序，返回 top_k all_results.sort(key=lambda x: x["score"], reverse=True) return all_results[:top_k] async def ingest( self, key: str, value: object, metadata: MetadataDict | None = None, ) -> dict[str, object] | None: """写入文档到知识库（可选操作） When contextual_chunking is enabled and llm_gateway is configured, the document content is enhanced with contextual prefixes before ingestion. Args: key: 文档标题或标识 value: 文档内容 metadata: 额外元数据 Returns: 写入结果，或 None 表示写入不可用 """ kb_ids = self._knowledge_base_ids if not kb_ids: logger.warning("HttpRAGService.ingest: no knowledge_base_ids configured") return None content = str(value) # Apply contextual chunking if enabled if self._contextual_chunking and self._llm_gateway: from agentkit.memory.contextual_retrieval import ContextualChunker chunker = ContextualChunker(llm_gateway=self._llm_gateway) # Simple chunking: split by paragraphs raw_chunks = [c.strip() for c in content.split("\n\n") if c.strip()] if raw_chunks: enhanced = await chunker.enhance_chunks( document=content, chunks=raw_chunks, metadata=metadata ) # Rejoin enhanced chunks content = "\n\n".join(chunk.enhanced_content for chunk in enhanced) payload = { "title": key, "content": content, "source_type": "text", "metadata": metadata or {}, } client = self._get_client() try: # 写入到第一个配置的知识库 kb_id = kb_ids[0] resp = await client.post(f"/bases/{kb_id}/documents", json=payload) resp.raise_for_status() return resp.json() except httpx.HTTPStatusError as e: logger.error(f"RAG ingest HTTP error: {e.response.status_code}") return None except Exception as e: logger.error(f"RAG ingest error: {e}") return None async def health_check(self) -> bool: """检查知识库服务是否可用""" client = self._get_client() try: resp = await client.get("/bases") return resp.status_code in (200, 401) # 401 = 服务在但需认证 except Exception: return False async def close(self) -> None: """关闭 HTTP 客户端""" if self._client and not self._client.is_closed: await self._client.aclose() self._client = None async def __aenter__(self) -> "HttpRAGService": return self async def __aexit__(self, *args: object) -> None: await self.close()