AI 知识库：基于 FastAPI 与 LLM 实现 RAG 流程基础功能

WebSocket 实时对话

系统提供了 websocket 接口用于模型对话应答，支持鉴权与流式输出。

# /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())

一次对话，阻塞响应

通过修改参数可切换为同步阻塞模式，适用于不需要流式输出的场景。

# /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())

向量化处理

向量化是将文本转换为高维向量空间中的点，以便进行相似度计算。

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 段落并构造知识库数据。你可以参考以下表结构和业务逻辑，得到最终向量化的数据结构后结合自身的业务需求进行调整。

核心流程

1. 接收文档字符串：通过 DocumentChunk 处理后的内容，或使用其他方式生成。
2. 初始化引擎：传入 LLMEngine 实例。
3. 生成 Chunk：利用 Prompt 将长文本切分为语义完整的段落。
4. 生成 QA 对：利用 Prompt 从段落中提取问题和答案。
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 流程的核心环节。当用户提问时，系统将问题向量化，在数据库中检索最相似的向量记录，获取上下文片段，最后由大模型生成回答。

检索逻辑

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())

总结

本指南详细阐述了构建 AI 知识库的技术路径。通过 FastAPI 提供异步接口，结合 LLM 引擎处理自然语言交互，利用向量化技术实现语义检索。开发者可根据实际业务调整文档解析策略、向量模型选择及数据库方案，以构建高性能的智能问答系统。