amorfo77/paperless-ngx
1
0
Fork 0
mirror of https://github.com/paperless-ngx/paperless-ngx.git synced 2026-09-27 03:40:30 +00:00
Code Issues Projects Releases Wiki Activity

Use guarded transports for AI LLM and embedding clients

Browse Source 
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
stumpylog
2026-09-22 12:48:21 -07:00
co-authored by Claude Opus 5
parent 526adbad1a
commit 6ad00ca55a
4 changed files with 155 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/paperless_ai/client.py
+8 -8
View File
@@ -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,
src/paperless_ai/embedding.py
+8 -8
View File
@@ -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,
),
)
src/paperless_ai/tests/test_client.py
+74
View File
@@ -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
src/paperless_ai/tests/test_embedding.py
+65
View File
@@ -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