diff --git a/templates/components/bedrock.jsonnet b/templates/components/bedrock.jsonnet index d519363e..11f112e1 100644 --- 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 index f06d143e..e2f31c52 100755 --- 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_input_queue = text_completion_request_queue 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 @@ class Processor(ConsumerProducer): 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 @@ class Processor(ConsumerProducer): 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 @@ class Processor(ConsumerProducer): 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' )