Add user/collection metadata to query

trustgraph-base/trustgraph/clients/graph_rag_client.py

@@ -38,9 +38,12 @@ class GraphRagClient(BaseClient):
             output_schema=GraphRagResponse,
         )
 
-    def request(self, query, timeout=500):
+    def request(
+            self, query, user="trustgraph", collection="default",
+            timeout=500
+    ):
 
         return self.call(
-            query=query, timeout=timeout
+            user=user, collection=collection, query=query, timeout=timeout
         ).response
 

trustgraph-base/trustgraph/schema/retrieval.py

@@ -9,6 +9,8 @@ from . types import Error, Value
 
 class GraphRagQuery(Record):
     query = String()
+    user = String()
+    collection = String()
 
 class GraphRagResponse(Record):
     error = Error()
@@ -27,6 +29,8 @@ graph_rag_response_queue = topic(
 
 class DocumentRagQuery(Record):
     query = String()
+    user = String()
+    collection = String()
 
 class DocumentRagResponse(Record):
     error = Error()