在 AlloyDB Omni 中注册和调用远程 AI 模型

如需使用模型调用预测或生成嵌入,请向模型端点管理注册模型端点。

在向模型端点管理服务注册模型端点之前,如果您的模型端点需要身份验证,您必须启用 google_ml_integration 扩展程序并根据模型提供程序设置身份验证。

确保您使用 postgres 默认用户名访问数据库。

启用扩展程序

您必须先添加并启用 google_ml_integration 扩展程序,然后才能开始使用关联的函数。若要进行模型端点管理,您需要安装 google_ml_integration 扩展程序。

  1. 验证实例的 google_ml_integration.enable_model_support 数据库标志是否设置为 on。如需详细了解如何设置数据库标志,请参阅配置数据库标志

  2. 可选:如果已安装 google_ml_integration 扩展程序,请对其进行更改以更新到最新版本:

        ALTER EXTENSION google_ml_integration UPDATE;
    
  3. 使用 psql 添加 google_ml_integration 扩展程序:

      CREATE EXTENSION IF NOT EXISTS google_ml_integration;
    
  4. 可选:向非超级 PostgreSQL 用户授予管理模型元数据的权限:

      GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA google_ml TO NON_SUPER_USER;
    

    NON_SUPER_USER 替换为非超级 PostgreSQL 用户名。

  5. 在数据库上启用模型端点管理:

      ALTER SYSTEM SET google_ml_integration.enable_model_support=on;
      SELECT pg_reload_conf();
    
  6. 确保已启用出站连接。

设置身份验证

以下部分介绍了如何在注册模型端点之前设置身份验证。

设置 Vertex AI 身份验证

如需使用 Google Vertex AI 模型端点,您必须向安装 AlloyDB Omni 时使用的服务账号添加 Vertex AI 权限。如需了解详情,请参阅配置 AlloyDB Omni 安装以查询基于云的模型

使用 Secret Manager 设置身份验证

如果您的模型端点不通过 Secret Manager 处理身份验证(例如,如果您的模型端点使用 HTTP 标头传递身份验证信息或根本不使用身份验证),则此步骤是可选的。

如需创建和使用 API 密钥或 Bearer 令牌,请完成以下步骤:

  1. 在 Secret Manager 中创建 Secret。如需了解详情,请参阅创建 Secret 并访问 Secret 版本

    该密钥路径用于 google_ml.create_sm_secret() SQL 函数。

  2. 向 AlloyDB 集群授予访问该 Secret 的权限。

      gcloud secrets add-iam-policy-binding 'SECRET_NAME' \
          --member="serviceAccount:SERVICE_ACCOUNT_ID" \
          --role="roles/secretmanager.secretAccessor"
    

    替换以下内容:

    • SECRET_NAME:Secret Manager 中的 Secret 名称。
    • SERVICE_ACCOUNT_ID:您在上一步中创建的服务账号的 ID。确保此账号与您在安装 AlloyDB Omni 时使用的账号相同。其中包括完整的 PROJECT_ID.iam.gserviceaccount.com 后缀。例如:my-service@my-project.iam.gserviceaccount.com

      您还可以在项目级别向服务账号授予此角色。如需了解详情,请参阅添加 Identity and Access Management 政策绑定

使用标头设置身份验证

以下示例展示了如何使用一个函数设置身份验证,该函数会返回一个 JSON 对象,其中包含向嵌入模型发出请求所需的标头。

  CREATE OR REPLACE FUNCTION HEADER_GEN_FUNCTION(
    model_id VARCHAR(100),
    input_text TEXT
  )
  RETURNS JSON
  LANGUAGE plpgsql
  AS $$
  #variable_conflict use_variable
  DECLARE
    api_key VARCHAR(255) := 'API_KEY';
    header_json JSON;
  BEGIN
    header_json := json_build_object(
      'Content-Type', 'application/json',
      'Authorization', 'Bearer ' || api_key
    );
    RETURN header_json;
  END;
  $$;

替换以下内容:

  • HEADER_GEN_FUNCTION:您可以在注册模型时使用的头文件生成函数的名称。
  • API_KEY:模型提供方的 API 密钥。

文本嵌入模型

本部分介绍了如何使用模型端点管理功能注册模型端点。

