# Create vector databases from BYO embeddings
URL: https://docs.datarobot.com/en/docs/agentic-ai/genai-code/create-vdb-byo-embedding.html

# Create vector databases from BYO embeddings

The following notebook outlines how you can build and validate a vector database from a bring-your-own (BYO) embedding and register the vector database to the DataRobot platform using the Python client. This notebook is designed for use with DataRobot notebooks; DataRobot recommends downloading this notebook and uploading it for use in the platform.

## Setup

The following steps outline the necessary configuration for integrating vector databases with the DataRobot platform.

1. This workflow uses the following feature flags. Contact your DataRobot representative or administrator for information on enabling these features.
2. Use a codespace, not a DataRobot Notebook, to ensure this notebook has access to a filesystem.
3. Set the notebook session timeout to 180 minutes.
4. Restart the notebook container using at least a "Medium" (16GB RAM) instance.
5. Optionally, upload your documents archive to the notebook filesystem.

### Install requirements

Import the following libraries and modules to interface with the DataRobot platform, prediction environments, and vector database functionality.

```
import datarobot as dr
from datarobot.enums import PredictionEnvironmentPlatform
from datarobot.enums import PredictionEnvironmentModelFormats
from datarobot.models.genai.custom_model_embedding_validation import CustomModelEmbeddingValidation
from datarobot.models.genai.vector_database import VectorDatabase
from datarobot.models.genai.vector_database import ChunkingParameters
```

### Connect to DataRobot

Provide a DataRobot `endpoint` and `token` to connect to DataRobot through the DataRobot Python client. Read more about options for [connecting to DataRobot from the Python client](https://docs.datarobot.com/en/docs/api/api-quickstart/api-qs.html).

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

## Select an environment

Using `dr.ExecutionEnvironment.list()`, iterate through the environments available to your organization, selecting the environment named `[DataRobot] Python 3.11 GenAI` to be the base environment of the vector database and assigning it to `base_environment`.

```
execution_environments = dr.ExecutionEnvironment.list()

base_environment = None
environment_versions = None

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

environment_version = environment_versions[0]
print(base_environment)
print(environment_version)
```

## Create a custom embedding model

Using `dr.CustomInferenceModel.list()` search the available custom models for `all-MiniLM-L6-v2-embedding-model`. If the custom model doesn't exist, create it as `custom_model` using `dr.CustomInferenceModel.create()`. If the custom model does exist, assign it to `custom_model`.

```
CUSTOM_MODEL_NAME = "all-MiniLM-L6-v2-embedding-model_20260218-01"
if CUSTOM_MODEL_NAME not in [c.name for c in dr.CustomInferenceModel.list()]:
    # Create a new custom model
    print("Creating new custom model")
    custom_model = dr.CustomInferenceModel.create(
        name=CUSTOM_MODEL_NAME,
        target_type=dr.TARGET_TYPE.UNSTRUCTURED,
        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()
```

### Write custom embedding model code

Create a directory called `custom_embedding_model` to write custom embedding model code into.

Write custom embedding model code into the `custom.py` file, creating an unstructured model from `all-MiniLM-L6-v2`.

```
import os
os.mkdir('custom_embedding_model')
```

```
%%writefile ./custom_embedding_model/custom.py
import os
os.environ["HF_HOME"] = "/tmp/hf_cache"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/tmp/hf_cache"

from sentence_transformers import SentenceTransformer


def load_model(input_dir):
    return SentenceTransformer("all-MiniLM-L6-v2")


def score_unstructured(model, data, query, **kwargs):
    import json

    data_dict = json.loads(data)
    outputs = model.encode(data_dict["input"])
    return json.dumps(
        {
            "result": outputs.tolist(),
            "device": str(model._target_device)
        }
    )
```

### Write a requirements file

Write the requirements for the custom embedding model into the `requirements.txt` file, ensuring the custom model environment includes the embedding model's dependencies.

```
%%writefile ./custom_embedding_model/requirements.txt
sentence-transformers==3.0.0
```

### Create a custom model version

Using `dr.CustomModelVersion.create_clean`, create a custom model version with the `custom_model`, `base_environment`, and `files` defined in previous steps. In addition, enable public network access using `dr.NETWORK_EGRESS_POLICY.PUBLIC`.

```
# Create a new custom model version in DataRobot
print("Upload new model version to DataRobot")
model_version = dr.CustomModelVersion.create_clean(
    custom_model_id=custom_model.id,
    base_environment_id=base_environment.id,
    files=[
        ("./custom_embedding_model/custom.py", "custom.py"),
        ("./custom_embedding_model/requirements.txt", "requirements.txt"),
    ],
    network_egress_policy=dr.NETWORK_EGRESS_POLICY.PUBLIC,
)
```

### Build a custom model environment

Using `dr.CustomModelVersionDependencyBuild`, build a custom model environment with the required dependencies installed.

```
# Build the custom model environment to ensure dependencies from `requirements.txt` are installed.
build_info = dr.CustomModelVersionDependencyBuild.start_build(
    custom_model_id=custom_model.id,
    custom_model_version_id=model_version.id,
    max_wait=60*10,  # Set a long timeout
)
```

## Register the custom model

Using `dr.RegisteredModel.list()`, search the available custom models for `all-MiniLM-L6-v2-embedding-model` (assigned to `CUSTOM_MODEL_NAME`). If the registered model doesn't exist, create a `registered_model_version` using `dr.RegisteredModelVersion.create_for_custom_model_version`. If the registered model does exist, assign it to `registered_model` and create a `registered_model_version` using `dr.RegisteredModelVersion.create_for_custom_model_version`.

```
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,
        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,
        registered_model_id=registered_model.id
    )
```

## Create a prediction environment for embedding models

Using `dr.PredictionEnvironment.list()`, search the available prediction environments for `Prediction environment for BYO embeddings models`. If the prediction environment doesn't exist, create a DataRobot Serverless `prediction_environment` using `dr.PredictionEnvironment.create`. If the prediction environment does exist, assign it to `prediction_environment`.

```
PREDICTION_ENVIRONMENT_NAME = "Prediction environment for BYO embeddings models"

prediction_environment = None
for _prediction_environment in dr.PredictionEnvironment.list():
    if _prediction_environment.name == PREDICTION_ENVIRONMENT_NAME:
        prediction_environment = _prediction_environment

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

## Deploy the custom embedding model

Using `dr.Deployment.list()`, search the available deployments for `Deployment for all-MiniLM-L6-v2`. If the prediction environment doesn't exist, create a `deployment`, deploying the `registered_model_version` created in a previous section with `dr.Deployment.create_from_registered_model_version`. If the deployment does exist, assign it to `deployment`.

```
MODEL_DEPLOYMENT_NAME = "Deployment for all-MiniLM-L6-v2"

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,
        max_wait=1000,
        prediction_environment_id=prediction_environment.id
    )
