109 lines
3.0 KiB
Python
109 lines
3.0 KiB
Python
"""数据库模型和管理器"""
|
|
import json
|
|
from datetime import datetime
|
|
from typing import Dict
|
|
|
|
from peewee import Model, CharField, TextField, IntegerField, DateTimeField, BooleanField, ForeignKeyField
|
|
from playhouse.sqlite_ext import SqliteExtDatabase
|
|
|
|
|
|
class BaseModel(Model):
|
|
"""基础模型类"""
|
|
pass
|
|
|
|
|
|
class Session(BaseModel):
|
|
"""聊天会话模型"""
|
|
name = CharField(max_length=255, unique=True)
|
|
model = CharField(max_length=50, default='deepseek-chat')
|
|
system_prompt = TextField(default='你是一个乐于助人的助手。')
|
|
created_at = DateTimeField(default=datetime.now)
|
|
updated_at = DateTimeField(default=datetime.now)
|
|
is_active = BooleanField(default=True)
|
|
metadata_json = TextField(default='{}')
|
|
|
|
@property
|
|
def metadata(self) -> Dict:
|
|
"""获取元数据"""
|
|
return json.loads(self.metadata_json) if self.metadata_json else {}
|
|
|
|
@metadata.setter
|
|
def metadata(self, value: Dict):
|
|
"""设置元数据"""
|
|
self.metadata_json = json.dumps(value, ensure_ascii=False)
|
|
|
|
def update_timestamp(self):
|
|
"""更新时间戳"""
|
|
self.updated_at = datetime.now()
|
|
self.save()
|
|
|
|
class Meta:
|
|
table_name = 'session'
|
|
|
|
|
|
class Message(BaseModel):
|
|
"""聊天消息模型"""
|
|
session = ForeignKeyField(Session, backref='messages', on_delete='CASCADE')
|
|
role = CharField(max_length=20, index=True)
|
|
content = TextField()
|
|
tokens = IntegerField(default=0)
|
|
timestamp = DateTimeField(default=datetime.now, index=True)
|
|
metadata_json = TextField(default='{}')
|
|
|
|
@property
|
|
def metadata(self) -> Dict:
|
|
"""获取元数据"""
|
|
return json.loads(self.metadata_json) if self.metadata_json else {}
|
|
|
|
@metadata.setter
|
|
def metadata(self, value: Dict):
|
|
"""设置元数据"""
|
|
self.metadata_json = json.dumps(value, ensure_ascii=False)
|
|
|
|
class Meta:
|
|
table_name = 'message'
|
|
indexes = (
|
|
(('session', 'timestamp'), False),
|
|
)
|
|
|
|
|
|
class DatabaseManager:
|
|
"""数据库管理器"""
|
|
|
|
def __init__(self, db_path: str = 'chat_sessions.db'):
|
|
"""
|
|
初始化数据库管理器
|
|
|
|
Args:
|
|
db_path: 数据库文件路径
|
|
"""
|
|
self.db = SqliteExtDatabase(
|
|
db_path,
|
|
pragmas={
|
|
'journal_mode': 'wal',
|
|
'cache_size': -1024 * 64,
|
|
'foreign_keys': 1,
|
|
'ignore_check_constraints': 0,
|
|
'synchronous': 1
|
|
}
|
|
)
|
|
self._init_database()
|
|
|
|
def _init_database(self):
|
|
"""初始化数据库连接和表"""
|
|
# 设置数据库
|
|
BaseModel._meta.database = self.db
|
|
Session._meta.database = self.db
|
|
Message._meta.database = self.db
|
|
|
|
# 连接数据库
|
|
self.db.connect()
|
|
|
|
# 创建表
|
|
self.db.create_tables([Session, Message], safe=True)
|
|
|
|
def close(self):
|
|
"""关闭数据库连接"""
|
|
if not self.db.is_closed():
|
|
self.db.close()
|