为大型表生成和管理自动向量嵌入

您可以为整个表列生成和管理向量嵌入,从而为大规模向量嵌入创建提供可扩缩的解决方案。此解决方案特别有助于对文本内容进行语义搜索和检索增强生成 (RAG),包括:

  • 为新表创建初始向量嵌入
  • 在大量数据导入后生成嵌入
  • 在数据发生重大变化后刷新嵌入
  • 以增量方式维护嵌入

了解自动向量嵌入

AlloyDB 中的自动向量嵌入提供了一种可扩缩的方式,可自动生成和维护数据的向量嵌入。您可以配置自动向量嵌入功能来处理此流程,而无需为每段新文本或更新后的文本手动生成嵌入。对于依赖最新嵌入进行语义搜索、检索增强生成 (RAG) 和其他 AI 赋能功能的应用,此功能特别有用。

借助自动向量嵌入,您可以执行以下操作:

  • 为整个表初始化嵌入:通过单个命令为表列中的所有现有数据生成嵌入。
  • 保持嵌入同步:在源数据发生变化时自动更新嵌入,确保 AI 应用始终使用最新信息。
  • 大规模生成嵌入:高效地为数百万行的大型表创建嵌入。
  • 通过为每个嵌入列调用管理函数,配置和管理同一表中多个列的嵌入。

此功能可抽象化向量嵌入创建和维护的复杂性,从而简化 AI 应用的开发和维护。

准备工作

在为大型表生成和管理向量嵌入之前,请执行以下操作:

  • 使用 psql 或 AlloyDB for PostgreSQL Studio 以 postgres 用户身份连接到您的数据库
  • 验证 google_ml_integration 扩展程序是否已安装
  • 验证 google_ml_integration.enable_model_support 标志是否已设置为 on
  • 验证 google_ml_integration 扩展程序的版本是否为 1.5.2 或更高版本,以及 google_ml_integration.enable_faster_embedding_generation 标志是否已设置为 on

    您可以使用以下命令查看扩展程序版本:

    SELECT extversion FROM pg_extension WHERE extname = 'google_ml_integration';
    

    如果您需要更新扩展程序,请使用 ALTER EXTENSION google_ml_integration UPDATE; 命令。

  • 您必须先将 AlloyDB 配置为与 Vertex AI 搭配使用,然后才能通过 AlloyDB 数据库生成嵌入。如需了解详情,请参阅将数据库与 Vertex AI 集成

  • 对于自动生成嵌入的管理和监控,用户默认拥有对 google_ml.embed_gen_progressgoogle_ml.embed_gen_settings 表的 Select 访问权限。

    如需允许用户管理自动嵌入生成,请在 google_ml.embed_gen_progressgoogle_ml.embed_gen_settings 表上授予 INSERTUPDATEDELETE 权限:

    GRANT INSERT, UPDATE, DELETE ON google_ml.embed_gen_progress TO 'USER_NAME';
    

    替换以下内容:

    • USER_NAME:为其授予权限的用户的名称。
  • 验证您使用的 PostgreSQL 客户端中是否已将 AUTOCOMMIT 设置为 ON

  • 确认您使用的嵌入模型有足够的配额用于自动向量嵌入。配额不足可能会导致自动嵌入操作缓慢或失败。例如,以下是 Vertex AI 嵌入模型的限制:

    • 文本嵌入限制:每个请求最多可包含 250 个输入文本(每个输入文本生成 1 个嵌入)和 2 万个token。系统仅使用每个输入文本中的前 2,048 个token来计算嵌入。
    • 每分钟请求数

      • base_model : text-embedding : 1500
      • base_model : gemini-embedding : 100000
    • Gemini 嵌入模型token限制:与其他主要受 RPM 配额限制的嵌入模型不同,Gemini 嵌入模型系列限制为每个项目每分钟 5,000,000 个token。

初始化表的嵌入

用于管理自动向量嵌入的函数位于 ai 架构中。此架构为 AlloyDB 中的最新 AI 功能提供了一个接口。

使用 ai.initialize_embeddings() SQL 函数为表的内容列生成嵌入。这是一个阻塞调用:

  • 如果函数返回成功,则表示向量嵌入创建完成。
  • 该函数会自动尝试从模型配额错误等暂时性问题中恢复。只有在这些恢复尝试失败时,系统才会向您返回失败消息。对于持续存在的问题(例如 batch_size 配置错误导致请求超出大小限制),或者如果操作是手动取消的,您必须手动重新发出调用。

