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_id of each reference based on the project_id (routing) - !246906 (merged)
  • Update the embeddings preprocessor to group embeddings generation by root_namespace_id - TBA
  • Update the Gitlab::Llm::Embeddings::* classes to accept an optional root_namespace_id, to be used for the AIGW request header - This MR

For details, see proposal.

This MR

  • updated the Gitlab::Llm::Embeddings::Client class to accept an option root_namespace_id. If this is given, it will be set as the x-gitlab-root-namespace-id header.
  • updated the Gitlab::Llm::Embeddings::CodeEmbeddings to accept an optional root_namespace_id, which is then passed on to the Client class
    • additional minor refactor: clarify the optionality of the user param (this is already set as nil by the indexing process)
    • this is the class invoked by the Semantic Search indexing pipeline to generate embeddings

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_id is not required information for billing or compliance
  • passing around a root_namespace ActiveRecord 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

  1. 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
  2. Verify that the header x-gitlab-root-namespace-id header 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

Edited by Pam Artiaga

Merge request reports

Loading