Implementation:Hpcaitech ColossalAI CustomRetriever

Knowledge Sources	ColossalAI
Domains	RAG, Information_Retrieval
Last Updated	2026-02-09 00:00 GMT

Overview

Multi-source vector retriever with incremental indexing and result caching, provided by ColossalQA.

Description

CustomRetriever manages multiple Chroma vector stores (one per document source), supports incremental document addition with SQL-backed record management, and provides top-k similarity search with optional score thresholding and result caching.

Usage

Create a CustomRetriever, add documents via add_documents(), then use as the retriever in a RetrievalQA chain.

Code Reference

Source Location

• Repository: ColossalAI
• File: applications/ColossalQA/colossalqa/retriever.py
• Lines: 22-179

Signature

class CustomRetriever(BaseRetriever):
    vector_stores: Dict[str, VectorStore] = {}
    sql_index_database: Dict[str, str] = {}
    record_managers: Dict[str, SQLRecordManager] = {}
    k: int = 3

    def add_documents(
        self,
        docs: Dict[str, Document] = [],
        cleanup: str = "incremental",
        mode: str = "by_source",
        embedding: Embeddings = None,
    ) -> None:
        """Add documents to vector stores with incremental indexing."""

    def _get_relevant_documents(
        self,
        query: str,
        run_manager=None,
        score_threshold=None,
        return_scores=False,
    ) -> List[Document]:
        """Retrieve top-k relevant documents across all stores."""

Import

from colossalqa.retriever import CustomRetriever
from langchain.embeddings import HuggingFaceEmbeddings

I/O Contract

Inputs

Name	Type	Required	Description
docs	List[Document]	Yes	Documents to index
embedding	Embeddings	Yes	Embedding model (e.g., HuggingFaceEmbeddings)
query	str	Yes	User query for retrieval
k	int	No	Number of results to return (default: 3)

Outputs

Name	Type	Description
documents	List[Document]	Top-k most relevant document chunks

Usage Examples

from colossalqa.retriever import CustomRetriever
from langchain.embeddings import HuggingFaceEmbeddings

embedding = HuggingFaceEmbeddings(model_name="moka-ai/m3e-base")
retriever = CustomRetriever(k=5)
retriever.add_documents(
    docs=chunked_documents,
    embedding=embedding,
    cleanup="incremental",
)

results = retriever._get_relevant_documents("What is ColossalAI?")

Related Pages

Implements Principle

• Principle:Hpcaitech_ColossalAI_Vector_Embedding_Retrieval

Environment and Heuristic Links

• Environment:Hpcaitech_ColossalAI_ColossalQA_RAG_Environment