Coding an AutoAI RAG experiment with a Milvus vector store

Grounding documents

Prepare and connect to the grounding documents you use to run the RAG experiment. For details, see Getting and preparing data in a project.

• Supported formats: PDF, HTML, DOCX, plain text
• Connect to data in a Cloud Object Storage bucket, a folder in a bucket, or specify up to 20 files.
• AutoAI samples the documentations for running the experiment

For example, to create a data connection when documents are stored in a Cloud Object Storage bucket:

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

The following example demonstrates how to use the data asset created in the project (or promoted to the space):

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)

input_data_references = [DataConnection(data_asset_id=asset_id)]

Evaluation data

Evaluation data must be in JSON format with a fixed schema with these fields: question, correct_answer, correct_answer_document_ids

For example:

[
    {
        "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"
        ]
    }
]

To prepare the evaluation data:

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

Connecting to a Milvus vector database

This code snippet demonstrates how to connect to a Milvus vector database.

from ibm_watsonx_ai.helpers import DataConnection

milvus_data_source_type_id = client.connections.get_datasource_type_uid_by_name("milvus")
details = client.connections.create(
    {
        client.connections.ConfigurationMetaNames.NAME: "Milvus Connection",
        client.connections.ConfigurationMetaNames.DATASOURCE_TYPE: milvus_data_source_type_id,
        client.connections.ConfigurationMetaNames.PROPERTIES: {
            "host": <PASTE MILVUS HOST HERE>,
            "port": <PASTE MILVUS PORT HERE>,
            "username": <PASTE MILVUS USERNAME HERE>,
            "password": <PASTE MILVUS PASSWORD HERE>,
            "ssl": True,
        },
    }
)

milvus_connection_id = client.connections.get_id(details)
vector_store_references = [DataConnection(connection_asset_id=milvus_connection_id)]

Select a pattern to test locally

The next step is to select a pattern and test it locally.

best_pattern = rag_optimizer.get_pattern()

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

resp = best_pattern.query(payload)
print(resp["predictions"][0]["values"][0][0])

Model's response:

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.

Testing the deployed pattern

This code sample demonstrates how to test the deployed solution. Enter test questions in the payload, by using the following format:

questions = ["How to use new approach of providing credentials to APIClient?"]

payload = {
    client.deployments.ScoringMetaNames.INPUT_DATA: [
        {
            "values": questions,
            "access_token": client.service_instance._get_token()
        }
    ]
}

resp = client.deployments.score(deployment_id, payload)
print(resp["predictions"][0]["values"][0][0])

Model's response:

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.

Get inference and indexing notebooks

To download the specified inference notebook from Service, use get_inference_notebook(). If you leave pattern_name empty, the method downloads the notebook for the highest-ranked pattern.

rag_optimizer.get_inference_notebook(pattern_name='Pattern3')