From c3c1d3fbcf843751964455b57aeae4bf400dd4f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=A6=82=E5=A8=81?= Date: Thu, 8 Jan 2026 13:20:09 +0800 Subject: [PATCH] feat: test rag --- src/pipeline/api/rag.py | 4 ++-- src/pipeline/main.py | 3 ++- tests/pipeline/api/test_rag.py | 14 ++++++++++++++ 3 files changed, 18 insertions(+), 3 deletions(-) diff --git a/src/pipeline/api/rag.py b/src/pipeline/api/rag.py index 9398396..1f316a3 100644 --- a/src/pipeline/api/rag.py +++ b/src/pipeline/api/rag.py @@ -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") diff --git a/src/pipeline/main.py b/src/pipeline/main.py index 6f27963..3304dd4 100644 --- a/src/pipeline/main.py +++ b/src/pipeline/main.py @@ -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) diff --git a/tests/pipeline/api/test_rag.py b/tests/pipeline/api/test_rag.py index 168f8c0..64f37a6 100644 --- a/tests/pipeline/api/test_rag.py +++ b/tests/pipeline/api/test_rag.py @@ -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