diff --git a/trustgraph/graph_rag.py b/trustgraph/graph_rag.py index b106e99e..aa24eef0 100644 --- a/trustgraph/graph_rag.py +++ b/trustgraph/graph_rag.py @@ -74,6 +74,8 @@ class GraphRag: limit=self.entity_limit ) + print("Obtained", len(res), "entities") + entities = set([ item["entity"]["entity"] for item in res diff --git a/trustgraph/triple_vectors.py b/trustgraph/triple_vectors.py index 1d0988cd..f80004b3 100644 --- a/trustgraph/triple_vectors.py +++ b/trustgraph/triple_vectors.py @@ -1,5 +1,6 @@ from pymilvus import MilvusClient, CollectionSchema, FieldSchema, DataType +import time class TripleVectors: @@ -13,11 +14,12 @@ class TripleVectors: # one are created. self.collections = {} -# self.collection = "edges" -# self.dimension = 384 + # Time between reloads + self.reload_time = 90 -# if not self.client.has_collection(collection_name=self.collection): -# self.init_collection() + # Next time to reload - this forces a reload at next window + self.next_reload = time.time() + self.reload_time + print("Reload at", self.next_reload) def init_collection(self, dimension): @@ -106,10 +108,13 @@ class TripleVectors: } } + print("Loading...") self.client.load_collection( collection_name=coll, ) + print("Searching...") + res = self.client.search( collection_name=coll, data=[embeds], @@ -118,11 +123,14 @@ class TripleVectors: search_params=search_params, )[0] - # FIXME: a lot of loading/unloading going on. How about using a - # time window? - self.client.release_collection( - collection_name=coll, - ) + + # If reload time has passed, unload collection + if time.time() > self.next_reload: + print("Unloading, reload at", self.next_reload) + self.client.release_collection( + collection_name=coll, + ) + self.next_reload = time.time() + self.reload_time return res diff --git a/trustgraph/trustgraph.py b/trustgraph/trustgraph.py index a7b53a06..cb313abc 100644 --- a/trustgraph/trustgraph.py +++ b/trustgraph/trustgraph.py @@ -39,7 +39,7 @@ class TrustGraph: s text, p text, o text, - PRIMARY KEY (s, p) + PRIMARY KEY (s, p, o) ); """);