Code Embeddings: add optional root_namespace_id for request header
What does this MR do and why?
Problem and Solution overview
For SaaS (GitLab.com), AIGW requests require the x-gitlab-root-namespace-id header for billing and compliance tracking. However, the semantic search indexing process does not set this header when invoking embeddings requests. This is because indexing is an asynchronous and batched process without a user and is namespace-agnostic.
In order to resolve this, we need to get the root_namespace_id from the projects being indexed, then set that root_namespace_id in the header. The implementation involves several steps:
- Add a preprocessor to set the
root_namespace_idof each reference based on the project_id (routing) - !246906 (merged) - Update the
embeddingspreprocessor to group embeddings generation byroot_namespace_id- TBA - Update the
Gitlab::Llm::Embeddings::*classes to accept an optionalroot_namespace_id, to be used for the AIGW request header - This MR
For details, see proposal.
This MR
- updated the
Gitlab::Llm::Embeddings::Clientclass to accept an optionroot_namespace_id. If this is given, it will be set as thex-gitlab-root-namespace-idheader. - updated the
Gitlab::Llm::Embeddings::CodeEmbeddingsto accept an optionalroot_namespace_id, which is then passed on to theClientclass- additional minor refactor: clarify the optionality of the
userparam (this is already set asnilby the indexing process) - this is the class invoked by the Semantic Search indexing pipeline to generate embeddings
- additional minor refactor: clarify the optionality of the
Consideration: passing root_namespace_id vs root_namespace
An alternative solution is to pass around a root_namespace record instead of a root_namespace_id. This will allow the Gitlab::Llm::Embeddings::Client to be able to set the organization_id header. The solution implementation passes the root_namespace_id because:
organization_idis not required information for billing or compliance- passing around a
root_namespaceActiveRecord object requires more memory; while GitLab may be able to handle this without issues, I would prefer we choose the solution without additional memory overhead since the additional information is not required
References
Screenshots or screen recordings
N/A
How to set up and validate locally
Validate root_namespace_id header
-
Run the following on the rails console:
model_definition = ::Gitlab::Llm::Embeddings::ModelDefinition.for_gitlab_provided_code_embeddings( identifier: 'text_embedding_005_vertex' ) namespace = Group.find_by(path: 'gitlab-duo') Gitlab::Llm::Embeddings::CodeEmbeddings.new( ["test content"], model_definition: model_definition, root_namespace_id: namespace.id, ).execute -
Verify that the header
x-gitlab-root-namespace-idheader is set in the request to AIGW# on the GDK directory tail -f log/gitlab-ai-gateway/gateway_debug.log # watch out for the `/v1/embeddings/code_embeddings/index` request logs 2026-07-23T09:30:16.221169Z [info ] 172.16.123.1:58284 - "POST /v1/embeddings/code_embeddings/index HTTP/1.1" 200 [api.access] <other fields> gitlab_organization_id=None gitlab_root_namespace_id=1000000 <other fields>
Additional validation for the indexing pipeline - verify no regression errors
Test the ::Ai::ActiveContext::Collections::Code.current_indexing_embedding_model specifically:
# with user not set
Ai::ActiveContext::Collections::Code.current_indexing_embedding_model.generate_embeddings("test input")
# with user set
Ai::ActiveContext::Collections::Code.current_indexing_embedding_model.generate_embeddings("test input", user: User.first)Note: we are not yet validating the the indexing process sending the root_namespace_id, since that will be done as a next step.
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 #606107