paperless-ngx/src/paperless_ai/tests/test_client.py

import json
from unittest.mock import ANY
from unittest.mock import MagicMock
from unittest.mock import patch

import httpx
import openai
import pytest
from llama_index.core.llms.llm import ToolSelection

from paperless_ai.client import LLM_SYSTEM_PROMPT
from paperless_ai.client import AIClient
from paperless_ai.exceptions import LLMTimeoutError


@pytest.fixture
def mock_ai_config():
    with patch("paperless_ai.client.AIConfig") as MockAIConfig:
        mock_config = MagicMock()
        mock_config.llm_allow_internal_endpoints = True
        mock_config.llm_context_size = 8192
        mock_config.llm_request_timeout = 120
        MockAIConfig.return_value = mock_config
        yield mock_config


@pytest.fixture
def mock_ollama_llm():
    with patch("llama_index.llms.ollama.Ollama") as MockOllama:
        yield MockOllama


@pytest.fixture
def mock_openai_llm():
    with patch("llama_index.llms.openai_like.OpenAILike") as MockOpenAILike:
        yield MockOpenAILike


def test_get_llm_ollama(mock_ai_config, mock_ollama_llm):
    mock_ai_config.llm_backend = "ollama"
    mock_ai_config.llm_model = "test_model"
    mock_ai_config.llm_endpoint = "http://test-url"

    client = AIClient()

    mock_ollama_llm.assert_called_once_with(
        model="test_model",
        base_url="http://test-url",
        context_window=8192,
        request_timeout=120,
        system_prompt=LLM_SYSTEM_PROMPT,
        client=ANY,
        async_client=ANY,
    )
    assert client.llm == mock_ollama_llm.return_value


def test_get_llm_openai(mock_ai_config, mock_openai_llm):
    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"

    client = AIClient()

    mock_openai_llm.assert_called_once_with(
        model="test_model",
        api_base="http://test-url",
        api_key="test_api_key",
        timeout=120,
        is_chat_model=True,
        is_function_calling_model=True,
        system_prompt=LLM_SYSTEM_PROMPT,
        http_client=ANY,
        async_http_client=ANY,
    )
    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-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"
    mock_ai_config.llm_allow_internal_endpoints = False

    with pytest.raises(ValueError, match="non-public address"):
        AIClient()


def test_get_llm_unsupported_backend(mock_ai_config):
    mock_ai_config.llm_backend = "unsupported"

    with pytest.raises(ValueError, match="Unsupported LLM backend: unsupported"):
        AIClient()


def test_run_llm_query_ollama_uses_structured_json(mock_ai_config, mock_ollama_llm):
    mock_ai_config.llm_backend = "ollama"
    mock_ai_config.llm_model = "test_model"
    mock_ai_config.llm_endpoint = "http://test-url"

    mock_llm_instance = mock_ollama_llm.return_value
    mock_llm_instance.chat.return_value = MagicMock()
    mock_llm_instance.chat.return_value.message.content = json.dumps(
        {
            "title": "Test Title",
            "tags": {"existing_ids": [1], "new_names": ["document"]},
            "correspondents": {"existing_ids": [], "new_names": ["John Doe"]},
            "document_types": {"existing_ids": [], "new_names": ["report"]},
            "storage_paths": {"existing_ids": [], "new_names": ["Reports"]},
            "dates": ["2023-01-01"],
        },
    )

    client = AIClient()
    result = client.run_llm_query("test_prompt")

    assert result["title"] == "Test Title"
    assert result["tags"] == {"existing_ids": [1], "new_names": ["document"]}
    mock_llm_instance.chat.assert_called_once_with(
        [ANY],
        format=ANY,
        think=False,
    )


def test_run_llm_query_openai_uses_tools(mock_ai_config, mock_openai_llm):
    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"

    mock_llm_instance = mock_openai_llm.return_value

    tool_selection = ToolSelection(
        tool_id="call_test",
        tool_name="DocumentClassifierSchema",
        tool_kwargs={
            "title": "Test Title",
            "tags": {"existing_ids": [1], "new_names": ["document"]},
            "correspondents": {"existing_ids": [], "new_names": ["John Doe"]},
            "document_types": {"existing_ids": [], "new_names": ["report"]},
            "storage_paths": {"existing_ids": [], "new_names": ["Reports"]},
            "dates": ["2023-01-01"],
        },
    )

    mock_llm_instance.chat_with_tools.return_value = MagicMock()
    mock_llm_instance.get_tool_calls_from_response.return_value = [tool_selection]

    client = AIClient()
    result = client.run_llm_query("test_prompt")

    assert result["title"] == "Test Title"
    assert result["tags"] == {"existing_ids": [1], "new_names": ["document"]}
    mock_llm_instance.chat_with_tools.assert_called_once()


def test_run_llm_query_openai_timeout_raises_local_error(
    mock_ai_config,
    mock_openai_llm,
):
    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"

    request = httpx.Request("POST", "http://test-url/v1/chat/completions")
    mock_openai_llm.return_value.chat_with_tools.side_effect = openai.APITimeoutError(
        request,
    )

    client = AIClient()

    with pytest.raises(LLMTimeoutError):
        client.run_llm_query("test_prompt")


def test_run_llm_query_httpx_timeout_raises_local_error(
    mock_ai_config,
    mock_ollama_llm,
):
    mock_ai_config.llm_backend = "ollama"
    mock_ai_config.llm_model = "test_model"
    mock_ai_config.llm_endpoint = "http://test-url"

    mock_llm_instance = mock_ollama_llm.return_value
    mock_llm_instance.chat.side_effect = httpx.ReadTimeout("timed out")

    client = AIClient()

    with pytest.raises(LLMTimeoutError):
        client.run_llm_query("test_prompt")