Project02/AIGC/DatabaseHandle.py
2025-07-03 19:01:27 +08:00

64 lines
2.1 KiB
Python

from contextlib import contextmanager
import sqlite3
class DatabaseHandle:
def __init__(self, db_path = "data.db"):
self.db_path = db_path
self._init_db()
def _init_db(self):
"""创建表"""
with self._get_connection() as conn:
#创建角色表
conn.execute('''
CREATE TABLE IF NOT EXISTS characters
(id INTEGER PRIMARY KEY,
name TEXT not NULL,
age INTEGER,
personality TEXT,
profession TEXT,
characterBackgrount TEXT,
chat_style TEXT
)
''')
#创建聊天记录表
conn.execute('''
CREATE TABLE IF NOT EXISTS chat_records
(id INTEGER PRIMARY KEY,
character_ids TEXT not NULL,
chat TEXT,
time DATETIME DEFAULT CURRENT_TIMESTAMP
)
''')
@contextmanager
def _get_connection(self):
conn = sqlite3.connect(self.db_path)
try:
yield conn
conn.commit()
except Exception:
conn.rollback()
raise
finally:
conn.close()
def add_character(self, data:dict):
"""添加角色数据"""
pass
# def add_employee(self, data: dict):
# """插入员工数据"""
# with self._get_connection() as conn:
# cursor = conn.cursor()
# cursor.execute('''
# INSERT INTO employees (
# name, age, personality, profession, background,
# conversation_scene, language_style
# ) VALUES (?, ?, ?, ?, ?, ?, ?)
# ''', (
# data["姓名"], data["年龄"], data["性格"],
# data["职业"], data["背景"],
# data["对话场景"], data["语言风格"]
# ))
# return cursor.lastrowid # 返回插入ID