Adding an SQLite database; incorporating a local demo Flask service.

Author: peng3307165

flask4modelcache_demo.py

@@ -0,0 +1,169 @@
+# -*- coding: utf-8 -*-
+import time
+from flask import Flask, request
+import logging
+import json
+from modelcache import cache
+from modelcache.adapter import adapter
+from modelcache.manager import CacheBase, VectorBase, get_data_manager
+from modelcache.similarity_evaluation.distance import SearchDistanceEvaluation
+from modelcache.processor.pre import query_multi_splicing
+from modelcache.processor.pre import insert_multi_splicing
+from concurrent.futures import ThreadPoolExecutor
+from modelcache.utils.model_filter import model_blacklist_filter
+from modelcache.embedding import Data2VecAudio
+
+# 创建一个Flask实例
+app = Flask(__name__)
+
+
+def response_text(cache_resp):
+    return cache_resp['data']
+
+
+def save_query_info(result, model, query, delta_time_log):
+    cache.data_manager.save_query_resp(result, model=model, query=json.dumps(query, ensure_ascii=False),
+                                       delta_time=delta_time_log)
+
+
+def response_hitquery(cache_resp):
+    return cache_resp['hitQuery']
+
+
+data2vec = Data2VecAudio()
+data_manager = get_data_manager(CacheBase("sqlite"), VectorBase("faiss", dimension=data2vec.dimension))
+
+
+cache.init(
+    embedding_func=data2vec.to_embeddings,
+    data_manager=data_manager,
+    similarity_evaluation=SearchDistanceEvaluation(),
+    query_pre_embedding_func=query_multi_splicing,
+    insert_pre_embedding_func=insert_multi_splicing,
+)
+
+# cache.set_openai_key()
+global executor
+executor = ThreadPoolExecutor(max_workers=6)
+
+
+@app.route('/welcome')
+def first_flask():  # 视图函数
+    return 'hello, modelcache!'
+
+
+@app.route('/modelcache', methods=['GET', 'POST'])
+def user_backend():
+    try:
+        if request.method == 'POST':
+            request_data = request.json
+        elif request.method == 'GET':
+            request_data = request.args
+        param_dict = json.loads(request_data)
+    except Exception as e:
+        result = {"errorCode": 101, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
+                  "answer": ''}
+        cache.data_manager.save_query_resp(result, model='', query='', delta_time=0)
+        return json.dumps(result)
+
+    # param parsing
+    try:
+        request_type = param_dict.get("type")
+        scope = param_dict.get("scope")
+        if scope is not None:
+            model = scope.get('model')
+            model = model.replace('-', '_')
+            model = model.replace('.', '_')
+        query = param_dict.get("query")
+        chat_info = param_dict.get("chat_info")
+        if request_type is None or request_type not in ['query', 'insert', 'detox', 'remove']:
+            result = {"errorCode": 102,
+                      "errorDesc": "type exception, should one of ['query', 'insert', 'detox', 'remove']",
+                      "cacheHit": False, "delta_time": 0, "hit_query": '', "answer": ''}
+            cache.data_manager.save_query_resp(result, model=model, query='', delta_time=0)
+            return json.dumps(result)
+    except Exception as e:
+        result = {"errorCode": 103, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
+                  "answer": ''}
+        return json.dumps(result)
+
+    # model filter
+    filter_resp = model_blacklist_filter(model, request_type)
+    if isinstance(filter_resp, dict):
+        return json.dumps(filter_resp)
+
+    if request_type == 'query':
+        try:
+            start_time = time.time()
+            response = adapter.ChatCompletion.create_query(
+                scope={"model": model},
+                query=query
+            )
+            delta_time = '{}s'.format(round(time.time() - start_time, 2))
+            if response is None:
+                result = {"errorCode": 0, "errorDesc": '', "cacheHit": False, "delta_time": delta_time, "hit_query": '',
+                          "answer": ''}
+            elif response in ['adapt_query_exception']:
+                result = {"errorCode": 201, "errorDesc": response, "cacheHit": False, "delta_time": delta_time,
+                          "hit_query": '', "answer": ''}
+            else:
+                answer = response_text(response)
+                hit_query = response_hitquery(response)
+                result = {"errorCode": 0, "errorDesc": '', "cacheHit": True, "delta_time": delta_time,
+                          "hit_query": hit_query, "answer": answer}
+            delta_time_log = round(time.time() - start_time, 2)
+            future = executor.submit(save_query_info, result, model, query, delta_time_log)
+        except Exception as e:
+            result = {"errorCode": 202, "errorDesc": e, "cacheHit": False, "delta_time": 0,
+                      "hit_query": '', "answer": ''}
+            logging.info('result: {}'.format(result))
+
+        return json.dumps(result, ensure_ascii=False)
+
+    if request_type == 'insert':
+        try:
+            try:
+                response = adapter.ChatCompletion.create_insert(
+                    model=model,
+                    chat_info=chat_info
+                )
+            except Exception as e:
+                result = {"errorCode": 303, "errorDesc": e, "writeStatus": "exception"}
+                return json.dumps(result, ensure_ascii=False)
+
+            if response in ['adapt_insert_exception']:
+                result = {"errorCode": 301, "errorDesc": response, "writeStatus": "exception"}
+            elif response == 'success':
+                result = {"errorCode": 0, "errorDesc": "", "writeStatus": "success"}
+            else:
+                result = {"errorCode": 302, "errorDesc": response,
+                          "writeStatus": "exception"}
+            return json.dumps(result, ensure_ascii=False)
+        except Exception as e:
+            result = {"errorCode": 304, "errorDesc": e, "writeStatus": "exception"}
+            return json.dumps(result, ensure_ascii=False)
+
+    if request_type == 'remove':
+        remove_type = param_dict.get("remove_type")
+        id_list = param_dict.get("id_list", [])
+
+        response = adapter.ChatCompletion.create_remove(
+            model=model,
+            remove_type=remove_type,
+            id_list=id_list
+        )
+
+        if not isinstance(response, dict):
+            result = {"errorCode": 401, "errorDesc": "", "response": response, "removeStatus": "exception"}
+            return json.dumps(result)
+
+        state = response.get('status')
+        if state == 'success':
+            result = {"errorCode": 0, "errorDesc": "", "response": response, "writeStatus": "success"}
+        else:
+            result = {"errorCode": 402, "errorDesc": "", "response": response, "writeStatus": "exception"}
+        return json.dumps(result)
+
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=5000, debug=True)

