Compare commits
2 Commits
cdad058c22
...
c3c1d3fbcf
| Author | SHA1 | Date |
|---|---|---|
|
|
c3c1d3fbcf | |
|
|
ad1b9c71fd |
|
|
@ -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")
|
||||||
|
|
|
||||||
|
|
@ -2,11 +2,11 @@ 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()
|
||||||
|
|
||||||
|
|
@ -15,18 +15,17 @@ async def lifespan(app:FastAPI):
|
||||||
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
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue