fix(retain): a zero retry budget must still perform the initial fact-extraction request

Author: chethanuk

hindsight-api-slim/hindsight_api/engine/retain/fact_extraction.py

@@ -1362,10 +1362,15 @@ async def _extract_facts_from_chunk(
     llm_max_retries = (
         config.retain_llm_max_retries if config.retain_llm_max_retries is not None else config.llm_max_retries
     )
+    # Legacy OUTER content-validation attempt count (re-prompts on malformed JSON),
+    # floored at 1 so a zero transport-retry budget still performs the initial request
+    # (#2731). The raw budget is still forwarded unchanged to llm_config.call(), which
+    # owns transport retries (`range(N + 1)` requests internally) — do NOT add +1 here.
+    outer_attempts = max(1, llm_max_retries)
     last_error: Exception | None = None
 
     usage = TokenUsage()  # Track cumulative usage across retries
-    for attempt in range(llm_max_retries):
+    for attempt in range(outer_attempts):
         try:
             initial_backoff = (
                 config.retain_llm_initial_backoff
@@ -1401,9 +1406,9 @@ async def _extract_facts_from_chunk(
             # Handle malformed LLM responses
             coerced_response_json = _coerce_fact_response(extraction_response_json)
             if coerced_response_json is None:
-                if attempt < llm_max_retries - 1:
+                if attempt < outer_attempts - 1:
                     logger.warning(
-                        f"LLM returned non-dict JSON on attempt {attempt + 1}/{llm_max_retries}: {type(extraction_response_json).__name__}. Retrying..."
+                        f"LLM returned non-dict JSON on attempt {attempt + 1}/{outer_attempts}: {type(extraction_response_json).__name__}. Retrying..."
                     )
                     continue
                 else:
@@ -1412,7 +1417,7 @@ async def _extract_facts_from_chunk(
                     # worker's retry machinery and ultimately fails loudly — never
                     # silently commit the document with 0 facts. See issue #1833.
                     raise RuntimeError(
-                        f"Fact extraction failed: LLM returned non-dict JSON after {llm_max_retries} attempts "
+                        f"Fact extraction failed: LLM returned non-dict JSON after {outer_attempts} attempts "
                         f"({type(extraction_response_json).__name__}). Raw: {str(extraction_response_json)[:500]}"
                     )
             extraction_response_json = coerced_response_json
@@ -1631,9 +1636,9 @@ def get_value(field_name):
                     continue
 
             # If we got malformed facts and haven't exhausted retries, try again
-            if has_malformed_facts and len(chunk_facts) < len(raw_facts) * 0.8 and attempt < llm_max_retries - 1:
+            if has_malformed_facts and len(chunk_facts) < len(raw_facts) * 0.8 and attempt < outer_attempts - 1:
                 logger.warning(
-                    f"Got {len(raw_facts) - len(chunk_facts)} malformed facts out of {len(raw_facts)} on attempt {attempt + 1}/{llm_max_retries}. Retrying..."
+                    f"Got {len(raw_facts) - len(chunk_facts)} malformed facts out of {len(raw_facts)} on attempt {attempt + 1}/{outer_attempts}. Retrying..."
                 )
                 continue
 
@@ -1672,7 +1677,7 @@ def get_value(field_name):
     # If we exhausted all retries, raise the last error or a descriptive fallback
     if last_error is not None:
         raise last_error
-    raise RuntimeError(f"Fact extraction failed after {llm_max_retries} attempts: LLM did not return valid JSON")
+    raise RuntimeError(f"Fact extraction failed after {outer_attempts} attempts: LLM did not return valid JSON")
 
 
 async def _extract_facts_with_auto_split(

hindsight-api-slim/tests/test_fact_extraction_retry.py

@@ -263,6 +263,40 @@ async def test_retain_llm_max_retries_overrides_global():
     assert llm_config.call.call_count == 5
 
 
+@pytest.mark.asyncio
+async def test_zero_retry_budget_performs_single_chunk_extraction_call():
+    """
+    Direct _extract_facts_from_chunk with a retry budget of 0 (issue #2731):
+    the outer loop must still run once, and the RAW budget (0) must reach the
+    provider so it stays the single owner of transport retries.
+    """
+    from hindsight_api.engine.retain.fact_extraction import _extract_facts_from_chunk
+
+    config = _make_config(llm_max_retries=3, retain_llm_max_retries=0)
+    llm_config = _make_llm_config(
+        mock_response={"facts": [{"what": "Alice visited Paris", "when": "2023", "who": "Alice", "why": "vacation"}]}
+    )
+
+    with patch(
+        "hindsight_api.engine.retain.fact_extraction._build_extraction_prompt_and_schema",
+        return_value=("system prompt", MagicMock()),
+    ):
+        facts, _usage = await _extract_facts_from_chunk(
+            chunk="Alice visited Paris in 2023.",
+            chunk_index=0,
+            total_chunks=1,
+            event_date=datetime(2023, 1, 1, tzinfo=timezone.utc),
+            context="",
+            llm_config=llm_config,
+            config=config,
+            agent_name="test-agent",
+        )
+
+    assert llm_config.call.call_count == 1
+    assert llm_config.call.call_args.kwargs["max_retries"] == 0
+    assert len(facts) == 1
+
+
 @pytest.mark.asyncio
 async def test_none_event_date_with_empty_facts_no_crash():
     """
@@ -376,3 +410,128 @@ def test_build_request_body_omits_temperature_when_none():
 
     body = _build_request_body(_make_batch_llm_config(), _make_batch_temp_config(None), "sys", "user", dict)
     assert "temperature" not in body
+
+
+# --- Retry budget semantics (issue #2731) -----------------------------------
+#
+# A retry BUDGET of N means N retries *after* the initial request — the meaning
+# every provider already implements (`for attempt in range(max_retries + 1)`) and
+# the meaning the OpenAI SDK documents for `max_retries=0` ("disable retries",
+# i.e. one request). The tests below drive the *public* extract_facts_from_text
+# entry point with the real HindsightConfig an operator gets from
+# HINDSIGHT_API_RETAIN_LLM_MAX_RETRIES, so they exercise the whole
+# env -> config -> chunk -> auto-split -> extraction -> provider chain rather
+# than a private helper.
+
+_VALID_EXTRACTION_RESPONSE = {
+    "facts": [{"what": "Alice visited Paris", "when": "2023", "who": "Alice", "why": "vacation"}]
+}
+
+
+def _make_recording_llm(mock_response):
+    """LLMProvider double returning ``mock_response`` and recording call kwargs."""
+    from hindsight_api.engine.llm_wrapper import LLMProvider
+    from hindsight_api.engine.response_models import TokenUsage
+
+    llm = MagicMock(spec=LLMProvider)
+    llm.provider = "mock"
+    llm.call = AsyncMock(return_value=(mock_response, TokenUsage()))
+    return llm
+
+
+@pytest.fixture
+def retain_config(monkeypatch):
+    """Factory for the real config an operator gets from the retry-budget env var.
+
+    Mirrors the reporter's setup: HINDSIGHT_API_RETAIN_LLM_MAX_RETRIES is the only
+    knob they touch. The global budget is pinned so the "unset" row provably
+    exercises the fallback.
+    """
+    from hindsight_api.config import _get_raw_config, clear_config_cache
+
+    def _build(retain_budget: str | None):
+        monkeypatch.setenv("HINDSIGHT_API_LLM_MAX_RETRIES", "3")
+        if retain_budget is None:
+            monkeypatch.delenv("HINDSIGHT_API_RETAIN_LLM_MAX_RETRIES", raising=False)
+        else:
+            monkeypatch.setenv("HINDSIGHT_API_RETAIN_LLM_MAX_RETRIES", retain_budget)
+        clear_config_cache()
+        return _get_raw_config()
+
+    yield _build
+    clear_config_cache()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "retain_budget, expected_forwarded_retries",
+    [
+        # The reported repro: a gateway owns transport retries, so the operator
+        # sets the budget to 0. This used to perform ZERO extraction requests and
+        # raise "Fact extraction failed after 0 attempts".
+        pytest.param("0", 0, id="zero_budget_gateway_owns_retries"),
+        pytest.param("1", 1, id="budget_one"),
+        pytest.param("3", 3, id="budget_three"),
+        pytest.param(None, 3, id="unset_falls_back_to_global"),
+    ],
+)
+async def test_retry_budget_always_performs_initial_extraction_request(
+    retain_config, retain_budget, expected_forwarded_retries
+):
+    """Any retry budget — including 0 — must still perform the initial request."""
+    from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
+
+    config = retain_config(retain_budget)
+    llm = _make_recording_llm(_VALID_EXTRACTION_RESPONSE)
+
+    facts, _chunks, _usage = await extract_facts_from_text(
+        text="Alice visited Paris in 2023.",
+        event_date=datetime(2023, 1, 1, tzinfo=timezone.utc),
+        llm_config=llm,
+        agent_name="test-agent",
+        config=config,
+        context="",
+    )
+
+    assert llm.call.call_count == 1
+    assert len(facts) == 1
+    assert "Alice visited Paris" in facts[0].fact
+    # The RAW budget reaches the provider — not the outer attempt count — so the
+    # provider stays the single retry owner (0 => gateway owns transport retries).
+    assert llm.call.call_args.kwargs["max_retries"] == expected_forwarded_retries
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "retain_budget, expected_calls",
+    [
+        # Budget 0 means "one request, no retry" — one attempt, then a real
+        # content error. Never zero attempts.
+        pytest.param("0", 1, id="zero_budget_one_attempt"),
+        pytest.param("1", 1, id="budget_one"),
+        pytest.param("3", 3, id="budget_three"),
+    ],
+)
+async def test_malformed_response_still_attempts_then_fails_loudly(retain_config, retain_budget, expected_calls):
+    """A malformed response must fail on content, never on a skipped request.
+
+    Guards the #1833 contract (raise, never silently return []) while proving the
+    zero budget spends its one attempt before failing.
+    """
+    from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
+
+    config = retain_config(retain_budget)
+    llm = _make_recording_llm(["not a dict"])
+
+    with pytest.raises(RuntimeError, match="non-dict JSON") as exc:
+        await extract_facts_from_text(
+            text="Alice visited Paris in 2023.",
+            event_date=datetime(2023, 1, 1, tzinfo=timezone.utc),
+            llm_config=llm,
+            agent_name="test-agent",
+            config=config,
+            context="",
+        )
+
+    assert llm.call.call_count == expected_calls
+    assert "after 0 attempts" not in str(exc.value)