此函数支持 Google 提供的模型(例如 Vertex AI 的 text-embedding-005),以及您注册的自定义模型。

在生成嵌入之前,请在表中创建列 content_embeddings。此列通常具有 vector(DIMENSION) 类型和 DEFAULT NULL 值。

ALTER TABLE user_reviews ADD COLUMN IF NOT EXISTS content_embeddings vector(768) DEFAULT NULL;

执行批量生成

默认情况下,AlloyDB 会使用批处理功能,通过单个请求为多个文本输入生成嵌入,从而提高效率。如果您未提供具体的批次大小,AlloyDB 会应用自动确定的默认值。

提示批次大小

ai.initialize_embeddings 中的 batch_size 参数可让您通过为直接支持的模型建议首选批次大小来指导 AlloyDB 的查询优化器。AlloyDB 可能会根据模型限制或配额动态减小此大小,但您的提示有助于影响查询执行计划。

CALL ai.initialize_embeddings(
    model_id => 'text-embedding-005',
    table_name => 'user_reviews',
    content_column => 'content',
    embedding_column => 'content_embeddings',
    batch_size => 50
);

使用支持批处理的自定义嵌入模型

如果您想使用支持批处理的自定义模型或外部支持的模型,请定义批处理转换函数,并在创建模型时将其指定为 model_batch_in_transform_fnmodel_batch_out_transform_fn。您还可以在 initialize_embeddings 调用中指定 batch_size。对于支持批处理的模型,我们建议您使用大于 1 的 batch_size 以获得更好的性能。

  1. 定义自定义模型的输入、输出和批处理转换函数。

    -- Scalar input transform functions
    CREATE OR REPLACE FUNCTION acme_text_input_transform(model_id TEXT, input TEXT) RETURNS JSON;
    CREATE OR REPLACE FUNCTION acme_text_output_transform(model_id TEXT, model_output JSON) RETURNS real[];
    CREATE OR REPLACE FUNCTION acme_generate_headers(model_id TEXT, input TEXT) RETURNS JSON;
    -- Batch input transform functions
    CREATE OR REPLACE FUNCTION acme_text_batch_input_transform(model_id TEXT, input TEXT[]) RETURNS JSON;
    CREATE OR REPLACE FUNCTION acme_text_batch_output_transform(model_id TEXT, model_output JSON) RETURNS real[][];
    
  2. 如需创建模型,请指定批处理转换函数。

    CALL
      ai.create_model(
        model_id => 'custom-embedding-model',
        model_request_url => 'https://acme.com/models/text/embeddings/v1',
        model_type => 'text_embedding',
        model_in_transform_fn => 'acme_text_input_transform',
        model_out_transform_fn => 'acme_text_output_transform',
        generate_headers_fn => 'acme_generate_headers',
        model_batch_in_transform_fn => 'acme_text_batch_input_transform',
        model_batch_out_transform_fn => 'acme_text_batch_output_transform'
      );
    
  3. 使用自定义模型生成向量嵌入。

    CALL
      ai.initialize_embeddings(
        model_id => 'custom-embedding-model',
        table_name => 'user_reviews',
        content_column => 'content',
        embedding_column => 'content_embeddings',
        batch_size => 10
    );
    

您还可以将自动嵌入功能与本身不支持批处理的自定义模型搭配使用。为此,您仍必须定义批处理转换函数 model_batch_in_transform_fnmodel_batch_out_transform_fn。对于非批处理模型,请定义这些函数以一次处理输入数组中的单个输入。为此模型调用 ai.initialize_embeddings 时,请将 batch_size 设置为 1

以增量方式刷新嵌入

刷新嵌入时,系统会根据输入内容列中的最新值重新生成嵌入。

