feat: add dynamic few-shot Lambda invocation to ExtractionService

Author: Daniel Lorch

lib/idp_common_pkg/idp_common/extraction/service.py

@@ -10,6 +10,7 @@
 
 from __future__ import annotations
 
+import base64
 import json
 import logging
 import os
@@ -286,7 +287,7 @@ def _build_prompt_content(
                 )
 
                 # Add few-shot examples
-                content.extend(self._build_few_shot_examples_content())
+                content.extend(self._build_few_shot_examples_content(image_content))
 
                 # Process after examples (only pass images if not already used)
                 image_for_after = (
@@ -379,24 +380,84 @@ def _prepare_image_attachments(self, image_content: Any) -> list[dict[str, Any]]
 
         return attachments
 
-    def _build_few_shot_examples_content(self) -> list[dict[str, Any]]:
+    def _build_few_shot_examples_content(
+        self, image_content: Any = None
+    ) -> list[dict[str, Any]]:
         """
         Build content items for few-shot examples from the configuration for a specific class.
 
+        Args:
+            image_content: Optional image content
+
         Returns:
             List of content items containing text and image content for examples
         """
         content: list[dict[str, Any]] = []
+        dynamic_few_shot_lambda_arn = self.config.extraction.dynamic_few_shot_lambda_arn
+
+        if dynamic_few_shot_lambda_arn and dynamic_few_shot_lambda_arn.strip():
+            logger.info(
+                f"Using Lambda to retrieve examples: {dynamic_few_shot_lambda_arn}"
+            )
+
+            # Create fully serializable payload using comprehensive helper
+            payload = {
+                "class_label": self._class_label,
+                "document_text": self._document_text,
+                "image_content": image_content,
+            }
 
-        # Use the stored class schema
-        if not self._class_schema:
-            logger.warning(
-                f"No class schema found for '{self._class_label}' for few-shot examples"
+            # Test JSON serialization before sending to Lambda to catch any remaining issues
+            try:
+                json.dumps(payload)
+                logger.info("Lambda payload successfully serialized")
+            except (TypeError, ValueError) as e:
+                logger.error(
+                    f"Lambda payload still contains non-serializable data: {e}"
+                )
+                logger.info("Using comprehensive serialization as fallback")
+                # Apply comprehensive serialization to entire payload
+                payload = self._make_json_serializable(payload)
+                try:
+                    json.dumps(payload)
+                    logger.info("Comprehensive serialization successful")
+                except (TypeError, ValueError) as e2:
+                    logger.error(f"Even comprehensive serialization failed: {e2}")
+                    # Ultimate fallback to minimal payload
+                    payload = {
+                        "class_label": self._class_label,
+                        "document_text": self._document_text,
+                    }
+
+            # Invoke dynamic few-shot Lambda and get result
+            examples = self._invoke_dynamic_few_shot_lambda(
+                dynamic_few_shot_lambda_arn, payload
             )
-            return content
 
-        # Get examples from the JSON Schema for this specific class
-        content = build_few_shot_extraction_examples_content(self._class_schema)
+            logger.info(f"Retrieved {len(examples)} examples from few-shot Lambda invocation")
+
+            # flatten examples content
+            for example in examples:
+                content.append({"text": example.get("attributes_prompt")})
+                for image_content in example.get("image_content", []):
+                    content.append(image.prepare_bedrock_image_attachment(image_content))
+
+            logger.info(content)
+
+        else:
+            logger.info(
+                f"No dynamic few-shot Lambda configured - add examples from config for class: {self._class_label}"
+            )
+
+            # Use the stored class schema
+            if not self._class_schema:
+                logger.warning(
+                    f"No class schema found for '{self._class_label}' for few-shot examples"
+                )
+                return content
+
+            # Get examples from the JSON Schema for this specific class
+            content = build_few_shot_extraction_examples_content(self._class_schema)
 
         return content
 
@@ -427,8 +488,8 @@ def _make_json_serializable(self, obj: Any) -> Any:
             # Handle objects with to_dict method
             return self._make_json_serializable(obj.to_dict())
         elif isinstance(obj, bytes):
-            # Convert bytes to base64 string or placeholder
-            return f"<bytes_object_{len(obj)}_bytes>"
+            # Convert bytes to base64 string
+            return base64.b64encode(obj).decode("utf-8")
         else:
             try:
                 # Test if it's already JSON serializable
@@ -498,6 +559,65 @@ def _invoke_custom_prompt_lambda(
             logger.error(error_msg)
             raise Exception(error_msg)
 
+    def _invoke_dynamic_few_shot_lambda(
+        self, lambda_arn: str, payload: dict
+    ) -> list[dict[str, Any]]:
+        """
+        Invoke dynamic few-shot Lambda function with JSON-serializable payload.
+
+        Args:
+            lambda_arn: ARN of the Lambda function to invoke
+            payload: Payload to send to Lambda function (must be JSON serializable)
+
+        Returns:
+            Dict containing example content with images and expected attributes prompt
+
+        Raises:
+            Exception: If Lambda invocation fails or returns invalid response
+        """
+        import boto3
+
+        lambda_client = boto3.client("lambda", region_name=self.region)
+
+        try:
+            logger.info(f"Invoking dynamic few-shot Lambda: {lambda_arn}")
+            response = lambda_client.invoke(
+                FunctionName=lambda_arn,
+                InvocationType="RequestResponse",
+                Payload=json.dumps(payload),
+            )
+
+            if response.get("FunctionError"):
+                error_payload = response.get("Payload", b"").read().decode()
+                error_msg = f"Dynamic few-shot Lambda failed: {error_payload}"
+                logger.error(error_msg)
+                raise Exception(error_msg)
+
+            result = json.loads(response["Payload"].read())
+            logger.info("Dynamic few-shot Lambda invoked successfully")
+
+            # Validate response structure
+            if not isinstance(result, list):
+                error_msg = f"Dynamic few-shot Lambda returned invalid response format: expected list, got {type(result)}"
+                logger.error(error_msg)
+                raise Exception(error_msg)
+
+            # Base64 decode images
+            for example in result:
+                decoded_images = []
+                for image_base64 in example.get("image_content", []):
+                    image_data = base64.b64decode(image_base64)
+                    decoded_images.append(image_data)
+
+                example["image_content"] = decoded_images
+
+            return result
+
+        except Exception as e:
+            error_msg = f"Failed to invoke custom prompt Lambda {lambda_arn}: {str(e)}"
+            logger.error(error_msg)
+            raise Exception(error_msg)
+
     def _reset_context(self) -> None:
         """Reset instance variables for clean state before processing."""
         self._document_text = ""
@@ -1078,7 +1198,6 @@ def process_document_section(self, document: Document, section_id: str) -> Docum
             content, system_prompt = self._build_extraction_content(
                 document, page_images
             )
-
             # Invoke model
             result = self._invoke_extraction_model(content, system_prompt, section_info)