zero-agent-core 0.2.0__py3-none-any.whl

agent_core/retrieval/chroma.py

@@ -0,0 +1,269 @@
+"""Chroma 本地持久化向量存储适配器。"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Iterable, Mapping, Sequence
+from pathlib import Path
+from typing import Any, cast
+
+from agent_core.access import AccessContext
+from agent_core.retrieval._common import (
+    decode_record_payload,
+    encode_record_payload,
+    metadata_matches,
+    normalize_record,
+    record_accessible,
+    validate_vector,
+)
+from agent_core.retrieval.core import VectorStore
+from agent_core.retrieval.types import VectorQuery, VectorRecord, VectorSearchResult
+
+_NAMESPACE_KEY = "_agent_core_namespace"
+_USER_KEY = "_agent_core_user_id"
+_TENANT_KEY = "_agent_core_tenant_id"
+_PAYLOAD_KEY = "_agent_core_payload"
+_PUBLIC_OWNER = ""
+
+
+class ChromaVectorStore(VectorStore):
+    """适合本地开发的 Chroma 持久化实现。
+
+    Chroma 的 Python API 是同步接口，因此所有读写都通过 ``asyncio.to_thread``
+    执行，避免阻塞 Agent 所在的事件循环。
+    """
+
+    def __init__(
+        self,
+        path: str | Path = "./chroma_db",
+        *,
+        collection_name: str = "agent_core",
+        require_access: bool = False,
+        client: Any | None = None,
+    ) -> None:
+        owns_client = client is None
+        if client is None:
+            try:
+                import chromadb
+            except ImportError as exc:
+                raise ImportError(
+                    "ChromaVectorStore 需要 chromadb，请安装 zero-agent-core[chroma]"
+                ) from exc
+            client = chromadb.PersistentClient(path=str(Path(path).resolve()))
+        # Chroma 各版本的 Client/Collection 泛型和重载差异较大，在 SDK 边界收口为 Any。
+        self._client: Any = client
+        self._owns_client = owns_client
+        self._collection: Any = client.get_or_create_collection(
+            name=collection_name,
+            metadata={"hnsw:space": "cosine"},
+        )
+        self._require_access = require_access
+
+    async def close(self) -> None:
+        """关闭由当前 Store 创建的 PersistentClient，释放本地文件句柄。"""
+        if not self._owns_client:
+            return
+        close = getattr(self._client, "close", None)
+        if close is not None:
+            await asyncio.to_thread(close)
+        self._owns_client = False
+
+    async def upsert(
+        self,
+        records: Sequence[VectorRecord],
+        *,
+        access: AccessContext | None = None,
+    ) -> list[str]:
+        values = [normalize_record(record, access, require_access=self._require_access) for record in records]
+        if not values:
+            return []
+        for record in values:
+            validate_vector(record.vector)
+        await self._check_existing_access([record.record_id for record in values], access)
+        metadatas = [self._metadata(record) for record in values]
+        await asyncio.to_thread(
+            self._collection.upsert,
+            ids=[record.record_id for record in values],
+            embeddings=[list(record.vector) for record in values],
+            documents=[record.document.content for record in values],
+            metadatas=metadatas,
+        )
+        return [record.record_id for record in values]
+
+    async def delete(
+        self,
+        record_ids: Sequence[str],
+        *,
+        access: AccessContext | None = None,
+    ) -> int:
+        self._check_required_access(access)
+        if not record_ids:
+            return 0
+        result = await asyncio.to_thread(
+            self._collection.get,
+            ids=list(record_ids),
+            include=["metadatas"],
+        )
+        ids = self._flat_list(result, "ids")
+        metadatas = self._flat_list(result, "metadatas")
+        allowed: list[str] = []
+        for index, record_id in enumerate(ids):
+            metadata = self._mapping_at(metadatas, index)
+            self._check_metadata_access(str(record_id), metadata, access)
+            allowed.append(str(record_id))
+        if allowed:
+            await asyncio.to_thread(self._collection.delete, ids=allowed)
+        return len(allowed)
+
+    async def search(self, query: VectorQuery) -> list[VectorSearchResult]:
+        self._check_required_access(query.access)
+        validate_vector(query.vector)
+        count = int(await asyncio.to_thread(self._collection.count))
+        if count <= 0:
+            return []
+        # 没有业务 metadata 过滤时，namespace 和访问范围已下推给 Chroma，
+        # 只需取 limit 条；带任意 metadata 过滤时仍全量读取以保持现有精确语义。
+        n_results = count if query.metadata_filter else min(count, query.limit)
+        result = await asyncio.to_thread(
+            self._collection.query,
+            query_embeddings=[list(query.vector)],
+            n_results=n_results,
+            where=self._where(query.namespace, query.access),
+            include=["metadatas", "distances", "embeddings"],
+        )
+        ids = self._first_batch(result, "ids")
+        metadatas = self._first_batch(result, "metadatas")
+        distances = self._first_batch(result, "distances")
+        embeddings = self._first_batch(result, "embeddings")
+        matches: list[VectorSearchResult] = []
+        for index, record_id in enumerate(ids):
+            metadata = self._mapping_at(metadatas, index)
+            raw_vector = embeddings[index] if index < len(embeddings) else ()
+            vector = tuple(float(cast(Any, value)) for value in cast(Sequence[object], raw_vector))
+            record = self._record_from_metadata(str(record_id), metadata, vector)
+            combined_metadata = {**record.document.metadata, **record.metadata}
+            if not metadata_matches(combined_metadata, query.metadata_filter):
+                continue
+            if not record_accessible(record, query.access):
+                continue
+            distance = float(cast(Any, distances[index])) if index < len(distances) else 1.0
+            score = 1.0 - distance
+            if query.min_score is not None and score < query.min_score:
+                continue
+            matches.append(VectorSearchResult(record=record, score=score))
+            if len(matches) >= query.limit:
+                break
+        return matches
+
+    async def _check_existing_access(
+        self,
+        record_ids: Sequence[str],
+        access: AccessContext | None,
+    ) -> None:
+        if access is None:
+            return
+        result = await asyncio.to_thread(
+            self._collection.get,
+            ids=list(record_ids),
+            include=["metadatas"],
+        )
+        ids = self._flat_list(result, "ids")
+        metadatas = self._flat_list(result, "metadatas")
+        for index, record_id in enumerate(ids):
+            self._check_metadata_access(
+                str(record_id),
+                self._mapping_at(metadatas, index),
+                access,
+            )
+
+    def _check_required_access(self, access: AccessContext | None) -> None:
+        if self._require_access and access is None:
+            raise PermissionError("该 VectorStore 要求提供 AccessContext")
+
+    @staticmethod
+    def _check_metadata_access(
+        record_id: str,
+        metadata: Mapping[str, object],
+        access: AccessContext | None,
+    ) -> None:
+        if access is None:
+            return
+        user_id = str(metadata[_USER_KEY]) if metadata.get(_USER_KEY) else None
+        tenant_id = str(metadata[_TENANT_KEY]) if metadata.get(_TENANT_KEY) else None
+        if not access.can_access(user_id, tenant_id):
+            raise PermissionError(f"无权访问向量记录：{record_id}")
+
+    @staticmethod
+    def _metadata(record: VectorRecord) -> dict[str, str]:
+        return {
+            _NAMESPACE_KEY: record.namespace,
+            _USER_KEY: record.user_id or _PUBLIC_OWNER,
+            _TENANT_KEY: record.tenant_id or _PUBLIC_OWNER,
+            _PAYLOAD_KEY: encode_record_payload(record),
+        }
+
+    @staticmethod
+    def _where(namespace: str, access: AccessContext | None) -> dict[str, object]:
+        clauses: list[dict[str, object]] = [{_NAMESPACE_KEY: {"$eq": namespace}}]
+        if access is not None:
+            clauses.append(
+                {
+                    "$or": [
+                        {_TENANT_KEY: {"$eq": _PUBLIC_OWNER}},
+                        {_TENANT_KEY: {"$eq": access.tenant_id}},
+                    ]
+                }
+            )
+            if not access.is_admin:
+                clauses.append(
+                    {
+                        "$or": [
+                            {_USER_KEY: {"$eq": _PUBLIC_OWNER}},
+                            {_USER_KEY: {"$eq": access.user_id}},
+                        ]
+                    }
+                )
+        return clauses[0] if len(clauses) == 1 else {"$and": clauses}
+
+    @staticmethod
+    def _record_from_metadata(
+        record_id: str,
+        metadata: Mapping[str, object],
+        vector: tuple[float, ...],
+    ) -> VectorRecord:
+        return decode_record_payload(
+            metadata.get(_PAYLOAD_KEY),
+            record_id=record_id,
+            vector=vector,
+            namespace=str(metadata.get(_NAMESPACE_KEY, "default")),
+            user_id=str(metadata[_USER_KEY]) if metadata.get(_USER_KEY) else None,
+            tenant_id=str(metadata[_TENANT_KEY]) if metadata.get(_TENANT_KEY) else None,
+        )
+
+    @staticmethod
+    def _flat_list(result: object, key: str) -> list[object]:
+        if not isinstance(result, Mapping):
+            return []
+        return ChromaVectorStore._to_list(result.get(key))
+
+    @staticmethod
+    def _first_batch(result: object, key: str) -> list[object]:
+        values = ChromaVectorStore._flat_list(result, key)
+        if not values:
+            return []
+        return ChromaVectorStore._to_list(values[0])
+
+    @staticmethod
+    def _to_list(value: object) -> list[object]:
+        if isinstance(value, (str, bytes, Mapping)) or not isinstance(value, Iterable):
+            return []
+        return list(value)
+
+    @staticmethod
+    def _mapping_at(values: Sequence[object], index: int) -> Mapping[str, object]:
+        if index >= len(values) or not isinstance(values[index], Mapping):
+            raise ValueError("Chroma 返回的向量记录缺少 metadata")
+        return cast(Mapping[str, object], values[index])
+
+
+__all__ = ["ChromaVectorStore"]

agent_core/retrieval/core.py

@@ -0,0 +1,63 @@
+"""检索依赖倒置协议。"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Protocol, runtime_checkable
+
+from agent_core.access import AccessContext
+from agent_core.retrieval.types import (
+    RetrievalQuery,
+    RetrievalResult,
+    VectorQuery,
+    VectorRecord,
+    VectorSearchResult,
+)
+
+
+@runtime_checkable
+class Embeddings(Protocol):
+    """文本向量化协议，厂商 SDK 由扩展适配器实现。"""
+
+    async def embed_documents(self, texts: Sequence[str]) -> list[list[float]]: ...
+
+    async def embed_query(self, text: str) -> list[float]: ...
+
+
+@runtime_checkable
+class Retriever(Protocol):
+    """根据查询返回排好序的文档。"""
+
+    async def retrieve(self, query: RetrievalQuery) -> list[RetrievalResult]: ...
+
+
+@runtime_checkable
+class VectorStore(Protocol):
+    """向量存储端口，不绑定 pgvector、Milvus、Qdrant 或其他实现。"""
+
+    async def upsert(
+        self,
+        records: Sequence[VectorRecord],
+        *,
+        access: AccessContext | None = None,
+    ) -> list[str]: ...
+
+    async def delete(
+        self,
+        record_ids: Sequence[str],
+        *,
+        access: AccessContext | None = None,
+    ) -> int: ...
+
+    async def search(self, query: VectorQuery) -> list[VectorSearchResult]: ...
+
+
+@runtime_checkable
+class Reranker(Protocol):
+    """可选的二阶段重排端口。"""
+
+    async def rerank(
+        self,
+        query: RetrievalQuery,
+        results: Sequence[RetrievalResult],
+    ) -> list[RetrievalResult]: ...

