conjurer/tests/unit/test_ai_provider_switch.py

"""Unit tests for the GPT/Claude provider switch.

The unit CI job installs only pytest, so the heavy runtime deps that
``ai_functions`` imports unguarded (``openai``, ``tiktoken``, ``other_functions``
-> ``discord``) are stubbed *only when genuinely absent*. Locally, where the
real packages exist, the stubs are skipped and the real modules are used.
"""
import sys
import types


def _stub_if_missing(name: str, build):
    if name in sys.modules:
        return
    try:  # real package present (local dev / bot image) -> use it
        __import__(name)
    except ImportError:
        sys.modules[name] = build()


def _build_tiktoken():
    mod = types.ModuleType("tiktoken")

    class _Enc:
        def encode(self, text):
            return list(text)

    mod.encoding_for_model = lambda _model: _Enc()
    return mod


def _build_other_functions():
    mod = types.ModuleType("other_functions")

    async def _noop(*_a, **_k):
        return None

    mod.discord_friendly_send = _noop
    mod.discord_friendly_reply = _noop
    return mod


def _build_openai():
    mod = types.ModuleType("openai")
    for cls_name in (
        "APITimeoutError",
        "APIConnectionError",
        "BadRequestError",
        "APIResponseValidationError",
        "AuthenticationError",
        "PermissionDeniedError",
        "RateLimitError",
        "UnprocessableEntityError",
        "APIError",
        "OpenAIError",
    ):
        setattr(mod, cls_name, type(cls_name, (Exception,), {}))
    return mod


_stub_if_missing("tiktoken", _build_tiktoken)
_stub_if_missing("other_functions", _build_other_functions)
_stub_if_missing("openai", _build_openai)

import ai_functions  # noqa: E402  (import after stubbing)
import openai  # noqa: E402  (real or stub, same object ai_functions uses)


def _reset_active(name="gpt"):
    ai_functions._ACTIVE_CONFIG_NAME = name


def test_split_extracts_system_and_starts_with_user():
    system, convo = ai_functions._to_anthropic_messages(
        [
            {"role": "system", "content": "SYS1"},
            {"role": "system", "content": "SYS2"},
            {"role": "assistant", "content": "leading-assistant-dropped"},
            {"role": "user", "content": "u1"},
            {"role": "assistant", "content": "a1"},
            {"role": "user", "content": "hi"},
        ]
    )
    assert system == "SYS1\n\nSYS2"
    assert convo[0] == {"role": "user", "content": "u1"}
    assert convo == [
        {"role": "user", "content": "u1"},
        {"role": "assistant", "content": "a1"},
        {"role": "user", "content": "hi"},
    ]


def test_split_synthesises_user_turn_when_only_system():
    _system, convo = ai_functions._to_anthropic_messages(
        [{"role": "system", "content": "only"}]
    )
    assert convo == [{"role": "user", "content": " "}]


def test_select_model_gpt_active():
    _reset_active("gpt")
    # legacy default auto-selects; MUSIC is cheap; explicit ids are honoured
    assert ai_functions.select_model("GENERAL", "gpt-4o") == "gpt-4o"
    assert ai_functions.select_model("MUSIC", "gpt-4o") == "gpt-4o-mini"
    assert ai_functions.select_model("MUSIC", "auto") == "gpt-4o-mini"
    assert ai_functions.select_model("GENERAL", "o1-preview") == "o1-preview"


def test_select_model_claude_active_maps_legacy_default():
    _reset_active("claude")
    try:
        # the old gpt-4o default must not leak to Claude - it auto-maps
        assert ai_functions.select_model("GENERAL", "gpt-4o") == "claude-opus-4-8"
        assert ai_functions.select_model("MUSIC", "gpt-4o") == "claude-haiku-4-5"
        # a real, deliberate model id is still honoured verbatim
        assert ai_functions.select_model("GENERAL", "claude-sonnet-5") == "claude-sonnet-5"
    finally:
        _reset_active("gpt")


def test_list_ai_configs_hides_templates():
    names = ai_functions.list_ai_configs()
    assert "_template" not in names
    assert {"gpt", "claude"}.issubset(set(names))


def _bare(cls):
    # Build an instance without invoking __init__ - the real openai SDK
    # exceptions require response/body kwargs, the CI stubs don't. isinstance
    # (all _map_openai_error cares about) works on __new__-created instances.
    return cls.__new__(cls)


