diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py index 98fe7915..fa933ef3 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py @@ -38,9 +38,11 @@ class Processor(Consumer): v = msg.value() - if v.entity.value != "": - for vec in v.vectors: - self.vecstore.insert(vec, v.entity.value) + for entity in v.entities: + + 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): diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py index b918c10b..868e866a 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py +++ b/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 = ( - "t-" + v.metadata.user + "-" + str(dim) - ) + dim = len(vec) - 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( - name = index_name, - dimension = dim, - metric = "cosine", - spec = ServerlessSpec( - cloud = self.cloud, - region = self.region, - ) - ) + try: - for i in range(0, 1000): + self.create_index(index_name, dim) - if self.pinecone.describe_index( - index_name - ).status["ready"]: - break + except Exception as e: + print("Pinecone index creation failed") + raise e - time.sleep(1) + print(f"Index {index_name} created", flush=True) - if not self.pinecone.describe_index( - index_name - ).status["ready"]: - raise RuntimeError( - "Gave up waiting for index creation" - ) + self.last_index_name = index_name - except Exception as e: - print("Pinecone index creation failed") - raise e + index = self.pinecone.Index(index_name) - print(f"Index {index_name} created", flush=True) + records = [ + { + "id": id, + "values": vec, + "metadata": { "entity": entity.value }, + } + ] - self.last_index_name = index_name - - index = self.pinecone.Index(index_name) - - records = [ - { - "id": id, - "values": vec, - "metadata": { "entity": v.entity.value }, - } - ] - - index.upsert( - vectors = records, - namespace = v.metadata.collection, - ) + index.upsert( + vectors = records, + namespace = v.metadata.collection, + ) @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py index bbea942f..c7c6b444 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py +++ b/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 + - "_" + str(dim) + + collection = self.get_collection( + 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=[