Coding an AutoAI RAG experiment with a Chroma vector store

Storing vectorized content in a Chroma database

When you set up your AutoAI RAG experiment and don't specify a connection to a vector store, the vectorized content is saved to the default, in-memory Chroma database. The content does not persist beyond the experiment, so this is not a viable production method for deploying a RAG pattern. However, it provides a fastpath for creating a RAG pattern.

The following sections expand on the annotated sample code provided with the Automating RAG pattern with Chroma database notebook.

The notebook uses the watsonx.ai Python client library (version 1.1.11 or later).

Follow these steps to code an AutoAI RAG experiment for your use case.

1. Prepare the prerequisites for preparing data and set up the experiment
2. Configure the RAG experiment
3. Run the experiment
4. Review the patterns and select the best one

Step 1: Prepare the prerequisites for preparing data and set up the experiment

Prepare the prerequisites for the experiment.

• Install and import the required modules and dependencies. For example:

  pip install 'ibm-watsonx-ai[rag]>=1.1.11'
  pip install langchain-community==0.2.4
  

  • Add task credentials. See Adding task credentials.
  • Add the watsonx.ai Runtime service. See Creating services.
  • Enter your API key. See Managing the user API key.

• Use these to initialize the client. For example:

  from ibm_watsonx_ai import APIClient, Credentials
  
  credentials = Credentials(
                  url = "https://us-south.ml.cloud.mydomain.com",
                  api_key = "***********"
                  )
  
  client = APIClient(credentials)
  

• Create a project or space for your work. See Creating a project or Creating a space.

• Get the ID for the project or space. See Finding the project ID.

• Set a default project or space:

client.set.default_project("<Project ID>")

client.set.default_space("<Space GUID>")

• Prepare the grounding documents
• Prepare the evaluation data

from ibm_watsonx_ai.helpers import DataConnection, S3Location

conn_meta_props= {
    client.connections.ConfigurationMetaNames.NAME: f"Connection to input data - {datasource_name} ",
    client.connections.ConfigurationMetaNames.DATASOURCE_TYPE: client.connections.get_datasource_type_id_by_name(datasource_name),
    client.connections.ConfigurationMetaNames.DESCRIPTION: "ibm-watsonx-ai SDK documentation",
    client.connections.ConfigurationMetaNames.PROPERTIES: {
        'bucket': <BUCKET_NAME>,
        'access_key': <ACCESS_KEY>,
        'secret_key': <SECRET_ACCESS_KEY>,
        'iam_url': 'https://iam.cloud.ibm.com/identity/token',
        'url': <ENDPOINT_URL>
    }
}

conn_details = client.connections.create(meta_props=conn_meta_props)
cos_connection_id = client.connections.get_id(conn_details)

input_data_references = [DataConnection(
    connection_asset_id=cos_connection_id,
    location=S3Location(
        bucket=<BACKET_NAME>,
        path=<BACKET_PATH>
    )
)]

import os, wget
from ibm_watsonx_ai.helpers import DataConnection

input_data_filename = "core_api.html"
input_data_path = f"https://ibm.github.io/watsonx-ai-python-sdk/{input_data_filename}"

if not os.path.isfile(input_data_filename): 
    wget.download(input_data_path, out=input_data_filename)
    
asset_details = client.data_assets.create(input_data_filename, input_data_filename)
asset_id = client.data_assets.get_id(asset_details)
asset_id

input_data_references = [DataConnection(data_asset_id=asset_id)]

[
    {
        "question": "What is the purpose of get_token()?",
        "correct_answer": "get_token() is used to retrieve an authentication token for secure API access.",
        "correct_answer_document_ids": [
            "core_api.html"
        ]
    },
    {
        "question": "How does the delete_model() function operate?",
        "correct_answer": "delete_model() method allows users to delete models they've created or managed.",
        "correct_answer_document_ids": [
            "core_api.html"
        ]
    }
]

import os, wget
from ibm_watsonx_ai.helpers import DataConnection

test_data_filename = "benchmarking_data_core_api.json"
test_data_path = f"https://github.com/IBM/watson-machine-learning-samples/raw/master/cloud/data/autoai_rag/{test_data_filename}"

if not os.path.isfile(test_data_filename): 
    wget.download(test_data_path, out=test_data_filename)

