feat(visualqa): fixed bugs

Author: Dinesh Sajwan

lambda/aws-qa-appsync-opensearch/question_answering/src/lambda.py

@@ -23,7 +23,7 @@
 tracer = Tracer(service="QUESTION_ANSWERING")
 metrics = Metrics(namespace="question_answering", service="QUESTION_ANSWERING")
 
-#@logger.inject_lambda_context(log_event=True)
+@logger.inject_lambda_context(log_event=True)
 @tracer.capture_lambda_handler
 @metrics.log_metrics(capture_cold_start_metric=True)
 def handler(event,  context: LambdaContext) -> dict:
@@ -41,32 +41,3 @@ def handler(event,  context: LambdaContext) -> dict:
 
     print(f"llm_response is {llm_response}")
     return llm_response
-
-input ={"detail": {
-            "jobid": "111",
-            "jobstatus": "",
-            "qa_model": {
-                "provider": "Bedrock",
-                "modelId": "anthropic.claude-3-sonnet-20240229-v1:0",
-                "streaming": True,
-                "modality": "Image"
-            },
-            "embeddings_model": {
-                "provider": "Bedrock",
-                "modelId": "amazon.titan-embed-image-v1",
-                "streaming": True
-            },
-            "retrieval": {
-                "max_docs": 1,
-                "index_name": "",
-                "filter_filename": ""
-            },
-            "filename": "two_cats.jpeg",
-            "presignedurl": "",
-            "question": "d2hhdCBhcmUgdGhlIGNhdHMgZG9pbmc/",
-            "verbose": False,
-            "responseGenerationMethod": "LONG_CONTEXT"
-        }
-    }
-
-handler(input, None)

lambda/aws-qa-appsync-opensearch/question_answering/src/llms/text_generation_llm_selector.py

@@ -27,6 +27,7 @@
 metrics = Metrics(namespace="question_answering", service="QUESTION_ANSWERING")
 
 
+
 def get_llm(callbacks=None,model_id="anthropic.claude-v2:1"):
     bedrock = boto3.client('bedrock-runtime')
 
@@ -63,7 +64,7 @@ def get_embeddings_llm(model_id,modality):
 def get_bedrock_fm(model_id,modality):
     bedrock_client = boto3.client('bedrock-runtime')
     validation_status= validate_model_id_in_bedrock(model_id,modality)
-    print(f' validation_status :: {validation_status}')
+    logger.info(f' validation_status :: {validation_status}')
     if(validation_status['status']):
         return bedrock_client
     else:
@@ -73,9 +74,16 @@ def get_bedrock_fm(model_id,modality):
 
 
 #TODO -add max token based on model id    
