Refactored qdrant write

This commit is contained in:
Cyber MacGeddon 2024-12-30 12:00:00 +00:00
parent 4a3bef9b6c
commit 59d1d6a752

View file

@ -44,45 +44,47 @@ class Processor(Consumer):
v = msg.value() v = msg.value()
if v.entity.value == "" or v.entity.value is None: return for entity in v.entities:
for vec in v.vectors: if entity.value == "" or entity.value is None: return
dim = len(vec) for vec in entity.vectors:
collection = (
"t_" + v.metadata.user + "_" + v.metadata.collection + "_" +
str(dim)
)
if collection != self.last_collection: dim = len(vec)
collection = (
"t_" + v.metadata.user + "_" + v.metadata.collection +
"_" + str(dim)
)
if not self.client.collection_exists(collection): if collection != self.last_collection:
try: if not self.client.collection_exists(collection):
self.client.create_collection(
collection_name=collection, try:
vectors_config=VectorParams( self.client.create_collection(
size=dim, distance=Distance.COSINE 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=[
PointStruct(
id=str(uuid.uuid4()),
vector=vec,
payload={
"entity": entity.value,
}
) )
except Exception as e: ]
print("Qdrant collection creation failed") )
raise e
self.last_collection = collection
self.client.upsert(
collection_name=collection,
points=[
PointStruct(
id=str(uuid.uuid4()),
vector=vec,
payload={
"entity": v.entity.value,
}
)
]
)
@staticmethod @staticmethod
def add_args(parser): def add_args(parser):