else:
    deployment = [d for d in dr.Deployment.list() if d.label == MODEL_DEPLOYMENT_NAME][0]
```

## Create a Use Case for BYO embeddings

Using `dr.UseCase.create`, create the Use Case to use the vector database with and assign it to `use_case`. When working with your own vector database, if [PostgreSQL](https://docs.datarobot.com/en/docs/gen-ai/vector-database/vector-dbs.html#connect-to-postgresql) is the connection method, the output dimension must be less than 2000.

```
use_case = dr.UseCase.create(name="For BYO embeddings")
```

### Upload a dataset to the Use Case

Using `dr.Dataset.create_from_url`, upload the example dataset for the vector database and assign it to `dataset`.

```
# this can be updated with any public URL that is pointing to a .zip file
# in the expected format
dataset_url = "https://s3.amazonaws.com/datarobot_public_datasets/genai/pirate_resumes.zip"

# We will use a vector database with our GenAI models. Let's upload a dataset with our documents.
# If you wish to use a local file as dataset, change this to
# `dataset = dr.Dataset.create_from_file(local_file_path)`
dataset = dr.Dataset.create_from_url(dataset_url)
```

Then, add the dataset to the `use_case` created in the previous section.

```
# Attach dataset to use case.
use_case.add(dataset)
```

## Validate and create the custom embedding model

The `CustomModelVectorDatabaseValidation.create` function executes the validation of the vector database, setting the required settings and associating the custom embedding model with the `use_case` and `deployment` created earlier in this notebook. This step stores the validation ID in `custom_model_embedding_validation`.

```
# Create BYO embeddings validation using prepared deployment
custom_model_embedding_validation = CustomModelEmbeddingValidation.create(
    prompt_column_name="input",
    target_column_name="result",
    deployment_id=deployment.id,
    use_case = use_case,
    name="BYO embeddings",
    wait_for_completion=True,
    prediction_timeout=300,
)
```

### Set chunking parameters and create a vector database

After validation completes, set the `ChunkingParameter()` and use `VectorDatabase.create()` to integrate the vector database. This step uses the `custom_model_embedding_validation`, `dataset`, and `use_case` defined in previous sections.

```
# Use created validation to set up chunking parameters
chunking_parameters = ChunkingParameters(
    embedding_validation=custom_model_embedding_validation,
    chunking_method="recursive",
    chunk_size=256,
    chunk_overlap_percentage=50,
    separators=["\n\n", "\n", " ", ""],
    embedding_model=None,
)


vdb = VectorDatabase.create(
    dataset_id=dataset.id,
    chunking_parameters=chunking_parameters,
    use_case=use_case
)
```

```
from time import sleep
vdb = VectorDatabase.get(vdb.id)
sleep(30)
assert vdb.execution_status == "COMPLETED"
```

---
