rag.search

class CrossEncoderReranker(model_name='cross-encoder/ms-marco-MiniLM-L-6-v2', max_length=512, fastembed_model_name=None)[source]

Bases: RerankerBase

Re-ranker using cross-encoder models.

Tries sentence-transformers first, then falls back to fastembed (ONNX-based). Requires at least one of:

pip install sentence-transformers pip install fastembed

Parameters:

model_name (str) – Cross-encoder model name.
max_length (int) – Maximum sequence length (sentence-transformers only).
fastembed_model_name (str | None) – Model name for fastembed fallback.

rerank(query, results, n_results=None)[source]

Re-rank results using cross-encoder scoring.

Parameters:

query (str) – Original search query.
results (list[dict[str, Any]]) – Search results to re-rank.
n_results (int | None) – Number of results to return.

Returns:

Re-ranked results with updated scores.

Return type:

list[dict[str, Any]]

class DiversityReranker(lambda_param=0.7)[source]

Bases: RerankerBase

Re-ranker that promotes result diversity.

Re-orders results to maximize diversity based on content similarity. Uses MMR (Maximal Marginal Relevance) style selection.

Parameters:: lambda_param (float) – Balance between relevance and diversity (0-1). Higher = more relevance, lower = more diversity.

rerank(query, results, n_results=None)[source]

Re-rank results for diversity.

Parameters:

query (str) – Original search query (not used, kept for interface).
results (list[dict[str, Any]]) – Search results to re-rank.
n_results (int | None) – Number of results to return.

Returns:

Diversified results.

Return type:

list[dict[str, Any]]

class HybridSearch(vector_search_fn=None, bm25_search_fn=None, alpha=0.6, rrf_k=60)[source]

Bases: object

Hybrid search combining vector similarity and BM25 keyword matching.

Uses reciprocal rank fusion (RRF) to combine results from multiple retrieval methods.

Parameters:

vector_search_fn (Callable[[str, int], list[dict[str, Any]]] | None) – Function that takes (query, n_results) and returns vector results.
bm25_search_fn (Callable[[str, int], list[dict[str, Any]]] | None) – Function that takes (query, n_results) and returns BM25 results.
alpha (float) – Weight for vector search (1-alpha for BM25). Default 0.6.
rrf_k (int) – RRF constant for rank fusion. Default 60.

search(query, n_results=5, mode='hybrid')[source]

Perform search using the specified mode.

Parameters:

query (str) – Search query.
n_results (int) – Maximum number of results.
mode (str) – Search mode - “vector”, “bm25”, or “hybrid”.

Returns:

List of search results with merged scores.

Return type:

list[dict[str, Any]]

set_alpha(alpha)[source]

Set the alpha parameter (vector weight).

Parameters:: alpha (float) – Weight for vector search (0.0 to 1.0).
Return type:: None

set_search_functions(vector_search_fn=None, bm25_search_fn=None)[source]

Set or update search functions.

Parameters:

vector_search_fn (Callable[[str, int], list[dict[str, Any]]] | None) – Vector search function.
bm25_search_fn (Callable[[str, int], list[dict[str, Any]]] | None) – BM25 search function.

Return type:

None

class NoopReranker[source]

Bases: RerankerBase

Pass-through re-ranker that does nothing.

Used when re-ranking is disabled but a reranker interface is expected.

rerank(query, results, n_results=None)[source]

Return results unchanged.

Parameters:

query (str) – Original search query (ignored).
results (list[dict[str, Any]]) – Search results.
n_results (int | None) – Number of results to return.

Returns:

Original results, optionally truncated.

Return type:

list[dict[str, Any]]

class QueryExpander(llm_generate_fn=None, n_variations=3, include_original=True)[source]

Bases: object

Expand queries using LLM for multi-query retrieval.

Generates multiple reformulations of a search query to improve recall by catching different phrasings and aspects of the topic.

Parameters:

llm_generate_fn (Callable[[str], Any] | None) – Function that takes a prompt string and returns LLM response.
n_variations (int) – Number of query variations to generate (default: 3).
include_original (bool) – Whether to include original query in results (default: True).

Example

>>> def mock_llm(prompt): return "1. What is auth?\n2. How to authenticate?"
>>> expander = QueryExpander(llm_generate_fn=mock_llm, n_variations=2)
>>> queries = expander.expand("authentication methods")
>>> # Returns: ["authentication methods", "What is auth?", "How to authenticate?"]

expand(query)[source]

Generate query variations for improved retrieval.

Parameters:: query (str) – Original search query.
Returns:: List of query variations including original (if include_original=True).
Return type:: list[str]

set_llm_function(fn)[source]

Set or update the LLM generate function.

Parameters:: fn (Callable[[str], Any]) – Function that takes a prompt and returns LLM response.
Return type:: None

class RerankerBase[source]

Bases: object

Base class for re-rankers.

rerank(query, results, n_results=None)[source]

Re-rank search results.

Parameters:

query (str) – Original search query.
results (list[dict[str, Any]]) – Search results to re-rank.
n_results (int | None) – Number of results to return (None = all).

Returns:

Re-ranked results.

Return type:

list[dict[str, Any]]

fuse_search_results(result_lists, n_results=5, dedupe_by='id')[source]

Fuse results from multiple searches with deduplication.

Combines results from multiple query variations, removing duplicates while preserving relevance ordering.

Parameters:

result_lists (list[list[dict[str, Any]]]) – List of search result lists from different queries.
n_results (int) – Maximum number of final results to return.
dedupe_by (str) – Field to use for deduplication (default: “id”).

Returns:

Fused and deduplicated list of results.

Return type:

list[dict[str, Any]]

get_reranker(reranker_type='none', **kwargs)[source]

Get a reranker by type name.

Parameters:

reranker_type (str) – Type of reranker (“cross_encoder”, “diversity”, “none”).
**kwargs (Any) – Additional arguments for the reranker.

Returns:

Configured reranker instance.

Return type:

RerankerBase

reciprocal_rank_fusion(result_lists, k=60, weights=None)[source]

Combine multiple result lists using reciprocal rank fusion.

Parameters:

result_lists (list[list[dict[str, Any]]]) – List of result lists to fuse.
k (int) – RRF constant.
weights (list[float] | None) – Optional weights for each result list.

Returns:

Fused and sorted results.

Return type:

list[dict[str, Any]]

Classes

class CrossEncoderReranker(model_name='cross-encoder/ms-marco-MiniLM-L-6-v2', max_length=512, fastembed_model_name=None)[source]

Bases: RerankerBase

Re-ranker using cross-encoder models.

Tries sentence-transformers first, then falls back to fastembed (ONNX-based). Requires at least one of:

pip install sentence-transformers pip install fastembed

Parameters:

model_name (str) – Cross-encoder model name.
max_length (int) – Maximum sequence length (sentence-transformers only).
fastembed_model_name (str | None) – Model name for fastembed fallback.

rerank(query, results, n_results=None)[source]

Re-rank results using cross-encoder scoring.

Parameters:

query (str) – Original search query.
results (list[dict[str, Any]]) – Search results to re-rank.
n_results (int | None) – Number of results to return.

Returns:

Re-ranked results with updated scores.

Return type:

list[dict[str, Any]]

class DiversityReranker(lambda_param=0.7)[source]

Bases: RerankerBase

Re-ranker that promotes result diversity.

Re-orders results to maximize diversity based on content similarity. Uses MMR (Maximal Marginal Relevance) style selection.

Parameters:: lambda_param (float) – Balance between relevance and diversity (0-1). Higher = more relevance, lower = more diversity.

rerank(query, results, n_results=None)[source]

Re-rank results for diversity.

Parameters:

query (str) – Original search query (not used, kept for interface).
results (list[dict[str, Any]]) – Search results to re-rank.
n_results (int | None) – Number of results to return.

Returns:

Diversified results.

Return type:

list[dict[str, Any]]

class HybridSearch(vector_search_fn=None, bm25_search_fn=None, alpha=0.6, rrf_k=60)[source]

Bases: object

Hybrid search combining vector similarity and BM25 keyword matching.

Uses reciprocal rank fusion (RRF) to combine results from multiple retrieval methods.

Parameters:

vector_search_fn (Callable[[str, int], list[dict[str, Any]]] | None) – Function that takes (query, n_results) and returns vector results.
bm25_search_fn (Callable[[str, int], list[dict[str, Any]]] | None) – Function that takes (query, n_results) and returns BM25 results.
alpha (float) – Weight for vector search (1-alpha for BM25). Default 0.6.
rrf_k (int) – RRF constant for rank fusion. Default 60.

