Enhancement: support ollama embeddings (#12753)

This commit is contained in:
shamoon
2026-05-09 00:06:14 +00:00
committed by GitHub
parent 177d81c8d4
commit 79d0a04df0
12 changed files with 211 additions and 10 deletions
@@ -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())
+4
View File
@@ -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",
),
),
]
+8
View File
@@ -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,
+2
View File
@@ -291,6 +291,8 @@ class ApplicationConfigurationSerializer(
return value
validate_llm_embedding_endpoint = validate_llm_endpoint
class Meta:
model = ApplicationConfiguration
fields = "__all__"
+2 -1
View File
@@ -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")
+25 -5
View File
@@ -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():
+65
View File
@@ -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"