79 lines
2.6 KiB
Python
79 lines
2.6 KiB
Python
from api.streaming import (
|
|
_live_usage_session_snapshot,
|
|
live_usage_prompt_estimate_after_tool_delta,
|
|
)
|
|
|
|
|
|
def test_live_usage_estimate_caps_tool_delta_against_previous_prompt():
|
|
usage = live_usage_prompt_estimate_after_tool_delta(
|
|
base_prompt_tokens=86_723,
|
|
exact_prompt_tokens=86_723,
|
|
messages=[{"role": "tool", "content": "x" * 80_000}],
|
|
)
|
|
|
|
assert usage["last_prompt_tokens"] <= 86_723 + 12_000
|
|
assert usage["last_prompt_tokens"] < 120_000
|
|
|
|
|
|
def test_live_usage_estimate_preserves_real_prompt_when_exact_prompt_advances():
|
|
usage = live_usage_prompt_estimate_after_tool_delta(
|
|
base_prompt_tokens=86_723,
|
|
exact_prompt_tokens=136_000,
|
|
messages=[{"role": "tool", "content": "x" * 80_000}],
|
|
)
|
|
|
|
assert usage["last_prompt_tokens"] == 136_000
|
|
assert usage["turn_tool_prompt_tokens"] == 0
|
|
|
|
|
|
def test_live_usage_estimate_caps_cumulative_tool_delta_per_turn():
|
|
from api import streaming
|
|
|
|
base_prompt_tokens = 86_723
|
|
turn_tool_prompt_tokens = 0
|
|
usage = None
|
|
original_delta = streaming._bounded_live_tool_prompt_delta
|
|
streaming._bounded_live_tool_prompt_delta = lambda messages, cap=12_000: int(cap or 0)
|
|
|
|
try:
|
|
for _ in range(20):
|
|
usage = live_usage_prompt_estimate_after_tool_delta(
|
|
base_prompt_tokens=base_prompt_tokens,
|
|
exact_prompt_tokens=base_prompt_tokens,
|
|
messages=[{"role": "tool", "content": "x" * 80_000}],
|
|
turn_tool_prompt_tokens=turn_tool_prompt_tokens,
|
|
)
|
|
turn_tool_prompt_tokens = usage["turn_tool_prompt_tokens"]
|
|
finally:
|
|
streaming._bounded_live_tool_prompt_delta = original_delta
|
|
|
|
assert usage is not None
|
|
assert usage["turn_tool_prompt_tokens"] == 24_000
|
|
assert usage["last_prompt_tokens"] == base_prompt_tokens + 24_000
|
|
assert usage["last_prompt_tokens"] < base_prompt_tokens + (20 * 12_000)
|
|
|
|
|
|
def test_live_usage_session_snapshot_prefers_current_stream_session():
|
|
current = object()
|
|
cache = [None]
|
|
|
|
def fail_loader(_sid):
|
|
raise AssertionError("hot live usage path should not reload the current session")
|
|
|
|
assert _live_usage_session_snapshot("sid", current, cache, loader=fail_loader) is current
|
|
assert cache[0] is current
|
|
|
|
|
|
def test_live_usage_session_snapshot_falls_back_once_then_uses_cache():
|
|
loaded = object()
|
|
calls = []
|
|
cache = [None]
|
|
|
|
def loader(sid):
|
|
calls.append(sid)
|
|
return loaded
|
|
|
|
assert _live_usage_session_snapshot("sid", None, cache, loader=loader) is loaded
|
|
assert _live_usage_session_snapshot("sid", None, cache, loader=loader) is loaded
|
|
assert calls == ["sid"]
|