Fix the pylint error and add the chromedb test

Signed-off-by: SimFG <bang.fu@zilliz.com>

Author: SimFG

gptcache/adapter/adapter.py

@@ -1,4 +1,5 @@
 import numpy as np
+
 from gptcache import cache
 from gptcache.processor.post import temperature_softmax
 from gptcache.utils.error import NotInitError
@@ -16,7 +17,6 @@ def adapt(llm_handler, cache_data_convert, update_cache_callback, *args, **kwarg
     :param kwargs: llm kwargs
     :return: llm result
     """
-    health_check_flag = kwargs.pop("health_check", False)
     search_only_flag = kwargs.pop("search_only", False)
     user_temperature = "temperature" in kwargs
     user_top_k = "top_k" in kwargs
@@ -114,7 +114,7 @@ def adapt(llm_handler, cache_data_convert, update_cache_callback, *args, **kwarg
                 continue
 
             # cache consistency check
-            if health_check_flag:
+            if chat_cache.config.data_check:
                 is_healthy = cache_health_check(
                     chat_cache.data_manager.v,
                     {
@@ -202,7 +202,7 @@ def post_process():
         kwargs["cache_context"] = context
         kwargs["cache_skip"] = cache_skip
         kwargs["cache_factor"] = cache_factor
-        kwargs["search_only_flag"] = search_only_flag
+        kwargs["search_only"] = search_only_flag
         llm_data = adapt(
             llm_handler, cache_data_convert, update_cache_callback, *args, **kwargs
         )
@@ -467,8 +467,8 @@ def update_cache_func(handled_llm_data, question=None):
             llm_data = update_cache_callback(
                 llm_data, update_cache_func, *args, **kwargs
             )
-        except Exception as e:  # pylint: disable=W0703
-            gptcache_log.warning("failed to save the data to cache, error: %s", e)
+        except Exception:  # pylint: disable=W0703
+            gptcache_log.error("failed to save the data to cache", exc_info=True)
     return llm_data
 
 

gptcache/config.py

@@ -44,7 +44,8 @@ def __init__(
             enable_token_counter: bool = True,
             input_summary_len: Optional[int] = None,
             context_len: Optional[int] = None,
-            skip_list: List[str] = None
+            skip_list: List[str] = None,
+            data_check: bool = False,
     ):
         if similarity_threshold < 0 or similarity_threshold > 1:
             raise CacheError(
@@ -61,3 +62,4 @@ def __init__(
         if skip_list is None:
             skip_list = ["system", "assistant"]
         self.skip_list = skip_list
+        self.data_check = data_check

gptcache/manager/vector_data/qdrant.py

@@ -1,36 +1,44 @@
 from typing import List, Optional
+
 import numpy as np
 
+from gptcache.manager.vector_data.base import VectorBase, VectorData
 from gptcache.utils import import_qdrant
 from gptcache.utils.log import gptcache_log
-from gptcache.manager.vector_data.base import VectorBase, VectorData
 
 import_qdrant()
 
-from qdrant_client import QdrantClient  # pylint: disable=C0413
-from qdrant_client.models import PointStruct, HnswConfigDiff, VectorParams, OptimizersConfigDiff, \
-    Distance  # pylint: disable=C0413
+# pylint: disable=C0413
+from qdrant_client import QdrantClient
+from qdrant_client.models import (
+    PointStruct,
+    HnswConfigDiff,
+    VectorParams,
+    OptimizersConfigDiff,
+    Distance,
+)
 
 
 class QdrantVectorStore(VectorBase):
+    """Qdrant Vector Store"""
 
     def __init__(
-            self,
-            url: Optional[str] = None,
-            port: Optional[int] = 6333,
-            grpc_port: int = 6334,
-            prefer_grpc: bool = False,
-            https: Optional[bool] = None,
-            api_key: Optional[str] = None,
-            prefix: Optional[str] = None,
-            timeout: Optional[float] = None,
-            host: Optional[str] = None,
-            collection_name: Optional[str] = "gptcache",
-            location: Optional[str] = "./qdrant",
-            dimension: int = 0,
-            top_k: int = 1,
-            flush_interval_sec: int = 5,
-            index_params: Optional[dict] = None,
+        self,
+        url: Optional[str] = None,
+        port: Optional[int] = 6333,
+        grpc_port: int = 6334,
+        prefer_grpc: bool = False,
+        https: Optional[bool] = None,
+        api_key: Optional[str] = None,
+        prefix: Optional[str] = None,
+        timeout: Optional[float] = None,
+        host: Optional[str] = None,
+        collection_name: Optional[str] = "gptcache",
+        location: Optional[str] = "./qdrant",
+        dimension: int = 0,
+        top_k: int = 1,
+        flush_interval_sec: int = 5,
+        index_params: Optional[dict] = None,
     ):
         if dimension <= 0:
             raise ValueError(
@@ -44,13 +52,17 @@ def __init__(
         if self._in_memory or location is not None:
             self._create_local(location)
         else:
-            self._create_remote(url, port, api_key, timeout, host, grpc_port, prefer_grpc, prefix, https)
+            self._create_remote(
+                url, port, api_key, timeout, host, grpc_port, prefer_grpc, prefix, https
+            )
         self._create_collection(collection_name, flush_interval_sec, index_params)
 
     def _create_local(self, location):
         self._client = QdrantClient(location=location)
 
-    def _create_remote(self, url, port, api_key, timeout, host, grpc_port, prefer_grpc, prefix, https):
+    def _create_remote(
+        self, url, port, api_key, timeout, host, grpc_port, prefer_grpc, prefix, https
+    ):
         self._client = QdrantClient(
             url=url,
             port=port,
@@ -63,45 +75,70 @@ def _create_remote(self, url, port, api_key, timeout, host, grpc_port, prefer_gr
             https=https,
         )
 
-    def _create_collection(self, collection_name: str, flush_interval_sec: int, index_params: Optional[dict] = None):
+    def _create_collection(
+        self,
+        collection_name: str,
+        flush_interval_sec: int,
+        index_params: Optional[dict] = None,
+    ):
         hnsw_config = HnswConfigDiff(**(index_params or {}))
-        vectors_config = VectorParams(size=self.dimension, distance=Distance.COSINE,
-                                      hnsw_config=hnsw_config)
-        optimizers_config = OptimizersConfigDiff(deleted_threshold=0.2, vacuum_min_vector_number=1000,
-                                                 flush_interval_sec=flush_interval_sec)
+        vectors_config = VectorParams(
+            size=self.dimension, distance=Distance.COSINE, hnsw_config=hnsw_config
+        )
+        optimizers_config = OptimizersConfigDiff(
+            deleted_threshold=0.2,
+            vacuum_min_vector_number=1000,
+            flush_interval_sec=flush_interval_sec,
+        )
         # check if the collection exists
         existing_collections = self._client.get_collections()
         for existing_collection in existing_collections.collections:
             if existing_collection.name == collection_name:
-                gptcache_log.warning("The %s collection already exists, and it will be used directly.", collection_name)
+                gptcache_log.warning(
+                    "The %s collection already exists, and it will be used directly.",
+                    collection_name,
+                )
                 break
         else:
-            self._client.create_collection(collection_name=collection_name, vectors_config=vectors_config,
-                                           optimizers_config=optimizers_config)
+            self._client.create_collection(
+                collection_name=collection_name,
+                vectors_config=vectors_config,
+                optimizers_config=optimizers_config,
+            )
 
     def mul_add(self, datas: List[VectorData]):
-        points = [PointStruct(id=d.id, vector=d.data.reshape(-1).tolist()) for d in datas]
-        self._client.upsert(collection_name=self._collection_name, points=points, wait=False)
+        points = [
+            PointStruct(id=d.id, vector=d.data.reshape(-1).tolist()) for d in datas
+        ]
+        self._client.upsert(
+            collection_name=self._collection_name, points=points, wait=False
+        )
 
     def search(self, data: np.ndarray, top_k: int = -1):
         if top_k == -1:
             top_k = self.top_k
         reshaped_data = data.reshape(-1).tolist()
-        search_result = self._client.search(collection_name=self._collection_name, query_vector=reshaped_data,
-                                            limit=top_k)
+        search_result = self._client.search(
+            collection_name=self._collection_name,
+            query_vector=reshaped_data,
+            limit=top_k,
+        )
         return list(map(lambda x: (x.score, x.id), search_result))
 
     def delete(self, ids: List[str]):
         self._client.delete(collection_name=self._collection_name, points_selector=ids)
 
     def rebuild(self, ids=None):  # pylint: disable=unused-argument
-        optimizers_config = OptimizersConfigDiff(deleted_threshold=0.2, vacuum_min_vector_number=1000)
-        self._client.update_collection(collection_name=self._collection_name, optimizer_config=optimizers_config)
+        optimizers_config = OptimizersConfigDiff(
+            deleted_threshold=0.2, vacuum_min_vector_number=1000
+        )
+        self._client.update_collection(
+            collection_name=self._collection_name, optimizer_config=optimizers_config
+        )
 
     def flush(self):
         # no need to flush manually as qdrant flushes automatically based on the optimizers_config for remote Qdrant
         pass
 
-
     def close(self):
         self.flush()