diff --git a/src/agentkit/memory/__init__.py b/src/agentkit/memory/__init__.py index 1d1ec20..e26e031 100644 --- a/src/agentkit/memory/__init__.py +++ b/src/agentkit/memory/__init__.py @@ -6,6 +6,7 @@ from agentkit.memory.episodic import EpisodicMemory from agentkit.memory.semantic import SemanticMemory from agentkit.memory.http_rag import HttpRAGService from agentkit.memory.retriever import MemoryRetriever +from agentkit.memory.multi_source_retriever import MultiSourceRetriever from agentkit.memory.query_transformer import ( QueryTransformerBase, LLMQueryTransformer, @@ -23,6 +24,7 @@ __all__ = [ "SemanticMemory", "HttpRAGService", "MemoryRetriever", + "MultiSourceRetriever", "QueryTransformerBase", "LLMQueryTransformer", "RuleQueryTransformer", diff --git a/src/agentkit/memory/adapters/__init__.py b/src/agentkit/memory/adapters/__init__.py new file mode 100644 index 0000000..d55c42e --- /dev/null +++ b/src/agentkit/memory/adapters/__init__.py @@ -0,0 +1,13 @@ +"""知识库适配器包""" + +from agentkit.memory.adapters.base import KBAdapter +from agentkit.memory.adapters.feishu import FeishuKBAdapter +from agentkit.memory.adapters.confluence import ConfluenceAdapter +from agentkit.memory.adapters.generic_http import GenericHTTPAdapter + +__all__ = [ + "KBAdapter", + "FeishuKBAdapter", + "ConfluenceAdapter", + "GenericHTTPAdapter", +] diff --git a/src/agentkit/memory/adapters/base.py b/src/agentkit/memory/adapters/base.py new file mode 100644 index 0000000..eddedec --- /dev/null +++ b/src/agentkit/memory/adapters/base.py @@ -0,0 +1,160 @@ +"""KBAdapter 抽象基类 - 知识库适配器的基础实现 + +实现 KnowledgeBase 协议,并提供通用扩展方法: +- search(): query() 的别名 +- get_document(): 按 ID 获取单个文档 +- authenticate(): 认证验证 +""" + +from __future__ import annotations + +import logging +from abc import ABC, abstractmethod +from typing import Any + +import httpx + +from agentkit.memory.knowledge_base import Document, QueryResult, SourceInfo + +logger = logging.getLogger(__name__) + + +class KBAdapter(ABC): + """知识库适配器抽象基类 + + 实现 KnowledgeBase 协议的所有方法,并提供额外便利方法。 + 子类需实现 _make_client() 和具体的 HTTP 调用逻辑。 + """ + + def __init__( + self, + source_id: str, + source_name: str, + source_type: str, + timeout: int = 30, + ): + self._source_id = source_id + self._source_name = source_name + self._source_type = source_type + self._timeout = timeout + self._client: httpx.AsyncClient | None = None + self._authenticated = False + + @property + def source_id(self) -> str: + return self._source_id + + @property + def source_name(self) -> str: + return self._source_name + + @property + def source_type(self) -> str: + return self._source_type + + def _get_client(self) -> httpx.AsyncClient: + """获取或创建 HTTP 客户端""" + if self._client is None or self._client.is_closed: + self._client = self._make_client() + return self._client + + @abstractmethod + def _make_client(self) -> httpx.AsyncClient: + """创建 HTTP 客户端(子类实现,配置 base_url、headers、auth 等)""" + ... + + # ------------------------------------------------------------------ + # KnowledgeBase 协议方法 + # ------------------------------------------------------------------ + + async def ingest(self, documents: list[Document]) -> list[str]: + """写入文档到知识库,返回文档 ID 列表 + + 默认实现:逐个调用 _ingest_one()。 + 子类可覆盖以实现批量写入。 + """ + ids: list[str] = [] + for doc in documents: + doc_id = await self._ingest_one(doc) + if doc_id: + ids.append(doc_id) + return ids + + async def _ingest_one(self, document: Document) -> str | None: + """写入单个文档(子类可覆盖)""" + logger.warning( + f"{self.__class__.__name__} does not support ingest; " + f"document '{document.doc_id}' skipped" + ) + return None + + async def query(self, text: str, top_k: int = 5) -> list[QueryResult]: + """语义检索知识库(委托给 search)""" + return await self.search(text, top_k=top_k) + + async def delete_by_id(self, id: str) -> bool: + """按文档 ID 删除(子类可覆盖)""" + logger.warning( + f"{self.__class__.__name__} does not support delete_by_id; " + f"id '{id}' skipped" + ) + return False + + async def list_sources(self) -> list[SourceInfo]: + """列出可用信息源""" + return [ + SourceInfo( + source_id=self._source_id, + source_name=self._source_name, + source_type=self._source_type, + ) + ] + + @abstractmethod + async def health_check(self) -> bool: + """检查知识库连接状态(子类实现)""" + ... + + # ------------------------------------------------------------------ + # 扩展方法 + # ------------------------------------------------------------------ + + @abstractmethod + async def search(self, query: str, top_k: int = 5) -> list[QueryResult]: + """语义检索知识库(子类实现)""" + ... + + async def get_document(self, doc_id: str) -> Document | None: + """按 ID 获取单个文档(子类可覆盖)""" + logger.warning( + f"{self.__class__.__name__} does not support get_document; " + f"doc_id '{doc_id}' not found" + ) + return None + + async def authenticate(self) -> bool: + """认证验证(子类可覆盖) + + 默认实现调用 health_check()。 + """ + try: + self._authenticated = await self.health_check() + except Exception: + self._authenticated = False + return self._authenticated + + # ------------------------------------------------------------------ + # 生命周期管理 + # ------------------------------------------------------------------ + + 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) -> KBAdapter: + return self + + async def __aexit__(self, *args: Any) -> None: + await self.close() diff --git a/src/agentkit/memory/adapters/confluence.py b/src/agentkit/memory/adapters/confluence.py new file mode 100644 index 0000000..0bed8da --- /dev/null +++ b/src/agentkit/memory/adapters/confluence.py @@ -0,0 +1,210 @@ +"""ConfluenceAdapter - Confluence 知识库适配器 + +对接 Confluence REST API,实现 KnowledgeBase 协议。 +通过 base_url + username + api_token 认证。 +""" + +from __future__ import annotations + +import logging +from typing import Any + +import httpx + +from agentkit.memory.adapters.base import KBAdapter +from agentkit.memory.knowledge_base import Document, QueryResult, SourceInfo + +logger = logging.getLogger(__name__) + + +class ConfluenceAdapter(KBAdapter): + """Confluence 知识库适配器 + + 对接 Confluence REST API,支持: + - CQL 搜索 + - 获取页面内容 + - 列出空间 + + 典型配置:: + + adapter = ConfluenceAdapter( + base_url="https://your-domain.atlassian.net/wiki", + username="user@example.com", + api_token="xxx", + ) + """ + + def __init__( + self, + base_url: str, + username: str, + api_token: str, + space_keys: list[str] | None = None, + timeout: int = 30, + ): + super().__init__( + source_id=f"confluence-{username.split('@')[0]}", + source_name="Confluence", + source_type="confluence", + timeout=timeout, + ) + self._base_url = base_url.rstrip("/") + self._username = username + self._api_token = api_token + self._space_keys = space_keys or [] + + def _make_client(self) -> httpx.AsyncClient: + """创建 Confluence API HTTP 客户端""" + import base64 + + credentials = base64.b64encode( + f"{self._username}:{self._api_token}".encode() + ).decode() + return httpx.AsyncClient( + base_url=self._base_url, + headers={ + "Content-Type": "application/json", + "Authorization": f"Basic {credentials}", + }, + timeout=self._timeout, + ) + + async def authenticate(self) -> bool: + """Confluence 认证验证""" + client = self._get_client() + try: + resp = await client.get("/rest/api/user/current") + self._authenticated = resp.status_code == 200 + return self._authenticated + except Exception as e: + logger.error(f"Confluence auth error: {e}") + self._authenticated = False + return False + + async def search(self, query: str, top_k: int = 5) -> list[QueryResult]: + """搜索 Confluence 页面 + + 使用 CQL (Confluence Query Language) 进行搜索。 + """ + client = self._get_client() + try: + cql = f'text ~ "{query}"' + if self._space_keys: + space_filter = " OR ".join( + f'space = "{key}"' for key in self._space_keys + ) + cql = f'{cql} AND ({space_filter})' + + resp = await client.get( + "/rest/api/content/search", + params={"cql": cql, "limit": top_k, "expand": "body.storage"}, + ) + resp.raise_for_status() + data = resp.json() + + results: list[QueryResult] = [] + for page in data.get("results", []): + body = page.get("body", {}).get("storage", {}).get("value", "") + # Strip HTML tags for plain text content + import re + content = re.sub(r"<[^>]+>", "", body) if body else page.get("title", "") + + results.append( + QueryResult( + content=content[:2000], # Limit content length + source_id=self._source_id, + source_name=self._source_name, + score=1.0, # Confluence doesn't return relevance score + metadata={ + "space_key": page.get("space", {}).get("key", ""), + "type": page.get("type", ""), + "status": page.get("status", ""), + }, + doc_id=page.get("id", ""), + title=page.get("title", ""), + ) + ) + return results[:top_k] + + except httpx.HTTPStatusError as e: + logger.error( + f"Confluence search HTTP error: {e.response.status_code} — " + f"{e.response.text[:200]}" + ) + return [] + except Exception as e: + logger.error(f"Confluence search error: {e}") + return [] + + async def get_document(self, doc_id: str) -> Document | None: + """获取 Confluence 页面内容""" + client = self._get_client() + try: + resp = await client.get( + f"/rest/api/content/{doc_id}", + params={"expand": "body.storage,space"}, + ) + resp.raise_for_status() + page = resp.json() + + body = page.get("body", {}).get("storage", {}).get("value", "") + import re + content = re.sub(r"<[^>]+>", "", body) if body else "" + + return Document( + doc_id=str(page.get("id", doc_id)), + content=content, + title=page.get("title", ""), + source_id=self._source_id, + metadata={ + "space_key": page.get("space", {}).get("key", ""), + "type": page.get("type", ""), + "version": page.get("version", {}).get("number", 0), + }, + ) + except Exception as e: + logger.error(f"Confluence get_document error: {e}") + return None + + async def list_sources(self) -> list[SourceInfo]: + """列出 Confluence 空间""" + client = self._get_client() + try: + resp = await client.get("/rest/api/space", params={"limit": 50}) + resp.raise_for_status() + data = resp.json() + + sources: list[SourceInfo] = [] + for space in data.get("results", []): + sources.append( + SourceInfo( + source_id=f"confluence-space-{space.get('key', '')}", + source_name=space.get("name", ""), + source_type="confluence", + ) + ) + return sources if sources else [ + SourceInfo( + source_id=self._source_id, + source_name=self._source_name, + source_type=self._source_type, + ) + ] + except Exception as e: + logger.error(f"Confluence list_sources error: {e}") + return [ + SourceInfo( + source_id=self._source_id, + source_name=self._source_name, + source_type=self._source_type, + ) + ] + + async def health_check(self) -> bool: + """检查 Confluence API 连接状态""" + client = self._get_client() + try: + resp = await client.get("/rest/api/space", params={"limit": 1}) + return resp.status_code == 200 + except Exception: + return False diff --git a/src/agentkit/memory/adapters/feishu.py b/src/agentkit/memory/adapters/feishu.py new file mode 100644 index 0000000..6672618 --- /dev/null +++ b/src/agentkit/memory/adapters/feishu.py @@ -0,0 +1,249 @@ +"""FeishuKBAdapter - 飞书知识库适配器 + +对接飞书知识库 API,实现 KnowledgeBase 协议。 +通过 app_id + app_secret 认证,调用飞书开放平台 API 检索知识库内容。 +""" + +from __future__ import annotations + +import logging +from typing import Any + +import httpx + +from agentkit.memory.adapters.base import KBAdapter +from agentkit.memory.knowledge_base import Document, QueryResult, SourceInfo + +logger = logging.getLogger(__name__) + + +class FeishuKBAdapter(KBAdapter): + """飞书知识库适配器 + + 对接飞书开放平台知识库 API,支持: + - 搜索知识库节点 + - 获取知识空间列表 + - 获取文档内容 + + 典型配置:: + + adapter = FeishuKBAdapter( + app_id="cli_xxx", + app_secret="xxx", + base_url="https://open.feishu.cn/open-apis", + ) + """ + + def __init__( + self, + app_id: str, + app_secret: str, + base_url: str = "https://open.feishu.cn/open-apis", + space_ids: list[str] | None = None, + timeout: int = 30, + ): + super().__init__( + source_id=f"feishu-{app_id[:8]}", + source_name="飞书知识库", + source_type="feishu", + timeout=timeout, + ) + self._app_id = app_id + self._app_secret = app_secret + self._base_url = base_url.rstrip("/") + self._space_ids = space_ids or [] + self._access_token: str | None = None + + def _make_client(self) -> httpx.AsyncClient: + """创建飞书 API HTTP 客户端""" + headers: dict[str, str] = {"Content-Type": "application/json"} + if self._access_token: + headers["Authorization"] = f"Bearer {self._access_token}" + return httpx.AsyncClient( + base_url=self._base_url, + headers=headers, + timeout=self._timeout, + ) + + async def _get_access_token(self) -> str | None: + """获取飞书 tenant_access_token""" + if self._access_token: + return self._access_token + + client = self._get_client() + try: + resp = await client.post( + "/auth/v3/tenant_access_token/internal", + json={ + "app_id": self._app_id, + "app_secret": self._app_secret, + }, + ) + resp.raise_for_status() + data = resp.json() + if data.get("code") == 0: + self._access_token = data.get("tenant_access_token") + # 重建客户端以携带 token + await self.close() + return self._access_token + else: + logger.error( + f"Feishu auth failed: code={data.get('code')}, " + f"msg={data.get('msg')}" + ) + return None + except Exception as e: + logger.error(f"Feishu auth error: {e}") + return None + + async def authenticate(self) -> bool: + """飞书认证""" + token = await self._get_access_token() + self._authenticated = token is not None + return self._authenticated + + async def search(self, query: str, top_k: int = 5) -> list[QueryResult]: + """搜索飞书知识库 + + 调用飞书搜索 API 检索知识库节点内容。 + """ + token = await self._get_access_token() + if not token: + logger.error("FeishuKBAdapter.search: not authenticated") + return [] + + client = self._get_client() + try: + payload: dict[str, Any] = { + "search_key": query, + "page_size": top_k, + } + if self._space_ids: + payload["wiki_space_ids"] = self._space_ids + + resp = await client.post( + "/search/v2/wiki", + json=payload, + ) + resp.raise_for_status() + data = resp.json() + + if data.get("code") != 0: + logger.error( + f"Feishu search failed: code={data.get('code')}, " + f"msg={data.get('msg')}" + ) + return [] + + results: list[QueryResult] = [] + for item in data.get("data", {}).get("items", []): + results.append( + QueryResult( + content=item.get("content", ""), + source_id=self._source_id, + source_name=self._source_name, + score=float(item.get("score", 0.0)), + metadata={ + "wiki_token": item.get("wiki_token", ""), + "space_id": item.get("space_id", ""), + }, + doc_id=item.get("wiki_token", ""), + title=item.get("title", ""), + ) + ) + return results[:top_k] + + except httpx.HTTPStatusError as e: + logger.error( + f"Feishu search HTTP error: {e.response.status_code} — " + f"{e.response.text[:200]}" + ) + return [] + except Exception as e: + logger.error(f"Feishu search error: {e}") + return [] + + async def get_document(self, doc_id: str) -> Document | None: + """获取飞书知识库文档内容""" + token = await self._get_access_token() + if not token: + return None + + client = self._get_client() + try: + resp = await client.get( + f"/wiki/v2/spaces/get_node", + params={"token": doc_id}, + ) + resp.raise_for_status() + data = resp.json() + + if data.get("code") != 0: + return None + + node = data.get("data", {}).get("node", {}) + return Document( + doc_id=doc_id, + content=node.get("content", ""), + title=node.get("title", ""), + source_id=self._source_id, + metadata={ + "space_id": node.get("space_id", ""), + "obj_type": node.get("obj_type", ""), + }, + ) + except Exception as e: + logger.error(f"Feishu get_document error: {e}") + return None + + async def list_sources(self) -> list[SourceInfo]: + """列出飞书知识空间""" + token = await self._get_access_token() + if not token: + return [ + SourceInfo( + source_id=self._source_id, + source_name=self._source_name, + source_type=self._source_type, + ) + ] + + client = self._get_client() + try: + resp = await client.get("/wiki/v2/spaces", params={"page_size": 50}) + resp.raise_for_status() + data = resp.json() + + sources: list[SourceInfo] = [] + for space in data.get("data", {}).get("items", []): + sources.append( + SourceInfo( + source_id=f"feishu-space-{space.get('space_id', '')}", + source_name=space.get("name", ""), + source_type="feishu", + ) + ) + return sources if sources else [ + SourceInfo( + source_id=self._source_id, + source_name=self._source_name, + source_type=self._source_type, + ) + ] + except Exception as e: + logger.error(f"Feishu list_sources error: {e}") + return [ + SourceInfo( + source_id=self._source_id, + source_name=self._source_name, + source_type=self._source_type, + ) + ] + + async def health_check(self) -> bool: + """检查飞书 API 连接状态""" + try: + token = await self._get_access_token() + return token is not None + except Exception: + return False diff --git a/src/agentkit/memory/adapters/generic_http.py b/src/agentkit/memory/adapters/generic_http.py new file mode 100644 index 0000000..7de5826 --- /dev/null +++ b/src/agentkit/memory/adapters/generic_http.py @@ -0,0 +1,289 @@ +"""GenericHTTPAdapter - 通用 HTTP 知识库适配器 + +配置 API endpoint + auth 即可对接任意 HTTP 知识库服务。 +实现 KnowledgeBase 协议,提供统一的检索接口。 +""" + +from __future__ import annotations + +import logging +from typing import Any + +import httpx + +from agentkit.memory.adapters.base import KBAdapter +from agentkit.memory.knowledge_base import Document, QueryResult, SourceInfo + +logger = logging.getLogger(__name__) + + +class GenericHTTPAdapter(KBAdapter): + """通用 HTTP 知识库适配器 + + 通过配置 API endpoint 和认证信息,对接任意提供 HTTP API 的知识库服务。 + + 期望的 API 接口: + - POST {endpoint_url}/search → 语义检索 + - POST {endpoint_url}/ingest → 文档写入(可选) + - GET {endpoint_url}/sources → 列出信息源(可选) + - GET {endpoint_url}/health → 健康检查(可选) + + 典型配置:: + + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/knowledge", + auth_config={"type": "bearer", "token": "sk-xxx"}, + headers={"X-Custom-Header": "value"}, + ) + """ + + def __init__( + self, + endpoint_url: str, + auth_config: dict[str, str] | None = None, + headers: dict[str, str] | None = None, + source_id: str | None = None, + source_name: str = "HTTP Knowledge Base", + timeout: int = 30, + ): + super().__init__( + source_id=source_id or f"http-{endpoint_url.rstrip('/').split('/')[-1]}", + source_name=source_name, + source_type="generic_http", + timeout=timeout, + ) + self._endpoint_url = endpoint_url.rstrip("/") + self._auth_config = auth_config or {} + self._extra_headers = headers or {} + + def _make_client(self) -> httpx.AsyncClient: + """创建通用 HTTP 客户端""" + headers: dict[str, str] = { + "Content-Type": "application/json", + **self._extra_headers, + } + + # 配置认证 + auth_type = self._auth_config.get("type", "") + if auth_type == "bearer": + token = self._auth_config.get("token", "") + if token: + headers["Authorization"] = f"Bearer {token}" + elif auth_type == "basic": + import base64 + username = self._auth_config.get("username", "") + password = self._auth_config.get("password", "") + if username and password: + credentials = base64.b64encode( + f"{username}:{password}".encode() + ).decode() + headers["Authorization"] = f"Basic {credentials}" + elif auth_type == "api_key": + key_name = self._auth_config.get("header_name", "X-API-Key") + key_value = self._auth_config.get("api_key", "") + if key_value: + headers[key_name] = key_value + + return httpx.AsyncClient( + base_url=self._endpoint_url, + headers=headers, + timeout=self._timeout, + ) + + async def search(self, query: str, top_k: int = 5) -> list[QueryResult]: + """搜索 HTTP 知识库 + + POST {endpoint_url}/search + Body: {"query": ..., "top_k": ...} + """ + client = self._get_client() + try: + payload = {"query": query, "top_k": top_k} + resp = await client.post("/search", json=payload) + resp.raise_for_status() + data = resp.json() + + # 兼容两种响应格式: + # 1. {"results": [...]} + # 2. [...] + if isinstance(data, dict) and "results" in data: + items = data["results"] + elif isinstance(data, list): + items = data + else: + logger.warning(f"Unexpected search response format: {type(data)}") + return [] + + results: list[QueryResult] = [] + for item in items: + if isinstance(item, dict): + results.append( + QueryResult( + content=item.get("content", ""), + source_id=self._source_id, + source_name=self._source_name, + score=float(item.get("score", 0.0)), + metadata=item.get("metadata", {}), + doc_id=item.get("doc_id", item.get("id", "")), + title=item.get("title", ""), + ) + ) + return results[:top_k] + + except httpx.HTTPStatusError as e: + logger.error( + f"GenericHTTP search HTTP error: {e.response.status_code} — " + f"{e.response.text[:200]}" + ) + return [] + except Exception as e: + logger.error(f"GenericHTTP search error: {e}") + return [] + + async def ingest(self, documents: list[Document]) -> list[str]: + """写入文档到 HTTP 知识库 + + POST {endpoint_url}/ingest + Body: {"documents": [...]} + """ + client = self._get_client() + try: + payload = { + "documents": [ + { + "doc_id": doc.doc_id, + "content": doc.content, + "title": doc.title, + "source_id": doc.source_id, + "metadata": doc.metadata, + } + for doc in documents + ] + } + resp = await client.post("/ingest", json=payload) + resp.raise_for_status() + data = resp.json() + + # 兼容响应格式 + if isinstance(data, dict) and "ids" in data: + return data["ids"] + elif isinstance(data, list): + return [str(item) for item in data] + else: + return [doc.doc_id for doc in documents] + + except httpx.HTTPStatusError as e: + logger.error( + f"GenericHTTP ingest HTTP error: {e.response.status_code} — " + f"{e.response.text[:200]}" + ) + return [] + except Exception as e: + logger.error(f"GenericHTTP ingest error: {e}") + return [] + + async def delete_by_id(self, id: str) -> bool: + """按文档 ID 删除 + + DELETE {endpoint_url}/documents/{id} + """ + client = self._get_client() + try: + resp = await client.delete(f"/documents/{id}") + if resp.status_code in (200, 204): + return True + return False + except Exception as e: + logger.error(f"GenericHTTP delete_by_id error: {e}") + return False + + async def get_document(self, doc_id: str) -> Document | None: + """获取单个文档 + + GET {endpoint_url}/documents/{doc_id} + """ + client = self._get_client() + try: + resp = await client.get(f"/documents/{doc_id}") + resp.raise_for_status() + data = resp.json() + + return Document( + doc_id=data.get("doc_id", data.get("id", doc_id)), + content=data.get("content", ""), + title=data.get("title", ""), + source_id=data.get("source_id", self._source_id), + metadata=data.get("metadata", {}), + ) + except Exception as e: + logger.error(f"GenericHTTP get_document error: {e}") + return None + + async def list_sources(self) -> list[SourceInfo]: + """列出信息源 + + GET {endpoint_url}/sources + """ + client = self._get_client() + try: + resp = await client.get("/sources") + resp.raise_for_status() + data = resp.json() + + if isinstance(data, list): + sources: list[SourceInfo] = [] + for item in data: + if isinstance(item, dict): + sources.append( + SourceInfo( + source_id=item.get("source_id", ""), + source_name=item.get("source_name", ""), + source_type=item.get("source_type", "generic_http"), + document_count=item.get("document_count", 0), + ) + ) + return sources if sources else [ + SourceInfo( + source_id=self._source_id, + source_name=self._source_name, + source_type=self._source_type, + ) + ] + except Exception as e: + logger.debug(f"GenericHTTP list_sources error (endpoint may not exist): {e}") + + return [ + SourceInfo( + source_id=self._source_id, + source_name=self._source_name, + source_type=self._source_type, + ) + ] + + async def health_check(self) -> bool: + """检查 HTTP 知识库服务连接状态""" + client = self._get_client() + try: + resp = await client.get("/health") + if resp.status_code == 200: + return True + except Exception: + pass + + # Fallback: try the base endpoint + try: + resp = await client.get("/") + return resp.status_code in (200, 401, 403) + except Exception: + return False + + async def authenticate(self) -> bool: + """认证验证 + + 对于 GenericHTTPAdapter,认证通过 health_check 验证。 + """ + try: + self._authenticated = await self.health_check() + except Exception: + self._authenticated = False + return self._authenticated diff --git a/src/agentkit/memory/chunking.py b/src/agentkit/memory/chunking.py new file mode 100644 index 0000000..ad0dbe4 --- /dev/null +++ b/src/agentkit/memory/chunking.py @@ -0,0 +1,330 @@ +"""Chunking - 文档分块策略 + +提供两种分块策略: +- TextChunker: 按字符数分块,带重叠 +- StructuralChunker: 按文档结构(标题/段落)分块,适用于 Markdown/HTML +""" + +from __future__ import annotations + +import logging +import re +import uuid +from dataclasses import dataclass, field +from typing import Any + +logger = logging.getLogger(__name__) + + +@dataclass +class Chunk: + """文档分块""" + + chunk_id: str + content: str + metadata: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if "source_doc" not in self.metadata: + self.metadata["source_doc"] = "" + if "position" not in self.metadata: + self.metadata["position"] = 0 + + def to_dict(self) -> dict[str, Any]: + return { + "chunk_id": self.chunk_id, + "content": self.content, + "metadata": self.metadata, + } + + +class TextChunker: + """按字符数分块,带重叠 + + 适用于纯文本文档,按固定字符数切分,相邻块之间有重叠区域。 + """ + + def __init__( + self, + chunk_size: int = 1000, + chunk_overlap: int = 200, + separator: str = "\n\n", + ): + """ + Args: + chunk_size: 每个块的最大字符数 + chunk_overlap: 相邻块之间的重叠字符数 + separator: 优先分割符 + """ + if chunk_overlap >= chunk_size: + raise ValueError(f"chunk_overlap ({chunk_overlap}) must be less than chunk_size ({chunk_size})") + self._chunk_size = chunk_size + self._chunk_overlap = chunk_overlap + self._separator = separator + + def chunk( + self, + text: str, + source_doc_id: str = "", + metadata: dict[str, Any] | None = None, + ) -> list[Chunk]: + """将文本分块 + + Args: + text: 待分块文本 + source_doc_id: 源文档 ID + metadata: 附加元数据 + + Returns: + Chunk 列表 + """ + if not text.strip(): + return [] + + # 先尝试按分隔符分割 + segments = self._split_by_separator(text) + + # 合并小段,切分大段 + chunks_text = self._merge_and_split(segments) + + base_meta = dict(metadata or {}) + base_meta["source_doc"] = source_doc_id + base_meta["chunking_strategy"] = "text" + + chunks = [] + for i, chunk_text in enumerate(chunks_text): + chunk_meta = dict(base_meta) + chunk_meta["position"] = i + chunk_meta["char_count"] = len(chunk_text) + chunks.append(Chunk( + chunk_id=str(uuid.uuid4()), + content=chunk_text, + metadata=chunk_meta, + )) + + return chunks + + def _split_by_separator(self, text: str) -> list[str]: + """按分隔符分割文本""" + segments = text.split(self._separator) + # 过滤空段 + return [s.strip() for s in segments if s.strip()] + + def _merge_and_split(self, segments: list[str]) -> list[str]: + """合并小段,切分大段""" + result: list[str] = [] + current: list[str] = [] + current_len = 0 + + for segment in segments: + seg_len = len(segment) + + # 如果单个段超过 chunk_size,需要进一步切分 + if seg_len > self._chunk_size: + # 先把当前累积的段输出 + if current: + result.append(self._separator.join(current)) + current = [] + current_len = 0 + + # 切分大段 + for sub in self._split_large_segment(segment): + result.append(sub) + continue + + # 如果加入当前段会超过 chunk_size,先输出当前累积 + if current_len + seg_len + len(self._separator) > self._chunk_size and current: + result.append(self._separator.join(current)) + # 保留重叠部分 + overlap_text = self._separator.join(current) + overlap_start = max(0, len(overlap_text) - self._chunk_overlap) + overlap_segments = self._get_overlap_segments( + overlap_text[overlap_start:], segments + ) + current = overlap_segments + current_len = sum(len(s) for s in current) + len(self._separator) * max(0, len(current) - 1) + + current.append(segment) + current_len += seg_len + len(self._separator) + + if current: + result.append(self._separator.join(current)) + + return result + + def _split_large_segment(self, segment: str) -> list[str]: + """切分超大段""" + result = [] + start = 0 + while start < len(segment): + end = start + self._chunk_size + # 尝试在句子边界切分 + if end < len(segment): + # 查找最近的句子结束符 + for sep in ["。", ".", "!", "!", "?", "?", "\n"]: + last_sep = segment.rfind(sep, start + self._chunk_size // 2, end) + if last_sep > start: + end = last_sep + len(sep) + break + result.append(segment[start:end].strip()) + start = end - self._chunk_overlap + if start <= 0 and end >= len(segment): + break + if start < 0: + start = 0 + return [r for r in result if r] + + def _get_overlap_segments(self, overlap_text: str, segments: list[str]) -> list[str]: + """从重叠文本中提取完整段""" + # 简化实现:将重叠文本作为一个段 + if overlap_text.strip(): + return [overlap_text.strip()] + return [] + + +class StructuralChunker: + """按文档结构分块 + + 适用于 Markdown 和 HTML 等有标题结构的文档。 + 按标题层级分块,每个标题下的内容作为一个块。 + 如果某个块超过 chunk_size,则回退到 TextChunker 继续切分。 + """ + + def __init__( + self, + chunk_size: int = 1000, + chunk_overlap: int = 200, + heading_levels: int = 3, + ): + """ + Args: + chunk_size: 每个块的最大字符数 + chunk_overlap: 回退 TextChunker 时的重叠字符数 + heading_levels: 识别的标题层级数(1-6 对应 # 到 ######) + """ + self._chunk_size = chunk_size + self._chunk_overlap = chunk_overlap + self._heading_levels = min(max(heading_levels, 1), 6) + self._text_chunker = TextChunker( + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + ) + + def chunk( + self, + text: str, + source_doc_id: str = "", + metadata: dict[str, Any] | None = None, + ) -> list[Chunk]: + """将文本按结构分块 + + Args: + text: 待分块文本(Markdown 格式) + source_doc_id: 源文档 ID + metadata: 附加元数据 + + Returns: + Chunk 列表 + """ + if not text.strip(): + return [] + + sections = self._split_by_headings(text) + + base_meta = dict(metadata or {}) + base_meta["source_doc"] = source_doc_id + base_meta["chunking_strategy"] = "structural" + + chunks = [] + position = 0 + + for section in sections: + heading = section["heading"] + content = section["content"] + level = section["level"] + + if not content.strip(): + continue + + # 如果内容超过 chunk_size,使用 TextChunker 继续切分 + if len(content) > self._chunk_size: + sub_chunks = self._text_chunker.chunk( + content, + source_doc_id=source_doc_id, + metadata=metadata, + ) + for sub in sub_chunks: + sub.metadata["position"] = position + sub.metadata["heading"] = heading + sub.metadata["heading_level"] = level + sub.metadata["chunking_strategy"] = "structural" + position += 1 + chunks.append(sub) + else: + chunk_meta = dict(base_meta) + chunk_meta["position"] = position + chunk_meta["heading"] = heading + chunk_meta["heading_level"] = level + chunk_meta["char_count"] = len(content) + chunks.append(Chunk( + chunk_id=str(uuid.uuid4()), + content=content, + metadata=chunk_meta, + )) + position += 1 + + return chunks + + def _split_by_headings(self, text: str) -> list[dict[str, Any]]: + """按标题分割 Markdown 文本 + + Returns: + 列表,每项包含 heading, content, level + """ + lines = text.split("\n") + sections: list[dict[str, Any]] = [] + current_heading = "" + current_level = 0 + current_lines: list[str] = [] + + heading_pattern = re.compile(r"^(#{1," + str(self._heading_levels) + r"})\s+(.+)$") + + for line in lines: + match = heading_pattern.match(line) + if match: + # 保存当前节 + if current_lines: + content = "\n".join(current_lines).strip() + if content: + sections.append({ + "heading": current_heading, + "content": content, + "level": current_level, + }) + + # 开始新节 + current_heading = match.group(2).strip() + current_level = len(match.group(1)) + current_lines = [line] + else: + current_lines.append(line) + + # 保存最后一节 + if current_lines: + content = "\n".join(current_lines).strip() + if content: + sections.append({ + "heading": current_heading, + "content": content, + "level": current_level, + }) + + # 如果没有标题结构,整体作为一个块 + if not sections: + sections.append({ + "heading": "", + "content": text.strip(), + "level": 0, + }) + + return sections diff --git a/src/agentkit/memory/document_loader.py b/src/agentkit/memory/document_loader.py new file mode 100644 index 0000000..d098b51 --- /dev/null +++ b/src/agentkit/memory/document_loader.py @@ -0,0 +1,330 @@ +"""DocumentLoader - 多格式文档解析器 + +支持 PDF(PyMuPDF/pdfplumber)、Word(python-docx)、Markdown(mistune)、 +HTML(BeautifulSoup)、纯文本。所有格式依赖均为可选(try/except ImportError)。 +""" + +from __future__ import annotations + +import logging +import uuid +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) + + +@dataclass +class Document: + """解析后的文档统一格式""" + + doc_id: str + title: str + content: str + metadata: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if "source" not in self.metadata: + self.metadata["source"] = "" + if "format" not in self.metadata: + self.metadata["format"] = "unknown" + if "page_count" not in self.metadata: + self.metadata["page_count"] = 0 + if "created_at" not in self.metadata: + self.metadata["created_at"] = datetime.now(timezone.utc).isoformat() + + def to_dict(self) -> dict[str, Any]: + return { + "doc_id": self.doc_id, + "title": self.title, + "content": self.content, + "metadata": self.metadata, + } + + +def _detect_format(filename: str) -> str: + """根据文件扩展名检测文档格式""" + ext = Path(filename).suffix.lower() + format_map = { + ".pdf": "pdf", + ".docx": "docx", + ".doc": "docx", + ".md": "markdown", + ".markdown": "markdown", + ".html": "html", + ".htm": "html", + ".txt": "text", + ".csv": "text", + ".json": "text", + ".xml": "text", + } + return format_map.get(ext, "text") + + +class DocumentLoader: + """多格式文档解析器 + + 支持格式: + - PDF: PyMuPDF (fitz) → pdfplumber → 纯文本回退 + - Word: python-docx → 纯文本回退 + - Markdown: mistune → 纯文本回退 + - HTML: BeautifulSoup → 纯文本回退 + - 纯文本: 直接读取 + """ + + def load(self, file_path: str | Path) -> Document: + """从文件路径加载文档 + + Args: + file_path: 文件路径 + + Returns: + 解析后的 Document 对象 + + Raises: + FileNotFoundError: 文件不存在 + ValueError: 不支持的格式 + """ + path = Path(file_path) + if not path.exists(): + raise FileNotFoundError(f"File not found: {path}") + + content = path.read_bytes() + return self.load_bytes(content, path.name) + + def load_bytes(self, content: bytes, filename: str) -> Document: + """从字节内容加载文档 + + Args: + content: 文件字节内容 + filename: 文件名(用于格式检测和元数据) + + Returns: + 解析后的 Document 对象 + """ + doc_format = _detect_format(filename) + doc_id = str(uuid.uuid4()) + + parsers = { + "pdf": self._parse_pdf, + "docx": self._parse_docx, + "markdown": self._parse_markdown, + "html": self._parse_html, + "text": self._parse_text, + } + + parser = parsers.get(doc_format) + if parser is None: + logger.warning(f"Unsupported format '{doc_format}' for {filename}, falling back to text") + parser = self._parse_text + + text, extra_meta = parser(content, filename) + + metadata: dict[str, Any] = { + "source": filename, + "format": doc_format, + "created_at": datetime.now(timezone.utc).isoformat(), + } + metadata.update(extra_meta) + + title = Path(filename).stem + if "title" in extra_meta: + title = extra_meta["title"] + + return Document( + doc_id=doc_id, + title=title, + content=text, + metadata=metadata, + ) + + def _parse_pdf(self, content: bytes, filename: str) -> tuple[str, dict[str, Any]]: + """解析 PDF 文件 + + 优先使用 PyMuPDF (fitz),回退到 pdfplumber,最终回退到纯文本。 + """ + # 尝试 PyMuPDF + try: + import fitz # PyMuPDF + + doc = fitz.open(stream=content, filetype="pdf") + pages = [] + for page in doc: + pages.append(page.get_text()) + text = "\n\n".join(pages) + meta = { + "page_count": len(doc), + "parser": "pymupdf", + } + # 提取 PDF 元数据中的标题 + pdf_meta = doc.metadata + if pdf_meta and pdf_meta.get("title"): + meta["title"] = pdf_meta["title"] + doc.close() + return text, meta + except ImportError: + pass + except Exception as e: + logger.warning(f"PyMuPDF parsing failed for {filename}: {e}") + + # 尝试 pdfplumber + try: + import pdfplumber + import io + + pdf = pdfplumber.open(io.BytesIO(content)) + pages = [] + for page in pdf.pages: + page_text = page.extract_text() + if page_text: + pages.append(page_text) + text = "\n\n".join(pages) + meta = { + "page_count": len(pdf.pages), + "parser": "pdfplumber", + } + pdf.close() + return text, meta + except ImportError: + pass + except Exception as e: + logger.warning(f"pdfplumber parsing failed for {filename}: {e}") + + # 回退到纯文本 + logger.warning(f"No PDF parser available for {filename}, falling back to text extraction") + return self._parse_text(content, filename) + + def _parse_docx(self, content: bytes, filename: str) -> tuple[str, dict[str, Any]]: + """解析 Word 文件 + + 使用 python-docx,回退到纯文本。 + """ + try: + from docx import Document as DocxDocument + import io + + doc = DocxDocument(io.BytesIO(content)) + paragraphs = [] + table_count = 0 + + # 提取段落文本 + for para in doc.paragraphs: + if para.text.strip(): + paragraphs.append(para.text.strip()) + + # 提取表格文本 + for table in doc.tables: + table_count += 1 + for row in table.rows: + row_text = " | ".join(cell.text.strip() for cell in row.cells) + if row_text.strip(" |"): + paragraphs.append(row_text) + + text = "\n\n".join(paragraphs) + meta = { + "parser": "python-docx", + "table_count": table_count, + } + + # 提取文档属性中的标题 + if doc.core_properties and doc.core_properties.title: + meta["title"] = doc.core_properties.title + + return text, meta + except ImportError: + logger.warning(f"python-docx not available for {filename}, falling back to text") + return self._parse_text(content, filename) + except Exception as e: + logger.warning(f"python-docx parsing failed for {filename}: {e}") + return self._parse_text(content, filename) + + def _parse_markdown(self, content: bytes, filename: str) -> tuple[str, dict[str, Any]]: + """解析 Markdown 文件 + + 使用 mistune(如果可用),否则直接读取文本。 + Markdown 原文保留,因为后续分块需要标题结构。 + """ + try: + text = content.decode("utf-8") + except UnicodeDecodeError: + text = content.decode("utf-8", errors="replace") + + # 提取第一个标题作为文档标题 + title = "" + for line in text.split("\n"): + line_stripped = line.strip() + if line_stripped.startswith("#"): + title = line_stripped.lstrip("#").strip() + break + + meta: dict[str, Any] = { + "parser": "markdown", + } + if title: + meta["title"] = title + + # 尝试用 mistune 提取结构信息(但保留原文用于分块) + try: + import mistune + + # 统计标题数量 + heading_count = 0 + for line in text.split("\n"): + if line.strip().startswith("#"): + heading_count += 1 + meta["heading_count"] = heading_count + except ImportError: + pass + + return text, meta + + def _parse_html(self, content: bytes, filename: str) -> tuple[str, dict[str, Any]]: + """解析 HTML 文件 + + 使用 BeautifulSoup 提取文本,回退到纯文本。 + """ + try: + from bs4 import BeautifulSoup + + try: + html_text = content.decode("utf-8") + except UnicodeDecodeError: + html_text = content.decode("utf-8", errors="replace") + + soup = BeautifulSoup(html_text, "html.parser") + + # 移除 script 和 style 标签 + for tag in soup(["script", "style"]): + tag.decompose() + + text = soup.get_text(separator="\n", strip=True) + + # 提取标题 + title = "" + if soup.title and soup.title.string: + title = soup.title.string.strip() + + meta: dict[str, Any] = { + "parser": "beautifulsoup", + } + if title: + meta["title"] = title + + return text, meta + except ImportError: + logger.warning(f"BeautifulSoup not available for {filename}, falling back to text") + return self._parse_text(content, filename) + except Exception as e: + logger.warning(f"BeautifulSoup parsing failed for {filename}: {e}") + return self._parse_text(content, filename) + + def _parse_text(self, content: bytes, filename: str) -> tuple[str, dict[str, Any]]: + """解析纯文本文件""" + try: + text = content.decode("utf-8") + except UnicodeDecodeError: + text = content.decode("utf-8", errors="replace") + + return text, {"parser": "text"} diff --git a/src/agentkit/memory/knowledge_base.py b/src/agentkit/memory/knowledge_base.py new file mode 100644 index 0000000..84ef46d --- /dev/null +++ b/src/agentkit/memory/knowledge_base.py @@ -0,0 +1,84 @@ +"""KnowledgeBase 协议定义 - 外部知识库统一接口 + +独立于 Memory 接口,提供语义检索模型的知识库协议。 +Memory 的 retrieve(key)/delete(key) 是精确 key-value 语义, +而 KnowledgeBase 的 query()/ingest() 更适合知识库的语义检索场景。 + +参见 KTD-7: KBAdapter 使用独立 KnowledgeBase 协议,不直接实现 Memory 接口。 +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any, Protocol, runtime_checkable + + +@dataclass +class Document: + """知识库文档""" + + doc_id: str + content: str + title: str = "" + source_id: str = "" + metadata: dict[str, Any] = field(default_factory=dict) + created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + + +@dataclass +class QueryResult: + """知识库检索结果""" + + content: str + source_id: str + source_name: str + score: float + metadata: dict[str, Any] = field(default_factory=dict) + doc_id: str = "" + title: str = "" + + +@dataclass +class SourceInfo: + """知识库信息源描述""" + + source_id: str + source_name: str + source_type: str # e.g. "feishu", "confluence", "generic_http" + document_count: int = 0 + last_updated: datetime | None = None + + +@runtime_checkable +class KnowledgeBase(Protocol): + """知识库协议 - 统一的外部知识库接口 + + 所有知识库适配器(飞书、Confluence、通用 HTTP 等)均实现此协议。 + 与 Memory 接口不同,KnowledgeBase 专注于语义检索模型: + - ingest() 批量写入文档 + - query() 语义检索 + - delete_by_id() 按文档 ID 删除 + - list_sources() 列出可用信息源 + - health_check() 检查连接状态 + """ + + async def ingest(self, documents: list[Document]) -> list[str]: + """写入文档到知识库,返回文档 ID 列表""" + ... + + async def query(self, text: str, top_k: int = 5) -> list[QueryResult]: + """语义检索知识库""" + ... + + async def delete_by_id(self, id: str) -> bool: + """按文档 ID 删除""" + ... + + async def list_sources(self) -> list[SourceInfo]: + """列出可用信息源""" + ... + + async def health_check(self) -> bool: + """检查知识库连接状态""" + ... diff --git a/src/agentkit/memory/local_rag.py b/src/agentkit/memory/local_rag.py new file mode 100644 index 0000000..6e9e2c5 --- /dev/null +++ b/src/agentkit/memory/local_rag.py @@ -0,0 +1,525 @@ +"""LocalRAGService - 本地文档 RAG 服务 + +实现 KnowledgeBase 协议,支持文档摄取、语义检索、删除和来源追溯。 +提供两种实现: +- LocalRAGService: 基于 pgvector + PostgreSQL(生产环境) +- InMemoryLocalRAGService: 基于内存(测试和开发环境) +""" + +from __future__ import annotations + +import json +import logging +import uuid +from datetime import datetime, timezone +from typing import Any + +from agentkit.memory.chunking import Chunk, StructuralChunker, TextChunker +from agentkit.memory.document_loader import Document as LoaderDocument +from agentkit.memory.embedder import Embedder +from agentkit.memory.knowledge_base import ( + Document, + KnowledgeBase, + QueryResult, + SourceInfo, +) + +logger = logging.getLogger(__name__) + + +def _loader_doc_to_kb_doc(loader_doc: LoaderDocument) -> Document: + """将 document_loader.Document 转换为 knowledge_base.Document""" + return Document( + doc_id=loader_doc.doc_id, + content=loader_doc.content, + title=loader_doc.title, + source_id=loader_doc.metadata.get("source", ""), + metadata=loader_doc.metadata, + ) + + +class LocalRAGService: + """基于 pgvector 的本地 RAG 服务 + + 实现 KnowledgeBase 协议,使用 pgvector 存储 + 检索。 + 复用 EpisodicMemory 的 pgvector 基础设施模式。 + + 摄取 Pipeline:上传 → 解析 → 分块 → 嵌入 → 写入 pgvector + """ + + def __init__( + self, + session_factory: Any, + embedder: Embedder, + chunk_size: int = 1000, + chunk_overlap: int = 200, + table_name: str = "knowledge_chunks", + pgvector_enabled: bool = True, + ): + """ + Args: + session_factory: 返回 async context manager 的工厂 + embedder: 嵌入器,用于生成向量 + chunk_size: 分块大小 + chunk_overlap: 分块重叠 + table_name: pgvector 查询使用的表名 + pgvector_enabled: 是否使用 pgvector 原生检索 + """ + self._session_factory = session_factory + self._embedder = embedder + self._chunk_size = chunk_size + self._chunk_overlap = chunk_overlap + self._table_name = table_name + self._pgvector_enabled = pgvector_enabled + self._text_chunker = TextChunker(chunk_size=chunk_size, chunk_overlap=chunk_overlap) + self._structural_chunker = StructuralChunker(chunk_size=chunk_size, chunk_overlap=chunk_overlap) + + async def ingest(self, documents: list[Document]) -> list[str]: + """摄取文档列表 + + Args: + documents: knowledge_base.Document 对象列表 + + Returns: + 成功摄取的文档 ID 列表 + """ + ingested_ids = [] + + for doc in documents: + try: + chunks = self._chunk_document(doc) + await self._store_chunks(doc, chunks) + ingested_ids.append(doc.doc_id) + logger.info(f"Ingested document '{doc.title}' with {len(chunks)} chunks") + except Exception as e: + logger.error(f"Failed to ingest document '{doc.title}': {e}") + + return ingested_ids + + async def query(self, text: str, top_k: int = 5) -> list[QueryResult]: + """语义检索 + + Args: + text: 查询文本 + top_k: 返回结果数量 + + Returns: + 检索结果列表 + """ + query_embedding = await self._embedder.embed(text) + + async with self._session_factory() as db: + try: + if self._pgvector_enabled: + return await self._query_pgvector(db, query_embedding, top_k) + return await self._query_client_side(db, query_embedding, top_k) + except Exception as e: + logger.error(f"Failed to query knowledge base: {e}") + return [] + + async def delete_by_id(self, id: str) -> bool: + """按文档 ID 删除 + + Args: + id: 文档 ID + + Returns: + 是否删除成功 + """ + async with self._session_factory() as db: + try: + from sqlalchemy import text as sql_text + + sql = sql_text( + f"DELETE FROM {self._table_name} WHERE source_doc_id = :doc_id" + ) + await db.execute(sql, {"doc_id": id}) + await db.commit() + return True + except Exception as e: + await db.rollback() + logger.error(f"Failed to delete document {id}: {e}") + return False + + async def list_sources(self) -> list[SourceInfo]: + """列出已摄取的文档 + + Returns: + 文档元信息列表 + """ + async with self._session_factory() as db: + try: + from sqlalchemy import text as sql_text + + sql = sql_text( + f"SELECT source_doc_id, source_title, doc_format, " + f"COUNT(*) as chunk_count, " + f"MIN(created_at) as created_at, " + f"MIN(doc_metadata) as doc_metadata " + f"FROM {self._table_name} " + f"GROUP BY source_doc_id, source_title, doc_format " + f"ORDER BY MIN(created_at) DESC" + ) + result = await db.execute(sql) + rows = result.mappings().all() + + sources = [] + for row in rows: + meta = {} + if row.get("doc_metadata"): + try: + meta = json.loads(row["doc_metadata"]) + except (json.JSONDecodeError, TypeError): + pass + + sources.append(SourceInfo( + source_id=row["source_doc_id"], + source_name=row.get("source_title", ""), + source_type=row.get("doc_format", "local"), + document_count=row.get("chunk_count", 0), + last_updated=row["created_at"] if row.get("created_at") else None, + )) + return sources + except Exception as e: + logger.error(f"Failed to list sources: {e}") + return [] + + async def health_check(self) -> bool: + """检查服务健康状态""" + async with self._session_factory() as db: + try: + from sqlalchemy import text as sql_text + + await db.execute(sql_text(f"SELECT 1 FROM {self._table_name} LIMIT 1")) + return True + except Exception as e: + logger.error(f"Health check failed: {e}") + return False + + def _chunk_document(self, doc: Document) -> list[Chunk]: + """将文档分块""" + doc_format = doc.metadata.get("format", "text") + + # Markdown 和 HTML 使用结构化分块 + if doc_format in ("markdown", "html"): + chunks = self._structural_chunker.chunk( + doc.content, + source_doc_id=doc.doc_id, + metadata=doc.metadata, + ) + else: + chunks = self._text_chunker.chunk( + doc.content, + source_doc_id=doc.doc_id, + metadata=doc.metadata, + ) + + return chunks + + async def _store_chunks(self, doc: Document, chunks: list[Chunk]) -> None: + """存储文档块到 pgvector""" + async with self._session_factory() as db: + try: + from sqlalchemy import text as sql_text + + for chunk in chunks: + # 生成嵌入 + embedding = await self._embedder.embed(chunk.content) + + sql = sql_text( + f"INSERT INTO {self._table_name} " + f"(chunk_id, source_doc_id, source_title, doc_format, " + f"content, embedding, chunk_metadata, doc_metadata, created_at) " + f"VALUES (:chunk_id, :doc_id, :title, :format, " + f":content, :embedding, :chunk_meta, :doc_meta, :created_at)" + ) + + await db.execute(sql, { + "chunk_id": chunk.chunk_id, + "doc_id": doc.doc_id, + "title": doc.title, + "format": doc.metadata.get("format", "unknown"), + "content": chunk.content, + "embedding": str(embedding), + "chunk_meta": json.dumps(chunk.metadata, ensure_ascii=False), + "doc_meta": json.dumps(doc.metadata, ensure_ascii=False), + "created_at": datetime.now(timezone.utc), + }) + + await db.commit() + except Exception as e: + await db.rollback() + logger.error(f"Failed to store chunks for document '{doc.title}': {e}") + raise + + async def _query_pgvector( + self, + db: Any, + query_embedding: list[float], + top_k: int, + ) -> list[QueryResult]: + """使用 pgvector <=> 算符检索""" + from sqlalchemy import text as sql_text + + sql = sql_text( + f"SELECT chunk_id, source_doc_id, source_title, content, " + f"chunk_metadata, embedding <=> :query_vec AS distance " + f"FROM {self._table_name} " + f"ORDER BY embedding <=> :query_vec " + f"LIMIT :lim" + ) + + result = await db.execute(sql, { + "query_vec": str(query_embedding), + "lim": top_k, + }) + rows = result.mappings().all() + + results = [] + for row in rows: + # 从 distance 计算 cosine similarity + distance = row.get("distance", 0.0) + # pgvector <=> 返回 cosine distance = 1 - cosine_similarity + cosine = max(0.0, 1.0 - float(distance)) + + chunk_meta = {} + if row.get("chunk_metadata"): + try: + chunk_meta = json.loads(row["chunk_metadata"]) + except (json.JSONDecodeError, TypeError): + pass + + results.append(QueryResult( + content=row["content"], + source_id=row["source_doc_id"], + source_name=row.get("source_title", ""), + score=cosine, + metadata=chunk_meta, + doc_id=row["source_doc_id"], + title=row.get("source_title", ""), + )) + + return results + + async def _query_client_side( + self, + db: Any, + query_embedding: list[float], + top_k: int, + ) -> list[QueryResult]: + """客户端 O(N) cosine similarity 检索(回退路径)""" + from sqlalchemy import text as sql_text + + sql = sql_text( + f"SELECT chunk_id, source_doc_id, source_title, content, " + f"chunk_metadata, embedding " + f"FROM {self._table_name} " + f"LIMIT 500" + ) + + result = await db.execute(sql) + rows = result.mappings().all() + + candidates = [] + for row in rows: + row_embedding = row.get("embedding") + if row_embedding is None: + continue + + # 解析存储的 embedding + try: + if isinstance(row_embedding, str): + stored_embedding = json.loads(row_embedding) + else: + stored_embedding = list(row_embedding) + except (json.JSONDecodeError, TypeError): + continue + + cosine = self._compute_cosine_similarity(query_embedding, stored_embedding) + if cosine < 0.1: + continue + + chunk_meta = {} + if row.get("chunk_metadata"): + try: + chunk_meta = json.loads(row["chunk_metadata"]) + except (json.JSONDecodeError, TypeError): + pass + + candidates.append(QueryResult( + content=row["content"], + source_id=row["source_doc_id"], + source_name=row.get("source_title", ""), + score=cosine, + metadata=chunk_meta, + doc_id=row["source_doc_id"], + title=row.get("source_title", ""), + )) + + candidates.sort(key=lambda x: x.score, reverse=True) + return candidates[:top_k] + + @staticmethod + def _compute_cosine_similarity(vec_a: list[float], vec_b: list[float]) -> float: + """计算两个向量的余弦相似度""" + if len(vec_a) != len(vec_b) or not vec_a: + return 0.0 + dot_product = sum(a * b for a, b in zip(vec_a, vec_b)) + magnitude_a = sum(a**2 for a in vec_a) ** 0.5 + magnitude_b = sum(b**2 for b in vec_b) ** 0.5 + if magnitude_a == 0.0 or magnitude_b == 0.0: + return 0.0 + return dot_product / (magnitude_a * magnitude_b) + + +class InMemoryLocalRAGService: + """基于内存的本地 RAG 服务 + + 用于测试和开发环境,无需 pgvector 依赖。 + 实现 KnowledgeBase 协议。 + """ + + def __init__( + self, + embedder: Embedder, + chunk_size: int = 1000, + chunk_overlap: int = 200, + ): + """ + Args: + embedder: 嵌入器 + chunk_size: 分块大小 + chunk_overlap: 分块重叠 + """ + self._embedder = embedder + self._text_chunker = TextChunker(chunk_size=chunk_size, chunk_overlap=chunk_overlap) + self._structural_chunker = StructuralChunker(chunk_size=chunk_size, chunk_overlap=chunk_overlap) + + # 内存存储 + self._chunks: dict[str, dict[str, Any]] = {} # chunk_id → {content, embedding, metadata} + self._documents: dict[str, dict[str, Any]] = {} # doc_id → {title, format, chunk_ids, metadata, created_at} + + async def ingest(self, documents: list[Document]) -> list[str]: + """摄取文档列表 + + 也支持传入 document_loader.Document,会自动转换为 knowledge_base.Document。 + """ + ingested_ids = [] + + for doc in documents: + # 支持 document_loader.Document 自动转换 + if isinstance(doc, LoaderDocument): + doc = _loader_doc_to_kb_doc(doc) + + try: + chunks = self._chunk_document(doc) + chunk_ids = [] + + for chunk in chunks: + embedding = await self._embedder.embed(chunk.content) + self._chunks[chunk.chunk_id] = { + "content": chunk.content, + "embedding": embedding, + "metadata": chunk.metadata, + "source_doc_id": doc.doc_id, + } + chunk_ids.append(chunk.chunk_id) + + self._documents[doc.doc_id] = { + "title": doc.title, + "source_id": doc.source_id, + "format": doc.metadata.get("format", "unknown"), + "chunk_ids": chunk_ids, + "metadata": doc.metadata, + "created_at": datetime.now(timezone.utc), + } + ingested_ids.append(doc.doc_id) + logger.info(f"Ingested document '{doc.title}' with {len(chunks)} chunks") + except Exception as e: + logger.error(f"Failed to ingest document '{doc.title}': {e}") + + return ingested_ids + + async def query(self, text: str, top_k: int = 5) -> list[QueryResult]: + """语义检索""" + query_embedding = await self._embedder.embed(text) + + candidates = [] + for chunk_id, chunk_data in self._chunks.items(): + stored_embedding = chunk_data["embedding"] + cosine = self._compute_cosine_similarity(query_embedding, stored_embedding) + if cosine < 0.1: + continue + + source_doc_id = chunk_data["source_doc_id"] + doc_info = self._documents.get(source_doc_id, {}) + + candidates.append(QueryResult( + content=chunk_data["content"], + source_id=source_doc_id, + source_name=doc_info.get("title", ""), + score=cosine, + metadata=chunk_data.get("metadata", {}), + doc_id=source_doc_id, + title=doc_info.get("title", ""), + )) + + candidates.sort(key=lambda x: x.score, reverse=True) + return candidates[:top_k] + + async def delete_by_id(self, id: str) -> bool: + """按文档 ID 删除""" + if id not in self._documents: + return False + + doc_info = self._documents[id] + for chunk_id in doc_info.get("chunk_ids", []): + self._chunks.pop(chunk_id, None) + + del self._documents[id] + return True + + async def list_sources(self) -> list[SourceInfo]: + """列出已摄取的文档""" + sources = [] + for doc_id, doc_info in self._documents.items(): + sources.append(SourceInfo( + source_id=doc_id, + source_name=doc_info["title"], + source_type=doc_info.get("format", "local"), + document_count=len(doc_info.get("chunk_ids", [])), + last_updated=doc_info.get("created_at"), + )) + return sources + + async def health_check(self) -> bool: + """检查服务健康状态""" + return True + + def _chunk_document(self, doc: Document) -> list[Chunk]: + """将文档分块""" + doc_format = doc.metadata.get("format", "text") + + if doc_format in ("markdown", "html"): + return self._structural_chunker.chunk( + doc.content, + source_doc_id=doc.doc_id, + metadata=doc.metadata, + ) + return self._text_chunker.chunk( + doc.content, + source_doc_id=doc.doc_id, + metadata=doc.metadata, + ) + + @staticmethod + def _compute_cosine_similarity(vec_a: list[float], vec_b: list[float]) -> float: + """计算两个向量的余弦相似度""" + if len(vec_a) != len(vec_b) or not vec_a: + return 0.0 + dot_product = sum(a * b for a, b in zip(vec_a, vec_b)) + magnitude_a = sum(a**2 for a in vec_a) ** 0.5 + magnitude_b = sum(b**2 for b in vec_b) ** 0.5 + if magnitude_a == 0.0 or magnitude_b == 0.0: + return 0.0 + return dot_product / (magnitude_a * magnitude_b) diff --git a/src/agentkit/memory/multi_source_retriever.py b/src/agentkit/memory/multi_source_retriever.py new file mode 100644 index 0000000..6ccf35d --- /dev/null +++ b/src/agentkit/memory/multi_source_retriever.py @@ -0,0 +1,230 @@ +"""MultiSourceRetriever - 多源混合检索器 + +管理多个 KnowledgeBase 协议实现,支持: +- 信息源指定:search(query, sources=["feishu", "local:合规文档"]) +- 并行查询多个源,按权重融合排序 +- 来源追溯:每个检索结果附带 source_id + document_title + chunk_location +- 基于 hash 的去重 +""" + +from __future__ import annotations + +import asyncio +import hashlib +import logging +from dataclasses import replace +from typing import Any + +from agentkit.memory.knowledge_base import KnowledgeBase, QueryResult, SourceInfo + +logger = logging.getLogger(__name__) + + +def _content_hash(content: str) -> str: + """计算内容哈希,用于去重""" + return hashlib.md5(content.encode("utf-8")).hexdigest() + + +class MultiSourceRetriever: + """多源混合检索器 + + 管理多个 KnowledgeBase 协议实现,支持按名称指定信息源、 + 并行查询、权重融合排序和来源追溯。 + + 用法:: + + retriever = MultiSourceRetriever() + retriever.register_source("feishu", feishu_adapter) + retriever.register_source("local:合规文档", local_rag) + + # 仅从指定源检索 + results = await retriever.search("合规要求", sources=["feishu", "local:合规文档"]) + + # 从所有可用源检索 + results = await retriever.search("合规要求") + + # 带权重检索 + results = await retriever.search("合规要求", weights={"feishu": 1.5, "local:合规文档": 0.8}) + """ + + def __init__(self, sources: dict[str, KnowledgeBase] | None = None): + """ + Args: + sources: 初始信息源映射,key 为源名称,value 为 KnowledgeBase 实现 + """ + self._sources: dict[str, KnowledgeBase] = {} + if sources: + for name, kb in sources.items(): + self._sources[name] = kb + + def register_source(self, name: str, knowledge_base: KnowledgeBase) -> None: + """注册信息源 + + Args: + name: 信息源名称,如 "feishu"、"local:合规文档" + knowledge_base: KnowledgeBase 协议实现 + """ + self._sources[name] = knowledge_base + logger.info(f"Registered knowledge source: {name}") + + def unregister_source(self, name: str) -> bool: + """注销信息源 + + Args: + name: 信息源名称 + + Returns: + 是否成功注销 + """ + if name in self._sources: + del self._sources[name] + logger.info(f"Unregistered knowledge source: {name}") + return True + return False + + async def search( + self, + query: str, + top_k: int = 5, + sources: list[str] | None = None, + weights: dict[str, float] | None = None, + ) -> list[QueryResult]: + """多源检索 + + Args: + query: 检索查询 + top_k: 返回最大结果数 + sources: 指定信息源列表,None 表示查询所有已注册源 + weights: 信息源权重映射,用于提升/降低特定源的分数 + + Returns: + 融合排序后的检索结果列表,每个结果包含来源追溯信息 + """ + # 确定要查询的源 + target_sources = self._resolve_sources(sources) + if not target_sources: + logger.warning("No knowledge sources available for search") + return [] + + # 并行查询所有目标源 + results = await self._query_sources(query, top_k, target_sources, weights) + + # 去重 + results = self._deduplicate(results) + + # 按 score 降序排序,截取 top_k + results.sort(key=lambda r: r.score, reverse=True) + return results[:top_k] + + async def list_all_sources(self) -> dict[str, SourceInfo]: + """列出所有已注册信息源 + + Returns: + 源名称到 SourceInfo 的映射 + """ + result: dict[str, SourceInfo] = {} + for name, kb in self._sources.items(): + try: + source_infos = await kb.list_sources() + if source_infos: + result[name] = source_infos[0] + else: + # 知识库未返回 source info,构造一个占位 + result[name] = SourceInfo( + source_id=name, + source_name=name, + source_type="unknown", + ) + except Exception as e: + logger.error(f"Failed to list sources for '{name}': {e}") + result[name] = SourceInfo( + source_id=name, + source_name=name, + source_type="error", + ) + return result + + def get_source_names(self) -> list[str]: + """获取所有已注册信息源名称""" + return list(self._sources.keys()) + + def _resolve_sources(self, sources: list[str] | None) -> dict[str, KnowledgeBase]: + """解析目标信息源 + + Args: + sources: 指定的源名称列表,None 表示所有源 + + Returns: + 源名称到 KnowledgeBase 的映射 + """ + if sources is None: + return dict(self._sources) + + resolved = {} + for name in sources: + if name in self._sources: + resolved[name] = self._sources[name] + else: + logger.warning(f"Knowledge source '{name}' not found, skipping") + return resolved + + async def _query_sources( + self, + query: str, + top_k: int, + target_sources: dict[str, KnowledgeBase], + weights: dict[str, float] | None, + ) -> list[QueryResult]: + """并行查询多个信息源 + + Args: + query: 查询文本 + top_k: 每个源返回的最大结果数 + target_sources: 目标源映射 + weights: 权重映射 + + Returns: + 所有源的检索结果列表(已应用权重) + """ + async def _query_one(name: str, kb: KnowledgeBase) -> list[QueryResult]: + try: + results = await kb.query(query, top_k=top_k) + # 应用权重 + weight = (weights or {}).get(name, 1.0) + return [ + replace(r, score=r.score * weight, source_name=name) + for r in results + ] + except Exception as e: + logger.error(f"Query failed for source '{name}': {e}") + return [] + + tasks = [_query_one(name, kb) for name, kb in target_sources.items()] + all_results = await asyncio.gather(*tasks, return_exceptions=True) + + merged: list[QueryResult] = [] + for result in all_results: + if isinstance(result, Exception): + logger.error(f"Source query raised exception: {result}") + continue + if isinstance(result, list): + merged.extend(result) + + return merged + + @staticmethod + def _deduplicate(results: list[QueryResult]) -> list[QueryResult]: + """基于内容哈希去重,保留分数最高的结果 + + Args: + results: 待去重的结果列表 + + Returns: + 去重后的结果列表 + """ + seen: dict[str, QueryResult] = {} + for r in results: + content_key = _content_hash(r.content) + if content_key not in seen or r.score > seen[content_key].score: + seen[content_key] = r + return list(seen.values()) diff --git a/src/agentkit/memory/retriever.py b/src/agentkit/memory/retriever.py index ebbc571..e1b36d2 100644 --- a/src/agentkit/memory/retriever.py +++ b/src/agentkit/memory/retriever.py @@ -19,6 +19,8 @@ from agentkit.memory.semantic import SemanticMemory from agentkit.memory.query_transformer import QueryTransformerBase from agentkit.memory.rag_loop import RAGSelfCorrectionLoop from agentkit.memory.relevance_scorer import RelevanceScorer +from agentkit.memory.knowledge_base import KnowledgeBase, QueryResult +from agentkit.memory.multi_source_retriever import MultiSourceRetriever from agentkit.tools.base import Tool logger = logging.getLogger(__name__) @@ -59,6 +61,7 @@ class MemoryRetriever: context_template: str = "structured", enable_self_correction: bool = False, max_correction_retries: int = 3, + knowledge_sources: dict[str, KnowledgeBase] | None = None, ): self._working = working_memory self._episodic = episodic_memory @@ -79,6 +82,7 @@ class MemoryRetriever: query_transformer=query_transformer, max_retries=max_correction_retries, ) + self._multi_source_retriever = MultiSourceRetriever(sources=knowledge_sources) async def retrieve( self, @@ -87,6 +91,8 @@ class MemoryRetriever: token_budget: int = 3000, filters: dict[str, Any] | None = None, _skip_correction: bool = False, + sources: list[str] | None = None, + source_weights: dict[str, float] | None = None, ) -> list[MemoryItem]: """混合检索三层记忆 @@ -96,7 +102,15 @@ class MemoryRetriever: token_budget: token 预算 filters: 过滤条件 _skip_correction: 内部参数,CRAG 循环内部调用时跳过自纠正 + sources: 指定信息源列表,如 ["feishu", "local:合规文档"]。 + 传入时仅从指定源检索,不查三层记忆。 + source_weights: 信息源权重映射,用于多源检索时调整分数 """ + # 多源检索路径:指定了 sources 时委托给 MultiSourceRetriever + if sources is not None: + return await self._retrieve_from_sources( + query, top_k, token_budget, sources, source_weights + ) # Self-correction loop (CRAG) if ( self._enable_self_correction @@ -199,6 +213,63 @@ class MemoryRetriever: return all_items + async def _retrieve_from_sources( + self, + query: str, + top_k: int = 5, + token_budget: int = 3000, + sources: list[str] | None = None, + source_weights: dict[str, float] | None = None, + ) -> list[MemoryItem]: + """从指定信息源检索,将 QueryResult 转换为 MemoryItem + + Args: + query: 检索查询 + top_k: 返回最大结果数 + token_budget: token 预算 + sources: 信息源名称列表 + source_weights: 信息源权重映射 + """ + kb_results = await self._multi_source_retriever.search( + query, top_k=top_k, sources=sources, weights=source_weights + ) + + # QueryResult → MemoryItem + items = [] + for r in kb_results: + items.append(MemoryItem( + key=r.source_id, + value=r.content, + metadata={ + **r.metadata, + "source": "rag", + "source_name": r.source_name, + "doc_id": r.doc_id, + "document_title": r.title, + }, + score=r.score, + )) + + # Token 预算管理 + selected = [] + total_tokens = 0 + for item in items: + text = str(item.value) + estimated_tokens = _estimate_tokens(text) + if total_tokens + estimated_tokens > token_budget: + continue + selected.append(item) + total_tokens += estimated_tokens + if len(selected) >= top_k: + break + + return selected + + @property + def multi_source_retriever(self) -> MultiSourceRetriever: + """获取多源检索器,用于直接注册/注销信息源""" + return self._multi_source_retriever + async def get_context_string( self, query: str, diff --git a/tests/unit/memory/__init__.py b/tests/unit/memory/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/memory/test_adapters.py b/tests/unit/memory/test_adapters.py new file mode 100644 index 0000000..ac7ca6b --- /dev/null +++ b/tests/unit/memory/test_adapters.py @@ -0,0 +1,1237 @@ +"""Tests for KnowledgeBase adapters — 飞书、Confluence、通用 HTTP 适配器""" + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from agentkit.memory.knowledge_base import Document, QueryResult, SourceInfo, KnowledgeBase +from agentkit.memory.adapters.base import KBAdapter +from agentkit.memory.adapters.feishu import FeishuKBAdapter +from agentkit.memory.adapters.confluence import ConfluenceAdapter +from agentkit.memory.adapters.generic_http import GenericHTTPAdapter + + +# --------------------------------------------------------------------------- +# KnowledgeBase Protocol tests +# --------------------------------------------------------------------------- + + +class TestKnowledgeBaseProtocol: + """KnowledgeBase 协议验证""" + + def test_document_creation(self): + doc = Document(doc_id="d1", content="测试内容", title="测试文档") + assert doc.doc_id == "d1" + assert doc.content == "测试内容" + assert doc.title == "测试文档" + assert doc.metadata == {} + + def test_query_result_creation(self): + result = QueryResult( + content="检索结果", + source_id="feishu-xxx", + source_name="飞书知识库", + score=0.92, + ) + assert result.content == "检索结果" + assert result.score == 0.92 + assert result.doc_id == "" + assert result.metadata == {} + + def test_source_info_creation(self): + info = SourceInfo( + source_id="feishu-xxx", + source_name="飞书知识库", + source_type="feishu", + document_count=100, + ) + assert info.source_id == "feishu-xxx" + assert info.source_type == "feishu" + assert info.document_count == 100 + + def test_knowledge_base_protocol_check(self): + """验证适配器满足 KnowledgeBase 协议""" + adapter = FeishuKBAdapter( + app_id="cli_test", + app_secret="secret", + ) + assert isinstance(adapter, KnowledgeBase) + + adapter2 = ConfluenceAdapter( + base_url="https://test.atlassian.net/wiki", + username="user@test.com", + api_token="token", + ) + assert isinstance(adapter2, KnowledgeBase) + + adapter3 = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + ) + assert isinstance(adapter3, KnowledgeBase) + + +# --------------------------------------------------------------------------- +# KBAdapter base class tests +# --------------------------------------------------------------------------- + + +class TestKBAdapterBase: + """KBAdapter 抽象基类测试""" + + def _make_concrete_adapter(self) -> KBAdapter: + """创建一个具体子类用于测试""" + + class ConcreteAdapter(KBAdapter): + def _make_client(self): + return MagicMock() + + async def search(self, query, top_k=5): + return [] + + async def health_check(self): + return True + + return ConcreteAdapter( + source_id="test-adapter", + source_name="Test Adapter", + source_type="test", + ) + + @pytest.mark.asyncio + async def test_query_delegates_to_search(self): + adapter = self._make_concrete_adapter() + adapter.search = AsyncMock(return_value=[ + QueryResult(content="result", source_id="test", source_name="test", score=0.9) + ]) + + results = await adapter.query("test query", top_k=3) + adapter.search.assert_called_once_with("test query", top_k=3) + assert len(results) == 1 + + @pytest.mark.asyncio + async def test_list_sources_default(self): + adapter = self._make_concrete_adapter() + sources = await adapter.list_sources() + assert len(sources) == 1 + assert sources[0].source_id == "test-adapter" + assert sources[0].source_name == "Test Adapter" + + @pytest.mark.asyncio + async def test_ingest_default_returns_empty(self): + adapter = self._make_concrete_adapter() + docs = [Document(doc_id="d1", content="test")] + ids = await adapter.ingest(docs) + assert ids == [] + + @pytest.mark.asyncio + async def test_delete_by_id_default_returns_false(self): + adapter = self._make_concrete_adapter() + result = await adapter.delete_by_id("d1") + assert result is False + + @pytest.mark.asyncio + async def test_get_document_default_returns_none(self): + adapter = self._make_concrete_adapter() + result = await adapter.get_document("d1") + assert result is None + + @pytest.mark.asyncio + async def test_authenticate_delegates_to_health_check(self): + adapter = self._make_concrete_adapter() + adapter.health_check = AsyncMock(return_value=True) + result = await adapter.authenticate() + assert result is True + adapter.health_check.assert_called_once() + + @pytest.mark.asyncio + async def test_authenticate_failure(self): + adapter = self._make_concrete_adapter() + adapter.health_check = AsyncMock(side_effect=Exception("connection error")) + result = await adapter.authenticate() + assert result is False + + @pytest.mark.asyncio + async def test_context_manager(self): + adapter = self._make_concrete_adapter() + adapter.close = AsyncMock() + async with adapter as a: + assert a is adapter + adapter.close.assert_called_once() + + +# --------------------------------------------------------------------------- +# FeishuKBAdapter tests +# --------------------------------------------------------------------------- + + +class TestFeishuKBAdapterInit: + """FeishuKBAdapter 初始化""" + + def test_basic_init(self): + adapter = FeishuKBAdapter( + app_id="cli_test1234", + app_secret="secret", + ) + assert adapter._app_id == "cli_test1234" + assert adapter._app_secret == "secret" + assert adapter._base_url == "https://open.feishu.cn/open-apis" + assert adapter._space_ids == [] + assert adapter._source_type == "feishu" + + def test_init_with_custom_base_url(self): + adapter = FeishuKBAdapter( + app_id="cli_test1234", + app_secret="secret", + base_url="https://internal.feishu.cn/open-apis/", + space_ids=["space-1", "space-2"], + timeout=60, + ) + assert adapter._base_url == "https://internal.feishu.cn/open-apis" + assert adapter._space_ids == ["space-1", "space-2"] + assert adapter._timeout == 60 + + +class TestFeishuKBAdapterAuth: + """FeishuKBAdapter 认证""" + + @pytest.mark.asyncio + async def test_authenticate_success(self): + adapter = FeishuKBAdapter( + app_id="cli_test", + app_secret="secret", + ) + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = { + "code": 0, + "msg": "ok", + "tenant_access_token": "t-xxx", + "expire": 7200, + } + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + result = await adapter.authenticate() + assert result is True + assert adapter._access_token == "t-xxx" + + @pytest.mark.asyncio + async def test_authenticate_failure(self): + adapter = FeishuKBAdapter( + app_id="cli_test", + app_secret="wrong_secret", + ) + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = { + "code": 10014, + "msg": "invalid app_id or app_secret", + } + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + result = await adapter.authenticate() + assert result is False + + +class TestFeishuKBAdapterSearch: + """FeishuKBAdapter 检索""" + + @pytest.fixture + def adapter(self): + return FeishuKBAdapter( + app_id="cli_test", + app_secret="secret", + space_ids=["space-1"], + ) + + @pytest.mark.asyncio + async def test_search_success(self, adapter): + # Mock authentication + adapter._access_token = "t-xxx" + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = { + "code": 0, + "data": { + "items": [ + { + "wiki_token": "wikcnxxx", + "title": "飞书知识库文档", + "content": "这是飞书知识库的内容", + "score": 0.92, + "space_id": "space-1", + }, + { + "wiki_token": "wikcnyyy", + "title": "另一个文档", + "content": "另一个文档的内容", + "score": 0.85, + "space_id": "space-1", + }, + ], + }, + } + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + results = await adapter.search("飞书知识库", top_k=5) + + assert len(results) == 2 + assert results[0].content == "这是飞书知识库的内容" + assert results[0].score == 0.92 + assert results[0].source_id.startswith("feishu-") + assert results[0].source_name == "飞书知识库" + assert results[0].doc_id == "wikcnxxx" + assert results[0].title == "飞书知识库文档" + + @pytest.mark.asyncio + async def test_search_not_authenticated(self, adapter): + adapter._access_token = None + adapter._get_access_token = AsyncMock(return_value=None) + + results = await adapter.search("test") + assert results == [] + + @pytest.mark.asyncio + async def test_search_api_error(self, adapter): + adapter._access_token = "t-xxx" + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = { + "code": 9999, + "msg": "internal error", + } + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + results = await adapter.search("test") + assert results == [] + + @pytest.mark.asyncio + async def test_search_http_error(self, adapter): + import httpx + + adapter._access_token = "t-xxx" + + mock_resp = MagicMock() + mock_resp.status_code = 500 + mock_resp.text = "Internal Server Error" + mock_resp.raise_for_status.side_effect = httpx.HTTPStatusError( + "500", request=MagicMock(), response=mock_resp + ) + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + results = await adapter.search("test") + assert results == [] + + +class TestFeishuKBAdapterHealthCheck: + """FeishuKBAdapter 健康检查""" + + @pytest.mark.asyncio + async def test_health_check_ok(self): + adapter = FeishuKBAdapter(app_id="cli_test", app_secret="secret") + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = {"code": 0, "tenant_access_token": "t-xxx"} + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + assert await adapter.health_check() is True + + @pytest.mark.asyncio + async def test_health_check_failure(self): + adapter = FeishuKBAdapter(app_id="cli_test", app_secret="secret") + + mock_client = AsyncMock() + mock_client.post = AsyncMock(side_effect=Exception("Connection refused")) + adapter._make_client = MagicMock(return_value=mock_client) + + assert await adapter.health_check() is False + + +class TestFeishuKBAdapterListSources: + """FeishuKBAdapter 列出信息源""" + + @pytest.mark.asyncio + async def test_list_sources_success(self): + adapter = FeishuKBAdapter(app_id="cli_test", app_secret="secret") + adapter._access_token = "t-xxx" + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = { + "data": { + "items": [ + {"space_id": "space-1", "name": "产品文档"}, + {"space_id": "space-2", "name": "技术文档"}, + ] + } + } + + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + sources = await adapter.list_sources() + assert len(sources) == 2 + assert sources[0].source_name == "产品文档" + assert sources[1].source_name == "技术文档" + + @pytest.mark.asyncio + async def test_list_sources_not_authenticated(self): + adapter = FeishuKBAdapter(app_id="cli_test", app_secret="secret") + adapter._access_token = None + adapter._get_access_token = AsyncMock(return_value=None) + + sources = await adapter.list_sources() + assert len(sources) == 1 + assert sources[0].source_id.startswith("feishu-") + + +class TestFeishuKBAdapterGetDocument: + """FeishuKBAdapter 获取文档""" + + @pytest.mark.asyncio + async def test_get_document_success(self): + adapter = FeishuKBAdapter(app_id="cli_test", app_secret="secret") + adapter._access_token = "t-xxx" + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = { + "code": 0, + "data": { + "node": { + "title": "测试文档", + "content": "文档内容", + "space_id": "space-1", + "obj_type": "doc", + } + }, + } + + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + doc = await adapter.get_document("wikcnxxx") + assert doc is not None + assert doc.doc_id == "wikcnxxx" + assert doc.title == "测试文档" + assert doc.content == "文档内容" + + @pytest.mark.asyncio + async def test_get_document_not_found(self): + adapter = FeishuKBAdapter(app_id="cli_test", app_secret="secret") + adapter._access_token = "t-xxx" + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = {"code": 10004, "msg": "node not found"} + + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + doc = await adapter.get_document("nonexistent") + assert doc is None + + +# --------------------------------------------------------------------------- +# ConfluenceAdapter tests +# --------------------------------------------------------------------------- + + +class TestConfluenceAdapterInit: + """ConfluenceAdapter 初始化""" + + def test_basic_init(self): + adapter = ConfluenceAdapter( + base_url="https://test.atlassian.net/wiki", + username="user@test.com", + api_token="token", + ) + assert adapter._base_url == "https://test.atlassian.net/wiki" + assert adapter._username == "user@test.com" + assert adapter._api_token == "token" + assert adapter._space_keys == [] + assert adapter._source_type == "confluence" + + def test_init_with_spaces(self): + adapter = ConfluenceAdapter( + base_url="https://test.atlassian.net/wiki/", + username="user@test.com", + api_token="token", + space_keys=["DEV", "DOC"], + timeout=60, + ) + assert adapter._base_url == "https://test.atlassian.net/wiki" + assert adapter._space_keys == ["DEV", "DOC"] + + +class TestConfluenceAdapterAuth: + """ConfluenceAdapter 认证""" + + @pytest.mark.asyncio + async def test_authenticate_success(self): + adapter = ConfluenceAdapter( + base_url="https://test.atlassian.net/wiki", + username="user@test.com", + api_token="token", + ) + + mock_resp = MagicMock() + mock_resp.status_code = 200 + + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + result = await adapter.authenticate() + assert result is True + + @pytest.mark.asyncio + async def test_authenticate_failure(self): + adapter = ConfluenceAdapter( + base_url="https://test.atlassian.net/wiki", + username="user@test.com", + api_token="wrong_token", + ) + + mock_resp = MagicMock() + mock_resp.status_code = 401 + + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + result = await adapter.authenticate() + assert result is False + + +class TestConfluenceAdapterSearch: + """ConfluenceAdapter 检索""" + + @pytest.fixture + def adapter(self): + return ConfluenceAdapter( + base_url="https://test.atlassian.net/wiki", + username="user@test.com", + api_token="token", + space_keys=["DEV"], + ) + + @pytest.mark.asyncio + async def test_search_success(self, adapter): + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = { + "results": [ + { + "id": "12345", + "title": "Confluence 页面", + "type": "page", + "status": "current", + "space": {"key": "DEV"}, + "body": { + "storage": { + "value": "

这是 Confluence 页面内容

" + } + }, + }, + ], + } + + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + results = await adapter.search("Confluence", top_k=5) + + assert len(results) == 1 + assert "Confluence 页面内容" in results[0].content + assert results[0].source_id.startswith("confluence-") + assert results[0].source_name == "Confluence" + assert results[0].doc_id == "12345" + assert results[0].title == "Confluence 页面" + + @pytest.mark.asyncio + async def test_search_with_space_filter(self, adapter): + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = {"results": []} + + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + await adapter.search("test") + + # Verify CQL includes space filter + call_args = mock_client.get.call_args + params = call_args[1].get("params", call_args[0][1] if len(call_args[0]) > 1 else {}) + cql = params.get("cql", "") + assert 'space = "DEV"' in cql + + @pytest.mark.asyncio + async def test_search_http_error(self, adapter): + import httpx + + mock_resp = MagicMock() + mock_resp.status_code = 500 + mock_resp.text = "Internal Server Error" + mock_resp.raise_for_status.side_effect = httpx.HTTPStatusError( + "500", request=MagicMock(), response=mock_resp + ) + + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + results = await adapter.search("test") + assert results == [] + + +class TestConfluenceAdapterHealthCheck: + """ConfluenceAdapter 健康检查""" + + @pytest.mark.asyncio + async def test_health_check_ok(self): + adapter = ConfluenceAdapter( + base_url="https://test.atlassian.net/wiki", + username="user@test.com", + api_token="token", + ) + + mock_resp = MagicMock() + mock_resp.status_code = 200 + + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + assert await adapter.health_check() is True + + @pytest.mark.asyncio + async def test_health_check_failure(self): + adapter = ConfluenceAdapter( + base_url="https://test.atlassian.net/wiki", + username="user@test.com", + api_token="token", + ) + + mock_client = AsyncMock() + mock_client.get = AsyncMock(side_effect=Exception("Connection refused")) + adapter._make_client = MagicMock(return_value=mock_client) + + assert await adapter.health_check() is False + + +class TestConfluenceAdapterGetDocument: + """ConfluenceAdapter 获取文档""" + + @pytest.mark.asyncio + async def test_get_document_success(self): + adapter = ConfluenceAdapter( + base_url="https://test.atlassian.net/wiki", + username="user@test.com", + api_token="token", + ) + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = { + "id": "12345", + "title": "测试页面", + "type": "page", + "space": {"key": "DEV"}, + "body": { + "storage": { + "value": "

页面内容

" + } + }, + "version": {"number": 3}, + } + + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + doc = await adapter.get_document("12345") + assert doc is not None + assert doc.doc_id == "12345" + assert doc.title == "测试页面" + assert "页面内容" in doc.content + assert doc.metadata["version"] == 3 + + +class TestConfluenceAdapterListSources: + """ConfluenceAdapter 列出信息源""" + + @pytest.mark.asyncio + async def test_list_sources_success(self): + adapter = ConfluenceAdapter( + base_url="https://test.atlassian.net/wiki", + username="user@test.com", + api_token="token", + ) + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = { + "results": [ + {"key": "DEV", "name": "Development"}, + {"key": "DOC", "name": "Documentation"}, + ] + } + + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + sources = await adapter.list_sources() + assert len(sources) == 2 + assert sources[0].source_name == "Development" + assert sources[1].source_name == "Documentation" + + +# --------------------------------------------------------------------------- +# GenericHTTPAdapter tests +# --------------------------------------------------------------------------- + + +class TestGenericHTTPAdapterInit: + """GenericHTTPAdapter 初始化""" + + def test_basic_init(self): + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + ) + assert adapter._endpoint_url == "http://localhost:8000/api/kb" + assert adapter._auth_config == {} + assert adapter._extra_headers == {} + assert adapter._source_type == "generic_http" + + def test_init_with_auth_bearer(self): + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb/", + auth_config={"type": "bearer", "token": "sk-test"}, + headers={"X-Custom": "value"}, + source_id="my-kb", + source_name="My KB", + timeout=60, + ) + assert adapter._endpoint_url == "http://localhost:8000/api/kb" + assert adapter._auth_config["type"] == "bearer" + assert adapter._extra_headers == {"X-Custom": "value"} + assert adapter._source_id == "my-kb" + assert adapter._source_name == "My KB" + + def test_client_bearer_auth_header(self): + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + auth_config={"type": "bearer", "token": "sk-test"}, + ) + client = adapter._make_client() + assert "Bearer sk-test" in str(client.headers.get("Authorization", "")) + + def test_client_basic_auth_header(self): + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + auth_config={"type": "basic", "username": "user", "password": "pass"}, + ) + client = adapter._make_client() + auth_header = str(client.headers.get("Authorization", "")) + assert auth_header.startswith("Basic ") + + def test_client_api_key_header(self): + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + auth_config={"type": "api_key", "header_name": "X-API-Key", "api_key": "key123"}, + ) + client = adapter._make_client() + assert client.headers.get("X-API-Key") == "key123" + + +class TestGenericHTTPAdapterSearch: + """GenericHTTPAdapter 检索""" + + @pytest.fixture + def adapter(self): + return GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + auth_config={"type": "bearer", "token": "sk-test"}, + ) + + @pytest.mark.asyncio + async def test_search_standard_response(self, adapter): + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = { + "results": [ + { + "content": "HTTP 知识库内容", + "score": 0.92, + "doc_id": "d1", + "title": "文档1", + "metadata": {"page": 1}, + }, + { + "content": "另一条结果", + "score": 0.85, + "doc_id": "d2", + "title": "文档2", + }, + ] + } + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + results = await adapter.search("知识查询", top_k=5) + + assert len(results) == 2 + assert results[0].content == "HTTP 知识库内容" + assert results[0].score == 0.92 + assert results[0].doc_id == "d1" + assert results[0].title == "文档1" + + # Verify payload + call_args = mock_client.post.call_args + assert call_args[0][0] == "/search" + payload = call_args[1]["json"] + assert payload["query"] == "知识查询" + assert payload["top_k"] == 5 + + @pytest.mark.asyncio + async def test_search_list_response(self, adapter): + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = [ + {"content": "直接列表结果", "score": 0.8, "id": "d1"}, + ] + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + results = await adapter.search("test") + assert len(results) == 1 + assert results[0].content == "直接列表结果" + + @pytest.mark.asyncio + async def test_search_http_error(self, adapter): + import httpx + + mock_resp = MagicMock() + mock_resp.status_code = 500 + mock_resp.text = "Internal Server Error" + mock_resp.raise_for_status.side_effect = httpx.HTTPStatusError( + "500", request=MagicMock(), response=mock_resp + ) + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + results = await adapter.search("test") + assert results == [] + + @pytest.mark.asyncio + async def test_search_unexpected_format(self, adapter): + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = {"error": "unexpected"} + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + results = await adapter.search("test") + assert results == [] + + +class TestGenericHTTPAdapterIngest: + """GenericHTTPAdapter 文档写入""" + + @pytest.mark.asyncio + async def test_ingest_success(self): + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + ) + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = {"ids": ["d1", "d2"]} + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + docs = [ + Document(doc_id="d1", content="内容1", title="文档1"), + Document(doc_id="d2", content="内容2", title="文档2"), + ] + ids = await adapter.ingest(docs) + + assert ids == ["d1", "d2"] + call_args = mock_client.post.call_args + assert call_args[0][0] == "/ingest" + payload = call_args[1]["json"] + assert len(payload["documents"]) == 2 + + @pytest.mark.asyncio + async def test_ingest_http_error(self): + import httpx + + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + ) + + mock_resp = MagicMock() + mock_resp.status_code = 500 + mock_resp.text = "Error" + mock_resp.raise_for_status.side_effect = httpx.HTTPStatusError( + "500", request=MagicMock(), response=mock_resp + ) + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + docs = [Document(doc_id="d1", content="test")] + ids = await adapter.ingest(docs) + assert ids == [] + + +class TestGenericHTTPAdapterDeleteById: + """GenericHTTPAdapter 按 ID 删除""" + + @pytest.mark.asyncio + async def test_delete_success(self): + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + ) + + mock_resp = MagicMock() + mock_resp.status_code = 200 + + mock_client = AsyncMock() + mock_client.delete = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + result = await adapter.delete_by_id("d1") + assert result is True + mock_client.delete.assert_called_once_with("/documents/d1") + + @pytest.mark.asyncio + async def test_delete_not_found(self): + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + ) + + mock_resp = MagicMock() + mock_resp.status_code = 404 + + mock_client = AsyncMock() + mock_client.delete = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + result = await adapter.delete_by_id("nonexistent") + assert result is False + + +class TestGenericHTTPAdapterGetDocument: + """GenericHTTPAdapter 获取文档""" + + @pytest.mark.asyncio + async def test_get_document_success(self): + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + ) + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = { + "id": "d1", + "content": "文档内容", + "title": "测试文档", + "metadata": {"page": 1}, + } + + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + doc = await adapter.get_document("d1") + assert doc is not None + assert doc.doc_id == "d1" + assert doc.content == "文档内容" + assert doc.title == "测试文档" + + @pytest.mark.asyncio + async def test_get_document_not_found(self): + import httpx + + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + ) + + mock_resp = MagicMock() + mock_resp.status_code = 404 + mock_resp.raise_for_status.side_effect = httpx.HTTPStatusError( + "404", request=MagicMock(), response=mock_resp + ) + + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + doc = await adapter.get_document("nonexistent") + assert doc is None + + +class TestGenericHTTPAdapterHealthCheck: + """GenericHTTPAdapter 健康检查""" + + @pytest.mark.asyncio + async def test_health_check_ok(self): + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + ) + + mock_resp = MagicMock() + mock_resp.status_code = 200 + + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + assert await adapter.health_check() is True + + @pytest.mark.asyncio + async def test_health_check_fallback_to_root(self): + """health endpoint 不存在时回退到根路径""" + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + ) + + import httpx + + # /health returns 404 + health_resp = MagicMock() + health_resp.status_code = 404 + health_resp.text = "Not Found" + health_resp.raise_for_status.side_effect = httpx.HTTPStatusError( + "404", request=MagicMock(), response=health_resp + ) + + # / returns 200 + root_resp = MagicMock() + root_resp.status_code = 200 + + mock_client = AsyncMock() + mock_client.get = AsyncMock(side_effect=[health_resp, root_resp]) + adapter._make_client = MagicMock(return_value=mock_client) + + assert await adapter.health_check() is True + + @pytest.mark.asyncio + async def test_health_check_connection_error(self): + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + ) + + mock_client = AsyncMock() + mock_client.get = AsyncMock(side_effect=Exception("Connection refused")) + adapter._make_client = MagicMock(return_value=mock_client) + + assert await adapter.health_check() is False + + +class TestGenericHTTPAdapterListSources: + """GenericHTTPAdapter 列出信息源""" + + @pytest.mark.asyncio + async def test_list_sources_success(self): + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + ) + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = [ + {"source_id": "src-1", "source_name": "知识库1", "source_type": "custom"}, + {"source_id": "src-2", "source_name": "知识库2", "document_count": 50}, + ] + + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_resp) + adapter._make_client = MagicMock(return_value=mock_client) + + sources = await adapter.list_sources() + assert len(sources) == 2 + assert sources[0].source_name == "知识库1" + assert sources[1].document_count == 50 + + @pytest.mark.asyncio + async def test_list_sources_endpoint_not_found(self): + """sources endpoint 不存在时返回默认信息源""" + adapter = GenericHTTPAdapter( + endpoint_url="http://localhost:8000/api/kb", + ) + + mock_client = AsyncMock() + mock_client.get = AsyncMock(side_effect=Exception("404")) + adapter._make_client = MagicMock(return_value=mock_client) + + sources = await adapter.list_sources() + assert len(sources) == 1 + assert sources[0].source_type == "generic_http" + + +# --------------------------------------------------------------------------- +# Cross-adapter integration tests +# --------------------------------------------------------------------------- + + +class TestCrossAdapterIntegration: + """跨适配器集成测试 — 验证统一 KnowledgeBase 接口""" + + @pytest.mark.asyncio + async def test_all_adapters_implement_knowledge_base_protocol(self): + """所有适配器都实现 KnowledgeBase 协议""" + adapters = [ + FeishuKBAdapter(app_id="cli_test", app_secret="secret"), + ConfluenceAdapter( + base_url="https://test.atlassian.net/wiki", + username="user@test.com", + api_token="token", + ), + GenericHTTPAdapter(endpoint_url="http://localhost:8000/api/kb"), + ] + for adapter in adapters: + assert isinstance(adapter, KnowledgeBase) + assert hasattr(adapter, "ingest") + assert hasattr(adapter, "query") + assert hasattr(adapter, "delete_by_id") + assert hasattr(adapter, "list_sources") + assert hasattr(adapter, "health_check") + + @pytest.mark.asyncio + async def test_all_adapters_have_search(self): + """所有适配器都有 search 方法(query 的别名)""" + adapters = [ + FeishuKBAdapter(app_id="cli_test", app_secret="secret"), + ConfluenceAdapter( + base_url="https://test.atlassian.net/wiki", + username="user@test.com", + api_token="token", + ), + GenericHTTPAdapter(endpoint_url="http://localhost:8000/api/kb"), + ] + for adapter in adapters: + assert hasattr(adapter, "search") + assert hasattr(adapter, "get_document") + assert hasattr(adapter, "authenticate") + + @pytest.mark.asyncio + async def test_unified_search_returns_query_result(self): + """统一检索接口返回 QueryResult 类型""" + # Feishu + feishu = FeishuKBAdapter(app_id="cli_test", app_secret="secret") + feishu._access_token = "t-xxx" + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.raise_for_status = MagicMock() + mock_resp.json.return_value = { + "code": 0, + "data": { + "items": [ + {"wiki_token": "w1", "title": "飞书文档", "content": "内容", "score": 0.9} + ] + }, + } + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_resp) + feishu._make_client = MagicMock(return_value=mock_client) + + results = await feishu.search("test") + assert all(isinstance(r, QueryResult) for r in results) + + # Confluence + confluence = ConfluenceAdapter( + base_url="https://test.atlassian.net/wiki", + username="user@test.com", + api_token="token", + ) + mock_resp2 = MagicMock() + mock_resp2.status_code = 200 + mock_resp2.raise_for_status = MagicMock() + mock_resp2.json.return_value = { + "results": [ + {"id": "123", "title": "Confluence 页面", "type": "page", + "body": {"storage": {"value": "

内容

"}}, "space": {"key": "DEV"}} + ] + } + mock_client2 = AsyncMock() + mock_client2.get = AsyncMock(return_value=mock_resp2) + confluence._make_client = MagicMock(return_value=mock_client2) + + results = await confluence.search("test") + assert all(isinstance(r, QueryResult) for r in results) + + # GenericHTTP + generic = GenericHTTPAdapter(endpoint_url="http://localhost:8000/api/kb") + mock_resp3 = MagicMock() + mock_resp3.status_code = 200 + mock_resp3.raise_for_status = MagicMock() + mock_resp3.json.return_value = { + "results": [ + {"content": "HTTP 内容", "score": 0.8, "doc_id": "d1", "title": "文档"} + ] + } + mock_client3 = AsyncMock() + mock_client3.post = AsyncMock(return_value=mock_resp3) + generic._make_client = MagicMock(return_value=mock_client3) + + results = await generic.search("test") + assert all(isinstance(r, QueryResult) for r in results) diff --git a/tests/unit/memory/test_document_loader.py b/tests/unit/memory/test_document_loader.py new file mode 100644 index 0000000..bff89c9 --- /dev/null +++ b/tests/unit/memory/test_document_loader.py @@ -0,0 +1,227 @@ +"""DocumentLoader 单元测试 - 多格式文档解析器""" + +import pytest + +from agentkit.memory.document_loader import Document, DocumentLoader, _detect_format + + +class TestDetectFormat: + """格式检测测试""" + + def test_pdf_format(self): + assert _detect_format("report.pdf") == "pdf" + + def test_docx_format(self): + assert _detect_format("document.docx") == "docx" + assert _detect_format("document.doc") == "docx" + + def test_markdown_format(self): + assert _detect_format("readme.md") == "markdown" + assert _detect_format("notes.markdown") == "markdown" + + def test_html_format(self): + assert _detect_format("page.html") == "html" + assert _detect_format("page.htm") == "html" + + def test_text_format(self): + assert _detect_format("data.txt") == "text" + assert _detect_format("data.csv") == "text" + assert _detect_format("data.json") == "text" + + def test_unknown_format_falls_back_to_text(self): + assert _detect_format("data.xyz") == "text" + + +class TestDocument: + """Document 数据类测试""" + + def test_default_metadata(self): + doc = Document(doc_id="1", title="Test", content="Hello") + assert doc.metadata["source"] == "" + assert doc.metadata["format"] == "unknown" + assert doc.metadata["page_count"] == 0 + assert "created_at" in doc.metadata + + def test_custom_metadata(self): + doc = Document( + doc_id="1", + title="Test", + content="Hello", + metadata={"source": "test.pdf", "format": "pdf", "page_count": 5}, + ) + assert doc.metadata["source"] == "test.pdf" + assert doc.metadata["format"] == "pdf" + assert doc.metadata["page_count"] == 5 + + def test_to_dict(self): + doc = Document(doc_id="1", title="Test", content="Hello", metadata={"format": "text"}) + d = doc.to_dict() + assert d["doc_id"] == "1" + assert d["title"] == "Test" + assert d["content"] == "Hello" + assert d["metadata"]["format"] == "text" + + +class TestDocumentLoaderText: + """纯文本解析测试""" + + def test_load_text_bytes(self): + loader = DocumentLoader() + content = "Hello, world!\nThis is a test document.".encode("utf-8") + doc = loader.load_bytes(content, "test.txt") + + assert doc.title == "test" + assert doc.content == "Hello, world!\nThis is a test document." + assert doc.metadata["format"] == "text" + assert doc.metadata["source"] == "test.txt" + assert doc.metadata["parser"] == "text" + assert doc.doc_id # 非空 UUID + + def test_load_text_file(self, tmp_path): + loader = DocumentLoader() + text_file = tmp_path / "sample.txt" + text_file.write_text("Sample text content", encoding="utf-8") + + doc = loader.load(text_file) + assert doc.content == "Sample text content" + assert doc.metadata["format"] == "text" + + def test_load_nonexistent_file(self): + loader = DocumentLoader() + with pytest.raises(FileNotFoundError): + loader.load("/nonexistent/path/file.txt") + + +class TestDocumentLoaderMarkdown: + """Markdown 解析测试""" + + def test_load_markdown_bytes(self): + loader = DocumentLoader() + md_content = """# Project Title + +## Introduction + +This is the introduction section. + +## Details + +Some details here. +""" + doc = loader.load_bytes(md_content.encode("utf-8"), "readme.md") + + assert doc.metadata["format"] == "markdown" + assert doc.metadata["title"] == "Project Title" + assert "Introduction" in doc.content + assert "Details" in doc.content + + def test_markdown_without_title(self): + loader = DocumentLoader() + md_content = "Just some text without a heading." + doc = loader.load_bytes(md_content.encode("utf-8"), "notes.md") + + assert doc.metadata["format"] == "markdown" + assert doc.content == "Just some text without a heading." + + +class TestDocumentLoaderHTML: + """HTML 解析测试""" + + def test_load_html_with_beautifulsoup(self): + """测试 BeautifulSoup 解析(如果可用)""" + loader = DocumentLoader() + html_content = """ + +Test Page + + + +

Hello

+

This is a paragraph.

+ +""" + doc = loader.load_bytes(html_content.encode("utf-8"), "page.html") + + assert doc.metadata["format"] == "html" + # BeautifulSoup 应该移除 script/style 标签 + # 如果 BeautifulSoup 不可用,则回退到文本 + if doc.metadata.get("parser") == "beautifulsoup": + assert "Test Page" in doc.metadata.get("title", "") or "Hello" in doc.content + assert "var x" not in doc.content + assert ".cls" not in doc.content + assert "Hello" in doc.content + else: + # 纯文本回退,内容可能包含 HTML 标签 + assert len(doc.content) > 0 + + def test_load_html_fallback_to_text(self): + """即使没有 BeautifulSoup,HTML 也能作为文本加载""" + loader = DocumentLoader() + html_content = "Simple content" + doc = loader.load_bytes(html_content.encode("utf-8"), "page.html") + + assert doc.metadata["format"] == "html" + assert len(doc.content) > 0 + + +class TestDocumentLoaderPDF: + """PDF 解析测试""" + + def test_load_pdf_without_parser(self): + """没有 PDF 解析器时回退到文本""" + loader = DocumentLoader() + # 传入一个非 PDF 二进制内容,模拟解析失败后的回退 + content = b"%PDF-1.4 fake pdf content" + doc = loader.load_bytes(content, "report.pdf") + + assert doc.metadata["format"] == "pdf" + # 即使解析失败,也应该返回文档对象(内容可能为空或乱码) + assert isinstance(doc, Document) + + +class TestDocumentLoaderDocx: + """Word 解析测试""" + + def test_load_docx_without_parser(self): + """没有 python-docx 时回退到文本""" + loader = DocumentLoader() + # 传入一个非 docx 二进制内容 + content = b"PK\x03\x04 fake docx content" + doc = loader.load_bytes(content, "document.docx") + + assert doc.metadata["format"] == "docx" + assert isinstance(doc, Document) + + +class TestDocumentLoaderEdgeCases: + """边界情况测试""" + + def test_empty_content(self): + loader = DocumentLoader() + doc = loader.load_bytes(b"", "empty.txt") + assert doc.content == "" + assert doc.metadata["format"] == "text" + + def test_unicode_content(self): + loader = DocumentLoader() + content = "中文内容测试\n日本語テスト\n한국어 테스트".encode("utf-8") + doc = loader.load_bytes(content, "unicode.txt") + assert "中文内容测试" in doc.content + assert "日本語テスト" in doc.content + + def test_large_content(self): + loader = DocumentLoader() + content = "A" * 1_000_000 # 1MB text + doc = loader.load_bytes(content.encode("utf-8"), "large.txt") + assert len(doc.content) == 1_000_000 + + def test_filename_with_spaces(self): + loader = DocumentLoader() + content = "Test content".encode("utf-8") + doc = loader.load_bytes(content, "my document.txt") + assert doc.title == "my document" + + def test_filename_with_path(self): + loader = DocumentLoader() + content = "Test content".encode("utf-8") + doc = loader.load_bytes(content, "reports/2024/summary.md") + assert doc.metadata["format"] == "markdown" diff --git a/tests/unit/memory/test_local_rag.py b/tests/unit/memory/test_local_rag.py new file mode 100644 index 0000000..8a38fa4 --- /dev/null +++ b/tests/unit/memory/test_local_rag.py @@ -0,0 +1,522 @@ +"""LocalRAGService 单元测试 - 本地文档 RAG 服务 + +使用 InMemoryLocalRAGService 进行测试,无需 pgvector 依赖。 +同时测试分块策略(TextChunker / StructuralChunker)。 +""" + +import pytest + +from agentkit.memory.chunking import Chunk, StructuralChunker, TextChunker +from agentkit.memory.document_loader import Document as LoaderDocument +from agentkit.memory.embedder import MockEmbedder +from agentkit.memory.knowledge_base import Document, KnowledgeBase, QueryResult, SourceInfo +from agentkit.memory.local_rag import InMemoryLocalRAGService + + +# ── Fixtures ────────────────────────────────────────────── + + +@pytest.fixture +def embedder(): + return MockEmbedder(dimension=128) + + +@pytest.fixture +def rag_service(embedder): + return InMemoryLocalRAGService(embedder=embedder, chunk_size=500, chunk_overlap=50) + + +@pytest.fixture +def sample_documents(): + """knowledge_base.Document 格式的测试文档""" + return [ + Document( + doc_id="doc-1", + content="Python 是一种通用编程语言。它支持多种编程范式,包括面向对象、命令式、函数式和过程式编程。Python 的设计哲学强调代码的可读性和简洁性。", + title="Python 入门指南", + source_id="python_intro.txt", + metadata={"source": "python_intro.txt", "format": "text"}, + ), + Document( + doc_id="doc-2", + content="机器学习是人工智能的一个分支,它使计算机系统能够从数据中学习和改进。常见的机器学习算法包括线性回归、决策树、支持向量机和神经网络。", + title="机器学习基础", + source_id="ml_basics.txt", + metadata={"source": "ml_basics.txt", "format": "text"}, + ), + ] + + +@pytest.fixture +def markdown_document(): + return Document( + doc_id="doc-md-1", + content="""# API 文档 + +## 认证 + +所有 API 请求需要 Bearer Token 认证。请在请求头中添加 Authorization 字段。 + +## 用户接口 + +### 获取用户信息 + +GET /api/users/{id} + +返回指定用户的详细信息。 + +### 创建用户 + +POST /api/users + +创建一个新用户。 + +## 数据接口 + +### 查询数据 + +POST /api/data/query + +根据条件查询数据。 +""", + title="API 文档", + source_id="api_doc.md", + metadata={"source": "api_doc.md", "format": "markdown"}, + ) + + +# ── TextChunker 测试 ────────────────────────────────────── + + +class TestTextChunker: + """TextChunker 单元测试""" + + def test_chunk_short_text(self): + chunker = TextChunker(chunk_size=1000, chunk_overlap=100) + chunks = chunker.chunk("Short text", source_doc_id="doc-1") + + assert len(chunks) == 1 + assert chunks[0].content == "Short text" + assert chunks[0].metadata["source_doc"] == "doc-1" + assert chunks[0].metadata["position"] == 0 + + def test_chunk_empty_text(self): + chunker = TextChunker(chunk_size=1000, chunk_overlap=100) + chunks = chunker.chunk("", source_doc_id="doc-1") + assert len(chunks) == 0 + + def test_chunk_whitespace_only(self): + chunker = TextChunker(chunk_size=1000, chunk_overlap=100) + chunks = chunker.chunk(" \n\n \t ", source_doc_id="doc-1") + assert len(chunks) == 0 + + def test_chunk_long_text(self): + chunker = TextChunker(chunk_size=100, chunk_overlap=20) + text = "A" * 300 + chunks = chunker.chunk(text, source_doc_id="doc-1") + + assert len(chunks) >= 2 + # 每个块不超过 chunk_size(允许少量超出用于句子边界) + for chunk in chunks: + assert len(chunk.content) <= 150 # 允许一些余量 + + def test_chunk_preserves_metadata(self): + chunker = TextChunker(chunk_size=1000, chunk_overlap=100) + chunks = chunker.chunk( + "Some content", + source_doc_id="doc-1", + metadata={"format": "pdf", "page_count": 5}, + ) + + assert len(chunks) == 1 + assert chunks[0].metadata["format"] == "pdf" + assert chunks[0].metadata["page_count"] == 5 + assert chunks[0].metadata["source_doc"] == "doc-1" + + def test_chunk_with_multiple_paragraphs(self): + chunker = TextChunker(chunk_size=200, chunk_overlap=20, separator="\n\n") + text = "第一段内容,包含一些文字。\n\n第二段内容,也有一些文字。\n\n第三段内容,同样有文字。" + chunks = chunker.chunk(text, source_doc_id="doc-1") + + assert len(chunks) >= 1 + for chunk in chunks: + assert len(chunk.content) > 0 + + def test_invalid_overlap(self): + with pytest.raises(ValueError): + TextChunker(chunk_size=100, chunk_overlap=100) + + def test_chunk_with_separator(self): + chunker = TextChunker(chunk_size=200, chunk_overlap=20, separator="\n\n") + text = "第一段内容\n\n第二段内容\n\n第三段内容" + chunks = chunker.chunk(text, source_doc_id="doc-1") + + assert len(chunks) >= 1 + for chunk in chunks: + assert len(chunk.content) > 0 + + +class TestStructuralChunker: + """StructuralChunker 单元测试""" + + def test_chunk_markdown_by_headings(self): + chunker = StructuralChunker(chunk_size=1000, chunk_overlap=50) + md = """# Title + +## Section A + +Content for section A. + +## Section B + +Content for section B. + +## Section C + +Content for section C.""" + chunks = chunker.chunk(md, source_doc_id="doc-1") + + assert len(chunks) >= 3 + # 每个块应该有标题元数据 + headings = [c.metadata.get("heading") for c in chunks] + assert "Section A" in headings + assert "Section B" in headings + assert "Section C" in headings + + def test_chunk_markdown_no_headings(self): + chunker = StructuralChunker(chunk_size=1000, chunk_overlap=50) + md = "Just some text without any headings." + chunks = chunker.chunk(md, source_doc_id="doc-1") + + assert len(chunks) == 1 + assert chunks[0].content == "Just some text without any headings." + + def test_chunk_empty_text(self): + chunker = StructuralChunker(chunk_size=1000, chunk_overlap=50) + chunks = chunker.chunk("", source_doc_id="doc-1") + assert len(chunks) == 0 + + def test_chunk_large_section_falls_back_to_text_chunker(self): + chunker = StructuralChunker(chunk_size=100, chunk_overlap=20) + md = """# Large Section + +""" + "A" * 300 + chunks = chunker.chunk(md, source_doc_id="doc-1") + + # 大段应被 TextChunker 进一步切分 + assert len(chunks) >= 2 + for chunk in chunks: + assert chunk.metadata.get("heading") == "Large Section" + + def test_heading_levels(self): + chunker = StructuralChunker(chunk_size=1000, heading_levels=2) + md = """# H1 + +Content 1. + +## H2 + +Content 2. + +### H3 + +This should be part of H2 section since heading_levels=2. +""" + chunks = chunker.chunk(md, source_doc_id="doc-1") + # H3 不应该作为独立标题分割 + assert len(chunks) >= 2 + + +# ── Chunk 数据类测试 ────────────────────────────────────── + + +class TestChunk: + """Chunk 数据类测试""" + + def test_default_metadata(self): + chunk = Chunk(chunk_id="c1", content="test") + assert chunk.metadata["source_doc"] == "" + assert chunk.metadata["position"] == 0 + + def test_to_dict(self): + chunk = Chunk( + chunk_id="c1", + content="test content", + metadata={"source_doc": "doc-1", "position": 0}, + ) + d = chunk.to_dict() + assert d["chunk_id"] == "c1" + assert d["content"] == "test content" + assert d["metadata"]["source_doc"] == "doc-1" + + +# ── InMemoryLocalRAGService 测试 ────────────────────────── + + +class TestInMemoryLocalRAGService: + """InMemoryLocalRAGService 单元测试""" + + @pytest.mark.asyncio + async def test_ingest_documents(self, rag_service, sample_documents): + ids = await rag_service.ingest(sample_documents) + + assert len(ids) == 2 + assert "doc-1" in ids + assert "doc-2" in ids + + @pytest.mark.asyncio + async def test_query_after_ingest(self, rag_service, sample_documents): + await rag_service.ingest(sample_documents) + + results = await rag_service.query("编程语言", top_k=2) + + assert len(results) >= 1 + assert all(isinstance(r, QueryResult) for r in results) + # 结果应该包含相关内容 + assert any("Python" in r.content or "编程" in r.content for r in results) + + @pytest.mark.asyncio + async def test_query_returns_source_info(self, rag_service, sample_documents): + await rag_service.ingest(sample_documents) + + results = await rag_service.query("机器学习", top_k=5) + + assert len(results) >= 1 + for r in results: + assert r.source_id != "" + assert r.source_name != "" + + @pytest.mark.asyncio + async def test_query_no_results_when_empty(self, rag_service): + results = await rag_service.query("anything", top_k=5) + assert len(results) == 0 + + @pytest.mark.asyncio + async def test_delete_by_id(self, rag_service, sample_documents): + await rag_service.ingest(sample_documents) + + deleted = await rag_service.delete_by_id("doc-1") + assert deleted is True + + # 删除后查询不应返回该文档的内容 + results = await rag_service.query("Python", top_k=5) + assert all(r.source_id != "doc-1" for r in results) + + @pytest.mark.asyncio + async def test_delete_nonexistent_id(self, rag_service): + deleted = await rag_service.delete_by_id("nonexistent") + assert deleted is False + + @pytest.mark.asyncio + async def test_list_sources(self, rag_service, sample_documents): + await rag_service.ingest(sample_documents) + + sources = await rag_service.list_sources() + + assert len(sources) == 2 + source_ids = {s.source_id for s in sources} + assert "doc-1" in source_ids + assert "doc-2" in source_ids + + for s in sources: + assert isinstance(s, SourceInfo) + assert s.source_name != "" + assert s.document_count > 0 + + @pytest.mark.asyncio + async def test_list_sources_empty(self, rag_service): + sources = await rag_service.list_sources() + assert len(sources) == 0 + + @pytest.mark.asyncio + async def test_health_check(self, rag_service): + assert await rag_service.health_check() is True + + @pytest.mark.asyncio + async def test_ingest_markdown_with_structural_chunking(self, rag_service, markdown_document): + ids = await rag_service.ingest([markdown_document]) + + assert len(ids) == 1 + sources = await rag_service.list_sources() + assert len(sources) == 1 + assert sources[0].source_type == "markdown" + + @pytest.mark.asyncio + async def test_query_markdown_by_section(self, rag_service, markdown_document): + await rag_service.ingest([markdown_document]) + + results = await rag_service.query("认证", top_k=3) + + # MockEmbedder 基于文本哈希,语义相关性不保证, + # 但应至少返回结果(因为文档已被摄取) + assert len(results) >= 0 # 可能因阈值过滤无结果 + # 使用与文档内容更相似的查询词来验证检索 + results = await rag_service.query("API 文档 认证", top_k=3) + assert len(results) >= 1 + + @pytest.mark.asyncio + async def test_ingest_empty_document(self, rag_service): + doc = Document( + doc_id="empty-doc", + content="", + title="Empty", + source_id="empty.txt", + metadata={"source": "empty.txt", "format": "text"}, + ) + ids = await rag_service.ingest([doc]) + + # 空文档应该被跳过(没有块生成) + assert len(ids) == 1 # doc_id 仍然返回 + sources = await rag_service.list_sources() + assert len(sources) == 1 + assert sources[0].document_count == 0 + + @pytest.mark.asyncio + async def test_ingest_large_document_chunking(self, embedder): + """大文件分块 → 块大小在配置范围内""" + rag = InMemoryLocalRAGService(embedder=embedder, chunk_size=200, chunk_overlap=20) + + large_content = "这是一段很长的文本。" * 200 # ~2000 字符 + doc = Document( + doc_id="large-doc", + content=large_content, + title="Large Document", + source_id="large.txt", + metadata={"source": "large.txt", "format": "text"}, + ) + ids = await rag.ingest([doc]) + + assert len(ids) == 1 + sources = await rag.list_sources() + assert sources[0].document_count > 1 # 应该被分成多个块 + + @pytest.mark.asyncio + async def test_query_result_has_score(self, rag_service, sample_documents): + await rag_service.ingest(sample_documents) + + results = await rag_service.query("编程", top_k=5) + + for r in results: + assert 0.0 <= r.score <= 1.0 + + @pytest.mark.asyncio + async def test_ingest_loader_document(self, rag_service): + """测试传入 document_loader.Document 时的自动转换""" + loader_doc = LoaderDocument( + doc_id="loader-doc-1", + title="Test Loader Doc", + content="This is content from document_loader.", + metadata={"source": "test.txt", "format": "text"}, + ) + ids = await rag_service.ingest([loader_doc]) + + assert len(ids) == 1 + results = await rag_service.query("content", top_k=3) + assert len(results) >= 1 + + @pytest.mark.asyncio + async def test_multiple_ingest_same_doc_id(self, rag_service): + """重复摄取相同 doc_id 的文档""" + doc1 = Document( + doc_id="same-id", + content="First version content", + title="Version 1", + source_id="v1.txt", + metadata={"source": "v1.txt", "format": "text"}, + ) + doc2 = Document( + doc_id="same-id", + content="Second version content with more text", + title="Version 2", + source_id="v2.txt", + metadata={"source": "v2.txt", "format": "text"}, + ) + + await rag_service.ingest([doc1]) + await rag_service.ingest([doc2]) + + # 第二次摄取会覆盖(内存实现中 doc_id 相同会覆盖) + sources = await rag_service.list_sources() + source_ids = [s.source_id for s in sources] + assert "same-id" in source_ids + + +# ── KnowledgeBase 协议测试 ──────────────────────────────── + + +class TestKnowledgeBaseProtocol: + """KnowledgeBase 协议兼容性测试""" + + @pytest.mark.asyncio + async def test_inmemory_service_implements_protocol(self, rag_service): + """InMemoryLocalRAGService 应该满足 KnowledgeBase 协议""" + assert isinstance(rag_service, KnowledgeBase) + + @pytest.mark.asyncio + async def test_protocol_methods_exist(self, rag_service): + """验证所有协议方法都存在""" + assert hasattr(rag_service, "ingest") + assert hasattr(rag_service, "query") + assert hasattr(rag_service, "delete_by_id") + assert hasattr(rag_service, "list_sources") + assert hasattr(rag_service, "health_check") + + # 验证方法可调用 + assert callable(rag_service.ingest) + assert callable(rag_service.query) + assert callable(rag_service.delete_by_id) + assert callable(rag_service.list_sources) + assert callable(rag_service.health_check) + + +# ── QueryResult / SourceInfo 测试 ───────────────────────── + + +class TestQueryResult: + """QueryResult 数据类测试""" + + def test_creation(self): + result = QueryResult( + content="test content", + source_id="doc-1", + source_name="Test Doc", + score=0.95, + ) + assert result.content == "test content" + assert result.source_id == "doc-1" + assert result.source_name == "Test Doc" + assert result.score == 0.95 + + def test_with_optional_fields(self): + result = QueryResult( + content="test content", + source_id="doc-1", + source_name="Test Doc", + score=0.95, + metadata={"position": 0}, + doc_id="doc-1", + title="Test Doc", + ) + assert result.doc_id == "doc-1" + assert result.title == "Test Doc" + assert result.metadata["position"] == 0 + + +class TestSourceInfo: + """SourceInfo 数据类测试""" + + def test_creation(self): + from datetime import datetime, timezone + + now = datetime.now(timezone.utc) + info = SourceInfo( + source_id="doc-1", + source_name="Test", + source_type="local", + document_count=5, + last_updated=now, + ) + assert info.source_id == "doc-1" + assert info.source_name == "Test" + assert info.source_type == "local" + assert info.document_count == 5 diff --git a/tests/unit/memory/test_multi_source_rag.py b/tests/unit/memory/test_multi_source_rag.py new file mode 100644 index 0000000..858af93 --- /dev/null +++ b/tests/unit/memory/test_multi_source_rag.py @@ -0,0 +1,610 @@ +"""MultiSourceRAG 单元测试 - 多源混合检索 + +测试场景: +- 指定单个信息源 → 仅从该源检索 +- 指定多个信息源 → 并行检索,结果融合排序 +- 不指定信息源 → 从所有可用源检索 +- 来源追溯 → 每个结果包含来源信息 +- AE4: 指定"合规文档库"和"法务知识库" → 仅从这两个源检索 +""" + +import pytest + +from agentkit.memory.embedder import MockEmbedder +from agentkit.memory.knowledge_base import Document, KnowledgeBase, QueryResult, SourceInfo +from agentkit.memory.local_rag import InMemoryLocalRAGService +from agentkit.memory.multi_source_retriever import MultiSourceRetriever +from agentkit.memory.retriever import MemoryRetriever + + +# ── Fixtures ────────────────────────────────────────────── + + +@pytest.fixture +def embedder(): + return MockEmbedder(dimension=128) + + +@pytest.fixture +def local_rag(embedder): + """本地合规文档库""" + return InMemoryLocalRAGService(embedder=embedder, chunk_size=500, chunk_overlap=50) + + +@pytest.fixture +def legal_rag(embedder): + """法务知识库""" + return InMemoryLocalRAGService(embedder=embedder, chunk_size=500, chunk_overlap=50) + + +@pytest.fixture +def tech_rag(embedder): + """技术文档库""" + return InMemoryLocalRAGService(embedder=embedder, chunk_size=500, chunk_overlap=50) + + +@pytest.fixture +def compliance_docs(): + """合规文档""" + return [ + Document( + doc_id="compliance-1", + content="数据保护合规要求:所有用户数据必须加密存储,访问需经授权审批。", + title="数据保护合规指南", + source_id="compliance_data_protection", + metadata={"source": "compliance_data_protection", "format": "text"}, + ), + Document( + doc_id="compliance-2", + content="跨境数据传输需遵守 GDPR 和中国网络安全法的相关规定。", + title="跨境数据传输合规", + source_id="compliance_cross_border", + metadata={"source": "compliance_cross_border", "format": "text"}, + ), + ] + + +@pytest.fixture +def legal_docs(): + """法务文档""" + return [ + Document( + doc_id="legal-1", + content="合同审查要点:注意违约责任条款、知识产权归属和保密义务。", + title="合同审查指南", + source_id="legal_contract_review", + metadata={"source": "legal_contract_review", "format": "text"}, + ), + Document( + doc_id="legal-2", + content="劳动法规定:员工加班需支付加班费,标准为平时工资的1.5倍至3倍。", + title="劳动法要点", + source_id="legal_labor_law", + metadata={"source": "legal_labor_law", "format": "text"}, + ), + ] + + +@pytest.fixture +def tech_docs(): + """技术文档""" + return [ + Document( + doc_id="tech-1", + content="API 网关配置:限流策略为每分钟 1000 次请求,超时设置 30 秒。", + title="API 网关配置手册", + source_id="tech_api_gateway", + metadata={"source": "tech_api_gateway", "format": "text"}, + ), + ] + + +# ── MultiSourceRetriever 核心测试 ───────────────────────── + + +class TestMultiSourceRetrieverBasic: + """MultiSourceRetriever 基础功能测试""" + + def test_register_source(self, local_rag, legal_rag): + retriever = MultiSourceRetriever() + retriever.register_source("local:合规文档", local_rag) + retriever.register_source("法务知识库", legal_rag) + + names = retriever.get_source_names() + assert "local:合规文档" in names + assert "法务知识库" in names + + def test_register_source_via_constructor(self, local_rag, legal_rag): + retriever = MultiSourceRetriever( + sources={"local:合规文档": local_rag, "法务知识库": legal_rag} + ) + + names = retriever.get_source_names() + assert len(names) == 2 + + def test_unregister_source(self, local_rag, legal_rag): + retriever = MultiSourceRetriever( + sources={"local:合规文档": local_rag, "法务知识库": legal_rag} + ) + + result = retriever.unregister_source("local:合规文档") + assert result is True + assert "local:合规文档" not in retriever.get_source_names() + + def test_unregister_nonexistent_source(self, local_rag): + retriever = MultiSourceRetriever(sources={"local:合规文档": local_rag}) + + result = retriever.unregister_source("不存在") + assert result is False + + @pytest.mark.asyncio + async def test_list_all_sources(self, local_rag, legal_rag, compliance_docs, legal_docs): + await local_rag.ingest(compliance_docs) + await legal_rag.ingest(legal_docs) + + retriever = MultiSourceRetriever( + sources={"local:合规文档": local_rag, "法务知识库": legal_rag} + ) + + sources = await retriever.list_all_sources() + assert "local:合规文档" in sources + assert "法务知识库" in sources + for name, info in sources.items(): + assert isinstance(info, SourceInfo) + + +class TestMultiSourceRetrieverSearch: + """MultiSourceRetriever 检索功能测试""" + + @pytest.mark.asyncio + async def test_search_single_source( + self, local_rag, legal_rag, compliance_docs, legal_docs + ): + """指定单个信息源 → 仅从该源检索""" + await local_rag.ingest(compliance_docs) + await legal_rag.ingest(legal_docs) + + retriever = MultiSourceRetriever( + sources={"local:合规文档": local_rag, "法务知识库": legal_rag} + ) + + # 仅从合规文档库检索 + results = await retriever.search("合规", top_k=5, sources=["local:合规文档"]) + + # 所有结果应来自合规文档库 + for r in results: + assert r.source_name == "local:合规文档" + + @pytest.mark.asyncio + async def test_search_multiple_sources( + self, local_rag, legal_rag, compliance_docs, legal_docs + ): + """指定多个信息源 → 并行检索,结果融合排序""" + await local_rag.ingest(compliance_docs) + await legal_rag.ingest(legal_docs) + + retriever = MultiSourceRetriever( + sources={"local:合规文档": local_rag, "法务知识库": legal_rag} + ) + + results = await retriever.search( + "合规 法务", top_k=10, sources=["local:合规文档", "法务知识库"] + ) + + # 结果应来自两个源 + source_names = {r.source_name for r in results} + assert source_names.issubset({"local:合规文档", "法务知识库"}) + + @pytest.mark.asyncio + async def test_search_all_sources_when_none_specified( + self, local_rag, legal_rag, tech_rag, compliance_docs, legal_docs, tech_docs + ): + """不指定信息源 → 从所有可用源检索""" + await local_rag.ingest(compliance_docs) + await legal_rag.ingest(legal_docs) + await tech_rag.ingest(tech_docs) + + retriever = MultiSourceRetriever( + sources={ + "local:合规文档": local_rag, + "法务知识库": legal_rag, + "技术文档库": tech_rag, + } + ) + + results = await retriever.search("合规 法务 技术", top_k=10) + + # 结果应来自所有三个源 + source_names = {r.source_name for r in results} + assert len(source_names) >= 1 # 至少有一个源返回结果 + + @pytest.mark.asyncio + async def test_search_no_sources_registered(self): + """无信息源注册时返回空结果""" + retriever = MultiSourceRetriever() + + results = await retriever.search("anything", top_k=5) + assert len(results) == 0 + + @pytest.mark.asyncio + async def test_search_nonexistent_source(self, local_rag, compliance_docs): + """指定不存在的源 → 跳过,返回空结果""" + await local_rag.ingest(compliance_docs) + + retriever = MultiSourceRetriever(sources={"local:合规文档": local_rag}) + + results = await retriever.search("合规", top_k=5, sources=["不存在的源"]) + assert len(results) == 0 + + @pytest.mark.asyncio + async def test_search_with_weights( + self, local_rag, legal_rag, compliance_docs, legal_docs + ): + """带权重检索 → 特定源分数被调整""" + await local_rag.ingest(compliance_docs) + await legal_rag.ingest(legal_docs) + + retriever = MultiSourceRetriever( + sources={"local:合规文档": local_rag, "法务知识库": legal_rag} + ) + + # 先不带权重检索 + results_no_weight = await retriever.search( + "合规", top_k=10, sources=["local:合规文档", "法务知识库"] + ) + + # 带权重检索:提升合规文档库 + results_with_weight = await retriever.search( + "合规", + top_k=10, + sources=["local:合规文档", "法务知识库"], + weights={"local:合规文档": 2.0}, + ) + + # 有权重时合规文档库的分数应更高 + compliance_scores_weighted = [ + r.score for r in results_with_weight if r.source_name == "local:合规文档" + ] + compliance_scores_unweighted = [ + r.score for r in results_no_weight if r.source_name == "local:合规文档" + ] + + if compliance_scores_weighted and compliance_scores_unweighted: + assert max(compliance_scores_weighted) >= max(compliance_scores_unweighted) + + +class TestMultiSourceRetrieverSourceTracing: + """来源追溯测试""" + + @pytest.mark.asyncio + async def test_result_contains_source_info( + self, local_rag, legal_rag, compliance_docs, legal_docs + ): + """每个检索结果包含来源追溯信息""" + await local_rag.ingest(compliance_docs) + await legal_rag.ingest(legal_docs) + + retriever = MultiSourceRetriever( + sources={"local:合规文档": local_rag, "法务知识库": legal_rag} + ) + + results = await retriever.search("合规", top_k=5) + + for r in results: + # source_id 应非空 + assert r.source_id != "" + # source_name 应为注册的源名称 + assert r.source_name in ("local:合规文档", "法务知识库") + # title 应非空 + assert r.title != "" + + @pytest.mark.asyncio + async def test_result_contains_document_title( + self, local_rag, compliance_docs + ): + """检索结果包含文档标题""" + await local_rag.ingest(compliance_docs) + + retriever = MultiSourceRetriever(sources={"local:合规文档": local_rag}) + + results = await retriever.search("数据保护", top_k=5) + + for r in results: + assert r.title != "" + assert r.doc_id != "" + + +class TestMultiSourceRetrieverDedup: + """去重测试""" + + @pytest.mark.asyncio + async def test_deduplicate_identical_content( + self, local_rag, legal_rag, embedder + ): + """相同内容从不同源返回时去重,保留高分""" + # 两个源包含相同内容 + same_doc = Document( + doc_id="same-doc", + content="这是一段完全相同的内容用于测试去重功能。", + title="重复文档", + source_id="same_source", + metadata={"source": "same_source", "format": "text"}, + ) + await local_rag.ingest([same_doc]) + await legal_rag.ingest([same_doc]) + + retriever = MultiSourceRetriever( + sources={"local:合规文档": local_rag, "法务知识库": legal_rag} + ) + + results = await retriever.search("去重", top_k=10) + + # 相同内容应去重,只保留一个 + content_counts: dict[str, int] = {} + for r in results: + content_counts[r.content] = content_counts.get(r.content, 0) + 1 + + for content, count in content_counts.items(): + assert count == 1, f"内容 '{content[:30]}...' 出现了 {count} 次,应去重为 1 次" + + +# ── AE4: 合规文档库 + 法务知识库指定检索 ────────────────── + + +class TestAE4ComplianceAndLegalSearch: + """AE4 场景:指定"合规文档库"和"法务知识库" → 仅从这两个源检索""" + + @pytest.mark.asyncio + async def test_search_compliance_and_legal_only( + self, local_rag, legal_rag, tech_rag, compliance_docs, legal_docs, tech_docs + ): + """指定合规和法务源 → 不从技术文档库检索""" + await local_rag.ingest(compliance_docs) + await legal_rag.ingest(legal_docs) + await tech_rag.ingest(tech_docs) + + retriever = MultiSourceRetriever( + sources={ + "合规文档库": local_rag, + "法务知识库": legal_rag, + "技术文档库": tech_rag, + } + ) + + results = await retriever.search( + "合规 法务", top_k=10, sources=["合规文档库", "法务知识库"] + ) + + # 结果不应来自技术文档库 + for r in results: + assert r.source_name != "技术文档库" + assert r.source_name in ("合规文档库", "法务知识库") + + @pytest.mark.asyncio + async def test_search_compliance_and_legal_results_merged( + self, local_rag, legal_rag, compliance_docs, legal_docs + ): + """合规和法务源的结果应合并排序""" + await local_rag.ingest(compliance_docs) + await legal_rag.ingest(legal_docs) + + retriever = MultiSourceRetriever( + sources={"合规文档库": local_rag, "法务知识库": legal_rag} + ) + + results = await retriever.search( + "合规 法务", top_k=10, sources=["合规文档库", "法务知识库"] + ) + + # 应有来自两个源的结果 + source_names = {r.source_name for r in results} + assert len(source_names) >= 1 + + # 结果应按 score 降序排列 + for i in range(len(results) - 1): + assert results[i].score >= results[i + 1].score + + +# ── MemoryRetriever 集成测试 ────────────────────────────── + + +class TestMemoryRetrieverIntegration: + """MemoryRetriever 与 MultiSourceRetriever 集成测试""" + + @pytest.mark.asyncio + async def test_retrieve_with_sources_parameter( + self, local_rag, legal_rag, compliance_docs, legal_docs + ): + """MemoryRetriever.retrieve(sources=...) 委托给 MultiSourceRetriever""" + await local_rag.ingest(compliance_docs) + await legal_rag.ingest(legal_docs) + + retriever = MemoryRetriever( + knowledge_sources={"local:合规文档": local_rag, "法务知识库": legal_rag} + ) + + items = await retriever.retrieve( + "合规", top_k=5, sources=["local:合规文档"] + ) + + # 结果应为 MemoryItem 类型 + for item in items: + assert hasattr(item, "key") + assert hasattr(item, "value") + assert hasattr(item, "score") + assert hasattr(item, "metadata") + + @pytest.mark.asyncio + async def test_retrieve_without_sources_keeps_current_behavior( + self, local_rag, compliance_docs + ): + """不指定 sources 时保持原有行为(三层记忆检索)""" + await local_rag.ingest(compliance_docs) + + retriever = MemoryRetriever( + knowledge_sources={"local:合规文档": local_rag} + ) + + # 不指定 sources → 走三层记忆路径 + items = await retriever.retrieve("合规", top_k=5) + # 三层记忆为空,应返回空结果 + assert isinstance(items, list) + + @pytest.mark.asyncio + async def test_retrieve_from_sources_with_source_tracing( + self, local_rag, legal_rag, compliance_docs, legal_docs + ): + """通过 MemoryRetriever 多源检索时,结果包含来源追溯""" + await local_rag.ingest(compliance_docs) + await legal_rag.ingest(legal_docs) + + retriever = MemoryRetriever( + knowledge_sources={"local:合规文档": local_rag, "法务知识库": legal_rag} + ) + + items = await retriever.retrieve( + "合规", top_k=5, sources=["local:合规文档", "法务知识库"] + ) + + for item in items: + assert item.metadata.get("source") == "rag" + assert "source_name" in item.metadata + assert "document_title" in item.metadata + + @pytest.mark.asyncio + async def test_multi_source_retriever_property( + self, local_rag, compliance_docs + ): + """通过 multi_source_retriever 属性直接访问""" + retriever = MemoryRetriever( + knowledge_sources={"local:合规文档": local_rag} + ) + + ms_retriever = retriever.multi_source_retriever + assert isinstance(ms_retriever, MultiSourceRetriever) + assert "local:合规文档" in ms_retriever.get_source_names() + + @pytest.mark.asyncio + async def test_register_source_via_property( + self, local_rag, legal_rag, compliance_docs, legal_docs + ): + """通过 multi_source_retriever 属性动态注册源""" + await local_rag.ingest(compliance_docs) + await legal_rag.ingest(legal_docs) + + retriever = MemoryRetriever( + knowledge_sources={"local:合规文档": local_rag} + ) + + # 动态注册法务知识库 + retriever.multi_source_retriever.register_source("法务知识库", legal_rag) + + # 现在可以从法务知识库检索 + items = await retriever.retrieve( + "合同", top_k=5, sources=["法务知识库"] + ) + + for item in items: + assert item.metadata.get("source_name") == "法务知识库" + + @pytest.mark.asyncio + async def test_retrieve_with_source_weights( + self, local_rag, legal_rag, compliance_docs, legal_docs + ): + """MemoryRetriever 支持 source_weights 参数""" + await local_rag.ingest(compliance_docs) + await legal_rag.ingest(legal_docs) + + retriever = MemoryRetriever( + knowledge_sources={"local:合规文档": local_rag, "法务知识库": legal_rag} + ) + + items = await retriever.retrieve( + "合规", + top_k=5, + sources=["local:合规文档", "法务知识库"], + source_weights={"local:合规文档": 1.5}, + ) + + assert isinstance(items, list) + + +# ── 边界情况测试 ────────────────────────────────────────── + + +class TestEdgeCases: + """边界情况测试""" + + @pytest.mark.asyncio + async def test_search_with_empty_source_list(self, local_rag, compliance_docs): + """sources=[] 空列表 → 不查询任何源,返回空""" + await local_rag.ingest(compliance_docs) + + retriever = MultiSourceRetriever(sources={"local:合规文档": local_rag}) + + results = await retriever.search("合规", top_k=5, sources=[]) + assert len(results) == 0 + + @pytest.mark.asyncio + async def test_search_top_k_limits_results( + self, local_rag, legal_rag, compliance_docs, legal_docs + ): + """top_k 限制返回结果数量""" + await local_rag.ingest(compliance_docs) + await legal_rag.ingest(legal_docs) + + retriever = MultiSourceRetriever( + sources={"local:合规文档": local_rag, "法务知识库": legal_rag} + ) + + results = await retriever.search("合规", top_k=1) + assert len(results) <= 1 + + @pytest.mark.asyncio + async def test_source_query_failure_graceful( + self, local_rag, compliance_docs + ): + """某个源查询失败时,其他源结果正常返回""" + await local_rag.ingest(compliance_docs) + + # 创建一个会抛异常的 mock 源 + class FailingSource: + async def ingest(self, documents): + return [] + + async def query(self, text, top_k=5): + raise ConnectionError("Service unavailable") + + async def delete_by_id(self, id): + return False + + async def list_sources(self): + return [SourceInfo(source_id="failing", source_name="Failing", source_type="mock")] + + async def health_check(self): + return False + + retriever = MultiSourceRetriever( + sources={"local:合规文档": local_rag, "failing_source": FailingSource()} + ) + + # 应不抛异常,返回合规文档库的结果 + results = await retriever.search("合规", top_k=5) + assert isinstance(results, list) + + @pytest.mark.asyncio + async def test_search_results_sorted_by_score( + self, local_rag, legal_rag, compliance_docs, legal_docs + ): + """结果按 score 降序排列""" + await local_rag.ingest(compliance_docs) + await legal_rag.ingest(legal_docs) + + retriever = MultiSourceRetriever( + sources={"local:合规文档": local_rag, "法务知识库": legal_rag} + ) + + results = await retriever.search("合规 法务", top_k=10) + + for i in range(len(results) - 1): + assert results[i].score >= results[i + 1].score