Allow for non-mutated user specified doc-ids (#977)

Co-authored-by: Copilot <[email protected]>

Author: mskarlin

paperqa/clients/journal_quality.py

@@ -45,12 +45,14 @@ async def _process(
         # remember, if both have docnames (i.e. key) they are
         # wiped and re-generated with resultant data
         return doc_details + DocDetails(
+            doc_id=doc_details.doc_id,  # ensure doc_id is preserved
+            dockey=doc_details.dockey,  # ensure dockey is preserved
             source_quality=max(
                 [
                     self.data.get(query.journal.casefold(), DocDetails.UNDEFINED_JOURNAL_QUALITY),  # type: ignore[union-attr]
                     self.data.get("the " + query.journal.casefold(), DocDetails.UNDEFINED_JOURNAL_QUALITY),  # type: ignore[union-attr]
                 ]
-            )
+            ),
         )
 
     def query_creator(self, doc_details: DocDetails, **kwargs) -> JournalQuery | None:

paperqa/clients/retractions.py

@@ -16,6 +16,9 @@
 
 
 class RetractionDataPostProcessor(MetadataPostProcessor[DOIQuery]):
+
+    RETRACTION_CACHE_DAYS: int = -1  # Number of days to cache, -1 is keep forever
+
     def __init__(self, retraction_data_path: os.PathLike | str | None = None) -> None:
 
         if retraction_data_path is None:
@@ -37,6 +40,8 @@ def __init__(self, retraction_data_path: os.PathLike | str | None = None) -> Non
         ]
 
     def _has_cache_expired(self) -> bool:
