diff --git a/ai4rag/components/optimization/rag_templates_optimization.py b/ai4rag/components/optimization/rag_templates_optimization.py index e14d1f6..0189db1 100644 --- a/ai4rag/components/optimization/rag_templates_optimization.py +++ b/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. diff --git a/ai4rag/components/optimization/search_space_preparation.py b/ai4rag/components/optimization/search_space_preparation.py index e473828..4dad72d 100644 --- a/ai4rag/components/optimization/search_space_preparation.py +++ b/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, diff --git a/tests/unit/ai4rag/components/optimization/test_search_space_prep.py b/tests/unit/ai4rag/components/optimization/test_search_space_prep.py index 980ce64..b674524 100644 --- a/tests/unit/ai4rag/components/optimization/test_search_space_prep.py +++ b/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")