feat(visualqa): implemented review comments

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,3 +41,32 @@ 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/__init__.py

@@ -1 +1 @@
-from .text_generation_llm_selector import get_llm, get_max_tokens, get_embeddings_llm
+from .text_generation_llm_selector import get_llm, get_max_tokens, get_embeddings_llm,get_bedrock_fm

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

@@ -10,6 +10,7 @@
 # OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
 # and limitations under the License.
 #
+from aiohttp import ClientError
 from langchain.llms.bedrock import Bedrock
 from langchain_community.embeddings import BedrockEmbeddings
 import os
@@ -26,7 +27,7 @@
 metrics = Metrics(namespace="question_answering", service="QUESTION_ANSWERING")
 
 
-def get_llm(callbacks=None):
+def get_llm(callbacks=None,model_id="anthropic.claude-v2:1"):
     bedrock = boto3.client('bedrock-runtime')
 
     params = {
@@ -39,7 +40,7 @@ def get_llm(callbacks=None):
 
     kwargs = {
         "client": bedrock,
-        "model_id": "anthropic.claude-v2:1",
+        "model_id": model_id,
         "model_kwargs": params,
         "streaming": False 
     }
@@ -50,10 +51,60 @@ def get_llm(callbacks=None):
 
     return Bedrock(**kwargs)
 
-def get_embeddings_llm(model_id):
+def get_embeddings_llm(model_id,modality):
     bedrock = boto3.client('bedrock-runtime')
-    return BedrockEmbeddings(client=bedrock, model_id=model_id)
-    
+    validation_status=validate_model_id_in_bedrock(model_id,modality)
+    if(validation_status['status']):
+        return BedrockEmbeddings(client=bedrock, model_id=model_id)
+    else:
+        return None
+
+
+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}')
+    if(validation_status['status']):
+        return bedrock_client
+    else:
+        logger.error(f"reason ::{validation_status['message']} ")
+        return None
+
+
+
+#TODO -add max token based on model id    
 def get_max_tokens():
     return 200000
-    
+
+def validate_model_id_in_bedrock(model_id,modality):
+        """
+        Validate if the listed model id is supported with given modality
+        in bedrock or not.
+        """
+        response={
+            "status":False,
+            "message":f"model {model_id} is not supported in bedrock."
+        }
+        try:
+            bedrock_client = boto3.client(service_name="bedrock")
+            bedrock_model_list = bedrock_client.list_foundation_models()
+            models = bedrock_model_list["modelSummaries"]
+            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}')
+            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/StreamingCallbackHandler.py

@@ -1,3 +1,15 @@
+#
+# 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.
+#
 from .helper import  send_job_status, JobStatus
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.schema import LLMResult

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

@@ -18,12 +18,10 @@
 import base64
 
 from langchain.chains import LLMChain
-from llms import get_llm, get_max_tokens
+from llms import  get_max_tokens
 from typing import Any, Dict, List, Union
-from langchain.prompts import PromptTemplate
 from .s3inmemoryloader import S3FileLoaderInMemory
-from .StreamingCallbackHandler import StreamingCallbackHandler
-from .helper import load_vector_db_opensearch, send_job_status, JobStatus
+from .helper import  send_job_status, JobStatus
 from .image_qa import  run_qa_agent_on_image_no_memory,run_qa_agent_rag_on_image_no_memory
 from .doc_qa import run_qa_agent_rag_no_memory, run_qa_agent_from_single_document_no_memory
 
@@ -45,8 +43,6 @@ def run_question_answering(arguments):
         filename = ''
         arguments['filename'] = ''
 
-    image_url = arguments['presignedurl']
-   
     #set deafult modality to text
     qa_model= arguments['qa_model']
     modality=qa_model.get('modality','Text')
@@ -57,7 +53,7 @@ def run_question_answering(arguments):
 
           # user didn't provide a image url as input, we use the RAG source against the entire knowledge base
          if response_generation_method == 'LONG_CONTEXT':
-                if not image_url: 
+                if not filename: 
                     warning = 'Error: Image presigned url is required for LONG_CONTEXT approach, defaulting to RAG.'
                     logger.warning(warning)
                     llm_response = run_qa_agent_rag_on_image_no_memory(arguments)

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

@@ -18,7 +18,7 @@
 import base64
 
 from langchain.chains import LLMChain
-from llms import get_llm, get_max_tokens
+from llms import get_llm
 from typing import Any, Dict, List, Union
 from langchain.prompts import PromptTemplate
 from .s3inmemoryloader import S3FileLoaderInMemory
@@ -41,7 +41,10 @@ def run_qa_agent_rag_no_memory(input_params):
     logger.info("starting qa agent with rag approach without memory :: {input_params}")
 
     base64_bytes = input_params['question'].encode("utf-8")
