From f3f4592206b467930b917eecba2c9bfe1199d3a8 Mon Sep 17 00:00:00 2001
From: Saeed Kasmani <amirgholipour@users.noreply.github.com>
Date: Sun, 9 Aug 2026 21:08:28 +1000
Subject: [PATCH] Improve image enhancer reliability

---
 deployment/lambdas/containers/image_enhancer/DEPLOYMENT.md | 4 ++--
 deployment/lambdas/containers/image_enhancer/Dockerfile    | 4 ++++
 deployment/lambdas/containers/image_enhancer/README.md     | 2 +-
 .../lambdas/containers/image_enhancer/agentic_enhancer.py  | 2 +-
 .../lambdas/containers/image_enhancer/lambda_handler.py    | 7 ++++++-
 5 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/deployment/lambdas/containers/image_enhancer/DEPLOYMENT.md b/deployment/lambdas/containers/image_enhancer/DEPLOYMENT.md
index 0eb91b4..4caeb69 100644
--- a/deployment/lambdas/containers/image_enhancer/DEPLOYMENT.md
+++ b/deployment/lambdas/containers/image_enhancer/DEPLOYMENT.md
@@ -78,7 +78,7 @@ Set in `lambda_stack.py` or override in AWS console:
 |----------|---------|---------|
 | `VISION_MODEL` | `us.anthropic.claude-sonnet-4-6` | Bedrock model ID |
 | `MAX_ITERATIONS` | `2` | Max agent iterations |
-| `MAX_IMAGE_DIMENSION` | `4000` | Max dimension for LLM |
+| `MAX_IMAGE_DIMENSION` | `2048` | Max dimension for LLM |
 | `JPEG_QUALITY` | `85` | LLM image encoding quality |
 | `OUTPUT_QUALITY` | `95` | Final output quality |
 | `OUTPUT_BUCKET` | (from CDK) | S3 bucket for enhanced images |
@@ -343,7 +343,7 @@ aws lambda update-function-configuration \
   --environment Variables="{
     VISION_MODEL=us.anthropic.claude-sonnet-4-6,
     MAX_ITERATIONS=3,
-    MAX_IMAGE_DIMENSION=4000
+    MAX_IMAGE_DIMENSION=2048
   }"
 ```
 
diff --git a/deployment/lambdas/containers/image_enhancer/Dockerfile b/deployment/lambdas/containers/image_enhancer/Dockerfile
index 70e8e29..338eb20 100644
--- a/deployment/lambdas/containers/image_enhancer/Dockerfile
+++ b/deployment/lambdas/containers/image_enhancer/Dockerfile
@@ -28,6 +28,10 @@ COPY lambda_handler.py ${LAMBDA_TASK_ROOT}/
 # into the build context from ./layer/python/foundation before docker build.
 COPY foundation/ ${LAMBDA_TASK_ROOT}/foundation/
 
+# Foundation imports the shared config package during initialization. The build
+# scripts stage it alongside foundation, so include it in the runtime image too.
+COPY config/ ${LAMBDA_TASK_ROOT}/config/
+
 # Fix file permissions
 RUN chmod 644 ${LAMBDA_TASK_ROOT}/*.py
 
diff --git a/deployment/lambdas/containers/image_enhancer/README.md b/deployment/lambdas/containers/image_enhancer/README.md
index 4dda1bd..7d3c697 100644
--- a/deployment/lambdas/containers/image_enhancer/README.md
+++ b/deployment/lambdas/containers/image_enhancer/README.md
@@ -154,7 +154,7 @@ Maps to MAX_ITERATIONS for the agent:
 |----------|---------|-------------|
 | `VISION_MODEL` | `us.anthropic.claude-sonnet-4-6` | Bedrock model ID |
 | `MAX_ITERATIONS` | `2` | Max agent iterations (overridden by enhancement_level at runtime) |
-| `MAX_IMAGE_DIMENSION` | `4000` | Max dimension for LLM submission |
+| `MAX_IMAGE_DIMENSION` | `2048` | Max dimension for LLM submission |
 | `JPEG_QUALITY` | `85` | Quality for LLM image encoding |
 | `OUTPUT_QUALITY` | `95` | Quality for final output |
 | `OUTPUT_BUCKET` | - | S3 bucket for enhanced images (if not set, returns base64) |
diff --git a/deployment/lambdas/containers/image_enhancer/agentic_enhancer.py b/deployment/lambdas/containers/image_enhancer/agentic_enhancer.py
index ebe313c..586fdb1 100644
--- a/deployment/lambdas/containers/image_enhancer/agentic_enhancer.py
+++ b/deployment/lambdas/containers/image_enhancer/agentic_enhancer.py
@@ -39,7 +39,7 @@
     "CLAUDE_OPUS_46_PROFILE_ARN", "us.anthropic.claude-opus-4-6-v1"
 )
 MAX_ITERATIONS = int(os.environ.get("MAX_ITERATIONS", "2"))
-MAX_IMAGE_DIMENSION = int(os.environ.get("MAX_IMAGE_DIMENSION", "4000"))
+MAX_IMAGE_DIMENSION = int(os.environ.get("MAX_IMAGE_DIMENSION", "2048"))
 JPEG_QUALITY = int(os.environ.get("JPEG_QUALITY", "85"))
 OUTPUT_QUALITY = int(os.environ.get("OUTPUT_QUALITY", "95"))
 
diff --git a/deployment/lambdas/containers/image_enhancer/lambda_handler.py b/deployment/lambdas/containers/image_enhancer/lambda_handler.py
index 20a9a7d..00d2d28 100644
--- a/deployment/lambdas/containers/image_enhancer/lambda_handler.py
+++ b/deployment/lambdas/containers/image_enhancer/lambda_handler.py
@@ -278,7 +278,12 @@ def _upload_to_s3(local_path: str, bucket: str, session_id: str, original: str)
     timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
     output_key = f"{session_id}/enhanced/{original_name}_enhanced_{timestamp}.jpg"
 
-    s3.upload_file(local_path, bucket, output_key)
+    s3.upload_file(
+        local_path,
+        bucket,
+        output_key,
+        ExtraArgs={"ContentType": "image/jpeg"},
+    )
     return f"s3://{bucket}/{output_key}"