feat: message 表
This commit is contained in:
parent
ef68e3cbd8
commit
b0b547825a
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from pydantic import MySQLDsn
|
|||
|
||||
class Settings(BaseSettings):
|
||||
logger_level: str
|
||||
env: str
|
||||
version: str
|
||||
port: int
|
||||
host: str
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
);
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
from pydantic import BaseModel
|
||||
|
||||
from typing import Optional
|
||||
class ChatRequest(BaseModel):
|
||||
query: str
|
||||
chat_uuid: Optional[str] = None
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Reference in New Issue