def test_map_openai_error_categories():
    assert ai_functions._map_openai_error(_bare(openai.AuthenticationError)).category == "auth"
    assert ai_functions._map_openai_error(_bare(openai.RateLimitError)).category == "rate_limit"
    assert ai_functions._map_openai_error(_bare(openai.APITimeoutError)).category == "timeout"
    assert ai_functions._map_openai_error(ValueError("x")).category == "api"


# ---------------------------------------------------------------- Ollama
# The self-hosted backend is wired through Ollama's OpenAI-compatible surface,
# so it reuses the message format and the error mapping above. What is new and
# worth pinning: it must appear in the picker even on an upgraded settings file,
# models come from the SERVER, and pinning one must stick.
import asyncio  # noqa: E402


class _FakeModel:
    def __init__(self, ident):
        self.id = ident


class _FakeModels:
    def __init__(self, ids, raises=None):
        self._ids = ids
        self._raises = raises

    async def list(self):
        if self._raises:
            raise self._raises
        return types.SimpleNamespace(data=[_FakeModel(i) for i in self._ids])


class _FakeOllamaClient:
    def __init__(self, ids=(), raises=None):
        self.models = _FakeModels(list(ids), raises)


def test_ollama_config_is_offered_in_the_picker():
    assert "ollama" in ai_functions.list_ai_configs()
    cfg = ai_functions.AI_CONFIGS["ollama"]
    assert cfg["provider"] == "ollama"


def test_select_model_uses_ollama_models_when_active(monkeypatch):
    monkeypatch.setitem(
        ai_functions.AI_CONFIGS,
        "ollama",
        {"provider": "ollama", "latest_model": "llama3.1:8b", "cheap_model": "qwen2.5:3b"},
    )
    _reset_active("ollama")
    try:
        # the legacy gpt-4o default must auto-map, not leak to Ollama
        assert ai_functions.select_model("GENERAL", "gpt-4o") == "llama3.1:8b"
        assert ai_functions.select_model("MUSIC", "gpt-4o") == "qwen2.5:3b"
        # an explicit id is still honoured verbatim
        assert ai_functions.select_model("GENERAL", "mistral:7b") == "mistral:7b"
    finally:
        _reset_active("gpt")


def test_list_provider_models_queries_the_ollama_server(monkeypatch):
    monkeypatch.setattr(
        ai_functions, "OLLAMACLIENT", _FakeOllamaClient(["b:2", "a:1", "a:1"])
    )
    models = asyncio.run(ai_functions.list_provider_models("ollama"))
    assert models == ["a:1", "b:2"]  # sorted + de-duplicated


def test_list_provider_models_for_hosted_provider_reports_configured_ids():
    models = asyncio.run(ai_functions.list_provider_models("gpt"))
    assert models == [
        ai_functions.AI_CONFIGS["gpt"]["latest_model"],
        ai_functions.AI_CONFIGS["gpt"]["cheap_model"],
    ]


def test_list_provider_models_wraps_server_failure(monkeypatch):
    monkeypatch.setattr(
        ai_functions, "OLLAMACLIENT", _FakeOllamaClient(raises=ValueError("boom"))
    )
    try:
        asyncio.run(ai_functions.list_provider_models("ollama"))
    except ai_functions.AIError as exc:
        assert exc.category == "api"
    else:
        raise AssertionError("a server failure must surface as AIError")


def test_list_provider_models_without_endpoint_is_an_auth_error(monkeypatch):
    # Contrast, so the assertion cannot pass vacuously: with a client present the
    # call succeeds, and ONLY setting it to None turns it into an auth error.
    monkeypatch.setattr(ai_functions, "OLLAMACLIENT", _FakeOllamaClient(["a:1"]))
    assert asyncio.run(ai_functions.list_provider_models("ollama")) == ["a:1"]

    monkeypatch.setattr(ai_functions, "OLLAMACLIENT", None)
    try:
        asyncio.run(ai_functions.list_provider_models("ollama"))
    except ai_functions.AIError as exc:
        assert exc.category == "auth"
    else:
        raise AssertionError("an unconfigured Ollama must surface as AIError")


