添加数据库的数据增加和查询功能

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 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
+    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_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()]

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对话生成器。请严格按以下要求生成两个角色的日常对话