Merge pull request #38 from codefuse-ai/modelcache_dev_mm

fix: sqlite does not correctly return the number of rows deleted

Author: peng3307165

modelcache/embedding/__init__.py

@@ -12,7 +12,7 @@ def Huggingface(model="sentence-transformers/all-mpnet-base-v2"):
     return huggingface.Huggingface(model)
 
 
-def Data2VecAudio(model="facebook/data2vec-audio-base-960h"):
+def Data2VecAudio(model="model/text2vec-base-chinese/"):
     return data2vec.Data2VecAudio(model)
 
 

modelcache/manager/scalar_data/sql_storage_sqlite.py

@@ -17,19 +17,6 @@ def __init__(
         self.create()
 
     def create(self):
-        # answer_table_sql = """CREATE TABLE IF NOT EXISTS `modelcache_llm_answer` (
-        # `id` bigint(20) NOT NULL AUTO_INCREMENT comment '主键',
-        # `gmt_create` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP comment '创建时间',
-        # `gmt_modified` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP comment '修改时间',
-        # `question` text NOT NULL comment 'question',
-        # `answer` text NOT NULL comment 'answer',
-        # `answer_type` int(11) NOT NULL comment 'answer_type',
-        # `hit_count` int(11) NOT NULL DEFAULT '0' comment 'hit_count',
-        # `model` varchar(1000) NOT NULL comment 'model',
-        # `embedding_data` blob NOT NULL comment 'embedding_data',
-        # PRIMARY KEY(`id`)
-        # ) AUTO_INCREMENT = 1 DEFAULT CHARSET = utf8mb4 COMMENT = 'modelcache_llm_answer';
-        # """
         answer_table_sql = """CREATE TABLE IF NOT EXISTS modelcache_llm_answer (
                 id INTEGER PRIMARY KEY AUTOINCREMENT,
                 gmt_create TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
@@ -43,21 +30,6 @@ def create(self):
                 );
                 """
 
-        # log_table_sql = """CREATE TABLE IF NOT EXISTS `modelcache_query_log` (
-        # `id` bigint(20) NOT NULL AUTO_INCREMENT comment '主键',
-        # `gmt_create` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP comment '创建时间',
-        # `gmt_modified` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP comment '修改时间',
-        # `error_code` int(11) NOT NULL comment 'errorCode',
-        # `error_desc` varchar(1000) NOT NULL comment 'errorDesc',
-        # `cache_hit` varchar(100) NOT NULL comment 'cacheHit',
-        # `delta_time` float NOT NULL comment 'delta_time',
-        # `model` varchar(1000) NOT NULL comment 'model',
-        # `query` text NOT NULL comment 'query',
-        # `hit_query` text NOT NULL comment 'hitQuery',
-        # `answer` text NOT NULL comment 'answer',
-        # PRIMARY KEY(`id`)
-        # ) AUTO_INCREMENT = 1 DEFAULT CHARSET = utf8mb4 COMMENT = 'modelcache_query_log';
-        # """
         log_table_sql = """CREATE TABLE IF NOT EXISTS modelcache_query_log (
                 id INTEGER PRIMARY KEY AUTOINCREMENT,
                 gmt_create TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
@@ -195,17 +167,21 @@ def mark_deleted(self, keys):
 
     def model_deleted(self, model_name):
         table_name = "modelcache_llm_answer"
-        delete_sql = "Delete from {} WHERE model='{}'".format(table_name, model_name)
+        delete_sql = "Delete from {} WHERE model=?".format(table_name)
         conn = sqlite3.connect(self._url)
         try:
             cursor = conn.cursor()
-            resp = cursor.execute(delete_sql)
+            cursor.execute(delete_sql, (model_name,))
             conn.commit()
+            # get delete rows
+            deleted_rows_count = cursor.rowcount
             cursor.close()
-            conn.close()
+        except sqlite3.Error as e:
+            print(f"SQLite error: {e}")
+            deleted_rows_count = 0  # if except, return 0
         finally:
             conn.close()
-        return resp
+        return deleted_rows_count
 
     def clear_deleted_data(self):
         pass