Experiment API¶

def run_pre_selection(
    self,
    foundation_models: list[BaseFoundationModel],
    embedding_models: list[BaseEmbeddingModel],
    n_records: int = 5,
    random_seed: int = 17,
) -> dict[str, list[BaseEmbeddingModel | BaseFoundationModel]]:
    """
    Run models pre-selection using ModelsPreSelector and sample
    of the data.

    Parameters
    ----------
    embedding_models : list[BaseEmbeddingModel]
        Embedding models to be considered during pre-selection process.

    foundation_models : list[BaseFoundationModel]
        Foundation models to be evaluated during pre-selection process.

    n_records : int, default=5
        Amount of records that should be used during models pre-selection.

    random_seed : int, default=17
        Random seed value used for sampling benchmark data records.

    Returns
    -------
    dict[str, list[BaseFoundationModel | EmbeddingModel]]
        Best embedding models and foundation models found in pre-selection.
    """
    _log_start_mps = (
        f"Starting foundation models pre-selection with following "
        f"foundation models: {[str(fm) for fm in foundation_models]} "
        f"and following embedding models: {[str(em) for em in embedding_models]}."
    )
    logger.info(_log_start_mps)
    self.event_handler.on_status_change(
        level=LogLevel.INFO,
        message=_log_start_mps,
        step=ExperimentStep.MODEL_SELECTION,
    )

    from ai4rag.core.experiment.mps import ModelsPreSelector

    mps = ModelsPreSelector(
        benchmark_data=self.benchmark_data.get_random_sample(n_records=n_records, random_seed=random_seed),
        documents=self.documents.copy(),
        foundation_models=foundation_models,
        embedding_models=embedding_models,
        metric=Metrics.OVERALL_SCORE,
    )
    mps.evaluate_patterns()

    selected_models = mps.select_models(
        n_embedding_models=self.n_mps_embedding_models, n_foundation_models=self.n_mps_foundation_models
    )

    logger.info(
        "Models pre-selection has been finished. Selected foundation models: %s and selected embedding models: %s.",
        [str(model) for model in selected_models["foundation_models"]],
        [str(model) for model in selected_models["embedding_models"]],
    )

    return selected_models

def run_single_evaluation(self, rag_params: RAGParamsType, publish_pattern: bool = True) -> float:
    """
    Evaluate a single RAG configuration and return its score using provided documents.

    Parameters
    ----------
    rag_params : RAGParamsType
        A dictionary containing rag parameters as keys and their values.

    publish_pattern : bool, default=True
        Whether to send the evaluated pattern to the event handler immediately.
        GAM optimization suppresses this only for warm-start candidates;
        GAM candidates are published as they complete.

    Returns
    -------
    float
        A single evaluation score obtained by the executed rag pattern.
    """
    start_time = time.time()

    chunking_params = get_chunking_params(rag_params)
    chunking_params["include_metadata"] = chunking_params.get(AI4RAGParamNames.CHUNKING_METHOD) == "hybrid"

    retrieval_params = get_retrieval_params(rag_params)

    foundation_model = rag_params.get(AI4RAGParamNames.FOUNDATION_MODEL)
    embedding_model = rag_params.get(AI4RAGParamNames.EMBEDDING_MODEL)

    embedding_params_dict = (
        asdict(embedding_model.params) if is_dataclass(embedding_model.params) else embedding_model.params
    )
    indexing_params = {
        "chunking": chunking_params,
        "embedding": {
            "model_id": embedding_model.model_id,
            "embedding_params": embedding_params_dict,
        },
    }

    logger.info("Using indexing params: %s", indexing_params)

    retrieval_method = retrieval_params[AI4RAGParamNames.RETRIEVAL_METHOD]
    number_of_chunks = retrieval_params[AI4RAGParamNames.NUMBER_OF_CHUNKS]

    search_mode = retrieval_params.get(AI4RAGParamNames.SEARCH_MODE, "vector")

    context_template_text = foundation_model.context_template_text
    system_message_text = foundation_model.system_message_text
    user_message_text = foundation_model.user_message_text

    rag_params = {
        "retrieval": retrieval_params,
        "generation": {
            "model_id": foundation_model.model_id,
            "temperature": foundation_model.params.temperature,
            "max_completion_tokens": foundation_model.params.max_completion_tokens,
            "context_template_text": context_template_text,
            "user_message_text": user_message_text,
            "system_message_text": system_message_text,
            "language": foundation_model.language.to_dict(),
        },
    }

    logger.info("Using retrieval and generation params: %s", rag_params)

    result_score = self.results.evaluation_explored_or_cached(
        indexing_params=indexing_params, rag_params=rag_params
    )
    if result_score is not None:
        return result_score

    pattern_name = self._create_pattern_name()
    logger.info("Using name '%s' for the currently evaluated pattern.", pattern_name)

    collection_name = self._get_reusable_collection_name(indexing_params=indexing_params)

    vector_store_config = self.vector_store_config
    if isinstance(vector_store_config, PGVectorConfig):
        # Size the connection pool to this run's actual query concurrency so a
        # fully concurrent query_rag() call never queues for a slot (see
        # PGVectorConfig.pool_max_size). Never shrink below a user-set ceiling:
        # a caller who deliberately raised pool_max_size (e.g. to share the store
        # with other concurrent work) must keep that headroom, so take the larger
        # of the configured size and this run's inference concurrency.
        pool_max_size = max(vector_store_config.pool_max_size, self.inference_max_threads)
        if pool_max_size != vector_store_config.pool_max_size:
            logger.info(
                "Raising PGVector pool_max_size from %d to %d to match inference_max_threads (%d).",
                vector_store_config.pool_max_size,
                pool_max_size,
                self.inference_max_threads,
            )
        else:
            logger.info(
                "Keeping configured PGVector pool_max_size %d (>= inference_max_threads %d).",
                vector_store_config.pool_max_size,
                self.inference_max_threads,
            )
        vector_store_config = replace(vector_store_config, pool_max_size=pool_max_size)

    try:
        vector_store = get_vector_store(
            embedding_model=embedding_model,
            collection_name=collection_name,
            config=vector_store_config,
        )
    except Exception as exc:
        raise VectorStoreInitializationError(
            exc,
            embedding_model_id=embedding_model.model_id,
            vector_store_provider_id=self.vector_store_config.provider,
        ) from exc

    collection_name = vector_store.collection_name

    with vector_store:
        if not self._collection_exists(collection_name=collection_name):
            chunking_method = chunking_params.get(AI4RAGParamNames.CHUNKING_METHOD)
            chunk_size = chunking_params.get(AI4RAGParamNames.CHUNK_SIZE)
            chunk_overlap = chunking_params.get(AI4RAGParamNames.CHUNK_OVERLAP)

            if chunking_method == "hybrid":
                chunker = DoclingChunker(max_tokens=chunk_size)
            else:
                chunker = LangChainChunker(
                    method=chunking_method, chunk_size=chunk_size, chunk_overlap=chunk_overlap
                )
            chunked_documents = chunker.split_documents(self.documents)

            if self.event_handler:
                self.event_handler.on_status_change(
                    level=LogLevel.INFO,
                    message=(
                        f"Chunking documents using the {chunking_method} method, chunk_size: {chunk_size} "
                        f"and chunk_overlap: {chunk_overlap}."
                    ),
                    step=ExperimentStep.CHUNKING,
                )

            self.event_handler.on_status_change(
                level=LogLevel.INFO,
                message=(
                    f"Embedding chunks using the {embedding_model.model_id} model. "
                    f"Building index: {collection_name}."
                ),
                step=ExperimentStep.EMBEDDING,
            )

            try:
                vector_store.add_documents(chunked_documents)
            except Exception as exc:
                raise IndexingError(exc, collection_name, embedding_model.model_id) from exc

        else:
            self.event_handler.on_status_change(
                level=LogLevel.INFO,
                message=f"Using index {collection_name}.",
                step=ExperimentStep.EMBEDDING,
            )

        logger.info("Using retriever with parameters: %s", retrieval_params)

        retriever = Retriever(
            vector_store=vector_store,
            number_of_chunks=number_of_chunks,
            method=retrieval_method,
            search_mode=search_mode,
            ranker_strategy=retrieval_params.get(AI4RAGParamNames.RANKER_STRATEGY),
            ranker_k=retrieval_params.get(AI4RAGParamNames.RANKER_K),
            ranker_alpha=retrieval_params.get(AI4RAGParamNames.RANKER_ALPHA),
        )

        rag_pattern = SimpleRAG(
            foundation_model=foundation_model,
            retriever=retriever,
        )

        _rag_log = (
            f"Retrieval and generation using collection: '{collection_name}' and "
            f"foundation model: '{foundation_model.model_id}'."
        )
        logger.info(_rag_log)
        self.event_handler.on_status_change(
            level=LogLevel.INFO,
            message=_rag_log,
            step=ExperimentStep.GENERATION,
        )

        inference_response = query_rag(
            rag=rag_pattern, questions=list(self.benchmark_data.questions), max_threads=self.inference_max_threads
        )

    result_scores, evaluation_data = self._evaluate_response(
        inference_response=inference_response,
        pattern_name=pattern_name,
    )

    stop_time = time.time()
    execution_time = stop_time - start_time

    final_score = self._resolve_optimization_score(result_scores, pattern_name)

    logger.info("Calculated optimization score for '%s': %s", pattern_name, final_score)

    evaluation_result = EvaluationResult(
        pattern_name=pattern_name,
        collection=collection_name,
        indexing_params=indexing_params,
        rag_params=rag_params,
        scores=result_scores,
        execution_time=execution_time,
        final_score=final_score,
    )

    evaluation_results_json = self.results.create_evaluation_results_json(
        evaluation_data=evaluation_data, evaluation_result=evaluation_result
    )

    logger.info(
        "Evaluation scores: %s",
        {el.get("question"): el.get("metrics") for el in evaluation_results_json if isinstance(el, dict)},
    )

    iteration = len(self.results) + (len(self.known_observations) if self.known_observations else 0)
    if publish_pattern:
        try:
            self._stream_finished_pattern(
                evaluation_result=evaluation_result,
                evaluation_results_json=evaluation_results_json,
                iteration=iteration,
            )
        except Exception as exc:
            raise AssetSaveError(exc) from exc

    self.results.add_evaluation(
        evaluation_data=evaluation_data,
        evaluation_result=evaluation_result,
    )

    return final_score

