Skip to content

Retrieval API

Retriever

Retriever(
    vector_store: BaseVectorStore,
    number_of_chunks: int,
    method: Literal["simple"] = "simple",
    search_mode: Literal["vector", "hybrid"] = "vector",
    ranker_strategy: str | None = None,
    ranker_k: int | None = None,
    ranker_alpha: float | None = None,
)

Class responsible for retrieving data from given vector store.

Parameters:

  • vector_store (BaseVectorStore) –

    Vector store / vector index to retrieve data from.

  • method (Literal['simple'], default: 'simple' ) –

    Method describing how data should be retrieved.

  • number_of_chunks (int) –

    Number of chunks to retrieve.

  • search_mode (Literal['vector', 'hybrid'], default: "vector" ) –

    Search mode passed to the vector store: "vector" or "hybrid".

  • ranker_strategy (str | None, default: None ) –

    Ranking strategy for hybrid search: "rrf", "weighted", or "normalized".

  • ranker_k (int | None, default: None ) –

    Parameter k for the ranking function.

  • ranker_alpha (float | None, default: None ) –

    Alpha parameter for weighted ranking strategy.

Source code in ai4rag/rag/retrieval/retriever.py
def __init__(
    self,
    vector_store: BaseVectorStore,
    number_of_chunks: int,
    method: Literal["simple"] = "simple",
    search_mode: Literal["vector", "hybrid"] = "vector",
    ranker_strategy: str | None = None,
    ranker_k: int | None = None,
    ranker_alpha: float | None = None,
):
    self._vector_store = vector_store
    self.method = method
    self.number_of_chunks = number_of_chunks
    self.search_mode = search_mode
    self.ranker_strategy = ranker_strategy
    self.ranker_k = ranker_k
    self.ranker_alpha = ranker_alpha

Functions

retrieve

retrieve(query: str, **kwargs) -> list[dict]

Retrieve relevant documents from vector store.

Parameters:

  • query (str) –

    question for which documents should be retrieved.

Returns:

  • list[Document]

    list of documents with their metadata corresponding to the query.

Source code in ai4rag/rag/retrieval/retriever.py
def retrieve(self, query: str, **kwargs) -> list[dict]:
    """Retrieve relevant documents from vector store.

    Parameters
    ----------
    query : str
        question for which documents should be retrieved.

    Returns
    -------
    list[Document]
        list of documents with their metadata corresponding to the query.
    """
    _number_of_chunks = kwargs.get("number_of_chunks", self.number_of_chunks)

    return self._vector_store.search(
        query,
        k=_number_of_chunks,
        search_mode=self.search_mode,
        ranker_strategy=self.ranker_strategy,
        ranker_k=self.ranker_k,
        ranker_alpha=self.ranker_alpha,
    )