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

Fix Bedrock (#119)

Browse source
This commit is contained in:
cybermaggedon 2024-10-15 19:21:43 +01:00 • committed by GitHub
parent f36c5e5d03
commit 25983d1557
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 14 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

2
templates/components/bedrock.jsonnet
View file

@ -16,7 +16,7 @@ local chunker = import "chunker-recursive.jsonnet";
local envSecrets = engine.envSecrets("bedrock-credentials") local envSecrets = engine.envSecrets("bedrock-credentials")
.with_env_var("AWS_ID_KEY", "aws-id-key") .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"); .with_env_var("AWS_REGION", "aws-region");
local container = local container =

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

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