feat: 调整 db init

This commit is contained in:
李如威 2026-01-09 13:27:37 +08:00
parent b0b547825a
commit 775469a37f
9 changed files with 203 additions and 32 deletions

View File

@ -10,6 +10,7 @@ aiofiles
pillow
loguru
httpx
asgi_lifespan
baidusearch
celery[redis]
tortoise-orm[asyncmy]

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -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:
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())