# -*- coding: utf-8 -*- """Runtime handle for a single knowledge base. A :class:`KnowledgeBase` instance is the **single algorithmic source of truth** for talking to one knowledge base: it pairs an embedding model with a vector-store collection (optionally scoped by a payload ``metadata_filter``) and exposes the four operations a caller ever needs — :meth:`search`, :meth:`insert_document`, :meth:`delete_document`, :meth:`list_documents`. The handle is *narrow on purpose* — it carries the resolved runtime state (embedding model + vector store + scope) and delegates every operation to the bound :class:`VectorStoreBase`. Document parsing, chunking, credential resolution, dimension policy validation, and persistence of knowledge-base records all belong one layer up (service-side :class:`KnowledgeBaseManagerBase` for hosted deployments; the caller directly otherwise). The backing collection is created on first use — each operation transparently calls :meth:`ensure_collection`, which is itself idempotent and memoised after the first success, so the only cost is one extra round-trip on the very first call against a fresh deployment. ``metadata_filter`` is the defense-in-depth scoping mechanism for co-locating multiple logical knowledge bases inside the same physical collection — typically multi-tenant deployments where every record carries a ``{"tenant_id": "..."}`` payload. It is set once at construction time and **always** applied: search/list never escape it, and insert forces it onto every chunk's metadata so a malicious or buggy parser cannot rebind a record into another scope. """ import asyncio from ._document import Chunk from ._vdb import VectorRecord, VectorSearchResult, VectorStoreBase from .._utils._common import _generate_id from ..embedding import EmbeddingModelBase from ..message import DataBlock, TextBlock from ._vdb import DocumentSummary class KnowledgeBase: """Runtime handle for one knowledge base. Binds an embedding model and a vector-store collection together so callers can retrieve / insert / delete / list documents without repeating the wiring. Cheap to construct (no I/O); the collection itself is created lazily on the first operation, so a fresh deployment "just works" without an explicit setup step. .. code-block:: python kb = KnowledgeBase( name="company-handbook", description="Internal HR and onboarding documents.", embedding_model=embedding_model, vector_store=vector_store, collection="handbook", ) await kb.insert_document(chunks) results = await kb.search(["What is the PTO policy?"]) """ name: str """Agent-oriented knowledge base name — used by tool descriptions and frontend rendering.""" description: str """Agent-oriented knowledge base description — what this knowledge base contains and when to retrieve from it.""" def __init__( self, name: str, description: str, embedding_model: EmbeddingModelBase, vector_store: VectorStoreBase, collection: str, metadata_filter: dict | None = None, ) -> None: """Initialize the runtime handle. Args: name (`str`): Agent-oriented knowledge base name. Surfaced to the LLM (via tool descriptions) and to the front-end. description (`str`): Agent-oriented description. Should answer "what is in this knowledge base and when should I search it?" — the LLM uses it to decide whether to call the search tool in agentic mode. embedding_model (`EmbeddingModelBase`): The embedding model used to embed both queries and inserted chunks. Must be the same model used at indexing time and at retrieval time, otherwise vectors will not be comparable. vector_store (`VectorStoreBase`): The shared vector-store connection. The store must already be entered (its own ``__aenter__`` already called) before any operation on this handle runs. collection (`str`): The physical collection backing this knowledge base. Created lazily on the first operation; see :meth:`ensure_collection`. metadata_filter (`dict | None`, optional): Defense-in-depth payload filter. When set: - :meth:`search` and :meth:`list_documents` restrict results to records whose payload matches every ``key == value`` pair; - :meth:`insert_document` forces these keys onto every inserted chunk's metadata, overriding caller-supplied values, so records cannot leak into another scope. ``None`` disables filtering — the default for deployments where every knowledge base owns its collection outright. """ self.name = name self.description = description self._embedding_model = embedding_model self._vector_store = vector_store self._collection = collection self._metadata_filter = metadata_filter # Memoise the "collection exists" check after the first # successful ensure_collection so subsequent operations avoid # the extra round-trip. self._collection_ready = False # ------------------------------------------------------------------ # Read-only accessors # ------------------------------------------------------------------ @property def embedding_model(self) -> EmbeddingModelBase: """The bound embedding model.""" return self._embedding_model @property def vector_store(self) -> VectorStoreBase: """The bound vector store.""" return self._vector_store @property def collection(self) -> str: """The physical collection backing this knowledge base.""" return self._collection @property def metadata_filter(self) -> dict | None: """The defense-in-depth payload filter, or ``None``.""" return self._metadata_filter # ------------------------------------------------------------------ # Lifecycle # ------------------------------------------------------------------ async def ensure_collection(self) -> None: """Idempotently create the backing collection if missing. Called transparently at the top of every public operation — callers should not need to invoke it themselves. Memoised on the instance after the first success, so subsequent calls are a single ``if`` check. Looks up the collection via :meth:`VectorStoreBase.has_collection` and creates it with the embedding model's :attr:`~EmbeddingModelBase.dimensions` when absent. Raises whatever the backend raises if the collection exists at an incompatible dimension (the backend is the authority on that; we do not double-check here). """ if self._collection_ready: return if not await self._vector_store.has_collection(self._collection): await self._vector_store.create_collection( self._collection, dimensions=self._embedding_model.dimensions, ) self._collection_ready = True # ------------------------------------------------------------------ # Search # ------------------------------------------------------------------ async def search( self, queries: list[str | TextBlock | DataBlock], top_k: int = 5, score_threshold: float | None = None, ) -> list[VectorSearchResult]: """Search the knowledge base with one or more queries. All queries are embedded in a single batch, then searched concurrently against the bound collection (with :attr:`metadata_filter` applied). Hits are deduplicated by ``(document_id, chunk_index)`` keeping the best score, optionally filtered by ``score_threshold``, sorted by descending score, and truncated to ``top_k``. Args: queries (`list[str | TextBlock | DataBlock]`): Query inputs. Text may be either bare ``str`` or :class:`TextBlock`; :class:`DataBlock` items are **silently dropped** when the bound embedding model does not declare ``supports_multimodal`` — text-only models would otherwise reject them. Callers can therefore pass a mixed list without per-KB filtering. top_k (`int`, defaults to ``5``): Maximum number of results returned across all queries (after dedup). score_threshold (`float | None`, optional): Minimum similarity score for a hit to be retained. Only meaningful for similarity metrics where higher is better (cosine / dot-product). ``None`` disables filtering. Returns: `list[VectorSearchResult]`: At most ``top_k`` deduplicated hits ordered by descending similarity score. Empty when there are no queries the bound embedding model can consume. """ if not queries: return [] if not self._embedding_model.supports_multimodal: queries = [q for q in queries if not isinstance(q, DataBlock)] if not queries: return [] await self.ensure_collection() response = await self._embedding_model(queries) results_per_query = await asyncio.gather( *( self._vector_store.search( collection=self._collection, query_vector=vector, top_k=top_k, metadata_filter=self._metadata_filter, ) for vector in response.embeddings ), ) best: dict[tuple[str, int], VectorSearchResult] = {} for results in results_per_query: for result in results: if ( score_threshold is not None and result.score < score_threshold ): continue # ``(document_id, chunk_index)`` is the stable identity # of a chunk: it survives reindex (block UUIDs do not) # and uniquely names "this slice of that document" # regardless of which query surfaced it. key = (result.document_id, result.chunk.chunk_index) if key not in best or result.score > best[key].score: best[key] = result merged = sorted( best.values(), key=lambda result: result.score, reverse=True, ) return merged[:top_k] # ------------------------------------------------------------------ # Document management # ------------------------------------------------------------------ async def insert_document( self, chunks: list[Chunk], document_id: str | None = None, document_metadata: dict | None = None, ) -> str: """Embed and insert a list of chunks as a single source document. All chunks share the resolved ``document_id``; :meth:`delete_document` later removes them as a unit. Each chunk's metadata is merged in this precedence (highest wins): 1. :attr:`metadata_filter` keys — defense-in-depth scoping, so a chunk can never be inserted with a payload that escapes the filter (any escape would silently disappear at retrieve time anyway, but failing closed at insert is clearer). 2. The chunk's pre-existing ``metadata`` — parser-supplied. 3. ``document_metadata`` — document-level fields propagated down (filename, media type, upload time, ...). Args: chunks (`list[Chunk]`): The pre-chunked document content (already produced by a parser + chunker pipeline). An empty list is a no-op. document_id (`str | None`, optional): The document identifier. When ``None`` a fresh UUID hex is generated and returned so the caller can record it for future :meth:`delete_document` calls. document_metadata (`dict | None`, optional): Document-level metadata (filename, media type, size, upload time, ...). Merged into each chunk's ``metadata``. Returns: `str`: The (possibly generated) document id. Raises: `RuntimeError`: If the embedding model returns a number of vectors that does not match the number of chunks. """ if not chunks: return document_id or _generate_id() document_id = document_id or _generate_id() await self.ensure_collection() # Precedence: metadata_filter wins (security boundary), then # chunk metadata, then document_metadata. See docstring. for chunk in chunks: chunk.metadata = { **(document_metadata or {}), **chunk.metadata, **(self._metadata_filter or {}), } response = await self._embedding_model( [chunk.content for chunk in chunks], ) if len(response.embeddings) != len(chunks): raise RuntimeError( f"Embedding model returned {len(response.embeddings)} " f"vectors for {len(chunks)} chunks.", ) records = [ VectorRecord( vector=vector, document_id=document_id, chunk=chunk, ) for vector, chunk in zip(response.embeddings, chunks) ] await self._vector_store.insert(self._collection, records) return document_id async def delete_document(self, document_id: str) -> None: """Remove every record for one source document. Args: document_id (`str`): The source document id whose records should be removed. """ await self.ensure_collection() await self._vector_store.delete( self._collection, document_id, ) async def list_documents(self) -> list["DocumentSummary"]: """List all distinct source documents in this knowledge base. Filtered by :attr:`metadata_filter` when set, so callers only ever see documents within their own scope. Returns: `list[DocumentSummary]`: One summary per indexed document, in unspecified order. """ await self.ensure_collection() return await self._vector_store.list_documents( self._collection, metadata_filter=self._metadata_filter, )