+
+ | Module |
+ Function |
+
+
+
+ | ModelCache |
+ GPTCache |
+
+
+ | Basic Interface |
+ Data query interface |
+ ☑ |
+ ☑ |
+
+
+ | Data writing interface |
+ ☑ |
+ ☑ |
+
+
+ | Embedding |
+ Embedding model configuration |
+ ☑ |
+ ☑ |
+
+
+ | Large model embedding layer |
+ ☑ |
+ |
+
+
+ | BERT model long text processing |
+ ☑ |
+ |
+
+
+ | Large model invocation |
+ Decoupling from large models |
+ ☑ |
+ |
+
+
+ | Local loading of embedding model |
+ ☑ |
+ |
+
+
+ | Data isolation |
+ Model data isolation |
+ ☑ |
+ ☑ |
+
+
+ | Hyperparameter isolation |
+ |
+ |
+
+
+ | Databases |
+ MySQL |
+ ☑ |
+ ☑ |
+
+
+ | Milvus |
+ ☑ |
+ ☑ |
+
+
+ | OceanBase |
+ ☑ |
+ |
+
+
+ | Session management |
+ Single-turn dialogue |
+ ☑ |
+ ☑ |
+
+
+ | System commands |
+ ☑ |
+ |
+
+
+ | Multi-turn dialogue |
+ ☑ |
+ |
+
+
+ | Data management |
+ Data persistence |
+ ☑ |
+ ☑ |
+
+
+ | One-click cache clearance |
+ ☑ |
+ |
+
+
+ | Tenant management |
+ Support for multi-tenancy |
+ ☑ |
+ |
+
+
+ | Milvus multi-collection capability |
+ ☑ |
+ |
+
+
+ | Other |
+ Long-short dialogue distinction |
+ ☑ |
+ |
+
+
diff --git a/docs/2.model-cache-features.md b/docs/2.model-cache-features.md
new file mode 100644
index 0000000..2294d03
--- /dev/null
+++ b/docs/2.model-cache-features.md
@@ -0,0 +1,29 @@
+# ModelCache features
+
+This topic describes ModelCache features. In ModelCache, we incorporated the core principles of GPTCache. ModelCache has four modules: adapter, embedding, similarity, and data_manager.
+
+- The adapter module orchestrates the business logic for various tasks, integrate the embedding, similarity, and data_manager modules.
+- The embedding module converts text into semantic vector representations, and transforms user queries into vectors.
+- The rank module ranks and evaluate the similarity of recalled vectors.
+- The data_manager module manages the databases.
+
+To make ModelCache more suitable for industrial use, we made several improvements to its architecture and functionality:
+
+- [x] Architectural adjustment (lightweight integration):
+ - Embedded into LLM products using a Redis-like caching mode.
+ - Provided semantic caching without interfering with LLM calls, security audits, and other functions.
+ - Compatible with all LLM services.
+- [x] Multiple model loading:
+ - Supported local embedding model loading, and resolved Hugging Face network connectivity issues.
+ - Supported loading embedding layers from various pre-trained models.
+- [x] Data isolation
+ - Environment isolation: Read different database configurations based on the environment. Isolate development, staging, and production environments.
+ - Multi-tenant data isolation: Dynamically create collections based on models for data isolation, addressing data separation issues in multi-model/service scenarios within large language model products.
+- [x] Supported system instruction: Adopted a concatenation approach to resolve issues with system instructions in the prompt paradigm.
+- [x] Long and short text differentiation: Long texts bring more challenges for similarity assessment. Added differentiation between long and short texts, allowing for separate threshold configurations.
+- [x] Milvus performance optimization: Adjusted Milvus consistency level to "Session" level for better performance.
+- [x] Data management:
+ - One-click cache clearing to enable easy data management after model upgrades.
+ - Recall of hit queries for subsequent data analysis and model iteration reference.
+ - Asynchronous log write-back for data analysis and statistics.
+ - Added model field and data statistics field to enhance features.
diff --git a/docs/3.model-cache-quick-start.md b/docs/3.model-cache-quick-start.md
new file mode 100644
index 0000000..7cd3903
--- /dev/null
+++ b/docs/3.model-cache-quick-start.md
@@ -0,0 +1,97 @@
+# Quick start
+
+This topic describes how to set up and use ModelCache.
+
+You can find the start script in `flask4modelcache.py` and `flask4modelcache_demo.py`.
+
+- `flask4modelcache_demo.py`: A quick test service that embeds SQLite and FAISS. No database configuration required.
+- `flask4modelcache.py`: The standard service that requires MySQL and Milvus configuration.
+
+## Dependencies
+
+- Python: V3.8 or above
+- Package installation
+
+ ```shell
+ pip install -r requirements.txt
+ ```
+
+## Start service
+
+### Start demo
+
+1. Download the embedding model bin file from [Hugging Face](https://huggingface.co/shibing624/text2vec-base-chinese/tree/main). Place it in the `model/text2vec-base-chinese` folder.
+2. Start the backend service:
+
+ ```shell
+ cd CodeFuse-ModelCache
+ ```
+
+ ```shell
+ python flask4modelcache_demo.py
+ ```
+
+### Start standard service
+
+Before you start standard service, do these steps:
+
+1. Install MySQL and import the SQL file from `reference_doc/create_table.sql`.
+2. Install vector database Milvus.
+3. Configure database access in:
+ - `modelcache/config/milvus_config.ini`
+ - `modelcache/config/mysql_config.ini`
+4. Download the embedding model bin file from [Hugging Face](https://huggingface.co/shibing624/text2vec-base-chinese/tree/main). Put it in `model/text2vec-base-chinese`.
+5. Start the backend service:
+
+ ```bash
+ python flask4modelcache.py
+ ```
+
+## Visit the service
+
+The service provides three core RESTful API functionalities: Cache-Writing, Cache-Querying, and Cache-Clearing.
+
+### Write cache
+
+```python
+import json
+import requests
+url = 'http://127.0.0.1:5000/modelcache'
+type = 'insert'
+scope = {"model": "CODEGPT-1008"}
+chat_info = [{"query": [{"role": "system", "content": "You are an AI code assistant and you must provide neutral and harmless answers to help users solve code-related problems."}, {"role": "user", "content": "你是谁?"}],
+ "answer": "Hello, I am an intelligent assistant. How can I assist you?"}]
+data = {'type': type, 'scope': scope, 'chat_info': chat_info}
+headers = {"Content-Type": "application/json"}
+res = requests.post(url, headers=headers, json=json.dumps(data))
+```
+
+### Query cache
+
+```python
+import json
+import requests
+url = 'http://127.0.0.1:5000/modelcache'
+type = 'query'
+scope = {"model": "CODEGPT-1008"}
+query = [{"role": "system", "content": "You are an AI code assistant and you must provide neutral and harmless answers to help users solve code-related problems."}, {"role": "user", "content": "Who are you?"}]
+data = {'type': type, 'scope': scope, 'query': query}
+
+headers = {"Content-Type": "application/json"}
+res = requests.post(url, headers=headers, json=json.dumps(data))
+```
+
+### Clear cache
+
+```python
+import json
+import requests
+url = 'http://127.0.0.1:5000/modelcache'
+type = 'remove'
+scope = {"model": "CODEGPT-1008"}
+remove_type = 'truncate_by_model'
+data = {'type': type, 'scope': scope, 'remove_type': remove_type}
+
+headers = {"Content-Type": "application/json"}
+res = requests.post(url, headers=headers, json=json.dumps(data))
+```
diff --git a/docs/4.create-cache.md b/docs/4.create-cache.md
new file mode 100644
index 0000000..e0506ba
--- /dev/null
+++ b/docs/4.create-cache.md
@@ -0,0 +1,42 @@
+# Create cache
+
+This topic describes how to create cache.
+
+## Default cache interface
+
+```py
+class Cache:
+ # ModelCache calls it whe you start the cache system
+ def __init__(self):
+ self.has_init = False
+ self.cache_enable_func = None
+ self.embedding_func = None
+ self.post_process_messages_func = None
+ self.config = Config()
+```
+
+This function embeds text into dense vectors for context similarity search. ModelCache supports these embedding context methods: Huggingface, ONNX, and SentenceTransformers. The default model is text2vec Hugging Face because it performs better for Chinese. Simply initialize your embedding function as `text2vec.to_embeddings`.
+
+```py
+data_manager = get_data_manager(CacheBase("mysql", config=mysql_config),
+ VectorBase("milvus", dimension=data2vec.dimension, milvus_config=milvus_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,
+)
+```
+
+data_manager CacheVase stores all scalar data, such as original questions, prompts, answers, and access times. ModelCache supports multiple cache storages like SQLite, MySQL, and OceanBase. NoSQL databases will be supported in the future.
+
+data_manager VectorBase stores and searches all embedding vectors to find semantically similar results. ModelCache supports using vector search libraries like FAISS or vector databases like Milvus. More vector database and cloud service will be supported in the future.
+
+## Examples
+
+```py
+data_manager = get_data_manager(CacheBase("sqlite"), VectorBase("faiss", dimension=data2vec.dimension))
+data_manager = get_data_manager(CacheBase("oceanbase"), VectorBase("milvus", dimension=data2vec.dimension))
+```
diff --git a/docs/cache-service-cost-time-distribution.webp b/docs/cache-service-cost-time-distribution.webp
new file mode 100644
index 0000000..d22a1f4
Binary files /dev/null and b/docs/cache-service-cost-time-distribution.webp differ
diff --git a/docs/time-cost-comparison.webp b/docs/time-cost-comparison.webp
new file mode 100644
index 0000000..e00c014
Binary files /dev/null and b/docs/time-cost-comparison.webp differ
diff --git a/examples/flask/llms_cache/data_insert.py b/examples/flask/llms_cache/data_insert.py
index 52d1405..e26bec7 100644
--- a/examples/flask/llms_cache/data_insert.py
+++ b/examples/flask/llms_cache/data_insert.py
@@ -14,6 +14,7 @@ def run():
res = requests.post(url, headers=headers, json=json.dumps(data))
res_text = res.text
+ print("data_insert:", res.status_code, res_text)
if __name__ == '__main__':
run()
diff --git a/examples/flask/llms_cache/data_query.py b/examples/flask/llms_cache/data_query.py
index aa59ae3..6fe8313 100644
--- a/examples/flask/llms_cache/data_query.py
+++ b/examples/flask/llms_cache/data_query.py
@@ -14,6 +14,7 @@ def run():
res = requests.post(url, headers=headers, json=json.dumps(data))
res_text = res.text
+ print("data_query:", res.status_code, res_text)
if __name__ == '__main__':
run()
diff --git a/examples/flask/llms_cache/data_query_long.py b/examples/flask/llms_cache/data_query_long.py
index e8c29f8..7aa4f14 100644
--- a/examples/flask/llms_cache/data_query_long.py
+++ b/examples/flask/llms_cache/data_query_long.py
@@ -19,6 +19,7 @@ def run():
res = requests.post(url, headers=headers, json=json.dumps(data))
res_text = res.text
+ print("data_query_long:", res.status_code, res_text)
if __name__ == '__main__':
run()
diff --git a/fastapi4modelcache.py b/fastapi4modelcache.py
new file mode 100644
index 0000000..b260eae
--- /dev/null
+++ b/fastapi4modelcache.py
@@ -0,0 +1,46 @@
+# -*- coding: utf-8 -*-
+import asyncio
+from contextlib import asynccontextmanager
+import uvicorn
+import json
+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:
+ 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:
+ 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)
+
+if __name__ == '__main__':
+ uvicorn.run(app, host='0.0.0.0', port=5000, loop="asyncio", http="httptools")
diff --git a/fastapi4modelcache_demo.py b/fastapi4modelcache_demo.py
new file mode 100644
index 0000000..c1f6a3f
--- /dev/null
+++ b/fastapi4modelcache_demo.py
@@ -0,0 +1,46 @@
+# -*- coding: utf-8 -*-
+import asyncio
+from contextlib import asynccontextmanager
+import uvicorn
+import json
+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="sqlite",
+ vector_storage="faiss",
+ 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:
+ 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:
+ 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)
+
+if __name__ == '__main__':
+ uvicorn.run(app, host='0.0.0.0', port=5000, loop="asyncio", http="httptools")
diff --git a/flask4modelcache.py b/flask4modelcache.py
index ee5955b..3d18518 100644
--- a/flask4modelcache.py
+++ b/flask4modelcache.py
@@ -1,191 +1,48 @@
# -*- coding: utf-8 -*-
-import time
-from flask import Flask, request
-import logging
-import configparser
-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
+import asyncio
-# 创建一个Flask实例
-app = Flask(__name__)
+from flask import Flask, request, jsonify
+from modelcache.cache import Cache
+from modelcache.embedding import EmbeddingModel
-def response_text(cache_resp):
- return cache_resp['data']
+async def main():
+ # 创建一个Flask实例
+ app = Flask(__name__)
-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)
+ cache,loop = await Cache.init(
+ sql_storage="mysql",
+ vector_storage="milvus",
+ embedding_model=EmbeddingModel.HUGGINGFACE_ALL_MPNET_BASE_V2,
+ embedding_workers_num=2
+ )
+ @app.route('/welcome')
+ def first_flask(): # 视图函数
+ return 'hello, modelcache!'
-def response_hitquery(cache_resp):
- return cache_resp['hitQuery']
-
-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')
-
-
-data_manager = get_data_manager(CacheBase("mysql", config=mysql_config),
- VectorBase("milvus", dimension=data2vec.dimension, milvus_config=milvus_config))
-
-# 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,
-)
-
-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', 'remove', 'register']:
- result = {"errorCode": 102,
- "errorDesc": "type exception, should one of ['query', 'insert', 'remove', 'register']",
- "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':
+ @app.post('/modelcache')
+ def user_backend():
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']:
- elif isinstance(response, str):
- 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": str(e), "cacheHit": False, "delta_time": 0,
- "hit_query": '', "answer": ''}
- logging.info('result: {}'.format(result))
+ param_dict = request.json
+ except Exception:
+ result = {"errorCode": 400, "errorDesc": "bad request", "cacheHit": False, "delta_time": 0, "hit_query": '',"answer": ''}
+ return jsonify(result), 400
- 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": 302, "errorDesc": str(e), "writeStatus": "exception"}
- return json.dumps(result, ensure_ascii=False)
-
- if response == 'success':
- result = {"errorCode": 0, "errorDesc": "", "writeStatus": "success"}
- else:
- result = {"errorCode": 301, "errorDesc": response, "writeStatus": "exception"}
- return json.dumps(result, ensure_ascii=False)
+ result = asyncio.run_coroutine_threadsafe(
+ cache.handle_request(param_dict), loop
+ ).result()
+ return jsonify(result), 200
except Exception as e:
- result = {"errorCode": 303, "errorDesc": str(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)
+ 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 jsonify(result), 500
- if request_type == 'register':
- response = adapter.ChatCompletion.create_register(
- model=model
- )
- if response in ['create_success', 'already_exists']:
- result = {"errorCode": 0, "errorDesc": "", "response": response, "writeStatus": "success"}
- else:
- result = {"errorCode": 502, "errorDesc": "", "response": response, "writeStatus": "exception"}
- return json.dumps(result)
+ await asyncio.to_thread(app.run, host='0.0.0.0', port=5000)
if __name__ == '__main__':
- app.run(host='0.0.0.0', port=5000, debug=True)
+ asyncio.run(main())
diff --git a/flask4modelcache_demo.py b/flask4modelcache_demo.py
index 54b9e8e..4251f94 100644
--- a/flask4modelcache_demo.py
+++ b/flask4modelcache_demo.py
@@ -1,169 +1,48 @@
# -*- 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
+import asyncio
-# 创建一个Flask实例
-app = Flask(__name__)
+from flask import Flask, request, jsonify
+from modelcache.cache import Cache
+from modelcache.embedding import EmbeddingModel
-def response_text(cache_resp):
- return cache_resp['data']
+async def main():
+ # 创建一个Flask实例
+ app = Flask(__name__)
-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)
+ cache,loop = await Cache.init(
+ sql_storage="sqlite",
+ vector_storage="faiss",
+ embedding_model=EmbeddingModel.HUGGINGFACE_ALL_MPNET_BASE_V2,
+ embedding_workers_num=2
+ )
+ @app.route('/welcome')
+ def first_flask(): # 视图函数
+ return 'hello, modelcache!'
-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':
+ @app.post('/modelcache')
+ def user_backend():
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)
+ param_dict = request.json
+ except Exception:
+ result = {"errorCode": 400, "errorDesc": "bad request", "cacheHit": False, "delta_time": 0, "hit_query": '',"answer": ''}
+ return jsonify(result), 400
- 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)
+ result = asyncio.run_coroutine_threadsafe(
+ cache.handle_request(param_dict), loop
+ ).result()
+ return jsonify(result), 200
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)
+ 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 jsonify(result), 500
- 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)
+ await asyncio.to_thread(app.run, host='0.0.0.0', port=5000)
if __name__ == '__main__':
- # app.run(host='0.0.0.0', port=5000, debug=True)
- app.run(host='0.0.0.0', port=5000)
+ asyncio.run(main())
diff --git a/model/.gitignore b/model/.gitignore
new file mode 100644
index 0000000..08ad955
--- /dev/null
+++ b/model/.gitignore
@@ -0,0 +1,2 @@
+*.tflite
+text2vec-base-chinese/*
\ No newline at end of file
diff --git a/model/text2vec-base-chinese/logs.txt b/model/text2vec-base-chinese/logs.txt
index 8b13789..1f2f0a8 100644
--- a/model/text2vec-base-chinese/logs.txt
+++ b/model/text2vec-base-chinese/logs.txt
@@ -1 +1,19 @@
-
+Epoch:0 Valid| corr: 0.794410
+Epoch:0 Valid| corr: 0.691819
+Epoch:1 Valid| corr: 0.722749
+Epoch:2 Valid| corr: 0.735054
+Epoch:3 Valid| corr: 0.738295
+Epoch:4 Valid| corr: 0.739411
+Test | corr: 0.679971
+Epoch:0 Valid| corr: 0.817416
+Epoch:1 Valid| corr: 0.832376
+Epoch:2 Valid| corr: 0.842308
+Epoch:3 Valid| corr: 0.843520
+Epoch:4 Valid| corr: 0.841837
+Test | corr: 0.793495
+Epoch:0 Valid| corr: 0.814648
+Epoch:1 Valid| corr: 0.831609
+Epoch:2 Valid| corr: 0.841678
+Epoch:3 Valid| corr: 0.842387
+Epoch:4 Valid| corr: 0.841435
+Test | corr: 0.794840
diff --git a/modelcache/__init__.py b/modelcache/__init__.py
index 9a82563..8bcb18a 100644
--- a/modelcache/__init__.py
+++ b/modelcache/__init__.py
@@ -1,5 +1,7 @@
# -*- coding: utf-8 -*-
-from modelcache.core import Cache
-from modelcache.core import cache
-from modelcache.config import Config
-import modelcache.adapter
+
+
+
+
+
+
diff --git a/modelcache/adapter/adapter.py b/modelcache/adapter/adapter.py
index 452be7b..2425539 100644
--- a/modelcache/adapter/adapter.py
+++ b/modelcache/adapter/adapter.py
@@ -10,47 +10,50 @@ class ChatCompletion(object):
"""Openai ChatCompletion Wrapper"""
@classmethod
- def create_query(cls, *args, **kwargs):
+ async def create_query(cls, *args, **kwargs):
def cache_data_convert(cache_data, cache_query):
return construct_resp_from_cache(cache_data, cache_query)
try:
- return adapt_query(
+ return await adapt_query(
cache_data_convert,
*args,
**kwargs
)
except Exception as e:
+ print(e)
return str(e)
@classmethod
- def create_insert(cls, *args, **kwargs):
+ async def create_insert(cls, *args, **kwargs):
try:
- return adapt_insert(
+ return await adapt_insert(
*args,
**kwargs
)
except Exception as e:
+ print(e)
return str(e)
@classmethod
- def create_remove(cls, *args, **kwargs):
+ async def create_remove(cls, *args, **kwargs):
try:
- return adapt_remove(
+ return await adapt_remove(
*args,
**kwargs
)
except Exception as e:
- logging.info('adapt_remove_e: {}'.format(e))
+ print(e)
return str(e)
@classmethod
- def create_register(cls, *args, **kwargs):
+ async def create_register(cls, *args, **kwargs):
try:
- return adapt_register(
+ return await adapt_register(
*args,
**kwargs
)
except Exception as e:
+ print(e)
return str(e)
diff --git a/modelcache/adapter/adapter_insert.py b/modelcache/adapter/adapter_insert.py
index 74aa619..a2f224a 100644
--- a/modelcache/adapter/adapter_insert.py
+++ b/modelcache/adapter/adapter_insert.py
@@ -1,39 +1,56 @@
# -*- coding: utf-8 -*-
-from modelcache import cache
+import asyncio
+
from modelcache.utils.error import NotInitError
from modelcache.utils.time import time_cal
-def adapt_insert(*args, **kwargs):
- chat_cache = kwargs.pop("cache_obj", cache)
+async def adapt_insert(*args, **kwargs):
+ chat_cache = kwargs.pop("cache_obj")
model = kwargs.pop("model", None)
require_object_store = kwargs.pop("require_object_store", False)
+
+ # Validate object store availability if required
if require_object_store:
assert chat_cache.data_manager.o, "Object store is required for adapter."
- if not chat_cache.has_init:
- raise NotInitError()
- cache_enable = chat_cache.cache_enable_func(*args, **kwargs)
+
context = kwargs.pop("cache_context", {})
- embedding_data = None
- pre_embedding_data = chat_cache.insert_pre_embedding_func(
- kwargs,
- extra_param=context.get("pre_embedding_func", None),
- prompts=chat_cache.config.prompts,
- )
chat_info = kwargs.pop("chat_info", [])
- llm_data = chat_info[-1]['answer']
- if cache_enable:
- embedding_data = time_cal(
+ # Initialize collections for parallel processing
+ pre_embedding_data_list = [] # Preprocessed data ready for embedding
+ embedding_futures_list = [] # Async embedding generation tasks
+ llm_data_list = [] # Extracted LLM response data
+
+ # Process each chat entry and prepare for parallel embedding generation
+ for row in chat_info:
+ # Preprocess chat data using configured preprocessing function
+ pre_embedding_data = chat_cache.insert_pre_embedding_func(
+ row,
+ extra_param=context.get("pre_embedding_func", None),
+ prompts=chat_cache.prompts,
+ )
+ pre_embedding_data_list.append(pre_embedding_data)
+ llm_data_list.append(row['answer']) # Extract answer text for storage
+
+ # Create async embedding generation task with performance monitoring
+ embedding_future = time_cal(
chat_cache.embedding_func,
func_name="embedding",
report_func=chat_cache.report.embedding,
+ cache_obj=chat_cache
)(pre_embedding_data)
+ embedding_futures_list.append(embedding_future)
+
+ # Wait for all embedding generation tasks to complete in parallel
+ embedding_data_list = await asyncio.gather(*embedding_futures_list)
- chat_cache.data_manager.save(
- pre_embedding_data,
- llm_data,
- embedding_data,
+ # Save all processed data to the data manager asynchronously
+ await asyncio.to_thread(
+ chat_cache.data_manager.save,
+ pre_embedding_data_list,
+ llm_data_list,
+ embedding_data_list,
model=model,
extra_param=context.get("save_func", None)
)
diff --git a/modelcache/adapter/adapter_query.py b/modelcache/adapter/adapter_query.py
index a8c90f1..d98e503 100644
--- a/modelcache/adapter/adapter_query.py
+++ b/modelcache/adapter/adapter_query.py
@@ -1,53 +1,72 @@
# -*- coding: utf-8 -*-
+import asyncio
import logging
-import time
-from modelcache import cache
-from modelcache.utils.error import NotInitError
+from modelcache.embedding import MetricType
from modelcache.utils.time import time_cal
-from modelcache.processor.pre import multi_analysis
+from FlagEmbedding import FlagReranker
+USE_RERANKER = False # 如果为 True 则启用 reranker，否则使用原有逻辑
-def adapt_query(cache_data_convert, *args, **kwargs):
- chat_cache = kwargs.pop("cache_obj", cache)
- scope = kwargs.pop("scope", None)
+async def adapt_query(cache_data_convert, *args, **kwargs):
+ # Extract query parameters
+ chat_cache = kwargs.pop("cache_obj")
+ scope = kwargs.pop("scope")
model = scope['model']
- if not chat_cache.has_init:
- raise NotInitError()
- cache_enable = chat_cache.cache_enable_func(*args, **kwargs)
context = kwargs.pop("cache_context", {})
- embedding_data = None
cache_factor = kwargs.pop("cache_factor", 1.0)
+
+ # Preprocess query for embedding generation
pre_embedding_data = chat_cache.query_pre_embedding_func(
kwargs,
extra_param=context.get("pre_embedding_func", None),
- prompts=chat_cache.config.prompts,
+ prompts=chat_cache.prompts,
)
- if cache_enable:
- embedding_data = time_cal(
- chat_cache.embedding_func,
- func_name="embedding",
- report_func=chat_cache.report.embedding,
- )(pre_embedding_data)
- if cache_enable:
- cache_data_list = time_cal(
- chat_cache.data_manager.search,
- func_name="vector_search",
- report_func=chat_cache.report.search,
- )(
- embedding_data,
- extra_param=context.get("search_func", None),
- top_k=kwargs.pop("top_k", -1),
- model=model
- )
- cache_answers = []
- cache_questions = []
- cache_ids = []
- similarity_threshold = chat_cache.config.similarity_threshold
- similarity_threshold_long = chat_cache.config.similarity_threshold_long
+
+ # Generate embedding with performance monitoring
+ embedding_data = await time_cal(
+ chat_cache.embedding_func,
+ func_name="embedding",
+ report_func=chat_cache.report.embedding,
+ cache_obj=chat_cache
+ )(pre_embedding_data)
+
+ search_time_cal = time_cal(
+ chat_cache.data_manager.search,
+ func_name="vector_search",
+ report_func=chat_cache.report.search,
+ cache_obj=chat_cache
+ )
+ cache_data_list = await asyncio.to_thread(
+ search_time_cal,
+ embedding_data,
+ extra_param=context.get("search_func", None),
+ top_k=kwargs.pop("top_k", -1),
+ model=model
+ )
+
+ # Initialize result containers
+ cache_answers = []
+ cache_questions = []
+ cache_ids = []
+ cosine_similarity = None
+
+ # Similarity evaluation based on metric type
+ if chat_cache.similarity_metric_type == MetricType.COSINE:
+ cosine_similarity = cache_data_list[0][0]
+ # This code uses the built-in cosine similarity evaluation in milvus
+ if cosine_similarity < chat_cache.similarity_threshold:
+ return None # No suitable match found
+
+ elif chat_cache.similarity_metric_type == MetricType.L2:
+ # this is the code that uses L2 for similarity evaluation
+ similarity_threshold = chat_cache.similarity_threshold
+ similarity_threshold_long = chat_cache.similarity_threshold_long
min_rank, max_rank = chat_cache.similarity_evaluation.range()
rank_threshold = (max_rank - min_rank) * similarity_threshold * cache_factor
rank_threshold_long = (max_rank - min_rank) * similarity_threshold_long * cache_factor
+
+ # Clamp thresholds to valid range
rank_threshold = (
max_rank
if rank_threshold > max_rank
@@ -62,6 +81,8 @@ def adapt_query(cache_data_convert, *args, **kwargs):
if rank_threshold_long < min_rank
else rank_threshold_long
)
+
+ # Evaluate similarity score
if cache_data_list is None or len(cache_data_list) == 0:
rank_pre = -1.0
else:
@@ -72,16 +93,26 @@ def adapt_query(cache_data_convert, *args, **kwargs):
extra_param=context.get("evaluation_func", None),
)
if rank_pre < rank_threshold:
- return
+ return None # Similarity too low
+ else:
+ raise ValueError(
+ f"Unsupported similarity metric type: {chat_cache.similarity_metric_type}"
+ )
+ # Process search results with optional reranking
+ if USE_RERANKER:
+ reranker = FlagReranker('BAAI/bge-reranker-v2-m3', use_fp16=False)
for cache_data in cache_data_list:
primary_id = cache_data[1]
- ret = chat_cache.data_manager.get_scalar_data(
- cache_data, extra_param=context.get("get_scalar_data", None)
+ ret = await asyncio.to_thread(
+ chat_cache.data_manager.get_scalar_data,
+ cache_data, extra_param=context.get("get_scalar_data", None), model=model
)
if ret is None:
continue
+ rank = reranker.compute_score([pre_embedding_data, ret[0]], normalize=True)[0]
+
if "deps" in context and hasattr(ret.question, "deps"):
eval_query_data = {
"question": context["deps"][0]["data"],
@@ -98,47 +129,109 @@ def adapt_query(cache_data_convert, *args, **kwargs):
"question": pre_embedding_data,
"embedding": embedding_data,
}
-
eval_cache_data = {
"question": ret[0],
"answer": ret[1],
"search_result": cache_data,
"embedding": None
}
- rank = chat_cache.similarity_evaluation.evaluation(
- eval_query_data,
- eval_cache_data,
- extra_param=context.get("evaluation_func", None),
- )
if len(pre_embedding_data) <= 256:
if rank_threshold <= rank:
- cache_answers.append((rank, ret[1]))
- cache_questions.append((rank, ret[0]))
+ cache_answers.append((rank, ret[0]))
+ cache_questions.append((rank, ret[1]))
cache_ids.append((rank, primary_id))
else:
if rank_threshold_long <= rank:
- cache_answers.append((rank, ret[1]))
- cache_questions.append((rank, ret[0]))
+ cache_answers.append((rank, ret[0]))
+ cache_questions.append((rank, ret[1]))
cache_ids.append((rank, primary_id))
- cache_answers = sorted(cache_answers, key=lambda x: x[0], reverse=True)
- cache_questions = sorted(cache_questions, key=lambda x: x[0], reverse=True)
- cache_ids = sorted(cache_ids, key=lambda x: x[0], reverse=True)
- if len(cache_answers) != 0:
- return_message = chat_cache.post_process_messages_func(
- [t[1] for t in cache_answers]
- )
- return_query = chat_cache.post_process_messages_func(
- [t[1] for t in cache_questions]
- )
- return_id = chat_cache.post_process_messages_func(
- [t[1] for t in cache_ids]
+ else:
+ # Original logic without reranking
+ for cache_data in cache_data_list:
+ primary_id = cache_data[1]
+ # Retrieve full cache entry data
+ ret = await asyncio.to_thread(
+ chat_cache.data_manager.get_scalar_data,
+ cache_data, extra_param=context.get("get_scalar_data", None), model=model
)
- # 更新命中次数
- try:
- chat_cache.data_manager.update_hit_count(return_id)
- except Exception:
- logging.info('update_hit_count except, please check!')
-
- chat_cache.report.hint_cache()
- return cache_data_convert(return_message, return_query)
+ if ret is None:
+ continue
+
+ if chat_cache.similarity_metric_type == MetricType.COSINE:
+ assert cosine_similarity is not None, "cosine_similarity should not be None"
+ cache_answers.append((cosine_similarity, ret[0]))
+ cache_questions.append((cosine_similarity, ret[1]))
+ cache_ids.append((cosine_similarity, primary_id))
+
+ elif chat_cache.similarity_metric_type == MetricType.L2:
+ if "deps" in context and hasattr(ret.question, "deps"):
+ eval_query_data = {
+ "question": context["deps"][0]["data"],
+ "embedding": None
+ }
+ eval_cache_data = {
+ "question": ret.question.deps[0].data,
+ "answer": ret.answers[0].answer,
+ "search_result": cache_data,
+ "embedding": None,
+ }
+ else:
+ eval_query_data = {
+ "question": pre_embedding_data,
+ "embedding": embedding_data,
+ }
+ eval_cache_data = {
+ "question": ret[0],
+ "answer": ret[1],
+ "search_result": cache_data,
+ "embedding": None
+ }
+
+ # Evaluate similarity for this specific result
+ rank = chat_cache.similarity_evaluation.evaluation(
+ eval_query_data,
+ eval_cache_data,
+ extra_param=context.get("evaluation_func", None),
+ )
+
+ if len(pre_embedding_data) <= 256:
+ if rank_threshold <= rank:
+ cache_answers.append((rank, ret[0]))
+ cache_questions.append((rank, ret[1]))
+ cache_ids.append((rank, primary_id))
+ else:
+ if rank_threshold_long <= rank:
+ cache_answers.append((rank, ret[0]))
+ cache_questions.append((rank, ret[1]))
+ cache_ids.append((rank, primary_id))
+ else:
+ raise ValueError(
+ f"Unsupported similarity metric type: {chat_cache.similarity_metric_type}"
+ )
+
+ # Sort results by similarity score (highest first)
+ cache_answers = sorted(cache_answers, key=lambda x: x[0], reverse=True)
+ cache_questions = sorted(cache_questions, key=lambda x: x[0], reverse=True)
+ cache_ids = sorted(cache_ids, key=lambda x: x[0], reverse=True)
+ if len(cache_answers) != 0:
+ return_message = chat_cache.post_process_messages_func(
+ [t[1] for t in cache_answers]
+ )
+ return_query = chat_cache.post_process_messages_func(
+ [t[1] for t in cache_questions]
+ )
+ return_id = chat_cache.post_process_messages_func(
+ [t[1] for t in cache_ids]
+ )
+
+ # Update hit count for analytics (async to avoid blocking)
+ try:
+ asyncio.create_task(asyncio.to_thread(chat_cache.data_manager.update_hit_count,return_id))
+ except Exception:
+ logging.info('update_hit_count except, please check!')
+
+ # Record cache hit for reporting
+ chat_cache.report.hint_cache()
+ return cache_data_convert(return_message, return_query)
+ return None
\ No newline at end of file
diff --git a/modelcache/adapter/adapter_register.py b/modelcache/adapter/adapter_register.py
index dafe597..c4d0317 100644
--- a/modelcache/adapter/adapter_register.py
+++ b/modelcache/adapter/adapter_register.py
@@ -1,12 +1,16 @@
# -*- coding: utf-8 -*-
-from modelcache import cache
+import asyncio
-def adapt_register(*args, **kwargs):
- chat_cache = kwargs.pop("cache_obj", cache)
+async def adapt_register(*args, **kwargs):
+ chat_cache = kwargs.pop("cache_obj")
model = kwargs.pop("model", None)
if model is None or len(model) == 0:
return ValueError('')
- register_resp = chat_cache.data_manager.create_index(model)
+ register_resp = await asyncio.to_thread(
+ chat_cache.data_manager.create_index,
+ model
+ )
+
return register_resp
diff --git a/modelcache/adapter/adapter_remove.py b/modelcache/adapter/adapter_remove.py
index 25f1ba3..aea0284 100644
--- a/modelcache/adapter/adapter_remove.py
+++ b/modelcache/adapter/adapter_remove.py
@@ -1,24 +1,29 @@
# -*- coding: utf-8 -*-
-from modelcache import cache
-from modelcache.utils.error import NotInitError, RemoveError
+import asyncio
+from modelcache.utils.error import RemoveError
-def adapt_remove(*args, **kwargs):
- chat_cache = kwargs.pop("cache_obj", cache)
+
+async def adapt_remove(*args, **kwargs):
+ chat_cache = kwargs.pop("cache_obj")
model = kwargs.pop("model", None)
remove_type = kwargs.pop("remove_type", None)
require_object_store = kwargs.pop("require_object_store", False)
if require_object_store:
assert chat_cache.data_manager.o, "Object store is required for adapter."
- if not chat_cache.has_init:
- raise NotInitError()
# delete data
if remove_type == 'delete_by_id':
id_list = kwargs.pop("id_list", [])
- resp = chat_cache.data_manager.delete(id_list, model=model)
+ resp = await asyncio.to_thread(
+ chat_cache.data_manager.delete,
+ id_list, model=model
+ )
elif remove_type == 'truncate_by_model':
- resp = chat_cache.data_manager.truncate(model)
+ resp = await asyncio.to_thread(
+ chat_cache.data_manager.truncate,
+ model
+ )
else:
# resp = "remove_type_error"
raise RemoveError()
diff --git a/modelcache/cache.py b/modelcache/cache.py
new file mode 100644
index 0000000..9149a3c
--- /dev/null
+++ b/modelcache/cache.py
@@ -0,0 +1,399 @@
+# -*- coding: utf-8 -*-
+import asyncio
+import atexit
+import json
+import logging
+import time
+from asyncio import AbstractEventLoop
+from typing import Callable, Optional, List, Any, Coroutine
+from modelcache.adapter import adapter
+from modelcache.embedding.embedding_dispatcher import EmbeddingDispatcher
+from modelcache.utils.model_filter import model_blacklist_filter
+from concurrent.futures import ThreadPoolExecutor, Future
+import configparser
+from modelcache.embedding.base import BaseEmbedding, EmbeddingModel, MetricType
+from modelcache.manager.scalar_data.sql_storage import SQLStorage
+from modelcache.manager.vector_data.base import VectorStorage
+from modelcache.processor.post import first
+from modelcache.processor.pre import query_with_role, query_multi_splicing, insert_multi_splicing
+from modelcache.similarity_evaluation.base import SimilarityEvaluation
+from modelcache.report import Report
+from modelcache.similarity_evaluation.distance import SearchDistanceEvaluation
+from modelcache.utils.error import CacheError
+from modelcache.utils.log import modelcache_log
+from modelcache.manager.data_manager import DataManager
+
+ #=====================================================================#
+ #==================== Cache class definition =========================#
+ #=====================================================================#
+
+
+# noinspection PyMethodMayBeStatic
+class Cache:
+ def __init__(
+ self,
+ embedding_model: EmbeddingModel,
+ similarity_metric_type: MetricType,
+ data_manager: DataManager,
+ query_pre_embedding_func: Callable,
+ insert_pre_embedding_func: Callable,
+ embedding_func: Callable[[str], Future],
+ report: Report, # TODO: figure out why this is needed
+ similarity_evaluation: Optional[SimilarityEvaluation],
+ post_process_messages_func: Callable,
+
+ similarity_threshold: float = 0.95,
+ similarity_threshold_long: float = 0.95,
+ prompts: Optional[List[str]] = None,
+ log_time_func: Callable[[str, float], None] = None,
+ ):
+ if similarity_threshold < 0 or similarity_threshold > 1:
+ raise CacheError(
+ "Invalid the similarity threshold param, reasonable range: 0-1"
+ )
+ self.data_manager: DataManager = data_manager
+ self.embedding_model: EmbeddingModel = embedding_model
+ self.similarity_metric_type: MetricType = similarity_metric_type
+ self.report: Report = report
+ self.query_pre_embedding_func: Callable = query_pre_embedding_func
+ self.insert_pre_embedding_func: Callable = insert_pre_embedding_func
+ self.embedding_func: Callable = embedding_func
+ self.similarity_evaluation: Optional[SimilarityEvaluation] = similarity_evaluation
+ self.post_process_messages_func: Callable = post_process_messages_func
+ self.similarity_threshold = similarity_threshold
+ self.similarity_threshold_long = similarity_threshold_long
+ self.prompts = prompts
+ self.log_time_func: Callable[[str, float], None] = log_time_func
+
+ @atexit.register
+ def close():
+ try:
+ self.data_manager.close()
+ except Exception as e:
+ modelcache_log.error(e)
+
+ def save_query_resp(self, query_resp_dict, **kwargs):
+ """
+ Save query response asynchronously to avoid blocking main thread.
+ Used for logging and analytics purposes.
+ """
+ # Execute save operation in a separate thread to maintain async performance
+ asyncio.create_task(asyncio.to_thread(
+ self.data_manager.save_query_resp,
+ query_resp_dict, **kwargs
+ ))
+
+ def save_query_info(self, result, model, query, delta_time_log):
+ """
+ Save query information with execution timing for performance analysis.
+ Serializes query data to JSON for storage.
+ """
+ # Convert query to JSON and save asynchronously
+ asyncio.create_task(asyncio.to_thread(
+ self.data_manager.save_query_resp,
+ result, model=model, query=json.dumps(query, ensure_ascii=False), delta_time=delta_time_log
+ ))
+
+ async def handle_request(self, param_dict: dict):
+ """
+ Main entry point for processing cache requests.
+
+ Routes requests to appropriate handlers based on request type.
+ Supports: query, insert, remove, register operations.
+
+ Args:
+ param_dict: Request parameters containing type, scope, query, etc.
+
+ Returns:
+ dict: Response dictionary with errorCode, result data, and metadata
+ """
+ # Parse and validate request parameters
+ try:
+ request_type = param_dict.get("type")
+
+ # Extract model information from scope
+ scope = param_dict.get("scope")
+ model = None
+ if scope is not None:
+ model = scope.get('model')
+ # Normalize model name for consistent storage (replace special chars)
+ model = model.replace('-', '_')
+ model = model.replace('.', '_')
+ query = param_dict.get("query")
+ chat_info = param_dict.get("chat_info")
+
+ # Validate request type against supported operations
+ if request_type is None 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": ''}
+ self.save_query_resp(result, model=model, query='', delta_time=0)
+ return result
+ except Exception as e:
+ # Return error response for parameter parsing failures
+ return {"errorCode": 103, "errorDesc": str(e), "cacheHit": False, "delta_time": 0, "hit_query": '',
+ "answer": ''}
+
+ # Apply model-based filtering (blacklist check)
+ filter_resp = model_blacklist_filter(model, request_type)
+ if isinstance(filter_resp, dict):
+ return filter_resp
+
+ # Route to appropriate handler based on request type
+ if request_type == 'query':
+ return await self.handle_query(model, query)
+ elif request_type == 'insert':
+ return await self.handle_insert(chat_info, model)
+ elif request_type == 'remove':
+ return await self.handle_remove(model, param_dict)
+ elif request_type == 'register':
+ return await self.handle_register(model)
+ else:
+ return {"errorCode": 400, "errorDesc": "bad request"}
+
+ async def handle_register(self, model):
+ response = await adapter.ChatCompletion.create_register(
+ model=model,
+ cache_obj=self
+ )
+ # Process registration response and return standardized result
+ if response in ['create_success', 'already_exists']:
+ result = {"errorCode": 0, "errorDesc": "", "response": response, "writeStatus": "success"}
+ else:
+ result = {"errorCode": 502, "errorDesc": "", "response": response, "writeStatus": "exception"}
+ return result
+
+ async def handle_remove(self, model, param_dict):
+ remove_type = param_dict.get("remove_type")
+ id_list = param_dict.get("id_list", [])
+
+ # Execute removal operation through adapter
+ response = await adapter.ChatCompletion.create_remove(
+ model=model,
+ remove_type=remove_type,
+ id_list=id_list,
+ cache_obj=self
+ )
+
+ # Process removal response and standardize result format
+ if not isinstance(response, dict):
+ return {"errorCode": 401, "errorDesc": "", "response": response, "removeStatus": "exception"}
+ 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 result
+
+ async def handle_insert(self, chat_info, model):
+ try:
+ try:
+ # Execute insertion through adapter with error handling
+ response = await adapter.ChatCompletion.create_insert(
+ model=model,
+ chat_info=chat_info,
+ cache_obj=self
+ )
+ except Exception as e:
+ return {"errorCode": 302, "errorDesc": str(e), "writeStatus": "exception"}
+
+ # Process insertion response
+ if response == 'success':
+ result = {"errorCode": 0, "errorDesc": "", "writeStatus": "success"}
+ else:
+ result = {"errorCode": 301, "errorDesc": response, "writeStatus": "exception"}
+ return result
+ except Exception as e:
+ return {"errorCode": 303, "errorDesc": str(e), "writeStatus": "exception"}
+
+ async def handle_query(self, model, query):
+ try:
+ start_time = time.time() # Start performance timer
+
+ # Execute query through adapter system
+ response = await adapter.ChatCompletion.create_query(
+ scope={"model": model},
+ query=query,
+ cache_obj=self
+ )
+
+ # Calculate query execution time
+ delta_time = '{}s'.format(round(time.time() - start_time, 2))
+
+ # Process different response types
+ if response is None:
+ # No cache hit found
+ result = {"errorCode": 0, "errorDesc": '', "cacheHit": False, "delta_time": delta_time, "hit_query": '',
+ "answer": ''}
+ elif isinstance(response, str):
+ # Error occurred during query processing
+ result = {"errorCode": 201, "errorDesc": response, "cacheHit": False, "delta_time": delta_time,
+ "hit_query": '', "answer": ''}
+ else:
+ # Cache hit found - extract response data
+ answer = response['data']
+ hit_query = response['hitQuery']
+ result = {"errorCode": 0, "errorDesc": '', "cacheHit": True, "delta_time": delta_time,
+ "hit_query": hit_query, "answer": answer}
+
+ # Log query performance data asynchronously
+ delta_time_log = round(time.time() - start_time, 2)
+ self.save_query_info(result, model, query, delta_time_log)
+ except Exception as e:
+ result = {"errorCode": 202, "errorDesc": str(e), "cacheHit": False, "delta_time": 0,
+ "hit_query": '', "answer": ''}
+ logging.info('result: {}'.format(result))
+ return result
+
+ def flush(self):
+ """Flush all cached data to persistent storage backends."""
+ self.data_manager.flush()
+
+ @staticmethod
+ async def init(
+ sql_storage: str,
+ vector_storage: str,
+ embedding_model: EmbeddingModel,
+ embedding_workers_num: int
+ ) -> tuple['Cache' , AbstractEventLoop]:
+ """
+ Initialize a complete Cache system with all required components.
+
+ Args:
+ sql_storage: SQL backend type ("mysql", "sqlite", "elasticsearch")
+ vector_storage: Vector backend type ("milvus", "faiss", "chromadb", "redis")
+ embedding_model: Embedding model enum value
+ embedding_workers_num: Number of parallel embedding worker processes
+
+ Returns:
+ tuple: (Cache instance, event loop) ready for async operations
+
+ Raises:
+ CacheError: If unsupported storage backends or model configuration issues
+
+ Example:
+ cache, loop = await Cache.init(
+ sql_storage="mysql",
+ vector_storage="milvus",
+ embedding_model=EmbeddingModel.HUGGINGFACE_ALL_MPNET_BASE_V2,
+ embedding_workers_num=4
+ )
+ """
+ #================= Database configuration loading ===================#
+
+ sql_config = configparser.ConfigParser()
+ vector_config = configparser.ConfigParser()
+
+ # Load SQL storage configuration from INI files
+ if sql_storage == "mysql":
+ sql_config.read('modelcache/config/mysql_config.ini')
+ elif sql_storage == "elasticsearch":
+ sql_config.read('modelcache/config/elasticsearch_config.ini')
+ elif sql_storage == "sqlite":
+ sql_config.read('modelcache/config/sqlite_config.ini')
+ else:
+ modelcache_log.error(f"Unsupported cache storage: {sql_storage}.")
+ raise CacheError(f"Unsupported cache storage: {sql_storage}.")
+
+ # Load vector storage configuration from INI files
+ if vector_storage == "milvus" :
+ vector_config.read('modelcache/config/milvus_config.ini')
+ elif vector_storage == "chromadb" :
+ vector_config.read('modelcache/config/chromadb_config.ini')
+ elif vector_storage == "redis" :
+ vector_config.read('modelcache/config/redis_config.ini')
+ elif vector_storage == "faiss" :
+ vector_config = None # FAISS does not require external configuration
+ else:
+ modelcache_log.error(f"Unsupported vector storage: {vector_storage}.")
+ raise CacheError(f"Unsupported vector storage: {vector_storage}.")
+
+
+ #=============== embedding-model-specific configuration =====================#
+
+ event_loop = asyncio.get_running_loop() # Get current async event loop
+ model_path = embedding_model.value['model_path']
+ dimension = embedding_model.value['dimension']
+
+ # Validate that embedding model has required configuration
+ if model_path is None or dimension is None:
+ modelcache_log.error(f"Please set the model_path and dimension for {embedding_model} in modelcache/embedding/base.py.")
+ raise CacheError(f"Please set the model_path and dimension for {embedding_model} in modelcache/embedding/base.py.")
+
+ # Initialize parallel embedding generation system
+ embedding_dispatcher = EmbeddingDispatcher(embedding_model, model_path, event_loop, embedding_workers_num)
+
+ #=== These will be used to initialize the cache ===#
+ query_pre_embedding_func: Callable = None
+ insert_pre_embedding_func: Callable = None
+ post_process_messages_func: Callable = None
+ similarity_evaluation: Optional[SimilarityEvaluation] = None
+ similarity_metric_type: MetricType = None
+ similarity_threshold: float = None
+ similarity_threshold_long: float = None
+ normalize: bool = None
+ #==================================================#
+
+ # Configure cache behavior based on embedding model type
+ if (embedding_model == EmbeddingModel.HUGGINGFACE_ALL_MPNET_BASE_V2
+ or embedding_model == EmbeddingModel.HUGGINGFACE_ALL_MINILM_L6_V2
+ or embedding_model == EmbeddingModel.HUGGINGFACE_ALL_MINILM_L12_V2):
+ query_pre_embedding_func = query_with_role
+ insert_pre_embedding_func = query_with_role
+ post_process_messages_func = first
+ similarity_evaluation = None # Uses the built-in cosine similarity evaluation in milvus
+ similarity_metric_type = MetricType.COSINE
+ similarity_threshold = 0.9
+ similarity_threshold_long = 0.9
+ normalize = False
+
+ elif embedding_model == EmbeddingModel.DATA2VEC_AUDIO:
+ query_pre_embedding_func = query_multi_splicing
+ insert_pre_embedding_func = insert_multi_splicing
+ post_process_messages_func = first
+ similarity_evaluation = SearchDistanceEvaluation()
+ similarity_metric_type = MetricType.L2
+ similarity_threshold = 0.95
+ similarity_threshold_long = 0.95
+ normalize = True
+
+ # Add configurations for additional embedding models as needed
+ else:
+ modelcache_log.error(f"Please add configuration for {embedding_model} in modelcache/cache.py.")
+ raise CacheError(f"Please add configuration for {embedding_model} in modelcache/cache.py.")
+
+ # ====================== Data manager ==============================#
+
+ # Create coordinated data manager with all storage backends
+ data_manager = DataManager.get(
+ SQLStorage.get(sql_storage, config=sql_config),
+ VectorStorage.get(
+ name=vector_storage,
+ dimension=dimension,
+ config=vector_config,
+ metric_type=similarity_metric_type,
+ ),
+ memory_cache_policy='ARC',
+ max_size=10000,
+ normalize=normalize,
+ )
+
+ #================== Cache Initialization ====================#
+
+ # Create fully configured Cache instance
+ cache = Cache(
+ embedding_model = embedding_model,
+ similarity_metric_type = similarity_metric_type,
+ data_manager = data_manager,
+ report = Report(),
+ embedding_func = embedding_dispatcher.embed,
+ query_pre_embedding_func = query_pre_embedding_func,
+ insert_pre_embedding_func = insert_pre_embedding_func,
+ similarity_evaluation = similarity_evaluation,
+ post_process_messages_func = post_process_messages_func,
+ similarity_threshold = similarity_threshold,
+ similarity_threshold_long = similarity_threshold_long,
+ prompts = None,
+ log_time_func = None,
+ )
+ return cache, event_loop
\ No newline at end of file
diff --git a/modelcache/config.py b/modelcache/config.py
deleted file mode 100644
index 69b3246..0000000
--- a/modelcache/config.py
+++ /dev/null
@@ -1,22 +0,0 @@
-# -*- coding: utf-8 -*-
-from typing import Optional, Callable, List
-from modelcache.utils.error import CacheError
-
-
-class Config:
-
- def __init__(
- self,
- log_time_func: Optional[Callable[[str, float], None]] = None,
- similarity_threshold: float = 0.95,
- similarity_threshold_long: float = 0.95,
- prompts: Optional[List[str]] = None
- ):
- if similarity_threshold < 0 or similarity_threshold > 1:
- raise CacheError(
- "Invalid the similarity threshold param, reasonable range: 0-1"
- )
- self.log_time_func = log_time_func
- self.similarity_threshold = similarity_threshold
- self.similarity_threshold_long = similarity_threshold_long
- self.prompts = prompts
diff --git a/modelcache/config/chromadb_config.ini b/modelcache/config/chromadb_config.ini
new file mode 100644
index 0000000..7400f80
--- /dev/null
+++ b/modelcache/config/chromadb_config.ini
@@ -0,0 +1,2 @@
+[chromadb]
+persist_directory=''
diff --git a/modelcache/config/elasticsearch_config.ini b/modelcache/config/elasticsearch_config.ini
new file mode 100644
index 0000000..ba87a5a
--- /dev/null
+++ b/modelcache/config/elasticsearch_config.ini
@@ -0,0 +1,5 @@
+[elasticsearch]
+host = ''
+port = ''
+user = ''
+password = ''
\ No newline at end of file
diff --git a/modelcache/config/milvus_config.ini b/modelcache/config/milvus_config.ini
index f5bd532..1e9aa66 100644
--- a/modelcache/config/milvus_config.ini
+++ b/modelcache/config/milvus_config.ini
@@ -1,5 +1,5 @@
[milvus]
-host = ''
-port = ''
+host = localhost
+port = 19530
user = ''
password = ''
\ No newline at end of file
diff --git a/modelcache/config/mysql_config.ini b/modelcache/config/mysql_config.ini
index 2c63f0e..3d92385 100644
--- a/modelcache/config/mysql_config.ini
+++ b/modelcache/config/mysql_config.ini
@@ -1,6 +1,6 @@
[mysql]
-host = ''
-port = ''
-username = ''
-password = ''
-database = ''
+host = localhost
+port = 3306
+username = modelcache
+password = modelcache
+database = modelcache
diff --git a/modelcache/core.py b/modelcache/core.py
deleted file mode 100644
index bd57029..0000000
--- a/modelcache/core.py
+++ /dev/null
@@ -1,76 +0,0 @@
-# -*- coding: utf-8 -*-
-import atexit
-from typing import Optional, List, Any
-from modelcache.processor.post import first
-from modelcache.similarity_evaluation import ExactMatchEvaluation
-from modelcache.similarity_evaluation import SimilarityEvaluation
-from modelcache.embedding.string_text import to_embeddings as string_embedding
-from modelcache.report import Report
-from modelcache.config import Config
-from modelcache.utils.cache_func import cache_all
-from modelcache.utils.log import modelcache_log
-from modelcache.manager import get_data_manager
-from modelcache.manager.data_manager import DataManager
-
-
-class Cache:
- def __init__(self):
- self.has_init = False
- self.cache_enable_func = None
- self.query_pre_embedding_func = None
- self.insert_pre_embedding_func = None
- self.mm_query_pre_embedding_func = None
- self.mm_insert_pre_embedding_func = None
- self.embedding_func = None
- self.embedding_concurrent_func = None
- self.data_manager: Optional[DataManager] = None
- self.similarity_evaluation: Optional[SimilarityEvaluation] = None
- self.post_process_messages_func = None
- self.config = Config()
- self.report = Report()
- self.next_cache = None
-
- def init(
- self,
- cache_enable_func=cache_all,
- query_pre_embedding_func=None,
- insert_pre_embedding_func=None,
- embedding_func=string_embedding,
- data_manager: DataManager = get_data_manager(),
- similarity_evaluation=ExactMatchEvaluation(),
- post_process_messages_func=first,
- config=Config(),
- next_cache=None,
- ):
- self.has_init = True
- self.cache_enable_func = cache_enable_func
- self.query_pre_embedding_func = query_pre_embedding_func
- self.insert_pre_embedding_func = insert_pre_embedding_func
- self.embedding_func = embedding_func
- self.data_manager: DataManager = data_manager
- self.similarity_evaluation = similarity_evaluation
- self.post_process_messages_func = post_process_messages_func
- self.config = config
- self.next_cache = next_cache
-
- @atexit.register
- def close():
- try:
- self.data_manager.close()
- except Exception as e:
- modelcache_log.error(e)
-
- def import_data(self, questions: List[Any], answers: List[Any]) -> None:
- self.data_manager.import_data(
- questions=questions,
- answers=answers,
- embedding_datas=[self.embedding_func(question) for question in questions],
- )
-
- def flush(self):
- self.data_manager.flush()
- if self.next_cache:
- self.next_cache.data_manager.flush()
-
-
-cache = Cache()
diff --git a/modelcache/embedding/__init__.py b/modelcache/embedding/__init__.py
index 5684a2d..19f9b94 100644
--- a/modelcache/embedding/__init__.py
+++ b/modelcache/embedding/__init__.py
@@ -1,36 +1,2 @@
# -*- coding: utf-8 -*-
-from modelcache.utils.lazy_import import LazyImport
-huggingface = LazyImport("huggingface", globals(), "modelcache.embedding.huggingface")
-data2vec = LazyImport("data2vec", globals(), "modelcache.embedding.data2vec")
-llmEmb = LazyImport("llmEmb", globals(), "modelcache.embedding.llmEmb")
-fasttext = LazyImport("fasttext", globals(), "modelcache.embedding.fasttext")
-paddlenlp = LazyImport("paddlenlp", globals(), "modelcache.embedding.paddlenlp")
-timm = LazyImport("timm", globals(), "modelcache.embedding.timm")
-huggingface_tei = LazyImport("huggingface_tei", globals(), "modelcache.embedding.huggingface_tei")
-
-
-def Huggingface(model="sentence-transformers/all-mpnet-base-v2"):
- return huggingface.Huggingface(model)
-
-
-def Data2VecAudio(model="model/text2vec-base-chinese/"):
- return data2vec.Data2VecAudio(model)
-
-
-def LlmEmb2vecAudio():
- return llmEmb.LlmEmb2Vec()
-
-
-def FastText(model="en", dim=None):
- return fasttext.FastText(model, dim)
-
-
-def PaddleNLP(model="ernie-3.0-medium-zh"):
- return paddlenlp.PaddleNLP(model)
-
-
-def Timm(model="resnet50", device="default"):
- return timm.Timm(model, device)
-
-def HuggingfaceTEI(base_url, model):
- return huggingface_tei.HuggingfaceTEI(base_url, model)
\ No newline at end of file
+from modelcache.embedding.base import EmbeddingModel, MetricType, BaseEmbedding
\ No newline at end of file
diff --git a/modelcache/embedding/base.py b/modelcache/embedding/base.py
index 4b316aa..70ae240 100644
--- a/modelcache/embedding/base.py
+++ b/modelcache/embedding/base.py
@@ -1,5 +1,46 @@
# -*- coding: utf-8 -*-
-from abc import ABCMeta, abstractmethod
+from abc import abstractmethod, ABCMeta
+
+from modelcache.utils.error import CacheError
+from modelcache.utils.lazy_import import LazyImport
+from enum import Enum
+
+from modelcache.utils.log import modelcache_log
+
+huggingface = LazyImport("huggingface", globals(), "modelcache.embedding.huggingface")
+data2vec = LazyImport("data2vec", globals(), "modelcache.embedding.data2vec")
+llmEmb = LazyImport("llmEmb", globals(), "modelcache.embedding.llmEmb")
+fasttext = LazyImport("fasttext", globals(), "modelcache.embedding.fasttext")
+paddlenlp = LazyImport("paddlenlp", globals(), "modelcache.embedding.paddlenlp")
+timm = LazyImport("timm", globals(), "modelcache.embedding.timm")
+huggingface_tei = LazyImport("huggingface_tei", globals(), "modelcache.embedding.huggingface_tei")
+bge_m3 = LazyImport("bge_m3", globals(), "modelcache.embedding.bge_m3")
+
+# define the embedding model enum
+class EmbeddingModel(Enum):
+ """
+ Enum for different embedding models.
+ """
+ # todo: fill in the dimension and model_path for each embedding model as needed
+ HUGGINGFACE_ALL_MPNET_BASE_V2 = {"dimension":768, "model_path":"sentence-transformers/all-mpnet-base-v2"}
+ HUGGINGFACE_ALL_MINILM_L6_V2 = {"dimension":384, "model_path":"sentence-transformers/all-MiniLM-L6-v2"}
+ HUGGINGFACE_ALL_MINILM_L12_V2 = {"dimension":384, "model_path":"sentence-transformers/all-MiniLM-L12-v2"}
+ DATA2VEC_AUDIO = {"dimension":768, "model_path":"model/text2vec-base-chinese/"}
+ LLM_EMB2VEC_AUDIO = {"dimension":None, "model_path":None}
+ FASTTEXT = {"dimension":None, "model_path":None}
+ PADDLE_NLP = {"dimension":None, "model_path":None}
+ TIMM = {"dimension":None, "model_path":None}
+ HUGGINGFACE_TEI = {"dimension":None, "model_path":None}
+ BGE_M3 = {"dimension":None, "model_path":None}
+
+
+class MetricType(Enum):
+ """
+ Enum for different metric types used in similarity evaluation.
+ Different models may require different metrics for optimal performance.
+ """
+ COSINE = "COSINE"
+ L2 = "L2"
class BaseEmbedding(metaclass=ABCMeta):
@@ -15,3 +56,62 @@ def to_embeddings(self, data, **kwargs):
@abstractmethod
def dimension(self) -> int:
return 0
+
+ @staticmethod
+ def get(model:EmbeddingModel, **kwargs):
+ """
+ Get the embedding model instance based on the specified model type.
+ :param model: The embedding model type.
+ :type model: EmbeddingModel
+ :param kwargs: Additional parameters for the model.
+ :return: An instance of the specified embedding model.
+ :rtype: BaseEmbedding
+ :raises ValueError: If the specified model type is not supported.
+ """
+ if model == EmbeddingModel.HUGGINGFACE_ALL_MPNET_BASE_V2:
+ model_path = kwargs.pop("model_path","sentence-transformers/all-mpnet-base-v2")
+ return huggingface.Huggingface(model_path)
+
+ elif model == EmbeddingModel.HUGGINGFACE_ALL_MINILM_L6_V2:
+ model_path = kwargs.pop("model_path","sentence-transformers/all-MiniLM-L6-v2")
+ return huggingface.Huggingface(model_path)
+
+ elif model == EmbeddingModel.HUGGINGFACE_ALL_MINILM_L12_V2:
+ model_path = kwargs.pop("model_path","sentence-transformers/all-MiniLM-L12-v2")
+ return huggingface.Huggingface(model_path)
+
+ elif model == EmbeddingModel.DATA2VEC_AUDIO:
+ model_path = kwargs.pop("model_path","model/text2vec-base-chinese/")
+ return data2vec.Data2VecAudio(model_path)
+
+ elif model == EmbeddingModel.LLM_EMB2VEC_AUDIO:
+ return llmEmb.LlmEmb2Vec()
+
+ elif model == EmbeddingModel.FASTTEXT:
+ model_path = kwargs.pop("model_path","en")
+ dim = kwargs.pop("dim", None)
+ return fasttext.FastText(model_path, dim)
+
+ elif model == EmbeddingModel.PADDLE_NLP:
+ model_path = kwargs.pop("model_path", "ernie-3.0-medium-zh")
+ return paddlenlp.PaddleNLP(model_path)
+
+ elif model == EmbeddingModel.TIMM:
+ model_path = kwargs.pop("model_path", "resnet50")
+ device = kwargs.pop("device", "default")
+ return timm.Timm(model_path, device)
+
+ elif model == EmbeddingModel.HUGGINGFACE_TEI:
+ base_url = kwargs.pop("base_url")
+ model_path = kwargs.pop("model_path")
+ return huggingface_tei.HuggingfaceTEI(base_url, model_path)
+
+ elif model == EmbeddingModel.BGE_M3:
+ model_path = kwargs.pop("model_path","model/bge-m3")
+ return bge_m3.BgeM3Embedding(model_path)
+
+ else:
+ modelcache_log.error(f"Please add configuration for {model} in modelcache/embedding/base.py.")
+ raise CacheError(f"Please add configuration for {model} in modelcache/embedding/base.py.")
+
+
diff --git a/modelcache/embedding/bge_m3.py b/modelcache/embedding/bge_m3.py
new file mode 100644
index 0000000..10cc15c
--- /dev/null
+++ b/modelcache/embedding/bge_m3.py
@@ -0,0 +1,28 @@
+# -*- coding: utf-8 -*-
+import numpy as np
+from modelcache.embedding.base import BaseEmbedding
+from transformers import AutoTokenizer, AutoModel
+from FlagEmbedding import BGEM3FlagModel
+
+class BgeM3Embedding(BaseEmbedding):
+ def __init__(self, model_path: str = "model/bge-m3"):
+ self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+ self.model = AutoModel.from_pretrained(model_path)
+
+ self.bge_model = BGEM3FlagModel(model_name_or_path=model_path,
+ model=self.model,
+ tokenizer=self.tokenizer,
+ use_fp16=False)
+
+ self.__dimension = 768
+
+ def to_embeddings(self, data, **_):
+ if not isinstance(data, list):
+ data = [data]
+
+ embeddings = self.bge_model.encode(data, batch_size=12, max_length=8192)['dense_vecs']
+ return np.array(embeddings).astype("float32")
+
+ @property
+ def dimension(self):
+ return self.__dimension
\ No newline at end of file
diff --git a/modelcache/embedding/embedding_dispatcher.py b/modelcache/embedding/embedding_dispatcher.py
new file mode 100644
index 0000000..16f4beb
--- /dev/null
+++ b/modelcache/embedding/embedding_dispatcher.py
@@ -0,0 +1,84 @@
+import multiprocessing
+import threading
+import uuid
+import asyncio
+import psutil
+from asyncio import Future, AbstractEventLoop
+
+from modelcache.embedding import EmbeddingModel
+from modelcache.embedding.base import BaseEmbedding
+
+
+def worker_func(embedding_model: EmbeddingModel, model_path, task_queue, result_queue, worker_id):
+ """Worker function that runs in separate processes to generate embeddings."""
+ base_embedding = BaseEmbedding.get(embedding_model, model_path=model_path)
+ print(f"Embedding worker {worker_id} started.")
+ try:
+ while True:
+ job_id, data = task_queue.get() # Get task from queue
+ try:
+ result = base_embedding.to_embeddings(data) # Generate embedding
+ except Exception as e:
+ result = e
+ result_queue.put((job_id, result)) # Send result back
+ except KeyboardInterrupt:
+ print(f"Embedding worker {worker_id} stopped.")
+ except Exception as e:
+ print(f"Embedding worker {worker_id} encountered an error: {e}")
+
+
+class EmbeddingDispatcher:
+ """Manages a pool of worker processes for parallel embedding generation."""
+
+ def __init__(
+ self,
+ embedding_model: EmbeddingModel,
+ model_path: str,
+ event_loop: AbstractEventLoop,
+ num_workers: int
+ ):
+ """Initialize the dispatcher with worker processes."""
+ if num_workers <= 0:
+ raise ValueError("Number of workers must be greater than 0.")
+
+ self.task_queue = multiprocessing.Queue() # Tasks to workers
+ self.result_queue = multiprocessing.Queue() # Results from workers
+ self.futures: dict[str, asyncio.Future] = {} # Pending futures
+ self.event_loop = event_loop
+ self._start_result_collector_thread() # Start result collection thread
+
+ # Start worker processes
+ self.workers = []
+ for i in range(num_workers):
+ p = multiprocessing.Process(
+ target=worker_func,
+ args=(embedding_model, model_path, self.task_queue, self.result_queue, i)
+ )
+ p.daemon = True
+ p.start()
+ psutil.Process(p.pid).nice(psutil.HIGH_PRIORITY_CLASS)
+ self.workers.append(p)
+
+ def _start_result_collector_thread(self):
+ """Start a thread to collect results from worker processes."""
+ def collect():
+ while True:
+ job_id, result = self.result_queue.get() # Get result from queue
+ future = self.futures.pop(job_id, None) # Retrieve future
+ if future:
+ self.event_loop.call_soon_threadsafe(
+ future.set_exception if isinstance(result, Exception) else future.set_result,
+ result
+ )
+
+ t = threading.Thread(target=collect, daemon=True)
+ t.start()
+
+ def embed(self, data: str) -> Future:
+ """Submit a task for embedding generation."""
+ job_id = str(uuid.uuid4()) # Generate unique job ID
+ future = asyncio.get_running_loop().create_future() # Create future
+ self.futures[job_id] = future # Store future
+ self.task_queue.put((job_id, data)) # Add task to queue
+ return future
+
diff --git a/modelcache/embedding/huggingface.py b/modelcache/embedding/huggingface.py
index 8c1434d..77ca08f 100644
--- a/modelcache/embedding/huggingface.py
+++ b/modelcache/embedding/huggingface.py
@@ -1,34 +1,21 @@
# -*- coding: utf-8 -*-
-import numpy as np
-
-from modelcache.utils import import_huggingface, import_torch
from modelcache.embedding.base import BaseEmbedding
-
-import_torch()
-import_huggingface()
-
-import torch # pylint: disable=C0413
-from transformers import AutoTokenizer, AutoModel # pylint: disable=C0413
-
+from sentence_transformers import SentenceTransformer
class Huggingface(BaseEmbedding):
- def __init__(self, model: str = "sentence-transformers/all-MiniLM-L6-v2"):
- self.model = AutoModel.from_pretrained(model, local_files_only=True)
- self.model.eval()
-
- # self.tokenizer = AutoTokenizer.from_pretrained(model)
- self.tokenizer = AutoTokenizer.from_pretrained(model, local_files_only=True)
- if not self.tokenizer.pad_token:
- self.tokenizer.pad_token = "[PAD]"
+ def __init__(self, model: str):
+ self.model = SentenceTransformer(model,tokenizer_kwargs={
+ "clean_up_tokenization_spaces":False
+ })
try:
self.__dimension = self.model.config.hidden_size
- except Exception: # pylint: disable=W0703
- from transformers import AutoConfig # pylint: disable=C0415
+ except Exception:
+ from transformers import AutoConfig
config = AutoConfig.from_pretrained(model)
self.__dimension = config.hidden_size
- def to_embeddings(self, data, **_):
+ def to_embeddings(self, data: str, **_):
"""Generate embedding given text input
:param data: text in string.
@@ -36,24 +23,11 @@ def to_embeddings(self, data, **_):
:return: a text embedding in shape of (dim,).
"""
- if not isinstance(data, list):
- data = [data]
- inputs = self.tokenizer(
- data, padding=True, truncation=True, return_tensors="pt"
- )
- outs = self.model(**inputs).last_hidden_state
- emb = self.post_proc(outs, inputs).squeeze(0).detach().numpy()
- return np.array(emb).astype("float32")
- def post_proc(self, token_embeddings, inputs):
- attention_mask = inputs["attention_mask"]
- input_mask_expanded = (
- attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
- )
- sentence_embs = torch.sum(
- token_embeddings * input_mask_expanded, 1
- ) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
- return sentence_embs
+ if not data:
+ raise ValueError("No data provided for embedding.")
+ embeddings = self.model.encode(data)
+ return embeddings[0] if len(data) == 1 else embeddings
@property
def dimension(self):
diff --git a/modelcache/embedding/string_text.py b/modelcache/embedding/string_text.py
deleted file mode 100644
index 4fd08e7..0000000
--- a/modelcache/embedding/string_text.py
+++ /dev/null
@@ -1,5 +0,0 @@
-# -*- coding: utf-8 -*-
-
-
-def to_embeddings(data, **_):
- return data
diff --git a/modelcache/manager/__init__.py b/modelcache/manager/__init__.py
index 59a5ffe..40a96af 100644
--- a/modelcache/manager/__init__.py
+++ b/modelcache/manager/__init__.py
@@ -1,5 +1 @@
# -*- coding: utf-8 -*-
-from modelcache.manager.scalar_data import CacheBase
-from modelcache.manager.vector_data import VectorBase
-from modelcache.manager.object_data import ObjectBase
-from modelcache.manager.factory import get_data_manager
diff --git a/modelcache/manager/data_manager.py b/modelcache/manager/data_manager.py
index 6734ef2..7f019cd 100644
--- a/modelcache/manager/data_manager.py
+++ b/modelcache/manager/data_manager.py
@@ -6,19 +6,13 @@
import numpy as np
import cachetools
from abc import abstractmethod, ABCMeta
-from typing import List, Any, Optional, Union
-from modelcache.manager.scalar_data.base import (
- CacheStorage,
- CacheData,
- DataType,
- Answer,
- Question
-)
+from typing import List, Any, Optional
+from typing import Union, Callable
+from modelcache.manager.scalar_data.base import CacheStorage,CacheData,DataType,Answer,Question
from modelcache.utils.error import CacheError, ParamError
-from modelcache.manager.vector_data.base import VectorBase, VectorData
+from modelcache.manager.vector_data.base import VectorStorage, VectorData
from modelcache.manager.object_data.base import ObjectBase
-from modelcache.manager.eviction import EvictionBase
-from modelcache.manager.eviction_manager import EvictionManager
+from modelcache.manager.eviction.memory_cache import MemoryCacheEviction
from modelcache.utils.log import modelcache_log
@@ -34,9 +28,7 @@ def save_query_resp(self, query_resp_dict, **kwargs):
pass
@abstractmethod
- def import_data(
- self, questions: List[Any], answers: List[Any], embedding_datas: List[Any], model:Any
- ):
+ def import_data(self, questions: List[Any], answers: List[Any], embedding_datas: List[Any], model:Any):
pass
@abstractmethod
@@ -58,9 +50,11 @@ def search(self, embedding_data, **kwargs):
def delete(self, id_list, **kwargs):
pass
+ @abstractmethod
def truncate(self, model_name):
pass
+ @abstractmethod
def flush(self):
pass
@@ -68,6 +62,30 @@ def flush(self):
def close(self):
pass
+ @staticmethod
+ def get(
+ cache_base: Union[CacheStorage, str] = None,
+ vector_base: Union[VectorStorage, str] = None,
+ object_base: Union[ObjectBase, str] = None,
+ max_size: int = 3,
+ clean_size: int = 1,
+ memory_cache_policy: str = "ARC",
+ data_path: str = "data_map.txt",
+ get_data_container: Callable = None,
+ normalize: bool = True
+ ):
+ if not cache_base and not vector_base:
+ return MapDataManager(data_path, max_size, get_data_container)
+
+ if isinstance(cache_base, str):
+ cache_base = CacheStorage.get(name=cache_base)
+ if isinstance(vector_base, str):
+ vector_base = VectorStorage.get(name=vector_base)
+ if isinstance(object_base, str):
+ object_base = ObjectBase.get(name=object_base)
+ assert cache_base and vector_base
+ return SSDataManager(cache_base, vector_base, object_base, max_size, clean_size,normalize, memory_cache_policy)
+
class MapDataManager(DataManager):
def __init__(self, data_path, max_size, get_data_container=None):
@@ -146,23 +164,33 @@ class SSDataManager(DataManager):
def __init__(
self,
s: CacheStorage,
- v: VectorBase,
+ v: VectorStorage,
o: Optional[ObjectBase],
max_size,
clean_size,
+ normalize: bool,
policy="LRU",
):
self.max_size = max_size
self.clean_size = clean_size
- self.s = s
- self.v = v
- self.o = o
-
- def save(self, question, answer, embedding_data, **kwargs):
+ self.s = s # SQL storage
+ self.v = v # Vector storage
+ self.o = o # Object storage (optional)
+ self.normalize = normalize
+
+ # Initialize memory cache with specified eviction policy
+ self.eviction_base = MemoryCacheEviction(
+ policy=policy,
+ maxsize=max_size,
+ clean_size=clean_size)
+
+ def save(self, questions: List[any], answers: List[any], embedding_datas: List[any], **kwargs):
+ """Save multiple questions, answers, and embeddings to storage."""
model = kwargs.pop("model", None)
- self.import_data([question], [answer], [embedding_data], model)
+ self.import_data(questions, answers, embedding_datas, model)
def save_query_resp(self, query_resp_dict, **kwargs):
+ """Save query response log to SQL storage for analytics."""
save_query_start_time = time.time()
self.s.insert_query_resp(query_resp_dict, **kwargs)
save_query_delta_time = '{}s'.format(round(time.time() - save_query_start_time, 2))
@@ -193,60 +221,97 @@ def _process_question_data(self, question: Union[str, Question]):
def import_data(
self, questions: List[Any], answers: List[Answer], embedding_datas: List[Any], model: Any
):
+ """
+ Add multiple cache entries into all storage backends.
+
+ Coordinates data insertion across SQL, vector, and object storage,
+ with memory cache population and optional vector normalization.
+ """
if len(questions) != len(answers) or len(questions) != len(embedding_datas):
raise ParamError("Make sure that all parameters have the same length")
cache_datas = []
- embedding_datas = [
- normalize(embedding_data) for embedding_data in embedding_datas
- ]
+ # Normalize embedding vectors if configured
+ if self.normalize:
+ embedding_datas = [
+ normalize(embedding_data) for embedding_data in embedding_datas
+ ]
- for i, embedding_data in enumerate(embedding_datas):
+ for embedding_data, answer, question in zip(embedding_datas,answers,questions):
if self.o is not None:
- ans = self._process_answer_data(answers[i])
- else:
- ans = answers[i]
+ answer = self._process_answer_data(answer)
- question = questions[i]
embedding_data = embedding_data.astype("float32")
- cache_datas.append([ans, question, embedding_data, model])
+ cache_datas.append([answer, question, embedding_data, model])
+ # Insert into SQL storage and get generated IDs
ids = self.s.batch_insert(cache_datas)
- self.v.mul_add(
- [
- VectorData(id=ids[i], data=embedding_data)
- for i, embedding_data in enumerate(embedding_datas)
- ],
- model
- )
+ # Prepare vector data and populate memory cache
+ datas = []
+ for _id,embedding_data,cache_data in zip(ids,embedding_datas,cache_datas):
+ datas.append(VectorData(id=_id, data=embedding_data.astype("float32")))
+ self.eviction_base.put([(_id, cache_data)],model=model)
+ self.v.mul_add(datas,model)
def get_scalar_data(self, res_data, **kwargs) -> Optional[CacheData]:
- cache_data = self.s.get_data_by_id(res_data[1])
+ """
+ Retrieve scalar data with multi-level caching strategy.
+
+ First checks memory cache, then falls back to SQL storage.
+ """
+ model = kwargs.pop("model")
+ _id = res_data[1]
+
+ # Try to get from memory cache first (fastest)
+ cache_hit = self.eviction_base.get(_id, model=model)
+ if cache_hit is not None:
+ return cache_hit
+ cache_data = self.s.get_data_by_id(_id)
if cache_data is None:
return None
+ self.eviction_base.put([(_id, cache_data)], model=model)
return cache_data
def update_hit_count(self, primary_id, **kwargs):
+ """Update hit count statistics in SQL storage."""
self.s.update_hit_count_by_id(primary_id)
def hit_cache_callback(self, res_data, **kwargs):
+ """Callback executed on cache hit to update memory cache."""
self.eviction_base.get(res_data[1])
def search(self, embedding_data, **kwargs):
+ """
+ Search for similar vectors in vector storage.
+
+ Applies normalization if configured and delegates to vector backend.
+ """
model = kwargs.pop("model", None)
- embedding_data = normalize(embedding_data)
+ if self.normalize:
+ embedding_data = normalize(embedding_data)
top_k = kwargs.get("top_k", -1)
return self.v.search(data=embedding_data, top_k=top_k, model=model)
def delete(self, id_list, **kwargs):
- model = kwargs.pop("model", None)
+ """
+ Delete cache entries from all storage backends.
+
+ Removes from memory cache, vector storage, and marks as deleted in SQL.
+ Returns detailed status of deletion operations.
+ """
+ model = kwargs.pop("model")
try:
+ # Remove from memory cache
+ for id in id_list:
+ self.eviction_base.get_cache(model).pop(id, None)
+ # Delete from vector storage
v_delete_count = self.v.delete(ids=id_list, model=model)
except Exception as e:
return {'status': 'failed', 'milvus': 'delete milvus data failed, please check! e: {}'.format(e),
'mysql': 'unexecuted'}
try:
+ # Mark as deleted in SQL storage
s_delete_count = self.s.mark_deleted(id_list)
except Exception as e:
return {'status': 'failed', 'milvus': 'success',
@@ -256,35 +321,43 @@ def delete(self, id_list, **kwargs):
'mysql': 'delete_count: '+str(s_delete_count)}
def create_index(self, model, **kwargs):
+ """Create vector index for a specific model."""
return self.v.create(model)
- def truncate(self, model_name):
- # drop vector base data
+ def truncate(self, model):
+ """
+ Truncate all data for a specific model across all storage backends.
+
+ Clears memory cache, rebuilds vector storage, and deletes SQL data.
+ Returns detailed status of truncation operations.
+ """
+ # Clear memory cache data
+ self.eviction_base.clear(model)
+
+ # Rebuild vector storage (drops and recreates collection)
try:
- vector_resp = self.v.rebuild_col(model_name)
+ vector_resp = self.v.rebuild_col(model)
except Exception as e:
return {'status': 'failed', 'VectorDB': 'truncate VectorDB data failed, please check! e: {}'.format(e),
'ScalarDB': 'unexecuted'}
if vector_resp:
return {'status': 'failed', 'VectorDB': vector_resp, 'ScalarDB': 'unexecuted'}
- # drop scalar base data
+
+ # Delete scalar data from SQL storage
try:
- delete_count = self.s.model_deleted(model_name)
+ delete_count = self.s.model_deleted(model)
except Exception as e:
return {'status': 'failed', 'VectorDB': 'rebuild',
'ScalarDB': 'truncate scalar data failed, please check! e: {}'.format(e)}
return {'status': 'success', 'VectorDB': 'rebuild', 'ScalarDB': 'delete_count: ' + str(delete_count)}
def flush(self):
+ """Flush all storage backends to ensure data persistence."""
self.s.flush()
self.v.flush()
def close(self):
+ """Close all storage connections and release resources."""
self.s.close()
self.v.close()
-
-# if __name__ == '__main__':
-# from modelcache.manager import CacheBase, VectorBase, get_data_manager
-# data_manager = get_data_manager(CacheBase('mysql'), VectorBase('milvus', dimension=128))
-# data_manager.save('hello', 'hi', np.random.random((128,)).astype('float32'), model='gptcode_6b')
diff --git a/modelcache/manager/eviction/__init__.py b/modelcache/manager/eviction/__init__.py
index 8ca7a3d..633f866 100644
--- a/modelcache/manager/eviction/__init__.py
+++ b/modelcache/manager/eviction/__init__.py
@@ -1,10 +1,2 @@
# -*- coding: utf-8 -*-
-from modelcache.utils.lazy_import import LazyImport
-eviction_manager = LazyImport(
- "eviction_manager", globals(), "modelcache.manager.eviction.manager"
-)
-
-
-def EvictionBase(name: str, **kwargs):
- return eviction_manager.EvictionBase.get(name, **kwargs)
diff --git a/modelcache/manager/eviction/arc_cache.py b/modelcache/manager/eviction/arc_cache.py
new file mode 100644
index 0000000..4f97231
--- /dev/null
+++ b/modelcache/manager/eviction/arc_cache.py
@@ -0,0 +1,133 @@
+from cachetools import Cache
+from collections import OrderedDict
+from readerwriterlock import rwlock
+
+_sentinel = object()
+
+class ARC(Cache):
+ """
+ Adaptive Replacement Cache (ARC) implementation.
+
+ ARC maintains four lists (T1, T2, B1, B2) to adaptively balance
+ between LRU and LFU eviction strategies based on access patterns.
+ """
+
+ def __init__(self, maxsize, getsizeof=None):
+ """Initialize ARC cache with maximum size."""
+ super().__init__(maxsize, getsizeof)
+ self.t1 = OrderedDict() # Recent items
+ self.t2 = OrderedDict() # Frequent items
+ self.b1 = OrderedDict() # Ghost entries for T1
+ self.b2 = OrderedDict() # Ghost entries for T2
+ self.p = 0 # Adaptive parameter
+ self._rw_lock = rwlock.RWLockWrite() # Thread safety
+
+ def __len__(self):
+ """Return total number of cached items."""
+ return len(self.t1) + len(self.t2)
+
+ def __contains__(self, key):
+ """Check if key exists in cache."""
+ return key in self.t1 or key in self.t2
+
+ def _evict_internal(self):
+ """Internal method to evict items when cache is full."""
+ # Evict from cache lists to ghost lists
+ while len(self.t1) + len(self.t2) > self.maxsize:
+ if len(self.t1) > self.p or (len(self.t1) == 0 and len(self.t2) > 0):
+ key, value = self.t1.popitem(last=False)
+ self.b1[key] = value
+ else:
+ key, value = self.t2.popitem(last=False)
+ self.b2[key] = value
+
+ # Maintain ghost list sizes
+ while len(self.b1) > (self.maxsize - self.p):
+ self.b1.popitem(last=False)
+ while len(self.b2) > self.p:
+ self.b2.popitem(last=False)
+
+ def __setitem__(self, key, value):
+ """Insert or update a cache entry."""
+ with self._rw_lock.gen_wlock():
+ # Remove key from all lists first
+ for l in (self.t1, self.t2, self.b1, self.b2):
+ l.pop(key, None)
+ # Add to recent list (T1)
+ self.t1[key] = value
+ self.t1.move_to_end(key)
+ self._evict_internal()
+
+ def __getitem__(self, key):
+ """Retrieve a cache entry and update access pattern."""
+ with self._rw_lock.gen_wlock():
+ if key in self.t1:
+ # Move from recent to frequent list
+ value = self.t1.pop(key)
+ self.t2[key] = value
+ self.t2.move_to_end(key)
+ self.p = max(0, self.p - 1) # Adjust adaptive parameter
+ self._evict_internal()
+ return value
+ if key in self.t2:
+ # Access frequent list
+ value = self.t2.pop(key)
+ self.t2[key] = value
+ self.t2.move_to_end(key)
+ self.p = min(self.maxsize, self.p + 1) # Adjust adaptive parameter
+ self._evict_internal()
+ return value
+ if key in self.b1:
+ # Promote from ghost list B1 to frequent list T2
+ self.b1.pop(key)
+ self.p = min(self.maxsize, self.p + 1) # Adjust adaptive parameter
+ self._evict_internal()
+ value = super().__missing__(key)
+ self.t2[key] = value
+ self.t2.move_to_end(key)
+ return value
+ if key in self.b2:
+ # Promote from ghost list B2 to frequent list T2
+ self.b2.pop(key)
+ self.p = max(0, self.p - 1) # Adjust adaptive parameter
+ self._evict_internal()
+ value = super().__missing__(key)
+ self.t2[key] = value
+ self.t2.move_to_end(key)
+ return value
+ return super().__getitem__(key)
+
+ def __missing__(self, key):
+ """Handle missing keys."""
+ raise KeyError(key)
+
+ def pop(self, key, default=_sentinel):
+ """Remove a cache entry."""
+ with self._rw_lock.gen_wlock():
+ for l in (self.t1, self.t2, self.b1, self.b2):
+ if key in l:
+ return l.pop(key)
+ if default is _sentinel:
+ raise KeyError(key)
+ return default
+
+ def clear(self):
+ """Clear all cache entries."""
+ with self._rw_lock.gen_wlock():
+ self.t1.clear()
+ self.t2.clear()
+ self.b1.clear()
+ self.b2.clear()
+ self.p = 0
+ super().clear()
+
+ def __iter__(self):
+ """Iterate over cache keys."""
+ yield from self.t1
+ yield from self.t2
+
+ def __repr__(self):
+ """Return string representation of the cache."""
+ return (f"ARC(maxsize={self.maxsize}, p={self.p}, len={len(self)}, "
+ f"t1_len={len(self.t1)}, t2_len={len(self.t2)}, "
+ f"b1_len={len(self.b1)}, b2_len={len(self.b2)})")
diff --git a/modelcache/manager/eviction/base.py b/modelcache/manager/eviction/base.py
index 352c9a3..56efa54 100644
--- a/modelcache/manager/eviction/base.py
+++ b/modelcache/manager/eviction/base.py
@@ -9,11 +9,11 @@ class EvictionBase(metaclass=ABCMeta):
"""
@abstractmethod
- def put(self, objs: List[Any]):
+ def put(self, objs: List[Any], model:str):
pass
@abstractmethod
- def get(self, obj: Any):
+ def get(self, obj: Any, model:str):
pass
@property
diff --git a/modelcache/manager/eviction/manager.py b/modelcache/manager/eviction/manager.py
deleted file mode 100644
index 61579f0..0000000
--- a/modelcache/manager/eviction/manager.py
+++ /dev/null
@@ -1,25 +0,0 @@
-# -*- coding: utf-8 -*-
-from typing import Callable, List, Any
-from modelcache.utils.error import NotFoundError
-
-
-class EvictionBase:
- """
- EvictionBase to evict the cache data.
- """
-
- def __init__(self):
- raise EnvironmentError(
- "EvictionBase is designed to be instantiated, "
- "please using the `EvictionBase.get(name, policy, maxsize, clean_size)`."
- )
-
- @staticmethod
- def get(name: str, policy: str, maxsize: int, clean_size: int, on_evict: Callable[[List[Any]], None], **kwargs):
- if name in "memory":
- from modelcache.manager.eviction.memory_cache import MemoryCacheEviction
-
- eviction_base = MemoryCacheEviction(policy, maxsize, clean_size, on_evict, **kwargs)
- else:
- raise NotFoundError("eviction base", name)
- return eviction_base
diff --git a/modelcache/manager/eviction/memory_cache.py b/modelcache/manager/eviction/memory_cache.py
index 3cb487f..838e8cd 100644
--- a/modelcache/manager/eviction/memory_cache.py
+++ b/modelcache/manager/eviction/memory_cache.py
@@ -1,8 +1,10 @@
# -*- coding: utf-8 -*-
-from typing import Any, Callable, List
+from typing import Any, Callable, List, Tuple
import cachetools
from modelcache.manager.eviction.base import EvictionBase
+from .arc_cache import ARC
+from .wtinylfu_cache import W2TinyLFU
def popitem_wrapper(func, wrapper_func, clean_size):
@@ -17,28 +19,52 @@ def wrapper(*args, **kwargs):
class MemoryCacheEviction(EvictionBase):
- def __init__(self, policy: str, maxsize: int, clean_size: int, on_evict: Callable[[List[Any]], None], **kwargs):
+ def __init__(self, policy: str, maxsize: int, clean_size: int, **kwargs):
self._policy = policy.upper()
+ self.model_to_cache = dict()
+ self.maxsize = maxsize
+ self.clean_size = clean_size
+ self.kwargs = kwargs
+
+ def create_cache(self, model: str):
if self._policy == "LRU":
- self._cache = cachetools.LRUCache(maxsize=maxsize, **kwargs)
+ cache = cachetools.LRUCache(maxsize=self.maxsize, **self.kwargs)
elif self._policy == "LFU":
- self._cache = cachetools.LFUCache(maxsize=maxsize, **kwargs)
+ cache = cachetools.LFUCache(maxsize=self.maxsize, **self.kwargs)
elif self._policy == "FIFO":
- self._cache = cachetools.FIFOCache(maxsize=maxsize, **kwargs)
+ cache = cachetools.FIFOCache(maxsize=self.maxsize, **self.kwargs)
elif self._policy == "RR":
- self._cache = cachetools.RRCache(maxsize=maxsize, **kwargs)
+ cache = cachetools.RRCache(maxsize=self.maxsize, **self.kwargs)
+ elif self._policy == "WTINYLFU":
+ cache = W2TinyLFU(maxsize=self.maxsize)
+ elif self._policy == "ARC":
+ cache = ARC(maxsize=self.maxsize)
else:
- raise ValueError(f"Unknown policy {policy}")
+ raise ValueError(f"Unknown policy {self.policy}")
+ return cache
+
+ def put(self, objs: List[Tuple[Any, Any]], model: str):
+ cache = self.get_cache(model)
+ for key, value in objs:
+ cache[key] = value
+
+
+ def get(self, obj: Any, model: str):
+ cache = self.get_cache(model)
+ return cache.get(obj)
- self._cache.popitem = popitem_wrapper(self._cache.popitem, on_evict, clean_size)
- def put(self, objs: List[Any]):
- for obj in objs:
- self._cache[obj] = True
+ def clear(self, model: str):
+ self.model_to_cache.pop(model, None)
+
+
+ def get_cache(self, model: str):
+ if not model in self.model_to_cache:
+ self.model_to_cache[model] = self.create_cache(model)
+ return self.model_to_cache[model]
- def get(self, obj: Any):
- return self._cache.get(obj)
@property
def policy(self) -> str:
return self._policy
+
diff --git a/modelcache/manager/eviction/wtinylfu_cache.py b/modelcache/manager/eviction/wtinylfu_cache.py
new file mode 100644
index 0000000..378b1ed
--- /dev/null
+++ b/modelcache/manager/eviction/wtinylfu_cache.py
@@ -0,0 +1,172 @@
+from cachetools import LRUCache, Cache, LFUCache
+from readerwriterlock import rwlock
+import random
+
+class CountMinSketch:
+ def __init__(self, width=1024, depth=4, decay_interval=10000):
+ """Initialize Count-Min Sketch with specified dimensions."""
+ self.width = width
+ self.depth = depth
+ self.tables = [[0]*width for _ in range(depth)] # Hash tables
+ self.seeds = [random.randrange(1<<30) for _ in range(depth)] # Hash seeds
+ self.ops = 0 # Operation counter for decay trigger
+ self.decay_interval = decay_interval
+
+ def _hash(self, x, seed):
+ """Hash function for mapping items to table positions."""
+ return hash((x, seed)) % self.width
+
+ def add(self, x):
+ """Add an item and increment its frequency estimate."""
+ self.ops += 1
+ est = self.estimate(x) # Get current estimate
+ # Update all hash tables
+ for i, seed in enumerate(self.seeds):
+ idx = self._hash(x, seed)
+ if self.tables[i][idx] <= est:
+ self.tables[i][idx] += 1
+
+ # Periodic decay to handle changing patterns
+ if self.ops >= self.decay_interval:
+ self.decay()
+ self.ops = 0
+
+ def estimate(self, x):
+ """Estimate frequency of an item (minimum across all tables)."""
+ return min(self.tables[i][self._hash(x, seed)]
+ for i, seed in enumerate(self.seeds))
+
+ def decay(self):
+ """Decay all frequency counts by half."""
+ for table in self.tables:
+ for i in range(len(table)):
+ table[i] >>= 1 # Right shift (divide by 2)
+
+class W2TinyLFU(Cache):
+ """
+ Window Tiny LFU cache implementation.
+
+ Combines a small LRU window cache with a main cache divided into
+ probation and protected segments, using frequency estimation for
+ admission control.
+ """
+
+ def __init__(self, maxsize, window_pct=0.01):
+ """
+ Initialize W-TinyLFU cache.
+
+ Args:
+ maxsize: Maximum size of the cache
+ window_pct: Percentage of cache size for the window (default 1%)
+ """
+ super().__init__(maxsize)
+ self.window_size = max(1, int(maxsize * window_pct))
+ rest = maxsize - self.window_size
+ self.probation_size = rest // 2
+ self.protected_size = rest - self.probation_size
+
+ # Three cache segments
+ self.window = LRUCache(maxsize=self.window_size) # Recent items
+ self.probation = LFUCache(maxsize=self.probation_size) # New main cache items
+ self.protected = LFUCache(maxsize=self.protected_size) # Frequently accessed items
+
+ self.cms = CountMinSketch() # Frequency estimator
+ self.data = {} # Cache data storage
+ self._rw_lock = rwlock.RWLockWrite() # Read-write lock for thread safety
+
+ def __setitem__(self, key, value):
+ """Add or update an item in the cache."""
+ with self._rw_lock.gen_wlock():
+ self.data[key] = value
+ self._put(key)
+
+ def __getitem__(self, key):
+ """Retrieve an item from the cache."""
+ val = self.get(key, default=None)
+ if val is None:
+ raise KeyError(key)
+ return val
+
+ def __contains__(self, key):
+ """Check if an item exists in the cache."""
+ return key in self.window or key in self.probation or key in self.protected
+
+ def __delitem__(self, key):
+ """Remove an item from the cache."""
+ with self._rw_lock.gen_wlock():
+ self.data.pop(key, None)
+ self.window.pop(key, None)
+ self.probation.pop(key, None)
+ self.protected.pop(key, None)
+
+ def get(self, key, default=None):
+ """
+ Retrieve an item from the cache, updating its position
+ in the cache hierarchy if necessary.
+ """
+ if key in self.window:
+ self.window[key] = True
+ return self.data.get(key, default)
+ if key in self.protected:
+ self.protected[key] = True
+ return self.data.get(key, default)
+ if key in self.probation:
+ self.probation.pop(key)
+ if len(self.protected) >= self.protected_size:
+ demoted = next(iter(self.protected))
+ self.protected.pop(demoted)
+ self.probation[demoted] = True
+ self.protected[key] = True
+ return self.data.get(key, default)
+ return default
+
+ def _put(self, key):
+ """
+ Add an item to the cache, using frequency-based admission
+ control and eviction policies.
+ """
+ self.cms.add(key)
+ if key in self:
+ return
+
+ if len(self.window) < self.window_size:
+ self.window[key] = True
+ return
+
+ victim = next(iter(self.window))
+ self.window.pop(victim)
+
+ if self.cms.estimate(key) >= self.cms.estimate(victim):
+ self._admit_to_main(victim)
+ self._admit_to_main(key)
+ else:
+ self._admit_to_main(victim)
+ self.data.pop(key, None)
+
+ def _admit_to_main(self, key):
+ """
+ Admit an item to the main cache (probation or protected segment).
+ """
+ if key in self.protected or key in self.probation:
+ return
+ if self.probation_size == 0:
+ self.data.pop(key, None)
+ return
+ if len(self.probation) < self.probation_size:
+ self.probation[key] = True
+ elif self.probation:
+ evicted = next(iter(self.probation))
+ self.probation.pop(evicted)
+ self.probation[key] = True
+ self.data.pop(evicted, None)
+ else:
+ self.data.pop(key, None)
+
+ def clear(self):
+ """Clear all items from the cache."""
+ with self._rw_lock.gen_wlock():
+ self.window.clear()
+ self.probation.clear()
+ self.protected.clear()
+ self.data.clear()
+
diff --git a/modelcache/manager/factory.py b/modelcache/manager/factory.py
deleted file mode 100644
index 7246cc6..0000000
--- a/modelcache/manager/factory.py
+++ /dev/null
@@ -1,27 +0,0 @@
-# -*- coding: utf-8 -*-
-from typing import Union, Callable
-from modelcache.manager import CacheBase, VectorBase, ObjectBase
-from modelcache.manager.data_manager import SSDataManager, MapDataManager
-
-
-def get_data_manager(
- cache_base: Union[CacheBase, str] = None,
- vector_base: Union[VectorBase, str] = None,
- object_base: Union[ObjectBase, str] = None,
- max_size: int = 1000,
- clean_size: int = None,
- eviction: str = "LRU",
- data_path: str = "data_map.txt",
- get_data_container: Callable = None,
-):
- if not cache_base and not vector_base:
- return MapDataManager(data_path, max_size, get_data_container)
-
- if isinstance(cache_base, str):
- cache_base = CacheBase(name=cache_base)
- if isinstance(vector_base, str):
- vector_base = VectorBase(name=vector_base)
- if isinstance(object_base, str):
- object_base = ObjectBase(name=object_base)
- assert cache_base and vector_base
- return SSDataManager(cache_base, vector_base, object_base, max_size, clean_size, eviction)
diff --git a/modelcache/manager/object_data/__init__.py b/modelcache/manager/object_data/__init__.py
index f1186c0..633f866 100644
--- a/modelcache/manager/object_data/__init__.py
+++ b/modelcache/manager/object_data/__init__.py
@@ -1,9 +1,2 @@
# -*- coding: utf-8 -*-
-from modelcache.utils.lazy_import import LazyImport
-object_manager = LazyImport(
- "object_manager", globals(), "modelcache.manager.object_data.manager"
-)
-
-def ObjectBase(name: str, **kwargs):
- return object_manager.ObjectBase.get(name, **kwargs)
diff --git a/modelcache/manager/object_data/base.py b/modelcache/manager/object_data/base.py
index 9ed00cf..41466f9 100644
--- a/modelcache/manager/object_data/base.py
+++ b/modelcache/manager/object_data/base.py
@@ -12,10 +12,6 @@ class ObjectBase(ABC):
def put(self, obj: Any) -> str:
pass
- @abstractmethod
- def get(self, obj: str) -> Any:
- pass
-
@abstractmethod
def get_access_link(self, obj: str) -> str:
pass
@@ -23,3 +19,7 @@ def get_access_link(self, obj: str) -> str:
@abstractmethod
def delete(self, to_delete: List[str]):
pass
+
+ @staticmethod
+ def get(name: str) -> Any:
+ pass
diff --git a/modelcache/manager/scalar_data/__init__.py b/modelcache/manager/scalar_data/__init__.py
index b63c430..40a96af 100644
--- a/modelcache/manager/scalar_data/__init__.py
+++ b/modelcache/manager/scalar_data/__init__.py
@@ -1,9 +1 @@
# -*- coding: utf-8 -*-
-from modelcache.utils.lazy_import import LazyImport
-scalar_manager = LazyImport(
- "scalar_manager", globals(), "modelcache.manager.scalar_data.manager"
-)
-
-
-def CacheBase(name: str, **kwargs):
- return scalar_manager.CacheBase.get(name, **kwargs)
diff --git a/modelcache/manager/scalar_data/base.py b/modelcache/manager/scalar_data/base.py
index 6db058d..5ed898b 100644
--- a/modelcache/manager/scalar_data/base.py
+++ b/modelcache/manager/scalar_data/base.py
@@ -5,6 +5,9 @@
from enum import IntEnum
import numpy as np
+from modelcache.utils import import_sql_client
+from modelcache.utils.error import NotFoundError
+
class DataType(IntEnum):
STR = 0
@@ -107,7 +110,7 @@ def mark_deleted(self, keys):
pass
@abstractmethod
- def model_deleted(self, model_name):
+ def model_deleted(self, model):
pass
@abstractmethod
@@ -122,9 +125,39 @@ def get_ids(self, deleted=True):
def count(self):
pass
+ @abstractmethod
def flush(self):
pass
@abstractmethod
def close(self):
pass
+
+ @abstractmethod
+ def batch_insert(self, all_data: List[CacheData]):
+ pass
+
+ @abstractmethod
+ def update_hit_count_by_id(self, primary_id):
+ pass
+
+ @staticmethod
+ def get(name, **kwargs):
+ if name in ["mysql", "oceanbase"]:
+ from modelcache.manager.scalar_data.sql_storage import SQLStorage
+ config = kwargs.get("config")
+ import_sql_client(name)
+ cache_base = SQLStorage(db_type=name, config=config)
+ elif name == 'sqlite':
+ SQL_URL = {"sqlite": "./sqlite.db"}
+ from modelcache.manager.scalar_data.sql_storage_sqlite import SQLStorage
+ sql_url = kwargs.get("sql_url", SQL_URL[name])
+ cache_base = SQLStorage(db_type=name, url=sql_url)
+ elif name == 'elasticsearch':
+ from modelcache.manager.scalar_data.sql_storage_es import SQLStorage
+ config = kwargs.get("config")
+ cache_base = SQLStorage(db_type=name, config=config)
+ else:
+ raise NotFoundError("cache store", name)
+ return cache_base
+
diff --git a/modelcache/manager/scalar_data/manager.py b/modelcache/manager/scalar_data/manager.py
deleted file mode 100644
index 4c02c45..0000000
--- a/modelcache/manager/scalar_data/manager.py
+++ /dev/null
@@ -1,32 +0,0 @@
-# -*- coding: utf-8 -*-
-from modelcache.utils import import_sql_client
-from modelcache.utils.error import NotFoundError
-
-SQL_URL = {"sqlite": "./sqlite.db"}
-
-
-class CacheBase:
- """
- CacheBase to manager the cache storage.
- """
-
- def __init__(self):
- raise EnvironmentError(
- "CacheBase is designed to be instantiated, please using the `CacheBase.get(name)`."
- )
-
- @staticmethod
- def get(name, **kwargs):
-
- if name in ["mysql", "oceanbase"]:
- from modelcache.manager.scalar_data.sql_storage import SQLStorage
- config = kwargs.get("config")
- import_sql_client(name)
- cache_base = SQLStorage(db_type=name, config=config)
- elif name == 'sqlite':
- from modelcache.manager.scalar_data.sql_storage_sqlite import SQLStorage
- sql_url = kwargs.get("sql_url", SQL_URL[name])
- cache_base = SQLStorage(db_type=name, url=sql_url)
- else:
- raise NotFoundError("cache store", name)
- return cache_base
diff --git a/modelcache/manager/scalar_data/sql_storage.py b/modelcache/manager/scalar_data/sql_storage.py
index 52fd559..20eba9b 100644
--- a/modelcache/manager/scalar_data/sql_storage.py
+++ b/modelcache/manager/scalar_data/sql_storage.py
@@ -1,16 +1,16 @@
# -*- coding: utf-8 -*-
-import os
-import time
+import uuid
import pymysql
import json
-import base64
+import numpy as np
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",
@@ -42,26 +42,59 @@ def _insert(self, data: List):
answer_type = 0
embedding_data = embedding_data.tobytes()
is_deleted = 0
+ _id = str(uuid.uuid4())
table_name = "modelcache_llm_answer"
- insert_sql = "INSERT INTO {} (question, answer, answer_type, model, embedding_data, is_deleted) VALUES (%s, %s, %s, %s, _binary%s, %s)".format(table_name)
+ insert_sql = f"""
+ INSERT INTO {table_name}
+ (id, question, answer, answer_type, model, embedding_data, is_deleted)
+ VALUES (%s, %s, %s, %s, %s, _binary%s, %s)
+ """
conn = self.pool.connection()
try:
with conn.cursor() as cursor:
# 执行插入数据操作
- values = (question, answer, answer_type, model, embedding_data, is_deleted)
+ values = (_id, question, answer, answer_type, model, embedding_data, is_deleted)
cursor.execute(insert_sql, values)
conn.commit()
- id = cursor.lastrowid
finally:
# 关闭连接，将连接返回给连接池
conn.close()
- return id
+ return _id
+
+ def batch_insert(self, all_data: List[List]):
+ table_name = "modelcache_llm_answer"
+ insert_sql = f"""
+ INSERT INTO {table_name}
+ (id, question, answer, answer_type, model, embedding_data, is_deleted)
+ VALUES (%s, %s, %s, %s, %s, %s, %s)
+ """
- def batch_insert(self, all_data: List[CacheData]):
+ values_list = []
ids = []
+
for data in all_data:
- ids.append(self._insert(data))
+ answer = data[0]
+ question = data[1]
+ embedding_data = data[2].tobytes()
+ model = data[3]
+ answer_type = 0
+ is_deleted = 0
+ _id = str(uuid.uuid4())
+ ids.append(_id)
+
+ values_list.append((
+ _id, question, answer, answer_type, model, embedding_data, is_deleted
+ ))
+
+ conn = self.pool.connection()
+ try:
+ with conn.cursor() as cursor:
+ cursor.executemany(insert_sql, values_list)
+ conn.commit()
+ finally:
+ conn.close()
+
return ids
def insert_query_resp(self, query_resp, **kwargs):
@@ -78,7 +111,11 @@ def insert_query_resp(self, query_resp, **kwargs):
hit_query = json.dumps(hit_query, ensure_ascii=False)
table_name = "modelcache_query_log"
- 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)
+ insert_sql = f"""
+ INSERT INTO {table_name}
+ (error_code, error_desc, cache_hit, model, query, delta_time, hit_query, answer)
+ VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
+ """
conn = self.pool.connection()
try:
with conn.cursor() as cursor:
@@ -92,35 +129,41 @@ def insert_query_resp(self, query_resp, **kwargs):
def get_data_by_id(self, key: int):
table_name = "modelcache_llm_answer"
- query_sql = "select question, answer, embedding_data, model from {} where id={}".format(table_name, key)
- conn_start = time.time()
+ query_sql = f"""
+ SELECT answer, question, embedding_data, model
+ FROM {table_name}
+ WHERE id = %s
+ """
conn = self.pool.connection()
-
- search_start = time.time()
try:
with conn.cursor() as cursor:
# 执行数据库操作
- cursor.execute(query_sql)
+ cursor.execute(query_sql, (key,))
resp = cursor.fetchone()
finally:
# 关闭连接，将连接返回给连接池
conn.close()
if resp is not None and len(resp) == 4:
- return resp
+ # parse the numpy array from bytes and return the data
+ return resp[0], resp[1], np.frombuffer(resp[2], dtype=np.float32), resp[3]
else:
return None
def update_hit_count_by_id(self, primary_id: int):
table_name = "modelcache_llm_answer"
- update_sql = "UPDATE {} SET hit_count = hit_count+1 WHERE id={}".format(table_name, primary_id)
+ update_sql = f"""
+ UPDATE {table_name}
+ SET hit_count = hit_count+1
+ WHERE id = %s
+ """
conn = self.pool.connection()
# 使用连接执行更新数据操作
try:
with conn.cursor() as cursor:
# 执行更新数据操作
- cursor.execute(update_sql)
+ cursor.execute(update_sql,(primary_id,))
conn.commit()
finally:
# 关闭连接，将连接返回给连接池
@@ -129,12 +172,16 @@ def update_hit_count_by_id(self, primary_id: int):
def get_ids(self, deleted=True):
table_name = "modelcache_llm_answer"
state = 1 if deleted else 0
- query_sql = "Select id FROM {} WHERE is_deleted = {}".format(table_name, state)
+ query_sql = f"""
+ SELECT id
+ FROM {table_name}
+ WHERE is_deleted = %s
+ """
conn = self.pool.connection()
try:
with conn.cursor() as cursor:
- cursor.execute(query_sql)
+ cursor.execute(query_sql, (state,))
ids = [row[0] for row in cursor.fetchall()]
finally:
conn.close()
@@ -143,37 +190,45 @@ def get_ids(self, deleted=True):
def mark_deleted(self, keys):
table_name = "modelcache_llm_answer"
- mark_sql = " update {} set is_deleted=1 WHERE id in ({})".format(table_name, ",".join([str(i) for i in keys]))
+ placeholders = ",".join(["%s"] * len(keys))
+ mark_sql = f"""
+ UPDATE {table_name}
+ SET is_deleted=1
+ WHERE id in ({placeholders})
+ """
- # 从连接池中获取连接
conn = self.pool.connection()
try:
with conn.cursor() as cursor:
- # 执行删除数据操作
- cursor.execute(mark_sql)
+ cursor.execute(mark_sql, keys)
delete_count = cursor.rowcount
conn.commit()
finally:
- # 关闭连接，将连接返回给连接池
conn.close()
return delete_count
def model_deleted(self, model_name):
table_name = "modelcache_llm_answer"
- delete_sql = "Delete from {} WHERE model='{}'".format(table_name, model_name)
+ delete_sql = f"""
+ Delete from {table_name}
+ WHERE model = %s
+ """
table_log_name = "modelcache_query_log"
- delete_log_sql = "Delete from {} WHERE model='{}'".format(table_log_name, model_name)
+ delete_log_sql = f"""
+ Delete from {table_log_name}
+ WHERE model = %s
+ """
conn = self.pool.connection()
# 使用连接执行删除数据操作
try:
with conn.cursor() as cursor:
# 执行删除数据操作
- resp = cursor.execute(delete_sql)
+ resp = cursor.execute(delete_sql, (model_name,))
conn.commit()
# 执行删除该模型对应日志操作 resp_log行数不返回
- resp_log = cursor.execute(delete_log_sql)
+ resp_log = cursor.execute(delete_log_sql, (model_name,))
conn.commit() # 分别提交事务
finally:
# 关闭连接，将连接返回给连接池
@@ -182,7 +237,10 @@ def model_deleted(self, model_name):
def clear_deleted_data(self):
table_name = "modelcache_llm_answer"
- delete_sql = "DELETE FROM {} WHERE is_deleted = 1".format(table_name)
+ delete_sql = f"""
+ DELETE FROM {table_name}
+ WHERE is_deleted = 1
+ """
conn = self.pool.connection()
try:
@@ -197,10 +255,15 @@ def clear_deleted_data(self):
def count(self, state: int = 0, is_all: bool = False):
table_name = "modelcache_llm_answer"
+
+ # we're not using prepared statements here, so we need to ensure state is an integer
+ if not isinstance(state, int):
+ raise ValueError("'state' must be an integer.")
+
if is_all:
- count_sql = "SELECT COUNT(*) FROM {}".format(table_name)
+ count_sql = f"SELECT COUNT(*) FROM {table_name}"
else:
- count_sql = "SELECT COUNT(*) FROM {} WHERE is_deleted = {}".format(table_name,state)
+ count_sql = f"SELECT COUNT(*) FROM {table_name} WHERE is_deleted = {state}"
conn = self.pool.connection()
try:
@@ -217,3 +280,6 @@ def close(self):
def count_answers(self):
pass
+
+ def flush(self):
+ pass
\ No newline at end of file
diff --git a/modelcache/manager/scalar_data/sql_storage_es.py b/modelcache/manager/scalar_data/sql_storage_es.py
new file mode 100644
index 0000000..3787623
--- /dev/null
+++ b/modelcache/manager/scalar_data/sql_storage_es.py
@@ -0,0 +1,199 @@
+# -*- coding: utf-8 -*-
+import json
+from typing import List
+from elasticsearch import Elasticsearch, helpers
+from modelcache.manager.scalar_data.base import CacheStorage, CacheData
+import time
+from snowflake import SnowflakeGenerator
+
+
+class SQLStorage(CacheStorage):
+
+ def __init__(
+ self,
+ db_type: str = "elasticsearch",
+ config=None
+ ):
+ self.host = config.get('elasticsearch', 'host')
+ self.port = int(config.get('elasticsearch', 'port'))
+ self.client = Elasticsearch(
+ hosts=[{"host": self.host, "port": self.port}],
+ timeout=30,
+ http_auth=('esuser', 'password')
+ )
+
+ self.log_index = "modelcache_query_log"
+ self.ans_index = "modelcache_llm_answer"
+ self.create()
+ self.instance_id = 1 # 雪花算法使用的机器id 使用同一套数据库的分布式系统需要配置不同id
+ # 生成雪花id
+ self.snowflake_id = SnowflakeGenerator(self.instance_id)
+
+ def create(self):
+ answer_index_body = {
+ "mappings": {
+ "properties": {
+ "gmt_create": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+ "gmt_modified": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+ "question": {"type": "text"},
+ "answer": {"type": "text"},
+ "answer_type": {"type": "integer"},
+ "hit_count": {"type": "integer"},
+ "model": {"type": "keyword"},
+ "embedding_data": {"type": "binary"},
+ "is_deleted": {"type": "integer"},
+ }
+ }
+ }
+
+ log_index_body = {
+ "mappings": {
+ "properties": {
+ "gmt_create": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+ "gmt_modified": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+ "error_code": {"type": "integer"},
+ "error_desc": {"type": "text"},
+ "cache_hit": {"type": "keyword"},
+ "delta_time": {"type": "float"},
+ "model": {"type": "keyword"},
+ "query": {"type": "text"},
+ "hit_query": {"type": "text"},
+ "answer": {"type": "text"}
+ }
+ }
+ }
+
+ if not self.client.indices.exists(index=self.ans_index):
+ self.client.indices.create(index=self.ans_index, body=answer_index_body)
+
+ if not self.client.indices.exists(index=self.log_index):
+ self.client.indices.create(index=self.log_index, body=log_index_body)
+
+ def _insert(self, data: List) -> str or None:
+ doc = {
+ "answer": data[0],
+ "question": data[1],
+ "embedding_data": data[2].tolist() if hasattr(data[2], "tolist") else data[2],
+ "model": data[3],
+ "answer_type": 0,
+ "hit_count": 0,
+ "is_deleted": 0
+ }
+
+ try:
+
+ response = self.client.index(
+ index=self.ans_index,
+ id=next(self.snowflake_id),
+ body=doc,
+ )
+ return int(response['_id'])
+ except Exception as e:
+
+ print(f"Failed to insert document: {e}")
+ return None
+
+ def batch_insert(self, all_data: List[List]) -> List[str]:
+ successful_ids = []
+ for data in all_data:
+ _id = self._insert(data)
+ if _id is not None:
+ successful_ids.append(_id)
+ self.client.indices.refresh(index=self.ans_index) # 批量插入后手动刷新
+
+ return successful_ids
+
+ def insert_query_resp(self, query_resp, **kwargs):
+ doc = {
+ "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": json.dumps(query_resp.get('hit_query'), ensure_ascii=False) if isinstance(
+ query_resp.get('hit_query'), list) else query_resp.get('hit_query'),
+ "answer": query_resp.get('answer'),
+ "hit_count": 0,
+ "is_deleted": 0
+
+ }
+ self.client.index(index=self.log_index, body=doc)
+
+ def get_data_by_id(self, key: int):
+ try:
+ response = self.client.get(index=self.ans_index, id=key, _source=['question', 'answer', 'embedding_data', 'model'])
+ source = response["_source"]
+ result = [
+ source.get('question'),
+ source.get('answer'),
+ source.get('embedding_data'),
+ source.get('model')
+ ]
+ return result
+ except Exception as e:
+ print(e)
+
+ def update_hit_count_by_id(self, primary_id: int):
+ self.client.update(
+ index=self.ans_index,
+ id=primary_id,
+ body={"script": {"source": "ctx._source.hit_count += 1"}}
+ )
+
+ def get_ids(self, deleted=True):
+ query = {
+ "query": {
+ "term": {"is_deleted": 1 if deleted else 0}
+ }
+ }
+ response = self.client.search(index=self.ans_index, body=query)
+ return [hit["_id"] for hit in response["hits"]["hits"]]
+
+ def mark_deleted(self, keys):
+ actions = [
+ {
+ "_op_type": "update",
+ "_index": self.ans_index,
+ "_id": key,
+ "doc": {"is_deleted": 1}
+ }
+ for key in keys
+ ]
+ responses = helpers.bulk(self.client, actions)
+ return responses[0] # 返回更新的文档数
+
+ def model_deleted(self, model_name):
+ query = {
+ "query": {
+ "term": {"model": model_name}
+ }
+ }
+
+ response = self.client.delete_by_query(index=self.ans_index, body=query)
+ return response["deleted"]
+
+ def clear_deleted_data(self):
+ query = {
+ "query": {
+ "term": {"is_deleted": 1}
+ }
+ }
+ response = self.client.delete_by_query(index=self.ans_index, body=query)
+ return response["deleted"]
+
+ def count(self, state: int = 0, is_all: bool = False):
+ query = {"query": {"match_all": {}}} if is_all else {"query": {"term": {"is_deleted": state}}}
+ response = self.client.count(index=self.ans_index, body=query)
+ return response["count"]
+
+ def close(self):
+ self.client.close()
+
+ def count_answers(self):
+ query = {"query": {"match_all": {}}}
+ response = self.client.count(index=self.ans_index, body=query)
+ return response["count"]
+
+ def flush(self):
+ pass
diff --git a/modelcache/manager/scalar_data/sql_storage_sqlite.py b/modelcache/manager/scalar_data/sql_storage_sqlite.py
index 3c68f9e..3b69979 100644
--- a/modelcache/manager/scalar_data/sql_storage_sqlite.py
+++ b/modelcache/manager/scalar_data/sql_storage_sqlite.py
@@ -100,8 +100,7 @@ def insert_query_resp(self, query_resp, **kwargs):
hit_query = json.dumps(hit_query, ensure_ascii=False)
table_name = "modelcache_query_log"
- 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)
-
+ insert_sql = "INSERT INTO {} (error_code, error_desc, cache_hit, model, query, delta_time, hit_query, answer) VALUES (?, ?, ?, ?, ?, ?, ?, ?)".format(table_name)
conn = sqlite3.connect(self._url)
try:
cursor = conn.cursor()
@@ -200,3 +199,6 @@ def close(self):
def count_answers(self):
pass
+
+ def flush(self):
+ pass
diff --git a/modelcache/manager/vector_data/__init__.py b/modelcache/manager/vector_data/__init__.py
index fdb98cb..40a96af 100644
--- a/modelcache/manager/vector_data/__init__.py
+++ b/modelcache/manager/vector_data/__init__.py
@@ -1,10 +1 @@
# -*- coding: utf-8 -*-
-from modelcache.utils.lazy_import import LazyImport
-
-vector_manager = LazyImport(
- "vector_manager", globals(), "modelcache.manager.vector_data.manager"
-)
-
-
-def VectorBase(name: str, **kwargs):
- return vector_manager.VectorBase.get(name, **kwargs)
diff --git a/modelcache/manager/vector_data/base.py b/modelcache/manager/vector_data/base.py
index 7bdf12a..689f833 100644
--- a/modelcache/manager/vector_data/base.py
+++ b/modelcache/manager/vector_data/base.py
@@ -4,6 +4,19 @@
from typing import List
from dataclasses import dataclass
+from modelcache.embedding import MetricType
+from modelcache.utils.error import ParamError, NotFoundError
+
+TOP_K = 1
+FAISS_INDEX_PATH = "faiss.index"
+DIMENSION = 0
+MILVUS_HOST = "localhost"
+MILVUS_PORT = 19530
+MILVUS_USER = ""
+MILVUS_PSW = ""
+MILVUS_SECURE = False
+
+COLLECTION_NAME = "modelcache"
@dataclass
class VectorData:
@@ -11,8 +24,7 @@ class VectorData:
data: np.ndarray
-class VectorBase(ABC):
- """VectorBase: base vector store interface"""
+class VectorStorage(ABC):
@abstractmethod
def mul_add(self, datas: List[VectorData], model=None):
@@ -34,8 +46,108 @@ def delete(self, ids) -> bool:
def rebuild_col(self, model):
pass
+ @abstractmethod
def flush(self):
pass
+ @abstractmethod
def close(self):
pass
+
+ @staticmethod
+ def get(name, **kwargs):
+ top_k = kwargs.get("top_k", TOP_K)
+ if name == "milvus":
+ from modelcache.manager.vector_data.milvus import Milvus
+ dimension = kwargs.get("dimension", DIMENSION)
+ milvus_config = kwargs.get("config")
+ check_dimension(dimension)
+ host = milvus_config.get('milvus', 'host')
+ port = milvus_config.get('milvus', 'port')
+ user = milvus_config.get('milvus', 'user')
+ password = milvus_config.get('milvus', 'password')
+
+ metric_type = kwargs.get("metric_type",MetricType.COSINE)
+ secure = kwargs.get("secure", MILVUS_SECURE)
+ collection_name = kwargs.get("collection_name", COLLECTION_NAME)
+ index_params = kwargs.get("index_params", None)
+ search_params = kwargs.get("search_params", None)
+ local_mode = kwargs.get("local_mode", False)
+ local_data = kwargs.get("local_data", "./milvus_data")
+ vector_base = Milvus(
+ host=host,
+ port=port,
+ user=user,
+ password=password,
+ secure=secure,
+ collection_name=collection_name,
+ dimension=dimension,
+ top_k=top_k,
+ index_params=index_params,
+ search_params=search_params,
+ local_mode=local_mode,
+ local_data=local_data,
+ metric_type=metric_type
+ )
+ elif name == "redis":
+ from modelcache.manager.vector_data.redis import RedisVectorStore
+ dimension = kwargs.get("dimension", DIMENSION)
+ check_dimension(dimension)
+
+ redis_config = kwargs.get("config")
+ host = redis_config.get('redis', 'host')
+ port = redis_config.get('redis', 'port')
+ user = redis_config.get('redis', 'user')
+ password = redis_config.get('redis', 'password')
+ namespace = kwargs.get("namespace", "")
+ # collection_name = kwargs.get("collection_name", COLLECTION_NAME)
+
+ vector_base = RedisVectorStore(
+ host=host,
+ port=port,
+ username=user,
+ password=password,
+ namespace=namespace,
+ top_k=top_k,
+ dimension=dimension,
+ )
+ elif name == "faiss":
+ from modelcache.manager.vector_data.faiss import Faiss
+
+ dimension = kwargs.get("dimension", DIMENSION)
+ index_path = kwargs.pop("index_path", FAISS_INDEX_PATH)
+ check_dimension(dimension)
+ vector_base = Faiss(
+ index_file_path=index_path, dimension=dimension, top_k=top_k
+ )
+ elif name == "chromadb":
+ from modelcache.manager.vector_data.chroma import Chromadb
+
+ chromadb_config = kwargs.get("config", None)
+ persist_directory = chromadb_config.get('chromadb','persist_directory')
+
+ vector_base = Chromadb(
+ persist_directory=persist_directory,
+ top_k=top_k,
+ )
+ elif name == "hnswlib":
+ from modelcache.manager.vector_data.hnswlib_store import Hnswlib
+
+ dimension = kwargs.get("dimension", DIMENSION)
+ index_path = kwargs.pop("index_path", "./hnswlib_index.bin")
+ max_elements = kwargs.pop("max_elements", 100000)
+ VectorStorage.check_dimension(dimension)
+ vector_base = Hnswlib(
+ index_file_path=index_path, dimension=dimension,
+ top_k=top_k, max_elements=max_elements
+ )
+ else:
+ raise NotFoundError("vector store", name)
+ return vector_base
+
+
+def check_dimension(dimension):
+ if dimension <= 0:
+ raise ParamError(f"the dimension should be greater than zero, current value: {dimension}.")
+
+
diff --git a/modelcache/manager/vector_data/chroma.py b/modelcache/manager/vector_data/chroma.py
new file mode 100644
index 0000000..b92ef3c
--- /dev/null
+++ b/modelcache/manager/vector_data/chroma.py
@@ -0,0 +1,92 @@
+from typing import List
+
+import numpy as np
+import logging
+from modelcache.manager.vector_data.base import VectorStorage, VectorData
+from modelcache.utils import import_chromadb, import_torch
+
+import_torch()
+import_chromadb()
+
+import chromadb
+
+
+class Chromadb(VectorStorage):
+
+ def __init__(
+ self,
+ persist_directory="./chromadb",
+ top_k: int = 1,
+ ):
+ self.collection_name = "modelcache"
+ self.top_k = top_k
+
+ self._client = chromadb.PersistentClient(path=persist_directory)
+ self._collection = None
+
+ def mul_add(self, datas: List[VectorData], model=None):
+ collection_name_model = self.collection_name + '_' + model
+ self._collection = self._client.get_or_create_collection(name=collection_name_model)
+
+ data_array, id_array = map(list, zip(*((data.data.tolist(), str(data.id)) for data in datas)))
+ self._collection.add(embeddings=data_array, ids=id_array)
+
+ def search(self, data: np.ndarray, top_k: int = -1, model=None):
+ collection_name_model = self.collection_name + '_' + model
+ self._collection = self._client.get_or_create_collection(name=collection_name_model)
+
+ if self._collection.count() == 0:
+ return []
+ if top_k == -1:
+ top_k = self.top_k
+ results = self._collection.query(
+ query_embeddings=[data.tolist()],
+ n_results=top_k,
+ include=["distances"],
+ )
+ return list(zip(results["distances"][0], [int(x) for x in results["ids"][0]]))
+
+ def rebuild(self, ids=None):
+ pass
+
+ def delete(self, ids, model=None):
+ try:
+ collection_name_model = self.collection_name + '_' + model
+ self._collection = self._client.get_or_create_collection(name=collection_name_model)
+ # 查询集合中实际存在的 ID
+ ids_str = [str(x) for x in ids]
+ existing_ids = set(self._collection.get(ids=ids_str).ids)
+
+ # 删除存在的 ID
+ if existing_ids:
+ self._collection.delete(list(existing_ids))
+
+ # 返回实际删除的条目数量
+ return len(existing_ids)
+
+ except Exception as e:
+ logging.error('Error during deletion: {}'.format(e))
+ raise ValueError(str(e))
+
+ def rebuild_col(self, model):
+ collection_name_model = self.collection_name + '_' + model
+
+ # 检查集合是否存在，如果存在则删除
+ collections = self._client.list_collections()
+ if any(col.name == collection_name_model for col in collections):
+ self._client.delete_collection(collection_name_model)
+ else:
+ return 'model collection not found, please check!'
+
+ try:
+ self._client.create_collection(collection_name_model)
+ except Exception as e:
+ logging.info(f'rebuild_collection: {e}')
+ raise ValueError(str(e))
+
+ def flush(self):
+ # chroma无flush方法
+ pass
+
+ def close(self):
+ pass
diff --git a/modelcache/manager/vector_data/faiss.py b/modelcache/manager/vector_data/faiss.py
index 0f8445c..77e7617 100644
--- a/modelcache/manager/vector_data/faiss.py
+++ b/modelcache/manager/vector_data/faiss.py
@@ -2,13 +2,13 @@
import os
from typing import List
import numpy as np
-from modelcache.manager.vector_data.base import VectorBase, VectorData
+from modelcache.manager.vector_data.base import VectorStorage, VectorData
from modelcache.utils import import_faiss
import_faiss()
import faiss # pylint: disable=C0413
-class Faiss(VectorBase):
+class Faiss(VectorStorage):
def __init__(self, index_file_path, dimension, top_k):
self._index_file_path = index_file_path
self._dimension = dimension
diff --git a/modelcache/manager/vector_data/manager.py b/modelcache/manager/vector_data/manager.py
deleted file mode 100644
index 70448b2..0000000
--- a/modelcache/manager/vector_data/manager.py
+++ /dev/null
@@ -1,127 +0,0 @@
-# -*- coding: utf-8 -*-
-from modelcache.utils.error import NotFoundError, ParamError
-
-TOP_K = 1
-FAISS_INDEX_PATH = "faiss.index"
-DIMENSION = 0
-MILVUS_HOST = "localhost"
-MILVUS_PORT = 19530
-MILVUS_USER = ""
-MILVUS_PSW = ""
-MILVUS_SECURE = False
-MILVUS_INDEX_PARAMS = {
- "metric_type": "L2",
- "index_type": "HNSW",
- "params": {"M": 8, "efConstruction": 64},
-}
-
-COLLECTION_NAME = "modelcache"
-
-
-class VectorBase:
- """
- VectorBase to manager the vector base.
- """
-
- def __init__(self):
- raise EnvironmentError(
- "VectorBase is designed to be instantiated, please using the `VectorBase.get(name)`."
- )
-
- @staticmethod
- def check_dimension(dimension):
- if dimension <= 0:
- raise ParamError(
- f"the dimension should be greater than zero, current value: {dimension}."
- )
-
- @staticmethod
- def get(name, **kwargs):
- top_k = kwargs.get("top_k", TOP_K)
- if name == "milvus":
- from modelcache.manager.vector_data.milvus import Milvus
- dimension = kwargs.get("dimension", DIMENSION)
- milvus_config = kwargs.get("milvus_config")
- VectorBase.check_dimension(dimension)
- host = milvus_config.get('milvus', 'host')
- port = milvus_config.get('milvus', 'port')
- user = milvus_config.get('milvus', 'user')
- password = milvus_config.get('milvus', 'password')
-
- secure = kwargs.get("secure", MILVUS_SECURE)
- collection_name = kwargs.get("collection_name", COLLECTION_NAME)
- index_params = kwargs.get("index_params", MILVUS_INDEX_PARAMS)
- search_params = kwargs.get("search_params", None)
- local_mode = kwargs.get("local_mode", False)
- local_data = kwargs.get("local_data", "./milvus_data")
- vector_base = Milvus(
- host=host,
- port=port,
- user=user,
- password=password,
- secure=secure,
- collection_name=collection_name,
- dimension=dimension,
- top_k=top_k,
- index_params=index_params,
- search_params=search_params,
- local_mode=local_mode,
- local_data=local_data
- )
- elif name == "redis":
- from modelcache.manager.vector_data.redis import RedisVectorStore
- dimension = kwargs.get("dimension", DIMENSION)
- VectorBase.check_dimension(dimension)
-
- redis_config = kwargs.get("redis_config")
- host = redis_config.get('redis', 'host')
- port = redis_config.get('redis', 'port')
- user = redis_config.get('redis', 'user')
- password = redis_config.get('redis', 'password')
- namespace = kwargs.get("namespace", "")
- # collection_name = kwargs.get("collection_name", COLLECTION_NAME)
-
- vector_base = RedisVectorStore(
- host=host,
- port=port,
- username=user,
- password=password,
- namespace=namespace,
- top_k=top_k,
- dimension=dimension,
- )
- elif name == "faiss":
- from modelcache.manager.vector_data.faiss import Faiss
-
- dimension = kwargs.get("dimension", DIMENSION)
- index_path = kwargs.pop("index_path", FAISS_INDEX_PATH)
- VectorBase.check_dimension(dimension)
- vector_base = Faiss(
- index_file_path=index_path, dimension=dimension, top_k=top_k
- )
- elif name == "chromadb":
- from modelcache.manager.vector_data.chroma import Chromadb
-
- client_settings = kwargs.get("client_settings", None)
- persist_directory = kwargs.get("persist_directory", None)
- collection_name = kwargs.get("collection_name", COLLECTION_NAME)
- vector_base = Chromadb(
- client_settings=client_settings,
- persist_directory=persist_directory,
- collection_name=collection_name,
- top_k=top_k,
- )
- elif name == "hnswlib":
- from modelcache.manager.vector_data.hnswlib_store import Hnswlib
-
- dimension = kwargs.get("dimension", DIMENSION)
- index_path = kwargs.pop("index_path", "./hnswlib_index.bin")
- max_elements = kwargs.pop("max_elements", 100000)
- VectorBase.check_dimension(dimension)
- vector_base = Hnswlib(
- index_file_path=index_path, dimension=dimension,
- top_k=top_k, max_elements=max_elements
- )
- else:
- raise NotFoundError("vector store", name)
- return vector_base
diff --git a/modelcache/manager/vector_data/milvus.py b/modelcache/manager/vector_data/milvus.py
index 50d6ab1..f5c7682 100644
--- a/modelcache/manager/vector_data/milvus.py
+++ b/modelcache/manager/vector_data/milvus.py
@@ -3,9 +3,11 @@
from typing import List
from uuid import uuid4
import numpy as np
+
+from modelcache.embedding import MetricType
from modelcache.utils import import_pymilvus
from modelcache.utils.log import modelcache_log
-from modelcache.manager.vector_data.base import VectorBase, VectorData
+from modelcache.manager.vector_data.base import VectorStorage, VectorData
import_pymilvus()
@@ -21,19 +23,7 @@
)
-class Milvus(VectorBase):
- SEARCH_PARAM = {
- "IVF_FLAT": {"metric_type": "L2", "params": {"nprobe": 10}},
- "IVF_SQ8": {"metric_type": "L2", "params": {"nprobe": 10}},
- "IVF_PQ": {"metric_type": "L2", "params": {"nprobe": 10}},
- "HNSW": {"metric_type": "L2", "params": {"ef": 10}},
- "RHNSW_FLAT": {"metric_type": "L2", "params": {"ef": 10}},
- "RHNSW_SQ": {"metric_type": "L2", "params": {"ef": 10}},
- "RHNSW_PQ": {"metric_type": "L2", "params": {"ef": 10}},
- "IVF_HNSW": {"metric_type": "L2", "params": {"nprobe": 10, "ef": 10}},
- "ANNOY": {"metric_type": "L2", "params": {"search_k": 10}},
- "AUTOINDEX": {"metric_type": "L2", "params": {}},
- }
+class Milvus(VectorStorage):
def __init__(
self,
@@ -48,7 +38,8 @@ def __init__(
index_params: dict = None,
search_params: dict = None,
local_mode: bool = False,
- local_data: str = "./milvus_data"
+ local_data: str = "./milvus_data",
+ metric_type: MetricType = MetricType.COSINE,
):
if dimension <= 0:
raise ValueError(
@@ -58,14 +49,29 @@ def __init__(
self._local_data = local_data
self.dimension = dimension
self.top_k = top_k
- self.index_params = index_params
if self._local_mode:
self._create_local(port, local_data)
self._connect(host, port, user, password, secure)
self.collection_name = collection_name
- self.search_params = (
- search_params or self.SEARCH_PARAM[self.index_params["index_type"]]
- )
+ self.search_params = {
+ "IVF_FLAT": {"metric_type": metric_type.value, "params": {"nprobe": 10}},
+ "IVF_SQ8": {"metric_type": metric_type.value, "params": {"nprobe": 10}},
+ "IVF_PQ": {"metric_type": metric_type.value, "params": {"nprobe": 10}},
+ "HNSW": {"metric_type": metric_type.value, "params": {"ef": 10}},
+ "RHNSW_FLAT": {"metric_type": metric_type.value, "params": {"ef": 10}},
+ "RHNSW_SQ": {"metric_type": metric_type.value, "params": {"ef": 10}},
+ "RHNSW_PQ": {"metric_type": metric_type.value, "params": {"ef": 10}},
+ "IVF_HNSW": {"metric_type": metric_type.value, "params": {"nprobe": 10, "ef": 10}},
+ "ANNOY": {"metric_type": metric_type.value, "params": {"search_k": 10}},
+ "AUTOINDEX": {"metric_type": metric_type.value, "params": {}},
+ }
+ self.index_params ={
+ "metric_type": metric_type.value,
+ "index_type": "HNSW",
+ "params": {"M": 16, "efConstruction": 64},
+ }
+ self.collections = dict()
+
def _connect(self, host, port, user, password, secure):
try:
@@ -87,12 +93,14 @@ def _connect(self, host, port, user, password, secure):
timeout=10
)
+
def _create_collection(self, collection_name):
if not utility.has_collection(collection_name, using=self.alias):
schema = [
FieldSchema(
name="id",
- dtype=DataType.INT64,
+ dtype=DataType.VARCHAR,
+ max_length=36,
is_primary=True,
auto_id=False,
),
@@ -101,7 +109,8 @@ def _create_collection(self, collection_name):
),
]
schema = CollectionSchema(schema)
- self.col = Collection(
+
+ new_collection = Collection(
collection_name,
schema=schema,
consistency_level="Session",
@@ -109,46 +118,48 @@ def _create_collection(self, collection_name):
)
else:
modelcache_log.warning("The %s collection already exists, and it will be used directly.", collection_name)
- self.col = Collection(
+ new_collection = Collection(
collection_name, consistency_level="Session", using=self.alias
)
- if len(self.col.indexes) == 0:
+ self.collections[collection_name] = new_collection
+
+ if len(new_collection.indexes) == 0:
try:
modelcache_log.info("Attempting creation of Milvus index.")
- self.col.create_index("embedding", index_params=self.index_params)
+ new_collection.create_index("embedding", index_params=self.index_params)
modelcache_log.info("Creation of Milvus index successful.")
except MilvusException as e:
modelcache_log.warning("Error with building index: %s, and attempting creation of default index.", e)
i_p = {"metric_type": "L2", "index_type": "AUTOINDEX", "params": {}}
- self.col.create_index("embedding", index_params=i_p)
+ new_collection.create_index("embedding", index_params=i_p)
self.index_params = i_p
else:
- self.index_params = self.col.indexes[0].to_dict()["index_param"]
+ self.index_params = new_collection.indexes[0].to_dict()["index_param"]
+
+ new_collection.load()
- self.col.load()
def _get_collection(self, collection_name):
- self.col = Collection(
- collection_name, consistency_level="Session", using=self.alias
- )
- self.col.load()
+ if collection_name not in self.collections:
+ self._create_collection(collection_name)
+ return self.collections[collection_name]
def mul_add(self, datas: List[VectorData], model=None):
collection_name_model = self.collection_name + '_' + model
- self._create_collection(collection_name_model)
-
+ col = self._get_collection(collection_name_model)
data_array, id_array = map(list, zip(*((data.data, data.id) for data in datas)))
np_data = np.array(data_array).astype("float32")
entities = [id_array, np_data]
- self.col.insert(entities)
+ col.insert(entities)
+
def search(self, data: np.ndarray, top_k: int = -1, model=None):
if top_k == -1:
top_k = self.top_k
collection_name_model = self.collection_name + '_' + model
- self._create_collection(collection_name_model)
- search_result = self.col.search(
+ col = self._get_collection(collection_name_model)
+ search_result = col.search(
data=data.reshape(1, -1).tolist(),
anns_field="embedding",
param=self.search_params,
@@ -156,12 +167,13 @@ def search(self, data: np.ndarray, top_k: int = -1, model=None):
)
return list(zip(search_result[0].distances, search_result[0].ids))
+
def delete(self, ids, model=None):
collection_name_model = self.collection_name + '_' + model
- self._get_collection(collection_name_model)
+ col = self._get_collection(collection_name_model)
- del_ids = ",".join([str(x) for x in ids])
- resp = self.col.delete(f"id in [{del_ids}]")
+ del_ids = ",".join([f'"{x}"' for x in ids])
+ resp = col.delete(f"id in [{del_ids}]")
delete_count = resp.delete_count
return delete_count
@@ -178,10 +190,12 @@ def rebuild_col(self, model):
logging.info('create_collection: {}'.format(e))
def rebuild(self, ids=None): # pylint: disable=unused-argument
- self.col.compact()
+ for col in self.collections.values():
+ col.compact()
def flush(self):
- self.col.flush(_async=True)
+ for col in self.collections.values():
+ col.flush(_async=True)
def close(self):
self.flush()
diff --git a/modelcache/manager/vector_data/redis.py b/modelcache/manager/vector_data/redis.py
index afa1088..6e51511 100644
--- a/modelcache/manager/vector_data/redis.py
+++ b/modelcache/manager/vector_data/redis.py
@@ -6,7 +6,7 @@
from redis.commands.search.field import TagField, VectorField, NumericField
from redis.client import Redis
-from modelcache.manager.vector_data.base import VectorBase, VectorData
+from modelcache.manager.vector_data.base import VectorStorage, VectorData
from modelcache.utils import import_redis
from modelcache.utils.log import modelcache_log
from modelcache.utils.index_util import get_index_name
@@ -14,7 +14,7 @@
import_redis()
-class RedisVectorStore(VectorBase):
+class RedisVectorStore(VectorStorage):
def __init__(
self,
host: str = "localhost",
diff --git a/modelcache/processor/pre.py b/modelcache/processor/pre.py
index 5875294..6a9d0f6 100644
--- a/modelcache/processor/pre.py
+++ b/modelcache/processor/pre.py
@@ -4,7 +4,7 @@
def insert_last_content(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
- return data.get("chat_info")[-1]["query"]
+ return data.get("query")[-1]["content"]
def query_last_content(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
@@ -64,9 +64,14 @@ def query_multi_splicing(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
def insert_multi_splicing(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
- insert_query_list = data.get("chat_info")[-1]['query']
+ insert_query_list = data['query']
return multi_splicing(insert_query_list)
+def query_with_role(data: Dict[str, Any], **_: Dict[str, Any]) -> Any:
+ query = data["query"][-1]
+ content = query["content"]
+ role = query["role"]
+ return role+": "+content
def multi_splicing(data_list) -> Any:
result_str = ""
diff --git a/modelcache/similarity_evaluation/__init__.py b/modelcache/similarity_evaluation/__init__.py
index 22dff8f..633f866 100644
--- a/modelcache/similarity_evaluation/__init__.py
+++ b/modelcache/similarity_evaluation/__init__.py
@@ -1,11 +1,2 @@
# -*- coding: utf-8 -*-
-from modelcache.similarity_evaluation.similarity_evaluation import SimilarityEvaluation
-from modelcache.utils.lazy_import import LazyImport
-exact_match = LazyImport(
- "exact_match", globals(), "modelcache.similarity_evaluation.exact_match"
-)
-
-
-def ExactMatchEvaluation():
- return exact_match.ExactMatchEvaluation()
diff --git a/modelcache/similarity_evaluation/similarity_evaluation.py b/modelcache/similarity_evaluation/base.py
similarity index 100%
rename from modelcache/similarity_evaluation/similarity_evaluation.py
rename to modelcache/similarity_evaluation/base.py
diff --git a/modelcache/similarity_evaluation/distance.py b/modelcache/similarity_evaluation/distance.py
index 44ca595..281bcd0 100644
--- a/modelcache/similarity_evaluation/distance.py
+++ b/modelcache/similarity_evaluation/distance.py
@@ -1,7 +1,6 @@
# -*- coding: utf-8 -*-
from typing import Tuple, Dict, Any
-from modelcache.similarity_evaluation import SimilarityEvaluation
-
+from modelcache.similarity_evaluation.base import SimilarityEvaluation
class SearchDistanceEvaluation(SimilarityEvaluation):
def __init__(self, max_distance=4.0, positive=False):
diff --git a/modelcache/similarity_evaluation/exact_match.py b/modelcache/similarity_evaluation/exact_match.py
index 553bd59..e207387 100644
--- a/modelcache/similarity_evaluation/exact_match.py
+++ b/modelcache/similarity_evaluation/exact_match.py
@@ -1,6 +1,6 @@
# -*- coding: utf-8 -*-
from typing import Tuple, Dict, Any
-from modelcache.similarity_evaluation.similarity_evaluation import SimilarityEvaluation
+from modelcache.similarity_evaluation.base import SimilarityEvaluation
class ExactMatchEvaluation(SimilarityEvaluation):
diff --git a/modelcache/utils/__init__.py b/modelcache/utils/__init__.py
index 147a56e..15d90ec 100644
--- a/modelcache/utils/__init__.py
+++ b/modelcache/utils/__init__.py
@@ -73,3 +73,7 @@ def import_pillow():
def import_redis():
_check_library("redis")
+
+
+def import_chromadb():
+ _check_library("chromadb", package="chromadb")
\ No newline at end of file
diff --git a/modelcache/utils/cache_func.py b/modelcache/utils/cache_func.py
deleted file mode 100644
index 461f542..0000000
--- a/modelcache/utils/cache_func.py
+++ /dev/null
@@ -1,3 +0,0 @@
-# -*- coding: utf-8 -*-
-def cache_all(*_, **__):
- return True
\ No newline at end of file
diff --git a/modelcache/utils/time.py b/modelcache/utils/time.py
index 7074b58..569af99 100644
--- a/modelcache/utils/time.py
+++ b/modelcache/utils/time.py
@@ -1,15 +1,14 @@
# -*- coding: utf-8 -*-
import time
-from modelcache import cache
-
-def time_cal(func, func_name=None, report_func=None):
+def time_cal(func, func_name=None, report_func=None, **kwargs):
+ cache = kwargs.pop("cache_obj")
def inner(*args, **kwargs):
time_start = time.time()
res = func(*args, **kwargs)
delta_time = time.time() - time_start
- if cache.config.log_time_func:
- cache.config.log_time_func(
+ if cache.log_time_func:
+ cache.log_time_func(
func.__name__ if func_name is None else func_name, delta_time
)
if report_func is not None:
diff --git a/modelcache_mm/config/chromadb_config.ini b/modelcache_mm/config/chromadb_config.ini
new file mode 100644
index 0000000..2a87f3b
--- /dev/null
+++ b/modelcache_mm/config/chromadb_config.ini
@@ -0,0 +1,2 @@
+[chromadb]
+persist_directory=./chromadb
diff --git a/modelcache_mm/config/elasticsearch_config.ini b/modelcache_mm/config/elasticsearch_config.ini
new file mode 100644
index 0000000..ba87a5a
--- /dev/null
+++ b/modelcache_mm/config/elasticsearch_config.ini
@@ -0,0 +1,5 @@
+[elasticsearch]
+host = ''
+port = ''
+user = ''
+password = ''
\ No newline at end of file
diff --git a/modelcache_mm/manager/scalar_data/sql_storage_es.py b/modelcache_mm/manager/scalar_data/sql_storage_es.py
new file mode 100644
index 0000000..562b349
--- /dev/null
+++ b/modelcache_mm/manager/scalar_data/sql_storage_es.py
@@ -0,0 +1,198 @@
+# -*- coding: utf-8 -*-
+import json
+from typing import List
+from elasticsearch import Elasticsearch, helpers
+from modelcache.manager.scalar_data.base import CacheStorage, CacheData
+import time
+from snowflake import SnowflakeGenerator
+
+
+class SQLStorage(CacheStorage):
+ def __init__(
+ self,
+ db_type: str = "elasticsearch",
+ config=None
+ ):
+ self.host = config.get('elasticsearch', 'host')
+ self.port = int(config.get('elasticsearch', 'port'))
+ self.client = Elasticsearch(
+ hosts=[{"host": self.host, "port": self.port}],
+ timeout=30,
+ http_auth=('esuser', 'password')
+ )
+
+ self.log_index = "open_cache_mm_query_log"
+ self.ans_index = "open_cache_mm_answer"
+ self.create()
+ self.instance_id = 1 # 雪花算法使用的机器id 使用同一套数据库的分布式系统需要配置不同id
+ # 生成雪花id
+ self.snowflake_id = SnowflakeGenerator(self.instance_id)
+
+ def create(self):
+ answer_index_body = {
+ "mappings": {
+ "properties": {
+ "gmt_create": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+ "gmt_modified": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+ "question": {"type": "text"},
+ "answer": {"type": "text"},
+ "answer_type": {"type": "integer"},
+ "hit_count": {"type": "integer"},
+ "model": {"type": "keyword"},
+ "image_url": {"type": "text"},
+ "image_id": {"type": "text"},
+ "is_deleted": {"type": "integer"},
+ }
+ }
+ }
+
+ log_index_body = {
+ "mappings": {
+ "properties": {
+ "gmt_create": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+ "gmt_modified": {"type": "date", "format": "strict_date_optional_time||epoch_millis"},
+ "error_code": {"type": "integer"},
+ "error_desc": {"type": "text"},
+ "cache_hit": {"type": "keyword"},
+ "delta_time": {"type": "float"},
+ "model": {"type": "keyword"},
+ "query": {"type": "text"},
+ "hit_query": {"type": "text"},
+ "answer": {"type": "text"}
+ }
+ }
+ }
+
+ if not self.client.indices.exists(index=self.ans_index):
+ self.client.indices.create(index=self.ans_index, body=answer_index_body)
+
+ if not self.client.indices.exists(index=self.log_index):
+ self.client.indices.create(index=self.log_index, body=log_index_body)
+
+ def _insert(self, data: List) -> str or None:
+ doc = {
+ "answer": data[0],
+ "question": data[1],
+ "image_url": data[2],
+ "image_id": data[3],
+ "model": data[4],
+ "answer_type": 0,
+ "hit_count": 0,
+ "is_deleted": 0
+ }
+
+ try:
+
+ response = self.client.index(
+ index=self.ans_index,
+ id=next(self.snowflake_id),
+ body=doc,
+ )
+ return int(response['_id'])
+ except Exception as e:
+
+ print(f"Failed to insert document: {e}")
+ return None
+
+ def batch_insert(self, all_data: List[List]) -> List[str]:
+ successful_ids = []
+ for data in all_data:
+ _id = self._insert(data)
+ if _id is not None:
+ successful_ids.append(_id)
+ self.client.indices.refresh(index=self.ans_index) # 批量插入后手动刷新
+
+ return successful_ids
+
+ def insert_query_resp(self, query_resp, **kwargs):
+ doc = {
+ "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": json.dumps(query_resp.get('hit_query'), ensure_ascii=False) if isinstance(
+ query_resp.get('hit_query'), list) else query_resp.get('hit_query'),
+ "answer": query_resp.get('answer'),
+ "hit_count": 0,
+ "is_deleted": 0
+
+ }
+ self.client.index(index=self.log_index, body=doc)
+
+ def get_data_by_id(self, key: int):
+ try:
+ response = self.client.get(index=self.ans_index, id=key, _source=['question', 'image_url','image_id', 'answer', 'model'])
+ source = response["_source"]
+ result = [
+ source.get('question'),
+ source.get('image_url'),
+ source.get('image_id'),
+ source.get('answer'),
+ source.get('model')
+ ]
+ return result
+ except Exception as e:
+ print(e)
+
+ def update_hit_count_by_id(self, primary_id: int):
+ self.client.update(
+ index=self.ans_index,
+ id=primary_id,
+ body={"script": {"source": "ctx._source.hit_count += 1"}}
+ )
+
+ def get_ids(self, deleted=True):
+ query = {
+ "query": {
+ "term": {"is_deleted": 1 if deleted else 0}
+ }
+ }
+ response = self.client.search(index=self.ans_index, body=query)
+ return [hit["_id"] for hit in response["hits"]["hits"]]
+
+ def mark_deleted(self, keys):
+ actions = [
+ {
+ "_op_type": "update",
+ "_index": self.ans_index,
+ "_id": key,
+ "doc": {"is_deleted": 1}
+ }
+ for key in keys
+ ]
+ responses = helpers.bulk(self.client, actions)
+ return responses[0] # 返回更新的文档数
+
+ def model_deleted(self, model_name):
+ query = {
+ "query": {
+ "term": {"model": model_name}
+ }
+ }
+
+ response = self.client.delete_by_query(index=self.ans_index, body=query)
+ return response["deleted"]
+
+ def clear_deleted_data(self):
+ query = {
+ "query": {
+ "term": {"is_deleted": 1}
+ }
+ }
+ response = self.client.delete_by_query(index=self.ans_index, body=query)
+ return response["deleted"]
+
+ def count(self, state: int = 0, is_all: bool = False):
+ query = {"query": {"match_all": {}}} if is_all else {"query": {"term": {"is_deleted": state}}}
+ response = self.client.count(index=self.ans_index, body=query)
+ return response["count"]
+
+ def close(self):
+ self.client.close()
+
+ def count_answers(self):
+ query = {"query": {"match_all": {}}}
+ response = self.client.count(index=self.ans_index, body=query)
+ return response["count"]
diff --git a/modelcache_mm/manager/vector_data/chroma.py b/modelcache_mm/manager/vector_data/chroma.py
new file mode 100644
index 0000000..7f5b78c
--- /dev/null
+++ b/modelcache_mm/manager/vector_data/chroma.py
@@ -0,0 +1,99 @@
+from typing import List
+
+import numpy as np
+import logging
+from modelcache_mm.manager.vector_data.base import VectorBase, VectorData
+from modelcache_mm.utils import import_chromadb, import_torch
+from modelcache_mm.utils.index_util import get_mm_index_name
+
+import_torch()
+import_chromadb()
+
+import chromadb
+
+
+class Chromadb(VectorBase):
+
+ def __init__(
+ self,
+ persist_directory="./chromadb",
+ top_k: int = 1,
+ ):
+ self.top_k = top_k
+
+ self._client = chromadb.PersistentClient(path=persist_directory)
+ self._collection = None
+
+ def create(self, model=None, mm_type=None):
+ try:
+ collection_name_model = get_mm_index_name(model, mm_type)
+ # collection_name_model = self.collection_name + '_' + model
+ self._client.get_or_create_collection(name=collection_name_model)
+ except Exception as e:
+ raise ValueError(str(e))
+
+ def add(self, datas: List[VectorData], model=None, mm_type=None):
+ collection_name_model = get_mm_index_name(model, mm_type)
+ self._collection = self._client.get_or_create_collection(name=collection_name_model)
+
+ data_array, id_array = map(list, zip(*((data.data.tolist(), str(data.id)) for data in datas)))
+ self._collection.add(embeddings=data_array, ids=id_array)
+
+ def search(self, data: np.ndarray, top_k: int = -1, model=None, mm_type='mm'):
+ collection_name_model = get_mm_index_name(model, mm_type)
+ self._collection = self._client.get_or_create_collection(name=collection_name_model)
+
+ if self._collection.count() == 0:
+ return []
+ if top_k == -1:
+ top_k = self.top_k
+ results = self._collection.query(
+ query_embeddings=[data.tolist()],
+ n_results=top_k,
+ include=["distances"],
+ )
+ return list(zip(results["distances"][0], [int(x) for x in results["ids"][0]]))
+
+ def delete(self, ids, model=None, mm_type=None):
+ try:
+ collection_name_model = get_mm_index_name(model, mm_type)
+ self._collection = self._client.get_or_create_collection(name=collection_name_model)
+ # 查询集合中实际存在的 ID
+ ids_str = [str(x) for x in ids]
+ existing_ids = set(self._collection.get(ids=ids_str).ids)
+
+ # 删除存在的 ID
+ if existing_ids:
+ self._collection.delete(list(existing_ids))
+
+ # 返回实际删除的条目数量
+ return len(existing_ids)
+
+ except Exception as e:
+ logging.error('Error during deletion: {}'.format(e))
+ raise ValueError(str(e))
+
+ def rebuild_idx(self, model, mm_type=None):
+ collection_name_model = get_mm_index_name(model, mm_type)
+
+ # 检查集合是否存在，如果存在则删除
+ collections = self._client.list_collections()
+ if any(col.name == collection_name_model for col in collections):
+ self._client.delete_collection(collection_name_model)
+ else:
+ return 'model collection not found, please check!'
+
+ try:
+ self._client.create_collection(collection_name_model)
+ except Exception as e:
+ logging.info(f'rebuild_collection: {e}')
+ raise ValueError(str(e))
+
+ def rebuild(self, ids=None):
+ pass
+
+ def flush(self):
+ pass
+
+ def close(self):
+ pass
diff --git a/modelcache_mm/manager/vector_data/manager.py b/modelcache_mm/manager/vector_data/manager.py
index c10decc..64cca45 100644
--- a/modelcache_mm/manager/vector_data/manager.py
+++ b/modelcache_mm/manager/vector_data/manager.py
@@ -108,6 +108,15 @@ def get(name, **kwargs):
dimension=dimension,
top_k=top_k
)
+ elif name == "chromadb":
+ from modelcache_mm.manager.vector_data.chroma import Chromadb
+
+ chromadb_config = kwargs.get("chromadb_config", None)
+ persist_directory = chromadb_config.get('chromadb', 'persist_directory')
+ vector_base = Chromadb(
+ persist_directory=persist_directory,
+ top_k=top_k,
+ )
else:
raise NotFoundError("vector store", name)
return vector_base
diff --git a/modelcache_mm/utils/__init__.py b/modelcache_mm/utils/__init__.py
index 147a56e..15d90ec 100644
--- a/modelcache_mm/utils/__init__.py
+++ b/modelcache_mm/utils/__init__.py
@@ -73,3 +73,7 @@ def import_pillow():
def import_redis():
_check_library("redis")
+
+
+def import_chromadb():
+ _check_library("chromadb", package="chromadb")
\ No newline at end of file
diff --git a/mulicache-readme-cn.md b/mulicache-readme-cn.md
new file mode 100644
index 0000000..9d21d83
--- /dev/null
+++ b/mulicache-readme-cn.md
@@ -0,0 +1,127 @@
+# MultiModal Cache
+
+为满足多模态的性能要求，我们在 LLModel Cache 的基础上，开发了 MultiModal Cache 系统。MultiModal Cache 增强了 ModelCache 功能，架优化架构，适应多种应用场景。
+
+- [MultiModal Cache](#multimodal-cache)
+ - [最新动态](#最新动态)
+ - [特性](#特性)
+ - [性能](#性能)
+ - [效果评估](#效果评估)
+ - [参与贡献](#参与贡献)
+
+## 最新动态
+
+- [2024.12.12] MultiModal Cache 系统正式发布。
+
+## 特性
+
+| 场景 | 数据类型 | 图像格式 | 数据隔离 |
+|------|----------|----------|----------|
+| 文本对话 | 文本 | 不适用 | 支持 |
+| 图文理解 | 文本+图像 | image_url/image_base64 | 支持 |
+
+- **兼容性**：支持文本和图片链接（image_url）和图片 Base64 编码三种数据格式及其组合。
+- **数据隔离**：支持多模型数据隔离，允许不同数据模型在同一系统中独立运行。
+- **模态隔离**：支持同一模型下不同模态数据（如文本和图像）的隔离处理。
+
+## 性能
+
+我们在生产环境中使用企业级数据库对 MultiModal Cache 进行了全面的性能评估。以下是详细的性能数据：
+
+