Using vector store knowledge bases in AutoAI RAG experiments

A vector store knowledge base is a pre-indexed vector store that contains document embeddings. Connecting vector stores to AutoAI RAG experiments speeds up retrieval and eliminates the need to reprocess documents. You can add a vector store knowledge base by using the UI or the AutoAI SDK.

Important:

Milvus and Elasticsearch vector stores do not support collections that contain only sparse vectors. You can use collections with dense vectors only or collections where each record includes both a dense vector and a sparse vector.

Adding a vector store as a knowledge base in the UI

You can add a vector store knowledge base during experiment setup by selecting a vector index from your project.

From the Select from project dropdown menu, click Vector index.
Select a connection to a Milvus, watsonx.data Milvus, or Elasticsearch vector index.
Define the details for each connection.
1. If a database is defined in your connection details, the database is selected automatically. If no database is defined, default is selected.
2. Select an index from the vector store.
3. Select an embedding model. If you're using a Milvus vector index connection, you can add only a dense embedding model from the setup window. To add a sparse embedding model, you have to use the API.
4. Add a detailed index description that explains what the data in the index represents and the type of information it contains. This is required to ensure that the data can be correctly processed and retrieved.
5. Specify the field mapping details. If you selected an AutoAI index that was created in a previous RAG experiment run, the field mapping is automatically generated.
To add more collections, repeat the steps.

Configuring a vector store knowledge base in code

To use a vector store as a knowledge base by using the AutoAI SDK, provide the following information:

vector_store_kb1 = VectorStoreKnowledgeBase(
    name="product_documentation",
    description="Vector store containing product documentation",
    connection=DataConnection(connection_asset_id="your-milvus-connection-id-1"),  # Replace with your connection asset ID
    settings={
        "index_name": "product_docs_index",
        "fields_mapping": [
            {
                "role": KnowledgeBaseFieldRole.DOCUMENT_NAME,
                "field_name": "doc_id"
            },
            {
                "role": KnowledgeBaseFieldRole.TEXT,
                "field_name": "content"
            },
            {
                "role": KnowledgeBaseFieldRole.DENSE_VECTOR_EMBEDDINGS,
                "field_name": "vector_embeddings"
            }
        ],
        "embeddings": {
            "model_id": "ibm/slate-125m-english-rtrvr" 
        }
    }
)

Important:

The embedding model that you specify must match the one used to create the vector store.

Map fields

Use the fields_mapping parameter to map the roles in the AutoAI RAG system to the actual field names in your vector store:

Role	Description	Example Field Name
`dense_vector_embeddings`	Dense vector embeddings field	`vector`
`sparse_vector_embeddings`	Sparse vector embeddings field (for hybrid search)	`sparse_embeddings`
`document_name`	Document identifier field	`document_id`
`text`	Text content field	`text`
`start_index`	Chunk starting position field	`start_index`
`sequence_number`	Chunk sequence number field	`sequence_number`

Important:

Always verify that your field names match the schema in your vector store. Incorrect mappings cause retrieval errors.
If you use the window retrieval method, you must include the chunk_sequence_number role in your fields_mapping. Experiments without this mapping fail during initialization.

Use hybrid search

You can enable hybrid search for vector store knowledge bases. Hybrid search combines dense and sparse vector retrieval to improve retrieval performance.

Select one of these hybrid search strategies:

Weighted: combines dense and sparse results with a weighted average.
Reciprocal Rank Fusion (RRF): merges rankings from various sources into one list.

Select a sparse vector model:

For Elasticsearch: an ELSER model (for example, .elser_model_2_linux-x86_64) or BM25
For Milvus: BM25

Using multiple vector store knowledge bases

You can connect up to 20 collections from a vector store database in a single experiment. Here’s an example of using two:

from ibm_watsonx_ai.utils.autoai.knowledge_base import VectorStoreKnowledgeBase
from ibm_watsonx_ai.utils.autoai.enums import KnowledgeBaseFieldRole
from ibm_watsonx_ai.helpers.connections import DataConnection

# Define the first vector store knowledge base reference
vector_store_kb1 = VectorStoreKnowledgeBase(
    name="product_documentation",
    description="Vector store containing product documentation",
    connection=DataConnection(connection_asset_id="your-milvus-connection-id-1"),  # Replace with your connection asset ID
    settings={
        "index_name": "product_docs_index",
        "fields_mapping": [
            {
                "role": KnowledgeBaseFieldRole.DOCUMENT_NAME,
                "field_name": "doc_id"
            },
            {
                "role": KnowledgeBaseFieldRole.TEXT,
                "field_name": "content"
            },
            {
                "role": KnowledgeBaseFieldRole.DENSE_VECTOR_EMBEDDINGS,
                "field_name": "vector_embeddings"
            }
        ],
        "embeddings": {
            "model_id": "ibm/slate-125m-english-rtrvr" 
        }
    }
)


# Define the second vector store knowledge base reference
vector_store_kb2 = VectorStoreKnowledgeBase(
    name="customer_support",
    description="Vector store containing customer support data",
    connection=DataConnection(connection_asset_id="your-milvus-connection-id-2"),  # Replace with your connection asset ID
    settings={
        "index_name": "support_data_index",
        "fields_mapping": [
            {
                "role": KnowledgeBaseFieldRole.DOCUMENT_NAME,
                "field_name": "support_id"
            },
            {
                "role": KnowledgeBaseFieldRole.TEXT,
                "field_name": "content"
            },
            {
                "role": KnowledgeBaseFieldRole.DENSE_VECTOR_EMBEDDINGS,
                "field_name": "vector_embeddings"
            },
            {
                "role": KnowledgeBaseFieldRole.SPARSE_VECTOR_EMBEDDINGS,
                "field_name": "sparse_embeddings"
            }
        ],
        "embeddings": {
            "model_id": "ibm/slate-125m-english-rtrvr"
        },
        # Optional: Configure hybrid search with sparse vectors
        "hybrid_ranker": {
            "sparse_vectors": {
                "model_id": "BM25" 
            }
        }
    }
)

Run the RAG optimizer with multiple knowledge base references:

run_details = rag_optimizer.run(
    knowledge_base_references=[vector_store_kb1, vector_store_kb2],
    test_data_references=[test_data_reference]
)

Limitations

The following features are not supported for vector store knowledge bases:

Chroma vector stores
Test data generation