feat: test rag
This commit is contained in:
parent
ad1b9c71fd
commit
c3c1d3fbcf
|
|
@ -12,9 +12,9 @@ async def stream_chat(body: ChatRequest):
|
|||
service = RagService()
|
||||
|
||||
async def event_generator():
|
||||
yield f"data: {json.dumps({'': "", 'status': 'start'})}\n\n"
|
||||
yield f"data: {json.dumps({'message': '', 'status': 'start'})}\n\n"
|
||||
async for tokens in service.stream_chat(body.query):
|
||||
yield f"data: {json.dumps({'message': tokens, 'status': 'process'})}\n\n"
|
||||
yield f"data: {json.dumps({'message': "", 'status': 'end'})}\n\n"
|
||||
yield f"data: {json.dumps({'message': '', 'status': 'end'})}\n\n"
|
||||
|
||||
return StreamingResponse(event_generator(), media_type="text/event-stream")
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ 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
|
||||
|
|
@ -18,6 +18,7 @@ async def lifespan(app: FastAPI):
|
|||
|
||||
app = FastAPI(title="AI Pipeline", description="轻量级 AI Pipeline", version=config.version, lifespan=lifespan)
|
||||
|
||||
|
||||
include_router(app)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue