Compare commits

..

No commits in common. "c3c1d3fbcf843751964455b57aeae4bf400dd4f2" and "cdad058c22b325ffab1607bc202cd4b60ebd0367" have entirely different histories.

3 changed files with 13 additions and 28 deletions

View File

@ -3,7 +3,6 @@ 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()
@ -12,9 +11,8 @@ 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):
yield f"data: {json.dumps({'message': tokens, 'status': 'process'})}\n\n" logger.debug(tokens)
yield f"data: {json.dumps({'message': '', 'status': 'end'})}\n\n" yield f"data: {tokens}\n\n"
return StreamingResponse(event_generator(), media_type="text/event-stream") return StreamingResponse(event_generator(), media_type="text/event-stream")

View File

@ -2,30 +2,31 @@ 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 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(
app = FastAPI(title="AI Pipeline", description="轻量级 AI Pipeline", version=config.version, lifespan=lifespan) title="AI Pipeline",
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,7 +4,6 @@ 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")
@ -26,16 +25,3 @@ 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