Use a flat list we change into our nested schema

src/paperless_ai/ai_classifier.py

@@ -1,4 +1,3 @@
-import json
 import logging
 
 from django.conf import settings
@@ -9,6 +8,7 @@ from documents.permissions import get_objects_for_user_owner_aware
 from paperless.config import AIConfig
 from paperless_ai.base_model import ClassificationSuggestions
 from paperless_ai.base_model import TaxonomyChoiceDict
+from paperless_ai.base_model import classification_suggestions_to_model
 from paperless_ai.client import AIClient
 from paperless_ai.db import db_connection_released
 from paperless_ai.indexing import _node_document_ids
@@ -124,20 +124,16 @@ def build_localization_prompt(
     suggestions: ClassificationSuggestions,
     output_language: str,
 ) -> str:
-    """``suggestions`` is the full nested-shape result of parse_ai_response
-    (each taxonomy field a ``{"existing_ids": [...], "new_names": [...]}``
-    dict) - passed through as-is so the model receives and returns the exact
-    DocumentClassifierSchema shape run_llm_query() always parses against.
-    Only each field's new_names (never existing_ids, which are plain
-    resolved-object IDs, not text) and title get used from the response; see
-    get_ai_document_classification's merge step, which always keeps the
-    *original* existing_ids regardless of what the model echoes back here.
+    """Render internal suggestions in the same flat shape the model returns.
+    Only the name fields and title are used from the localized response; the
+    merge step always keeps the original ID fields.
     """
     language_name = get_language_name(output_language)
+    model_suggestions = classification_suggestions_to_model(suggestions)
     return render_prompt(
         LocalizationPromptContext(
             language_name=language_name,
-            suggestions_json=json.dumps(suggestions, ensure_ascii=False),
+            suggestions_json=model_suggestions.model_dump_json(),
         ),
     )
 
@@ -208,11 +204,9 @@ def get_taxonomy_context(
 
 
 def parse_ai_response(raw: dict) -> ClassificationSuggestions:
-    """``raw`` is AIClient.run_llm_query()'s return value - already a
-    DocumentClassifierSchema.model_dump(), so every key below is always
-    present with the right shape; this only exists to give the rest of the
-    module a named, typed boundary instead of passing the client's bare dict
-    straight through everywhere.
+    """``raw`` is AIClient.run_llm_query()'s validated internal-shape result.
+    This gives the rest of the module a named, typed boundary instead of
+    passing the client's bare dict straight through everywhere.
     """
 
     def _choice(value: dict | None) -> TaxonomyChoiceDict:

src/paperless_ai/base_model.py

@@ -6,7 +6,6 @@ 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"
@@ -32,59 +31,9 @@ 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 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,
-        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."
-        ),
-    )
-
-    @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])
-
-
+# This model is serialized into the schema handed to the LLM, so its docstring
+# and field descriptions are instructions for the model. Keep implementation
+# details in code comments instead.
 class DocumentClassifierSchema(BaseModel):
     """Classification suggestions for a single document."""
 
@@ -95,36 +44,79 @@ class DocumentClassifierSchema(BaseModel):
             f"{MAX_TITLE_LENGTH} characters."
         ),
     )
