feat: message 表

This commit is contained in:
李如威 2026-01-08 17:44:18 +08:00
parent ef68e3cbd8
commit b0b547825a
9 changed files with 138 additions and 30 deletions

View File

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

View File

@ -4,6 +4,7 @@ from pydantic import MySQLDsn
class Settings(BaseSettings):
logger_level: str
env: str
version: str
port: int
host: str

View File

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

View File

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

View File

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

View File

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

View File

@ -1,4 +1,5 @@
from pydantic import BaseModel
from typing import Optional
class ChatRequest(BaseModel):
query: str
chat_uuid: Optional[str] = None

View File

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