chat-python/rag_server.py

# 请求的入口

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"

from fastapi import FastAPI, File, UploadFile, Form, Request, Response, WebSocket, WebSocketDisconnect, Depends, APIRouter, Body, HTTPException
from fastapi.responses import JSONResponse, FileResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from sse_starlette import EventSourceResponse
import uvicorn
from utils.get_logger import setup_logger
from rag.vector_db.milvus_vector import HybridRetriever
from response_info import generate_message, generate_response
from rag.db import MilvusOperate, MysqlOperate
from rag.file_process import ParseFile
from rag.documents_process import ProcessDocuments
from rag.task_registry import task_registry
from rag.chat_message import ChatRetrieverRag


logger = setup_logger(__name__)
app = FastAPI()

# 设置跨域
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/health")
async def health_check():
    return {"status": "healthy"}

@app.post("/upload_knowledge")
async def upload_file_to_db(file_json: dict):
    logger.info(f"上传文件请求参数：{file_json}")
    parse_file = ProcessDocuments(file_json)
    resp = await parse_file.process_documents(file_json)
    # parse_file = ParseFile(file_json)
    # resp = parse_file.save_file_to_db()

    logger.info(f"上传文件响应结果：{resp}")
    return JSONResponse(resp)


# @app.post("/network/search")
# async def chat_with_rag(request: Request, chat_json:dict):
#     retriever = ChatRetrieverRag(chat_json)

#     return EventSourceResponse(retriever.generate_event(chat_json, request), ping=300)

@app.post("/rag/chat")
async def chat_with_rag(request: Request, chat_json:dict):
    retriever = ChatRetrieverRag(chat_json)

    return EventSourceResponse(retriever.generate_event(chat_json, request), ping=300)

@app.post("/rag/chat/sync")
async def chat_with_rag_sync(chat_json: dict):
    """非流式RAG聊天接口"""
    logger.info(f"非流式RAG请求参数：{chat_json}")
    retriever = ChatRetrieverRag(chat_json)
    result = await retriever.generate_sync_response(chat_json)
    logger.info(f"非流式RAG响应结果：{result.get('code')}")
    return JSONResponse(result)

@app.post("/rag/query")
async def generate_query(request: Request, query_json:dict):
    logger.info(f"请求参数：{query_json}")
    relevant_query = ChatRetrieverRag(query_json)
    relevant_json = await relevant_query.generate_relevant_query(query_json)

    return JSONResponse(relevant_json)

@app.get("/rag/slice/search/{chat_id}")
async def generate_query(request: Request, chat_id:str = None):
    chat = ChatRetrieverRag(chat_id=chat_id)
    chunk_json = await chat.search_slice()

    return JSONResponse(chunk_json)


@app.delete("/rag/delete_slice/{slice_id}/{knowledge_id}/{document_id}")
async def delete_by_chunk_id(slice_id:str=None, knowledge_id:str=None, document_id:str=None):
    logger.info(f"删除切片接口中，知识库：{knowledge_id}, 切片id：{slice_id}")
    resp = MilvusOperate(collection_name=knowledge_id)._delete_by_chunk_id(slice_id, knowledge_id, document_id)
    logger.info(f"删除切片信息的结果：{resp}")
    return JSONResponse(resp)

@app.delete("/rag/delete_doc/{doc_id}/{knowledge_id}")
async def delete_by_doc_id(doc_id:str=None, knowledge_id:str=None):
    logger.info(f"删除文档id接口，知识库：{knowledge_id}, 文档id：{doc_id}")
    resp = MilvusOperate(collection_name=knowledge_id)._delete_by_doc_id(doc_id=doc_id)
    logger.info(f"删除文档的结果：{resp}")
    return JSONResponse(resp)

@app.put("/rag/update_slice")
async def put_by_id(slice_json:dict):
    logger.info(f"更新切片信息的请求参数：{slice_json}")
    collection_name = slice_json.get("knowledge_id")
    resp = MilvusOperate(collection_name=collection_name)._put_by_id(slice_json)

    logger.info(f"更新切片信息的结果：{resp}")
    return JSONResponse(resp)

@app.post("/rag/insert_slice")
async def insert_slice_text(slice_json:dict):
    logger.info(f"新增切片信息的请求参数：{slice_json}")
    collection_name = slice_json.get("knowledge_id")
    resp = MilvusOperate(collection_name=collection_name)._insert_slice(slice_json)

    logger.info(f"新增切片信息的结果：{resp}")
    return JSONResponse(resp)

@app.get("/rag/search/{knowledge_id}/{slice_id}")
async def search_by_doc_id(knowledge_id:str=None, slice_id:str=None):
    # 根据切片id查询切片信息
    # print(f"知识库：{knowledge_id}, 切片：{slice_id}")
    logger.info(f"根据切片id查询的数据库名：{knowledge_id}，切片id：{slice_id}")
    collection_name = knowledge_id  # 根据传过来的id处理对应知识库
    resp = MilvusOperate(collection_name=collection_name)._search_by_chunk_id(slice_id)

    logger.info(f"根据切片id查询结果：{resp}")
    return JSONResponse(resp)

@app.post("/rag/search_word")
async def search_by_key_word(search_json:dict):
    # 根据doc_id 查询切片列表信息
    collection_name = search_json.get("knowledge_id")
    logger.info(f"根据关键字请求的参数：{search_json}")
    resp = MilvusOperate(collection_name=collection_name)._search_by_key_word(search_json)

    logger.info(f"根据关键字查询的结果：{resp}")
    return JSONResponse(resp)

@app.delete("/rag/delete_knowledge/{knowledge_id}")
async def delete_collection(knowledge_id: str = None):
    logger.info(f"删除数据库请求的参数：{knowledge_id}")
    resp = MilvusOperate(collection_name=knowledge_id)._delete_collection()
    logger.info(f"删除向量库结果：{resp}")

    return JSONResponse(resp)

@app.post("/rag/create_collection")
async def create_collection(collection: dict):
    collection_name = collection.get("knowledge_id")
    embedding_name = collection.get("embedding_id")
    logger.info(f"创建向量库的库名：{collection_name}，向量名称：{embedding_name}")
    resp = MilvusOperate(collection_name=collection_name, embedding_name=embedding_name)._create_collection()
    logger.info(f"创建向量库结果：{resp}")
    
    return JSONResponse(resp)


# ==================== 任务取消接口 ====================

@app.post("/rag/cancel_task/{task_id}")
async def cancel_task(task_id: str):
    """
    取消正在执行的文档处理任务
    
    参数:
    - task_id: 任务ID（即 document_id）
    
    返回:
    - code: 200 成功, 404 任务不存在, 400 任务已取消
    """
    mysql_client = MysqlOperate()
    # 首先删除 bm_document 表中的记录
    del_success, del_info = mysql_client.delete_document(task_id)

    success, message = task_registry.cancel(task_id)
    
    if success:
        if del_success:
            logger.info(f"任务 {task_id} 已被取消，bm_document 记录已删除")
        else:
            logger.warning(f"任务 {task_id} 取消成功，但删除 bm_document 记录失败: {del_info}")
        
        return JSONResponse({"code": 200, "message": "任务取消成功"})
    
    if "不存在" in message:
        return JSONResponse({"code": 404, "message": message})
    
    return JSONResponse({"code": 400, "message": message})


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=6666)