-    tags: TaxonomyChoice = Field(
-        default_factory=TaxonomyChoice,
+    tags: list[str] = Field(
+        default_factory=list,
+        max_length=MAX_NEW_NAMES,
         description=(
-            "Topic labels describing what this document is about. A document "
-            "may have several, e.g. 'Insurance', 'Car', 'Warranty'."
+            "Names of topic labels describing what this document is about, "
+            "e.g. 'Insurance', 'Car', 'Warranty'. When an available tag "
+            "represents the same label, use its ID in tag_ids instead."
         ),
     )
-    correspondents: TaxonomyChoice = Field(
-        default_factory=TaxonomyChoice,
+    tag_ids: list[int] = Field(
+        default_factory=list,
+        max_length=MAX_EXISTING_IDS,
         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."
+            "IDs of available tags that clearly apply to this document. Only "
+            "use IDs shown in the prompt; never invent one or choose a weak match."
         ),
     )
-    document_types: TaxonomyChoice = Field(
-        default_factory=TaxonomyChoice,
+    correspondents: list[str] = Field(
+        default_factory=list,
+        max_length=MAX_NEW_NAMES,
         description=(
-            "What kind of document this is, e.g. 'Invoice', 'Contract', "
-            "'Bank Statement', 'Letter'. Never its subject matter and never "
-            "who sent it."
+            "Names of people, institutions or companies this document is from "
+            "or was sent to, not every party merely mentioned. When an "
+            "available correspondent is the same entity, use its ID in "
+            "correspondent_ids instead."
         ),
     )
-    storage_paths: TaxonomyChoice = Field(
-        default_factory=TaxonomyChoice,
+    correspondent_ids: list[int] = Field(
+        default_factory=list,
+        max_length=MAX_EXISTING_IDS,
         description=(
-            "A folder-style filing location for this document, e.g. "
+            "IDs of available correspondents that clearly apply to this "
+            "document. Only use IDs shown in the prompt; never invent one or "
+            "choose a weak match."
+        ),
+    )
+    document_types: list[str] = Field(
+        default_factory=list,
+        max_length=MAX_NEW_NAMES,
+        description=(
+            "Names describing what kind of document this is, e.g. 'Invoice', "
+            "'Contract', 'Bank Statement', 'Letter'. Never use its subject or "
+            "sender as a document type. When an available document type is the "
+            "same kind, use its ID in document_type_ids instead."
+        ),
+    )
+    document_type_ids: list[int] = Field(
+        default_factory=list,
+        max_length=MAX_EXISTING_IDS,
+        description=(
+            "IDs of available document types that clearly apply to this "
+            "document. Only use IDs shown in the prompt; never invent one or "
+            "choose a weak match."
+        ),
+    )
+    storage_paths: list[str] = Field(
+        default_factory=list,
+        max_length=MAX_NEW_NAMES,
+        description=(
+            "Names of folder-style filing locations, e.g. "
             "'Finance/Invoices'. Leave empty unless a filing location is "
             "clearly implied - never put tags, document types or "
-            "correspondents here."
+            "correspondents here. When an available storage path is the same "
+            "location, use its ID in storage_path_ids instead."
+        ),
+    )
+    storage_path_ids: list[int] = Field(
+        default_factory=list,
+        max_length=MAX_EXISTING_IDS,
+        description=(
+            "IDs of available storage paths that clearly apply to this "
+            "document. Only use IDs shown in the prompt; never invent one or "
+            "choose a weak match."
         ),
     )
     dates: list[str] = Field(
@@ -137,41 +129,33 @@ class DocumentClassifierSchema(BaseModel):
         ),
     )
 
-    @field_validator("title", "dates", mode="before")
+    @field_validator(
+        "title",
+        "tags",
+        "tag_ids",
+        "correspondents",
+        "correspondent_ids",
+        "document_types",
+        "document_type_ids",
+        "storage_paths",
+        "storage_path_ids",
+        "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
-    TaxonomyChoice.model_dump() actually produces, typed for callers that
-    work with the dumped dict rather than the pydantic instance."""
+    """Internal representation of names and existing IDs for one taxonomy."""
 
     existing_ids: list[int]
     new_names: list[str]
 
 
 class ClassificationSuggestions(TypedDict):
-    """Plain-dict counterpart of DocumentClassifierSchema.model_dump() -
-    the shape threaded through parse_ai_response, build_localization_prompt,
-    get_ai_document_classification, and the ai_suggestions view."""
+    """Internal shape used after the flat LLM response is validated."""
 
     title: str
     tags: TaxonomyChoiceDict
@@ -179,3 +163,47 @@ class ClassificationSuggestions(TypedDict):
     document_types: TaxonomyChoiceDict
     storage_paths: TaxonomyChoiceDict
     dates: list[str]
+
+
+def model_to_classification_suggestions(
+    model: DocumentClassifierSchema,
+) -> ClassificationSuggestions:
+    """Convert the flat, model-friendly response to the internal shape."""
+    return ClassificationSuggestions(
+        title=model.title,
+        tags=TaxonomyChoiceDict(
+            existing_ids=model.tag_ids,
+            new_names=model.tags,
+        ),
+        correspondents=TaxonomyChoiceDict(
+            existing_ids=model.correspondent_ids,
+            new_names=model.correspondents,
+        ),
+        document_types=TaxonomyChoiceDict(
+            existing_ids=model.document_type_ids,
+            new_names=model.document_types,
+        ),
+        storage_paths=TaxonomyChoiceDict(
+            existing_ids=model.storage_path_ids,
+            new_names=model.storage_paths,
+        ),
+        dates=model.dates,
+    )
+
+
+def classification_suggestions_to_model(
+    suggestions: ClassificationSuggestions,
+) -> DocumentClassifierSchema:
+    """Convert internal suggestions to the flat shape used for localization."""
+    return DocumentClassifierSchema(
+        title=suggestions["title"],
+        tags=suggestions["tags"]["new_names"],
+        tag_ids=suggestions["tags"]["existing_ids"],
+        correspondents=suggestions["correspondents"]["new_names"],
+        correspondent_ids=suggestions["correspondents"]["existing_ids"],
+        document_types=suggestions["document_types"]["new_names"],
+        document_type_ids=suggestions["document_types"]["existing_ids"],
+        storage_paths=suggestions["storage_paths"]["new_names"],
+        storage_path_ids=suggestions["storage_paths"]["existing_ids"],
+        dates=suggestions["dates"],
+    )

src/paperless_ai/client.py

@@ -19,7 +19,9 @@ from paperless.network import PinnedHostHTTPTransport
 from paperless.network import create_pinned_async_httpx_client
 from paperless.network import create_pinned_httpx_client
 from paperless.network import validate_outbound_http_url
+from paperless_ai.base_model import ClassificationSuggestions
 from paperless_ai.base_model import DocumentClassifierSchema
+from paperless_ai.base_model import model_to_classification_suggestions
 from paperless_ai.exceptions import LLMTimeoutError
 
 logger = logging.getLogger("paperless_ai.client")
@@ -115,7 +117,7 @@ class AIClient:
         else:
             raise ValueError(f"Unsupported LLM backend: {self.settings.llm_backend}")
 
-    def run_llm_query(self, prompt: str) -> str:
+    def run_llm_query(self, prompt: str) -> ClassificationSuggestions:
         logger.debug(
             "Running LLM query against %s with model %s",
             self.settings.llm_backend,
@@ -134,7 +136,7 @@ class AIClient:
                 )
             logger.debug("LLM query result: %s", result)
             parsed = DocumentClassifierSchema(**json.loads(result.message.content))
-            return parsed.model_dump()
+            return model_to_classification_suggestions(parsed)
 
         from llama_index.core.program.function_program import get_function_tool
 
@@ -153,7 +155,7 @@ class AIClient:
             )
         logger.debug("LLM query result: %s", tool_calls)
         parsed = DocumentClassifierSchema(**tool_calls[0].tool_kwargs)
-        return parsed.model_dump()
+        return model_to_classification_suggestions(parsed)
 
     @contextmanager
     def _normalize_timeouts(self) -> Iterator[None]:

src/paperless_ai/prompts/classification.j2

@@ -13,10 +13,10 @@ Analyze the following document and fill in these fields:
 - dates: up to 3 relevant dates in YYYY-MM-DD format
 {% if has_candidates %}
 
-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.
+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 the matching tag_ids, correspondent_ids, document_type_ids, or storage_path_ids field instead of repeating its name. If no candidate represents the suggestion, put its name in tags, correspondents, document_types, or storage_paths. 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.
+No candidates are shown for this document, so leave every field ending in "_ids" empty and put suggestions in the corresponding name fields.
 {% endif %}
 
 Filename:

src/paperless_ai/prompts/localization.j2

@@ -1,6 +1,6 @@
 You are localizing document classification suggestions for display in Paperless-ngx.
 
-Rewrite only the "title" field and each taxonomy field's "new_names" list in {{ language_name }}. Leave every "existing_ids" list exactly as given - these are database identifiers, not text, and are not used from your response even if changed.
+Rewrite only the "title", "tags", "document_types", and "storage_paths" fields in {{ language_name }}. Leave every field ending in "_ids" exactly as given - these are database identifiers, not text, and are not used from your response even if changed.
 
 Do not translate correspondents or dates.
 Preserve proper nouns, organization names, product names, and exact official document names. Translate generic category words when a {{ language_name }} equivalent exists.

src/paperless_ai/tests/test_ai_classifier.py

@@ -245,7 +245,7 @@ def test_prompt_with_without_rag(mock_document):
     THEN:
         - build_prompt_without_rag() has no similar-documents section
         - build_prompt_with_rag() includes the similar-documents context
-        - build_localization_prompt() asks to rewrite only new_names/title and
+        - build_localization_prompt() asks to rewrite only names/title and
           not to translate correspondents or dates
     """
     config = AIConfig()
@@ -264,6 +264,7 @@ def test_prompt_with_without_rag(mock_document):
     prompt = build_localization_prompt(NESTED_SUGGESTIONS, output_language="de-de")
     assert "Rewrite only the" in prompt
     assert "Do not translate correspondents or dates" in prompt
+    assert '"tag_ids":[]' in prompt
 
 
 def test_get_language_name_falls_back_to_language_code():
@@ -607,7 +608,7 @@ 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/new_names instructions appear
+        - The candidate's id and the flat name/ID instructions appear
         - Candidates are presented as deduplication options, not requirements
     """
     document = DocumentFactory.create(content="Some content")
@@ -633,8 +634,8 @@ 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 "tag_ids" in prompt
+    assert "correspondent_ids" in prompt
     assert "not requirements" in prompt
     assert "weak candidate" in prompt
 
@@ -651,7 +652,7 @@ def test_build_prompt_without_rag_identical_when_no_hints():
         - Both prompts are identical
         - Neither carries the "Available ..." candidate block or the
           id-vs-name routing instruction
-        - Both still tell the model to leave existing_ids empty
+        - Both still tell the model to leave every ID field empty
     """
     document = DocumentFactory.create(content="Some content")
     config = AIConfig()
@@ -678,8 +679,8 @@ def test_build_prompt_without_rag_identical_when_no_hints():
 
     assert with_empty_hints == 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
+    assert "put its id in the matching" not in with_no_hints
+    assert 'leave every field ending in "_ids" empty' in with_no_hints
 
 
 @pytest.mark.django_db
@@ -691,9 +692,9 @@ def test_build_prompt_without_rag_tells_model_to_skip_ids_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 prompt tells the model to leave existing_ids empty
+        - The prompt tells the model to leave every ID field empty
 
-    Staying silent about existing_ids here is not enough: the response schema
+    Staying silent about ID fields 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).
     """
@@ -721,7 +722,7 @@ def test_build_prompt_without_rag_tells_model_to_skip_ids_when_no_candidates():
 
     assert "already assigned" in prompt
     assert "No candidates are shown" in prompt
-    assert "leave every existing_ids list empty" in prompt
+    assert 'leave every field ending in "_ids" empty' in prompt
 
 
 @pytest.mark.django_db

src/paperless_ai/tests/test_base_model.py

@@ -6,8 +6,9 @@ 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
 from paperless_ai.base_model import TaxonomyChoiceDict
+from paperless_ai.base_model import classification_suggestions_to_model
+from paperless_ai.base_model import model_to_classification_suggestions
 
 
 def test_document_classifier_schema_declared_defaults():
@@ -18,66 +19,67 @@ def test_document_classifier_schema_declared_defaults():
     WHEN:
         - The schema is dumped to a dict via model_dump()
     THEN:
-        - Every taxonomy field dumps as an empty existing_ids/new_names
-          dict, and dates dumps as an empty list
+        - Every name and ID field, and dates, dump as empty lists
 
-    This is the one project-owned fact worth pinning down here: which
-    defaults this schema declares for a partial LLM response (see
-    client.py's DocumentClassifierSchema(**json.loads(...)) call sites,
-    which construct from whatever subset of fields the backend actually
-    returned). It deliberately hardcodes the expected literal rather than
-    re-deriving it from TaxonomyChoice()/[] - pydantic's own
-    default_factory machinery is not this project's to re-test, and a
-    test that recomputes the expected value from the model under test
-    can't ever catch a wrong default.
+    The model may omit optional fields, so the schema must provide the complete
+    empty shape expected by the conversion and matching pipeline.
     """
     schema = DocumentClassifierSchema(title="Test Title")
 
     dumped = schema.model_dump()
 
-    empty_choice = {"existing_ids": [], "new_names": []}
-    assert dumped["tags"] == empty_choice
-    assert dumped["correspondents"] == empty_choice
-    assert dumped["document_types"] == empty_choice
-    assert dumped["storage_paths"] == empty_choice
-    assert dumped["dates"] == []
+    assert dumped == {
+        "title": "Test Title",
+        "tags": [],
+        "tag_ids": [],
+        "correspondents": [],
+        "correspondent_ids": [],
+        "document_types": [],
+        "document_type_ids": [],
+        "storage_paths": [],
+        "storage_path_ids": [],
+        "dates": [],
+    }
 
 
-def test_flat_taxonomy_lists_are_normalized_for_legacy_model_responses():
+def test_flat_model_response_converts_to_internal_taxonomy_choices():
     """
     GIVEN:
-        - A model response using the flat taxonomy lists accepted before 3.1
-        - Strings, integer candidate IDs, and a mixture of both
+        - A flat model response with separate name and candidate-ID fields
     WHEN:
-        - DocumentClassifierSchema validates the response
+        - It is converted to Paperless' internal suggestion representation
     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.
+        - Names and IDs are paired under their taxonomy category
     """
     parsed = DocumentClassifierSchema(
         title="Electricity Bill",
         tags=["Utilities", "Electricity"],
-        correspondents=[12],
-        document_types=[34, "Utility Bill"],
+        tag_ids=[12],
+        correspondents=["Power Company"],
+        correspondent_ids=[23],
+        document_types=["Utility Bill"],
+        document_type_ids=[34],
         storage_paths=["Finance/Utilities"],
+        storage_path_ids=[45],
     )
+    suggestions = model_to_classification_suggestions(parsed)
 
-    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"],
-    )
+    assert suggestions["tags"] == {
+        "existing_ids": [12],
+        "new_names": ["Utilities", "Electricity"],
+    }
+    assert suggestions["correspondents"] == {
+        "existing_ids": [23],
+        "new_names": ["Power Company"],
+    }
+    assert suggestions["document_types"] == {
+        "existing_ids": [34],
+        "new_names": ["Utility Bill"],
+    }
+    assert suggestions["storage_paths"] == {
+        "existing_ids": [45],
+        "new_names": ["Finance/Utilities"],
+    }
 
 
 def test_document_classifier_schema_json_schema_is_self_contained():
@@ -87,23 +89,20 @@ def test_document_classifier_schema_json_schema_is_self_contained():
     WHEN:
         - Its JSON schema is generated via model_json_schema()
     THEN:
-        - No $defs section and no $ref at any depth survives in the schema
-        - Each taxonomy property carries existing_ids/new_names inline
+        - The schema contains no definitions, references, or nested objects
+        - Every response field is a scalar or flat array
 
-    Regression guard: Google's function-declaration schema rejects the $ref
-    Pydantic normally emits for the nested TaxonomyChoice model.
+    This keeps the function declaration compatible with backends that reject
+    JSON Schema references and with smaller models that struggle with nesting.
     """
     schema = DocumentClassifierSchema.model_json_schema()
 
     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",
-        }
+    assert all(
+        field_schema.get("type") != "object"
+        for field_schema in schema["properties"].values()
+    )
 
 
 def test_every_field_describes_itself_to_the_model():
@@ -113,8 +112,7 @@ def test_every_field_describes_itself_to_the_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
+        - Every property 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
@@ -123,48 +121,14 @@ def test_every_field_describes_itself_to_the_model():
     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()
+        name
+        for name, prop in schema["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():
     """
     GIVEN:
@@ -172,22 +136,13 @@ 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 each
-          inlined TaxonomyChoice, carries a maxItems
+        - Every array property in the schema carries a maxItems
     """
     schema = DocumentClassifierSchema.model_json_schema()
 
     unbounded = [
-        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()
+        name
+        for name, prop in schema["properties"].items()
         if prop.get("type") == "array" and "maxItems" not in prop
     ]
 
@@ -219,17 +174,15 @@ def test_over_long_response_is_truncated_rather_than_rejected():
     """
     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),
-        ),
+        tags=["n"] * (MAX_NEW_NAMES + 20),
+        tag_ids=list(range(MAX_EXISTING_IDS + 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
+    assert len(parsed.tag_ids) == MAX_EXISTING_IDS
+    assert len(parsed.tags) == MAX_NEW_NAMES
 
 
 def test_truncation_keeps_the_earliest_entries():
@@ -249,24 +202,48 @@ def test_truncation_keeps_the_earliest_entries():
     assert parsed.dates == ["2016-10-01", "2016-09-01", "2016-08-01"]
 
 
-def test_model_dump_matches_typed_dict_keys():
+def test_model_conversion_matches_internal_typed_dict_keys():
     """
     GIVEN:
         - A DocumentClassifierSchema instance
     WHEN:
-        - It is dumped to a dict via model_dump()
+        - It is converted to ClassificationSuggestions
     THEN:
-        - The dumped dict's keys exactly match ClassificationSuggestions'
+        - The converted dict's keys exactly match ClassificationSuggestions'
           declared keys
-        - The dumped tags dict's keys exactly match TaxonomyChoiceDict's
+        - The converted tags dict's keys exactly match TaxonomyChoiceDict's
           declared keys
     """
-    # TaxonomyChoiceDict/ClassificationSuggestions are the static-typing
-    # counterparts of TaxonomyChoice/DocumentClassifierSchema - this pins
-    # down that .model_dump()'s actual runtime keys are exactly what the
-    # TypedDicts declare, so the two don't silently drift apart.
-    schema = DocumentClassifierSchema(title="T", tags=TaxonomyChoice(existing_ids=[1]))
-    dumped = schema.model_dump()
+    schema = DocumentClassifierSchema(title="T", tags=["Tag"], tag_ids=[1])
+    suggestions = model_to_classification_suggestions(schema)
 
-    assert set(dumped.keys()) == set(ClassificationSuggestions.__annotations__.keys())
-    assert set(dumped["tags"].keys()) == set(TaxonomyChoiceDict.__annotations__.keys())
+    assert set(suggestions.keys()) == set(
+        ClassificationSuggestions.__annotations__.keys(),
+    )
+    assert set(suggestions["tags"].keys()) == set(
+        TaxonomyChoiceDict.__annotations__.keys(),
+    )
+
+
+def test_internal_suggestions_round_trip_through_flat_model():
+    suggestions = ClassificationSuggestions(
+        title="Electricity Bill",
+        tags=TaxonomyChoiceDict(existing_ids=[1], new_names=["Utilities"]),
+        correspondents=TaxonomyChoiceDict(
+            existing_ids=[2],
+            new_names=["Power Company"],
+        ),
+        document_types=TaxonomyChoiceDict(
+            existing_ids=[3],
+            new_names=["Utility Bill"],
+        ),
+        storage_paths=TaxonomyChoiceDict(
+            existing_ids=[4],
+            new_names=["Finance/Utilities"],
+        ),
+        dates=["2026-08-30"],
+    )
+
+    model = classification_suggestions_to_model(suggestions)
+
+    assert model_to_classification_suggestions(model) == suggestions

src/paperless_ai/tests/test_client.py

@@ -123,10 +123,14 @@ def test_run_llm_query_ollama_uses_structured_json(mock_ai_config, mock_ollama_l
     mock_llm_instance.chat.return_value.message.content = json.dumps(
         {
             "title": "Test Title",
-            "tags": {"existing_ids": [1], "new_names": ["document"]},
-            "correspondents": {"existing_ids": [], "new_names": ["John Doe"]},
-            "document_types": {"existing_ids": [], "new_names": ["report"]},
-            "storage_paths": {"existing_ids": [], "new_names": ["Reports"]},
+            "tags": ["document"],
+            "tag_ids": [1],
+            "correspondents": ["John Doe"],
+            "correspondent_ids": [],
+            "document_types": ["report"],
+            "document_type_ids": [],
+            "storage_paths": ["Reports"],
+            "storage_path_ids": [],
             "dates": ["2023-01-01"],
         },
     )
@@ -156,10 +160,14 @@ def test_run_llm_query_openai_uses_tools(mock_ai_config, mock_openai_llm):
         tool_name="DocumentClassifierSchema",
         tool_kwargs={
             "title": "Test Title",
-            "tags": {"existing_ids": [1], "new_names": ["document"]},
-            "correspondents": {"existing_ids": [], "new_names": ["John Doe"]},
-            "document_types": {"existing_ids": [], "new_names": ["report"]},
-            "storage_paths": {"existing_ids": [], "new_names": ["Reports"]},
+            "tags": ["document"],
+            "tag_ids": [1],
+            "correspondents": ["John Doe"],
+            "correspondent_ids": [],
+            "document_types": ["report"],
+            "document_type_ids": [],
+            "storage_paths": ["Reports"],
+            "storage_path_ids": [],
             "dates": ["2023-01-01"],
         },
     )