Help us learn about your current experience with the documentation. Take the survey.

嵌入

嵌入是一种将数据以向量化格式表示的方法,使得查找相似文档变得简单高效。

目前,嵌入仅针对问题生成,这支持了以下功能:

架构

嵌入存储在 Elasticsearch 中,该系统也用于高级搜索

graph LR
  A[数据库记录] --> B[ActiveRecord 回调]
  B --> C[构建嵌入引用]
  C -->|添加到队列| N[队列]
  E[每分钟运行的 cron worker] <-->|从队列提取| N
  E --> G[反序列化引用]
  G --> H[生成嵌入]
  H <--> I[AI 网关]
  I <--> J[Vertex API]
  H --> K[使用嵌入更新文档]
  K --> L[Elasticsearch]

该过程由 Search::Elastic::ProcessEmbeddingBookkeepingService 驱动,它负责向 Redis 队列添加和提取任务。

添加到嵌入队列

以下过程描述以问题为例。

问题嵌入从内容 "issue with title '#{issue.title}' and description '#{issue.description}'" 生成。

使用 Search::Elastic::IssuesSearch 中定义的 ActiveRecord 回调,如果问题被创建或标题或描述被更新,并且该问题支持嵌入生成,则会将嵌入引用添加到嵌入队列中。

从嵌入队列提取

一个 Search::ElasticIndexEmbeddingBulkCronWorker cron 工作程序每分钟运行一次,执行以下操作:

graph LR
  A[cron] --> B{端点是否限流?}
  B -->|否| C[调度 16 个 worker]
  C ..->|每个 worker| D{端点是否限流?}
  D -->|否| E[从队列提取 19 个引用]
  E ..->|每个引用| F[增加端点计数]
  F --> G{端点是否限流?}
  G -->|否| H[调用 AI 网关生成嵌入]

因此,我们始终确保即使有 16 个并发进程同时生成嵌入,也不会超过每分钟 450 个嵌入的速率限制。

回填

使用高级搜索迁移来执行回填。它本质上将引用批量添加到队列中,然后由上述的 cron 工作程序处理。

添加新的嵌入类型

以下过程概述了生成嵌入并将其存储到 Elasticsearch 的步骤。

  1. 进行成本和资源计算,查看 Elasticsearch 集群能否处理嵌入生成,或者是否需要额外资源。
  2. 决定嵌入的存储位置。查看Elasticsearch 中的现有索引,如果没有合适的现有索引,则创建新索引
  3. 向索引添加嵌入字段:示例
  4. 更新内容的生成方式以适应新类型。
  5. 添加新的单元原语:这里这里
  6. 使用 Elastic::ApplicationVersionedSearch 访问回调并添加必要的检查,确定何时生成嵌入。示例参见 Search::Elastic::IssuesSearch
  7. 回填嵌入:示例

在本地添加工作项嵌入

先决条件

  1. 确保 Elasticsearch 正在运行

  2. 如果您已有 Elasticsearch 设置,请通过执行以下命令直到返回结果,确保 AddEmbeddingToWorkItems 迁移已完成:

    Elastic::MigrationWorker.new.perform
  3. 确保您可以在本地环境中运行 GitLab Duo 功能

  4. 确保在 rails 控制台中运行以下命令能输出嵌入(一个 768 维的向量)。如果没有,说明 AI 设置有问题。

    Gitlab::Llm::VertexAi::Embeddings::Text.new('text', user: nil, tracking_context: {}, unit_primitive: 'semantic_search_issue').execute

运行回填

要为项目的工作项回填工作项嵌入,在 rails 控制台中运行以下命令:

Gitlab::Duo::Developments::BackfillWorkItemEmbeddings.execute(project_id: project_id)

该任务将工作项添加到队列中并批量处理,将嵌入索引到 Elasticsearch 中。 它遵循每分钟 450 个嵌入的速率限制。如有任何问题,请在 Slack 中联系 #g_global_search

验证

如果以下命令返回 0,则该项目所有的工作项都已生成嵌入:

curl "http://localhost:9200/gitlab-development-work_items/_count" \
--header "Content-Type: application/json" \
--data '{"query": {"bool": {"filter": [{"term": {"project_id": PROJECT_ID}}], "must_not": [{"exists": {"field": "embedding_0"}}]}}}' | jq '.count'

PROJECT_ID 替换为您的项目 ID。