chore: use working bucket from GenAIIDP for dataset + adapt threshold

Author: Daniel Lorch

plugins/dynamic-few-shot-lambda/README.md

@@ -1,64 +1,77 @@
-# Dynamic-Few Shot Prompting - Complete Guide
+# Dynamic Few-Shot Prompting Lambda - Complete Guide
 
-This directory contains the **complete implementation and demonstration** of the dynamic-few shot prompting feature for GenAI IDP Accelerator. This feature enables users to dynamically retrieve few-shot examples using S3 Vectors similarity search to improve extraction accuracy for Pattern 2.
+This directory contains the **complete implementation** of the dynamic few-shot prompting Lambda function for GenAI IDP Accelerator. This Lambda function integrates with Pattern 2 extraction as a custom prompt generator, dynamically retrieving similar examples using S3 Vectors similarity search to improve extraction accuracy.
 
 ## 🎯 Overview
 
-The dynamic-few shot prompting feature allows you to:
+The dynamic few-shot prompting Lambda function allows you to:
 
 - **Dynamically retrieve similar examples** based on document content using vector similarity search
-- **Provide few-shot examples** to improve extraction accuracy through example-based prompting
+- **Automatically inject few-shot examples** into extraction prompts using the `{FEW_SHOT_EXAMPLES}` placeholder
 - **Leverage S3 Vectors** for efficient similarity search across large example datasets
 - **Integrate multimodal embeddings** using Amazon Nova models for image-based similarity
-- **Customize example selection** based on document characteristics and business rules
+- **Seamlessly integrate** with existing IDP extraction workflows as a custom prompt Lambda
 
 ## 📁 Files in This Directory
 
-- **`GENAIIDP-dynamic-few-shot.py`** - Dynamic few-shot Lambda function with S3 Vectors lookup
-- **`template.yml`** - CloudFormation SAM template to deploy the complete stack
-- **`requirements.txt`** - Python dependencies for the Lambda function
+- **`src/GENAIIDP-dynamic-few-shot.py`** - Dynamic few-shot Lambda function with S3 Vectors lookup
+- **`src/requirements.txt`** - Python dependencies for the Lambda function
+- **`template.yml`** - CloudFormation SAM template to deploy the Lambda function
 - **`README.md`** - This comprehensive documentation and guide
 
 ## 🏗️ Architecture
 
 ```mermaid
 flowchart TD
-    A[Document Processing] --> B{Dynamic-few shot configured?}
-    B -->|No| C[Use Default Extraction]
-    B -->|Yes| D[Invoke Dynamic-few shot Lambda]
-
-    subgraph Lambda
-        D --> E[Receive Document Images]
-        E --> F[Generate Embeddings with Nova]
-        F --> G[Query S3 Vectors Index]
-        G --> H[Retrieve Similar Examples]
-        H --> I[Load Example Images from S3]
-        I --> J[Format Examples for Bedrock]
+    A[IDP Document Processing] --> B{Custom Prompt Lambda ARN configured?}
+    B -->|No| C[Use Default Task Prompt]
+    B -->|Yes| D[Invoke Dynamic Few-Shot Lambda]
+
+    subgraph "Lambda Function: GENAIIDP-dynamic-few-shot"
+        D --> E[Receive IDP Context & Placeholders]
+        E --> F[Extract Document Images from DOCUMENT_IMAGE]
+        F --> G[Generate Nova Multimodal Embeddings]
+        G --> H[Query S3 Vectors Index]
+        H --> I[Filter by Distance Threshold]
+        I --> J[Merge & Deduplicate Results]
+        J --> K[Load Example Images from S3]
+        K --> L[Build Prompt Content Array]
+        L --> M[Replace FEW_SHOT_EXAMPLES Placeholder]
     end
 
-    J --> K[Use Examples in Extraction Prompt]
-    C --> L[Continue with Standard Extraction]
-    K --> L
+    M --> N[Return Modified Task Prompt Content]
+    C --> O[Continue with Bedrock Extraction]
+    N --> O
 
-    subgraph Input
-        M[Document Class]
-        N[Document Text]
-        O[Document Images]
+    subgraph "Input Payload"
+        P[config: IDP Configuration]
+        Q[prompt_placeholders: DOCUMENT_TEXT, DOCUMENT_CLASS, etc.]
+        R[default_task_prompt_content: Original prompt]
+        S[serialized_document: Document metadata]
     end
 
-    subgraph Output
-        P[Example Attributes Prompts]
-        Q[Example Images]
-        R[Similarity Distances]
+    subgraph "Output Payload"
+        T[system_prompt: Unchanged]
+        U[task_prompt_content: Array with Prompt segments and Example images]
     end
 
-    D -.-> M
-    D -.-> N
-    D -.-> O
+    D -.-> P
+    D -.-> Q
+    D -.-> R
+    D -.-> S
 
-    J -.-> P
-    J -.-> Q
-    J -.-> R
+    N -.-> T
+    N -.-> U
+
+    subgraph "S3 Vectors Infrastructure"
+        X[Vector Bucket: Encrypted storage]
+        Y[Vector Index: 3072-dim cosine similarity]
+        Z[Metadata: classPrompt, attributesPrompt, imagePath]
+    end
+
+    H -.-> X
+    H -.-> Y
+    H -.-> Z
 ```
 
 ## Quick Start
@@ -88,7 +101,7 @@ aws cloudformation describe-stacks \
 
 Use the [fewshot_dataset_import.ipynb](notebooks/fewshot_dataset_import.ipynb) notebook to import a dataset into S3 Vectors, or manually upload your example documents and metadata to the S3 bucket and vector index created by the stack.
 
-### Step 4: Configure IDP to Use Dynamic-few shot
+### Step 4: Configure IDP to Use Dynamic Few-Shot
 
 Add the Lambda ARN to your IDP extraction configuration:
 
@@ -97,42 +110,81 @@ extraction:
   custom_prompt_lambda_arn: "arn:aws:lambda:region:account:function:GENAIIDP-dynamic-few-shot"
 ```
 
+**Important**: Your extraction task prompt must include the `{FEW_SHOT_EXAMPLES}` placeholder where you want the dynamic examples to be inserted.
+
+### Step 5: Run the Demo Notebook
+
+0. Run `notebooks/examples` steps 0, 1, 2
+1. Open `plugins/dynamic-few-shot-lambda/notebooks/step3_extraction_with_custom_lambda.ipynb`
+2. Run all cells to see the comparison
+
 ## Lambda Interface
 
 ### Input Payload Structure
+
+The Lambda receives the full IDP context as a custom prompt Lambda:
+
 ```json
 {
-  "class_label": "invoice",
-  "document_text": "Text or markdown from section 1 (pages 1-3)...",
-  "image_content": [
-    "base64_encoded_image_1",
-    "base64_encoded_image_2"
-  ]
+  "config": {
+    "extraction": {...},
+    "classes": [...],
+    ...
+  },
+  "prompt_placeholders": {
+    "DOCUMENT_TEXT": "Full OCR text from all pages",
+    "DOCUMENT_CLASS": "invoice", 
+    "ATTRIBUTE_NAMES_AND_DESCRIPTIONS": "LineItems: List of line items in the invoice...",
+    "DOCUMENT_IMAGE": ["s3://bucket/document/page1.jpg", "s3://bucket/document/page2.jpg"]
+  },
+  "default_task_prompt_content": [
+    {"text": "Resolved default task prompt..."},
+    {"image_uri": "s3://..."}, // if images present
+    {"cachePoint": true} // if cache points present
+  ],
+  "serialized_document": {
+    "id": "document-123",
+    "input_bucket": "my-bucket",
+    "pages": {...},
+    "sections": [...],
+    ...
+  }
 }
 ```
 
 ### Output Payload Structure
+
+The Lambda returns modified prompt content with dynamic few-shot examples:
+
 ```json
