Skip to content

feat(reflect): step-by-step context caching for the Gemini tool loop#2540

Open
nicoloboschi wants to merge 1 commit into
mainfrom
feat/reflect-step-cache
Open

feat(reflect): step-by-step context caching for the Gemini tool loop#2540
nicoloboschi wants to merge 1 commit into
mainfrom
feat/reflect-step-cache

Conversation

@nicoloboschi

Copy link
Copy Markdown
Collaborator

Summary

Reflect's LLM cost is dominated by the tool-result context re-sent on every turn, not the static system prefix. Caching only the fixed system+tools prefix (Gemini CachedContent) recovered ~9% of input on a short reflect; Gemini implicit caching gave 0% on short loops (and is unreliable in general).

This PR rolls a step-by-step context cache forward through the agent's tool loop: after each turn the CachedContent is extended to cover that turn's full input, so the next auto turn reuses the entire prior conversation at the cached-input rate and sends only its own new tool results as the delta. Each retrieved payload is billed at full price exactly once, then cached thereafter.

Measured (gemini-2.5-flash-lite)

Loop Implicit only Old static prefix This PR
Short (3 calls) 0% 9% ~29%
Deep (8-11 calls) ~70% (unreliable) ~10% ~74-81%

On deep loops the per-call ratio ramps to ~99% on the deepest turns (each turn reuses nearly the whole prior context). Unlike implicit caching, this is deterministic and works on short loops too.

How it works

  • GeminiCacheManager: caches contents (the growing conversation), with session-tracked create/delete for ephemeral per-reflect caches.
  • call_with_tools gains cached_prefix_message_count so only the un-cached tail is sent. Forced (tool_config) retrieval turns stay uncached — Gemini rejects cached_content + tool_config — but the cache still advances underneath them, so the first auto turn inherits a cache covering the whole forced phase.
  • Latency hidden: the next-turn cache create is scheduled to run concurrently with tool execution and only awaited right before the next auto call.
  • Ephemeral: the reflect wrapper deletes the whole cache session on every exit path (answer / error / cancellation); a short TTL is the backstop. Verified no cache leak.
  • Flag: HINDSIGHT_API_REFLECT_PROMPT_CACHE_ENABLED (default true) disables reflect caching independently of the global prompt cache.

Tests

  • Deterministic step-by-step contract test (each auto turn reuses the previous turn's full input; caches torn down once) — no real LLM.
  • GeminiCacheManager unit tests (create/track/delete/soft-fallback/error-swallow), no network.
  • Real-Gemini cached-ratio measurement test updated with a real floor (gated behind HINDSIGHT_RUN_GEMINI_EVALS).

Not an LLM-behaviour change (no prompt/classification changes), so no judge test needed.

Reflect's cost is dominated by the tool-result context re-sent on every
turn, not the static system prefix. Caching only the fixed system+tools
prefix (Gemini CachedContent) recovered ~9% of input on a short reflect
and 0% via implicit caching.

Instead, roll a cache forward one step at a time: after each turn extend
the CachedContent to cover that turn's full input, so the next `auto`
turn reuses the entire prior conversation at the cached-input rate and
sends only its own new tool results as the delta. Each retrieved payload
is billed fresh exactly once, then cached thereafter. Measured on
gemini-2.5-flash-lite: short loops ~29%, deep loops ~74-81% cached input
(deepest turns ~99%).

- GeminiCacheManager: cache `contents` (growing conversation), session-
  tracked create/delete for ephemeral per-reflect caches.
- call_with_tools gains `cached_prefix_message_count` so only the un-cached
  tail is sent; forced (tool_config) turns stay uncached as Gemini requires.
- The next-turn cache create is scheduled to run concurrently with tool
  execution, hiding its latency; the reflect wrapper deletes the session
  (all caches) on every exit path, with a short TTL as backstop.
- New HINDSIGHT_API_REFLECT_PROMPT_CACHE_ENABLED flag (default true) to
  disable reflect caching independently of the global prompt cache.

Tests: deterministic step-by-step contract test, cache-manager units,
and a real-Gemini cached-ratio measurement test.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant