嵌入
嵌入是一种将数据以向量化格式表示的方法,使得查找相似文档变得简单高效。
目前,嵌入仅针对问题生成,这支持了以下功能:
架构
嵌入存储在 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 的步骤。
- 进行成本和资源计算,查看 Elasticsearch 集群能否处理嵌入生成,或者是否需要额外资源。
- 决定嵌入的存储位置。查看Elasticsearch 中的现有索引,如果没有合适的现有索引,则创建新索引。
- 向索引添加嵌入字段:示例。
- 更新内容的生成方式以适应新类型。
- 添加新的单元原语:这里和这里。
- 使用
Elastic::ApplicationVersionedSearch访问回调并添加必要的检查,确定何时生成嵌入。示例参见Search::Elastic::IssuesSearch。 - 回填嵌入:示例。
在本地添加工作项嵌入
先决条件
-
如果您已有 Elasticsearch 设置,请通过执行以下命令直到返回结果,确保
AddEmbeddingToWorkItems迁移已完成:Elastic::MigrationWorker.new.perform -
确保您可以在本地环境中运行 GitLab Duo 功能。
-
确保在 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。