Fix Cassandra index, cache loaded vector index from Milvus

This commit is contained in:
Cyber MacGeddon 2024-07-16 23:20:30 +01:00
parent bdb42dac84
commit d1066d8e12
3 changed files with 20 additions and 10 deletions

View file

@ -74,6 +74,8 @@ class GraphRag:
limit=self.entity_limit limit=self.entity_limit
) )
print("Obtained", len(res), "entities")
entities = set([ entities = set([
item["entity"]["entity"] item["entity"]["entity"]
for item in res for item in res

View file

@ -1,5 +1,6 @@
from pymilvus import MilvusClient, CollectionSchema, FieldSchema, DataType from pymilvus import MilvusClient, CollectionSchema, FieldSchema, DataType
import time
class TripleVectors: class TripleVectors:
@ -13,11 +14,12 @@ class TripleVectors:
# one are created. # one are created.
self.collections = {} self.collections = {}
# self.collection = "edges" # Time between reloads
# self.dimension = 384 self.reload_time = 90
# if not self.client.has_collection(collection_name=self.collection): # Next time to reload - this forces a reload at next window
# self.init_collection() self.next_reload = time.time() + self.reload_time
print("Reload at", self.next_reload)
def init_collection(self, dimension): def init_collection(self, dimension):
@ -106,10 +108,13 @@ class TripleVectors:
} }
} }
print("Loading...")
self.client.load_collection( self.client.load_collection(
collection_name=coll, collection_name=coll,
) )
print("Searching...")
res = self.client.search( res = self.client.search(
collection_name=coll, collection_name=coll,
data=[embeds], data=[embeds],
@ -118,11 +123,14 @@ class TripleVectors:
search_params=search_params, search_params=search_params,
)[0] )[0]
# FIXME: a lot of loading/unloading going on. How about using a
# time window? # If reload time has passed, unload collection
self.client.release_collection( if time.time() > self.next_reload:
collection_name=coll, print("Unloading, reload at", self.next_reload)
) self.client.release_collection(
collection_name=coll,
)
self.next_reload = time.time() + self.reload_time
return res return res

View file

@ -39,7 +39,7 @@ class TrustGraph:
s text, s text,
p text, p text,
o text, o text,
PRIMARY KEY (s, p) PRIMARY KEY (s, p, o)
); );
"""); """);