# Vector Database Customization
* [Available Vector Databases](#available-vector-databases)
* [Configuring Milvus with GPU Acceleration](#configuring-milvus-with-gpu-acceleration)
* [Configuring pgvector as the Vector Database](#configuring-pgvector-as-the-vector-database)
* [Configuring Support for an External Milvus or pgvector database](#configuring-support-for-an-external-milvus-or-pgvector-database)
* [Adding a New Vector Store](#adding-a-new-vector-store)
* [LlamaIndex Framework](#llamaindex-framework)
* [LangChain Framework](#langchain-framework)
## Available Vector Databases
By default, the Docker Compose files for the examples deploy Milvus as the vector database with CPU-only support.
You must install the NVIDIA Container Toolkit to use Milvus with GPU acceleration.
The available vector databases in the examples are shown in the following list:
- LlamaIndex: Milvus, pgvector
- LangChain: FAISS, Milvus, pgvector
The following customizations are common:
- Use Milvus with GPU acceleration.
- Use pgvector as an alternative to Milvus.
pgvector uses CPU only.
- Use your own vector database and prevent deploying a vector database with each RAG example.
## Configuring Milvus with GPU Acceleration
1. Edit the `RAG/examples/local_deploy/docker-compose-vectordb.yaml` file and make the following changes to the Milvus service.
- Change the image tag to include the `-gpu` suffix:
```yaml
milvus:
container_name: milvus-standalone
image: milvusdb/milvus:v2.4.5-gpu
...
```
- Add the GPU resource reservation:
```yaml
...
depends_on:
- "etcd"
- "minio"
deploy:
resources:
reservations:
devices:
- driver: nvidia
capabilities: ["gpu"]
device_ids: ['${VECTORSTORE_GPU_DEVICE_ID:-0}']
profiles: ["nemo-retriever", "milvus", ""]
```
1. Stop and start the containers:
```console
docker compose down
docker compose up -d --build
```
Note: when deploying milvus with `local-nim` you have to use `milvus` profile to deploy the vectorstore
```
docker compose --profile local-nim --profile milvus up -d --build
```
1. Optional: View the chain server logs to confirm the vector database is operational.
1. View the logs:
```console
docker logs -f chain-server
```
1. Upload a document to the knowledge base.
Refer to [Use Unstructured Documents as a Knowledge Base](./using-sample-web-application.md#use-unstructured-documents-as-a-knowledge-base) for more information.
1. Confirm the log output includes the vector database:
```output
INFO:RAG.src.chain_server.utils:Using milvus collection: nvidia_api_catalog
INFO:RAG.src.chain_server.utils:Vector store created and saved.
```
## Configuring pgvector as the Vector Database
1. Export the following environment variables in your terminal:
```console
export POSTGRES_PASSWORD=password
export POSTGRES_USER=postgres
export POSTGRES_DB=api
```
1. Edit the `docker-compose.yaml` file for the RAG example and set the following environment variables for the Chain Server:
```yaml
environment:
APP_VECTORSTORE_URL: "pgvector:5432"
APP_VECTORSTORE_NAME: "pgvector"
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-password}
POSTGRES_USER: ${POSTGRES_USER:-postgres}
POSTGRES_DB: ${POSTGRES_DB:-api}
...
```
1. Start the containers:
```console
docker compose --profile pgvector up -d --build
```
1. Optional: View the chain server logs to confirm the vector database is operational.
1. View the logs:
```console
docker logs -f chain-server
```
1. Upload a document to the knowledge base.
Refer to [Use Unstructured Documents as a Knowledge Base](./using-sample-web-application.md#use-unstructured-documents-as-a-knowledge-base) for more information.
1. Confirm the log output includes the vector database:
```output
INFO:RAG.src.chain_server.utils:Using PGVector collection: nvidia_api_catalog
INFO:RAG.src.chain_server.utils:Vector store created and saved.
```
To stop pgvector and the other containers run `docker compose --profile pgvector down`.
## Configuring Support for an External Milvus or pgvector database
1. Edit the `docker-compose.yaml` file for the RAG example and make the following edits.
- Remove or comment the `include` path to the `docker-compose-vectordb.yaml` file:
```yaml
include:
- path:
# - ../../local_deploy/docker-compose-vectordb.yaml
- ../../local_deploy/docker-compose-nim-ms.yaml
```
- To use an external Milvus server, specify the connection information:
```yaml
environment:
APP_VECTORSTORE_URL: "http://:19530"
APP_VECTORSTORE_NAME: "milvus"
...
```
- To use an external pgvector server, specify the connection information:
```yaml
environment:
APP_VECTORSTORE_URL: ":5432"
APP_VECTORSTORE_NAME: "pgvector"
...
```
Also export the `POSTGRES_PASSWORD`, `POSTGRES_USER`, and `POSTGRES_DB` environment variables in your terminal.
1. Start the containers:
```console
docker compose up -d --build
```
## Adding a New Vector Store
You can extend the code to add support for any vector store.
### LlamaIndex Framework
1. Navigate to the file `RAG/src/chain_server/utils.py` from the project's root directory. This file contains the utility functions used for vector store interactions.
2. Modify the `get_vector_index` function to handle your new vector store. Implement the logic for creating your vector store object within this function.
```python
def get_vector_index():
# existing code
elif config.vector_store.name == "chromadb":
import chromadb
from llama_index.vector_stores.chroma import ChromaVectorStore
if not collection_name:
collection_name = os.getenv('COLLECTION_NAME', "vector_db")
logger.info(f"Using Chroma collection: {collection_name}")
chroma_client = chromadb.EphemeralClient()
chroma_collection = chroma_client.create_collection(collection_name)
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
```
3. Modify the `get_docs_vectorstore_llamaindex` function to retrieve the list of files stored in your new vector store.
```python
def get_docs_vectorstore_llamaindex():
# existing code
elif settings.vector_store.name == "chromadb":
ref_doc_info = index.ref_doc_info
# iterate over all the document in vectorstore and return unique filename
for _ , ref_doc_value in ref_doc_info.items():
metadata = ref_doc_value.metadata
if 'filename' in metadata:
filename = metadata['filename']
decoded_filenames.append(filename)
decoded_filenames = list(set(decoded_filenames))
```
4. Update the `del_docs_vectorstore_llamaindex` function to handle document deletion in your new vector store.
```python
def del_docs_vectorstore_llamaindex(filenames: List[str]):
# existing code
elif settings.vector_store.name == "chromadb":
ref_doc_info = index.ref_doc_info
# Iterate over all the filenames and if filename present in metadata of doc delete it
for filename in filenames:
for ref_doc_id, doc_info in ref_doc_info.items():
if 'filename' in doc_info.metadata and doc_info.metadata['filename'] == filename:
index.delete_ref_doc(ref_doc_id, delete_from_docstore=True)
logger.info(f"Deleted documents with filenames {filename}")
```
5. In your custom `chains.py` implementation, import the functions from `utils.py`.
The sample `chains.py` in `RAG/examples/basic_rag/llamaindex` already imports the functions.
```python
from RAG.src.chain_server.utils import (
get_vector_index,
get_docs_vectorstore_llamaindex,
del_docs_vectorstore_llamaindex,
)
```
6. Update `RAG/src/chain_server/requirements.txt` with any additional package required for the vector store.
```text
# existing dependency
llama-index-vector-stores-chroma
```
7. Build and start the containers.
1. Navigate to the example directory.
```console
cd RAG/examples/basic_rag/llamaindex
```
1. Set the `APP_VECTORSTORE_NAME` environment variable for the `chain-server` microservice in your `docker-compose.yaml` file.
Set it to the name of your newly added vector store.
```yaml
APP_VECTORSTORE_NAME: "chromadb"
```
1. Build and deploy the microservice.
```console
docker compose up -d --build chain-server rag-playground
```
### LangChain Framework
1. Navigate to the file `RAG/src/chain_server/utils.py` in the project's root directory.
2. Modify the `create_vectorstore_langchain` function to handle your new vector store. Implement the logic for creating your vector store object within it.
```python
def create_vectorstore_langchain(document_embedder, collection_name: str = "") -> VectorStore:
# existing code
elif config.vector_store.name == "chromadb":
from langchain_chroma import Chroma
import chromadb
logger.info(f"Using Chroma collection: {collection_name}")
persistent_client = chromadb.PersistentClient()
vectorstore = Chroma(
client=persistent_client,
collection_name=collection_name,
embedding_function=document_embedder,
)
```
3. Update the `get_docs_vectorstore_langchain` function to retrieve a list of documents from your new vector store. Implement your retrieval logic within it.
```python
def get_docs_vectorstore_langchain(vectorstore: VectorStore) -> List[str]:
# Existing code
elif settings.vector_store.name == "chromadb":
chroma_data = vectorstore.get()
filenames = set([extract_filename(metadata) for metadata in chroma_data.get("metadatas", [])])
return filenames
```
4. Update the `del_docs_vectorstore_langchain` function to handle document deletion in your new vector store.
```python
def del_docs_vectorstore_langchain(vectorstore: VectorStore, filenames: List[str]) -> bool:
# Existing code
elif settings.vector_store.name == "chromadb":
chroma_data = vectorstore.get()
for filename in filenames:
ids_list = [chroma_data.get("ids")[idx] for idx, metadata in enumerate(chroma_data.get("metadatas", [])) if extract_filename(metadata) == filename]
vectorstore.delete(ids_list)
return True
```
5. In your custom `chains.py` implementation, import the preceding functions from `utils.py`.
The sample `chains.py` in `RAG/examples/basic_rag/langchain` already imports the functions.
```python
from RAG.src.chain_server.utils import (
create_vectorstore_langchain,
get_docs_vectorstore_langchain,
del_docs_vectorstore_langchain,
get_vectorstore
)
```
6. Update `RAG/src/chain_server/requirements.txt` with any additional package required for the vector store.
```text
# existing dependency
langchain-core==0.1.40 # Update this dependency as there is conflict with existing one
langchain-chroma
```
7. Build and start the containers.
1. Navigate to the example directory.
```console
cd RAG/examples/basic_rag/langchain
```
1. Set the `APP_VECTORSTORE_NAME` environment variable for the `chain-server` microservice in your `docker-compose.yaml` file.
Set it to the name of your newly added vector store.
```yaml
APP_VECTORSTORE_NAME: "chromadb"
```
1. Build and deploy the microservices.
```console
docker compose up -d --build chain-server rag-playground
```