用 FastAPI 和大模型搭一个 RAG 知识库

WebSocket 实时对话

如果要做在线问答，websocket 比普通 HTTP 更合适，至少流式输出更自然。这里还带了鉴权，没通过就直接断开，不需要绕弯。

# /ExileChat/app/api/chat_ws/chat_ws.py

from fastapi import APIRouter, WebSocket, Depends
from starlette.websockets import WebSocketDisconnect

chat_ws_router = APIRouter()

@chat_ws_router.websocket("/{token}/{chat_id}")
async def chat(websocket: WebSocket, token: str, chat_id: str, user: dict = Depends(check_user)):
    """WebSocket 对话接口"""
    await websocket.accept()

    if not user:
        await websocket.close(code=1008)
        raise Exception("鉴权验证失败")

    try:
        while True:
            data = await websocket.receive_text()
            
            llm_engine = LLMEngine(model_name='azure_open_ai', api_key=api_key)
            llm_engine.system_prompt = "你是一名 Python 专家"
            response_generator = llm_engine.chat(input=data)

            async for chunk in response_generator:
                if isinstance(chunk, str):
                    await websocket.send_text(chunk)
                else:
                    print(type(chunk), chunk)

            await websocket.close()
            break
    except Exception as e:
        print(f"WebSocket 连接发生异常：{e}")
        await websocket.close()

多轮对话与流式响应

异步迭代是这里比较顺手的地方。模型返回 token 流时，前端可以边收边显示，不用等整句结束才看见结果。

# /ExileChat/test/test_ai/test_llm_engine.py

import asyncio
from utils.ai.llm_engine import LLMEngine

async def main():
    new_engine = LLMEngine(model_name='azure_open_ai', api_key="your_api_key")
    new_engine.system_prompt = "你是一名 Python 专家"
    response_generator = new_engine.chat(input="Python 是什么时候诞生的")

    async for chunk in response_generator:
        if isinstance(chunk, str):
            print(f"Received: {chunk}")
        else:
            print(f"Received chunk type: {type(chunk)}, value: {chunk}")

if __name__ == "__main__":
    asyncio.run(main())

一次对话，阻塞响应

如果不需要流式输出，chat_only 会简单很多。它牺牲的是实时性，换来的是调用方式更直接。

# /ExileChat/test/test_ai/test_llm_engine_only.py

import asyncio
from utils.ai.llm_engine import LLMEngine

async def main():
    new_engine = LLMEngine(model_name='azure_open_ai', api_key="your_api_key")
    generated_message = await new_engine.chat_only(prompt="你是强大的人工智能", input="你是谁？")
    print(generated_message)

if __name__ == "__main__":
    asyncio.run(main())

向量化处理

向量化就是把文本映射到高维空间里，后面靠相似度去找近邻。RAG 能不能稳，很多时候就看这里的 embedding 质量和切块方式是否靠谱。

import asyncio
from utils.ai.llm_engine import LLMEngine

async def main():
    llm_engine = LLMEngine(model_name='azure_open_ai', api_key="your_api_key")

    question = "1+1 等于多少"
    answer = "等于 2"
    question_embedding = await llm_engine.embedding(text=question)
    answer_embedding = await llm_engine.embedding(text=answer)
    return question_embedding, answer_embedding

if __name__ == "__main__":
    asyncio.run(main())

生成向量化数据结构

接下来是把文档内容变成知识库记录。做法不复杂：先分块，再抽 QA，再分别生成向量。难点不在代码量，而在每一步输出是否足够稳定，尤其是分块和 QA 生成，太松会把上下文切碎，太紧又会让召回效果发散。

核心流程

1. 接收文档字符串：通过 DocumentChunk 处理后的内容，或其他方式生成。
2. 初始化引擎：传入 LLMEngine 实例。
3. 生成 Chunk：用 Prompt 把长文本切成语义完整的段落。
4. 生成 QA 对：从段落中提取问题和答案。
5. 生成向量：对问题和答案分别做 Embedding。

# /ExileChat/test/test_ai/test_document_vector.py

import asyncio
from utils.ai.document_chunk import DocumentChunk
from utils.ai.document_vector import DocumentVector
from utils.ai.llm_engine import LLMEngine

