dsProject/dsLightRag/Start.py

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)