Vector Stores API¶

@abstractmethod
def search(self, query: str, k: int, **kwargs) -> list[AI4RAGChunk]:
    """
    Search for the chunks relevant to the query.
    The method used will be simple similarity search.

    Parameters
    ----------
    query : str
        Question / query for which the similarity search will be executed.

    k : int
        Number of chunks to be returned as a result of similarity search

    **kwargs : Any
        Backend-specific search options (e.g. metadata filters or hybrid
        search parameters). Ignored by backends that do not support them.

    Returns
    -------
    list[AI4RAGChunk]
        List of chunks with content and metadata.
    """

@abstractmethod
def add_documents(self, documents: Sequence[AI4RAGChunk]) -> None:
    """
    Add documents to the collection.

    Parameters
    ----------
    documents : Sequence[AI4RAGChunk]
        Chunks to add to the collection.
    """

def close(self) -> None:
    """Release backend resources held by this store (connections, clients).

    A no-op by default. Concrete stores that hold a real connection or
    client (e.g. :class:`~ai4rag.rag.vector_store.pgvector.PGVectorStore`,
    :class:`~ai4rag.rag.vector_store.milvus.MilvusVectorStore`) override this
    to release it; callers should call ``close()`` (or use the store as a
    context manager) once they are done searching or indexing, since the
    store performs no automatic cleanup on garbage collection. Idempotent:
    safe to call more than once.
    """

@classmethod
@abstractmethod
def from_env(cls) -> "BaseVectorStoreConfig":
    """Create config from environment variables."""

def __post_init__(self) -> None:
    """Reject any ``uri`` that is not an explicit Milvus server URL.

    Guards against the footgun where an incorrect ``uri`` (a typo, a bare
    hostname, or a stray path) is silently interpreted by ``MilvusClient`` as
    a local Milvus Lite database file, creating a throwaway store instead of
    connecting to the intended server. Local, embedded use must go through
    :class:`MilvusLiteConfig`.
    """
    if not _is_server_url(self.uri):
        raise ValueError(
            f"MilvusConfig.uri must be a Milvus server URL starting with 'http://' or 'https://', "
            f"got {self.uri!r}. For a local, embedded database use MilvusLiteConfig(db_path=...) "
            "(provider 'milvus_lite') instead."
        )

@classmethod
def from_env(cls) -> "MilvusConfig":
    """Build config from ``MILVUS_*`` environment variables.

    Reads ``MILVUS_URI`` (required), plus the optional ``MILVUS_TOKEN`` and
    ``MILVUS_SERVER_CERT``. ``MILVUS_SERVER_CERT`` holds the PEM certificate
    text itself, not a filesystem path.

    Returns
    -------
    MilvusConfig
        Config populated from the ``MILVUS_*`` environment variables.

    Raises
    ------
    KeyError
        If the required ``MILVUS_URI`` variable is not set.
    ValueError
        If ``MILVUS_URI`` is not an ``http://``/``https://`` URL.
    """
    return cls(
        uri=os.environ["MILVUS_URI"],
        token=os.environ.get("MILVUS_TOKEN"),
        server_cert=os.environ.get("MILVUS_SERVER_CERT"),
    )

def __post_init__(self) -> None:
    """Reject a ``db_path`` that is blank or is actually a server URL.

    The symmetric guard to :meth:`MilvusConfig.__post_init__`: a value like
    ``https://host:19530`` is a server endpoint, not a local database file,
    and belongs in :class:`MilvusConfig`. An empty or whitespace-only path
    is rejected here too, rather than being handed to ``MilvusClient`` where
    it would surface as an opaque, hard-to-trace pymilvus error.
    """
    if not isinstance(self.db_path, str) or not self.db_path.strip():
        raise ValueError(
            f"MilvusLiteConfig.db_path must be a non-empty local filesystem path, got {self.db_path!r}."
        )
    if _is_server_url(self.db_path):
        raise ValueError(
            f"MilvusLiteConfig.db_path must be a local filesystem path, not a server URL, "
            f"got {self.db_path!r}. For a remote Milvus server use MilvusConfig(uri=...) "
            "(provider 'milvus') instead."
        )

