cleanup

Author: kazmer97

lib/idp_common_pkg/idp_common/assessment/example_usage.py

@@ -142,14 +142,15 @@ def example_granular_assessment():
 
     # Demonstrate task creation (this would normally be done internally)
     if hasattr(assessment_service, "_create_assessment_tasks"):
-        attributes = assessment_service._get_class_attributes("Bank Statement")
-        tasks = assessment_service._create_assessment_tasks(
-            example_extraction_results, attributes, 0.9
+        class_schema = assessment_service._get_class_schema("Bank Statement")
+        properties = class_schema.get("properties", {})
+        tasks, assessment_structure = assessment_service._create_assessment_tasks(
+            example_extraction_results, properties, 0.9
         )
 
         logger.info(f"\nCreated {len(tasks)} assessment tasks:")
         for task in tasks:
-            logger.info(f"  - {task.task_id}: {task.task_type} for {task.attributes}")
+            logger.info(f"  - {task.task_id}: {task.task_type} for {task.field_name}")
 
     return assessment_service, config
 

lib/idp_common_pkg/idp_common/assessment/granular_service.py

@@ -309,7 +309,6 @@ def _traverse(
                                 field_name=_convert_field_path_to_string(item_path),
                                 field_schema=items_schema,
                                 confidence_threshold=threshold,
-                                parent_assessment_dict=assessment_list,  # type: ignore
                             )
                             tasks.append(task)
                             task_counter[0] += 1
@@ -332,7 +331,6 @@ def _traverse(
                         field_name=_convert_field_path_to_string(field_path),
                         field_schema=prop_schema,
                         confidence_threshold=threshold,
-                        parent_assessment_dict=parent_dict,
                     )
                     tasks.append(task)
                     task_counter[0] += 1
@@ -536,14 +534,40 @@ def _is_throttling_exception(self, exception: Exception) -> bool:
             for throttle_term in self.throttling_exceptions
         )
 
+    def _insert_at_field_path(
+        self,
+        structure: dict[str, Any],
+        field_path: tuple[str | int, ...],
+        value: Any,
+    ) -> None:
+        """
+        Navigate through structure using field_path and insert value at the end.
+
+        Args:
+            structure: The assessment structure to navigate
+            field_path: Tuple path like ("Account Holder Address", "City") or ("Transactions", 0, "Amount")
+            value: The assessment data to insert
+
+        Example:
+            field_path = ("Account Holder Address", "City")
+            -> structure["Account Holder Address"]["City"] = value
+
+            field_path = ("Transactions", 0, "Amount")
+            -> structure["Transactions"][0]["Amount"] = value
+        """
+        parent = structure
+        for key in field_path[:-1]:
+            parent = parent[key]
+        parent[field_path[-1]] = value
+
     def _aggregate_assessment_results(
         self,
         tasks: list[AssessmentTask],
         results: list[AssessmentResult],
         assessment_structure: dict[str, Any],
     ) -> tuple[dict[str, Any], list[dict[str, Any]], dict[str, Any]]:
         """
-        Aggregate individual task results into assessment structure using direct parent insertion.
+        Aggregate individual task results into assessment structure using field_path navigation.
 
         Args:
             tasks: list of assessment tasks
@@ -553,13 +577,12 @@ def _aggregate_assessment_results(
         Returns:
             Tuple of (assessment_structure, confidence_alerts, aggregated_metering)
         """
-        all_confidence_alerts = []
-        aggregated_metering = {}
+        all_confidence_alerts: list[dict[str, Any]] = []
+        aggregated_metering: dict[str, Any] = {}
 
         # Create a mapping from task_id to result
         result_map = {result.task_id: result for result in results}
 
-        # Process each task result - direct O(1) insertion using parent reference
         for task in tasks:
             result = result_map.get(task.task_id)
             if not result or not result.success:
@@ -575,45 +598,31 @@ def _aggregate_assessment_results(
             # Add confidence alerts
             all_confidence_alerts.extend(result.confidence_alerts)
 
-            # Get assessment data from result - should be a single assessment object
-            # The Strands agent returns the assessment in result.assessment_data
-            assessment_obj = result.assessment_data
+            # Get assessment data directly from result
+            # strands_service returns flat assessment dict: {confidence, value, reasoning, ...}
+            field_assessment = result.assessment_data
 
-            if not isinstance(assessment_obj, dict):
+            if not isinstance(field_assessment, dict):
                 logger.warning(
-                    f"Task {task.task_id}: expected dict assessment, got {type(assessment_obj)}"
+                    f"Task {task.task_id}: expected dict assessment, got {type(field_assessment)}"
                 )
                 continue
 
-            # Add confidence_threshold to the assessment object
-            assessment_obj["confidence_threshold"] = task.confidence_threshold
-
-            # Direct insertion using parent reference - O(1) operation!
-            parent = task.parent_assessment_dict
-            field_name = task.field_name
-
-            if isinstance(parent, dict):
-                # Regular field - insert into parent dict
-                parent[field_name] = assessment_obj
-            elif isinstance(parent, list):
-                # Array item - get index from field_path
-                # field_path is like ("items", 0, "price") - second-to-last is the index
-                if len(task.field_path) >= 2 and isinstance(task.field_path[-2], int):
-                    idx = task.field_path[-2]
-                    # Replace the None placeholder we created during structure building
-                    if idx < len(parent):
-                        parent[idx] = assessment_obj
-                    else:
-                        logger.warning(
-                            f"Task {task.task_id}: index {idx} out of range for list of length {len(parent)}"
-                        )
-                else:
-                    logger.warning(
-                        f"Task {task.task_id}: cannot determine array index from path {task.field_path}"
-                    )
-            else:
-                logger.warning(
-                    f"Task {task.task_id}: unexpected parent type {type(parent)}"
+            # Add confidence_threshold if not already present
+            if "confidence_threshold" not in field_assessment:
+                field_assessment["confidence_threshold"] = task.confidence_threshold
+
+            # Insert directly at field_path - no unwrapping needed
+            try:
+                self._insert_at_field_path(
+                    assessment_structure, task.field_path, field_assessment
+                )
+                logger.debug(
+                    f"Task {task.task_id}: Inserted assessment at {task.field_path}"
+                )
+            except (KeyError, IndexError, TypeError) as e:
+                logger.error(
+                    f"Task {task.task_id}: Failed to insert at path {task.field_path}: {e}"
                 )
 
         return assessment_structure, all_confidence_alerts, aggregated_metering

lib/idp_common_pkg/idp_common/assessment/models.py

@@ -78,10 +78,6 @@ class AssessmentTask(BaseModel):
     # Confidence threshold for this field
     confidence_threshold: float
 
-    # Direct reference to parent container in assessment structure (for O(1) insertion)
-    # Can be Dict for regular fields or list for array items
-    parent_assessment_dict: dict[str, Any] | list[Any]
-
 
 class AssessmentResult(BaseModel):
     """Result of a single assessment task (used by both granular and strands services)."""

lib/idp_common_pkg/idp_common/assessment/strands_service.py

@@ -405,8 +405,12 @@ def _convert_to_assessment_result(
     metering: dict[str, Any],
     processing_time: float,
 ) -> AssessmentResult:
-    """Convert Strands AssessmentOutput to AssessmentResult with standardized geometry format."""
-    # Single field assessment
+    """
+    Convert Strands AssessmentOutput to AssessmentResult with standardized geometry format.
+
+    The assessment_data is returned as a flat dict (not wrapped by field name) because
+    the aggregation step uses task.field_path for insertion into the final structure.
+    """
     field_name = output.field_name
     assessment = output.assessment
 
@@ -429,8 +433,9 @@ def _convert_to_assessment_result(
         page_num=assessment.bounding_box.page if assessment.bounding_box else None,
     )
 
-    # Convert to explainability format
-    assessment_data = {field_name: field_data.to_explainability_format()}
+    # Return assessment data directly (not wrapped by field name)
+    # The aggregation step uses task.field_path for proper insertion
+    assessment_data = field_data.to_explainability_format()
 
     # Check for confidence threshold violations
     confidence_alerts = []

lib/idp_common_pkg/idp_common/assessment/strands_tools.py

@@ -37,22 +37,6 @@ class ViewImageInput(BaseModel):
 
 @tool
 def submit_assessment(assessment: AssessmentOutput, agent: Agent) -> str:
-    """
-    Submit your final confidence assessment.
-
-    Use this tool when you have:
-    1. Located the values in the document images
-    2. Determined precise bounding box coordinates using ruler markings
-    3. Assessed the confidence based on clarity and accuracy
-
-    Args:
-        assessment: Dictionary with:
-            - assessments: dict mapping attribute names to ConfidenceAssessment
-            - alerts: list of any threshold alerts (optional)
-
-    Returns:
-        Success confirmation message or validation error details
-    """
     # Validate assessment structure and return helpful errors
     validated_assessment = AssessmentOutput.model_validate(assessment)
 

lib/idp_common_pkg/idp_common/image/__init__.py

@@ -79,7 +79,7 @@ def resize_image(
         logger.info(
             f"Resizing image from {current_width}x{current_height} to {new_width}x{new_height} (scale: {scale_factor:.3f})"
         )
-        image = image.resize((new_width, new_height), Image.LANCZOS)
+        image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
 
         # Save in original format if possible
         img_byte_array = io.BytesIO()

lib/idp_common_pkg/idp_common/utils/pdf_helpers.py

@@ -75,7 +75,7 @@ def pdf_page_to_image(
 
         # Render page to pixmap
         mat = fitz.Matrix(dpi_scale, dpi_scale)
-        pix = page.get_pixmap(matrix=mat)
+        pix = page.get_pixmap(matrix=mat)  # pyright: ignore[reportAttributeAccessIssue]
 
         # Convert to PNG bytes
         png_bytes = pix.tobytes("png")
@@ -144,7 +144,7 @@ def pdf_to_images(
 
             # Render page to pixmap
             mat = fitz.Matrix(dpi_scale, dpi_scale)
-            pix = page.get_pixmap(matrix=mat)
+            pix = page.get_pixmap(matrix=mat)  # pyright: ignore[reportAttributeAccessIssue]
 
             # Convert to PNG bytes
             png_bytes = pix.tobytes("png")