-    embedding_model_id = input_params['embeddings_model']['modelId']
+    embedding_model = input_params['embeddings_model']
+    embedding_model_id = embedding_model['modelId']
+    modality=embedding_model.get("modality", "Text")
+
     qa_model_id = input_params['qa_model']['modelId']
     sample_string_bytes = base64.b64decode(base64_bytes)
     decoded_question = sample_string_bytes.decode("utf-8")
@@ -69,7 +72,8 @@ def run_qa_agent_rag_no_memory(input_params):
                                               os.environ.get('OPENSEARCH_DOMAIN_ENDPOINT'),
                                               os.environ.get('OPENSEARCH_INDEX'),
                                               os.environ.get('OPENSEARCH_SECRET_ID'),
-                                              embedding_model_id)
+                                              embedding_model_id,
+                                              modality)
 
     else:
         logger.info("_retriever already exists")
@@ -106,7 +110,7 @@ def run_qa_agent_rag_no_memory(input_params):
     # 2 : load llm using the selector
     streaming = input_params.get("streaming", False)
     callback_manager = [StreamingCallbackHandler(status_variables)] if streaming else None
-    _qa_llm = get_llm(callback_manager)
+    _qa_llm = get_llm(callback_manager,qa_model_id)
 
     if (_qa_llm is None):
         logger.error('llm is None, returning')
@@ -154,6 +158,7 @@ def run_qa_agent_from_single_document_no_memory(input_params):
     logger.info("starting qa agent without memory single document")
 
     base64_bytes = input_params['question'].encode("utf-8")
+    qa_model_id = input_params['qa_model']['modelId']
 
     sample_string_bytes = base64.b64decode(base64_bytes)
     decoded_question = sample_string_bytes.decode("utf-8")
@@ -200,7 +205,7 @@ def run_qa_agent_from_single_document_no_memory(input_params):
     # 2 : run the question
     streaming = input_params.get("streaming", False)
     callback_manager = [StreamingCallbackHandler(status_variables)] if streaming else None
-    _qa_llm = get_llm(callback_manager)
+    _qa_llm = get_llm(callback_manager,qa_model_id)
 
     if (_qa_llm is None):
         logger.info('llm is None, returning')

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

@@ -10,8 +10,10 @@
 # OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
 # and limitations under the License.
 #
+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
@@ -57,7 +59,7 @@ class JobStatus(Enum):
         base64.b64encode("Sorry, it seems an issue happened on my end, and I'm not able to answer your question. Please contact an administrator to understand why !".encode('utf-8'))
     )
     ERROR_SEMANTIC_SEARCH = (
-        'Exception during simialirty search, Please verify model for the selected modality', 
+        'Exception during similarity search, Please verify model for the selected modality', 
         base64.b64encode("Sorry, it seems an issue happened on my end, and I'm not able to answer your question. Please contact an administrator to understand why !".encode('utf-8'))
     )
 
@@ -97,7 +99,8 @@ def load_vector_db_opensearch(region: str,
                               opensearch_domain_endpoint: str,
                               opensearch_index: str,
                               secret_id: str,
-                              model_id: str) -> OpenSearchVectorSearch:
+                              model_id: str,
+                              modality: str) -> OpenSearchVectorSearch:
     print(f"load_vector_db_opensearch, region={region}, "
                 f"opensearch_domain_endpoint={opensearch_domain_endpoint}, opensearch_index={opensearch_index}")
 
@@ -114,16 +117,17 @@ def load_vector_db_opensearch(region: str,
             opensearch_api_name,
             session_token=credentials.token,
         )
-    embedding_function = get_embeddings_llm(model_id)
+    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}")
     return vector_db
 
@@ -155,7 +159,8 @@ def send_job_status(variables):
 
     print(request)
 
-    GRAPHQL_URL = os.environ['GRAPHQL_URL']
+    #GRAPHQL_URL = os.environ['GRAPHQL_URL']
+    GRAPHQL_URL ="https://j2uzmlvujbhbzoduvpctgkpu2e.appsync-api.us-east-1.amazonaws.com/graphql"
     HEADERS={
         "Content-Type": "application/json",
         }
@@ -180,4 +185,21 @@ def get_presigned_url(bucket,key) -> str:
              return url
         except Exception as exception:
             print(f"Reason: {exception}")
-            return ""
+            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}")
+            return file_path
+        except ClientError as client_err:
+            print(f"Couldn\'t download file {client_err.response['Error']['Message']}")
+        
+        except Exception as exp:
+            print(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:
+            b64_image = base64.b64encode(image_file.read()).decode('utf8')       
+        return b64_image