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.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({'message': '', '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

@ -2,31 +2,30 @@ from fastapi import FastAPI
from src.pipeline.api import include_router from src.pipeline.api import include_router
from src.pipeline.config import config 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
@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

View File

@ -4,6 +4,7 @@ from src.pipeline.utils import logger
from src.pipeline.core import llm from src.pipeline.core import llm
import pytest import pytest
import pytest_asyncio import pytest_asyncio
import json
@pytest_asyncio.fixture(scope="session") @pytest_asyncio.fixture(scope="session")
@ -25,3 +26,16 @@ async def app_client():
async def test_stream_chat(app_client, init_llm): async def test_stream_chat(app_client, init_llm):
res = await app_client.post("/api/v1/rag/stream-chat", json={"query": "hello"}) res = await app_client.post("/api/v1/rag/stream-chat", json={"query": "hello"})
assert res.status_code == 200 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