Upload tool

tool.py

@@ -1,7 +1,7 @@
 from smolagents.tools import Tool
 import pronouncing
-import json
 import difflib
+import json
 import string
 
 class ParodyWordSuggestionTool(Tool):
@@ -11,6 +11,41 @@ class ParodyWordSuggestionTool(Tool):
     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}}
     output_type = "string"
 
+    def _are_vowels_similar(self, v1: str, v2: str) -> bool:
+        """Check if two vowel sounds are similar enough to rhyme."""
+        # Strip stress markers
+        v1 = v1.rstrip('012')
+        v2 = v2.rstrip('012')
+    
+        # Define groups of similar vowel sounds
+        similar_vowels = [
+            {'AH', 'UH'},  # Short u sounds
+            {'AE', 'EH'},  # Short e/a sounds
+            {'IY', 'IH'},  # Long/short i sounds
+            {'AO', 'AA'},  # Open o/a sounds
+            {'UW', 'UH'},  # Long/short oo sounds
+        ]
+    
+        # Direct match
+        if v1 == v2:
+            return True
+        
+        # Check if they're in the same similarity group
+        for group in similar_vowels:
+            if v1 in group and v2 in group:
+                return True
+            
+        return False
+
+
+    def _contains_vowel(self, phone: str, vowels: list) -> bool:
+        """Helper function to check if a phone contains any vowel from the list."""
+        for v in vowels:
+            if v in phone:
+                return True
+        return False
+
+
     def forward(self, target: str, word_list_str: str, min_similarity: str = "0.5") -> str:
         """Get rhyming word suggestions."""
         import pronouncing
@@ -18,6 +53,8 @@ class ParodyWordSuggestionTool(Tool):
         import json
         from difflib import SequenceMatcher
     
+        VOWEL_LETTERS = ['A', 'E', 'I', 'O', 'U']
+    
         target = target.lower().strip(string.punctuation)
         min_similarity = float(min_similarity)
         suggestions = []
@@ -42,6 +79,14 @@ class ParodyWordSuggestionTool(Tool):
         target_phones = target_phones[0]
         target_phone_list = target_phones.split()
     
+        # Find the last vowel in target
+        target_last_vowel = None
+        target_last_vowel_idx = -1
+        for i, phone in enumerate(target_phone_list):
+            if self._contains_vowel(phone, VOWEL_LETTERS):
+                target_last_vowel = phone
+                target_last_vowel_idx = i
+    
         # Check each word
         for word in words:
             word = word.lower().strip(string.punctuation)
@@ -50,39 +95,46 @@ class ParodyWordSuggestionTool(Tool):
                 word_phones = phones[0]
                 word_phone_list = word_phones.split()
             
+                # Find last vowel in word
+                word_last_vowel = None
+                word_last_vowel_idx = -1
+                for i, phone in enumerate(word_phone_list):
+                    if self._contains_vowel(phone, VOWEL_LETTERS):
+                        word_last_vowel = phone
+                        word_last_vowel_idx = i
+            
                 # 1. Rhyme score (most important - 60%)
                 rhyme_score = 0
-                if len(word_phone_list) > 1 and len(target_phone_list) > 1:
-                    # Check if words share the same ending (vowel + final consonants)
-                    vowel_plus_end = -2  # Index of the vowel in final syllable
-                    while vowel_plus_end < -1:
-                        if 'A' in word_phone_list[vowel_plus_end] or 'E' in word_phone_list[vowel_plus_end] or 'I' in word_phone_list[vowel_plus_end] or 'O' in word_phone_list[vowel_plus_end] or 'U' in word_phone_list[vowel_plus_end]:
-                            break
-                        vowel_plus_end += 1
-                
-                    if vowel_plus_end == -1:
-                        vowel_plus_end = -2  # Fall back if no vowel found
-                    
-                    # Check if the ending (from vowel onwards) matches
-                    if word_phone_list[vowel_plus_end:] == target_phone_list[vowel_plus_end:]:
-                        rhyme_score = 1.0
+                if word_last_vowel and target_last_vowel:
+                    # Check if the vowels are similar
+                    if self._are_vowels_similar(word_last_vowel, target_last_vowel):
+                        # Check if the endings after the vowel match
+                        if (word_phone_list[word_last_vowel_idx:] == 
+                            target_phone_list[target_last_vowel_idx:]):
+                            rhyme_score = 1.0
+                        # Partial match for similar endings
+                        elif (len(word_phone_list) > word_last_vowel_idx + 1 and
+                              len(target_phone_list) > target_last_vowel_idx + 1 and
+                              word_phone_list[-1] == target_phone_list[-1]):
+                            rhyme_score = 0.8
             
                 # 2. Syllable match (25%)
                 target_syl = pronouncing.syllable_count(target_phones)
                 word_syl = pronouncing.syllable_count(word_phones)
                 syllable_score = 1.0 if target_syl == word_syl else 0.0
             
-                # 3. Overall similarity (15%) - using string similarity
+                # 3. Overall similarity (15%)
                 string_similarity = SequenceMatcher(None, target, word).ratio()
             
-                # Combined score (60% rhyme, 25% syllables, 15% similarity)
+                # Combined score
                 similarity = (rhyme_score * 0.6) + (syllable_score * 0.25) + (string_similarity * 0.15)
             
                 if similarity >= min_similarity:
                     suggestions.append({
                         "word": word,
                         "similarity": round(similarity, 3),
-                        "rhyme_match": rhyme_score == 1.0,
+                        "rhyme_match": rhyme_score > 0,
+                        "rhyme_score": round(rhyme_score, 3),
                         "syllable_match": syllable_score == 1.0,
                         "string_similarity": round(string_similarity, 3),
                         "syllables": word_syl,