From ad74a7ca7d27d705ece202db70cb07577a1625d9 Mon Sep 17 00:00:00 2001
From: 997146918 <997146918@qq.com>
Date: Thu, 3 Jul 2025 19:01:27 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E6=95=B0=E6=8D=AE=E5=BA=93?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 AIGC/AICore.py                       |  4 +-
 AIGC/DatabaseHandle.py               | 64 ++++++++++++++++++++++++++++
 AIGC/main.py                         |  3 ++
 AIGC/数据库工具/SQLiteStudio/qt.conf |  2 +
 4 files changed, 71 insertions(+), 2 deletions(-)
 create mode 100644 AIGC/DatabaseHandle.py
 create mode 100644 AIGC/数据库工具/SQLiteStudio/qt.conf

diff --git a/AIGC/AICore.py b/AIGC/AICore.py
index 378e2ca..b01ae16 100644
--- a/AIGC/AICore.py
+++ b/AIGC/AICore.py
@@ -11,8 +11,8 @@ class AICore:
 
     def __init__(self, model):
          #初始化ollama客户端
-        ollamaClient = Client(host='http://localhost:11434', headers={'x-some-header': 'some-value'})
-        response = ollamaClient.show(model)
+        self.ollamaClient = Client(host='http://localhost:11434', headers={'x-some-header': 'some-value'})
+        response = self.ollamaClient.show(model)
         modelMaxTokens = response.modelinfo['qwen2.context_length']
         
     def getPromptToken(self, prompt)-> int: 
diff --git a/AIGC/DatabaseHandle.py b/AIGC/DatabaseHandle.py
new file mode 100644
index 0000000..0f3d689
--- /dev/null
+++ b/AIGC/DatabaseHandle.py
@@ -0,0 +1,64 @@
+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
\ No newline at end of file
diff --git a/AIGC/main.py b/AIGC/main.py
index 1fdb06c..e3c9287 100644
--- a/AIGC/main.py
+++ b/AIGC/main.py
@@ -12,6 +12,7 @@ from fastapi.websockets import WebSocketState
 from h11 import ConnectionClosed
 import uvicorn
 from AICore import AICore
+from DatabaseHandle import DatabaseHandle
 from Utils.AIGCLog import AIGCLog
 
 
@@ -238,6 +239,8 @@ if __name__ == "__main__":
 
     # Test
     aicore.getPromptToken("测试功能")
+    database = DatabaseHandle()
+    
     asyncio.run(
         generateAIChat(promptStr = f"""
             #你是一个游戏NPC对话生成器。请严格按以下要求生成两个角色的日常对话
diff --git a/AIGC/数据库工具/SQLiteStudio/qt.conf b/AIGC/数据库工具/SQLiteStudio/qt.conf
new file mode 100644
index 0000000..01f4e8c
--- /dev/null
+++ b/AIGC/数据库工具/SQLiteStudio/qt.conf
@@ -0,0 +1,2 @@
+[Paths]
+Prefix=..