DataRobotのベクターデータベースを作成およびデプロイする¶

import datarobot as dr
from datarobot.models.genai.vector_database import VectorDatabase
from datarobot.models.genai.vector_database import ChunkingParameters
from datarobot.enums import VectorDatabaseEmbeddingModel
from datarobot.enums import VectorDatabaseChunkingMethod
from datarobot.enums import PredictionEnvironmentPlatform
from datarobot.enums import PredictionEnvironmentModelFormats
import time
import requests

# Option 1: Use environment variables (recommended)
# The client will automatically use DATAROBOT_ENDPOINT and DATAROBOT_API_TOKEN
dr.Client()

# Option 2: Explicitly provide endpoint and token
# endpoint = "https://app.datarobot.com/api/v2"
# token = "<ADD_YOUR_TOKEN_HERE>"
# dr.Client(endpoint=endpoint, token=token)

# Create the Use Case
USE_CASE_NAME = "VDB Example Use Case"
use_case = dr.UseCase.create(name=USE_CASE_NAME)
print(f"Created Use Case: {use_case.name}")

# Get or create the dataset
DATASET_NAME = "pirate_resumes.zip"
dataset = None

# Search for existing dataset by exact name
try:
    all_datasets = dr.Dataset.list(filter_failed=True)
    for d in all_datasets:
        if d.name == DATASET_NAME:
            dataset = d
            print(f"Found existing dataset: {dataset.name}")
            break
except Exception:
    pass

# Upload if not found
if not dataset:
    dataset_url = "https://s3.amazonaws.com/datarobot_public_datasets/genai/pirate_resumes.zip"
    dataset = dr.Dataset.create_from_url(dataset_url)
    print(f"Uploaded new dataset: {dataset.name}")

# Add dataset to Use Case if not already added
try:
    use_case_datasets = use_case.get_datasets()
    dataset_ids = [d.id for d in use_case_datasets]
    if dataset.id not in dataset_ids:
        use_case.add(dataset)
        print(f"Added dataset to Use Case")
except Exception:
    pass  # Already added or error adding

chunking_parameters = ChunkingParameters(
    embedding_model=VectorDatabaseEmbeddingModel.JINA_EMBEDDING_T_EN_V1,
    chunking_method=VectorDatabaseChunkingMethod.RECURSIVE,
    chunk_size=256,
    chunk_overlap_percentage=25,
    separators=["\n\n", "\n", " ", ""],
)

vdb = VectorDatabase.create(
    dataset_id=dataset.id,
    chunking_parameters=chunking_parameters,
    use_case=use_case,
    name="My Vector Database"
)

max_wait_time = 600
check_interval = 5
start_time = time.time()

print("Waiting for vector database creation...")
while time.time() - start_time < max_wait_time:
    vdb = VectorDatabase.get(vdb.id)
    status = vdb.execution_status
    if status == "COMPLETED":
        print(f"Vector database created: {vdb.name}")
        break
    elif status == "FAILED":
        error_msg = getattr(vdb, 'error_message', 'Unknown error')
        raise Exception(f"Vector database creation failed: {error_msg}")
    else:
        # Show progress if available
        percentage = getattr(vdb, 'percentage', None)
        if percentage is not None:
            print(f"  Status: {status} ({percentage}%)")
        else:
            print(f"  Status: {status}...")
    time.sleep(check_interval)
else:
    raise Exception(f"Vector database creation timed out after {max_wait_time} seconds")

assert vdb.execution_status == "COMPLETED", f"Vector database creation failed with status: {vdb.execution_status}"

PREDICTION_ENVIRONMENT_NAME = "Vector Database Prediction Environment"

# Get or create prediction environment
prediction_environment = None
for env in dr.PredictionEnvironment.list():
    if env.name == PREDICTION_ENVIRONMENT_NAME:
        prediction_environment = env
        break

if prediction_environment is None:
    prediction_environment = dr.PredictionEnvironment.create(
        name=PREDICTION_ENVIRONMENT_NAME,
        platform=PredictionEnvironmentPlatform.DATAROBOT_SERVERLESS,
        supported_model_formats=[
            PredictionEnvironmentModelFormats.DATAROBOT,
            PredictionEnvironmentModelFormats.CUSTOM_MODEL
        ],
    )
    print(f"Created prediction environment: {prediction_environment.name}")
else:
    print(f"Using existing prediction environment: {prediction_environment.name}")

