assesment update

Author: kazmer97

lib/idp_common_pkg/idp_common/assessment/granular_service.py

@@ -14,9 +14,10 @@
 import json
 import os
 import time
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
 
 from aws_lambda_powertools import Logger
+from mypy_boto3_dynamodb.service_resource import DynamoDBServiceResource
 
 from idp_common import image, metrics, s3, utils
 from idp_common.assessment.models import AssessmentResult, AssessmentTask
@@ -31,7 +32,7 @@
     X_AWS_IDP_CONFIDENCE_THRESHOLD,
     X_AWS_IDP_DOCUMENT_TYPE,
 )
-from idp_common.extraction.models import ExtractionData, ExtractionMetadata
+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

lib/idp_common_pkg/idp_common/assessment/strands_service.py

@@ -12,7 +12,10 @@
 
 from aws_lambda_powertools import Logger
 from botocore.config import Config
-from strands import Agent, tool
+from idp_common.utils.bedrock_utils import (
+    async_exponential_backoff_retry,
+)
+from strands import Agent
 from strands.agent.conversation_manager import SummarizingConversationManager
 from strands.models.bedrock import BedrockModel
 from strands.types.content import ContentBlock, Message
@@ -24,48 +27,6 @@
 logger = Logger(service="assessment", level=os.getenv("LOG_LEVEL", "INFO"))
 
 
-def create_submit_assessment_tool():
-    """
-    Create a tool for submitting assessment results.
-
-    Returns:
-        A Strands tool function for submitting assessments
-    """
-
-    @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(**assessment)  # pyright: ignore[reportCallIssue]
-
-        # Store in agent state
-        agent.state.set("assessment_output", validated_assessment.model_dump())
-
-        logger.info(
-            "Assessment submitted successfully",
-            extra={"assessment": validated_assessment.model_dump()},
-        )
-
-        return "Assessment submitted successfully. You can now finish the task."
-
-    return submit_assessment
-
-
 async def assess_attribute_with_strands(
     task: AssessmentTask,
     extraction_results: dict[str, Any],
@@ -105,9 +66,7 @@ async def assess_attribute_with_strands(
     try:
         # 1. Create tools (image viewer + todo list + submit assessment)
         base_tools = create_strands_tools(page_images, sorted_page_ids)
-        submit_tool = create_submit_assessment_tool()
-        tools = base_tools + [submit_tool]
-
+        tools = base_tools
         # 2. Build enhanced system prompt with schema and extraction (for caching)
         enhanced_system_prompt = _build_system_prompt_with_context(
             system_prompt, document_schema, extraction_results, len(page_images)
@@ -161,8 +120,17 @@ async def assess_attribute_with_strands(
             },
         )
 
-        response = await agent.invoke_async([user_message])
+        @async_exponential_backoff_retry(
+            max_retries=30,
+            initial_delay=5,
+            exponential_base=2,
+            jitter=0.5,
+            max_delay=900,
+        )
+        async def invoke_agent_with_retry():
+            return await agent.invoke_async([user_message])
 
+        response = await invoke_agent_with_retry()
         logger.debug("Agent response received", extra={"task_id": task.task_id})
 
         # 7. Extract assessment from agent state

lib/idp_common_pkg/idp_common/assessment/strands_tools.py

@@ -11,10 +11,8 @@
 from aws_lambda_powertools import Logger
 from pydantic import BaseModel, Field
 from strands import Agent, tool
-from strands.types.content import ContentBlock, ImageContent
-from strands.types.tools import ToolResult, ToolResultContent, ToolResultStatus
 
-from idp_common.assessment.strands_models import BoundingBox
+from idp_common.assessment.strands_models import AssessmentOutput, BoundingBox
 from idp_common.utils.grid_overlay import draw_bounding_boxes
 from idp_common.utils.strands_agent_tools.todo_list import (
     create_todo_list,
@@ -37,6 +35,38 @@ class ViewImageInput(BaseModel):
     label: str | None = Field(None, description="Optional label for the bounding box")
 
 
+@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(**assessment)  # pyright: ignore[reportCallIssue]
+
+    # Store in agent state
+    agent.state.set("assessment_output", validated_assessment.model_dump())
+
+    logger.info(
+        "Assessment submitted successfully",
+        extra={"assessment": validated_assessment.model_dump()},
+    )
+
+    return "Assessment submitted successfully. You can now finish the task."
+
+
 def create_view_image_tool(page_images: list[bytes], sorted_page_ids: list[str]) -> Any:
     """
     Create a view_image tool that has access to page images.
@@ -163,4 +193,5 @@ def create_strands_tools(
         create_todo_list,
         update_todo,
         view_todo_list,
+        submit_assessment,
     ]

lib/idp_common_pkg/idp_common/config/models.py

@@ -256,7 +256,7 @@ class AssessmentConfig(BaseModel):
     image: ImageConfig = Field(default_factory=ImageConfig)
     # Granular assessment settings (always enabled, no longer nested)
     max_workers: int = Field(
-        default=20, gt=0, description="Max concurrent workers for parallel assessment"
+        default=50, gt=0, description="Max concurrent workers for parallel assessment"
     )
 
     @field_validator(

lib/idp_common_pkg/pyproject.toml

@@ -21,10 +21,11 @@ description = "Common utilities for GenAI IDP Accelerator patterns"
 authors = [{ name = "AWS", email = "[email protected]" }]
 requires-python = ">=3.12,<3.14"
 dependencies = [
-  "boto3==1.39.7",      # Core dependency for AWS services
+  "boto3==1.39.7", # Core dependency for AWS services
   "jsonschema>=4.25.1",
   "pydantic>=2.12.0",
   "deepdiff>=6.0.0",
+  "mypy-boto3-bedrock-runtime>=1.40.76",
 ]
 readme = "README.md"
 classifiers = [
@@ -73,6 +74,8 @@ extraction = [
 assessment = [
   "Pillow==11.2.1",               # For image handling
   "aws-lambda-powertools>=3.2.0", # Structured logging and observability
+  "mypy-boto3-dynamodb>=1.40.56",
+  "strands-agents==1.17.0",
 ]
 
 # Evaluation module dependencies
@@ -98,15 +101,15 @@ appsync = ["requests==2.32.4"]
 
 # Agents module dependencies
 agents = [
-  "strands-agents==1.14.0; python_version>='3.10'",
+  "strands-agents==1.17.0; python_version>='3.10'",
   "strands-agents-tools==0.2.13; python_version>='3.10'",
   "bedrock-agentcore>=0.1.1; python_version>='3.10'",     # Specifically for the code interpreter tool
 ]
 
 # Code intelligence module dependencies
 code_intel = [
   "requests==2.32.4",
-  "strands-agents==1.14.0",
+  "strands-agents==1.17.0",
   "strands-agents-tools==0.2.13",
   "bedrock-agentcore>=0.1.1",
   "PyYAML>=6.0.0",
@@ -158,14 +161,14 @@ all = [
   "pyarrow==20.0.0",
   "openpyxl==3.1.5",
   "python-docx==1.2.0",
-  "strands-agents==1.14.0; python_version>='3.10'",
+  "strands-agents==1.17.0; python_version>='3.10'",
   "strands-agents-tools==0.2.13; python_version>='3.10'",
   "bedrock-agentcore>=0.1.1; python_version>='3.10'",
   # "s3fs==2023.12.2" - - disabled till we fix package dependencies
 ]
 agentic-extraction = [
   "jsonpatch==1.33",
-  "strands-agents==1.14.0 ; python_full_version >= '3.10'",
+  "strands-agents==1.17.0 ; python_full_version >= '3.10'",
   "pandas>=2.2.3",
   "pymupdf==1.25.5",                                        # Pinned to 1.25.5 - has pre-built ARM64 wheels, 1.26.x requires compilation
   "email-validator>=2.3.0",