fewfefffwefweef

ai_main/database.py

@@ -0,0 +1,108 @@
+"""数据库模型和管理器"""
+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()