添加数据库

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: 

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

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

AIGC/数据库工具/SQLiteStudio/qt.conf

@@ -0,0 +1,2 @@
+[Paths]
+Prefix=..