diff --git a/AIGC/DatabaseHandle.py b/AIGC/DatabaseHandle.py index 0f3d689..9c23d28 100644 --- a/AIGC/DatabaseHandle.py +++ b/AIGC/DatabaseHandle.py @@ -13,11 +13,11 @@ class DatabaseHandle: conn.execute(''' CREATE TABLE IF NOT EXISTS characters (id INTEGER PRIMARY KEY, - name TEXT not NULL, + name TEXT not NULL UNIQUE, age INTEGER, personality TEXT, profession TEXT, - characterBackgrount TEXT, + characterBackground TEXT, chat_style TEXT ) ''') @@ -45,20 +45,70 @@ class DatabaseHandle: def add_character(self, data:dict): """添加角色数据""" - pass + with self._get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + INSERT INTO characters ( + name, age, personality, profession, characterBackground, + chat_style + ) VALUES (?, ?, ?, ?, ?, ?) + ON CONFLICT(name) DO NOTHING + ''', ( + data["name"], data["age"], data["personality"], + data["profession"], data["characterBackground"], + data["chat_style"] + )) + conn.commit() + return cursor.lastrowid + + def get_character_byname(self, name:str) ->list: + """ + 根据角色名称查询数据(name为空时返回全部数据) + :param name: 角色名称(精确匹配,None 或空字符串时返回全部) + :return: 角色数据字典列表 + """ + with self._get_connection() as conn: + cursor = conn.cursor() + + # 动态构建SQL语句 + sql = "SELECT * FROM characters" + params = () + + if name: # 当name非空时添加条件 + sql += " WHERE name = ?" + params = (name,) + + cursor.execute(sql, params) + + # 转换结果为字典列表 + columns = [col[0] for col in cursor.description] + return [dict(zip(columns, row)) for row in cursor.fetchall()] - # 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 \ No newline at end of file + def add_chat(self, data:dict): + """添加聊天数据""" + with self._get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + INSERT INTO chat_records ( + character_ids, chat + ) VALUES (?, ?) + ''', ( + data["character_ids"], data["chat"] + )) + conn.commit() + return cursor.lastrowid + + def get_chats_by_character_id(self, character_id: int) -> list: + """ + 根据角色ID查询聊天记录(target_id为空时返回全部数据) + :param target_id: 目标角色ID(None时返回全部记录) + :return: 聊天记录字典列表 + """ + with self._get_connection() as conn: + cursor = conn.cursor() + sql = "SELECT * FROM chat_records WHERE ',' || character_ids || ',' LIKE '%,' || ? || ',%'" + params = (str(character_id)) + cursor.execute(sql, params) + # 转换结果为字典列表 + columns = [col[0] for col in cursor.description] + return [dict(zip(columns, row)) for row in cursor.fetchall()] \ No newline at end of file diff --git a/AIGC/main.py b/AIGC/main.py index e3c9287..591f78d 100644 --- a/AIGC/main.py +++ b/AIGC/main.py @@ -237,10 +237,22 @@ if __name__ == "__main__": server_thread.daemon = True # 设为守护线程(主程序退出时自动终止) server_thread.start() - # Test - aicore.getPromptToken("测试功能") - database = DatabaseHandle() + #Test database + database = DatabaseHandle() + id = database.add_character({"name":"张三","age":35,"personality":"成熟稳重/惜字如金","profession":"阿里巴巴算法工程师" + ,"characterBackground":"浙大计算机系毕业,专注AI优化项目","chat_style":"请在对话中表现出专业、冷静、惜字如金。用口语化的方式简短回答"}) + + id = database.add_character({"name":"李明","age":30,"personality":"活泼健谈","profession":"产品经理" + ,"characterBackground":"公司资深产品经理","chat_style":"热情"}) + + characters = database.get_character_byname("") + #chat_id = database.add_chat({"character_ids":"1,2","chat":"张三:[第一轮发言] 李明:[第一轮回应] 张三:[第二轮发言] 李明:[第二轮回应"}) + chat = database.get_chats_by_character_id(3) + if id == 0: + logger.log(logging.ERROR, f"角色 张三已经添加到数据库") + # Test AI + aicore.getPromptToken("测试功能") asyncio.run( generateAIChat(promptStr = f""" #你是一个游戏NPC对话生成器。请严格按以下要求生成两个角色的日常对话