Fixing crash on possessive query (#992)

Author: jamesbraza

paperqa/agents/search.py

@@ -46,7 +46,7 @@
 from paperqa.docs import Docs
 from paperqa.settings import IndexSettings, get_settings
 from paperqa.types import VAR_MATCH_LOOKUP, DocDetails
-from paperqa.utils import ImpossibleParsingError, hexdigest
+from paperqa.utils import ImpossibleParsingError, clean_possessives, hexdigest
 
 from .models import JSONType, SupportsPickle
 
@@ -408,9 +408,8 @@ async def get_saved_object(
                 return self.storage.read_from_string(content)
         return None
 
-    def clean_query(self, query: str) -> str:
-        # SEE: https://regex101.com/r/DoLMoa/3
-        return re.sub(r'[*\[\]:(){}~^><+"\\]', "", query)
+    # Remove these characters, SEE: https://regex101.com/r/DoLMoa/3
+    CLEAN_QUERY_REGEX: ClassVar[re.Pattern] = re.compile(r'[*\[\]:(){}~^><+"\\]')
 
     async def query(
         self,
@@ -424,13 +423,17 @@ async def query(
         query_fields = list(field_subset or self.fields)
         searcher = await self.searcher
         index = await self.index
+        cleaned_query = self.CLEAN_QUERY_REGEX.sub("", query)
+        try:
+            parsed_query = index.parse_query(cleaned_query, query_fields)
+        except ValueError:  # Rejected by tantivy
+            # Retry with more aggressive cleaning
+            parsed_query = index.parse_query(
+                clean_possessives(cleaned_query), query_fields
+            )
         addresses = [
             s[1]
-            for s in searcher.search(
-                index.parse_query(self.clean_query(query), query_fields),
-                top_n,
-                offset=offset,
-            ).hits
+            for s in searcher.search(parsed_query, top_n, offset=offset).hits
             if s[0] > min_score
         ]
         search_index_docs = [searcher.doc(address) for address in addresses]

paperqa/utils.py

@@ -564,3 +564,20 @@ def maybe_get_date(date: str | datetime | None) -> datetime | None:
                 continue
         return None
     return date
+
+
+def clean_possessives(text: str) -> str:
+    """Remove possessive apostrophes from text (e.g. "X's Y" to "Xs Y")."""
+    # Handle apostrophes after word characters
+    # (possessive 's or trailing apostrophes)
+    text = re.sub(
+        r"(?<=\w)'(?:s\b|(?=\s|$))",
+        lambda m: "s" if m.group().endswith("s") else "",
+        text,
+    )
+    # Remove standalone 's patterns
+    text = re.sub(r"\s+'s\b", "", text)
+    text = re.sub(r"^'s\s*", "", text)
+    # Remove standalone apostrophes
+    text = re.sub(r"\s+'\s+", " ", text)
+    return re.sub(r"(?<!\w)'\s*", "", text)

tests/test_agents.py

@@ -105,6 +105,14 @@ async def test_get_directory_index(
                 f" {[d.formatted_citation for d in results[0].docs.values()]}."
             )
 
+            # Check single quoted text in the query doesn't crash us
+            results = await index.query(query="Who is 'Bates'")
+            assert results
+
+            # Check possessive in the query doesn't crash us
+            results = await index.query(query="What is Bates' first name")
+            assert results
+
         with subtests.test(msg="check-md-query"):
             results = await index.query(query="what is a gravity hill?", min_score=5)
             assert results

tests/test_paperqa.py

@@ -50,6 +50,7 @@
 from paperqa.readers import parse_pdf_to_pages, read_doc
 from paperqa.types import ChunkMetadata
 from paperqa.utils import (
+    clean_possessives,
     encode_id,
     extract_score,
     maybe_get_date,
@@ -1980,3 +1981,29 @@ def test_maybe_get_date():
     assert maybe_get_date(datetime(2023, 1, 1)) == datetime(2023, 1, 1)
     assert maybe_get_date("foo") is None
     assert maybe_get_date("") is None
+
+
+@pytest.mark.parametrize(
+    ("raw_text", "cleaned_text"),
+    [
+        ("name", "name"),
+        (" name", " name"),
+        ("name ", "name "),
+        (" ", " "),
+        ("Bates name", "Bates name"),
+        ("Bate's name", "Bates name"),
+        ("Bate's name Bate's name", "Bates name Bates name"),
+        ("Bates' name", "Bates name"),
+        ("X's Y", "Xs Y"),
+        ("' name", "name"),
+        (" ' name", " name"),
+        ("name ' name", "name name"),
+        ("'s name", "name"),
+        (" 's name", " name"),
+        ("s' name", "s name"),
+        ("S' name", "S name"),
+        ("Bates 's name", "Bates name"),
+    ],
+)
+def test_clean_possessives(raw_text: str, cleaned_text: str) -> None:
+    assert clean_possessives(raw_text) == cleaned_text