diff --git a/src/documents/tests/test_views.py b/src/documents/tests/test_views.py index 82208fd77..376cb8e93 100644 --- a/src/documents/tests/test_views.py +++ b/src/documents/tests/test_views.py @@ -5,7 +5,6 @@ from pathlib import Path from unittest.mock import MagicMock from unittest.mock import patch -import httpx from django.conf import settings from django.contrib.auth.models import Group from django.contrib.auth.models import Permission @@ -31,7 +30,7 @@ from documents.signals.handlers import update_llm_suggestions_cache from documents.tests.utils import DirectoriesMixin from documents.tests.utils import read_streaming_response from paperless.models import ApplicationConfiguration -from paperless_ai.client import LLMTimeoutError +from paperless_ai.exceptions import LLMTimeoutError class TestViews(DirectoriesMixin, TestCase): @@ -486,33 +485,6 @@ class TestAISuggestions(DirectoriesMixin, TestCase): def test_ai_suggestions_with_llm_timeout( self, mock_get_ai_classification, - ) -> None: - mock_get_ai_classification.side_effect = httpx.ReadTimeout("timed out") - - 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_503_SERVICE_UNAVAILABLE) - self.assertEqual( - response.json(), - { - "ai": ["AI backend request timed out."], - }, - ) - self.assertIsNone( - 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_openai_timeout( - self, - mock_get_ai_classification, ) -> None: mock_get_ai_classification.side_effect = LLMTimeoutError() diff --git a/src/documents/views.py b/src/documents/views.py index e70da35c6..0b8ce0f2a 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -241,7 +241,7 @@ from paperless.serialisers import UserSerializer from paperless.views import StandardPagination from paperless_ai.ai_classifier import get_ai_document_classification from paperless_ai.chat import stream_chat_with_documents -from paperless_ai.client import LLMTimeoutError +from paperless_ai.exceptions import LLMTimeoutError from paperless_ai.matching import extract_unmatched_names from paperless_ai.matching import match_correspondents_by_name from paperless_ai.matching import match_document_types_by_name @@ -1511,7 +1511,7 @@ class DocumentViewSet( exc_info=True, ) raise ValidationError({"ai": [_("Invalid AI configuration.")]}) from exc - except (httpx.TimeoutException, LLMTimeoutError) as exc: + except LLMTimeoutError as exc: logger.exception( "AI backend timed out while generating suggestions for document %s: %s", doc.pk, diff --git a/src/paperless_ai/client.py b/src/paperless_ai/client.py index 70235b71d..bb3054081 100644 --- a/src/paperless_ai/client.py +++ b/src/paperless_ai/client.py @@ -1,11 +1,14 @@ import json import logging +from collections.abc import Iterator +from contextlib import contextmanager from typing import TYPE_CHECKING +import httpx + from paperless.models import LLMBackend if TYPE_CHECKING: - from llama_index.core.llms import ChatMessage from llama_index.llms.ollama import Ollama from llama_index.llms.openai_like import OpenAILike @@ -16,6 +19,7 @@ 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 +from paperless_ai.exceptions import LLMTimeoutError logger = logging.getLogger("paperless_ai.client") @@ -31,10 +35,6 @@ LLM_SYSTEM_PROMPT = ( ) -class LLMTimeoutError(Exception): - pass - - class AIClient: """ A client for interacting with an LLM backend. @@ -120,11 +120,12 @@ class AIClient: user_msg = ChatMessage(role="user", content=prompt) if self.settings.llm_backend == LLMBackend.OLLAMA: - result = self.llm.chat( - [user_msg], - format=DocumentClassifierSchema.model_json_schema(), - think=False, - ) + with self._normalize_timeouts(): + result = self.llm.chat( + [user_msg], + format=DocumentClassifierSchema.model_json_schema(), + think=False, + ) logger.debug("LLM query result: %s", result) parsed = DocumentClassifierSchema(**json.loads(result.message.content)) return parsed.model_dump() @@ -132,7 +133,7 @@ class AIClient: from llama_index.core.program.function_program import get_function_tool tool = get_function_tool(DocumentClassifierSchema) - try: + with self._normalize_timeouts(): result = self.llm.chat_with_tools( tools=[tool], user_msg=user_msg, @@ -143,34 +144,27 @@ class AIClient: result, error_on_no_tool_call=True, ) - except Exception as exc: - self._raise_llm_timeout_if_openai_timeout(exc) - raise logger.debug("LLM query result: %s", tool_calls) parsed = DocumentClassifierSchema(**tool_calls[0].tool_kwargs) return parsed.model_dump() - def run_chat(self, messages: list["ChatMessage"]) -> str: - logger.debug( - "Running chat query against %s with model %s", - self.settings.llm_backend, - self.settings.llm_model, - ) + @contextmanager + def _normalize_timeouts(self) -> Iterator[None]: try: - result = self.llm.chat(messages) + yield + except httpx.TimeoutException as exc: + raise LLMTimeoutError from exc except Exception as exc: - self._raise_llm_timeout_if_openai_timeout(exc) + if self._is_openai_timeout(exc): + raise LLMTimeoutError from exc raise - logger.debug("Chat result: %s", result) - return result - def _raise_llm_timeout_if_openai_timeout(self, exc: Exception) -> None: + def _is_openai_timeout(self, exc: Exception) -> bool: if self.settings.llm_backend != LLMBackend.OPENAI_LIKE: - return + return False # Keep OpenAI imports out of module import paths and only load the SDK # when translating an error from an OpenAI-backed request. from openai import APITimeoutError - if isinstance(exc, APITimeoutError): - raise LLMTimeoutError from exc + return isinstance(exc, APITimeoutError) diff --git a/src/paperless_ai/exceptions.py b/src/paperless_ai/exceptions.py new file mode 100644 index 000000000..48ac1db5d --- /dev/null +++ b/src/paperless_ai/exceptions.py @@ -0,0 +1,2 @@ +class LLMTimeoutError(Exception): + pass diff --git a/src/paperless_ai/tests/test_client.py b/src/paperless_ai/tests/test_client.py index 8d68e7060..b51c22acb 100644 --- a/src/paperless_ai/tests/test_client.py +++ b/src/paperless_ai/tests/test_client.py @@ -6,12 +6,11 @@ from unittest.mock import patch import httpx import openai import pytest -from llama_index.core.llms import ChatMessage from llama_index.core.llms.llm import ToolSelection from paperless_ai.client import LLM_SYSTEM_PROMPT from paperless_ai.client import AIClient -from paperless_ai.client import LLMTimeoutError +from paperless_ai.exceptions import LLMTimeoutError @pytest.fixture @@ -176,17 +175,18 @@ def test_run_llm_query_openai_timeout_raises_local_error( client.run_llm_query("test_prompt") -def test_run_chat(mock_ai_config, mock_ollama_llm): +def test_run_llm_query_httpx_timeout_raises_local_error( + mock_ai_config, + mock_ollama_llm, +): mock_ai_config.llm_backend = "ollama" mock_ai_config.llm_model = "test_model" mock_ai_config.llm_endpoint = "http://test-url" mock_llm_instance = mock_ollama_llm.return_value - mock_llm_instance.chat.return_value = "test_chat_result" + mock_llm_instance.chat.side_effect = httpx.ReadTimeout("timed out") client = AIClient() - messages = [ChatMessage(role="user", content="Hello")] - result = client.run_chat(messages) - mock_llm_instance.chat.assert_called_once_with(messages) - assert result == "test_chat_result" + with pytest.raises(LLMTimeoutError): + client.run_llm_query("test_prompt")