Generative AI documentation > Code-based GenAI > Create external LLMs with code

Create external LLMs with code¶

The following, designed for use with DataRobot Notebooks, outlines how you can build and validate an external LLM using the DataRobot Python client. DataRobot recommends downloading this notebook and uploading it for use in the platform.

Note: For self-managed users, when code samples reference app.datarobot.com, change them to the appropriate URL for your instance.

Setup¶

The following steps outline the configuration necessary for integrating an external LLM with the DataRobot platform.

This workflow requires that the following feature flags are enabled. Contact your DataRobot representative or administrator for information on enabling these features.
- Enable MLOps
- Enable Public Network Access for all Custom Models
- Enable Monitoring Support for Generative Models
- Enable Custom Inference Models
- Enable GenAI Experimentation
Create a new credential in the DataRobot Credentials Management tool:
- Create it as an API Token type credential.
- Set the display name as openai-api-key.
- Place your OpenAI API key in the Token field.
Use a codespace, not a DataRobot Notebook, to ensure this notebook has access to a filesystem.
Set the notebook session timeout to 180 minutes.
Restart the notebook container using at least a "Medium" (16GB RAM) instance.

Install libraries¶

Install the following libraries:

In [ ]:

Copied!

!pip install openai datarobot-drum datarobot-predict
!pip install openai datarobot-drum datarobot-predict

In [ ]:

Copied!

import datarobot as dr
from datarobot.models.genai.custom_model_llm_validation import CustomModelLLMValidation
import datarobot as dr
from datarobot.models.genai.custom_model_llm_validation import CustomModelLLMValidation

Connect to DataRobot¶

Read more about different options for connecting to DataRobot from the Python client.

In [ ]:

Copied!

dr.Client()

token = dr.Client().token
endpoint = f"{dr.Client().endpoint}"

dr.Client(endpoint=endpoint, token=token)
dr.Client()

token = dr.Client().token
endpoint = f"{dr.Client().endpoint}"

dr.Client(endpoint=endpoint, token=token)

Create a directory for custom code¶

Create a directory called custom_model that will hold your OpenAI wrapper code.

In [ ]:

Copied!

!mkdir custom_model
!mkdir custom_model

Define environment variables¶

In the cell below, define each environment variable required to run this notebook. These environment variables are added as runtime parameters to the custom model in the workshop.

In [ ]:

Copied!





%%writefile .env

# Define required environment variables for the codespace environment
OPENAI_API_KEY=<OPENAI_API_KEY> # For codespace testing, the API key for the Azure OpenAI API
OPENAI_API_VERSION=<OPENAI_API_VERSION>
OPENAI_API_BASE=<OPENAI_API_BASE>
OPENAI_DEPLOYMENT_NAME=<OPENAI_DEPLOYMENT_NAME>
DATAROBOT_CREDENTIAL_OPENAI=<DATAROBOT_CREDENTIAL_OPENAI> # For the deployed model, the ObjectId of the Azure OpenAI API Token credential
%%writefile .env

# Define required environment variables for the codespace environment
OPENAI_API_KEY= # For codespace testing, the API key for the Azure OpenAI API
OPENAI_API_VERSION=
OPENAI_API_BASE=
OPENAI_DEPLOYMENT_NAME=
DATAROBOT_CREDENTIAL_OPENAI= # For the deployed model, the ObjectId of the Azure OpenAI API Token credential

The value for DATAROBOT_CREDENTIAL_OPENAI environment variable must be the ObjectId of the Azure OpenAI API Token credential. You can find the ID of the credential using the DataRobot Python API client's credentials list method.

In [ ]:

Copied!

dr.Credential.list()
dr.Credential.list()

Define hooks¶

The following cell defines the methods used to deploy a text generation custom model. These include loading the custom model and using the model for scoring.

In [ ]:

Copied!





import os
import pandas as pd
from typing import Any, Iterator, Union, Dict, List
from openai import AzureOpenAI
from openai.types.chat import ChatCompletion, ChatCompletionChunk

# Try to load dotenv for codespace testing
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass  # Don't load the .env file

CompletionCreateParams = Dict[str, Any]


def get_config():
    # Get configuration from runtime parameters or environment variables
    try:
        # Try DataRobot runtime parameters first
        from datarobot_drum import RuntimeParameters
        return {
            "api_key": RuntimeParameters.get("OPENAI_API_KEY")["apiToken"],
            "api_base": RuntimeParameters.get("OPENAI_API_BASE"),
            "api_version": RuntimeParameters.get("OPENAI_API_VERSION"),
            "deployment_name": RuntimeParameters.get("OPENAI_DEPLOYMENT_NAME")
        }
    except Exception:
        # Fallback to environment variables for codespace testing
        return {
            "api_key": os.environ.get("OPENAI_API_KEY", ""),
            "api_base": os.environ.get("OPENAI_API_BASE", ""),
            "api_version": os.environ.get("OPENAI_API_VERSION", ""),
            "deployment_name": os.environ.get("OPENAI_DEPLOYMENT_NAME", "")
        }


# Implement the load_model hook.
def load_model(*args, **kwargs):
    config = get_config()
    return AzureOpenAI(
        api_key=config["api_key"],
        azure_endpoint=config["api_base"],
        api_version=config["api_version"]
    )


# Load the Azure OpenAI client
def load_client(*args, **kwargs):
    return load_model(*args, **kwargs)


# Get supported LLM models
def get_supported_llm_models(client: AzureOpenAI) -> List[Dict[str, Any]]:
    azure_models = client.models.list()
    model_ids = [m.id for m in azure_models]
    return model_ids if model_ids else ["datarobot-deployed-llm"]


# On-demand chat requests
def chat(
    completion_create_params: CompletionCreateParams, client: AzureOpenAI
) -> Union[ChatCompletion, Iterator[ChatCompletionChunk], Dict]:

    try:
        if completion_create_params.get("model") == "datarobot-deployed-llm":
            config = get_config()
            completion_create_params["model"] = config["deployment_name"]

        return client.chat.completions.create(**completion_create_params)
    except Exception as e:
        return {
            "error": f"{e.__class__.__name__}: {str(e)}"
        }


# Batch chat requests
PROMPT_COLUMN_NAME = "promptText"
COMPLETION_COLUMN_NAME = "resultText"
ERROR_COLUMN_NAME = "error"


def score(data, client, **kwargs):
    prompts = data["promptText"].tolist()
    responses = []
    errors = []

    for prompt in prompts:
        try:
            # Get model config
            config = get_config()

            # Attempt to get a completion from the client
            response = client.chat.completions.create(
                model=config["deployment_name"],
                messages=[{"role": "user", "content": f"{prompt}"},],
                max_tokens=20,
                temperature=0
            )
            # On success, append the content and a null error
            responses.append(response.choices[0].message.content or "")
            errors.append("")
        except Exception as e:
            # On failure, format the error message
            error = f"{e.__class__.__name__}: {str(e)}"
            responses.append("")
            errors.append(error)

    return pd.DataFrame({
        PROMPT_COLUMN_NAME: prompts,
        COMPLETION_COLUMN_NAME: responses,
        ERROR_COLUMN_NAME: errors
    })
import os
import pandas as pd
from typing import Any, Iterator, Union, Dict, List
from openai import AzureOpenAI
from openai.types.chat import ChatCompletion, ChatCompletionChunk

# Try to load dotenv for codespace testing
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass  # Don't load the .env file

CompletionCreateParams = Dict[str, Any]


def get_config():
    # Get configuration from runtime parameters or environment variables
    try:
        # Try DataRobot runtime parameters first
        from datarobot_drum import RuntimeParameters
        return {
            "api_key": RuntimeParameters.get("OPENAI_API_KEY")["apiToken"],
            "api_base": RuntimeParameters.get("OPENAI_API_BASE"),
            "api_version": RuntimeParameters.get("OPENAI_API_VERSION"),
            "deployment_name": RuntimeParameters.get("OPENAI_DEPLOYMENT_NAME")
        }
    except Exception:
        # Fallback to environment variables for codespace testing
        return {
            "api_key": os.environ.get("OPENAI_API_KEY", ""),
            "api_base": os.environ.get("OPENAI_API_BASE", ""),
            "api_version": os.environ.get("OPENAI_API_VERSION", ""),
            "deployment_name": os.environ.get("OPENAI_DEPLOYMENT_NAME", "")
        }


# Implement the load_model hook.
def load_model(*args, **kwargs):
    config = get_config()
    return AzureOpenAI(
        api_key=config["api_key"],
        azure_endpoint=config["api_base"],
        api_version=config["api_version"]
    )


# Load the Azure OpenAI client
def load_client(*args, **kwargs):
    return load_model(*args, **kwargs)


# Get supported LLM models
def get_supported_llm_models(client: AzureOpenAI) -> List[Dict[str, Any]]:
    azure_models = client.models.list()
    model_ids = [m.id for m in azure_models]
    return model_ids if model_ids else ["datarobot-deployed-llm"]


