ai_query
function
Applies to: Databricks SQL
Important
This feature is in Public Preview.
Invokes an existing Azure Databricks Model Serving endpoint and parses and returns its response.
Requirements
- This function is not available on Azure Databricks SQL Classic.
- Querying Foundation Model APIs is enabled by default. To query endpoints that serve custom models or external models, you must enroll in the public preview. Please populate and submit the AI Functions Public Preview enrollment form.
Syntax
To query an endpoint that serves an external model or a foundation model:
ai_query(endpointName, request)
To query a custom model serving endpoint:
ai_query(endpointName, request, returnType)
Arguments
endpointName
: A STRING literal, the name of the existing Mosaic AI 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 endpoint is an external model serving endpoint or Databricks Foundation Model APIs endpoint, the request must be a STRING.
- If the endpoint is a custom model serving endpoint, the request can be a single column or a struct expression. The struct field names should match the input feature names expected by the endpoint.
returnType
: 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. Required for querying a custom Model Serving Endpoint.
Returns
The parsed response from the endpoint.
Examples
To query an External Model Serving Endpoint or Databricks Foundation Model:
> 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."
> CREATE FUNCTION correct_grammar(text STRING)
RETURNS STRING
RETURN ai_query(
'databricks-llama-2-70b-chat',
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;
To query a custom Model Serving Endpoint:
> 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."]}
Related functions
Feedback
https://aka.ms/ContentUserFeedback.
Coming soon: Throughout 2024 we will be phasing out GitHub Issues as the feedback mechanism for content and replacing it with a new feedback system. For more information see:Submit and view feedback for