modelcache/manager/scalar_data/sql_storage_sqlite.py

@@ -0,0 +1,179 @@
+# -*- coding: utf-8 -*-
+import os
+import time
+
+import pymysql
+import json
+import base64
+from typing import List
+from modelcache.manager.scalar_data.base import CacheStorage, CacheData
+from DBUtils.PooledDB import PooledDB
+
+
+class SQLStorage(CacheStorage):
+    def __init__(
+        self,
+        db_type: str = "mysql",
+        config=None,
+        url="sqlite:///./sqlite.db"
+    ):
+        if db_type in ["mysql", "oceanbase"]:
+            self.host = config.get('mysql', 'host')
+            self.port = int(config.get('mysql', 'port'))
+            self.username = config.get('mysql', 'username')
+            self.password = config.get('mysql', 'password')
+            self.database = config.get('mysql', 'database')
+
+            self.pool = PooledDB(
+                creator=pymysql,
+                host=self.host,
+                user=self.username,
+                password=self.password,
+                port=self.port,
+                database=self.database
+            )
+        elif db_type == 'sqlite':
+            self._url = url
+
+    def create(self):
+        pass
+
+    def _insert(self, data: List):
+        answer = data[0]
+        question = data[1]
+        embedding_data = data[2]
+        model = data[3]
+        answer_type = 0
+        embedding_data = embedding_data.tobytes()
+
+        table_name = "cache_codegpt_answer"
+        insert_sql = "INSERT INTO {} (question, answer, answer_type, model, embedding_data) VALUES (%s, %s, %s, %s, _binary%s)".format(table_name)
+
+        conn = self.pool.connection()
+        try:
+            with conn.cursor() as cursor:
+                # 执行插入数据操作
+                values = (question, answer, answer_type, model, embedding_data)
+                cursor.execute(insert_sql, values)
+                conn.commit()
+                id = cursor.lastrowid
+        finally:
+            # 关闭连接，将连接返回给连接池
+            conn.close()
+        return id
+
+    def batch_insert(self, all_data: List[CacheData]):
+        ids = []
+        for data in all_data:
+            ids.append(self._insert(data))
+        return ids
+
+    def insert_query_resp(self, query_resp, **kwargs):
+        error_code = query_resp.get('errorCode')
+        error_desc = query_resp.get('errorDesc')
+        cache_hit = query_resp.get('cacheHit')
+        model = kwargs.get('model')
+        query = kwargs.get('query')
+        delta_time = kwargs.get('delta_time')
+        hit_query = query_resp.get('hit_query')
+        answer = query_resp.get('answer')
+
+        if isinstance(hit_query, list):
+            hit_query = json.dumps(hit_query, ensure_ascii=False)
+
+        table_name = "cache_query_log_info"
+        insert_sql = "INSERT INTO {} (error_code, error_desc, cache_hit, model, query, delta_time, hit_query, answer) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)".format(table_name)
+        conn = self.pool.connection()
+        try:
+            with conn.cursor() as cursor:
+                # 执行插入数据操作
+                values = (error_code, error_desc, cache_hit, model, query, delta_time, hit_query, answer)
+                cursor.execute(insert_sql, values)
+                conn.commit()
+        finally:
+            # 关闭连接，将连接返回给连接池
+            conn.close()
+        return id
+
+    def get_data_by_id(self, key: int):
+        table_name = "cache_codegpt_answer"
+        query_sql = "select question, answer, embedding_data, model from {} where id={}".format(table_name, key)
+        conn_start = time.time()
+        conn = self.pool.connection()
+
+        search_start = time.time()
+        try:
+            with conn.cursor() as cursor:
+                # 执行数据库操作
+                cursor.execute(query_sql)
+                resp = cursor.fetchone()
+        finally:
+            # 关闭连接，将连接返回给连接池
+            conn.close()
+
+        if resp is not None and len(resp) == 4:
+            return resp
+        else:
+            return None
+
+    def update_hit_count_by_id(self, primary_id: int):
+        table_name = "cache_codegpt_answer"
+        update_sql = "UPDATE {} SET hit_count = hit_count+1 WHERE id={}".format(table_name, primary_id)
+        conn = self.pool.connection()
+
+        # 使用连接执行更新数据操作
+        try:
+            with conn.cursor() as cursor:
+                # 执行更新数据操作
+                cursor.execute(update_sql)
+                conn.commit()
+        finally:
+            # 关闭连接，将连接返回给连接池
+            conn.close()
+
+    def get_ids(self, deleted=True):
+        pass
+
+    def mark_deleted(self, keys):
+        table_name = "cache_codegpt_answer"
+        delete_sql = "Delete from {} WHERE id in ({})".format(table_name, ",".join([str(i) for i in keys]))
+
+        # 从连接池中获取连接
+        conn = self.pool.connection()
+        try:
+            with conn.cursor() as cursor:
+                # 执行删除数据操作
+                cursor.execute(delete_sql)
+                delete_count = cursor.rowcount
+                conn.commit()
+        finally:
+            # 关闭连接，将连接返回给连接池
+            conn.close()
+        return delete_count
+
+    def model_deleted(self, model_name):
+        table_name = "cache_codegpt_answer"
+        delete_sql = "Delete from {} WHERE model='{}'".format(table_name, model_name)
+        conn = self.pool.connection()
+        # 使用连接执行删除数据操作
+        try:
+            with conn.cursor() as cursor:
+                # 执行删除数据操作
+                resp = cursor.execute(delete_sql)
+                conn.commit()
+        finally:
+            # 关闭连接，将连接返回给连接池
+            conn.close()
+        return resp
+
+    def clear_deleted_data(self):
+        pass
+
+    def count(self, state: int = 0, is_all: bool = False):
+        pass
+
+    def close(self):
+        pass
+
+    def count_answers(self):
+        pass