encapsulate ruler

Author: kazmer97

lib/idp_common_pkg/idp_common/assessment/granular_service.py

@@ -36,7 +36,6 @@
 from idp_common.extraction.models import ExtractionData
 from idp_common.models import Document, Status
 from idp_common.utils import check_token_limit
-from idp_common.utils.grid_overlay import add_ruler_edges
 
 logger = Logger(service="assessment", level=os.getenv("LOG_LEVEL", "INFO"))
 
@@ -845,28 +844,21 @@ def process_document_section(self, document: Document, section_id: str) -> Docum
                     f"Found {len(cached_task_results)} cached assessment task results, processing {len(tasks_to_process)} remaining tasks"
                 )
 
-                # Apply grid overlay to page images for assessment
-                grid_page_images = []
-                for idx, page_img in enumerate(page_images):
-                    grid_img = add_ruler_edges(page_img)
-                    logger.info(
-                        f"Added ruler overlay to page {idx}: {len(page_img):,} bytes -> {len(grid_img):,} bytes"
-                    )
-                    grid_page_images.append(grid_img)
-
                 # Execute tasks using Strands-based parallel executor
+                # Note: ruler overlay is added internally by strands_service/strands_tools
                 logger.info(
                     f"Processing {len(tasks_to_process)} assessment tasks with Strands executor (max_concurrent={self.max_workers})"
                 )
 
                 request_start_time = time.time()
 
                 # Call Strands executor - handles both parallel and sequential based on max_concurrent
+                # Pass raw page_images - ruler overlay is added internally when needed
                 task_results, task_metering, processing_time = (
                     execute_assessment_tasks_parallel(
                         tasks=tasks_to_process,
                         extraction_results=extraction_results,
-                        page_images=grid_page_images,
+                        page_images=page_images,
                         sorted_page_ids=sorted_page_ids,
                         model_id=self.config.assessment.model,
                         system_prompt=self.config.assessment.system_prompt,

lib/idp_common_pkg/idp_common/assessment/models.py

@@ -114,7 +114,7 @@ def from_corners(
         scale: float = 1000.0,
     ) -> "BoundingBoxCoordinates":
         """
-        Create from corner coordinates in document space.
+        Create from corner coordinates.
 
         Args:
             x1, y1: Top-left corner in 0-scale range
@@ -228,7 +228,16 @@ def from_llm_response(
         bbox_coords: list[float] | None = None,
         page_num: int | None = None,
     ) -> "FieldAssessmentData":
-        """Create from LLM response data."""
+        """
+        Create from LLM response data.
+
+        Args:
+            confidence: Confidence score (0-1)
+            reasoning: Explanation for confidence
+            confidence_threshold: Threshold for confidence alerts
+            bbox_coords: Optional bounding box coordinates [x1, y1, x2, y2]
+            page_num: Optional page number (1-indexed)
+        """
         geometry = None
         if bbox_coords is not None and page_num is not None:
             geom = Geometry.from_bbox_list(bbox_coords, page_num)

lib/idp_common_pkg/idp_common/assessment/strands_executor.py

@@ -40,7 +40,7 @@ async def execute_tasks_async(
     Args:
         tasks: List of assessment tasks to execute
         extraction_results: Full extraction results
-        page_images: List of page images (with grid overlay)
+        page_images: List of raw page images (ruler overlay added internally)
         sorted_page_ids: List of page IDs
         model_id: Model to use
         system_prompt: System prompt
@@ -166,7 +166,7 @@ def execute_assessment_tasks_parallel(
     Args:
         tasks: List of assessment tasks
         extraction_results: Full extraction results
-        page_images: List of page images (with grid overlay already applied)
+        page_images: List of raw page images (ruler overlay added internally)
         sorted_page_ids: List of page IDs in sorted order
         model_id: Model ID
         system_prompt: System prompt

lib/idp_common_pkg/idp_common/assessment/strands_service.py

@@ -26,6 +26,7 @@
 from idp_common.assessment.strands_tools import create_strands_tools
 from idp_common.bedrock import build_model_config
 from idp_common.utils.bedrock_utils import async_exponential_backoff_retry
+from idp_common.utils.grid_overlay import add_ruler_edges
 
 logger = Logger(service="assessment", level=os.getenv("LOG_LEVEL", "INFO"))
 
@@ -49,9 +50,8 @@ async def assess_attribute_with_strands(
 
     Args:
         task: Assessment task to process
-        base_content: Base prompt content (includes images)
         extraction_results: Full extraction results
-        page_images: List of page images (with grid overlay already applied)
+        page_images: List of raw page images (ruler overlay added internally)
         sorted_page_ids: List of page IDs in sorted order
         model_id: Bedrock model ID
         system_prompt: System prompt for assessment
@@ -364,16 +364,21 @@ def _build_task_prompt(
 
     Args:
         task: Assessment task for one specific field
-        page_images: List of page images to include in the prompt
+        page_images: List of raw page images (ruler will be added here)
 
     Returns:
         List of content blocks with images and task text
     """
     field_path_str = _convert_field_path_to_string(task.field_path)
 
-    # Create image content blocks
+    # Create image content blocks with ruler overlay
+    # Rulers are added here so the LLM can see coordinate reference marks
     image_blocks = [
-        ContentBlock(image=ImageContent(format="png", source=ImageSource(bytes=img)))
+        ContentBlock(
+            image=ImageContent(
+                format="png", source=ImageSource(bytes=add_ruler_edges(img))
+            )
+        )
         for img in page_images
     ]
 
@@ -414,6 +419,8 @@ def _convert_to_assessment_result(
     assessment = output.assessment
 
     # Create standardized field assessment data
+    # Note: bounding box coordinates are already adjusted for ruler offset
+    # by the submit_assessment tool in strands_tools.py
     field_data = FieldAssessmentData.from_llm_response(
         confidence=assessment.confidence,
         reasoning=assessment.reasoning,

lib/idp_common_pkg/idp_common/assessment/strands_tools.py

@@ -13,7 +13,7 @@
 from strands import Agent, tool
 
 from idp_common.assessment.strands_models import AssessmentOutput, BoundingBox
-from idp_common.utils.grid_overlay import draw_bounding_boxes
+from idp_common.utils.grid_overlay import add_ruler_and_draw_boxes, add_ruler_edges
 from idp_common.utils.strands_agent_tools.todo_list import (
     create_todo_list,
     update_todo,
@@ -37,10 +37,13 @@ class ViewImageInput(BaseModel):
 
 @tool
 def submit_assessment(assessment: AssessmentOutput, agent: Agent) -> str:
+    """Submit the final assessment for a field."""
     # Validate assessment structure and return helpful errors
     validated_assessment = AssessmentOutput.model_validate(assessment)
 
-    # Store in agent state
+    # Store in agent state - coordinates are in 0-1000 document space
+    # The ruler shows 0-1000 scale mapping to the document, so LLM coordinates
+    # are already in document space and need no adjustment
     agent.state.set("assessment_output", validated_assessment.model_dump(mode="json"))
 
     logger.info(
@@ -56,15 +59,15 @@ def create_view_image_tool(page_images: list[bytes], sorted_page_ids: list[str])
     Create a view_image tool that has access to page images.
 
     Args:
-        page_images: List of page image bytes (with grid overlay already applied)
+        page_images: List of raw page image bytes (without ruler overlay)
         sorted_page_ids: List of page IDs in sorted order
 
     Returns:
         A Strands tool function for viewing images
     """
 
     @tool
-    def view_image(input_data: ViewImageInput, agent: Agent) -> dict:
+    def view_image(input_data: ViewImageInput, agent: Agent) -> dict[str, Any]:
         """
         View a specific page image, optionally highlighting a bounding box area.
 
@@ -97,13 +100,13 @@ def view_image(input_data: ViewImageInput, agent: Agent) -> dict:
                 f"Valid range: 0-{len(page_images) - 1}"
             )
 
-        # Get the base image (already has grid overlay)
-        img_bytes = page_images[view_input.image_index]
+        # Get the raw image (no ruler overlay yet)
+        raw_img_bytes = page_images[view_input.image_index]
         page_id = sorted_page_ids[view_input.image_index]
 
-        # If bounding box is specified, draw it on the image
+        # Add ruler and optionally draw bounding box
         if view_input.bounding_box:
-            # Convert BoundingBox to dict format for draw_bounding_boxes
+            # Convert BoundingBox to dict format
             bbox_dict = {
                 "bbox": [
                     view_input.bounding_box.x1,
@@ -115,21 +118,19 @@ def view_image(input_data: ViewImageInput, agent: Agent) -> dict:
                 "color": "red",
             }
 
-            # Draw the bounding box on the image (which has 30px margin for ruler)
-            # Let drawing errors propagate - if we can't draw, something is wrong
-            img_bytes = draw_bounding_boxes(
-                img_bytes,
-                [bbox_dict],
-                margin_offset=30,
-            )
+            # Add ruler overlay and draw bounding box in one step
+            img_bytes = add_ruler_and_draw_boxes(raw_img_bytes, [bbox_dict])
 
             logger.debug(
-                "Drew bounding box on image",
+                "Added ruler and drew bounding box on image",
                 extra={
                     "image_index": view_input.image_index,
                     "bbox": bbox_dict["bbox"],
                 },
             )
+        else:
+            # Just add ruler overlay (no bounding box)
+            img_bytes = add_ruler_edges(raw_img_bytes)
 
         logger.info(
             "Returning image to agent",