core/agent: calibrate soft-compact with last-request prompt_tokens

After the first model call, scale char-based token estimates by
API/resolved prompt size so budget checks track near-real tokens.
/compact uses the same calibration when available.
This commit is contained in:
2026-07-15 11:11:15 +08:00
parent c9e805f38d
commit b1223c8189
6 changed files with 128 additions and 11 deletions
+49 -8
View File
@@ -11,12 +11,14 @@ from plyngent.lmproto.openai_compatible.model import (
)
if TYPE_CHECKING:
from collections.abc import Sequence
from collections.abc import Callable, Sequence
from plyngent.lmproto.openai_compatible.model import AnyChatMessage
type TokenMeasure = Callable[[Sequence[AnyChatMessage]], int]
DEFAULT_TOOL_RESULT_MAX_CHARS = 32_000
# Soft context budget in estimated tokens (~4 chars/token); not a hard model limit.
# Soft context budget in tokens (API-calibrated when possible; else ~4 chars/token).
DEFAULT_CONTEXT_MAX_TOKENS = 200_000
DEFAULT_OLD_TOOL_RESULT_CHARS = 800
DEFAULT_RECENT_TOOL_RESULTS = 4
@@ -64,12 +66,37 @@ def estimate_messages_chars(messages: Sequence[AnyChatMessage]) -> int:
def estimate_messages_tokens(messages: Sequence[AnyChatMessage]) -> int:
"""Char-based token estimate for soft context budget checks."""
"""Char-based token estimate (fallback when no API calibration is available)."""
from plyngent.agent.usage import chars_to_tokens
return chars_to_tokens(estimate_messages_chars(messages))
def measure_messages_tokens(
messages: Sequence[AnyChatMessage],
*,
prompt_tokens_hint: int | None = None,
sent_estimate_tokens: int | None = None,
) -> int:
"""Token size for budget checks.
When ``prompt_tokens_hint`` is the last request's API (or resolved) prompt
size and ``sent_estimate_tokens`` is the char-estimate of that same payload,
scale the current char-estimate by ``hint / sent_estimate`` so soft-compact
tracks near-real tokens after the first model call.
"""
est = estimate_messages_tokens(messages)
if (
prompt_tokens_hint is not None
and prompt_tokens_hint > 0
and sent_estimate_tokens is not None
and sent_estimate_tokens > 0
):
scaled = round(est * (prompt_tokens_hint / sent_estimate_tokens))
return max(1, scaled) if est > 0 else 0
return est
def _shrink_tool(message: ToolChatMessage, max_chars: int) -> ToolChatMessage:
if len(message.content) <= max_chars:
return message
@@ -109,13 +136,14 @@ def _shrink_largest(
*,
max_tokens: int,
shrink_cap: int,
measure: TokenMeasure,
) -> None:
def tool_len(i: int) -> int:
msg = messages[i]
return len(msg.content) if isinstance(msg, ToolChatMessage) else 0
for idx in sorted(tool_indices, key=tool_len, reverse=True):
if estimate_messages_tokens(messages) <= max_tokens:
if measure(messages) <= max_tokens:
return
tool_msg = messages[idx]
if isinstance(tool_msg, ToolChatMessage):
@@ -128,17 +156,29 @@ def compact_messages_for_request(
max_tokens: int = DEFAULT_CONTEXT_MAX_TOKENS,
old_tool_result_chars: int = DEFAULT_OLD_TOOL_RESULT_CHARS,
keep_recent_tool_results: int = DEFAULT_RECENT_TOOL_RESULTS,
# Deprecated alias: treated as token budget if max_tokens not overridden via callers.
prompt_tokens_hint: int | None = None,
sent_estimate_tokens: int | None = None,
# Deprecated alias: treated as token budget.
max_chars: int | None = None,
) -> list[AnyChatMessage]:
"""Return a request-time copy with older tool dumps shrunk if over budget.
Budget is in **estimated tokens** (char/4). Does not mutate the original history.
Budget is in tokens. Prefer API-calibrated measurement via
``prompt_tokens_hint`` / ``sent_estimate_tokens`` (last request); otherwise
fall back to char/4. Does not mutate the original history.
``max_tokens < 1`` disables compacting.
"""
budget = max_tokens if max_chars is None else max_chars
def measure(msgs: Sequence[AnyChatMessage]) -> int:
return measure_messages_tokens(
msgs,
prompt_tokens_hint=prompt_tokens_hint,
sent_estimate_tokens=sent_estimate_tokens,
)
out: list[AnyChatMessage] = list(messages)
if budget < 1 or estimate_messages_tokens(out) <= budget:
if budget < 1 or measure(out) <= budget:
return out
indices = _tool_indices(out)
@@ -147,7 +187,7 @@ def compact_messages_for_request(
protect = _protect_indices(indices, keep_recent_tool_results)
_shrink_except(out, indices, protect, old_tool_result_chars)
if estimate_messages_tokens(out) <= budget:
if measure(out) <= budget:
return out
_shrink_largest(
@@ -155,5 +195,6 @@ def compact_messages_for_request(
indices,
max_tokens=budget,
shrink_cap=max(64, old_tool_result_chars // 2),
measure=measure,
)
return out
+16 -2
View File
@@ -68,9 +68,16 @@ def soft_compact_transcript(
messages: Sequence[AnyChatMessage],
*,
max_tokens: int = DEFAULT_CONTEXT_MAX_TOKENS,
prompt_tokens_hint: int | None = None,
sent_estimate_tokens: int | None = None,
) -> str:
"""Soft-compact tool dumps then format as transcript text."""
compacted = compact_messages_for_request(messages, max_tokens=max_tokens)
compacted = compact_messages_for_request(
messages,
max_tokens=max_tokens,
prompt_tokens_hint=prompt_tokens_hint,
sent_estimate_tokens=sent_estimate_tokens,
)
return format_transcript(compacted)
@@ -81,12 +88,19 @@ async def summarize_messages(
model: str,
max_context_tokens: int = DEFAULT_CONTEXT_MAX_TOKENS,
temperature: float | None = 0.2,
prompt_tokens_hint: int | None = None,
sent_estimate_tokens: int | None = None,
) -> str:
"""Soft-compact history and ask the model for a dense summary (no tools)."""
if not messages:
msg = "nothing to compact"
raise ValueError(msg)
transcript = soft_compact_transcript(messages, max_tokens=max_context_tokens)
transcript = soft_compact_transcript(
messages,
max_tokens=max_context_tokens,
prompt_tokens_hint=prompt_tokens_hint,
sent_estimate_tokens=sent_estimate_tokens,
)
if not transcript.strip():
msg = "nothing to compact"
raise ValueError(msg)
+11
View File
@@ -23,6 +23,7 @@ from .budget import (
DEFAULT_CONTEXT_MAX_TOKENS,
DEFAULT_TOOL_RESULT_MAX_CHARS,
compact_messages_for_request,
estimate_messages_tokens,
truncate_tool_result,
)
from .events import (
@@ -240,6 +241,9 @@ async def run_chat_loop(
rounds_used = 0
allowance = max_rounds
# Calibrate soft-compact from last model call's prompt_tokens (API preferred).
prompt_tokens_hint: int | None = None
sent_estimate_tokens: int | None = None
while True:
while rounds_used < allowance:
@@ -247,7 +251,10 @@ async def run_chat_loop(
request_messages = compact_messages_for_request(
messages,
max_tokens=max_context_tokens,
prompt_tokens_hint=prompt_tokens_hint,
sent_estimate_tokens=sent_estimate_tokens,
)
sent_est = estimate_messages_tokens(request_messages)
param = ChatCompletionsParam(
messages=request_messages,
model=model,
@@ -257,6 +264,10 @@ async def run_chat_loop(
pre_len = len(messages)
async for event in _assistant_round(client, param, messages, stream=stream):
if isinstance(event, UsageEvent):
# Next rounds scale char-estimates by real/resolved prompt size.
prompt_tokens_hint = event.usage.prompt_tokens
sent_estimate_tokens = sent_est
yield event
assistant = _last_assistant(messages, pre_len)
tool_calls = assistant.tool_calls
+11
View File
@@ -170,11 +170,22 @@ class ReplState:
msg = "nothing to compact (empty history)"
raise ValueError(msg)
# Prefer last API prompt_tokens to drive soft-compact toward real size.
hint: int | None = None
sent_est: int | None = None
if not self.agent.last_request_usage.is_zero():
from plyngent.agent.budget import estimate_messages_tokens
hint = self.agent.last_request_usage.prompt_tokens
# Approximate: calibrate against current full history char-est.
sent_est = estimate_messages_tokens(messages)
summary = await summarize_messages(
self.client,
messages,
model=self.model,
max_context_tokens=self.agent.max_context_tokens,
prompt_tokens_hint=hint,
sent_estimate_tokens=sent_est,
)
session_name = name or f"compact-from-{old_id}"
await self.new_session(name=session_name)