通过

ai_query 函数

  • 项目

适用于：勾选“是” Databricks SQL 勾选“是” Databricks Runtime

重要

此功能目前以公共预览版提供。

调用现有的 Azure Databricks 模型服务终结点，然后分析并返回其响应。

请参阅以下指南，了解如何 ai_query 用于不同的用例：

要求

注意

  • 在 Databricks Runtime 14.2 及更高版本中，Databricks 笔记本（包括作为任务在 Databricks 工作流中运行的笔记本）支持此函数。
  • 在 Databricks Runtime 14.1 及更低版本中，Databricks 笔记本不支持此函数。

语法

若要查询为外部模型或基础模型提供服务的终结点，请执行以下操作：

ai_query(endpointName, request)

若要查询具有模型架构的自定义模型服务终结点，请执行以下操作：

ai_query(endpointName, request)

若要查询没有模型架构的自定义模型服务终结点，请执行以下操作：

ai_query(endpointName, request, returnType, failOnError)

参数

  • endpointName：字符串字面量、Databricks 基础模型服务终结点的名称、同一工作区中用于调用的外部模型服务终结点或自定义模型终结点。 定义者必须对终结点具有“可查询”权限。
  • request：一个表达式，是用于调用终结点的请求。
    • 如果终结点是一个外部模型服务终结点或 Databricks 基础模型 API 终结点，则请求必须是一个字符串。
    • 如果终结点是一个自定义模型服务终结点，则请求可以是单个列或结构表达式。 结构字段名称应与终结点所需的输入特征名称匹配。
  • returnType：一个表达式，即来自终结点的预期 returnType。 这与 from_json 函数中的架构参数类似，它接受 STRING 表达式或 schema_of_json 函数的调用。
    • 在 Databricks Runtime 14.2 及更高版本中，如果未提供此表达式，则 ai_query() 会自动从自定义模型服务终结点的模型架构推断返回类型。
    • 在 Databricks Runtime 14.1 及更低版本中，查询自定义模型服务终结点需要此表达式。
  • failOnError：（可选） 默认为 true 的布尔文本。 此标志指示是否在响应中包含 ai_query 错误状态。 当此格式设置为 false时，响应格式不同。
  • modelParameters （可选）：包含聊天、完成和嵌入模型参数的结构字段，用于为基础模型或外部模型提供服务。 这些模型参数必须是常量参数，而不是依赖于数据。 如果未指定或设置为 null 默认值，则使用这些模型参数。 除了默认值为 0.0 之外temperature，这些模型参数的默认值与基础模型 REST API 参考中列出的默认值相同。

返回

来自终结点的已分析响应。

  • 如果 failOnError => true，该函数将返回与现有行为相同的结果，即终结点分析的响应。 分析的响应的数据类型是从模型类型、模型架构终结点或 returnType 函数中的参数推断的 ai_query 。
  • 如果 failOnError => false，该函数将分析的响应和错误状态字符串作为 STRUCT 对象返回。
    • 如果行的推理成功，则该 errorStatus 字段为 null。
    • 如果由于模型终结点错误导致行推理失败，则该 response 字段为 null。
    • 如果由于其他错误导致行的推理失败，则整个查询将失败。

示例

下面是一个批处理推理示例，该示例使用 failOnError and modelParameters with max_tokens 和 temperature. 此示例还演示如何使用 concat()连接模型和推理列的提示符。 有多种方法可以执行串联，例如使用 ||、 concat()或 format_string()。


CREATE OR REPLACE TABLE ${output_table_name} AS (
  SELECT
      ${input_column_name},
      AI_QUERY(
        "${endpoint}",
        CONCAT("${prompt}", ${input_column_name}),
        failOnError => false,
        modelParameters => named_struct('max_tokens', ${num_output_tokens},'temperature', ${temperature})
      ) as response
    FROM ${input_table_name}
    LIMIT ${input_num_rows}
)

若要查询创建外部模型服务终结点：

> SELECT ai_query(
    'my-external-model-openai-chat',
    'Describe Databricks SQL in 30 words.'
  ) AS summary

  "Databricks SQL is a cloud-based platform for data analytics and machine learning, providing a unified workspace for collaborative data exploration, analysis, and visualization using SQL queries."

查询 Databricks Foundation 模型 API 支持的基础模型：

> SELECT *,
  ai_query(
    'databricks-meta-llama-3-1-70b-instruct',
    "Can you tell me the name of the US state that serves the provided ZIP code? zip code: " || pickup_zip
    )
  FROM samples.nyctaxi.trips
  LIMIT 10

（可选）还可以在 UDF 中包装对 ai_query() 的调用，以便调用函数，如下所示：

> CREATE FUNCTION correct_grammar(text STRING)
  RETURNS STRING
  RETURN ai_query(
    'databricks-meta-llama-3-1-70b-instruct',
    CONCAT('Correct this to standard English:\n', text));
> GRANT EXECUTE ON correct_grammar TO ds;
- DS fixes grammar issues in a batch.
> SELECT
    * EXCEPT text,
    correct_grammar(text) AS text
  FROM articles;

若要查询自定义模型服务终结点，请执行以下操作：


> SELECT text, ai_query(
    endpoint => 'spam-classification-endpoint',
    request => named_struct(
      'timestamp', timestamp,
      'sender', from_number,
      'text', text),
    returnType => 'BOOLEAN') AS is_spam
  FROM messages

> SELECT ai_query(
    'weekly-forecast',
    request => struct(*),
    returnType => 'FLOAT') AS predicted_revenue
  FROM retail_revenue

> SELECT ai_query(
    'custom-llama-2-7b-chat',
    request => named_struct("messages",
        ARRAY(named_struct("role", "user", "content", "What is ML?"))),
    returnType => 'STRUCT<candidates:ARRAY<STRING>>')

  {"candidates":["ML stands for Machine Learning. It's a subfield of Artificial Intelligence that involves the use of algorithms and statistical models to enable machines to learn from data, make decisions, and improve their performance on a specific task over time."]}

将 DLT 通道设置为预览的示例查询：

> create or replace materialized view
    ai_query_mv
    TBLPROPERTIES('pipelines.channel' = 'PREVIEW') AS
  SELECT
    ai_query("databricks-dbrx-instruct", text) as response
  FROM
    messages