File size: 4,985 Bytes
0f9caed
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
"""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)