Semantic Search: persist batch_size for embeddings indexing

What does this MR do and why?

When configuring embedding models in Self-Managed instances, we want the users to be able to specify the indexing batch size.

The indexing batch size refers to the number of individual inputs that is sent to an embedding model's bulk request API. This value depends on the the embedding model's Bulk API limits, and the configured indexing chunk size. On SaaS, we have calculated a batch size of 30 for the text_embedding_005_vertex. However, Self-Managed users may need to configure their own batch size, especially if they use their own Self-hosted models with different Bulk API limits.

In this MR

  • add a embeddings_request_batch_size field to the model_metadata schema
  • update Ai::ActiveContext::Embedding.embeddings_request_batch_size to consider the model_metadata[:embeddings_request_batch_size]
  • update affected callers and tests

Next Steps

  • Refactor: update the "Semantic search embeddings UI" to split sections into partials - !246671 (merged)
  • Allow administrators to update the indexing batch size through the Admin UI - TBA, see PoC MR

References

Screenshots or screen recordings

TBA

How to set up and validate locally

As we are simply introducing a new field in a jsonb column and updating a simple method (Ai::ActiveContext::Embedding.embeddings_request_batch_size), the unit tests should cover all required validations.

However, this can also be tested locally by:

  1. Test for embeddings_request_batch_size specified:

    # set a `current_indexing_embedding_model` with batch size
    Ai::ActiveContext::Collections::Code.collection_record.update_metadata!(
      current_indexing_embedding_model: {
        field: 'embeddings_v1',
        model_type: 'gitlab_managed',
        model_ref: 'text_embedding_005_vertex',
        embeddings_request_batch_size: 15
      }
    )
    
    # check the params that would be sent to the embeddings request
    Ai::ActiveContext::Collections::Code.current_indexing_embedding_model.llm_params[:batch_size]
    => 15
    
    # if you have semantic search setup, you can further test the actual embeddings generation
    Ai::ActiveContext::Collections::Code.current_indexing_embedding_model.generate_embeddings("test")
  2. Test for embeddings_request_batch_size not specified

    # set a `current_indexing_embedding_model` without batch size
    Ai::ActiveContext::Collections::Code.collection_record.update_metadata!(
      current_indexing_embedding_model: {
        field: 'embeddings_v1',
        model_type: 'gitlab_managed',
        model_ref: 'text_embedding_005_vertex',
      }
    )
    
    # check the params that would be sent to the embeddings request
    Ai::ActiveContext::Collections::Code.current_indexing_embedding_model.llm_params[:batch_size]
    => 30 # default batch size for "text_embedding_005_vertex" model

MR acceptance checklist

Evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.

Related to #600634 (closed)

Edited by Pam Artiaga

Merge request reports

Loading
Loading