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=..