diff --git a/src/documents/tests/test_tasks.py b/src/documents/tests/test_tasks.py index 247b6a529..1afeb1d8a 100644 --- a/src/documents/tests/test_tasks.py +++ b/src/documents/tests/test_tasks.py @@ -18,6 +18,7 @@ from documents.models import WorkflowAction from documents.sanity_checker import SanityCheckFailedException from documents.sanity_checker import SanityCheckMessages from documents.tests.helpers import dummy_preprocess +from paperless_ai.exceptions import LLMBlockedError from paperless_testing.assertions import FileSystemAssertsMixin from paperless_testing.dirs import DirectoriesMixin @@ -555,3 +556,37 @@ class TestApplyAISuggestionsTask(DirectoriesMixin, TestCase): apply_suggestions.assert_not_called() self.assertIn("no longer exists", "".join(cm.output)) + + @override_settings(AI_ENABLED=True) + def test_blocked_request_fails_without_retry(self) -> None: + """ + GIVEN: + - AI enabled and a document with content + - The AI classification call blocked by the outbound request policy + WHEN: + - The task runs through Celery + THEN: + - The workflow code does not swallow the block + - The task fails with LLMBlockedError and is never retried + """ + with ( + mock.patch( + "documents.workflows.ai.get_ai_document_classification", + side_effect=LLMBlockedError( + "AI backend request was blocked by the outbound request " + "policy: detail", + ), + ), + mock.patch.object( + tasks.apply_ai_suggestions, + "retry", + wraps=tasks.apply_ai_suggestions.retry, + ) as retry, + ): + result = tasks.apply_ai_suggestions.apply( + args=(self.action.pk, self.doc.pk), + ) + + self.assertTrue(result.failed()) + self.assertIsInstance(result.result, LLMBlockedError) + retry.assert_not_called() diff --git a/src/documents/tests/test_views.py b/src/documents/tests/test_views.py index df2601120..398aed4e0 100644 --- a/src/documents/tests/test_views.py +++ b/src/documents/tests/test_views.py @@ -29,6 +29,7 @@ from documents.models import Tag from documents.models import UiSettings from documents.signals.handlers import update_llm_suggestions_cache from paperless.models import ApplicationConfiguration +from paperless_ai.exceptions import LLMBlockedError from paperless_ai.exceptions import LLMProviderError from paperless_ai.exceptions import LLMTimeoutError from paperless_testing.dirs import DirectoriesMixin @@ -770,6 +771,48 @@ class TestAISuggestions(DirectoriesMixin, TestCase): get_llm_suggestion_cache(self.document.pk, backend="openai-like"), ) + @patch("documents.views.get_ai_document_classification") + @override_settings( + AI_ENABLED=True, + LLM_BACKEND="openai-like", + ) + def test_ai_suggestions_with_blocked_llm_request( + self, + mock_get_ai_classification, + ) -> None: + """ + GIVEN: + - An AI backend request blocked by the outbound request policy + WHEN: + - AI suggestions are requested + THEN: + - 502 is returned with a generic message and nothing is cached + """ + mock_get_ai_classification.side_effect = LLMBlockedError( + "AI backend request was blocked by the outbound request policy: detail", + ) + + self.client.force_login(user=self.user) + response = self.client.get( + f"/api/documents/{self.document.pk}/ai_suggestions/", + ) + + self.assertEqual(response.status_code, status.HTTP_502_BAD_GATEWAY) + self.assertEqual( + response.json(), + { + "ai": [ + ( + "AI backend request was blocked by the outbound request " + "policy. Check logs for details." + ), + ], + }, + ) + self.assertIsNone( + get_llm_suggestion_cache(self.document.pk, backend="openai-like"), + ) + @patch("documents.views.get_ai_document_classification") @override_settings( AI_ENABLED=True, diff --git a/src/documents/views.py b/src/documents/views.py index cc09709a4..a98be332e 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -256,6 +256,7 @@ from paperless.views import StandardPagination from paperless_ai.ai_classifier import get_ai_document_classification from paperless_ai.ai_classifier import get_llm_output_language from paperless_ai.chat import stream_chat_with_documents +from paperless_ai.exceptions import LLMBlockedError from paperless_ai.exceptions import LLMProviderError from paperless_ai.exceptions import LLMTimeoutError from paperless_ai.matching import extract_unmatched_names @@ -1697,6 +1698,23 @@ class DocumentViewSet( }, status=status.HTTP_502_BAD_GATEWAY, ) + except LLMBlockedError as exc: + logger.warning( + "AI backend request for document %s was blocked: %s", + doc.pk, + exc, + ) + return Response( + { + "ai": [ + _( + "AI backend request was blocked by the outbound " + "request policy. Check logs for details.", + ), + ], + }, + status=status.HTTP_502_BAD_GATEWAY, + ) set_llm_suggestions_cache( doc.pk, llm_suggestions, diff --git a/src/paperless_ai/client.py b/src/paperless_ai/client.py index 2402ba4eb..b4b1dfb4e 100644 --- a/src/paperless_ai/client.py +++ b/src/paperless_ai/client.py @@ -16,12 +16,14 @@ if TYPE_CHECKING: from paperless.config import AIConfig from paperless.network import GuardedAsyncHTTPTransport from paperless.network import GuardedHTTPTransport +from paperless.network import OutboundRequestBlockedError from paperless.network import create_guarded_async_httpx_client from paperless.network import create_guarded_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 LLMBlockedError from paperless_ai.exceptions import LLMProviderError from paperless_ai.exceptions import LLMTimeoutError @@ -43,6 +45,19 @@ LLM_SYSTEM_PROMPT = ( PLACEHOLDER_API_KEY: Final = "fake" +def _find_blocked_cause(exc: BaseException) -> OutboundRequestBlockedError | None: + # The openai SDK wraps transport errors in APIConnectionError, so the + # block can sit anywhere in the __cause__ chain. + current: BaseException | None = exc + seen: set[int] = set() + while current is not None and id(current) not in seen: + if isinstance(current, OutboundRequestBlockedError): + return current + seen.add(id(current)) + current = current.__cause__ + return None + + class AIClient: """ A client for interacting with an LLM backend. @@ -179,6 +194,12 @@ class AIClient: except httpx.TimeoutException as exc: raise LLMTimeoutError from exc except Exception as exc: + blocked = _find_blocked_cause(exc) + if blocked is not None: + raise LLMBlockedError( + "AI backend request was blocked by the outbound request " + f"policy: {blocked}", + ) from exc if self._is_openai_timeout(exc): raise LLMTimeoutError from exc if self._is_provider_error(exc): diff --git a/src/paperless_ai/exceptions.py b/src/paperless_ai/exceptions.py index ab3415e85..1fa4e979e 100644 --- a/src/paperless_ai/exceptions.py +++ b/src/paperless_ai/exceptions.py @@ -4,3 +4,7 @@ class LLMTimeoutError(Exception): class LLMProviderError(Exception): """The LLM backend rejected the request.""" + + +class LLMBlockedError(Exception): + """The outbound request policy refused the connection to the LLM backend.""" diff --git a/src/paperless_ai/tests/test_client.py b/src/paperless_ai/tests/test_client.py index d3d0d8115..729787203 100644 --- a/src/paperless_ai/tests/test_client.py +++ b/src/paperless_ai/tests/test_client.py @@ -1,3 +1,4 @@ +import ipaddress import json from unittest.mock import ANY from unittest.mock import MagicMock @@ -9,9 +10,12 @@ import openai import pytest from llama_index.core.llms.llm import ToolSelection +from paperless.network import BlockReason +from paperless.network import OutboundRequestBlockedError from paperless_ai.client import LLM_SYSTEM_PROMPT from paperless_ai.client import PLACEHOLDER_API_KEY from paperless_ai.client import AIClient +from paperless_ai.exceptions import LLMBlockedError from paperless_ai.exceptions import LLMProviderError from paperless_ai.exceptions import LLMTimeoutError from paperless_testing.outbound import guard_of @@ -351,3 +355,69 @@ class TestGuardedLLMClients: kwargs = mock_openai_llm.call_args.kwargs assert guard_of(kwargs["http_client"])._allow_internal is allow_internal assert guard_of(kwargs["async_http_client"])._allow_internal is allow_internal + + +def _block() -> OutboundRequestBlockedError: + return OutboundRequestBlockedError( + host="llm.example", + port=443, + reason=BlockReason.NON_PUBLIC_ADDRESS, + address=ipaddress.ip_address("10.0.0.1"), + ) + + +class TestBlockedLLMRequests: + def test_ollama_block_becomes_llm_blocked_error( + self, + mock_ai_config: MagicMock, + mock_ollama_llm: MagicMock, + ) -> None: + """ + GIVEN: + - The Ollama backend and a connection blocked by policy + WHEN: + - An LLM query runs + THEN: + - LLMBlockedError is raised with a message, chained to the block + - The message, which tracked tasks store, names the destination but + not the resolved internal address + """ + mock_ai_config.llm_backend = "ollama" + mock_ai_config.llm_model = "test_model" + mock_ai_config.llm_endpoint = "http://test-url" + block = _block() + mock_ollama_llm.return_value.chat.side_effect = block + + with pytest.raises(LLMBlockedError) as exc_info: + AIClient().run_llm_query("test_prompt") + + assert exc_info.value.__cause__ is block + assert "llm.example:443" in str(exc_info.value) + assert "10.0.0.1" not in str(exc_info.value) + + def test_openai_wrapped_block_becomes_llm_blocked_error( + self, + mock_ai_config: MagicMock, + mock_openai_llm: MagicMock, + ) -> None: + """ + GIVEN: + - The OpenAI-like backend, whose SDK wraps the block in + APIConnectionError + WHEN: + - An LLM query runs + THEN: + - LLMBlockedError is raised + """ + mock_ai_config.llm_backend = "openai-like" + mock_ai_config.llm_model = "test_model" + mock_ai_config.llm_api_key = "key" + mock_ai_config.llm_endpoint = "http://test-url" + wrapped = openai.APIConnectionError( + request=httpx.Request("POST", "http://test-url/v1/chat/completions"), + ) + wrapped.__cause__ = _block() + mock_openai_llm.return_value.chat_with_tools.side_effect = wrapped + + with pytest.raises(LLMBlockedError): + AIClient().run_llm_query("test_prompt")