From a2c02636fad4fdeb8baf48d7f979ed31d318582d Mon Sep 17 00:00:00 2001
From: w_xiaolizu <w_xiaolizu@kingsoft.com>
Date: Tue, 6 Jun 2023 12:28:02 +0800
Subject: [PATCH] =?UTF-8?q?=E5=B0=86=E6=95=B0=E6=8D=AE=E5=BA=93=E7=9A=84pr?=
 =?UTF-8?q?ompt=E6=8D=A2=E4=B8=BAUNIQUE=E5=94=AF=E4=B8=80=E5=B1=9E?=
 =?UTF-8?q?=E6=80=A7=EF=BC=8C=E5=B9=B6=E4=B8=94=E6=8F=92=E5=85=A5=E8=AF=AD?=
 =?UTF-8?q?=E5=8F=A5=E6=8D=A2=E6=88=90REPLACE=20INTO?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 func_box.py         |  8 ++++----
 prompt_generator.py | 11 +++++++----
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/func_box.py b/func_box.py
index b146647..225c183 100644
--- a/func_box.py
+++ b/func_box.py
@@ -287,8 +287,8 @@ def diff_list(txt='', percent=0.70, switch: list = None, lst: list = None, sp=15
     import difflib
     count_dict = {}
     if not lst:
-        lst = SqliteHandle('ai_common').get_prompt_value()
-        lst.update(SqliteHandle(f"ai_private_{hosts}").get_prompt_value())
+        lst = SqliteHandle('ai_common').get_prompt_value(txt)
+        lst.update(SqliteHandle(f"ai_private_{hosts}").get_prompt_value(txt))
     # diff 数据，根据precent系数归类数据
     for i in lst:
         found = False
@@ -378,10 +378,10 @@ def prompt_retrieval(is_all, hosts='', search=False):
     if '所有人' in is_all:
         for tab in SqliteHandle('ai_common').get_tables():
             if tab.startswith('prompt'):
-                data = SqliteHandle(tab).get_prompt_value()
+                data = SqliteHandle(tab).get_prompt_value(None)
                 if data: count_dict.update(data)
     elif '个人' in is_all:
-        data = SqliteHandle(f'prompt_{hosts}').get_prompt_value()
+        data = SqliteHandle(f'prompt_{hosts}').get_prompt_value(None)
         if data: count_dict.update(data)
     retrieval = []
     if count_dict != {}:
diff --git a/prompt_generator.py b/prompt_generator.py
index b814160..0cf1415 100644
--- a/prompt_generator.py
+++ b/prompt_generator.py
@@ -41,7 +41,7 @@ class SqliteHandle:
         self.__connect.close()
 
     def create_tab(self):
-        self.__cursor.execute(f"CREATE TABLE `{self.__table}` ( 'prompt' TEXT, 'result' TEXT)")
+        self.__cursor.execute(f"CREATE TABLE `{self.__table}` ('prompt' TEXT UNIQUE, 'result' TEXT)")
 
     def get_tables(self):
         all_tab = []
@@ -50,16 +50,19 @@ class SqliteHandle:
             all_tab.append(tab[0])
         return all_tab
 
-    def get_prompt_value(self):
+    def get_prompt_value(self, find):
         temp_all = {}
-        result = self.__cursor.execute(f"SELECT prompt, result FROM `{self.__table}`").fetchall()
+        if find:
+            result = self.__cursor.execute(f"SELECT prompt, result FROM `{self.__table}` WHERE prompt LIKE '%{find}%'").fetchall()
+        else:
+            result = self.__cursor.execute(f"SELECT prompt, result FROM `{self.__table}`").fetchall()
         for row in result:
             temp_all[row[0]] = row[1]
         return temp_all
 
     def inset_prompt(self, prompt: dict):
         for key in prompt:
-            self.__cursor.execute(f"INSERT INTO `{self.__table}` (prompt, result) VALUES (?, ?);", (str(key), str(prompt[key])))
+            self.__cursor.execute(f"REPLACE INTO `{self.__table}` (prompt, result) VALUES (?, ?);", (str(key), str(prompt[key])))
         self.__connect.commit()
 
     def delete_prompt(self):