Skip to content

Commit

Permalink
Updated bedrock.py to handle case of deploying it over Lambda functio…
Browse files Browse the repository at this point in the history
…n with attached role access.

Bedrock client access was working fine when used locally but it was not working when deployed the same solution over AWS Lambda function or EC2 Instance or ECS Task, when no access credentials being passed/set as environment variables. 

I have made changes in this to build bedrock-runtime client directly using boto3 which uses attached Role to managed AWS service. 

This is more secure way to running code over AWS. Tested it with AWS Lambda function and ECS container task.
  • Loading branch information
irshadc authored Nov 25, 2024
1 parent 1946d0d commit 7b48b68
Showing 1 changed file with 14 additions and 8 deletions.
22 changes: 14 additions & 8 deletions autogen/oai/bedrock.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,14 +79,20 @@ def __init__(self, **kwargs: Any):
retries={"max_attempts": self._retries, "mode": "standard"},
)

session = boto3.Session(
aws_access_key_id=self._aws_access_key,
aws_secret_access_key=self._aws_secret_key,
aws_session_token=self._aws_session_token,
profile_name=self._aws_profile_name,
)

self.bedrock_runtime = session.client(service_name="bedrock-runtime", config=bedrock_config)
#if haven't got any access_key or secret_key in environment variable or via arugments then
if self._aws_access_key is None or self._aws_access_key == "" or self._aws_secret_key is None or self._aws_secret_key == "":

# attempts to get client from attached role of mananged service (lambda, ec2, ecs, etc.)
self.bedrock_runtime = boto3.client(service_name="bedrock-runtime", config=bedrock_config)
else:
session = boto3.Session(
aws_access_key_id=self._aws_access_key,
aws_secret_access_key=self._aws_secret_key,
aws_session_token=self._aws_session_token,
profile_name=self._aws_profile_name,
)

self.bedrock_runtime = session.client(service_name="bedrock-runtime", config=bedrock_config)

def message_retrieval(self, response):
"""Retrieve the messages from the response."""
Expand Down

0 comments on commit 7b48b68

Please sign in to comment.