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 pillow
loguru loguru
httpx httpx
asgi_lifespan
baidusearch baidusearch
celery[redis] celery[redis]
tortoise-orm[asyncmy] tortoise-orm[asyncmy]

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 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 import APIRouter
from fastapi.responses import StreamingResponse from fastapi.responses import StreamingResponse
from src.pipeline.services.rag import RagService from src.pipeline.services.rag import rag_service
from src.pipeline.schemas.schemas import ChatRequest 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.utils import logger
from src.pipeline.models import User as TUser from src.pipeline.models import (
User as TUser,
Chat as TChat
)
import json import json
router = APIRouter() router = APIRouter()
@router.post("/stream-chat") @router.post("/stream-chat")
async def stream_chat(body: ChatRequest): async def stream_chat(body: StreamChatRequest):
logger.debug(await TUser.all())
service = RagService()
async def event_generator(): async def event_generator():
chat_uuid, message_uuid = "", "" chat_uuid, message_uuid = "", ""
yield f"data: {json.dumps({'status': 'start'})}\n\n" 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 chat_uuid = uuid_0
message_uuid = uuid_1 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({'message': tokens, 'status': 'process', "chat": chat_uuid, "message": message_uuid})}\n\n"
yield f"data: {json.dumps({'status': 'end'})}\n\n" yield f"data: {json.dumps({'status': 'end'})}\n\n"
return StreamingResponse(event_generator(), media_type="text/event-stream") 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( async def init_tortoise(
app: FastAPI,
db_url: str, db_url: str,
modules: dict = DEFAULT_MODELS, modules: dict = DEFAULT_MODELS,
schema_sql: str = DEFAULT_INIT_SQL, schema_sql: str = DEFAULT_INIT_SQL,
@ -32,7 +31,10 @@ async def init_tortoise(
""" """
config = generate_config(db_url, app_modules=modules) 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: if schema_sql:
conn = Tortoise.get_connection("default") conn = Tortoise.get_connection("default")
@ -46,3 +48,11 @@ async def init_tortoise(
logger.debug("初始化数据库") logger.debug("初始化数据库")
logger.debug(f"sql:\n{sql}") logger.debug(f"sql:\n{sql}")
await conn.execute_script(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.api import include_router
from src.pipeline.config import config from src.pipeline.config import config
from src.pipeline.core import llm, es 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 from contextlib import asynccontextmanager
@ -10,12 +10,13 @@ from contextlib import asynccontextmanager
async def lifespan(app: FastAPI): async def lifespan(app: FastAPI):
await llm.init_client() await llm.init_client()
await es.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 yield
await llm.close_client() await llm.close_client()
await es.close_client() await es.close_client()
await close_tortoise()
app = FastAPI(title="AI Pipeline", description="轻量级 AI Pipeline", version=config.version, lifespan=lifespan) 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: 页数不设置返回全部 :param page_size: 页数不设置返回全部
:type page_size: int | None :type page_size: int | None
""" """
datas = cls.filter(**params).order_by(*order_by)
if page and page_size: 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: 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): class User(BaseModel):

View File

@ -1,5 +1,87 @@
from mimetypes import init
from pydantic import BaseModel from pydantic import BaseModel
from typing import Optional from typing import Optional, List
class ChatRequest(BaseModel): 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 query: str
chat_uuid: Optional[str] = None 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 "question": query, "answer": content
} }
await t_message.save() 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 httpx import AsyncClient, ASGITransport
from src.pipeline.main import app from src.pipeline.main import app
from src.pipeline.utils import logger from src.pipeline.utils import logger
from src.pipeline.core import llm from fastapi.testclient import TestClient
import pytest import pytest
import pytest_asyncio import pytest_asyncio
import json import json
# ----------------------
@pytest_asyncio.fixture(scope="session") # async
async def init_llm(): # ----------------------
logger.debug('init_llm')
await llm.init_client()
yield
await llm.close_client()
@pytest_asyncio.fixture(scope="session") @pytest_asyncio.fixture(scope="session")
async def app_client(): async def app_client_async():
transport = ASGITransport(app=app) async with LifespanManager(app):
async with AsyncClient(transport=transport, base_url="http://test") as client: async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
yield client yield client
@pytest.mark.asyncio @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"}) res = await app_client.post("/api/v1/rag/stream-chat", json={"query": "hello"})
assert res.status_code == 200 assert res.status_code == 200
sse_messages = [] sse_messages = []
@ -39,3 +36,27 @@ async def test_stream_chat(app_client, init_llm):
logger.debug(obj) logger.debug(obj)
sse_messages.append(obj) sse_messages.append(obj)
assert sse_messages 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())