diff --git a/ai4rag/components/optimization/search_space_preparation.py b/ai4rag/components/optimization/search_space_preparation.py index ca1d0a5e..165434af 100644 --- a/ai4rag/components/optimization/search_space_preparation.py +++ b/ai4rag/components/optimization/search_space_preparation.py @@ -113,6 +113,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 +154,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 @@ -171,27 +176,47 @@ def prepare_search_space_report( # pylint: disable=too-many-locals,too-many-arg If *metric* is not one of the supported values. TypeError 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) # 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_with_ogx( + payload, + 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 @@ -230,16 +255,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, @@ -255,16 +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.") - - -def _validate_chunking_methods(methods: list[str] | None) -> None: - """Validate that chunking methods, if provided, are non-empty strings.""" - 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.") diff --git a/ai4rag/search_space/prepare/input_payload_types.py b/ai4rag/search_space/prepare/input_payload_types.py index 2575fb63..fb93f14c 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 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") @@ -36,3 +35,40 @@ 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, 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("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_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: + raise ValueError( + 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 75018b78..e384fa98 100644 --- a/ai4rag/search_space/prepare/prepare_search_space.py +++ b/ai4rag/search_space/prepare/prepare_search_space.py @@ -18,104 +18,179 @@ 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"] -def prepare_search_space_with_ogx( - payload: dict[str, Any], +def _resolve_models_from_payload( + validated_payload: AI4RAGConstraints, client: OgxClient, - vector_store_type: str = "ogx", - benchmark_data: pd.DataFrame | None = None, -) -> AI4RAGSearchSpace: - """ - Prepare AutoRAGSearchSpace. +) -> tuple[list, list]: + """Retrieve and validate foundation and embedding models from OGX. Parameters ---------- - payload : dict[str, Any] - A mapping between parameter name and its associated values. - + validated_payload : AI4RAGConstraints + Validated constraint payload specifying which models to include. 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. + Authenticated OGX client used for model discovery and validation. Returns ------- - AI4RAGSearchSpace - A valid AI4RAGSearchSpace used in RAG optimization process. + tuple[list, list] + A (foundation_models, embedding_models) pair of instantiated model objects. Raises ------ SearchSpaceValueError - Raised when payload contains non-recognized parameter name. + When client is not an OgxClient instance. """ - logger.info("Preparing search space based on provided constraints: %s.", payload) - - 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: + 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=registered_foundation_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=registered_foundation_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=registered_embedding_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=registered_embedding_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 prepare_search_space_with_ogx( + payload: dict[str, Any], + client: OgxClient, + vector_store_type: str = "ogx", + benchmark_data: pd.DataFrame | None = None, +) -> AI4RAGSearchSpace: + """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 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. + + 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 + 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 + ------ + pydantic.ValidationError + Raised when payload contains a non-recognized parameter name, + 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 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: - 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) - - fms_param = Parameter(name="foundation_model", values=foundation_models) - ems_param = Parameter(name="embedding_model", values=embedding_models) + _apply_language_detection(foundation_models, benchmark_data) + 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: + 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], + 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..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,7 +12,6 @@ from ai4rag.components.optimization.search_space_preparation import ( SUPPORTED_METRICS, SearchSpaceReport, - _validate_chunking_methods, _validate_model_list, prepare_search_space_report, ) @@ -172,107 +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] - - -# --------------------------------------------------------------------------- -# 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 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"], - ) 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..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) @@ -401,3 +344,97 @@ 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 _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_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_with_ogx({}, 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_with_ogx({"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_with_ogx({"chunk_sizes": [256, 512]}, mock_client) + + assert set(result["chunk_size"].values) == {256, 512} + + 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_with_ogx( + {"chunking_methods": ["hybrid"], "chunk_sizes": [1024]}, + mock_client, + ) + + assert result["chunking_method"].values == ("hybrid",) + assert result["chunk_size"].values == (1024,) + + def test_duplicate_chunking_methods_are_deduplicated(self, mocker): + """Duplicate chunking_methods are silently deduplicated, preserving order.""" + mock_client = self._setup_mock_client(mocker) + + 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_with_ogx({"chunk_sizes": [128, 129, 128]}, mock_client) + + assert result["chunk_size"].values == (128, 129) + + def test_unsupported_chunking_method_raises_error(self): + """An unsupported chunking method raises ValidationError before any I/O.""" + with pytest.raises(ValidationError, match="Unsupported chunking methods"): + prepare_search_space_with_ogx({"chunking_methods": ["semantic"]}, MagicMock()) + + def test_chunk_size_below_min_raises_error(self): + """A chunk size below MIN_CHUNK_SIZE raises ValidationError before any I/O.""" + with pytest.raises(ValidationError): + prepare_search_space_with_ogx({"chunk_sizes": [1]}, MagicMock()) + + def test_chunk_size_above_max_raises_error(self): + """A chunk size above MAX_CHUNK_SIZE raises ValidationError before any I/O.""" + with pytest.raises(ValidationError): + prepare_search_space_with_ogx({"chunk_sizes": [99999]}, MagicMock()) + + def test_bool_chunk_size_raises_error(self): + """A bool value in chunk_sizes raises ValidationError before any I/O.""" + with pytest.raises(ValidationError, match="must be a positive integer"): + prepare_search_space_with_ogx({"chunk_sizes": [True]}, MagicMock())