From 668ddf57d45d6a63e21420ae1915488f299eba19 Mon Sep 17 00:00:00 2001 From: Matthew Printz Date: Thu, 2 Jul 2026 09:23:06 -0600 Subject: [PATCH] Fix: Summarization - handle oversized messages and track statetool tokens MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two summarization-related bugs: Issue #85 — oversized messages stalled summarization. When a single record at the head of the history was larger than the summarization threshold, the record-selection walk chose nothing, so summarization fired on every query but never made progress. Selection and context-fitting are now separate gates: get_records_up_to_threshold selects records (and falls back to the minimal oversized prefix when the head record alone exceeds the threshold), and fit_records_to_context middle-truncates only the copies sent to the summarizer when the selected records overflow the context window. A record above the threshold but below the window is summarized in full; originals are never mutated. Issue #86 — statetool injections skewed token accounting. Statetool input/output pairs are ephemeral tail injections that never land in the persisted records, so their size was being absorbed into base_tokens and destabilizing estimates across loop steps. ChatHistory now tracks the last-known size of each statetool's injection separately (statetool_token_estimates) and folds it into token_overhead. Entries are updated when a statetool fires and retained (not zeroed) when it doesn't. Also clarifies summarizers.py: a module docstring describing the select-then-fit pipeline, helper extraction for the tool-response boundary walks, and explanatory comments on the fit budget. Adds tests/fake_model.py plus coverage for both fixes. --- archytas/agent.py | 21 ++ archytas/chat_history.py | 23 ++ archytas/summarizers.py | 242 +++++++++++++++++- tests/fake_model.py | 50 ++++ tests/test_oversized_message_summarization.py | 220 ++++++++++++++++ tests/test_statetool_token_tracking.py | 141 ++++++++++ 6 files changed, 685 insertions(+), 12 deletions(-) create mode 100644 tests/fake_model.py create mode 100644 tests/test_oversized_message_summarization.py create mode 100644 tests/test_statetool_token_tracking.py diff --git a/archytas/agent.py b/archytas/agent.py index 12ed656..9b2f575 100644 --- a/archytas/agent.py +++ b/archytas/agent.py @@ -441,6 +441,27 @@ async def build_tail_injections(self) -> list[BaseMessage]: fired.append((tool_name, output)) if fired: + # Record the last-known token size of each fired statetool's output + # (issue #86). These injections are ephemeral and never persisted, + # so ChatHistory tracks them separately (statetool_token_estimates) + # and folds them into token_overhead. Entries are only ever + # updated here — a statetool that doesn't fire this turn keeps its + # previous size rather than dropping to zero — which keeps token + # estimates stable across loop steps and prevents the actual-usage + # delta from being wrongly absorbed into base_tokens. + for tool_name, output in fired: + try: + estimate = await self.model.get_num_tokens_from_messages( + [HumanMessage(content=output)] + ) + except Exception as err: + logger.debug( + "Token estimation for statetool %r failed: %r", tool_name, err + ) + continue + if estimate: + self.chat_history.statetool_token_estimates[tool_name] = estimate + state_messages = await self.build_state_injection(fired) tail.extend(state_messages) if self.verbose: diff --git a/archytas/chat_history.py b/archytas/chat_history.py index c61fe1c..05acb02 100644 --- a/archytas/chat_history.py +++ b/archytas/chat_history.py @@ -415,6 +415,9 @@ class OutboundChatHistory: summary_token_count: Optional[int] overhead_token_count: Optional[int] summarization_threshold: Optional[int] + # Last-known per-statetool injection sizes (issue #86). Optional with a + # default so existing constructors that don't pass it keep working. + statetool_token_estimates: Optional[dict[str, int]] = None class ChatHistory: @@ -430,6 +433,13 @@ class ChatHistory: history_summarizer: Optional[callable] summarization_threshold: int tool_token_estimate: int + # Last-known token size of each statetool's injected output, keyed by tool + # name (issue #86). Statetool pairs are ephemeral tail injections, so they + # never appear in the persisted records; tracking their last observed size + # here keeps the token overhead stable across loop steps where a statetool + # does or doesn't fire. Entries are updated when a statetool fires and + # retained (not zeroed) when it doesn't. + statetool_token_estimates: dict[str, int] _tool_hash: str _token_estimate: int|None history_summarization_task: Optional[asyncio.Task] @@ -459,6 +469,7 @@ def __init__( self.system_message = None self.model = model self.tool_token_estimate = 0 + self.statetool_token_estimates = {} self._tool_hash = "" self._token_estimate = None self.loop_summarizer = loop_summarizer @@ -612,11 +623,20 @@ def instruction_token_estimate(self) -> int: return self.instruction.token_count return 0 + @property + def statetool_token_estimate(self) -> int: + """Sum of the last-known token sizes of all statetool injections.""" + return sum( + value for value in self.statetool_token_estimates.values() + if isinstance(value, int) + ) + @property def token_overhead(self): return sum((value for value in ( self.base_tokens, self.tool_token_estimate, + self.statetool_token_estimate, self.instruction_token_estimate, ) if isinstance(value, int))) @@ -928,6 +948,7 @@ def to_dict(self) -> dict[str, Any]: "current_loop_id": self.current_loop_id, "summarization_threshold": self.summarization_threshold, "tool_token_estimate": self.tool_token_estimate, + "statetool_token_estimates": dict(self.statetool_token_estimates), "token_estimate": self._token_estimate, }, "system_message": self.system_message.to_dict() if self.system_message else None, @@ -983,6 +1004,8 @@ def from_dict( history.summarization_threshold = metadata["summarization_threshold"] if metadata.get("tool_token_estimate") is not None: history.tool_token_estimate = metadata["tool_token_estimate"] + if metadata.get("statetool_token_estimates"): + history.statetool_token_estimates = dict(metadata["statetool_token_estimates"]) if metadata.get("token_estimate") is not None: history._token_estimate = metadata["token_estimate"] diff --git a/archytas/summarizers.py b/archytas/summarizers.py index f2ab110..f8763c6 100644 --- a/archytas/summarizers.py +++ b/archytas/summarizers.py @@ -1,3 +1,25 @@ +""" +Summarization pipeline for context management. + +Two independent gates shrink the outgoing chat history, and keeping them +separate is what makes the flow tractable: + +1. Selection (threshold): ``get_records_up_to_threshold`` decides *which* + records are eligible for summarization, based on ``token_threshold``. It + never truncates content — it only chooses where the summarize/keep boundary + sits. +2. Fitting (context window): ``fit_records_to_context`` runs on the records the + selection step chose, just before the summarization request, and + middle-truncates copies of any records that, together, overflow the model's + context window. Records that fit are passed through untouched, and the + original records are never mutated. + +The distinction matters: a record larger than the summarization threshold but +smaller than the context window is summarized in full and never truncated. +Truncation happens only when the *selected* records don't fit the context +window (see ``fit_records_to_context`` for the exact budget). +""" + import asyncio import json import logging @@ -27,6 +49,23 @@ MESSAGE_SUMMARIZATION_THRESHOLD: int = 1000 MESSAGE_SUMMARIZATION_SNIPPET_SIZE: int = 1000 +# Headroom reserved out of the model's context window when sizing a +# summarization request: covers the system/user template scaffolding plus the +# generated summary output. +SUMMARIZATION_TOKEN_RESERVE: int = 8192 +# Below this per-record token target, middle-truncation would leave too little +# content to be meaningful; fall back to a placeholder instead. +MIN_TRUNCATION_TOKENS: int = 256 + +TRUNCATION_NOTICE_TEMPLATE = ( + "\n\n[... {removed} characters removed from the middle of this message " + "because it was too large to summarize in full ...]\n\n" +) +TRUNCATION_PLACEHOLDER = ( + "[Message content omitted from summarization because it was too large to " + "fit in the context window.]" +) + async def get_summarizable_records(chat_history: "ChatHistory") -> "list[RecordType]": from .chat_history import RecordType, SummaryRecord, SystemMessage, AIMessage @@ -44,33 +83,206 @@ async def get_summarizable_records(chat_history: "ChatHistory") -> "list[RecordT return summarizable_records -async def get_records_up_to_threshold(chat_history: "ChatHistory", model: BaseArchytasModel, token_threshold: int) -> "list[RecordType]": +def _has_pending_tool_calls(message: "BaseMessage") -> bool: + """True if `message` is an AIMessage that issued tool calls (whose + ToolMessage responses therefore appear in later records).""" from .chat_history import AIMessage + return isinstance(message, AIMessage) and bool(message.tool_calls) + + +def retreat_past_pending_tool_calls(records: "list[RecordType]", idx: int) -> int: + """Move `idx` backward so that ``records[:idx+1]`` doesn't end on an + AIMessage whose ToolMessage responses would be left out. Returns the new + boundary index (may be -1 if no complete boundary exists at or before + `idx`).""" + while idx >= 0 and _has_pending_tool_calls(records[idx].message): + idx -= 1 + return idx + + +def extend_to_include_tool_responses(records: "list[RecordType]", idx: int) -> int: + """Move `idx` forward so that ``records[:idx+1]`` includes the ToolMessage + responses following an AIMessage's tool calls, avoiding a split between an + AIMessage and its responses. Returns the new boundary index.""" + from .chat_history import ToolMessage + while idx < len(records) - 1 and ( + _has_pending_tool_calls(records[idx].message) + or isinstance(records[idx + 1].message, ToolMessage) + ): + idx += 1 + return idx + + +async def get_records_up_to_threshold(chat_history: "ChatHistory", model: BaseArchytasModel, token_threshold: int) -> "list[RecordType]": + """Select the prefix of summarizable records whose cumulative token count + stays under ``token_threshold``. + + Note: the returned prefix may extend *past* the threshold. When the very + first summarizable record is already larger than the whole threshold, the + normal walk would select nothing and summarization could never make + progress, so the fallback branch selects the minimal oversized prefix + instead. This function never truncates content — fitting the selected + records into the context window (including any truncation) is handled + separately by ``fit_records_to_context``. + """ token_count = chat_history.base_tokens or 0 target_idx: int = -1 + threshold_exceeded = False summarizable_records = await get_summarizable_records(chat_history) for idx, record in enumerate(summarizable_records): - token_count += record.token_count + token_count += record.token_count or 0 if token_count >= token_threshold: + threshold_exceeded = True break target_idx = idx else: + # The whole history fits under the threshold; nothing to summarize. target_idx = -1 - # Don't split AIMessages and their tool call - done = False - while target_idx >= 0 and not done: - target_record = summarizable_records[target_idx] - target_message = target_record.message - if isinstance(target_message, AIMessage) and target_message.tool_calls: - target_idx -= 1 - else: - done = True + # Don't end the selection on an AIMessage whose tool responses would be split off. + target_idx = retreat_past_pending_tool_calls(summarizable_records, target_idx) if target_idx >= 0: return summarizable_records[:target_idx+1] + + # Fallback (issue #85): a single record at the head of the history is + # larger than the whole summarization threshold, so the walk above selected + # nothing. Without this, summarization triggers on every query but never + # makes progress. Select the minimal prefix instead — the head record plus + # however many records are needed to avoid splitting an AIMessage from its + # ToolMessage responses — and let the summarizer shrink any content that + # doesn't fit in the context window (see fit_records_to_context). + if threshold_exceeded and summarizable_records: + end_idx = extend_to_include_tool_responses(summarizable_records, 0) + selected = summarizable_records[:end_idx + 1] + logger.warning( + "A single message (%s tokens) exceeds the summarization threshold " + "(%s tokens); summarizing %d record(s) beyond the threshold, " + "truncating content as needed to fit the context window.", + selected[0].token_count, token_threshold, len(selected), + ) + return selected + + return [] + + +def middle_truncate_text(text: str, keep_chars: int) -> str: + """Cut characters out of the middle of `text` so roughly `keep_chars` + characters remain, keeping the head and tail and inserting a notice where + content was removed.""" + if keep_chars <= 0: + return TRUNCATION_PLACEHOLDER + if len(text) <= keep_chars: + return text + head = keep_chars // 2 + tail = keep_chars - head + notice = TRUNCATION_NOTICE_TEMPLATE.format(removed=len(text) - keep_chars) + return text[:head] + notice + text[-tail:] + + +async def shrink_record_for_summarization( + record: "MessageRecord", + model: BaseArchytasModel, + target_tokens: int, +) -> "MessageRecord": + """ + Return a copy of `record` whose message content has been middle-truncated + to fit within `target_tokens` (issue #85 remediation b, with placeholder + replacement as the fallback). The copy keeps the original record's uuid so + the resulting summary excludes the original — which is never mutated — from + future outgoing history. + """ + from langchain_core.messages import HumanMessage + from .chat_history import MessageRecord + + message = record.message.model_copy(deep=True) + content = message.content + new_content: str + token_count: int + + if isinstance(content, str) and target_tokens >= MIN_TRUNCATION_TOKENS: + token_count = record.token_count or 0 + new_content = content + # Estimate a keep-size from the observed chars-per-token ratio, then + # re-estimate and tighten a few times if the cut wasn't deep enough. + for _ in range(4): + chars_per_token = len(new_content) / max(token_count, 1) + keep_chars = int(target_tokens * chars_per_token * 0.9) + new_content = middle_truncate_text(content, keep_chars) + token_count = await model.get_num_tokens_from_messages( + [HumanMessage(content=new_content)] + ) + if token_count <= target_tokens: + break + else: + new_content = TRUNCATION_PLACEHOLDER + token_count = await model.get_num_tokens_from_messages( + [HumanMessage(content=new_content)] + ) else: - return [] + # Non-string content (multimodal/structured blocks) has no + # well-defined middle to cut; replace it outright. + new_content = TRUNCATION_PLACEHOLDER + token_count = await model.get_num_tokens_from_messages( + [HumanMessage(content=new_content)] + ) + + message.content = new_content + return MessageRecord( + message=message, + uuid=record.uuid, + token_count=token_count, + metadata=dict(record.metadata), + react_loop_id=record.react_loop_id, + ) + + +async def fit_records_to_context( + records: "list[MessageRecord]", + model: BaseArchytasModel, +) -> "list[MessageRecord]": + """ + Ensure the records destined for a summarization request fit within the + model's context window (minus SUMMARIZATION_TOKEN_RESERVE headroom). + Oversized records are replaced by middle-truncated copies (originals are + left untouched); records that fit are passed through unchanged. + """ + context_window = model.contextsize() + if not context_window or not records: + return records + + # Reserve headroom for the summarization prompt scaffolding + generated + # summary, but never let that reserve claim more than half the window: on a + # small context window SUMMARIZATION_TOKEN_RESERVE could otherwise dominate + # (or drive the budget negative), leaving no room to summarize at all. + budget = max(context_window - SUMMARIZATION_TOKEN_RESERVE, context_window // 2) + fitted = list(records) + total = sum(record.token_count or 0 for record in fitted) + + attempts = 0 + while total > budget and attempts < len(fitted) * 2: + attempts += 1 + idx, largest = max( + enumerate(fitted), key=lambda pair: pair[1].token_count or 0 + ) + largest_tokens = largest.token_count or 0 + if largest_tokens <= 0: + break + overage = total - budget + target_tokens = max(largest_tokens - overage, 0) + logger.warning( + "Record %s (%s tokens) is too large to summarize in full; " + "truncating its copy to ~%s tokens for the summarization request.", + largest.uuid, largest_tokens, target_tokens, + ) + shrunk = await shrink_record_for_summarization(largest, model, target_tokens) + if (shrunk.token_count or 0) >= largest_tokens: + # No progress is possible on this record; give up rather than loop. + break + fitted[idx] = shrunk + total = sum(record.token_count or 0 for record in fitted) + + return fitted async def get_records_up_to_loop(chat_history: "ChatHistory", model: BaseArchytasModel) -> "list[RecordType]": @@ -132,6 +344,12 @@ async def default_history_summarizer( case MessageRecord(): records_to_summarize.append(record) + # Shrink any content too large to fit the summarization request into the + # context window (issue #85). Only the copies sent to the summarizer are + # truncated; the original records are untouched and are excluded from + # future outgoing history once the summary lands. + records_to_summarize = await fit_records_to_context(records_to_summarize, model) + uuids = [record.uuid for record in records_to_summarize] jinja_globals = { diff --git a/tests/fake_model.py b/tests/fake_model.py new file mode 100644 index 0000000..53214d7 --- /dev/null +++ b/tests/fake_model.py @@ -0,0 +1,50 @@ +""" +A minimal offline model implementation for unit tests that must not hit a +provider API. Token counts are estimated as ``len(content) // 4``. +""" +from types import SimpleNamespace +from typing import Optional, Sequence + +from archytas.models.base import BaseArchytasModel, ModelConfig + + +class FakeChatModel: + """Stands in for the underlying LangChain chat model.""" + + def __init__(self, response_text: str = "SUMMARY TEXT"): + self.response_text = response_text + self.calls: list = [] + + async def ainvoke(self, input=None, config=None, **kwargs): + self.calls.append(input) + return SimpleNamespace( + content=self.response_text, + usage_metadata={"input_tokens": 0, "output_tokens": 10, "total_tokens": 10}, + ) + + +class FakeModel(BaseArchytasModel): + DEFAULT_MODEL = "fake-model" + + def __init__(self, context_window: int = 10000, **kwargs): + self._context_window = context_window + super().__init__(ModelConfig(model_name="fake-model"), **kwargs) + + def initialize_model(self, **kwargs): + return FakeChatModel() + + def contextsize(self, model_name: Optional[str] = None) -> int | None: + return self._context_window + + async def get_num_tokens_from_messages( + self, + messages, + tools: Optional[Sequence] = None, + ) -> int: + total = 0 + for message in messages: + content = message.content + if not isinstance(content, str): + content = str(content) + total += max(len(content) // 4, 1) + return total diff --git a/tests/test_oversized_message_summarization.py b/tests/test_oversized_message_summarization.py new file mode 100644 index 0000000..5cab43d --- /dev/null +++ b/tests/test_oversized_message_summarization.py @@ -0,0 +1,220 @@ +""" +Tests for issue #85: summarization stalls when a single message is larger than +the summarization threshold. + +All tests run offline against FakeModel (tokens ~= chars // 4). +""" +import pytest +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage + +from archytas.chat_history import ChatHistory, MessageRecord +from archytas.summarizers import ( + TRUNCATION_PLACEHOLDER, + default_history_summarizer, + fit_records_to_context, + get_records_up_to_threshold, + middle_truncate_text, + shrink_record_for_summarization, +) + +from .fake_model import FakeModel + + +def make_history(model, messages_with_tokens) -> ChatHistory: + """Build a ChatHistory whose records carry preset token counts.""" + history = ChatHistory(model=model) + for message, tokens in messages_with_tokens: + record = history.add_message(message) + record.token_count = tokens + return history + + +def text_of_tokens(tokens: int) -> str: + """Text sized so FakeModel estimates it at roughly `tokens` tokens.""" + return "x" * (tokens * 4) + + +class TestRecordSelectionFallback: + + @pytest.mark.asyncio + async def test_single_oversized_record_is_selected(self): + """A lone record larger than the threshold is selected instead of + returning an empty set (the issue #85 stall).""" + model = FakeModel() + history = make_history(model, [ + (HumanMessage(content=text_of_tokens(5000)), 5000), + (AIMessage(content="ok"), 20), + ]) + + selected = await get_records_up_to_threshold(history, model, token_threshold=1000) + + assert len(selected) == 1 + assert selected[0].uuid == history.raw_records[0].uuid + + @pytest.mark.asyncio + async def test_oversized_tool_message_keeps_pair_together(self): + """When the oversized record is a ToolMessage, its calling AIMessage + and sibling ToolMessages are included so the pair is never split.""" + model = FakeModel() + tool_calls = [ + {"id": "call_1", "name": "big_tool", "args": {}, "type": "tool_call"}, + {"id": "call_2", "name": "big_tool", "args": {}, "type": "tool_call"}, + ] + history = make_history(model, [ + (AIMessage(content="calling tool", tool_calls=tool_calls), 50), + (ToolMessage(content=text_of_tokens(5000), tool_call_id="call_1"), 5000), + (ToolMessage(content="small result", tool_call_id="call_2"), 10), + (HumanMessage(content="next question"), 20), + ]) + + selected = await get_records_up_to_threshold(history, model, token_threshold=1000) + + assert [record.uuid for record in selected] == [ + record.uuid for record in history.raw_records[:3] + ] + + @pytest.mark.asyncio + async def test_history_under_threshold_still_selects_nothing(self): + """The fallback must not fire when the history fits under the + threshold; that path intentionally summarizes nothing.""" + model = FakeModel() + history = make_history(model, [ + (HumanMessage(content="hello"), 100), + (AIMessage(content="hi"), 100), + ]) + + selected = await get_records_up_to_threshold(history, model, token_threshold=1000) + + assert selected == [] + + @pytest.mark.asyncio + async def test_normal_threshold_selection_unchanged(self): + """The regular path (several records under the threshold) behaves as + before: select the prefix that fits.""" + model = FakeModel() + history = make_history(model, [ + (HumanMessage(content=f"message {i}"), 300) for i in range(5) + ]) + + selected = await get_records_up_to_threshold(history, model, token_threshold=1000) + + assert len(selected) == 3 + + +class TestMiddleTruncation: + + def test_keeps_head_and_tail(self): + text = "HEAD" + ("m" * 10000) + "TAIL" + truncated = middle_truncate_text(text, keep_chars=1000) + assert truncated.startswith("HEAD") + assert truncated.endswith("TAIL") + assert "removed from the middle" in truncated + assert len(truncated) < len(text) + + def test_short_text_untouched(self): + assert middle_truncate_text("short", keep_chars=1000) == "short" + + def test_zero_budget_gives_placeholder(self): + assert middle_truncate_text("anything", keep_chars=0) == TRUNCATION_PLACEHOLDER + + @pytest.mark.asyncio + async def test_shrink_record_fits_target_and_preserves_original(self): + model = FakeModel() + original_text = "HEAD" + ("m" * 40000) + "TAIL" + record = MessageRecord(message=HumanMessage(content=original_text), token_count=10000) + + shrunk = await shrink_record_for_summarization(record, model, target_tokens=1000) + + assert shrunk.uuid == record.uuid + assert shrunk.token_count <= 1000 + assert "removed from the middle" in shrunk.message.content + # Original record is left completely untouched. + assert record.message.content == original_text + assert record.token_count == 10000 + + @pytest.mark.asyncio + async def test_non_string_content_replaced_with_placeholder(self): + model = FakeModel() + content = [{"type": "text", "text": "m" * 40000}] + record = MessageRecord(message=HumanMessage(content=content), token_count=10000) + + shrunk = await shrink_record_for_summarization(record, model, target_tokens=1000) + + assert shrunk.message.content == TRUNCATION_PLACEHOLDER + assert record.message.content == content + + @pytest.mark.asyncio + async def test_tiny_target_replaced_with_placeholder(self): + model = FakeModel() + record = MessageRecord(message=HumanMessage(content="m" * 40000), token_count=10000) + + shrunk = await shrink_record_for_summarization(record, model, target_tokens=10) + + assert shrunk.message.content == TRUNCATION_PLACEHOLDER + + +class TestFitRecordsToContext: + + @pytest.mark.asyncio + async def test_records_within_budget_pass_through(self): + model = FakeModel(context_window=100000) + records = [ + MessageRecord(message=HumanMessage(content="small"), token_count=10), + ] + fitted = await fit_records_to_context(records, model) + assert fitted[0] is records[0] + + @pytest.mark.asyncio + async def test_oversized_record_truncated_to_fit(self): + # context 2000 -> budget = max(2000 - 8192, 1000) = 1000 + model = FakeModel(context_window=2000) + small = MessageRecord(message=HumanMessage(content="small question"), token_count=10) + huge = MessageRecord( + message=HumanMessage(content=text_of_tokens(5000)), token_count=5000, + ) + + fitted = await fit_records_to_context([small, huge], model) + + assert fitted[0] is small + assert fitted[1] is not huge + assert fitted[1].uuid == huge.uuid + total = sum(record.token_count for record in fitted) + assert total <= 1000 + # Original untouched. + assert huge.token_count == 5000 + + +class TestEndToEndSummarization: + + @pytest.mark.asyncio + async def test_oversized_message_gets_summarized(self): + """The full path: an oversized record is selected, truncated for the + request, summarized, and thereafter excluded from outgoing history.""" + model = FakeModel(context_window=2000) # threshold = 1000 + history = make_history(model, [ + (HumanMessage(content="HEAD" + text_of_tokens(5000) + "TAIL"), 5001), + ]) + original_content = history.raw_records[0].message.content + + recordset = await get_records_up_to_threshold(history, model, token_threshold=1000) + assert recordset, "oversized record must be selected for summarization" + + await default_history_summarizer( + chat_history=history, agent=None, recordset=recordset, model=model, + ) + + assert len(history.summaries) == 1 + summary = history.summaries[0] + assert summary.summarized_messages == {history.raw_records[0].uuid} + + # The original record was not mutated... + assert history.raw_records[0].message.content == original_content + # ...but is excluded from the outgoing history now that it is summarized. + outgoing = await history.records() + assert history.raw_records[0].uuid not in [record.uuid for record in outgoing] + + # The request actually sent to the model was truncated to fit. + sent = model._model.calls[0] + user_prompt = sent[-1].content + assert "removed from the middle" in user_prompt + assert len(user_prompt) < len(original_content) diff --git a/tests/test_statetool_token_tracking.py b/tests/test_statetool_token_tracking.py new file mode 100644 index 0000000..ed4a60d --- /dev/null +++ b/tests/test_statetool_token_tracking.py @@ -0,0 +1,141 @@ +""" +Tests for issue #86: statetool outputs are counted sporadically in token +estimates because their tail injections are ephemeral. ChatHistory now tracks +the last-known size of each statetool's output separately and folds it into +the token overhead, updating on fire and retaining (not zeroing) otherwise. + +All tests run offline against FakeModel (tokens ~= chars // 4). +""" +import pytest + +from archytas.agent import Agent +from archytas.chat_history import ChatHistory +from archytas.tool_utils import statetool, AgentRef + +from .fake_model import FakeModel + + +def make_gated_statetool(flag: dict, output: dict, name: str = "gated_state"): + """A statetool that fires when flag['fire'] and returns output['content'].""" + + def gate_condition(agent: AgentRef) -> bool: + return bool(flag.get("fire")) + + @statetool(condition=gate_condition, name=name) + def gated_state() -> str: + """ + ** INTERNAL ** Return canned state content for tests. + + Returns: + str: The canned content. + """ + return output["content"] + + return gated_state + + +class TestChatHistoryStatetoolAccounting: + + def test_statetool_estimate_sums_entries(self): + history = ChatHistory(model=FakeModel()) + assert history.statetool_token_estimate == 0 + + history.statetool_token_estimates["tool_a"] = 100 + history.statetool_token_estimates["tool_b"] = 250 + assert history.statetool_token_estimate == 350 + + def test_token_overhead_includes_statetool_estimate(self): + history = ChatHistory(model=FakeModel()) + history.base_tokens = 50 + history.tool_token_estimate = 200 + baseline = history.token_overhead + + history.statetool_token_estimates["tool_a"] = 100 + assert history.token_overhead == baseline + 100 + + def test_serde_roundtrip_preserves_estimates(self): + history = ChatHistory(model=FakeModel()) + history.statetool_token_estimates = {"tool_a": 100, "tool_b": 250} + + data = history.to_dict() + restored = ChatHistory.from_dict(data, model=FakeModel()) + + assert restored.statetool_token_estimates == {"tool_a": 100, "tool_b": 250} + + +class TestStatetoolEstimateTracking: + + @pytest.mark.asyncio + async def test_estimate_recorded_when_fired(self): + flag = {"fire": True} + output = {"content": "s" * 400} # ~100 tokens under FakeModel + tool_fn = make_gated_statetool(flag, output) + + agent = Agent(model=FakeModel(), spinner=None) + agent.statetools = {"gated_state": tool_fn} + + await agent.build_tail_injections() + + estimate = agent.chat_history.statetool_token_estimates.get("gated_state") + assert estimate == 100 + + @pytest.mark.asyncio + async def test_estimate_retained_when_not_fired(self): + """A statetool that doesn't fire keeps its last-known size instead of + dropping back to zero — the 'jumping around' from the issue.""" + flag = {"fire": True} + output = {"content": "s" * 400} + tool_fn = make_gated_statetool(flag, output) + + agent = Agent(model=FakeModel(), spinner=None) + agent.statetools = {"gated_state": tool_fn} + + await agent.build_tail_injections() + overhead_after_fire = agent.chat_history.token_overhead + assert agent.chat_history.statetool_token_estimates["gated_state"] == 100 + + flag["fire"] = False + tail = await agent.build_tail_injections() + + assert tail == [] + assert agent.chat_history.statetool_token_estimates["gated_state"] == 100 + assert agent.chat_history.token_overhead == overhead_after_fire + + @pytest.mark.asyncio + async def test_estimate_updated_on_refire_with_new_size(self): + flag = {"fire": True} + output = {"content": "s" * 400} + tool_fn = make_gated_statetool(flag, output) + + agent = Agent(model=FakeModel(), spinner=None) + agent.statetools = {"gated_state": tool_fn} + + await agent.build_tail_injections() + assert agent.chat_history.statetool_token_estimates["gated_state"] == 100 + + output["content"] = "s" * 1200 # ~300 tokens + await agent.build_tail_injections() + assert agent.chat_history.statetool_token_estimates["gated_state"] == 300 + + @pytest.mark.asyncio + async def test_multiple_statetools_tracked_independently(self): + flag_a = {"fire": True} + flag_b = {"fire": True} + tool_a = make_gated_statetool(flag_a, {"content": "a" * 400}, name="state_a") + tool_b = make_gated_statetool(flag_b, {"content": "b" * 800}, name="state_b") + + agent = Agent(model=FakeModel(), spinner=None) + agent.statetools = {"state_a": tool_a, "state_b": tool_b} + + await agent.build_tail_injections() + assert agent.chat_history.statetool_token_estimates == { + "state_a": 100, "state_b": 200, + } + + # Only state_b fires next turn; state_a's size is retained. + flag_a["fire"] = False + await agent.build_tail_injections() + assert agent.chat_history.statetool_token_estimates == { + "state_a": 100, "state_b": 200, + } + assert agent.chat_history.statetool_token_estimate == 300