コードでの外部LLMの作成¶

!pip install openai datarobot-drum datarobot-predict

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

dr.Client()

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

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

!mkdir custom_model

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

dr.Credential.list()

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
    })

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

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

# 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 score()
score(test_data, client=load_client())

%%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/requirements.txt
openai
datarobot-drum
pandas
python-dotenv

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

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

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

# 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)

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()

# 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", "")),
    ]
)

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

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

# 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"
)

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,
        )
    )

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

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
    ]

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

# 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)

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,
)

assert external_llm_validation.validation_status == "PASSED"

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