From 38df49776053ed74c4d34f74116087d4b1d2cc16 Mon Sep 17 00:00:00 2001 From: shamoon <4887959+shamoon@users.noreply.github.com> Date: Sun, 16 Aug 2026 19:55:03 -0700 Subject: [PATCH] Fix: DocumentClassifierSchema bounds --- src/paperless_ai/base_model.py | 46 +++++++++++-- src/paperless_ai/tests/test_base_model.py | 84 +++++++++++++++++++++++ 2 files changed, 126 insertions(+), 4 deletions(-) diff --git a/src/paperless_ai/base_model.py b/src/paperless_ai/base_model.py index 3df311035..c6d24bc5d 100644 --- a/src/paperless_ai/base_model.py +++ b/src/paperless_ai/base_model.py @@ -1,7 +1,32 @@ +from typing import Any +from typing import Final from typing import TypedDict from pydantic import BaseModel from pydantic import Field +from pydantic import ValidationInfo +from pydantic import field_validator +from pydantic.fields import FieldInfo + +# taxonomy.py MAX_TAG_CANDIDATES = 10, prompt is "up to 3 relevant dates" +MAX_EXISTING_IDS: Final = 10 +MAX_NEW_NAMES: Final = 8 +MAX_DATES: Final = 3 +# Matches documents.models.Document.title's CharField(max_length=128). +MAX_TITLE_LENGTH: Final = 128 + + +def _truncate_to_field_limit(value: Any, field: FieldInfo) -> Any: + """ + Clip down to its it's declared maximum. Run as a `mode="before"` validator. + """ + limit = next( + (m.max_length for m in field.metadata if hasattr(m, "max_length")), + None, + ) + if limit is None or not isinstance(value, (list, str)): + return value + return value[:limit] class TaxonomyChoice(BaseModel): @@ -14,19 +39,32 @@ class TaxonomyChoice(BaseModel): TaxonomyChoiceDict below. """ - existing_ids: list[int] = Field(default_factory=list) - new_names: list[str] = Field(default_factory=list) + existing_ids: list[int] = Field( + default_factory=list, + max_length=MAX_EXISTING_IDS, + ) + new_names: list[str] = Field(default_factory=list, max_length=MAX_NEW_NAMES) + + @field_validator("existing_ids", "new_names", mode="before") + @classmethod + def _truncate(cls, value: Any, info: ValidationInfo) -> Any: + return _truncate_to_field_limit(value, cls.model_fields[info.field_name]) class DocumentClassifierSchema(BaseModel): """Schema for document classification suggestions.""" - title: str + title: str = Field(max_length=MAX_TITLE_LENGTH) tags: TaxonomyChoice = Field(default_factory=TaxonomyChoice) correspondents: TaxonomyChoice = Field(default_factory=TaxonomyChoice) document_types: TaxonomyChoice = Field(default_factory=TaxonomyChoice) storage_paths: TaxonomyChoice = Field(default_factory=TaxonomyChoice) - dates: list[str] = Field(default_factory=list) + dates: list[str] = Field(default_factory=list, max_length=MAX_DATES) + + @field_validator("title", "dates", mode="before") + @classmethod + def _truncate(cls, value: Any, info: ValidationInfo) -> Any: + return _truncate_to_field_limit(value, cls.model_fields[info.field_name]) class TaxonomyChoiceDict(TypedDict): diff --git a/src/paperless_ai/tests/test_base_model.py b/src/paperless_ai/tests/test_base_model.py index 0ffda52cf..339fef0bd 100644 --- a/src/paperless_ai/tests/test_base_model.py +++ b/src/paperless_ai/tests/test_base_model.py @@ -1,3 +1,7 @@ +from paperless_ai.base_model import MAX_DATES +from paperless_ai.base_model import MAX_EXISTING_IDS +from paperless_ai.base_model import MAX_NEW_NAMES +from paperless_ai.base_model import MAX_TITLE_LENGTH from paperless_ai.base_model import ClassificationSuggestions from paperless_ai.base_model import DocumentClassifierSchema from paperless_ai.base_model import TaxonomyChoice @@ -63,6 +67,86 @@ def test_document_classifier_schema_json_schema_is_self_contained(): assert set(taxonomy_choice_properties.keys()) == {"existing_ids", "new_names"} +def test_every_sequence_in_the_emitted_schema_is_bounded(): + """ + GIVEN: + - The DocumentClassifierSchema pydantic model + WHEN: + - Its JSON schema is generated via model_json_schema() + THEN: + - Every array property in the schema, including those on the + referenced TaxonomyChoice definition, carries a maxItems + """ + schema = DocumentClassifierSchema.model_json_schema() + + unbounded = [ + f"{owner}.{name}" + for owner, definition in [ + ("DocumentClassifierSchema", schema), + *schema.get("$defs", {}).items(), + ] + for name, prop in definition.get("properties", {}).items() + if prop.get("type") == "array" and "maxItems" not in prop + ] + + assert unbounded == [] + + +def test_dates_bound_matches_what_the_prompt_asks_for(): + """ + GIVEN: + - The DocumentClassifierSchema pydantic model + WHEN: + - The emitted maxItems for dates is inspected + THEN: + - It equals the 3 that build_prompt_without_rag asks the model for + """ + dates_schema = DocumentClassifierSchema.model_json_schema()["properties"]["dates"] + + assert dates_schema["maxItems"] == MAX_DATES == 3 + + +def test_over_long_response_is_truncated_rather_than_rejected(): + """ + GIVEN: + - An LLM response overshooting every declared bound + WHEN: + - DocumentClassifierSchema is constructed from it + THEN: + - Each field is clipped to its maximum, with no ValidationError + """ + parsed = DocumentClassifierSchema( + title="T" * (MAX_TITLE_LENGTH + 50), + tags=TaxonomyChoice( + existing_ids=list(range(MAX_EXISTING_IDS + 20)), + new_names=["n"] * (MAX_NEW_NAMES + 20), + ), + dates=[f"2016-{month:02d}-01" for month in range(1, 13)], + ) + + assert len(parsed.title) == MAX_TITLE_LENGTH + assert len(parsed.dates) == MAX_DATES + assert len(parsed.tags.existing_ids) == MAX_EXISTING_IDS + assert len(parsed.tags.new_names) == MAX_NEW_NAMES + + +def test_truncation_keeps_the_earliest_entries(): + """ + GIVEN: + - An over-long dates list from an LLM response + WHEN: + - DocumentClassifierSchema is constructed from it + THEN: + - The kept entries are the first ones the model emitted + """ + parsed = DocumentClassifierSchema( + title="T", + dates=["2016-10-01", "2016-09-01", "2016-08-01", "2016-07-01", "2016-06-01"], + ) + + assert parsed.dates == ["2016-10-01", "2016-09-01", "2016-08-01"] + + def test_model_dump_matches_typed_dict_keys(): """ GIVEN: