babylon.rag.retrieval

Query and retrieval interface for the RAG system.

Uses VectorStoreProtocol from babylon.persistence for backend-agnostic vector storage. Implementations include PgVectorStore (Feature 037).

Classes

QueryResponse(**data)	Represents the complete response to a query.
QueryResult(**data)	Represents a single query result with similarity score.
Retriever(vector_store, embedding_manager)	High-level retrieval interface for RAG queries.

class babylon.rag.retrieval.QueryResult(**data)

Bases: BaseModel

Represents a single query result with similarity score.

Parameters:

• chunk (DocumentChunk)

• similarity_score (float)

• distance (float)

• metadata (dict[str, Any] | None)

model_config: ClassVar[ConfigDict] = {'validate_assignment': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

chunk: DocumentChunk

similarity_score: float

distance: float

metadata: dict[str, Any] | None

convert_similarity_to_distance()

Convert similarity score to distance if not provided.

Return type:

QueryResult

class babylon.rag.retrieval.QueryResponse(**data)

Bases: BaseModel

Represents the complete response to a query.

Parameters:

• query (str)

• results (list[QueryResult])

• total_results (int)

• processing_time_ms (float)

• embedding_time_ms (float)

• search_time_ms (float)

• metadata (dict[str, Any] | None)

model_config: ClassVar[ConfigDict] = {'validate_assignment': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

query: str

results: list[QueryResult]

total_results: int

processing_time_ms: float

embedding_time_ms: float

search_time_ms: float

metadata: dict[str, Any] | None

get_top_k(k)

Get the top k results by similarity score.

Return type:

list[QueryResult]

Parameters:

k (int)

get_combined_context(max_length=4000, separator='\\n\\n')

Combine result chunks into a single context string.

Return type:

str

Parameters:

• max_length (int)

• separator (str)

class babylon.rag.retrieval.Retriever(vector_store, embedding_manager)

Bases: object

High-level retrieval interface for RAG queries.

Uses VectorStoreProtocol for backend-agnostic similarity search.

Parameters:

• vector_store (VectorStoreProtocol)

• embedding_manager (EmbeddingManager)

__init__(vector_store, embedding_manager)

Initialize the retriever.

Parameters:

• vector_store (VectorStoreProtocol) – Any VectorStoreProtocol implementation

• embedding_manager (EmbeddingManager) – EmbeddingManager for query embedding

async aquery(query, k=10, similarity_threshold=0.0, metadata_filter=None)

Asynchronously query for relevant document chunks.

Parameters:

• query (str) – Query text to search for

• k (int) – Number of results to return

• similarity_threshold (float) – Minimum similarity score for results

• metadata_filter (dict[str, Any] | None) – Optional filters for chunk metadata

Return type:

QueryResponse

Returns:

QueryResponse with results and timing information

Raises:

RagError – If query processing fails

query(query, k=10, similarity_threshold=0.0, metadata_filter=None)

Synchronously query for relevant document chunks.

Parameters:

• query (str) – Query text to search for

• k (int) – Number of results to return

• similarity_threshold (float) – Minimum similarity score for results

• metadata_filter (dict[str, Any] | None) – Optional filters for chunk metadata

Return type:

QueryResponse

Returns:

QueryResponse with results and timing information

Raises:

RagError – If query processing fails