butler/tools/budget_config.py
Rama main · réplica de lectura del repositorio original.
"""Immutable budget configuration for multi-layer tool result persistence.
This module provides a frozen configuration object that controls how
much tool output is retained inline vs. persisted to disk per result
and per turn. The companion :class:`TurnBudgetTracker` is a lightweight
runtime helper that accumulates the size of tool results produced in a
single agent turn and reports whether the aggregate budget has been
exceeded so that downstream persistence logic can flush / summarize /
drop data accordingly.
Design principles:
* Immutability — :class:`BudgetConfig` is a frozen dataclass so that
budget decisions cannot be mutated mid-turn (which would be unsafe
when the config is shared across threads / tasks).
* Proportional scaling — every numeric threshold can be rescaled to
the target model context window via :meth:`BudgetConfig.scale_to_context`
so the same code base can serve 8 k BPE models and 200 k BPE models
without hand-tuning.
* Pin protection — certain tools (``read_file``, etc.) must *never* be
subjected to a result-size cap, because capping them would defeat
the whole purpose of reading a file and lead to silent data loss or
``persist → read → persist`` feedback loops. See ``PINNED_THRESHOLDS``.
"""
from __future__ import annotations
import math
from dataclasses import dataclass, field
from typing import Dict, Mapping, Optional
# ── constants ────────────────────────────────────────────────────
DEFAULT_RESULT_SIZE_CHARS: int = 100_000
"""Default per-result size threshold in characters."""
DEFAULT_TURN_BUDGET_CHARS: int = 200_000
"""Default aggregate per-turn budget in characters."""
DEFAULT_PREVIEW_SIZE_CHARS: int = 1_500
"""Default size of the inline snippet kept after persistence."""
_CHARS_PER_TOKEN: int = 4
"""Rough token → character conversion factor used for budget scaling."""
_PER_RESULT_WINDOW_FRACTION: float = 0.15
"""Fraction of the model context window reserved for a single result."""
_PER_TURN_WINDOW_FRACTION: float = 0.30
"""Fraction of the model context window reserved for the whole turn."""
_MIN_RESULT_SIZE_CHARS: int = 8_000
"""Floor applied when scaling down so tiny models still get a usable budget."""
# Tools whose thresholds MUST NEVER be overridden.
# Mapping is ``tool_name → threshold_in_chars``. ``math.inf`` means
# "never cap the result" which is the correct behaviour for tools like
# ``read_file`` — capping a read would silently truncate source code
# viewed by the model, producing incorrect downstream decisions.
PINNED_THRESHOLDS: Dict[str, float] = {
"read_file": math.inf,
"read_multiple_files": math.inf,
"glob": math.inf,
"grep": math.inf,
"list_dir": math.inf,
"list_project_files": math.inf,
"web_fetch": math.inf,
"web_search": math.inf,
}
# ── helpers ───────────────────────────────────────────────────────
def _scale_result_size(context_window_tokens: int) -> int:
"""Scale the per-result budget proportionally to *context_window_tokens*."""
context_chars = context_window_tokens * _CHARS_PER_TOKEN
scaled = math.ceil(context_chars * _PER_RESULT_WINDOW_FRACTION)
return max(_MIN_RESULT_SIZE_CHARS, scaled)
def _scale_turn_budget(context_window_tokens: int) -> int:
"""Scale the per-turn budget proportionally to *context_window_tokens*."""
context_chars = context_window_tokens * _CHARS_PER_TOKEN
scaled = math.ceil(context_chars * _PER_TURN_WINDOW_FRACTION)
return max(_MIN_RESULT_SIZE_CHARS * 2, scaled)
# ── BudgetConfig ─────────────────────────────────────────────────
@dataclass(frozen=True)
class BudgetConfig:
"""Immutable configuration governing tool-result persistence budgets.
Attributes:
default_result_size:
Default cap (chars) applied to any tool result that has no
more specific override. ``math.inf`` disables the cap.
turn_budget:
Aggregate per-turn cap (chars) across *all* tool results
produced in a single agent step. When this is exceeded the
:class:`TurnBudgetTracker` reports ``exceeded() == True`` so
that callers can flush / summarize / drop data.
preview_size:
Number of characters retained inline after a tool result has
been persisted to disk (``=`` the size of the snippet the
model sees in the follow-up message).
tool_overrides:
Tool-specific overrides keyed by tool name. These take
precedence over the registry's ``max_result_size`` value
(but *not* over the :data:`PINNED_THRESHOLDS`).
"""
default_result_size: int = DEFAULT_RESULT_SIZE_CHARS
turn_budget: int = DEFAULT_TURN_BUDGET_CHARS
preview_size: int = DEFAULT_PREVIEW_SIZE_CHARS
tool_overrides: Dict[str, int] = field(default_factory=dict)
# ── resolution ────────────────────────────────────────────────
def resolve_threshold(
self,
tool_name: str,
registry: Optional[Mapping[str, int]] = None,
) -> float:
"""Return the effective per-result threshold for *tool_name*.
Resolution priority (highest → lowest):
1. :data:`PINNED_THRESHOLDS` — pinned tools always return the
pinned value (typically ``math.inf``) regardless of any
caller-supplied override.
2. ``tool_overrides`` set on this config instance.
3. ``registry[tool_name]`` — tool-specific threshold registered
externally (e.g. via the :class:`ToolRegistry`).
4. ``default_result_size`` — the config-wide default.
:param tool_name:
Exact tool identifier as registered with the tool system.
:param registry:
Optional mapping ``{tool_name → threshold}``. If supplied
and *tool_name* is missing from both the pinned set and the
local overrides, this mapping is consulted as the third
source of truth before falling back to the default.
:returns:
The effective threshold in chars. May be ``math.inf`` when
the tool is pinned or when an override was explicitly set
to infinity. The return type is ``float`` because it may
be ``math.inf``; consumers that expect a finite integer
should check ``math.isfinite(value)`` first.
"""
if tool_name in PINNED_THRESHOLDS:
return PINNED_THRESHOLDS[tool_name]
override = self.tool_overrides.get(tool_name)
if override is not None:
return float(override)
if registry is not None:
registry_value = registry.get(tool_name)
if registry_value is not None:
return float(registry_value)
return float(self.default_result_size)
# ── scaling ───────────────────────────────────────────────────
def scale_to_context(self, context_window_tokens: int) -> "BudgetConfig":
"""Return a new config with thresholds rescaled to *context_window_tokens*.
The scaling preserves the relative proportions documented in the
module-level constants but never produces a value below
``_MIN_RESULT_SIZE_CHARS`` so that very small models still get
a usable budget.
:param context_window_tokens:
Model context window expressed in tokens (BPE).
:returns:
A new :class:`BudgetConfig` instance with rescaled
``default_result_size`` and ``turn_budget``. ``preview_size``
and ``tool_overrides`` are preserved verbatim — they encode
user intent, not window size.
"""
scaled_result = _scale_result_size(context_window_tokens)
scaled_turn = _scale_turn_budget(context_window_tokens)
return BudgetConfig(
default_result_size=scaled_result,
turn_budget=scaled_turn,
preview_size=self.preview_size,
tool_overrides=dict(self.tool_overrides),
)
# ── construction helpers ──────────────────────────────────────
def with_overrides(self, overrides: Mapping[str, int]) -> "BudgetConfig":
"""Return a new config merging *overrides* into ``tool_overrides``.
Existing entries are preserved; entries in *overrides* take
precedence when keys collide. The original instance is never
mutated (frozen dataclass contract).
"""
merged: Dict[str, int] = dict(self.tool_overrides)
for key, value in overrides.items():
merged[key] = value
return BudgetConfig(
default_result_size=self.default_result_size,
turn_budget=self.turn_budget,
preview_size=self.preview_size,
tool_overrides=merged,
)
# ── default singleton ─────────────────────────────────────────────
DEFAULT_BUDGET: BudgetConfig = BudgetConfig()
"""Module-wide default :class:`BudgetConfig` instance.
Consumers that do not need custom thresholds can import this singleton
directly:
.. code-block:: python
from butler.tools.budget_config import DEFAULT_BUDGET
threshold = DEFAULT_BUDGET.resolve_threshold("read_file")
"""
# ── TurnBudgetTracker ─────────────────────────────────────────────
class TurnBudgetTracker:
"""Accumulates tool-result sizes for a single agent turn.
Usage::
tracker = TurnBudgetTracker(budget)
for result in tool_results:
tracker.record_result(result.tool_name, len(result.content))
if tracker.exceeded():
flush_persistence_buffers()
break
remaining = tracker.remaining()
if tracker.exceeded():
# trigger summarization / fallback persistence
...
The tracker is *not* thread-safe — budget accounting is an
intrinsically per-turn, single-threaded concern.
"""
__slots__ = ("_budget", "_used", "_tool_totals")
def __init__(self, budget: BudgetConfig = DEFAULT_BUDGET) -> None:
self._budget = budget
self._used: int = 0
self._tool_totals: Dict[str, int] = {}
# ── mutation ──────────────────────────────────────────────────
def record_result(self, tool_name: str, char_count: int) -> None:
"""Record a single tool result of *char_count* chars.
:param tool_name:
Tool identifier (used for per-tool accounting breakdown).
:param char_count:
Size of the result in characters. Negative values are
treated as zero — callers may safely pass ``len(content)``
for any string-like payload.
"""
if char_count < 0:
char_count = 0
self._used += char_count
self._tool_totals[tool_name] = self._tool_totals.get(tool_name, 0) + char_count
def reset(self) -> None:
"""Reset all accumulators (start of a new turn)."""
self._used = 0
self._tool_totals.clear()
# ── queries ───────────────────────────────────────────────────
@property
def used(self) -> int:
"""Total characters recorded so far this turn."""
return self._used
@property
def tool_totals(self) -> Dict[str, int]:
"""Per-tool character totals (defensive copy)."""
return dict(self._tool_totals)
def remaining(self) -> int:
"""Characters left before the aggregate budget is exhausted.
If the budget has already been overrun the returned value is
clamped to ``0`` so callers can safely use it as a capacity
hint without additional guard checks.
"""
remaining = self._budget.turn_budget - self._used
return remaining if remaining > 0 else 0
def exceeded(self) -> bool:
"""Return ``True`` when the turn budget has been exhausted."""
return self._used >= self._budget.turn_budget
def tool_count(self) -> int:
"""Number of distinct tool names recorded this turn."""
return len(self._tool_totals)
# ── tests (``python -m butler.tools.budget_config``) ─────────────
if __name__ == "__main__": # pragma: no cover - manual smoke tests
def _assert(condition: bool, message: str) -> None:
if not condition:
raise AssertionError(message)
config = BudgetConfig()
# ── 1. Priority resolution ──────────────────────────────────
# Pinned tool wins over everything
pinned_threshold = config.resolve_threshold("read_file")
_assert(math.isinf(pinned_threshold), "read_file must be pinned to inf")
# Tool override wins over registry
override_cfg = config.with_overrides({"my_tool": 42})
registry_map = {"my_tool": 9999, "another_tool": 8888}
_assert(
override_cfg.resolve_threshold("my_tool", registry_map) == 42,
"tool_overrides must take precedence over registry",
)
# Registry consulted when no override
_assert(
override_cfg.resolve_threshold("another_tool", registry_map) == 8888,
"registry must be consulted when no override exists",
)
# Falls back to default when neither
_assert(
config.resolve_threshold("unknown_tool") == DEFAULT_RESULT_SIZE_CHARS,
"must fall back to default_result_size",
)
# ── 2. Context window scaling ───────────────────────────────
tiny_cfg = config.scale_to_context(4_000)
_assert(tiny_cfg.default_result_size >= _MIN_RESULT_SIZE_CHARS,
"scaled result must not drop below MIN_RESULT_SIZE_CHARS")
_assert(tiny_cfg.turn_budget >= _MIN_RESULT_SIZE_CHARS * 2,
"scaled turn budget must not drop below 2x MIN_RESULT_SIZE_CHARS")
large_cfg = config.scale_to_context(200_000)
_assert(
large_cfg.default_result_size > config.default_result_size,
"scaling to a larger window must increase result size",
)
_assert(
large_cfg.turn_budget > config.turn_budget,
"scaling to a larger window must increase turn budget",
)
_assert(
large_cfg.preview_size == DEFAULT_PREVIEW_SIZE_CHARS,
"preview_size must be preserved by scale_to_context",
)
# Proportional sanity check for a 200k context
expected_result = max(
_MIN_RESULT_SIZE_CHARS,
math.ceil(200_000 * _CHARS_PER_TOKEN * _PER_RESULT_WINDOW_FRACTION),
)
_assert(
large_cfg.default_result_size == expected_result,
"200k scaling must match documented fraction",
)
# ── 3. Pin protection ────────────────────────────────────────
# Override must NOT be able to override a pinned tool
attacked_cfg = config.with_overrides({"read_file": 1000})
_assert(
math.isinf(attacked_cfg.resolve_threshold("read_file")),
"pinned tools must be immune to tool_overrides",
)
attacked_cfg2 = config.with_overrides({"read_file": 1000})
registry_pinned = {"read_file": 500}
_assert(
math.isinf(attacked_cfg2.resolve_threshold("read_file", registry_pinned)),
"pinned tools must be immune to registry overrides",
)
# Glob/grep/web_fetch are also pinned
for pinned_tool in ("glob", "grep", "web_fetch", "list_dir"):
_assert(
math.isinf(config.resolve_threshold(pinned_tool)),
f"{pinned_tool} must be pinned",
)
# ── 4. Turn budget tracking ─────────────────────────────────
tracker = TurnBudgetTracker(budget=BudgetConfig(turn_budget=1_000))
_assert(tracker.used == 0, "initial used must be 0")
_assert(tracker.remaining() == 1_000, "initial remaining must equal turn_budget")
_assert(tracker.exceeded() is False, "must not be exceeded at start")
_assert(tracker.tool_count() == 0, "tool_count must be 0 before any record")
tracker.record_result("bash", 300)
tracker.record_result("bash", 200)
tracker.record_result("read_file", 400)
_assert(tracker.used == 900, "used must sum recorded chars")
_assert(tracker.tool_count() == 2, "must count distinct tools")
_assert(tracker.remaining() == 100, "remaining must equal budget - used")
_assert(tracker.exceeded() is False, "must not be exceeded below budget")
tracker.record_result("bash", 200) # pushes us over
_assert(tracker.used == 1_100, "used must keep growing")
_assert(tracker.remaining() == 0, "remaining must clamp to 0 after overrun")
_assert(tracker.exceeded() is True, "must be exceeded once budget overrun")
# Negative char counts are treated as zero
tracker.record_result("bash", -50)
_assert(tracker.used == 1_100, "negative char_count must be treated as zero")
tool_totals = tracker.tool_totals
_assert(tool_totals.get("bash") == 700, "per-tool bash total must be 700")
_assert(tool_totals.get("read_file") == 400, "per-tool read_file total must be 400")
# reset() wipes state
tracker.reset()
_assert(tracker.used == 0, "reset must zero used")
_assert(tracker.remaining() == 1_000, "reset must restore full remaining")
_assert(tracker.exceeded() is False, "reset must clear exceeded flag")
_assert(tracker.tool_count() == 0, "reset must clear tool_totals")
# with_overrides produces a new frozen instance
merged = config.with_overrides({"custom_tool": 5_000})
_assert(
merged.resolve_threshold("custom_tool") == 5_000,
"with_overrides must add a new threshold",
)
_assert(
merged.resolve_threshold("read_file") == math.inf,
"with_overrides must not affect pinned tools",
)
_assert(
config.resolve_threshold("custom_tool") == DEFAULT_RESULT_SIZE_CHARS,
"original config must remain unchanged (immutable)",
)
# DEFAULT_BUDGET singleton is usable directly
_assert(
DEFAULT_BUDGET.resolve_threshold("unknown") == DEFAULT_RESULT_SIZE_CHARS,
"DEFAULT_BUDGET must behave like default config",
)
print("budget_config: all smoke tests passed ✓")