aiecs 1.0.1__py3-none-any.whl → 1.7.17__py3-none-any.whl

aiecs/tools/task_tools/image_tool.py

@@ -1,136 +1,83 @@
 import os
 import logging
 import subprocess
-import uuid
 import tempfile
 from typing import Dict, Any, List, Optional
 from dataclasses import dataclass
 from dataclasses import field
 
-from pydantic import BaseModel, ValidationError, field_validator, ConfigDict
-from pydantic_settings import BaseSettings
+from pydantic import (
+    BaseModel,
+    ValidationError,
+    field_validator,
+    Field,
+)
+from pydantic_settings import BaseSettings, SettingsConfigDict
 from PIL import Image, ExifTags, ImageFilter
 from queue import Queue
 
 from aiecs.tools.base_tool import BaseTool
 from aiecs.tools import register_tool
 
-# Configuration for ImageTool
-class ImageSettings(BaseSettings):
-    """
-    Configuration for ImageTool.
+# Module-level default configuration for validators
+_DEFAULT_MAX_FILE_SIZE_MB = 50
+_DEFAULT_ALLOWED_EXTENSIONS = [
+    ".jpg",
+    ".jpeg",
+    ".png",
+    ".bmp",
+    ".tiff",
+    ".gif",
+]
 
-    Attributes:
-        max_file_size_mb (int): Maximum file size in megabytes.
-        allowed_extensions (List[str]): Allowed image file extensions.
-        tesseract_pool_size (int): Number of Tesseract processes for OCR.
-        env_prefix (str): Environment variable prefix for settings.
-    """
-    max_file_size_mb: int = 50
-    allowed_extensions: List[str] = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.gif']
-    tesseract_pool_size: int = 2
-    env_prefix: str = 'IMAGE_TOOL_'
+# Exceptions
 
-    model_config = ConfigDict(env_prefix='IMAGE_TOOL_')
 
-# Exceptions
 class ImageToolError(Exception):
     """Base exception for ImageTool errors."""
-    pass
+
 
 class FileOperationError(ImageToolError):
     """Raised when file operations fail."""
-    pass
+
 
 class SecurityError(ImageToolError):
     """Raised for security-related issues."""
-    pass
+
 
 # Base schema for common fields
+
+
 class BaseFileSchema(BaseModel):
     file_path: str
     _mtime: Optional[float] = None  # Internal use for cache
 
-    @field_validator('file_path')
+    @field_validator("file_path")
     @classmethod
     def validate_file_path(cls, v: str) -> str:
         """Validate file path for existence, size, and extension."""
-        settings = ImageSettings()
         abs_path = os.path.abspath(os.path.normpath(v))
         ext = os.path.splitext(abs_path)[1].lower()
-        if ext not in settings.allowed_extensions:
-            raise SecurityError(f"Extension '{ext}' not allowed, expected {settings.allowed_extensions}")
+        if ext not in _DEFAULT_ALLOWED_EXTENSIONS:
+            raise SecurityError(f"Extension '{ext}' not allowed, expected {_DEFAULT_ALLOWED_EXTENSIONS}")
         if not os.path.isfile(abs_path):
             raise FileOperationError(f"File not found: {abs_path}")
         size_mb = os.path.getsize(abs_path) / (1024 * 1024)
-        if size_mb > settings.max_file_size_mb:
-            raise FileOperationError(f"File too large: {size_mb:.1f}MB, max {settings.max_file_size_mb}MB")
+        if size_mb > _DEFAULT_MAX_FILE_SIZE_MB:
+            raise FileOperationError(f"File too large: {size_mb:.1f}MB, max {_DEFAULT_MAX_FILE_SIZE_MB}MB")
         return abs_path
 
-# Schemas for operations
-class LoadSchema(BaseFileSchema):
-    """Schema for load operation."""
-    pass
-
-class OCRSchema(BaseFileSchema):
-    """Schema for OCR operation."""
-    lang: Optional[str] = None
-
-
-class MetadataSchema(BaseFileSchema):
-    """Schema for metadata extraction operation."""
-    include_exif: bool = False
 
-class ResizeSchema(BaseFileSchema):
-    """Schema for resize operation."""
-    output_path: str
-    width: int
-    height: int
+# Schemas for operations - moved to ImageTool class as inner classes
 
