diff --git a/tests/e2e/perf/session_anomaly.py b/tests/e2e/perf/session_anomaly.py
index e105b076943..b68482e965f 100644
--- a/tests/e2e/perf/session_anomaly.py
+++ b/tests/e2e/perf/session_anomaly.py
@@ -70,24 +70,37 @@ class AnomalyReport:
@property
def warm_cache_read_share(self) -> float:
- billed = self.warm_uncached_input_tokens + self.warm_cache_read_tokens + self.warm_cache_creation_tokens
+ billed = (
+ self.warm_uncached_input_tokens
+ + self.warm_cache_read_tokens
+ + self.warm_cache_creation_tokens
+ )
return self.warm_cache_read_tokens / billed if billed else 0.0
def _system_prefix_block(marker: str) -> TextBlock:
- text = " ".join(f"Project context paragraph {index} for session {marker}." for index in range(300))
+ text = " ".join(
+ f"Project context paragraph {index} for session {marker}." for index in range(300)
+ )
return TextBlock(text=text, cache_control=CacheControl())
def _user_turn_text(marker: str, turn_index: int) -> str:
- notes = " ".join(f"Working note {index} of turn {turn_index} in session {marker}." for index in range(80))
+ notes = " ".join(
+ f"Working note {index} of turn {turn_index} in session {marker}."
+ for index in range(80)
+ )
return f"Reply with one short sentence.\n{notes}"
def _reminder_turn() -> RichMessage:
return RichMessage(
role="system",
- content=[TextBlock(text="Keep the answer to one short sentence.")],
+ content=[
+ TextBlock(
+ text="Keep the answer to one short sentence."
+ )
+ ],
)
@@ -114,7 +127,9 @@ def retried(
return retried(call, attempts - 1, backoff_seconds, sleep)
-def _metric(result: Result[SessionMessagesResponse], turn_index: int, latency_seconds: float) -> TurnMetric:
+def _metric(
+ result: Result[SessionMessagesResponse], turn_index: int, latency_seconds: float
+) -> TurnMetric:
if isinstance(result, Success):
usage = result.data.usage
return TurnMetric(
@@ -152,7 +167,11 @@ def _drive_turns(
return ()
user_turn = RichMessage(
role="user",
- content=[TextBlock(text=_user_turn_text(marker, turn_index), cache_control=CacheControl())],
+ content=[
+ TextBlock(
+ text=_user_turn_text(marker, turn_index), cache_control=CacheControl()
+ )
+ ],
)
started = time.monotonic()
result = retried(
@@ -171,7 +190,9 @@ def _drive_turns(
turn = _metric(result, turn_index, time.monotonic() - started)
if not isinstance(result, Success):
return (turn,)
- assistant_turn = RichMessage(role="assistant", content=[TextBlock(text=result.data.text or "Understood.")])
+ assistant_turn = RichMessage(
+ role="assistant", content=[TextBlock(text=result.data.text or "Understood.")]
+ )
return (
turn,
*_drive_turns(
@@ -221,7 +242,9 @@ def run_concurrent_sessions(
) -> tuple[TurnMetric, ...]:
with ThreadPoolExecutor(max_workers=sessions) as pool:
futures = [
- pool.submit(run_session, transport, key, model, turns_per_session, attempts_per_turn)
+ pool.submit(
+ run_session, transport, key, model, turns_per_session, attempts_per_turn
+ )
for _ in range(sessions)
]
return tuple(turn for future in futures for turn in future.result())
@@ -273,5 +296,7 @@ def summarize(turns: tuple[TurnMetric, ...], planned_turns: int) -> AnomalyRepor
warm_uncached_input_tokens=sum(turn.uncached_input_tokens for turn in warm),
warm_cache_read_tokens=sum(turn.cache_read_tokens for turn in warm),
warm_cache_creation_tokens=sum(turn.cache_creation_tokens for turn in warm),
- p95_turn_seconds=_p95(tuple(turn.latency_seconds for turn in turns if turn.ok)),
+ p95_turn_seconds=_p95(
+ tuple(turn.latency_seconds for turn in turns if turn.ok)
+ ),
)