def test_set_active_model_pins_latest_and_keeps_cheap(monkeypatch):
    monkeypatch.setitem(
        ai_functions.AI_CONFIGS,
        "ollama",
        {"provider": "ollama", "latest_model": "old:1", "cheap_model": "cheap:1"},
    )
    written = {}
    monkeypatch.setattr(
        ai_functions,
        "_persist_active_ai_config",
        lambda name, model_for=None: written.update(name=name, model_for=model_for),
    
    )
    ai_functions.set_active_model("mistral:7b", "ollama")
    # Assert on the shared registry, not on the returned object - that object IS
    # the mutated dict, so asserting on it would pass even if nothing was stored.
    stored = ai_functions.AI_CONFIGS["ollama"]
    assert stored["latest_model"] == "mistral:7b"
    assert stored["cheap_model"] == "cheap:1"   # MUSIC path untouched
    assert written["name"]                      # the choice was persisted...
    assert written["model_for"] == "ollama"     # ...scoped to the config we changed


def test_set_active_model_rejects_blank_and_unknown_config(monkeypatch):
    monkeypatch.setattr(
        ai_functions, "_persist_active_ai_config", lambda _n, model_for=None: None
    )
    for bad in ("", "   "):
        try:
            ai_functions.set_active_model(bad, "gpt")
        except ValueError:
            pass
        else:
            raise AssertionError("a blank model id must be rejected")
    try:
        ai_functions.set_active_model("x", "nie-ma-takiego")
    except KeyError:
        pass
    else:
        raise AssertionError("an unknown config must be rejected")


def test_switching_to_ollama_without_endpoint_explains_itself(monkeypatch):
    monkeypatch.setattr(ai_functions, "OLLAMACLIENT", None)
    try:
        ai_functions.set_active_ai_config("ollama")
    except RuntimeError as exc:
        assert "CONJURER_OLLAMA_URL" in str(exc)
    else:
        raise AssertionError("switching to an unconfigured Ollama must raise")
    finally:
        _reset_active("gpt")


def test_provider_generate_routes_to_ollama(monkeypatch):
    monkeypatch.setitem(
        ai_functions.AI_CONFIGS,
        "ollama",
        {"provider": "ollama", "latest_model": "m:1", "cheap_model": "m:1"},
    )
    _reset_active("ollama")
    seen = {}

    async def _fake_ollama_call(messages, model, cfg):
        seen.update(model=model, messages=messages)
        return "odpowiedź z domu"

    monkeypatch.setattr(ai_functions, "_ollama_call", _fake_ollama_call)
    try:
        out = asyncio.run(
            ai_functions.provider_generate([{"role": "user", "content": "hej"}], "m:1")
        )
    finally:
        _reset_active("gpt")
    assert out == "odpowiedź z domu"
    assert seen["model"] == "m:1"


# ------------------------------------------------- persistence (real disk path)
# This path had NO coverage, which is exactly how a config-clobbering regression
# got in: persisting the whole in-memory AI_CONFIGS (built-in defaults merged
# under the file) overwrote operator hand-edits and resurrected deleted configs.
import json  # noqa: E402


def _settings_file(tmp_path, configs, active="gpt"):
    path = tmp_path / "system_gpt_settings.json"
    path.write_text(
        json.dumps(
            [
                {"role": "system", "content": "sys"},
                {"someuser": [1, "a", "b", "c", "asst_x"]},
                {"active": active, "configs": configs},
            ]
        ),
        encoding="utf-8",
    )
    return path


def test_persist_writes_the_pin_without_clobbering_operator_edits(tmp_path, monkeypatch):
    # The file is authoritative for everything the bot does not itself change:
    # a hand-tuned cheap_model, and a config deliberately deleted from it.
    settings = _settings_file(
        tmp_path,
        {"gpt": {"provider": "openai", "latest_model": "gpt-4.1", "cheap_model": "hand-tuned"}},
    )
    monkeypatch.setattr(ai_functions, "SYSTEM_GPT_SETTINGS", str(settings))
    monkeypatch.setitem(
        ai_functions.AI_CONFIGS,
        "ollama",
        {"provider": "ollama", "latest_model": "mistral:7b", "cheap_model": "c:1"},
    )

    ai_functions._persist_active_ai_config("ollama", model_for="ollama")

    data = json.loads(settings.read_text(encoding="utf-8"))
    configs = data[2]["configs"]
    assert data[2]["active"] == "ollama"
    assert configs["ollama"]["latest_model"] == "mistral:7b"   # the pin landed
    assert configs["gpt"]["latest_model"] == "gpt-4.1"         # edit survived
    assert configs["gpt"]["cheap_model"] == "hand-tuned"       # edit survived
    assert "claude" not in configs   # a deleted config is NOT resurrected
    assert data[0]["content"] == "sys" and "someuser" in data[1]  # rest intact


def test_plain_switch_leaves_the_configs_block_untouched(tmp_path, monkeypatch):
    original = {"gpt": {"provider": "openai", "latest_model": "gpt-4.1", "cheap_model": "hand-tuned"}}
    settings = _settings_file(tmp_path, original, active="claude")
    monkeypatch.setattr(ai_functions, "SYSTEM_GPT_SETTINGS", str(settings))

    # Switching backend without pinning a model must only move "active".
    ai_functions._persist_active_ai_config("gpt")

    data = json.loads(settings.read_text(encoding="utf-8"))
    assert data[2]["active"] == "gpt"
    assert data[2]["configs"] == original


def test_pinned_model_survives_a_restart(tmp_path, monkeypatch):
    # The whole point of persisting: re-reading the file must yield the pin.
    settings = _settings_file(tmp_path, {"ollama": {"provider": "ollama", "latest_model": "old:1"}})
    monkeypatch.setattr(ai_functions, "SYSTEM_GPT_SETTINGS", str(settings))
    monkeypatch.setitem(
        ai_functions.AI_CONFIGS,
        "ollama",
        {"provider": "ollama", "latest_model": "new:2", "cheap_model": "c:1"},
    )

    ai_functions._persist_active_ai_config("ollama", model_for="ollama")

    reread = json.loads(settings.read_text(encoding="utf-8"))[2]
    assert reread["configs"]["ollama"]["latest_model"] == "new:2"


def test_persist_survives_an_unreadable_settings_file(tmp_path, monkeypatch):
    # Best-effort by contract: a broken file must not raise into the command.
    broken = tmp_path / "broken.json"
    broken.write_text("{ not json", encoding="utf-8")
    monkeypatch.setattr(ai_functions, "SYSTEM_GPT_SETTINGS", str(broken))
    ai_functions._persist_active_ai_config("gpt", model_for="gpt")  # must not raise


# ----------------------------------------------- keep-warm (Ollama ONLY) ----
# The money guard: preloading a self-hosted model is free, but firing the same
# thing at a metered API would burn tokens for nothing. These pin that it can
# only ever happen for Ollama.


def test_warm_active_model_is_a_noop_for_paid_providers(monkeypatch):
    called = []
    monkeypatch.setattr(
        ai_functions, "_ollama_preload", lambda *a, **k: called.append(a) or True
    )
    for paid in ("gpt", "claude"):
        _reset_active(paid)
        try:
            assert asyncio.run(ai_functions.warm_active_model()) is False
        finally:
            _reset_active("gpt")
    assert called == [], "a paid backend must never be preloaded"


def test_warm_active_model_preloads_when_ollama_is_active(monkeypatch):
    monkeypatch.setitem(
        ai_functions.AI_CONFIGS,
        "ollama",
        {"provider": "ollama", "latest_model": "qwen2.5:7b", "cheap_model": "c"},
    )
    seen = {}
    monkeypatch.setattr(
        ai_functions, "_ollama_preload", lambda model, *a, **k: seen.update(model=model) or True
    )
    _reset_active("ollama")
    try:
        assert asyncio.run(ai_functions.warm_active_model()) is True
    finally:
        _reset_active("gpt")
    assert seen["model"] == "qwen2.5:7b"


def test_active_provider_reports_the_switch():
    _reset_active("gpt")
    assert ai_functions.active_provider() == "openai"
    _reset_active("claude")
    try:
        assert ai_functions.active_provider() == "anthropic"
    finally:
        _reset_active("gpt")


