apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-17 01:01:03 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

- More AWS Boto3 settings (profile and session key)

Browse source 
- Align environment variable and profile setting names with AWS
  conventions.
Hopefully this should be able to run from an EC2 instance just with role
setting.
This commit is contained in:
Cyber MacGeddon 2025-01-29 12:48:52 +00:00
parent 2c5289765d
commit 6b779f7817
1 changed files with 31 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

46
trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py
View file

@ -24,9 +24,14 @@ default_subscriber = module
default_model = 'mistral.mistral-large-2407-v1:0'
default_temperature = 0.0
default_max_output = 2048
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')
# Actually, these could all just be None, no need to get environment
# variables, as Boto3 would pick all these up if not passed in as args
default_access_key_id = os.getenv("AWS_ACCESS_KEY_ID", None)
default_secret_access_key = os.getenv("AWS_SECRET_ACCESS_KEY", None)
default_session_token = os.getenv("AWS_SESSION_TOKEN", None)
default_profile = os.getenv("AWS_PROFILE", None)
default_region = os.getenv("AWS_DEFAULT_REGION", None)
class Processor(ConsumerProducer):
@ -35,21 +40,30 @@ class Processor(ConsumerProducer):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
subscriber = params.get("subscriber", default_subscriber)
model = params.get("model", default_model)
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_aws_region)
temperature = params.get("temperature", default_temperature)
max_output = params.get("max_output", default_max_output)
if aws_id_key is None:
raise RuntimeError("AWS ID not specified")
aws_access_key_id = params.get(
"aws_access_key_id", default_access_key_id
)
if aws_secret is None:
raise RuntimeError("AWS secret not specified")
aws_secret_access_key = params.get(
"aws_secret_access_key", default_secret_access_key
)
if aws_region is None:
raise RuntimeError("AWS region not specified")
aws_session_token = params.get(
"aws_session_token", default_session_token
)
aws_region = params.get(
"aws_region", default_region
)
aws_profile = params.get(
"aws_profile", default_profile
)
super(Processor, self).__init__(
**params | {
@ -82,9 +96,11 @@ class Processor(ConsumerProducer):
self.max_output = max_output
self.session = boto3.Session(
aws_access_key_id=aws_id_key,
aws_secret_access_key=aws_secret,
region_name=aws_region
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
profile_name=aws_profile,
region_name=aws_region,
)
self.bedrock = self.session.client(service_name='bedrock-runtime')