From 7ddc1c98015c620a530cc286c7cbd049fbee3e04 Mon Sep 17 00:00:00 2001
From: shamoon <4887959+shamoon@users.noreply.github.com>
Date: Sat, 29 Aug 2026 10:56:05 -0700
Subject: [PATCH] Fix: 3.1 LLM suggestions fix rank ordering (#13848)

---
 src/paperless_ai/ai_classifier.py            | 11 ++-
 src/paperless_ai/tests/test_ai_classifier.py | 71 ++++++++++++++++++++
 2 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/src/paperless_ai/ai_classifier.py b/src/paperless_ai/ai_classifier.py
index 528300fc9..c9be94453 100644
--- a/src/paperless_ai/ai_classifier.py
+++ b/src/paperless_ai/ai_classifier.py
@@ -183,9 +183,14 @@ def get_taxonomy_context(
 
         candidates = build_taxonomy_candidates(nodes, user)
 
-        similar_docs = list(
-            Document.objects.filter(pk__in=_node_document_ids(nodes))[:max_docs],
-        )
+        # ``nodes`` are already ordered by descending vector similarity; don't lose it.
+        similar_document_ids = list(dict.fromkeys(_node_document_ids(nodes)))
+        similar_documents_by_id = Document.objects.in_bulk(similar_document_ids)
+        similar_docs = [
+            similar_documents_by_id[document_id]
+            for document_id in similar_document_ids
+            if document_id in similar_documents_by_id
+        ][:max_docs]
         context_blocks = []
         for similar in similar_docs:
             text = similar.content[:1000] or ""
diff --git a/src/paperless_ai/tests/test_ai_classifier.py b/src/paperless_ai/tests/test_ai_classifier.py
index 6159f78b9..d0d455dfb 100644
--- a/src/paperless_ai/tests/test_ai_classifier.py
+++ b/src/paperless_ai/tests/test_ai_classifier.py
@@ -1,3 +1,4 @@
+import datetime
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 from unittest.mock import patch
@@ -343,6 +344,76 @@ def test_get_taxonomy_context_assembles_rag_text_and_candidates():
     }
 
 
+@pytest.mark.django_db
+def test_get_taxonomy_context_preserves_similarity_order_and_distinct_documents():
+    """
+    GIVEN:
+        - Ranked nodes whose similarity order conflicts with Document's
+          newest-created-first default ordering
+        - Two chunks belonging to the most similar document
+        - A stale node whose document no longer exists
+    WHEN:
+        - get_taxonomy_context() builds a two-document RAG context
+    THEN:
+        - The two most similar distinct documents are used in ranked order
+        - The duplicate chunk does not consume a context slot
+        - The missing document does not consume a context slot
+    """
+    most_similar = DocumentFactory.create(
+        created=datetime.date(2020, 1, 1),
+        content="Most similar content",
+        title="Most Similar",
+    )
+    second_most_similar = DocumentFactory.create(
+        created=datetime.date(2021, 1, 1),
+        content="Second most similar content",
+        title="Second Most Similar",
+    )
+    newest_but_least_similar = DocumentFactory.create(
+        created=datetime.date(2026, 1, 1),
+        content="Least similar content",
+        title="Newest But Least Similar",
+    )
+    document = DocumentFactory.create(content="Some content")
+    fake_nodes = [
+        SimpleNamespace(
+            metadata={"document_id": str(most_similar.pk)},
+            score=0.9,
+        ),
+        SimpleNamespace(
+            metadata={"document_id": str(most_similar.pk)},
+            score=0.8,
+        ),
+        SimpleNamespace(
+            metadata={"document_id": "999999999"},
+            score=0.75,
+        ),
+        SimpleNamespace(
+            metadata={"document_id": str(second_most_similar.pk)},
+            score=0.7,
+        ),
+        SimpleNamespace(
+            metadata={"document_id": str(newest_but_least_similar.pk)},
+            score=0.6,
+        ),
+    ]
+
+    with patch(
+        "paperless_ai.ai_classifier.retrieve_similar_nodes",
+        return_value=fake_nodes,
+    ):
+        _candidates, _assigned, context = get_taxonomy_context(
+            document,
+            user=None,
+            max_docs=2,
+        )
+
+    assert context == (
+        "TITLE: Most Similar\nMost similar content\n\n"
+        "TITLE: Second Most Similar\nSecond most similar content"
+    )
+
+
 @pytest.mark.django_db
 def test_get_taxonomy_context_no_similar_docs():
     """