diff --git a/src/documents/admin.py b/src/documents/admin.py index d0a43ce9c..719d33c16 100644 --- a/src/documents/admin.py +++ b/src/documents/admin.py @@ -198,6 +198,7 @@ class ShareLinksAdmin(GuardedModelAdmin): class ShareLinkBundleAdmin(GuardedModelAdmin): list_display = ("created", "status", "expiration", "owner", "slug") list_filter = ("status", "created", "expiration", "owner") + readonly_fields = ("file_path",) search_fields = ("slug",) def get_queryset(self, request): # pragma: no cover diff --git a/src/documents/models.py b/src/documents/models.py index 0819b9984..a77447512 100644 --- a/src/documents/models.py +++ b/src/documents/models.py @@ -1019,7 +1019,17 @@ class ShareLinkBundle(models.Model): def absolute_file_path(self) -> Path | None: if not self.file_path: return None - return (settings.SHARE_LINK_BUNDLE_DIR / Path(self.file_path)).resolve() + relative_path = Path(self.file_path) + if relative_path.is_absolute(): + return None + + bundle_dir = settings.SHARE_LINK_BUNDLE_DIR.resolve() + absolute_path = (bundle_dir / relative_path).resolve() + try: + absolute_path.relative_to(bundle_dir) + except ValueError: + return None + return absolute_path def remove_file(self) -> None: if self.absolute_file_path is not None and self.absolute_file_path.exists(): diff --git a/src/documents/tests/test_share_link_bundles.py b/src/documents/tests/test_share_link_bundles.py index 4f0eb057e..0040c5030 100644 --- a/src/documents/tests/test_share_link_bundles.py +++ b/src/documents/tests/test_share_link_bundles.py @@ -124,7 +124,7 @@ class ShareLinkBundleAPITests(DirectoriesMixin, APITestCase): self.assertIn("document_ids", response.data) def test_download_ready_bundle_streams_file(self) -> None: - bundle_file = Path(self.dirs.media_dir) / "bundles" / "ready.zip" + bundle_file = settings.SHARE_LINK_BUNDLE_DIR / "bundles" / "ready.zip" bundle_file.parent.mkdir(parents=True, exist_ok=True) bundle_file.write_bytes(b"binary-zip-content") @@ -132,7 +132,7 @@ class ShareLinkBundleAPITests(DirectoriesMixin, APITestCase): slug="readyslug", file_version=ShareLink.FileVersion.ARCHIVE, status=ShareLinkBundle.Status.READY, - file_path=str(bundle_file), + file_path=str(bundle_file.relative_to(settings.SHARE_LINK_BUNDLE_DIR)), ) bundle.documents.set([self.document]) @@ -199,11 +199,11 @@ class ShareLinkBundleTaskTests(DirectoriesMixin, APITestCase): self.document = DocumentFactory.create() def test_cleanup_expired_share_link_bundles(self) -> None: - expired_path = Path(self.dirs.media_dir) / "expired.zip" + expired_path = settings.SHARE_LINK_BUNDLE_DIR / "expired.zip" expired_path.parent.mkdir(parents=True, exist_ok=True) expired_path.write_bytes(b"expired") - active_path = Path(self.dirs.media_dir) / "active.zip" + active_path = settings.SHARE_LINK_BUNDLE_DIR / "active.zip" active_path.write_bytes(b"active") expired_bundle = ShareLinkBundle.objects.create( @@ -211,7 +211,7 @@ class ShareLinkBundleTaskTests(DirectoriesMixin, APITestCase): file_version=ShareLink.FileVersion.ARCHIVE, status=ShareLinkBundle.Status.READY, expiration=timezone.now() - timedelta(days=1), - file_path=str(expired_path), + file_path=expired_path.name, ) expired_bundle.documents.set([self.document]) @@ -220,7 +220,7 @@ class ShareLinkBundleTaskTests(DirectoriesMixin, APITestCase): file_version=ShareLink.FileVersion.ARCHIVE, status=ShareLinkBundle.Status.READY, expiration=timezone.now() + timedelta(days=1), - file_path=str(active_path), + file_path=active_path.name, ) active_bundle.documents.set([self.document]) @@ -424,7 +424,7 @@ class ShareLinkBundleFilterSetTests(DirectoriesMixin, APITestCase): class ShareLinkBundleModelTests(DirectoriesMixin, APITestCase): - def test_absolute_file_path_handles_relative_and_absolute(self) -> None: + def test_absolute_file_path_handles_relative_path(self) -> None: relative_path = Path("relative.zip") bundle = ShareLinkBundle.objects.create( slug="relative-bundle", @@ -437,10 +437,23 @@ class ShareLinkBundleModelTests(DirectoriesMixin, APITestCase): (settings.SHARE_LINK_BUNDLE_DIR / relative_path).resolve(), ) - absolute_path = Path(self.dirs.media_dir) / "absolute.zip" - bundle.file_path = str(absolute_path) + def test_absolute_file_path_rejects_absolute_path(self) -> None: + bundle = ShareLinkBundle.objects.create( + slug="absolute-bundle", + file_version=ShareLink.FileVersion.ORIGINAL, + file_path=str(Path(self.dirs.media_dir) / "absolute.zip"), + ) - self.assertEqual(bundle.absolute_file_path.resolve(), absolute_path.resolve()) + self.assertIsNone(bundle.absolute_file_path) + + def test_absolute_file_path_rejects_traversal_outside_bundle_dir(self) -> None: + bundle = ShareLinkBundle.objects.create( + slug="traversal-bundle", + file_version=ShareLink.FileVersion.ORIGINAL, + file_path="../escaped.zip", + ) + + self.assertIsNone(bundle.absolute_file_path) def test_str_returns_translated_slug(self) -> None: bundle = ShareLinkBundle.objects.create( diff --git a/src/documents/tests/test_views.py b/src/documents/tests/test_views.py index 4876ec6ab..b8ab42256 100644 --- a/src/documents/tests/test_views.py +++ b/src/documents/tests/test_views.py @@ -437,8 +437,14 @@ class TestAIChatStreamingView(DirectoriesMixin, TestCase): ) super().setUp() + def grant_view_document_permission(self) -> None: + self.user.user_permissions.add( + *Permission.objects.filter(codename="view_document"), + ) + @override_settings(AI_ENABLED=False) def test_post_ai_disabled(self) -> None: + self.grant_view_document_permission() response = self.client.post( self.ENDPOINT, data='{"q": "question"}', @@ -451,6 +457,7 @@ class TestAIChatStreamingView(DirectoriesMixin, TestCase): @patch("documents.views.get_objects_for_user_owner_aware") @override_settings(AI_ENABLED=True) def test_post_no_document_id(self, mock_get_objects, mock_stream_chat) -> None: + self.grant_view_document_permission() mock_get_objects.return_value = [self.document] mock_stream_chat.return_value = iter([b"data"]) response = self.client.post( @@ -464,6 +471,7 @@ class TestAIChatStreamingView(DirectoriesMixin, TestCase): @patch("documents.views.stream_chat_with_documents") @override_settings(AI_ENABLED=True) def test_post_with_document_id(self, mock_stream_chat) -> None: + self.grant_view_document_permission() mock_stream_chat.return_value = iter([b"data"]) response = self.client.post( self.ENDPOINT, @@ -475,6 +483,7 @@ class TestAIChatStreamingView(DirectoriesMixin, TestCase): @override_settings(AI_ENABLED=True) def test_post_with_invalid_document_id(self) -> None: + self.grant_view_document_permission() response = self.client.post( self.ENDPOINT, data='{"q": "question", "document_id": 999999}', @@ -486,6 +495,7 @@ class TestAIChatStreamingView(DirectoriesMixin, TestCase): @patch("documents.views.has_perms_owner_aware") @override_settings(AI_ENABLED=True) def test_post_with_document_id_no_permission(self, mock_has_perms) -> None: + self.grant_view_document_permission() mock_has_perms.return_value = False response = self.client.post( self.ENDPOINT, @@ -494,3 +504,31 @@ class TestAIChatStreamingView(DirectoriesMixin, TestCase): ) self.assertEqual(response.status_code, 403) self.assertIn(b"Insufficient permissions", response.content) + + @patch("documents.views.stream_chat_with_documents") + @override_settings(AI_ENABLED=True) + def test_post_no_document_id_requires_view_document_permission( + self, + mock_stream_chat, + ) -> None: + response = self.client.post( + self.ENDPOINT, + data='{"q": "question"}', + content_type="application/json", + ) + self.assertEqual(response.status_code, 403) + mock_stream_chat.assert_not_called() + + @patch("documents.views.stream_chat_with_documents") + @override_settings(AI_ENABLED=True) + def test_post_with_document_id_requires_view_document_permission( + self, + mock_stream_chat, + ) -> None: + response = self.client.post( + self.ENDPOINT, + data=f'{{"q": "question", "document_id": {self.document.pk}}}', + content_type="application/json", + ) + self.assertEqual(response.status_code, 403) + mock_stream_chat.assert_not_called() diff --git a/src/documents/views.py b/src/documents/views.py index d93c23c16..cc508012e 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -2150,7 +2150,7 @@ class ChatStreamingSerializer(serializers.Serializer[dict[str, Any]]): name="dispatch", ) class ChatStreamingView(GenericAPIView[Any]): - permission_classes = (IsAuthenticated,) + permission_classes = (IsAuthenticated, ViewDocumentsPermissions) serializer_class = ChatStreamingSerializer def post(self, request, *args, **kwargs): diff --git a/src/documents/workflows/webhooks.py b/src/documents/workflows/webhooks.py index 0c510a35d..e2eff558d 100644 --- a/src/documents/workflows/webhooks.py +++ b/src/documents/workflows/webhooks.py @@ -4,70 +4,12 @@ import httpx from celery import shared_task from django.conf import settings -from paperless.network import format_host_for_url -from paperless.network import is_public_ip -from paperless.network import resolve_hostname_ips +from paperless.network import PinnedHostHTTPTransport from paperless.network import validate_outbound_http_url logger = logging.getLogger("paperless.workflows.webhooks") -class WebhookTransport(httpx.HTTPTransport): - """ - Transport that resolves/validates hostnames and rewrites to a vetted IP - while keeping Host/SNI as the original hostname. - """ - - def __init__( - self, - hostname: str, - *args, - allow_internal: bool = False, - **kwargs, - ) -> None: - super().__init__(*args, **kwargs) - self.hostname = hostname - self.allow_internal = allow_internal - - def handle_request(self, request: httpx.Request) -> httpx.Response: - hostname = request.url.host - - if not hostname: - raise httpx.ConnectError("No hostname in request URL") - - try: - ips = resolve_hostname_ips(hostname) - except ValueError as e: - raise httpx.ConnectError(str(e)) from e - - if not self.allow_internal: - for ip_str in ips: - if not is_public_ip(ip_str): - raise httpx.ConnectError( - f"Connection blocked: {hostname} resolves to a non-public address", - ) - - ip_str = ips[0] - formatted_ip = format_host_for_url(ip_str) - - new_headers = httpx.Headers(request.headers) - if "host" in new_headers: - del new_headers["host"] - new_headers["Host"] = hostname - new_url = request.url.copy_with(host=formatted_ip) - - request = httpx.Request( - method=request.method, - url=new_url, - headers=new_headers, - content=request.stream, - extensions=request.extensions, - ) - request.extensions["sni_hostname"] = hostname - - return super().handle_request(request) - - @shared_task( retry_backoff=True, autoretry_for=(httpx.HTTPStatusError,), @@ -83,7 +25,7 @@ def send_webhook( as_json: bool = False, ): try: - parsed = validate_outbound_http_url( + validate_outbound_http_url( url, allowed_schemes=settings.WEBHOOKS_ALLOWED_SCHEMES, allowed_ports=settings.WEBHOOKS_ALLOWED_PORTS, @@ -94,12 +36,7 @@ def send_webhook( logger.warning("Webhook blocked: %s", e) raise - hostname = parsed.hostname - if hostname is None: # pragma: no cover - raise ValueError("Invalid URL scheme or hostname.") - - transport = WebhookTransport( - hostname=hostname, + transport = PinnedHostHTTPTransport( allow_internal=settings.WEBHOOKS_ALLOW_INTERNAL_REQUESTS, ) diff --git a/src/paperless/network.py b/src/paperless/network.py index 4af00052f..8a5b28a91 100644 --- a/src/paperless/network.py +++ b/src/paperless/network.py @@ -4,6 +4,8 @@ from collections.abc import Collection from urllib.parse import ParseResult from urllib.parse import urlparse +import httpx + def is_public_ip(ip: str | int) -> bool: try: @@ -74,3 +76,121 @@ def validate_outbound_http_url( ) return parsed + + +def _rewrite_request_to_pinned_ip( + request: httpx.Request, + *, + allow_internal: bool, +) -> httpx.Request: + hostname = request.url.host + + if not hostname: + raise httpx.ConnectError("No hostname in request URL") + + try: + ips = resolve_hostname_ips(hostname) + except ValueError as e: + raise httpx.ConnectError(str(e)) from e + + if not allow_internal: + for ip_str in ips: + if not is_public_ip(ip_str): + raise httpx.ConnectError( + f"Connection blocked: {hostname} resolves to a non-public address", + ) + + ip_str = ips[0] + formatted_ip = format_host_for_url(ip_str) + + new_headers = httpx.Headers(request.headers) + if "host" in new_headers: + del new_headers["host"] + host_header = format_host_for_url(hostname) + default_port = 443 if request.url.scheme == "https" else 80 + if request.url.port and request.url.port != default_port: + host_header = f"{host_header}:{request.url.port}" + new_headers["Host"] = host_header + new_url = request.url.copy_with(host=formatted_ip) + + rewritten_request = httpx.Request( + method=request.method, + url=new_url, + headers=new_headers, + content=request.stream, + extensions=request.extensions, + ) + rewritten_request.extensions["sni_hostname"] = hostname + + return rewritten_request + + +class PinnedHostHTTPTransport(httpx.HTTPTransport): + """ + HTTP transport that resolves/validates hostnames per request and connects to + a vetted IP while preserving the original Host header and TLS SNI hostname. + """ + + def __init__( + self, + *args, + allow_internal: bool = False, + **kwargs, + ) -> None: + super().__init__(*args, **kwargs) + self.allow_internal = allow_internal + + def handle_request(self, request: httpx.Request) -> httpx.Response: + request = _rewrite_request_to_pinned_ip( + request, + allow_internal=self.allow_internal, + ) + return super().handle_request(request) + + +class PinnedHostAsyncHTTPTransport(httpx.AsyncHTTPTransport): + """ + Async variant of PinnedHostHTTPTransport. + """ + + def __init__( + self, + *args, + allow_internal: bool = False, + **kwargs, + ) -> None: + super().__init__(*args, **kwargs) + self.allow_internal = allow_internal + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + request = _rewrite_request_to_pinned_ip( + request, + allow_internal=self.allow_internal, + ) + return await super().handle_async_request(request) + + +def create_pinned_httpx_client( + url: str, + *, + allow_internal: bool = False, + **kwargs, +) -> httpx.Client: + validate_outbound_http_url(url, allow_internal=allow_internal) + return httpx.Client( + transport=PinnedHostHTTPTransport(allow_internal=allow_internal), + **kwargs, + ) + + +def create_pinned_async_httpx_client( + url: str, + *, + allow_internal: bool = False, + **kwargs, +) -> httpx.AsyncClient: + validate_outbound_http_url(url, allow_internal=allow_internal) + return httpx.AsyncClient( + transport=PinnedHostAsyncHTTPTransport(allow_internal=allow_internal), + **kwargs, + ) diff --git a/src/paperless/tests/test_network.py b/src/paperless/tests/test_network.py new file mode 100644 index 000000000..a306c52b4 --- /dev/null +++ b/src/paperless/tests/test_network.py @@ -0,0 +1,50 @@ +from unittest import mock + +import httpx +import pytest + +from paperless.network import PinnedHostHTTPTransport + + +def test_pinned_host_transport_blocks_internal_rebinding(): + transport = PinnedHostHTTPTransport(allow_internal=False) + request = httpx.Request("GET", "http://example.com/test") + + with ( + mock.patch( + "paperless.network.resolve_hostname_ips", + return_value=["127.0.0.1"], + ), + pytest.raises(httpx.ConnectError, match="non-public address"), + ): + transport.handle_request(request) + + +def test_pinned_host_transport_rewrites_to_vetted_ip(): + transport = PinnedHostHTTPTransport(allow_internal=False) + request = httpx.Request("GET", "https://example.com:8443/test") + + def assert_rewritten_request( + self, + rewritten_request, + ): + assert str(rewritten_request.url) == "https://93.184.216.34:8443/test" + assert rewritten_request.headers["Host"] == "example.com:8443" + assert rewritten_request.extensions["sni_hostname"] == "example.com" + return httpx.Response(200, request=rewritten_request) + + with ( + mock.patch( + "paperless.network.resolve_hostname_ips", + return_value=["93.184.216.34"], + ), + mock.patch.object( + httpx.HTTPTransport, + "handle_request", + autospec=True, + side_effect=assert_rewritten_request, + ), + ): + response = transport.handle_request(request) + + assert response.status_code == 200 diff --git a/src/paperless_ai/ai_classifier.py b/src/paperless_ai/ai_classifier.py index e60ca37ff..c522a89f9 100644 --- a/src/paperless_ai/ai_classifier.py +++ b/src/paperless_ai/ai_classifier.py @@ -60,11 +60,14 @@ def get_context_for_document( if user else None ) + visible_document_ids = ( + list(visible_documents.values_list("pk", flat=True)) + if visible_documents is not None + else None + ) similar_docs = query_similar_documents( document=doc, - document_ids=[document.pk for document in visible_documents] - if visible_documents - else None, + document_ids=visible_document_ids, )[:max_docs] context_blocks = [] for similar in similar_docs: diff --git a/src/paperless_ai/client.py b/src/paperless_ai/client.py index ce1874461..d4bcef0c8 100644 --- a/src/paperless_ai/client.py +++ b/src/paperless_ai/client.py @@ -9,6 +9,10 @@ if TYPE_CHECKING: from llama_index.llms.openai_like import OpenAILike from paperless.config import AIConfig +from paperless.network import PinnedHostAsyncHTTPTransport +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 DocumentClassifierSchema @@ -27,23 +31,47 @@ class AIClient: def get_llm(self) -> "Ollama | OpenAILike": if self.settings.llm_backend == LLMBackend.OLLAMA: from llama_index.llms.ollama import Ollama + from ollama import AsyncClient + from ollama import Client endpoint = self.settings.llm_endpoint or "http://localhost:11434" validate_outbound_http_url( endpoint, allow_internal=self.settings.llm_allow_internal_endpoints, ) + transport = PinnedHostHTTPTransport( + allow_internal=self.settings.llm_allow_internal_endpoints, + ) + async_transport = PinnedHostAsyncHTTPTransport( + allow_internal=self.settings.llm_allow_internal_endpoints, + ) return Ollama( model=self.settings.llm_model or "llama3.1", base_url=endpoint, request_timeout=120, + client=Client( + host=endpoint, + timeout=120, + transport=transport, + ), + async_client=AsyncClient( + host=endpoint, + timeout=120, + transport=async_transport, + ), ) elif self.settings.llm_backend == LLMBackend.OPENAI_LIKE: from llama_index.llms.openai_like import OpenAILike endpoint = self.settings.llm_endpoint or None + http_client = None + async_http_client = None if endpoint: - validate_outbound_http_url( + http_client = create_pinned_httpx_client( + endpoint, + allow_internal=self.settings.llm_allow_internal_endpoints, + ) + async_http_client = create_pinned_async_httpx_client( endpoint, allow_internal=self.settings.llm_allow_internal_endpoints, ) @@ -53,6 +81,8 @@ class AIClient: api_key=self.settings.llm_api_key, is_chat_model=True, is_function_calling_model=True, + http_client=http_client, + async_http_client=async_http_client, ) else: raise ValueError(f"Unsupported LLM backend: {self.settings.llm_backend}") diff --git a/src/paperless_ai/embedding.py b/src/paperless_ai/embedding.py index 7ef841e4b..407dd4c0e 100644 --- a/src/paperless_ai/embedding.py +++ b/src/paperless_ai/embedding.py @@ -13,6 +13,10 @@ from documents.models import Document from documents.models import Note from paperless.config import AIConfig from paperless.models import LLMEmbeddingBackend +from paperless.network import PinnedHostAsyncHTTPTransport +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 OCR_LEADER_REGEX = re.compile(r"[._\-\u00b7]{4,}") @@ -27,8 +31,14 @@ def get_embedding_model() -> "BaseEmbedding": from llama_index.embeddings.openai_like import OpenAILikeEmbedding endpoint = config.llm_embedding_endpoint or config.llm_endpoint or None + http_client = None + async_http_client = None if endpoint: - validate_outbound_http_url( + http_client = create_pinned_httpx_client( + endpoint, + allow_internal=config.llm_allow_internal_endpoints, + ) + async_http_client = create_pinned_async_httpx_client( endpoint, allow_internal=config.llm_allow_internal_endpoints, ) @@ -36,6 +46,8 @@ def get_embedding_model() -> "BaseEmbedding": model_name=config.llm_embedding_model or "text-embedding-3-small", api_key=config.llm_api_key, api_base=endpoint, + http_client=http_client, + async_http_client=async_http_client, ) case LLMEmbeddingBackend.HUGGINGFACE: from llama_index.embeddings.huggingface import HuggingFaceEmbedding @@ -47,6 +59,8 @@ def get_embedding_model() -> "BaseEmbedding": ) case LLMEmbeddingBackend.OLLAMA: from llama_index.embeddings.ollama import OllamaEmbedding + from ollama import AsyncClient + from ollama import Client endpoint = ( config.llm_embedding_endpoint @@ -57,10 +71,23 @@ def get_embedding_model() -> "BaseEmbedding": endpoint, allow_internal=config.llm_allow_internal_endpoints, ) - return OllamaEmbedding( + embedding = OllamaEmbedding( model_name=config.llm_embedding_model or "embeddinggemma", base_url=endpoint, ) + embedding._client = Client( + host=endpoint, + transport=PinnedHostHTTPTransport( + allow_internal=config.llm_allow_internal_endpoints, + ), + ) + embedding._async_client = AsyncClient( + host=endpoint, + transport=PinnedHostAsyncHTTPTransport( + allow_internal=config.llm_allow_internal_endpoints, + ), + ) + return embedding case _: raise ValueError( f"Unsupported embedding backend: {config.llm_embedding_backend}", diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index d596805fa..28be2c94b 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -1,5 +1,6 @@ import logging import shutil +from collections.abc import Iterable from datetime import timedelta from pathlib import Path from typing import TYPE_CHECKING @@ -327,14 +328,24 @@ def truncate_content(content: str) -> str: return " ".join(truncated_chunks) +def normalize_document_ids(document_ids: Iterable[int | str] | None) -> set[str] | None: + if document_ids is None: + return None + return {str(document_id) for document_id in document_ids} + + def query_similar_documents( document: Document, top_k: int = 5, - document_ids: list[int] | None = None, + document_ids: Iterable[int | str] | None = None, ) -> list[Document]: """ Runs a similarity query and returns top-k similar Document objects. """ + allowed_document_ids = normalize_document_ids(document_ids) + if allowed_document_ids is not None and not allowed_document_ids: + return [] + if not vector_store_file_exists(): queue_llm_index_update_if_needed( rebuild=False, @@ -349,11 +360,13 @@ def query_similar_documents( [ node.node_id for node in index.docstore.docs.values() - if node.metadata.get("document_id") in document_ids + if node.metadata.get("document_id") in allowed_document_ids ] - if document_ids + if allowed_document_ids is not None else None ) + if doc_node_ids is not None and not doc_node_ids: + return [] from llama_index.core.retrievers import VectorIndexRetriever @@ -368,10 +381,23 @@ def query_similar_documents( ) results = retriever.retrieve(query_text) - document_ids = [ - int(node.metadata["document_id"]) - for node in results - if "document_id" in node.metadata - ] + retrieved_document_ids: list[int] = [] + for node in results: + document_id = node.metadata.get("document_id") + if document_id is None: + continue + normalized_document_id = str(document_id) + if ( + allowed_document_ids is not None + and normalized_document_id not in allowed_document_ids + ): + continue + try: + retrieved_document_ids.append(int(normalized_document_id)) + except ValueError: + logger.warning( + "Skipping LLM index result with invalid document_id %r.", + document_id, + ) - return list(Document.objects.filter(pk__in=document_ids)) + return list(Document.objects.filter(pk__in=retrieved_document_ids)) diff --git a/src/paperless_ai/tests/test_ai_indexing.py b/src/paperless_ai/tests/test_ai_indexing.py index aba9e2f7f..356619549 100644 --- a/src/paperless_ai/tests/test_ai_indexing.py +++ b/src/paperless_ai/tests/test_ai_indexing.py @@ -3,6 +3,7 @@ from unittest.mock import MagicMock from unittest.mock import patch import pytest +from django.contrib.auth.models import User from django.test import override_settings from django.utils import timezone from llama_index.core.base.embeddings.base import BaseEmbedding @@ -428,3 +429,78 @@ def test_query_similar_documents_triggers_update_when_index_missing( ) mock_load.assert_not_called() assert result == [] + + +@pytest.mark.django_db +def test_query_similar_documents_normalizes_and_post_filters_allowed_ids( + real_document, +) -> None: + real_document.owner = User.objects.create_user(username="rag-owner") + real_document.save() + private_owner = User.objects.create_user(username="rag-private-owner") + private_document = Document.objects.create( + title="Private similar document", + content="Similar private content that must not reach RAG.", + owner=private_owner, + added=timezone.now(), + ) + + with ( + patch( + "paperless_ai.indexing.vector_store_file_exists", + return_value=True, + ), + patch("paperless_ai.indexing.load_or_build_index") as mock_load_or_build_index, + patch("llama_index.core.retrievers.VectorIndexRetriever") as mock_retriever_cls, + ): + allowed_node = MagicMock() + allowed_node.node_id = "allowed-node" + allowed_node.metadata = {"document_id": str(real_document.pk)} + private_node = MagicMock() + private_node.node_id = "private-node" + private_node.metadata = {"document_id": str(private_document.pk)} + + mock_index = MagicMock() + mock_index.docstore.docs.values.return_value = [allowed_node, private_node] + mock_load_or_build_index.return_value = mock_index + + mock_retriever = MagicMock() + mock_retriever.retrieve.return_value = [private_node, allowed_node] + mock_retriever_cls.return_value = mock_retriever + + result = indexing.query_similar_documents( + real_document, + top_k=2, + document_ids=[real_document.pk], + ) + + mock_retriever_cls.assert_called_once_with( + index=mock_index, + similarity_top_k=2, + doc_ids=["allowed-node"], + ) + assert result == [real_document] + assert private_document not in result + + +@pytest.mark.django_db +def test_query_similar_documents_empty_allow_list_fails_closed( + real_document, +) -> None: + with ( + patch( + "paperless_ai.indexing.vector_store_file_exists", + return_value=True, + ) as mock_vector_store_exists, + patch("paperless_ai.indexing.load_or_build_index") as mock_load_or_build_index, + patch("llama_index.core.retrievers.VectorIndexRetriever") as mock_retriever_cls, + ): + result = indexing.query_similar_documents( + real_document, + document_ids=[], + ) + + assert result == [] + mock_vector_store_exists.assert_not_called() + mock_load_or_build_index.assert_not_called() + mock_retriever_cls.assert_not_called() diff --git a/src/paperless_ai/tests/test_client.py b/src/paperless_ai/tests/test_client.py index 35a881400..ae903b8a0 100644 --- a/src/paperless_ai/tests/test_client.py +++ b/src/paperless_ai/tests/test_client.py @@ -1,3 +1,4 @@ +from unittest.mock import ANY from unittest.mock import MagicMock from unittest.mock import patch @@ -40,6 +41,8 @@ def test_get_llm_ollama(mock_ai_config, mock_ollama_llm): model="test_model", base_url="http://test-url", request_timeout=120, + client=ANY, + async_client=ANY, ) assert client.llm == mock_ollama_llm.return_value @@ -58,6 +61,8 @@ def test_get_llm_openai(mock_ai_config, mock_openai_llm): api_key="test_api_key", is_chat_model=True, is_function_calling_model=True, + http_client=ANY, + async_http_client=ANY, ) assert client.llm == mock_openai_llm.return_value diff --git a/src/paperless_ai/tests/test_embedding.py b/src/paperless_ai/tests/test_embedding.py index f8bc98f7d..d3eff080e 100644 --- a/src/paperless_ai/tests/test_embedding.py +++ b/src/paperless_ai/tests/test_embedding.py @@ -1,4 +1,5 @@ import json +from unittest.mock import ANY from unittest.mock import MagicMock from unittest.mock import patch @@ -70,6 +71,8 @@ def test_get_embedding_model_openai(mock_ai_config): model_name="text-embedding-3-small", api_key="test_api_key", api_base="http://test-url", + http_client=ANY, + async_http_client=ANY, ) assert model == MockOpenAIEmbedding.return_value @@ -89,6 +92,8 @@ def test_get_embedding_model_openai_prefers_embedding_endpoint(mock_ai_config): model_name="text-embedding-3-small", api_key="test_api_key", api_base="http://embedding-url", + http_client=ANY, + async_http_client=ANY, ) assert model == MockOpenAIEmbedding.return_value