fix(claude_sdk,conversations,rotation): context size from the last api call, new-day text by rotation reason
This commit is contained in:
@@ -391,7 +391,7 @@ class ClaudeSdkBackend:
|
|||||||
live.session_id = turn.result.session_id
|
live.session_id = turn.result.session_id
|
||||||
if conversation_id is None:
|
if conversation_id is None:
|
||||||
self._rekey(live.key, fingerprint([*history, *turn.synthesized]))
|
self._rekey(live.key, fingerprint([*history, *turn.synthesized]))
|
||||||
usage = _usage_of(turn.result)
|
usage = _usage_of(turn.result, context_tokens=turn.context_tokens)
|
||||||
interrupted = live.interrupt_requested
|
interrupted = live.interrupt_requested
|
||||||
live.interrupt_requested = False
|
live.interrupt_requested = False
|
||||||
if capture is not None:
|
if capture is not None:
|
||||||
@@ -455,6 +455,7 @@ class ClaudeSdkBackend:
|
|||||||
event = message.event
|
event = message.event
|
||||||
if event.get("type") == "message_start":
|
if event.get("type") == "message_start":
|
||||||
offset = next_index
|
offset = next_index
|
||||||
|
turn.context_tokens = _context_of(event.get("message"))
|
||||||
continue
|
continue
|
||||||
index = event.get("index")
|
index = event.get("index")
|
||||||
if isinstance(index, int):
|
if isinstance(index, int):
|
||||||
@@ -740,6 +741,9 @@ class _SessionSpec:
|
|||||||
class _Turn:
|
class _Turn:
|
||||||
events: int = 0
|
events: int = 0
|
||||||
"""Wire events already yielded to the caller."""
|
"""Wire events already yielded to the caller."""
|
||||||
|
|
||||||
|
context_tokens: int = 0
|
||||||
|
"""Input size of the latest API call (``message_start`` usage)."""
|
||||||
synthesized: list[dict[str, Any]] = field(default_factory=list)
|
synthesized: list[dict[str, Any]] = field(default_factory=list)
|
||||||
result: ResultMessage | None = None
|
result: ResultMessage | None = None
|
||||||
stop_reason: StopReason = "end_turn"
|
stop_reason: StopReason = "end_turn"
|
||||||
@@ -895,11 +899,23 @@ def _block_to_dict(block: Any) -> dict[str, Any]:
|
|||||||
raise TypeError(msg)
|
raise TypeError(msg)
|
||||||
|
|
||||||
|
|
||||||
def _usage_of(result: ResultMessage | None) -> TurnUsage:
|
def _context_of(message: Any) -> int:
|
||||||
|
usage = message.get("usage") if isinstance(message, dict) else None
|
||||||
|
if not isinstance(usage, dict):
|
||||||
|
return 0
|
||||||
|
return (
|
||||||
|
_int(usage.get("input_tokens"))
|
||||||
|
+ _int(usage.get("cache_read_input_tokens"))
|
||||||
|
+ _int(usage.get("cache_creation_input_tokens"))
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _usage_of(result: ResultMessage | None, *, context_tokens: int = 0) -> TurnUsage:
|
||||||
if result is None:
|
if result is None:
|
||||||
return TurnUsage()
|
return TurnUsage(context_tokens=context_tokens)
|
||||||
usage = result.usage or {}
|
usage = result.usage or {}
|
||||||
return TurnUsage(
|
return TurnUsage(
|
||||||
|
context_tokens=context_tokens,
|
||||||
input_tokens=_int(usage.get("input_tokens")),
|
input_tokens=_int(usage.get("input_tokens")),
|
||||||
output_tokens=_int(usage.get("output_tokens")),
|
output_tokens=_int(usage.get("output_tokens")),
|
||||||
cache_read_tokens=_int(usage.get("cache_read_input_tokens")),
|
cache_read_tokens=_int(usage.get("cache_read_input_tokens")),
|
||||||
|
|||||||
@@ -428,6 +428,7 @@ async def _build_backends(
|
|||||||
output_tokens=event.usage.output_tokens,
|
output_tokens=event.usage.output_tokens,
|
||||||
cache_read_tokens=event.usage.cache_read_tokens,
|
cache_read_tokens=event.usage.cache_read_tokens,
|
||||||
cache_creation_tokens=event.usage.cache_creation_tokens,
|
cache_creation_tokens=event.usage.cache_creation_tokens,
|
||||||
|
context_tokens=event.usage.context_tokens,
|
||||||
cost_usd=event.usage.cost_usd,
|
cost_usd=event.usage.cost_usd,
|
||||||
duration_ms=event.usage.duration_ms,
|
duration_ms=event.usage.duration_ms,
|
||||||
num_turns=event.usage.num_turns,
|
num_turns=event.usage.num_turns,
|
||||||
|
|||||||
@@ -23,7 +23,7 @@ import logging
|
|||||||
import re
|
import re
|
||||||
import uuid
|
import uuid
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from datetime import UTC, datetime, timedelta
|
from datetime import UTC, date, datetime, timedelta
|
||||||
from typing import TYPE_CHECKING, Any, cast
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
|
|
||||||
from claude_agent_sdk import (
|
from claude_agent_sdk import (
|
||||||
@@ -134,6 +134,19 @@ class SeedContext:
|
|||||||
title: str | None
|
title: str | None
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class NewDayContext:
|
||||||
|
"""What the new master hears first.
|
||||||
|
|
||||||
|
``reason`` is ``ночь`` / ``возраст`` / ``транскрипт`` - only the first
|
||||||
|
one is actually a new day.
|
||||||
|
"""
|
||||||
|
|
||||||
|
day: date
|
||||||
|
reason: str
|
||||||
|
moved: int
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
@dataclass(frozen=True, slots=True)
|
||||||
class ConversationTexts:
|
class ConversationTexts:
|
||||||
"""Texts the gateway cannot invent for a setup.
|
"""Texts the gateway cannot invent for a setup.
|
||||||
@@ -155,7 +168,10 @@ class ConversationTexts:
|
|||||||
"Этот мастер закрывается ({reason}). Напиши хендаут за {day}: справку "
|
"Этот мастер закрывается ({reason}). Напиши хендаут за {day}: справку "
|
||||||
"на утро, не задание - прошедшее время, без повелительного наклонения."
|
"на утро, не задание - прошедшее время, без повелительного наклонения."
|
||||||
)
|
)
|
||||||
new_day: str = "Новый день: мастер сменился, хендаут за {day} записан."
|
new_day: Callable[[NewDayContext], Awaitable[str] | str] | str = (
|
||||||
|
"Мастер сменился ({reason}), хендаут за {day} записан."
|
||||||
|
)
|
||||||
|
"""First inject of the new master; a callable sees the rotation reason."""
|
||||||
distill: Callable[[DistillContext], Awaitable[str] | str] | None = None
|
distill: Callable[[DistillContext], Awaitable[str] | str] | None = None
|
||||||
"""The distiller fork's first message (§8.4): which chat, what day, where
|
"""The distiller fork's first message (§8.4): which chat, what day, where
|
||||||
the digest goes; ``None`` uses a path-less default."""
|
the digest goes; ``None`` uses a path-less default."""
|
||||||
@@ -553,9 +569,7 @@ class Conversations:
|
|||||||
.limit(1)
|
.limit(1)
|
||||||
)
|
)
|
||||||
).first()
|
).first()
|
||||||
if row is None:
|
return context_of(row)
|
||||||
return 0
|
|
||||||
return row.input_tokens + row.cache_read_tokens + row.cache_creation_tokens
|
|
||||||
|
|
||||||
async def usage_tokens(self, since: datetime) -> int:
|
async def usage_tokens(self, since: datetime) -> int:
|
||||||
async with self._db.session() as session:
|
async with self._db.session() as session:
|
||||||
@@ -1213,9 +1227,18 @@ class Conversations:
|
|||||||
_log.exception("%s could not mark %s", fe.name, conv.external_id)
|
_log.exception("%s could not mark %s", fe.name, conv.external_id)
|
||||||
return marked
|
return marked
|
||||||
|
|
||||||
async def new_day(self, conv: Conversation, *, moved: int = 0) -> InjectQueueItem:
|
async def new_day(
|
||||||
day = datetime.now(UTC).astimezone().date().isoformat()
|
self, conv: Conversation, *, reason: str = "ночь", moved: int = 0
|
||||||
text = self._texts.new_day.format(day=day)
|
) -> InjectQueueItem:
|
||||||
|
ctx = NewDayContext(
|
||||||
|
day=datetime.now(UTC).astimezone().date(), reason=reason, moved=moved
|
||||||
|
)
|
||||||
|
source = self._texts.new_day
|
||||||
|
if isinstance(source, str):
|
||||||
|
text = source.format(day=ctx.day.isoformat(), reason=ctx.reason)
|
||||||
|
else:
|
||||||
|
produced: Any = source(ctx)
|
||||||
|
text = await produced if inspect.isawaitable(produced) else produced
|
||||||
if moved:
|
if moved:
|
||||||
text += f" Инжектов переехало из старого мастера: {moved}."
|
text += f" Инжектов переехало из старого мастера: {moved}."
|
||||||
return await self.inject(
|
return await self.inject(
|
||||||
@@ -1902,3 +1925,17 @@ def _usage_dict(capture: TurnCapture) -> dict[str, Any] | None:
|
|||||||
"cost_usd": usage.cost_usd,
|
"cost_usd": usage.cost_usd,
|
||||||
"duration_ms": usage.duration_ms,
|
"duration_ms": usage.duration_ms,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def context_of(row: Usage | None) -> int:
|
||||||
|
"""Context size of a turn.
|
||||||
|
|
||||||
|
The last API call's input, or, for rows written before it was recorded,
|
||||||
|
the per-call average of the turn's input sums.
|
||||||
|
"""
|
||||||
|
if row is None:
|
||||||
|
return 0
|
||||||
|
if row.context_tokens:
|
||||||
|
return row.context_tokens
|
||||||
|
total = row.input_tokens + row.cache_read_tokens + row.cache_creation_tokens
|
||||||
|
return round(total / max(row.num_turns or 1, 1))
|
||||||
|
|||||||
@@ -115,7 +115,7 @@ class Rotation:
|
|||||||
elif not branch.running_turn:
|
elif not branch.running_turn:
|
||||||
await conversations.mark_closed(branch)
|
await conversations.mark_closed(branch)
|
||||||
moved = await conversations.queue.move(old, new, priority="normal")
|
moved = await conversations.queue.move(old, new, priority="normal")
|
||||||
await conversations.new_day(new, moved=moved)
|
await conversations.new_day(new, reason=reason, moved=moved)
|
||||||
conversations.bus.publish(
|
conversations.bus.publish(
|
||||||
"conversation.rotated",
|
"conversation.rotated",
|
||||||
conversation_id=new.external_id,
|
conversation_id=new.external_id,
|
||||||
|
|||||||
@@ -27,6 +27,12 @@ class TurnUsage:
|
|||||||
model_usage: dict[str, Any] | None = None
|
model_usage: dict[str, Any] | None = None
|
||||||
"""``ResultMessage.model_usage`` verbatim: per-model tokens, cost, web searches."""
|
"""``ResultMessage.model_usage`` verbatim: per-model tokens, cost, web searches."""
|
||||||
|
|
||||||
|
context_tokens: int = 0
|
||||||
|
"""Input of the last API call in the turn (fresh + cached + written to
|
||||||
|
cache) - the context size the model actually ran with, what Claude Code
|
||||||
|
shows as the context. The token fields above are sums over every API
|
||||||
|
call of the turn and grow with the number of tool calls."""
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class TurnCapture:
|
class TurnCapture:
|
||||||
|
|||||||
@@ -299,6 +299,8 @@ class Usage(SQLModel, table=True):
|
|||||||
output_tokens: int = 0
|
output_tokens: int = 0
|
||||||
cache_read_tokens: int = 0
|
cache_read_tokens: int = 0
|
||||||
cache_creation_tokens: int = 0
|
cache_creation_tokens: int = 0
|
||||||
|
context_tokens: int | None = Field(default=None)
|
||||||
|
"""Input of the turn's last API call - the real context size."""
|
||||||
cost_usd: float | None = Field(default=None)
|
cost_usd: float | None = Field(default=None)
|
||||||
duration_ms: int | None = Field(default=None)
|
duration_ms: int | None = Field(default=None)
|
||||||
num_turns: int | None = Field(default=None)
|
num_turns: int | None = Field(default=None)
|
||||||
|
|||||||
@@ -98,7 +98,12 @@ class FakeClient:
|
|||||||
self.interrupted = True
|
self.interrupted = True
|
||||||
|
|
||||||
async def receive_response(self):
|
async def receive_response(self):
|
||||||
start = {"type": "message_start", "message": {}}
|
start = {
|
||||||
|
"type": "message_start",
|
||||||
|
"message": {
|
||||||
|
"usage": {"input_tokens": 5, "cache_read_input_tokens": 20_000}
|
||||||
|
},
|
||||||
|
}
|
||||||
yield StreamEvent(uuid="u", session_id="s", event=start)
|
yield StreamEvent(uuid="u", session_id="s", event=start)
|
||||||
for e in _stream(0, "calling "):
|
for e in _stream(0, "calling "):
|
||||||
yield e
|
yield e
|
||||||
@@ -119,7 +124,17 @@ class FakeClient:
|
|||||||
yield UserMessage(
|
yield UserMessage(
|
||||||
content=[ToolResultBlock(tool_use_id="toolu_1", content="42")]
|
content=[ToolResultBlock(tool_use_id="toolu_1", content="42")]
|
||||||
)
|
)
|
||||||
yield StreamEvent(uuid="u", session_id="s", event=start)
|
second = {
|
||||||
|
"type": "message_start",
|
||||||
|
"message": {
|
||||||
|
"usage": {
|
||||||
|
"input_tokens": 7,
|
||||||
|
"cache_read_input_tokens": 20_000,
|
||||||
|
"cache_creation_input_tokens": 3_000,
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
yield StreamEvent(uuid="u", session_id="s", event=second)
|
||||||
for e in _stream(0, "done"):
|
for e in _stream(0, "done"):
|
||||||
yield e
|
yield e
|
||||||
yield AssistantMessage(content=[TextBlock(text="done")], model="m")
|
yield AssistantMessage(content=[TextBlock(text="done")], model="m")
|
||||||
@@ -194,6 +209,7 @@ async def test_stream_envelope_and_index_rebase(cwd: Path) -> None:
|
|||||||
|
|
||||||
assert capture.session_id == "fresh-session"
|
assert capture.session_id == "fresh-session"
|
||||||
assert capture.usage is not None and capture.usage.cost_usd == 0.01
|
assert capture.usage is not None and capture.usage.cost_usd == 0.01
|
||||||
|
assert capture.usage.context_tokens == 23_007
|
||||||
assert capture.synthesized_messages == [
|
assert capture.synthesized_messages == [
|
||||||
{
|
{
|
||||||
"role": "assistant",
|
"role": "assistant",
|
||||||
|
|||||||
+28
-2
@@ -100,7 +100,7 @@ async def test_rotation_order_handout_close_marks_moves_and_new_day(
|
|||||||
return f"напиши хендаут за {ctx.day}"
|
return f"напиши хендаут за {ctx.day}"
|
||||||
|
|
||||||
world.conversations._texts = ConversationTexts( # noqa: SLF001
|
world.conversations._texts = ConversationTexts( # noqa: SLF001
|
||||||
handout=handout, new_day="Новый день {day}."
|
handout=handout, new_day=lambda ctx: f"Новый день {ctx.day} ({ctx.reason})."
|
||||||
)
|
)
|
||||||
old = await world.conversations.spawn(kind="master", agent="a", seed="clean")
|
old = await world.conversations.spawn(kind="master", agent="a", seed="clean")
|
||||||
old = await age(world, old, datetime(2026, 8, 27, 9, 0, tzinfo=UTC))
|
old = await age(world, old, datetime(2026, 8, 27, 9, 0, tzinfo=UTC))
|
||||||
@@ -140,7 +140,7 @@ async def test_rotation_order_handout_close_marks_moves_and_new_day(
|
|||||||
prompt = new_client.prompts[0]
|
prompt = new_client.prompts[0]
|
||||||
assert prompt.startswith("[сид: morning] master")
|
assert prompt.startswith("[сид: morning] master")
|
||||||
assert "[инжект: ротация" in prompt
|
assert "[инжект: ротация" in prompt
|
||||||
assert "Новый день" in prompt
|
assert "Новый день 20" in prompt and "(ночь)" in prompt
|
||||||
assert "переехало из старого мастера: 1" in prompt
|
assert "переехало из старого мастера: 1" in prompt
|
||||||
moved = await world.conversations.queue.pending(new.id)
|
moved = await world.conversations.queue.pending(new.id)
|
||||||
assert [(i.priority, i.text) for i in moved] == [("normal", "later")]
|
assert [(i.priority, i.text) for i in moved] == [("normal", "later")]
|
||||||
@@ -172,3 +172,29 @@ async def test_due_uses_last_usage_row_for_context_size(world: World) -> None:
|
|||||||
rotation = Rotation(world.conversations, RotationPolicy(max_context_tokens=5))
|
rotation = Rotation(world.conversations, RotationPolicy(max_context_tokens=5))
|
||||||
(pair,) = await rotation.due()
|
(pair,) = await rotation.due()
|
||||||
assert pair[0].id == conv.id and pair[1] == "транскрипт"
|
assert pair[0].id == conv.id and pair[1] == "транскрипт"
|
||||||
|
|
||||||
|
|
||||||
|
def test_context_of_prefers_last_call_and_averages_old_rows() -> None:
|
||||||
|
from beaver_gateway.core.conversations import context_of
|
||||||
|
from beaver_gateway.storage.models import Usage
|
||||||
|
|
||||||
|
fresh = Usage(
|
||||||
|
agent_name="a",
|
||||||
|
model="m",
|
||||||
|
input_tokens=12,
|
||||||
|
cache_read_tokens=142_984,
|
||||||
|
cache_creation_tokens=29_584,
|
||||||
|
num_turns=6,
|
||||||
|
context_tokens=29_000,
|
||||||
|
)
|
||||||
|
assert context_of(fresh) == 29_000
|
||||||
|
old = Usage(
|
||||||
|
agent_name="a",
|
||||||
|
model="m",
|
||||||
|
input_tokens=12,
|
||||||
|
cache_read_tokens=142_984,
|
||||||
|
cache_creation_tokens=29_584,
|
||||||
|
num_turns=6,
|
||||||
|
)
|
||||||
|
assert context_of(old) == 28_763
|
||||||
|
assert context_of(None) == 0
|
||||||
|
|||||||
Reference in New Issue
Block a user