# On-demand chat requests
def chat(
    completion_create_params: CompletionCreateParams, client: AzureOpenAI
) -> Union[ChatCompletion, Iterator[ChatCompletionChunk], Dict]:

    try:
        if completion_create_params.get("model") == "datarobot-deployed-llm":
            config = get_config()
            completion_create_params["model"] = config["deployment_name"]

        return client.chat.completions.create(**completion_create_params)
    except Exception as e:
        return {
            "error": f"{e.__class__.__name__}: {str(e)}"
        }


# Batch chat requests
PROMPT_COLUMN_NAME = "promptText"
COMPLETION_COLUMN_NAME = "resultText"
ERROR_COLUMN_NAME = "error"


def score(data, client, **kwargs):
    prompts = data["promptText"].tolist()
    responses = []
    errors = []

    for prompt in prompts:
        try:
            # Get model config
            config = get_config()

            # Attempt to get a completion from the client
            response = client.chat.completions.create(
                model=config["deployment_name"],
                messages=[{"role": "user", "content": f"{prompt}"},],
                max_tokens=20,
                temperature=0
            )
            # On success, append the content and a null error
            responses.append(response.choices[0].message.content or "")
            errors.append("")
        except Exception as e:
            # On failure, format the error message
            error = f"{e.__class__.__name__}: {str(e)}"
            responses.append("")
            errors.append(error)

    return pd.DataFrame({
        PROMPT_COLUMN_NAME: prompts,
        COMPLETION_COLUMN_NAME: responses,
        ERROR_COLUMN_NAME: errors
    })

Test hooks locally¶

Before proceeding with the deployment, use the cell below to test that the custom model hooks function correctly.

In [ ]:

Copied!

# Provide test data
test_data = pd.DataFrame({PROMPT_COLUMN_NAME: ["What is a large language model (LLM)?"]})
# Provide test data
test_data = pd.DataFrame({PROMPT_COLUMN_NAME: ["What is a large language model (LLM)?"]})

In [ ]:

Copied!

# Test get_supported_llm_models()
models = get_supported_llm_models(load_client())
print(f"Available models: {models}")
# Test get_supported_llm_models()
models = get_supported_llm_models(load_client())
print(f"Available models: {models}")

In [ ]:

Copied!





# Test chat()
chat(
    {
        "model": "datarobot-deployed-llm",
        "messages": [{"role": "user", "content": "What is a large language model (LLM)?"}],
        "max_tokens": 20,
        "temperature": 0,
    },
    client=load_client(),
)
# Test chat()
chat(
    {
        "model": "datarobot-deployed-llm",
        "messages": [{"role": "user", "content": "What is a large language model (LLM)?"}],
        "max_tokens": 20,
        "temperature": 0,
    },
    client=load_client(),
)

In [ ]:

Copied!

# Test score()
score(test_data, client=load_client())
# Test score()
score(test_data, client=load_client())

Save the custom model code¶

Next, save the hooks above as custom_model/custom.py. This python file will be executed by DataRobot using your credentials. The following is a copy of the cell where you previously defined the hooks.

In [ ]:

Copied!





%%writefile custom_model/custom.py

import os
import pandas as pd
from typing import Any, Iterator, Union, Dict, List
from openai import AzureOpenAI
from openai.types.chat import ChatCompletion, ChatCompletionChunk

# Try to load dotenv for codespace testing
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass  # Don't load the .env file

CompletionCreateParams = Dict[str, Any]


def get_config():
    # Get configuration from runtime parameters or environment variables 
    try:
        # Try DataRobot runtime parameters
        from datarobot_drum import RuntimeParameters
        return {
            "api_key": RuntimeParameters.get("OPENAI_API_KEY")["apiToken"],
            "api_base": RuntimeParameters.get("OPENAI_API_BASE"),
            "api_version": RuntimeParameters.get("OPENAI_API_VERSION"),
            "deployment_name": RuntimeParameters.get("OPENAI_DEPLOYMENT_NAME")
        }
    except Exception:
        # Fallback to environment variables for codespace testing
        return {
            "api_key": os.environ.get("OPENAI_API_KEY", ""),
            "api_base": os.environ.get("OPENAI_API_BASE", ""),
            "api_version": os.environ.get("OPENAI_API_VERSION", ""),
            "deployment_name": os.environ.get("OPENAI_DEPLOYMENT_NAME", "")
        }


# Implement the load_model hook.
def load_model(*args, **kwargs):
    config = get_config()
    return AzureOpenAI(
        api_key=config["api_key"],
        azure_endpoint=config["api_base"],
        api_version=config["api_version"]
    )

# Load the Azure OpenAI client.
def load_client(*args, **kwargs):
    return load_model(*args, **kwargs)


# Get supported LLM models
def get_supported_llm_models(client: AzureOpenAI) -> List[Dict[str, Any]]:
    azure_models = client.models.list()
    model_ids = [m.id for m in azure_models]
    return model_ids if model_ids else ["datarobot-deployed-llm"]


# On-demand chat requests
def chat(
    completion_create_params: CompletionCreateParams, client: AzureOpenAI
) -> Union[ChatCompletion, Iterator[ChatCompletionChunk], Dict]:

    try:
        if completion_create_params.get("model") == "datarobot-deployed-llm":
            config = get_config()
            completion_create_params["model"] = config["deployment_name"]
            
        return client.chat.completions.create(**completion_create_params)
    except Exception as e:
        return {
            "error": f"{e.__class__.__name__}: {str(e)}"
        }


# Batch chat requests
PROMPT_COLUMN_NAME = "promptText"
COMPLETION_COLUMN_NAME = "resultText"
ERROR_COLUMN_NAME = "error"


def score(data, client, **kwargs):
    prompts = data["promptText"].tolist()
    responses = []
    errors = []

    for prompt in prompts:
        try:
            # Get model config
            config = get_config()
            
            # Attempt to get a completion from the client
            response = client.chat.completions.create(
                model=config["deployment_name"],
                messages=[{"role": "user", "content": f"{prompt}"},],
                max_tokens=20,
                temperature=0
            )
            # On success, append the content and a null error
            responses.append(response.choices[0].message.content or "")
            errors.append("")
        except Exception as e:
            # On failure, format the error message
            error = f"{e.__class__.__name__}: {str(e)}"
            responses.append("")
            errors.append(error)

    return pd.DataFrame({
        PROMPT_COLUMN_NAME: prompts,
        COMPLETION_COLUMN_NAME: responses,
        ERROR_COLUMN_NAME: errors
    })
%%writefile custom_model/custom.py

import os
import pandas as pd
from typing import Any, Iterator, Union, Dict, List
from openai import AzureOpenAI
from openai.types.chat import ChatCompletion, ChatCompletionChunk

# Try to load dotenv for codespace testing
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass  # Don't load the .env file

CompletionCreateParams = Dict[str, Any]


def get_config():
    # Get configuration from runtime parameters or environment variables 
    try:
        # Try DataRobot runtime parameters
        from datarobot_drum import RuntimeParameters
        return {
            "api_key": RuntimeParameters.get("OPENAI_API_KEY")["apiToken"],
            "api_base": RuntimeParameters.get("OPENAI_API_BASE"),
            "api_version": RuntimeParameters.get("OPENAI_API_VERSION"),
            "deployment_name": RuntimeParameters.get("OPENAI_DEPLOYMENT_NAME")
        }
    except Exception:
        # Fallback to environment variables for codespace testing
        return {
            "api_key": os.environ.get("OPENAI_API_KEY", ""),
            "api_base": os.environ.get("OPENAI_API_BASE", ""),
            "api_version": os.environ.get("OPENAI_API_VERSION", ""),
            "deployment_name": os.environ.get("OPENAI_DEPLOYMENT_NAME", "")
        }


# Implement the load_model hook.
def load_model(*args, **kwargs):
    config = get_config()
    return AzureOpenAI(
        api_key=config["api_key"],
        azure_endpoint=config["api_base"],
        api_version=config["api_version"]
    )

# Load the Azure OpenAI client.
def load_client(*args, **kwargs):
    return load_model(*args, **kwargs)


# Get supported LLM models
def get_supported_llm_models(client: AzureOpenAI) -> List[Dict[str, Any]]:
    azure_models = client.models.list()
    model_ids = [m.id for m in azure_models]
    return model_ids if model_ids else ["datarobot-deployed-llm"]


# On-demand chat requests
def chat(
    completion_create_params: CompletionCreateParams, client: AzureOpenAI
) -> Union[ChatCompletion, Iterator[ChatCompletionChunk], Dict]:

    try:
        if completion_create_params.get("model") == "datarobot-deployed-llm":
            config = get_config()
            completion_create_params["model"] = config["deployment_name"]
            
        return client.chat.completions.create(**completion_create_params)
    except Exception as e:
        return {
            "error": f"{e.__class__.__name__}: {str(e)}"
        }


# Batch chat requests
PROMPT_COLUMN_NAME = "promptText"
COMPLETION_COLUMN_NAME = "resultText"
ERROR_COLUMN_NAME = "error"


