Merge pull request #87 from Yuval-Roth/main

feat: big performance improvements, new features and refactorings

Author: peng3307165

.gitignore

@@ -93,7 +93,7 @@ celerybeat.pid
 
 # Environments
 .env
-.venv
+.venv*
 env/
 venv/
 ENV/
@@ -142,7 +142,7 @@ dmypy.json
 **/multicache_serving.py
 **/modelcache_serving.py
 
-**/model/
+**/model/text2vec-base-chinese
 
 /data/milvus/db
 /data/mysql/db

data/mysql/init/init.sql

@@ -3,7 +3,7 @@ CREATE DATABASE IF NOT EXISTS `modelcache`;
 USE `modelcache`;
 
 CREATE TABLE IF NOT EXISTS `modelcache_llm_answer` (
-  `id` bigint(20) unsigned NOT NULL AUTO_INCREMENT comment '主键',
+  `id` CHAR(36) 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',

docker-compose.yaml

@@ -1,4 +1,4 @@
-version: 'Beta'
+name: "modelcache"
 services:
   mysql:
     image: mysql:8.0.23
@@ -14,12 +14,12 @@ services:
       - ./data/mysql/db:/var/lib/mysql
       - ./data/mysql/my.cnf:/etc/mysql/conf.d/my.cnf
       - ./data/mysql/init:/docker-entrypoint-initdb.d
-    restart: on-failure
+#    restart: on-failure
     networks:
       - modelcache
 
   milvus:
-    image: milvusdb/milvus:v2.5.0-beta
+    image: milvusdb/milvus:v2.5.10
     container_name: milvus
     security_opt:
       - seccomp:unconfined
@@ -36,35 +36,35 @@ services:
       - 19530:19530
       - 9091:9091
       - 2379:2379
-    healthcheck:
-      test: ["CMD", "curl", "-f", "http://localhost:9091/healthz"]
-      interval: 30s
-      start_period: 90s
-      timeout: 20s
-      retries: 3
+#    healthcheck:
+#      test: ["CMD", "curl", "-f", "http://localhost:9091/healthz"]
+#      interval: 30s
+#      start_period: 90s
+#      timeout: 20s
+#      retries: 3
     networks:
       - modelcache
-    restart: on-failure  
+#    restart: on-failure
     command: milvus run standalone
 
-  modelcache:
-    build: 
-      context: .
-      dockerfile: Dockerfile
-    container_name: modelcache
-    image: modelcache:0.1.0
-    ports:
-      - 5000:5000
-    volumes:
-      - ./model:/home/user/model
-      - ./modelcache:/home/user/modelcache
-      - ./modelcache_mm:/home/user/modelcache_mm
-      - ./fastapi4modelcache.py:/home/user/fastapi4modelcache.py
-    networks:
-      - modelcache
-    restart: on-failure
-    command: sh -c "uvicorn fastapi4modelcache:app --reload --reload-dir /home/user --port=5000 --host=0.0.0.0"
+#  modelcache:
+#    build:
+#      context: .
+#      dockerfile: Dockerfile
+#    container_name: modelcache
+#    image: modelcache:0.1.0
+#    ports:
+#      - 5000:5000
+#    volumes:
+#      - ./model:/home/user/model
+#      - ./modelcache:/home/user/modelcache
+#      - ./modelcache_mm:/home/user/modelcache_mm
+#      - ./fastapi4modelcache.py:/home/user/fastapi4modelcache.py
+#    networks:
+#      - modelcache
+#    restart: on-failure
+#    command: sh -c "uvicorn fastapi4modelcache:app --reload --reload-dir /home/user --port=5000 --host=0.0.0.0"
 
 networks:
   modelcache: 
-    external: true
+    driver: bridge

fastapi4modelcache.py

@@ -1,193 +1,46 @@
 # -*- coding: utf-8 -*-
-import time
-import uvicorn
 import asyncio
-import logging
-import configparser
+from contextlib import asynccontextmanager
+import uvicorn
 import json
-from fastapi import FastAPI, Request, HTTPException
-from pydantic import BaseModel
-from concurrent.futures import ThreadPoolExecutor
-from starlette.responses import PlainTextResponse  
-import functools
-
-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 modelcache.utils.model_filter import model_blacklist_filter
-from modelcache.embedding import Data2VecAudio
-
-#创建一个FastAPI实例
-app = FastAPI()
-
-class RequestData(BaseModel):
-    type: str
-    scope: dict = None
-    query: str = None
-    chat_info: dict = None
-    remove_type: str = None
-    id_list: list = []
-
-data2vec = Data2VecAudio()
-mysql_config = configparser.ConfigParser()
-mysql_config.read('modelcache/config/mysql_config.ini')
-
-milvus_config = configparser.ConfigParser()
-milvus_config.read('modelcache/config/milvus_config.ini')
-
-# redis_config = configparser.ConfigParser()
-# redis_config.read('modelcache/config/redis_config.ini')
-
-# 初始化datamanager
-data_manager = get_data_manager(
-    CacheBase("mysql", config=mysql_config),
-    VectorBase("milvus", dimension=data2vec.dimension, milvus_config=milvus_config)
-)
-
-# # 使用redis初始化datamanager
-# data_manager = get_data_manager(
-#     CacheBase("mysql", config=mysql_config),
-#     VectorBase("redis", dimension=data2vec.dimension, redis_config=redis_config)
-# )
-
-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,
-)
-
-executor = ThreadPoolExecutor(max_workers=6)
-
-# 异步保存查询信息
-async def save_query_info(result, model, query, delta_time_log):
-    loop = asyncio.get_running_loop()
-    func = functools.partial(cache.data_manager.save_query_resp, result, model=model, query=json.dumps(query, ensure_ascii=False), delta_time=delta_time_log)
-    await loop.run_in_executor(None, func)
-
-
-
-@app.get("/welcome", response_class=PlainTextResponse)
+from fastapi.responses import JSONResponse
+from fastapi import FastAPI, Request
+from modelcache.cache import Cache
+from modelcache.embedding import EmbeddingModel
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global cache
+    cache, _ = await Cache.init(
+        sql_storage="mysql",
+        vector_storage="milvus",
+        embedding_model=EmbeddingModel.HUGGINGFACE_ALL_MPNET_BASE_V2,
+        embedding_workers_num=2
+    )
+    yield
+
+app = FastAPI(lifespan=lifespan)
+cache: Cache = None
+
+@app.get("/welcome")
 async def first_fastapi():
     return "hello, modelcache!"
 
 @app.post("/modelcache")
 async def user_backend(request: Request):
