64 lines
2.1 KiB
Python
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
|