-    @field_validator('output_path')
-    @classmethod
-    def validate_output_path(cls, v: str) -> str:
-        """Validate output path for existence and extension."""
-        settings = ImageSettings()
-        abs_path = os.path.abspath(os.path.normpath(v))
-        ext = os.path.splitext(abs_path)[1].lower()
-        if ext not in settings.allowed_extensions:
-            raise SecurityError(f"Output extension '{ext}' not allowed, expected {settings.allowed_extensions}")
-        if os.path.exists(abs_path):
-            raise FileOperationError(f"Output file already exists: {abs_path}")
-        return abs_path
 
-class FilterSchema(BaseFileSchema):
-    """Schema for filter operation."""
-    output_path: str
-    filter_type: str = 'blur'
+# Tesseract process manager
 
-    @field_validator('filter_type')
-    @classmethod
-    def validate_filter_type(cls, v: str) -> str:
-        """Validate filter type."""
-        valid_filters = ['blur', 'sharpen', 'edge_enhance']
-        if v not in valid_filters:
-            raise ValueError(f"Invalid filter_type '{v}', expected {valid_filters}")
-        return v
-
-    @field_validator('output_path')
-    @classmethod
-    def validate_output_path(cls, v: str) -> str:
-        """Validate output path for existence and extension."""
-        settings = ImageSettings()
-        abs_path = os.path.abspath(os.path.normpath(v))
-        ext = os.path.splitext(abs_path)[1].lower()
-        if ext not in settings.allowed_extensions:
-            raise SecurityError(f"Output extension '{ext}' not allowed, expected {settings.allowed_extensions}")
-        if os.path.exists(abs_path):
-            raise FileOperationError(f"Output file already exists: {abs_path}")
-        return abs_path
 
-# Tesseract process manager
 @dataclass
 class TesseractManager:
     """Manages a pool of Tesseract processes for OCR."""
+
     pool_size: int
     processes: List[subprocess.Popen] = field(default_factory=list)
     queue: Queue = field(default_factory=lambda: Queue())
@@ -140,11 +87,11 @@ class TesseractManager:
         for _ in range(self.pool_size):
             try:
                 proc = subprocess.Popen(
-                    ['tesseract', '--oem', '1', '-', 'stdout', '-l', 'eng'],
+                    ["tesseract", "--oem", "1", "-", "stdout", "-l", "eng"],
                     stdin=subprocess.PIPE,
                     stdout=subprocess.PIPE,
                     stderr=subprocess.PIPE,
-                    text=True
+                    text=True,
                 )
                 self.queue.put(proc)
                 self.processes.append(proc)
@@ -171,7 +118,8 @@ class TesseractManager:
             except (subprocess.TimeoutExpired, OSError) as e:
                 logging.getLogger(__name__).warning(f"Error terminating Tesseract process: {e}")
 
-@register_tool('image')
+
+@register_tool("image")
 class ImageTool(BaseTool):
     """
     Image processing tool supporting:
@@ -183,38 +131,142 @@ class ImageTool(BaseTool):
 
     Inherits from BaseTool to leverage ToolExecutor for caching, concurrency, and error handling.
     """
-    def __init__(self, config: Dict[Any, Any] = None):
+
+    # Configuration schema
+    class Config(BaseSettings):
+        """Configuration for the image tool
+        
+        Automatically reads from environment variables with IMAGE_TOOL_ prefix.
+        Example: IMAGE_TOOL_MAX_FILE_SIZE_MB -> max_file_size_mb
+        """
+
+        model_config = SettingsConfigDict(env_prefix="IMAGE_TOOL_")
+
+        max_file_size_mb: int = Field(default=50, description="Maximum file size in megabytes")
+        allowed_extensions: List[str] = Field(
+            default=[".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".gif"],
+            description="Allowed image file extensions",
+        )
+        tesseract_pool_size: int = Field(default=2, description="Number of Tesseract processes for OCR")
+        default_ocr_language: str = Field(
+            default="eng",
+            description="Default OCR language code (e.g., 'eng', 'chi_sim'). Supports multi-language format like 'eng+chi_sim'"
+        )
+
+    # Schema definitions
+    class LoadSchema(BaseFileSchema):
+        """Schema for load operation"""
+
+        file_path: str = Field(description="Path to the image file to load")
+
+    class OcrSchema(BaseFileSchema):
+        """Schema for ocr operation"""
+
+        file_path: str = Field(description="Path to the image file for OCR text extraction")
+        lang: Optional[str] = Field(
+            default=None,
+            description="Optional language code for OCR (e.g., 'eng', 'chi_sim', 'eng+chi_sim'). If not specified, uses the configured default_ocr_language"
+        )
+
+    class MetadataSchema(BaseFileSchema):
+        """Schema for metadata operation"""
+
+        file_path: str = Field(description="Path to the image file to extract metadata from")
+        include_exif: bool = Field(default=False, description="Whether to include EXIF data in the metadata. If False, only basic info (size, mode) is returned")
+
+    class ResizeSchema(BaseFileSchema):
+        """Schema for resize operation"""
+
+        file_path: str = Field(description="Path to the source image file")
+        output_path: str = Field(description="Path where the resized image will be saved")
+        width: int = Field(description="Target width in pixels for the resized image")
+        height: int = Field(description="Target height in pixels for the resized image")
+
+        @field_validator("output_path")
+        @classmethod
+        def validate_output_path(cls, v: str) -> str:
+            """Validate output path for existence and extension."""
+            abs_path = os.path.abspath(os.path.normpath(v))
+            ext = os.path.splitext(abs_path)[1].lower()
+            if ext not in _DEFAULT_ALLOWED_EXTENSIONS:
+                raise SecurityError(f"Output extension '{ext}' not allowed, expected {_DEFAULT_ALLOWED_EXTENSIONS}")
+            if os.path.exists(abs_path):
+                raise FileOperationError(f"Output file already exists: {abs_path}")
+            return abs_path
+
+    class FilterSchema(BaseFileSchema):
+        """Schema for filter operation"""
+
+        file_path: str = Field(description="Path to the source image file")
+        output_path: str = Field(description="Path where the filtered image will be saved")
+        filter_type: str = Field(default="blur", description="Type of filter to apply: 'blur', 'sharpen', or 'edge_enhance'")
+
+        @field_validator("filter_type")
+        @classmethod
+        def validate_filter_type(cls, v: str) -> str:
+            """Validate filter type."""
+            valid_filters = ["blur", "sharpen", "edge_enhance"]
+            if v not in valid_filters:
+                raise ValueError(f"Invalid filter_type '{v}', expected {valid_filters}")
+            return v
+
+        @field_validator("output_path")
+        @classmethod
+        def validate_output_path(cls, v: str) -> str:
+            """Validate output path for existence and extension."""
+            abs_path = os.path.abspath(os.path.normpath(v))
+            ext = os.path.splitext(abs_path)[1].lower()
+            if ext not in _DEFAULT_ALLOWED_EXTENSIONS:
+                raise SecurityError(f"Output extension '{ext}' not allowed, expected {_DEFAULT_ALLOWED_EXTENSIONS}")
+            if os.path.exists(abs_path):
+                raise FileOperationError(f"Output file already exists: {abs_path}")
+            return abs_path
+
+    def __init__(self, config: Optional[Dict[str, Any]] = None, **kwargs):
         """
-        Initialize ImageTool with settings and resources.
+        Initialize ImageTool with configuration and resources.
 
         Args:
-            config (Dict, optional): Configuration overrides for ImageSettings.
+            config (Dict, optional): Configuration overrides for ImageTool.
+            **kwargs: Additional arguments passed to BaseTool (e.g., tool_name)
 
         Raises:
             ValueError: If config contains invalid settings.
+
+        Configuration is automatically loaded by BaseTool from:
+        1. Explicit config dict (highest priority)
+        2. YAML config files (config/tools/image.yaml, config/tools/image_tool.yaml, or config/tools/ImageTool.yaml)
+        3. Environment variables (via dotenv from .env files with IMAGE_TOOL_ prefix)
+        4. Tool defaults (lowest priority)
+        
+        YAML configuration files are automatically discovered by ToolConfigLoader using multiple naming conventions.
+        See examples/config/tools/image_tool.yaml.example for a configuration template.
         """
-        super().__init__(config)
-        self.settings = ImageSettings()
-        if config:
-            try:
-                self.settings = self.settings.model_validate({**self.settings.model_dump(), **config})
-            except ValidationError as e:
-                raise ValueError(f"Invalid configuration: {e}")
+        # Pass tool_name="image" to BaseTool so it can find config/tools/image.yaml
+        if "tool_name" not in kwargs:
+            kwargs["tool_name"] = "image"
+        super().__init__(config, **kwargs)
+
+        # Configuration is automatically loaded by BaseTool into self._config_obj
+        # Access config via self._config_obj (BaseSettings instance)
+        self.config = self._config_obj if self._config_obj else self.Config()
+
         self.logger = logging.getLogger(__name__)
         if not self.logger.handlers:
             handler = logging.StreamHandler()
-            handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(message)s'))
+            handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
             self.logger.addHandler(handler)
         self.logger.setLevel(logging.INFO)
+
         # Initialize Tesseract manager
-        self._tesseract_manager = TesseractManager(self.settings.tesseract_pool_size)
+        self._tesseract_manager = TesseractManager(self.config.tesseract_pool_size)
         self._tesseract_manager.initialize()
 
     def __del__(self):
         """Clean up Tesseract processes on destruction."""
         self._tesseract_manager.cleanup()
 
-    def update_settings(self, config: Dict):
+    def update_config(self, config: Dict) -> None:
         """
         Update configuration settings dynamically.
 
@@ -225,11 +277,11 @@ class ImageTool(BaseTool):
             ValueError: If config contains invalid settings.
         """
         try:
-            self.settings = self.settings.model_validate({**self.settings.model_dump(), **config})
+            self.config = self.Config(**{**self.config.model_dump(), **config})
             # Reinitialize Tesseract if pool size changes
-            if 'tesseract_pool_size' in config:
+            if "tesseract_pool_size" in config:
                 self._tesseract_manager.cleanup()
-                self._tesseract_manager = TesseractManager(self.settings.tesseract_pool_size)
+                self._tesseract_manager = TesseractManager(self.config.tesseract_pool_size)
                 self._tesseract_manager.initialize()
         except ValidationError as e:
             raise ValueError(f"Invalid configuration: {e}")
@@ -248,22 +300,23 @@ class ImageTool(BaseTool):
             FileOperationError: If file is invalid or inaccessible.
         """
         # Validate input using schema
-        validated_input = LoadSchema(file_path=file_path)
-        
+        validated_input = self.LoadSchema(file_path=file_path)
+
         try:
             with Image.open(validated_input.file_path) as img:
                 img.load()
-                return {'size': img.size, 'mode': img.mode}
+                return {"size": img.size, "mode": img.mode}
         except Exception as e:
             raise FileOperationError(f"load: Failed to load image '{file_path}': {e}")
 
     def ocr(self, file_path: str, lang: Optional[str] = None) -> str:
         """
-        Extract text from an image using a pooled Tesseract process.
+        Extract text from an image using Tesseract OCR.
 
         Args:
             file_path (str): Path to the image file.
-            lang (Optional[str]): Language code for OCR (e.g., 'eng').
+            lang (Optional[str]): Language code for OCR (e.g., 'eng', 'chi_sim', 'eng+chi_sim').
+                If not specified, uses the configured default_ocr_language.
 
         Returns:
             str: Extracted text.
@@ -272,31 +325,51 @@ class ImageTool(BaseTool):
             FileOperationError: If OCR fails or Tesseract is unavailable.
         """
         # Validate input using schema
-        validated_input = OCRSchema(file_path=file_path, lang=lang)
+        validated_input = self.OcrSchema(file_path=file_path, lang=lang)
         
-        proc = self._tesseract_manager.get_process()
-        if not proc:
-            raise FileOperationError(f"ocr: No Tesseract processes available (lang: {lang or 'eng'})")
-        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file:
+        # Use configured default language if lang is not specified
+        ocr_lang = lang if lang is not None else self.config.default_ocr_language
+
+        # Prepare temporary file for image processing
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
             temp_path = temp_file.name
         try:
-            img = Image.open(validated_input.file_path).convert('L').filter(ImageFilter.SHARPEN)
+            # Preprocess image for better OCR results
+            img = Image.open(validated_input.file_path).convert("L").filter(ImageFilter.SHARPEN)
             img.save(temp_path)
-            stdout, stderr = proc.communicate(input=temp_path, timeout=30)
-            if proc.returncode != 0:
-                raise FileOperationError(f"ocr: Tesseract failed for '{file_path}' (lang: {lang or 'eng'}): {stderr}")
-            return stdout.strip()
+            
+            # Call Tesseract with dynamic language parameter
+            # Use subprocess.run instead of process pool to support dynamic languages
+            try:
+                result = subprocess.run(
+                    ["tesseract", "--oem", "1", temp_path, "stdout", "-l", ocr_lang],
+                    capture_output=True,
+                    text=True,
+                    timeout=30,
+                    check=False,  # Don't raise on non-zero return code, handle manually
+                )
+            except subprocess.TimeoutExpired:
+                raise FileOperationError(f"ocr: Tesseract timeout for '{file_path}' (lang: {ocr_lang})")
+            except FileNotFoundError:
+                raise FileOperationError("ocr: Tesseract not found. Please install Tesseract OCR.")
+            
+            if result.returncode != 0:
+                raise FileOperationError(
+                    f"ocr: Tesseract failed for '{file_path}' (lang: {ocr_lang}): {result.stderr}"
+                )
+            
+            return result.stdout.strip()
+        except FileOperationError:
+            raise  # Re-raise FileOperationError as-is
         except Exception as e:
-            raise FileOperationError(f"ocr: Failed to process '{file_path}' (lang: {lang or 'eng'}): {e}")
+            raise FileOperationError(f"ocr: Failed to process '{file_path}' (lang: {ocr_lang}): {e}")
         finally:
-            self._tesseract_manager.return_process(proc)
             if os.path.exists(temp_path):
                 try:
                     os.unlink(temp_path)
                 except Exception as e:
                     self.logger.warning(f"Failed to remove temporary file {temp_path}: {e}")
 
-
     def metadata(self, file_path: str, include_exif: bool = False) -> Dict[str, Any]:
         """
         Retrieve metadata (size, mode, EXIF) from an image.
@@ -312,19 +385,19 @@ class ImageTool(BaseTool):
             FileOperationError: If metadata extraction fails.
         """
         # Validate input using schema
-        validated_input = MetadataSchema(file_path=file_path, include_exif=include_exif)
-        
+        validated_input = self.MetadataSchema(file_path=file_path, include_exif=include_exif)
+
         try:
             with Image.open(validated_input.file_path) as img:
                 img.load()
-                info = {'size': img.size, 'mode': img.mode}
+                info = {"size": img.size, "mode": img.mode}
                 if include_exif:
                     exif = {}
                     raw = img._getexif() or {}
                     for tag, val in raw.items():
                         decoded = ExifTags.TAGS.get(tag, tag)
                         exif[decoded] = val
-                    info['exif'] = exif
+                    info["exif"] = exif
                 return info
         except Exception as e:
             raise FileOperationError(f"metadata: Failed to process '{file_path}': {e}")
@@ -346,18 +419,21 @@ class ImageTool(BaseTool):
             FileOperationError: If resizing fails.
         """
         # Validate input using schema
-        validated_input = ResizeSchema(
-            file_path=file_path, 
-            output_path=output_path, 
-            width=width, 
-            height=height
+        validated_input = self.ResizeSchema(
+            file_path=file_path,
+            output_path=output_path,
+            width=width,
+            height=height,
         )
-        
+
         try:
             with Image.open(validated_input.file_path) as img:
                 img = img.resize((width, height), Image.Resampling.LANCZOS)
                 img.save(validated_input.output_path)
-            return {'success': True, 'output_path': validated_input.output_path}
+            return {
+                "success": True,
+                "output_path": validated_input.output_path,
+            }
         except Exception as e:
             raise FileOperationError(f"resize: Failed to process '{file_path}' (output_path: {output_path}): {e}")
 
@@ -377,21 +453,24 @@ class ImageTool(BaseTool):
             FileOperationError: If filtering fails.
         """
         # Validate input using schema
-        validated_input = FilterSchema(
-            file_path=file_path, 
-            output_path=output_path, 
-            filter_type=filter_type
+        validated_input = self.FilterSchema(
+            file_path=file_path,
+            output_path=output_path,
+            filter_type=filter_type,
         )
-        
+
         try:
             filter_map = {
-                'blur': ImageFilter.BLUR,
-                'sharpen': ImageFilter.SHARPEN,
-                'edge_enhance': ImageFilter.EDGE_ENHANCE
+                "blur": ImageFilter.BLUR,
+                "sharpen": ImageFilter.SHARPEN,
+                "edge_enhance": ImageFilter.EDGE_ENHANCE,
             }
             with Image.open(validated_input.file_path) as img:
                 img = img.filter(filter_map[filter_type])
                 img.save(validated_input.output_path)
-            return {'success': True, 'output_path': validated_input.output_path}
+            return {
+                "success": True,
+                "output_path": validated_input.output_path,
+            }
         except Exception as e:
             raise FileOperationError(f"filter: Failed to process '{file_path}' (output_path: {output_path}, filter_type: {filter_type}): {e}")