chore: add validation for chunk sizes and chunking methods

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.")

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

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,
     )

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