Use a seperate exception for timeout, context manager to determine type

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()
 

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,

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)

src/paperless_ai/exceptions.py

@@ -0,0 +1,2 @@
+class LLMTimeoutError(Exception):
+    pass

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")