mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-09-14 21:58:00 +00:00
Enhancement (dev): Use OpenAI-like backend (#12668)
This commit is contained in:
@@ -1,10 +1,12 @@
|
||||
import logging
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from paperless.models import LLMBackend
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from llama_index.core.llms import ChatMessage
|
||||
from llama_index.llms.ollama import Ollama
|
||||
from llama_index.llms.openai import OpenAI
|
||||
from llama_index.llms.openai_like import OpenAILike
|
||||
|
||||
from paperless.config import AIConfig
|
||||
from paperless.network import validate_outbound_http_url
|
||||
@@ -22,8 +24,8 @@ class AIClient:
|
||||
self.settings = AIConfig()
|
||||
self.llm = self.get_llm()
|
||||
|
||||
def get_llm(self) -> "Ollama | OpenAI":
|
||||
if self.settings.llm_backend == "ollama":
|
||||
def get_llm(self) -> "Ollama | OpenAILike":
|
||||
if self.settings.llm_backend == LLMBackend.OLLAMA:
|
||||
from llama_index.llms.ollama import Ollama
|
||||
|
||||
endpoint = self.settings.llm_endpoint or "http://localhost:11434"
|
||||
@@ -36,8 +38,8 @@ class AIClient:
|
||||
base_url=endpoint,
|
||||
request_timeout=120,
|
||||
)
|
||||
elif self.settings.llm_backend == "openai":
|
||||
from llama_index.llms.openai import OpenAI
|
||||
elif self.settings.llm_backend == LLMBackend.OPENAI_LIKE:
|
||||
from llama_index.llms.openai_like import OpenAILike
|
||||
|
||||
endpoint = self.settings.llm_endpoint or None
|
||||
if endpoint:
|
||||
@@ -45,10 +47,12 @@ class AIClient:
|
||||
endpoint,
|
||||
allow_internal=self.settings.llm_allow_internal_endpoints,
|
||||
)
|
||||
return OpenAI(
|
||||
return OpenAILike(
|
||||
model=self.settings.llm_model or "gpt-3.5-turbo",
|
||||
api_base=endpoint,
|
||||
api_key=self.settings.llm_api_key,
|
||||
is_chat_model=True,
|
||||
is_function_calling_model=True,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported LLM backend: {self.settings.llm_backend}")
|
||||
|
||||
@@ -19,8 +19,8 @@ def get_embedding_model() -> "BaseEmbedding":
|
||||
config = AIConfig()
|
||||
|
||||
match config.llm_embedding_backend:
|
||||
case LLMEmbeddingBackend.OPENAI:
|
||||
from llama_index.embeddings.openai import OpenAIEmbedding
|
||||
case LLMEmbeddingBackend.OPENAI_LIKE:
|
||||
from llama_index.embeddings.openai_like import OpenAILikeEmbedding
|
||||
|
||||
endpoint = config.llm_endpoint or None
|
||||
if endpoint:
|
||||
@@ -28,8 +28,8 @@ def get_embedding_model() -> "BaseEmbedding":
|
||||
endpoint,
|
||||
allow_internal=config.llm_allow_internal_endpoints,
|
||||
)
|
||||
return OpenAIEmbedding(
|
||||
model=config.llm_embedding_model or "text-embedding-3-small",
|
||||
return OpenAILikeEmbedding(
|
||||
model_name=config.llm_embedding_model or "text-embedding-3-small",
|
||||
api_key=config.llm_api_key,
|
||||
api_base=endpoint,
|
||||
)
|
||||
@@ -54,7 +54,7 @@ def get_embedding_dim() -> int:
|
||||
config = AIConfig()
|
||||
model = config.llm_embedding_model or (
|
||||
"text-embedding-3-small"
|
||||
if config.llm_embedding_backend == "openai"
|
||||
if config.llm_embedding_backend == LLMEmbeddingBackend.OPENAI_LIKE
|
||||
else "sentence-transformers/all-MiniLM-L6-v2"
|
||||
)
|
||||
|
||||
|
||||
@@ -98,7 +98,7 @@ def test_update_llm_index_removes_meta(
|
||||
config = AIConfig()
|
||||
expected_model = config.llm_embedding_model or (
|
||||
"text-embedding-3-small"
|
||||
if config.llm_embedding_backend == "openai"
|
||||
if config.llm_embedding_backend == "openai-like"
|
||||
else "sentence-transformers/all-MiniLM-L6-v2"
|
||||
)
|
||||
assert meta == {"embedding_model": expected_model, "dim": 384}
|
||||
|
||||
@@ -25,8 +25,8 @@ def mock_ollama_llm():
|
||||
|
||||
@pytest.fixture
|
||||
def mock_openai_llm():
|
||||
with patch("llama_index.llms.openai.OpenAI") as MockOpenAI:
|
||||
yield MockOpenAI
|
||||
with patch("llama_index.llms.openai_like.OpenAILike") as MockOpenAILike:
|
||||
yield MockOpenAILike
|
||||
|
||||
|
||||
def test_get_llm_ollama(mock_ai_config, mock_ollama_llm):
|
||||
@@ -45,7 +45,7 @@ def test_get_llm_ollama(mock_ai_config, mock_ollama_llm):
|
||||
|
||||
|
||||
def test_get_llm_openai(mock_ai_config, mock_openai_llm):
|
||||
mock_ai_config.llm_backend = "openai"
|
||||
mock_ai_config.llm_backend = "openai-like"
|
||||
mock_ai_config.llm_model = "test_model"
|
||||
mock_ai_config.llm_api_key = "test_api_key"
|
||||
mock_ai_config.llm_endpoint = "http://test-url"
|
||||
@@ -56,12 +56,14 @@ def test_get_llm_openai(mock_ai_config, mock_openai_llm):
|
||||
model="test_model",
|
||||
api_base="http://test-url",
|
||||
api_key="test_api_key",
|
||||
is_chat_model=True,
|
||||
is_function_calling_model=True,
|
||||
)
|
||||
assert client.llm == mock_openai_llm.return_value
|
||||
|
||||
|
||||
def test_get_llm_openai_blocks_internal_endpoint_when_disallowed(mock_ai_config):
|
||||
mock_ai_config.llm_backend = "openai"
|
||||
mock_ai_config.llm_backend = "openai-like"
|
||||
mock_ai_config.llm_model = "test_model"
|
||||
mock_ai_config.llm_api_key = "test_api_key"
|
||||
mock_ai_config.llm_endpoint = "http://127.0.0.1:1234"
|
||||
|
||||
@@ -54,15 +54,17 @@ def mock_document():
|
||||
|
||||
|
||||
def test_get_embedding_model_openai(mock_ai_config):
|
||||
mock_ai_config.return_value.llm_embedding_backend = LLMEmbeddingBackend.OPENAI
|
||||
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_endpoint = "http://test-url"
|
||||
|
||||
with patch("llama_index.embeddings.openai.OpenAIEmbedding") as MockOpenAIEmbedding:
|
||||
with patch(
|
||||
"llama_index.embeddings.openai_like.OpenAILikeEmbedding",
|
||||
) as MockOpenAIEmbedding:
|
||||
model = get_embedding_model()
|
||||
MockOpenAIEmbedding.assert_called_once_with(
|
||||
model="text-embedding-3-small",
|
||||
model_name="text-embedding-3-small",
|
||||
api_key="test_api_key",
|
||||
api_base="http://test-url",
|
||||
)
|
||||
@@ -72,7 +74,7 @@ def test_get_embedding_model_openai(mock_ai_config):
|
||||
def test_get_embedding_model_openai_blocks_internal_endpoint_when_disallowed(
|
||||
mock_ai_config,
|
||||
):
|
||||
mock_ai_config.return_value.llm_embedding_backend = LLMEmbeddingBackend.OPENAI
|
||||
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_endpoint = "http://127.0.0.1:11434"
|
||||
@@ -109,7 +111,7 @@ def test_get_embedding_model_invalid_backend(mock_ai_config):
|
||||
|
||||
|
||||
def test_get_embedding_dim_infers_and_saves(temp_llm_index_dir, mock_ai_config):
|
||||
mock_ai_config.return_value.llm_embedding_backend = "openai"
|
||||
mock_ai_config.return_value.llm_embedding_backend = "openai-like"
|
||||
mock_ai_config.return_value.llm_embedding_model = None
|
||||
|
||||
class DummyEmbedding:
|
||||
@@ -129,7 +131,7 @@ def test_get_embedding_dim_infers_and_saves(temp_llm_index_dir, mock_ai_config):
|
||||
|
||||
|
||||
def test_get_embedding_dim_reads_existing_meta(temp_llm_index_dir, mock_ai_config):
|
||||
mock_ai_config.return_value.llm_embedding_backend = "openai"
|
||||
mock_ai_config.return_value.llm_embedding_backend = "openai-like"
|
||||
mock_ai_config.return_value.llm_embedding_model = None
|
||||
|
||||
(temp_llm_index_dir / "meta.json").write_text(
|
||||
@@ -142,7 +144,7 @@ def test_get_embedding_dim_reads_existing_meta(temp_llm_index_dir, mock_ai_confi
|
||||
|
||||
|
||||
def test_get_embedding_dim_raises_on_model_change(temp_llm_index_dir, mock_ai_config):
|
||||
mock_ai_config.return_value.llm_embedding_backend = "openai"
|
||||
mock_ai_config.return_value.llm_embedding_backend = "openai-like"
|
||||
mock_ai_config.return_value.llm_embedding_model = None
|
||||
|
||||
(temp_llm_index_dir / "meta.json").write_text(
|
||||
|
||||
Reference in New Issue
Block a user