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