Spaces:
Sleeping
Sleeping
File size: 3,382 Bytes
0bb4dfa | 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 | from app.services.context_budget import (
ContextSummary,
DEFAULT_CONTEXT,
DEFAULT_REPLY_BUDGET,
SUMMARIZE_THRESHOLD,
TRIM_THRESHOLD,
build_compressed_messages,
context_window_for,
estimate_messages_tokens,
select_summarizer_model_id,
should_summarize,
)
def test_context_window_known_model():
assert context_window_for("gpt-4.1-mini") == 128_000
def test_context_window_unknown_falls_back():
assert context_window_for("totally-unknown") == DEFAULT_CONTEXT
def test_context_window_neon_falls_back():
assert context_window_for("neon:Foo/Bar@2025.10.01:Researcher") == DEFAULT_CONTEXT
def test_estimate_tokens_grows_with_text():
a = estimate_messages_tokens([{"role": "user", "content": "hi"}])
b = estimate_messages_tokens([{"role": "user", "content": "hi" * 1000}])
assert b > a
def test_should_summarize_below_threshold():
"""A small ~10-token prompt against a 1M-token gemini window: never
should we trigger summarize.
"""
summary = ContextSummary()
api = [{"role": "user", "content": "hello"}]
needs_sum, needs_trim, _ = should_summarize("gemini-2.5-flash", api, summary)
assert needs_sum is False
assert needs_trim is False
def test_should_summarize_above_threshold_triggers_summarize():
summary = ContextSummary()
big = "x" * 10_000 # ~2500 estimated tokens
api = [
{"role": "system", "content": big},
{"role": "user", "content": big},
]
# Small 8K-window model -> 6K input budget -> 55% = 3300; we're way over.
needs_sum, _, budget = should_summarize("totally-unknown", api, summary)
assert needs_sum is True
assert budget >= 2_048
def test_should_trim_only_when_summary_exists():
"""Even at 70%, trim should only happen once we already have a
running summary - otherwise we'd drop history with no replacement.
"""
summary = ContextSummary()
big = "x" * 30_000
api = [
{"role": "system", "content": big},
{"role": "user", "content": big},
]
_, needs_trim_no_sum, _ = should_summarize("totally-unknown", api, summary)
assert needs_trim_no_sum is False
summary.summary_text = "Previously, the group discussed X."
_, needs_trim_yes_sum, _ = should_summarize("totally-unknown", api, summary)
assert needs_trim_yes_sum is True
def test_build_compressed_keeps_system_and_recent():
summary = ContextSummary(summary_text="condensed history")
msgs = [
{"role": "system", "content": "you are X"},
{"role": "user", "content": "old1"},
{"role": "assistant", "content": "oldA"},
{"role": "user", "content": "old2"},
{"role": "assistant", "content": "oldB"},
{"role": "user", "content": "recent"},
]
out = build_compressed_messages(msgs, summary, needs_trim=True)
# head + summary + last KEEP_RECENT_MESSAGES (=6) <= len(msgs) so we
# might end up with everything; the contract is just that the
# original system and the running summary are preserved.
assert out[0] == msgs[0]
assert any("Summary of earlier" in m["content"] for m in out)
def test_select_summarizer_falls_back():
# Override wins
assert select_summarizer_model_id("custom", "orch") == "custom"
# Falls back to orchestrator if no override
assert select_summarizer_model_id(None, "orch") == "orch"
|