fix(recall): add timeout/fallback guard to expand_observations entity CTE - #3724
fix(recall): add timeout/fallback guard to expand_observations entity CTE#3724jordigilh wants to merge 2 commits into
Conversation
… CTE _expand_combined() got a per-entity LATERAL cap and an asyncio.wait_for timeout+fallback in vectorize-io#911, but its sibling _expand_observations() only got the cap. On a large/dense bank the entity CTE can still exceed the configured link_expansion_timeout, and with no guard the raw TimeoutError propagates and fails the whole recall() call instead of degrading to semantic+causal results. Wraps just the entity_rows fetch in expand_observations() with the same asyncio.wait_for/except asyncio.TimeoutError pattern, falling back to an empty entity result set (the semantic+causal query still runs). Observed in production: RuntimeError: Failed to search memories (TimeoutError): TimeoutError() on a bank with ~410K entity_cooccurrences rows.
Strix Security ReviewWarning This pull request has 1 commit after the last Strix review ( No security issues found. Updated for Reviewed by Strix |
Sanderhoff-alt
left a comment
Sanderhoff-alt
left a comment
There was a problem hiding this comment.
Two items to address before merge: the new regression test also passes without the change, and the Oracle implementation of the same query is left unguarded, so Oracle deployments keep the failure this fixes.
| # Recall succeeds even though entity expansion timed out; the | ||
| # seed observation itself is still found via semantic search. | ||
| assert result.results is not None | ||
| assert len(result.results) > 0 |
There was a problem hiding this comment.
This test passes on main as well, so it does not guard the fix. Before this change link_expansion_timeout is read in exactly one place (hindsight_api/engine/search/link_expansion_retrieval.py:345, inside _expand_combined), and fact_type=["observation"] routes to _expand_observations, which never consulted it — so lowering the timeout cannot raise TimeoutError on the old code path. With the guard in place the assertion here is satisfied by the semantic seed regardless of whether entity expansion ran. Please make the test discriminate, for example assert via caplog that the new [ExpandObservations] Entity expansion timed out warning fired and that no entity-expanded rows came back.
There was a problem hiding this comment.
Good catch, confirmed. Added a caplog check for the [ExpandObservations] Entity expansion timed out warning, which only fires once the guard exists and actually runs. Ran the test against main and this branch to confirm: fails on main (no warning), passes here. Pushed in 5403261.
|
|
||
| config = get_config() | ||
| try: | ||
| entity_rows = await asyncio.wait_for( |
There was a problem hiding this comment.
The guard is added to the PostgreSQL ops class only. hindsight_api/engine/db/ops_oracle.py:717 has the same unguarded entity_rows = await conn.fetch(...) inside its expand_observations, so on Oracle 23ai a slow entity CTE still propagates TimeoutError and fails the entire recall (memory_engine.py:7850). The wrapper is plain Python around the fetch with no dialect-specific SQL, so it can be mirrored in the Oracle implementation without an Oracle environment, or factored into a helper both ops classes call.
There was a problem hiding this comment.
Confirmed, same gap. Mirrored the wait_for/except TimeoutError guard into ops_oracle.py's expand_observations, same log message format as the PG fix. Left the two PG-only optimizations you noted alone, they're already explained inline. Pushed in 5403261.
…criminate Address review feedback on vectorize-io#3724: - ops_oracle.py's expand_observations() had the same unguarded entity CTE the PG backend was fixed for, so Oracle recalls could still fail on timeout instead of degrading to semantic+causal results. - The new regression test asserted only that recall succeeds, which also held true before this fix (the seed observation is found via direct semantic match regardless of the guard). It now also asserts via caplog that the timeout warning actually fired, which only happens once the guard exists and is taken.
2 participants
Fixes #3723.
Problem
_expand_combined()has anasyncio.wait_for/fallback guard around its entity CTE (added in #911); its sibling_expand_observations()only got that PR's per-entity LATERAL cap, not the timeout guard. On a large/dense bank the entity query insideexpand_observations()can still exceed budget, and with no guard the rawTimeoutErrorfails the wholerecall()call instead of degrading to semantic+causal results.Fix
Wraps just the
entity_rowsfetch inPostgreSQLOps.expand_observations()with the sameasyncio.wait_for(..., timeout=config.link_expansion_timeout)/except asyncio.TimeoutErrorpattern already used by_expand_combined(), falling back toentity_rows = [](the semantic+causal query still runs independently, unaffected).Test plan
test_expand_observations_falls_back_gracefully_on_timeout, mirroring the existingtest_entity_expansion_timeout_fallbacktechnique (forces the real query to exceed an artificially lowlink_expansion_timeout) for the observation path. Fails with the exact production error (RuntimeError: Failed to search memories (TimeoutError): TimeoutError()) without this fix, passes with it.test_expand_combined_gracefully_falls_back_on_timeoutalongside it as a contrast/sanity check that the sibling path's existing behavior is unchanged.uv run pytest tests/test_expand_observations_timeout_gap.py tests/test_link_expansion_retrieval.py tests/test_observation_expansion_scoring.py tests/test_graph_entity_fanout_cap.pyall green locally (one unrelated pre-existing error in a test requiring a real LLM API key I don't have configured).ruff check --fix/ruff formatclean.Not touched: Oracle's
expand_observations()(ops_oracle.py) has the same shape but I don't have an Oracle test environment to validate a change there — flagging in case a maintainer wants the same guard applied there.