feat: 调整 db init
This commit is contained in:
parent
b0b547825a
commit
775469a37f
|
|
@ -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]
|
||||||
|
|
@ -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
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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))
|
||||||
|
|
|
||||||
|
|
@ -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")
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
|
|
@ -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):
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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())
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue