mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-08-30 06:27:14 +00:00
Enhancement: support ollama embeddings (#12753)
This commit is contained in:
@@ -74,6 +74,7 @@ class TestApiAppConfig(DirectoriesMixin, APITestCase):
|
||||
"ai_enabled": False,
|
||||
"llm_embedding_backend": None,
|
||||
"llm_embedding_model": None,
|
||||
"llm_embedding_endpoint": None,
|
||||
"llm_backend": None,
|
||||
"llm_model": None,
|
||||
"llm_api_key": None,
|
||||
@@ -868,3 +869,19 @@ class TestApiAppConfig(DirectoriesMixin, APITestCase):
|
||||
)
|
||||
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
|
||||
self.assertIn("non-public address", str(response.data).lower())
|
||||
|
||||
@override_settings(LLM_ALLOW_INTERNAL_ENDPOINTS=False)
|
||||
def test_update_llm_embedding_endpoint_blocks_internal_endpoint_when_disallowed(
|
||||
self,
|
||||
) -> None:
|
||||
response = self.client.patch(
|
||||
f"{self.ENDPOINT}1/",
|
||||
json.dumps(
|
||||
{
|
||||
"llm_embedding_endpoint": "http://127.0.0.1:11434",
|
||||
},
|
||||
),
|
||||
content_type="application/json",
|
||||
)
|
||||
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
|
||||
self.assertIn("non-public address", str(response.data).lower())
|
||||
|
||||
@@ -194,6 +194,7 @@ class AIConfig(BaseConfig):
|
||||
ai_enabled: bool = dataclasses.field(init=False)
|
||||
llm_embedding_backend: str = dataclasses.field(init=False)
|
||||
llm_embedding_model: str = dataclasses.field(init=False)
|
||||
llm_embedding_endpoint: str = dataclasses.field(init=False)
|
||||
llm_backend: str = dataclasses.field(init=False)
|
||||
llm_model: str = dataclasses.field(init=False)
|
||||
llm_api_key: str = dataclasses.field(init=False)
|
||||
@@ -210,6 +211,9 @@ class AIConfig(BaseConfig):
|
||||
self.llm_embedding_model = (
|
||||
app_config.llm_embedding_model or settings.LLM_EMBEDDING_MODEL
|
||||
)
|
||||
self.llm_embedding_endpoint = (
|
||||
app_config.llm_embedding_endpoint or settings.LLM_EMBEDDING_ENDPOINT
|
||||
)
|
||||
self.llm_backend = app_config.llm_backend or settings.LLM_BACKEND
|
||||
self.llm_model = app_config.llm_model or settings.LLM_MODEL
|
||||
self.llm_api_key = app_config.llm_api_key or settings.LLM_API_KEY
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
# Generated by Django 5.2.6 on 2026-05-08 00:00
|
||||
|
||||
from django.db import migrations
|
||||
from django.db import models
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
dependencies = [
|
||||
("paperless", "0009_alter_applicationconfiguration_options"),
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.AlterField(
|
||||
model_name="applicationconfiguration",
|
||||
name="llm_embedding_backend",
|
||||
field=models.CharField(
|
||||
blank=True,
|
||||
choices=[
|
||||
("openai-like", "OpenAI-compatible"),
|
||||
("huggingface", "Huggingface"),
|
||||
("ollama", "Ollama"),
|
||||
],
|
||||
max_length=128,
|
||||
null=True,
|
||||
verbose_name="Sets the LLM embedding backend",
|
||||
),
|
||||
),
|
||||
migrations.AddField(
|
||||
model_name="applicationconfiguration",
|
||||
name="llm_embedding_endpoint",
|
||||
field=models.CharField(
|
||||
blank=True,
|
||||
max_length=256,
|
||||
null=True,
|
||||
verbose_name="Sets the LLM embedding endpoint, optional",
|
||||
),
|
||||
),
|
||||
]
|
||||
@@ -77,6 +77,7 @@ class ColorConvertChoices(models.TextChoices):
|
||||
class LLMEmbeddingBackend(models.TextChoices):
|
||||
OPENAI_LIKE = ("openai-like", _("OpenAI-compatible"))
|
||||
HUGGINGFACE = ("huggingface", _("Huggingface"))
|
||||
OLLAMA = ("ollama", _("Ollama"))
|
||||
|
||||
|
||||
class LLMBackend(models.TextChoices):
|
||||
@@ -310,6 +311,13 @@ class ApplicationConfiguration(AbstractSingletonModel):
|
||||
max_length=128,
|
||||
)
|
||||
|
||||
llm_embedding_endpoint = models.CharField(
|
||||
verbose_name=_("Sets the LLM embedding endpoint, optional"),
|
||||
blank=True,
|
||||
null=True,
|
||||
max_length=256,
|
||||
)
|
||||
|
||||
llm_backend = models.CharField(
|
||||
verbose_name=_("Sets the LLM backend"),
|
||||
blank=True,
|
||||
|
||||
@@ -291,6 +291,8 @@ class ApplicationConfigurationSerializer(
|
||||
|
||||
return value
|
||||
|
||||
validate_llm_embedding_endpoint = validate_llm_endpoint
|
||||
|
||||
class Meta:
|
||||
model = ApplicationConfiguration
|
||||
fields = "__all__"
|
||||
|
||||
@@ -1178,8 +1178,9 @@ REMOTE_OCR_ENDPOINT = os.getenv("PAPERLESS_REMOTE_OCR_ENDPOINT")
|
||||
AI_ENABLED = get_bool_from_env("PAPERLESS_AI_ENABLED", "NO")
|
||||
LLM_EMBEDDING_BACKEND = os.getenv(
|
||||
"PAPERLESS_AI_LLM_EMBEDDING_BACKEND",
|
||||
) # "huggingface" or "openai-like"
|
||||
) # "huggingface", "openai-like", or "ollama"
|
||||
LLM_EMBEDDING_MODEL = os.getenv("PAPERLESS_AI_LLM_EMBEDDING_MODEL")
|
||||
LLM_EMBEDDING_ENDPOINT = os.getenv("PAPERLESS_AI_LLM_EMBEDDING_ENDPOINT")
|
||||
LLM_BACKEND = os.getenv("PAPERLESS_AI_LLM_BACKEND") # "ollama" or "openai-like"
|
||||
LLM_MODEL = os.getenv("PAPERLESS_AI_LLM_MODEL")
|
||||
LLM_API_KEY = os.getenv("PAPERLESS_AI_LLM_API_KEY")
|
||||
|
||||
@@ -22,7 +22,7 @@ def get_embedding_model() -> "BaseEmbedding":
|
||||
case LLMEmbeddingBackend.OPENAI_LIKE:
|
||||
from llama_index.embeddings.openai_like import OpenAILikeEmbedding
|
||||
|
||||
endpoint = config.llm_endpoint or None
|
||||
endpoint = config.llm_embedding_endpoint or config.llm_endpoint or None
|
||||
if endpoint:
|
||||
validate_outbound_http_url(
|
||||
endpoint,
|
||||
@@ -40,6 +40,22 @@ def get_embedding_model() -> "BaseEmbedding":
|
||||
model_name=config.llm_embedding_model
|
||||
or "sentence-transformers/all-MiniLM-L6-v2",
|
||||
)
|
||||
case LLMEmbeddingBackend.OLLAMA:
|
||||
from llama_index.embeddings.ollama import OllamaEmbedding
|
||||
|
||||
endpoint = (
|
||||
config.llm_embedding_endpoint
|
||||
or config.llm_endpoint
|
||||
or "http://localhost:11434"
|
||||
)
|
||||
validate_outbound_http_url(
|
||||
endpoint,
|
||||
allow_internal=config.llm_allow_internal_endpoints,
|
||||
)
|
||||
return OllamaEmbedding(
|
||||
model_name=config.llm_embedding_model or "embeddinggemma",
|
||||
base_url=endpoint,
|
||||
)
|
||||
case _:
|
||||
raise ValueError(
|
||||
f"Unsupported embedding backend: {config.llm_embedding_backend}",
|
||||
@@ -52,11 +68,15 @@ def get_embedding_dim() -> int:
|
||||
from a dummy embedding and stores it for future use.
|
||||
"""
|
||||
config = AIConfig()
|
||||
model = config.llm_embedding_model or (
|
||||
"text-embedding-3-small"
|
||||
if config.llm_embedding_backend == LLMEmbeddingBackend.OPENAI_LIKE
|
||||
else "sentence-transformers/all-MiniLM-L6-v2"
|
||||
default_model = {
|
||||
LLMEmbeddingBackend.OPENAI_LIKE: "text-embedding-3-small",
|
||||
LLMEmbeddingBackend.HUGGINGFACE: "sentence-transformers/all-MiniLM-L6-v2",
|
||||
LLMEmbeddingBackend.OLLAMA: "embeddinggemma",
|
||||
}.get(
|
||||
config.llm_embedding_backend,
|
||||
"sentence-transformers/all-MiniLM-L6-v2",
|
||||
)
|
||||
model = config.llm_embedding_model or default_model
|
||||
|
||||
meta_path: Path = settings.LLM_INDEX_DIR / "meta.json"
|
||||
if meta_path.exists():
|
||||
|
||||
@@ -14,6 +14,7 @@ from paperless_ai.embedding import get_embedding_model
|
||||
@pytest.fixture
|
||||
def mock_ai_config():
|
||||
with patch("paperless_ai.embedding.AIConfig") as MockAIConfig:
|
||||
MockAIConfig.return_value.llm_embedding_endpoint = None
|
||||
MockAIConfig.return_value.llm_allow_internal_endpoints = True
|
||||
yield MockAIConfig
|
||||
|
||||
@@ -71,6 +72,25 @@ def test_get_embedding_model_openai(mock_ai_config):
|
||||
assert model == MockOpenAIEmbedding.return_value
|
||||
|
||||
|
||||
def test_get_embedding_model_openai_prefers_embedding_endpoint(mock_ai_config):
|
||||
mock_ai_config.return_value.llm_embedding_backend = LLMEmbeddingBackend.OPENAI_LIKE
|
||||
mock_ai_config.return_value.llm_embedding_model = "text-embedding-3-small"
|
||||
mock_ai_config.return_value.llm_api_key = "test_api_key"
|
||||
mock_ai_config.return_value.llm_embedding_endpoint = "http://embedding-url"
|
||||
mock_ai_config.return_value.llm_endpoint = "http://test-url"
|
||||
|
||||
with patch(
|
||||
"llama_index.embeddings.openai_like.OpenAILikeEmbedding",
|
||||
) as MockOpenAIEmbedding:
|
||||
model = get_embedding_model()
|
||||
MockOpenAIEmbedding.assert_called_once_with(
|
||||
model_name="text-embedding-3-small",
|
||||
api_key="test_api_key",
|
||||
api_base="http://embedding-url",
|
||||
)
|
||||
assert model == MockOpenAIEmbedding.return_value
|
||||
|
||||
|
||||
def test_get_embedding_model_openai_blocks_internal_endpoint_when_disallowed(
|
||||
mock_ai_config,
|
||||
):
|
||||
@@ -100,6 +120,51 @@ def test_get_embedding_model_huggingface(mock_ai_config):
|
||||
assert model == MockHuggingFaceEmbedding.return_value
|
||||
|
||||
|
||||
def test_get_embedding_model_ollama(mock_ai_config):
|
||||
mock_ai_config.return_value.llm_embedding_backend = LLMEmbeddingBackend.OLLAMA
|
||||
mock_ai_config.return_value.llm_embedding_model = "embeddinggemma"
|
||||
mock_ai_config.return_value.llm_endpoint = "http://test-url"
|
||||
|
||||
with patch(
|
||||
"llama_index.embeddings.ollama.OllamaEmbedding",
|
||||
) as MockOllamaEmbedding:
|
||||
model = get_embedding_model()
|
||||
MockOllamaEmbedding.assert_called_once_with(
|
||||
model_name="embeddinggemma",
|
||||
base_url="http://test-url",
|
||||
)
|
||||
assert model == MockOllamaEmbedding.return_value
|
||||
|
||||
|
||||
def test_get_embedding_model_ollama_prefers_embedding_endpoint(mock_ai_config):
|
||||
mock_ai_config.return_value.llm_embedding_backend = LLMEmbeddingBackend.OLLAMA
|
||||
mock_ai_config.return_value.llm_embedding_model = "embeddinggemma"
|
||||
mock_ai_config.return_value.llm_embedding_endpoint = "http://embedding-url"
|
||||
mock_ai_config.return_value.llm_endpoint = "http://test-url"
|
||||
|
||||
with patch(
|
||||
"llama_index.embeddings.ollama.OllamaEmbedding",
|
||||
) as MockOllamaEmbedding:
|
||||
model = get_embedding_model()
|
||||
MockOllamaEmbedding.assert_called_once_with(
|
||||
model_name="embeddinggemma",
|
||||
base_url="http://embedding-url",
|
||||
)
|
||||
assert model == MockOllamaEmbedding.return_value
|
||||
|
||||
|
||||
def test_get_embedding_model_ollama_blocks_internal_endpoint_when_disallowed(
|
||||
mock_ai_config,
|
||||
):
|
||||
mock_ai_config.return_value.llm_embedding_backend = LLMEmbeddingBackend.OLLAMA
|
||||
mock_ai_config.return_value.llm_embedding_model = "embeddinggemma"
|
||||
mock_ai_config.return_value.llm_endpoint = "http://127.0.0.1:11434"
|
||||
mock_ai_config.return_value.llm_allow_internal_endpoints = False
|
||||
|
||||
with pytest.raises(ValueError, match="non-public address"):
|
||||
get_embedding_model()
|
||||
|
||||
|
||||
def test_get_embedding_model_invalid_backend(mock_ai_config):
|
||||
mock_ai_config.return_value.llm_embedding_backend = "INVALID_BACKEND"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user