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-ai · Feb 8, 2025 · e99c0ac · e99c0ac
1 parent 41ccb6c
commit e99c0ac
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 6 deletions.
diff --git a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py
@@ -30,6 +30,8 @@ def __init__(self, **params):
         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 @@ def __init__(self, **params):
                 "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 @@ def add_args(parser):
         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():

diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py
@@ -30,6 +30,7 @@ def __init__(self, **params):
         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 @@ def __init__(self, **params):
                 "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 @@ def add_args(parser):
         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():

diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py
@@ -26,13 +26,15 @@ def __init__(self, **params):
         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 | {
                 "input_queue": input_queue,
                 "subscriber": subscriber,
                 "input_schema": DocumentEmbeddings,
                 "store_uri": store_uri,
+                "api_key": api_key,
             }
         )
 
@@ -97,7 +99,13 @@ def add_args(parser):
         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():

diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py
@@ -26,19 +26,21 @@ def __init__(self, **params):
         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 | {
                 "input_queue": input_queue,
                 "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 @@ def add_args(parser):
             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():