-    try:
-        raw_body = await request.body()
-        # 解析字符串为JSON对象
-        if isinstance(raw_body, bytes):
-            raw_body = raw_body.decode("utf-8")
-        if isinstance(raw_body, str):
-            try:
-                # 尝试将字符串解析为JSON对象
-                request_data = json.loads(raw_body)
-            except json.JSONDecodeError as e:
-                # 如果无法解析，返回格式错误
-                result = {"errorCode": 101, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
-                  "answer": ''}
-                asyncio.create_task(save_query_info(result, model='', query='', delta_time_log=0))
-                raise HTTPException(status_code=101, detail="Invalid JSON format")
-        else:
-            request_data = raw_body
-
-        # 确保request_data是字典对象
-        if isinstance(request_data, str):
-            try:
-                request_data = json.loads(request_data)
-            except json.JSONDecodeError:
-                raise HTTPException(status_code=101, detail="Invalid JSON format")
-
-        request_type = request_data.get('type')
-        model = None
-        if 'scope' in request_data:
-            model = request_data['scope'].get('model', '').replace('-', '_').replace('.', '_')
-        query = request_data.get('query')
-        chat_info = request_data.get('chat_info')
 
-        if not request_type or request_type not in ['query', 'insert', 'remove', 'register']:
-            result = {"errorCode": 102,
-                      "errorDesc": "type exception, should one of ['query', 'insert', 'remove', 'register']",
-                      "cacheHit": False, "delta_time": 0, "hit_query": '', "answer": ''}
-            asyncio.create_task(save_query_info(result, model=model, query='', delta_time_log=0))
-            raise HTTPException(status_code=102, detail="Type exception, should be one of ['query', 'insert', 'remove', 'register']")
+    try:
+        request_data = await request.json()
+    except Exception:
+        result = {"errorCode": 400, "errorDesc": "bad request", "cacheHit": False, "delta_time": 0, "hit_query": '', "answer": ''}
+        return JSONResponse(status_code=400, content=result)
 
