mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
core/agent+config: soft context budget in estimated tokens
Replace max_context_chars with max_context_tokens (default 200k est. tokens via char/4); compact and /status use the token budget.
This commit is contained in:
@@ -59,7 +59,7 @@ Async SQLAlchemy + aiosqlite. `MemoryStore`: schema init (+ lightweight SQLite `
|
||||
- **`/compact`**: soft-compact tool dumps → model summary (no tools) → **new** session seeded with summary message.
|
||||
- Events: text_delta, assistant_message, tool_call/result, max_rounds, **error** (`retryable`/`source`), **cancelled** (`reason`), **usage** (`TokenUsage`).
|
||||
- Usage: API `usage` from completions (stream with `include_usage`); **char≈token fallback** (~4 chars/token) when omitted; `ChatAgent.session_usage` / `last_turn_usage`; CLI end-of-turn + `/status` (marks `(est)` / `(api+est)`).
|
||||
- Config ``[agent]``: `system_prompt`, `max_tool_result_chars`, `parallel_tools`, `confirm_destructive`, `path_denylist`, `max_context_chars`.
|
||||
- Config ``[agent]``: `system_prompt`, `max_tool_result_chars`, `parallel_tools`, `confirm_destructive`, `path_denylist`, `max_context_tokens` (default 200k est. tokens).
|
||||
|
||||
### Tools (`tools/`)
|
||||
|
||||
|
||||
@@ -16,10 +16,14 @@ if TYPE_CHECKING:
|
||||
from plyngent.lmproto.openai_compatible.model import AnyChatMessage
|
||||
|
||||
DEFAULT_TOOL_RESULT_MAX_CHARS = 32_000
|
||||
DEFAULT_CONTEXT_MAX_CHARS = 200_000
|
||||
# Soft context budget in estimated tokens (~4 chars/token); not a hard model limit.
|
||||
DEFAULT_CONTEXT_MAX_TOKENS = 200_000
|
||||
DEFAULT_OLD_TOOL_RESULT_CHARS = 800
|
||||
DEFAULT_RECENT_TOOL_RESULTS = 4
|
||||
|
||||
# Backward-compat alias (older code/docs may still import this name).
|
||||
DEFAULT_CONTEXT_MAX_CHARS = DEFAULT_CONTEXT_MAX_TOKENS * 4
|
||||
|
||||
|
||||
def truncate_tool_result(text: str, max_chars: int = DEFAULT_TOOL_RESULT_MAX_CHARS) -> str:
|
||||
"""Cap tool output so huge dumps do not flood model context."""
|
||||
@@ -59,6 +63,13 @@ def estimate_messages_chars(messages: Sequence[AnyChatMessage]) -> int:
|
||||
return sum(estimate_message_chars(m) for m in messages)
|
||||
|
||||
|
||||
def estimate_messages_tokens(messages: Sequence[AnyChatMessage]) -> int:
|
||||
"""Char-based token estimate for soft context budget checks."""
|
||||
from plyngent.agent.usage import chars_to_tokens
|
||||
|
||||
return chars_to_tokens(estimate_messages_chars(messages))
|
||||
|
||||
|
||||
def _shrink_tool(message: ToolChatMessage, max_chars: int) -> ToolChatMessage:
|
||||
if len(message.content) <= max_chars:
|
||||
return message
|
||||
@@ -96,7 +107,7 @@ def _shrink_largest(
|
||||
messages: list[AnyChatMessage],
|
||||
tool_indices: Sequence[int],
|
||||
*,
|
||||
max_chars: int,
|
||||
max_tokens: int,
|
||||
shrink_cap: int,
|
||||
) -> None:
|
||||
def tool_len(i: int) -> int:
|
||||
@@ -104,7 +115,7 @@ def _shrink_largest(
|
||||
return len(msg.content) if isinstance(msg, ToolChatMessage) else 0
|
||||
|
||||
for idx in sorted(tool_indices, key=tool_len, reverse=True):
|
||||
if estimate_messages_chars(messages) <= max_chars:
|
||||
if estimate_messages_tokens(messages) <= max_tokens:
|
||||
return
|
||||
tool_msg = messages[idx]
|
||||
if isinstance(tool_msg, ToolChatMessage):
|
||||
@@ -114,17 +125,20 @@ def _shrink_largest(
|
||||
def compact_messages_for_request(
|
||||
messages: Sequence[AnyChatMessage],
|
||||
*,
|
||||
max_chars: int = DEFAULT_CONTEXT_MAX_CHARS,
|
||||
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.
|
||||
max_chars: int | None = None,
|
||||
) -> list[AnyChatMessage]:
|
||||
"""Return a request-time copy with older tool dumps shrunk if over budget.
|
||||
|
||||
Does not mutate the original history (full results stay for persistence/UI).
|
||||
``max_chars < 1`` disables compacting.
|
||||
Budget is in **estimated tokens** (char/4). Does not mutate the original history.
|
||||
``max_tokens < 1`` disables compacting.
|
||||
"""
|
||||
budget = max_tokens if max_chars is None else max_chars
|
||||
out: list[AnyChatMessage] = list(messages)
|
||||
if max_chars < 1 or estimate_messages_chars(out) <= max_chars:
|
||||
if budget < 1 or estimate_messages_tokens(out) <= budget:
|
||||
return out
|
||||
|
||||
indices = _tool_indices(out)
|
||||
@@ -133,13 +147,13 @@ 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_chars(out) <= max_chars:
|
||||
if estimate_messages_tokens(out) <= budget:
|
||||
return out
|
||||
|
||||
_shrink_largest(
|
||||
out,
|
||||
indices,
|
||||
max_chars=max_chars,
|
||||
max_tokens=budget,
|
||||
shrink_cap=max(64, old_tool_result_chars // 2),
|
||||
)
|
||||
return out
|
||||
|
||||
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
|
||||
|
||||
from plyngent.lmproto.openai_compatible.model import SystemChatMessage, UserChatMessage
|
||||
|
||||
from .budget import DEFAULT_CONTEXT_MAX_CHARS, DEFAULT_TOOL_RESULT_MAX_CHARS
|
||||
from .budget import DEFAULT_CONTEXT_MAX_TOKENS, DEFAULT_TOOL_RESULT_MAX_CHARS
|
||||
from .events import UsageEvent
|
||||
from .loop import DEFAULT_MAX_ROUNDS, run_chat_loop
|
||||
from .usage import TokenUsage
|
||||
@@ -37,7 +37,7 @@ class ChatAgent:
|
||||
system_prompt: str | None
|
||||
max_tool_result_chars: int
|
||||
parallel_tools: bool
|
||||
max_context_chars: int
|
||||
max_context_tokens: int
|
||||
messages: list[AnyChatMessage]
|
||||
pending_retry_text: str | None
|
||||
session_usage: TokenUsage
|
||||
@@ -59,7 +59,7 @@ class ChatAgent:
|
||||
system_prompt: str | None = None,
|
||||
max_tool_result_chars: int = DEFAULT_TOOL_RESULT_MAX_CHARS,
|
||||
parallel_tools: bool = True,
|
||||
max_context_chars: int = DEFAULT_CONTEXT_MAX_CHARS,
|
||||
max_context_tokens: int = DEFAULT_CONTEXT_MAX_TOKENS,
|
||||
) -> None:
|
||||
self.client = client
|
||||
self.model = model
|
||||
@@ -73,7 +73,7 @@ class ChatAgent:
|
||||
self.system_prompt = system_prompt
|
||||
self.max_tool_result_chars = max_tool_result_chars
|
||||
self.parallel_tools = parallel_tools
|
||||
self.max_context_chars = max_context_chars
|
||||
self.max_context_tokens = max_context_tokens
|
||||
self.messages = list(messages) if messages is not None else []
|
||||
self.pending_retry_text = None
|
||||
self.session_usage = TokenUsage()
|
||||
@@ -135,7 +135,7 @@ class ChatAgent:
|
||||
stream=self.stream,
|
||||
max_tool_result_chars=self.max_tool_result_chars,
|
||||
parallel_tools=self.parallel_tools,
|
||||
max_context_chars=self.max_context_chars,
|
||||
max_context_tokens=self.max_context_tokens,
|
||||
):
|
||||
if isinstance(event, UsageEvent):
|
||||
turn_usage = turn_usage.add(event.usage)
|
||||
|
||||
@@ -13,7 +13,7 @@ from plyngent.lmproto.openai_compatible.model import (
|
||||
UserChatMessage,
|
||||
)
|
||||
|
||||
from .budget import DEFAULT_CONTEXT_MAX_CHARS, compact_messages_for_request
|
||||
from .budget import DEFAULT_CONTEXT_MAX_TOKENS, compact_messages_for_request
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
@@ -67,10 +67,10 @@ def format_transcript(messages: Sequence[AnyChatMessage]) -> str:
|
||||
def soft_compact_transcript(
|
||||
messages: Sequence[AnyChatMessage],
|
||||
*,
|
||||
max_chars: int = DEFAULT_CONTEXT_MAX_CHARS,
|
||||
max_tokens: int = DEFAULT_CONTEXT_MAX_TOKENS,
|
||||
) -> str:
|
||||
"""Soft-compact tool dumps then format as transcript text."""
|
||||
compacted = compact_messages_for_request(messages, max_chars=max_chars)
|
||||
compacted = compact_messages_for_request(messages, max_tokens=max_tokens)
|
||||
return format_transcript(compacted)
|
||||
|
||||
|
||||
@@ -79,14 +79,14 @@ async def summarize_messages(
|
||||
messages: Sequence[AnyChatMessage],
|
||||
*,
|
||||
model: str,
|
||||
max_context_chars: int = DEFAULT_CONTEXT_MAX_CHARS,
|
||||
max_context_tokens: int = DEFAULT_CONTEXT_MAX_TOKENS,
|
||||
temperature: float | None = 0.2,
|
||||
) -> 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_chars=max_context_chars)
|
||||
transcript = soft_compact_transcript(messages, max_tokens=max_context_tokens)
|
||||
if not transcript.strip():
|
||||
msg = "nothing to compact"
|
||||
raise ValueError(msg)
|
||||
|
||||
@@ -20,7 +20,7 @@ from plyngent.lmproto.openai_compatible.model import (
|
||||
from plyngent.typedef import Unset # noqa: TC001
|
||||
|
||||
from .budget import (
|
||||
DEFAULT_CONTEXT_MAX_CHARS,
|
||||
DEFAULT_CONTEXT_MAX_TOKENS,
|
||||
DEFAULT_TOOL_RESULT_MAX_CHARS,
|
||||
compact_messages_for_request,
|
||||
truncate_tool_result,
|
||||
@@ -225,14 +225,14 @@ async def run_chat_loop(
|
||||
stream: bool = True,
|
||||
max_tool_result_chars: int = DEFAULT_TOOL_RESULT_MAX_CHARS,
|
||||
parallel_tools: bool = True,
|
||||
max_context_chars: int = DEFAULT_CONTEXT_MAX_CHARS,
|
||||
max_context_tokens: int = DEFAULT_CONTEXT_MAX_TOKENS,
|
||||
) -> AsyncIterator[AgentEvent]:
|
||||
"""Multi-round chat/tool loop; mutates ``messages`` in place and yields events.
|
||||
|
||||
When ``stream=True``, uses ``chat_completions(..., stream=True)`` and yields
|
||||
text deltas as chunks arrive; tool calls are merged from stream deltas.
|
||||
Multiple tool calls in one round run in parallel when ``parallel_tools``.
|
||||
Request payloads may shrink older tool results when over ``max_context_chars``.
|
||||
Request payloads may shrink older tool results when over ``max_context_tokens``.
|
||||
"""
|
||||
tool_items: Sequence[AnyToolItem] | None = None
|
||||
if tools is not None and len(tools) > 0:
|
||||
@@ -246,7 +246,7 @@ async def run_chat_loop(
|
||||
rounds_used += 1
|
||||
request_messages = compact_messages_for_request(
|
||||
messages,
|
||||
max_chars=max_context_chars,
|
||||
max_tokens=max_context_tokens,
|
||||
)
|
||||
param = ChatCompletionsParam(
|
||||
messages=request_messages,
|
||||
|
||||
@@ -28,7 +28,7 @@ _MINIMAL_CONFIG = """\
|
||||
# parallel_tools = true
|
||||
# confirm_destructive = true
|
||||
# path_denylist = ["/secrets/", ".ssh/"]
|
||||
# max_context_chars = 200000
|
||||
# max_context_tokens = 200000
|
||||
|
||||
# [providers.example]
|
||||
# preset = "openai-compatible"
|
||||
|
||||
@@ -54,25 +54,23 @@ _CONTENT_PREVIEW = 200
|
||||
|
||||
|
||||
def _cmd_status(state: ReplState) -> None:
|
||||
from plyngent.agent.budget import estimate_messages_chars
|
||||
from plyngent.agent.budget import estimate_messages_chars, estimate_messages_tokens
|
||||
|
||||
pending = state.agent.pending_retry_text
|
||||
pending_disp = "yes" if pending else "no"
|
||||
ctx_chars = estimate_messages_chars(state.agent.messages)
|
||||
ctx_budget = state.agent.max_context_chars
|
||||
from plyngent.agent.usage import chars_to_tokens
|
||||
|
||||
ctx_tokens = estimate_messages_tokens(state.agent.messages)
|
||||
ctx_budget = state.agent.max_context_tokens
|
||||
session_u = state.agent.session_usage
|
||||
last_u = state.agent.last_turn_usage
|
||||
ctx_est_tokens = chars_to_tokens(ctx_chars)
|
||||
click.echo(
|
||||
f"provider={state.provider_name} model={state.model}\n"
|
||||
f"session={state.session_id} messages={len(state.agent.messages)} "
|
||||
f"pending_retry={pending_disp}\n"
|
||||
f"tools={'on' if state.tools_enabled else 'off'} "
|
||||
f"rounds={state.max_rounds} stream={'on' if state.agent.stream else 'off'}\n"
|
||||
f"context_chars={ctx_chars}/{ctx_budget} "
|
||||
f"context_tokens~={ctx_est_tokens} (est) "
|
||||
f"context_tokens~={ctx_tokens}/{ctx_budget} (est) "
|
||||
f"context_chars={ctx_chars} "
|
||||
f"tool_result_max={state.agent.max_tool_result_chars}\n"
|
||||
f"usage_session={session_u.format_line()}\n"
|
||||
f"usage_last_turn={last_u.format_line()}\n"
|
||||
|
||||
@@ -78,7 +78,7 @@ class ReplState:
|
||||
system_prompt=system_prompt,
|
||||
max_tool_result_chars=agent_cfg.max_tool_result_chars,
|
||||
parallel_tools=agent_cfg.parallel_tools,
|
||||
max_context_chars=agent_cfg.max_context_chars,
|
||||
max_context_tokens=agent_cfg.max_context_tokens,
|
||||
)
|
||||
|
||||
def rebuild_client(self) -> None:
|
||||
@@ -174,7 +174,7 @@ class ReplState:
|
||||
self.client,
|
||||
messages,
|
||||
model=self.model,
|
||||
max_context_chars=self.agent.max_context_chars,
|
||||
max_context_tokens=self.agent.max_context_tokens,
|
||||
)
|
||||
session_name = name or f"compact-from-{old_id}"
|
||||
await self.new_session(name=session_name)
|
||||
|
||||
@@ -18,7 +18,7 @@ class AgentConfig(Struct, omit_defaults=True):
|
||||
parallel_tools: bool = True
|
||||
confirm_destructive: bool = True
|
||||
path_denylist: list[str] = field(default_factory=list)
|
||||
max_context_chars: int = 200_000
|
||||
max_context_tokens: int = 200_000
|
||||
|
||||
|
||||
class ModelConfig(Struct, omit_defaults=True):
|
||||
|
||||
@@ -6,7 +6,7 @@ from msgspec import UNSET
|
||||
|
||||
from plyngent.agent.budget import (
|
||||
compact_messages_for_request,
|
||||
estimate_messages_chars,
|
||||
estimate_messages_tokens,
|
||||
truncate_tool_result,
|
||||
)
|
||||
from plyngent.agent.loop import run_chat_loop
|
||||
@@ -75,7 +75,7 @@ def test_compact_shrinks_old_tool_results() -> None:
|
||||
|
||||
compacted = compact_messages_for_request(
|
||||
messages,
|
||||
max_chars=estimate_messages_chars(messages) - 1,
|
||||
max_tokens=max(1, estimate_messages_tokens(messages) - 1),
|
||||
old_tool_result_chars=40,
|
||||
keep_recent_tool_results=1,
|
||||
)
|
||||
@@ -90,11 +90,11 @@ def test_compact_shrinks_old_tool_results() -> None:
|
||||
assert compacted[5].content == "NEW" * 50
|
||||
|
||||
|
||||
def test_compact_disabled_when_max_chars_zero() -> None:
|
||||
def test_compact_disabled_when_max_tokens_zero() -> None:
|
||||
messages: list[AnyChatMessage] = [
|
||||
ToolChatMessage(content="x" * 500, tool_call_id="1"),
|
||||
]
|
||||
out = compact_messages_for_request(messages, max_chars=0)
|
||||
out = compact_messages_for_request(messages, max_tokens=0)
|
||||
assert out[0] is messages[0] or (
|
||||
isinstance(out[0], ToolChatMessage) and out[0].content == "x" * 500
|
||||
)
|
||||
@@ -211,7 +211,7 @@ async def test_loop_sends_compacted_request_not_history() -> None:
|
||||
history,
|
||||
model="m",
|
||||
stream=False,
|
||||
max_context_chars=200,
|
||||
max_context_tokens=50,
|
||||
max_tool_result_chars=50,
|
||||
)
|
||||
]
|
||||
|
||||
@@ -43,7 +43,7 @@ def test_soft_compact_transcript_shrinks_tools() -> None:
|
||||
ToolChatMessage(content=big, tool_call_id="1"),
|
||||
ToolChatMessage(content="recent", tool_call_id="2"),
|
||||
]
|
||||
out = soft_compact_transcript(messages, max_chars=500)
|
||||
out = soft_compact_transcript(messages, max_tokens=100)
|
||||
assert "truncated" in out or len(out) < len(big) + 50
|
||||
assert "recent" in out
|
||||
|
||||
|
||||
@@ -142,7 +142,7 @@ async def test_rounds(state: ReplState) -> None:
|
||||
assert state.agent.max_rounds == 40
|
||||
|
||||
|
||||
async def test_status_shows_context_chars(
|
||||
async def test_status_shows_context_tokens(
|
||||
state: ReplState, capsys: pytest.CaptureFixture[str]
|
||||
) -> None:
|
||||
from plyngent.lmproto.openai_compatible.model import UserChatMessage
|
||||
@@ -150,6 +150,7 @@ async def test_status_shows_context_chars(
|
||||
state.agent.messages = [UserChatMessage(content="hello")]
|
||||
assert await handle_slash(state, "/status") is True
|
||||
out = capsys.readouterr().out
|
||||
assert "context_tokens~=" in out
|
||||
assert "context_chars=" in out
|
||||
assert "tool_result_max=" in out
|
||||
assert str(state.workspace) in out
|
||||
|
||||
@@ -17,7 +17,7 @@ def test_agent_section_defaults(tmp_path: Path) -> None:
|
||||
assert store.agent_config.parallel_tools is True
|
||||
assert store.agent_config.confirm_destructive is True
|
||||
assert store.agent_config.path_denylist == []
|
||||
assert store.agent_config.max_context_chars == 200_000
|
||||
assert store.agent_config.max_context_tokens == 200_000
|
||||
|
||||
|
||||
def test_agent_section_parse(tmp_path: Path) -> None:
|
||||
@@ -30,7 +30,7 @@ max_tool_result_chars = 100
|
||||
parallel_tools = false
|
||||
confirm_destructive = false
|
||||
path_denylist = ["/secrets/", ".ssh/"]
|
||||
max_context_chars = 5000
|
||||
max_context_tokens = 5000
|
||||
""",
|
||||
encoding="utf-8",
|
||||
)
|
||||
@@ -40,4 +40,5 @@ max_context_chars = 5000
|
||||
assert store.agent_config.parallel_tools is False
|
||||
assert store.agent_config.confirm_destructive is False
|
||||
assert store.agent_config.path_denylist == ["/secrets/", ".ssh/"]
|
||||
assert store.agent_config.max_context_chars == 5000
|
||||
assert store.agent_config.max_context_tokens == 5000
|
||||
|
||||
|
||||
Reference in New Issue
Block a user