Refactoring of all vector stores in place

trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py

@@ -38,9 +38,11 @@ class Processor(Consumer):
 
         v = msg.value()
 
-        if v.entity.value != "":
+        for entity in v.entities:
-            for vec in v.vectors:
+
-                self.vecstore.insert(vec, v.entity.value)
+            if entity.value != "" and entity.value is not None:
+                for vec in entity.vectors:
+                    self.vecstore.insert(vec, entity.value)
 
     @staticmethod
     def add_args(parser):

trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py

@@ -60,76 +60,82 @@ class Processor(Consumer):
 
         self.last_index_name = None
 
+    def create_index(self, index_name, dim):
+
+        self.pinecone.create_index(
+            name = index_name,
+            dimension = dim,
+            metric = "cosine",
+            spec = ServerlessSpec(
+                cloud = self.cloud,
+                region = self.region,
+            )
+        )
+
+        for i in range(0, 1000):
+
+            if self.pinecone.describe_index(
+                    index_name
+            ).status["ready"]:
+                break
+
+            time.sleep(1)
+
+        if not self.pinecone.describe_index(
+                index_name
+        ).status["ready"]:
+            raise RuntimeError(
+                "Gave up waiting for index creation"
+            )
+
     def handle(self, msg):
 
         v = msg.value()
 
         id = str(uuid.uuid4())
 
-        if v.entity.value == "" or v.entity.value is None: return
+        for entity in v.entities:
 
-        for vec in v.vectors:
+            if v.entity.value == "" or v.entity.value is None: continue
 
-            dim = len(vec)
+            for vec in entity.vectors:
 
-            index_name = (
+                dim = len(vec)
-                "t-" + v.metadata.user + "-" + str(dim)
-            )
 
-            if index_name != self.last_index_name:
+                index_name = (
+                    "t-" + v.metadata.user + "-" + str(dim)
+                )
 
-                if not self.pinecone.has_index(index_name):
+                if index_name != self.last_index_name:
 
-                    try:
+                    if not self.pinecone.has_index(index_name):
 
-                        self.pinecone.create_index(
+                        try:
-                            name = index_name,
-                            dimension = dim,
-                            metric = "cosine",
-                            spec = ServerlessSpec(
-                                cloud = self.cloud,
-                                region = self.region,
-                            )
-                        )
 
-                        for i in range(0, 1000):
+                            self.create_index(index_name, dim)
 
-                            if self.pinecone.describe_index(
+                        except Exception as e:
-                                    index_name
+                            print("Pinecone index creation failed")
-                            ).status["ready"]:
+                            raise e
-                                break
 
-                            time.sleep(1)
+                        print(f"Index {index_name} created", flush=True)
 
-                        if not self.pinecone.describe_index(
+                    self.last_index_name = index_name
-                                index_name
-                        ).status["ready"]:
-                            raise RuntimeError(
-                                "Gave up waiting for index creation"
-                            )
 
-                    except Exception as e:
+                index = self.pinecone.Index(index_name)
-                        print("Pinecone index creation failed")
-                        raise e
 
-                    print(f"Index {index_name} created", flush=True)
+                records = [
+                    {
+                        "id": id,
+                        "values": vec,
+                        "metadata": { "entity": entity.value },
+                    }
+                ]
 
-                self.last_index_name = index_name
+                index.upsert(
-
+                    vectors = records,
-            index = self.pinecone.Index(index_name)
+                    namespace = v.metadata.collection,
-
+                )
-            records = [
-                {
-                    "id": id,
-                    "values": vec,
-                    "metadata": { "entity": v.entity.value },
-                }
-            ]
-
-            index.upsert(
-                vectors = records,
-                namespace = v.metadata.collection,
-            )
 
     @staticmethod
     def add_args(parser):

trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py

@@ -40,6 +40,31 @@ class Processor(Consumer):
 
         self.client = QdrantClient(url=store_uri)
 
+    def get_collection(self, dim, user, collection):
+
+        cname = (
+            "t_" + user + "_" + collection + "_" + str(dim)
+        )
+
+        if cname != self.last_collection:
+
+            if not self.client.collection_exists(cname):
+
+                try:
+                    self.client.create_collection(
+                        collection_name=cname,
+                        vectors_config=VectorParams(
+                            size=dim, distance=Distance.COSINE
+                        ),
+                    )
+                except Exception as e:
+                    print("Qdrant collection creation failed")
+                    raise e
+
+            self.last_collection = cname
+
+        return cname
+
     def handle(self, msg):
 
         v = msg.value()
@@ -51,28 +76,11 @@ class Processor(Consumer):
             for vec in entity.vectors:
 
                 dim = len(vec)
-                collection = (
+
-                    "t_" + v.metadata.user + "_" + v.metadata.collection +
+                collection = self.get_collection(
-                    "_" + str(dim)
+                    dim, v.metadata.user, v.metadata.collection
                 )
 
-                if collection != self.last_collection:
-
-                    if not self.client.collection_exists(collection):
-
-                        try:
-                            self.client.create_collection(
-                                collection_name=collection,
-                                vectors_config=VectorParams(
-                                    size=dim, distance=Distance.COSINE
-                                ),
-                            )
-                        except Exception as e:
-                            print("Qdrant collection creation failed")
-                            raise e
-
-                    self.last_collection = collection
-
                 self.client.upsert(
                     collection_name=collection,
                     points=[