def score(data, client, **kwargs):
    prompts = data["promptText"].tolist()
    responses = []
    errors = []

    for prompt in prompts:
        try:
            # Get model config
            config = get_config()
            
            # Attempt to get a completion from the client
            response = client.chat.completions.create(
                model=config["deployment_name"],
                messages=[{"role": "user", "content": f"{prompt}"},],
                max_tokens=20,
                temperature=0
            )
            # On success, append the content and a null error
            responses.append(response.choices[0].message.content or "")
            errors.append("")
        except Exception as e:
            # On failure, format the error message
            error = f"{e.__class__.__name__}: {str(e)}"
            responses.append("")
            errors.append(error)

    return pd.DataFrame({
        PROMPT_COLUMN_NAME: prompts,
        COMPLETION_COLUMN_NAME: responses,
        ERROR_COLUMN_NAME: errors
    })

Save requirements and metadata files to help describe the model's environment and usage.

In [ ]:

Copied!





%%writefile custom_model/requirements.txt
openai
datarobot-drum
pandas
python-dotenv
%%writefile custom_model/requirements.txt
openai
datarobot-drum
pandas
python-dotenv

In [ ]:

Copied!





%%writefile custom_model/model-metadata.yaml
---
name: OpenAI gpt-4o
type: inference
targetType: textgeneration
runtimeParameterDefinitions:
  - fieldName: OPENAI_API_KEY
    type: credential
    credentialType: api_token
    description: OpenAI API key stored in DataRobot
    allowEmpty: false
  - fieldName: OPENAI_API_VERSION
    type: string
    description: OpenAI API version string
    allowEmpty: false
  - fieldName: OPENAI_API_BASE
    type: string
    description: OpenAI API base URL string
    allowEmpty: false
  - fieldName: OPENAI_DEPLOYMENT_NAME
    type: string
    description: OpenAI API deployment ID string
    allowEmpty: false
%%writefile custom_model/model-metadata.yaml
---
name: OpenAI gpt-4o
type: inference
targetType: textgeneration
runtimeParameterDefinitions:
  - fieldName: OPENAI_API_KEY
    type: credential
    credentialType: api_token
    description: OpenAI API key stored in DataRobot
    allowEmpty: false
  - fieldName: OPENAI_API_VERSION
    type: string
    description: OpenAI API version string
    allowEmpty: false
  - fieldName: OPENAI_API_BASE
    type: string
    description: OpenAI API base URL string
    allowEmpty: false
  - fieldName: OPENAI_DEPLOYMENT_NAME
    type: string
    description: OpenAI API deployment ID string
    allowEmpty: false

Test the code locally¶

The DataRobot DRUM library allows you to test the code as if DataRobot were running it via a simple CLI. To do this, supply a test file and then run it.

In [ ]:

Copied!

# Create the test file
test_data.to_csv("custom_model/test_data.csv", index=False)
# Create the test file
test_data.to_csv("custom_model/test_data.csv", index=False)

In [ ]:

Copied!

os.putenv("TARGET_NAME", COMPLETION_COLUMN_NAME)
!drum score --code-dir custom_model/ --target-type textgeneration --input custom_model/test_data.csv
os.putenv("TARGET_NAME", COMPLETION_COLUMN_NAME)
!drum score --code-dir custom_model/ --target-type textgeneration --input custom_model/test_data.csv

Create a custom model in DataRobot¶

The code below performs a few steps to register your code with DataRobot:

Creates a custom model to contain the versioned code.
Creates a custom model version with the code in the custom_model folder and adds the required Runtime Parameters.
Builds the environment to run the model by installing the requirements.txt file.
Tests the entire setup.

In [ ]:

Copied!





# List all existing base environments
execution_environments = dr.ExecutionEnvironment.list()
execution_environments

BASE_ENVIRONMENT = None
for execution_environment in execution_environments:
    if execution_environment.name == "[DataRobot] Python 3.11 GenAI":
        BASE_ENVIRONMENT = execution_environment
        environment_versions = dr.ExecutionEnvironmentVersion.list(
            execution_environment.id
        )
        break

if BASE_ENVIRONMENT is None:
    raise ValueError(
        "Required execution environment '[DataRobot] Python 3.11 GenAI' not found. Please check your DataRobot instance."
    )

BASE_ENVIRONMENT_VERSION = environment_versions[0]

print(BASE_ENVIRONMENT)
print(BASE_ENVIRONMENT_VERSION)
print(BASE_ENVIRONMENT.id)
# List all existing base environments
execution_environments = dr.ExecutionEnvironment.list()
execution_environments

BASE_ENVIRONMENT = None
for execution_environment in execution_environments:
    if execution_environment.name == "[DataRobot] Python 3.11 GenAI":
        BASE_ENVIRONMENT = execution_environment
        environment_versions = dr.ExecutionEnvironmentVersion.list(
            execution_environment.id
        )
        break

