RAG+Vllm的知识库部署检索方案

克里斯蒂亚诺更新2026-04-04 8:08

轻量级快速方案（推荐起步）

架构组件

组件	选型建议	说明
嵌入模型	`BAAI/bge-large-zh-v1.5` 或 `Qwen3-Embedding-0.6B`	中文支持好
向量数据库	Chroma（轻量）或 FAISS（高性能）	本地部署，无需额外服务
文档切分	RecursiveCharacterTextSplitter	按段落/句子智能切分

核心代码实现

python

复制代码

# 1. 知识库准备（一次性）
from langchain.document_loaders import TextLoader, DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma

# 加载控规文档
loader = DirectoryLoader("./控规文档/", glob="*.txt", loader_cls=TextLoader)
documents = loader.load()

# 智能切分（适合控规文件特点）
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=800,          # 每块800字符
    chunk_overlap=200,       # 重叠200，保留上下文
    separators=["\n\n", "\n", "。", "；", "，", " "]
)
chunks = text_splitter.split_documents(documents)

# 中文嵌入模型
embeddings = HuggingFaceEmbeddings(
    model_name="BAAI/bge-large-zh-v1.5",
    model_kwargs={"device": "cuda"}
)

# 存入向量库
vectorstore = Chroma.from_documents(chunks, embeddings, persist_directory="./chroma_db")

python

复制代码

# 2. 查询接口（每次请求）
def rag_query(user_query: str, system_prompt: str) -> str:
    # 检索相关文档片段
    retriever = vectorstore.as_retriever(search_kwargs={"k": 5})
    relevant_docs = retriever.get_relevant_documents(user_query)
    
    # 构建上下文
    context = "\n\n---\n\n".join([doc.page_content for doc in relevant_docs])
    
    # 增强后的提示词
    augmented_prompt = f"""【参考知识库内容】
{context}

【用户问题】
{user_query}

【要求】
请严格依据上述知识库内容进行回答。如果知识库中没有相关信息，请明确说明。"""

    # 调用 vLLM
    import requests
    response = requests.post(
        "http://192.168.3.215:7899/v1/chat/completions",
        headers={"Authorization": "Bearer sk-*****"},
        json={
            "model": "Qwen/Qwen3-8B",
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": augmented_prompt}
            ],
            "temperature": 0.0,
            "max_tokens": 8192
        }
    )
    return response.json()["choices"][0]["message"]["content"]