-[
-  {
-    "attributes_prompt": "Expected attributes are: invoice_number [Unique identifier], invoice_date [Invoice date], total_amount [Total amount]...",
-    "class_prompt": "This is an example of the class 'invoice'",
-    "distance": 0.122344521145, # lower is more similar
-    "image_content": ["<base64_image_content_1>", "<base64_image_content_2>", ...]
-  }
-]
+{
+  "system_prompt": "Custom system prompt text",
+  "task_prompt_content": [
+    {"text": "Extract the following attributes from this invoice document:\n\nLineItems: List of line items in the invoice...\n\n<few_shot_examples>"},
+    {"text": "expected attributes are:\n    \"invoice_number\": \"INV-2024-001\",\n    \"total_amount\": \"$1,250.00\""},
+    {"image_uri": "s3://examples-bucket/invoices/example-001/page1.jpg"},
+    {"text": "</few_shot_examples>\n\n<<CACHEPOINT>>\n\nDocument content:\nINVOICE\nInvoice #: INV-2024-002..."}
+  ]
+}
 ```
 
 ## Core Functionality
 
-### 1. Vector Similarity Search
+### 1. Custom Prompt Integration
+
+The Lambda integrates with IDP's custom prompt system by:
+- Receiving the full extraction context and configuration
+- Processing the `{FEW_SHOT_EXAMPLES}` placeholder in task prompts
+- Returning modified prompt content with dynamically retrieved examples
+
+### 2. Vector Similarity Search
 
 The Lambda uses Amazon Nova multimodal embeddings to find similar examples:
 
 ```python
 # Generate embedding from document image
