From 77defd96e9069c775eb5e740557e9503f75bf0aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Mon, 2 Mar 2026 12:03:12 +0100 Subject: [PATCH] fix(reflect): prevent context_length_exceeded on large memory banks (#462) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(reflect): prevent context_length_exceeded on large memory banks (#457) The reflect agent's agentic loop accumulated tool-call messages across iterations with no upper bound on token count, causing context_length_exceeded errors on banks with 19K+ nodes. Changes: - Add proactive token-budget guard: before each call_with_tools, count accumulated message tokens via tiktoken; if >= max_context_tokens and evidence has been gathered, immediately synthesize from what was found - Detect context-overflow errors specifically (_is_context_overflow_error) and skip the retry path — retrying after overflow only makes it worse - Truncate context_history in build_final_prompt to a 60K-token budget so the fallback synthesis prompt itself cannot overflow - Add HINDSIGHT_API_REFLECT_MAX_CONTEXT_TOKENS config (default 100000) wired through config.py → main.py → memory_engine → run_reflect_agent - Tests: unit tests for helpers + mock-LLM behavior tests + an end-to-end integration test using a real LLM with max_context_tokens=1 * fix(reflect): derive final prompt context budget from max_context_tokens Replace the hardcoded _FINAL_PROMPT_CONTEXT_BUDGET (60K tokens) with a fraction of max_context_tokens (80%), so the fallback synthesis prompt automatically scales with whatever context window is configured. --- hindsight-api/hindsight_api/config.py | 6 + .../hindsight_api/engine/memory_engine.py | 2 + .../hindsight_api/engine/reflect/agent.py | 114 +++++++++- .../hindsight_api/engine/reflect/prompts.py | 31 ++- hindsight-api/hindsight_api/main.py | 1 + hindsight-api/tests/test_reflect_agent.py | 199 +++++++++++++++++- .../docs/developer/configuration.md | 1 + 7 files changed, 342 insertions(+), 12 deletions(-) diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index 10d6780d..e705430b 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/hindsight-api/hindsight_api/config.py @@ -314,6 +314,7 @@ ENV_WORKER_CONSOLIDATION_MAX_SLOTS = "HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLO # Reflect agent settings ENV_REFLECT_MAX_ITERATIONS = "HINDSIGHT_API_REFLECT_MAX_ITERATIONS" +ENV_REFLECT_MAX_CONTEXT_TOKENS = "HINDSIGHT_API_REFLECT_MAX_CONTEXT_TOKENS" ENV_REFLECT_MISSION = "HINDSIGHT_API_REFLECT_MISSION" # Disposition settings @@ -453,6 +454,7 @@ DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS = 2 # Max concurrent consolidation tasks # Reflect agent settings DEFAULT_REFLECT_MAX_ITERATIONS = 10 # Max tool call iterations before forcing response +DEFAULT_REFLECT_MAX_CONTEXT_TOKENS = 100_000 # Max accumulated context tokens before forcing final prompt # Disposition defaults (None = not set, fall back to bank DB value or 3) DEFAULT_DISPOSITION_SKEPTICISM = None @@ -720,6 +722,7 @@ class HindsightConfig: # Reflect agent settings reflect_max_iterations: int + reflect_max_context_tokens: int # OpenTelemetry tracing configuration otel_traces_enabled: bool @@ -1134,6 +1137,9 @@ class HindsightConfig: ), # Reflect agent settings reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))), + reflect_max_context_tokens=int( + os.getenv(ENV_REFLECT_MAX_CONTEXT_TOKENS, str(DEFAULT_REFLECT_MAX_CONTEXT_TOKENS)) + ), reflect_mission=os.getenv(ENV_REFLECT_MISSION) or None, # Disposition settings (None = fall back to DB value) disposition_skepticism=int(os.getenv(ENV_DISPOSITION_SKEPTICISM)) diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index 484d4809..22e888f5 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -4572,6 +4572,7 @@ class MemoryEngine(MemoryEngineInterface): budget_multipliers = {Budget.LOW: 0.5, Budget.MID: 1.0, Budget.HIGH: 2.0} effective_budget = budget or Budget.LOW max_iterations = max(1, int(base_max_iterations * budget_multipliers.get(effective_budget, 1.0))) + max_context_tokens = config.reflect_max_context_tokens # Run agentic loop - acquire connections only when needed for DB operations # (not held during LLM calls which can be slow) @@ -4681,6 +4682,7 @@ class MemoryEngine(MemoryEngineInterface): directives=directives, has_mental_models=has_mental_models, budget=effective_budget, + max_context_tokens=max_context_tokens, ) total_time = time.time() - reflect_start diff --git a/hindsight-api/hindsight_api/engine/reflect/agent.py b/hindsight-api/hindsight_api/engine/reflect/agent.py index fd3353a6..cdac3cdd 100644 --- a/hindsight-api/hindsight_api/engine/reflect/agent.py +++ b/hindsight-api/hindsight_api/engine/reflect/agent.py @@ -14,6 +14,8 @@ import re import time from typing import TYPE_CHECKING, Any, Awaitable, Callable +import tiktoken + from .models import DirectiveInfo, LLMCall, ReflectAgentResult, TokenUsageSummary, ToolCall from .prompts import FINAL_SYSTEM_PROMPT, _extract_directive_rules, build_final_prompt, build_system_prompt_for_tools from .tools_schema import get_reflect_tools @@ -259,6 +261,46 @@ OUTPUT:""" return None, 0, 0 +_TIKTOKEN_ENCODING = tiktoken.get_encoding("cl100k_base") + + +def _count_messages_tokens(messages: list[dict[str, Any]]) -> int: + """Estimate the token count of the messages list using cl100k_base encoding.""" + total = 0 + for msg in messages: + content = msg.get("content") or "" + if isinstance(content, str): + total += len(_TIKTOKEN_ENCODING.encode(content)) + elif isinstance(content, list): + for part in content: + if isinstance(part, dict) and isinstance(part.get("text"), str): + total += len(_TIKTOKEN_ENCODING.encode(part["text"])) + # Tool call arguments and results also count + for tc in msg.get("tool_calls") or []: + if isinstance(tc, dict): + func = tc.get("function", {}) + total += len(_TIKTOKEN_ENCODING.encode(func.get("arguments", ""))) + return total + + +def _is_context_overflow_error(exc: Exception) -> bool: + """Return True if the exception signals the LLM context window was exceeded.""" + msg = str(exc).lower() + return any( + phrase in msg + for phrase in ( + "context_length_exceeded", + "context length exceeded", + "maximum context length", + "prompt_too_long", + "prompt is too long", + "resource_exhausted", + "input is too long", + "too many tokens", + ) + ) + + async def run_reflect_agent( llm_config: "LLMProvider", bank_id: str, @@ -275,6 +317,7 @@ async def run_reflect_agent( directives: list[dict[str, Any]] | None = None, has_mental_models: bool = False, budget: str | None = None, + max_context_tokens: int = 100_000, ) -> ReflectAgentResult: """ Execute the reflect agent loop using native tool calling. @@ -388,7 +431,7 @@ async def run_reflect_agent( if is_last: # Force text response on last iteration - no tools - prompt = build_final_prompt(query, context_history, bank_profile, context) + prompt = build_final_prompt(query, context_history, bank_profile, context, max_context_tokens=max_context_tokens) llm_start = time.time() response, usage = await llm_config.call( messages=[ @@ -433,6 +476,60 @@ async def run_reflect_agent( directives_applied=directives_applied, ) + # Proactive context-window guard: if accumulated messages would exceed the + # configured token budget, bail out early and synthesize from what we have. + estimated_tokens = _count_messages_tokens(messages) + if estimated_tokens >= max_context_tokens and ( + bool(available_memory_ids) or bool(available_mental_model_ids) or bool(available_observation_ids) + ): + logger.warning( + f"[REFLECT {reflect_id}] Context budget exceeded on iteration {iteration + 1}: " + f"~{estimated_tokens} tokens >= {max_context_tokens} limit. Forcing final synthesis." + ) + prompt = build_final_prompt(query, context_history, bank_profile, context, max_context_tokens=max_context_tokens) + llm_start = time.time() + response, usage = await llm_config.call( + messages=[ + {"role": "system", "content": FINAL_SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ], + scope="reflect", + max_completion_tokens=max_tokens, + return_usage=True, + ) + llm_duration = int((time.time() - llm_start) * 1000) + total_input_tokens += usage.input_tokens + total_output_tokens += usage.output_tokens + llm_trace.append( + { + "scope": "final", + "duration_ms": llm_duration, + "input_tokens": usage.input_tokens, + "output_tokens": usage.output_tokens, + } + ) + answer = _clean_answer_text(response.strip()) + + structured_output = None + if response_schema and answer: + structured_output, struct_in, struct_out = await _generate_structured_output( + answer, response_schema, llm_config, reflect_id + ) + total_input_tokens += struct_in + total_output_tokens += struct_out + + _log_completion(answer, iteration + 1, forced=True) + return ReflectAgentResult( + text=answer, + structured_output=structured_output, + iterations=iteration + 1, + tools_called=total_tools_called, + tool_trace=tool_trace, + llm_trace=_get_llm_trace(), + usage=_get_usage(), + directives_applied=directives_applied, + ) + # Call LLM with tools llm_start = time.time() @@ -478,13 +575,20 @@ async def run_reflect_agent( consecutive_errors += 1 logger.warning(f"[REFLECT {reflect_id}] LLM error on iteration {iteration + 1}: {e} ({err_duration}ms)") llm_trace.append({"scope": f"agent_{iteration + 1}_err", "duration_ms": err_duration}) - # Guardrail: If no evidence gathered yet, retry (but cap consecutive errors to avoid long hangs) has_gathered_evidence = ( bool(available_memory_ids) or bool(available_mental_model_ids) or bool(available_observation_ids) ) - if not has_gathered_evidence and iteration < max_iterations - 1 and consecutive_errors < 2: + # Context overflow errors must never be retried — retrying would only make them worse. + # Skip straight to final synthesis with whatever evidence we have. + if _is_context_overflow_error(e): + logger.warning( + f"[REFLECT {reflect_id}] Context window exceeded on iteration {iteration + 1}, " + "forcing final synthesis from gathered evidence." + ) + # For other errors: retry if no evidence yet (but cap consecutive errors to avoid long hangs) + elif not has_gathered_evidence and iteration < max_iterations - 1 and consecutive_errors < 2: continue - prompt = build_final_prompt(query, context_history, bank_profile, context) + prompt = build_final_prompt(query, context_history, bank_profile, context, max_context_tokens=max_context_tokens) llm_start = time.time() response, usage = await llm_config.call( messages=[ @@ -555,7 +659,7 @@ async def run_reflect_agent( directives_applied=directives_applied, ) # Empty response, force final - prompt = build_final_prompt(query, context_history, bank_profile, context) + prompt = build_final_prompt(query, context_history, bank_profile, context, max_context_tokens=max_context_tokens) llm_start = time.time() response, usage = await llm_config.call( messages=[ diff --git a/hindsight-api/hindsight_api/engine/reflect/prompts.py b/hindsight-api/hindsight_api/engine/reflect/prompts.py index 08aebee8..43faffe9 100644 --- a/hindsight-api/hindsight_api/engine/reflect/prompts.py +++ b/hindsight-api/hindsight_api/engine/reflect/prompts.py @@ -10,6 +10,14 @@ The reflect agent uses hierarchical retrieval: import json from typing import Any +import tiktoken + +_TIKTOKEN_ENCODING = tiktoken.get_encoding("cl100k_base") + +# Fraction of max_context_tokens reserved for tool results in the final synthesis prompt. +# The remainder covers the system prompt, question, bank context, and output tokens. +_FINAL_PROMPT_CONTEXT_FRACTION = 0.8 + def _extract_directive_rules(directives: list[dict[str, Any]]) -> list[str]: """Extract directive rules as a list of strings.""" @@ -394,6 +402,7 @@ def build_final_prompt( context_history: list[dict], bank_profile: dict, additional_context: str | None = None, + max_context_tokens: int = 100_000, ) -> str: """Build the final prompt when forcing a text response (no tools).""" parts = [] @@ -423,18 +432,32 @@ def build_final_prompt( if additional_context: parts.append(f"\n## Additional Context\n{additional_context}") - # Tool call history + # Tool call history — include as many entries as fit within the token budget, + # preferring the most recent calls (they tend to be the most targeted). if context_history: parts.append("\n## Retrieved Data (synthesize and reason from this data)") - for entry in context_history: + token_budget = int(max_context_tokens * _FINAL_PROMPT_CONTEXT_FRACTION) + # Render entries newest-first, then reverse so the prompt reads chronologically. + rendered: list[str] = [] + truncated = False + for entry in reversed(context_history): tool = entry["tool"] output = entry["output"] - # Format as proper JSON for LLM readability try: output_str = json.dumps(output, indent=2, default=str) except (TypeError, ValueError): output_str = str(output) - parts.append(f"\n### From {tool}:\n```json\n{output_str}\n```") + block = f"\n### From {tool}:\n```json\n{output_str}\n```" + block_tokens = len(_TIKTOKEN_ENCODING.encode(block)) + if block_tokens > token_budget: + truncated = True + break + rendered.append(block) + token_budget -= block_tokens + for block in reversed(rendered): + parts.append(block) + if truncated: + parts.append("\n*Note: Some earlier tool results were omitted to stay within the context window.*") else: parts.append("\n## Retrieved Data\nNo data was retrieved.") diff --git a/hindsight-api/hindsight_api/main.py b/hindsight-api/hindsight_api/main.py index 17bf4d46..6ff01b31 100644 --- a/hindsight-api/hindsight_api/main.py +++ b/hindsight-api/hindsight_api/main.py @@ -292,6 +292,7 @@ def main(): worker_max_slots=config.worker_max_slots, worker_consolidation_max_slots=config.worker_consolidation_max_slots, reflect_max_iterations=config.reflect_max_iterations, + reflect_max_context_tokens=config.reflect_max_context_tokens, reflect_mission=config.reflect_mission, disposition_skepticism=config.disposition_skepticism, disposition_literalism=config.disposition_literalism, diff --git a/hindsight-api/tests/test_reflect_agent.py b/hindsight-api/tests/test_reflect_agent.py index d867b44d..4acee0ac 100644 --- a/hindsight-api/tests/test_reflect_agent.py +++ b/hindsight-api/tests/test_reflect_agent.py @@ -7,14 +7,17 @@ These tests verify: 3. Recovery from tool execution errors """ +from unittest.mock import AsyncMock, MagicMock + import pytest -from unittest.mock import AsyncMock, MagicMock, patch from hindsight_api.engine.reflect.agent import ( - _normalize_tool_name, - _is_done_tool, _clean_answer_text, _clean_done_answer, + _count_messages_tokens, + _is_context_overflow_error, + _is_done_tool, + _normalize_tool_name, run_reflect_agent, ) from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage @@ -412,3 +415,193 @@ class TestReflectAgentMocked: # Should have a result even if no memories found assert result is not None assert result.iterations == 3 + + +class TestContextOverflowHelpers: + """Unit tests for context-overflow detection helpers.""" + + def test_count_messages_tokens_basic(self): + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is the capital of France?"}, + ] + count = _count_messages_tokens(messages) + assert count > 0 + # Rough sanity check: ~10 tokens for each message + assert count < 100 + + def test_count_messages_tokens_with_tool_result(self): + """A large tool result should substantially increase the count.""" + small_messages = [{"role": "user", "content": "hi"}] + large_messages = [ + {"role": "user", "content": "hi"}, + { + "role": "tool", + "tool_call_id": "x", + "name": "recall", + "content": '{"memories": [' + ', '.join([f'{{"id": "m{i}", "content": "A long memory fact about some topic that goes on and on."}}' for i in range(50)]) + ']}', + }, + ] + small = _count_messages_tokens(small_messages) + large = _count_messages_tokens(large_messages) + assert large > small + 200 + + def test_is_context_overflow_error_openai(self): + assert _is_context_overflow_error(Exception("context_length_exceeded: too many tokens")) + assert _is_context_overflow_error(Exception("This model's maximum context length is 128000 tokens. However, your messages resulted in 142164 tokens.")) + + def test_is_context_overflow_error_anthropic(self): + assert _is_context_overflow_error(Exception("prompt_too_long")) + assert _is_context_overflow_error(Exception("prompt is too long for this model")) + + def test_is_context_overflow_error_gemini(self): + assert _is_context_overflow_error(Exception("RESOURCE_EXHAUSTED: quota exceeded")) + + def test_is_context_overflow_error_generic(self): + assert _is_context_overflow_error(Exception("input is too long to process")) + assert _is_context_overflow_error(Exception("too many tokens in the request")) + + def test_is_context_overflow_error_unrelated(self): + assert not _is_context_overflow_error(Exception("connection timeout")) + assert not _is_context_overflow_error(Exception("rate limit exceeded")) + assert not _is_context_overflow_error(ValueError("invalid argument")) + + +class TestContextOverflowBehavior: + """Test that the reflect agent handles context overflow gracefully.""" + + @pytest.fixture + def mock_llm(self): + llm = MagicMock() + llm.call_with_tools = AsyncMock() + llm.call = AsyncMock( + return_value=("Synthesized answer from gathered evidence.", TokenUsage(input_tokens=50, output_tokens=20, total_tokens=70)) + ) + return llm + + @pytest.fixture + def mock_functions_with_large_output(self): + """Mock functions that return a large enough payload to exceed a tiny token budget.""" + large_memories = [ + {"id": f"mem-{i}", "content": f"Memory fact number {i}: " + "A" * 200} + for i in range(20) + ] + return { + "search_mental_models_fn": AsyncMock(return_value={"mental_models": []}), + "search_observations_fn": AsyncMock(return_value={"observations": []}), + "recall_fn": AsyncMock(return_value={"memories": large_memories}), + "expand_fn": AsyncMock(return_value={"memories": []}), + } + + @pytest.mark.asyncio + async def test_proactive_guard_fires_when_budget_exceeded(self, mock_llm, mock_functions_with_large_output): + """When token count exceeds max_context_tokens after a tool call, the agent + should immediately synthesize from gathered evidence instead of making + another LLM call that would overflow.""" + # First call: LLM calls recall (forced by iter 0 with no mental models) + mock_llm.call_with_tools.return_value = LLMToolCallResult( + tool_calls=[LLMToolCall(id="1", name="recall", arguments={"query": "test"})], + finish_reason="tool_calls", + ) + + # Set a tiny token budget — the recall result alone will blow past it + result = await run_reflect_agent( + llm_config=mock_llm, + bank_id="test-bank", + query="What do you know?", + bank_profile={"name": "Test", "mission": "Testing"}, + max_context_tokens=100, + **mock_functions_with_large_output, + ) + + assert result.text == "Synthesized answer from gathered evidence." + # call_with_tools was called once (for the forced recall), then the guard + # kicked in — no further tool-call iterations + assert mock_llm.call_with_tools.call_count == 1 + # llm.call() was invoked to generate the final synthesis + mock_llm.call.assert_called_once() + + @pytest.mark.asyncio + async def test_context_overflow_error_skips_retry(self, mock_llm, mock_functions_with_large_output): + """A context_length_exceeded error from the LLM should NOT be retried — + it should immediately fall back to final synthesis.""" + mock_llm.call_with_tools.side_effect = Exception( + "context_length_exceeded: messages resulted in 150000 tokens." + ) + + result = await run_reflect_agent( + llm_config=mock_llm, + bank_id="test-bank", + query="What do you know?", + bank_profile={"name": "Test", "mission": "Testing"}, + max_iterations=5, + **mock_functions_with_large_output, + ) + + assert result is not None + # Should have attempted only 1 iteration (no retry on overflow error) + assert mock_llm.call_with_tools.call_count == 1 + # Final synthesis was called + mock_llm.call.assert_called_once() + + +class TestContextOverflowIntegration: + """Integration test: real LLM with a very small max_context_tokens. + + The agent will make one real LLM call (forced tool choice), receive a large + tool result that exceeds the tiny budget, then synthesize from it via a second + real LLM call — all without raising a context_length_exceeded error. + """ + + @pytest.mark.asyncio + async def test_reflect_completes_with_tiny_context_budget(self, memory, request_context): + """End-to-end: reflect on a bank with max_context_tokens=1 (tiny budget). + + Setting max_context_tokens=1 guarantees the proactive guard fires as soon + as the first tool result is received and evidence is available. + The result must be a non-empty string with no exception raised. + """ + import uuid + from unittest.mock import patch + + bank_id = f"test-ctx-overflow-{uuid.uuid4().hex[:8]}" + try: + # Retain a handful of facts so the recall tool has something to return + await memory.retain_async( + bank_id=bank_id, + content="Alice is a software engineer who enjoys hiking on weekends.", + request_context=request_context, + ) + await memory.retain_async( + bank_id=bank_id, + content="Bob is a designer who loves cooking Italian food.", + request_context=request_context, + ) + + # Patch get_config where memory_engine uses it, injecting a tiny + # max_context_tokens. Everything else delegates to the real config. + real_config = memory._get_raw_config() if hasattr(memory, "_get_raw_config") else None + from hindsight_api.config import get_config as _real_get_config + + class _TinyContextProxy: + """Forwards all attribute access to the real config proxy except + reflect_max_context_tokens which is forced to 1.""" + _real = _real_get_config() + + def __getattr__(self, name: str): + if name == "reflect_max_context_tokens": + return 1 + return getattr(self._real, name) + + with patch("hindsight_api.engine.memory_engine.get_config", return_value=_TinyContextProxy()): + result = await memory.reflect_async( + bank_id=bank_id, + query="Tell me about the people you know.", + request_context=request_context, + ) + + assert result.text, "reflect must return a non-empty answer" + assert result.usage.total_tokens > 0 + + finally: + await memory.delete_bank(bank_id, request_context=request_context) diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index cb5aea1b..f003ae9b 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -780,6 +780,7 @@ export HINDSIGHT_API_OBSERVATIONS_MISSION="Observations are recurring patterns i | Variable | Description | Default | |----------|-------------|---------| | `HINDSIGHT_API_REFLECT_MAX_ITERATIONS` | Max tool call iterations before forcing a response | `10` | +| `HINDSIGHT_API_REFLECT_MAX_CONTEXT_TOKENS` | Max accumulated context tokens in the reflect loop before forcing final synthesis. Prevents `context_length_exceeded` errors on large banks. Lower this if your LLM has a context window smaller than 128K. | `100000` | | `HINDSIGHT_API_REFLECT_MISSION` | Global reflect mission (identity and reasoning framing). Overridden per bank via config API. | - | #### Disposition