diff --git a/src/paperless_ai/client.py b/src/paperless_ai/client.py index 64a3be45c..8c2a9abf4 100644 --- a/src/paperless_ai/client.py +++ b/src/paperless_ai/client.py @@ -182,17 +182,22 @@ class AIClient: except Exception as exc: if self._is_openai_timeout(exc): raise LLMTimeoutError from exc - if self._is_openai_status_error(exc): + if self._is_provider_error(exc): raise LLMProviderError from exc raise - def _is_openai_status_error(self, exc: Exception) -> bool: - if self.settings.llm_backend != LLMBackend.OPENAI_LIKE: - return False + def _is_provider_error(self, exc: Exception) -> bool: + if self.settings.llm_backend == LLMBackend.OLLAMA: + from ollama import ResponseError - from openai import APIStatusError + return isinstance(exc, ResponseError) - return isinstance(exc, APIStatusError) + if self.settings.llm_backend == LLMBackend.OPENAI_LIKE: + from openai import APIStatusError + + return isinstance(exc, APIStatusError) + + return False def _is_openai_timeout(self, exc: Exception) -> bool: if self.settings.llm_backend != LLMBackend.OPENAI_LIKE: diff --git a/src/paperless_ai/tests/test_client.py b/src/paperless_ai/tests/test_client.py index 8d6c52c29..79bb6ad44 100644 --- a/src/paperless_ai/tests/test_client.py +++ b/src/paperless_ai/tests/test_client.py @@ -4,6 +4,7 @@ from unittest.mock import MagicMock from unittest.mock import patch import httpx +import ollama import openai import pytest from llama_index.core.llms.llm import ToolSelection @@ -239,6 +240,28 @@ def test_run_llm_query_openai_status_error_raises_provider_error( assert isinstance(exc_info.value.__cause__, openai.BadRequestError) +def test_run_llm_query_ollama_response_error_raises_provider_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" + + response_error = ollama.ResponseError( + "confidential provider response", + status_code=400, + ) + mock_ollama_llm.return_value.chat.side_effect = response_error + + client = AIClient() + + with pytest.raises(LLMProviderError) as exc_info: + client.run_llm_query("test_prompt") + assert str(exc_info.value) == "" + assert exc_info.value.__cause__ is response_error + + def test_run_llm_query_httpx_timeout_raises_local_error( mock_ai_config, mock_ollama_llm,