Upload tool

requirements.txt

@@ -1,2 +1,2 @@
-smolagents
 pronouncing
+smolagents

tool.py

@@ -1,14 +1,14 @@
 from smolagents.tools import Tool
 import pronouncing
+import difflib
 import json
 import string
-import difflib
 
 class ParodyWordSuggestionTool(Tool):
     name = "parody_word_suggester"
     description = """Suggests rhyming funny words using CMU dictionary pronunciations.
     Returns similar-sounding words that rhyme, especially focusing on common vowel sounds."""
-    inputs = {'target': {'type': 'string', 'description': 'The word you want to find rhyming alternatives for'}, 'word_list_str': {'type': 'string', 'description': 'JSON string of word list (e.g. \'["word1", "word2"]\')'}, 'min_similarity': {'type': 'string', 'description': 'Minimum similarity threshold (0.0-1.0)', 'nullable': True, 'default': '0.5'}}
+    inputs = {'target': {'type': 'string', 'description': 'The word you want to find rhyming alternatives for'}, 'word_list_str': {'type': 'string', 'description': 'JSON string of word list (e.g. \'["word1", "word2"]\')'}, 'min_similarity': {'type': 'string', 'description': 'Minimum similarity threshold (0.0-1.0)', 'nullable': True, 'default': '0.5'}, 'custom_phones': {'type': 'object', 'description': 'Optional dictionary of custom word pronunciations', 'nullable': True, 'default': None}}
     output_type = "string"
     VOWEL_REF = "AH,AX|UH|AE,EH|IY,IH|AO,AA|UW|AY,EY|OW,AO|AW,AO|OY,OW|ER,AXR"
     CONSONANT_REF = "M,N,NG|P,B|T,D|K,G|F,V|TH,DH|S,Z|SH,ZH|L,R|W,Y"
@@ -29,6 +29,16 @@ class ParodyWordSuggestionTool(Tool):
         return groups
 
 
+    def _get_word_phones(self, word, custom_phones=None):
+        """Get phones for a word, checking custom dictionary first."""
+        if custom_phones and word in custom_phones:
+            return custom_phones[word]["primary_phones"]
+        
+        import pronouncing
+        phones = pronouncing.phones_for_word(word)
+        return phones[0] if phones else None
+
+
     def _get_last_syllable(self, phones: list) -> tuple:
         """Extract the last syllable (vowel + remaining consonants)."""
         last_vowel_idx = -1
@@ -173,7 +183,7 @@ class ParodyWordSuggestionTool(Tool):
         }
 
 
-    def forward(self, target: str, word_list_str: str, min_similarity: str = "0.5") -> str:
+    def forward(self, target: str, word_list_str: str, min_similarity: str = "0.5", custom_phones: dict = None) -> str:
         import pronouncing
         import string
         import json
@@ -205,40 +215,38 @@ class ParodyWordSuggestionTool(Tool):
                 "suggestions": []
             }, indent=2)
         
-        # Get target pronunciation
-        target_phones = pronouncing.phones_for_word(target)
+        # Get target pronunciation using custom phones
+        target_phones = self._get_word_phones(target, custom_phones)
         if not target_phones:
             return json.dumps({
-                "error": f"Target word '{target}' not found in CMU dictionary",
+                "error": f"Target word '{target}' not found in dictionary or custom phones",
                 "suggestions": []
             }, indent=2)
     
-        # Filter word list to only words in CMU dictionary
+        # Filter word list checking both CMU and custom phones
         valid_words = []
         invalid_words = []
         for word in words:
             word = word.lower().strip(string.punctuation)
-            if pronouncing.phones_for_word(word):
+            if self._get_word_phones(word, custom_phones):
                 valid_words.append(word)
             else:
                 invalid_words.append(word)
     
         if not valid_words:
             return json.dumps({
-                "error": "No valid words found in CMU dictionary",
+                "error": "No valid words found in dictionary or custom phones",
                 "invalid_words": invalid_words,
                 "suggestions": []
             }, indent=2)
     
-        target_phones = target_phones[0]
         target_phone_list = target_phones.split()
         target_vowel, target_end = self._get_last_syllable(target_phone_list)
     
         # Check each word
         for word in valid_words:
-            phones = pronouncing.phones_for_word(word)
-            if phones:
-                word_phones = phones[0]
+            word_phones = self._get_word_phones(word, custom_phones)
+            if word_phones:
                 similarity_result = self._calculate_similarity(word, word_phones, target, target_phones)
             
                 if similarity_result["similarity"] >= min_similarity:
@@ -254,6 +262,7 @@ class ParodyWordSuggestionTool(Tool):
                         "phones": word_phones,
                         "last_vowel": word_vowel,
                         "ending": " ".join(word_end) if word_end else "",
+                        "is_custom": word in custom_phones if custom_phones else False,
                         "details": similarity_result["details"]
                     })