Skip to content

Client for generating vertex embedding

Madelein van Niekerk requested to merge 451431-vertex-embedding-client into master

What does this MR do and why?

Adds a Vertex embedding client which can be called as:

Gitlab::Llm::VertexAi::Embeddings::Text.new(text, user: user, tracking_context: tracking_context).execute

This is done by moving the logic from https://gitlab.com/gitlab-org/gitlab/-/blob/master/ee/app/workers/llm/embedding/gitlab_documentation/set_embeddings_on_the_record_worker.rb#L38-45 so that it can be reused.

groupglobal search will use this to generate embeddings for issues and store it in Elasticsearch.

groupai framework are planning to add the vertex embedding API to the AI Gateway but until that is done, we will continue calling the API directly.

MR acceptance checklist

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

How to set up and validate locally

  1. Ensure you have access to vertex locally either by creating a GCP project or following the latest way to access vertex
  2. Ensure the following returns embeddings of 768 dimensions: Gitlab::Llm::VertexAi::Embeddings::Text.new('some text', user: nil, tracking_context: {}).execute
  3. Ensure the following returns nil (when the text is blank): Gitlab::Llm::VertexAi::Embeddings::Text.new('', user: nil, tracking_context: {}).execute

Related to #451431 (closed)

Edited by Madelein van Niekerk

Merge request reports