tests/test_context.py

"""Testy context.py — odhad tokenů, TokenStats, kompakce kontextu."""

from types import SimpleNamespace

from context import (TRUNCATION_MARK, TokenStats, compact_messages,
                     estimate_tokens)


def _msgs(*parts):
    return list(parts)


def sys_msg(text="system prompt"):
    return {"role": "system", "content": text}


def user_msg(text):
    return {"role": "user", "content": text}


def asst_tool_call(name, call_id):
    return {"role": "assistant", "content": "",
            "tool_calls": [{"id": call_id, "type": "function",
                            "function": {"name": name, "arguments": "{}"}}]}


def tool_msg(call_id, content):
    return {"role": "tool", "tool_call_id": call_id, "content": content}


# ---------------------------------------------------------------- odhad

def test_estimate_tokens_scales_with_length():
    short = estimate_tokens([user_msg("abcd" * 10)])
    long = estimate_tokens([user_msg("abcd" * 1000)])
    assert long > short * 10
    # ~4 znaky/token + režie
    assert abs(estimate_tokens([user_msg("x" * 4000)]) - 1004) < 50


def test_estimate_counts_tool_calls():
    plain = estimate_tokens([{"role": "assistant", "content": "hi"}])
    with_calls = estimate_tokens([asst_tool_call("run_command", "c1")])
    assert with_calls > plain


# ---------------------------------------------------------------- TokenStats

def test_token_stats_object_and_dict_usage():
    stats = TokenStats()
    stats.add("local", SimpleNamespace(prompt_tokens=100, completion_tokens=20))
    stats.add("api", {"prompt_tokens": 50, "completion_tokens": 5})
    stats.add("local", None)  # cache hit / bez usage — nepočítá se
    snap = stats.snapshot()
    assert snap["calls"] == 2
    assert snap["prompt_tokens"] == 150
    assert snap["completion_tokens"] == 25
    assert snap["total_tokens"] == 175
    assert snap["last_context_tokens"] == 50
    assert snap["by_source"]["local"]["calls"] == 1
    assert snap["by_source"]["api"]["prompt_tokens"] == 50


def test_token_stats_saved():
    stats = TokenStats()
    stats.add_saved(500)
    stats.add_saved(-10)  # záporné se ignoruje
    assert stats.snapshot()["saved_by_compaction_tokens"] == 500


# ---------------------------------------------------------------- kompakce

def _history(steps, result_chars=4000):
    msgs = [sys_msg(), user_msg("Oprav bug v projektu")]
    for i in range(steps):
        msgs.append(asst_tool_call(f"tool_{i}", f"c{i}"))
        msgs.append(tool_msg(f"c{i}", "x" * result_chars))
    return msgs


def test_under_budget_untouched():
    msgs = _history(2)
    out, saved = compact_messages(msgs, budget_tokens=10**6)
    assert out == msgs and saved == 0


def test_phase1_ages_old_tool_results():
    msgs = _history(10)
    out, saved = compact_messages(msgs, budget_tokens=8000,
                                  keep_last_steps=2, aged_chars=500)
    assert saved > 0
    # staré tool výsledky zkráceny, poslední 2 bloky netknuté
    aged = [m for m in out if m.get("role") == "tool" and TRUNCATION_MARK in m["content"]]
    full = [m for m in out if m.get("role") == "tool" and TRUNCATION_MARK not in m["content"]]
    assert len(aged) >= 1
    assert len(full) >= 2
    assert out[-1]["content"].endswith("x" * 100)  # poslední výsledek celý
    # vstup nemodifikován
    assert TRUNCATION_MARK not in msgs[3]["content"]


def test_phase2_drops_oldest_blocks_with_summary():
    msgs = _history(20)
    out, saved = compact_messages(msgs, budget_tokens=4000,
                                  keep_last_steps=3, aged_chars=300)
    assert saved > 0
    assert estimate_tokens(out) <= 4000
    # kotva zachována
    assert out[0]["role"] == "system"
    assert out[1]["role"] == "user" and "Oprav bug" in out[1]["content"]
    # souhrnná poznámka o vypuštění
    summary = [m for m in out if "Kompakce kontextu" in str(m.get("content"))]
    assert len(summary) == 1
    assert "tool_0" in summary[0]["content"]
    # struktura validní: každá tool zpráva má před sebou assistant s tool_calls
    for i, m in enumerate(out):
        if m.get("role") == "tool":
            prev = next(p for p in reversed(out[:i])
                        if p.get("role") in ("assistant", "user", "system"))
            assert prev.get("tool_calls"), "tool zpráva bez tool_calls předchůdce"
    # poslední 3 bloky přežily
    tool_names = [tc["function"]["name"] for m in out
                  for tc in (m.get("tool_calls") or [])]
    assert tool_names[-3:] == ["tool_17", "tool_18", "tool_19"]


def test_best_effort_when_budget_impossible():
    msgs = _history(20)
    out, saved = compact_messages(msgs, budget_tokens=10,
                                  keep_last_steps=3, aged_chars=300)
    # nevejde se ani po maximální kompakci — vrací best effort, nikoli prázdno
    assert saved > 0
    assert any(m.get("role") == "user" for m in out)
    assert estimate_tokens(out) < estimate_tokens(msgs)