Optimization Components¶

def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arguments,too-many-positional-arguments
    test_data_path: str | Path,
    extracted_text_path: str | Path,
    ogx_client: OgxClient,
    embedding_models: list[str] | None = None,
    generation_models: list[str] | None = None,
    top_n_generation: int = _DEFAULT_TOP_N_GENERATION,
    top_k_embedding: int = _DEFAULT_TOP_K_EMBEDDING,
    sample_size: int = _DEFAULT_SAMPLE_SIZE,
    random_seed: int = _DEFAULT_SEED,
    chunking_methods: list[str] | None = None,
    chunk_sizes: list[int] | None = None,
    chunk_overlaps: list[int] | None = None,
    inference_max_threads: int = 10,
) -> SearchSpaceReport:
    """Run model pre-selection and prepare a search-space report.

    Builds an :class:`AI4RAGSearchSpace` from the given model lists, runs
    :class:`ModelsPreSelector` when the number of models exceeds the
    configured caps, detects the benchmark language, and returns a
    structured report.

    Parameters
    ----------
    test_data_path
        Path to a JSON file containing benchmark questions and expected
        answers.
    extracted_text_path
        Path to a single DoclingDocument JSON file or a directory of such
        files.
    ogx_client
        An authenticated :class:`OgxClient` instance.
    embedding_models
        Embedding model identifiers.  ``None`` uses the server defaults.
    generation_models
        Generation model identifiers.  ``None`` uses the server defaults.
    top_n_generation
        Maximum number of generation models to retain.
    top_k_embedding
        Maximum number of embedding models to retain.
    sample_size
        Number of benchmark records sampled for model pre-selection.
    random_seed
        Seed for reproducible sampling.
    chunking_methods
        When provided, constrains the ``chunking_method`` dimension of the
        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.
    chunk_overlaps
        When provided, constrains the ``chunk_overlap`` dimension of the
        search space to only these values (e.g. ``[0, 128]``).  ``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
        per-request concurrency (useful when each request carries more
        retrieved context).  Defaults to ``10``.

    Returns
    -------
    SearchSpaceReport
        Structured report containing the verbose search space, selected
        models, and detected language.

    Raises
    ------
    ValueError
        If *metric* is not one of the supported values.
    TypeError
        If *embedding_models* or *generation_models* contain invalid entries.
    pydantic.ValidationError
        If *chunking_methods*, *chunk_sizes*, or *chunk_overlaps* 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]``,
        or *chunk_overlaps* contains values outside
        ``[ChunkingConstraints.MIN_CHUNK_OVERLAP, ChunkingConstraints.MAX_CHUNK_OVERLAP]``.
    """
    _validate_model_list(embedding_models, "embedding_models")
    _validate_model_list(generation_models, "generation_models")

    # Build payload and create search space via OGX
    payload: dict[str, Any] = {}
    if generation_models:
        payload["foundation_models"] = [{"model_id": gm} for gm in generation_models]
    if embedding_models:
        payload["embedding_models"] = [{"model_id": em} for em in embedding_models]
    if chunking_methods is not None:
        payload["chunking_methods"] = chunking_methods
    if chunk_sizes is not None:
        payload["chunk_sizes"] = chunk_sizes
    if chunk_overlaps is not None:
        payload["chunk_overlaps"] = chunk_overlaps

    # Load benchmark data and documents
    benchmark_df = pd.read_json(Path(test_data_path))
    benchmark_data = BenchmarkData(benchmark_df)
    documents = load_docling_documents(extracted_text_path)

    search_space = prepare_search_space_with_ogx(
        payload,
        client=ogx_client,
        benchmark_data=benchmark_df,
    )
    _logger.info(
        "Search space chunking_method=%s chunk_size=%s chunk_overlap=%s",
        list(search_space["chunking_method"].values),
        list(search_space["chunk_size"].values),
        list(search_space["chunk_overlap"].values),
    )

    # Run model pre-selection when the number of models exceeds the caps
    fm_values = search_space["foundation_model"].values
    em_values = search_space["embedding_model"].values

    if len(fm_values) > top_n_generation or len(em_values) > top_k_embedding:
        mps = ModelsPreSelector(
            benchmark_data=benchmark_data.get_random_sample(n_records=sample_size, random_seed=random_seed),
            documents=documents,
            foundation_models=search_space._search_space["foundation_model"].values,  # pylint: disable=protected-access
            embedding_models=search_space._search_space["embedding_model"].values,  # pylint: disable=protected-access
            max_threads=inference_max_threads,
        )
        mps.evaluate_patterns()
        selected = mps.select_models(
            n_embedding_models=top_k_embedding,
            n_foundation_models=top_n_generation,
        )
        selected_models = {
            "foundation_model": selected["foundation_models"],
            "embedding_model": selected["embedding_models"],
        }
    else:
        selected_models = {
            "foundation_model": list(fm_values),
            "embedding_model": list(em_values),
        }

    # Build verbose representation from valid (rule-filtered) combinations only
    valid_combinations = search_space.combinations
    if not valid_combinations:
        _logger.warning("No valid combinations remain after applying search space rules.")
    non_model_keys = [p.name for p in search_space.params if p.name not in ("foundation_model", "embedding_model")]
    verbose_repr: dict[str, Any] = {
        key: list(dict.fromkeys(combo[key] for combo in valid_combinations)) for key in non_model_keys
    }
    verbose_repr["foundation_model"] = [_serialize_model(m) for m in selected_models["foundation_model"]]
    verbose_repr["embedding_model"] = [_serialize_model(m) for m in selected_models["embedding_model"]]

    return SearchSpaceReport(
        search_space=verbose_repr,
        selected_models=selected_models,
    )

