apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-18 09:41:03 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Refactoring of all vector stores in place

Browse source
This commit is contained in:
Cyber MacGeddon 2024-12-30 12:09:37 +00:00
parent 59d1d6a752
commit caf58bb5f2
3 changed files with 89 additions and 73 deletions
Download patch file Download diff file Expand all files Collapse all files

8
trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py
View file

@ -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):

106
trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py
View file

@ -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):

48
trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py
View file

@ -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=[