async def main():
    is_debug = True
    file_path = "./data/测试文档分段.docx"
    
    # 1. 文档分块
    dc = DocumentChunk(image_base_path="./data/images", is_debug=is_debug)
    document_content = dc.process_file(file_path)

    # 2. 初始化引擎
    llm_engine = LLMEngine(model_name='azure_open_ai', api_key="your_api_key", is_debug=is_debug)
    
    # 3. 构建向量对象
    dv = DocumentVector(document_content=document_content, llm_engine=llm_engine, is_debug=is_debug)
    
    # 4. 执行生成任务
    await dv.gen_chunks()      # 生成段落
    await dv.gen_qa()          # 生成 QA 对
    await dv.gen_qa_vector()   # 生成向量

if __name__ == '__main__':
    asyncio.run(main())

数据结构说明

Chunks 结构：

[
    {"index": 1, "chunk": "段落 1"},
    {"index": 2, "chunk": "段落 2"}
]

QA 结构：

[
    {
        "Q": "什么是 Python 中的协程？",
        "A": "Python 中的协程是一种用于实现异步编程的机制...",
        "chunks": ["Python 中的协程是一种用于实现异步编程的机制..."]
    }
]

最终向量对象：

{
    'able_id': 0,
    'document_id': 0,
    'answer': '1+1 等于 2',
    'question': '1+1 等于几？',
    'chunks': ["1+1=2", "..."],
    'answer_embedding': [0.027..., -0.010..., ...],
    'question_embedding': [0.027..., -0.010..., ...]
}

向量化应答（RAG 检索）

真正用起来时，RAG 的重点不是'能不能查到'，而是'查回来的东西有没有用'。问题先向量化，再去向量库里做相似度搜索，拿到 Top-K 片段后组装 Prompt，让模型基于上下文回答。

检索逻辑

1. 问题嵌入：把用户输入的问题转换成向量。
2. 相似度搜索：在向量数据库（如 Milvus、Chroma、Faiss）中计算余弦相似度或欧氏距离。
3. Top-K 召回：取最近的 K 个文档片段作为上下文。
4. Prompt 组装：把原始问题和检索结果拼成新的 Prompt。
5. 模型生成：调用 LLM 根据上下文回答问题。

代码示例

# /ExileChat/utils/ai/rag_retriever.py

import asyncio
from typing import List, Dict
from utils.ai.llm_engine import LLMEngine
from database.vector_store import VectorStore  # 假设存在的向量存储类

class RagRetriever:
    def __init__(self, llm_engine: LLMEngine, vector_store: VectorStore):
        self.llm_engine = llm_engine
        self.vector_store = vector_store

    async def retrieve(self, query: str, top_k: int = 3) -> List[Dict]:
        """检索相关片段"""
        # 1. 问题向量化
        query_embedding = await self.llm_engine.embedding(text=query)
        
        # 2. 向量检索
        results = await self.vector_store.search(query_embedding, limit=top_k)
        return results

    async def generate_answer(self, query: str, context: List[str]) -> str:
        """基于上下文生成回答"""
        prompt_template = f"""
        请根据以下提供的背景信息回答问题。如果背景信息不足以回答问题，请直接说明。
        
        背景信息：
        {chr(10).join(context)}
        
        问题：{query}
        
        回答：
        """
        
        response = await self.llm_engine.chat_only(prompt=prompt_template, input="")
        return response

    async def process_query(self, query: str) -> str:
        """完整流程"""
        context_list = []
        results = await self.retrieve(query)
        
        for item in results:
            context_list.append(item['answer'])
            context_list.extend(item['chunks'])
        
        return await self.generate_answer(query, context_list)

# 使用示例
async def test_rag():
    llm_engine = LLMEngine(model_name='azure_open_ai', api_key="your_api_key")
    # 初始化向量库连接
    vector_store = VectorStore(collection_name="knowledge_base") 
    retriever = RagRetriever(llm_engine, vector_store)
    
    final_answer = await retriever.process_query("1+1 等于多少？")
    print(final_answer)

if __name__ == "__main__":
    asyncio.run(test_rag())

收尾

整条链路看下来，FastAPI 负责接口层，LLMEngine 负责模型接入，DocumentChunk 和 DocumentVector 负责把文档变成可检索数据，最后由向量库和检索器把答案拼回来。它不是一套特别花哨的方案，但足够把一个基础知识库跑起来。后面真要继续做，优先补的是 PDF 解析、向量库索引策略和检索结果的重排，这几块比再多包一层接口更有价值。