From db83705c3dfa5201e8bb4d8c48d0d23fadc0a0ea Mon Sep 17 00:00:00 2001
From: Joel Balcaen <jobalcaen+github@gmail.com>
Date: Thu, 25 Apr 2024 16:25:00 -0300
Subject: [PATCH] udate logging

---
 lambdas/bedrock_invoker/src/index.py | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/lambdas/bedrock_invoker/src/index.py b/lambdas/bedrock_invoker/src/index.py
index 1260aee..cea9b2c 100644
--- a/lambdas/bedrock_invoker/src/index.py
+++ b/lambdas/bedrock_invoker/src/index.py
@@ -1,10 +1,13 @@
 import boto3
 import json
 from botocore.exceptions import BotoCoreError, ClientError
+from aws_lambda_powertools import Logger
 
+logger = Logger()
 s3 = boto3.client('s3')
 bedrock = boto3.client('bedrock-runtime')
 
+@logger.inject_lambda_context
 def lambda_handler(event, context):
     """
     Invokes a bedrock model with the given parameters and s3 text object
@@ -12,11 +15,14 @@ def lambda_handler(event, context):
     s3_arn = event['s3_arn']
     prompt = event['prompt']
 
+    logger.info(f"Invoking claude with prompt: ",prompt)
+
     # Parse the S3 ARN to get the bucket and key
     s3_path = s3_arn.replace("arn:aws:s3:::", "")
     bucket, key = s3_path.split('/', 1)
 
     # Download the file from S3
+    logger.info(f"Fetching file bucket: {bucket}, key: {key}")
     try:
         s3_object = s3.get_object(Bucket=bucket, Key=key)
     except ClientError as e:
@@ -27,7 +33,8 @@ def lambda_handler(event, context):
 
     # Extract text from the S3 object
     extracted_text = s3_object['Body'].read().decode('utf-8')
-
+    logger.info(f"Extracted text what is {len(extracted_text)} characters long")
+    
     claude_body = {
         "anthropic_version": "bedrock-2023-05-31",
         "max_tokens": 4096,
@@ -44,6 +51,18 @@ def lambda_handler(event, context):
             }
         ]
     }
+    # Create a copy of claude_body
+    redacted_claude_body = claude_body.copy()
+
+    # Replace all the text values in messages with "redacted"
+    for message in redacted_claude_body["messages"]:
+        for content in message["content"]:
+            if 'text' in content:
+                content['text'] = "redacted"
+
+    logger.info(f"Invoke bedrock with this body: ", redacted_claude_body)
+    bedrock_model = 'anthropic.claude-3-sonnet-20240229-v1:0'
+    logger.info(f"Invoke bedock with this model: ", bedrock_model)
 
     # Invoke the Bedrock model with the extracted text and the provided parameters
     try:
@@ -51,7 +70,7 @@ def lambda_handler(event, context):
             body=json.dumps(claude_body),
             contentType='application/json',
             accept='application/json',
-            modelId='anthropic.claude-3-sonnet-20240229-v1:0',
+            modelId=bedrock_model,
         )
     except BotoCoreError as e:
         return {