def search(self, **kwargs) -> None:
    """
    Prepare and execute experiment to find the best RAG parameters.

    Result of the search() can be reviewed via self.results as this object
    stores results of each evaluation or via self.event_handler with custom
    implementation.
    """

    logger.info("Starting RAG optimization process...")

    optimizer_class: type[BaseOptimizer] = kwargs.get("optimizer", GAMOptimizer)
    is_gam_optimizer = issubclass(optimizer_class, GAMOptimizer)

    def objective_function(space: RAGParamsType) -> float | None:
        """Function passed to the optimizer."""
        self._optimization_phase = getattr(optimizer, "current_phase", None)
        try:
            return self.run_single_evaluation(
                space,
                publish_pattern=not is_gam_optimizer or self._optimization_phase != "warm_start",
            )
        except AI4RAGError as err:
            msg = self._exception_handler.handle_exception(err)
            raise FailedIterationError(msg) from err
        finally:
            self._optimization_phase = None

    # MPS - models pre-selection based on sample evaluation.
    # Run if there are more than 3 foundation models or more than 2 embedding models.
    foundation_models = list(self.search_space[AI4RAGParamNames.FOUNDATION_MODEL].values)
    embedding_models = list(self.search_space[AI4RAGParamNames.EMBEDDING_MODEL].values)

    if (
        len(embedding_models) > self.n_mps_embedding_models or len(foundation_models) > self.n_mps_foundation_models
    ) and not kwargs.get("skip_mps", False):
        selected_models = self.run_pre_selection(
            foundation_models=foundation_models, embedding_models=embedding_models
        )
        self.search_space[AI4RAGParamNames.FOUNDATION_MODEL] = Parameter(
            name=AI4RAGParamNames.FOUNDATION_MODEL, param_type="C", values=selected_models["foundation_models"]
        )
        self.search_space[AI4RAGParamNames.EMBEDDING_MODEL] = Parameter(
            name=AI4RAGParamNames.EMBEDDING_MODEL, param_type="C", values=selected_models["embedding_models"]
        )

    optimizer_kwargs = {}
    if self.known_observations is not None and is_gam_optimizer:
        optimizer_kwargs["known_observations"] = self.known_observations

    optimizer = optimizer_class(
        objective_function=objective_function,
        search_space=self.search_space,
        settings=self.optimizer_settings,
        **optimizer_kwargs,
    )
    logger.info(
        "Using optimizer: %s with optimizer settings: %s",
        optimizer_class.__name__,
        self.optimizer_settings.to_dict(),
    )

    try:
        _ = optimizer.search()
    except OptimizationError as err:
        final_error_msg = self._exception_handler.get_final_error_msg()
        raise RAGExperimentError(final_error_msg) from err

    if is_gam_optimizer:
        self._publish_best_warm_start_pattern()

    self.event_handler.on_status_change(
        level=LogLevel.INFO,
        message="Experiment optimization process finished.",
    )