Added the get_manifest_kwargs function to get more than doi and title… (#848)

Co-authored-by: pre-commit-ci-lite[bot] <117423508+pre-commit-ci-lite[bot]@users.noreply.github.com>

Author: JoaquinPolonuer

paperqa/agents/search.py

@@ -10,7 +10,7 @@
 import warnings
 import zlib
 from collections import Counter
-from collections.abc import AsyncIterator, Callable, Collection, Sequence
+from collections.abc import AsyncIterator, Callable, Sequence
 from datetime import datetime
 from enum import StrEnum, auto
 from typing import TYPE_CHECKING, Any, ClassVar
@@ -44,7 +44,7 @@
 
 from paperqa.docs import Docs
 from paperqa.settings import IndexSettings, get_settings
-from paperqa.types import DocDetails
+from paperqa.types import VAR_MATCH_LOOKUP, DocDetails
 from paperqa.utils import ImpossibleParsingError, hexdigest
 
 from .models import SupportsPickle
@@ -108,14 +108,12 @@ def read_from_string(self, data: str | bytes) -> BaseModel | SupportsPickle:
         return pickle.loads(data)  # type: ignore[arg-type] # noqa: S301
 
 
-ENV_VAR_MATCH: Collection[str] = {"1", "true"}
-
 # Cache keys are a two-tuple of index name and absolute index directory
 # Cache values are a two-tuple of an opened Index instance and the count
 # of SearchIndex instances currently referencing that Index
 _OPENED_INDEX_CACHE: dict[tuple[str, str], tuple[Index, int]] = {}
 DONT_USE_OPENED_INDEX_CACHE = (
-    os.environ.get("PQA_INDEX_DONT_CACHE_INDEXES", "").lower() in ENV_VAR_MATCH
+    os.environ.get("PQA_INDEX_DONT_CACHE_INDEXES", "").lower() in VAR_MATCH_LOOKUP
 )
 
 
@@ -429,6 +427,17 @@ async def query(
         ]
 
 
+def fetch_kwargs_from_manifest(
+    file_location: str, manifest: dict[str, Any], manifest_fallback_location: str
+) -> dict[str, Any]:
+    manifest_entry: DocDetails | None = manifest.get(file_location) or manifest.get(
+        manifest_fallback_location
+    )
+    if manifest_entry:
+        return manifest_entry.model_dump()
+    return {}
+
+
 async def maybe_get_manifest(
     filename: anyio.Path | None = None,
 ) -> dict[str, DocDetails]:
@@ -491,23 +500,17 @@ async def process_file(
         if not await search_index.filecheck(filename=file_location):
             logger.info(f"New file to index: {file_location}...")
 
-            doi, title = None, None
-            if file_location in manifest:
-                manifest_entry = manifest[file_location]
-                doi, title = manifest_entry.doi, manifest_entry.title
-            elif manifest_fallback_location in manifest:
-                # Perhaps manifest used the opposite pathing scheme
-                manifest_entry = manifest[manifest_fallback_location]
-                doi, title = manifest_entry.doi, manifest_entry.title
+            kwargs = fetch_kwargs_from_manifest(
+                file_location, manifest, manifest_fallback_location
+            )
 
             tmp_docs = Docs()
             try:
                 await tmp_docs.aadd(
                     path=abs_file_path,
-                    title=title,
-                    doi=doi,
                     fields=["title", "author", "journal", "year"],
                     settings=settings,
+                    **kwargs,
                 )
             except Exception as e:
                 # We handle any exception here because we want to save_index so we
@@ -569,10 +572,10 @@ def _make_progress_bar_update(
 ) -> tuple[contextlib.AbstractContextManager, Callable[[], Any] | None]:
     # Disable should override enable
     env_var_disable = (
-        os.environ.get("PQA_INDEX_DISABLE_PROGRESS_BAR", "").lower() in ENV_VAR_MATCH
+        os.environ.get("PQA_INDEX_DISABLE_PROGRESS_BAR", "").lower() in VAR_MATCH_LOOKUP
     )
     env_var_enable = (
-        os.environ.get("PQA_INDEX_ENABLE_PROGRESS_BAR", "").lower() in ENV_VAR_MATCH
+        os.environ.get("PQA_INDEX_ENABLE_PROGRESS_BAR", "").lower() in VAR_MATCH_LOOKUP
     )
     try:
         is_cli = is_running_under_cli()  # pylint: disable=used-before-assignment

paperqa/clients/__init__.py

@@ -201,7 +201,7 @@ async def upgrade_doc_to_doc_details(self, doc: Doc, **kwargs) -> DocDetails:
         # note we have some extra fields which may have come from reading the doc text,
         # but aren't in the doc object, we add them here too.
         extra_fields = {
-            k: v for k, v in kwargs.items() if k in {"title", "authors", "doi"}
+            k: v for k, v in kwargs.items() if k in set(DocDetails.model_fields)
         }
         # abuse our doc_details object to be an int if it's empty
         # our __add__ operation supports int by doing nothing
@@ -210,18 +210,21 @@ async def upgrade_doc_to_doc_details(self, doc: Doc, **kwargs) -> DocDetails:
         )
 
         if doc_details := await self.query(**kwargs):
-            if doc.overwrite_fields_from_metadata:
-                return extra_doc + doc_details
 
             # hard overwrite the details from the prior object
-            doc_details.dockey = doc.dockey
-            doc_details.doc_id = doc.dockey
-            doc_details.docname = doc.docname
-            doc_details.key = doc.docname
-            doc_details.citation = doc.citation
+            if "dockey" in doc.fields_to_overwrite_from_metadata:
+                doc_details.dockey = doc.dockey
+            if "doc_id" in doc.fields_to_overwrite_from_metadata:
+                doc_details.doc_id = doc.dockey
+            if "docname" in doc.fields_to_overwrite_from_metadata:
+                doc_details.docname = doc.docname
+            if "key" in doc.fields_to_overwrite_from_metadata:
+                doc_details.key = doc.docname
+            if "citation" in doc.fields_to_overwrite_from_metadata:
+                doc_details.citation = doc.citation
             return extra_doc + doc_details
 
         # if we can't get metadata, just return the doc, but don't overwrite any fields
         prior_doc = doc.model_dump()
-        prior_doc["overwrite_fields_from_metadata"] = False
+        prior_doc["fields_to_overwrite_from_metadata"] = set()
         return DocDetails(**(prior_doc | extra_fields))

paperqa/docs.py

@@ -347,7 +347,7 @@ async def aadd(  # noqa: PLR0912
                 )
         # see if we can upgrade to DocDetails
         # if not, we can progress with a normal Doc
-        # if "overwrite_fields_from_metadata" is used:
+        # if "fields_to_overwrite_from_metadata" is used:
         # will map "docname" to "key", and "dockey" to "doc_id"
         if (title or doi) and parse_config.use_doc_details:
             if kwargs.get("metadata_client"):

paperqa/sources/clinical_trials.py

@@ -170,7 +170,7 @@ def format_to_doc_details(trial_data: dict) -> DocDetails:
         year=year or None,
         citation=citation,
         other={"client_source": [CLINICAL_TRIALS_BASE]},
-        overwrite_fields_from_metadata=False,
+        fields_to_overwrite_from_metadata=set(),
     )
 
 
