Share via

Query an external model with ai_query()

  • Article

Note

This feature is in Public Preview. To query endpoints that serve external models, enable AI_Query for Custom Models and External Models in the Databricks Previews UI.

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 Azure 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

Additional resources