feat: api

This commit is contained in:
李如威 2026-01-08 11:28:22 +08:00
parent 63d875bdc5
commit cdad058c22
10 changed files with 99 additions and 19 deletions

View File

@ -1,4 +1,4 @@
#!/usr/bin/env bash #!/usr/bin/env bash
export $(cat .env | xargs) 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

View File

@ -1,5 +1,6 @@
from fastapi import FastAPI from fastapi import FastAPI
from src.pipeline.api.rag import router as rag_router
def include_router(app:FastAPI): def include_router(app:FastAPI):
pass app.include_router(rag_router, prefix="/api/v1/rag")

18
src/pipeline/api/rag.py Normal file
View File

@ -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")

View File

@ -11,8 +11,8 @@ class AsyncES:
max_connections: int = 50, max_connections: int = 50,
max_keepalive: int = 10, max_keepalive: int = 10,
): ):
self.base = f"{config['es_host']}:{config['es_port']}".rstrip("/") self.base = f"{config.es_host}:{config.es_port}".rstrip("/")
self.auth = (config["es_user"], config["es_password"]) self.auth = (config.es_user, config.es_password)
self.client = httpx.AsyncClient( self.client = httpx.AsyncClient(
http2=False, http2=False,
@ -48,7 +48,7 @@ class AsyncES:
"content": {"type": "text", "analyzer": "ik_smart"}, "content": {"type": "text", "analyzer": "ik_smart"},
"embedding": { "embedding": {
"type": "dense_vector", "type": "dense_vector",
"dims": config['embedding_dims'], "dims": config.embedding_dims,
"index": True, "index": True,
"similarity": "cosine", "similarity": "cosine",
}, },
@ -175,7 +175,6 @@ class AsyncES:
for h in hits for h in hits
] ]
async def close(self): async def close(self):
await self.client.aclose() await self.client.aclose()

View File

@ -12,9 +12,9 @@ class AsyncLLm:
max_keepalive: int = 10, max_keepalive: int = 10,
): ):
# chat # chat
self.chat_api = config["llm_api_host"].rstrip("/") + "/chat/completions" self.chat_api = config.llm_api_host.rstrip("/") + "/chat/completions"
self.chat_model = config["llm_model"] self.chat_model = config.llm_model
self.chat_api_key = config["llm_api_key"] self.chat_api_key = config.llm_api_key
self.chat_client = httpx.AsyncClient( self.chat_client = httpx.AsyncClient(
http2=False, http2=False,
trust_env=False, trust_env=False,
@ -29,9 +29,9 @@ class AsyncLLm:
}, },
) )
# embedding # embedding
self.embedding_api = config["embedding_api_host"].rstrip("/") + "/embeddings" self.embedding_api = config.embedding_api_host.rstrip("/") + "/embeddings"
self.embedding_model = config["embedding_model"] self.embedding_model = config.embedding_model
self.embedding_api_key = config["embedding_api_key"] self.embedding_api_key = config.embedding_api_key
self.embedding_client = httpx.AsyncClient( self.embedding_client = httpx.AsyncClient(
http2=False, http2=False,
trust_env=False, trust_env=False,
@ -46,9 +46,9 @@ class AsyncLLm:
}, },
) )
# rerank # rerank
self.rerank_api = config["rerank_api_host"].rstrip("/") + "/score" self.rerank_api = config.rerank_api_host.rstrip("/") + "/score"
self.rerank_model = config["rerank_model"] self.rerank_model = config.rerank_model
self.rerank_api_key = config["rerank_api_key"] self.rerank_api_key = config.rerank_api_key
self.rerank_client = httpx.AsyncClient( self.rerank_client = httpx.AsyncClient(
http2=False, http2=False,
trust_env=False, trust_env=False,

View File

@ -1,17 +1,32 @@
from fastapi import FastAPI 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 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( app = FastAPI(
title="AI Pipeline", title="AI Pipeline",
description="轻量级 AI Pipeline", description="轻量级 AI Pipeline",
version=config["version"], 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

@ -0,0 +1,4 @@
from pydantic import BaseModel
class ChatRequest(BaseModel):
query: str

View File

@ -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

View File

@ -17,10 +17,10 @@ from baidusearch.baidusearch import search
# ----------------------------- # -----------------------------
logger.remove() # 清除已有 handler包括 pytest 导致的重复加载 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( logger.add(
"logs/pipeline.log", "logs/pipeline.log",
level=config["logger_level"], level=config.logger_level,
rotation="10 MB", # 自动分割 rotation="10 MB", # 自动分割
retention="7 days", # 保留时间 retention="7 days", # 保留时间
compression="zip", # 自动压缩 compression="zip", # 自动压缩

View File

@ -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