if BASE_ENVIRONMENT is None:
    raise ValueError(
        "Required execution environment '[DataRobot] Python 3.11 GenAI' not found. Please check your DataRobot instance."
    )

BASE_ENVIRONMENT_VERSION = environment_versions[0]

print(BASE_ENVIRONMENT)
print(BASE_ENVIRONMENT_VERSION)
print(BASE_ENVIRONMENT.id)

In [ ]:

Copied!





CUSTOM_MODEL_NAME = "External LLM OpenAI Wrapper Model"
if CUSTOM_MODEL_NAME not in [c.name for c in dr.CustomInferenceModel.list()]:
    # Create a new custom model
    print("Creating a new custom model")
    custom_model = dr.CustomInferenceModel.create(
        name=CUSTOM_MODEL_NAME,
        target_type=dr.TARGET_TYPE.TEXT_GENERATION,
        target_name=COMPLETION_COLUMN_NAME,
        description="Wrapper for OpenAI completion",
        language="Python",
        is_training_data_for_versions_permanently_enabled=True,
    )
else:
    print("Custom model exists")
    custom_model = [
        c for c in dr.CustomInferenceModel.list() if c.name == CUSTOM_MODEL_NAME
    ].pop()
CUSTOM_MODEL_NAME = "External LLM OpenAI Wrapper Model"
if CUSTOM_MODEL_NAME not in [c.name for c in dr.CustomInferenceModel.list()]:
    # Create a new custom model
    print("Creating a new custom model")
    custom_model = dr.CustomInferenceModel.create(
        name=CUSTOM_MODEL_NAME,
        target_type=dr.TARGET_TYPE.TEXT_GENERATION,
        target_name=COMPLETION_COLUMN_NAME,
        description="Wrapper for OpenAI completion",
        language="Python",
        is_training_data_for_versions_permanently_enabled=True,
    )
else:
    print("Custom model exists")
    custom_model = [
        c for c in dr.CustomInferenceModel.list() if c.name == CUSTOM_MODEL_NAME
    ].pop()

In [ ]:

Copied!





# Create a new custom model version in DataRobot with required runtime parameters
print("Upload new version of model to DataRobot")
model_version = dr.CustomModelVersion.create_clean(
    custom_model_id=custom_model.id,
    base_environment_id=BASE_ENVIRONMENT.id,
    files=[
        "./custom_model/custom.py",
        "./custom_model/requirements.txt",
        "./custom_model/model-metadata.yaml",
        
    ],
    network_egress_policy=dr.NETWORK_EGRESS_POLICY.PUBLIC,
    runtime_parameter_values=[
        dr.models.custom_model_version.RuntimeParameterValue(field_name="OPENAI_API_KEY", type="credential", value=os.environ.get("DATAROBOT_CREDENTIAL_OPENAI", "")),
        dr.models.custom_model_version.RuntimeParameterValue(field_name="OPENAI_API_VERSION", type="string", value=os.environ.get("OPENAI_API_VERSION", "")),
        dr.models.custom_model_version.RuntimeParameterValue(field_name="OPENAI_API_BASE", type="string", value=os.environ.get("OPENAI_API_BASE", "")),
        dr.models.custom_model_version.RuntimeParameterValue(field_name="OPENAI_DEPLOYMENT_NAME", type="string", value=os.environ.get("OPENAI_DEPLOYMENT_NAME", "")),
    ]
)
# Create a new custom model version in DataRobot with required runtime parameters
print("Upload new version of model to DataRobot")
model_version = dr.CustomModelVersion.create_clean(
    custom_model_id=custom_model.id,
    base_environment_id=BASE_ENVIRONMENT.id,
    files=[
        "./custom_model/custom.py",
        "./custom_model/requirements.txt",
        "./custom_model/model-metadata.yaml",
        
    ],
    network_egress_policy=dr.NETWORK_EGRESS_POLICY.PUBLIC,
    runtime_parameter_values=[
        dr.models.custom_model_version.RuntimeParameterValue(field_name="OPENAI_API_KEY", type="credential", value=os.environ.get("DATAROBOT_CREDENTIAL_OPENAI", "")),
        dr.models.custom_model_version.RuntimeParameterValue(field_name="OPENAI_API_VERSION", type="string", value=os.environ.get("OPENAI_API_VERSION", "")),
        dr.models.custom_model_version.RuntimeParameterValue(field_name="OPENAI_API_BASE", type="string", value=os.environ.get("OPENAI_API_BASE", "")),
        dr.models.custom_model_version.RuntimeParameterValue(field_name="OPENAI_DEPLOYMENT_NAME", type="string", value=os.environ.get("OPENAI_DEPLOYMENT_NAME", "")),
    ]
)

In [ ]:

Copied!





try:
    build_info = dr.CustomModelVersionDependencyBuild.start_build(
        custom_model_id=custom_model.id,
        custom_model_version_id=model_version.id,
        max_wait=3600,
    )
    print("Finished new dependency build")
except dr.errors.ClientError as e:
    if "already has a dependency image" in str(e):
        print("Dependency build already exists, skipping build step")
        build_info = None
    else:
        raise e
try:
    build_info = dr.CustomModelVersionDependencyBuild.start_build(
        custom_model_id=custom_model.id,
        custom_model_version_id=model_version.id,
        max_wait=3600,
    )
    print("Finished new dependency build")
except dr.errors.ClientError as e:
    if "already has a dependency image" in str(e):
        print("Dependency build already exists, skipping build step")
        build_info = None
    else:
        raise e

Test the custom model in DataRobot¶

Next, use the environment to run the model with prediction test data to verify that the custom model is functional before deployment. To do this, upload the inference dataset for testing predictions.

In [ ]:

Copied!

pred_test_dataset = dr.Dataset.create_from_in_memory_data(test_data)
pred_test_dataset.modify(name="LLM Test Data")
pred_test_dataset.update()
pred_test_dataset = dr.Dataset.create_from_in_memory_data(test_data)
pred_test_dataset.modify(name="LLM Test Data")
pred_test_dataset.update()

After uploading the inference dataset, you can test the custom model.

In [ ]:

Copied!





# Test a new version in DataRobot
print("Run test of new version in DataRobot")
custom_model_test = dr.CustomModelTest.create(
    custom_model_id=custom_model.id,
    custom_model_version_id=model_version.id,
    dataset_id=pred_test_dataset.id,
    max_wait=3600,  # 1 hour timeout
)
custom_model_test.overall_status

HOST = "https://app.datarobot.com"
for name, test in custom_model_test.detailed_status.items():
    print("Test: {}".format(name))
    print("Status: {}".format(test["status"]))
    print("Message: {}".format(test["message"]))

print(
    "Finished testing: "
    + HOST
    + "model-registry/custom-models/"
    + custom_model.id
    + "/assemble"
)
# Test a new version in DataRobot
print("Run test of new version in DataRobot")
custom_model_test = dr.CustomModelTest.create(
    custom_model_id=custom_model.id,
    custom_model_version_id=model_version.id,
    dataset_id=pred_test_dataset.id,
    max_wait=3600,  # 1 hour timeout
)
custom_model_test.overall_status

HOST = "https://app.datarobot.com"
for name, test in custom_model_test.detailed_status.items():
    print("Test: {}".format(name))
    print("Status: {}".format(test["status"]))
    print("Message: {}".format(test["message"]))

print(
    "Finished testing: "
    + HOST
    + "model-registry/custom-models/"
    + custom_model.id
    + "/assemble"
)

Register and deploy the LLM¶

Next, register the model in the Registry. The model registry contains entries from all models (predictive, generative, built in DataRobot, and externally hosted).

In [ ]:

Copied!





if CUSTOM_MODEL_NAME not in [m.name for m in dr.RegisteredModel.list()]:
    print("Creating New Registered Model")
    registered_model_version = (
        dr.RegisteredModelVersion.create_for_custom_model_version(
            model_version.id,
            name=CUSTOM_MODEL_NAME,
            description="LLM Wrapper Example from DataRobot Docs",
            registered_model_name=CUSTOM_MODEL_NAME,
        )
    )
else:
    print("Using Existing Model")
    registered_model = [
        m for m in dr.RegisteredModel.list() if m.name == CUSTOM_MODEL_NAME
    ].pop()
    registered_model_version = (
        dr.RegisteredModelVersion.create_for_custom_model_version(
            model_version.id,
            name=CUSTOM_MODEL_NAME,
            description="LLM Wrapper Example from DataRobot Docs",
            registered_model_id=registered_model.id,
        )
    )
if CUSTOM_MODEL_NAME not in [m.name for m in dr.RegisteredModel.list()]:
    print("Creating New Registered Model")
    registered_model_version = (
        dr.RegisteredModelVersion.create_for_custom_model_version(
            model_version.id,
            name=CUSTOM_MODEL_NAME,
            description="LLM Wrapper Example from DataRobot Docs",
            registered_model_name=CUSTOM_MODEL_NAME,
        )
    )