def run_rag_optimization(  # pylint: disable=too-many-locals,too-many-arguments,too-many-positional-arguments
    extracted_text_path: str | Path,
    test_data_path: str | Path,
    search_space_report_path: str | Path,
    output_dir: str | Path,
    ogx_client: OgxClient,
    vector_io_provider_id: str,
    test_data_key: str = "",
    input_data_key: str = "",
    optimization_settings: dict | None = None,
    inference_max_threads: int = 10,
    indexing_pipeline_params: dict | None = None,
    judge_enabled: bool = True,
) -> OptimizationResult:
    """Run a full AI4RAG optimization experiment and generate output artefacts.

    Orchestrates the end-to-end workflow: load documents, reconstruct the
    search space from a JSON report, run the experiment, then generate
    per-pattern outputs (``pattern.json``, notebooks, evaluation results).

    Parameters
    ----------
    extracted_text_path
        Path to a folder of DoclingDocument JSON files (or a single file).
    test_data_path
        Path to a benchmark JSON file with questions and expected answers.
    search_space_report_path
        Path to the JSON report produced by the search-space preparation step.
    output_dir
        Root directory where per-pattern output folders are written.
    ogx_client
        An authenticated :class:`OgxClient` instance.
    vector_io_provider_id
        Vector I/O provider identifier registered in OGX.
    test_data_key
        Object-storage key for the test data file, embedded into generated
        notebooks.
    input_data_key
        Object-storage key for the documents directory, embedded into
        generated notebooks.
    optimization_settings
        Optional dictionary with ``"metric"`` and/or
        ``"max_number_of_rag_patterns"`` overrides.
    inference_max_threads
        Maximum number of concurrent threads used when querying the
        RAG service during benchmark evaluation.  Lower values reduce
        per-request concurrency (useful when each request carries more
        retrieved context).  Defaults to ``10``.
    indexing_pipeline_params : dict | None, default=None
        Parameters required to enhance pattern.json with indexing pipeline
        settings.
    judge_enabled : bool, default=True
        Whether LLM as a Judge metrics should be calculated.

    Returns
    -------
    OptimizationResult
        Contains the list of pattern definitions, raw evaluations, and the
        total number of parameter combinations explored.

    Raises
    ------
    ValueError
        If ``test_data_key`` does not point to a JSON file,
        ``vector_io_provider_id`` is empty, or the optimization metric is
        not supported.
    TypeError
        If ``optimization_settings`` has invalid types.
    """
    # --- Input validation ---
    if not isinstance(test_data_key, str) or not test_data_key.strip() or not test_data_key.lower().endswith(".json"):
        raise ValueError("test_data_key must point to a JSON file.")

    if not isinstance(vector_io_provider_id, str) or not vector_io_provider_id.strip():
        raise ValueError("vector_io_provider_id must be a non-empty string.")
    vector_io_provider_id = vector_io_provider_id.strip()

    settings = _validate_optimization_settings(optimization_settings)
    optimization_metric = settings.get("metric") or DEFAULT_METRIC
    if optimization_metric not in SUPPORTED_OPTIMIZATION_METRICS:
        raise ValueError(
            f"Optimization metric {optimization_metric} is not supported. "
            f"Select one of {sorted(SUPPORTED_OPTIMIZATION_METRICS)}."
        )

    documents = load_docling_documents(extracted_text_path)
    benchmark_data = pd.read_json(Path(test_data_path))
    benchmark_data_obj = BenchmarkData(benchmark_data)

    # --- Reconstruct search space from report ---
    with open(search_space_report_path, "r", encoding="utf-8") as f:
        search_space_raw: dict[str, Any] = json.load(f)

    foundation_models: list[OGXFoundationModel] = []
    embedding_models: list[OGXEmbeddingModel] = []
    params: list[Parameter] = []

    for param_name, values in search_space_raw.items():
        if param_name == "foundation_model":
            values = [_deserialize_model(m, ogx_client) for m in values]
            foundation_models = values
        elif param_name == "embedding_model":
            values = [_deserialize_model(m, ogx_client) for m in values]
            embedding_models = values
        params.append(Parameter(param_name, "C", values=values))

    search_space = AI4RAGSearchSpace(params=params)

    if judge_enabled:
        # --- Select judge model and build evaluators ---
        judge_model = select_judge_model(
            generation_models=foundation_models,
            embedding_models=embedding_models,
            benchmark_data=benchmark_data_obj,
            documents=documents,
            max_threads=inference_max_threads,
        )
        _logger.info("Judge model selected: %s", judge_model.model_id)

        evaluators = [UnitxtEvaluator(), LLMaJEvaluator(model=judge_model)]
    else:
        evaluators = [UnitxtEvaluator()]

    # --- Configure experiment ---
    max_rag_patterns = settings.get("max_number_of_rag_patterns", DEFAULT_MAX_RAG_PATTERNS)
    if isinstance(max_rag_patterns, str):
        max_rag_patterns = int(max_rag_patterns.strip())
    optimizer_settings = GAMOptSettings(max_evals=max_rag_patterns)

    event_handler = KFPEventHandler()

    rag_exp = AI4RAGExperiment(
        client=ogx_client,
        event_handler=event_handler,
        optimizer_settings=optimizer_settings,
        search_space=search_space,
        benchmark_data=benchmark_data,
        vector_store_type="ogx",
        documents=documents,
        optimization_metric=optimization_metric,
        ogx_vector_io_provider_id=vector_io_provider_id,
        inference_max_threads=inference_max_threads,
        evaluators=evaluators,
    )

    # --- Run the optimization loop ---
    rag_exp.search()

    # --- Generate output artefacts ---
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    patterns = _generate_output_artifacts(
        patterns_raw=event_handler.patterns,
        output_dir=output_dir,
        input_data_key=input_data_key,
        test_data_key=test_data_key,
        indexing_pipeline_params=indexing_pipeline_params,
    )

    return OptimizationResult(
        patterns=patterns,
        evaluations=list(rag_exp.results.evaluations),
    )