somew1nd
5d92939d91
Preserve user guidance through history compression, stabilize Flash index targets, and retire withdrawn checks without applying stale verdicts. Improve runtime diagnostics, tool responses, and history rendering. Simplify documentation and comments. Validation: 2076 tests passed in the full suite; 381 affected tests and 26 frontend tests passed after cleanup. Ruff and the quality ratchet passed.
112 lines
4.2 KiB
Python
112 lines
4.2 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Session and per-source cache-hit ratios.
|
|
|
|
A session whose operator calls never hit the provider's prompt cache must be
|
|
visible without reading every ``llm_usage`` trace: the session snapshot and
|
|
each source's totals carry ``cached_ratio`` (cached / prompt tokens).
|
|
"""
|
|
|
|
import logging
|
|
from unittest.mock import Mock
|
|
from uuid import uuid4
|
|
|
|
from langchain_core.messages import AIMessage
|
|
|
|
from artemis.services.token_meter import (
|
|
SessionTokenMeter,
|
|
cached_ratio,
|
|
format_session_summary,
|
|
get_meter,
|
|
log_session_summary,
|
|
record_llm_usage,
|
|
)
|
|
|
|
|
|
def _msg(prompt: int, cached: int = 0) -> AIMessage:
|
|
usage = {"input_tokens": prompt, "output_tokens": 10, "total_tokens": prompt + 10}
|
|
if cached:
|
|
usage["input_token_details"] = {"cache_read": cached}
|
|
return AIMessage(content="ok", usage_metadata=usage)
|
|
|
|
|
|
def test_cached_ratio_is_zero_without_prompt_tokens():
|
|
assert cached_ratio(0, 0) == 0.0
|
|
assert cached_ratio(50, 0) == 0.0
|
|
assert cached_ratio(250, 400) == 0.625
|
|
|
|
|
|
def test_meter_tracks_cached_ratio_per_session_and_per_source():
|
|
meter = SessionTokenMeter("s")
|
|
meter.record({"prompt_tokens": 1000, "cached_tokens": 0}, source="google:operator")
|
|
meter.record({"prompt_tokens": 1000, "cached_tokens": 0}, source="google:operator")
|
|
snapshot = meter.record(
|
|
{"prompt_tokens": 200, "cached_tokens": 150},
|
|
source="lens:step_capsule:gemini",
|
|
update_last_prompt=False,
|
|
)
|
|
|
|
# Session ratio blends every call; the call's own source rides along.
|
|
assert snapshot["session_cached_ratio"] == round(150 / 2200, 3)
|
|
assert snapshot["source_cached_ratio"] == 0.75
|
|
assert meter.last_prompt_tokens == 1000
|
|
|
|
summary = meter.summary()
|
|
assert summary["session_id"] == "s"
|
|
assert summary["sources"]["google:operator"] == {
|
|
"llm_calls": 2,
|
|
"prompt_tokens": 2000,
|
|
"cached_tokens": 0,
|
|
"cache_hit_calls": 0,
|
|
"cached_ratio": 0.0, # the operator never hit the cache: visible at a glance
|
|
}
|
|
assert summary["sources"]["lens:step_capsule:gemini"]["cached_ratio"] == 0.75
|
|
|
|
|
|
def test_record_llm_usage_payload_carries_both_ratios():
|
|
engine = Mock()
|
|
engine.current_session_id = uuid4()
|
|
engine.current_step_id = None
|
|
|
|
record_llm_usage(engine, _msg(1000), source="google:op")
|
|
payload = record_llm_usage(engine, _msg(1000, cached=800), source="google:op")
|
|
|
|
assert payload["session_cached_ratio"] == 0.4
|
|
assert payload["source_cached_ratio"] == 0.4
|
|
# A source-less call still reports the session ratio, never a stale source one.
|
|
payload = record_llm_usage(engine, _msg(1000))
|
|
assert payload["session_cached_ratio"] == round(800 / 3000, 3)
|
|
assert "source_cached_ratio" not in payload
|
|
|
|
|
|
def test_session_summary_logs_once_per_metered_session(caplog):
|
|
session_id = f"ratio-{uuid4()}"
|
|
assert log_session_summary(session_id) is None # never metered: nothing to say
|
|
|
|
meter = get_meter(session_id)
|
|
meter.record({"prompt_tokens": 500, "cached_tokens": 0}, source="google:op")
|
|
meter.record({"prompt_tokens": 100, "cached_tokens": 90}, source="lens:visual")
|
|
|
|
line = format_session_summary(session_id)
|
|
assert f"LLM usage for session {session_id}: 2 calls, 600 prompt tokens, 90 cached" in line
|
|
assert "cached_ratio=0.150" in line
|
|
assert "google:op: 1 calls, cached_ratio=0.000" in line
|
|
assert "lens:visual: 1 calls, cached_ratio=0.900" in line
|
|
|
|
with caplog.at_level(logging.INFO, logger="artemis.services.token_meter"):
|
|
summary = log_session_summary(session_id)
|
|
assert summary is not None and summary["session_cached_ratio"] == 0.15
|
|
assert any(line in record.getMessage() for record in caplog.records)
|