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

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