From 4c9cc18165ffb8161a58a3047cab9e82e728a3c5 Mon Sep 17 00:00:00 2001
From: shamoon <4887959+shamoon@users.noreply.github.com>
Date: Sun, 30 Aug 2026 10:32:02 -0700
Subject: [PATCH] Update ai_classifier.py

---
 src/paperless_ai/ai_classifier.py | 21 +++++++++++++++++----
 1 file changed, 17 insertions(+), 4 deletions(-)

diff --git a/src/paperless_ai/ai_classifier.py b/src/paperless_ai/ai_classifier.py
index 5e59b13c7..9e5f9b662 100644
--- a/src/paperless_ai/ai_classifier.py
+++ b/src/paperless_ai/ai_classifier.py
@@ -264,6 +264,22 @@ def _restrict_to_shown_candidates(
     )
 
 
+def _candidate_id_allowlist(
+    candidates: TaxonomyCandidates,
+) -> dict[str, set[int]]:
+    """Candidate IDs grouped by category for validating model mappings."""
+    return {
+        "tags": {candidate["id"] for candidate in candidates["tags"]},
+        "document_types": {
+            candidate["id"] for candidate in candidates["document_types"]
+        },
+        "correspondents": {
+            candidate["id"] for candidate in candidates["correspondents"]
+        },
+        "storage_paths": {candidate["id"] for candidate in candidates["storage_paths"]},
+    }
+
+
 def get_ai_document_classification(
     document: Document,
     user: User | None = None,
@@ -289,10 +305,7 @@ def get_ai_document_classification(
     with db_connection_released():
         result = client.run_llm_query(
             prompt,
-            allowed_candidate_ids={
-                category: {candidate["id"] for candidate in values}
-                for category, values in candidates.items()
-            },
+            allowed_candidate_ids=_candidate_id_allowlist(candidates),
         )
         suggestions = _restrict_to_shown_candidates(
             parse_ai_response(result),