Compare commits

...

2 Commits

Author SHA1 Message Date
李如威 c3c1d3fbcf feat: test rag 2026-01-08 13:20:09 +08:00
李如威 ad1b9c71fd feat: 调整返回结构 2026-01-08 13:09:35 +08:00
3 changed files with 28 additions and 13 deletions

View File

@ -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({'message': '', '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")

View File

@ -2,11 +2,11 @@ from fastapi import FastAPI
from src.pipeline.api import include_router
from src.pipeline.config import config
from contextlib import asynccontextmanager
from src.pipeline.core import llm, es, nodes
from src.pipeline.core import llm, es
@asynccontextmanager
async def lifespan(app: FastAPI):
print("------")
await llm.init_client()
await es.init_client()
@ -15,18 +15,17 @@ async def lifespan(app:FastAPI):
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

View File

@ -4,6 +4,7 @@ from src.pipeline.utils import logger
from src.pipeline.core import llm
import pytest
import pytest_asyncio
import json
@pytest_asyncio.fixture(scope="session")
@ -25,3 +26,16 @@ async def app_client():
async def test_stream_chat(app_client, init_llm):
res = await app_client.post("/api/v1/rag/stream-chat", json={"query": "hello"})
assert res.status_code == 200
sse_messages = []
async for line in res.aiter_lines():
if not line :
continue
if line.startswith("data:"):
payload = line[len("data:") :].strip()
try:
obj = json.loads(payload)
except Exception:
obj = {"raw": payload}
logger.debug(obj)
sse_messages.append(obj)
assert sse_messages