+        if self.RETRACTION_CACHE_DAYS < 0:
+            return False
         creation_time = os.path.getctime(self.retraction_data_path)
         file_creation_date = datetime.datetime.fromtimestamp(creation_time).replace(
             tzinfo=datetime.UTC
@@ -45,7 +50,7 @@ def _has_cache_expired(self) -> bool:
         current_time = datetime.datetime.now(datetime.UTC)
         time_difference = current_time - file_creation_date
 
-        return time_difference > datetime.timedelta(days=30)
+        return time_difference > datetime.timedelta(days=self.RETRACTION_CACHE_DAYS)
 
     def _is_csv_cached(self) -> bool:
         return os.path.exists(self.retraction_data_path)
@@ -71,7 +76,11 @@ async def _process(self, query: DOIQuery, doc_details: DocDetails) -> DocDetails
         if not self.doi_set:
             await self.load_data()
 
-        return doc_details + DocDetails(is_retracted=query.doi in self.doi_set)
+        return doc_details + DocDetails(
+            doc_id=doc_details.doc_id,  # ensure doc_id is preserved
+            dockey=doc_details.dockey,  # ensure dockey is preserved
+            is_retracted=query.doi in self.doi_set,
+        )
 
     def query_creator(self, doc_details: DocDetails, **kwargs) -> DOIQuery | None:
         try:

paperqa/types.py

@@ -431,14 +431,15 @@ def lowercase_doi_and_populate_doc_id(cls, data: dict[str, Any]) -> dict[str, An
                 if doi.startswith(url_prefix_to_remove):
                     doi = doi.replace(url_prefix_to_remove, "")
             data["doi"] = doi.lower()
-            data["doc_id"] = encode_id(doi.lower())
-        else:
+            if "doc_id" not in data or not data["doc_id"]:  # keep user defined doc_ids
+                data["doc_id"] = encode_id(doi.lower())
+        elif "doc_id" not in data or not data["doc_id"]:  # keep user defined doc_ids
             data["doc_id"] = encode_id(uuid4())
 
         if "dockey" in data.get(
             "fields_to_overwrite_from_metadata",
             DEFAULT_FIELDS_TO_OVERWRITE_FROM_METADATA,
-        ):
+        ) and ("dockey" not in data or not data["dockey"]):
             data["dockey"] = data["doc_id"]
 
         return data

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 (
+    encode_id,
     extract_score,
     get_citenames,
     maybe_get_date,
@@ -1488,6 +1489,150 @@ def test_docdetails_deserialization() -> None:
     ), "Deserialization should not mutate input"
 
 
+def test_docdetails_doc_id_roundtrip() -> None:
+    """Test that DocDetails can be initialized with doc_id or doi inputs."""
+    test_doi = "10.1234/test.doi"
+    test_doi_doc_id = encode_id(test_doi.lower())
+    test_specified_doc_id = "abc123"
+    # first we test without a doc_id or doi, ensure it's still valid
+    doc_details_no_doi_no_doc_id = DocDetails(
+        docname="test_doc",
+        citation="Test Citation",
+        dockey="test_dockey",
+        embedding=None,
+        formatted_citation="Formatted Test Citation",
+    )
+
+    assert (
+        doc_details_no_doi_no_doc_id.doc_id != test_doi_doc_id
+    ), "DocDetails without doc_id should not match test_doi_doc_id"
+    assert (
+        doc_details_no_doi_no_doc_id.doi is None
+    ), "DocDetails without doi should have None doi"
+    assert doc_details_no_doi_no_doc_id.dockey == doc_details_no_doi_no_doc_id.doc_id
+
+    # now round-trip serializaiton should keep the same doc_id
+    new_no_doi_no_doc_id = DocDetails(
+        **doc_details_no_doi_no_doc_id.model_dump(exclude_none=True)
+    )
+    assert (
+        new_no_doi_no_doc_id.doc_id == doc_details_no_doi_no_doc_id.doc_id
+    ), "DocDetails without doc_id should keep the same doc_id after serialization"
+
+    # since validation runs on assignment, make sure we can assign correctly
+    doc_details_no_doi_no_doc_id.doc_id = test_specified_doc_id
+    assert (
+        doc_details_no_doi_no_doc_id.doc_id == test_specified_doc_id
+    ), "DocDetails with doc_id should match test_specified_doc_id"
+    assert doc_details_no_doi_no_doc_id.dockey == doc_details_no_doi_no_doc_id.doc_id
+
+    # now let's do this with a doi
+    doc_details_with_doi_no_doc_id = DocDetails(
+        doi=test_doi,
+        docname="test_doc",
+        citation="Test Citation",
+        dockey="test_dockey",
+        embedding=None,
+        formatted_citation="Formatted Test Citation",
+    )
+    assert (
+        doc_details_with_doi_no_doc_id.doc_id == test_doi_doc_id
+    ), "DocDetails with doc_id should not match test_doi_doc_id"
+    assert (
+        doc_details_with_doi_no_doc_id.doi == test_doi
+    ), "DocDetails with doi should match test_doi"
+    assert (
+        doc_details_with_doi_no_doc_id.dockey == doc_details_with_doi_no_doc_id.doc_id
+    )
+
+    # round-trip serializaiton should keep the same doc_id
+    new_with_doi_no_doc_id = DocDetails(
+        **doc_details_with_doi_no_doc_id.model_dump(exclude_none=True)
+    )
+    assert (
+        new_with_doi_no_doc_id.doc_id == doc_details_with_doi_no_doc_id.doc_id
+    ), "DocDetails with doc_id should keep the same doc_id after serialization"
+
+    # since validation runs on assignment, make sure we can assign correctly
+    doc_details_with_doi_no_doc_id.doc_id = test_specified_doc_id
+    assert (
+        doc_details_with_doi_no_doc_id.doc_id == test_specified_doc_id
+    ), "DocDetails with doc_id should match test_specified_doc_id"
+    assert (
+        doc_details_with_doi_no_doc_id.dockey == doc_details_with_doi_no_doc_id.doc_id
+    )
+
+    # let's specify the doc_id directly
+    doc_details_no_doi_with_doc_id = DocDetails(
+        doc_id=test_specified_doc_id,
+        docname="test_doc",
+        citation="Test Citation",
+        dockey="test_dockey",
+        embedding=None,
+        formatted_citation="Formatted Test Citation",
+    )
+    assert (
+        doc_details_no_doi_with_doc_id.doc_id == test_specified_doc_id
+    ), "DocDetails with doc_id should not match test_specified_doc_id"
+    assert (
+        doc_details_no_doi_with_doc_id.doi is None
+    ), "DocDetails without doi should be None"
+    assert (
+        doc_details_no_doi_with_doc_id.dockey == doc_details_no_doi_with_doc_id.doc_id
+    ), "DocDetails dockey should match doc_id for the same object"
+
+    # round-trip serializaiton should keep the same doc_id
+    new_no_doi_with_doc_id = DocDetails(
+        **doc_details_no_doi_with_doc_id.model_dump(exclude_none=True)
+    )
+    assert (
+        new_no_doi_with_doc_id.doc_id == doc_details_with_doi_no_doc_id.doc_id
+    ), "DocDetails with doc_id should keep the same doc_id after serialization"
+
+    # since validation runs on assignment, make sure we can assign correctly
+    new_no_doi_with_doc_id.doc_id = test_doi_doc_id
+    assert (
+        new_no_doi_with_doc_id.doc_id == test_doi_doc_id
+    ), "DocDetails with doc_id should match test_specified_doc_id"
+    assert new_no_doi_with_doc_id.dockey == new_no_doi_with_doc_id.doc_id
+
+    # now we specify both doi and doc_id, ensuring doc_id takes precedence
+    doc_details_with_doi_with_doc_id = DocDetails(
+        doc_id=test_specified_doc_id,
+        doi=test_doi,
+        docname="test_doc",
+        citation="Test Citation",
+        dockey="test_dockey",
+        embedding=None,
+        formatted_citation="Formatted Test Citation",
+    )
+    assert (
+        doc_details_with_doi_with_doc_id.doc_id == test_specified_doc_id
+    ), "DocDetails with doc_id should not match test_specified_doc_id"
+    assert (
+        doc_details_with_doi_with_doc_id.doi == test_doi
+    ), "DocDetails without doi should match test_doi"
+    assert (
+        doc_details_with_doi_with_doc_id.dockey
+        == doc_details_with_doi_with_doc_id.doc_id
+    )
+
+    # round-trip serializaiton should keep the same doc_id
+    new_with_doi_with_doc_id = DocDetails(
+        **doc_details_with_doi_with_doc_id.model_dump(exclude_none=True)
+    )
+    assert (
+        new_with_doi_with_doc_id.doc_id == doc_details_with_doi_with_doc_id.doc_id
+    ), "DocDetails with doc_id should keep the same doc_id after serialization"
+
+    # since validation runs on assignment, make sure we can assign correctly
+    new_with_doi_with_doc_id.doc_id = test_doi_doc_id
+    assert (
+        new_with_doi_with_doc_id.doc_id == test_doi_doc_id
+    ), "DocDetails with doc_id should match test_specified_doc_id"
+    assert new_with_doi_with_doc_id.dockey == new_with_doi_with_doc_id.doc_id
+
+
 @pytest.mark.vcr
 @pytest.mark.parametrize("use_partition", [True, False])
 @pytest.mark.asyncio