feat: message 表
This commit is contained in:
parent
ef68e3cbd8
commit
b0b547825a
|
|
@ -16,9 +16,12 @@ async def stream_chat(body: ChatRequest):
|
||||||
service = RagService()
|
service = RagService()
|
||||||
|
|
||||||
async def event_generator():
|
async def event_generator():
|
||||||
yield f"data: {json.dumps({'message': '', 'status': 'start'})}\n\n"
|
chat_uuid, message_uuid = "", ""
|
||||||
async for tokens in service.stream_chat(body.query):
|
yield f"data: {json.dumps({'status': 'start'})}\n\n"
|
||||||
yield f"data: {json.dumps({'message': tokens, 'status': 'process'})}\n\n"
|
async for tokens, uuid_0, uuid_1 in service.stream_chat(body.query, chat_uuid=body.chat_uuid):
|
||||||
yield f"data: {json.dumps({'message': '', 'status': 'end'})}\n\n"
|
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")
|
return StreamingResponse(event_generator(), media_type="text/event-stream")
|
||||||
|
|
|
||||||
|
|
@ -4,6 +4,7 @@ from pydantic import MySQLDsn
|
||||||
|
|
||||||
class Settings(BaseSettings):
|
class Settings(BaseSettings):
|
||||||
logger_level: str
|
logger_level: str
|
||||||
|
env: str
|
||||||
version: str
|
version: str
|
||||||
port: int
|
port: int
|
||||||
host: str
|
host: str
|
||||||
|
|
|
||||||
|
|
@ -9,9 +9,30 @@ DEFAULT_MODELS = {"models": ["src.pipeline.models"]}
|
||||||
DEFAULT_INIT_SQL = "src/pipeline/db/init.sql"
|
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)
|
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:
|
if schema_sql:
|
||||||
conn = Tortoise.get_connection("default")
|
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}")
|
logger.debug(f"数据库初始化文件: {schema_sql}")
|
||||||
sql_file = Path(schema_sql)
|
sql_file = Path(schema_sql)
|
||||||
if sql_file.exists():
|
if sql_file.exists():
|
||||||
logger.debug("初始化数据库")
|
|
||||||
sql = sql_file.read_text()
|
sql = sql_file.read_text()
|
||||||
logger.debug(f'sql:\n{sql}')
|
if sql:
|
||||||
|
logger.debug("初始化数据库")
|
||||||
|
logger.debug(f"sql:\n{sql}")
|
||||||
await conn.execute_script(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):
|
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))
|
await init_tortoise(app, db_url=str(config.mysql_dsn), generate_schemas=config.env == "DEBUG")
|
||||||
|
|
||||||
yield
|
yield
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,16 +1,48 @@
|
||||||
from email.policy import default
|
|
||||||
from tortoise import fields
|
from tortoise import fields
|
||||||
|
from tortoise.transactions import in_transaction
|
||||||
from tortoise.models import Model
|
from tortoise.models import Model
|
||||||
|
from uuid import uuid4
|
||||||
|
|
||||||
|
|
||||||
class BaseModel(Model):
|
class BaseModel(Model):
|
||||||
id = fields.IntField(pk=True)
|
id = fields.IntField(pk=True)
|
||||||
create_at = fields.DatetimeField(auto_now_add=True)
|
create_at = fields.DatetimeField(auto_now_add=True)
|
||||||
update_at = fields.DatetimeField(auto_now=True)
|
update_at = fields.DatetimeField(auto_now=True)
|
||||||
|
is_deleted = fields.IntField(default=0)
|
||||||
|
|
||||||
class Meta:
|
class Meta:
|
||||||
abstract = True
|
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):
|
class User(BaseModel):
|
||||||
username = fields.CharField(max_length=50, unique=True)
|
username = fields.CharField(max_length=50, unique=True)
|
||||||
|
|
@ -19,3 +51,30 @@ class User(BaseModel):
|
||||||
|
|
||||||
class Meta:
|
class Meta:
|
||||||
table = "tb_users"
|
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 pydantic import BaseModel
|
||||||
|
from typing import Optional
|
||||||
class ChatRequest(BaseModel):
|
class ChatRequest(BaseModel):
|
||||||
query: str
|
query: str
|
||||||
|
chat_uuid: Optional[str] = None
|
||||||
|
|
|
||||||
|
|
@ -1,16 +1,47 @@
|
||||||
from src.pipeline.core import llm
|
from src.pipeline.core import llm
|
||||||
from src.pipeline.utils import logger
|
from src.pipeline.utils import logger
|
||||||
|
from src.pipeline.models import Chat as TChat, Message as TMessage, User as TUser
|
||||||
|
|
||||||
|
|
||||||
class RagService:
|
class RagService:
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
async def stream_chat(self, query:str):
|
async def stream_chat(
|
||||||
logger.debug(query)
|
self,
|
||||||
request = llm.client.stream_chat(
|
query: str,
|
||||||
messages=[
|
chat_uuid: str | None = None,
|
||||||
{"role": "user", "content": "query"},
|
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:
|
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