From 9b241b24ea11854797fb26dcd8752f38401f010c Mon Sep 17 00:00:00 2001 From: kigland Date: Sun, 6 Sep 2026 10:38:35 +0800 Subject: [PATCH] Enforce reflect length after forced synthesis --- .../hindsight_api/engine/reflect/agent.py | 143 +++++++++++------- .../tests/test_reflect_agent.py | 48 ++++++ 2 files changed, 140 insertions(+), 51 deletions(-) diff --git a/hindsight-api-slim/hindsight_api/engine/reflect/agent.py b/hindsight-api-slim/hindsight_api/engine/reflect/agent.py index b28ad1bb74..52820b3bd9 100644 --- a/hindsight-api-slim/hindsight_api/engine/reflect/agent.py +++ b/hindsight-api-slim/hindsight_api/engine/reflect/agent.py @@ -18,6 +18,7 @@ from ..llm_interface import LLM_TOOL_CHOICE_AUTO, LLMToolChoice from ..llm_trace import LLMQueueWait, reset_queue_wait_sink, set_queue_wait_sink from ..llm_transport import describe_llm_error +from ..response_models import TokenUsage from .models import DirectiveInfo, LLMCall, ReflectAgentResult, StructuredOutputResult, TokenUsageSummary, ToolCall from .prompts import ( _SPLIT_SYNTHESIS_WARN_CHUNKS, @@ -836,6 +837,26 @@ async def _forced_final_synthesis(iterations_completed: int) -> ReflectAgentResu f"over {len(chunks)} context chunk(s)." ) + answer, _, rewrite_usage, rewrite_trace = await _rewrite_final_answer( + answer, + llm_config, + max_tokens, + ) + if rewrite_usage is not None: + 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 + if rewrite_trace is not None: + llm_trace.append( + { + "scope": rewrite_trace.scope, + "duration_ms": rewrite_trace.duration_ms, + "input_tokens": rewrite_trace.input_tokens, + "output_tokens": rewrite_trace.output_tokens, + } + ) + structured_output = None # ``answer`` is non-empty past the guard above, so only the schema gates this. if response_schema: @@ -1340,6 +1361,68 @@ def _document_from_rewrite(rewritten: str, previous_answer: str) -> CanonicalDoc return CanonicalDocument(markdown=text, structure=split_markdown(text)) +async def _rewrite_final_answer( + answer: str, + llm_config: "LLMProvider | None", + max_tokens: int | None, + document: StructuredDocument | None = None, +) -> tuple[str, StructuredDocument | None, TokenUsage | None, LLMCall | None]: + """Shorten an over-budget final answer, regardless of how reflect completed.""" + if llm_config is None or max_tokens is None or count_prompt_tokens(answer) <= max_tokens: + return answer, document, None, None + + rewrite_start = time.time() + # In document mode the trim is asked for as a document too. Asking for + # prose here would put the model back in the business of writing the + # markdown that gets stored — on the one path where the answer is long + # enough that its structure matters most. + if document is not None: + rewrite_system = ( + "Shorten the user's document so it fits within the requested token budget. " + "Preserve the key facts and the document's structure; drop lower-priority detail. " + 'Respond ONLY with JSON: {"sections": [{"heading": "...", "level": 2, ' + '"blocks": ["...", "..."]}]}. A heading carries no "#", and each block is one ' + "paragraph, list, table or code fence." + ) + rewrite_user = f"Target budget: {max_tokens} tokens.\n\nDocument to shorten:\n{answer}" + else: + # The token budget is enforced via the prompt, not a hard provider cap: + # on thinking models a hard cap is eaten by reasoning tokens and would + # truncate the rewrite mid-word (#3365). Cost is bounded by the separate + # reflect_max_completion_tokens config (uncapped by default). + rewrite_system = ( + "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." + ) + rewrite_user = f"Target budget: {max_tokens} tokens.\n\nText to rewrite:\n{answer}" + + call_result = await llm_config.call( + messages=[ + {"role": "system", "content": rewrite_system}, + {"role": "user", "content": rewrite_user}, + ], + scope="reflect", + temperature=get_config().llm_temperature_reflect, + max_completion_tokens=get_config().reflect_max_completion_tokens, + ) + rewritten = call_result.content + if document is not None: + trimmed = _document_from_rewrite(rewritten, answer) + document, answer = trimmed.structure, trimmed.markdown + else: + answer = rewritten.strip() + + usage = call_result.usage + trace = LLMCall( + scope="final_rewrite", + duration_ms=int((time.time() - rewrite_start) * 1000), + input_tokens=usage.input_tokens, + output_tokens=usage.output_tokens, + ) + return answer, document, usage, trace + + async def _process_done_tool( done_call: "LLMToolCall", available_memory_ids: set[str], @@ -1389,49 +1472,14 @@ async def _process_done_tool( ) final_usage = usage - if llm_config and max_tokens is not None and count_prompt_tokens(answer) > max_tokens: - rewrite_start = time.time() - # In document mode the trim is asked for as a document too. Asking for - # prose here would put the model back in the business of writing the - # markdown that gets stored — on the one path where the answer is long - # enough that its structure matters most. - if document is not None: - rewrite_system = ( - "Shorten the user's document so it fits within the requested token budget. " - "Preserve the key facts and the document's structure; drop lower-priority detail. " - 'Respond ONLY with JSON: {"sections": [{"heading": "...", "level": 2, ' - '"blocks": ["...", "..."]}]}. A heading carries no "#", and each block is one ' - "paragraph, list, table or code fence." - ) - rewrite_user = f"Target budget: {max_tokens} tokens.\n\nDocument to shorten:\n{answer}" - else: - # The token budget is enforced via the prompt, not a hard provider cap: - # on thinking models a hard cap is eaten by reasoning tokens and would - # truncate the rewrite mid-word (#3365). Cost is bounded by the separate - # reflect_max_completion_tokens config (uncapped by default). - rewrite_system = ( - "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." - ) - rewrite_user = f"Target budget: {max_tokens} tokens.\n\nText to rewrite:\n{answer}" - - call_result = await llm_config.call( - messages=[ - {"role": "system", "content": rewrite_system}, - {"role": "user", "content": rewrite_user}, - ], - scope="reflect", - temperature=get_config().llm_temperature_reflect, - max_completion_tokens=get_config().reflect_max_completion_tokens, - ) - rewritten = call_result.content - rewrite_usage = call_result.usage - if document is not None: - trimmed = _document_from_rewrite(rewritten, answer) - document, answer = trimmed.structure, trimmed.markdown - else: - answer = rewritten.strip() + answer, document, rewrite_usage, rewrite_trace = await _rewrite_final_answer( + answer, + llm_config, + max_tokens, + document, + ) + if rewrite_usage is not None: + assert rewrite_trace is not None final_usage = TokenUsageSummary( input_tokens=usage.input_tokens + rewrite_usage.input_tokens, output_tokens=usage.output_tokens + rewrite_usage.output_tokens, @@ -1439,14 +1487,7 @@ async def _process_done_tool( cached_tokens=usage.cached_tokens + (getattr(rewrite_usage, "cached_tokens", 0) or 0), thoughts_tokens=usage.thoughts_tokens + (getattr(rewrite_usage, "thoughts_tokens", 0) or 0), ) - llm_trace.append( - LLMCall( - scope="final_rewrite", - duration_ms=int((time.time() - rewrite_start) * 1000), - input_tokens=rewrite_usage.input_tokens, - output_tokens=rewrite_usage.output_tokens, - ) - ) + llm_trace.append(rewrite_trace) # 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] diff --git a/hindsight-api-slim/tests/test_reflect_agent.py b/hindsight-api-slim/tests/test_reflect_agent.py index 747a1cc973..8372368b86 100644 --- a/hindsight-api-slim/tests/test_reflect_agent.py +++ b/hindsight-api-slim/tests/test_reflect_agent.py @@ -1012,6 +1012,54 @@ async def test_stop_after_evidence_uses_forced_final_synthesis(self, mock_llm, m final_prompt = mock_llm.call.await_args.kwargs["messages"][1]["content"] assert f"approximately {cap} tokens" in final_prompt + @pytest.mark.asyncio + async def test_forced_final_synthesis_rewrites_over_budget_answer(self, mock_llm, mock_functions, monkeypatch): + config = MagicMock( + reflect_prompt_cache_enabled=False, + reflect_max_completion_tokens=None, + llm_temperature_reflect=0.17, + ) + monkeypatch.setattr("hindsight_api.engine.reflect.agent.get_config", lambda: config) + mock_functions["search_mental_models_fn"].return_value = { + "mental_models": [{"id": "mm-1", "name": "Prefs", "content": "Fresh content.", "is_stale": False}] + } + mock_llm.call_with_tools.side_effect = [ + self._mm_call(), + LLMToolCallResult(tool_calls=[], content="I have enough to answer.", finish_reason="stop"), + ] + mock_llm.call = AsyncMock( + side_effect=[ + LLMCallResult( + content="important detail " * 100, + usage=TokenUsage(input_tokens=40, output_tokens=100, total_tokens=140), + ), + LLMCallResult( + content="Concise final answer.", + usage=TokenUsage(input_tokens=110, output_tokens=4, total_tokens=114), + ), + ] + ) + + result = await run_reflect_agent( + llm_config=mock_llm, + bank_id="test-bank", + query="test query", + bank_profile={"name": "Test", "mission": "Testing"}, + has_mental_models=True, + budget="low", + max_tokens=8, + **mock_functions, + ) + + assert result.text == "Concise final answer." + assert mock_llm.call.await_count == 2 + rewrite_call = mock_llm.call.await_args_list[1] + assert "Target budget: 8 tokens" in rewrite_call.kwargs["messages"][1]["content"] + assert rewrite_call.kwargs["max_completion_tokens"] is None + assert result.llm_trace[-1].scope == "final_rewrite" + assert result.usage.input_tokens == 150 + assert result.usage.output_tokens == 104 + @pytest.mark.asyncio async def test_max_iterations_reached(self, mock_llm, mock_functions): """Test that agent stops after max iterations even with errors."""