agent_core/retrieval/factory.py

@@ -0,0 +1,72 @@
+"""按环境和显式配置创建向量存储。"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Literal
+
+from agent_core.retrieval.chroma import ChromaVectorStore
+from agent_core.retrieval.core import VectorStore
+from agent_core.retrieval.memory import InMemoryVectorStore
+from agent_core.retrieval.pgvector import PgVectorStore
+
+VectorStoreBackend = Literal["memory", "chroma", "pgvector"]
+
+
+def create_vector_store(
+    env: str = "development",
+    *,
+    backend: str | None = None,
+    chroma_path: str | Path = "./chroma_db",
+    chroma_client: Any | None = None,
+    postgres_url: str | None = None,
+    dimension: int | None = None,
+    collection_name: str = "agent_core",
+    table_name: str = "agent_vector_records",
+    require_access: bool = False,
+    pg_create_extension: bool = True,
+    pg_create_index: bool = True,
+) -> VectorStore:
+    """创建向量存储；显式 backend 优先于环境默认值。
+
+    默认映射：测试使用内存，开发使用 Chroma，生产使用 pgvector。
+    PostgreSQL 实现返回后仍需调用 ``initialize()``。
+    """
+    normalized_env = env.strip().lower()
+    selected = (backend or _default_backend(normalized_env)).strip().lower()
+    selected = {"postgres": "pgvector", "postgresql": "pgvector"}.get(selected, selected)
+    if selected == "memory":
+        return InMemoryVectorStore(require_access=require_access)
+    if selected == "chroma":
+        return ChromaVectorStore(
+            chroma_path,
+            collection_name=collection_name,
+            require_access=require_access,
+            client=chroma_client,
+        )
+    if selected == "pgvector":
+        if not postgres_url:
+            raise ValueError("pgvector 后端需要配置 PostgreSQL 连接地址")
+        if dimension is None:
+            raise ValueError("pgvector 后端需要配置 embedding dimension")
+        return PgVectorStore(
+            postgres_url,
+            dimension=dimension,
+            collection_name=collection_name,
+            table_name=table_name,
+            require_access=require_access,
+            create_extension=pg_create_extension,
+            create_index=pg_create_index,
+        )
+    raise ValueError(f"不支持的向量存储后端：{selected}")
+
+
+def _default_backend(env: str) -> VectorStoreBackend:
+    if env in {"test", "testing"}:
+        return "memory"
+    if env in {"production", "prod"}:
+        return "pgvector"
+    return "chroma"
+
+
+__all__ = ["VectorStoreBackend", "create_vector_store"]

agent_core/retrieval/memory.py

@@ -0,0 +1,232 @@
+"""用于开发、测试和小数据集的零依赖检索实现。"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Iterable, Sequence
+from dataclasses import replace
+from typing import Any
+
+from agent_core.access import AccessContext
+from agent_core.documents import Document
+from agent_core.retrieval._common import check_record_access, normalize_record, validate_vector
+from agent_core.retrieval.core import Embeddings, VectorStore
+from agent_core.retrieval.types import (
+    RetrievalQuery,
+    RetrievalResult,
+    SearchType,
+    VectorQuery,
+    VectorRecord,
+    VectorSearchResult,
+)
+
+
+def _metadata_matches(metadata: dict[str, Any], expected: dict[str, Any]) -> bool:
+    return all(metadata.get(key) == value for key, value in expected.items())
+
+
+def _cosine_similarity(left: tuple[float, ...], right: tuple[float, ...]) -> float:
+    if len(left) != len(right):
+        raise ValueError(f"向量维度不一致：{len(left)} != {len(right)}")
+    dot = sum(a * b for a, b in zip(left, right, strict=True))
+    left_norm = math.sqrt(sum(value * value for value in left))
+    right_norm = math.sqrt(sum(value * value for value in right))
+    if left_norm == 0 or right_norm == 0:
+        return 0.0
+    return dot / (left_norm * right_norm)
+
+
+class InMemoryVectorStore(VectorStore):
+    """带 namespace、metadata 和用户/租户过滤的内存向量存储。"""
+
+    def __init__(self, *, require_access: bool = False) -> None:
+        self._records: dict[str, VectorRecord] = {}
+        self._dimension: int | None = None
+        self._require_access = require_access
+
+    async def upsert(
+        self,
+        records: Sequence[VectorRecord],
+        *,
+        access: AccessContext | None = None,
+    ) -> list[str]:
+        self._check_required_access(access)
+        saved: list[str] = []
+        for record in records:
+            self._check_dimension(record.vector)
+            value = normalize_record(record, access, require_access=self._require_access)
+            existing = self._records.get(value.record_id)
+            if existing is not None:
+                check_record_access(existing, access)
+            self._records[value.record_id] = value
+            saved.append(value.record_id)
+        return saved
+
+    async def delete(
+        self,
+        record_ids: Sequence[str],
+        *,
+        access: AccessContext | None = None,
+    ) -> int:
+        self._check_required_access(access)
+        existing = [self._records[record_id] for record_id in record_ids if record_id in self._records]
+        # 先完成整批权限校验，避免中途失败时前面的记录已经被删除。
+        for record in existing:
+            self._check_record_access(record, access)
+        deleted = 0
+        for record_id in record_ids:
+            target = self._records.get(record_id)
+            if target is None:
+                continue
+            del self._records[record_id]
+            deleted += 1
+        return deleted
+
+    async def search(self, query: VectorQuery) -> list[VectorSearchResult]:
+        self._check_required_access(query.access)
+        if self._dimension is not None and len(query.vector) != self._dimension:
+            raise ValueError(f"向量维度不一致：期望 {self._dimension}，实际 {len(query.vector)}")
+        matches: list[VectorSearchResult] = []
+        for record in self._records.values():
+            if record.namespace != query.namespace:
+                continue
+            if not _metadata_matches({**record.document.metadata, **record.metadata}, query.metadata_filter):
+                continue
+            if not self._record_accessible(record, query.access):
+                continue
+            score = _cosine_similarity(query.vector, record.vector)
+            if query.min_score is not None and score < query.min_score:
+                continue
+            matches.append(VectorSearchResult(record=record, score=score))
+        matches.sort(key=lambda item: (-item.score, item.record.record_id))
+        return matches[: query.limit]
+
+    def _check_dimension(self, vector: tuple[float, ...]) -> None:
+        validate_vector(vector)
+        if self._dimension is None:
+            self._dimension = len(vector)
+        elif len(vector) != self._dimension:
+            raise ValueError(f"向量维度不一致：期望 {self._dimension}，实际 {len(vector)}")
+
+    def _check_required_access(self, access: AccessContext | None) -> None:
+        if self._require_access and access is None:
+            raise PermissionError("该 VectorStore 要求提供 AccessContext")
+
+    @staticmethod
+    def _record_accessible(record: VectorRecord, access: AccessContext | None) -> bool:
+        return access is None or access.can_access(record.user_id, record.tenant_id)
+
+    @classmethod
+    def _check_record_access(cls, record: VectorRecord, access: AccessContext | None) -> None:
+        if not cls._record_accessible(record, access):
+            raise PermissionError(f"无权访问向量记录：{record.record_id}")
+
+
+class EmbeddingRetriever:
+    """组合 Embeddings 和 VectorStore 的通用语义 Retriever。"""
+
+    def __init__(self, embeddings: Embeddings, vector_store: VectorStore) -> None:
+        self._embeddings = embeddings
+        self._vector_store = vector_store
+
+    async def add_documents(
+        self,
+        documents: Sequence[Document],
+        *,
+        namespace: str = "default",
+        access: AccessContext | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> list[str]:
+        vectors = await self._embeddings.embed_documents([item.content for item in documents])
+        if len(vectors) != len(documents):
+            raise ValueError("Embedding 返回数量与文档数量不一致")
+        records = [
+            VectorRecord(
+                record_id=document.document_id,
+                vector=tuple(float(value) for value in vector),
+                document=document,
+                namespace=namespace,
+                user_id=access.user_id if access else None,
+                tenant_id=access.tenant_id if access else None,
+                metadata=dict(metadata or {}),
+            )
+            for document, vector in zip(documents, vectors, strict=True)
+        ]
+        return await self._vector_store.upsert(records, access=access)
+
+    async def retrieve(self, query: RetrievalQuery) -> list[RetrievalResult]:
+        if query.search_type != SearchType.SIMILARITY:
+            raise ValueError("EmbeddingRetriever 当前只支持 similarity 检索")
+        vector = await self._embeddings.embed_query(query.text)
+        matches = await self._vector_store.search(
+            VectorQuery(
+                vector=tuple(float(value) for value in vector),
+                limit=query.limit,
+                namespace=query.namespace,
+                metadata_filter=query.metadata_filter,
+                min_score=query.min_score,
+                access=query.access,
+            )
+        )
+        return [
+            RetrievalResult(
+                document=item.record.document,
+                score=item.score,
+                rank=index,
+                metadata={"record_id": item.record.record_id, **item.record.metadata},
+            )
+            for index, item in enumerate(matches, start=1)
+        ]
+
+
+class KeywordRetriever:
+    """零依赖关键词 Retriever，适合测试和小型静态语料。"""
+
+    def __init__(
+        self,
+        documents: Iterable[Document] | None = None,
+        *,
+        require_access: bool = False,
+    ) -> None:
+        self._documents = list(documents or ())
+        self._require_access = require_access
+
+    def add_documents(self, documents: Iterable[Document]) -> None:
+        self._documents.extend(documents)
+
+    async def retrieve(self, query: RetrievalQuery) -> list[RetrievalResult]:
+        if query.search_type != SearchType.KEYWORD:
+            raise ValueError("KeywordRetriever 只支持 keyword 检索")
+        if self._require_access and query.access is None:
+            raise PermissionError("该 Retriever 要求提供 AccessContext")
+        terms = [term.casefold() for term in query.text.split() if term.strip()]
+        if not terms and query.text.strip():
+            terms = [query.text.strip().casefold()]
+        results: list[RetrievalResult] = []
+        for document in self._documents:
+            metadata = document.metadata
+            if metadata.get("namespace", "default") != query.namespace:
+                continue
+            if not _metadata_matches(metadata, query.metadata_filter):
+                continue
+            if query.access is not None and not query.access.can_access(
+                metadata.get("user_id"), metadata.get("tenant_id")
+            ):
+                continue
+            content = document.content.casefold()
+            matches = sum(content.count(term) for term in terms)
+            if matches == 0:
+                continue
+            score = matches / max(len(terms), 1)
+            if query.min_score is not None and score < query.min_score:
+                continue
+            results.append(
+                RetrievalResult(
+                    document=document,
+                    score=score,
+                    rank=0,
+                    metadata={"matches": matches},
+                )
+            )
+        results.sort(key=lambda item: (-item.score, item.document.document_id))
+        return [replace(item, rank=index) for index, item in enumerate(results[: query.limit], start=1)]