Fixed duplicate keys

Author: whitead

paperqa/docs.py

@@ -39,14 +39,34 @@ class Docs:
     def __init__(self, chunk_size_limit: int = 3000) -> None:
         self.docs = dict()
         self.chunk_size_limit = chunk_size_limit
+        self.keys = set()
 
     def add(self, path: str, citation: str, key: Optional[str] = None) -> bool:
         """Add a document to the collection."""
+        if path in self.docs:
+            return False
         if key is None:
             # get first name and year from citation
-            author = re.search(r"([A-Z][a-z]+)", citation).group(1)
-            year = re.search(r"(\d{4})", citation).group(1)
+            try:
+                author = re.search(r"([A-Z][a-z]+)", citation).group(1)
+            except AttributeError:
+                # panicking - no word??
+                return False
+            try:
+                year = re.search(r"(\d{4})", citation).group(1)
+            except AttributeError:
+                year = ""
             key = f"{author}{year}"
+        suffix = ""
+        while key + suffix in self.keys:
+            # move suffix to next letter
+            if suffix == "":
+                suffix = "a"
+            else:
+                suffix = chr(ord(suffix) + 1)
+        key += suffix
+        self.keys.add(key)
+
         data = {"citation": citation, "key": key}
         d = gpt_index.SimpleDirectoryReader(input_files=[path]).load_data()
         # loose check to see if document was loaded
@@ -124,7 +144,8 @@ def query(self, query: str, k: int = 3, max_sources: int = 5):
         else:
             answer = qa_chain.run(question=query, context_str=context_str)[1:]
         for data in self.docs.values():
-            if data["key"] in answer:
+            # do check for whole key (so we don't catch Callahan2019a with Callahan2019)
+            if data["key"] + " " in answer or data["key"] + ")" in answer:
                 bib.append(f'{data["key"]}: {data["citation"]}')
         bib_str = "\n\n".join(bib)
         formatted_answer = f"Question: {query}\n\n{answer}\n"

paperqa/qaprompts.py

@@ -32,14 +32,14 @@
 
 _qa_prompt = prompts.PromptTemplate(
     input_variables=["question", "context_str"],
-    template="Generate a comprehensive answer (about 100 words) "
+    template="Write a comprehensive answer (about 100 words) "
     "for the question below solely based on the provided context. "
-    "If the context is insufficient to "
+    "If the context is insufficient "
     'answer, reply "I cannot answer". '
     "For each sentence in your answer, indicate which sources most support it "
-    "via valid citation markers at the end of sentences like (Foo2012). "
+    "via valid citation markers at the end of sentences, like (Foo2012). "
     "Answer in an unbiased, balanced, and scientific tone. "
-    # "To answer, start by writing out the reasoning and then "
+    "Try to use the direct quotes, if present, from the context. "
     # "write a complete unbiased answer prefixed by \"Answer:\""
     "\n--------------------\n"
     "{context_str}\n"

paperqa/version.py

@@ -1 +1 @@
-__version__ = "0.0.2"
+__version__ = "0.0.3"

tests/test_paperqa.py

@@ -33,7 +33,7 @@ def test_docs():
 def test_evidence():
     doc_path = "example.txt"
     with open(doc_path, "w", encoding="utf-8") as f:
-        # get wiki page about Obama
+        # get wiki page about politician
         r = requests.get("https://en.wikipedia.org/wiki/Frederick_Bates_(politician)")
         f.write(r.text)
     docs = paperqa.Docs()
@@ -48,7 +48,7 @@ def test_evidence():
 def test_query():
     doc_path = "example.txt"
     with open(doc_path, "w", encoding="utf-8") as f:
-        # get wiki page about Obama
+        # get wiki page about politician
         r = requests.get("https://en.wikipedia.org/wiki/Frederick_Bates_(politician)")
         f.write(r.text)
     docs = paperqa.Docs()
@@ -73,3 +73,46 @@ def test_docs_pickle():
         "What is today?", k=1, max_sources=1
     ) == docs2.get_evidence("What is today?", k=1, max_sources=1)
     os.remove(doc_path)
+
+
+def test_bad_context():
+    doc_path = "example.txt"
+    with open(doc_path, "w", encoding="utf-8") as f:
+        # get wiki page about politician
+        r = requests.get("https://en.wikipedia.org/wiki/Frederick_Bates_(politician)")
+        f.write(r.text)
+    docs = paperqa.Docs()
+    docs.add(doc_path, "WikiMedia Foundation, 2023, Accessed now")
+    answer = docs.query("What year was Barack Obama born greatest accomplishment?")
+    assert (
+        answer.answer
+        == "I cannot answer this question due to insufficient information."
+    )
+    os.remove(doc_path)
+
+
+def test_repeat_keys():
+    doc_path = "example.txt"
+    with open(doc_path, "w", encoding="utf-8") as f:
+        # get wiki page about politician
+        r = requests.get("https://en.wikipedia.org/wiki/Frederick_Bates_(politician)")
+        f.write(r.text)
+    docs = paperqa.Docs()
+    docs.add(doc_path, "WikiMedia Foundation, 2023, Accessed now")
+    docs.add(doc_path, "WikiMedia Foundation, 2023, Accessed now")
+    assert len(docs.docs) == 1
+
+    # now with different paths
+    doc_path2 = "example2.txt"
+    with open(doc_path2, "w", encoding="utf-8") as f:
+        # get wiki page about politician
+        f.write(r.text)
+    docs.add(doc_path2, "WikiMedia Foundation, 2023, Accessed now")
+    assert len(docs.docs) == 2
+
+    # check keys
+    assert docs.docs[doc_path]["key"] == "Wiki2023"
+    assert docs.docs[doc_path2]["key"] == "Wiki2023a"
+
+    os.remove(doc_path)
+    os.remove(doc_path2)