odysseus/tests/test_model_context.py

"""Tests for model_context.py — local endpoint detection, token estimation, known model lookup."""

import sys
import types

import pytest

import src.model_context as model_context
from src.model_context import is_local_endpoint, estimate_tokens, _lookup_known


class _Column:
    def __init__(self, name):
        self.name = name

    def __eq__(self, value):
        return ("eq", self.name, value)


class _ModelEndpoint:
    is_enabled = _Column("is_enabled")


class _Query:
    def __init__(self, rows):
        self.rows = list(rows)

    def filter(self, *conditions):
        for condition in conditions:
            if isinstance(condition, tuple) and condition[0] == "eq":
                _, field, value = condition
                self.rows = [row for row in self.rows if getattr(row, field) == value]
        return self

    def all(self):
        return list(self.rows)


class _Db:
    def __init__(self, rows):
        self.rows = rows

    def query(self, model):
        return _Query(self.rows)

    def close(self):
        pass


def _install_endpoint_db(monkeypatch, rows):
    mod = types.ModuleType("core.database")
    mod.ModelEndpoint = _ModelEndpoint
    mod.SessionLocal = lambda: _Db(rows)
    monkeypatch.setitem(sys.modules, "core.database", mod)


class TestIsLocalEndpoint:
    def test_localhost(self):
        assert is_local_endpoint("http://localhost:5000/v1/chat/completions") is True

    def test_loopback_ipv4(self):
        assert is_local_endpoint("http://127.0.0.1:8080/v1/chat/completions") is True

    def test_private_192_168(self):
        assert is_local_endpoint("http://192.168.1.1:11434/v1/chat/completions") is True

    def test_private_10(self):
        assert is_local_endpoint("http://10.0.0.5:8000/v1/chat/completions") is True

    @pytest.mark.parametrize("host", [
        "10.example-cloud.com",
        "172.16.example-cloud.com",
        "192.168.example-cloud.com",
    ])
    def test_private_prefix_dns_names_are_remote(self, host):
        assert is_local_endpoint(f"https://{host}/v1/chat/completions") is False

    def test_tailscale_100(self):
        # 100.64.0.0/10 is the CGNAT range Tailscale uses.
        assert is_local_endpoint("http://100.64.0.1:5000/v1/chat/completions") is True

    def test_configured_tailscale_proxy_is_remote(self, monkeypatch):
        _install_endpoint_db(monkeypatch, [
            types.SimpleNamespace(
                base_url="http://100.117.136.97:34521/v1",
                endpoint_kind="proxy",
                api_key="fake-key",
                is_enabled=True,
            )
        ])

        assert is_local_endpoint("http://100.117.136.97:34521/v1/chat/completions") is False

    def test_openai_is_remote(self):
        assert is_local_endpoint("https://api.openai.com/v1/chat/completions") is False

    def test_anthropic_is_remote(self):
        assert is_local_endpoint("https://api.anthropic.com/v1/messages") is False

    def test_empty_url(self):
        assert is_local_endpoint("") is False

    def test_malformed_url(self):
        assert is_local_endpoint("not-a-url") is False


class TestEstimateTokens:
    def test_empty_list(self):
        assert estimate_tokens([]) == 0

    def test_single_short_message(self):
        messages = [{"role": "user", "content": "Hello"}]
        tokens = estimate_tokens(messages)
        # 4 overhead + int(5 * 0.3) = 4 + 1 = 5
        assert tokens == 5

    def test_multiple_messages(self):
        messages = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "Hi there"},
        ]
        tokens = estimate_tokens(messages)
        assert tokens > 0
        # Each message adds 4 overhead + chars * 0.3
        assert tokens == 4 + int(16 * 0.3) + 4 + int(8 * 0.3)

    def test_multimodal_content_list(self):
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "Describe this image"},
                    {"type": "image_url", "image_url": {"url": "data:..."}},
                ],
            }
        ]
        tokens = estimate_tokens(messages)
        # 4 overhead + int(19 * 0.3) for the text item; image_url is ignored
        assert tokens == 4 + int(19 * 0.3)

    def test_missing_content_key(self):
        messages = [{"role": "assistant"}]
        tokens = estimate_tokens(messages)
        # 4 overhead + 0 content
        assert tokens == 4

    def test_scales_with_length(self):
        short = estimate_tokens([{"role": "user", "content": "short"}])
        long_text = "a" * 10000
        long = estimate_tokens([{"role": "user", "content": long_text}])
        assert long > short * 10


class TestLookupKnown:
    def test_claude_sonnet(self):
        assert _lookup_known("claude-sonnet-4-5") == 200000

    def test_gpt4o(self):
        assert _lookup_known("gpt-4o") == 128000

    def test_deepseek_r1(self):
        assert _lookup_known("deepseek-r1") == 64000

    def test_gemini_pro(self):
        assert _lookup_known("gemini-2.5-pro") == 1048576

    def test_unknown_model(self):
        assert _lookup_known("totally-unknown-model-xyz") is None

    def test_namespaced_model(self):
        """Models prefixed with provider/ should still match."""
        result = _lookup_known("openrouter/deepseek-r1")
        assert result == 64000

    def test_model_with_tag(self):
        """Models with :free or :extended suffixes should still match."""
        result = _lookup_known("deepseek-r1:free")
        assert result == 64000

    def test_o1_mini_not_shadowed_by_o1(self):
        """'o1' (200k) precedes 'o1-mini' (128k) in the table; longest match wins."""
        assert _lookup_known("o1-mini") == 128000

    def test_o1_full(self):
        assert _lookup_known("o1") == 200000

    def test_gpt4o_mini_not_shadowed_by_gpt4(self):
        assert _lookup_known("gpt-4o-mini") == 128000

    def test_gpt4_base(self):
        assert _lookup_known("gpt-4") == 8192


class _FakeResp:
    def __init__(self, payload, ok=True):
        self._payload = payload
        self.is_success = ok

    def json(self):
        return self._payload


class TestGetContextLength:
    def setup_method(self):
        model_context._context_cache.clear()
        model_context._catalog_ctx_cache.clear()
        model_context._ollama_runtime_endpoints.clear()

    def test_local_endpoint_requeries_same_model_after_restart(self, monkeypatch):
        calls = []

        def fake_query(endpoint_url, model):
            calls.append((endpoint_url, model))
            return (8192, True) if len(calls) == 1 else (27000, True)

        monkeypatch.setattr(model_context, "_query_context_length", fake_query)

        endpoint = "http://127.0.0.1:8000/v1/chat/completions"
        model = "Qwen/Qwen3-14B"

        first = model_context.get_context_length(endpoint, model)
        second = model_context.get_context_length(endpoint, model)

        assert first == 8192
        assert second == 27000
        assert len(calls) == 2

    def test_remote_endpoint_keeps_cached_context(self, monkeypatch):
        calls = []

        def fake_query(endpoint_url, model):
            calls.append((endpoint_url, model))
            return (200000, True) if len(calls) == 1 else (12345, True)

        monkeypatch.setattr(model_context, "_query_context_length", fake_query)

        endpoint = "https://api.openai.com/v1/chat/completions"
        model = "gpt-5"

        first = model_context.get_context_length(endpoint, model)
        second = model_context.get_context_length(endpoint, model)

        assert first == 200000
        assert second == 200000
        assert len(calls) == 1

    def _proxy_db(self, monkeypatch):
        _install_endpoint_db(monkeypatch, [
            types.SimpleNamespace(
                base_url="http://100.117.136.97:34521/v1",
                endpoint_kind="proxy",
                api_key="fake-key",
                is_enabled=True,
            )
        ])

    def _api_db(self, monkeypatch, base_url="http://127.0.0.1:11434/v1"):
        _install_endpoint_db(monkeypatch, [
            types.SimpleNamespace(
                base_url=base_url,
                endpoint_kind="api",
                api_key=None,
                is_enabled=True,
            )
        ])

    def test_configured_api_ollama_uses_loaded_runtime_allocation(self, monkeypatch):
        self._api_db(monkeypatch)
        fetches = []

        def fake_get(url, *args, **kwargs):
            fetches.append(url)
            assert url == "http://127.0.0.1:11434/api/ps"
            return _FakeResp({
                "models": [
                    {"model": "qwen3:latest", "context_length": 65536},
                ]
            })

        monkeypatch.setattr(model_context.httpx, "get", fake_get)

        endpoint = "http://127.0.0.1:11434/v1/chat/completions"
        assert model_context.get_context_length(endpoint, "qwen3") == 65536
        assert fetches == ["http://127.0.0.1:11434/api/ps"]

    def test_configured_api_ollama_retries_after_model_load(self, monkeypatch):
        self._api_db(monkeypatch)
        responses = [
            {"models": []},
            {"models": [{"model": "qwen3", "context_length": 32768}]},
        ]

        def fake_get(url, *args, **kwargs):
            assert url == "http://127.0.0.1:11434/api/ps"
            return _FakeResp(responses.pop(0))

        monkeypatch.setattr(model_context.httpx, "get", fake_get)

        endpoint = "http://127.0.0.1:11434/v1/chat/completions"
        assert model_context.get_context_length_known(endpoint, "qwen3") == (
            model_context.DEFAULT_CONTEXT,
            False,
        )
        assert model_context.get_context_length_known(endpoint, "qwen3") == (32768, True)
        assert responses == []

    @pytest.mark.parametrize(
        ("endpoint", "expected"),
        [
            ("http://localhost:11434", "http://localhost:11434/api/ps"),
            ("http://ollama:11434/v1", "http://ollama:11434/api/ps"),
            ("http://localhost:11434/api/chat", "http://localhost:11434/api/ps"),
            ("http://localhost:11434/v1/chat/completions", "http://localhost:11434/api/ps"),
            ("http://localhost:11434/ollama/v1", "http://localhost:11434/ollama/api/ps"),
        ],
    )
    def test_ollama_ps_url_preserves_supported_prefixes(self, endpoint, expected):
        assert model_context._ollama_ps_url(endpoint, "api") == expected

    def test_ollama_ps_url_excludes_remote_and_configured_proxy(self):
        assert model_context._ollama_ps_url("https://api.openai.com/v1", "api") is None
        assert model_context._ollama_ps_url("https://ollama.example.com/v1", "api") is None
        assert model_context._ollama_ps_url("http://127.0.0.1:11434/v1", "proxy") is None
        assert model_context._ollama_ps_url("http://127.0.0.1:8080/v1", None) is None

    def test_configured_proxy_known_model_skips_model_listing(self, monkeypatch):
        # A model covered by the known-context table must still resolve without
        # touching /models — the cheap path the proxy short-circuit exists for.
        self._proxy_db(monkeypatch)

        def fake_get(*args, **kwargs):
            raise AssertionError("/models must not be queried for a known proxy model")

        monkeypatch.setattr(model_context.httpx, "get", fake_get)

        endpoint = "http://100.117.136.97:34521/v1/chat/completions"
        assert model_context.get_context_length(endpoint, "gpt-4o") == 128000

    def test_configured_proxy_unknown_model_reads_catalog_context(self, monkeypatch):
        # A model missing from the known table (e.g. a new OpenRouter model)
        # must report the catalog's real window, not the bare default (#4886).
        # The catalog is fetched once per endpoint and reused for other models.
        self._proxy_db(monkeypatch)
        fetches = []

        def fake_get(url, *args, **kwargs):
            fetches.append(url)
            return _FakeResp({"data": [
                {"id": "owl-alpha", "context_length": 1048576},
                {"id": "tiny-proxy-model", "context_length": 8192},
            ]})

        monkeypatch.setattr(model_context.httpx, "get", fake_get)

        endpoint = "http://100.117.136.97:34521/v1/chat/completions"
        assert model_context.get_context_length(endpoint, "owl-alpha") == 1048576
        # A second unknown model on the same endpoint reuses the cached catalog.
        assert model_context.get_context_length(endpoint, "tiny-proxy-model") == 8192
        assert len(fetches) == 1

    def test_configured_proxy_unknown_model_falls_back_to_default(self, monkeypatch):
        # If the catalog can be read but doesn't list the model, keep the
        # conservative default rather than guessing.
        self._proxy_db(monkeypatch)

        def fake_get(url, *args, **kwargs):
            return _FakeResp({"data": [{"id": "some-other-model", "context_length": 4096}]})

        monkeypatch.setattr(model_context.httpx, "get", fake_get)

        endpoint = "http://100.117.136.97:34521/v1/chat/completions"
        assert model_context.get_context_length(endpoint, "absent-model") == model_context.DEFAULT_CONTEXT

    def test_configured_proxy_catalog_fetch_failure_uses_default(self, monkeypatch):
        # A failed/unreachable catalog must not raise — fall back to the default.
        self._proxy_db(monkeypatch)

        def fake_get(url, *args, **kwargs):
            raise RuntimeError("network down")

        monkeypatch.setattr(model_context.httpx, "get", fake_get)

        endpoint = "http://100.117.136.97:34521/v1/chat/completions"
        assert model_context.get_context_length(endpoint, "unknown-proxy-model") == model_context.DEFAULT_CONTEXT