在 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 密钥或不记名令牌,请完成以下步骤:

  1. 在 Secret Manager 中创建 Secret。如需了解详情,请参阅创建 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) 角色,以便可以访问托管在其他项目中的模型。

打开 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 添加到 Secret Manager 以进行身份验证
  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 时在 Secret Manager 中设置的 Secret ID。
    • PROJECT_ID:您的 Google Cloud 项目的 ID。
    • VERSION_NUMBER:Secret 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 添加到 Secret Manager 以进行身份验证

  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 时在 Secret Manager 中设置的 Secret ID。
    • PROJECT_ID:您的 Google Cloud 项目的 ID。
    • VERSION_NUMBER:Secret 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 添加到 Secret Manager 以进行身份验证。如果您已为任何其他 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:您设置的 Secret ID,随后在注册模型端点时会用到。
    • SECRET_MANAGER_SECRET_ID:您创建 Secret 时在 Secret Manager 中设置的 Secret ID。
    • PROJECT_ID:您的 Google Cloud 项目的 ID。
    • VERSION_NUMBER:Secret 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 模型端点。模型端点管理提供了注册 Anthropic 模型所需的标头函数。

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

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

    Secret Manager

    1. 将不记名令牌作为 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:您设置的 Secret ID,随后在注册模型端点时会用到。
      • SECRET_MANAGER_SECRET_ID:您创建 Secret 时在 Secret Manager 中设置的 Secret ID。
      • PROJECT_ID:您的 Google Cloud 项目的 ID。
      • VERSION_NUMBER:Secret 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

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

后续步骤