RAG/main.py

import asyncio
import os
import threading
from contextlib import asynccontextmanager

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

# 导入配置
from core.config import logger

# 导入API路由
from api.search import router as search_router
from api.health import router as health_router
from utils.keywords_utils import KeywordSummaryTask

# 创建 FastAPI 应用
app = FastAPI(
    title="RAG",
    description="RAG数据检索",
    version="1.0.0"
)

# 添加 CORS 中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 注册路由
app.include_router(health_router, tags=["健康检查"])
app.include_router(search_router, tags=["数据检索"])


# 定义生命周期事件处理器
@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup 逻辑
    print("正在启动数据胶囊 API 服务...")
    # 初始化一些资源，例如数据库连接
    # app.state.db = await init_database()  # 假设是异步初始化
    yield  # 应用在此处运行，处理请求

    # Shutdown 逻辑
    print("数据胶囊 API 服务正在关闭...")


if __name__ == "__main__":
    import uvicorn

    os.environ['RAG_ENV'] = 'prod'
    keyword_summary_task = KeywordSummaryTask()

    # 启动单独的线程来执行 process_texts_concurrently 方法
    threading.Thread(target=keyword_summary_task.process_texts_concurrently, daemon=True).start()
    uvicorn.run(app, host="0.0.0.0", port=5000)