基于 FastAPI 的 AI 知识库 RAG 流程基础功能实现

当然也提供了 websocket 与模型对话应答。

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

# 代码片段
# 把鉴权删除，替换模型与提示词即可。
@chat_ws_router.websocket("/{token}/{chat_id}")
async def chat(websocket: WebSocket, token: str, chat_id: str, user: dict = Depends(check_user)):
    """对话"""

    await websocket.accept()

    if not user:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        print("鉴权验证失败 ws 关闭...")
        raise CustomException(status_code=403, detail="鉴权验证失败...", custom_code=10005)

    try:
        while True:
            data = await websocket.receive_text()
            # await websocket.send_text(f"chat_id: {chat_id} 用户：{user} token: {token} 消息：{data}")

            # 测试代码
            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()

多轮对话，流式响应

# /ExileChat/test/test_ai/test_llm_engine.py

import asyncio
from utils.ai.llm_engine import LLMEngine
from api_key import api_key

async def main():
    """main"""

    new_engine = LLMEngine(model_name='azure_open_ai', api_key=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())

一次对话，阻塞响应

你可以通过修改 stream 改变响应性方式。

# /ExileChat/test/test_ai/test_llm_engine_only.py

import asyncio
from utils.ai.llm_engine import LLMEngine
from api_key import api_key

async def main():
    """main"""

    new_engine = LLMEngine(model_name='azure_open_ai', api_key=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
from api_key import api_key

async def main():
    """main"""

    llm_engine = LLMEngine(model_name='azure_open_ai', api_key=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 段落并且向量化构造知识库数据。你完全可以不需要理会我项目的表结构和业务逻辑，你得到最终向量化的数据结构，结合自身的业务即可。

# /ExileChat/test/test_ai/test_document_vector.py

import asyncio

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

"""
文档向量
1.接收已经处理好的文档字符串。

    通过 `from utils.ai.document_chunk import DocumentChunk` 文档规整生成文档字符串，当然你可以使用其他更好方式生成它。

2.接收大模型引擎 `LLMEngine` 实例，例如以下使用了 `azure_open_ai` 作为模型。
    
    llm_engine = LLMEngine(model_name='azure_open_ai', api_key=api_key, is_debug=is_debug)
    
3.使用 `LLMEngine` 结合 `prompt` 文档生成 `chunks`，同时 `self.chunks` 也会被赋值，如果你需要自定义可以修改它 `self.chunks=[...]`

    dv = DocumentVector(...)
    dv.gen_chunks()
    
    *.得到例如以下结构的数据，自定义时需要按照以下数据结构。
    [
        {
            "index": 1,
            "chunk": "段落 1"
        },
        {
            "index": 2,
            "chunk": "段落 2"
        }
        ...
    ]
    
4.使用 `LLMEngine` 结合 `prompt` 对段落生成 `QA`，同时 `self.qa` 也会被赋值，如果你需要自定义可以修改它 `self.qa=[...]`

    dv.gen_qa()
    
    *.得到例如以下结构的数据，自定义时需要按照以下数据结构。
    [
        {
            "Q": "什么是 Python 中的协程？",
            "A": "Python 中的协程是一种用于实现异步编程的机制，它允许程序在等待操作完成时挂起执行，从而可以执行其他任务。",
            "chunks": [
                "Python 中的协程是一种用于实现异步编程的机制，它允许程序在等待操作完成时挂起执行，从而可以执行其他任务。"
            ]
        },
        {
            "Q": "协程在 Python 中的作用是什么？",
            "A": "协程是 Python 语言中的一个特性，用于编写更高效和响应更快的代码。",
            "chunks": [
                "协程是 Python 语言中的一个特性，用于编写更高效和响应更快的代码。"
            ]
        },
        ...
    ]
    
5.生成 `QA` 向量。

    dv.gen_qa_vector()
    
    *.最终会得到一个组装好的对象列表，例如以下结构的数据。
    [
        {
            'able_id': 0,
            'document_id': 0,
            'answer': '1+1 等于 2',
            'question': '1+1 等于几？',
            'chunks': ["1+1=2", "..."],
            'answer_embedding': [0.02715362422168255, -0.010939446277916431, 0.006153851747512817, -0.009505090303719044, ...],
            'question_embedding': [0.02715362422168255, -0.010939446277916431, 0.006153851747512817, -0.009505090303719044, ...],
        },
        ...
    ]
"""
if __name__ == '__main__':
    is_debug = True
    # 为何测试问答的精准，我使用了我的毕业论文作为测试，hhh。
    # 你可以把下面 document_content = "1+1=2" 注释打开来进行测试。
    # file_path = "./test/基于 Python+Vue 自动化测试平台的设计与实现.docx"
    file_path = "./test/测试文档分段.docx"
    dc = DocumentChunk(image_base_path="./test/test_ai", is_debug=is_debug)
    document_content = dc.process_file(file_path)

    # document_content = "1+1=2"

    llm_engine = LLMEngine(model_name='azure_open_ai', api_key=api_key, is_debug=is_debug)
    dv = DocumentVector(document_content=document_content, llm_engine=llm_engine, is_debug=is_debug)
    asyncio.run(dv.gen_chunks())
    asyncio.run(dv.gen_qa())
    asyncio.run(dv.gen_qa_vector())

向量化应答

问题向量化通过数据库检索得到向量距离最相近的答案，结合大模型生成应答。以下是完整的检索与生成流程示例。

1. 接收用户提问。
2. 将问题转换为向量（Embedding）。
3. 在向量数据库中检索相似度最高的 Top-K 个知识片段。
4. 将检索到的上下文与用户问题组合成 Prompt。
5. 调用大模型生成最终回答。

# /ExileChat/utils/ai/retrieval.py

import asyncio
from typing import List, Dict
from utils.ai.llm_engine import LLMEngine

async def retrieve_and_answer(question: str, top_k: int = 3):
    """
    检索增强生成 (RAG) 核心函数
    :param question: 用户问题
    :param top_k: 检索数量
    :return: 生成的回答
    """
    # 1. 初始化引擎
    llm_engine = LLMEngine(model_name='azure_open_ai', api_key=api_key)
    
    # 2. 向量化问题
    question_embedding = await llm_engine.embedding(text=question)
    
    # 3. 检索相似向量 (此处模拟数据库查询)
    # 实际场景中应调用 VectorDB.search(query_vector=question_embedding, limit=top_k)
    retrieved_docs = await search_database_simulated(question_embedding, top_k)
    
    # 4. 构建上下文
    context = "\n".join([doc['answer'] for doc in retrieved_docs])
    prompt = f"""基于以下知识库内容回答问题：
    知识库：
    {context}
    
    问题：{question}
    
    如果知识库中没有相关信息，请说明无法回答。
    """
    
    # 5. 生成回答
    response = await llm_engine.chat_only(prompt=prompt, input=question)
    return response

async def search_database_simulated(embedding: List[float], limit: int):
    """模拟数据库检索"""
    # 这里应替换为真实的向量数据库查询逻辑
    # 例如 Milvus, Chroma, Pinecone 等
    return []

if __name__ == "__main__":
    asyncio.run(retrieve_and_answer("1+1 等于多少"))

通过以上步骤，即可完成从文档解析、向量化存储到检索问答的完整 RAG 流程闭环。