From a596922bd5be1616139aeee2f16149a0d586cf72 Mon Sep 17 00:00:00 2001 From: Kevin Xia Date: Wed, 5 Aug 2026 14:32:46 +0800 Subject: [PATCH] fix: bound conversation compaction budgets --- deeptutor/services/session/context_builder.py | 56 ++++++--- .../services/session/test_context_builder.py | 114 ++++++++++++++++++ 2 files changed, 155 insertions(+), 15 deletions(-) diff --git a/deeptutor/services/session/context_builder.py b/deeptutor/services/session/context_builder.py index 7aaed07239..ff2c7bd9d9 100644 --- a/deeptutor/services/session/context_builder.py +++ b/deeptutor/services/session/context_builder.py @@ -1,6 +1,4 @@ -""" -Build bounded conversation history for unified chat sessions. -""" +"""Build budgeted conversation history for unified chat sessions.""" from __future__ import annotations @@ -11,7 +9,10 @@ from deeptutor.core.stream import StreamEvent, StreamEventType from deeptutor.core.trace import build_trace_metadata, merge_trace_metadata, new_call_id from deeptutor.services.llm.config import LLMConfig -from deeptutor.services.llm.context_window import resolve_effective_context_window +from deeptutor.services.llm.context_window import ( + coerce_positive_int, + resolve_effective_context_window, +) from .protocol import SessionStoreProtocol @@ -19,6 +20,16 @@ #: cap, assume the provider cut it mid-sentence and trim the partial tail. TRUNCATION_GUARD_RATIO = 0.95 +# Ratio-only budgets become impractical for models with very large context +# windows: a 1M-token window would otherwise reserve hundreds of thousands of +# tokens for chat history and ask the summarizer for a six-figure response. +# Keep the rolling-summary strategy, but bound its history plan, summary +# output, and raw-rebuild eligibility threshold. The summary ceiling is also +# constrained by the active generation limit when that setting is lower. +MAX_HISTORY_PLAN_TOKENS = 131_072 +MAX_SUMMARY_OUTPUT_TOKENS = 16_384 +MAX_RAW_REBUILD_TOKENS = 131_072 + def count_tokens(text: str) -> int: """Estimate token count with tiktoken when available.""" @@ -102,7 +113,11 @@ async def process(self, *_args, **_kwargs) -> dict[str, Any]: class ContextBuilder: - """Construct a bounded conversation history plus optional summary trace.""" + """Construct history against a bounded plan plus optional summary trace. + + The budget is a planning target, not destructive truncation: the newest + non-empty message is retained even when that single message exceeds it. + """ def __init__( self, @@ -123,18 +138,29 @@ def _effective_context_window(self, llm_config: LLMConfig) -> int: def _history_budget(self, llm_config: LLMConfig) -> int: effective_context_window = self._effective_context_window(llm_config) - return max(256, int(effective_context_window * self.history_budget_ratio)) - - def _summary_budget(self, budget: int) -> int: - return max(96, int(budget * self.summary_target_ratio)) + ratio_budget = max(256, int(effective_context_window * self.history_budget_ratio)) + return min(ratio_budget, MAX_HISTORY_PLAN_TOKENS) + + def _summary_budget(self, budget: int, llm_config: LLMConfig | None = None) -> int: + ratio_budget = max(96, int(budget * self.summary_target_ratio)) + output_cap = MAX_SUMMARY_OUTPUT_TOKENS + if llm_config is not None: + generation_limit = coerce_positive_int(getattr(llm_config, "max_tokens", None)) + if generation_limit is not None: + output_cap = min(output_cap, generation_limit) + return min(ratio_budget, output_cap) def _recent_budget(self, budget: int) -> int: - return max(128, budget - self._summary_budget(budget)) + # Keep the original ratio-based split independent of the summarizer's + # output cap. Otherwise lowering that cap expands the verbatim tail and + # leaves almost no headroom before the next compaction. + return max(128, int(budget * (1 - self.summary_target_ratio))) def _rebuild_source_budget(self, llm_config: LLMConfig) -> int: - # Raw-rebuild input may use up to half the effective context window; - # beyond that we degrade to fold-in (existing summary + new turns). - return max(1024, self._effective_context_window(llm_config) // 2) + # A raw prefix is eligible for drift-free rebuild up to this threshold; + # beyond it we degrade to fold-in (existing summary + new turns). + ratio_budget = max(1024, self._effective_context_window(llm_config) // 2) + return min(ratio_budget, MAX_RAW_REBUILD_TOKENS) def _build_history(self, summary: str, messages: list[dict[str, Any]]) -> list[dict[str, Any]]: history: list[dict[str, Any]] = [] @@ -282,7 +308,7 @@ async def _trace_bridge(update: dict[str, Any]) -> None: ) # The instruction targets ~80% of the hard cap so the model's own # length control — not the max_tokens cut — is the binding limit. - target_tokens = max(96, int(summary_budget * 0.8)) + target_tokens = max(1, int(summary_budget * 0.8)) system_prompt = ( "You maintain a running summary of a conversation so future turns can " "continue seamlessly. Rewrite the summary from the material provided, " @@ -365,7 +391,7 @@ async def build( return ContextBuildResult([], "", "", [], 0, self._history_budget(llm_config)) budget = self._history_budget(llm_config) - summary_budget = self._summary_budget(budget) + summary_budget = self._summary_budget(budget, llm_config) recent_budget = self._recent_budget(budget) stored_summary = str(session.get("compressed_summary", "") or "").strip() diff --git a/tests/services/session/test_context_builder.py b/tests/services/session/test_context_builder.py index 9b4f8ca83c..2d0e4b1a95 100644 --- a/tests/services/session/test_context_builder.py +++ b/tests/services/session/test_context_builder.py @@ -8,6 +8,9 @@ import pytest from deeptutor.services.session.context_builder import ( + MAX_HISTORY_PLAN_TOKENS, + MAX_RAW_REBUILD_TOKENS, + MAX_SUMMARY_OUTPUT_TOKENS, ContextBuilder, ContextBuildResult, build_history_text, @@ -152,6 +155,11 @@ def test_history_budget_uses_large_context_model_heuristic(self) -> None: budget = builder._history_budget(self._make_llm_config(4096, model="gpt-4o-mini")) assert budget == int(65536 * 0.35) + def test_history_budget_has_absolute_cap(self) -> None: + builder = ContextBuilder(store=MagicMock(), history_budget_ratio=0.35) + budget = builder._history_budget(self._make_llm_config(4096, context_window=983_616)) + assert budget == MAX_HISTORY_PLAN_TOKENS + def test_history_budget_minimum(self) -> None: builder = ContextBuilder(store=MagicMock(), history_budget_ratio=0.01) budget = builder._history_budget(self._make_llm_config(100, model="unknown-local-model")) @@ -165,6 +173,20 @@ def test_summary_budget_minimum(self) -> None: builder = ContextBuilder(store=MagicMock(), summary_target_ratio=0.01) assert builder._summary_budget(100) >= 96 + def test_summary_budget_has_absolute_cap(self) -> None: + builder = ContextBuilder(store=MagicMock(), summary_target_ratio=0.40) + assert builder._summary_budget(MAX_HISTORY_PLAN_TOKENS) == (MAX_SUMMARY_OUTPUT_TOKENS) + + def test_summary_budget_respects_configured_output_limit(self) -> None: + builder = ContextBuilder(store=MagicMock(), summary_target_ratio=0.40) + cfg = self._make_llm_config(4096, context_window=983_616) + assert builder._summary_budget(MAX_HISTORY_PLAN_TOKENS, cfg) == 4096 + + def test_summary_budget_can_follow_small_generation_limit(self) -> None: + builder = ContextBuilder(store=MagicMock(), summary_target_ratio=0.40) + cfg = self._make_llm_config(50, context_window=983_616) + assert builder._summary_budget(MAX_HISTORY_PLAN_TOKENS, cfg) == 50 + def test_recent_budget(self) -> None: builder = ContextBuilder(store=MagicMock(), summary_target_ratio=0.40) recent = builder._recent_budget(1000) @@ -174,6 +196,20 @@ def test_recent_budget_minimum(self) -> None: builder = ContextBuilder(store=MagicMock(), summary_target_ratio=0.99) assert builder._recent_budget(200) >= 128 + def test_recent_budget_preserves_headroom_when_summary_is_capped(self) -> None: + builder = ContextBuilder(store=MagicMock(), summary_target_ratio=0.40) + cfg = self._make_llm_config(4096, context_window=983_616) + recent_budget = builder._recent_budget(MAX_HISTORY_PLAN_TOKENS) + summary_budget = builder._summary_budget(MAX_HISTORY_PLAN_TOKENS, cfg) + + assert recent_budget == int(MAX_HISTORY_PLAN_TOKENS * 0.60) + assert recent_budget + summary_budget < MAX_HISTORY_PLAN_TOKENS + + def test_rebuild_source_budget_has_absolute_cap(self) -> None: + builder = ContextBuilder(store=MagicMock()) + cfg = self._make_llm_config(4096, context_window=983_616) + assert builder._rebuild_source_budget(cfg) == MAX_RAW_REBUILD_TOKENS + # --------------------------------------------------------------------------- # ContextBuilder._build_history @@ -328,6 +364,50 @@ async def test_large_context_model_avoids_premature_summarize(self) -> None: assert result.events == [] store.update_summary.assert_not_called() + @pytest.mark.asyncio + async def test_large_context_build_uses_capped_plan_with_headroom(self) -> None: + store = AsyncMock() + store.get_session = AsyncMock( + return_value={ + "id": "s1", + "compressed_summary": "", + "summary_up_to_msg_id": 0, + } + ) + store.get_messages_for_context = AsyncMock( + return_value=[ + {"id": 1, "role": "user", "content": "old turn"}, + {"id": 2, "role": "assistant", "content": "recent turn"}, + ] + ) + + builder = ContextBuilder(store=store) + builder._summarize = AsyncMock(return_value=("SUMMARY", [])) + select_recent = MagicMock(wraps=builder._select_recent_messages) + builder._select_recent_messages = select_recent + cfg = MagicMock() + cfg.max_tokens = 4096 + cfg.model = "qwen3.8-max" + cfg.context_window = 983_616 + + def _fake_count_tokens(text: str) -> int: + if text in {"old turn", "recent turn"}: + return 70_000 + if "User: old turn" in text and "Assistant: recent turn" in text: + return 140_000 + return 100 + + with patch( + "deeptutor.services.session.context_builder.count_tokens", + side_effect=_fake_count_tokens, + ): + result = await builder.build(session_id="s1", llm_config=cfg) + + assert result.budget == MAX_HISTORY_PLAN_TOKENS + assert builder._summarize.call_args.kwargs["summary_budget"] == 4096 + assert select_recent.call_args.args[1] == int(MAX_HISTORY_PLAN_TOKENS * 0.60) + store.update_summary.assert_awaited_once_with("s1", "SUMMARY", 1) + # --------------------------------------------------------------------------- # trim_incomplete_tail @@ -342,6 +422,40 @@ def test_single_line_kept(self) -> None: assert trim_incomplete_tail("only one line, keep it") == "only one line, keep it" +# --------------------------------------------------------------------------- +# ContextBuilder._summarize +# --------------------------------------------------------------------------- + + +class TestContextBuilderSummarize: + @pytest.mark.asyncio + async def test_target_stays_below_small_output_cap(self) -> None: + captured: dict[str, Any] = {} + + async def _stream_llm(**kwargs: Any): + captured.update(kwargs) + yield "short summary" + + agent = MagicMock() + agent.stream_llm = _stream_llm + builder = ContextBuilder(store=MagicMock()) + + with patch( + "deeptutor.services.session.context_builder._ContextSummaryAgent", + return_value=agent, + ): + summary, _events = await builder._summarize( + session_id="s1", + language="en", + source_text="User: hello", + summary_budget=50, + ) + + assert summary == "short summary" + assert captured["max_tokens"] == 50 + assert "under 40 tokens" in captured["user_prompt"] + + # --------------------------------------------------------------------------- # ContextBuilder.build — summarize paths (rebuild / fold-in / failure / branch) # ---------------------------------------------------------------------------