diff --git a/requirements.txt b/requirements.txt index ac3b269..fea759f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,6 +10,7 @@ aiofiles pillow loguru httpx +asgi_lifespan baidusearch celery[redis] tortoise-orm[asyncmy] \ No newline at end of file diff --git a/scripts/run_test.sh b/scripts/run_test.sh index 6bb70cf..a9b4483 100644 --- a/scripts/run_test.sh +++ b/scripts/run_test.sh @@ -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 diff --git a/src/pipeline/api/rag.py b/src/pipeline/api/rag.py index 1b4a961..cdc2e73 100644 --- a/src/pipeline/api/rag.py +++ b/src/pipeline/api/rag.py @@ -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)) diff --git a/src/pipeline/db/__init__.py b/src/pipeline/db/__init__.py index b158d68..af574ea 100644 --- a/src/pipeline/db/__init__.py +++ b/src/pipeline/db/__init__.py @@ -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") diff --git a/src/pipeline/main.py b/src/pipeline/main.py index 45ed23e..d679b99 100644 --- a/src/pipeline/main.py +++ b/src/pipeline/main.py @@ -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) diff --git a/src/pipeline/models/__init__.py b/src/pipeline/models/__init__.py index 5b73795..be236bd 100644 --- a/src/pipeline/models/__init__.py +++ b/src/pipeline/models/__init__.py @@ -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): diff --git a/src/pipeline/schemas/schemas.py b/src/pipeline/schemas/schemas.py index 82ca543..6d4d796 100644 --- a/src/pipeline/schemas/schemas.py +++ b/src/pipeline/schemas/schemas.py @@ -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 diff --git a/src/pipeline/services/rag.py b/src/pipeline/services/rag.py index 06f60f9..30b03ae 100644 --- a/src/pipeline/services/rag.py +++ b/src/pipeline/services/rag.py @@ -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() diff --git a/tests/pipeline/api/test_rag.py b/tests/pipeline/api/test_rag.py index 64f37a6..55de439 100644 --- a/tests/pipeline/api/test_rag.py +++ b/tests/pipeline/api/test_rag.py @@ -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())