6krrt/tests/test_dispatcher_helpers.py

"""Unit tests for the new availability/relevance helpers in dispatcher.py.

Offline. Mocks ``requests.post`` so no model/network is reached. Covers the
fail-safe ``_embed_for_relevance``, the min_candidates/disabled gates in
``_relevance_order_for``, the circuit-breaker exclusion set ``_open_circuits``,
and the connection-opening helper ``_open_upstream``.
"""

from __future__ import annotations

import dispatcher
import circuit_breaker


class _FakeResp:
    def __init__(self, status_code=200, text="", json_payload=None, lines=None):
        self.status_code = status_code
        self._text = text
        self._json = json_payload
        self._lines = lines or []
        self.closed = False

    @property
    def text(self):
        return self._text

    def json(self):
        if self._json is None:
            raise ValueError("no json")
        return self._json

    def close(self):
        self.closed = True

    @property
    def encoding(self):
        """The charset requests would pick, by requests' own rule.

        Derived rather than asserted, so this fake tracks requests instead of
        restating a belief about it. The default is a charset-less
        ``text/event-stream`` on purpose: that is what OpenRouter actually
        sends, and `get_encoding_from_headers` answers ISO-8859-1 for any
        ``text/*`` without a charset. Decoding a UTF-8 stream with that and
        re-encoding it produced mojibake on every streamed non-ASCII
        character -- so the hostile case is the DEFAULT here, and a fake can
        no longer make a broken passthrough look correct.
        """
        from requests.structures import CaseInsensitiveDict
        from requests.utils import get_encoding_from_headers

        headers = CaseInsensitiveDict(
            getattr(self, "headers", None) or {"Content-Type": "text/event-stream"}
        )
        return get_encoding_from_headers(headers) or "utf-8"

    def iter_lines(self, decode_unicode=False):
        # Faithful to requests in both modes: BYTES unless decode_unicode is
        # set, and when it is set the charset comes from Content-Type.
        #
        # The `wire` line is the load-bearing one. Whatever a test wrote into
        # `lines`, what a provider actually puts on the wire is UTF-8 bytes,
        # so that is what gets decoded. An earlier version of this fake passed
        # str lines straight through under decode_unicode, which modelled a
        # stream that had ALREADY been decoded correctly -- and a fake that
        # hands back the right answer cannot reproduce a charset bug. It let
        # the whole streaming suite pass against a proxy that was mangling
        # every non-ASCII character.
        for line in self._lines:
            wire = line.encode("utf-8") if isinstance(line, str) else line
            yield wire.decode(self.encoding) if decode_unicode else wire


def _cfg_with_relevance(enabled=True, min_candidates=2):
    class _Rel:
        enabled = True
        model = "nomic-embed-text"
        base_url = "http://localhost:11434/v1"
        timeout_seconds = 10
        min_candidates = 2

    if not enabled:
        _Rel.enabled = False
    if min_candidates != 2:
        _Rel.min_candidates = min_candidates

    class _Pinch:
        enabled = True
        keep_last_turns = 1
        budget_tokens = 100
        max_summarize_chars = 4000
        relevance = _Rel()

    class _CB:
        enabled = False

    class _Cfg:
        pinch = _Pinch()
        circuit_breaker = _CB()

    return _Cfg()


def _embedding_data(n_vectors, dim=4):
    data = []
    for i in range(n_vectors):
        vec = [0.0] * dim
        vec[i % dim] = 1.0
        data.append({"embedding": vec, "index": i})
    return data


def test_embed_for_relevance_happy_path_returns_order(monkeypatch):
    captured = {}

    def fake_post(url, headers=None, json=None, timeout=None):
        captured["json"] = json
        # Return vectors: query (idx0), candidate0 = same as query (most
        # relevant), candidate1 = orthogonal (least relevant).
        payload = {
            "data": [
                {"embedding": [1.0, 0.0, 0.0, 0.0]},
                {"embedding": [1.0, 0.0, 0.0, 0.0]},
                {"embedding": [0.0, 1.0, 0.0, 0.0]},
            ]
        }
        return _FakeResp(status_code=200, json_payload=payload)

    monkeypatch.setattr(dispatcher.requests, "post", fake_post)
    order = dispatcher._embed_for_relevance("hi", ["a", "b"], _cfg_with_relevance())
    # candidate1 (b) is least relevant, so it is compressed first -> index 1.
    assert order == [1, 0]
    assert captured["json"]["model"] == "nomic-embed-text"


def test_embed_for_relevance_failure_paths_return_none(monkeypatch):
    cases = [
        _FakeResp(status_code=500, text="err"),            # non-200
        _FakeResp(status_code=200),                        # unparseable (no json)
        _FakeResp(status_code=200, json_payload={"data": None}),  # malformed
        _FakeResp(status_code=200, json_payload={"data": [{"embedding": []}]}),
    ]
    for resp in cases:
        monkeypatch.setattr(dispatcher.requests, "post", lambda *a, _r=resp, **k: _r)
        assert dispatcher._embed_for_relevance("q", ["a"], _cfg_with_relevance()) is None


def test_embed_for_relevance_request_exception_returns_none(monkeypatch):
    def boom(*a, **k):
        raise dispatcher.requests.RequestException("timeout")

    monkeypatch.setattr(dispatcher.requests, "post", boom)
    assert dispatcher._embed_for_relevance("q", ["a"], _cfg_with_relevance()) is None


def test_relevance_order_for_returns_none_when_disabled():
    cfg = _cfg_with_relevance(enabled=False)
    assert dispatcher._relevance_order_for([], cfg, budget_tokens=100) is None


def test_relevance_order_for_below_min_candidates_skips_embedding(monkeypatch):
    called = {"n": 0}

    def fake_post(*a, **k):
        called["n"] += 1
        return _FakeResp(status_code=200, json_payload={"data": None})

    monkeypatch.setattr(dispatcher.requests, "post", fake_post)
    cfg = _cfg_with_relevance(enabled=True, min_candidates=2)
    # One candidate only -> below min_candidates -> no embed call.
    messages = [
        {"role": "user", "content": "q"},
        {"role": "tool", "name": "read", "content": "x" * 6000},
        {"role": "assistant", "content": "a"},
    ]
    assert dispatcher._relevance_order_for(messages, cfg, budget_tokens=100) is None
    assert called["n"] == 0


def test_relevance_order_for_skips_embedding_when_under_budget(monkeypatch):
    called = {"n": 0}

    def fake_post(*a, **k):
        called["n"] += 1
        return _FakeResp(status_code=200, json_payload={"data": None})

    monkeypatch.setattr(dispatcher.requests, "post", fake_post)
    cfg = _cfg_with_relevance(enabled=True, min_candidates=2)
    # Two short tool-results -> token estimate under budget -> no embed call.
    messages = [
        {"role": "user", "content": "q"},
        {"role": "tool", "name": "read", "content": "short"},
        {"role": "tool", "name": "read", "content": "also short"},
        {"role": "assistant", "content": "a"},
    ]
    assert dispatcher._relevance_order_for(messages, cfg, budget_tokens=50000) is None
    assert called["n"] == 0


def test_open_circuits_empty_when_disabled():
    cfg = _cfg_with_relevance()
    cfg.circuit_breaker.enabled = False
    rows = [{"model_id": "m1", "provider": "p1"}]
    assert dispatcher._open_circuits(rows, cfg) == set()


def test_open_circuits_lists_down_models(monkeypatch):
    cfg = _cfg_with_relevance()
    cfg.circuit_breaker.enabled = True
    circuit_breaker.clear()
    try:
        monkeypatch.setattr(circuit_breaker.time, "time", lambda: 1000.0)
        circuit_breaker.record_failure("m1", "p1", 30, 600, 2.0)
        rows = [
            {"model_id": "m1", "provider": "p1"},   # down
            {"model_id": "m2", "provider": "p1"},   # fine
        ]
        assert dispatcher._open_circuits(rows, cfg) == {"m1"}
    finally:
        circuit_breaker.clear()


def test_open_upstream_posts_stream_with_candidate_model(monkeypatch):
    captured = {}

    def fake_post(url, headers=None, json=None, stream=None, timeout=None):
        captured["model"] = json["model"]
        captured["stream"] = stream
        return _FakeResp(status_code=200)

    monkeypatch.setattr(dispatcher.requests, "post", fake_post)
    dispatcher._open_upstream("candidate-model", "http://x", {"authorization": "Bearer k"},
                              {"model": "ignore", "messages": []})
    assert captured["model"] == "candidate-model"
    assert captured["stream"] is True