为了让您能够控制一致性和性能,AlloyDB 支持多种增量嵌入刷新模式。您可以使用 ai.initialize_embeddings() 中的 incremental_refresh_mode 枚举参数选择模式。以下是可能模式的列表:

  • transactional:在更新内容列的事务中刷新嵌入。此过程通常使用类似于数据库触发器的机制,在内容列更新时自动生成嵌入,但可能会产生开销并减慢更新操作的速度。引入的开销是为了保持事务语义并确保嵌入与内容同步而做出的权衡。此模式依赖于模型的标量转换函数,因此您必须在创建模型时定义 model_in_transform_fnmodel_out_transform_fn。如需使用 transactional 模式,您必须拥有表的“所有者”角色

    CALL
      ai.initialize_embeddings(
        model_id => 'text-embedding-005',
        table_name => 'user_reviews',
        content_column => 'content',
        embedding_column => 'content_embeddings',
        batch_size => 10,
        incremental_refresh_mode => 'transactional'
    );
    
  • none:此为默认模式。在此模式下,AlloyDB 不会自动更新嵌入。系统不会跟踪增量更改,因此调用 ai.refresh_embeddings() 函数会为整个表重新生成嵌入。此模式提供完全控制权。

刷新表的所有嵌入

成功为表运行 ai.initialize_embeddings() 后,您可以使用 ai.refresh_embeddings() 函数重新生成嵌入。您可以使用刷新操作来更新在初始 initialize_embeddings 调用期间并发修改的行的嵌入,或执行定期完整刷新。

刷新函数会重复使用初始调用中的设置,因此您只需指定表和嵌入列。您还可以提供可选的 batch_size 来替换默认值。

CALL ai.refresh_embeddings(
    table_name => 'user_reviews',
    embedding_column => 'content_embeddings',
    batch_size => 50  -- Optional override
);

在创建向量嵌入期间处理表数据

虽然 ai.initialize_embeddings() 对于运行它的会话来说是一个阻塞调用,但其他连接可以继续使用该表。自动向量嵌入流程会分批更新行,这涉及标准的行级锁定。这意味着,如果并发 DML 操作(例如 UPDATEDELETE)尝试修改嵌入进程正在更新的相同行,则只会短暂受阻。来自其他连接的查询不会被阻止。请注意,initialize_embeddings 会跳过同时修改的行。如果 incremental_refresh_modenone,则在调用后续刷新之前,不会更新这些已修改行的嵌入。

-- connection1 (starts embedding generation)
SELECT
  ai.initialize_embeddings(
    model_id => 'text-embedding-005',
    table_name => 'user_reviews',
    content_column => 'content',
    embedding_column => 'content_embeddings'
  );

-- connection2 (performs DMLs/queries without blocking)
INSERT INTO user_reviews(id, review_time, is_edited, content)
VALUES (48290, now(), false, 'I really liked the product functionality, but wish it came in orange color');

UPDATE user_reviews
SET is_edited = TRUE, content = 'Changing to 5 star. My issue is resolved by the support'
WHERE id = 700;

如果您使用 pg_cancel 取消 initialize_embeddings,或者 initialize_embeddings 因内部错误而失败,则会返回失败状态。任何已成功创建的向量嵌入都不会回滚。如需从失败中恢复并完成向量嵌入创建,您必须先使用 ai.drop_embedding_config() 函数清理配置,然后重新发出 ai.initialize_embeddings() 调用。

如需刷新并发修改的行的嵌入,请在 ai.initialize_embeddings 调用完成后调用 ai.refresh_embeddings。此刷新调用会重新生成整个表的嵌入。

删除自动向量嵌入设置

如果您需要移除特定表和嵌入列组合的自动向量嵌入配置,请使用 ai.drop_embedding_config() 函数。此函数可用于清理或在您重新配置列的嵌入管理时使用。

CALL
  ai.drop_embedding_config(
    table_name => 'user_reviews',
    embedding_column => 'content_embeddings');

自动生成嵌入的示例

本部分提供了使用已注册模型端点自动生成嵌入的示例。

OpenAI 嵌入模型

如需使用 OpenAI 提供的已注册的 text-embedding-3-small 模型端点生成嵌入,请运行以下语句:

CALL ai.initialize_embeddings(
    model_id => 'text-embedding-3-small',
    table_name => 'user_reviews',
    chunk_column => 'content',
    embedding_column => 'content_embeddings'
);

自定义嵌入模型

对于您自己的模型或外部支持的模型,您必须定义输入和输出转换函数,并使用 ai.create_model 注册这些函数。如果您计划使用自动嵌入功能,则必须同时指定标量转换函数(例如 acme_text_input_transformacme_text_output_transform)和批处理转换函数(例如 acme_text_batch_input_transformacme_text_batch_output_transform)。

后续步骤