search(query, n_results=5, mode='hybrid')[source]

Perform search using the specified mode.

Parameters:

query (str) – Search query.
n_results (int) – Maximum number of results.
mode (str) – Search mode - “vector”, “bm25”, or “hybrid”.

Returns:

List of search results with merged scores.

Return type:

list[dict[str, Any]]

set_alpha(alpha)[source]

Set the alpha parameter (vector weight).

Parameters:: alpha (float) – Weight for vector search (0.0 to 1.0).
Return type:: None

set_search_functions(vector_search_fn=None, bm25_search_fn=None)[source]

Set or update search functions.

Parameters:

vector_search_fn (Callable[[str, int], list[dict[str, Any]]] | None) – Vector search function.
bm25_search_fn (Callable[[str, int], list[dict[str, Any]]] | None) – BM25 search function.

Return type:

None

class NoopReranker[source]

Bases: RerankerBase

Pass-through re-ranker that does nothing.

Used when re-ranking is disabled but a reranker interface is expected.

rerank(query, results, n_results=None)[source]

Return results unchanged.

Parameters:

query (str) – Original search query (ignored).
results (list[dict[str, Any]]) – Search results.
n_results (int | None) – Number of results to return.

Returns:

Original results, optionally truncated.

Return type:

list[dict[str, Any]]

class QueryExpander(llm_generate_fn=None, n_variations=3, include_original=True)[source]

Bases: object

Expand queries using LLM for multi-query retrieval.

Generates multiple reformulations of a search query to improve recall by catching different phrasings and aspects of the topic.

Parameters:

llm_generate_fn (Callable[[str], Any] | None) – Function that takes a prompt string and returns LLM response.
n_variations (int) – Number of query variations to generate (default: 3).
include_original (bool) – Whether to include original query in results (default: True).

Example

>>> def mock_llm(prompt): return "1. What is auth?\n2. How to authenticate?"
>>> expander = QueryExpander(llm_generate_fn=mock_llm, n_variations=2)
>>> queries = expander.expand("authentication methods")
>>> # Returns: ["authentication methods", "What is auth?", "How to authenticate?"]

expand(query)[source]

Generate query variations for improved retrieval.

Parameters:: query (str) – Original search query.
Returns:: List of query variations including original (if include_original=True).
Return type:: list[str]

set_llm_function(fn)[source]

Set or update the LLM generate function.

Parameters:: fn (Callable[[str], Any]) – Function that takes a prompt and returns LLM response.
Return type:: None

class RerankerBase[source]

Bases: object

Base class for re-rankers.

rerank(query, results, n_results=None)[source]

Re-rank search results.

Parameters:

query (str) – Original search query.
results (list[dict[str, Any]]) – Search results to re-rank.
n_results (int | None) – Number of results to return (None = all).

Returns:

Re-ranked results.

Return type:

list[dict[str, Any]]

Functions

fuse_search_results(result_lists, n_results=5, dedupe_by='id')[source]

Fuse results from multiple searches with deduplication.

Combines results from multiple query variations, removing duplicates while preserving relevance ordering.

Parameters:

result_lists (list[list[dict[str, Any]]]) – List of search result lists from different queries.
n_results (int) – Maximum number of final results to return.
dedupe_by (str) – Field to use for deduplication (default: “id”).

Returns:

Fused and deduplicated list of results.

Return type:

list[dict[str, Any]]

get_reranker(reranker_type='none', **kwargs)[source]

Get a reranker by type name.

Parameters:

reranker_type (str) – Type of reranker (“cross_encoder”, “diversity”, “none”).
**kwargs (Any) – Additional arguments for the reranker.

Returns:

Configured reranker instance.

Return type:

RerankerBase

reciprocal_rank_fusion(result_lists, k=60, weights=None)[source]

Combine multiple result lists using reciprocal rank fusion.

Parameters:

result_lists (list[list[dict[str, Any]]]) – List of result lists to fuse.
k (int) – RRF constant.
weights (list[float] | None) – Optional weights for each result list.

Returns:

Fused and sorted results.

Return type:

list[dict[str, Any]]