-def get_max_tokens():
-    return 200000
-
+def get_max_tokens(model_id):
+    match model_id:
+        case "anthropic.claude-v2:1":
+            return 200000
+        case "anthropic.claude-3-sonnet-20240229-v1:0":
+            return 200000
+        case _:
+            return 4096
+
+        
 def validate_model_id_in_bedrock(model_id,modality):
         """
         Validate if the listed model id is supported with given modality
@@ -92,19 +100,16 @@ def validate_model_id_in_bedrock(model_id,modality):
             for model in models:
                 if model["modelId"].lower() == model_id.lower():   
                     response["message"]=f"model {model_id} does not support modality {modality} "                 
-                    print(f' modality :: {model["inputModalities"]}')
                     for inputModality in model["inputModalities"]:
                         if inputModality.lower() == modality.lower():
-                            print(f' modality supported')
                             response["message"]=f"model {model_id} with modality {modality} is supported with bedrock "                 
                             response["status"] = True
 
-            print(f' response :: {response}')
+            logger.info(f' response :: {response}')
             return response         
         except ClientError as ce:
             message=f"error occured while validating model in bedrock {ce}"
             logger.error(message)
             response["status"] = False
             response["message"] = message
-            print(f' response :: {response}')
             return response     

lambda/aws-qa-appsync-opensearch/question_answering/src/qa_agent/helper.py

@@ -10,19 +10,25 @@
 # OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
 # and limitations under the License.
 #
+import os
+import boto3
+import json
+import base64
 from pathlib import Path
 from aiohttp import ClientError
 from langchain_community.vectorstores import OpenSearchVectorSearch
-#from opensearchpy import RequestsHttpConnection
+from opensearchpy import RequestsHttpConnection
 from llms import get_embeddings_llm
 import requests
-import os
-import boto3
-import json
-import base64
 from enum import Enum
 from requests_aws4auth import AWS4Auth
 s3 = boto3.client('s3')
+from aws_lambda_powertools import Logger, Tracer, Metrics
+
+
+logger = Logger(service="QUESTION_ANSWERING")
+tracer = Tracer(service="QUESTION_ANSWERING")
+metrics = Metrics(namespace="question_answering", service="QUESTION_ANSWERING")
 
 
 class JobStatus(Enum):
@@ -101,7 +107,7 @@ def load_vector_db_opensearch(region: str,
                               secret_id: str,
                               model_id: str,
                               modality: str) -> OpenSearchVectorSearch:
-    print(f"load_vector_db_opensearch, region={region}, "
+    logger.info(f"load_vector_db_opensearch, region={region}, "
                 f"opensearch_domain_endpoint={opensearch_domain_endpoint}, opensearch_index={opensearch_index}")
 
     # if the secret id is not provided
@@ -120,15 +126,16 @@ def load_vector_db_opensearch(region: str,
     embedding_function = get_embeddings_llm(model_id,modality)
 
     opensearch_url = opensearch_domain_endpoint if opensearch_domain_endpoint.startswith("https://") else f"https://{opensearch_domain_endpoint}"
-    # vector_db = OpenSearchVectorSearch(index_name=opensearch_index,
-    #                                    embedding_function=embedding_function,
-    #                                    opensearch_url=opensearch_url,
-    #                                    http_auth=http_auth,
-    #                                    use_ssl = True,
-    #                                    verify_certs = True,
-    #                                    connection_class = RequestsHttpConnection)
+    
+    vector_db = OpenSearchVectorSearch(index_name=opensearch_index,
+                                        embedding_function=embedding_function,
+                                        opensearch_url=opensearch_url,
+                                        http_auth=http_auth,
+                                        use_ssl = True,
+                                        verify_certs = True,
+                                        connection_class = RequestsHttpConnection)
     vector_db=""
-    print(f"returning handle to OpenSearchVectorSearch, vector_db={vector_db}")
+    logger.info(f"returning handle to OpenSearchVectorSearch, vector_db={vector_db}")
     return vector_db
 
 def send_job_status(variables):
@@ -159,8 +166,7 @@ def send_job_status(variables):
 
     print(request)
 
-    #GRAPHQL_URL = os.environ['GRAPHQL_URL']
-    GRAPHQL_URL ="https://j2uzmlvujbhbzoduvpctgkpu2e.appsync-api.us-east-1.amazonaws.com/graphql"
+    GRAPHQL_URL = os.environ['GRAPHQL_URL']
     HEADERS={
         "Content-Type": "application/json",
         }
@@ -172,7 +178,7 @@ def send_job_status(variables):
         auth=aws_auth_appsync,
         timeout=10
     )
-    print('res :: {}',responseJobstatus)
+    logger.info('res :: {}',responseJobstatus)
 
 def get_presigned_url(bucket,key) -> str:
         try:
@@ -181,23 +187,23 @@ def get_presigned_url(bucket,key) -> str:
                 Params={'Bucket': bucket, 'Key': key},
                 ExpiresIn=900
                 )
-             print(f"presigned url generated for {key} from {bucket}")
+             logger.info(f"presigned url generated for {key} from {bucket}")
              return url
         except Exception as exception:
-            print(f"Reason: {exception}")
+            logger.error(f"Reason: {exception}")
             return ""
 
 def download_file(bucket,key )-> str:
         try: 
             file_path = "/tmp/" + os.path.basename(key)
             s3.download_file(bucket, key,file_path)
-            print(f"file downloaded {file_path}")
+            logger.info(f"file downloaded {file_path}")
             return file_path
         except ClientError as client_err:
-            print(f"Couldn\'t download file {client_err.response['Error']['Message']}")
+            logger.error(f"Couldn\'t download file {client_err.response['Error']['Message']}")
 
         except Exception as exp:
-            print(f"Couldn\'t download file : {exp}")
+            logger.error(f"Couldn\'t download file : {exp}")
 
 def encode_image_to_base64(image_file_path,image_file) -> str:
         with open(image_file_path, "rb") as image_file:

lambda/aws-qa-appsync-opensearch/question_answering/src/qa_agent/image_qa.py

@@ -35,8 +35,7 @@
 
 
 
-#bucket_name = os.environ['INPUT_BUCKET']
-bucket_name="persistencestack-inputassets7d1d3f52-qert2sgpwhtu"
+bucket_name = os.environ['INPUT_BUCKET']
 
 def run_qa_agent_on_image_no_memory(input_params):
     logger.info("starting qa agent without memory on uploaded image")

lambda/aws-qa-appsync-opensearch/question_answering/src/qa_agent/sagemaker_endpoint.py

@@ -1,11 +1,24 @@
-
-from langchain.llms.sagemaker_endpoint import LLMContentHandler, SagemakerEndpoint
-from aws_lambda_powertools import Logger, Tracer, Metrics
-
+#
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+#
 import json
 import os
+from langchain.llms.sagemaker_endpoint import LLMContentHandler, SagemakerEndpoint
+from aws_lambda_powertools import Logger, Tracer, Metrics
 logger = Logger(service="QUESTION_ANSWERING")
 
+
+sageMakerEndpoint= os.environ['SAGEMAKER_ENDPOINT']
+
 class ContentHandler(LLMContentHandler):
     content_type = "application/json"
     accepts = "application/json"
@@ -35,17 +48,22 @@ class MultiModal():
 
     @classmethod
     def sagemakerendpoint_llm(self,model_id):
-        try: 
-            endpoint= SagemakerEndpoint(
-            endpoint_name=model_id,
-            region_name=os.environ["AWS_REGION"],
-            model_kwargs=self.parameters,
-            content_handler=content_handler,
-            )
-            return endpoint
-        except Exception as err:
-            logger.error(' Error when accessing sagemaker endpoint :: {model_id} , returning...')
-            return ''
+        if(sageMakerEndpoint ==model_id):
+            try: 
+                endpoint= SagemakerEndpoint(
+                endpoint_name=model_id,
+                region_name=os.environ["AWS_REGION"],
+                model_kwargs=self.parameters,
+                content_handler=content_handler,
+                )
+                return endpoint
+            except Exception as err:
+                logger.error(f' Error when accessing sagemaker endpoint :: {model_id} , returning...')
+                return ''
+        else:
+                logger.error(f" The sagemaker model Id {model_id} do not match a sagemaker endpoint name {sageMakerEndpoint}")
+                return ''
+            
 
 
 

lambda/aws-rag-appsync-stepfn-opensearch/s3_file_transformer/src/helpers/utils.py

@@ -47,9 +47,14 @@ def transform_pdf_document(input_bucket: str,file_name: str,output_bucket: str,o
         if not document_content:
                 return 'Unable to load document'             
         else:
+            try:
                 encoded_string = document_content.encode("utf-8")
                 s3.Bucket(output_bucket).put_object(Key=output_file_name, Body=encoded_string) 
                 return 'File transformed' 
+            except Exception as e: 
+                  logger.error(f'Error in uploading {output_file_name} to  {output_bucket} :: {e}')
+                  return 'File transformed Failed'
+
 
 @tracer.capture_method
 def transform_image_document(input_bucket: str,file_name: str,output_bucket: str):  

lambda/aws-rag-appsync-stepfn-opensearch/s3_file_transformer/src/lambda.py

@@ -87,7 +87,6 @@ def handler(event,  context: LambdaContext) -> dict:
             response['name'] = output_file_name
             if(extension == '.pdf'):
                 response['status'] = transform_pdf_document(input_bucket,file_name,output_bucket,output_file_name)
-                print(f' pdf processed ::' )
             elif(extension == '.jpg'or extension == '.jpeg' or extension == '.png' or extension == '.svg'):
                 response['status'] = transform_image_document(input_bucket,file_name,output_bucket)
             #TODO add csv, doc, docx file type support as well.

lambda/aws-rag-appsync-stepfn-opensearch/s3_file_transformer/src/requirements.txt

@@ -1,9 +1,9 @@
 aws-lambda-powertools
-aws-xray-sdk
-fastjsonschema
-typing-extensions
+aws-xray-sdk==3.5.4
+fastjsonschema==2.19.1
+typing-extensions====4.7.0
 boto3>=1.34.29
-requests
+requests==2.31.0
 langchain==0.1.4
 pypdf2==3.0.1
 Pillow==10.2.0

src/patterns/gen-ai/aws-qa-appsync-opensearch/index.ts

@@ -152,6 +152,12 @@ export interface QaAppsyncOpensearchProps {
    * and settings instead of the existing
    */
   readonly customDockerLambdaProps?: DockerLambdaCustomProps | undefined;
+
+  /**
+   * Optional. Allows to provide custom lambda code
+   * and settings instead of the existing
+   */
+  readonly sagemakerEndpointName?: string
 }
 
 /**
@@ -466,14 +472,16 @@ export class QaAppsyncOpensearch extends Construct {
         resources: ['*'],
       }),
     );
-    question_answering_function_role.addToPolicy(
-      new iam.PolicyStatement({
-        effect: iam.Effect.ALLOW,
-        actions: ['sagemaker:InvokeEndpoint'],
-        resources: ['*'],
-      }),
-    );
-
+        
+    if(props.sagemakerEndpointName){
+      question_answering_function_role.addToPolicy(
+        new iam.PolicyStatement({
+          effect: iam.Effect.ALLOW,
+          actions: ['sagemaker:InvokeEndpoint'],
+          resources: ['arn:'+ Aws.PARTITION +':sagemaker:' + Aws.ACCOUNT_ID + ':endpoint/*' ],
+        }),
+      );
+    }
     // The lambda will access the opensearch credentials
     if (props.openSearchSecret) {
       props.openSearchSecret.grantRead(question_answering_function_role);
@@ -553,6 +561,7 @@ export class QaAppsyncOpensearch extends Construct {
       true,
     );
 
+    const sagemakerEndpointNamestr = props.sagemakerEndpointName || ""
     const construct_docker_lambda_props = {
       code: lambda.DockerImageCode.fromImageAsset(
         path.join(
@@ -576,6 +585,7 @@ export class QaAppsyncOpensearch extends Construct {
         OPENSEARCH_DOMAIN_ENDPOINT: opensearch_helper.getOpenSearchEndpoint(props),
         OPENSEARCH_INDEX: props.openSearchIndexName,
         OPENSEARCH_SECRET_ID: SecretId,
+        SAGEMAKER_ENDPOINT:sagemakerEndpointNamestr
       },
       ...(props.lambdaProvisionedConcurrency && {
         currentVersionOptions: {

src/patterns/gen-ai/aws-rag-appsync-stepfn-opensearch/index.ts

@@ -555,19 +555,9 @@ export class RagAppsyncStepfnOpensearch extends Construct {
     s3_transformer_job_function_role.addToPolicy(new iam.PolicyStatement({
       effect: iam.Effect.ALLOW,
       actions: [
-        'rekognition:CompareFaces',
-        'rekognition:DetectFaces',
-        'rekognition:DetectLabels',
-        'rekognition:ListFaces',
-        'rekognition:SearchFaces',
-        'rekognition:SearchFacesByImage',
-        'rekognition:DetectText',
-        'rekognition:GetCelebrityInfo',
-        'rekognition:RecognizeCelebrities',
         'rekognition:DetectModerationLabels',
       ],
-      //TODO: change the resource to specific arn
-      resources: ['*'],
+      resources: ['arn:'+ Aws.PARTITION +':rekognition:' + Aws.ACCOUNT_ID + ':project/*' ],
     }));
 
     s3_transformer_job_function_role.addToPolicy(new iam.PolicyStatement({