Use the DataRobot LLM gateway¶
The DataRobot LLM gateway service provides a DataRobot API endpoint to interface with LLMs hosted by external LLM providers. To request LLM responses from the DataRobot LLM gateway, you can use any API client that supports OpenAI-compatible chat completion API, for example, the OpenAI Python API library.
Setup¶
The DataRobot LLM gateway is a premium feature; contact your DataRobot representative or administrator for information on enabling it. To use the DataRobot LLM gateway with the OpenAI Python API library, first, make sure the OpenAI client package is installed and imported. You must also import the DataRobot Python client.
%pip install openai
import datarobot as dr
from openai import OpenAI
Connect to the LLM gateway¶
Next, initialize the OpenAI client. The base URL is the DataRobot LLM gateway endpoint. The example below assembles this URL by combining your DataRobot API endpoint (retrieved from dr_client) and /genai/llmgw. Usually this results in https://app.datarobot.com/api/v2/genai/llmgw, https://app.eu.datarobot.com/api/v2/genai/llmgw, https://app.jp.datarobot.com/api/v2/genai/llmgw, or your organization's DataRobot API endpoint URL with /genai/llmgw appended to it. The API key is your DataRobot API key.
dr_client = dr.Client()
DR_API_TOKEN = dr_client.token
LLM_GATEWAY_BASE_URL = f"{dr_client.endpoint}/genai/llmgw"
client = OpenAI(
base_url=LLM_GATEWAY_BASE_URL,
api_key=DR_API_TOKEN,
)
print(f"Your LLM gateway URL is {LLM_GATEWAY_BASE_URL}.")
Select models and make requests¶
In your code, you can specify any supported provider LLM and set up the message to send to the LLM as a prompt. An optional argument is client_id, where you can specify the caller service to use for metering: genai-playground, custom-model, or moderations.
This example calls the LLM gateway catalog endpoint to get one of the latest supported LLMs from each provider.
response = dr_client.get(url="genai/llmgw/catalog/")
catalog = response.json()["data"]
first_model_by_provider = {}
# Iterate through the catalog to select the first supported LLM from each provider
for supported_llm in catalog:
# Get the provider name from the provider key
provider_name = supported_llm['provider']
# If the provider isn't already recorded, store the full model string
if provider_name not in first_model_by_provider:
model_name = supported_llm['model']
first_model_by_provider[provider_name] = model_name
# Store the list of models and a message
models = list(first_model_by_provider.values())
message = [{"role": "user", "content": "Hello! What is your name and who made you?"}]
print(models)
After you define the client, models and message, you can make chat completion requests to the LLM gateway. The authentication uses DataRobot-provided credentials.
from IPython.display import display, Markdown
for model in models:
response = client.chat.completions.create(
model=model,
messages=message,
)
response_text = response.choices[0].message.content
output_as_markdown = f"""
**{model}:**
{response_text}
---
"""
display(Markdown(output_as_markdown));
To further configure your chat completion request when making direct calls to an LLM gateway, specify LLM parameter settings like temperature, max_completion_tokens, and more. These parameters are also supported for custom models. For more information on the available parameters, see the OpenAI chat completion documentation.
model2 = models[2]
message2 = [{"role": "user", "content": "Hello! What is your name and who made you? How do you feel about Agentic AI"}]
extra_body2 = {
"temperature": 0.8,
"max_completion_tokens": 200,
}
response2 = client.chat.completions.create(
model=model2,
messages=message2,
extra_body=extra_body2,
)
response_text2 = response2.choices[0].message.content
output_as_markdown2 = f"""
**{model2}:**
{response_text2}
---
"""
display(Markdown(output_as_markdown2));
Identify supported LLMs¶
To provide a list of LLMs supported by the LLM gateway, this example runs a get call to the catalog endpoint and prints the returned list of LLMs.
from pprint import pprint
response = dr_client.get(url="genai/llmgw/catalog/")
data = response.json()["data"]
supported_llms = [llm_model["model"] for llm_model in data]
pprint(supported_llms)
If you try to use an unsupported LLM, the LLM gateway returns an error message, relaying that the specified LLM is not in the LLM catalog.
messages3 = [
{"role": "user", "content": "Hello!"}
]
response = client.chat.completions.create(
model="unsupported-provider/random-llm",
messages=messages3,
)
print(response.choices[0].message.content)