+    try:
+        return await cache.handle_request(request_data)
     except Exception as e:
-        request_data = raw_body if 'raw_body' in locals() else None
-        result = {
-            "errorCode": 103,
-            "errorDesc": str(e),
-            "cacheHit": False,
-            "delta_time": 0,
-            "hit_query": '',
-            "answer": '',
-            "para_dict": request_data
-        }
-        return result
-
-
-    # model filter
-    filter_resp = model_blacklist_filter(model, request_type)
-    if isinstance(filter_resp, dict):
-        return filter_resp
-
-    if request_type == 'query':
-        try:
-            start_time = time.time()
-            response = adapter.ChatCompletion.create_query(scope={"model": model}, query=query)
-            delta_time = f"{round(time.time() - start_time, 2)}s"
-
-            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['data']
-                hit_query = response['hitQuery']
-                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)
-            asyncio.create_task(save_query_info(result, model, query, delta_time_log))
-            return result
-        except Exception as e:
-            result = {"errorCode": 202, "errorDesc": str(e), "cacheHit": False, "delta_time": 0,
-                      "hit_query": '', "answer": ''}
-            logging.info(f'result: {str(result)}')
-            return result
-
-    if request_type == 'insert':
-        try:
-            response = adapter.ChatCompletion.create_insert(model=model, chat_info=chat_info)
-            if response == 'success':
-                return {"errorCode": 0, "errorDesc": "", "writeStatus": "success"}
-            else:
-                return {"errorCode": 301, "errorDesc": response, "writeStatus": "exception"}
-        except Exception as e:
-            return {"errorCode": 303, "errorDesc": str(e), "writeStatus": "exception"}
-
-    if request_type == 'remove':
-        response = adapter.ChatCompletion.create_remove(model=model, remove_type=request_data.get("remove_type"), id_list=request_data.get("id_list"))
-        if not isinstance(response, dict):
-            return {"errorCode": 401, "errorDesc": "", "response": response, "removeStatus": "exception"}
-
-        state = response.get('status')
-        if state == 'success':
-            return {"errorCode": 0, "errorDesc": "", "response": response, "writeStatus": "success"}
-        else:
-            return {"errorCode": 402, "errorDesc": "", "response": response, "writeStatus": "exception"}
-
-    if request_type == 'register':
-        response = adapter.ChatCompletion.create_register(model=model)
-        if response in ['create_success', 'already_exists']:
-            return {"errorCode": 0, "errorDesc": "", "response": response, "writeStatus": "success"}
-        else:
-            return {"errorCode": 502, "errorDesc": "", "response": response, "writeStatus": "exception"}
+        result = {"errorCode": 500, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '', "answer": ''}
+        cache.save_query_resp(result, model='', query='', delta_time=0)
+        return JSONResponse(status_code=500, content=result)
 
-# TODO: 可以修改为在命令行中使用`uvicorn your_module_name:app --host 0.0.0.0 --port 5000 --reload`的命令启动
 if __name__ == '__main__':
-    uvicorn.run(app, host='0.0.0.0', port=5000)
+    uvicorn.run(app, host='0.0.0.0', port=5000, loop="asyncio", http="httptools")