@classmethod
def from_env(cls) -> "MilvusLiteConfig":
    """Build config from the ``MILVUS_LITE_DB_PATH`` environment variable.

    An unset variable falls back to :data:`DEFAULT_MILVUS_LITE_DB_PATH`.

    Returns
    -------
    MilvusLiteConfig
        Config populated from ``MILVUS_LITE_DB_PATH`` (or the default path).
    """
    return cls(db_path=os.environ.get("MILVUS_LITE_DB_PATH", DEFAULT_MILVUS_LITE_DB_PATH))

@classmethod
def from_env(cls) -> "PGVectorConfig":
    """Build config from ``PGVECTOR_*`` environment variables.

    Reads ``PGVECTOR_HOST``, ``PGVECTOR_PORT``, ``PGVECTOR_DB``,
    ``PGVECTOR_USER`` and ``PGVECTOR_PASSWORD``. Unset variables fall back to
    the local-PostgreSQL defaults; ``PGVECTOR_PASSWORD`` defaults to ``None``
    for trust/peer authentication.

    Returns
    -------
    PGVectorConfig
        Config populated from the ``PGVECTOR_*`` environment variables.
    """
    return cls(
        host=os.environ.get("PGVECTOR_HOST", "localhost"),
        port=int(os.environ.get("PGVECTOR_PORT", "5432")),
        dbname=os.environ.get("PGVECTOR_DB", "postgres"),
        user=os.environ.get("PGVECTOR_USER", "postgres"),
        password=os.environ.get("PGVECTOR_PASSWORD"),
    )

def search(
    self,
    query: str,
    k: int = 5,
    include_scores: bool = False,
    search_mode: str = "vector",
    ranker_strategy: str | None = None,
    ranker_k: int | None = None,
    ranker_alpha: float | None = None,
    **kwargs,
) -> list[AI4RAGChunk] | list[tuple[AI4RAGChunk, float]]:
    """Search for chunks relevant to *query*.

    Parameters
    ----------
    query : str
        Search query text.
    k : int, default=5
        Number of results to return.
    include_scores : bool, default=False
        Whether to include similarity scores in the results.
    search_mode : str, default="vector"
        ``"vector"`` for dense-only search or ``"hybrid"`` for dense + BM25
        sparse search.
    ranker_strategy : str | None, default=None
        Hybrid ranker: ``"rrf"``, ``"weighted"``, or ``"normalized"``.
    ranker_k : int | None, default=None
        RRF smoothing constant (``k``).
    ranker_alpha : float | None, default=None
        Weighted blend factor (``0`` = keyword, ``1`` = vector).
    **kwargs : Any
        Accepted for interface compatibility; ignored by this backend.

    Returns
    -------
    list[AI4RAGChunk] | list[tuple[AI4RAGChunk, float]]
        Matched chunks, optionally paired with their scores.
    """
    validate_search_params(search_mode, ranker_strategy, ranker_k, ranker_alpha)

    if search_mode == "hybrid":
        return self._search_hybrid(query, k, include_scores, ranker_strategy, ranker_k, ranker_alpha)
    return self._search_vector(query, k, include_scores)

def add_documents(self, documents: list[AI4RAGChunk], **kwargs) -> None:
    """Embed, deduplicate, and upsert chunks into Milvus.

    Duplicate ``chunk_id`` values within *documents* are skipped (first
    occurrence wins) and logged. Rows are upserted in batches.

    Parameters
    ----------
    documents : list[AI4RAGChunk]
        Chunks to be embedded and stored.
    **kwargs : Any
        Optional overrides. ``batch_size`` (int) sets the upsert batch size
        (default :attr:`_BATCH_SIZE`).
    """
    if not documents:
        return

    embeddings = self.embedding_model.embed_documents([doc.text for doc in documents])

    data: list[dict[str, Any]] = []
    for doc, embedding in iter_unique_chunks(documents, embeddings):
        data.append(
            {
                "chunk_id": doc.chunk_id,
                "content": doc.text,
                "vector": embedding,
                "metadata": doc.metadata,
            }
        )

    batch_size = kwargs.get("batch_size", self._BATCH_SIZE)
    for idx in range(0, len(data), batch_size):
        self._client.upsert(self._collection_name, data=data[idx : idx + batch_size])

