feat: api
This commit is contained in:
parent
63d875bdc5
commit
cdad058c22
|
|
@ -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
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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")
|
||||||
|
|
@ -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")
|
||||||
|
|
@ -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()
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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,
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,4 @@
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
class ChatRequest(BaseModel):
|
||||||
|
query: str
|
||||||
|
|
@ -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
|
||||||
|
|
@ -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", # 自动压缩
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
Loading…
Reference in New Issue