Query a Chat Completion Model

1. Query a Chat Completion Model

Now, let's learn about some of the AI functionality that Databricks provides. Specifically, we will dive into how to use Databricks to query Chat Completion Models.

2. Query serving endpoint

The Serving Endpoint API exposes a query function that can be used to query the AI model being served by a specified serving endpoint. Some of the parameters it accepts are name, messages, and max_tokens. name is the name of the model representing the serving endpoint. messages is a list of ChatMessage objects and max_tokens tells the model the maximum number of "tokens" or words that the model can include in its response to queries, allowing us to limit the length of responses from the model.

3. ChatMessage

We mentioned that we pass in a list of ChatMessage objects into the messages parameter when querying a serving endpoint. A ChatMessage requires content, which is the content of the query or directive we want to send to the LLM. The role represents the ChatMessageRole.

4. ChatMessageRole

If we want to provide instructions to the LLM about how to respond to future queries, we send a ChatMessage with a SYSTEM ChatMessageRole. If we want to ask the LLM a query then we specify the role as a USER ChatMessageRole.

5. Query chat completion Large Language Model (LLM)

Let's look put this all together with an example. First we use the system role to tell the LLM to act like a helpful assistant when responding to future queries. Next, we ask it to explain what a fibonacci sequence is, using the "user" role. The response from the query is a list of choices objects. We can use response.choices[0].message.content to retrieve the content of the first chat message response from the LLM.

6. Query response structure

If we ask a single question, we can parse the content of the AI model response as shown in the code snippet.

7. Example: summarize text

LLMs in Databricks can do more than just provide responses to queries,they can also summarize large amounts of content succinctly and accurately. To demonstrate the Chat model's ability to summarize text, we query the Databricks Meta Llama model to summarize a long string of text, stored in a variable some_long_text.

8. Example: generate content

Here is an example of using the service_endpoints.query() method on the WorkspaceClient to query the Meta Llama 3 Instruct Model. To demonstrate the Chat model's ability to generate content, we first tell the AI agent to act like a ghost writer for a famous country singer, using the "SYSTEM" ChatMessageRole. We also send a message with the "USER" ChatMessageRole, asking it to write the lyrics to a country song that takes place in Mississippi.

9. Let's practice!

Let's practice using the Databricks SDK to interact with AI models and ask them to do useful tasks for us.

Create Your Free Account

By continuing, you accept our Terms of Use, our Privacy Policy and that your data is stored in the USA.