paperless-ngx/src/paperless/tests/test_ai_client.py

import json
from unittest.mock import patch

import pytest
from django.conf import settings

from paperless.ai.client import AIClient


@pytest.fixture
def mock_settings():
    settings.LLM_BACKEND = "openai"
    settings.LLM_MODEL = "gpt-3.5-turbo"
    settings.LLM_API_KEY = "test-api-key"
    yield settings


@pytest.mark.django_db
@patch("paperless.ai.client.AIClient._run_openai_query")
@patch("paperless.ai.client.AIClient._run_ollama_query")
def test_run_llm_query_openai(mock_ollama_query, mock_openai_query, mock_settings):
    mock_settings.LLM_BACKEND = "openai"
    mock_openai_query.return_value = "OpenAI response"
    client = AIClient()
    result = client.run_llm_query("Test prompt")
    assert result == "OpenAI response"
    mock_openai_query.assert_called_once_with("Test prompt")
    mock_ollama_query.assert_not_called()


@pytest.mark.django_db
@patch("paperless.ai.client.AIClient._run_openai_query")
@patch("paperless.ai.client.AIClient._run_ollama_query")
def test_run_llm_query_ollama(mock_ollama_query, mock_openai_query, mock_settings):
    mock_settings.LLM_BACKEND = "ollama"
    mock_ollama_query.return_value = "Ollama response"
    client = AIClient()
    result = client.run_llm_query("Test prompt")
    assert result == "Ollama response"
    mock_ollama_query.assert_called_once_with("Test prompt")
    mock_openai_query.assert_not_called()


@pytest.mark.django_db
def test_run_llm_query_unsupported_backend(mock_settings):
    mock_settings.LLM_BACKEND = "unsupported"
    client = AIClient()
    with pytest.raises(ValueError, match="Unsupported LLM backend: unsupported"):
        client.run_llm_query("Test prompt")


@pytest.mark.django_db
def test_run_openai_query(httpx_mock, mock_settings):
    mock_settings.LLM_BACKEND = "openai"
    httpx_mock.add_response(
        url="https://api.openai.com/v1/chat/completions",
        json={
            "choices": [{"message": {"content": "OpenAI response"}}],
        },
    )

    client = AIClient()
    result = client.run_llm_query("Test prompt")
    assert result == "OpenAI response"

    request = httpx_mock.get_request()
    assert request.method == "POST"
    assert request.headers["Authorization"] == f"Bearer {mock_settings.LLM_API_KEY}"
    assert request.headers["Content-Type"] == "application/json"
    assert json.loads(request.content) == {
        "model": mock_settings.LLM_MODEL,
        "messages": [{"role": "user", "content": "Test prompt"}],
        "temperature": 0.3,
    }


@pytest.mark.django_db
def test_run_ollama_query(httpx_mock, mock_settings):
    mock_settings.LLM_BACKEND = "ollama"
    httpx_mock.add_response(
        url="http://localhost:11434/api/chat",
        json={"message": {"content": "Ollama response"}},
    )

    client = AIClient()
    result = client.run_llm_query("Test prompt")
    assert result == "Ollama response"

    request = httpx_mock.get_request()
    assert request.method == "POST"
    assert json.loads(request.content) == {
        "model": mock_settings.LLM_MODEL,
        "messages": [{"role": "user", "content": "Test prompt"}],
        "stream": False,
    }