diff --git a/src/paperless_ai/base_model.py b/src/paperless_ai/base_model.py index 8af146ee5..da7ebcb5f 100644 --- a/src/paperless_ai/base_model.py +++ b/src/paperless_ai/base_model.py @@ -6,6 +6,7 @@ from pydantic import BaseModel from pydantic import Field from pydantic import ValidationInfo from pydantic import field_validator +from pydantic import model_validator from pydantic.fields import FieldInfo # taxonomy.py MAX_TAG_CANDIDATES = 10, prompt is "up to 3 relevant dates" @@ -37,6 +38,27 @@ def _truncate_to_field_limit(value: Any, field: FieldInfo) -> Any: class TaxonomyChoice(BaseModel): """One field's suggestions: existing values to reuse, plus new ones to create.""" + @model_validator(mode="before") + @classmethod + def _normalize_flat_list(cls, value: Any) -> Any: + """Accept the flat list shape used before 3.1 and still emitted by + some smaller models despite the nested tool schema. Strings are new + names and integers are candidate IDs; the latter remain subject to + the shown-candidate allowlist in ai_classifier.py. + """ + if not isinstance(value, list): + return value + if not all( + isinstance(item, str) + or (isinstance(item, int) and not isinstance(item, bool)) + for item in value + ): + return value + return { + "existing_ids": [item for item in value if isinstance(item, int)], + "new_names": [item for item in value if isinstance(item, str)], + } + existing_ids: list[int] = Field( default_factory=list, max_length=MAX_EXISTING_IDS, diff --git a/src/paperless_ai/tests/test_base_model.py b/src/paperless_ai/tests/test_base_model.py index caf6ea32a..775938ce9 100644 --- a/src/paperless_ai/tests/test_base_model.py +++ b/src/paperless_ai/tests/test_base_model.py @@ -43,6 +43,43 @@ def test_document_classifier_schema_declared_defaults(): assert dumped["dates"] == [] +def test_flat_taxonomy_lists_are_normalized_for_legacy_model_responses(): + """ + GIVEN: + - A model response using the flat taxonomy lists accepted before 3.1 + - Strings, integer candidate IDs, and a mixture of both + WHEN: + - DocumentClassifierSchema validates the response + THEN: + - Strings become new_names and integers become existing_ids + + Some smaller models emit the old flat shape even when shown the nested + tool schema. Candidate IDs are still restricted to the IDs actually shown + to the model later in ai_classifier.py. + """ + parsed = DocumentClassifierSchema( + title="Electricity Bill", + tags=["Utilities", "Electricity"], + correspondents=[12], + document_types=[34, "Utility Bill"], + storage_paths=["Finance/Utilities"], + ) + + assert parsed.tags == TaxonomyChoice( + existing_ids=[], + new_names=["Utilities", "Electricity"], + ) + assert parsed.correspondents == TaxonomyChoice(existing_ids=[12], new_names=[]) + assert parsed.document_types == TaxonomyChoice( + existing_ids=[34], + new_names=["Utility Bill"], + ) + assert parsed.storage_paths == TaxonomyChoice( + existing_ids=[], + new_names=["Finance/Utilities"], + ) + + def test_document_classifier_schema_json_schema_is_self_contained(): """ GIVEN: