ai_query function

Applies to: check marked yes Databricks SQL

Preview

This feature is in Public Preview.

Invokes an existing Databricks Model Serving endpoint and parses and returns its response.

Requirements

  • This function is only available on Databricks SQL Pro and Serverless.

  • This feature is in public preview. To enroll in the public preview, please populate and submit the AI Functions Public Preview enrollment form.

Syntax 

ai_query(endpointName, request[, param1, value1] [...])

Arguments

  • endpointName: A STRING literal, the name of the existing Databricks Model Serving endpoint in the same workspace for invocations. The definer must have “Can Query” permission on the endpoint.

  • request: An expression, the request used to invoke the endpoint.

    • If the input is not a struct expression, the endpoint should take only one input.

    • If the input is a struct expression, the endpoint should take named inputs. The struct field names should match the input names expected by the endpoint.

  • paramN and valueN: key-value pairs to provide required information on the selected endpoint. The keys must be string literals and are case-sensitive. The types of the values depend on the following keys:

    • returnType: Required. An expression, the expected returnType from the endpoint. This is similar to the schema parameter in from_json function, which accepts both A STRING expression or invocation of schema_of_json function.

Returns

The parsed response from the endpoint.

Examples 

> SELECT ai_query(
    'qa_bot',
    'What gets stored in the metastore of Unity Catalog?',
    'returnType', 'STRING')

  "The metastore of Unity Catalog stores metadata about data assets (tables and views) and the permissions that govern access to them."


> SELECT ai_query(
    'qa_bot',
    named_struct(
      'prompt', 'Write me a tweet about the launch of Dolly 2.0.',
      'temperature', 0.5),
    'returnType', 'STRING')
  "We've upgraded our LLM, making it more efficient, more powerful, and more accessible to a broader audience."


> CREATE FUNCTION correct_grammar(text STRING)
  RETURNS STRING
  RETURN ai_query(
    'dolly',
    named_struct(
      'prompt',
      CONCAT(
        'Correct this to standard English:\n',
        text),
      'temperature', 0.5)
    'returnType', 'STRING');
> GRANT EXECUTE ON correct_grammer TO ds;
- DS fixes grammar issues in a batch.
> SELECT
    * EXCEPT text,
    correct_grammar(text) AS text
  FROM articles;