def clean_collection(self) -> None:
    """Drop the Milvus collection."""
    self._client.drop_collection(self._collection_name)

def close(self) -> None:
    """Close the underlying Milvus client connection.

    The temporary TLS certificate file (when one was materialized) is
    intentionally *not* removed here: pymilvus can reconnect an idle channel
    from a background thread and re-read ``server_pem_path`` after ``close()``,
    so the file is kept for the process lifetime and cleaned at interpreter
    exit by :func:`_cleanup_server_certs`.
    """
    self._client.close()

def search(
    self,
    query: str,
    k: int,
    include_scores: bool = False,
    search_mode: str = "vector",
    ranker_strategy: str | None = None,
    ranker_k: int | None = None,
    ranker_alpha: float | None = None,
    **kwargs,
) -> list[AI4RAGChunk] | list[tuple[AI4RAGChunk, float]]:
    """Search for chunks relevant to *query*.

    Parameters
    ----------
    query : str
        Search query text.
    k : int
        Number of results to return.
    include_scores : bool, default=False
        Whether to include similarity scores.
    search_mode : str, default="vector"
        ``"vector"`` for dense-only search or ``"hybrid"`` for dense +
        full-text search.
    ranker_strategy : str | None, default=None
        Hybrid ranker: ``"rrf"``, ``"weighted"``, or ``"normalized"``.
    ranker_k : int | None, default=None
        RRF smoothing constant (``k``).
    ranker_alpha : float | None, default=None
        Weighted blend factor (``0`` = keyword, ``1`` = vector).
    **kwargs : Any
        Accepted for interface compatibility; ignored by this backend.

    Returns
    -------
    list[AI4RAGChunk] | list[tuple[AI4RAGChunk, float]]
        Matched chunks, optionally paired with their scores.
    """
    validate_search_params(search_mode, ranker_strategy, ranker_k, ranker_alpha)
    self._ensure_indexes()

    if search_mode == "hybrid":
        return self._search_hybrid(query, k, include_scores, ranker_strategy, ranker_k, ranker_alpha)
    return self._search_vector(query, k, include_scores)

def add_documents(self, documents: list[AI4RAGChunk], **kwargs) -> None:
    """Embed, deduplicate, and upsert chunks into PGVector.

    Duplicate ``chunk_id`` values within *documents* are skipped (first
    occurrence wins) and logged. Rows are upserted in batches, each with a
    one-shot retry on a dropped connection.

    Parameters
    ----------
    documents : list[AI4RAGChunk]
        Chunks to be embedded and stored.
    **kwargs : Any
        Optional overrides. ``batch_size`` (int) sets the insert batch size
        (default :attr:`_BATCH_SIZE`).
    """
    if not documents:
        return

    # embed_documents() is a synchronous, network-bound call: it runs here, on
    # the caller's own thread, before anything is dispatched to the shared
    # event loop (see the matching note in _search_vector).
    embeddings = self.embedding_model.embed_documents([doc.text for doc in documents])
    pool = self._ensure_db()

    values: list[tuple[str, dict, list[float], str, str]] = []
    for doc, embedding in iter_unique_chunks(documents, embeddings):
        values.append((doc.chunk_id, doc.metadata, embedding, doc.text, doc.text))

    batch_size = kwargs.get("batch_size", self._BATCH_SIZE)
    for idx in range(0, len(values), batch_size):
        self._insert_batch_with_retry(pool, values[idx : idx + batch_size])

def clean_collection(self) -> None:
    """Drop the PostgreSQL table."""
    pool = self._ensure_pool()
    self._run_with_retry(lambda: self._drop_table(pool))

