diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index af20669e5..07c20b526 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -721,20 +721,3 @@ def retrieve_similar_nodes( continue filtered.append(node) return filtered - - -def _node_document_ids(nodes: list["NodeWithScore"]) -> list[int]: - document_ids: list[int] = [] - for node in nodes: - document_id = node.metadata.get("document_id") - if document_id is None: # pragma: no cover - # See the matching guard in retrieve_similar_nodes() above. - continue - try: - document_ids.append(int(document_id)) - except ValueError: # pragma: no cover - logger.warning( - "Skipping LLM index result with invalid document_id %r.", - document_id, - ) - return document_ids diff --git a/src/paperless_ai/tests/test_ai_indexing.py b/src/paperless_ai/tests/test_ai_indexing.py index 659eaa2f5..3866e6345 100644 --- a/src/paperless_ai/tests/test_ai_indexing.py +++ b/src/paperless_ai/tests/test_ai_indexing.py @@ -1188,9 +1188,7 @@ class TestRetrieveSimilarNodesAgainstRealIndex: nodes = indexing.retrieve_similar_nodes(a, document_ids=[b.id]) - assert all( - document_id == b.id for document_id in indexing._node_document_ids(nodes) - ) + assert all(int(node.metadata["document_id"]) == b.id for node in nodes) def test_excludes_self( self, @@ -1212,7 +1210,7 @@ class TestRetrieveSimilarNodesAgainstRealIndex: nodes = indexing.retrieve_similar_nodes(a, top_k=5) - assert set(indexing._node_document_ids(nodes)) == {b.id} + assert {int(node.metadata["document_id"]) for node in nodes} == {b.id} def test_excludes_self_with_multiple_chunks( self, @@ -1235,4 +1233,4 @@ class TestRetrieveSimilarNodesAgainstRealIndex: nodes = indexing.retrieve_similar_nodes(a, top_k=3) - assert set(indexing._node_document_ids(nodes)) == {b.id} + assert {int(node.metadata["document_id"]) for node in nodes} == {b.id}