dsProject/dsRag/Start.py

import asyncio
import logging
from contextlib import asynccontextmanager
from logging.handlers import RotatingFileHandler

import uvicorn
from fastapi import FastAPI, UploadFile, File

from Dao.KbDao import KbDao
from Util.MySQLUtil import init_mysql_pool

# 初始化日志
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
handler = RotatingFileHandler('Logs/start.log', maxBytes=1024*1024, backupCount=5)
handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
logger.addHandler(handler)

@asynccontextmanager
async def lifespan(app: FastAPI):
    # 初始化数据库连接池
    app.state.kb_dao = KbDao(await init_mysql_pool())
    
    # 启动文档处理任务
    async def document_processor():
        while True:
            try:
                # 获取未处理文档
                # 处理文档
                # 保存到ES
                await asyncio.sleep(10)
            except Exception as e:
                logger.error(f"文档处理出错: {e}")
                await asyncio.sleep(10)
    
    task = asyncio.create_task(document_processor())
    yield
    # 关闭时取消任务
    task.cancel()
    # 关闭数据库连接池
    await app.state.kb_dao.mysql_pool.close()

app = FastAPI(lifespan=lifespan)

# 知识库CRUD接口
@app.post("/kb")
async def create_kb(kb: dict):
    """创建知识库"""
    return await app.state.kb_dao.create_kb(kb)

@app.get("/kb/{kb_id}")
async def read_kb(kb_id: int):
    """获取知识库详情"""
    return await app.state.kb_dao.get_kb(kb_id)

@app.post("/kb/update/{kb_id}")
async def update_kb(kb_id: int, kb: dict):
    """更新知识库信息"""
    return await app.state.kb_dao.update_kb(kb_id, kb)

@app.delete("/kb/{kb_id}")
async def delete_kb(kb_id: int):
    """删除知识库"""
    return await app.state.kb_dao.delete_kb(kb_id)

# 知识库文件CRUD接口
@app.post("/kb_file")
async def create_kb_file(file: dict):
    """创建知识库文件记录"""
    return await app.state.kb_dao.create_kb_file(file)

@app.get("/kb_files/{file_id}")
async def read_kb_file(file_id: int):
    """获取文件详情"""
    return await app.state.kb_dao.get_kb_file(file_id)

@app.post("/kb_files/update/{file_id}")
async def update_kb_file(file_id: int, file: dict):
    """更新文件信息"""
    return await app.state.kb_dao.update_kb_file(file_id, file)

@app.delete("/kb_files/{file_id}")
async def delete_kb_file(file_id: int):
    """删除文件记录"""
    return await app.state.kb_dao.delete_kb_file(file_id)

# 文件上传接口
@app.post("/upload")
async def upload_file(kb_id: int, file: UploadFile = File(...)):
    """文件上传接口"""
    return await app.state.kb_dao.handle_upload(kb_id, file)

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