# How DB-GPT's RAG Framework Handles Vector Storage and Retrieval for Knowledge Bases

> Discover how DB-GPT's RAG framework efficiently handles vector storage and retrieval for knowledge bases using a pluggable architecture supporting Chroma PGVector and Milvus.

- Repository: [eosphoros/DB-GPT](https://github.com/eosphoros-ai/db-gpt)
- Tags: how-to-guide
- Published: 2026-02-23

---

**DB-GPT's RAG framework manages vector storage and retrieval through a pluggable architecture that abstracts vector databases behind a unified API, using `VectorStoreConfig` for configuration, `StorageManager` for connection handling, and `EmbeddingRetriever` for semantic search across Chroma, PGVector, and Milvus backends.**

DB-GPT is an open-source AI database framework that implements Retrieval-Augmented Generation (RAG) for knowledge base interactions. The system's approach to **vector storage and retrieval** follows a modular design pattern that separates embedding logic from storage backends, enabling seamless swapping of vector databases without changing application code.

## Core Architecture Components

The framework treats a knowledge base as a collection of text chunks embedded into dense vectors and persisted in a **vector store**. This workflow relies on three primary abstractions defined in the `dbgpt` package structure.

### Configuration Layer with VectorStoreConfig

The base configuration class `VectorStoreConfig` in [`packages/dbgpt-core/src/dbgpt/storage/vector_store/base.py`](https://github.com/eosphoros-ai/DB-GPT/blob/main/packages/dbgpt-core/src/dbgpt/storage/vector_store/base.py) (lines 68-106) defines the contract for all vector store implementations. It holds the vector database type, connection parameters, and the embedding function. Concrete implementations like `ChromaVectorConfig` or `PGVectorConfig` subclass this base to provide database-specific initialization logic.

### Storage Management and Caching

The `StorageManager` class in [`packages/dbgpt-serve/src/dbgpt_serve/rag/storage_manager.py`](https://github.com/eosphoros-ai/DB-GPT/blob/main/packages/dbgpt-serve/src/dbgpt_serve/rag/storage_manager.py) (lines 89-107) acts as a factory and cache for vector store instances. When a RAG request arrives, it looks up the configured vector store type, builds the concrete store class via `VectorStoreConfig.create_store()`, and maintains a per-collection cache (`self._store_cache`) guarded by locks to prevent repeated connection overhead.

### Retrieval Interface via EmbeddingRetriever

The `EmbeddingRetriever` class in [`packages/dbgpt-core/src/dbgpt/rag/retriever/embedding.py`](https://github.com/eosphoros-ai/DB-GPT/blob/main/packages/dbgpt-core/src/dbgpt/rag/retriever/embedding.py) (lines 95-99) implements the `BaseRetriever` interface. It delegates similarity searches to the underlying vector store's `similar_search()` or `similar_search_with_scores()` methods, returning `Chunk` objects ordered by relevance. This retriever supports optional query rewriting and re-ranking before returning results.

## Vector Store Implementations

DB-GPT provides concrete implementations for multiple vector databases, all subclassing `VectorStoreBase`. Each implementation translates generic API calls to native database clients.

- **ChromaStore**: Located in [`packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/chroma_store.py`](https://github.com/eosphoros-ai/DB-GPT/blob/main/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/chroma_store.py) (lines 58-66), this implementation builds collections, upserts embeddings, and executes cosine-based similarity queries.
- **PGVectorStore**: Found in [`packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/pgvector_store.py`](https://github.com/eosphoros-ai/DB-GPT/blob/main/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/pgvector_store.py), this adapter enables PostgreSQL-backed vector storage using the `pgvector` extension.
- **MilvusStore**: Supports distributed vector storage for enterprise-scale deployments.

The `VectorStoreConnector` in [`packages/dbgpt-serve/src/dbgpt_serve/rag/connector.py`](https://github.com/eosphoros-ai/DB-GPT/blob/main/packages/dbgpt-serve/src/dbgpt_serve/rag/connector.py) (lines 24-70) provides a thin wrapper that validates store types and exposes uniform methods including `load_document()`, `similar_search()`, and `delete_vector_name()`.

## End-to-End Retrieval Workflow

The complete vector storage and retrieval pipeline follows these steps:

1. **Application startup** reads the RAG storage configuration (e.g., `rag.storage.vector.type = "chroma"`).
2. **Query initiation** triggers `StorageManager.get_storage_connector()` with the collection name and storage type.
3. **Store instantiation** occurs via `create_store()`, utilizing automatic discovery through `VectorStoreConfig.__subclasses__()`.
4. **Semantic search** executes when `EmbeddingRetriever.retrieve()` calls the vector store's `similar_search()` method.
5. **Post-processing** applies optional score thresholds or reranking before returning chunks to the LLM.

This design ensures that adding a new vector database requires only implementing a new `VectorStoreBase` subclass and corresponding config class, with automatic registration via Python's subclass introspection.

## Practical Implementation Examples

### Creating a Chroma Vector Store and Loading Documents

```python
from dbgpt.storage.vector_store.connector import VectorStoreConnector
from dbgpt_ext.storage.vector_store.chroma_store import ChromaVectorConfig
from dbgpt.rag.embedding.embedding_factory import DefaultEmbeddingFactory
from dbgpt.core import Chunk

# Build the embedding function

embedding_factory = DefaultEmbeddingFactory()
embedding_fn = embedding_factory.create(model_name="text-embedding-ada-002")

# Configure Chroma

cfg = ChromaVectorConfig(name="my_collection", embedding_fn=embedding_fn)

# Connect to the vector store

connector = VectorStoreConnector(vector_store_type="Chroma", vector_store_config=cfg)

# Prepare chunks

chunks = [
    Chunk(content="DB-GPT is an LLM-augmented database assistant.", metadata={"source": "doc1"}),
    Chunk(content="It supports RAG with multiple vector back-ends.", metadata={"source": "doc2"}),
]

# Load into the store

ids = connector.load_document(chunks)
print("Loaded IDs:", ids)

```

The `VectorStoreConnector` initialization handles store creation through the connector's `__init__` method, as implemented in [`packages/dbgpt-serve/src/dbgpt_serve/rag/connector.py`](https://github.com/eosphoros-ai/DB-GPT/blob/main/packages/dbgpt-serve/src/dbgpt_serve/rag/connector.py).

### Retrieving Similar Chunks

```python
from dbgpt.rag.retriever.embedding import EmbeddingRetriever

# Reuse the connector from previous example

retriever = EmbeddingRetriever(index_store=connector, top_k=3)

# Perform semantic search

result_chunks = retriever.retrieve("What does DB-GPT do?")
for c in result_chunks:
    print("- ", c.content, "(score:", c.score, ")")

```

The `EmbeddingRetriever._retrieve()` method internally calls `self._index_store.similar_search()` to fetch relevant documents.

### Switching to PGVector Without Code Changes

```python
from dbgpt_ext.storage.vector_store.pgvector_store import PGVectorConfig

pg_cfg = PGVectorConfig(
    name="pg_collection",
    host="localhost",
    port=5432,
    user="postgres",
    password="******",
    embedding_fn=embedding_fn,
)

pg_connector = VectorStoreConnector(vector_store_type="PGVector", vector_store_config=pg_cfg)

# Use pg_connector identically to the Chroma connector

```

The `PGVectorConfig` class implements `create_store()` to instantiate the PostgreSQL-backed vector store, demonstrating the framework's storage-agnostic design.

## Summary

- **Pluggable backends** allow switching between Chroma, PGVector, and Milvus by changing configuration rather than application code, with automatic discovery via `VectorStoreConfig.__subclasses__()`.
- **Unified API** across all vector stores exposes consistent methods like `similar_search()`, `load_document()`, and `truncate()`, keeping RAG components storage-agnostic.
- **Connection caching** in `StorageManager` maintains a `_store_cache` dictionary to reuse client instances and minimize connection overhead for repeated queries.
- **Embedding-as-service** architecture injects embedding functions through configuration, enabling model reuse across different storage backends.
- **Modular retrieval** through `EmbeddingRetriever` separates similarity search logic from storage implementation, supporting query rewriting and reranking pipelines.

## Frequently Asked Questions

### What vector databases does DB-GPT support?

DB-GPT supports Chroma, PostgreSQL with PGVector, Milvus, and Weaviate through concrete implementations of `VectorStoreBase`. Each implementation resides in the `dbgpt_ext.storage.vector_store` package and translates generic API calls to database-specific operations. New backends require only subclassing `VectorStoreConfig` and `VectorStoreBase` to integrate automatically.

### How does DB-GPT cache vector store connections?

The `StorageManager` maintains a `_store_cache` dictionary that maps collection names to instantiated store objects, guarded by thread locks. When `get_storage_connector()` receives a request for an existing collection, it returns the cached instance rather than creating a new connection, significantly reducing latency for repeated RAG queries against the same knowledge base.

### Can I retrieve vectors with similarity scores?

Yes. The `VectorStoreBase` interface exposes both `similar_search()` for basic retrieval and `similar_search_with_scores()` for ranked results with distance metrics. The `EmbeddingRetriever` can utilize either method, and returned `Chunk` objects include score attributes that enable threshold filtering or reranking before passing context to the LLM.

### How does the framework handle different embedding models?

Embedding functions are injected through `VectorStoreConfig` and shared across store instances via the `Embeddings` interface. The `DefaultEmbeddingFactory` creates embedding models that get passed during vector store initialization, ensuring consistent vector representations for both document indexing and query encoding regardless of the underlying storage backend.