chore: ruff format

Author: Daniel Lorch

plugins/dynamic-few-shot-lambda/GENAIIDP-dynamic-few-shot.py

@@ -26,20 +26,21 @@
 
 # Parse environment variables with error handling
 try:
-    S3VECTOR_BUCKET = os.environ['S3VECTOR_BUCKET']
-    S3VECTOR_INDEX = os.environ['S3VECTOR_INDEX']
-    S3VECTOR_DIMENSIONS = int(os.environ['S3VECTOR_DIMENSIONS'])
-    MODEL_ID = os.environ['MODEL_ID']
-    TOP_K = int(os.environ['TOP_K'])
-    THRESHOLD = float(os.environ['THRESHOLD'])
+    S3VECTOR_BUCKET = os.environ["S3VECTOR_BUCKET"]
+    S3VECTOR_INDEX = os.environ["S3VECTOR_INDEX"]
+    S3VECTOR_DIMENSIONS = int(os.environ["S3VECTOR_DIMENSIONS"])
+    MODEL_ID = os.environ["MODEL_ID"]
+    TOP_K = int(os.environ["TOP_K"])
+    THRESHOLD = float(os.environ["THRESHOLD"])
 except (KeyError, ValueError, IndexError) as e:
     logger.error(f"Failed to parse environment variables: {e}")
     raise
 
 # Initialize clients
-s3vectors = boto3.client('s3vectors')
+s3vectors = boto3.client("s3vectors")
 bedrock_client = bedrock.BedrockClient()
 
+
 def lambda_handler(event, context):
     """
     Process a document to find similar examples using S3 Vectors similarity search.
@@ -61,7 +62,7 @@ def lambda_handler(event, context):
         }
     ]
     """
-    
+
     try:
         logger.info("=== DYNAMIC FEW-SHOT LAMBDA INVOKED ===")
         logger.debug(f"Complete input event: {json.dumps(event, indent=2)}")
@@ -73,7 +74,9 @@ def lambda_handler(event, context):
 
         logger.info(f"=== INPUT VALUES ===")
         logger.info(f"Class label: {class_label if class_label else 'Not specified'}")
-        logger.info(f"Document text: {len(document_text) if document_text else "0"} bytes")
+        logger.info(
+            f"Document text: {len(document_text) if document_text else '0'} bytes"
+        )
         logger.info(f"Image content: {len(image_content)} images")
 
         # Decode input data
@@ -91,15 +94,18 @@ def lambda_handler(event, context):
 
         logger.info("=== DYNAMIC FEW-SHOT LAMBDA COMPLETED ===")
         return result
-            
+
     except Exception as e:
         logger.error(f"=== DYNAMIC FEW-SHOT LAMBDA ERROR ===")
         logger.error(f"Error type: {type(e).__name__}")
         logger.error(f"Error message: {str(e)}")
-        logger.error(f"Input event keys: {list(event.keys()) if 'event' in locals() else 'Unknown'}")
+        logger.error(
+            f"Input event keys: {list(event.keys()) if 'event' in locals() else 'Unknown'}"
+        )
         # In demo, we'll fail gracefully with detailed error info
         raise Exception(f"Dynamic few-shot Lambda failed: {str(e)}")
 
+
 def _decode_images(image_content):
     """Base64 decode image content to bytes"""
     result = []
@@ -108,14 +114,16 @@ def _decode_images(image_content):
         result.append(image_data)
     return result
 
+
 def _encode_images(image_content):
     """Base64 encode image content to JSON-serializable string"""
     result = []
     for image_bytes in image_content:
-       image_base64 = base64.b64encode(image_bytes).decode("utf-8")
-       result.append(image_base64)
+        image_base64 = base64.b64encode(image_bytes).decode("utf-8")
+        result.append(image_base64)
     return result
 
+
 def _s3vectors_find_similar_items(image_data):
     """Find similar items for input"""
     # find similar items based on image similarity only
@@ -133,21 +141,21 @@ def _s3vectors_find_similar_items(image_data):
 
         # Only process this example if it has a non-empty attributesPrompt
         if not attributes_prompt or not attributes_prompt.strip():
-            logger.info(
-                f"Skipping example with empty attributesPrompt: {key}"
-            )
+            logger.info(f"Skipping example with empty attributesPrompt: {key}")
             continue
 
         attributes = _extract_metadata(metadata, distance)
         result.append(attributes)
 
     # sort results by distance score (lowest to highest - lower is more similar)
-    sorted_result = sorted(result, key=lambda example: example['distance'], reverse=False)
+    sorted_result = sorted(
+        result, key=lambda example: example["distance"], reverse=False
+    )
 
     # filter result by distance score
     filtered_result = []
     for example in sorted_result:
-        if example['distance'] > THRESHOLD:
+        if example["distance"] > THRESHOLD:
             logger.info(
                 f"Skipping example with distance {example['distance']} above threshold {THRESHOLD}: {key}"
             )
@@ -156,6 +164,7 @@ def _s3vectors_find_similar_items(image_data):
 
     return filtered_result
 
+
 def _s3vectors_find_similar_items_from_image(page_image):
     """Search for similar items using image query"""
     embedding = bedrock_client.generate_embedding(
@@ -169,11 +178,12 @@ def _s3vectors_find_similar_items_from_image(page_image):
         queryVector={"float32": embedding},
         topK=TOP_K,
         returnDistance=True,
-        returnMetadata=True
+        returnMetadata=True,
     )
     logger.debug(f"S3 vectors lookup result: {response['vectors']}")
     return response["vectors"]
 
+
 def _merge_examples(examples, new_examples):
     """
     Merge in-place new examples into the result list, avoiding duplicates.
@@ -185,7 +195,7 @@ def _merge_examples(examples, new_examples):
     for new_example in new_examples:
         key = new_example["key"]
         new_distance = new_example.get("distance", 1.0)
-        
+
         # update example
         if examples.get(key):
             existing_distance = examples[key].get("distance", 1.0)
@@ -195,9 +205,10 @@ def _merge_examples(examples, new_examples):
         else:
             examples[key] = {
                 "distance": new_distance,
-                "metadata": new_example.get("metadata")
+                "metadata": new_example.get("metadata"),
             }
 
+
 def _extract_metadata(metadata, distance):
     """Create result object from S3 vectors metadata"""
     # Result object attributes
@@ -215,6 +226,7 @@ def _extract_metadata(metadata, distance):
 
     return attributes
 
+
 def _get_image_data_from_s3_path(image_path):
     """
     Load images from image path
@@ -237,9 +249,7 @@ def _get_image_data_from_s3_path(image_path):
                 # Direct S3 URI
                 image_bytes = s3.get_binary_content(image_file_path)
             else:
-                raise ValueError(
-                    f"Invalid file path {image_path} - expecting S3 path"
-                )
+                raise ValueError(f"Invalid file path {image_path} - expecting S3 path")
 
             image_content.append(image_bytes)
         except Exception as e:
@@ -248,6 +258,7 @@ def _get_image_data_from_s3_path(image_path):
 
     return image_content
 
+
 def _get_image_files_from_s3_path(image_path):
     """
     Get list of image files from an S3 path.
@@ -260,9 +271,7 @@ def _get_image_files_from_s3_path(image_path):
     """
     # Handle S3 URIs
     if not image_path.startswith("s3://"):
-        raise ValueError(
-            f"Invalid file path {image_path} - expecting S3 URI"
-        )
+        raise ValueError(f"Invalid file path {image_path} - expecting S3 URI")
 
     # Check if it's a direct file or a prefix
     if image_path.endswith(