Implementation:Hpcaitech ColossalAI RAG ChatBot

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

Overview

Orchestrator class for deploying a complete RAG chatbot with REST API and web UI, provided by ColossalQA.

Description

RAG_ChatBot manages the complete RAG lifecycle: embedding model setup, text splitter configuration, memory management, retriever setup, chain configuration, document loading/splitting/indexing, and query processing. It exposes run() for QA and load_doc_from_files() for dynamic document management.

Usage

Create with an LLM and configuration dict, then use with FastAPI server or Gradio UI.

Code Reference

Source Location

• Repository: ColossalAI
• File: applications/ColossalQA/examples/webui_demo/RAG_ChatBot.py
• Lines: 16-172

Signature

class RAG_ChatBot:
    def __init__(self, llm, rag_config) -> None:
        """
        Args:
            llm: Language model instance
            rag_config: Configuration dict with embed/model/splitter/retrieval/chain sections
        """

    def load_doc_from_files(self, files, data_name="default_kb", json_parse_args={}):
        """Load documents from file paths."""

    def run(self, user_input: str, memory) -> Tuple[str, Memory]:
        """Process a user query through the RAG pipeline."""

    def clear_docs(self):
        """Clear all loaded documents and reset retriever."""

Import

from RAG_ChatBot import RAG_ChatBot

I/O Contract

Inputs

Name	Type	Required	Description
llm	LLM	Yes	Language model instance
rag_config	Dict	Yes	Full configuration with embed, splitter, retrieval, chain sections
user_input	str	Yes	User question for run()
files	List	Yes	Document files for load_doc_from_files()

Outputs

Name	Type	Description
answer	str	Generated answer
memory	ConversationBufferWithSummary	Updated conversation memory

Usage Examples

from colossalqa.local.llm import ColossalAPI, ColossalLLM

# Initialize LLM
api = ColossalAPI("chatglm2", "/models/chatglm2-6b")
llm = ColossalLLM(n=1, api=api)

# Configure RAG
rag_config = {
    "embed": {"model_name": "moka-ai/m3e-base"},
    "splitter": {"chunk_size": 100, "chunk_overlap": 20},
    "retrieval": {"k": 3},
    "chain": {"chain_type": "stuff"},
}

# Create chatbot
chatbot = RAG_ChatBot(llm=llm, rag_config=rag_config)
chatbot.load_doc_from_files(["/docs/manual.pdf"])

# Query
answer, memory = chatbot.run("How to install?", chatbot.memory)

Related Pages

Implements Principle

• Principle:Hpcaitech_ColossalAI_RAG_Deployment

Environment and Heuristic Links

• Environment:Hpcaitech_ColossalAI_ColossalQA_RAG_Environment