Merge branch 'main' into jd/ci-workflow-split

Author: jakubduda-dsai

docs/api_reference/document_search/index.md

@@ -1,3 +1,5 @@
 # Document Search
 
-::: ragbits.document_search.DocumentSearch
+::: ragbits.document_search.DocumentSearchOptions
+
+::: ragbits.document_search.DocumentSearch

examples/document-search/chroma.py

@@ -39,7 +39,7 @@
 from ragbits.core.embeddings.dense import LiteLLMEmbedder, LiteLLMEmbedderOptions
 from ragbits.core.vector_stores.base import VectorStoreOptions
 from ragbits.core.vector_stores.chroma import ChromaVectorStore
-from ragbits.document_search import DocumentSearch, SearchConfig
+from ragbits.document_search import DocumentSearch, DocumentSearchOptions
 from ragbits.document_search.documents.document import DocumentMeta
 
 set_trace_handlers("cli")
@@ -101,14 +101,12 @@ async def main() -> None:
     print([doc.metadata["content"] for doc in all_documents])
 
     query = "I'm boiling my water and I need a joke"
-    vector_store_kwargs = {
-        "k": 2,
-        "score_threshold": 0.4,
-    }
-    results = await document_search.search(
-        query,
-        config=SearchConfig(vector_store_kwargs=vector_store_kwargs),
+    vector_store_options = VectorStoreOptions(
+        k=2,
+        score_threshold=0.4,
     )
+    options = DocumentSearchOptions(vector_store_options=vector_store_options)
+    results = await document_search.search(query, options)
 
     print()
     print(f"Documents similar to: {query}")

examples/document-search/qdrant.py

@@ -37,8 +37,9 @@
 
 from ragbits.core.audit import set_trace_handlers
 from ragbits.core.embeddings.dense import LiteLLMEmbedder
+from ragbits.core.vector_stores.base import VectorStoreOptions
 from ragbits.core.vector_stores.qdrant import QdrantVectorStore
-from ragbits.document_search import DocumentSearch, SearchConfig
+from ragbits.document_search import DocumentSearch, DocumentSearchOptions
 from ragbits.document_search.documents.document import DocumentMeta
 
 set_trace_handlers("cli")
@@ -92,14 +93,12 @@ async def main() -> None:
     print([doc.metadata["content"] for doc in all_documents])
 
     query = "I'm boiling my water and I need a joke"
-    vector_store_kwargs = {
-        "k": 2,
-        "score_threshold": 0.6,
-    }
-    results = await document_search.search(
-        query,
-        config=SearchConfig(vector_store_kwargs=vector_store_kwargs),
+    vector_store_options = VectorStoreOptions(
+        k=2,
+        score_threshold=0.6,
     )
+    options = DocumentSearchOptions(vector_store_options=vector_store_options)
+    results = await document_search.search(query, options=options)
 
     print()
     print(f"Documents similar to: {query}")

packages/ragbits-core/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Adjust typing for DocumentSearch (#554)
 - Add Prometheus & Grafana Monitoring for LLMs Using OpenTelemetry (#427)
 - Restructure audit module (#427)
 

packages/ragbits-core/tests/integration/vector_stores/test_vector_store.py

@@ -233,7 +233,7 @@ async def test_handling_document_ingestion_with_different_content_and_verifying_
     document_1 = DocumentMeta.create_text_document_from_literal(document_1_content)
     document_2 = DocumentMeta.create_text_document_from_literal(document_2_content)
 
-    document_search = DocumentSearch(
+    document_search: DocumentSearch = DocumentSearch(
         vector_store=text_vector_store,
     )
     await document_search.ingest([document_1, document_2])

packages/ragbits-document-search/CHANGELOG.md

@@ -3,6 +3,9 @@
 ## Unreleased
 
 - Update audit imports (#427)
+- BREAKING CHANGE: Adjust document search configurable interface (#554)
+- BREAKING CHANGE: Rename SearchConfig to DocumentSearchOptions (#554)
+- BREAKING CHANGE: Improve typing for SearchConfig (#554)
 
 ## 0.17.1 (2025-05-09)
 

packages/ragbits-document-search/src/ragbits/document_search/__init__.py

@@ -1,3 +1,3 @@
-from ._main import DocumentSearch, SearchConfig
+from ragbits.document_search._main import DocumentSearch, DocumentSearchOptions
 
-__all__ = ["DocumentSearch", "SearchConfig"]
+__all__ = ["DocumentSearch", "DocumentSearchOptions"]

packages/ragbits-document-search/src/ragbits/document_search/__version__.py

@@ -1,51 +1,52 @@
 from collections.abc import Iterable, Sequence
 from pathlib import Path
 from types import ModuleType
-from typing import Any, ClassVar
+from typing import ClassVar, Generic
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 from typing_extensions import Self
 
 from ragbits import document_search
 from ragbits.core.audit.traces import trace, traceable
 from ragbits.core.config import CoreConfig
+from ragbits.core.options import Options
 from ragbits.core.sources.base import Source, SourceResolver
+from ragbits.core.types import NOT_GIVEN, NotGiven
 from ragbits.core.utils._pyproject import get_config_from_yaml
 from ragbits.core.utils.config_handling import (
+    ConfigurableComponent,
     NoPreferredConfigError,
     ObjectConstructionConfig,
-    WithConstructionConfig,
 )
-from ragbits.core.vector_stores import VectorStore
-from ragbits.core.vector_stores.base import VectorStoreOptions
+from ragbits.core.vector_stores.base import VectorStore, VectorStoreOptionsT
 from ragbits.document_search.documents.document import Document, DocumentMeta
 from ragbits.document_search.documents.element import Element
 from ragbits.document_search.ingestion.enrichers.router import ElementEnricherRouter
 from ragbits.document_search.ingestion.parsers.router import DocumentParserRouter
-from ragbits.document_search.ingestion.strategies import (
-    IngestStrategy,
-    SequentialIngestStrategy,
-)
+from ragbits.document_search.ingestion.strategies import IngestStrategy, SequentialIngestStrategy
 from ragbits.document_search.ingestion.strategies.base import IngestExecutionError, IngestExecutionResult
 from ragbits.document_search.retrieval.rephrasers.base import QueryRephraser
 from ragbits.document_search.retrieval.rephrasers.noop import NoopQueryRephraser
-from ragbits.document_search.retrieval.rerankers.base import Reranker, RerankerOptions
+from ragbits.document_search.retrieval.rerankers.base import Reranker, RerankerOptionsT
 from ragbits.document_search.retrieval.rerankers.noop import NoopReranker
 
 
-class SearchConfig(BaseModel):
+class DocumentSearchOptions(Options, Generic[VectorStoreOptionsT, RerankerOptionsT]):
     """
-    Configuration for the search process.
+    Object representing the options for the document search.
+
+    Attributes:
+        vector_store_options: The options for the vector store.
+        reranker_options: The options for the reranker.
     """
 
-    reranker_kwargs: dict[str, Any] = Field(default_factory=dict)
-    vector_store_kwargs: dict[str, Any] = Field(default_factory=dict)
-    embedder_kwargs: dict[str, Any] = Field(default_factory=dict)
+    vector_store_options: VectorStoreOptionsT | None | NotGiven = NOT_GIVEN
+    reranker_options: RerankerOptionsT | None | NotGiven = NOT_GIVEN
 
 
 class DocumentSearchConfig(BaseModel):
     """
-    Schema for the dict taken by DocumentSearch.from_config method.
+    Schema for the document search config.
     """
 
     vector_store: ObjectConstructionConfig
@@ -56,39 +57,49 @@ class DocumentSearchConfig(BaseModel):
     enricher_router: dict[str, ObjectConstructionConfig] = {}
 
 
-class DocumentSearch(WithConstructionConfig):
+class DocumentSearch(ConfigurableComponent[DocumentSearchOptions[VectorStoreOptionsT, RerankerOptionsT]]):
     """
-    A main entrypoint to the DocumentSearch functionality.
-
-    It provides methods for both ingestion and retrieval.
+    Main entrypoint to the document search functionality. It provides methods for document retrieval and ingestion.
 
     Retrieval:
-
         1. Uses QueryRephraser to rephrase the query.
-        2. Uses VectorStore to retrieve the most relevant chunks.
-        3. Uses Reranker to rerank the chunks.
+        2. Uses VectorStore to retrieve the most relevant elements.
+        3. Uses Reranker to rerank the elements.
+
+    Ingestion:
+        1. Uses IngestStrategy to orchestrate ingestion process.
+        2. Uses DocumentParserRouter to route the document to the appropriate DocumentParser to parse the content.
+        3. Uses ElementEnricherRouter to redirect the element to the appropriate ElementEnricher to enrich the element.
     """
 
+    options_cls: type[DocumentSearchOptions] = DocumentSearchOptions
     default_module: ClassVar[ModuleType | None] = document_search
     configuration_key: ClassVar[str] = "document_search"
 
-    vector_store: VectorStore
-    query_rephraser: QueryRephraser
-    reranker: Reranker
-
-    ingest_strategy: IngestStrategy
-    parser_router: DocumentParserRouter
-    enricher_router: ElementEnricherRouter
-
     def __init__(
         self,
-        vector_store: VectorStore,
+        vector_store: VectorStore[VectorStoreOptionsT],
+        *,
         query_rephraser: QueryRephraser | None = None,
-        reranker: Reranker | None = None,
+        reranker: Reranker[RerankerOptionsT] | None = None,
+        default_options: DocumentSearchOptions[VectorStoreOptionsT, RerankerOptionsT] | None = None,
         ingest_strategy: IngestStrategy | None = None,
         parser_router: DocumentParserRouter | None = None,
         enricher_router: ElementEnricherRouter | None = None,
     ) -> None:
+        """
+        Initialize the DocumentSearch instance.
+
+        Args:
+            vector_store: The vector store to use for retrieval.
+            query_rephraser: The query rephraser to use for retrieval.
+            reranker: The reranker to use for retrieval.
+            default_options: The default options for the search.
+            ingest_strategy: The ingestion strategy to use for ingestion.
+            parser_router: The document parser router to use for ingestion.
+            enricher_router: The element enricher router to use for ingestion.
+        """
+        super().__init__(default_options=default_options)
         self.vector_store = vector_store
         self.query_rephraser = query_rephraser or NoopQueryRephraser()
         self.reranker = reranker or NoopReranker()
@@ -178,39 +189,47 @@ def preferred_subclass(
 
         raise NoPreferredConfigError(f"Could not find preferred factory or configuration for {cls.configuration_key}")
 
-    async def search(self, query: str, config: SearchConfig | None = None) -> Sequence[Element]:
+    async def search(
+        self,
+        query: str,
+        options: DocumentSearchOptions[VectorStoreOptionsT, RerankerOptionsT] | None = None,
+    ) -> Sequence[Element]:
         """
         Search for the most relevant chunks for a query.
 
         Args:
             query: The query to search for.
-            config: The search configuration.
+            options: The document search retrieval options.
 
         Returns:
             A list of chunks.
         """
-        config = config or SearchConfig()
-        queries = await self.query_rephraser.rephrase(query)
-        with trace(queries=queries, config=config, vectore_store=self.vector_store, reranker=self.reranker) as outputs:
-            elements = []
-
-            for rephrased_query in queries:
-                results = await self.vector_store.retrieve(
-                    text=rephrased_query,
-                    options=VectorStoreOptions(**config.vector_store_kwargs),
-                )
-                elements.append([Element.from_vector_db_entry(result.entry, result.score) for result in results])
-
-            outputs.search_results = await self.reranker.rerank(
+        merged_options = (self.default_options | options) if options else self.default_options
+        vector_store_options = merged_options.vector_store_options or None
+        reranker_options = merged_options.reranker_options or None
+
+        with trace(query=query, options=merged_options) as outputs:
+            queries = await self.query_rephraser.rephrase(query)
+            elements = [
+                [
+                    Element.from_vector_db_entry(result.entry, result.score)
+                    for result in await self.vector_store.retrieve(query, vector_store_options)
+                ]
+                for query in queries
+            ]
+            outputs.results = await self.reranker.rerank(
                 elements=elements,
                 query=query,
-                options=RerankerOptions(**config.reranker_kwargs),
+                options=reranker_options,
             )
-            return outputs.search_results
+
+        return outputs.results
 
     @traceable
     async def ingest(
-        self, documents: str | Iterable[DocumentMeta | Document | Source], fail_on_error: bool = True
+        self,
+        documents: str | Iterable[DocumentMeta | Document | Source],
+        fail_on_error: bool = True,
     ) -> IngestExecutionResult:
         """
         Ingest documents into the search index.

packages/ragbits-document-search/src/ragbits/document_search/_main.py

@@ -8,7 +8,8 @@
 
 from ragbits.cli._utils import get_instance_or_exit
 from ragbits.cli.state import print_output
-from ragbits.document_search._main import DocumentSearch, SearchConfig
+from ragbits.core.vector_stores.base import VectorStoreOptions
+from ragbits.document_search._main import DocumentSearch, DocumentSearchOptions
 
 ds_app = typer.Typer(no_args_is_help=True)
 
@@ -84,8 +85,8 @@ async def run() -> None:
         if state.document_search is None:
             raise ValueError("Document search not initialized")
 
-        search_config = SearchConfig(vector_store_kwargs={"k": k})
-        entries = await state.document_search.search(query, config=search_config)
+        options: DocumentSearchOptions = DocumentSearchOptions(vector_store_options=VectorStoreOptions(k=k))
+        entries = await state.document_search.search(query, options)
         print_output(entries, columns=columns)
 
     asyncio.run(run())