From b0b547825aad3e6b34c2f37eb2cb5f820344e0aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=A6=82=E5=A8=81?= Date: Thu, 8 Jan 2026 17:44:18 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20message=20=E8=A1=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pipeline/api/rag.py | 11 +++--- src/pipeline/config.py | 1 + src/pipeline/db/__init__.py | 32 ++++++++++++++--- src/pipeline/db/init.sql | 9 ----- src/pipeline/main.py | 2 +- src/pipeline/models/__init__.py | 63 +++++++++++++++++++++++++++++++-- src/pipeline/rag/__init__.py | 0 src/pipeline/schemas/schemas.py | 5 +-- src/pipeline/services/rag.py | 45 +++++++++++++++++++---- 9 files changed, 138 insertions(+), 30 deletions(-) delete mode 100644 src/pipeline/rag/__init__.py diff --git a/src/pipeline/api/rag.py b/src/pipeline/api/rag.py index 1c83817..1b4a961 100644 --- a/src/pipeline/api/rag.py +++ b/src/pipeline/api/rag.py @@ -16,9 +16,12 @@ async def stream_chat(body: ChatRequest): service = RagService() async def event_generator(): - yield f"data: {json.dumps({'message': '', 'status': 'start'})}\n\n" - async for tokens in service.stream_chat(body.query): - yield f"data: {json.dumps({'message': tokens, 'status': 'process'})}\n\n" - yield f"data: {json.dumps({'message': '', 'status': 'end'})}\n\n" + 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): + 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") diff --git a/src/pipeline/config.py b/src/pipeline/config.py index 3985649..9ff94fe 100644 --- a/src/pipeline/config.py +++ b/src/pipeline/config.py @@ -4,6 +4,7 @@ from pydantic import MySQLDsn class Settings(BaseSettings): logger_level: str + env: str version: str port: int host: str diff --git a/src/pipeline/db/__init__.py b/src/pipeline/db/__init__.py index 9d19ffa..b158d68 100644 --- a/src/pipeline/db/__init__.py +++ b/src/pipeline/db/__init__.py @@ -9,9 +9,30 @@ DEFAULT_MODELS = {"models": ["src.pipeline.models"]} 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): +async def init_tortoise( + app: FastAPI, + db_url: str, + modules: dict = DEFAULT_MODELS, + schema_sql: str = DEFAULT_INIT_SQL, + generate_schemas: bool = False, +): + """ + 初始化数据库 + + :param app: fastapi 实例 + :type app: FastAPI + :param db_url: 数据库配置地址 + :type db_url: str + :param modules: 模型定位目录 + :type modules: dict + :param schema_sql: 数据库初始化 sql 文件 + :type schema_sql: str + :param generate_schemas: 根据模型创建数据表,生产环境建议关闭 + :type generate_schemas: bool + """ config = generate_config(db_url, app_modules=modules) - await RegisterTortoise(app, config=config) + + await RegisterTortoise(app, config=config, generate_schemas=generate_schemas) if schema_sql: conn = Tortoise.get_connection("default") @@ -20,7 +41,8 @@ async def init_tortoise(app: FastAPI, db_url: str, modules: dict = DEFAULT_MODEL logger.debug(f"数据库初始化文件: {schema_sql}") sql_file = Path(schema_sql) if sql_file.exists(): - logger.debug("初始化数据库") sql = sql_file.read_text() - logger.debug(f'sql:\n{sql}') - await conn.execute_script(sql) + if sql: + logger.debug("初始化数据库") + logger.debug(f"sql:\n{sql}") + await conn.execute_script(sql) diff --git a/src/pipeline/db/init.sql b/src/pipeline/db/init.sql index 92d5dcd..e69de29 100644 --- a/src/pipeline/db/init.sql +++ b/src/pipeline/db/init.sql @@ -1,9 +0,0 @@ --- tb_users table -CREATE TABLE IF NOT EXISTS tb_users ( - id INT AUTO_INCREMENT PRIMARY KEY, - username VARCHAR(50) NOT NULL UNIQUE, - password VARCHAR(128) NOT NULL, - user_status TINYINT(1) NOT NULL DEFAULT 1, -- 0=禁用, 1=启用 - create_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - update_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP -); \ No newline at end of file diff --git a/src/pipeline/main.py b/src/pipeline/main.py index b05ac31..45ed23e 100644 --- a/src/pipeline/main.py +++ b/src/pipeline/main.py @@ -10,7 +10,7 @@ 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)) + await init_tortoise(app, db_url=str(config.mysql_dsn), generate_schemas=config.env == "DEBUG") yield diff --git a/src/pipeline/models/__init__.py b/src/pipeline/models/__init__.py index 3827fba..5b73795 100644 --- a/src/pipeline/models/__init__.py +++ b/src/pipeline/models/__init__.py @@ -1,21 +1,80 @@ -from email.policy import default from tortoise import fields +from tortoise.transactions import in_transaction from tortoise.models import Model +from uuid import uuid4 class BaseModel(Model): id = fields.IntField(pk=True) create_at = fields.DatetimeField(auto_now_add=True) update_at = fields.DatetimeField(auto_now=True) + is_deleted = fields.IntField(default=0) class Meta: abstract = True + @classmethod + def generate_uuid(cls): + return str(uuid4()) + + @classmethod + async def select( + cls, + params: dict, + order_by: list[str] = ["-id"], + page: int = 1, + page_size: int | None = None, + ): + """ + SELECT 语句 + + :param cls: + :param params: 筛选条件 + :type params: dict + :param order_by: + :type order_by: list[str] + :param page: 页码 1 开始 + :type page: int | None + :param page_size: 页数,不设置返回全部 + :type page_size: int | None + """ + if page and page_size: + return await cls.filter(**params).order_by(*order_by).limit(page_size).offset((page - 1) * page_size) + else: + return await cls.filter(**params).order_by(*order_by).all() + class User(BaseModel): username = fields.CharField(max_length=50, unique=True) password = fields.CharField(max_length=128) - user_status = fields.IntField(default=1) # 0=禁用, 1=启用 + user_status = fields.IntField(default=1) # 0=禁用, 1=启用 class Meta: table = "tb_users" + + +class Chat(BaseModel): + user_id = fields.IntField(null=True) + title = fields.CharField(max_length=50) + uuid = fields.UUIDField(index=True) + + class Meta: + table = "tb_chats" + + +class Message(BaseModel): + user_id = fields.IntField(null=True) + chat_id = fields.IntField(null=True, index=True) + uuid = fields.UUIDField(index=True) + content = fields.JSONField(null=True) + + class Meta: + table = "tb_messages" + + def to_llm_messages(self) -> list[dict]: + llm_messages = [] + if "question" in self.content and self.content["question"]: + llm_messages.append({"role": "user", "content": self.content["question"]}) + if "answer" in self.content and self.content["answer"]: + llm_messages.append({"role": "assistant", "content": self.content["answer"]}) + return llm_messages diff --git a/src/pipeline/rag/__init__.py b/src/pipeline/rag/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/pipeline/schemas/schemas.py b/src/pipeline/schemas/schemas.py index 20b2647..82ca543 100644 --- a/src/pipeline/schemas/schemas.py +++ b/src/pipeline/schemas/schemas.py @@ -1,4 +1,5 @@ from pydantic import BaseModel - +from typing import Optional class ChatRequest(BaseModel): - query: str \ No newline at end of file + query: str + chat_uuid: Optional[str] = None diff --git a/src/pipeline/services/rag.py b/src/pipeline/services/rag.py index 6ace158..06f60f9 100644 --- a/src/pipeline/services/rag.py +++ b/src/pipeline/services/rag.py @@ -1,16 +1,47 @@ from src.pipeline.core import llm from src.pipeline.utils import logger +from src.pipeline.models import Chat as TChat, Message as TMessage, User as TUser + class RagService: def __init__(self): pass - async def stream_chat(self, query:str): - logger.debug(query) - request = llm.client.stream_chat( - messages=[ - {"role": "user", "content": "query"}, - ] + async def stream_chat( + self, + query: str, + chat_uuid: str | None = None, + top_n: int = 10, + user_id: int = 0, + ): + # 构造数据 + t_chat, _ = await TChat.get_or_create( + { + "title": query.strip()[:50], + "user_id": user_id, + "uuid": chat_uuid or TChat.generate_uuid(), + } ) + t_message = TMessage( + user_id=user_id, + chat_id=t_chat.id, + uuid=TMessage.generate_uuid(), + role="assistant", + ) + t_messages_list = await TMessage.select({"chat_id": t_chat.id}, page_size=top_n - 1) + + messages = [m for x in t_messages_list[::-1] for m in x.to_llm_messages()] + messages.append({"role": "user", "content": query}) + + # 对话 + request = llm.client.stream_chat(messages=messages) + content = "" async for chunk in request: - yield chunk + content += chunk + yield chunk, str(t_chat.uuid), str(t_message.uuid) + + # 保存数据 + t_message.content = { + "question": query, "answer": content + } + await t_message.save()