From 071c638cfd4028759704196e2aa535d7fb01bdf3 Mon Sep 17 00:00:00 2001 From: Tomasz Guzik Date: Thu, 6 Aug 2026 07:21:09 +0200 Subject: [PATCH 1/6] =?UTF-8?q?feat:=203=20files=20modified:=20=20=201.=20?= =?UTF-8?q?ai4rag/utils/validators.py=20=E2=80=94=20added=20validate=5Fmod?= =?UTF-8?q?el=5Flist()=20function=20(moved=20from=20search=5Fspace=5Fprepa?= =?UTF-8?q?ration.py)=20=20=202.=20ai4rag/components/optimization/search?= =?UTF-8?q?=5Fspace=5Fpreparation.py=20=E2=80=94=20=20=20=20=20-=20Imports?= =?UTF-8?q?=20validate=5Fmodel=5Flist=20from=20shared=20location=20+=20AI4?= =?UTF-8?q?RAGSearchSpace=20=20=20=20=20-=20Adds=20=5Fvalidate=5Fmodel=5Fl?= =?UTF-8?q?ist=20=3D=20validate=5Fmodel=5Flist=20alias=20for=20backward=20?= =?UTF-8?q?compat=20=20=20=20=20-=20Adds=20pre=5Fvalidated=5Fsearch=5Fspac?= =?UTF-8?q?e:=20AI4RAGSearchSpace=20|=20None=20=3D=20None=20parameter=20to?= =?UTF-8?q?=20prepare=5Fsearch=5Fspace=5Freport()=20=20=20=20=20-=20When?= =?UTF-8?q?=20provided:=20skips=20model=20validation,=20payload=20construc?= =?UTF-8?q?tion,=20and=20prepare=5Fsearch=5Fspace=5Fwith=5Fogx()=20call=20?= =?UTF-8?q?=20=20=20=20-=20When=20None:=20existing=20behavior=20unchanged?= =?UTF-8?q?=20=20=203.=20tests/unit/ai4rag/components/optimization/test=5F?= =?UTF-8?q?search=5Fspace=5Fprep.py=20=E2=80=94=20=20=20=20=20-=20Added=20?= =?UTF-8?q?TestValidateModelListShared=20(verifies=20alias=20identity)=20?= =?UTF-8?q?=20=20=20=20-=20Added=20TestPrepareSearchSpaceReportPreValidate?= =?UTF-8?q?d=20with=203=20tests:=20skips=20OGX=20call,=20uses=20provided?= =?UTF-8?q?=20search=20space,=20still=20loads=20documents?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Tomasz Guzik --- .../optimization/search_space_preparation.py | 68 +++++----- ai4rag/utils/validators.py | 11 ++ .../optimization/test_search_space_prep.py | 119 ++++++++++++++++++ 3 files changed, 167 insertions(+), 31 deletions(-) diff --git a/ai4rag/components/optimization/search_space_preparation.py b/ai4rag/components/optimization/search_space_preparation.py index e4738287..ee351b7e 100644 --- a/ai4rag/components/optimization/search_space_preparation.py +++ b/ai4rag/components/optimization/search_space_preparation.py @@ -18,10 +18,14 @@ 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 +from ai4rag.utils.validators import validate_model_list _logger = logging.getLogger("search-space-preparation") _logger.addHandler(handler) +_validate_model_list = validate_model_list + _DEFAULT_METRIC = "faithfulness" _DEFAULT_TOP_N_GENERATION = 3 _DEFAULT_TOP_K_EMBEDDING = 2 @@ -113,6 +117,7 @@ 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. @@ -161,6 +166,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 + :func:`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 ------- @@ -185,32 +197,35 @@ def prepare_search_space_report( # pylint: disable=too-many-locals,too-many-arg 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), @@ -262,12 +277,3 @@ def prepare_search_space_report( # pylint: disable=too-many-locals,too-many-arg ) -def _validate_model_list(models: list[str] | None, name: str) -> None: - """Validate that a model list, if provided, contains only non-empty strings.""" - if models is None: - return - if not isinstance(models, list): - raise TypeError(f"{name} must be a list.") - for i, m in enumerate(models): - if not m: - raise TypeError(f"{name}[{i}] must be a non-empty string.") diff --git a/ai4rag/utils/validators.py b/ai4rag/utils/validators.py index 8d449312..c6dad897 100644 --- a/ai4rag/utils/validators.py +++ b/ai4rag/utils/validators.py @@ -66,3 +66,14 @@ def validate(self, obj, value: T): f"on attribute {self.private_name[1:]}." ) return value + + +def validate_model_list(models: list[str] | None, name: str) -> None: + """Validate that a model list, if provided, contains only non-empty strings.""" + if models is None: + return + if not isinstance(models, list): + raise TypeError(f"{name} must be a list.") + for i, m in enumerate(models): + if not m: + raise TypeError(f"{name}[{i}] must be a non-empty string.") 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 980ce647..f2deda24 100644 --- a/tests/unit/ai4rag/components/optimization/test_search_space_prep.py +++ b/tests/unit/ai4rag/components/optimization/test_search_space_prep.py @@ -14,6 +14,7 @@ _validate_model_list, prepare_search_space_report, ) +from ai4rag.utils.validators import validate_model_list from ai4rag.search_space.src.parameter import Parameter from ai4rag.search_space.src.search_space import AI4RAGSearchSpace from ai4rag.utils.constants import AI4RAGParamNames @@ -216,3 +217,121 @@ 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"] == [] + + +# --------------------------------------------------------------------------- +# validate_model_list shared location +# --------------------------------------------------------------------------- + + +class TestValidateModelListShared: + """Verify that the shared validator is accessible from both locations.""" + + def test_shared_validator_is_same_function(self): + """The alias in search_space_preparation must point to the shared function.""" + assert _validate_model_list is validate_model_list + + +# --------------------------------------------------------------------------- +# 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") From b60977aa4f44694bf904383b79460c3fba1ee3a1 Mon Sep 17 00:00:00 2001 From: Tomasz Guzik Date: Thu, 6 Aug 2026 07:36:06 +0200 Subject: [PATCH 2/6] chore: removing NL Signed-off-by: Tomasz Guzik --- ai4rag/components/optimization/search_space_preparation.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/ai4rag/components/optimization/search_space_preparation.py b/ai4rag/components/optimization/search_space_preparation.py index ee351b7e..9b1a17bc 100644 --- a/ai4rag/components/optimization/search_space_preparation.py +++ b/ai4rag/components/optimization/search_space_preparation.py @@ -275,5 +275,3 @@ def prepare_search_space_report( # pylint: disable=too-many-locals,too-many-arg search_space=verbose_repr, selected_models=selected_models, ) - - From cfeac1542d2c8395c8c5fe9d2fbaad1df9223614 Mon Sep 17 00:00:00 2001 From: Tomasz Guzik Date: Thu, 6 Aug 2026 08:11:24 +0200 Subject: [PATCH 3/6] chore: cleanups lint, isort Signed-off-by: Tomasz Guzik --- .../ai4rag/components/optimization/test_search_space_prep.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 f2deda24..721cb137 100644 --- a/tests/unit/ai4rag/components/optimization/test_search_space_prep.py +++ b/tests/unit/ai4rag/components/optimization/test_search_space_prep.py @@ -14,10 +14,10 @@ _validate_model_list, prepare_search_space_report, ) -from ai4rag.utils.validators import validate_model_list from ai4rag.search_space.src.parameter import Parameter from ai4rag.search_space.src.search_space import AI4RAGSearchSpace from ai4rag.utils.constants import AI4RAGParamNames +from ai4rag.utils.validators import validate_model_list # --------------------------------------------------------------------------- # Fixtures From 88dd0a411796853ef1488ce3e93604d9a4d0060d Mon Sep 17 00:00:00 2001 From: Tomasz Guzik Date: Thu, 6 Aug 2026 12:09:08 +0200 Subject: [PATCH 4/6] chore: cleanups lint, isort, docs Signed-off-by: Tomasz Guzik --- .../ogx_inference_template.ipynb | 4 +--- .../optimization/search_space_preparation.py | 24 ++++++++++++------- ai4rag/utils/validators.py | 11 --------- .../optimization/test_search_space_prep.py | 14 ----------- tests/unit/ai4rag/core/hpo/test_gam_opt.py | 4 +--- 5 files changed, 18 insertions(+), 39 deletions(-) diff --git a/ai4rag/components/assets_generator/notebook_templates/ogx_inference_template.ipynb b/ai4rag/components/assets_generator/notebook_templates/ogx_inference_template.ipynb index f7594413..8a24561d 100644 --- a/ai4rag/components/assets_generator/notebook_templates/ogx_inference_template.ipynb +++ b/ai4rag/components/assets_generator/notebook_templates/ogx_inference_template.ipynb @@ -172,9 +172,7 @@ "\n", "# Build retriever kwargs, only including non-None hybrid search parameters\n", "retriever_kwargs = dict(\n", - " ((\"vector_store\", ogx_vector_store),\n", - " (\"method\", method),\n", - " (\"number_of_chunks\", number_of_chunks))\n", + " ((\"vector_store\", ogx_vector_store), (\"method\", method), (\"number_of_chunks\", number_of_chunks))\n", ")\n", "if search_mode is not None:\n", " retriever_kwargs[\"search_mode\"] = search_mode\n", diff --git a/ai4rag/components/optimization/search_space_preparation.py b/ai4rag/components/optimization/search_space_preparation.py index 9b1a17bc..95785135 100644 --- a/ai4rag/components/optimization/search_space_preparation.py +++ b/ai4rag/components/optimization/search_space_preparation.py @@ -19,13 +19,10 @@ 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 -from ai4rag.utils.validators import validate_model_list _logger = logging.getLogger("search-space-preparation") _logger.addHandler(handler) -_validate_model_list = validate_model_list - _DEFAULT_METRIC = "faithfulness" _DEFAULT_TOP_N_GENERATION = 3 _DEFAULT_TOP_K_EMBEDDING = 2 @@ -121,8 +118,8 @@ def prepare_search_space_report( # pylint: disable=too-many-locals,too-many-arg ) -> 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. @@ -135,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 @@ -169,7 +166,7 @@ def prepare_search_space_report( # pylint: disable=too-many-locals,too-many-arg pre_validated_search_space When provided, the function skips model-list validation, payload construction, and the - :func:`prepare_search_space_with_ogx` call and uses this + ``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. @@ -191,7 +188,7 @@ 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 @@ -275,3 +272,14 @@ def prepare_search_space_report( # pylint: disable=too-many-locals,too-many-arg search_space=verbose_repr, selected_models=selected_models, ) + + +def _validate_model_list(models: list[str] | None, name: str) -> None: + """Validate that a model list, if provided, contains only non-empty strings.""" + if models is None: + return + if not isinstance(models, list): + raise TypeError(f"{name} must be a list.") + for i, m in enumerate(models): + if not m: + raise TypeError(f"{name}[{i}] must be a non-empty string.") diff --git a/ai4rag/utils/validators.py b/ai4rag/utils/validators.py index c6dad897..8d449312 100644 --- a/ai4rag/utils/validators.py +++ b/ai4rag/utils/validators.py @@ -66,14 +66,3 @@ def validate(self, obj, value: T): f"on attribute {self.private_name[1:]}." ) return value - - -def validate_model_list(models: list[str] | None, name: str) -> None: - """Validate that a model list, if provided, contains only non-empty strings.""" - if models is None: - return - if not isinstance(models, list): - raise TypeError(f"{name} must be a list.") - for i, m in enumerate(models): - if not m: - raise TypeError(f"{name}[{i}] must be a non-empty string.") 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 721cb137..ee5b59a4 100644 --- a/tests/unit/ai4rag/components/optimization/test_search_space_prep.py +++ b/tests/unit/ai4rag/components/optimization/test_search_space_prep.py @@ -17,7 +17,6 @@ from ai4rag.search_space.src.parameter import Parameter from ai4rag.search_space.src.search_space import AI4RAGSearchSpace from ai4rag.utils.constants import AI4RAGParamNames -from ai4rag.utils.validators import validate_model_list # --------------------------------------------------------------------------- # Fixtures @@ -219,19 +218,6 @@ def test_recursive_with_too_small_chunk_sizes_yields_empty_search_space(self, mo assert result.search_space["chunking_method"] == [] -# --------------------------------------------------------------------------- -# validate_model_list shared location -# --------------------------------------------------------------------------- - - -class TestValidateModelListShared: - """Verify that the shared validator is accessible from both locations.""" - - def test_shared_validator_is_same_function(self): - """The alias in search_space_preparation must point to the shared function.""" - assert _validate_model_list is validate_model_list - - # --------------------------------------------------------------------------- # pre_validated_search_space parameter # --------------------------------------------------------------------------- diff --git a/tests/unit/ai4rag/core/hpo/test_gam_opt.py b/tests/unit/ai4rag/core/hpo/test_gam_opt.py index eb49ce4f..f8f2113a 100644 --- a/tests/unit/ai4rag/core/hpo/test_gam_opt.py +++ b/tests/unit/ai4rag/core/hpo/test_gam_opt.py @@ -840,9 +840,7 @@ def test_already_seen_reduces_stratified_set(self): {"search_mode": "vector", "chunk_size": 512}, ] # "hybrid" already covered by warm-start; stratification should pull vector first. - result = GAMOptimizer._get_stratified_combinations( - combinations, already_seen={"search_mode": {"hybrid"}} - ) + result = GAMOptimizer._get_stratified_combinations(combinations, already_seen={"search_mode": {"hybrid"}}) assert result[0]["search_mode"] == "vector" def test_already_seen_all_values_skips_stratification(self): From 89fdd13907bdc8f49ec0eca1435819589055b6ad Mon Sep 17 00:00:00 2001 From: Tomasz Guzik Date: Thu, 6 Aug 2026 12:47:05 +0200 Subject: [PATCH 5/6] chore: revert irrelevant changes Signed-off-by: Tomasz Guzik --- .../notebook_templates/ogx_inference_template.ipynb | 4 +++- tests/unit/ai4rag/core/hpo/test_gam_opt.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/ai4rag/components/assets_generator/notebook_templates/ogx_inference_template.ipynb b/ai4rag/components/assets_generator/notebook_templates/ogx_inference_template.ipynb index 8a24561d..f7594413 100644 --- a/ai4rag/components/assets_generator/notebook_templates/ogx_inference_template.ipynb +++ b/ai4rag/components/assets_generator/notebook_templates/ogx_inference_template.ipynb @@ -172,7 +172,9 @@ "\n", "# Build retriever kwargs, only including non-None hybrid search parameters\n", "retriever_kwargs = dict(\n", - " ((\"vector_store\", ogx_vector_store), (\"method\", method), (\"number_of_chunks\", number_of_chunks))\n", + " ((\"vector_store\", ogx_vector_store),\n", + " (\"method\", method),\n", + " (\"number_of_chunks\", number_of_chunks))\n", ")\n", "if search_mode is not None:\n", " retriever_kwargs[\"search_mode\"] = search_mode\n", diff --git a/tests/unit/ai4rag/core/hpo/test_gam_opt.py b/tests/unit/ai4rag/core/hpo/test_gam_opt.py index f8f2113a..eb49ce4f 100644 --- a/tests/unit/ai4rag/core/hpo/test_gam_opt.py +++ b/tests/unit/ai4rag/core/hpo/test_gam_opt.py @@ -840,7 +840,9 @@ def test_already_seen_reduces_stratified_set(self): {"search_mode": "vector", "chunk_size": 512}, ] # "hybrid" already covered by warm-start; stratification should pull vector first. - result = GAMOptimizer._get_stratified_combinations(combinations, already_seen={"search_mode": {"hybrid"}}) + result = GAMOptimizer._get_stratified_combinations( + combinations, already_seen={"search_mode": {"hybrid"}} + ) assert result[0]["search_mode"] == "vector" def test_already_seen_all_values_skips_stratification(self): From 04b5ffcd4ad3d30010b3dc56d28b1325f1be9048 Mon Sep 17 00:00:00 2001 From: Tomasz Guzik Date: Thu, 6 Aug 2026 13:34:55 +0200 Subject: [PATCH 6/6] chore: cr changes Signed-off-by: Tomasz Guzik --- .../components/optimization/rag_templates_optimization.py | 8 ++++---- .../components/optimization/search_space_preparation.py | 6 +++--- .../components/optimization/test_search_space_prep.py | 8 ++++---- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/ai4rag/components/optimization/rag_templates_optimization.py b/ai4rag/components/optimization/rag_templates_optimization.py index e14d1f6e..0189db1d 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 95785135..4dad72d6 100644 --- a/ai4rag/components/optimization/search_space_preparation.py +++ b/ai4rag/components/optimization/search_space_preparation.py @@ -30,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, @@ -265,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 ee5b59a4..b6745240 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"}, ) @@ -255,7 +255,7 @@ def test_skips_ogx_call_when_pre_validated(self, mocker): ) 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"}, ) @@ -281,7 +281,7 @@ def test_uses_provided_search_space(self, mocker): ) 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"}, ) @@ -309,7 +309,7 @@ def test_still_loads_documents(self, mocker): ) 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"}, )