@@ -311,7 +311,7 @@ async def add_clinical_trials_to_docs(
         year=datetime.now().year,
         citation=f"Clinical Trials Search via ClinicalTrials.gov: {query}",
         other={"client_source": [CLINICAL_TRIALS_BASE]},
-        overwrite_fields_from_metadata=False,
+        fields_to_overwrite_from_metadata=set(),
     )
 
     await docs.aadd_texts(

paperqa/types.py

@@ -40,17 +40,26 @@
 logger = logging.getLogger(__name__)
 
 
+VAR_MATCH_LOOKUP: Collection[str] = {"1", "true"}
+VAR_MISMATCH_LOOKUP: Collection[str] = {"0", "false"}
+DEFAULT_FIELDS_TO_OVERWRITE_FROM_METADATA: Collection[str] = {
+    "key",
+    "doc_id",
+    "docname",
+    "dockey",
+    "citation",
+}
+
+
 class Doc(Embeddable):
     model_config = ConfigDict(extra="forbid")
 
     docname: str
     dockey: DocKey
     citation: str
-    overwrite_fields_from_metadata: bool = Field(
-        default=True,
-        description=(
-            "flag to overwrite fields from metadata when upgrading to a DocDetails"
-        ),
+    fields_to_overwrite_from_metadata: set[str] = Field(
+        default_factory=lambda: set(DEFAULT_FIELDS_TO_OVERWRITE_FROM_METADATA),
+        description="fields from metadata to overwrite when upgrading to a DocDetails",
     )
 
     @model_validator(mode="before")
@@ -410,7 +419,10 @@ def lowercase_doi_and_populate_doc_id(cls, data: dict[str, Any]) -> dict[str, An
         else:
             data["doc_id"] = encode_id(uuid4())
 
-        if data.get("overwrite_fields_from_metadata", True):
+        if "dockey" in data.get(
+            "fields_to_overwrite_from_metadata",
+            DEFAULT_FIELDS_TO_OVERWRITE_FROM_METADATA,
+        ):
             data["dockey"] = data["doc_id"]
 
         return data
@@ -502,10 +514,13 @@ def overwrite_docname_dockey_for_compatibility_w_doc(
     ) -> dict[str, Any]:
         """Overwrite fields from metadata if specified."""
         overwrite_fields = {"key": "docname", "doc_id": "dockey"}
-        if data.get("overwrite_fields_from_metadata", True):
-            for field, old_field in overwrite_fields.items():
-                if data.get(field):
-                    data[old_field] = data[field]
+        fields_to_overwrite = data.get(
+            "fields_to_overwrite_from_metadata",
+            DEFAULT_FIELDS_TO_OVERWRITE_FROM_METADATA,
+        )
+        for field in overwrite_fields.keys() & fields_to_overwrite:
+            if data.get(field):
+                data[overwrite_fields[field]] = data[field]
         return data
 
     @classmethod
@@ -516,7 +531,7 @@ def populate_bibtex_key_citation(  # noqa: PLR0912
 
         Missing values, 'unknown' keys, and incomplete bibtex entries are regenerated.
 
-        When overwrite_fields_from_metadata:
+        When fields_to_overwrite_from_metadata:
             If bibtex is regenerated, the citation field is also regenerated.
 
             Otherwise we keep the citation field as is.
@@ -529,7 +544,10 @@ def populate_bibtex_key_citation(  # noqa: PLR0912
                 data.get("year") or CITATION_FALLBACK_DATA["year"],  # type: ignore[arg-type]
                 data.get("title") or CITATION_FALLBACK_DATA["title"],  # type: ignore[arg-type]
             )
-            if data.get("overwrite_fields_from_metadata", True):
+            if "docname" in data.get(
+                "fields_to_overwrite_from_metadata",
+                DEFAULT_FIELDS_TO_OVERWRITE_FROM_METADATA,
+            ):
                 data["docname"] = data["key"]
 
         # even if we have a bibtex, it may not be complete, thus we need to add to it
@@ -591,26 +609,35 @@ def populate_bibtex_key_citation(  # noqa: PLR0912
                     entries={data["key"]: new_entry}
                 ).to_string("bibtex")
                 # clear out the citation, since it will be regenerated
-                if data.get("overwrite_fields_from_metadata", True):
+                if "citation" in data.get(
+                    "fields_to_overwrite_from_metadata",
+                    DEFAULT_FIELDS_TO_OVERWRITE_FROM_METADATA,
+                ):
                     data["citation"] = None
             except Exception:
                 logger.warning(
                     "Failed to generate bibtex for"
                     f" {data.get('docname') or data.get('citation')}"
                 )
-        if not data.get("citation") and data.get("bibtex") is not None:
+        if data.get("citation") is None and data.get("bibtex") is not None:
             data["citation"] = format_bibtex(
                 data["bibtex"], missing_replacements=CITATION_FALLBACK_DATA  # type: ignore[arg-type]
             )
-        elif not data.get("citation"):
+        elif data.get("citation") is None:
             data["citation"] = data.get("title") or CITATION_FALLBACK_DATA["title"]
         return data
 
     @model_validator(mode="before")
     @classmethod
     def validate_all_fields(cls, data: Mapping[str, Any]) -> dict[str, Any]:
+
         data = deepcopy(data)  # Avoid mutating input
         data = dict(data)
+        if isinstance(data.get("fields_to_overwrite_from_metadata"), str):
+            data["fields_to_overwrite_from_metadata"] = {
+                s.strip()
+                for s in data.get("fields_to_overwrite_from_metadata", "").split(",")
+            }
         data = cls.lowercase_doi_and_populate_doc_id(data)
         data = cls.remove_invalid_authors(data)
         data = cls.misc_string_cleaning(data)

tests/stub_data/stub_manifest_nocitation.csv

@@ -0,0 +1,2 @@
+file_location,doi,title,citation,fields_to_overwrite_from_metadata
+"bates.txt",,"Frederick Bates (Wikipedia article)","","key, doc_id, docname, dockey"

tests/test_agents.py

@@ -243,10 +243,31 @@ async def test_get_directory_index_w_manifest(agent_test_settings: Settings) ->
 
         results = await index.query(query="who is Frederick Bates?")
         top_result = next(iter(results[0].docs.values()))
-        assert top_result.dockey == md5sum(abs_paper_dir / "bates.txt")
+
+        # note: we get every possible field from the manifest constructed in maybe_get_manifest,
+        # and then DocDetails construction sets the dockey to the doc_id.
+        assert top_result.dockey == top_result.doc_id
         # note: this title comes from the manifest, so we know it worked
         assert top_result.title == "Frederick Bates (Wikipedia article)"
 
+        assert "wikipedia article" in top_result.citation.lower(), (
+            "Other tests check we can override citation,"
+            " so here we check here it's actually populated"
+        )
+
+
+@pytest.mark.asyncio
+async def test_get_directory_index_w_no_citations(
+    agent_test_settings: Settings,
+) -> None:
+    agent_test_settings.agent.index.manifest_file = "stub_manifest_nocitation.csv"
+    index = await get_directory_index(settings=agent_test_settings)
+
+    results = await index.query(query="who is Frederick Bates?")
+    top_result = next(iter(results[0].docs.values()))
+
+    assert not top_result.citation
+
 
 @pytest.mark.flaky(reruns=2, only_rerun=["AssertionError", "httpx.RemoteProtocolError"])
 @pytest.mark.parametrize("agent_type", [FAKE_AGENT_TYPE, ToolSelector, SimpleAgent])

tests/test_paperqa.py

@@ -1324,7 +1324,6 @@ def test_docdetails_deserialization() -> None:
         "docname": "Stub",
         "embedding": None,
         "formatted_citation": "stub",
-        "overwrite_fields_from_metadata": True,
     }
     deepcopy_deserialize_to_doc = deepcopy(deserialize_to_doc)
     doc = Doc(**deserialize_to_doc)
@@ -1338,7 +1337,6 @@ def test_docdetails_deserialization() -> None:
     for key, value in {
         "docname": "unknownauthorsUnknownyearunknowntitle",
         "citation": "Unknown authors. Unknown title. Unknown journal, Unknown year.",
-        "overwrite_fields_from_metadata": True,
         "key": "unknownauthorsUnknownyearunknowntitle",
         "bibtex": (
             '@article{unknownauthorsUnknownyearunknowntitle,\n    author = "authors,'