diff --git a/src/paperless_ai/client.py b/src/paperless_ai/client.py index 682a0e86c..2402ba4eb 100644 --- a/src/paperless_ai/client.py +++ b/src/paperless_ai/client.py @@ -14,10 +14,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 GuardedAsyncHTTPTransport +from paperless.network import GuardedHTTPTransport +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 @@ -63,10 +63,10 @@ class AIClient: endpoint, allow_internal=self.settings.llm_allow_internal_endpoints, ) - transport = PinnedHostHTTPTransport( + transport = GuardedHTTPTransport( allow_internal=self.settings.llm_allow_internal_endpoints, ) - async_transport = PinnedHostAsyncHTTPTransport( + async_transport = GuardedAsyncHTTPTransport( allow_internal=self.settings.llm_allow_internal_endpoints, ) return Ollama( @@ -93,12 +93,12 @@ class AIClient: http_client = None async_http_client = None if endpoint: - http_client = create_pinned_httpx_client( + http_client = create_guarded_httpx_client( endpoint, allow_internal=self.settings.llm_allow_internal_endpoints, timeout=self.settings.llm_request_timeout, ) - async_http_client = create_pinned_async_httpx_client( + async_http_client = create_guarded_async_httpx_client( endpoint, allow_internal=self.settings.llm_allow_internal_endpoints, timeout=self.settings.llm_request_timeout, diff --git a/src/paperless_ai/embedding.py b/src/paperless_ai/embedding.py index dbbb2b3be..0a83b900a 100644 --- a/src/paperless_ai/embedding.py +++ b/src/paperless_ai/embedding.py @@ -9,10 +9,10 @@ if TYPE_CHECKING: from documents.models import Document 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 GuardedAsyncHTTPTransport +from paperless.network import GuardedHTTPTransport +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.client import PLACEHOLDER_API_KEY @@ -29,12 +29,12 @@ def get_embedding_model(config: AIConfig) -> "BaseEmbedding": http_client = None async_http_client = None if endpoint: - http_client = create_pinned_httpx_client( + http_client = create_guarded_httpx_client( endpoint, allow_internal=config.llm_allow_internal_endpoints, timeout=config.llm_request_timeout, ) - async_http_client = create_pinned_async_httpx_client( + async_http_client = create_guarded_async_httpx_client( endpoint, allow_internal=config.llm_allow_internal_endpoints, timeout=config.llm_request_timeout, @@ -77,14 +77,14 @@ def get_embedding_model(config: AIConfig) -> "BaseEmbedding": embedding._client = Client( host=endpoint, timeout=config.llm_request_timeout, - transport=PinnedHostHTTPTransport( + transport=GuardedHTTPTransport( allow_internal=config.llm_allow_internal_endpoints, ), ) embedding._async_client = AsyncClient( host=endpoint, timeout=config.llm_request_timeout, - transport=PinnedHostAsyncHTTPTransport( + transport=GuardedAsyncHTTPTransport( allow_internal=config.llm_allow_internal_endpoints, ), ) diff --git a/src/paperless_ai/tests/test_client.py b/src/paperless_ai/tests/test_client.py index 79bb6ad44..d3d0d8115 100644 --- a/src/paperless_ai/tests/test_client.py +++ b/src/paperless_ai/tests/test_client.py @@ -14,6 +14,7 @@ from paperless_ai.client import PLACEHOLDER_API_KEY from paperless_ai.client import AIClient from paperless_ai.exceptions import LLMProviderError from paperless_ai.exceptions import LLMTimeoutError +from paperless_testing.outbound import guard_of @pytest.fixture @@ -277,3 +278,76 @@ def test_run_llm_query_httpx_timeout_raises_local_error( with pytest.raises(LLMTimeoutError): client.run_llm_query("test_prompt") + + +class TestGuardedLLMClients: + @pytest.mark.parametrize( + ("endpoint", "allow_internal"), + [ + pytest.param("http://test-url", True, id="internal-allowed"), + pytest.param("http://93.184.216.34:11434", False, id="internal-blocked"), + ], + ) + def test_ollama_clients_are_guarded( + self, + mock_ai_config: MagicMock, + mock_ollama_llm: MagicMock, + endpoint: str, + *, + allow_internal: bool, + ) -> None: + """ + GIVEN: + - The Ollama backend + WHEN: + - The LLM is built + THEN: + - Its sync and async clients use guarded transports with the setting + """ + mock_ai_config.llm_backend = "ollama" + mock_ai_config.llm_model = "test_model" + mock_ai_config.llm_endpoint = endpoint + mock_ai_config.llm_allow_internal_endpoints = allow_internal + + AIClient() + + kwargs = mock_ollama_llm.call_args.kwargs + assert guard_of(kwargs["client"]._client)._allow_internal is allow_internal + assert ( + guard_of(kwargs["async_client"]._client)._allow_internal is allow_internal + ) + + @pytest.mark.parametrize( + ("endpoint", "allow_internal"), + [ + pytest.param("http://test-url", True, id="internal-allowed"), + pytest.param("http://93.184.216.34:8080", False, id="internal-blocked"), + ], + ) + def test_openai_like_clients_are_guarded( + self, + mock_ai_config: MagicMock, + mock_openai_llm: MagicMock, + endpoint: str, + *, + allow_internal: bool, + ) -> None: + """ + GIVEN: + - The OpenAI-like backend with an endpoint + WHEN: + - The LLM is built + THEN: + - Its sync and async http clients use guarded transports + """ + 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 = endpoint + mock_ai_config.llm_allow_internal_endpoints = allow_internal + + AIClient() + + 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 diff --git a/src/paperless_ai/tests/test_embedding.py b/src/paperless_ai/tests/test_embedding.py index 1373a37bd..e843d194e 100644 --- a/src/paperless_ai/tests/test_embedding.py +++ b/src/paperless_ai/tests/test_embedding.py @@ -1,9 +1,12 @@ +from typing import TYPE_CHECKING +from typing import cast from unittest.mock import ANY from unittest.mock import MagicMock from unittest.mock import patch import pytest from django.conf import settings +from pytest_mock import MockerFixture from documents.models import Document from paperless.models import LLMEmbeddingBackend @@ -12,6 +15,10 @@ from paperless_ai.embedding import _normalize_llm_index_text from paperless_ai.embedding import build_llm_index_text from paperless_ai.embedding import get_configured_model_name from paperless_ai.embedding import get_embedding_model +from paperless_testing.outbound import guard_of + +if TYPE_CHECKING: + from llama_index.embeddings.ollama import OllamaEmbedding @pytest.fixture @@ -283,3 +290,61 @@ def test_normalize_llm_index_text_collapses_ocr_leaders_without_joining_lines(): def test_normalize_llm_index_text_collapses_non_breaking_spaces(): assert _normalize_llm_index_text("A\u00a0........\u00a0B") == "A B" + + +class TestGuardedEmbeddingClients: + def test_ollama_embedding_clients_are_guarded( + self, + mocker: MockerFixture, + mock_ai_config: MagicMock, + ) -> None: + """ + GIVEN: + - The Ollama embedding backend + WHEN: + - The embedding model is built + THEN: + - The clients swapped onto it use guarded transports + """ + config = mock_ai_config.return_value + config.llm_embedding_backend = LLMEmbeddingBackend.OLLAMA + config.llm_embedding_model = "embeddinggemma" + config.llm_endpoint = "http://93.184.216.34:11434" + config.llm_allow_internal_endpoints = False + + mocker.patch("llama_index.embeddings.ollama.OllamaEmbedding") + + model = cast("OllamaEmbedding", get_embedding_model(config)) + + assert guard_of(model._client._client)._allow_internal is False + assert guard_of(model._async_client._client)._allow_internal is False + + def test_openai_like_embedding_clients_are_guarded( + self, + mocker: MockerFixture, + mock_ai_config: MagicMock, + ) -> None: + """ + GIVEN: + - The OpenAI-like embedding backend with an endpoint + WHEN: + - The embedding model is built + THEN: + - Its http clients use guarded transports + """ + config = mock_ai_config.return_value + config.llm_embedding_backend = LLMEmbeddingBackend.OPENAI_LIKE + config.llm_embedding_model = "text-embedding-3-small" + config.llm_api_key = "key" + config.llm_endpoint = "http://93.184.216.34:8080" + config.llm_allow_internal_endpoints = False + + embedding_class = mocker.patch( + "llama_index.embeddings.openai_like.OpenAILikeEmbedding", + ) + + get_embedding_model(config) + + kwargs = embedding_class.call_args.kwargs + assert guard_of(kwargs["http_client"])._allow_internal is False + assert guard_of(kwargs["async_http_client"])._allow_internal is False