End-to-end code-first generative AI experimentation¶
This notebook outlines a comprehensive overview of the generative AI features DataRobot has to offer.
Note: For self-managed users, code samples that reference app.datarobot.com
need to be changed to the appropriate URL for your instance.
import datarobot as dr
from datarobot.models.genai.vector_database import VectorDatabase
from datarobot.models.genai.vector_database import ChunkingParameters
from datarobot.enums import PromptType
from datarobot.enums import VectorDatabaseEmbeddingModel
from datarobot.enums import VectorDatabaseChunkingMethod
from datarobot.models.genai.playground import Playground
from datarobot.models.genai.llm import LLMDefinition
from datarobot.models.genai.llm_blueprint import LLMBlueprint
from datarobot.models.genai.llm_blueprint import VectorDatabaseSettings
from datarobot.models.genai.chat import Chat
from datarobot.models.genai.chat_prompt import ChatPrompt
from datarobot.models.genai.vector_database import CustomModelVectorDatabaseValidation
from datarobot.models.genai.comparison_chat import ComparisonChat
from datarobot.models.genai.comparison_prompt import ComparisonPrompt
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.
endpoint="https://app.datarobot.com/api/v2"
token="<ADD_VALUE_HERE>"
dr.Client(endpoint=endpoint, token=token)
Create a Use Case¶
Generative AI works within a DataRobot Use Case, so to begin, create a new use case.
use_case = dr.UseCase.create()
This workflow uses a vector database with GenAI models. Next, you will upload a dataset with a set of documents.
# 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)
Create a vector database¶
chunking_parameters = ChunkingParameters(
embedding_model=VectorDatabaseEmbeddingModel.JINA_EMBEDDING_T_EN_V1,
chunking_method=VectorDatabaseChunkingMethod.RECURSIVE,
chunk_size=128,
chunk_overlap_percentage=25,
separators=[],
)
vdb = VectorDatabase.create(dataset.id, chunking_parameters, use_case)
Use the following cell to retrieve the vector database and make sure it has completed. Approximate completion time is 30-60 seconds.
vdb = VectorDatabase.get(vdb.id)
assert vdb.execution_status == "COMPLETED"
Create an LLM playground¶
Create an LLM playground that will host large language models (LLMs).
playground = Playground.create(name="New playground for example", use_case=use_case)
Use the following cell to see what kind of LLMs are available. By default this method returns as a dict for easy readability. It includes a list of allowed settings for each LLM along with any constraints on the values.
llms_dict = LLMDefinition.list(use_case=use_case)
print(llms_dict)
As an example, use GPT 3.5.
llms = LLMDefinition.list(use_case=use_case, as_dict=False)
gpt = llms[0]
To interact with GPT, you need to create an LLM blueprint with the settings that you want. Since the allowed settings depend on the LLM type, take a generic dict with those values that will be validated during the creation of the blueprint. Review the allowed settings for GPT.
print([setting.to_dict() for setting in gpt.settings])
Create an LLM blueprint¶
llm_settings = {
"system_prompt": (
"You are a pirate who begins each response with 'Arrr' "
"and ends each response with 'Matey!'"
),
"max_completion_length": 256, # Let's ask for short responses
"temperature": 2.0, # With a high degree of variability
}
# PromptType.ONE_TIME_PROMPT is an alternative if you don't wish
# for previous chat prompts (history) to be included in each subsequent prompt
prompting_strategy = PromptType.CHAT_HISTORY_AWARE
Use some vector database settings that could be different from the defaults when the database was created.
vector_database_settings = VectorDatabaseSettings(
max_documents_retrieved_per_prompt=2,
max_tokens=128,
)
llm_blueprint = LLMBlueprint.create(
playground=playground,
name="GPT",
llm=gpt,
prompt_type=prompting_strategy,
llm_settings=llm_settings,
vector_database=vdb,
vector_database_settings=vector_database_settings,
)
Create a chat¶
Now create a chat and associate it with the LLM blueprint you just created. A chat is the entity that groups together a set of prompt messages. It can be thought of as a conversation with a particular LLM blueprint.
chat = Chat.create(
name="Resume review chat with a pirate",
llm_blueprint=llm_blueprint
)
Chat with the LLM blueprint¶
Use the following cell to chat with the LLM blueprint.
prompt1 = ChatPrompt.create(
chat=chat,
text="How can a pirate resume be evaluated?",
wait_for_completion=True,
)
print(prompt1.result_text)
The example cell above returns a truncated output, confirmed by result_metadata
.
print(prompt1.result_metadata.error_message)
Try to chat again with a larger max_completion_length
. LLM settings can be updated within a chat prompt as long as the LLM blueprint hasn't yet been saved. Doing so updates the underlying settings of the LLM blueprint.
llm_settings["max_completion_length"] = 1024
prompt2 = ChatPrompt.create(
text="Please answer the previous question again.",
chat=chat,
llm_settings=llm_settings,
wait_for_completion=True,
)
print(prompt2.result_text)
You can keep the new max_completion_length
and confirm that the LLM blueprint has been updated.
llm_blueprint = LLMBlueprint.get(llm_blueprint.id)
print(llm_blueprint.llm_settings)
Next, lower the temperature to modify the reply, as the response may be unsatisfactory.
llm_settings["temperature"] = 0.8
prompt3 = ChatPrompt.create(
text="Please summarize the best pirate resume.",
chat=chat,
llm_settings=llm_settings,
wait_for_completion=True,
)
print(prompt3.result_text)
Confirm that the LLM is using information from the vector database.
print([citation.text for citation in prompt3.citations])
Next, create a new LLM blueprint from the existing one so that you can change some settings and compare the two.
new_llm_blueprint = LLMBlueprint.create_from_llm_blueprint(llm_blueprint, name="new blueprint")
Remove the vector database from the new blueprint, modify the settings as shown below, and save the blueprint.
llm_settings["system_prompt"] = "You are an AI assistant who helps to evaluate resumes."
new_llm_blueprint = new_llm_blueprint.update(
llm_settings=llm_settings,
remove_vector_database=True,
)
Once you are satisfied with the blueprint, in the UI you can perform an LLM blueprint comparison or deploy it from the playground.
Add an external vector database¶
Now add an external vector database from a deployment.
Note that creating the deployment is outside the scope of this example.
To continue with this portion of the example, first create an external vector database deployment, validate and deploy it, and update the external_vdb_id
value here.
You can create an external vector database and deploy it by first following this guide
Once you have done that, simply update your external vector database ID here to use it.
external_vdb_id = "<ADD_VALUE_HERE>"
# Now create a vector database
external_model_vdb = VectorDatabase.get(
vector_database_id=external_vdb_id
)
assert external_model_vdb.execution_status == "COMPLETED"
Add an external model LLM¶
Now add an external LLM blueprint from a deployment. Because creating the deployment is outside the scope of this example, it instead assumes a previously deployed text generation model.
To continue with this portion of the example, first create an external LLM deployment, validate it, and update the external_llm_validation_id
value here.
You can create an external LLM and validate it by first following this guide.
Once you have done that, update your external LLM validation ID here to use it.
external_llm_validation_id = "<ADD_VALUE_HERE>"
external_llm_validation = CustomModelLLMValidation.get(validation_id=external_llm_validation_id)
assert external_llm_validation.validation_status == "PASSED"
Now you can create an LLM blueprint using the custom model LLM and the custom model vector database.
custom_model_llm_blueprint = LLMBlueprint.create(
playground=playground,
name="custom model LLM with custom model vdb",
llm="custom-model",
llm_settings={
"validation_id": external_llm_validation.id,
# NOTE - update this value based on the context size in tokens
# of the external LLM that you are deploying
"external_llm_context_size": 4096
},
vector_database=external_model_vdb,
)
Create a comparison chat¶
Finally create a comparison_chat
and associate it with a playground.
A comparison chat, analogous to the chat for a chat prompt, is the entity that groups together a set of comparison prompt messages.
It can be thought of as a conversation with one or more LLM blueprints.
comparison_chat = ComparisonChat.create(
name="Resume review comparison chat with a pirate",
playground=playground
)
Now compare the LLM blueprints.
The response of the custom_model_llm_blueprint
as it is setup in the documentation will not be related to the pirate resume and that is expected. Feel free to modify the external VDB creation and/or external LLM to suit your use case.
comparison_prompt = ComparisonPrompt.create(
llm_blueprints=[llm_blueprint, new_llm_blueprint, custom_model_llm_blueprint],
text="Summarize the best resume",
comparison_chat=comparison_chat,
wait_for_completion=True,
)
print([result.result_text for result in comparison_prompt.results])