diff --git a/hindsight-api-slim/hindsight_api/engine/reflect/agent.py b/hindsight-api-slim/hindsight_api/engine/reflect/agent.py index 2fe1e93620..4be0780039 100644 --- a/hindsight-api-slim/hindsight_api/engine/reflect/agent.py +++ b/hindsight-api-slim/hindsight_api/engine/reflect/agent.py @@ -15,7 +15,15 @@ from typing import TYPE_CHECKING, Any, Awaitable, Callable from ...config import get_config -from .models import DirectiveInfo, LLMCall, ReflectAgentResult, StructuredOutputResult, TokenUsageSummary, ToolCall +from .models import ( + DirectiveInfo, + LLMCall, + MaxTokensCapResult, + ReflectAgentResult, + StructuredOutputResult, + TokenUsageSummary, + ToolCall, +) from .prompts import ( _extract_directive_rules, build_final_prompt, @@ -358,6 +366,60 @@ def _json_schema_type_to_python(field_schema: dict) -> type: return StructuredOutputResult() +async def _enforce_answer_token_cap( + answer: str, + max_tokens: int | None, + llm_config: "LLMProvider | None", + reflect_id: str, +) -> MaxTokensCapResult: + """Rewrite ``answer`` to fit within ``max_tokens`` if it overshoots. + + The reflect agent's tool-driven completion is intentionally uncapped so the + LLM has headroom to emit tool-call JSON and intermediate reasoning. But once + the agent's text becomes the user-visible final answer it must respect + max_tokens. Every completion path (forced-final, direct-text short-circuit, + and the done tool) routes through this helper so the cap is enforced + uniformly. See issue #2756. + + When the answer is already within budget (or no cap/LLM is configured), no + call is made and the original answer is returned with ``rewritten=False``. + """ + if max_tokens is None or llm_config is None or not answer: + return MaxTokensCapResult(answer=answer) + if count_cl100k_tokens(answer) <= max_tokens: + return MaxTokensCapResult(answer=answer) + + rewrite_start = time.time() + rewritten, usage = await llm_config.call( + messages=[ + { + "role": "system", + "content": ( + "Rewrite the user's text so it fits within the requested token " + "budget. Preserve the key facts and structure; drop lower-priority " + "detail. Respond with the rewritten text only, no preamble." + ), + }, + { + "role": "user", + "content": f"Target budget: {max_tokens} tokens.\n\nText to rewrite:\n{answer}", + }, + ], + scope="reflect", + max_completion_tokens=max_tokens, + return_usage=True, + ) + return MaxTokensCapResult( + answer=_clean_answer_text(rewritten.strip()), + rewritten=True, + input_tokens=usage.input_tokens, + output_tokens=usage.output_tokens, + cached_tokens=getattr(usage, "cached_tokens", 0) or 0, + thoughts_tokens=getattr(usage, "thoughts_tokens", 0) or 0, + duration_ms=int((time.time() - rewrite_start) * 1000), + ) + + def _count_messages_tokens(messages: list[dict[str, Any]]) -> int: """Estimate the token count of the messages list using cl100k_base encoding.""" total = 0 @@ -870,40 +932,21 @@ def _log_completion(answer: str, iterations: int, forced: bool = False): # must respect max_tokens like the forced-final paths do. If it # overshoots, run one extra capped call to rewrite it within # the cap. - if max_tokens is not None and count_cl100k_tokens(answer) > max_tokens: - rewrite_start = time.time() - rewritten, rewrite_usage = await llm_config.call( - messages=[ - { - "role": "system", - "content": ( - "Rewrite the user's text so it fits within the requested token " - "budget. Preserve the key facts and structure; drop lower-priority " - "detail. Respond with the rewritten text only, no preamble." - ), - }, - { - "role": "user", - "content": f"Target budget: {max_tokens} tokens.\n\nText to rewrite:\n{answer}", - }, - ], - scope="reflect", - max_completion_tokens=max_tokens, - return_usage=True, - ) - total_input_tokens += rewrite_usage.input_tokens - total_output_tokens += rewrite_usage.output_tokens - total_cached_tokens += getattr(rewrite_usage, "cached_tokens", 0) or 0 - total_thoughts_tokens += getattr(rewrite_usage, "thoughts_tokens", 0) or 0 + cap = await _enforce_answer_token_cap(answer, max_tokens, llm_config, reflect_id) + if cap.rewritten: + answer = cap.answer + total_input_tokens += cap.input_tokens + total_output_tokens += cap.output_tokens + total_cached_tokens += cap.cached_tokens + total_thoughts_tokens += cap.thoughts_tokens llm_trace.append( { "scope": "final_rewrite", - "duration_ms": int((time.time() - rewrite_start) * 1000), - "input_tokens": rewrite_usage.input_tokens, - "output_tokens": rewrite_usage.output_tokens, + "duration_ms": cap.duration_ms, + "input_tokens": cap.input_tokens, + "output_tokens": cap.output_tokens, } ) - answer = _clean_answer_text(rewritten.strip()) # Generate structured output if schema provided structured_output = None @@ -1035,6 +1078,7 @@ def _log_completion(answer: str, iterations: int, forced: bool = False): directives_applied=directives_applied, llm_config=llm_config, response_schema=response_schema, + max_tokens=max_tokens, ) # Execute other tools in parallel (exclude done tool in all its format variants) @@ -1244,6 +1288,7 @@ async def _process_done_tool( directives_applied: list[DirectiveInfo], llm_config: "LLMProvider | None" = None, response_schema: dict | None = None, + max_tokens: int | None = None, ) -> ReflectAgentResult: """Process the done tool call and return the result.""" args = done_call.arguments @@ -1254,6 +1299,30 @@ async def _process_done_tool( if not answer: answer = "No answer provided." + # The done tool is the normal completion path, and its answer comes verbatim + # from the (intentionally uncapped) tool-call arguments. Enforce max_tokens + # here too, exactly like the forced-final and direct-text short-circuit paths, + # so a per-model cap is honoured no matter how the agent finishes (#2756). + cap = await _enforce_answer_token_cap(answer, max_tokens, llm_config, reflect_id) + if cap.rewritten: + answer = cap.answer + usage = TokenUsageSummary( + input_tokens=usage.input_tokens + cap.input_tokens, + output_tokens=usage.output_tokens + cap.output_tokens, + total_tokens=usage.total_tokens + cap.input_tokens + cap.output_tokens, + cached_tokens=usage.cached_tokens + cap.cached_tokens, + thoughts_tokens=usage.thoughts_tokens + cap.thoughts_tokens, + ) + llm_trace = [ + *llm_trace, + LLMCall( + scope="final_rewrite", + duration_ms=cap.duration_ms, + input_tokens=cap.input_tokens, + output_tokens=cap.output_tokens, + ), + ] + # Validate IDs (only include IDs that were actually retrieved) used_memory_ids = [mid for mid in (args.get("memory_ids") or []) if mid in available_memory_ids] used_mental_model_ids = [mid for mid in (args.get("mental_model_ids") or []) if mid in available_mental_model_ids] diff --git a/hindsight-api-slim/hindsight_api/engine/reflect/models.py b/hindsight-api-slim/hindsight_api/engine/reflect/models.py index cc54af4e8a..0cabef067c 100644 --- a/hindsight-api-slim/hindsight_api/engine/reflect/models.py +++ b/hindsight-api-slim/hindsight_api/engine/reflect/models.py @@ -106,6 +106,24 @@ class StructuredOutputResult(BaseModel): thoughts_tokens: int = Field(default=0, description="Reasoning/thinking tokens, when reported by the provider") +class MaxTokensCapResult(BaseModel): + """Result of enforcing a max_tokens cap on a final answer. + + Carries the (possibly rewritten) answer plus the token usage and timing of + the rewrite call, so the caller can fold it into its own accounting. When + ``rewritten`` is False the answer was already within budget and no LLM call + was made (all token counts are zero). + """ + + answer: str = Field(description="The answer, rewritten within budget if it overshot") + rewritten: bool = Field(default=False, description="Whether a capped rewrite call was actually made") + input_tokens: int = Field(default=0, description="Input tokens used by the rewrite call") + output_tokens: int = Field(default=0, description="Visible output tokens used by the rewrite call") + cached_tokens: int = Field(default=0, description="Cached prefix tokens. Subset of input_tokens.") + thoughts_tokens: int = Field(default=0, description="Reasoning/thinking tokens, when reported by the provider") + duration_ms: int = Field(default=0, description="Wall-clock duration of the rewrite call") + + class ReflectAgentResult(BaseModel): """Result from the reflect agent.""" diff --git a/hindsight-api-slim/tests/test_reflect_agent.py b/hindsight-api-slim/tests/test_reflect_agent.py index 63d49a2586..09985b61de 100644 --- a/hindsight-api-slim/tests/test_reflect_agent.py +++ b/hindsight-api-slim/tests/test_reflect_agent.py @@ -1263,3 +1263,96 @@ async def test_real_stale_mental_model_forces_and_grounds_in_deeper_evidence(sel context="A stale mental model claimed the launch was still pending, but the freshly retrieved raw " "fact (deploy log A-1029) shows it shipped on Friday. The agent should correct the stale summary.", ) + + +class TestDoneToolMaxTokensCap: + """Regression tests for #2756: max_tokens must be enforced on the done-tool + completion path (the normal path), not only the forced-final and direct-text + short-circuit paths.""" + + @staticmethod + def _done_call(answer: str) -> "LLMToolCall": + return LLMToolCall(id="done-1", name="done", arguments={"answer": answer}) + + @staticmethod + def _process(done_call, llm_config, max_tokens, llm_trace=None): + from hindsight_api.engine.reflect.agent import _process_done_tool + from hindsight_api.engine.reflect.models import TokenUsageSummary + + return _process_done_tool( + done_call, + set(), + set(), + set(), + 2, # iterations + 1, # total_tools_called + [], # tool_trace + list(llm_trace or []), # llm_trace + TokenUsageSummary(input_tokens=100, output_tokens=50, total_tokens=150), + lambda *a, **k: None, # log_completion + "test-reflect", # reflect_id + directives_applied=[], + llm_config=llm_config, + response_schema=None, + max_tokens=max_tokens, + ) + + @pytest.mark.asyncio + async def test_done_tool_over_budget_answer_is_rewritten(self): + """An over-budget done answer triggers a single capped rewrite call.""" + long_answer = "This is a sentence with several words. " * 60 # well over 50 tokens + llm = MagicMock() + llm.call = AsyncMock( + return_value=("Short capped answer.", TokenUsage(input_tokens=200, output_tokens=20, total_tokens=220)) + ) + + result = await self._process(self._done_call(long_answer), llm, max_tokens=50) + + # Cap enforced: the rewritten (short) answer is returned, not the raw one. + assert result.text == "Short capped answer." + # Exactly one rewrite call, with the budget forwarded. + llm.call.assert_awaited_once() + kwargs = llm.call.await_args.kwargs + assert kwargs["max_completion_tokens"] == 50 + assert kwargs["scope"] == "reflect" + # Rewrite usage is folded into the result and traced. + assert result.usage.input_tokens == 100 + 200 + assert result.usage.output_tokens == 50 + 20 + assert any(c.scope == "final_rewrite" for c in result.llm_trace) + + @pytest.mark.asyncio + async def test_done_tool_under_budget_answer_is_untouched(self): + """An answer already within budget makes no rewrite call and is returned as-is.""" + llm = MagicMock() + llm.call = AsyncMock() + + result = await self._process(self._done_call("Concise."), llm, max_tokens=50) + + assert result.text == "Concise." + llm.call.assert_not_awaited() + assert not any(c.scope == "final_rewrite" for c in result.llm_trace) + + @pytest.mark.asyncio + async def test_done_tool_no_cap_configured_is_untouched(self): + """With max_tokens=None the answer is never rewritten, however long.""" + long_answer = "This is a sentence with several words. " * 60 + llm = MagicMock() + llm.call = AsyncMock() + + result = await self._process(self._done_call(long_answer), llm, max_tokens=None) + + assert result.text == _clean_done_answer(long_answer.strip()) + llm.call.assert_not_awaited() + + @pytest.mark.asyncio + async def test_enforce_cap_helper_skips_when_no_cap(self): + """The shared helper is a no-op (no LLM call) when max_tokens is None.""" + from hindsight_api.engine.reflect.agent import _enforce_answer_token_cap + + llm = MagicMock() + llm.call = AsyncMock() + cap = await _enforce_answer_token_cap("some long text " * 50, None, llm, "rid") + + assert cap.rewritten is False + assert cap.answer == "some long text " * 50 + llm.call.assert_not_awaited()