Fix Bedrock (#119)

trustgraph-ai · Oct 15, 2024 · 25983d1 · 25983d1
1 parent f36c5e5
commit 25983d1
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/templates/components/bedrock.jsonnet b/templates/components/bedrock.jsonnet
@@ -16,7 +16,7 @@ local chunker = import "chunker-recursive.jsonnet";
 
             local envSecrets = engine.envSecrets("bedrock-credentials")
                 .with_env_var("AWS_ID_KEY", "aws-id-key")
-                .with_env_var("AWS_SECRET_KEY", "aws-secret-key")
+                .with_env_var("AWS_SECRET", "aws-secret")
                 .with_env_var("AWS_REGION", "aws-region");
 
             local container =

diff --git a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py
@@ -22,11 +22,11 @@
 default_output_queue = text_completion_response_queue
 default_subscriber = module
 default_model = 'mistral.mistral-large-2407-v1:0'
-default_region = os.getenv("AWS_REGION", 'us-west-2')
 default_temperature = 0.0
 default_max_output = 2048
-default_aws_id = os.getenv("AWS_ID_KEY")
-default_aws_secret = os.getenv("AWS_SECRET_KEY")
+default_aws_id_key = os.getenv("AWS_ID_KEY", None)
+default_aws_secret = os.getenv("AWS_SECRET", None)
+default_aws_region = os.getenv("AWS_REGION", 'us-west-2')
 
 class Processor(ConsumerProducer):
 
@@ -36,18 +36,21 @@ def __init__(self, **params):
         output_queue = params.get("output_queue", default_output_queue)
         subscriber = params.get("subscriber", default_subscriber)
         model = params.get("model", default_model)
-        aws_id = params.get("aws_id_key", default_aws_id)
+        aws_id_key = params.get("aws_id_key", default_aws_id_key)
         aws_secret = params.get("aws_secret", default_aws_secret)
-        aws_region = params.get("aws_region", default_region)
+        aws_region = params.get("aws_region", default_aws_region)
         temperature = params.get("temperature", default_temperature)
         max_output = params.get("max_output", default_max_output)
 
-        if aws_id is None:
+        if aws_id_key is None:
             raise RuntimeError("AWS ID not specified")
 
         if aws_secret is None:
             raise RuntimeError("AWS secret not specified")
 
+        if aws_region is None:
+            raise RuntimeError("AWS region not specified")
+
         super(Processor, self).__init__(
             **params | {
                 "input_queue": input_queue,
@@ -79,7 +82,7 @@ def __init__(self, **params):
         self.max_output = max_output
 
         self.session = boto3.Session(
-            aws_access_key_id=aws_id,
+            aws_access_key_id=aws_id_key,
             aws_secret_access_key=aws_secret,
             region_name=aws_region
         )
@@ -297,16 +300,19 @@ def add_args(parser):
 
         parser.add_argument(
             '-z', '--aws-id-key',
+            default=default_aws_id_key,
             help=f'AWS ID Key'
         )
 
         parser.add_argument(
             '-k', '--aws-secret',
+            default=default_aws_secret,
             help=f'AWS Secret Key'
         )
 
         parser.add_argument(
             '-r', '--aws-region',
+            default=default_aws_region,
             help=f'AWS Region'
         )