diff --git a/CLAUDE.md b/CLAUDE.md index d253eed..5e1a99d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -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/`) diff --git a/src/plyngent/agent/budget.py b/src/plyngent/agent/budget.py index 0397200..04b3c7a 100644 --- a/src/plyngent/agent/budget.py +++ b/src/plyngent/agent/budget.py @@ -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 diff --git a/src/plyngent/agent/chat.py b/src/plyngent/agent/chat.py index ec79948..972c099 100644 --- a/src/plyngent/agent/chat.py +++ b/src/plyngent/agent/chat.py @@ -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) diff --git a/src/plyngent/agent/compact.py b/src/plyngent/agent/compact.py index 77bba6a..4f5084d 100644 --- a/src/plyngent/agent/compact.py +++ b/src/plyngent/agent/compact.py @@ -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) diff --git a/src/plyngent/agent/loop.py b/src/plyngent/agent/loop.py index 3ddeaba..8364e2a 100644 --- a/src/plyngent/agent/loop.py +++ b/src/plyngent/agent/loop.py @@ -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, diff --git a/src/plyngent/cli/editor.py b/src/plyngent/cli/editor.py index 71e9452..0dc7e3b 100644 --- a/src/plyngent/cli/editor.py +++ b/src/plyngent/cli/editor.py @@ -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" diff --git a/src/plyngent/cli/repl.py b/src/plyngent/cli/repl.py index c43fe33..0263ccc 100644 --- a/src/plyngent/cli/repl.py +++ b/src/plyngent/cli/repl.py @@ -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" diff --git a/src/plyngent/cli/state.py b/src/plyngent/cli/state.py index fcbc686..023c4d7 100644 --- a/src/plyngent/cli/state.py +++ b/src/plyngent/cli/state.py @@ -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) diff --git a/src/plyngent/config/models.py b/src/plyngent/config/models.py index ef4ad6f..c8a03b6 100644 --- a/src/plyngent/config/models.py +++ b/src/plyngent/config/models.py @@ -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): diff --git a/tests/test_agent/test_compact.py b/tests/test_agent/test_compact.py index d152779..dea3ad4 100644 --- a/tests/test_agent/test_compact.py +++ b/tests/test_agent/test_compact.py @@ -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, ) ] diff --git a/tests/test_agent/test_compact_summarize.py b/tests/test_agent/test_compact_summarize.py index c9a9db4..dca01db 100644 --- a/tests/test_agent/test_compact_summarize.py +++ b/tests/test_agent/test_compact_summarize.py @@ -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 diff --git a/tests/test_cli/test_repl_commands.py b/tests/test_cli/test_repl_commands.py index bf1e5bd..481262d 100644 --- a/tests/test_cli/test_repl_commands.py +++ b/tests/test_cli/test_repl_commands.py @@ -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 diff --git a/tests/test_config/test_agent_section.py b/tests/test_config/test_agent_section.py index 9032bae..4f4f07d 100644 --- a/tests/test_config/test_agent_section.py +++ b/tests/test_config/test_agent_section.py @@ -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 +