Merge pull request #112 from aws-samples/spy_dev

add llm support ak/sk connect
aws-samples · Jun 24, 2024 · 31c5376 · 31c5376
2 parents e0ce07a + e6d9d74
commit 31c5376
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 4 deletions.
diff --git a/application/.env.cntemplate b/application/.env.cntemplate
@@ -26,3 +26,7 @@ SAGEMAKER_ENDPOINT_SQL=sql-sqlcoder-7b-2-7e5b6
 SAGEMAKER_ENDPOINT_EXPLAIN=llm-internlm2-chat-7b-3ab71
 
 EMBEDDING_DIMENSION=1024
+
+BEDROCK_SECRETS_AK_SK=bedrock-ak-sk
+OPENSEARCH_SECRETS_URL_HOST=opensearch-host-url
+OPENSEARCH_SECRETS_USERNAME_PASSWORD=opensearch-master-user
diff --git a/application/.env.template b/application/.env.template
@@ -22,3 +22,7 @@ DYNAMODB_AWS_REGION=us-west-2
 
 EMBEDDING_DIMENSION=1536
 BEDROCK_EMBEDDING_MODEL=amazon.titan-embed-text-v1
+
+BEDROCK_SECRETS_AK_SK=bedrock-ak-sk
+OPENSEARCH_SECRETS_URL_HOST=opensearch-host-url
+OPENSEARCH_SECRETS_USERNAME_PASSWORD=opensearch-master-user
diff --git a/application/utils/env_var.py b/application/utils/env_var.py
@@ -20,7 +20,7 @@
 OPENSEARCH_REGION = os.getenv('AOS_AWS_REGION')
 
 AOS_HOST = os.getenv('AOS_HOST')
-AOS_PORT = os.getenv('AOS_PORT')
+AOS_PORT = int(os.getenv('AOS_PORT'))
 AOS_USER = os.getenv('AOS_USER')
 AOS_PASSWORD = os.getenv('AOS_PASSWORD')
 AOS_DOMAIN = os.getenv('AOS_DOMAIN')
@@ -35,19 +35,25 @@
 
 OPENSEARCH_TYPE = os.getenv('OPENSEARCH_TYPE')
 
+OPENSEARCH_SECRETS_URL_HOST = os.getenv('OPENSEARCH_SECRETS_URL_HOST', 'opensearch-host-url')
+
+OPENSEARCH_SECRETS_USERNAME_PASSWORD = os.getenv('OPENSEARCH_SECRETS_USERNAME_PASSWORD', 'opensearch-master-user')
+
+BEDROCK_SECRETS_AK_SK = os.getenv('BEDROCK_SECRETS_AK_SK')
+
 
 def get_opensearch_parameter():
     try:
         session = boto3.session.Session()
         sm_client = session.client(service_name='secretsmanager', region_name=AWS_DEFAULT_REGION)
-        master_user = sm_client.get_secret_value(SecretId='opensearch-host-url')['SecretString']
+        master_user = sm_client.get_secret_value(SecretId=OPENSEARCH_SECRETS_URL_HOST)['SecretString']
         data = json.loads(master_user)
         es_host_name = data.get('host')
         # cluster endpoint, for example: my-test-domain.us-east-1.es.amazonaws.com/
         host = es_host_name + '/' if es_host_name[-1] != '/' else es_host_name
 
         sm_client = session.client(service_name='secretsmanager', region_name=AWS_DEFAULT_REGION)
-        master_user = sm_client.get_secret_value(SecretId='opensearch-master-user')['SecretString']
+        master_user = sm_client.get_secret_value(SecretId=OPENSEARCH_SECRETS_USERNAME_PASSWORD)['SecretString']
         data = json.loads(master_user)
         username = data.get('username')
         password = data.get('password')
@@ -59,6 +65,25 @@ def get_opensearch_parameter():
         raise e
 
 
+def get_bedrock_parameter():
+    bedrock_ak_sk_info = {}
+    try:
+        session = boto3.session.Session()
+        sm_client = session.client(service_name='secretsmanager', region_name=AWS_DEFAULT_REGION)
+        if BEDROCK_SECRETS_AK_SK is not None and BEDROCK_SECRETS_AK_SK != "":
+            bedrock_info = sm_client.get_secret_value(SecretId=BEDROCK_SECRETS_AK_SK)['SecretString']
+            data = json.loads(bedrock_info)
+            access_key = data.get('access_key_id')
+            secret_key = data.get('secret_access_key')
+            bedrock_ak_sk_info['access_key_id'] = access_key
+            bedrock_ak_sk_info['secret_access_key'] = secret_key
+            return bedrock_ak_sk_info
+    except ClientError as e:
+        # For a list of exceptions thrown, see
+        # https://docs.aws.amazon.com/secretsmanager/latest/apireference/API_GetSecretValue.html
+        return bedrock_ak_sk_info
+
+
 if OPENSEARCH_TYPE == "service":
     opensearch_host, opensearch_port, opensearch_username, opensearch_password = get_opensearch_parameter()
     AOS_HOST = opensearch_host
@@ -78,3 +103,5 @@ def get_opensearch_parameter():
     'agent_index': AOS_INDEX_AGENT,
     'embedding_dimension': EMBEDDING_DIMENSION
 }
+
+bedrock_ak_sk_info = get_bedrock_parameter()
diff --git a/application/utils/llm.py b/application/utils/llm.py
@@ -13,6 +13,7 @@
     generate_agent_analyse_prompt, generate_data_summary_prompt, generate_suggest_question_prompt, \
     generate_query_rewrite_prompt
 
+from utils.env_var import bedrock_ak_sk_info
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
@@ -38,7 +39,13 @@
 def get_bedrock_client():
     global bedrock
     if not bedrock:
-        bedrock = boto3.client(service_name='bedrock-runtime', config=config)
+        if len(bedrock_ak_sk_info) == 0:
+            bedrock = boto3.client(service_name='bedrock-runtime', config=config)
+        else:
+            bedrock = boto3.client(
+                service_name='bedrock-runtime',  config=config,
+                aws_access_key_id=bedrock_ak_sk_info['access_key_id'],
+                aws_secret_access_key=bedrock_ak_sk_info['secret_access_key'])
     return bedrock