Add support for Qdrant API Auth (#300)

Added the necessary changes to support API Key in Qdrant Client Query+Storage
- Doc Embeddings
- Graph Embeddings

trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py

@@ -30,6 +30,8 @@ class Processor(ConsumerProducer):
         output_queue = params.get("output_queue", default_output_queue)
         subscriber = params.get("subscriber", default_subscriber)
         store_uri = params.get("store_uri", default_store_uri)
+        #optional api key
+        api_key = params.get("api_key", None)
 
         super(Processor, self).__init__(
             **params | {
@@ -39,10 +41,11 @@ class Processor(ConsumerProducer):
                 "input_schema": DocumentEmbeddingsRequest,
                 "output_schema": DocumentEmbeddingsResponse,
                 "store_uri": store_uri,
+                "api_key": api_key,
             }
         )
 
-        self.client = QdrantClient(url=store_uri)
+        self.client = QdrantClient(url=store_uri, api_key=api_key)
 
     def handle(self, msg):
 
@@ -111,7 +114,13 @@ class Processor(ConsumerProducer):
         parser.add_argument(
             '-t', '--store-uri',
             default=default_store_uri,
-            help=f'Milvus store URI (default: {default_store_uri})'
+            help=f'Qdrant store URI (default: {default_store_uri})'
+        )
+        
+        parser.add_argument(
+            '-k', '--api-key',
+            default=None,
+            help=f'API key for qdrant (default: None)'
         )
 
 def run():

trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py

@@ -30,6 +30,7 @@ class Processor(ConsumerProducer):
         output_queue = params.get("output_queue", default_output_queue)
         subscriber = params.get("subscriber", default_subscriber)
         store_uri = params.get("store_uri", default_store_uri)
+        api_key = params.get("api_key", None)
 
         super(Processor, self).__init__(
             **params | {
@@ -39,10 +40,11 @@ class Processor(ConsumerProducer):
                 "input_schema": GraphEmbeddingsRequest,
                 "output_schema": GraphEmbeddingsResponse,
                 "store_uri": store_uri,
+                "api_key": api_key,
             }
         )
 
-        self.client = QdrantClient(url=store_uri)
+        self.client = QdrantClient(url=store_uri, api_key=api_key)
 
     def create_value(self, ent):
         if ent.startswith("http://") or ent.startswith("https://"):
@@ -137,7 +139,13 @@ class Processor(ConsumerProducer):
         parser.add_argument(
             '-t', '--store-uri',
             default=default_store_uri,
-            help=f'Milvus store URI (default: {default_store_uri})'
+            help=f'Qdrant store URI (default: {default_store_uri})'
+        )
+        
+        parser.add_argument(
+            '-k', '--api-key',
+            default=None,
+            help=f'API key for qdrant (default: None)'
         )
 
 def run():

trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py

@@ -26,6 +26,7 @@ class Processor(Consumer):
         input_queue = params.get("input_queue", default_input_queue)
         subscriber = params.get("subscriber", default_subscriber)
         store_uri = params.get("store_uri", default_store_uri)
+        api_key = params.get("api_key", None)
 
         super(Processor, self).__init__(
             **params | {
@@ -33,6 +34,7 @@ class Processor(Consumer):
                 "subscriber": subscriber,
                 "input_schema": DocumentEmbeddings,
                 "store_uri": store_uri,
+                "api_key": api_key,
             }
         )
 
@@ -97,7 +99,13 @@ class Processor(Consumer):
         parser.add_argument(
             '-t', '--store-uri',
             default=default_store_uri,
-            help=f'Qdrant store URI (default: {default_store_uri})'
+            help=f'Qdrant URI (default: {default_store_uri})'
+        )
+        
+        parser.add_argument(
+            '-k', '--api-key',
+            default=None,
+            help=f'Qdrant API key (default: None)'
         )
 
 def run():

trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py

@@ -26,6 +26,7 @@ class Processor(Consumer):
         input_queue = params.get("input_queue", default_input_queue)
         subscriber = params.get("subscriber", default_subscriber)
         store_uri = params.get("store_uri", default_store_uri)
+        api_key = params.get("api_key", None)
 
         super(Processor, self).__init__(
             **params | {
@@ -33,12 +34,13 @@ class Processor(Consumer):
                 "subscriber": subscriber,
                 "input_schema": GraphEmbeddings,
                 "store_uri": store_uri,
+                "api_key": api_key,
             }
         )
 
         self.last_collection = None
 
-        self.client = QdrantClient(url=store_uri)
+        self.client = QdrantClient(url=store_uri, api_key=api_key)
 
     def get_collection(self, dim, user, collection):
 
@@ -106,6 +108,12 @@ class Processor(Consumer):
             default=default_store_uri,
             help=f'Qdrant store URI (default: {default_store_uri})'
         )
+        
+        parser.add_argument(
+            '-k', '--api-key',
+            default=None,
+            help=f'Qdrant API key'
+        )
 
 def run():