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_sizefield to themodel_metadataschema - update
Ai::ActiveContext::Embedding.embeddings_request_batch_sizeto consider themodel_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:
-
Test for
embeddings_request_batch_sizespecified:# 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") -
Test for
embeddings_request_batch_sizenot 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)