模型端点管理支持将某些文本嵌入和通用 Vertex AI 模型作为预注册的模型端点。您可以根据模型类型直接使用模型 ID 生成嵌入或调用预测。如需详细了解受支持的预注册模型,请参阅预注册的 Vertex AI 模型

例如,如需调用预注册的 textembedding-gecko 模型,您可以使用嵌入函数直接调用该模型:

SELECT
      google_ml.embedding(
        model_id => 'textembedding-gecko',
        content => 'AlloyDB is a managed, cloud-hosted SQL database service');

同样,如需调用预先注册的 gemini-1.5-pro:generateContent 模型,您可以使用预测函数直接调用该模型:

 SELECT
        json_array_elements(
        google_ml.predict_row(
            model_id => 'gemini-1.5-pro:generateContent',
            request_body => '{
        "contents": [
            {
                "role": "user",
                "parts": [
                    {
                        "text": "For TPCH database schema as mentioned here https://www.tpc.org/TPC_Documents_Current_Versions/pdf/TPC-H_v3.0.1.pdf , generate a SQL query to find all supplier names which are located in the India nation."
                    }
                ]
            }
        ]
        }'))-> 'candidates' -> 0 -> 'content' -> 'parts' -> 0 -> 'text'; 

如需生成嵌入,请参阅如何为预注册的模型端点生成嵌入。如需调用预测,请参阅如何为预注册的模型端点调用预测

内置支持的文本嵌入模型

模型端点管理功能为 Vertex AI 和 OpenAI 的部分模型提供了内置支持。如需查看内置支持的模型列表,请参阅内置支持的模型

对于内置支持的模型,您可以将限定名称设置为模型限定名称,并指定请求网址。模型端点管理会自动识别模型并设置默认的转换函数。

Vertex AI 嵌入模型

以下步骤展示了如何注册具有内置支持的 Vertex AI 模型。以 text-embedding-005text-multilingual-embedding-002 模型端点为例。

对于 AlloyDB Omni,请确保您已设置 AlloyDB Omni 以查询基于云的 Vertex AI 模型

  1. 创建并启用 google_ml_integration 扩展程序

  2. 使用 psql 连接到数据库。

  3. 调用 create model 函数以添加模型端点:

    text-embedding-005

      CALL
        google_ml.create_model(
          model_id => 'text-embedding-005',
          model_request_url => 'publishers/google/models/text-embedding-005',
          model_provider => 'google',
          model_qualified_name => 'text-embedding-005',
          model_type => 'text_embedding',
          model_auth_type => 'alloydb_service_agent_iam');
    

    text-multilingual-embedding-002

      CALL
        google_ml.create_model(
          model_id => 'text-multilingual-embedding-002',
          model_request_url => 'publishers/google/models/text-multilingual-embedding-002',
          model_provider => 'google',
          model_qualified_name => 'text-multilingual-embedding-002',
          model_type => 'text_embedding',
          model_auth_type => 'alloydb_service_agent_iam'
          model_in_transform_fn => 'google_ml.vertexai_text_embedding_input_transform',
          model_out_transform_fn => 'google_ml.vertexai_text_embedding_output_transform');
    

该函数生成的请求网址是指与 AlloyDB Omni 服务账号关联的项目。如果您要引用其他项目,请务必明确指定 model_request_url

如果模型存储在与 AlloyDB 集群不同的项目和区域中,请将请求网址设置为 projects/PROJECT_ID/locations/REGION_ID/publishers/google/models/MODEL_ID,其中 REGION_ID 是模型的托管区域,MODEL_ID 是限定的模型名称。

此外,请向安装 AlloyDB Omni 的项目的服务账号授予 Vertex AI User (roles/aiplatform.user) 角色,以便访问托管在其他项目中的模型。

Open AI 文本嵌入模型

google_ml_integration 扩展程序会自动设置默认转换函数,并调用对远程 OpenAI 模型的调用。如需查看内置支持的 OpenAI 模型列表,请参阅内置支持的模型

以下示例添加了 text-embedding-ada-002 OpenAI 模型端点。您可以使用相同的步骤注册 OpenAI text-embedding-3-smalltext-embedding-3-large 模型端点,并设置特定于模型的模型限定名称。

  1. 使用 psql 连接到数据库。
  2. 创建并启用 google_ml_integration 扩展程序
  3. 将 OpenAI API 密钥添加为 Secret Manager 中的 Secret 以进行身份验证
  4. 调用存储在 Secret Manager 中的 Secret:

    CALL
    google_ml.create_sm_secret(
      secret_id => 'SECRET_ID',
      secret_path => 'projects/PROJECT_ID/secrets/SECRET_MANAGER_SECRET_ID/versions/VERSION_NUMBER');
    

    替换以下内容:

    • SECRET_ID:您设置的 Secret ID,之后会在注册模型端点时使用,例如 key1
    • SECRET_MANAGER_SECRET_ID:您在 Secret Manager 中创建 Secret 时设置的 Secret ID。
    • PROJECT_ID:您的 Google Cloud 项目的 ID。
    • VERSION_NUMBER:密文 ID 的版本号。
  5. 调用 create model 函数以注册 text-embedding-ada-002 模型端点:

    CALL
      google_ml.create_model(
        model_id => 'MODEL_ID',
        model_provider => 'open_ai',
        model_type => 'text_embedding',
        model_qualified_name => 'text-embedding-ada-002',
        model_auth_type => 'secret_manager',
        model_auth_id => 'SECRET_ID');
    

    替换以下内容:

    • MODEL_ID:您定义的模型端点的唯一 ID。此模型 ID 用于引用模型端点生成嵌入或调用预测所需的元数据。
    • SECRET_ID:您之前在 google_ml.create_sm_secret() 过程中使用的 Secret ID。

如需生成嵌入,请参阅如何为内置支持的模型端点生成嵌入

自定义托管的文本嵌入模型

本部分介绍了如何注册自托管模型端点,以及如何创建转换函数和(可选)自定义 HTTP 标头。无论托管在何处,所有自定义托管模型端点都受支持。

以下示例添加了由 Cymbal 托管的 custom-embedding-model 自定义模型端点。cymbal_text_input_transformcymbal_text_output_transform 转换函数用于将模型的输入和输出格式转换为预测函数的输入和输出格式。

如需注册自托管文本嵌入模型端点,请完成以下步骤:

  1. 使用 psql 连接到数据库。

  2. 创建并启用 google_ml_integration 扩展程序

  3. 可选:将 API 密钥添加为 Secret Manager 中的 Secret 以进行身份验证

  4. 调用存储在 Secret Manager 中的 Secret:

    CALL
      google_ml.create_sm_secret(
        secret_id => 'SECRET_ID',
        secret_path => 'projects/project-id/secrets/SECRET_MANAGER_SECRET_ID/versions/VERSION_NUMBER');
    

    替换以下内容:

    • SECRET_ID:您设置的 Secret ID,之后会在注册模型端点时使用,例如 key1
    • SECRET_MANAGER_SECRET_ID:您在 Secret Manager 中创建 Secret 时设置的 Secret ID。
    • PROJECT_ID:您的 Google Cloud 项目的 ID。
    • VERSION_NUMBER:密文 ID 的版本号。
  5. 根据文本嵌入模型端点的预测函数的以下签名,创建输入和输出转换函数。如需详细了解如何创建转换函数,请参阅转换函数示例

    以下是特定于 custom-embedding-model 文本嵌入模型端点的转换函数示例:

    -- Input Transform Function corresponding to the custom model endpoint
    CREATE OR REPLACE FUNCTION cymbal_text_input_transform(model_id VARCHAR(100), input_text TEXT)
    RETURNS JSON
    LANGUAGE plpgsql
    AS $$
    DECLARE
      transformed_input JSON;
      model_qualified_name TEXT;
    BEGIN
      SELECT json_build_object('prompt', json_build_array(input_text))::JSON INTO transformed_input;
      RETURN transformed_input;
    END;
    $$;
    -- Output Transform Function corresponding to the custom model endpoint
    CREATE OR REPLACE FUNCTION cymbal_text_output_transform(model_id VARCHAR(100), response_json JSON)
    RETURNS REAL[]
    LANGUAGE plpgsql
    AS $$
    DECLARE
      transformed_output REAL[];
    BEGIN
      SELECT ARRAY(SELECT json_array_elements_text(response_json->0)) INTO transformed_output;
      RETURN transformed_output;
    END;
    $$;
    
  6. 调用 create model 函数以注册自定义嵌入模型端点:

    CALL
      google_ml.create_model(
        model_id => 'MODEL_ID',
        model_request_url => 'REQUEST_URL',
        model_provider => 'custom',
        model_type => 'text_embedding',
        model_auth_type => 'secret_manager',
        model_auth_id => 'SECRET_ID',
        model_qualified_name => 'MODEL_QUALIFIED_NAME',
        model_in_transform_fn => 'cymbal_text_input_transform',
        model_out_transform_fn => 'cymbal_text_output_transform');
    

    替换以下内容:

    • MODEL_ID:必填。您定义的模型端点的唯一 ID,例如 custom-embedding-model。此模型 ID 用于引用模型端点生成嵌入或调用预测所需的元数据。
    • REQUEST_URL:必填。添加自定义文本嵌入和通用模型端点时,模型专用端点,例如 https://cymbal.com/models/text/embeddings/v1
    • MODEL_QUALIFIED_NAME:如果您的模型端点使用限定名称,则必填。完全限定名称(如果模型端点有多个版本)。
    • SECRET_ID:您之前在 google_ml.create_sm_secret() 过程中使用的 Secret ID。

通用模型

本部分介绍了如何注册托管式模型提供商(例如 Hugging Face、OpenAI、Vertex AI、Anthropic 或任何其他提供商)上提供的任何通用模型端点。本部分将通过示例介绍如何注册托管在 Hugging Face 上的通用模型端点、Vertex AI Model Garden 中的通用 gemini-pro 模型以及 claude-haiku 模型端点。

只要输入和输出采用 JSON 格式,您就可以注册任何通用模型端点。根据模型端点元数据,您可能需要生成 HTTP 标头或定义请求网址。

如需详细了解预注册的通用模型和内置支持的模型,请参阅支持的模型

Gemini 模型

由于某些 gemini-pro 模型是预先注册的,因此您可以直接调用模型 ID 来调用预测。确保您已设置 AlloyDB Omni 以查询基于云的 Vertex AI 模型

以下示例使用 Vertex AI Model Garden 中的 gemini-1.5-pro:generateContent 模型端点。

  1. 使用 psql 连接到数据库。
  2. 创建并启用 google_ml_integration 扩展程序
  3. 使用预注册的模型 ID 调用预测:

    SELECT
        json_array_elements(
        google_ml.predict_row(
            model_id => 'gemini-1.5-pro:generateContent',
            request_body => '{
        "contents": [
            {
                "role": "user",
                "parts": [
                    {
                        "text": "For TPCH database schema as mentioned here https://www.tpc.org/TPC_Documents_Current_Versions/pdf/TPC-H_v3.0.1.pdf , generate a SQL query to find all supplier names which are located in the India nation."
                    }
                ]
            }
        ]
        }'))-> 'candidates' -> 0 -> 'content' -> 'parts' -> 0 -> 'text';
    

Hugging Face 上的通用模型

以下示例展示了如何添加托管在 Hugging Face 上的 facebook/bart-large-mnli 自定义分类模型端点。

  1. 使用 psql 连接到数据库。
  2. 创建并启用 google_ml_integration 扩展程序
  3. 将 OpenAI API 密钥添加为 Secret Manager 中的 Secret 以进行身份验证。如果您已为任何其他 OpenAI 模型创建 Secret,则可以重复使用同一 Secret。
  4. 调用存储在 Secret Manager 中的 Secret:

    CALL
      google_ml.create_sm_secret(
        secret_id => 'SECRET_ID',
        secret_path => 'projects/project-id/secrets/SECRET_MANAGER_SECRET_ID/versions/VERSION_NUMBER');
    

    替换以下内容:

    • SECRET_ID:您设置的密钥 ID,之后会在注册模型端点时使用。
    • SECRET_MANAGER_SECRET_ID:您在 Secret Manager 中创建 Secret 时设置的 Secret ID。
    • PROJECT_ID:您的 Google Cloud 项目的 ID。
    • VERSION_NUMBER:密文 ID 的版本号。
  5. 调用 create model 函数以注册 facebook/bart-large-mnli 模型端点:

    CALL
      google_ml.create_model(
        model_id => 'MODEL_ID',
        model_provider => 'hugging_face',
        model_request_url => 'REQUEST_URL',
        model_qualified_name => 'MODEL_QUALIFIED_NAME',
        model_auth_type => 'secret_manager',
        model_auth_id => 'SECRET_ID');
    

    替换以下内容:

    • MODEL_ID:您定义的模型端点的唯一 ID,例如 custom-classification-model。此模型 ID 用于引用模型端点生成嵌入或调用预测所需的元数据。
    • REQUEST_URL:添加自定义文本嵌入和通用模型端点时,模型专用端点,例如 https://api-inference.huggingface.co/models/facebook/bart-large-mnli
    • MODEL_QUALIFIED_NAME:模型端点版本的完全限定名称,例如 facebook/bart-large-mnli
    • SECRET_ID:您之前在 google_ml.create_sm_secret() 过程中使用的 Secret ID。

Anthropic 通用模型

以下示例添加了 claude-3-opus-20240229 模型端点。模型端点管理功能提供了注册 Anthropoic 模型所需的头文件函数。

确保您已设置 AlloyDB Omni 以查询基于云的 Vertex AI 模型

  1. 使用 psql 连接到数据库。
  2. 创建并启用 google_ml_integration 扩展程序

    Secret Manager

    1. 将 Bearer 令牌作为 Secret 添加到 Secret Manager 以进行身份验证
    2. 调用存储在 Secret Manager 中的 Secret:

      CALL
        google_ml.create_sm_secret(
          secret_id => 'SECRET_ID',
          secret_path => 'projects/project-id/secrets/SECRET_MANAGER_SECRET_ID/versions/VERSION_NUMBER');
      

      替换以下内容:

      • SECRET_ID:您设置的密钥 ID,之后会在注册模型端点时使用。
      • SECRET_MANAGER_SECRET_ID:您在 Secret Manager 中创建 Secret 时设置的 Secret ID。
      • PROJECT_ID:您的 Google Cloud 项目的 ID。
      • VERSION_NUMBER:密文 ID 的版本号。
    3. 调用 create model 函数以注册 claude-3-opus-20240229 模型端点。

      CALL
        google_ml.create_model(
          model_id => 'MODEL_ID',
          model_provider => 'anthropic',
          model_request_url => 'REQUEST_URL',
          model_auth_type => 'secret_manager',
          model_auth_id => 'SECRET_ID',
          generate_headers_fn => 'google_ml.anthropic_claude_header_gen_fn');
      

      替换以下内容:

      • MODEL_ID:您定义的模型端点的唯一 ID,例如 anthropic-opus。此模型 ID 用于引用模型端点生成嵌入或调用预测所需的元数据。
      • REQUEST_URL:添加自定义文本嵌入和通用模型端点时,模型专用端点,例如 https://api.anthropic.com/v1/messages

    身份验证标头

    1. 使用 google_ml.anthropic_claude_header_gen_fn 默认标头生成函数或创建标头生成函数。

        CREATE OR REPLACE FUNCTION anthropic_sample_header_gen_fn(model_id VARCHAR(100), request_body JSON)
        RETURNS JSON
        LANGUAGE plpgsql
        AS $$
        #variable_conflict use_variable
        BEGIN
              RETURN json_build_object('x-api-key', 'ANTHROPIC_API_KEY', 'anthropic-version', 'ANTHROPIC_VERSION')::JSON;
        END;
        $$;
      

      替换以下内容:

      • ANTHROPIC_API_KEY:Anthropic API 密钥。
      • ANTHROPIC_VERSION(可选):您要使用的特定模型版本,例如 2023-06-01
    2. 调用 create model 函数以注册 claude-3-opus-20240229 模型端点。

      CALL
        google_ml.create_model(
          model_id => 'MODEL_ID',
          model_provider => 'anthropic',
          model_request_url => 'REQUEST_URL',
          generate_headers_fn => 'google_ml.anthropic_claude_header_gen_fn');
      

      替换以下内容:

      • MODEL_ID:您定义的模型端点的唯一 ID,例如 anthropic-opus。此模型 ID 用于引用模型端点生成嵌入或调用预测所需的元数据。
      • REQUEST_URL:添加自定义文本嵌入和通用模型端点时,模型专用端点,例如 https://api.anthropic.com/v1/messages

如需了解详情,请参阅如何为通用模型端点调用预测

后续步骤