diff --git a/scripts/run_test.sh b/scripts/run_test.sh index 7bbcfdd..6bb70cf 100644 --- a/scripts/run_test.sh +++ b/scripts/run_test.sh @@ -1,4 +1,4 @@ #!/usr/bin/env bash export $(cat .env | xargs) -pytest -s -W ignore::DeprecationWarning src/tests/test_nodes.py::test_agent +pytest -s -W ignore::DeprecationWarning tests/pipeline/api/test_rag.py diff --git a/src/pipeline/api/__init__.py b/src/pipeline/api/__init__.py index 668b471..763f6c3 100644 --- a/src/pipeline/api/__init__.py +++ b/src/pipeline/api/__init__.py @@ -1,5 +1,6 @@ from fastapi import FastAPI +from src.pipeline.api.rag import router as rag_router def include_router(app:FastAPI): - pass \ No newline at end of file + app.include_router(rag_router, prefix="/api/v1/rag") \ No newline at end of file diff --git a/src/pipeline/api/rag.py b/src/pipeline/api/rag.py new file mode 100644 index 0000000..a7bae94 --- /dev/null +++ b/src/pipeline/api/rag.py @@ -0,0 +1,18 @@ +from fastapi import APIRouter +from fastapi.responses import StreamingResponse +from src.pipeline.services.rag import RagService +from src.pipeline.schemas.schemas import ChatRequest +from src.pipeline.utils import logger + +router = APIRouter() + +@router.post("/stream-chat") +async def stream_chat(body: ChatRequest): + service = RagService() + + async def event_generator(): + async for tokens in service.stream_chat(body.query): + logger.debug(tokens) + yield f"data: {tokens}\n\n" + + return StreamingResponse(event_generator(), media_type="text/event-stream") diff --git a/src/pipeline/core/es.py b/src/pipeline/core/es.py index 178fafe..5cf57dc 100644 --- a/src/pipeline/core/es.py +++ b/src/pipeline/core/es.py @@ -11,8 +11,8 @@ class AsyncES: max_connections: int = 50, max_keepalive: int = 10, ): - self.base = f"{config['es_host']}:{config['es_port']}".rstrip("/") - self.auth = (config["es_user"], config["es_password"]) + self.base = f"{config.es_host}:{config.es_port}".rstrip("/") + self.auth = (config.es_user, config.es_password) self.client = httpx.AsyncClient( http2=False, @@ -48,7 +48,7 @@ class AsyncES: "content": {"type": "text", "analyzer": "ik_smart"}, "embedding": { "type": "dense_vector", - "dims": config['embedding_dims'], + "dims": config.embedding_dims, "index": True, "similarity": "cosine", }, @@ -175,7 +175,6 @@ class AsyncES: for h in hits ] - async def close(self): await self.client.aclose() diff --git a/src/pipeline/core/llm.py b/src/pipeline/core/llm.py index 60fbb83..808fa31 100644 --- a/src/pipeline/core/llm.py +++ b/src/pipeline/core/llm.py @@ -12,9 +12,9 @@ class AsyncLLm: max_keepalive: int = 10, ): # chat - self.chat_api = config["llm_api_host"].rstrip("/") + "/chat/completions" - self.chat_model = config["llm_model"] - self.chat_api_key = config["llm_api_key"] + self.chat_api = config.llm_api_host.rstrip("/") + "/chat/completions" + self.chat_model = config.llm_model + self.chat_api_key = config.llm_api_key self.chat_client = httpx.AsyncClient( http2=False, trust_env=False, @@ -29,9 +29,9 @@ class AsyncLLm: }, ) # embedding - self.embedding_api = config["embedding_api_host"].rstrip("/") + "/embeddings" - self.embedding_model = config["embedding_model"] - self.embedding_api_key = config["embedding_api_key"] + self.embedding_api = config.embedding_api_host.rstrip("/") + "/embeddings" + self.embedding_model = config.embedding_model + self.embedding_api_key = config.embedding_api_key self.embedding_client = httpx.AsyncClient( http2=False, trust_env=False, @@ -46,9 +46,9 @@ class AsyncLLm: }, ) # rerank - self.rerank_api = config["rerank_api_host"].rstrip("/") + "/score" - self.rerank_model = config["rerank_model"] - self.rerank_api_key = config["rerank_api_key"] + self.rerank_api = config.rerank_api_host.rstrip("/") + "/score" + self.rerank_model = config.rerank_model + self.rerank_api_key = config.rerank_api_key self.rerank_client = httpx.AsyncClient( http2=False, trust_env=False, diff --git a/src/pipeline/main.py b/src/pipeline/main.py index cada61f..c92dcab 100644 --- a/src/pipeline/main.py +++ b/src/pipeline/main.py @@ -1,17 +1,32 @@ 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 + +@asynccontextmanager +async def lifespan(app:FastAPI): + print("------") + await llm.init_client() + await es.init_client() + + yield + + await llm.close_client() + await es.close_client() + app = FastAPI( title="AI Pipeline", description="轻量级 AI Pipeline", - version=config["version"], + version=config.version, + lifespan=lifespan ) include_router(app) @app.get("/") async def healthz(): - return {"status": "running", "version": config["version"]} + return {"status": "running", "version": config.version} if __name__ == '__main__': pass diff --git a/src/pipeline/schemas/schemas.py b/src/pipeline/schemas/schemas.py new file mode 100644 index 0000000..20b2647 --- /dev/null +++ b/src/pipeline/schemas/schemas.py @@ -0,0 +1,4 @@ +from pydantic import BaseModel + +class ChatRequest(BaseModel): + query: str \ No newline at end of file diff --git a/src/pipeline/services/rag.py b/src/pipeline/services/rag.py new file mode 100644 index 0000000..6ace158 --- /dev/null +++ b/src/pipeline/services/rag.py @@ -0,0 +1,16 @@ +from src.pipeline.core import llm +from src.pipeline.utils import logger + +class RagService: + def __init__(self): + pass + + async def stream_chat(self, query:str): + logger.debug(query) + request = llm.client.stream_chat( + messages=[ + {"role": "user", "content": "query"}, + ] + ) + async for chunk in request: + yield chunk diff --git a/src/pipeline/utils.py b/src/pipeline/utils.py index 5cb4e2a..6c65cad 100644 --- a/src/pipeline/utils.py +++ b/src/pipeline/utils.py @@ -17,10 +17,10 @@ from baidusearch.baidusearch import search # ----------------------------- logger.remove() # 清除已有 handler(包括 pytest 导致的重复加载 -logger.add(sys.stdout, level=config["logger_level"], colorize=True) +logger.add(sys.stdout, level=config.logger_level, colorize=True) logger.add( "logs/pipeline.log", - level=config["logger_level"], + level=config.logger_level, rotation="10 MB", # 自动分割 retention="7 days", # 保留时间 compression="zip", # 自动压缩 diff --git a/tests/pipeline/api/test_rag.py b/tests/pipeline/api/test_rag.py new file mode 100644 index 0000000..168f8c0 --- /dev/null +++ b/tests/pipeline/api/test_rag.py @@ -0,0 +1,27 @@ +from httpx import AsyncClient, ASGITransport +from src.pipeline.main import app +from src.pipeline.utils import logger +from src.pipeline.core import llm +import pytest +import pytest_asyncio + + +@pytest_asyncio.fixture(scope="session") +async def init_llm(): + logger.debug('init_llm') + await llm.init_client() + yield + await llm.close_client() + + +@pytest_asyncio.fixture(scope="session") +async def app_client(): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as client: + yield client + + +@pytest.mark.asyncio +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