feat: support passing metadata (#53) (#56)

* feat: support passing metadata

* pass credentials via metadata call credentials

* Update common proto submodule

Co-authored-by: Viraj Karandikar <[email protected]>

Author: rmittal-github

common

@@ -57,4 +57,5 @@ def add_connection_argparse_parameters(parser: argparse.ArgumentParser) -> argpa
     parser.add_argument(
         "--use-ssl", action='store_true', help="Boolean to control if SSL/TLS encryption should be used."
     )
+    parser.add_argument("--metadata", action='append', nargs='+', help="Send HTTP Header(s) to server")
     return parser

riva/client/argparse_utils.py

@@ -4,20 +4,26 @@
 import os
 from pathlib import Path
 from typing import List, Optional, Tuple, Union
-
 import grpc
 
 
 def create_channel(
-    ssl_cert: Optional[Union[str, os.PathLike]] = None, use_ssl: bool = False, uri: str = "localhost:50051",
+    ssl_cert: Optional[Union[str, os.PathLike]] = None, use_ssl: bool = False, uri: str = "localhost:50051", metadata: Optional[List[Tuple[str, str]]] = None,
 ) -> grpc.Channel:
+
+    def metadata_callback(context, callback):
+        callback(metadata, None)
+
     if ssl_cert is not None or use_ssl:
         root_certificates = None
         if ssl_cert is not None:
             ssl_cert = Path(ssl_cert).expanduser()
             with open(ssl_cert, 'rb') as f:
                 root_certificates = f.read()
         creds = grpc.ssl_channel_credentials(root_certificates)
+        if metadata:
+            auth_creds = grpc.metadata_call_credentials(metadata_callback)
+            creds = grpc.composite_channel_credentials(creds, auth_creds)
         channel = grpc.secure_channel(uri, creds)
     else:
         channel = grpc.insecure_channel(uri)
@@ -30,6 +36,7 @@ def __init__(
         ssl_cert: Optional[Union[str, os.PathLike]] = None,
         use_ssl: bool = False,
         uri: str = "localhost:50051",
+        metadata_args: List[List[str]] = None,
     ) -> None:
         """
         A class responsible for establishing connection with a server and providing security metadata.
@@ -44,7 +51,13 @@ def __init__(
         self.ssl_cert: Optional[Path] = None if ssl_cert is None else Path(ssl_cert).expanduser()
         self.uri: str = uri
         self.use_ssl: bool = use_ssl
-        self.channel: grpc.Channel = create_channel(self.ssl_cert, self.use_ssl, self.uri)
+        self.metadata = []
+        if metadata_args:
+            for meta in metadata_args:
+                if len(meta) != 2:
+                    raise ValueError(f"Metadata should have 2 parameters in \"key\" \"value\" pair. Receieved {len(meta)} parameters.")
+                self.metadata.append(tuple(meta))
+        self.channel: grpc.Channel = create_channel(self.ssl_cert, self.use_ssl, self.uri, self.metadata)
 
     def get_auth_metadata(self) -> List[Tuple[str, str]]:
         """

riva/client/auth.py

@@ -50,7 +50,7 @@ def streaming_transcription_worker(
 ) -> None:
     output_file = Path(output_file).expanduser()
     try:
-        auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+        auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
         asr_service = riva.client.ASRService(auth)
         config = riva.client.StreamingRecognitionConfig(
             config=riva.client.RecognitionConfig(

scripts/asr/riva_streaming_asr_client.py

@@ -66,7 +66,7 @@ def main() -> None:
     if args.list_devices:
         riva.client.audio_io.list_output_devices()
         return
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     asr_service = riva.client.ASRService(auth)
     config = riva.client.StreamingRecognitionConfig(
         config=riva.client.RecognitionConfig(

scripts/asr/transcribe_file.py

@@ -26,7 +26,7 @@ def parse_args() -> argparse.Namespace:
 
 def main() -> None:
     args = parse_args()
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     asr_service = riva.client.ASRService(auth)
     config = riva.client.RecognitionConfig(
         language_code=args.language_code,

scripts/asr/transcribe_file_offline.py

@@ -41,7 +41,7 @@ def main() -> None:
     if args.list_devices:
         riva.client.audio_io.list_input_devices()
         return
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     asr_service = riva.client.ASRService(auth)
     config = riva.client.StreamingRecognitionConfig(
         config=riva.client.RecognitionConfig(

scripts/asr/transcribe_mic.py

@@ -340,7 +340,7 @@ def parse_args() -> argparse.Namespace:
 
 def main() -> None:
     args = parse_args()
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     service = riva.client.NLPService(auth)
     intent_report, slot_report = intent_slots_classification_report(
         args.input_file,

scripts/nlp/eval_intent_slot.py

@@ -44,7 +44,7 @@ def pretty_print_result(
 
 def main() -> None:
     args = parse_args()
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     service = riva.client.NLPService(auth)
     if args.interactive:
         while True:

scripts/nlp/intentslot_client.py

@@ -31,7 +31,7 @@ def parse_args() -> argparse.Namespace:
 
 def main() -> None:
     args = parse_args()
-    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server)
+    auth = riva.client.Auth(args.ssl_cert, args.use_ssl, args.server, args.metadata)
     service = riva.client.NLPService(auth)
     tokens, slots, slot_confidences, starts, ends = riva.client.extract_most_probable_token_classification_predictions(
         service.classify_tokens(input_strings=args.query, model_name=args.model)