feat: 调整返回结构

This commit is contained in:
李如威 2026-01-08 13:09:35 +08:00
parent cdad058c22
commit ad1b9c71fd
2 changed files with 12 additions and 12 deletions

View File

@ -3,6 +3,7 @@ from fastapi.responses import StreamingResponse
from src.pipeline.services.rag import RagService from src.pipeline.services.rag import RagService
from src.pipeline.schemas.schemas import ChatRequest from src.pipeline.schemas.schemas import ChatRequest
from src.pipeline.utils import logger from src.pipeline.utils import logger
import json
router = APIRouter() router = APIRouter()
@ -11,8 +12,9 @@ async def stream_chat(body: ChatRequest):
service = RagService() service = RagService()
async def event_generator(): async def event_generator():
yield f"data: {json.dumps({'': "", 'status': 'start'})}\n\n"
async for tokens in service.stream_chat(body.query): async for tokens in service.stream_chat(body.query):
logger.debug(tokens) yield f"data: {json.dumps({'message': tokens, 'status': 'process'})}\n\n"
yield f"data: {tokens}\n\n" yield f"data: {json.dumps({'message': "", 'status': 'end'})}\n\n"
return StreamingResponse(event_generator(), media_type="text/event-stream") return StreamingResponse(event_generator(), media_type="text/event-stream")

View File

@ -4,29 +4,27 @@ from src.pipeline.config import config
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from src.pipeline.core import llm, es, nodes from src.pipeline.core import llm, es, nodes
@asynccontextmanager @asynccontextmanager
async def lifespan(app:FastAPI): async def lifespan(app: FastAPI):
print("------")
await llm.init_client() await llm.init_client()
await es.init_client() await es.init_client()
yield yield
await llm.close_client() await llm.close_client()
await es.close_client() await es.close_client()
app = FastAPI(
title="AI Pipeline", app = FastAPI(title="AI Pipeline", description="轻量级 AI Pipeline", version=config.version, lifespan=lifespan)
description="轻量级 AI Pipeline",
version=config.version,
lifespan=lifespan
)
include_router(app) include_router(app)
@app.get("/") @app.get("/")
async def healthz(): async def healthz():
return {"status": "running", "version": config.version} return {"status": "running", "version": config.version}
if __name__ == '__main__':
if __name__ == "__main__":
pass pass