Merge branch 'feature/chat-with-document' into 'develop'

Guardrails Protection

See merge request genaiic-reusable-assets/engagement-artifacts/genaiic-idp-accelerator!310

Author: rstrahan

src/lambda/chat_with_document_resolver/index.py

@@ -9,15 +9,28 @@
 import base64
 import hashlib
 import os
+import re 
 from urllib.parse import urlparse
 from botocore.exceptions import ClientError
 from idp_common.bedrock.client import BedrockClient
 
 # Set up logging
 logger = logging.getLogger()
 logger.setLevel(os.environ.get("LOG_LEVEL", "INFO"))
-# Get LOG_LEVEL from environment variable with INFO as default
 
+def remove_text_between_brackets(text):
+    # Find position of first opening bracket
+    start = text.find('{')
+    # Find position of last closing bracket
+    end = text.rfind('}')
+    
+    # If both brackets exist, remove text between them including brackets
+    if start != -1 and end != -1:
+        return text[:start] + text[end+1:]
+    # If brackets not found, return original string
+    return text
+
+# Get LOG_LEVEL from environment variable with INFO as default
 def s3_object_exists(bucket, key):
     try:
         s3 = boto3.client('s3')
@@ -66,7 +79,7 @@ def get_full_text(bucket, key):
 
     except Exception as e:
         logger.error(f"Error getting document pages: {str(e)}")
-        raise
+        raise Exception(f"Error getting document pages: {str(e)}")
 
 
 def get_summarization_model():
@@ -98,7 +111,6 @@ def handler(event, context):
 
     try:
         # logger.info(f"Received event: {json.dumps(event)}")
-
         objectKey = event['arguments']['s3Uri']
         prompt = event['arguments']['prompt']
         history = event['arguments']['history']
@@ -115,8 +127,6 @@ def handler(event, context):
 
         output_bucket = os.environ['OUTPUT_BUCKET']
 
-        bedrock_runtime = boto3.client('bedrock-runtime', region_name=os.environ['AWS_REGION'])
-
         if (len(objectKey)):
             fulltext_key = objectKey + '/summary/fulltext.txt'
             content_str = ""
@@ -140,7 +150,6 @@ def handler(event, context):
             logger.info(f"Output Bucket: {output_bucket}")
             logger.info(f"Full Text Key: {fulltext_key}")
 
-            client = BedrockClient()
             # Content with cachepoint tags
             content = [
                 {
@@ -150,16 +159,32 @@ def handler(event, context):
                 }
             ]
 
-            model_response = client.invoke_model(
-                model_id="us.amazon.nova-pro-v1:0",
+            client = BedrockClient(
+                region=os.environ['AWS_REGION'],
+                max_retries=5,
+                initial_backoff=1.5,
+                max_backoff=300,
+                metrics_enabled=True
+            )
+
+            # Invoke a model
+            response = client.invoke_model(
+                model_id=selectedModelId,
                 system_prompt="You are an assistant that's responsible for getting details from document text attached here based on questions from the user.\n\nIf you don't know the answer, just say that you don't know. Don't try to make up an answer.\n\nAdditionally, use the user and assistant responses in the following JSON object to see what's been asked and what the resposes were in the past.\n\n",
                 content=content,
                 temperature=0.0
             )
 
-            text = client.extract_text_from_response(model_response)
+            text = client.extract_text_from_response(response)
+            logger.info(f"Full response: {text}")
+
+            # right now, there is a JSON object before the full response when a guardrail is tripped
+            # need to remove that JSON object first
+            logger.info(f"New response: {remove_text_between_brackets(text).strip("\n")}")
+            cleaned_up_text = remove_text_between_brackets(text).strip("\n")
+            
+            chat_response = {"cr": {"content": [{"text": cleaned_up_text}]}}
 
-            chat_response = {"cr": {"content": [{"text": text}]}}
             return json.dumps(chat_response)
 
     except ClientError as e:
@@ -175,7 +200,7 @@ def handler(event, context):
             raise Exception(error_message)
 
     except Exception as e:
-        logger.error(f"Unexpected error: {str(e)}")
-        raise Exception(f"Unexpected error: {str(e)}")
+        logger.error(f"{str(e)}")
+        raise Exception(f"{str(e)}")
 
     return response_data

template.yaml

@@ -6107,6 +6107,7 @@ Resources:
           OUTPUT_BUCKET: !Ref OutputBucket
           CONFIGURATION_TABLE_NAME: !Ref ConfigurationTable
           TRACKING_TABLE_NAME: !Ref TrackingTable
+          GUARDRAIL_ID_AND_VERSION: !If [HasGuardrailConfig, !Sub "${BedrockGuardrailId}:${BedrockGuardrailVersion}", ""]
       LoggingConfig:
         LogGroup: !Ref ChatWithDocumentResolverFunctionLogGroup
       Policies: