Update scripts/text_cleaner_tool.py

scripts/text_cleaner_tool.py

@@ -1,111 +1,62 @@
 from smolagents import Tool
-from typing import Dict, Optional, Any
+from typing import Dict, Any, Optional
 import logging
 
 logger = logging.getLogger(__name__)
 
 
 class TextCleanerTool(Tool):
+    """A simplified text cleaner tool that avoids typing issues."""
+
     name = "clean_text"
     description = (
-        "Cleans and normalizes text by removing or replacing unwanted elements"
+        "Cleans and normalizes text using the cleantext library. "
+        "Example usage: clean_text(text='Your text here', options={'lower': True, 'no_urls': True})"
     )
     inputs = {
         "text": {"type": "string", "description": "The input text to clean"},
-        "fix_unicode": {
-            "type": "boolean",
-            "description": "Fix broken unicode characters and mojibake",
-            "default": True,
-            # Removed nullable: True - this parameter doesn't accept None
-        },
-        "to_ascii": {
-            "type": "boolean",
-            "description": "Convert non-ASCII characters to their closest ASCII equivalents",
-            "default": True,
-        },
-        "lower": {
-            "type": "boolean",
-            "description": "Convert text to lowercase",
-            "default": True,
-        },
-        "no_line_breaks": {
-            "type": "boolean",
-            "description": "Replace line breaks with spaces",
-            "default": False,
-        },
-        "no_urls": {
-            "type": "boolean",
-            "description": "Replace URLs with a token",
-            "default": False,
-        },
-        "no_emails": {
-            "type": "boolean",
-            "description": "Replace email addresses with a token",
-            "default": False,
-        },
-        "no_phone_numbers": {
-            "type": "boolean",
-            "description": "Replace phone numbers with a token",
-            "default": False,
-        },
-        "no_numbers": {
-            "type": "boolean",
-            "description": "Replace all numbers with a token",
-            "default": False,
-        },
-        "no_digits": {
-            "type": "boolean",
-            "description": "Replace all digits with 0",
-            "default": False,
-        },
-        "no_currency_symbols": {
-            "type": "boolean",
-            "description": "Replace currency symbols with a token",
-            "default": False,
-        },
-        "no_punct": {
-            "type": "boolean",
-            "description": "Remove all punctuation",
-            "default": False,
-        },
-        "no_emoji": {
-            "type": "boolean",
-            "description": "Remove all emoji characters",
-            "default": False,
-        },
-        "lang": {
-            "type": "string",
-            "description": "Language code for special handling ('en' or 'de' supported)",
-            "default": "en",
-        },
-        "custom_replacements": {
+        "options": {
             "type": "object",
-            "description": "Dictionary of custom string replacements to apply",
+            "description": (
+                "Optional parameters for text cleaning. Available options: "
+                "fix_unicode, to_ascii, lower, no_line_breaks, no_urls, no_emails, "
+                "no_phone_numbers, no_numbers, no_digits, no_currency_symbols, "
+                "no_punct, no_emoji, lang"
+            ),
             "optional": True,
+            "nullable": True,
         },
     }
     output_type = "string"
 
-    def forward(
-        self,
-        text: str,
-        fix_unicode: bool = True,  # No Optional - this doesn't accept None
-        to_ascii: bool = True,
-        lower: bool = True,
-        no_line_breaks: bool = False,
-        no_urls: bool = False,
-        no_emails: bool = False,
-        no_phone_numbers: bool = False,
-        no_numbers: bool = False,
-        no_digits: bool = False,
-        no_currency_symbols: bool = False,
-        no_punct: bool = False,
-        no_emoji: bool = False,
-        lang: str = "en",
-        custom_replacements: Optional[Dict[str, str]] = None,
-    ) -> str:
-        """Clean and normalize text by removing or replacing unwanted elements."""
-        # Basic input validation
+    def forward(self, text: str, options: Optional[Dict[str, Any]] = None) -> str:
+        """
+        Clean text using the cleantext library with flexible options.
+
+        Example API:
+        clean("some input",
+            fix_unicode=True,               # fix various unicode errors
+            to_ascii=True,                  # transliterate to closest ASCII representation
+            lower=True,                     # lowercase text
+            no_line_breaks=False,           # fully strip line breaks as opposed to only normalizing them
+            no_urls=False,                  # replace all URLs with a special token
+            no_emails=False,                # replace all email addresses with a special token
+            no_phone_numbers=False,         # replace all phone numbers with a special token
+            no_numbers=False,               # replace all numbers with a special token
+            no_digits=False,                # replace all digits with a special token
+            no_currency_symbols=False,      # replace all currency symbols with a special token
+            no_punct=False,                 # remove punctuations
+            replace_with_punct="",          # instead of removing punctuations you may replace them
+            replace_with_url="<URL>",
+            replace_with_email="<EMAIL>",
+            replace_with_phone_number="<PHONE>",
+            replace_with_number="<NUMBER>",
+            replace_with_digit="0",
+            replace_with_currency_symbol="<CUR>",
+            lang="en"                       # set to 'de' for German special handling
+        )
+        """
+        # Input validation
         if not text:
             return ""
 
@@ -116,47 +67,53 @@ class TextCleanerTool(Tool):
                 logger.error(f"Failed to convert input to string: {e}")
                 return f"Error: Could not process input of type {type(text)}"
 
+        # Import cleantext safely
         try:
-            # Import cleantext
             from cleantext import clean
-
-            # Directly use the clean function with replacement parameters
-            cleaned_text = clean(
-                text,
-                fix_unicode=fix_unicode,
-                to_ascii=to_ascii,
-                lower=lower,
-                no_line_breaks=no_line_breaks,
-                no_urls=no_urls,
-                no_emails=no_emails,
-                no_phone_numbers=no_phone_numbers,
-                no_numbers=no_numbers,
-                no_digits=no_digits,
-                no_currency_symbols=no_currency_symbols,
-                no_punct=no_punct,
-                no_emoji=no_emoji,
-                lang=lang,
-                replace_with_url="<URL>",
-                replace_with_email="<EMAIL>",
-                replace_with_phone_number="<PHONE>",
-                replace_with_number="<NUMBER>",
-                replace_with_digit="0",
-                replace_with_currency_symbol="<CUR>",
-                replace_with_punct="",
-            )
-
-            # Apply any custom replacements
-            if custom_replacements:
-                for old, new in custom_replacements.items():
-                    cleaned_text = cleaned_text.replace(old, new)
-
-            return cleaned_text
-
         except ImportError:
             logger.error(
                 "cleantext package not installed. Install with: pip install clean-text"
             )
             return "Error: Required dependency 'clean-text' is not installed."
+
+        # Default replacement tokens
+        replacements = {
+            "replace_with_url": "<URL>",
+            "replace_with_email": "<EMAIL>",
+            "replace_with_phone_number": "<PHONE>",
+            "replace_with_number": "<NUMBER>",
+            "replace_with_digit": "0",
+            "replace_with_currency_symbol": "<CUR>",
+            "replace_with_punct": "",
+        }
+
+        # Default options
+        default_options = {
+            "fix_unicode": True,
+            "to_ascii": True,
+            "lower": True,
+            "no_line_breaks": False,
+            "no_urls": False,
+            "no_emails": False,
+            "no_phone_numbers": False,
+            "no_numbers": False,
+            "no_digits": False,
+            "no_currency_symbols": False,
+            "no_punct": False,
+            "no_emoji": False,
+            "lang": "en",
+        }
+
+        # Merge user options with defaults
+        if options:
+            default_options.update(options)
+
+        # Merge all parameters
+        params = {**default_options, **replacements}
+
+        try:
+            # Apply cleantext with parameters
+            return clean(text, **params)
         except Exception as e:
             logger.error(f"Error cleaning text: {e}")
             return f"Error during text cleaning: {str(e)}"