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

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

    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")
    if search_mode != "vector" and self.vector_store_type == "chroma":
        raise RAGExperimentError(
            f"Search mode '{search_mode}' is not supported with chroma vector store. "
            "Only 'vector' mode is supported for chroma."
        )

    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)

    reuse_collection_name = self._get_reusable_collection_name(indexing_params=indexing_params)

    try:
        vector_store = get_vector_store(
            vs_type=self.vector_store_type,
            embedding_model=embedding_model,
            reuse_collection_name=reuse_collection_name,
            client=self.client,
            ogx_vector_io_provider_id=self.ogx_vector_io_provider_id,
        )
    except Exception as exc:
        raise VectorStoreInitializationError(
            exc,
            embedding_model_id=embedding_model.model_id,
            vector_store_provider_id=self.ogx_vector_io_provider_id or "local_chroma",
        ) from exc

    collection_name = vector_store.collection_name

    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 = next(
        (r["scores"]["mean"] for r in result_scores["metrics"] if r["name"] == self.optimization_metric.name),
        None,
    )
    if final_score is None:
        raise RAGExperimentError(
            f"Optimization metric '{self.optimization_metric.name}' not found in evaluation results. "
            f"Available: {[m['name'] for m in result_scores['metrics']]}."
        )

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

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

    try:
        self._stream_finished_pattern(
            evaluation_result=evaluation_result,
            evaluation_results_json=evaluation_results_json,
        )
    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...")

    def objective_function(space: RAGParamsType) -> float | None:
        """Function passed to the optimizer."""
        try:
            return self.run_single_evaluation(space)
        except AI4RAGError as err:
            msg = self._exception_handler.handle_exception(err)
            raise FailedIterationError(msg) from err

    # 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_class: type[BaseOptimizer] = kwargs.get("optimizer", GAMOptimizer)

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

    # In the search kwargs user may pass different optimizer class for testing purposes
    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

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