diff --git a/src/pipeline/api/rag.py b/src/pipeline/api/rag.py index a7bae94..9398396 100644 --- a/src/pipeline/api/rag.py +++ b/src/pipeline/api/rag.py @@ -3,6 +3,7 @@ from fastapi.responses import StreamingResponse from src.pipeline.services.rag import RagService from src.pipeline.schemas.schemas import ChatRequest from src.pipeline.utils import logger +import json router = APIRouter() @@ -11,8 +12,9 @@ async def stream_chat(body: ChatRequest): service = RagService() async def event_generator(): + yield f"data: {json.dumps({'': "", 'status': 'start'})}\n\n" async for tokens in service.stream_chat(body.query): - logger.debug(tokens) - yield f"data: {tokens}\n\n" + yield f"data: {json.dumps({'message': tokens, 'status': 'process'})}\n\n" + yield f"data: {json.dumps({'message': "", 'status': 'end'})}\n\n" return StreamingResponse(event_generator(), media_type="text/event-stream") diff --git a/src/pipeline/main.py b/src/pipeline/main.py index c92dcab..6f27963 100644 --- a/src/pipeline/main.py +++ b/src/pipeline/main.py @@ -4,29 +4,27 @@ from src.pipeline.config import config from contextlib import asynccontextmanager from src.pipeline.core import llm, es, nodes + @asynccontextmanager -async def lifespan(app:FastAPI): - print("------") +async def lifespan(app: FastAPI): await llm.init_client() await es.init_client() - + yield await llm.close_client() await es.close_client() -app = FastAPI( - title="AI Pipeline", - description="轻量级 AI Pipeline", - version=config.version, - lifespan=lifespan -) + +app = FastAPI(title="AI Pipeline", description="轻量级 AI Pipeline", version=config.version, lifespan=lifespan) include_router(app) + @app.get("/") async def healthz(): return {"status": "running", "version": config.version} -if __name__ == '__main__': + +if __name__ == "__main__": pass