Escape \ and " in typeahead query (#1306)

Author: papa99do

src/marqo/core/typeahead/typeahead.py

@@ -4,7 +4,7 @@
 
 import blake3
 
-from marqo.core.constants import MARQO_TYPEAHEAD_SCHEMA_MINIMUM_VERSION
+from marqo.core.constants import MARQO_TYPEAHEAD_SCHEMA_MINIMUM_VERSION, CHARACTERS_TO_BE_ESCAPED_IN_VESPA
 from marqo.core.index_management.index_management import IndexManagement
 from marqo.core.models.typeahead import (
     TypeaheadRequest, TypeaheadResponse, TypeaheadSuggestion,
@@ -66,19 +66,20 @@ def get_suggestions(self, index_name: str, request: TypeaheadRequest) -> Typeahe
             retrieval_terms = []
             ranking_terms = []
             for token in tokens:
+                escaped_token = self._escape_token(token)
                 if len(token) < request.min_fuzzy_match_length:
                     # Use exact prefix matching for short tokens
                     retrieval_terms.append(
-                        f"query_words contains ({{prefix:true}}\"{token}\")"
+                        f"query_words contains ({{prefix:true}}\"{escaped_token}\")"
                     )
                 else:
                     # Use fuzzy matching for longer tokens
                     retrieval_terms.append(
                         f"query_words contains "
-                        f"({{maxEditDistance:{request.fuzzy_edit_distance}, prefix:true}}fuzzy(\"{token}\"))"
+                        f"({{maxEditDistance:{request.fuzzy_edit_distance}, prefix:true}}fuzzy(\"{escaped_token}\"))"
                     )
 
-                ranking_terms.append(f"query_index contains \"{token}\"")
+                ranking_terms.append(f"query_index contains \"{escaped_token}\"")
 
             # Create single YQL query that ORs all token conditions
             yql_retrieval = " OR ".join(retrieval_terms)
@@ -309,6 +310,23 @@ def get_queries(self, index_name: str, queries: List[str]) -> TypeaheadGetQuerie
 
         return TypeaheadGetQueriesResponse(queries=query_results)
 
+    def _escape_token(self, token: str) -> str:
+        """Escape special characters in a token for Vespa YQL queries.
+
+        Args:
+            token: The token to escape
+
+        Returns:
+            The escaped token
+        """
+        escaped = []
+        for char in token:
+            if char in CHARACTERS_TO_BE_ESCAPED_IN_VESPA:
+                escaped.append('\\' + char)
+            else:
+                escaped.append(char)
+        return ''.join(escaped)
+
     def _generate_query_hash(self, query: str) -> str:
         """Generate a 128-bit blake3 hash for a query string.
 

tests/integ_tests/core/typeahead/test_typeahead_integration.py

@@ -456,6 +456,47 @@ def test_typeahead_version_check_with_old_version_index(self):
 
         self._assert_version_error_message(context.exception, self.index_220_name, "2.22.0")
 
+    # G. Special Character Handling Tests
+    def test_typeahead_with_special_characters_in_user_input(self):
+        """Test typeahead search handles special characters in user input without causing Vespa 500 errors."""
+        self._index_test_queries()
+
+        # Test user input queries with special characters that should NOT cause Vespa 500 errors
+        user_input_test_cases = [
+            # Single special characters
+            '"',
+            '\\',
+            # Queries with quotes
+            'a"b"c',
+            'a"b"',
+            '"b"c',
+            '"bc',
+            'bc"',
+            # Queries with backslashes
+            'Path\\to\\file',
+            '\\',
+            '\\\\',
+            'a\\b',
+            '\\a',
+            'b\\',
+            # Mixed special characters
+            'Program "with spaces"\\folder',
+            '"\\',
+            '\\"',
+        ]
+
+        for user_query in user_input_test_cases:
+            with self.subTest(user_query=user_query):
+                request = TypeaheadRequest(q=user_query, limit=10)
+                # This should NOT raise a 500 error from Vespa (main goal of the fix)
+                try:
+                    response = self.config.typeahead.get_suggestions(self.test_index_name, request)
+                    # Just verify we get a response without error
+                    self.assertIsNotNone(response)
+                    self.assertGreaterEqual(len(response.suggestions), 0)
+                except Exception as e:
+                    self.fail(f"User input query '{user_query}' caused an error: {e}")
+
     def _assert_version_error_message(self, exception: UnsupportedFeatureError, index_name: str, version: str):
         """Helper method to verify the error message contains expected information."""
         error_message = str(exception)

tests/unit_tests/marqo/core/typeahead/test_typeahead.py

@@ -782,5 +782,102 @@ def test_get_queries_mixed_found_not_found(self):
         self.assertEqual(result.queries[0].query, "found query")
 
 
+class TestTypeaheadEscaping(unittest.TestCase):
+    """Test cases for the special character escaping functionality."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        self.mock_vespa_client = Mock(spec=VespaClient)
+        self.mock_index_management = Mock(spec=IndexManagement)
+        self.typeahead = Typeahead(
+            vespa_client=self.mock_vespa_client,
+            index_management=self.mock_index_management
+        )
+
+    def test_escape_token_no_special_chars(self):
+        """Test escaping tokens with no special characters."""
+        # Simple strings should not be modified
+        test_cases = [
+            "hello",
+            "world",
+            "123",
+            "test_query",
+            "simple-text",
+            "with spaces",
+            "",
+        ]
+
+        for token in test_cases:
+            with self.subTest(token=token):
+                result = self.typeahead._escape_token(token)
+                self.assertEqual(result, token)
+
+    def test_escape_token_quotes(self):
+        """Test escaping tokens with double quotes."""
+        test_cases = [
+            ('"', '\\"'),
+            ('hello"world', 'hello\\"world'),
+            ('"start', '\\"start'),
+            ('end"', 'end\\"'),
+            ('""', '\\"\\"'),
+            ('say "hello"', 'say \\"hello\\"'),
+        ]
+
+        for input_token, expected in test_cases:
+            with self.subTest(input_token=input_token):
+                result = self.typeahead._escape_token(input_token)
+                self.assertEqual(result, expected)
+
+    def test_escape_token_backslashes(self):
+        """Test escaping tokens with backslashes."""
+        test_cases = [
+            ('\\', '\\\\'),
+            ('hello\\world', 'hello\\\\world'),
+            ('\\start', '\\\\start'),
+            ('end\\', 'end\\\\'),
+            ('\\\\', '\\\\\\\\'),
+            ('path\\to\\file', 'path\\\\to\\\\file'),
+        ]
+
+        for input_token, expected in test_cases:
+            with self.subTest(input_token=input_token):
+                result = self.typeahead._escape_token(input_token)
+                self.assertEqual(result, expected)
+
+    def test_escape_token_mixed_special_chars(self):
+        """Test escaping tokens with both quotes and backslashes."""
+        test_cases = [
+            ('"\\', '\\"\\\\'),
+            ('"hello\\world"', '\\"hello\\\\world\\"'),
+            ('C:\\Program Files\\"test"', 'C:\\\\Program Files\\\\\\"test\\"'),
+            ('a"b\\c"d', 'a\\"b\\\\c\\"d'),
+        ]
+
+        for input_token, expected in test_cases:
+            with self.subTest(input_token=input_token):
+                result = self.typeahead._escape_token(input_token)
+                self.assertEqual(result, expected)
+
+    def test_escape_token_preserves_other_chars(self):
+        """Test that escaping preserves other special characters."""
+        # Characters that are NOT in CHARACTERS_TO_BE_ESCAPED_IN_VESPA should be preserved
+        test_cases = [
+            "hello!world",
+            "[email protected]",
+            "price$100",
+            "50%off",
+            "a+b=c",
+            "question?",
+            "array[0]",
+            "function()",
+            "hash#tag",
+        ]
+
+        for token in test_cases:
+            with self.subTest(token=token):
+                result = self.typeahead._escape_token(token)
+                self.assertEqual(result, token)
+
+
 if __name__ == "__main__":
     unittest.main()