You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

92 lines
2.9 KiB

import inspect
import json
import logging
import os
import warnings
from logging.handlers import RotatingFileHandler
import fastapi
import uvicorn
from fastapi import FastAPI
from sse_starlette import EventSourceResponse
from starlette.staticfiles import StaticFiles
from Util.LightRagUtil import initialize_rag
from lightrag import QueryParam
# 初始化日志
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
# 配置日志处理器
log_file = os.path.join(os.path.dirname(__file__), 'Logs', 'app.log')
os.makedirs(os.path.dirname(log_file), exist_ok=True)
# 文件处理器
file_handler = RotatingFileHandler(
log_file, maxBytes=1024 * 1024, backupCount=5, encoding='utf-8')
file_handler.setFormatter(logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
# 控制台处理器
console_handler = logging.StreamHandler()
console_handler.setFormatter(logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
logger.addHandler(file_handler)
logger.addHandler(console_handler)
async def lifespan(app: FastAPI):
# 抑制HTTPS相关警告
warnings.filterwarnings('ignore', message='Connecting to .* using TLS with verify_certs=False is insecure')
warnings.filterwarnings('ignore', message='Unverified HTTPS request is being made to host')
yield
app = FastAPI(lifespan=lifespan)
# 挂载静态文件目录
app.mount("/static", StaticFiles(directory="Static"), name="static")
@app.post("/api/rag", response_model=None)
async def rag(request: fastapi.Request):
data = await request.json()
query = data.get('query', '')
working_dir = "./Test/Math" # 使用与T2_Query.py相同的目录
async def generate_response_stream():
rag = None
try:
logger.info(f"Initializing RAG with working_dir: {working_dir}")
rag = await initialize_rag(working_dir)
logger.info("RAG initialized successfully")
resp = await rag.aquery(
query,
param=QueryParam(mode="hybrid", stream=True),
)
logger.info("Query execution started")
async for chunk in resp:
if not chunk: # Skip empty chunks
continue
try:
yield f"data: {json.dumps({'reply': chunk}, ensure_ascii=False)}\n\n"
except Exception as e:
logger.error(f"Error processing chunk: {e}")
yield f"data: {json.dumps({'error': str(e)})}\n\n"
except Exception as e:
yield f"data: {json.dumps({'error': str(e)})}\n\n"
finally:
if rag:
await rag.finalize_storages()
return EventSourceResponse(generate_response_stream())
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000)