feat: add search space validation step for fail-fast on misconfiguration

ai4rag/components/optimization/rag_templates_optimization.py

@@ -161,10 +161,10 @@ def run_rag_optimization(  # pylint: disable=too-many-locals,too-many-arguments,
 
     for param_name, values in search_space_raw.items():
         if param_name == "foundation_model":
-            values = [_deserialize_model(m, ogx_client) for m in values]
+            values = [deserialize_model(m, ogx_client) for m in values]
             foundation_models = values
         elif param_name == "embedding_model":
-            values = [_deserialize_model(m, ogx_client) for m in values]
+            values = [deserialize_model(m, ogx_client) for m in values]
             embedding_models = values
         params.append(Parameter(param_name, "C", values=values))
 
@@ -274,13 +274,13 @@ def _generate_output_artifacts(
     return patterns
 
 
-def _deserialize_model(data: dict[str, Any], ogx_client: OgxClient) -> OGXEmbeddingModel | OGXFoundationModel:
+def deserialize_model(data: dict[str, Any], ogx_client: OgxClient) -> OGXEmbeddingModel | OGXFoundationModel:
     """Reconstruct a model instance from its serialized dictionary.
 
     Parameters
     ----------
     data
-        Dictionary produced by :func:`_serialize_model` in the search-space
+        Dictionary produced by ``serialize_model`` in the search-space
         preparation step.
     ogx_client
         Client bound to the reconstructed model instance.

ai4rag/components/optimization/search_space_preparation.py

@@ -18,6 +18,7 @@
 from ai4rag.rag.embedding.base_model import BaseEmbeddingModel
 from ai4rag.rag.foundation_models.base_model import BaseFoundationModel
 from ai4rag.search_space.prepare.prepare_search_space import prepare_search_space_with_ogx
+from ai4rag.search_space.src.search_space import AI4RAGSearchSpace
 
 _logger = logging.getLogger("search-space-preparation")
 _logger.addHandler(handler)
@@ -29,7 +30,7 @@
 _DEFAULT_SEED = 17
 
 
-def _serialize_model(model: BaseFoundationModel | BaseEmbeddingModel) -> dict[str, Any]:
+def serialize_model(model: BaseFoundationModel | BaseEmbeddingModel) -> dict[str, Any]:
     """Convert a model instance to a plain dictionary with all its settings.
 
     Captures model identifier, type discriminator, inference parameters,
@@ -113,11 +114,12 @@ def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arg
     chunk_sizes: list[int] | None = None,
     chunk_overlaps: list[int] | None = None,
     inference_max_threads: int = 10,
+    pre_validated_search_space: AI4RAGSearchSpace | None = None,
 ) -> SearchSpaceReport:
     """Run model pre-selection and prepare a search-space report.
 
-    Builds an :class:`AI4RAGSearchSpace` from the given model lists, runs
-    :class:`ModelsPreSelector` when the number of models exceeds the
+    Builds an ``AI4RAGSearchSpace`` from the given model lists, runs
+    ``ModelsPreSelector`` when the number of models exceeds the
     configured caps, detects the benchmark language, and returns a
     structured report.
 
@@ -130,7 +132,7 @@ def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arg
         Path to a single DoclingDocument JSON file or a directory of such
         files.
     ogx_client
-        An authenticated :class:`OgxClient` instance.
+        An authenticated ``OgxClient`` instance.
     embedding_models
         Embedding model identifiers.  ``None`` uses the server defaults.
     generation_models
@@ -161,6 +163,13 @@ def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arg
         RAG service during benchmark evaluation.  Lower values reduce
         per-request concurrency (useful when each request carries more
         retrieved context).  Defaults to ``10``.
+    pre_validated_search_space
+        When provided, the function skips model-list validation,
+        payload construction, and the
+        ``prepare_search_space_with_ogx`` call and uses this
+        search space directly.  Pass the result of an earlier
+        validation step to avoid redundant OGX API calls.
+        ``None`` (default) preserves the original behaviour.
 
     Returns
     -------
@@ -179,38 +188,41 @@ def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arg
         validation (wrong type, empty list, or invalid element types).
     SearchSpaceValueError
         If *chunking_methods* contains values not in
-        :attr:`~ai4rag.utils.constants.ChunkingConstraints.METHODS`, or
+        ``ChunkingConstraints.METHODS``, or
         *chunk_sizes* contains values outside
         ``[ChunkingConstraints.MIN_CHUNK_SIZE, ChunkingConstraints.MAX_CHUNK_SIZE]``,
         or *chunk_overlaps* contains values outside
         ``[ChunkingConstraints.MIN_CHUNK_OVERLAP, ChunkingConstraints.MAX_CHUNK_OVERLAP]``.
     """
-    _validate_model_list(embedding_models, "embedding_models")
-    _validate_model_list(generation_models, "generation_models")
-
-    # Build payload and create search space via OGX
-    payload: dict[str, Any] = {}
-    if generation_models:
-        payload["foundation_models"] = [{"model_id": gm} for gm in generation_models]
-    if embedding_models:
-        payload["embedding_models"] = [{"model_id": em} for em in embedding_models]
-    if chunking_methods is not None:
-        payload["chunking_methods"] = chunking_methods
-    if chunk_sizes is not None:
-        payload["chunk_sizes"] = chunk_sizes
-    if chunk_overlaps is not None:
-        payload["chunk_overlaps"] = chunk_overlaps
-
-    # Load benchmark data and documents
+    if pre_validated_search_space is not None:
+        search_space = pre_validated_search_space
+    else:
+        _validate_model_list(embedding_models, "embedding_models")
+        _validate_model_list(generation_models, "generation_models")
+
+        payload: dict[str, Any] = {}
+        if generation_models:
+            payload["foundation_models"] = [{"model_id": gm} for gm in generation_models]
+        if embedding_models:
+            payload["embedding_models"] = [{"model_id": em} for em in embedding_models]
+        if chunking_methods is not None:
+            payload["chunking_methods"] = chunking_methods
+        if chunk_sizes is not None:
+            payload["chunk_sizes"] = chunk_sizes
+        if chunk_overlaps is not None:
+            payload["chunk_overlaps"] = chunk_overlaps
+
+        benchmark_df = pd.read_json(Path(test_data_path))
+        search_space = prepare_search_space_with_ogx(
+            payload,
+            client=ogx_client,
+            benchmark_data=benchmark_df,
+        )
+
     benchmark_df = pd.read_json(Path(test_data_path))
     benchmark_data = BenchmarkData(benchmark_df)
     documents = load_docling_documents(extracted_text_path)
 
-    search_space = prepare_search_space_with_ogx(
-        payload,
-        client=ogx_client,
-        benchmark_data=benchmark_df,
-    )
     _logger.info(
         "Search space chunking_method=%s chunk_size=%s chunk_overlap=%s",
         list(search_space["chunking_method"].values),
@@ -253,8 +265,8 @@ def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arg
     verbose_repr: dict[str, Any] = {
         key: list(dict.fromkeys(combo[key] for combo in valid_combinations)) for key in non_model_keys
     }
-    verbose_repr["foundation_model"] = [_serialize_model(m) for m in selected_models["foundation_model"]]
-    verbose_repr["embedding_model"] = [_serialize_model(m) for m in selected_models["embedding_model"]]
+    verbose_repr["foundation_model"] = [serialize_model(m) for m in selected_models["foundation_model"]]
+    verbose_repr["embedding_model"] = [serialize_model(m) for m in selected_models["embedding_model"]]
 
     return SearchSpaceReport(
         search_space=verbose_repr,

tests/unit/ai4rag/components/optimization/test_search_space_prep.py

@@ -201,7 +201,7 @@ def test_recursive_with_too_small_chunk_sizes_yields_empty_search_space(self, mo
         )
         mocker.patch("ai4rag.components.optimization.search_space_preparation.BenchmarkData")
         mocker.patch(
-            "ai4rag.components.optimization.search_space_preparation._serialize_model",
+            "ai4rag.components.optimization.search_space_preparation.serialize_model",
             return_value={"model_id": "mock"},
         )
 
@@ -216,3 +216,108 @@ def test_recursive_with_too_small_chunk_sizes_yields_empty_search_space(self, mo
         assert result.search_space["chunk_size"] == []
         assert result.search_space["chunk_overlap"] == []
         assert result.search_space["chunking_method"] == []
+
+
+# ---------------------------------------------------------------------------
+# pre_validated_search_space parameter
+# ---------------------------------------------------------------------------
+
+
+class TestPrepareSearchSpaceReportPreValidated:
+    """Test the pre_validated_search_space bypass path."""
+
+    def _make_search_space(self) -> AI4RAGSearchSpace:
+        mock_em = MagicMock()
+        mock_em.params.context_length = None
+        return AI4RAGSearchSpace(
+            params=[
+                Parameter(name=AI4RAGParamNames.FOUNDATION_MODEL, values=(MagicMock(),)),
+                Parameter(name=AI4RAGParamNames.EMBEDDING_MODEL, values=(mock_em,)),
+                Parameter(name=AI4RAGParamNames.CHUNKING_METHOD, values=("recursive",)),
+                Parameter(name=AI4RAGParamNames.CHUNK_SIZE, values=(512,)),
+                Parameter(name=AI4RAGParamNames.CHUNK_OVERLAP, values=(128,)),
+            ]
+        )
+
+    def test_skips_ogx_call_when_pre_validated(self, mocker):
+        """prepare_search_space_with_ogx must not be called when pre_validated_search_space is given."""
+        search_space = self._make_search_space()
+        mock_prepare = mocker.patch(
+            "ai4rag.components.optimization.search_space_preparation.prepare_search_space_with_ogx",
+        )
+        mocker.patch(
+            "ai4rag.components.optimization.search_space_preparation.pd.read_json",
+            return_value=MagicMock(),
+        )
+        mocker.patch(
+            "ai4rag.components.optimization.search_space_preparation.load_docling_documents",
+            return_value=[],
+        )
+        mocker.patch("ai4rag.components.optimization.search_space_preparation.BenchmarkData")
+        mocker.patch(
+            "ai4rag.components.optimization.search_space_preparation.serialize_model",
+            return_value={"model_id": "mock"},
+        )
+
+        prepare_search_space_report(
+            test_data_path="dummy.json",
+            extracted_text_path="dummy_dir",
+            ogx_client=MagicMock(),
+            pre_validated_search_space=search_space,
+        )
+
+        mock_prepare.assert_not_called()
+
+    def test_uses_provided_search_space(self, mocker):
+        """The report must reflect parameters from the pre-validated search space."""
+        search_space = self._make_search_space()
+        mocker.patch(
+            "ai4rag.components.optimization.search_space_preparation.pd.read_json",
+            return_value=MagicMock(),
+        )
+        mocker.patch(
+            "ai4rag.components.optimization.search_space_preparation.load_docling_documents",
+            return_value=[],
+        )
+        mocker.patch("ai4rag.components.optimization.search_space_preparation.BenchmarkData")
+        mocker.patch(
+            "ai4rag.components.optimization.search_space_preparation.serialize_model",
+            return_value={"model_id": "mock"},
+        )
+
+        result = prepare_search_space_report(
+            test_data_path="dummy.json",
+            extracted_text_path="dummy_dir",
+            ogx_client=MagicMock(),
+            pre_validated_search_space=search_space,
+        )
+
+        assert result.search_space["chunk_size"] == [512]
+        assert result.search_space["chunk_overlap"] == [128]
+        assert result.search_space["chunking_method"] == ["recursive"]
+
+    def test_still_loads_documents(self, mocker):
+        """Documents must still be loaded even when pre_validated_search_space is given (needed for MPS)."""
+        search_space = self._make_search_space()
+        mocker.patch(
+            "ai4rag.components.optimization.search_space_preparation.pd.read_json",
+            return_value=MagicMock(),
+        )
+        mock_load_docs = mocker.patch(
+            "ai4rag.components.optimization.search_space_preparation.load_docling_documents",
+            return_value=[],
+        )
+        mocker.patch("ai4rag.components.optimization.search_space_preparation.BenchmarkData")
+        mocker.patch(
+            "ai4rag.components.optimization.search_space_preparation.serialize_model",
+            return_value={"model_id": "mock"},
+        )
+
+        prepare_search_space_report(
+            test_data_path="dummy.json",
+            extracted_text_path="dummy_dir",
+            ogx_client=MagicMock(),
+            pre_validated_search_space=search_space,
+        )
+
+        mock_load_docs.assert_called_once_with("dummy_dir")