else:
    print("Using Existing Model")
    registered_model = [
        m for m in dr.RegisteredModel.list() if m.name == CUSTOM_MODEL_NAME
    ].pop()
    registered_model_version = (
        dr.RegisteredModelVersion.create_for_custom_model_version(
            model_version.id,
            name=CUSTOM_MODEL_NAME,
            description="LLM Wrapper Example from DataRobot Docs",
            registered_model_id=registered_model.id,
        )
    )

Now, deploy the model. If you are a DataRobot multitenant SaaS user, you must select a prediction environment.

In [ ]:

Copied!

pred_server = [s for s in dr.PredictionServer.list()][0]
print(f"Prediction server ID: {pred_server}")
pred_server = [s for s in dr.PredictionServer.list()][0]
print(f"Prediction server ID: {pred_server}")

In [ ]:

Copied!





MODEL_DEPLOYMENT_NAME = "LLM Wrapper Deployment"

if MODEL_DEPLOYMENT_NAME not in [d.label for d in dr.Deployment.list()]:
    deployment = dr.Deployment.create_from_registered_model_version(
        registered_model_version.id,
        label=MODEL_DEPLOYMENT_NAME,
        description="Your new deployment",
        max_wait=1000,
        # Only needed for DataRobot Managed AI Platform
        default_prediction_server_id=pred_server.id,
    )
else:
    deployment = [d for d in dr.Deployment.list() if d.label == MODEL_DEPLOYMENT_NAME][
        0
    ]
MODEL_DEPLOYMENT_NAME = "LLM Wrapper Deployment"

if MODEL_DEPLOYMENT_NAME not in [d.label for d in dr.Deployment.list()]:
    deployment = dr.Deployment.create_from_registered_model_version(
        registered_model_version.id,
        label=MODEL_DEPLOYMENT_NAME,
        description="Your new deployment",
        max_wait=1000,
        # Only needed for DataRobot Managed AI Platform
        default_prediction_server_id=pred_server.id,
    )
else:
    deployment = [d for d in dr.Deployment.list() if d.label == MODEL_DEPLOYMENT_NAME][
        0
    ]

Test the deployment¶

Test that the deployment can successfully provide responses to prompts.

In [ ]:

Copied!





from datarobot_predict.deployment import predict

input_df = pd.DataFrame(
    {
        PROMPT_COLUMN_NAME: [
            "Give me some context on large language models and their applications?",
            "What is AutoML?",
            "Tell me a joke",
        ],
    }
)


result_df, response_headers = predict(deployment, input_df)
result_df
from datarobot_predict.deployment import predict

input_df = pd.DataFrame(
    {
        PROMPT_COLUMN_NAME: [
            "Give me some context on large language models and their applications?",
            "What is AutoML?",
            "Tell me a joke",
        ],
    }
)


result_df, response_headers = predict(deployment, input_df)
result_df

Validate the external LLM¶

The following methods execute and validate the external LLM.

This example associates a Use Case with the validation and creates the vector database within that Use Case. Set the use_case_id to specify an existing Use Case or create a new one with that name.

In [ ]:

Copied!





# Option 1: Create a new Use Case (default approach)
use_case = dr.UseCase.create()

# Option 2: Use an existing Use Case (replace with a Use Case ID)
# use_case_id = <use_case_id>
# use_case = dr.UseCase.get(use_case_id)
# Option 1: Create a new Use Case (default approach)
use_case = dr.UseCase.create()

# Option 2: Use an existing Use Case (replace with a Use Case ID)
# use_case_id = 
# use_case = dr.UseCase.get(use_case_id)

CustomModelLLMValidation.create executes the validation of the external LLM. Be sure to provide the deployment ID.

In [ ]:

Copied!





external_llm_validation = CustomModelLLMValidation.create(
    prompt_column_name=PROMPT_COLUMN_NAME,
    target_column_name=COMPLETION_COLUMN_NAME,
    deployment_id=deployment.id,
    name="My External LLM",
    use_case=use_case,
    wait_for_completion=True,
)
external_llm_validation = CustomModelLLMValidation.create(
    prompt_column_name=PROMPT_COLUMN_NAME,
    target_column_name=COMPLETION_COLUMN_NAME,
    deployment_id=deployment.id,
    name="My External LLM",
    use_case=use_case,
    wait_for_completion=True,
)

In [ ]:

Copied!

assert external_llm_validation.validation_status == "PASSED"
assert external_llm_validation.validation_status == "PASSED"

In [ ]:

Copied!

print(f"External LLM Validation ID: {external_llm_validation.id}")
print(f"External LLM Validation ID: {external_llm_validation.id}")

This external LLM can now be used in the GenAI E2E walkthrough, for example to create the LLM blueprint.