From 1f78bdbb3c3192db850c48f7b6b5a5f8b3a803fb Mon Sep 17 00:00:00 2001 From: Mateusz Switala Date: Thu, 2 Jul 2026 21:34:46 +0200 Subject: [PATCH 01/13] add validation for chunk szies and chunking methods Signed-off-by: Mateusz Switala Assisted-by: Claude Code --- .../optimization/search_space_preparation.py | 63 +++-- ai4rag/search_space/prepare/__init__.py | 5 +- .../prepare/input_payload_types.py | 4 +- .../prepare/prepare_search_space.py | 239 ++++++++++++++---- ai4rag/utils/constants.py | 2 +- .../optimization/test_search_space_prep.py | 152 ++++++++--- .../prepare/test_prepare_search_space.py | 86 ++++++- 7 files changed, 450 insertions(+), 101 deletions(-) diff --git a/ai4rag/components/optimization/search_space_preparation.py b/ai4rag/components/optimization/search_space_preparation.py index ca1d0a5e..126abd24 100644 --- a/ai4rag/components/optimization/search_space_preparation.py +++ b/ai4rag/components/optimization/search_space_preparation.py @@ -17,7 +17,8 @@ from ai4rag.core.experiment.mps import ModelsPreSelector 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.prepare.prepare_search_space import prepare_search_space_custom +from ai4rag.utils.constants import ChunkingConstraints _logger = logging.getLogger("search-space-preparation") _logger.addHandler(handler) @@ -113,6 +114,7 @@ def prepare_search_space_report( # pylint: disable=too-many-locals,too-many-arg sample_size: int = _DEFAULT_SAMPLE_SIZE, random_seed: int = _DEFAULT_SEED, chunking_methods: list[str] | None = None, + chunk_sizes: list[int] | None = None, inference_max_threads: int = 10, ) -> SearchSpaceReport: """Run model pre-selection and prepare a search-space report. @@ -153,6 +155,10 @@ def prepare_search_space_report( # pylint: disable=too-many-locals,too-many-arg search space to only these methods (e.g. ``["recursive"]`` or ``["hybrid"]``). ``None`` uses the platform defaults (both ``"recursive"`` and ``"hybrid"``). + chunk_sizes + When provided, constrains the ``chunk_size`` dimension of the + search space to only these sizes (e.g. ``[256, 512]``). ``None`` + uses the platform defaults. inference_max_threads Maximum number of concurrent threads used when querying the RAG service during benchmark evaluation. Lower values reduce @@ -168,9 +174,11 @@ def prepare_search_space_report( # pylint: disable=too-many-locals,too-many-arg Raises ------ ValueError - If *metric* is not one of the supported values. + If *metric* is not one of the supported values, or if + *chunking_methods* or *chunk_sizes* contain unsupported entries. TypeError - If *embedding_models* or *generation_models* contain invalid entries. + If *embedding_models*, *generation_models*, *chunking_methods*, or + *chunk_sizes* contain invalid entries. """ if metric not in SUPPORTED_METRICS: raise ValueError(f"Metric {metric!r} is not supported. Supported metrics are {list(SUPPORTED_METRICS)}.") @@ -178,20 +186,29 @@ def prepare_search_space_report( # pylint: disable=too-many-locals,too-many-arg _validate_model_list(embedding_models, "embedding_models") _validate_model_list(generation_models, "generation_models") _validate_chunking_methods(chunking_methods) + _validate_chunk_sizes(chunk_sizes) # Build payload and create search space via OGX - payload: dict[str, list[dict[str, str]]] = {} + 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 # Load benchmark data and documents 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) + search_space = prepare_search_space_custom( + payload, + client=ogx_client, + benchmark_data=benchmark_df, + ) # Run model pre-selection when the number of models exceeds the caps fm_values = search_space["foundation_model"].values @@ -230,16 +247,6 @@ def prepare_search_space_report( # pylint: disable=too-many-locals,too-many-arg 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"]] - if chunking_methods is not None: - available = set(verbose_repr["chunking_method"]) - unsupported = [m for m in chunking_methods if m not in available] - if unsupported: - raise ValueError( - f"Unsupported chunking methods: {unsupported!r}. " f"Available methods: {sorted(available)!r}." - ) - verbose_repr["chunking_method"] = chunking_methods - _logger.info("Chunking methods constrained to: %s", verbose_repr["chunking_method"]) - return SearchSpaceReport( search_space=verbose_repr, selected_models=selected_models, @@ -258,7 +265,7 @@ def _validate_model_list(models: list[str] | None, name: str) -> None: def _validate_chunking_methods(methods: list[str] | None) -> None: - """Validate that chunking methods, if provided, are non-empty strings.""" + """Validate that chunking methods, if provided, are non-empty strings and supported.""" if methods is None: return if not isinstance(methods, list): @@ -268,3 +275,27 @@ def _validate_chunking_methods(methods: list[str] | None) -> None: for i, m in enumerate(methods): if not isinstance(m, str) or not m.strip(): raise TypeError(f"chunking_methods[{i}] must be a non-empty string.") + unsupported = [m for m in methods if m not in ChunkingConstraints.METHODS] + if unsupported: + raise ValueError( + f"Unsupported chunking methods: {unsupported!r}. " + f"Supported methods: {ChunkingConstraints.METHODS!r}." + ) + + +def _validate_chunk_sizes(chunk_sizes: list[int] | None) -> None: + """Validate that chunk sizes, if provided, are integers within ChunkingConstraints bounds.""" + if chunk_sizes is None: + return + if not isinstance(chunk_sizes, list): + raise TypeError("chunk_sizes must be a list.") + if not chunk_sizes: + raise ValueError("chunk_sizes must not be empty when provided.") + for i, s in enumerate(chunk_sizes): + if not isinstance(s, int) or s <= 0: + raise TypeError(f"chunk_sizes[{i}] must be a positive integer.") + if not (ChunkingConstraints.MIN_CHUNK_SIZE <= s <= ChunkingConstraints.MAX_CHUNK_SIZE): + raise ValueError( + f"chunk_sizes[{i}]={s} is out of range " + f"[{ChunkingConstraints.MIN_CHUNK_SIZE}, {ChunkingConstraints.MAX_CHUNK_SIZE}]." + ) diff --git a/ai4rag/search_space/prepare/__init__.py b/ai4rag/search_space/prepare/__init__.py index 09b93d8b..bac8837a 100644 --- a/ai4rag/search_space/prepare/__init__.py +++ b/ai4rag/search_space/prepare/__init__.py @@ -3,4 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 # ----------------------------------------------------------------------------- -from ai4rag.search_space.prepare.prepare_search_space import prepare_search_space_with_ogx +from ai4rag.search_space.prepare.prepare_search_space import ( + prepare_search_space_custom, + prepare_search_space_with_ogx, +) diff --git a/ai4rag/search_space/prepare/input_payload_types.py b/ai4rag/search_space/prepare/input_payload_types.py index 2575fb63..c53cb38f 100644 --- a/ai4rag/search_space/prepare/input_payload_types.py +++ b/ai4rag/search_space/prepare/input_payload_types.py @@ -4,7 +4,7 @@ # ----------------------------------------------------------------------------- from typing import Annotated, Optional -from annotated_types import MinLen +from annotated_types import Gt, MinLen from pydantic import ( BaseModel, ConfigDict, @@ -36,3 +36,5 @@ class AI4RAGConstraints(BaseModel): embedding_models: Optional[Annotated[list[AI4RAGEmbeddingModel], MinLen(1)]] = None foundation_models: Optional[Annotated[list[AI4RAGFoundationModel], MinLen(1)]] = None + chunking_methods: Optional[Annotated[list[Annotated[str, MinLen(1)]], MinLen(1)]] = None + chunk_sizes: Optional[Annotated[list[Annotated[int, Gt(0)]], MinLen(1)]] = None diff --git a/ai4rag/search_space/prepare/prepare_search_space.py b/ai4rag/search_space/prepare/prepare_search_space.py index 75018b78..e5d3f9a4 100644 --- a/ai4rag/search_space/prepare/prepare_search_space.py +++ b/ai4rag/search_space/prepare/prepare_search_space.py @@ -18,8 +18,113 @@ from ai4rag.search_space.src.exceptions import SearchSpaceValueError from ai4rag.search_space.src.parameter import Parameter from ai4rag.search_space.src.search_space import AI4RAGSearchSpace +from ai4rag.utils.constants import AI4RAGParamNames -__all__ = ["prepare_search_space_with_ogx"] +__all__ = ["prepare_search_space_with_ogx", "prepare_search_space_custom"] + + +def _resolve_models_from_payload( + validated_payload: AI4RAGConstraints, + client: OgxClient, +) -> tuple[list, list]: + """Retrieve and validate foundation and embedding models from OGX. + + Parameters + ---------- + validated_payload : AI4RAGConstraints + Validated constraint payload specifying which models to include. + client : OgxClient + Authenticated OGX client used for model discovery and validation. + + Returns + ------- + tuple[list, list] + A ``(foundation_models, embedding_models)`` pair of instantiated + model objects ready for use in the search space. + + Raises + ------ + SearchSpaceValueError + When *client* is not an :class:`OgxClient` instance. + """ + if not isinstance(client, OgxClient): + raise SearchSpaceValueError(f"Unrecognized client type: '{client.__class__.__name__}'") + + models = _get_default_ogx_models(client) + + if validated_payload.foundation_models: + foundation_models = _validate_availability_and_create_models( + registered_models=models["foundation_models"], + models_type="llm", + client=client, + provided_models_ids=[m.model_id for m in validated_payload.foundation_models], + ) + else: + foundation_models = _validate_availability_and_create_models( + registered_models=models["foundation_models"], + models_type="llm", + client=client, + ) + + if validated_payload.embedding_models: + embedding_models = _validate_availability_and_create_models( + registered_models=models["embedding_models"], + models_type="embedding", + client=client, + provided_models_ids=[m.model_id for m in validated_payload.embedding_models], + ) + else: + embedding_models = _validate_availability_and_create_models( + registered_models=models["embedding_models"], + models_type="embedding", + client=client, + ) + + return foundation_models, embedding_models + + +def _apply_language_detection(foundation_models: list, benchmark_data: pd.DataFrame) -> None: + """Detect language from benchmark questions and set it on each foundation model in-place. + + Parameters + ---------- + foundation_models : list + Foundation model objects to update with detected language. + benchmark_data : pd.DataFrame + Benchmark data whose ``question`` column is sampled for detection. + """ + for fm in foundation_models: + lang = detect_language_with_llm( + questions=[str(q) for q in benchmark_data["question"][:10]], + generation_model=fm, + ) + if lang is not None: + fm.language = Language(**lang) + logger.info("Model %s: language set to %s (%s).", fm.model_id, lang["name"], lang["code"]) + else: + logger.warning("Model %s: language detection failed, falling back to auto-detect.", fm.model_id) + + +def _build_model_params(foundation_models: list, embedding_models: list) -> tuple[Parameter, Parameter]: + """Create :class:`Parameter` objects for model lists and log selections. + + Parameters + ---------- + foundation_models : list + Foundation model instances to wrap in a Parameter. + embedding_models : list + Embedding model instances to wrap in a Parameter. + + Returns + ------- + tuple[Parameter, Parameter] + ``(fms_param, ems_param)`` ready for inclusion in a search space. + """ + fms_param = Parameter(name="foundation_model", values=foundation_models) + ems_param = Parameter(name="embedding_model", values=embedding_models) + logger.info("Selected foundation models for the experiment: %s.", [m.model_id for m in fms_param.values]) + logger.info("Selected embedding models for the experiment: %s.", [m.model_id for m in ems_param.values]) + return fms_param, ems_param def prepare_search_space_with_ogx( @@ -62,60 +167,98 @@ def prepare_search_space_with_ogx( validated_payload = AI4RAGConstraints(**payload) - if isinstance(client, OgxClient): - models = _get_default_ogx_models(client) - registered_foundation_models = models["foundation_models"] - registered_embedding_models = models["embedding_models"] - else: - raise SearchSpaceValueError(f"Unrecognized client type: '{client.__class__.__name__}'") - - if validated_payload.foundation_models: - foundation_models = _validate_availability_and_create_models( - registered_models=registered_foundation_models, - models_type="llm", - client=client, - provided_models_ids=[m.model_id for m in validated_payload.foundation_models], - ) - else: - foundation_models = _validate_availability_and_create_models( - registered_models=registered_foundation_models, - models_type="llm", - client=client, + _CHUNKING_FIELDS = ("chunking_methods", "chunk_sizes") + skipped = [f for f in _CHUNKING_FIELDS if getattr(validated_payload, f) is not None] + if skipped: + logger.warning( + "Fields %s are not used by prepare_search_space_with_ogx and will be skipped. " + "Use prepare_search_space_custom to apply chunking constraints.", + skipped, ) - if validated_payload.embedding_models: - embedding_models = _validate_availability_and_create_models( - registered_models=registered_embedding_models, - models_type="embedding", - client=client, - provided_models_ids=[m.model_id for m in validated_payload.embedding_models], - ) - else: - embedding_models = _validate_availability_and_create_models( - registered_models=registered_embedding_models, - models_type="embedding", - client=client, - ) + foundation_models, embedding_models = _resolve_models_from_payload(validated_payload, client) if benchmark_data is not None: - for fm in foundation_models: - lang = detect_language_with_llm( - questions=[str(q) for q in benchmark_data["question"][:10]], - generation_model=fm, - ) - if lang is not None: - fm.language = Language(**lang) - logger.info("Model %s: language set to %s (%s).", fm.model_id, lang["name"], lang["code"]) - else: - logger.warning("Model %s: language detection failed, falling back to auto-detect.", fm.model_id) + _apply_language_detection(foundation_models, benchmark_data) - fms_param = Parameter(name="foundation_model", values=foundation_models) - ems_param = Parameter(name="embedding_model", values=embedding_models) - - logger.info("Selected foundation models for the experiment: %s.", [m.model_id for m in fms_param.values]) - logger.info("Selected embedding models for the experiment: %s.", [m.model_id for m in ems_param.values]) + fms_param, ems_param = _build_model_params(foundation_models, embedding_models) return AI4RAGSearchSpace( params=[fms_param, ems_param], vector_store_type=vector_store_type, ) + + +def prepare_search_space_custom( + payload: dict[str, Any], + client: OgxClient, + vector_store_type: str = "ogx", + benchmark_data: pd.DataFrame | None = None, +) -> AI4RAGSearchSpace: + """Prepare an :class:`AI4RAGSearchSpace` with optional chunking customization. + + Extends :func:`prepare_search_space_with_ogx` by allowing the caller to + override the default ``chunking_method`` and ``chunk_size`` dimensions of + the search space via the *payload* dict. All other parameters and rules + remain unchanged. + + Parameters + ---------- + payload : dict[str, Any] + A mapping of constraint names to their values. Supports all keys + accepted by :func:`prepare_search_space_with_ogx` plus: + + - ``"chunking_methods"`` *(list[str])* — overrides the default + ``chunking_method`` dimension (e.g. ``["recursive"]``). + - ``"chunk_sizes"`` *(list[int])* — overrides the default + ``chunk_size`` dimension (e.g. ``[256, 512]``). + + Omitting a key keeps the platform default for that dimension. + + client : OgxClient + Client instance for listing and validating available models. + + vector_store_type : str, default="ogx" + Type of vector store. Supported values: ``"ogx"`` and ``"chroma"``. + When ``"chroma"``, hybrid search parameters are excluded from the + default search space since ChromaDB does not support hybrid search. + + benchmark_data : pd.DataFrame | None, default=None + Benchmark data used for language detection. + If not given, models will use automatic language detection per session. + + Returns + ------- + AI4RAGSearchSpace + A valid AI4RAGSearchSpace used in RAG optimization process. + + Raises + ------ + SearchSpaceValueError + Raised when payload contains a non-recognized parameter name or + when *client* is not an :class:`OgxClient`. + """ + logger.info("Preparing custom search space based on provided constraints: %s.", payload) + + validated_payload = AI4RAGConstraints(**payload) + foundation_models, embedding_models = _resolve_models_from_payload(validated_payload, client) + + if benchmark_data is not None: + _apply_language_detection(foundation_models, benchmark_data) + + fms_param, ems_param = _build_model_params(foundation_models, embedding_models) + + extra_params: list[Parameter] = [] + if validated_payload.chunking_methods is not None: + extra_params.append( + Parameter(name=AI4RAGParamNames.CHUNKING_METHOD, values=tuple(validated_payload.chunking_methods)) + ) + if validated_payload.chunk_sizes is not None: + extra_params.append( + Parameter(name=AI4RAGParamNames.CHUNK_SIZE, values=tuple(validated_payload.chunk_sizes)) + ) + + return AI4RAGSearchSpace( + params=[fms_param, ems_param, *extra_params], + vector_store_type=vector_store_type, + ) diff --git a/ai4rag/utils/constants.py b/ai4rag/utils/constants.py index 2ae93e0c..ed97fb1b 100644 --- a/ai4rag/utils/constants.py +++ b/ai4rag/utils/constants.py @@ -106,7 +106,7 @@ class ChatGenerationConstants(metaclass=ConstantMeta): class ChunkingConstraints(metaclass=ConstantMeta): """Constants used to define chunking constraints on what below parameters can be.""" - MIN_CHUNK_SIZE = 512 + MIN_CHUNK_SIZE = 128 MAX_CHUNK_SIZE = 2048 MIN_CHUNK_OVERLAP = 0 MAX_CHUNK_OVERLAP = 256 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 2897e552..a8a1b753 100644 --- a/tests/unit/ai4rag/components/optimization/test_search_space_prep.py +++ b/tests/unit/ai4rag/components/optimization/test_search_space_prep.py @@ -12,6 +12,7 @@ from ai4rag.components.optimization.search_space_preparation import ( SUPPORTED_METRICS, SearchSpaceReport, + _validate_chunk_sizes, _validate_chunking_methods, _validate_model_list, prepare_search_space_report, @@ -243,36 +244,121 @@ class TestUnsupportedChunkingMethods: """Test that providing chunking methods not in the search space raises.""" def test_unsupported_method_raises_value_error(self, mock_ogx_client): - """A chunking method absent from the search space must raise ValueError.""" - from unittest.mock import patch - - from ai4rag.components.optimization import search_space_preparation as mod - - search_space_items = { - "chunking_method": MagicMock( - values=["recursive", "hybrid"], all_values=MagicMock(return_value=["recursive", "hybrid"]) - ), - "chunk_size": MagicMock(values=[256], all_values=MagicMock(return_value=[256])), - "foundation_model": MagicMock(values=[MagicMock()]), - "embedding_model": MagicMock(values=[MagicMock()]), - } - fake_search_space = MagicMock() - fake_search_space._search_space = search_space_items - fake_search_space.__getitem__ = lambda self, key: search_space_items[key] - - fake_benchmark_df = MagicMock(spec=mod.pd.DataFrame) - fake_benchmark_df.__len__ = lambda self: 1 - - with ( - patch.object(mod, "prepare_search_space_with_ogx", return_value=fake_search_space), - patch.object(mod.pd, "read_json", return_value=fake_benchmark_df), - patch.object(mod, "load_docling_documents", return_value=[]), - patch.object(mod, "BenchmarkData"), - ): - with pytest.raises(ValueError, match="Unsupported chunking methods"): - prepare_search_space_report( - test_data_path="dummy.json", - extracted_text_path="dummy_dir", - ogx_client=mock_ogx_client, - chunking_methods=["semantic"], - ) + """A chunking method not in ChunkingConstraints.METHODS must raise ValueError. + + Validation is now performed eagerly before any I/O so no mocking is needed. + """ + with pytest.raises(ValueError, match="Unsupported chunking methods"): + prepare_search_space_report( + test_data_path="dummy.json", + extracted_text_path="dummy_dir", + ogx_client=mock_ogx_client, + chunking_methods=["semantic"], + ) + + +# --------------------------------------------------------------------------- +# _validate_chunk_sizes +# --------------------------------------------------------------------------- + + +class TestValidateChunkSizes: + """Tests for the ``_validate_chunk_sizes`` helper.""" + + def test_none_is_valid(self): + """None means 'use defaults' and must not raise.""" + _validate_chunk_sizes(None) + + def test_valid_list_passes(self): + """A list of positive integers must be accepted.""" + _validate_chunk_sizes([256, 512, 1024]) + + def test_single_element_passes(self): + """A single-element list must be accepted.""" + _validate_chunk_sizes([512]) + + def test_non_list_raises_type_error(self): + """A non-list value must raise TypeError.""" + with pytest.raises(TypeError, match="must be a list"): + _validate_chunk_sizes(512) # type: ignore[arg-type] + + def test_empty_list_raises_value_error(self): + """An empty list must raise ValueError.""" + with pytest.raises(ValueError, match="must not be empty"): + _validate_chunk_sizes([]) + + def test_zero_raises_type_error(self): + """Zero is not a positive integer and must raise TypeError.""" + with pytest.raises(TypeError, match=r"chunk_sizes\[0\] must be a positive integer"): + _validate_chunk_sizes([0]) + + def test_negative_raises_type_error(self): + """A negative value must raise TypeError.""" + with pytest.raises(TypeError, match=r"chunk_sizes\[1\] must be a positive integer"): + _validate_chunk_sizes([512, -1]) + + def test_float_raises_type_error(self): + """A float value must raise TypeError even if it looks like an integer.""" + with pytest.raises(TypeError, match=r"chunk_sizes\[0\] must be a positive integer"): + _validate_chunk_sizes([512.0]) # type: ignore[list-item] + + def test_below_min_raises_value_error(self): + """A size below ChunkingConstraints.MIN_CHUNK_SIZE must raise ValueError.""" + from ai4rag.utils.constants import ChunkingConstraints + + below_min = ChunkingConstraints.MIN_CHUNK_SIZE - 1 + with pytest.raises(ValueError, match="out of range"): + _validate_chunk_sizes([below_min]) + + def test_above_max_raises_value_error(self): + """A size above ChunkingConstraints.MAX_CHUNK_SIZE must raise ValueError.""" + from ai4rag.utils.constants import ChunkingConstraints + + above_max = ChunkingConstraints.MAX_CHUNK_SIZE + 1 + with pytest.raises(ValueError, match="out of range"): + _validate_chunk_sizes([above_max]) + + def test_boundary_values_pass(self): + """MIN_CHUNK_SIZE and MAX_CHUNK_SIZE themselves must be accepted.""" + from ai4rag.utils.constants import ChunkingConstraints + + _validate_chunk_sizes([ChunkingConstraints.MIN_CHUNK_SIZE, ChunkingConstraints.MAX_CHUNK_SIZE]) + + +# --------------------------------------------------------------------------- +# prepare_search_space_report — chunk_sizes input validation +# --------------------------------------------------------------------------- + + +class TestPrepareSearchSpaceReportChunkSizesValidation: + """Test chunk_sizes input validation in prepare_search_space_report.""" + + def test_non_list_chunk_sizes_raises_type_error(self, mock_ogx_client): + """A non-list chunk_sizes must raise TypeError before any I/O.""" + with pytest.raises(TypeError, match="must be a list"): + prepare_search_space_report( + test_data_path="dummy.json", + extracted_text_path="dummy_dir", + ogx_client=mock_ogx_client, + chunk_sizes=512, # type: ignore[arg-type] + ) + + def test_empty_chunk_sizes_raises_value_error(self, mock_ogx_client): + """An empty chunk_sizes list must raise ValueError before any I/O.""" + with pytest.raises(ValueError, match="must not be empty"): + prepare_search_space_report( + test_data_path="dummy.json", + extracted_text_path="dummy_dir", + ogx_client=mock_ogx_client, + chunk_sizes=[], + ) + + def test_non_positive_chunk_size_raises_type_error(self, mock_ogx_client): + """A non-positive integer in chunk_sizes must raise TypeError before any I/O.""" + with pytest.raises(TypeError, match="must be a positive integer"): + prepare_search_space_report( + test_data_path="dummy.json", + extracted_text_path="dummy_dir", + ogx_client=mock_ogx_client, + chunk_sizes=[0], + ) diff --git a/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py b/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py index 2e76e070..3453313e 100644 --- a/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py +++ b/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py @@ -10,7 +10,7 @@ from ogx_client import OgxClient from pydantic import ValidationError -from ai4rag.search_space.prepare import prepare_search_space_with_ogx +from ai4rag.search_space.prepare import prepare_search_space_custom, prepare_search_space_with_ogx from ai4rag.search_space.src.exceptions import SearchSpaceValueError @@ -401,3 +401,87 @@ def test_user_picks_available_model_while_others_fail(self, mocker): fm_ids = [m.model_id for m in result["foundation_model"].values] assert fm_ids == ["llm-ok"] + + +class TestPrepareSearchSpaceCustom: + """Test prepare_search_space_custom function.""" + + def _setup_mock_client(self, mocker): + mock_client = MagicMock(spec=OgxClient) + + mock_llm = Mock() + mock_llm.id = "default-llm" + mock_llm.custom_metadata = {"model_type": "llm"} + + mock_embedding = Mock() + mock_embedding.id = "default-embedding" + mock_embedding.custom_metadata = {"model_type": "embedding", "embedding_dimension": 768} + + mock_client.models.list.return_value.data = [mock_llm, mock_embedding] + + mocker.patch("ai4rag.search_space.prepare.ogx_utils._validate_foundation_model", return_value=True) + mocker.patch("ai4rag.search_space.prepare.ogx_utils._validate_embedding_model", return_value=True) + + return mock_client + + def test_no_extra_params_behaves_like_with_ogx(self, mocker): + """Without chunking_methods or chunk_sizes the result matches prepare_search_space_with_ogx.""" + mock_client = self._setup_mock_client(mocker) + + result = prepare_search_space_custom({}, mock_client) + + param_names = [p.name for p in result.params] + assert "foundation_model" in param_names + assert "embedding_model" in param_names + assert "chunking_method" in param_names + assert "chunk_size" in param_names + + def test_custom_chunking_methods_override_defaults(self, mocker): + """chunking_methods in payload overrides the default chunking_method dimension.""" + mock_client = self._setup_mock_client(mocker) + + result = prepare_search_space_custom({"chunking_methods": ["recursive"]}, mock_client) + + assert result["chunking_method"].values == ("recursive",) + + def test_custom_chunk_sizes_override_defaults(self, mocker): + """chunk_sizes in payload overrides the default chunk_size dimension.""" + mock_client = self._setup_mock_client(mocker) + + result = prepare_search_space_custom({"chunk_sizes": [256, 512]}, mock_client) + + assert set(result["chunk_size"].values) == {256, 512} + + def test_both_custom_params_applied_together(self, mocker): + """Both chunking_methods and chunk_sizes in payload can be set simultaneously.""" + mock_client = self._setup_mock_client(mocker) + + result = prepare_search_space_custom( + {"chunking_methods": ["hybrid"], "chunk_sizes": [1024]}, + mock_client, + ) + + assert result["chunking_method"].values == ("hybrid",) + assert result["chunk_size"].values == (1024,) + + def test_non_ogx_client_raises_error(self): + """Non-OgxClient raises SearchSpaceValueError.""" + with pytest.raises(SearchSpaceValueError, match="Unrecognized client type"): + prepare_search_space_custom({}, MagicMock(spec=object)) + + def test_invalid_payload_raises_validation_error(self, mocker): + """Invalid payload fields raise a pydantic ValidationError.""" + mock_client = self._setup_mock_client(mocker) + with pytest.raises(ValidationError): + prepare_search_space_custom({"invalid_field": "value"}, mock_client) + + def test_chroma_vector_store_excludes_hybrid_params(self, mocker): + """chroma vector_store_type excludes ranker parameters.""" + mock_client = self._setup_mock_client(mocker) + + result = prepare_search_space_custom({}, mock_client, vector_store_type="chroma") + + param_names = [p.name for p in result.params] + assert "ranker_strategy" not in param_names + assert "ranker_k" not in param_names + assert "ranker_alpha" not in param_names From 2643742e26c63d99949ca80294b0d4b25c8b2b88 Mon Sep 17 00:00:00 2001 From: Mateusz Switala Date: Thu, 2 Jul 2026 22:08:24 +0200 Subject: [PATCH 02/13] fix black formatting Signed-off-by: Mateusz Switala --- ai4rag/components/optimization/search_space_preparation.py | 3 +-- ai4rag/search_space/prepare/prepare_search_space.py | 4 +--- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/ai4rag/components/optimization/search_space_preparation.py b/ai4rag/components/optimization/search_space_preparation.py index 126abd24..fc351273 100644 --- a/ai4rag/components/optimization/search_space_preparation.py +++ b/ai4rag/components/optimization/search_space_preparation.py @@ -278,8 +278,7 @@ def _validate_chunking_methods(methods: list[str] | None) -> None: unsupported = [m for m in methods if m not in ChunkingConstraints.METHODS] if unsupported: raise ValueError( - f"Unsupported chunking methods: {unsupported!r}. " - f"Supported methods: {ChunkingConstraints.METHODS!r}." + f"Unsupported chunking methods: {unsupported!r}. " f"Supported methods: {ChunkingConstraints.METHODS!r}." ) diff --git a/ai4rag/search_space/prepare/prepare_search_space.py b/ai4rag/search_space/prepare/prepare_search_space.py index e5d3f9a4..1e0a207d 100644 --- a/ai4rag/search_space/prepare/prepare_search_space.py +++ b/ai4rag/search_space/prepare/prepare_search_space.py @@ -254,9 +254,7 @@ def prepare_search_space_custom( Parameter(name=AI4RAGParamNames.CHUNKING_METHOD, values=tuple(validated_payload.chunking_methods)) ) if validated_payload.chunk_sizes is not None: - extra_params.append( - Parameter(name=AI4RAGParamNames.CHUNK_SIZE, values=tuple(validated_payload.chunk_sizes)) - ) + extra_params.append(Parameter(name=AI4RAGParamNames.CHUNK_SIZE, values=tuple(validated_payload.chunk_sizes))) return AI4RAGSearchSpace( params=[fms_param, ems_param, *extra_params], From 7c8414936137dfc8324ed90fdfb24bd33de17f78 Mon Sep 17 00:00:00 2001 From: Mateusz Switala Date: Thu, 2 Jul 2026 22:50:32 +0200 Subject: [PATCH 03/13] resolve pylint check Signed-off-by: Mateusz Switala --- ai4rag/components/optimization/search_space_preparation.py | 2 +- ai4rag/search_space/prepare/prepare_search_space.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/ai4rag/components/optimization/search_space_preparation.py b/ai4rag/components/optimization/search_space_preparation.py index fc351273..4e77de97 100644 --- a/ai4rag/components/optimization/search_space_preparation.py +++ b/ai4rag/components/optimization/search_space_preparation.py @@ -293,7 +293,7 @@ def _validate_chunk_sizes(chunk_sizes: list[int] | None) -> None: for i, s in enumerate(chunk_sizes): if not isinstance(s, int) or s <= 0: raise TypeError(f"chunk_sizes[{i}] must be a positive integer.") - if not (ChunkingConstraints.MIN_CHUNK_SIZE <= s <= ChunkingConstraints.MAX_CHUNK_SIZE): + if not ChunkingConstraints.MIN_CHUNK_SIZE <= s <= ChunkingConstraints.MAX_CHUNK_SIZE: raise ValueError( f"chunk_sizes[{i}]={s} is out of range " f"[{ChunkingConstraints.MIN_CHUNK_SIZE}, {ChunkingConstraints.MAX_CHUNK_SIZE}]." diff --git a/ai4rag/search_space/prepare/prepare_search_space.py b/ai4rag/search_space/prepare/prepare_search_space.py index 1e0a207d..2197b286 100644 --- a/ai4rag/search_space/prepare/prepare_search_space.py +++ b/ai4rag/search_space/prepare/prepare_search_space.py @@ -167,8 +167,7 @@ def prepare_search_space_with_ogx( validated_payload = AI4RAGConstraints(**payload) - _CHUNKING_FIELDS = ("chunking_methods", "chunk_sizes") - skipped = [f for f in _CHUNKING_FIELDS if getattr(validated_payload, f) is not None] + skipped = [f for f in ("chunking_methods", "chunk_sizes") if getattr(validated_payload, f) is not None] if skipped: logger.warning( "Fields %s are not used by prepare_search_space_with_ogx and will be skipped. " From fe78e5207956918b2abf86cdb49a17ee2f15187a Mon Sep 17 00:00:00 2001 From: Mateusz Switala Date: Fri, 3 Jul 2026 09:24:39 +0200 Subject: [PATCH 04/13] fixes after claude review Signed-off-by: Mateusz Switala Assisted-by: Claude Code --- .../optimization/search_space_preparation.py | 22 ++++----- .../prepare/prepare_search_space.py | 32 ++++++++++--- .../optimization/test_search_space_prep.py | 47 ++----------------- .../prepare/test_prepare_search_space.py | 46 ++++++++++++++++++ 4 files changed, 85 insertions(+), 62 deletions(-) diff --git a/ai4rag/components/optimization/search_space_preparation.py b/ai4rag/components/optimization/search_space_preparation.py index 4e77de97..ffcb8fbc 100644 --- a/ai4rag/components/optimization/search_space_preparation.py +++ b/ai4rag/components/optimization/search_space_preparation.py @@ -18,7 +18,6 @@ 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_custom -from ai4rag.utils.constants import ChunkingConstraints _logger = logging.getLogger("search-space-preparation") _logger.addHandler(handler) @@ -209,6 +208,11 @@ def prepare_search_space_report( # pylint: disable=too-many-locals,too-many-arg client=ogx_client, benchmark_data=benchmark_df, ) + _logger.info( + "Search space chunking_method=%s chunk_size=%s", + list(search_space["chunking_method"].values), + list(search_space["chunk_size"].values), + ) # Run model pre-selection when the number of models exceeds the caps fm_values = search_space["foundation_model"].values @@ -265,6 +269,8 @@ def _validate_model_list(models: list[str] | None, name: str) -> None: def _validate_chunking_methods(methods: list[str] | None) -> None: + # TODO: remove — structural validation is now covered by AI4RAGConstraints (pydantic) and + # supported-methods check is covered by prepare_search_space_custom. """Validate that chunking methods, if provided, are non-empty strings and supported.""" if methods is None: return @@ -275,14 +281,11 @@ def _validate_chunking_methods(methods: list[str] | None) -> None: for i, m in enumerate(methods): if not isinstance(m, str) or not m.strip(): raise TypeError(f"chunking_methods[{i}] must be a non-empty string.") - unsupported = [m for m in methods if m not in ChunkingConstraints.METHODS] - if unsupported: - raise ValueError( - f"Unsupported chunking methods: {unsupported!r}. " f"Supported methods: {ChunkingConstraints.METHODS!r}." - ) def _validate_chunk_sizes(chunk_sizes: list[int] | None) -> None: + # TODO: remove — structural validation is now covered by AI4RAGConstraints (pydantic) and + # bounds check is covered by prepare_search_space_custom. """Validate that chunk sizes, if provided, are integers within ChunkingConstraints bounds.""" if chunk_sizes is None: return @@ -291,10 +294,5 @@ def _validate_chunk_sizes(chunk_sizes: list[int] | None) -> None: if not chunk_sizes: raise ValueError("chunk_sizes must not be empty when provided.") for i, s in enumerate(chunk_sizes): - if not isinstance(s, int) or s <= 0: + if isinstance(s, bool) or not isinstance(s, int) or s <= 0: raise TypeError(f"chunk_sizes[{i}] must be a positive integer.") - if not ChunkingConstraints.MIN_CHUNK_SIZE <= s <= ChunkingConstraints.MAX_CHUNK_SIZE: - raise ValueError( - f"chunk_sizes[{i}]={s} is out of range " - f"[{ChunkingConstraints.MIN_CHUNK_SIZE}, {ChunkingConstraints.MAX_CHUNK_SIZE}]." - ) diff --git a/ai4rag/search_space/prepare/prepare_search_space.py b/ai4rag/search_space/prepare/prepare_search_space.py index 2197b286..e59a4e22 100644 --- a/ai4rag/search_space/prepare/prepare_search_space.py +++ b/ai4rag/search_space/prepare/prepare_search_space.py @@ -18,7 +18,7 @@ from ai4rag.search_space.src.exceptions import SearchSpaceValueError 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.constants import AI4RAGParamNames, ChunkingConstraints __all__ = ["prepare_search_space_with_ogx", "prepare_search_space_custom"] @@ -120,8 +120,8 @@ def _build_model_params(foundation_models: list, embedding_models: list) -> tupl tuple[Parameter, Parameter] ``(fms_param, ems_param)`` ready for inclusion in a search space. """ - fms_param = Parameter(name="foundation_model", values=foundation_models) - ems_param = Parameter(name="embedding_model", values=embedding_models) + fms_param = Parameter(name=AI4RAGParamNames.FOUNDATION_MODEL, values=foundation_models) + ems_param = Parameter(name=AI4RAGParamNames.EMBEDDING_MODEL, values=embedding_models) logger.info("Selected foundation models for the experiment: %s.", [m.model_id for m in fms_param.values]) logger.info("Selected embedding models for the experiment: %s.", [m.model_id for m in ems_param.values]) return fms_param, ems_param @@ -169,10 +169,9 @@ def prepare_search_space_with_ogx( skipped = [f for f in ("chunking_methods", "chunk_sizes") if getattr(validated_payload, f) is not None] if skipped: - logger.warning( - "Fields %s are not used by prepare_search_space_with_ogx and will be skipped. " - "Use prepare_search_space_custom to apply chunking constraints.", - skipped, + raise SearchSpaceValueError( + f"Fields {skipped} are not supported by prepare_search_space_with_ogx. " + "Use prepare_search_space_custom to apply chunking constraints." ) foundation_models, embedding_models = _resolve_models_from_payload(validated_payload, client) @@ -240,6 +239,25 @@ def prepare_search_space_custom( logger.info("Preparing custom search space based on provided constraints: %s.", payload) validated_payload = AI4RAGConstraints(**payload) + + if validated_payload.chunking_methods is not None: + unsupported = [m for m in validated_payload.chunking_methods if m not in ChunkingConstraints.METHODS] + if unsupported: + raise SearchSpaceValueError( + f"Unsupported chunking methods: {unsupported!r}. " + f"Supported methods: {ChunkingConstraints.METHODS!r}." + ) + + if validated_payload.chunk_sizes is not None: + for i, s in enumerate(validated_payload.chunk_sizes): + if isinstance(s, bool): + raise SearchSpaceValueError(f"chunk_sizes[{i}] must be a positive integer, got bool.") + if not ChunkingConstraints.MIN_CHUNK_SIZE <= s <= ChunkingConstraints.MAX_CHUNK_SIZE: + raise SearchSpaceValueError( + f"chunk_sizes[{i}]={s} is out of range " + f"[{ChunkingConstraints.MIN_CHUNK_SIZE}, {ChunkingConstraints.MAX_CHUNK_SIZE}]." + ) + foundation_models, embedding_models = _resolve_models_from_payload(validated_payload, client) if benchmark_data is not None: 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 a8a1b753..cfa525b9 100644 --- a/tests/unit/ai4rag/components/optimization/test_search_space_prep.py +++ b/tests/unit/ai4rag/components/optimization/test_search_space_prep.py @@ -235,28 +235,6 @@ def test_non_string_raises_type_error(self): _validate_chunking_methods(["recursive", 42]) # type: ignore[list-item] -# --------------------------------------------------------------------------- -# prepare_search_space_report — unsupported chunking_methods -# --------------------------------------------------------------------------- - - -class TestUnsupportedChunkingMethods: - """Test that providing chunking methods not in the search space raises.""" - - def test_unsupported_method_raises_value_error(self, mock_ogx_client): - """A chunking method not in ChunkingConstraints.METHODS must raise ValueError. - - Validation is now performed eagerly before any I/O so no mocking is needed. - """ - with pytest.raises(ValueError, match="Unsupported chunking methods"): - prepare_search_space_report( - test_data_path="dummy.json", - extracted_text_path="dummy_dir", - ogx_client=mock_ogx_client, - chunking_methods=["semantic"], - ) - - # --------------------------------------------------------------------------- # _validate_chunk_sizes # --------------------------------------------------------------------------- @@ -302,27 +280,10 @@ def test_float_raises_type_error(self): with pytest.raises(TypeError, match=r"chunk_sizes\[0\] must be a positive integer"): _validate_chunk_sizes([512.0]) # type: ignore[list-item] - def test_below_min_raises_value_error(self): - """A size below ChunkingConstraints.MIN_CHUNK_SIZE must raise ValueError.""" - from ai4rag.utils.constants import ChunkingConstraints - - below_min = ChunkingConstraints.MIN_CHUNK_SIZE - 1 - with pytest.raises(ValueError, match="out of range"): - _validate_chunk_sizes([below_min]) - - def test_above_max_raises_value_error(self): - """A size above ChunkingConstraints.MAX_CHUNK_SIZE must raise ValueError.""" - from ai4rag.utils.constants import ChunkingConstraints - - above_max = ChunkingConstraints.MAX_CHUNK_SIZE + 1 - with pytest.raises(ValueError, match="out of range"): - _validate_chunk_sizes([above_max]) - - def test_boundary_values_pass(self): - """MIN_CHUNK_SIZE and MAX_CHUNK_SIZE themselves must be accepted.""" - from ai4rag.utils.constants import ChunkingConstraints - - _validate_chunk_sizes([ChunkingConstraints.MIN_CHUNK_SIZE, ChunkingConstraints.MAX_CHUNK_SIZE]) + def test_bool_raises_type_error(self): + """A bool value must raise TypeError even though bool is a subclass of int.""" + with pytest.raises(TypeError, match=r"chunk_sizes\[0\] must be a positive integer"): + _validate_chunk_sizes([True]) # --------------------------------------------------------------------------- diff --git a/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py b/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py index 3453313e..7091600c 100644 --- a/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py +++ b/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py @@ -402,6 +402,20 @@ def test_user_picks_available_model_while_others_fail(self, mocker): fm_ids = [m.model_id for m in result["foundation_model"].values] assert fm_ids == ["llm-ok"] + def test_chunking_fields_in_payload_raise_error(self): + """chunking_methods or chunk_sizes in payload must raise SearchSpaceValueError.""" + mock_client = MagicMock(spec=OgxClient) + + with pytest.raises(SearchSpaceValueError, match="not supported by prepare_search_space_with_ogx"): + prepare_search_space_with_ogx({"chunking_methods": ["recursive"]}, mock_client) + + def test_chunk_sizes_in_payload_raise_error(self): + """chunk_sizes in payload must raise SearchSpaceValueError.""" + mock_client = MagicMock(spec=OgxClient) + + with pytest.raises(SearchSpaceValueError, match="not supported by prepare_search_space_with_ogx"): + prepare_search_space_with_ogx({"chunk_sizes": [512]}, mock_client) + class TestPrepareSearchSpaceCustom: """Test prepare_search_space_custom function.""" @@ -485,3 +499,35 @@ def test_chroma_vector_store_excludes_hybrid_params(self, mocker): assert "ranker_strategy" not in param_names assert "ranker_k" not in param_names assert "ranker_alpha" not in param_names + + def test_unsupported_chunking_method_raises_error(self, mocker): + """An unsupported chunking method raises SearchSpaceValueError before any I/O.""" + mock_client = self._setup_mock_client(mocker) + + with pytest.raises(SearchSpaceValueError, match="Unsupported chunking methods"): + prepare_search_space_custom({"chunking_methods": ["semantic"]}, mock_client) + + def test_chunk_size_below_min_raises_error(self, mocker): + """A chunk size below MIN_CHUNK_SIZE raises SearchSpaceValueError before any I/O.""" + mock_client = self._setup_mock_client(mocker) + + with pytest.raises(SearchSpaceValueError, match="out of range"): + prepare_search_space_custom({"chunk_sizes": [1]}, mock_client) + + def test_chunk_size_above_max_raises_error(self, mocker): + """A chunk size above MAX_CHUNK_SIZE raises SearchSpaceValueError before any I/O.""" + mock_client = self._setup_mock_client(mocker) + + with pytest.raises(SearchSpaceValueError, match="out of range"): + prepare_search_space_custom({"chunk_sizes": [99999]}, mock_client) + + def test_bool_chunk_size_raises_error(self, mocker): + """A bool value in chunk_sizes raises SearchSpaceValueError. + + Pydantic coerces True→1 before our validator runs, so the error is + "out of range" (1 < MIN_CHUNK_SIZE) rather than "must be a positive integer". + """ + mock_client = self._setup_mock_client(mocker) + + with pytest.raises(SearchSpaceValueError, match="out of range"): + prepare_search_space_custom({"chunk_sizes": [True]}, mock_client) From bb325aeb8731a3023c27fe87cdb2ed2b5b09212f Mon Sep 17 00:00:00 2001 From: Mateusz Switala Date: Fri, 3 Jul 2026 09:26:47 +0200 Subject: [PATCH 05/13] fix pylint Signed-off-by: Mateusz Switala --- ai4rag/components/optimization/search_space_preparation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ai4rag/components/optimization/search_space_preparation.py b/ai4rag/components/optimization/search_space_preparation.py index ffcb8fbc..28d74f9f 100644 --- a/ai4rag/components/optimization/search_space_preparation.py +++ b/ai4rag/components/optimization/search_space_preparation.py @@ -269,7 +269,7 @@ def _validate_model_list(models: list[str] | None, name: str) -> None: def _validate_chunking_methods(methods: list[str] | None) -> None: - # TODO: remove — structural validation is now covered by AI4RAGConstraints (pydantic) and + # TODO: remove — structural validation is now covered by AI4RAGConstraints (pydantic) and # pylint: disable=fixme # supported-methods check is covered by prepare_search_space_custom. """Validate that chunking methods, if provided, are non-empty strings and supported.""" if methods is None: @@ -284,7 +284,7 @@ def _validate_chunking_methods(methods: list[str] | None) -> None: def _validate_chunk_sizes(chunk_sizes: list[int] | None) -> None: - # TODO: remove — structural validation is now covered by AI4RAGConstraints (pydantic) and + # TODO: remove — structural validation is now covered by AI4RAGConstraints (pydantic) and # pylint: disable=fixme # bounds check is covered by prepare_search_space_custom. """Validate that chunk sizes, if provided, are integers within ChunkingConstraints bounds.""" if chunk_sizes is None: From dfd3a8a15405248816dbc2630644144c2bd7e4ab Mon Sep 17 00:00:00 2001 From: Mateusz Switala Date: Fri, 3 Jul 2026 15:54:26 +0200 Subject: [PATCH 06/13] remove chunking_methods and chunk_sizes validation Signed-off-by: Mateusz Switala Assisted-by: Claude Code --- .../optimization/search_space_preparation.py | 44 ++--- .../optimization/test_search_space_prep.py | 152 ------------------ 2 files changed, 10 insertions(+), 186 deletions(-) diff --git a/ai4rag/components/optimization/search_space_preparation.py b/ai4rag/components/optimization/search_space_preparation.py index 28d74f9f..2794d2cd 100644 --- a/ai4rag/components/optimization/search_space_preparation.py +++ b/ai4rag/components/optimization/search_space_preparation.py @@ -173,19 +173,23 @@ def prepare_search_space_report( # pylint: disable=too-many-locals,too-many-arg Raises ------ ValueError - If *metric* is not one of the supported values, or if - *chunking_methods* or *chunk_sizes* contain unsupported entries. + If *metric* is not one of the supported values. TypeError - If *embedding_models*, *generation_models*, *chunking_methods*, or - *chunk_sizes* contain invalid entries. + If *embedding_models* or *generation_models* contain invalid entries. + pydantic.ValidationError + If *chunking_methods* or *chunk_sizes* fail structural validation + (wrong type, empty list, or invalid element types). + SearchSpaceValueError + If *chunking_methods* contains values not in + :attr:`~ai4rag.utils.constants.ChunkingConstraints.METHODS`, or + *chunk_sizes* contains values outside + ``[ChunkingConstraints.MIN_CHUNK_SIZE, ChunkingConstraints.MAX_CHUNK_SIZE]``. """ if metric not in SUPPORTED_METRICS: raise ValueError(f"Metric {metric!r} is not supported. Supported metrics are {list(SUPPORTED_METRICS)}.") _validate_model_list(embedding_models, "embedding_models") _validate_model_list(generation_models, "generation_models") - _validate_chunking_methods(chunking_methods) - _validate_chunk_sizes(chunk_sizes) # Build payload and create search space via OGX payload: dict[str, Any] = {} @@ -268,31 +272,3 @@ def _validate_model_list(models: list[str] | None, name: str) -> None: raise TypeError(f"{name}[{i}] must be a non-empty string.") -def _validate_chunking_methods(methods: list[str] | None) -> None: - # TODO: remove — structural validation is now covered by AI4RAGConstraints (pydantic) and # pylint: disable=fixme - # supported-methods check is covered by prepare_search_space_custom. - """Validate that chunking methods, if provided, are non-empty strings and supported.""" - if methods is None: - return - if not isinstance(methods, list): - raise TypeError("chunking_methods must be a list.") - if not methods: - raise ValueError("chunking_methods must not be empty when provided.") - for i, m in enumerate(methods): - if not isinstance(m, str) or not m.strip(): - raise TypeError(f"chunking_methods[{i}] must be a non-empty string.") - - -def _validate_chunk_sizes(chunk_sizes: list[int] | None) -> None: - # TODO: remove — structural validation is now covered by AI4RAGConstraints (pydantic) and # pylint: disable=fixme - # bounds check is covered by prepare_search_space_custom. - """Validate that chunk sizes, if provided, are integers within ChunkingConstraints bounds.""" - if chunk_sizes is None: - return - if not isinstance(chunk_sizes, list): - raise TypeError("chunk_sizes must be a list.") - if not chunk_sizes: - raise ValueError("chunk_sizes must not be empty when provided.") - for i, s in enumerate(chunk_sizes): - if isinstance(s, bool) or not isinstance(s, int) or s <= 0: - raise TypeError(f"chunk_sizes[{i}] must be a positive integer.") 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 cfa525b9..f25fbb4b 100644 --- a/tests/unit/ai4rag/components/optimization/test_search_space_prep.py +++ b/tests/unit/ai4rag/components/optimization/test_search_space_prep.py @@ -12,8 +12,6 @@ from ai4rag.components.optimization.search_space_preparation import ( SUPPORTED_METRICS, SearchSpaceReport, - _validate_chunk_sizes, - _validate_chunking_methods, _validate_model_list, prepare_search_space_report, ) @@ -173,153 +171,3 @@ def test_non_list_models_raises_type_error(self, mock_ogx_client): embedding_models="not-a-list", # type: ignore[arg-type] ) - def test_invalid_chunking_methods_raises_type_error(self, mock_ogx_client): - """A non-list value for chunking_methods must raise TypeError.""" - with pytest.raises(TypeError, match="must be a list"): - prepare_search_space_report( - test_data_path="dummy.json", - extracted_text_path="dummy_dir", - ogx_client=mock_ogx_client, - chunking_methods="recursive", # type: ignore[arg-type] - ) - - def test_empty_chunking_methods_raises_value_error(self, mock_ogx_client): - """An empty list for chunking_methods must raise ValueError.""" - with pytest.raises(ValueError, match="must not be empty"): - prepare_search_space_report( - test_data_path="dummy.json", - extracted_text_path="dummy_dir", - ogx_client=mock_ogx_client, - chunking_methods=[], - ) - - -# --------------------------------------------------------------------------- -# _validate_chunking_methods -# --------------------------------------------------------------------------- - - -class TestValidateChunkingMethods: - """Tests for the ``_validate_chunking_methods`` helper.""" - - def test_none_is_valid(self): - """None means 'use defaults' and must not raise.""" - _validate_chunking_methods(None) - - def test_valid_list_passes(self): - """A list of non-empty strings must be accepted.""" - _validate_chunking_methods(["recursive", "hybrid"]) - - def test_single_element_passes(self): - """A single-element list must be accepted.""" - _validate_chunking_methods(["recursive"]) - - def test_non_list_raises_type_error(self): - """A non-list value must raise TypeError.""" - with pytest.raises(TypeError, match="must be a list"): - _validate_chunking_methods("recursive") # type: ignore[arg-type] - - def test_empty_list_raises_value_error(self): - """An empty list must raise ValueError.""" - with pytest.raises(ValueError, match="must not be empty"): - _validate_chunking_methods([]) - - def test_empty_string_raises_type_error(self): - """An empty string element must raise TypeError.""" - with pytest.raises(TypeError, match=r"chunking_methods\[0\] must be a non-empty string"): - _validate_chunking_methods([""]) - - def test_non_string_raises_type_error(self): - """A non-string element must raise TypeError.""" - with pytest.raises(TypeError, match=r"chunking_methods\[1\] must be a non-empty string"): - _validate_chunking_methods(["recursive", 42]) # type: ignore[list-item] - - -# --------------------------------------------------------------------------- -# _validate_chunk_sizes -# --------------------------------------------------------------------------- - - -class TestValidateChunkSizes: - """Tests for the ``_validate_chunk_sizes`` helper.""" - - def test_none_is_valid(self): - """None means 'use defaults' and must not raise.""" - _validate_chunk_sizes(None) - - def test_valid_list_passes(self): - """A list of positive integers must be accepted.""" - _validate_chunk_sizes([256, 512, 1024]) - - def test_single_element_passes(self): - """A single-element list must be accepted.""" - _validate_chunk_sizes([512]) - - def test_non_list_raises_type_error(self): - """A non-list value must raise TypeError.""" - with pytest.raises(TypeError, match="must be a list"): - _validate_chunk_sizes(512) # type: ignore[arg-type] - - def test_empty_list_raises_value_error(self): - """An empty list must raise ValueError.""" - with pytest.raises(ValueError, match="must not be empty"): - _validate_chunk_sizes([]) - - def test_zero_raises_type_error(self): - """Zero is not a positive integer and must raise TypeError.""" - with pytest.raises(TypeError, match=r"chunk_sizes\[0\] must be a positive integer"): - _validate_chunk_sizes([0]) - - def test_negative_raises_type_error(self): - """A negative value must raise TypeError.""" - with pytest.raises(TypeError, match=r"chunk_sizes\[1\] must be a positive integer"): - _validate_chunk_sizes([512, -1]) - - def test_float_raises_type_error(self): - """A float value must raise TypeError even if it looks like an integer.""" - with pytest.raises(TypeError, match=r"chunk_sizes\[0\] must be a positive integer"): - _validate_chunk_sizes([512.0]) # type: ignore[list-item] - - def test_bool_raises_type_error(self): - """A bool value must raise TypeError even though bool is a subclass of int.""" - with pytest.raises(TypeError, match=r"chunk_sizes\[0\] must be a positive integer"): - _validate_chunk_sizes([True]) - - -# --------------------------------------------------------------------------- -# prepare_search_space_report — chunk_sizes input validation -# --------------------------------------------------------------------------- - - -class TestPrepareSearchSpaceReportChunkSizesValidation: - """Test chunk_sizes input validation in prepare_search_space_report.""" - - def test_non_list_chunk_sizes_raises_type_error(self, mock_ogx_client): - """A non-list chunk_sizes must raise TypeError before any I/O.""" - with pytest.raises(TypeError, match="must be a list"): - prepare_search_space_report( - test_data_path="dummy.json", - extracted_text_path="dummy_dir", - ogx_client=mock_ogx_client, - chunk_sizes=512, # type: ignore[arg-type] - ) - - def test_empty_chunk_sizes_raises_value_error(self, mock_ogx_client): - """An empty chunk_sizes list must raise ValueError before any I/O.""" - with pytest.raises(ValueError, match="must not be empty"): - prepare_search_space_report( - test_data_path="dummy.json", - extracted_text_path="dummy_dir", - ogx_client=mock_ogx_client, - chunk_sizes=[], - ) - - def test_non_positive_chunk_size_raises_type_error(self, mock_ogx_client): - """A non-positive integer in chunk_sizes must raise TypeError before any I/O.""" - with pytest.raises(TypeError, match="must be a positive integer"): - prepare_search_space_report( - test_data_path="dummy.json", - extracted_text_path="dummy_dir", - ogx_client=mock_ogx_client, - chunk_sizes=[0], - ) From 43700f953cf139bca8baad133ae96007369681b6 Mon Sep 17 00:00:00 2001 From: Mateusz Switala Date: Fri, 3 Jul 2026 15:56:51 +0200 Subject: [PATCH 07/13] fix pylint Signed-off-by: Mateusz Switala --- 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 2794d2cd..14743704 100644 --- a/ai4rag/components/optimization/search_space_preparation.py +++ b/ai4rag/components/optimization/search_space_preparation.py @@ -270,5 +270,3 @@ def _validate_model_list(models: list[str] | None, name: str) -> None: for i, m in enumerate(models): if not m: raise TypeError(f"{name}[{i}] must be a non-empty string.") - - From ce423d50c45e69791db8f8ccbe35ec60ad943a9b Mon Sep 17 00:00:00 2001 From: Mateusz Switala Date: Mon, 6 Jul 2026 09:49:37 +0200 Subject: [PATCH 08/13] remove addtional function Signed-off-by: Mateusz Switala Assisted-by: Claude Code --- .../optimization/search_space_preparation.py | 4 +- ai4rag/search_space/prepare/__init__.py | 5 +- .../prepare/prepare_search_space.py | 104 +++++------------- .../prepare/test_prepare_search_space.py | 68 ++++-------- 4 files changed, 51 insertions(+), 130 deletions(-) diff --git a/ai4rag/components/optimization/search_space_preparation.py b/ai4rag/components/optimization/search_space_preparation.py index 14743704..165434af 100644 --- a/ai4rag/components/optimization/search_space_preparation.py +++ b/ai4rag/components/optimization/search_space_preparation.py @@ -17,7 +17,7 @@ from ai4rag.core.experiment.mps import ModelsPreSelector 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_custom +from ai4rag.search_space.prepare.prepare_search_space import prepare_search_space_with_ogx _logger = logging.getLogger("search-space-preparation") _logger.addHandler(handler) @@ -207,7 +207,7 @@ def prepare_search_space_report( # pylint: disable=too-many-locals,too-many-arg benchmark_data = BenchmarkData(benchmark_df) documents = load_docling_documents(extracted_text_path) - search_space = prepare_search_space_custom( + search_space = prepare_search_space_with_ogx( payload, client=ogx_client, benchmark_data=benchmark_df, diff --git a/ai4rag/search_space/prepare/__init__.py b/ai4rag/search_space/prepare/__init__.py index bac8837a..09b93d8b 100644 --- a/ai4rag/search_space/prepare/__init__.py +++ b/ai4rag/search_space/prepare/__init__.py @@ -3,7 +3,4 @@ # SPDX-License-Identifier: Apache-2.0 # ----------------------------------------------------------------------------- -from ai4rag.search_space.prepare.prepare_search_space import ( - prepare_search_space_custom, - prepare_search_space_with_ogx, -) +from ai4rag.search_space.prepare.prepare_search_space import prepare_search_space_with_ogx diff --git a/ai4rag/search_space/prepare/prepare_search_space.py b/ai4rag/search_space/prepare/prepare_search_space.py index e59a4e22..88d56f98 100644 --- a/ai4rag/search_space/prepare/prepare_search_space.py +++ b/ai4rag/search_space/prepare/prepare_search_space.py @@ -20,7 +20,7 @@ from ai4rag.search_space.src.search_space import AI4RAGSearchSpace from ai4rag.utils.constants import AI4RAGParamNames, ChunkingConstraints -__all__ = ["prepare_search_space_with_ogx", "prepare_search_space_custom"] +__all__ = ["prepare_search_space_with_ogx"] def _resolve_models_from_payload( @@ -133,88 +133,31 @@ def prepare_search_space_with_ogx( vector_store_type: str = "ogx", benchmark_data: pd.DataFrame | None = None, ) -> AI4RAGSearchSpace: - """ - Prepare AutoRAGSearchSpace. - - Parameters - ---------- - payload : dict[str, Any] - A mapping between parameter name and its associated values. - - client : OgxClient - Client instance for listing and validating available models. - - vector_store_type : str, default="ogx" - Type of vector store. Supported values: ``"ogx"`` and ``"chroma"``. - When ``"chroma"``, hybrid search parameters are excluded from the - default search space since ChromaDB does not support hybrid search. - - benchmark_data : pd.DataFrame | None, default=None - Benchmark data used for language detection. - If not given, models with use automatic language detection per session. - - Returns - ------- - AI4RAGSearchSpace - A valid AI4RAGSearchSpace used in RAG optimization process. - - Raises - ------ - SearchSpaceValueError - Raised when payload contains non-recognized parameter name. - """ - logger.info("Preparing search space based on provided constraints: %s.", payload) - - validated_payload = AI4RAGConstraints(**payload) - - skipped = [f for f in ("chunking_methods", "chunk_sizes") if getattr(validated_payload, f) is not None] - if skipped: - raise SearchSpaceValueError( - f"Fields {skipped} are not supported by prepare_search_space_with_ogx. " - "Use prepare_search_space_custom to apply chunking constraints." - ) - - foundation_models, embedding_models = _resolve_models_from_payload(validated_payload, client) - - if benchmark_data is not None: - _apply_language_detection(foundation_models, benchmark_data) - - fms_param, ems_param = _build_model_params(foundation_models, embedding_models) - - return AI4RAGSearchSpace( - params=[fms_param, ems_param], - vector_store_type=vector_store_type, - ) + """Prepare an :class:`AI4RAGSearchSpace` using OGX for model validation. - -def prepare_search_space_custom( - payload: dict[str, Any], - client: OgxClient, - vector_store_type: str = "ogx", - benchmark_data: pd.DataFrame | None = None, -) -> AI4RAGSearchSpace: - """Prepare an :class:`AI4RAGSearchSpace` with optional chunking customization. - - Extends :func:`prepare_search_space_with_ogx` by allowing the caller to - override the default ``chunking_method`` and ``chunk_size`` dimensions of - the search space via the *payload* dict. All other parameters and rules - remain unchanged. + Foundation and embedding models are discovered and validated via the OGX + platform. Chunking parameters (``chunking_methods``, ``chunk_sizes``) are + validated locally against :class:`~ai4rag.utils.constants.ChunkingConstraints` + and, when provided, override the platform defaults for those dimensions. Parameters ---------- payload : dict[str, Any] - A mapping of constraint names to their values. Supports all keys - accepted by :func:`prepare_search_space_with_ogx` plus: + A mapping of constraint names to their values. Supported keys: + - ``"foundation_models"`` *(list[dict])* — foundation model identifiers + to include; ``None`` uses all OGX defaults. + - ``"embedding_models"`` *(list[dict])* — embedding model identifiers + to include; ``None`` uses all OGX defaults. - ``"chunking_methods"`` *(list[str])* — overrides the default ``chunking_method`` dimension (e.g. ``["recursive"]``). + ``None`` keeps the platform default. - ``"chunk_sizes"`` *(list[int])* — overrides the default ``chunk_size`` dimension (e.g. ``[256, 512]``). - - Omitting a key keeps the platform default for that dimension. + ``None`` keeps the platform default. client : OgxClient - Client instance for listing and validating available models. + Authenticated OGX client used for model discovery and validation. vector_store_type : str, default="ogx" Type of vector store. Supported values: ``"ogx"`` and ``"chroma"``. @@ -233,10 +176,11 @@ def prepare_search_space_custom( Raises ------ SearchSpaceValueError - Raised when payload contains a non-recognized parameter name or - when *client* is not an :class:`OgxClient`. + Raised when payload contains a non-recognized parameter name, + when *client* is not an :class:`OgxClient`, when *chunking_methods* + contains unsupported values, or when *chunk_sizes* are out of range. """ - logger.info("Preparing custom search space based on provided constraints: %s.", payload) + logger.info("Preparing search space based on provided constraints: %s.", payload) validated_payload = AI4RAGConstraints(**payload) @@ -267,11 +211,15 @@ def prepare_search_space_custom( extra_params: list[Parameter] = [] if validated_payload.chunking_methods is not None: - extra_params.append( - Parameter(name=AI4RAGParamNames.CHUNKING_METHOD, values=tuple(validated_payload.chunking_methods)) - ) + deduped_methods = list(dict.fromkeys(validated_payload.chunking_methods)) + if len(deduped_methods) < len(validated_payload.chunking_methods): + logger.warning("Duplicate chunking_methods detected and removed: %s.", validated_payload.chunking_methods) + extra_params.append(Parameter(name=AI4RAGParamNames.CHUNKING_METHOD, values=tuple(deduped_methods))) if validated_payload.chunk_sizes is not None: - extra_params.append(Parameter(name=AI4RAGParamNames.CHUNK_SIZE, values=tuple(validated_payload.chunk_sizes))) + deduped_sizes = list(dict.fromkeys(validated_payload.chunk_sizes)) + if len(deduped_sizes) < len(validated_payload.chunk_sizes): + logger.warning("Duplicate chunk_sizes detected and removed: %s.", validated_payload.chunk_sizes) + extra_params.append(Parameter(name=AI4RAGParamNames.CHUNK_SIZE, values=tuple(deduped_sizes))) return AI4RAGSearchSpace( params=[fms_param, ems_param, *extra_params], diff --git a/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py b/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py index 7091600c..84b96a33 100644 --- a/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py +++ b/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py @@ -10,7 +10,7 @@ from ogx_client import OgxClient from pydantic import ValidationError -from ai4rag.search_space.prepare import prepare_search_space_custom, prepare_search_space_with_ogx +from ai4rag.search_space.prepare import prepare_search_space_with_ogx from ai4rag.search_space.src.exceptions import SearchSpaceValueError @@ -402,24 +402,6 @@ def test_user_picks_available_model_while_others_fail(self, mocker): fm_ids = [m.model_id for m in result["foundation_model"].values] assert fm_ids == ["llm-ok"] - def test_chunking_fields_in_payload_raise_error(self): - """chunking_methods or chunk_sizes in payload must raise SearchSpaceValueError.""" - mock_client = MagicMock(spec=OgxClient) - - with pytest.raises(SearchSpaceValueError, match="not supported by prepare_search_space_with_ogx"): - prepare_search_space_with_ogx({"chunking_methods": ["recursive"]}, mock_client) - - def test_chunk_sizes_in_payload_raise_error(self): - """chunk_sizes in payload must raise SearchSpaceValueError.""" - mock_client = MagicMock(spec=OgxClient) - - with pytest.raises(SearchSpaceValueError, match="not supported by prepare_search_space_with_ogx"): - prepare_search_space_with_ogx({"chunk_sizes": [512]}, mock_client) - - -class TestPrepareSearchSpaceCustom: - """Test prepare_search_space_custom function.""" - def _setup_mock_client(self, mocker): mock_client = MagicMock(spec=OgxClient) @@ -438,11 +420,11 @@ def _setup_mock_client(self, mocker): return mock_client - def test_no_extra_params_behaves_like_with_ogx(self, mocker): - """Without chunking_methods or chunk_sizes the result matches prepare_search_space_with_ogx.""" + def test_no_chunking_params_returns_default_chunking_dimensions(self, mocker): + """Without chunking overrides the result includes default chunking_method and chunk_size dimensions.""" mock_client = self._setup_mock_client(mocker) - result = prepare_search_space_custom({}, mock_client) + result = prepare_search_space_with_ogx({}, mock_client) param_names = [p.name for p in result.params] assert "foundation_model" in param_names @@ -454,7 +436,7 @@ def test_custom_chunking_methods_override_defaults(self, mocker): """chunking_methods in payload overrides the default chunking_method dimension.""" mock_client = self._setup_mock_client(mocker) - result = prepare_search_space_custom({"chunking_methods": ["recursive"]}, mock_client) + result = prepare_search_space_with_ogx({"chunking_methods": ["recursive"]}, mock_client) assert result["chunking_method"].values == ("recursive",) @@ -462,15 +444,15 @@ def test_custom_chunk_sizes_override_defaults(self, mocker): """chunk_sizes in payload overrides the default chunk_size dimension.""" mock_client = self._setup_mock_client(mocker) - result = prepare_search_space_custom({"chunk_sizes": [256, 512]}, mock_client) + result = prepare_search_space_with_ogx({"chunk_sizes": [256, 512]}, mock_client) assert set(result["chunk_size"].values) == {256, 512} - def test_both_custom_params_applied_together(self, mocker): + def test_both_chunking_params_applied_together(self, mocker): """Both chunking_methods and chunk_sizes in payload can be set simultaneously.""" mock_client = self._setup_mock_client(mocker) - result = prepare_search_space_custom( + result = prepare_search_space_with_ogx( {"chunking_methods": ["hybrid"], "chunk_sizes": [1024]}, mock_client, ) @@ -478,48 +460,42 @@ def test_both_custom_params_applied_together(self, mocker): assert result["chunking_method"].values == ("hybrid",) assert result["chunk_size"].values == (1024,) - def test_non_ogx_client_raises_error(self): - """Non-OgxClient raises SearchSpaceValueError.""" - with pytest.raises(SearchSpaceValueError, match="Unrecognized client type"): - prepare_search_space_custom({}, MagicMock(spec=object)) - - def test_invalid_payload_raises_validation_error(self, mocker): - """Invalid payload fields raise a pydantic ValidationError.""" + def test_duplicate_chunking_methods_are_deduplicated(self, mocker): + """Duplicate chunking_methods are silently deduplicated, preserving order.""" mock_client = self._setup_mock_client(mocker) - with pytest.raises(ValidationError): - prepare_search_space_custom({"invalid_field": "value"}, mock_client) - def test_chroma_vector_store_excludes_hybrid_params(self, mocker): - """chroma vector_store_type excludes ranker parameters.""" + result = prepare_search_space_with_ogx({"chunking_methods": ["recursive", "recursive"]}, mock_client) + + assert result["chunking_method"].values == ("recursive",) + + def test_duplicate_chunk_sizes_are_deduplicated(self, mocker): + """Duplicate chunk_sizes are silently deduplicated, preserving order.""" mock_client = self._setup_mock_client(mocker) - result = prepare_search_space_custom({}, mock_client, vector_store_type="chroma") + result = prepare_search_space_with_ogx({"chunk_sizes": [128, 129, 128]}, mock_client) - param_names = [p.name for p in result.params] - assert "ranker_strategy" not in param_names - assert "ranker_k" not in param_names - assert "ranker_alpha" not in param_names + assert result["chunk_size"].values == (128, 129) def test_unsupported_chunking_method_raises_error(self, mocker): """An unsupported chunking method raises SearchSpaceValueError before any I/O.""" mock_client = self._setup_mock_client(mocker) with pytest.raises(SearchSpaceValueError, match="Unsupported chunking methods"): - prepare_search_space_custom({"chunking_methods": ["semantic"]}, mock_client) + prepare_search_space_with_ogx({"chunking_methods": ["semantic"]}, mock_client) def test_chunk_size_below_min_raises_error(self, mocker): """A chunk size below MIN_CHUNK_SIZE raises SearchSpaceValueError before any I/O.""" mock_client = self._setup_mock_client(mocker) with pytest.raises(SearchSpaceValueError, match="out of range"): - prepare_search_space_custom({"chunk_sizes": [1]}, mock_client) + prepare_search_space_with_ogx({"chunk_sizes": [1]}, mock_client) def test_chunk_size_above_max_raises_error(self, mocker): """A chunk size above MAX_CHUNK_SIZE raises SearchSpaceValueError before any I/O.""" mock_client = self._setup_mock_client(mocker) with pytest.raises(SearchSpaceValueError, match="out of range"): - prepare_search_space_custom({"chunk_sizes": [99999]}, mock_client) + prepare_search_space_with_ogx({"chunk_sizes": [99999]}, mock_client) def test_bool_chunk_size_raises_error(self, mocker): """A bool value in chunk_sizes raises SearchSpaceValueError. @@ -530,4 +506,4 @@ def test_bool_chunk_size_raises_error(self, mocker): mock_client = self._setup_mock_client(mocker) with pytest.raises(SearchSpaceValueError, match="out of range"): - prepare_search_space_custom({"chunk_sizes": [True]}, mock_client) + prepare_search_space_with_ogx({"chunk_sizes": [True]}, mock_client) From a5e9547629fa7d33d61d0afab48f914e4121d635 Mon Sep 17 00:00:00 2001 From: Mateusz Switala Date: Mon, 6 Jul 2026 12:19:27 +0200 Subject: [PATCH 09/13] remove sphinx rest markup Signed-off-by: Mateusz Switala Assisted-by: Claude Code --- .../prepare/prepare_search_space.py | 51 +++++++++---------- 1 file changed, 25 insertions(+), 26 deletions(-) diff --git a/ai4rag/search_space/prepare/prepare_search_space.py b/ai4rag/search_space/prepare/prepare_search_space.py index 88d56f98..c8a2eb91 100644 --- a/ai4rag/search_space/prepare/prepare_search_space.py +++ b/ai4rag/search_space/prepare/prepare_search_space.py @@ -39,13 +39,12 @@ def _resolve_models_from_payload( Returns ------- tuple[list, list] - A ``(foundation_models, embedding_models)`` pair of instantiated - model objects ready for use in the search space. + A (foundation_models, embedding_models) pair of instantiated model objects. Raises ------ SearchSpaceValueError - When *client* is not an :class:`OgxClient` instance. + When client is not an OgxClient instance. """ if not isinstance(client, OgxClient): raise SearchSpaceValueError(f"Unrecognized client type: '{client.__class__.__name__}'") @@ -91,7 +90,7 @@ def _apply_language_detection(foundation_models: list, benchmark_data: pd.DataFr foundation_models : list Foundation model objects to update with detected language. benchmark_data : pd.DataFrame - Benchmark data whose ``question`` column is sampled for detection. + Benchmark data whose "question" column is sampled for detection. """ for fm in foundation_models: lang = detect_language_with_llm( @@ -106,7 +105,7 @@ def _apply_language_detection(foundation_models: list, benchmark_data: pd.DataFr def _build_model_params(foundation_models: list, embedding_models: list) -> tuple[Parameter, Parameter]: - """Create :class:`Parameter` objects for model lists and log selections. + """Create Parameter objects for model lists and log selections. Parameters ---------- @@ -118,7 +117,7 @@ def _build_model_params(foundation_models: list, embedding_models: list) -> tupl Returns ------- tuple[Parameter, Parameter] - ``(fms_param, ems_param)`` ready for inclusion in a search space. + (fms_param, ems_param) ready for inclusion in a search space. """ fms_param = Parameter(name=AI4RAGParamNames.FOUNDATION_MODEL, values=foundation_models) ems_param = Parameter(name=AI4RAGParamNames.EMBEDDING_MODEL, values=embedding_models) @@ -133,35 +132,35 @@ def prepare_search_space_with_ogx( vector_store_type: str = "ogx", benchmark_data: pd.DataFrame | None = None, ) -> AI4RAGSearchSpace: - """Prepare an :class:`AI4RAGSearchSpace` using OGX for model validation. + """Prepare an AI4RAGSearchSpace using OGX for model validation. Foundation and embedding models are discovered and validated via the OGX - platform. Chunking parameters (``chunking_methods``, ``chunk_sizes``) are - validated locally against :class:`~ai4rag.utils.constants.ChunkingConstraints` - and, when provided, override the platform defaults for those dimensions. + platform. Chunking parameters (chunking_methods, chunk_sizes) are validated + locally against ChunkingConstraints and, when provided, override the platform + defaults for those dimensions. Parameters ---------- payload : dict[str, Any] - A mapping of constraint names to their values. Supported keys: - - - ``"foundation_models"`` *(list[dict])* — foundation model identifiers - to include; ``None`` uses all OGX defaults. - - ``"embedding_models"`` *(list[dict])* — embedding model identifiers - to include; ``None`` uses all OGX defaults. - - ``"chunking_methods"`` *(list[str])* — overrides the default - ``chunking_method`` dimension (e.g. ``["recursive"]``). - ``None`` keeps the platform default. - - ``"chunk_sizes"`` *(list[int])* — overrides the default - ``chunk_size`` dimension (e.g. ``[256, 512]``). - ``None`` keeps the platform default. + A mapping of constraint names to their values. Supported keys: + + - "foundation_models" (list[dict]) — foundation model identifiers + to include; None uses all OGX defaults. + - "embedding_models" (list[dict]) — embedding model identifiers + to include; None uses all OGX defaults. + - "chunking_methods" (list[str]) — overrides the default + chunking_method dimension (e.g. ["recursive"]). + None keeps the platform default. + - "chunk_sizes" (list[int]) — overrides the default + chunk_size dimension (e.g. [256, 512]). + None keeps the platform default. client : OgxClient Authenticated OGX client used for model discovery and validation. vector_store_type : str, default="ogx" - Type of vector store. Supported values: ``"ogx"`` and ``"chroma"``. - When ``"chroma"``, hybrid search parameters are excluded from the + Type of vector store. Supported values: "ogx" and "chroma". + When "chroma", hybrid search parameters are excluded from the default search space since ChromaDB does not support hybrid search. benchmark_data : pd.DataFrame | None, default=None @@ -177,8 +176,8 @@ def prepare_search_space_with_ogx( ------ SearchSpaceValueError Raised when payload contains a non-recognized parameter name, - when *client* is not an :class:`OgxClient`, when *chunking_methods* - contains unsupported values, or when *chunk_sizes* are out of range. + when client is not an OgxClient, when chunking_methods contains + unsupported values, or when chunk_sizes are out of range. """ logger.info("Preparing search space based on provided constraints: %s.", payload) From d02f09aff5ff41ab5ca47eff7c0aa7ca066e215a Mon Sep 17 00:00:00 2001 From: Mateusz Switala Date: Mon, 6 Jul 2026 12:30:18 +0200 Subject: [PATCH 10/13] improve chunking related params validation Signed-off-by: Mateusz Switala Assisted-by: Claude Code --- .../prepare/input_payload_types.py | 37 ++++++++++++++++--- .../prepare/prepare_search_space.py | 28 +++----------- .../prepare/test_prepare_search_space.py | 20 ++++------ 3 files changed, 45 insertions(+), 40 deletions(-) diff --git a/ai4rag/search_space/prepare/input_payload_types.py b/ai4rag/search_space/prepare/input_payload_types.py index c53cb38f..481460b1 100644 --- a/ai4rag/search_space/prepare/input_payload_types.py +++ b/ai4rag/search_space/prepare/input_payload_types.py @@ -4,11 +4,10 @@ # ----------------------------------------------------------------------------- from typing import Annotated, Optional -from annotated_types import Gt, MinLen -from pydantic import ( - BaseModel, - ConfigDict, -) +from annotated_types import Ge, Le, MinLen +from pydantic import BaseModel, ConfigDict, field_validator + +from ai4rag.utils.constants import ChunkingConstraints CONFIG = ConfigDict(extra="forbid") @@ -37,4 +36,30 @@ class AI4RAGConstraints(BaseModel): embedding_models: Optional[Annotated[list[AI4RAGEmbeddingModel], MinLen(1)]] = None foundation_models: Optional[Annotated[list[AI4RAGFoundationModel], MinLen(1)]] = None chunking_methods: Optional[Annotated[list[Annotated[str, MinLen(1)]], MinLen(1)]] = None - chunk_sizes: Optional[Annotated[list[Annotated[int, Gt(0)]], MinLen(1)]] = None + chunk_sizes: Optional[ + Annotated[ + list[Annotated[int, Ge(ChunkingConstraints.MIN_CHUNK_SIZE), Le(ChunkingConstraints.MAX_CHUNK_SIZE)]], + MinLen(1), + ] + ] = None + + @field_validator("chunk_sizes", mode="before") + @classmethod + def _reject_bool_chunk_sizes(cls, v): + if isinstance(v, list): + for i, s in enumerate(v): + if isinstance(s, bool): + raise ValueError(f"chunk_sizes[{i}] must be a positive integer, got bool.") + return v + + @field_validator("chunking_methods", mode="after") + @classmethod + def _validate_chunking_methods(cls, v): + if v is not None: + unsupported = [m for m in v if m not in ChunkingConstraints.METHODS] + if unsupported: + raise ValueError( + f"Unsupported chunking methods: {unsupported!r}. " + f"Supported methods: {ChunkingConstraints.METHODS!r}." + ) + return v diff --git a/ai4rag/search_space/prepare/prepare_search_space.py b/ai4rag/search_space/prepare/prepare_search_space.py index c8a2eb91..b4856b4f 100644 --- a/ai4rag/search_space/prepare/prepare_search_space.py +++ b/ai4rag/search_space/prepare/prepare_search_space.py @@ -18,7 +18,7 @@ from ai4rag.search_space.src.exceptions import SearchSpaceValueError from ai4rag.search_space.src.parameter import Parameter from ai4rag.search_space.src.search_space import AI4RAGSearchSpace -from ai4rag.utils.constants import AI4RAGParamNames, ChunkingConstraints +from ai4rag.utils.constants import AI4RAGParamNames __all__ = ["prepare_search_space_with_ogx"] @@ -174,33 +174,17 @@ def prepare_search_space_with_ogx( Raises ------ - SearchSpaceValueError + pydantic.ValidationError Raised when payload contains a non-recognized parameter name, - when client is not an OgxClient, when chunking_methods contains - unsupported values, or when chunk_sizes are out of range. + when chunking_methods contains unsupported values, or when + chunk_sizes are out of range or contain bool values. + SearchSpaceValueError + Raised when client is not an OgxClient. """ logger.info("Preparing search space based on provided constraints: %s.", payload) validated_payload = AI4RAGConstraints(**payload) - if validated_payload.chunking_methods is not None: - unsupported = [m for m in validated_payload.chunking_methods if m not in ChunkingConstraints.METHODS] - if unsupported: - raise SearchSpaceValueError( - f"Unsupported chunking methods: {unsupported!r}. " - f"Supported methods: {ChunkingConstraints.METHODS!r}." - ) - - if validated_payload.chunk_sizes is not None: - for i, s in enumerate(validated_payload.chunk_sizes): - if isinstance(s, bool): - raise SearchSpaceValueError(f"chunk_sizes[{i}] must be a positive integer, got bool.") - if not ChunkingConstraints.MIN_CHUNK_SIZE <= s <= ChunkingConstraints.MAX_CHUNK_SIZE: - raise SearchSpaceValueError( - f"chunk_sizes[{i}]={s} is out of range " - f"[{ChunkingConstraints.MIN_CHUNK_SIZE}, {ChunkingConstraints.MAX_CHUNK_SIZE}]." - ) - foundation_models, embedding_models = _resolve_models_from_payload(validated_payload, client) if benchmark_data is not None: diff --git a/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py b/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py index 84b96a33..7b1a0bb6 100644 --- a/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py +++ b/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py @@ -477,33 +477,29 @@ def test_duplicate_chunk_sizes_are_deduplicated(self, mocker): assert result["chunk_size"].values == (128, 129) def test_unsupported_chunking_method_raises_error(self, mocker): - """An unsupported chunking method raises SearchSpaceValueError before any I/O.""" + """An unsupported chunking method raises ValidationError before any I/O.""" mock_client = self._setup_mock_client(mocker) - with pytest.raises(SearchSpaceValueError, match="Unsupported chunking methods"): + with pytest.raises(ValidationError, match="Unsupported chunking methods"): prepare_search_space_with_ogx({"chunking_methods": ["semantic"]}, mock_client) def test_chunk_size_below_min_raises_error(self, mocker): - """A chunk size below MIN_CHUNK_SIZE raises SearchSpaceValueError before any I/O.""" + """A chunk size below MIN_CHUNK_SIZE raises ValidationError before any I/O.""" mock_client = self._setup_mock_client(mocker) - with pytest.raises(SearchSpaceValueError, match="out of range"): + with pytest.raises(ValidationError): prepare_search_space_with_ogx({"chunk_sizes": [1]}, mock_client) def test_chunk_size_above_max_raises_error(self, mocker): - """A chunk size above MAX_CHUNK_SIZE raises SearchSpaceValueError before any I/O.""" + """A chunk size above MAX_CHUNK_SIZE raises ValidationError before any I/O.""" mock_client = self._setup_mock_client(mocker) - with pytest.raises(SearchSpaceValueError, match="out of range"): + with pytest.raises(ValidationError): prepare_search_space_with_ogx({"chunk_sizes": [99999]}, mock_client) def test_bool_chunk_size_raises_error(self, mocker): - """A bool value in chunk_sizes raises SearchSpaceValueError. - - Pydantic coerces True→1 before our validator runs, so the error is - "out of range" (1 < MIN_CHUNK_SIZE) rather than "must be a positive integer". - """ + """A bool value in chunk_sizes raises ValidationError before any I/O.""" mock_client = self._setup_mock_client(mocker) - with pytest.raises(SearchSpaceValueError, match="out of range"): + with pytest.raises(ValidationError, match="must be a positive integer"): prepare_search_space_with_ogx({"chunk_sizes": [True]}, mock_client) From f741cfaf5daa2d7b89c3b39384b061db105ffdf9 Mon Sep 17 00:00:00 2001 From: Mateusz Switala Date: Mon, 6 Jul 2026 12:32:20 +0200 Subject: [PATCH 11/13] remove redundant function Signed-off-by: Mateusz Switala Assisted-by: Claude Code --- .../prepare/prepare_search_space.py | 27 +++---------------- 1 file changed, 4 insertions(+), 23 deletions(-) diff --git a/ai4rag/search_space/prepare/prepare_search_space.py b/ai4rag/search_space/prepare/prepare_search_space.py index b4856b4f..a93fe0b4 100644 --- a/ai4rag/search_space/prepare/prepare_search_space.py +++ b/ai4rag/search_space/prepare/prepare_search_space.py @@ -104,28 +104,6 @@ def _apply_language_detection(foundation_models: list, benchmark_data: pd.DataFr logger.warning("Model %s: language detection failed, falling back to auto-detect.", fm.model_id) -def _build_model_params(foundation_models: list, embedding_models: list) -> tuple[Parameter, Parameter]: - """Create Parameter objects for model lists and log selections. - - Parameters - ---------- - foundation_models : list - Foundation model instances to wrap in a Parameter. - embedding_models : list - Embedding model instances to wrap in a Parameter. - - Returns - ------- - tuple[Parameter, Parameter] - (fms_param, ems_param) ready for inclusion in a search space. - """ - fms_param = Parameter(name=AI4RAGParamNames.FOUNDATION_MODEL, values=foundation_models) - ems_param = Parameter(name=AI4RAGParamNames.EMBEDDING_MODEL, values=embedding_models) - logger.info("Selected foundation models for the experiment: %s.", [m.model_id for m in fms_param.values]) - logger.info("Selected embedding models for the experiment: %s.", [m.model_id for m in ems_param.values]) - return fms_param, ems_param - - def prepare_search_space_with_ogx( payload: dict[str, Any], client: OgxClient, @@ -190,7 +168,10 @@ def prepare_search_space_with_ogx( if benchmark_data is not None: _apply_language_detection(foundation_models, benchmark_data) - fms_param, ems_param = _build_model_params(foundation_models, embedding_models) + fms_param = Parameter(name=AI4RAGParamNames.FOUNDATION_MODEL, values=foundation_models) + ems_param = Parameter(name=AI4RAGParamNames.EMBEDDING_MODEL, values=embedding_models) + logger.info("Selected foundation models for the experiment: %s.", [m.model_id for m in fms_param.values]) + logger.info("Selected embedding models for the experiment: %s.", [m.model_id for m in ems_param.values]) extra_params: list[Parameter] = [] if validated_payload.chunking_methods is not None: From 2a0f36ffb977c7c66b16776511dc49bcbdb8428f Mon Sep 17 00:00:00 2001 From: Mateusz Switala Date: Mon, 6 Jul 2026 12:42:23 +0200 Subject: [PATCH 12/13] improve chunking params validation and deduplication Signed-off-by: Mateusz Switala Assisted-by: Claude Code --- .../prepare/input_payload_types.py | 11 +++++++++- .../prepare/prepare_search_space.py | 21 ++++++++++++------- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/ai4rag/search_space/prepare/input_payload_types.py b/ai4rag/search_space/prepare/input_payload_types.py index 481460b1..fb93f14c 100644 --- a/ai4rag/search_space/prepare/input_payload_types.py +++ b/ai4rag/search_space/prepare/input_payload_types.py @@ -52,9 +52,16 @@ def _reject_bool_chunk_sizes(cls, v): raise ValueError(f"chunk_sizes[{i}] must be a positive integer, got bool.") return v + @field_validator("chunk_sizes", mode="after") + @classmethod + def _deduplicate_chunk_sizes(cls, v): + if v is not None: + return list(dict.fromkeys(v)) + return v + @field_validator("chunking_methods", mode="after") @classmethod - def _validate_chunking_methods(cls, v): + def _validate_and_deduplicate_chunking_methods(cls, v): if v is not None: unsupported = [m for m in v if m not in ChunkingConstraints.METHODS] if unsupported: @@ -62,4 +69,6 @@ def _validate_chunking_methods(cls, v): f"Unsupported chunking methods: {unsupported!r}. " f"Supported methods: {ChunkingConstraints.METHODS!r}." ) + # use dict.fromkeys to deduplicate chunking_methods + return list(dict.fromkeys(v)) return v diff --git a/ai4rag/search_space/prepare/prepare_search_space.py b/ai4rag/search_space/prepare/prepare_search_space.py index a93fe0b4..e384fa98 100644 --- a/ai4rag/search_space/prepare/prepare_search_space.py +++ b/ai4rag/search_space/prepare/prepare_search_space.py @@ -163,6 +163,15 @@ def prepare_search_space_with_ogx( validated_payload = AI4RAGConstraints(**payload) + if validated_payload.chunking_methods is not None and len(validated_payload.chunking_methods) < len( + payload.get("chunking_methods", []) + ): + logger.warning("Duplicate chunking_methods removed: %s.", payload["chunking_methods"]) + if validated_payload.chunk_sizes is not None and len(validated_payload.chunk_sizes) < len( + payload.get("chunk_sizes", []) + ): + logger.warning("Duplicate chunk_sizes removed: %s.", payload["chunk_sizes"]) + foundation_models, embedding_models = _resolve_models_from_payload(validated_payload, client) if benchmark_data is not None: @@ -175,15 +184,11 @@ def prepare_search_space_with_ogx( extra_params: list[Parameter] = [] if validated_payload.chunking_methods is not None: - deduped_methods = list(dict.fromkeys(validated_payload.chunking_methods)) - if len(deduped_methods) < len(validated_payload.chunking_methods): - logger.warning("Duplicate chunking_methods detected and removed: %s.", validated_payload.chunking_methods) - extra_params.append(Parameter(name=AI4RAGParamNames.CHUNKING_METHOD, values=tuple(deduped_methods))) + extra_params.append( + Parameter(name=AI4RAGParamNames.CHUNKING_METHOD, values=tuple(validated_payload.chunking_methods)) + ) if validated_payload.chunk_sizes is not None: - deduped_sizes = list(dict.fromkeys(validated_payload.chunk_sizes)) - if len(deduped_sizes) < len(validated_payload.chunk_sizes): - logger.warning("Duplicate chunk_sizes detected and removed: %s.", validated_payload.chunk_sizes) - extra_params.append(Parameter(name=AI4RAGParamNames.CHUNK_SIZE, values=tuple(deduped_sizes))) + extra_params.append(Parameter(name=AI4RAGParamNames.CHUNK_SIZE, values=tuple(validated_payload.chunk_sizes))) return AI4RAGSearchSpace( params=[fms_param, ems_param, *extra_params], From 4e99c4435b5026cb9c959884411aab2145075555 Mon Sep 17 00:00:00 2001 From: Mateusz Switala Date: Mon, 6 Jul 2026 12:48:36 +0200 Subject: [PATCH 13/13] improve tests Signed-off-by: Mateusz Switala Assisted-by: Claude Code --- .../prepare/test_prepare_search_space.py | 85 +++---------------- 1 file changed, 10 insertions(+), 75 deletions(-) diff --git a/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py b/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py index 7b1a0bb6..1e5fedae 100644 --- a/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py +++ b/tests/unit/ai4rag/search_space/prepare/test_prepare_search_space.py @@ -135,36 +135,10 @@ def test_payload_with_custom_embedding_models(self, mocker): assert len(embedding_param.values) == 1 assert embedding_param.values[0].model_id == "custom-embedding" - def test_invalid_payload_raises_validation_error(self, mocker): + def test_invalid_payload_raises_validation_error(self): """Test that invalid payload raises validation error.""" - mock_client = MagicMock(spec=OgxClient) - - # Mock model list - mock_llm = Mock() - mock_llm.id = "default-llm" - mock_llm.custom_metadata = {"model_type": "llm"} - - mock_embedding = Mock() - mock_embedding.id = "default-embedding" - mock_embedding.custom_metadata = {"model_type": "embedding", "embedding_dimension": 768} - - mock_client.models.list.return_value.data = [mock_llm, mock_embedding] - - # Mock validation functions to always return True - mocker.patch( - "ai4rag.search_space.prepare.ogx_utils._validate_foundation_model", - return_value=True, - ) - mocker.patch( - "ai4rag.search_space.prepare.ogx_utils._validate_embedding_model", - return_value=True, - ) - - # Invalid payload with unrecognized parameter - payload = {"invalid_parameter": "value"} - with pytest.raises(ValidationError, match="Unknown validation error|invalid_parameter"): - prepare_search_space_with_ogx(payload, mock_client) + prepare_search_space_with_ogx({"invalid_parameter": "value"}, MagicMock()) def test_non_ogx_client_raises_error(self): """Test that non-OgxClient raises error.""" @@ -244,37 +218,6 @@ def test_ogx_vector_store_includes_hybrid_params_by_default(self, mocker): assert "vector" in search_mode_param.values assert "hybrid" in search_mode_param.values - def test_default_vector_store_type_is_ogx(self, mocker): - """Test that default vector_store_type is ogx (includes hybrid params by default).""" - mock_client = MagicMock(spec=OgxClient) - - mock_llm = Mock() - mock_llm.id = "default-llm" - mock_llm.custom_metadata = {"model_type": "llm"} - - mock_embedding = Mock() - mock_embedding.id = "default-embedding" - mock_embedding.custom_metadata = {"model_type": "embedding", "embedding_dimension": 768} - - mock_client.models.list.return_value.data = [mock_llm, mock_embedding] - - mocker.patch( - "ai4rag.search_space.prepare.ogx_utils._validate_foundation_model", - return_value=True, - ) - mocker.patch( - "ai4rag.search_space.prepare.ogx_utils._validate_embedding_model", - return_value=True, - ) - - result = prepare_search_space_with_ogx({}, mock_client) - - param_names = [p.name for p in result.params] - assert "search_mode" in param_names - assert "ranker_strategy" in param_names - assert "ranker_k" in param_names - assert "ranker_alpha" in param_names - def test_user_specifies_not_responding_foundation_model(self, mocker): """Error when user requests a foundation model that is registered but not responding.""" mock_client = MagicMock(spec=OgxClient) @@ -476,30 +419,22 @@ def test_duplicate_chunk_sizes_are_deduplicated(self, mocker): assert result["chunk_size"].values == (128, 129) - def test_unsupported_chunking_method_raises_error(self, mocker): + def test_unsupported_chunking_method_raises_error(self): """An unsupported chunking method raises ValidationError before any I/O.""" - mock_client = self._setup_mock_client(mocker) - with pytest.raises(ValidationError, match="Unsupported chunking methods"): - prepare_search_space_with_ogx({"chunking_methods": ["semantic"]}, mock_client) + prepare_search_space_with_ogx({"chunking_methods": ["semantic"]}, MagicMock()) - def test_chunk_size_below_min_raises_error(self, mocker): + def test_chunk_size_below_min_raises_error(self): """A chunk size below MIN_CHUNK_SIZE raises ValidationError before any I/O.""" - mock_client = self._setup_mock_client(mocker) - with pytest.raises(ValidationError): - prepare_search_space_with_ogx({"chunk_sizes": [1]}, mock_client) + prepare_search_space_with_ogx({"chunk_sizes": [1]}, MagicMock()) - def test_chunk_size_above_max_raises_error(self, mocker): + def test_chunk_size_above_max_raises_error(self): """A chunk size above MAX_CHUNK_SIZE raises ValidationError before any I/O.""" - mock_client = self._setup_mock_client(mocker) - with pytest.raises(ValidationError): - prepare_search_space_with_ogx({"chunk_sizes": [99999]}, mock_client) + prepare_search_space_with_ogx({"chunk_sizes": [99999]}, MagicMock()) - def test_bool_chunk_size_raises_error(self, mocker): + def test_bool_chunk_size_raises_error(self): """A bool value in chunk_sizes raises ValidationError before any I/O.""" - mock_client = self._setup_mock_client(mocker) - with pytest.raises(ValidationError, match="must be a positive integer"): - prepare_search_space_with_ogx({"chunk_sizes": [True]}, mock_client) + prepare_search_space_with_ogx({"chunk_sizes": [True]}, MagicMock())