Query an external model with ai_query()
Note
This feature is in Public Preview. To query endpoints that serve external models, you must enroll in the public preview. Please populate and submit the AI Functions Public Preview enrollment form.
This article illustrates how to set up and query an external model endpoint using the built-in Databricks SQL function ai_query()
. The example uses external model support in Mosaic AI Model Serving to query gpt-4
provided by OpenAI and accomplish chat tasks. See AI Functions on Databricks for more detail about this AI function.
Requirements
See the requirements of ai_query SQL function.
An OpenAI API key.
Store the key in a Databricks secret. In this example you store the API key in scope
my-external-model
and secretopenai
.
Create an external model endpoint
The following creates an external model serving endpoint that serves OpenAI gpt-4
for a chat task.
To create a personal access token, see Authentication for Databricks automation.
import requests
import json
personal_access_token = "your-personal-access-token"
headers = {
"Authorization": "Bearer " + personal_access_token,
}
host = "https://oregon.cloud.databricks.com/"
url = host + "api/2.0/serving-endpoints"
data = {
"name": "my-external-openai-chat",
"config": {
"served_entities": [
{
"name": "my_entity",
"external_model": {
"name": "gpt-4",
"provider": "openai",
"openai_config": {
"openai_api_key": "{{secrets/my-external-model/openai}}",
},
"task": "llm/v1/chat",
},
}
],
},
}
response = requests.post(url, headers=headers, json=data)
print("Status Code", response.status_code)
print("JSON Response ", json.dumps(json.loads(response.text), indent=4))
Query the external model with ai_query()
In the Databricks SQL query editor, you can write SQL queries to query the external model serving endpoint.
Example queries:
SELECT ai_query(
"my-external-openai-chat",
"What is a large language model?"
)
SELECT question, ai_query(
"my-external-openai-chat",
"You are a customer service agent. Answer the customer's question in 100 words: " || question
) AS answer
FROM
uc_catalog.schema.customer_questions
SELECT
sku_id,
product_name,
ai_query(
"my-external-openai-chat",
"You are a marketing expert for a winter holiday promotion targeting GenZ. Generate a promotional text in 30 words mentioning a 50% discount for product: " || product_name
)
FROM
uc_catalog.schema.retail_products
WHERE
inventory > 2 * forecasted_sales