-embedding = bedrock.generate_embedding(
-    image_source=image_data,
+embedding = bedrock_client.generate_embedding(
+    image_source=page_image,
     model_id=MODEL_ID,
     dimensions=S3VECTOR_DIMENSIONS,
 )
@@ -148,34 +200,36 @@ response = s3vectors.query_vectors(
 )
 ```
 
-### 2. Example Merging and Deduplication
+### 3. Example Merging and Deduplication
 
 Multiple document images are processed and results are merged to avoid duplicates:
 
 ```python
-def merge_examples(combined_examples, new_examples):
+def _merge_examples(examples, new_examples):
     """Merge examples, keeping the best similarity score for duplicates"""
     for new_example in new_examples:
         key = new_example["key"]
-        if combined_examples.get(key):
-            # Keep the better (lower) distance score
-            combined_examples[key]["distance"] = min(
-                new_example.get("distance"),
-                combined_examples[key]["distance"]
-            )
+        new_distance = new_example.get("distance", 1.0)
+        
+        if examples.get(key):
+            existing_distance = examples[key].get("distance", 1.0)
+            examples[key]["distance"] = min(new_distance, existing_distance)
 ```
 
-### 3. Example Image Loading
+### 4. Prompt Content Building
 
-The Lambda loads example images from S3 paths stored in vector metadata:
+The Lambda builds structured prompt content handling multiple placeholders:
 
 ```python
-def get_image_files_from_s3_path(image_path: str) -> List[str]:
-    """Get list of image files from S3 path or prefix"""
-    if image_path.endswith((".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".tif", ".webp")):
-        return [image_path]  # Direct file
-    else:
-        return s3.list_images_from_path(image_path)  # Directory/prefix
+def _build_prompt_content(prompt_template, substitutions, image_content):
+    """
+    Build prompt content array handling FEW_SHOT_EXAMPLES and DOCUMENT_IMAGE placeholders.
+    
+    Handles:
+    - {FEW_SHOT_EXAMPLES}: Inserts few-shot examples from S3 Vectors
+    - {DOCUMENT_IMAGE}: Inserts images at specific location
+    - Regular text placeholders: DOCUMENT_TEXT, DOCUMENT_CLASS, etc.
+    """
 ```
 
 ## Configuration
@@ -188,7 +242,9 @@ The Lambda function uses these environment variables (set by the CloudFormation
 - `S3VECTOR_INDEX` - Name of the S3 Vectors index
 - `S3VECTOR_DIMENSIONS` - Embedding dimensions (e.g. `3072` for Nova Multimodal Embedding model)
 - `MODEL_ID` - Bedrock model ID for embeddings (e.g. `amazon.nova-2-multimodal-embeddings-v1:0`)
-- `TOP_K` - Number of similar examples to retrieve
+- `TOP_K` - Number of similar examples to retrieve (default: 3)
+- `THRESHOLD` - Maximum distance threshold for filtering results (default: 0.5)
+- `LOG_LEVEL` - Logging level (default: INFO)
 
 ### S3 Vectors Configuration
 
@@ -208,16 +264,22 @@ Monitor the Lambda function logs:
 
 **Successful Operation:**
 ```
-Processing document ID: document-123
-Document class: invoice
-Response contains 2 elements
+=== DYNAMIC FEW-SHOT LAMBDA INVOKED ===
+=== EXTRACTION CONFIG ===
+Model: anthropic.claude-3-5-sonnet-20241022-v2:0
+=== HANDLE INPUT DOCUMENT ===
+=== OUTPUT ANALYSIS ===
+Output keys: ['system_prompt', 'task_prompt_content']
+Task prompt content items: 5
+=== DYNAMIC FEW-SHOT LAMBDA COMPLETED ===
 ```
 
 **Error Conditions:**
 ```
-No class_label found in event
-No document_texts found in event or not in list format
-Failed to load example images from s3://bucket/path: error
+Failed to parse environment variables: KeyError('S3VECTOR_BUCKET')
+Skipping example with empty attributesPrompt: example_key
+Skipping example with distance 0.8 above threshold 0.5: example_key
+Invalid file path /local/path - expecting S3 URI
 ```
 
 ### Performance Monitoring
@@ -331,22 +393,53 @@ aws cloudformation delete-stack --stack-name GENAIIDP-dynamic-few-shot-stack
 
 ### Configuration in IDP Stack
 
-Add the dynamic-few shot Lambda ARN to your IDP configuration:
+Add the dynamic few-shot Lambda ARN to your IDP extraction configuration:
 
 ```yaml
-# In your IDP stack parameters or configuration
 extraction:
-  dynamic_few_shot_lambda_arn: "arn:aws:lambda:region:account:function:GENAIIDP-dynamic-few-shot"
+  custom_prompt_lambda_arn: "arn:aws:lambda:region:account:function:GENAIIDP-dynamic-few-shot"
 ```
 
+### Required Task Prompt Configuration
+
+**Critical**: Your extraction task prompt must include the `{FEW_SHOT_EXAMPLES}` placeholder where you want the dynamic examples to be inserted. The Lambda specifically looks for this placeholder and replaces it with retrieved examples.
+
 ### Expected Behavior
 
 When configured:
 1. IDP processes document and extracts images/text
-2. Dynamic few-shot Lambda is invoked with document data
-3. Lambda returns similar examples with prompts and images
-4. IDP includes examples in extraction prompt to Bedrock
-5. Bedrock uses examples to improve extraction accuracy
+2. IDP invokes the dynamic few-shot Lambda with full extraction context
+3. Lambda generates embeddings from document images using Amazon Nova
+4. Lambda queries S3 Vectors to find similar examples
+5. Lambda loads example images and metadata from S3
+6. Lambda builds modified prompt content with examples inserted at `{FEW_SHOT_EXAMPLES}` location
+7. IDP uses the modified prompt content for Bedrock extraction
+8. Bedrock uses the dynamic examples to improve extraction accuracy
+
+### Prompt Flow Example
+
+**Original Task Prompt:**
+```
+Extract attributes from this invoice:
+{ATTRIBUTE_NAMES_AND_DESCRIPTIONS}
+{FEW_SHOT_EXAMPLES}
+<<CACHEPOINT>>
+Document: {DOCUMENT_TEXT}
+```
+
+**After Lambda Processing:**
+```
+Extract attributes from this invoice:
+invoice_number [Unique identifier]...
+
+expected attributes are:
+    "invoice_number": "INV-2024-001",
+    "total_amount": "$1,250.00"
+[Example image content]
+
+<<CACHEPOINT>>
+Document: INVOICE #INV-2024-002...
+```
 
 ## Next Steps