From d22754f3ae926d72fdec8ab5b796d32eb496c6c6 Mon Sep 17 00:00:00 2001 From: Tyler O <4535788+toliver38@users.noreply.github.com> Date: Fri, 7 Feb 2025 18:24:28 +0000 Subject: [PATCH] Add support for Qdrant API Auth Added the necessary changes to support API Key in Qdrant Client Query+Storage - Doc Embeddings - Graph Embeddings --- .../query/doc_embeddings/qdrant/service.py | 13 +++++++++++-- .../query/graph_embeddings/qdrant/service.py | 12 ++++++++++-- .../storage/doc_embeddings/qdrant/write.py | 10 +++++++++- .../storage/graph_embeddings/qdrant/write.py | 10 +++++++++- 4 files changed, 39 insertions(+), 6 deletions(-) diff --git a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py index 7bb5133a..dd53862f 100755 --- a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py +++ b/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(): diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py index c2dcaa4c..4e9492c6 100755 --- a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py +++ b/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(): diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py index 813c4f29..66a0a5da 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py +++ b/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": ChunkEmbeddings, "store_uri": store_uri, + "api_key": api_key, } ) @@ -96,7 +98,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(): diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py index 47b53979..2292097d 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py +++ b/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 handle(self, msg): @@ -96,6 +98,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():