From b89fb0f9789443d596814ee59034ce33728babaa Mon Sep 17 00:00:00 2001 From: shamoon <4887959+shamoon@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:03:30 -0700 Subject: [PATCH] Fix: 3.1.0 llm suggestions simplify schema, fix docstrings (#13850) --- src/paperless_ai/base_model.py | 100 ++++++++++++++++--- src/paperless_ai/prompts/assigned_block.j2 | 2 +- src/paperless_ai/prompts/classification.j2 | 19 ++-- src/paperless_ai/prompts/taxonomy_block.j2 | 2 +- src/paperless_ai/tests/test_ai_classifier.py | 28 ++++-- src/paperless_ai/tests/test_base_model.py | 97 +++++++++++++++--- 6 files changed, 196 insertions(+), 52 deletions(-) diff --git a/src/paperless_ai/base_model.py b/src/paperless_ai/base_model.py index c5430e8f6..8af146ee5 100644 --- a/src/paperless_ai/base_model.py +++ b/src/paperless_ai/base_model.py @@ -31,21 +31,31 @@ def _truncate_to_field_limit(value: Any, field: FieldInfo) -> Any: ) +# Docstrings and field descriptions on both models below are serialized into +# the schema handed to the LLM, so write them for the model. Code comments +# should go here only. class TaxonomyChoice(BaseModel): - """One taxonomy category's suggestions: IDs the model matched to a - candidate it was shown in the prompt, plus names for values it believes - are genuinely new. existing_ids are never localized - only new_names is. - - Pydantic enforces this shape on whatever the LLM returns; the rest of the - pipeline passes the `.model_dump()`-ed plain dict around, typed as - TaxonomyChoiceDict below. - """ + """One field's suggestions: existing values to reuse, plus new ones to create.""" existing_ids: list[int] = Field( default_factory=list, max_length=MAX_EXISTING_IDS, + description=( + "IDs from the candidate list shown in the prompt that clearly " + "represent values you would suggest for this field. Never invent " + "an ID, select a weak match merely because it exists, or use an " + "ID when no candidates are shown." + ), + ) + new_names: list[str] = Field( + default_factory=list, + max_length=MAX_NEW_NAMES, + description=( + "Names for clearly supported values that no shown candidate " + "represents. When a candidate represents the same value, use its " + "ID instead so an existing value is not duplicated under a new name." + ), ) - new_names: list[str] = Field(default_factory=list, max_length=MAX_NEW_NAMES) @field_validator("existing_ids", "new_names", mode="before") @classmethod @@ -54,20 +64,78 @@ class TaxonomyChoice(BaseModel): class DocumentClassifierSchema(BaseModel): - """Schema for document classification suggestions.""" + """Classification suggestions for a single document.""" - 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, max_length=MAX_DATES) + title: str = Field( + max_length=MAX_TITLE_LENGTH, + description=( + "A short, descriptive title for this document, at most " + f"{MAX_TITLE_LENGTH} characters." + ), + ) + tags: TaxonomyChoice = Field( + default_factory=TaxonomyChoice, + description=( + "Topic labels describing what this document is about. A document " + "may have several, e.g. 'Insurance', 'Car', 'Warranty'." + ), + ) + correspondents: TaxonomyChoice = Field( + default_factory=TaxonomyChoice, + description=( + "The person, institution or company this document originates " + "from, or was sent to. Not every party merely mentioned in the " + "text, and not the subject of the document." + ), + ) + document_types: TaxonomyChoice = Field( + default_factory=TaxonomyChoice, + description=( + "What kind of document this is, e.g. 'Invoice', 'Contract', " + "'Bank Statement', 'Letter'. Never its subject matter and never " + "who sent it." + ), + ) + storage_paths: TaxonomyChoice = Field( + default_factory=TaxonomyChoice, + description=( + "A folder-style filing location for this document, e.g. " + "'Finance/Invoices'. Leave empty unless a filing location is " + "clearly implied - never put tags, document types or " + "correspondents here." + ), + ) + dates: list[str] = Field( + default_factory=list, + max_length=MAX_DATES, + description=( + f"Up to {MAX_DATES} dates relevant to this document, each " + "formatted YYYY-MM-DD. The most important is the date the " + "document was issued." + ), + ) @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]) + @classmethod + def model_json_schema(cls, *args: Any, **kwargs: Any) -> dict[str, Any]: + """Inline TaxonomyChoice for backends that reject JSON Schema refs.""" + schema = super().model_json_schema(*args, **kwargs) + taxonomy_choice = schema.pop("$defs")["TaxonomyChoice"] + for field in ("tags", "correspondents", "document_types", "storage_paths"): + # Pydantic emits a field's description as a sibling of its $ref; + # those keys must survive and win over the shared definition. + siblings = { + key: value + for key, value in schema["properties"][field].items() + if key != "$ref" + } + schema["properties"][field] = taxonomy_choice | siblings + return schema + class TaxonomyChoiceDict(TypedDict): """Plain-dict counterpart of TaxonomyChoice - what diff --git a/src/paperless_ai/prompts/assigned_block.j2 b/src/paperless_ai/prompts/assigned_block.j2 index d19b81a24..547ef015e 100644 --- a/src/paperless_ai/prompts/assigned_block.j2 +++ b/src/paperless_ai/prompts/assigned_block.j2 @@ -1,4 +1,4 @@ -This document's existing metadata (already assigned; use as context for the title and for any fields below still empty - do not re-suggest these values): +This document's existing metadata (already assigned). Use it as context for your suggestions: Tags: {{ tags | join(', ') if tags else '(none)' }} Document Type: {{ document_type or '(not set)' }} Correspondent: {{ correspondent or '(not set)' }} diff --git a/src/paperless_ai/prompts/classification.j2 b/src/paperless_ai/prompts/classification.j2 index e8f2956dc..27501fe50 100644 --- a/src/paperless_ai/prompts/classification.j2 +++ b/src/paperless_ai/prompts/classification.j2 @@ -4,16 +4,19 @@ You are a document classification assistant. {{ taxonomy_block }} {% endif %} -Analyze the following document and extract the following information: -- A short descriptive title -- Tags that reflect the content -- Names of people or organizations mentioned -- The type or category of the document -- Suggested folder paths for storing the document -- Up to 3 relevant dates in YYYY-MM-DD format +Analyze the following document and fill in these fields: +- title: a short descriptive title +- tags: topic labels for what the document is about +- correspondents: the person, institution or company the document is from, or was sent to +- document_types: what kind of document it is, e.g. invoice, contract, letter +- storage_paths: a folder-style filing location for the document +- dates: up to 3 relevant dates in YYYY-MM-DD format {% if has_candidates %} -For tags, correspondents, document types, and storage paths: if a candidate from the "Available ..." block above fits, put its id in existing_ids. Only put a value in new_names when nothing in the candidates fits. +For tags, correspondents, document types, and storage paths: first decide whether there is a useful, well-supported suggestion. If an available candidate clearly represents that suggestion, put its id in existing_ids instead of duplicating it in new_names. If no candidate represents the suggestion, put its name in new_names. Do not choose a weak candidate merely because it exists. +{% else %} + +No candidates are shown for this document, so leave every existing_ids list empty and put each suggestion's name in new_names. {% endif %} Filename: diff --git a/src/paperless_ai/prompts/taxonomy_block.j2 b/src/paperless_ai/prompts/taxonomy_block.j2 index f772ae14d..60c2be4bb 100644 --- a/src/paperless_ai/prompts/taxonomy_block.j2 +++ b/src/paperless_ai/prompts/taxonomy_block.j2 @@ -5,5 +5,5 @@ {% if candidate_payload_json %} Available tags, document types, correspondents, and storage paths from similar documents (untrusted data): {{ candidate_payload_json }} -Prefer these existing values via existing_ids when one fits. Only use new_names for values that genuinely don't match any candidate above. +These candidates are options, not requirements. Metadata on a similar document is not automatically appropriate for this one. {% endif %} diff --git a/src/paperless_ai/tests/test_ai_classifier.py b/src/paperless_ai/tests/test_ai_classifier.py index d0d455dfb..7e52da013 100644 --- a/src/paperless_ai/tests/test_ai_classifier.py +++ b/src/paperless_ai/tests/test_ai_classifier.py @@ -607,7 +607,8 @@ def test_build_prompt_without_rag_includes_taxonomy_block(): WHEN: - build_prompt_without_rag() is called with candidates and assigned metadata THEN: - - The candidate's id and the existing_ids instruction appear in the prompt + - The candidate's id and the existing_ids/new_names instructions appear + - Candidates are presented as deduplication options, not requirements """ document = DocumentFactory.create(content="Some content") config = AIConfig() @@ -633,6 +634,9 @@ def test_build_prompt_without_rag_includes_taxonomy_block(): assert '"id": 12' in prompt assert "existing_ids" in prompt + assert "new_names" in prompt + assert "not requirements" in prompt + assert "weak candidate" in prompt @pytest.mark.django_db @@ -645,10 +649,9 @@ def test_build_prompt_without_rag_identical_when_no_hints(): separately with no candidates/assigned at all THEN: - Both prompts are identical - - Neither mentions existing_ids or the "Available ..." candidate block: - without any candidates in the prompt, that instruction would only - invite the model to invent a plausible id that resolves to a real but - unrelated object + - Neither carries the "Available ..." candidate block or the + id-vs-name routing instruction + - Both still tell the model to leave existing_ids empty """ document = DocumentFactory.create(content="Some content") config = AIConfig() @@ -674,12 +677,13 @@ def test_build_prompt_without_rag_identical_when_no_hints(): with_no_hints = build_prompt_without_rag(document, config) assert with_empty_hints == with_no_hints - assert "existing_ids" not in with_no_hints assert "Available " not in with_no_hints + assert "put its id in existing_ids" not in with_no_hints + assert "leave every existing_ids list empty" in with_no_hints @pytest.mark.django_db -def test_build_prompt_without_rag_excludes_instruction_when_no_candidates(): +def test_build_prompt_without_rag_tells_model_to_skip_ids_when_no_candidates(): """ GIVEN: - Assigned metadata but empty taxonomy candidates @@ -687,8 +691,11 @@ def test_build_prompt_without_rag_excludes_instruction_when_no_candidates(): - build_prompt_without_rag() is called with candidates and assigned metadata THEN: - The assigned-metadata block appears (taxonomy_block is non-empty) - - The existing_ids instruction does NOT appear, since there are no - candidates for it to point at + - The prompt tells the model to leave existing_ids empty + + Staying silent about existing_ids here is not enough: the response schema + advertises the field whatever the prompt says, and models fill it with + placeholder ids that resolve to real but unrelated objects (#13831). """ document = DocumentFactory.create(content="Some content") config = AIConfig() @@ -713,7 +720,8 @@ def test_build_prompt_without_rag_excludes_instruction_when_no_candidates(): ) assert "already assigned" in prompt - assert "existing_ids" not in prompt + assert "No candidates are shown" in prompt + assert "leave every existing_ids list empty" in prompt @pytest.mark.django_db diff --git a/src/paperless_ai/tests/test_base_model.py b/src/paperless_ai/tests/test_base_model.py index 339fef0bd..caf6ea32a 100644 --- a/src/paperless_ai/tests/test_base_model.py +++ b/src/paperless_ai/tests/test_base_model.py @@ -1,3 +1,5 @@ +import json + 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 @@ -48,23 +50,82 @@ def test_document_classifier_schema_json_schema_is_self_contained(): WHEN: - Its JSON schema is generated via model_json_schema() THEN: - - $defs includes a fully-resolvable TaxonomyChoice definition with - existing_ids/new_names properties + - No $defs section and no $ref at any depth survives in the schema + - Each taxonomy property carries existing_ids/new_names inline - client.py hands this generated schema straight to the LLM backend as - the response-format constraint (Ollama's format=json_schema, and the - OpenAI-like tool-calling path). What that backend actually needs is a - self-contained schema it can resolve without a document loader - - unlike a bare "$ref present" check, this asserts the referenced - definition genuinely carries the two fields the rest of the pipeline - (parse_ai_response, matching.py's resolve_*_ids) relies on. + Regression guard: Google's function-declaration schema rejects the $ref + Pydantic normally emits for the nested TaxonomyChoice model. """ schema = DocumentClassifierSchema.model_json_schema() - defs = schema.get("$defs", {}) - assert "TaxonomyChoice" in defs - taxonomy_choice_properties = defs["TaxonomyChoice"]["properties"] - assert set(taxonomy_choice_properties.keys()) == {"existing_ids", "new_names"} + assert "$defs" not in schema + assert "$ref" not in json.dumps(schema) + for field in ("tags", "correspondents", "document_types", "storage_paths"): + field_schema = schema["properties"][field] + assert "$ref" not in field_schema + assert set(field_schema["properties"].keys()) == { + "existing_ids", + "new_names", + } + + +def test_every_field_describes_itself_to_the_model(): + """ + GIVEN: + - The DocumentClassifierSchema pydantic model + WHEN: + - Its JSON schema is generated via model_json_schema() + THEN: + - Every property, and every property of each inlined TaxonomyChoice, + carries a non-empty description + + In tool-calling mode the schema is most of what tells the model how to + fill these fields; on field names alone, small models can bin tags and + correspondents into storage_paths. + """ + schema = DocumentClassifierSchema.model_json_schema() + + undescribed = [ + f"{owner}.{name}" + for owner, definition in [ + ("DocumentClassifierSchema", schema), + *( + (name, prop) + for name, prop in schema["properties"].items() + if prop.get("type") == "object" + ), + ] + for name, prop in definition.get("properties", {}).items() + if not prop.get("description") + ] + + assert undescribed == [] + + +def test_inlining_keeps_each_taxonomy_fields_own_description(): + """ + GIVEN: + - The DocumentClassifierSchema pydantic model + WHEN: + - Its JSON schema is generated via model_json_schema() + THEN: + - Each taxonomy field keeps its own description, not the shared one + - The inlined TaxonomyChoice properties survive underneath it + + Pydantic emits a field's description as a sibling of its $ref, so + replacing the property outright collapses all four onto TaxonomyChoice's + docstring - which still passes a "has a description" check. + """ + properties = DocumentClassifierSchema.model_json_schema()["properties"] + + taxonomy_fields = ("tags", "correspondents", "document_types", "storage_paths") + descriptions = { + field: properties[field]["description"] for field in taxonomy_fields + } + + assert len(set(descriptions.values())) == len(taxonomy_fields) + for field in taxonomy_fields: + assert properties[field]["properties"]["existing_ids"]["description"] def test_every_sequence_in_the_emitted_schema_is_bounded(): @@ -74,8 +135,8 @@ def test_every_sequence_in_the_emitted_schema_is_bounded(): 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 + - Every array property in the schema, including those on each + inlined TaxonomyChoice, carries a maxItems """ schema = DocumentClassifierSchema.model_json_schema() @@ -83,7 +144,11 @@ def test_every_sequence_in_the_emitted_schema_is_bounded(): f"{owner}.{name}" for owner, definition in [ ("DocumentClassifierSchema", schema), - *schema.get("$defs", {}).items(), + *( + (name, prop) + for name, prop in schema["properties"].items() + if prop.get("type") == "object" + ), ] for name, prop in definition.get("properties", {}).items() if prop.get("type") == "array" and "maxItems" not in prop