Log every model call's time and tokens #175
No files matched your search
@@ -31,6 +31,7 @@ dangerous direction.
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
import json
|
import json
|
||||||
|
import time
|
||||||
from collections.abc import Callable, Sequence
|
from collections.abc import Callable, Sequence
|
||||||
from dataclasses import dataclass, replace
|
from dataclasses import dataclass, replace
|
||||||
from typing import Any
|
from typing import Any
|
||||||
@@ -297,6 +298,31 @@ class BudgetedClient:
|
|||||||
def last_rate_limit_wait_s(self) -> float:
|
def last_rate_limit_wait_s(self) -> float:
|
||||||
return float(getattr(self._inner, "last_rate_limit_wait_s", 0.0) or 0.0)
|
return float(getattr(self._inner, "last_rate_limit_wait_s", 0.0) or 0.0)
|
||||||
|
|
||||||
|
def _log_call(self, model: str, started: float, error: str | None = None) -> None:
|
||||||
|
"""One event per model call: how long it took and what it used, so a
|
||||||
|
slow turn shows whether reading or writing took the time."""
|
||||||
|
if self._oplog is None:
|
||||||
|
return
|
||||||
|
fields: dict[str, Any] = {
|
||||||
|
"model": model,
|
||||||
|
"seconds": round(time.monotonic() - started, 1),
|
||||||
|
}
|
||||||
|
if error is not None:
|
||||||
|
self._oplog.emit("llm.call", cause="failed", error=error, **fields)
|
||||||
|
return
|
||||||
|
usage = self.last_usage
|
||||||
|
answered = self.last_model
|
||||||
|
if answered and answered != model:
|
||||||
|
fields["answered_by"] = answered
|
||||||
|
self._oplog.emit(
|
||||||
|
"llm.call",
|
||||||
|
cause="answered",
|
||||||
|
input_tokens=int(usage.get("input", 0)),
|
||||||
|
output_tokens=int(usage.get("output", 0)),
|
||||||
|
cached_tokens=int(usage.get("cached", 0)),
|
||||||
|
**fields,
|
||||||
|
)
|
||||||
|
|
||||||
def budget_for(self, model: str) -> Budget:
|
def budget_for(self, model: str) -> Budget:
|
||||||
window, model_max_output = self._limits_for(model)
|
window, model_max_output = self._limits_for(model)
|
||||||
ceiling = self._max_output_tokens
|
ceiling = self._max_output_tokens
|
||||||
@@ -360,17 +386,23 @@ class BudgetedClient:
|
|||||||
)
|
)
|
||||||
if require_tool_call is None:
|
if require_tool_call is None:
|
||||||
require_tool_call = bool(tools) and self._enforces_tool_choice(model)
|
require_tool_call = bool(tools) and self._enforces_tool_choice(model)
|
||||||
produced = self._inner.complete(
|
started = time.monotonic()
|
||||||
fitted,
|
try:
|
||||||
tools,
|
produced = self._inner.complete(
|
||||||
model,
|
fitted,
|
||||||
system=system,
|
tools,
|
||||||
max_tokens=budget.max_output,
|
model,
|
||||||
# Every caller that offers tools needs one called: the executor
|
system=system,
|
||||||
# finishes with finish, the planner with submit_plan or
|
max_tokens=budget.max_output,
|
||||||
# ask_human, a delegation with hand_back.
|
# Every caller that offers tools needs one called: the executor
|
||||||
require_tool_call=require_tool_call,
|
# finishes with finish, the planner with submit_plan or
|
||||||
)
|
# ask_human, a delegation with hand_back.
|
||||||
|
require_tool_call=require_tool_call,
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
self._log_call(model, started, error=type(e).__name__)
|
||||||
|
raise
|
||||||
|
self._log_call(model, started)
|
||||||
reported = int(self.last_usage.get("input", 0))
|
reported = int(self.last_usage.get("input", 0))
|
||||||
if reported > 0:
|
if reported > 0:
|
||||||
estimated = estimate_tokens(fitted, system, tools)
|
estimated = estimate_tokens(fitted, system, tools)
|
||||||
|
|||||||
@@ -244,6 +244,10 @@ class LiteLLMClient:
|
|||||||
self.last_usage = {
|
self.last_usage = {
|
||||||
"input": int(usage.get("prompt_tokens") or 0),
|
"input": int(usage.get("prompt_tokens") or 0),
|
||||||
"output": int(usage.get("completion_tokens") or 0),
|
"output": int(usage.get("completion_tokens") or 0),
|
||||||
|
# Prompt tokens the server reused from its cache, when it says.
|
||||||
|
"cached": int(
|
||||||
|
(usage.get("prompt_tokens_details") or {}).get("cached_tokens") or 0
|
||||||
|
),
|
||||||
}
|
}
|
||||||
out: list[CanonicalEvent] = []
|
out: list[CanonicalEvent] = []
|
||||||
if message.get("content"):
|
if message.get("content"):
|
||||||
@@ -371,6 +375,7 @@ class AnthropicPassthroughClient:
|
|||||||
self.last_usage = {
|
self.last_usage = {
|
||||||
"input": int(usage.get("input_tokens") or 0),
|
"input": int(usage.get("input_tokens") or 0),
|
||||||
"output": int(usage.get("output_tokens") or 0),
|
"output": int(usage.get("output_tokens") or 0),
|
||||||
|
"cached": int(usage.get("cache_read_input_tokens") or 0),
|
||||||
}
|
}
|
||||||
out: list[CanonicalEvent] = []
|
out: list[CanonicalEvent] = []
|
||||||
for block in content:
|
for block in content:
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
"""The context budget (architecture L6)."""
|
"""The context budget (architecture L6)."""
|
||||||
|
|
||||||
import io
|
import io
|
||||||
|
import json
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
@@ -25,6 +26,7 @@ from nisse_core.models.context import (
|
|||||||
fit,
|
fit,
|
||||||
head_and_tail,
|
head_and_tail,
|
||||||
)
|
)
|
||||||
|
from nisse_core.models.llm import LLMError
|
||||||
from nisse_core.oplog import OpLog
|
from nisse_core.oplog import OpLog
|
||||||
from nisse_core.state.store import StateStore
|
from nisse_core.state.store import StateStore
|
||||||
|
|
||||||
@@ -190,7 +192,7 @@ def test_reported_usage_only_ever_raises_the_estimate() -> None:
|
|||||||
assert under_then_over._factors["m"] == 2.0 # never loosened
|
assert under_then_over._factors["m"] == 2.0 # never loosened
|
||||||
|
|
||||||
|
|
||||||
def test_a_request_that_fits_is_sent_whole_and_nothing_is_logged() -> None:
|
def test_a_request_that_fits_is_sent_whole_and_no_cut_is_logged() -> None:
|
||||||
stream = io.StringIO()
|
stream = io.StringIO()
|
||||||
client = BudgetedClient(
|
client = BudgetedClient(
|
||||||
RecordingClient(), lambda m: (100_000, None), oplog=OpLog("t", stream)
|
RecordingClient(), lambda m: (100_000, None), oplog=OpLog("t", stream)
|
||||||
@@ -198,7 +200,49 @@ def test_a_request_that_fits_is_sent_whole_and_nothing_is_logged() -> None:
|
|||||||
events = transcript(3, 100)
|
events = transcript(3, 100)
|
||||||
client.complete(events, None, "m")
|
client.complete(events, None, "m")
|
||||||
assert client.last_trim is None
|
assert client.last_trim is None
|
||||||
assert stream.getvalue() == ""
|
assert "context.trimmed" not in stream.getvalue()
|
||||||
|
|
||||||
|
|
||||||
|
def logged(stream: io.StringIO, event: str) -> list[dict[str, Any]]:
|
||||||
|
records = [json.loads(line) for line in stream.getvalue().splitlines()]
|
||||||
|
return [r for r in records if r["event"] == event]
|
||||||
|
|
||||||
|
|
||||||
|
def test_every_call_logs_its_time_and_tokens() -> None:
|
||||||
|
# A slow turn should show whether reading or writing took the time (#174).
|
||||||
|
stream = io.StringIO()
|
||||||
|
inner = RecordingClient(reported_input=900)
|
||||||
|
inner.last_usage["cached"] = 800
|
||||||
|
inner.last_model = "fallback-model"
|
||||||
|
client = BudgetedClient(inner, lambda m: (100_000, None), oplog=OpLog("t", stream))
|
||||||
|
client.complete(transcript(1, 10), None, "m")
|
||||||
|
[call] = logged(stream, "llm.call")
|
||||||
|
assert call["cause"] == "answered"
|
||||||
|
assert call["model"] == "m"
|
||||||
|
assert call["answered_by"] == "fallback-model"
|
||||||
|
assert (call["input_tokens"], call["output_tokens"], call["cached_tokens"]) == (
|
||||||
|
900,
|
||||||
|
5,
|
||||||
|
800,
|
||||||
|
)
|
||||||
|
assert isinstance(call["seconds"], float)
|
||||||
|
|
||||||
|
|
||||||
|
def test_a_failed_call_logs_how_long_it_ran_and_still_raises() -> None:
|
||||||
|
class TimesOut(RecordingClient):
|
||||||
|
def complete(self, *args, **kwargs):
|
||||||
|
raise LLMError("LLM request returned 408: timed out")
|
||||||
|
|
||||||
|
stream = io.StringIO()
|
||||||
|
client = BudgetedClient(
|
||||||
|
TimesOut(), lambda m: (100_000, None), oplog=OpLog("t", stream)
|
||||||
|
)
|
||||||
|
with pytest.raises(LLMError):
|
||||||
|
client.complete(transcript(1, 10), None, "m")
|
||||||
|
[call] = logged(stream, "llm.call")
|
||||||
|
assert call["cause"] == "failed"
|
||||||
|
assert call["error"] == "LLMError"
|
||||||
|
assert "input_tokens" not in call
|
||||||
|
|
||||||
|
|
||||||
def test_a_cut_is_recorded_and_logged() -> None:
|
def test_a_cut_is_recorded_and_logged() -> None:
|
||||||
|
|||||||
@@ -484,6 +484,48 @@ def test_the_answering_model_is_recorded() -> None:
|
|||||||
assert client.last_model == "gpt-test-1"
|
assert client.last_model == "gpt-test-1"
|
||||||
|
|
||||||
|
|
||||||
|
def test_usage_includes_what_the_server_reused_from_its_cache() -> None:
|
||||||
|
client = make_client(
|
||||||
|
lambda request: httpx.Response(
|
||||||
|
200,
|
||||||
|
json={
|
||||||
|
"model": "qwen",
|
||||||
|
"choices": [{"message": {"role": "assistant", "content": "hi"}}],
|
||||||
|
"usage": {
|
||||||
|
"prompt_tokens": 900,
|
||||||
|
"completion_tokens": 40,
|
||||||
|
"prompt_tokens_details": {"cached_tokens": 800},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
client.complete(EVENTS, tools=None, model="qwen")
|
||||||
|
assert client.last_usage == {"input": 900, "output": 40, "cached": 800}
|
||||||
|
|
||||||
|
passthrough = AnthropicPassthroughClient(
|
||||||
|
"https://litellm.test",
|
||||||
|
"tok",
|
||||||
|
transport=httpx.MockTransport(
|
||||||
|
lambda request: httpx.Response(
|
||||||
|
200,
|
||||||
|
json={
|
||||||
|
"model": "claude-opus-5-5",
|
||||||
|
"stop_reason": "end_turn",
|
||||||
|
"content": [{"type": "text", "text": "hi"}],
|
||||||
|
"usage": {
|
||||||
|
"input_tokens": 50,
|
||||||
|
"output_tokens": 7,
|
||||||
|
"cache_read_input_tokens": 30,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
)
|
||||||
|
),
|
||||||
|
sleep=lambda _s: None,
|
||||||
|
)
|
||||||
|
passthrough.complete(EVENTS, tools=None, model="claude-opus-5-5")
|
||||||
|
assert passthrough.last_usage == {"input": 50, "output": 7, "cached": 30}
|
||||||
|
|
||||||
|
|
||||||
# --- outages delay, never fail (L5) --------------------------------------------
|
# --- outages delay, never fail (L5) --------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user