Add integration for Marqo vector store

examples/data_manager/vector_store.py

@@ -20,6 +20,7 @@ def run():
         'docarray',
         'redis',
         'weaviate',
+        'marqo',
     ]
     for vector_store in vector_stores:
         cache_base = CacheBase('sqlite')

gptcache/manager/vector_data/manager.py

@@ -289,6 +289,20 @@ def get(name, **kwargs):
                 class_schema=class_schema,
                 top_k=top_k,
             )
+        elif name=="marqo":
+            from gptcache.manager.vector_data.marqo_vectorstore import MarqoVectorStore
+
+            marqo_url = kwargs.get("marqo_url", None)
+            index_name = kwargs.get("index_name", COLLECTION_NAME)
+            dimension = kwargs.get("dimension", DIMENSION)
+            top_k = kwargs.get("top_k", TOP_K)
+
+            vector_base = MarqoVectorStore(
+                marqo_url=marqo_url,
+                index_name=index_name,
+                dimension=dimension,
+                top_k=top_k,
+            )
         else:
             raise NotFoundError("vector store", name)
         return vector_base

gptcache/manager/vector_data/marqo_vectorstore.py

@@ -0,0 +1,108 @@
+"""A vector store integration for Marqo"""
+
+from typing import List, Optional
+
+import numpy as np
+
+from gptcache.manager.vector_data.base import VectorBase, VectorData
+from gptcache.utils import import_marqo
+from gptcache.utils.log import gptcache_log
+
+import_marqo()
+from marqo import Client # pylint: disable=C0413
+
+class MarqoVectorStore(VectorBase):
+    """
+    Marqo Vector Store - https://github.com/marqo-ai/marqo
+    """
+
+    def __init__(self,
+                 marqo_url,
+                 index_name: Optional[str] = "gptcache",
+                 dimension: Optional[int] = 0,
+                 top_k: int = 1,
+                ):
+        self.top_k = top_k
+        self._client = Client(
+            url=marqo_url
+        )
+        self._index_name = index_name
+        self._dimension = dimension
+        self._index_settings_dict = {
+                    'index_defaults': {
+                        'model': 'no_model',
+                        'model_properties': {
+                            'dimensions': self._dimension 
+                        },
+                        'ann_parameters':{
+                            'space_type': 'cosinesimil'
+                        }
+                    }
+                }
+
+        self._client.create_index(
+            index_name=self._index_name,
+            settings_dict=self._index_settings_dict)
+
+
+    def mul_add(self, datas: List[VectorData]):
+
+        vecs = []
+        for d in datas:
+            vecs.append(
+                {
+                    '_id': str(d.id),
+                    'gptcachevec':{
+                        'vector': d.data.tolist()
+                    }
+                }
+            )
+        self._client.index(self._index_name).add_documents(
+            documents=vecs,
+            mappings={
+                'gptcachevec': 
+                    {
+                    'type': 'custom_vector'
+                    }
+            },
+            tensor_fields=['gptcachevec'],
+            auto_refresh=True
+        )
+
+    def search(self, data: np.ndarray, top_k: int = -1):
+        if self._client.index(self._index_name).get_stats()['numberOfDocuments']==0:
+            return []
+        if top_k == -1:
+            top_k = self.top_k
+
+        search_results = self._client.index(self._index_name).search(
+                context={
+                    'tensor':[{'vector': data.tolist(), 'weight' : 1}]
+                },
+                limit=top_k,
+            )['hits']
+        return [(r['_id'], r['_score']) for r in search_results]
+
+    def get_embeddings(self, data_id: int | str) -> np.ndarray | None:
+
+        doc = self._client.index(self._index_name).get_document(
+            document_id=str(data_id),
+            expose_facets=True
+        )
+        embedding = doc['_tensor_facets'][0]['_embedding']
+        return embedding
+
+    def delete(self, ids) -> bool:
+        self._client.index(self._index_name).delete_documents(
+                ids=[str(_id) for _id in ids]
+            )
+
+    def rebuild(self, ids=None) -> bool:
+        return
+
+    def flush(self):
+        pass
+
+    def close(self):
+        return self.flush()
+

gptcache/utils/__init__.py

@@ -42,6 +42,7 @@
     "import_redis",
     "import_qdrant",
     "import_weaviate",
+    "import_marqo",
     ]
 
 import importlib.util
@@ -59,6 +60,8 @@ def _check_library(libname: str, prompt: bool = True, package: Optional[str] = N
         prompt_install(package if package else libname)
     return is_avail
 
+def import_marqo():
+    _check_library("marqo")
 
 def import_pymilvus():
     _check_library("pymilvus")

tests/unit_tests/manager/test_marqo.py

@@ -0,0 +1,17 @@
+import unittest
+
+import numpy as np
+
+from gptcache.manager.vector_data import VectorBase
+from gptcache.manager.vector_data.base import VectorData
+
+
+class TestMarqo(unittest.TestCase):
+    def test_normal(self):
+        marqo_url = "http://0.0.0.0:8882"
+        db = VectorBase("marqo", marqo_url=marqo_url, dimension=10, top_k=3)
+        db.mul_add([VectorData(id=i, data=np.random.sample(10)) for i in range(100)])
+        search_res = db.search(np.random.sample(10))
+        self.assertEqual(len(search_res), 3)
+        db.delete(["1", "10", "50", "70"])
+        self.assertEqual(db._collection.count(), 96)