Use user/collection in RAG

trustgraph-base/trustgraph/clients/graph_embeddings_client.py

@@ -38,8 +38,12 @@ class GraphEmbeddingsClient(BaseClient):
             output_schema=GraphEmbeddingsResponse,
         )
 
-    def request(self, vectors, limit=10, timeout=300):
+    def request(
+            self, vectors, user="trustgraph", collection="default",
+            limit=10, timeout=300
+    ):
         return self.call(
+            user=user, collection=collection,
             vectors=vectors, limit=limit, timeout=timeout
         ).entities
 

trustgraph-base/trustgraph/clients/triples_query_client.py

@@ -48,11 +48,18 @@ class TriplesQueryClient(BaseClient):
 
         return Value(value=ent, is_uri=False)
 
-    def request(self, s, p, o, limit=10, timeout=60):
+    def request(
+            self, 
+            s, p, o,
+            user="trustgraph", collection="default",
+            limit=10, timeout=60,
+    ):
         return self.call(
             s=self.create_value(s),
             p=self.create_value(p),
             o=self.create_value(o),
+            user=user,
+            collection=collection,
             limit=limit,
             timeout=timeout,
         ).triples

trustgraph-base/trustgraph/schema/graph.py

@@ -23,6 +23,8 @@ graph_embeddings_store_queue = topic('graph-embeddings-store')
 class GraphEmbeddingsRequest(Record):
     vectors = Array(Array(Double()))
     limit = Integer()
+    user = String()
+    collection = String()
 
 class GraphEmbeddingsResponse(Record):
     error = Error()
@@ -56,6 +58,8 @@ class TriplesQueryRequest(Record):
     p = Value()
     o = Value()
     limit = Integer()
+    user = String()
+    collection = String()
 
 class TriplesQueryResponse(Record):
     error = Error()