def test_preload_sends_no_prompt_so_it_generates_nothing(monkeypatch):
    # Ollama's documented preload: a model and keep_alive, and NO prompt. If a
    # prompt ever crept in, every warm-up would silently generate tokens.
    sent = {}

    class _Resp:
        status_code = 200

    monkeypatch.setattr(ai_functions, "OLLAMA_URL", "http://ollama:11434")
    monkeypatch.setattr(
        ai_functions.requests, "post",
        lambda url, json=None, timeout=None: sent.update(url=url, body=json) or _Resp(),
    )
    assert ai_functions._ollama_preload("qwen2.5:7b") is True
    assert sent["url"].endswith("/api/generate")
    assert sent["body"]["model"] == "qwen2.5:7b"
    assert "keep_alive" in sent["body"]
    assert "prompt" not in sent["body"], "a preload must not generate"


def test_preload_without_endpoint_is_a_noop(monkeypatch):
    monkeypatch.setattr(ai_functions, "OLLAMA_URL", "")
    assert ai_functions._ollama_preload("x") is False


# ------------------------------------------- personal assistants (per user) --
# Replaces the sunset OpenAI Assistants API. The two properties that matter:
# each user's DM history is ISOLATED (private DMs must not leak into another
# user's context or the bar's shared memory), and it stays BOUNDED.


def _fresh_assistant_memory(tmp_path, monkeypatch, turns=40):
    monkeypatch.setattr(
        ai_functions, "ASSISTANT_MEMORY_FILE", str(tmp_path / "assistant_memory.json")
    )
    monkeypatch.setattr(ai_functions, "ASSISTANT_MEMORY_TURNS", turns)
    monkeypatch.setattr(ai_functions, "_ASSISTANT_MEMORY", None)


def test_assistant_history_is_isolated_per_user(tmp_path, monkeypatch):
    _fresh_assistant_memory(tmp_path, monkeypatch)
    ai_functions.remember_assistant_turn(111, "sekret Anny", "ok Anna")
    ai_functions.remember_assistant_turn(222, "sekret Bartka", "ok Bartek")

    anna = ai_functions.assistant_history(111)
    bartek = ai_functions.assistant_history(222)
    assert [m["content"] for m in anna] == ["sekret Anny", "ok Anna"]
    assert [m["content"] for m in bartek] == ["sekret Bartka", "ok Bartek"]
    assert "sekret Anny" not in str(bartek)   # no cross-user bleed


def test_assistant_history_is_trimmed_to_the_bound(tmp_path, monkeypatch):
    _fresh_assistant_memory(tmp_path, monkeypatch, turns=4)
    for i in range(10):
        ai_functions.remember_assistant_turn(1, f"u{i}", f"a{i}")
    history = ai_functions.assistant_history(1)
    assert len(history) == 4                          # bounded
    assert history[-1]["content"] == "a9"             # newest kept
    assert all("u0" != m["content"] for m in history)  # oldest dropped


def test_assistant_history_survives_a_restart(tmp_path, monkeypatch):
    _fresh_assistant_memory(tmp_path, monkeypatch)
    ai_functions.remember_assistant_turn(7, "pamietaj", "pamietam")
    # Simulate a restart: drop the in-memory cache, re-read from disk.
    monkeypatch.setattr(ai_functions, "_ASSISTANT_MEMORY", None)
    assert [m["content"] for m in ai_functions.assistant_history(7)] == [
        "pamietaj",
        "pamietam",
    ]


def test_assistant_messages_carry_persona_history_and_new_turn(tmp_path, monkeypatch):
    _fresh_assistant_memory(tmp_path, monkeypatch)
    ai_functions.remember_assistant_turn(5, "wczoraj", "odpowiedz")
    msgs = ai_functions.build_assistant_messages(
        5, "Towarzysz Młotek", "Mówisz po polsku.", "dzisiaj"
    )
    assert msgs[0]["role"] == "system"
    assert "Towarzysz Młotek" in msgs[0]["content"]
    assert "Mówisz po polsku." in msgs[0]["content"]
    assert [m["content"] for m in msgs[1:]] == ["wczoraj", "odpowiedz", "dzisiaj"]


def test_corrupt_assistant_memory_starts_empty_instead_of_crashing(tmp_path, monkeypatch):
    path = tmp_path / "assistant_memory.json"
    path.write_text("{ not json", encoding="utf-8")
    monkeypatch.setattr(ai_functions, "ASSISTANT_MEMORY_FILE", str(path))
    monkeypatch.setattr(ai_functions, "_ASSISTANT_MEMORY", None)
    assert ai_functions.assistant_history(1) == []