# Select 3XL resource bundle
resource_bundle_id = None
try:
    dr_client = dr.Client()
    bundles_url = f"{dr_client.endpoint}/mlops/compute/bundles/"
    headers = {"Authorization": f"Bearer {dr_client.token}"}
    bundles_response = requests.get(bundles_url, headers=headers, params={"useCases": "customModel"})
    
    if bundles_response.status_code == 200:
        bundles_data = bundles_response.json()
        if bundles_data.get("data"):
            bundles = bundles_data["data"]
            # Look for 3XL bundle
            bundle_3xl = next((b for b in bundles if "3XL" in b.get("name", "").upper()), None)
            if bundle_3xl:
                resource_bundle_id = bundle_3xl["id"]
                print(f"Selected 3XL bundle: {bundle_3xl['name']}")
            else:
                # Fallback to largest available
                sorted_bundles = sorted(bundles, key=lambda b: b.get("memoryBytes", 0), reverse=True)
                if sorted_bundles:
                    resource_bundle_id = sorted_bundles[0]["id"]
                    print(f"Warning: 3XL bundle not found. Using largest available: {sorted_bundles[0]['name']}")
        else:
            print("Using memory settings (no resource bundles available)")
    else:
        print("Using memory settings (resource bundles not enabled)")
except (ImportError, KeyError) as e:
    print(f"Using memory settings (error checking bundles): {e}")
except requests.RequestException as e:
    print(f"Using memory settings (network error checking bundles): {e}")

assert vdb.execution_status == "COMPLETED", f"Vector database must be completed. Current status: {vdb.execution_status}"

# Send to workshop with 3XL bundle or memory settings
if resource_bundle_id:
    custom_model_version = vdb.send_to_custom_model_workshop(
        resource_bundle_id=resource_bundle_id,
        replicas=1,
        network_egress_policy=dr.NETWORK_EGRESS_POLICY.PUBLIC,
    )
else:
    custom_model_version = vdb.send_to_custom_model_workshop(
        maximum_memory=4096*1024*1024,
        replicas=1,
        network_egress_policy=dr.NETWORK_EGRESS_POLICY.PUBLIC,
    )

print(f"Custom model version created: {custom_model_version}")

REGISTERED_MODEL_NAME = f"Vector Database - {vdb.name}"

# Register model (adds new version if model already exists)
existing_models = [m for m in dr.RegisteredModel.list() if m.name == REGISTERED_MODEL_NAME]

if existing_models:
    registered_model_version = dr.RegisteredModelVersion.create_for_custom_model_version(
        custom_model_version_id=custom_model_version.id,
        registered_model_id=existing_models[0].id,
    )
    print(f"Added new version to existing registered model: {REGISTERED_MODEL_NAME}")
else:
    registered_model_version = dr.RegisteredModelVersion.create_for_custom_model_version(
        custom_model_version_id=custom_model_version.id,
        registered_model_name=REGISTERED_MODEL_NAME,
    )
    print(f"Created new registered model: {REGISTERED_MODEL_NAME}")

registered_model = dr.RegisteredModel.get(registered_model_version.registered_model_id)
max_wait_time = 600
check_interval = 10
start_time = time.time()

print("Waiting for model build to complete...")
while time.time() - start_time < max_wait_time:
    version = registered_model.get_version(registered_model_version.id)
    build_status = getattr(version, 'build_status', None) or getattr(version, 'buildStatus', None)
    
    if build_status in ('READY', 'complete', 'COMPLETE'):
        print(f"Model build completed (status: {build_status})")
        break
    elif build_status in ('FAILED', 'ERROR', 'error'):
        raise Exception(f"Model build failed. Status: {build_status}")
    else:
        print(f"  Build status: {build_status}...")
    time.sleep(check_interval)
else:
    version = registered_model.get_version(registered_model_version.id)
    build_status = getattr(version, 'build_status', None) or getattr(version, 'buildStatus', None)
    raise Exception(f"Model build timed out. Current status: {build_status}")

# Verify ready status
version = registered_model.get_version(registered_model_version.id)
final_status = getattr(version, 'build_status', None) or getattr(version, 'buildStatus', None)
if final_status not in ('READY', 'complete', 'COMPLETE'):
    raise Exception(f"Model not ready for deployment. Status: {final_status}")

deployment = dr.Deployment.create_from_registered_model_version(
    registered_model_version.id,
    label=f"Vector Database Deployment - {vdb.name}",
    description="Vector database deployment for RAG applications",
    prediction_environment_id=prediction_environment.id,
    max_wait=600,
)

print(f"Deployment created: {deployment.id}")

# List all vector databases in a Use Case
vdbs = VectorDatabase.list(use_case=use_case)
print(f"Found {len(vdbs)} vector database(s) in Use Case '{use_case.name}'")