def close(self) -> None:
    """Close the connection pool and stop the store's background event loop.

    Idempotent: a second call sees :attr:`_db` already ``None`` and returns
    immediately, rather than dispatching onto a loop that has already
    stopped — which would hang forever waiting for a callback the stopped
    loop can never run.

    Safe to call while another thread is mid-:meth:`_run`: :attr:`_db` is
    cleared up front (under :attr:`_inflight`) so no *new* call can start,
    but the loop and pool are only stopped once every call that already
    started has finished — see :meth:`_run`. A call that started before
    this method clears :attr:`_db` still completes normally; one that
    starts after gets a clear ``RuntimeError`` instead of racing the teardown.
    """
    with self._inflight.cond:
        if self._db is None:
            return
        db = self._db
        self._db = None
        while self._inflight.count > 0:
            self._inflight.cond.wait()
    self._dispatch(db.loop, db.pool.close())
    db.loop.call_soon_threadsafe(db.loop.stop)
    db.thread.join()
    db.loop.close()

@staticmethod
def weighted_rerank(
    vector_scores: dict[str, float],
    keyword_scores: dict[str, float],
    alpha: float = 0.5,
) -> dict[str, float]:
    """Weighted average of normalized vector and keyword scores.

    Parameters
    ----------
    vector_scores : dict[str, float]
        Scores from vector search keyed by chunk id.
    keyword_scores : dict[str, float]
        Scores from keyword search keyed by chunk id.
    alpha : float
        Blend factor: ``0`` = keyword only, ``1`` = vector only.

    Returns
    -------
    dict[str, float]
        Combined scores keyed by chunk id.
    """
    all_ids = set(vector_scores) | set(keyword_scores)
    norm_vec = WeightedInMemoryAggregator._normalize_scores(vector_scores)
    norm_kw = WeightedInMemoryAggregator._normalize_scores(keyword_scores)
    return {
        doc_id: (1 - alpha) * norm_kw.get(doc_id, 0.0) + alpha * norm_vec.get(doc_id, 0.0) for doc_id in all_ids
    }

@staticmethod
def rrf_rerank(
    vector_scores: dict[str, float],
    keyword_scores: dict[str, float],
    k: float = 60.0,
) -> dict[str, float]:
    """Reciprocal Rank Fusion of vector and keyword search results.

    Parameters
    ----------
    vector_scores : dict[str, float]
        Scores from vector search keyed by chunk id.
    keyword_scores : dict[str, float]
        Scores from keyword search keyed by chunk id.
    k : float
        RRF smoothing constant (default ``60.0``).

    Returns
    -------
    dict[str, float]
        Fused RRF scores keyed by chunk id.
    """
    vector_ranks = {
        doc_id: i + 1
        for i, (doc_id, _) in enumerate(sorted(vector_scores.items(), key=lambda x: x[1], reverse=True))
    }
    keyword_ranks = {
        doc_id: i + 1
        for i, (doc_id, _) in enumerate(sorted(keyword_scores.items(), key=lambda x: x[1], reverse=True))
    }

    all_ids = set(vector_scores) | set(keyword_scores)
    return {
        doc_id: 1.0 / (k + vector_ranks.get(doc_id, float("inf")))
        + 1.0 / (k + keyword_ranks.get(doc_id, float("inf")))
        for doc_id in all_ids
    }

@staticmethod
def combine_search_results(
    vector_scores: dict[str, float],
    keyword_scores: dict[str, float],
    reranker_type: str = "rrf",
    reranker_params: dict[str, Any] | None = None,
) -> dict[str, float]:
    """Dispatch to the appropriate reranking strategy.

    Parameters
    ----------
    vector_scores : dict[str, float]
        Scores from vector search keyed by chunk id.
    keyword_scores : dict[str, float]
        Scores from keyword search keyed by chunk id.
    reranker_type : str
        ``"rrf"``, ``"weighted"``, or ``"normalized"`` (falls through to RRF).
    reranker_params : dict[str, Any] | None
        Strategy-specific params: ``{"k": float}`` for RRF,
        ``{"alpha": float}`` for weighted.

    Returns
    -------
    dict[str, float]
        Combined scores keyed by chunk id, produced by the selected strategy.
    """
    if reranker_params is None:
        reranker_params = {}

    if reranker_type == "weighted":
        alpha = reranker_params.get("alpha", 0.5)
        return WeightedInMemoryAggregator.weighted_rerank(vector_scores, keyword_scores, alpha)

    k = reranker_params.get("k", 60.0)
    return WeightedInMemoryAggregator.rrf_rerank(vector_scores, keyword_scores, k)