arc53--docsgpt/tests/llm/test_premai.py

"""Unit tests for application/llm/premai.py — PremAILLM.

Covers:
  - Constructor
  - _raw_gen: API call and return value
  - _raw_gen_stream: streaming with delta content filtering
"""

import sys
import types

import pytest


# ---------------------------------------------------------------------------
# Fake premai module
# ---------------------------------------------------------------------------


class _FakeMessage:
    def __init__(self, content):
        self.message = {"content": content}


class _FakeDelta:
    def __init__(self, content):
        self.delta = {"content": content}


class _FakeChoice:
    def __init__(self, content):
        self.message = {"content": content}


class _FakeStreamChoice:
    def __init__(self, content):
        self.delta = {"content": content}


class _FakeResponse:
    def __init__(self, content="result_text"):
        self.choices = [_FakeChoice(content)]


class _FakeStreamLine:
    def __init__(self, content):
        self.choices = [_FakeStreamChoice(content)]


class _FakeChatCompletions:
    def __init__(self):
        self.last_kwargs = None

    def create(self, **kwargs):
        self.last_kwargs = kwargs
        if kwargs.get("stream"):
            return [
                _FakeStreamLine("chunk1"),
                _FakeStreamLine("chunk2"),
                _FakeStreamLine(None),  # None content should be filtered
            ]
        return _FakeResponse()


class _FakeChat:
    def __init__(self):
        self.completions = _FakeChatCompletions()


class _FakePrem:
    def __init__(self, api_key=None):
        self.api_key = api_key
        self.chat = _FakeChat()


@pytest.fixture(autouse=True)
def patch_premai(monkeypatch):
    fake_mod = types.ModuleType("premai")
    fake_mod.Prem = _FakePrem
    sys.modules["premai"] = fake_mod

    if "application.llm.premai" in sys.modules:
        del sys.modules["application.llm.premai"]
    yield
    sys.modules.pop("premai", None)
    if "application.llm.premai" in sys.modules:
        del sys.modules["application.llm.premai"]


@pytest.fixture
def llm():
    from application.llm.premai import PremAILLM

    return PremAILLM(api_key="test-key")


# ---------------------------------------------------------------------------
# Constructor
# ---------------------------------------------------------------------------


@pytest.mark.unit
class TestPremAIConstructor:

    def test_sets_api_key(self, llm):
        assert llm.api_key == "test-key"

    def test_sets_user_api_key_none(self, llm):
        assert llm.user_api_key is None

    def test_client_created(self, llm):
        assert isinstance(llm.client, _FakePrem)

    def test_project_id_from_settings(self, llm):
        from application.core.settings import settings

        assert llm.project_id == settings.PREMAI_PROJECT_ID


# ---------------------------------------------------------------------------
# _raw_gen
# ---------------------------------------------------------------------------


@pytest.mark.unit
class TestRawGen:

    def test_returns_content(self, llm):
        msgs = [{"role": "user", "content": "hi"}]
        result = llm._raw_gen(llm, model="model-1", messages=msgs)
        assert result == "result_text"

    def test_passes_model_and_project_id(self, llm):
        msgs = [{"role": "user", "content": "hi"}]
        llm._raw_gen(llm, model="my-model", messages=msgs)
        kwargs = llm.client.chat.completions.last_kwargs
        assert kwargs["model"] == "my-model"
        assert kwargs["project_id"] == llm.project_id
        assert kwargs["stream"] is False

    def test_passes_messages(self, llm):
        msgs = [{"role": "user", "content": "hello"}]
        llm._raw_gen(llm, model="m", messages=msgs)
        kwargs = llm.client.chat.completions.last_kwargs
        assert kwargs["messages"] == msgs

    def test_extra_kwargs_forwarded(self, llm):
        msgs = [{"role": "user", "content": "hi"}]
        llm._raw_gen(llm, model="m", messages=msgs, temperature=0.5)
        kwargs = llm.client.chat.completions.last_kwargs
        assert kwargs["temperature"] == 0.5


# ---------------------------------------------------------------------------
# _raw_gen_stream
# ---------------------------------------------------------------------------


@pytest.mark.unit
class TestRawGenStream:

    def test_yields_non_none_content(self, llm):
        msgs = [{"role": "user", "content": "hi"}]
        chunks = list(
            llm._raw_gen_stream(llm, model="m", messages=msgs, stream=True)
        )
        assert chunks == ["chunk1", "chunk2"]

    def test_filters_none_content(self, llm):
        msgs = [{"role": "user", "content": "hi"}]
        chunks = list(
            llm._raw_gen_stream(llm, model="m", messages=msgs, stream=True)
        )
        assert None not in chunks

    def test_passes_stream_true(self, llm):
        msgs = [{"role": "user", "content": "hi"}]
        list(llm._raw_gen_stream(llm, model="m", messages=msgs))
        kwargs = llm.client.chat.completions.last_kwargs
        assert kwargs["stream"] is True

    def test_passes_extra_kwargs(self, llm):
        msgs = [{"role": "user", "content": "hi"}]
        list(
            llm._raw_gen_stream(
                llm, model="m", messages=msgs, max_tokens=100
            )
        )
        kwargs = llm.client.chat.completions.last_kwargs
        assert kwargs["max_tokens"] == 100