feat: 调整 db init
This commit is contained in:
parent
b0b547825a
commit
775469a37f
|
|
@ -10,6 +10,7 @@ aiofiles
|
|||
pillow
|
||||
loguru
|
||||
httpx
|
||||
asgi_lifespan
|
||||
baidusearch
|
||||
celery[redis]
|
||||
tortoise-orm[asyncmy]
|
||||
|
|
@ -1,4 +1,4 @@
|
|||
#!/usr/bin/env bash
|
||||
export $(cat .env | xargs)
|
||||
pytest -s -W ignore::DeprecationWarning tests/pipeline/api/test_rag.py
|
||||
pytest -s -W ignore::DeprecationWarning tests/pipeline/api/test_rag.py::test_chat_message_list
|
||||
|
||||
|
|
|
|||
|
|
@ -1,27 +1,47 @@
|
|||
from pydoc import pager
|
||||
from typing import Optional
|
||||
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.services.rag import rag_service
|
||||
from src.pipeline.schemas.schemas import ChatListRequest, ChatListResponse, MessageListResponse, StreamChatRequest, mapping_chat_item, mapping_message_items
|
||||
from src.pipeline.utils import logger
|
||||
from src.pipeline.models import User as TUser
|
||||
from src.pipeline.models import (
|
||||
User as TUser,
|
||||
Chat as TChat
|
||||
)
|
||||
import json
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.post("/stream-chat")
|
||||
async def stream_chat(body: ChatRequest):
|
||||
|
||||
logger.debug(await TUser.all())
|
||||
|
||||
service = RagService()
|
||||
|
||||
async def stream_chat(body: StreamChatRequest):
|
||||
async def event_generator():
|
||||
chat_uuid, message_uuid = "", ""
|
||||
yield f"data: {json.dumps({'status': 'start'})}\n\n"
|
||||
async for tokens, uuid_0, uuid_1 in service.stream_chat(body.query, chat_uuid=body.chat_uuid):
|
||||
async for tokens, uuid_0, uuid_1 in rag_service.stream_chat(body.query, chat_uuid=body.chat_uuid):
|
||||
chat_uuid = uuid_0
|
||||
message_uuid = uuid_1
|
||||
yield f"data: {json.dumps({'message': tokens, 'status': 'process', "chat": chat_uuid, "message": message_uuid})}\n\n"
|
||||
yield f"data: {json.dumps({'status': 'end'})}\n\n"
|
||||
|
||||
return StreamingResponse(event_generator(), media_type="text/event-stream")
|
||||
|
||||
|
||||
@router.get("/chat-list")
|
||||
async def chat_list(
|
||||
user_id: int,
|
||||
page: Optional[int] = None,
|
||||
page_size: Optional[int] = None,
|
||||
) -> ChatListResponse:
|
||||
data, total = await rag_service.chat_list(user_id, page, page_size)
|
||||
return ChatListResponse(
|
||||
data=[mapping_chat_item(x) for x in data],
|
||||
total=total,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/message-list/{chat_uuid}")
|
||||
async def chat_message_list(chat_uuid: str) -> MessageListResponse:
|
||||
res = await rag_service.chat_message_list(chat_uuid=chat_uuid)
|
||||
return MessageListResponse(total=len(res), data=mapping_message_items(TChat(uuid=chat_uuid), res))
|
||||
|
|
|
|||
|
|
@ -10,7 +10,6 @@ DEFAULT_INIT_SQL = "src/pipeline/db/init.sql"
|
|||
|
||||
|
||||
async def init_tortoise(
|
||||
app: FastAPI,
|
||||
db_url: str,
|
||||
modules: dict = DEFAULT_MODELS,
|
||||
schema_sql: str = DEFAULT_INIT_SQL,
|
||||
|
|
@ -32,7 +31,10 @@ async def init_tortoise(
|
|||
"""
|
||||
config = generate_config(db_url, app_modules=modules)
|
||||
|
||||
await RegisterTortoise(app, config=config, generate_schemas=generate_schemas)
|
||||
await Tortoise.init(config=config)
|
||||
if generate_schemas:
|
||||
await Tortoise.generate_schemas()
|
||||
# await RegisterTortoise(app, config=config, generate_schemas=generate_schemas)
|
||||
|
||||
if schema_sql:
|
||||
conn = Tortoise.get_connection("default")
|
||||
|
|
@ -46,3 +48,11 @@ async def init_tortoise(
|
|||
logger.debug("初始化数据库")
|
||||
logger.debug(f"sql:\n{sql}")
|
||||
await conn.execute_script(sql)
|
||||
|
||||
logger.debug("init_tortoise")
|
||||
|
||||
|
||||
async def close_tortoise():
|
||||
await Tortoise.close_connections()
|
||||
|
||||
logger.debug("close_tortoise")
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ from fastapi import FastAPI
|
|||
from src.pipeline.api import include_router
|
||||
from src.pipeline.config import config
|
||||
from src.pipeline.core import llm, es
|
||||
from src.pipeline.db import init_tortoise
|
||||
from src.pipeline.db import init_tortoise, close_tortoise
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
|
||||
|
|
@ -10,12 +10,13 @@ from contextlib import asynccontextmanager
|
|||
async def lifespan(app: FastAPI):
|
||||
await llm.init_client()
|
||||
await es.init_client()
|
||||
await init_tortoise(app, db_url=str(config.mysql_dsn), generate_schemas=config.env == "DEBUG")
|
||||
await init_tortoise(db_url=str(config.mysql_dsn), generate_schemas=config.env == "DEBUG")
|
||||
|
||||
yield
|
||||
|
||||
await llm.close_client()
|
||||
await es.close_client()
|
||||
await close_tortoise()
|
||||
|
||||
|
||||
app = FastAPI(title="AI Pipeline", description="轻量级 AI Pipeline", version=config.version, lifespan=lifespan)
|
||||
|
|
|
|||
|
|
@ -38,10 +38,30 @@ class BaseModel(Model):
|
|||
:param page_size: 页数,不设置返回全部
|
||||
:type page_size: int | None
|
||||
"""
|
||||
datas = cls.filter(**params).order_by(*order_by)
|
||||
if page and page_size:
|
||||
return await cls.filter(**params).order_by(*order_by).limit(page_size).offset((page - 1) * page_size)
|
||||
datas = datas.limit(page_size).offset((page - 1) * page_size)
|
||||
else:
|
||||
return await cls.filter(**params).order_by(*order_by).all()
|
||||
datas = datas.all()
|
||||
return await datas
|
||||
|
||||
@classmethod
|
||||
async def count(
|
||||
cls,
|
||||
params: dict,
|
||||
):
|
||||
"""
|
||||
SELECT 语句
|
||||
|
||||
:param cls:
|
||||
:param params: 筛选条件
|
||||
:type params: dict
|
||||
"""
|
||||
return await cls.filter(**params).count()
|
||||
|
||||
@classmethod
|
||||
async def find(cls, params:dict):
|
||||
return await cls.filter(**params).first()
|
||||
|
||||
|
||||
class User(BaseModel):
|
||||
|
|
|
|||
|
|
@ -1,5 +1,87 @@
|
|||
from mimetypes import init
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional
|
||||
class ChatRequest(BaseModel):
|
||||
from typing import Optional, List
|
||||
from src.pipeline.models import Chat as TChat, User as TUser, Message as TMessage
|
||||
|
||||
# ----------------------------------------
|
||||
# 对话相关
|
||||
# ----------------------------------------
|
||||
|
||||
|
||||
class ChatItem(BaseModel):
|
||||
user_name: str
|
||||
user_id: int
|
||||
title: str
|
||||
uuid: str
|
||||
create_at: str
|
||||
|
||||
|
||||
def mapping_chat_item(chat: TChat, user: TUser | None = None):
|
||||
return ChatItem(
|
||||
user_name=(user.username or "") if user else "",
|
||||
user_id=user.id if user else 0,
|
||||
title=chat.title,
|
||||
uuid=str(chat.uuid),
|
||||
create_at=chat.create_at.isoformat(),
|
||||
)
|
||||
|
||||
|
||||
class StreamChatRequest(BaseModel):
|
||||
query: str
|
||||
chat_uuid: Optional[str] = None
|
||||
|
||||
|
||||
class ChatListRequest(BaseModel):
|
||||
user_id: int
|
||||
page: int = 1
|
||||
page_size: Optional[int] = None
|
||||
|
||||
|
||||
class ChatListResponse(BaseModel):
|
||||
total: int
|
||||
data: List[ChatItem]
|
||||
|
||||
|
||||
class MessageItem(BaseModel):
|
||||
user_id: int
|
||||
user_name: str
|
||||
uuid: str
|
||||
role: str
|
||||
content: str
|
||||
create_at: str
|
||||
chat_uuid: str
|
||||
|
||||
|
||||
def mapping_message_items(chat: TChat, t_message_list: list[TMessage], user: TUser | None = None) -> list[MessageItem]:
|
||||
res: list[MessageItem] = []
|
||||
for obj in t_message_list:
|
||||
if "question" in obj.content and obj.content["question"]:
|
||||
res.append(
|
||||
MessageItem(
|
||||
user_id=user.id if user else 0,
|
||||
user_name=user.username if user else "",
|
||||
chat_uuid=str(chat.uuid) if chat else "",
|
||||
uuid=str(obj.uuid),
|
||||
content=obj.content["question"],
|
||||
role="user",
|
||||
create_at=str(obj.create_at),
|
||||
)
|
||||
)
|
||||
if "answer" in obj.content and obj.content["answer"]:
|
||||
res.append(
|
||||
MessageItem(
|
||||
user_id=user.id if user else 0,
|
||||
user_name=user.username if user else "",
|
||||
chat_uuid=str(chat.uuid) if chat else "",
|
||||
uuid=str(obj.uuid),
|
||||
content=obj.content["answer"],
|
||||
role="assistant",
|
||||
create_at=str(obj.create_at),
|
||||
)
|
||||
)
|
||||
return res
|
||||
|
||||
|
||||
class MessageListResponse(BaseModel):
|
||||
data: List[MessageItem]
|
||||
total: int
|
||||
|
|
|
|||
|
|
@ -45,3 +45,19 @@ class RagService:
|
|||
"question": query, "answer": content
|
||||
}
|
||||
await t_message.save()
|
||||
|
||||
async def chat_list(self, user_id:int, page:int|None = None, page_size:int|None = None):
|
||||
params = {"user_id": user_id}
|
||||
t_chat_list = await TChat.select(params, page=page, page_size=page_size)
|
||||
t_count = await TChat.count(params)
|
||||
return t_chat_list, t_count
|
||||
|
||||
async def chat_message_list(self, chat_uuid: str):
|
||||
res:list[TMessage] = []
|
||||
t_chat = await TChat.find({"uuid": chat_uuid})
|
||||
if not t_chat:
|
||||
return res
|
||||
res = await TMessage.select({"chat_id": t_chat.id}, order_by=["id"])
|
||||
return res
|
||||
|
||||
rag_service = RagService()
|
||||
|
|
|
|||
|
|
@ -1,29 +1,26 @@
|
|||
from asgi_lifespan import LifespanManager
|
||||
from httpx import AsyncClient, ASGITransport
|
||||
from src.pipeline.main import app
|
||||
from src.pipeline.utils import logger
|
||||
from src.pipeline.core import llm
|
||||
from fastapi.testclient import TestClient
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import json
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope="session")
|
||||
async def init_llm():
|
||||
logger.debug('init_llm')
|
||||
await llm.init_client()
|
||||
yield
|
||||
await llm.close_client()
|
||||
# ----------------------
|
||||
# async
|
||||
# ----------------------
|
||||
|
||||
|
||||
@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
|
||||
async def app_client_async():
|
||||
async with LifespanManager(app):
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_chat(app_client, init_llm):
|
||||
async def test_stream_chat(app_client_async):
|
||||
res = await app_client.post("/api/v1/rag/stream-chat", json={"query": "hello"})
|
||||
assert res.status_code == 200
|
||||
sse_messages = []
|
||||
|
|
@ -39,3 +36,27 @@ async def test_stream_chat(app_client, init_llm):
|
|||
logger.debug(obj)
|
||||
sse_messages.append(obj)
|
||||
assert sse_messages
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_message_list_async(app_client_async):
|
||||
res = await app_client_async.get("/api/v1/rag/message-list/a1c9108c-9201-4e60-a436-505edec3f47e")
|
||||
assert res.status_code == 200
|
||||
logger.debug(res.json())
|
||||
|
||||
|
||||
# ----------------------
|
||||
# normal
|
||||
# ----------------------
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def app_client():
|
||||
client = TestClient(app=app)
|
||||
return client
|
||||
|
||||
|
||||
def test_chat_message_list(app_client):
|
||||
res = app_client.get("/api/v1/rag/message-list/a1c9108c-9201-4e60-a436-505edec3f47e")
|
||||
assert res.status_code == 200
|
||||
logger.debug(res.json())
|
||||
|
|
|
|||
Loading…
Reference in New Issue