Skip to content

chore: add validation for chunk sizes and chunking methods#89

Merged
Mateusz-Switala merged 13 commits into
IBM:mainfrom
Mateusz-Switala:chore-refactor-prepare-search-space
Jul 6, 2026
Merged

chore: add validation for chunk sizes and chunking methods#89
Mateusz-Switala merged 13 commits into
IBM:mainfrom
Mateusz-Switala:chore-refactor-prepare-search-space

Conversation

@Mateusz-Switala

@Mateusz-Switala Mateusz-Switala commented Jul 2, 2026

Copy link
Copy Markdown
Collaborator

Assisted-by: Claude Code

Description

Extends prepare_search_space_with_ogx to accept optional chunking_method and chunk_size
overrides directly via the payload, and adds chunk_sizes as a new parameter to
prepare_search_space_report. Also lowers the minimum allowed chunk size to 128 tokens and
deduplicates any repeated values in the chunking inputs.

Motivation

The existing prepare_search_space_with_ogx always used platform-default chunking methods
(recursive, hybrid) and chunk sizes (512, 1024, 2048). Users running experiments on
shorter documents or with a single chunking strategy had no way to constrain these dimensions
without post-processing the search space manually. This change exposes that control at the
preparation step, in the single function callers already use.

Changes

  • ai4rag/utils/constants.py — lowered ChunkingConstraints.MIN_CHUNK_SIZE from 512 to 128
  • ai4rag/search_space/prepare/input_payload_types.py — added optional chunking_methods and
    chunk_sizes fields to AI4RAGConstraints; both are pydantic-validated (non-empty list,
    non-empty strings / positive integers respectively)
  • ai4rag/search_space/prepare/prepare_search_space.py
    • Extracted three private helpers to eliminate duplication: _resolve_models_from_payload,
      _apply_language_detection, _build_model_params
    • Extended prepare_search_space_with_ogx to accept and validate chunking_methods and
      chunk_sizes in the payload; when provided they override the platform defaults for those
      dimensions; duplicate values are deduplicated with a warning
    • Updated docstring to clarify that only model validation goes through OGX; chunking is
      validated locally against ChunkingConstraints
  • ai4rag/search_space/prepare/__init__.py — removed prepare_search_space_custom
    (functionality merged into prepare_search_space_with_ogx)
  • ai4rag/components/optimization/search_space_preparation.py
    • Added chunk_sizes: list[int] | None = None parameter to prepare_search_space_report
    • Updated _validate_chunking_methods to validate values against ChunkingConstraints.METHODS,
      providing an early, clear error
    • prepare_search_space_report now passes chunking_methods and chunk_sizes into the payload
      and calls prepare_search_space_with_ogx

Testing

  • Added tests for chunking overrides, individual and combined, via prepare_search_space_with_ogx
  • Added tests for duplicate deduplication in chunking_methods and chunk_sizes
  • Added TestValidateChunkSizes (11 tests) covering: None, valid list, single element,
    non-list, empty list, zero, negative, float, below-min, above-max, and exact boundary values
  • Added TestPrepareSearchSpaceReportChunkSizesValidation (3 tests) covering: non-list, empty
    list, and non-positive integer
  • Simplified TestUnsupportedChunkingMethods — validation fires before any I/O so complex
    mocking is no longer needed

Checklist

  • Tests added/updated
  • Documentation updated
  • Code follows style gu

@Mateusz-Switala Mateusz-Switala marked this pull request as ready for review July 3, 2026 08:24
@Mateusz-Switala Mateusz-Switala force-pushed the chore-refactor-prepare-search-space branch from 1f00377 to 3f9c2fc Compare July 3, 2026 08:27
Comment thread ai4rag/components/optimization/search_space_preparation.py Outdated
Comment thread ai4rag/components/optimization/search_space_preparation.py Outdated
Comment thread ai4rag/components/optimization/search_space_preparation.py Outdated
Comment thread ai4rag/search_space/prepare/prepare_search_space.py Outdated
Comment thread ai4rag/search_space/prepare/prepare_search_space.py
Signed-off-by: Mateusz Switala <mswitala@redhat.com>
Assisted-by: Claude Code
Signed-off-by: Mateusz Switala <mswitala@redhat.com>
Signed-off-by: Mateusz Switala <mswitala@redhat.com>
Signed-off-by: Mateusz Switala <mswitala@redhat.com>
Assisted-by: Claude Code
Signed-off-by: Mateusz Switala <mswitala@redhat.com>
Signed-off-by: Mateusz Switala <mswitala@redhat.com>
Assisted-by: Claude Code
Signed-off-by: Mateusz Switala <mswitala@redhat.com>
@Mateusz-Switala Mateusz-Switala force-pushed the chore-refactor-prepare-search-space branch from 4a1841d to 43700f9 Compare July 3, 2026 13:57
Comment thread tests/unit/ai4rag/components/optimization/test_search_space_prep.py

@filip-komarzyniec filip-komarzyniec left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🆗 /lgtm

Signed-off-by: Mateusz Switala <mswitala@redhat.com>
Assisted-by: Claude Code
Comment thread ai4rag/search_space/prepare/prepare_search_space.py Outdated
Comment thread ai4rag/search_space/prepare/prepare_search_space.py Outdated
Comment thread ai4rag/search_space/prepare/prepare_search_space.py Outdated
Comment thread ai4rag/search_space/prepare/prepare_search_space.py Outdated
Comment thread ai4rag/search_space/prepare/prepare_search_space.py Outdated

@filip-komarzyniec filip-komarzyniec left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

some comments to think through, everything else totally fine 👌🏻

Signed-off-by: Mateusz Switala <mswitala@redhat.com>
Assisted-by: Claude Code
Signed-off-by: Mateusz Switala <mswitala@redhat.com>
Assisted-by: Claude Code
Signed-off-by: Mateusz Switala <mswitala@redhat.com>
Assisted-by: Claude Code
Signed-off-by: Mateusz Switala <mswitala@redhat.com>
Assisted-by: Claude Code
Signed-off-by: Mateusz Switala <mswitala@redhat.com>
Assisted-by: Claude Code
@Mateusz-Switala Mateusz-Switala merged commit aa5f02f into IBM:main Jul 6, 2026
4 checks passed
@Mateusz-Switala Mateusz-Switala deleted the chore-refactor-prepare-search-space branch July 6, 2026 12:58
@Mateusz-Switala Mateusz-Switala restored the chore-refactor-prepare-search-space branch July 7, 2026 15:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants