remove image tools

not throughly tested yet

Author: kazmer97

lib/idp_common_pkg/idp_common/extraction/agentic_idp.py

@@ -17,7 +17,6 @@
     Any,
     Dict,
     List,
-    Literal,
     Optional,
     Tuple,
     Type,
@@ -28,7 +27,7 @@
 
 import jsonpatch
 from botocore.config import Config
-from PIL import Image, ImageEnhance, ImageOps
+from PIL import Image
 from pydantic import BaseModel, Field
 from strands import Agent, tool
 from strands.models.bedrock import BedrockModel
@@ -215,81 +214,6 @@ def apply_json_patches(
     return extraction_tool, apply_json_patches
 
 
-@tool
-def load_image(image_index: int, agent: Agent) -> Image.Image:
-    """
-    Load an image from state as a PIL Image for viewing and analysis.
-
-    Args:
-        image_index: Index of the image to load (0-based)
-
-    Returns:
-        PIL Image object
-    """
-    images = agent.state.get("images")
-    if not images or str(image_index) not in images:
-        raise ValueError(f"Image index {image_index} does not exist")
-
-    image_data = images[str(image_index)]
-    if isinstance(image_data, str):
-        # Convert binary string back to PIL Image
-        import base64
-
-        img_bytes = base64.b64decode(image_data)
-        return Image.open(io.BytesIO(img_bytes))
-    else:
-        # Legacy support for direct PIL images
-        return image_data
-
-
-@tool
-def enhance_image(
-    image_index: int,
-    enhancement_type: Literal["brightness", "contrast", "sharpen", "grayscale"],
-    agent: Agent,
-    factor: float = 1.5,
-) -> str:
-    """
-    Apply image enhancement for better extraction accuracy.
-
-    Args:
-        image_index: Index of the image to enhance (0-based)
-        enhancement_type: Type of enhancement (brightness, contrast, sharpen, grayscale)
-        factor: Enhancement factor (for brightness/contrast)
-    """
-    try:
-        # Load the image
-        image = load_image(image_index, agent)
-
-        # Apply enhancement
-        if enhancement_type == "brightness":
-            enhancer = ImageEnhance.Brightness(image)
-            enhanced = enhancer.enhance(factor)
-        elif enhancement_type == "contrast":
-            enhancer = ImageEnhance.Contrast(image)
-            enhanced = enhancer.enhance(factor)
-        elif enhancement_type == "sharpen":
-            enhancer = ImageEnhance.Sharpness(image)
-            enhanced = enhancer.enhance(factor)
-        elif enhancement_type == "grayscale":
-            enhanced = ImageOps.grayscale(image)
-
-        # Convert enhanced image back to binary string and store
-        img_buffer = io.BytesIO()
-        enhanced.save(img_buffer, format="PNG")
-
-        img_binary = base64.b64encode(img_buffer.getvalue()).decode("utf-8")
-
-        images = agent.state.get("images")
-        images[str(image_index)] = img_binary
-        agent.state.set(key="images", value=images)
-
-        return f"Applied {enhancement_type} to image {image_index}"
-
-    except Exception as e:
-        return f"Error enhancing image {image_index}: {str(e)}"
-
-
 SYSTEM_PROMPT = """
 You are a useful assistant that helps turn unstructured data into structured data using the provided tools.
 
@@ -317,7 +241,6 @@ def enhance_image(
 4. Verify that all characters and formatting are preserved exactly as they appear
 5. When fixing errors, use JSON patches to target specific problems
 
-Use load_image tool to view images when needed for better accuracy.
 
 FINAL REVIEW (CRITICAL):
 After successfully using the extraction tool, you MUST:
@@ -334,7 +257,6 @@ async def structured_output_async(
     model_id: str,
     data_format: Type[TargetModel],
     prompt: Union[str, Message, Image.Image],
-    enable_image_tools: bool = True,
     existing_data: Optional[BaseModel] = None,
     system_prompt: str | None = None,
     custom_instruction: str | None = None,
@@ -425,9 +347,7 @@ async def structured_output_async(
     )
 
     # Prepare tools list
-    tools = [extraction_tool, apply_json_patches, load_image]
-    if enable_image_tools:
-        tools.append(enhance_image)
+    tools = [extraction_tool, apply_json_patches]
 
     # Create agent with system prompt and tools
     schema_json = json.dumps(data_format.model_json_schema(), indent=2)
@@ -654,7 +574,6 @@ def structured_output(
     model_id: str,
     data_format: Type[BaseModel],
     prompt: Union[str, Message, Image.Image],
-    enable_image_tools: bool = True,
     existing_data: Optional[BaseModel] = None,
     system_prompt: str | None = None,
     custom_instruction: str | None = None,
@@ -682,7 +601,6 @@ def structured_output(
         model_id: Model identifier (e.g., "us.anthropic.claude-sonnet-4-20250514-v1:0")
         data_format: Pydantic model class defining the expected structure
         prompt: Input content (text, image, or content blocks)
-        enable_image_tools: Whether to enable image enhancement tools (default: True)
         existing_data: Optional existing data to update via patches
         system_prompt: **DISCOURAGED** - Custom system prompt. Only use if the default
                       SYSTEM_PROMPT is completely unsuitable for your use case.
@@ -738,7 +656,6 @@ def run_in_new_loop():
                         model_id=model_id,
                         data_format=data_format,
                         prompt=prompt,
-                        enable_image_tools=enable_image_tools,
                         existing_data=existing_data,
                         system_prompt=system_prompt,
                         custom_instruction=custom_instruction,
@@ -770,7 +687,6 @@ def run_in_new_loop():
                 model_id=model_id,
                 data_format=data_format,
                 prompt=prompt,
-                enable_image_tools=enable_image_tools,
                 existing_data=existing_data,
                 system_prompt=system_prompt,
                 custom_instruction=custom_instruction,

lib/idp_common_pkg/idp_common/extraction/service.py

@@ -1348,7 +1348,6 @@ def process_document_section(self, document: Document, section_id: str) -> Docum
                     model_id=model_id,
                     data_format=dynamic_model,
                     prompt=message_prompt,  # pyright: ignore[reportArgumentType]
-                    enable_image_tools=True,
                     custom_instruction=system_prompt,
                     review_agent=agentic_config.get("review_agent", False),
                     context="Extraction",