test_asset_details = client.data_assets.create(name=test_data_filename, file_path=test_data_filename)
test_asset_id = client.data_assets.get_id(test_asset_details)

test_data_references = [DataConnection(data_asset_id=test_asset_id)]

Step 2: Configure the RAG optimizer

The rag_optimizer object provides a set of methods for working with the AutoAI RAG experiment. In this step, enter the details to define the experiment. These are the available configuration options:

This sample code shows the configuration options for running the experiment with the ibm-watsonx-ai SDK documentation:

from ibm_watsonx_ai.experiment import AutoAI

experiment = AutoAI(credentials, project_id=project_id)

rag_optimizer = experiment.rag_optimizer(
    name='DEMO - AutoAI RAG ibm-watsonx-ai SDK documentation',
    description="AutoAI RAG experiment grounded with the ibm-watsonx-ai SDK documentation",
    max_number_of_rag_patterns=5,
    optimization_metrics=[AutoAI.RAGMetrics.ANSWER_CORRECTNESS]
)

rag_optimizer = experiment.rag_optimizer(
    name='DEMO - AutoAI RAG ibm-watsonx-ai SDK documentation',
    description="AutoAI RAG experiment grounded with the ibm-watsonx-ai SDK documentation",
    embedding_models=["ibm/slate-125m-english-rtrvr"],
    foundation_models=["ibm/granite-13b-chat-v2","mistralai/mixtral-8x7b-instruct-v01"],
    max_number_of_rag_patterns=5,
    optimization_metrics=[AutoAI.RAGMetrics.ANSWER_CORRECTNESS]
)

Step 3: Run the experiment

Run the optimizer to create the RAG patterns using the specified configuration options. In this code sample for running a Chroma experiment, the task is run in interactive mode. You can run the task in the background by changing the background_mode to True.

run_details = rag_optimizer.run(
    input_data_references=input_data_references,
    test_data_references=test_data_references,
    background_mode=False
)

Step 4: Review the patterns and select the best one

After the AutoAI RAG experiment completes successfully, you can review the patterns. Use the summary method to list completed patterns and evaluation metrics information in the form of a Pandas DataFrame so you can review the patterns, ranked according to performance against the optimized metric.

summary = rag_optimizer.summary()
summary

For example, pattern results display like this:

from langchain_community.document_loaders import WebBaseLoader

best_pattern = rag_optimizer.get_pattern()

urls = [
    "https://ibm.github.io/watsonx-ai-python-sdk/core_api.html",
    "https://ibm.github.io/watsonx-ai-python-sdk/fm_embeddings.html",
]
docs_list = WebBaseLoader(urls).load()
doc_splits = best_pattern.chunker.split_documents(docs_list)
best_pattern.indexing_function(doc_splits)

payload = {
    client.deployments.ScoringMetaNames.INPUT_DATA: [
        {
            "values": ["How to use new approach of providing credentials to APIClient?"],
        }
    ]
}

best_pattern.query(payload)

According to the document, the new approach to provide credentials to APIClient is by using the Credentials class. Here's an example:


from ibm_watsonx_ai import APIClient
from ibm_watsonx_ai import Credentials

credentials = Credentials(
                   url = "https://us-south.ml.cloud.ibm.com",
                   token = "***********",
                  )

client = APIClient(credentials)


This replaces the old approach of passing a dictionary with credentials to the APIClient constructor.

Reviewing experiment results in Cloud Object Storage

If the final status of the experiment is failed or error, use rag_optimizer.get_logs() or refer to experiment results to understand what went wrong. Experiment results and logs are stored in the default Cloud Object Storage instance linked to your account. By default, results are saved in the default_autoai_rag_out directory.

Results are organized by pattern. For example:

|-- Pattern1
|      | -- evaluation_results.json
|      | -- indexing_inference_notebook.ipynb (Chroma)
|-- Pattern2
|    ...
|-- training_status.json

Each pattern contains these results:

• The evaluation_results.json file contains evaluation results for each benchmark question.
• The indexing_inference_notebook.ipynb contains the python code for building vector database index as well as building retrieval and generation function. The notebook introduces commands for retrieving data, chunking, and embeddings creation as well as for retrieving chunks, building prompts and generating answers.

rag_optimizer.get_inference_notebook(pattern_name='Pattern3')