Compare commits

...
Author SHA1 Message Date
Nicolò Boschi 2ff5d7acb2 Merge remote-tracking branch 'origin/main' into test-gemini-cache-1936 2026-06-04 13:56:24 +02:00
Nicolò Boschi 1910a5fb50 test: fix consolidation-scope-parallelism mock + metrics counter count
- test_consolidation_scope_parallelism.py: the inline mock read facts from
  messages[0], which is now the (cached) system message after the consolidation
  prompt split — read the user message(s) instead.
- test_metrics.py: mock_meter provided 5 counter mocks but MetricsCollector now
  creates 7 (the cached_input + thoughts counters), so create_counter.side_effect
  ran out (StopIteration at setup). Bump both fixtures to 7.
2026-06-04 13:56:01 +02:00
Nicolò Boschi 7dfebb1147 style: ruff-format the prompt-cache config line (fixes verify-generated-files) 2026-06-04 13:49:37 +02:00
Nicolò Boschi f64a816361 fix(gemini): bound the cache-create call with a timeout
get_or_create holds the manager lock across the caches.create network call, which
correctly dedups concurrent callers (a 10-chunk retain batch produces exactly one
create, not ten). But with no timeout, a hung create would block every waiting
chunk indefinitely. Wrap the create in asyncio.wait_for (30s default, configurable
via create_timeout_seconds); on timeout it soft-fails to None and callers proceed
uncached instead of stalling the batch. Unit test covers the timeout path.
2026-06-04 13:36:11 +02:00
Nicolò Boschi 2954f0d2a5 fix(gemini): fall back to uncached on a cached-request 400
A 400 from a generate request that references a CachedContent (expired/deleted
cache, cross-project mismatch, cache+tool_config incompatibility, ...) was treated
as a generic retryable error: the same cached request was retried, 400'd again,
and the whole operation failed. The soft-fallback only covered cache *creation*,
not the call that *uses* the cache.

Now, on the first 400 while a cache is in use, call()/call_with_tools():
- drop the cache and rebuild the request inline (re-send system prefix + schema/
  tools) so the request still succeeds,
- invalidate the dead cache name (GeminiCacheManager.invalidate) so the next
  operation recreates it instead of reusing the bad name,
- retry immediately (no backoff — it's a config switch, not a transient error).

If the uncached retry also 400s it's a genuine bad request and errors normally.

Supporting fix: system_instruction is now ALWAYS captured from the messages (it
was skipped when cached), so the fallback has the prefix to inline; the config
builder still omits it from the request while the cache carries it. New unit test
covers the 400 → uncached-retry → invalidate path. Cached success path unchanged
(real Gemini retain still 90.8%).
2026-06-04 11:40:32 +02:00
Nicolò Boschi b5da24a1bf feat(llm): generic, default-on prompt caching knob
Rename the Gemini-specific opt-in flag to a provider-agnostic, default-on knob,
modelled on HINDSIGHT_API_RETAIN_BATCH_ENABLED:

- HINDSIGHT_API_LLM_GEMINI_PROMPT_CACHE_ENABLED → HINDSIGHT_API_LLM_PROMPT_CACHE_ENABLED
  (config field llm_gemini_prompt_cache_enabled → llm_prompt_cache_enabled, kwarg
  gemini_prompt_cache_enabled → prompt_cache_enabled), single global knob (not per-op).
- DEFAULT_LLM_PROMPT_CACHE_ENABLED = True. Safe to default on: the cached prefix is
  bank-agnostic (one shared cache) and creation soft-fails to an uncached call, so
  it never breaks a request. Providers that don't implement caching ignore the flag.
- Resolve the flag for every provider (drop the gemini/vertexai restriction) so any
  future provider that implements supports_prompt_caching() picks it up.

Docs: models page now says "on by default; disable with
HINDSIGHT_API_LLM_PROMPT_CACHE_ENABLED=false". The per-operation ratio test sets the
flag explicitly in both modes since the default is now on. Includes the regenerated
skills/hindsight-docs mirror.
2026-06-04 11:27:26 +02:00
Nicolò Boschi 477342ce94 docs(models): drive provider capability table from llmProviders.json
Replace the hand-written capability table with a data-driven one so adding a
provider stays a single-file edit. The capability flags (batchApi, promptCaching)
live in llmProviders.json — the existing single source of truth for the provider
grid and default-models table — and a new LLMProviderCapabilities component (plus
a matching renderer in generate-docs-skill.sh) renders them. Tool-calling dropped
(not differentiating here). Keep flags aligned with supports_batch_api() /
supports_prompt_caching() on the provider classes.
2026-06-04 11:10:34 +02:00
Nicolò Boschi f1fa1904e0 docs(models): add per-provider capability table (batch API, prompt caching)
Adds a "Provider Capabilities" table to the LLM section of the models page
showing which providers support the Batch API (OpenAI/Groq/Fireworks) and
explicit prompt-prefix caching (Gemini/Vertex via CachedContent), with notes on
OpenAI's automatic prefix caching and the bank-agnostic shared-cache design.
Includes the regenerated skills/hindsight-docs mirror.
2026-06-04 11:04:01 +02:00
Nicolò Boschi e6a519dc06 refactor(llm): make prompt-prefix caching a provider-interface feature
Hoist caching out of Gemini-specific duck-typing into the LLMInterface contract,
mirroring supports_batch_api():
- LLMInterface.supports_prompt_caching() -> bool (default False) and
  get_or_create_cached_prefix(...) -> str | None (default None), with docs on how
  explicit-cache (Gemini handle), automatic-cache (OpenAI), and inline-marker
  (Anthropic cache_control) providers each map onto the hook.
- call()/call_with_tools() gain a provider-neutral cached_prefix handle (renamed
  from the Gemini-flavoured cached_content_name); the wrapper forwards it only
  when set so non-caching providers' signatures are untouched.
- GeminiLLM implements supports_prompt_caching(); the retain/consolidation/reflect
  call sites gate on it instead of hasattr().

The engine already decides WHAT is cacheable (bank-agnostic system prefix), so a
new provider only implements HOW — e.g. OpenAI can benefit with no code (stable
leading prefix is auto-cached) or a thin override.
2026-06-04 10:52:35 +02:00
Nicolò Boschi e1ff77954a test(retain): cacheable prefix invariant to per-bank free-text (concise/verbose)
Parametrized over the concise and verbose modes: the cached system prompt must be
byte-identical regardless of the retain mission (any value, incl. JSON/unicode/
long text) and custom instructions, so per-bank free-text can never fragment the
shared Gemini cache. Structural toggles (causal/labels/language) are intentionally
out of scope — they legitimately partition the cache via the fingerprint.
2026-06-04 10:52:04 +02:00
Nicolò Boschi 2fbcf35b54 test(retain): assert different missions yield one shared cache prefix
Extend the mission-relocation test to prove the payoff directly: two banks with
different retain missions produce a byte-identical system prompt → the same cache
fingerprint → a single shared CachedContent instead of one per mission.
2026-06-03 18:20:40 +02:00
Nicolò Boschi 2d56782287 feat(gemini): make cached prefix bank-agnostic (mission → user message)
The retain and consolidation system prompts embedded the per-bank mission, so
each distinct mission produced a different cache fingerprint → one CachedContent
per bank. With many banks/missions that multiplies create + storage cost and
cached-object count, and makes default-on uneconomical.

Move the mission out of the cached prefix into the per-request user message:
- retain: _build_extraction_prompt_and_schema now returns a bank-agnostic prompt;
  the mission rides in the user message via _retain_mission_preamble().
- consolidation: build_consolidation_system_prompt drops the mission param; the
  mission moves into build_consolidation_input (the user message).

Result: the cached prefix is identical across all banks, so a single shared
CachedContent serves every bank — cardinality drops from O(missions) to O(1) per
operation, and the cost-inversion for many-low-volume-bank workloads goes away.

Behavioral note: the mission now appears in the user turn rather than the system
prompt. Validate mission-adherence against the accuracy benchmarks before flipping
the global default on. Tests updated to assert the new location + cross-bank
prefix sharing.
2026-06-03 18:17:43 +02:00
Nicolò Boschi 9b9769baeb perf(consolidation): move stable observation-format note into cached prefix
The "## INPUT FORMAT" boilerplate (the explanation of the observation JSON
shape: id/text/proof_count/occurred_*/source_memories) was re-sent in every
per-batch user message. It's stable, so move it into the cached system prefix
(build_consolidation_system_prompt); the per-batch user message now carries
only the variable facts + observations data. Lifts the cached/input ratio a
couple of points without changing what the model sees.
2026-06-03 17:40:43 +02:00
Nicolò Boschi fc173a6c5e test(gemini): per-operation cached-ratio test + consolidation split coverage
- New tests/test_gemini_implicit_cache_ratio.py: measures cached/input token
  ratio per operation (retain, reflect, consolidation) against real Gemini via
  the LLM-request tracer. Dual mode: default records the implicit-cache baseline
  (~0% for this access pattern); HINDSIGHT_GEMINI_EXPLICIT_CACHE=1 asserts the
  explicit cache engages (cached_tokens > 0, per-op ratio floor). Gated behind
  HINDSIGHT_RUN_GEMINI_EVALS=1 + a Gemini key.
- test_consolidation.py: unit test for the system/user prompt split (cacheable
  byte-stable prefix; data only in the user message). Fix the inline mock LLM
  callbacks to read facts from the user message(s) rather than messages[0], now
  that the stable instructions are a separate system message.
2026-06-03 17:32:42 +02:00
Nicolò Boschi 7a3e42ae74 feat(gemini): cache consolidation prefix + gate reflect cache to auto turns
- Consolidation: split the batch prompt into a stable system instruction
  (mission + rules + decision guide + output format) and a per-batch user
  message (facts + existing observations + capacity note). The system prefix
  is byte-identical across batches in a run, so it is cached and reused; the
  variable data and the per-batch response_schema stay out of the cached
  surface so it never busts. Measures ~30-40% cached/input per batch (the
  remainder is irreducible per-batch data).
- Reflect: Gemini rejects cached_content alongside a per-request tool_config
  ("CachedContent can not be used with ... tool_config"). The forced-retrieval
  iterations set tool_config, so only the `auto` iterations can reference the
  cache. Gate cached_content_name on tool_choice == "auto"; forced iterations
  send the prefix inline.
2026-06-03 17:32:19 +02:00
Nicolò Boschi ec0920e435 fix(gemini): make explicit prompt caching actually work end-to-end
The caching paths could never produce a cache hit:

- CreateCachedContentConfig was given response_schema/response_mime_type,
  which the google-genai SDK forbids (extra_forbidden) — so every cache
  create raised and soft-fell-back to an uncached call. Cache only holds
  system_instruction (+ tools); response_schema is a generation-time
  constraint and stays on the per-request GenerateContentConfig.
- call() dropped response_schema when a cache was in use (assuming the
  schema lived in the cache — impossible). Keep it on the request; only
  system_instruction moves into the cache. Structured output is preserved.
- cached_content_name was plumbed into the leaf GeminiLLM.call /
  call_with_tools but NOT through the LLMProvider wrapper, so the real call
  path raised "unexpected keyword argument 'cached_content_name'". Thread it
  through both wrappers, forwarding only when set (other providers untouched).

With these, retain extraction caches the ~1.7k-token prefix at ~90%.
2026-06-03 17:31:54 +02:00
Nicolò Boschi 08a16f1dde Merge remote-tracking branch 'origin/main' into test-gemini-cache-1936
# Conflicts:
#	hindsight-api-slim/hindsight_api/engine/providers/gemini_llm.py
2026-06-03 16:14:51 +02:00
Chris Bartholomew 9d9ac2b722 feat(gemini): extend context caching to the tool-calling reflect loop
Adds caching support to the agentic tool-loop path. The reflect agent's
``system_prompt + tools`` is stable for the duration of a single reflect
(and across reflects against the same bank), so caching them once and
reusing the cache name across every iteration of the loop collapses the
dominant input cost — the prefix repeated on every turn.

Mechanism
---------
1. ``GeminiCacheManager.fingerprint(...)`` now accepts ``tools`` and
   includes the OpenAI-style tool list in the hash. A loop that swaps a
   tool gets a fresh cache automatically; a loop that doesn't, hits the
   cache deterministically. The tool list is serialised with sort_keys
   so upstream dict-reordering doesn't cause phantom cache misses.

2. ``GeminiCacheManager.get_or_create(...)`` accepts ``tools`` and
   converts the OpenAI-style entries into Gemini ``Tool`` /
   ``FunctionDeclaration`` shapes inside ``CreateCachedContentConfig``.
   The cached prefix now holds system_instruction + tools, so the
   subsequent ``call_with_tools(cached_content_name=...)`` invocation
   skips resending both.

3. ``GeminiLLM.call_with_tools(...)`` gains ``cached_content_name``.
   When set, ``system_instruction`` and ``tools`` are dropped from the
   per-request config (the SDK rejects re-sending them alongside
   ``cached_content``); ``tool_config`` (mode / allowed_function_names)
   stays per-request as it must.

4. ``GeminiLLM.get_or_create_cached_prefix(...)`` accepts ``tools``
   and forwards them to the cache manager.

5. ``reflect/agent.py:run_reflect_agent`` looks up (or creates) the
   cached prefix ONCE per reflect — right after the ``system_prompt``
   and ``tools`` are built — and reuses the returned cache name across
   every iteration of the agentic loop. The lookup is wrapped in a
   try/except so a cache-side failure can never block a reflect.

6. ``call_with_tools`` now extracts ``cached_content_token_count``
   and ``thoughts_token_count`` from ``usage_metadata`` and threads them
   through ``metrics.record_llm_call`` — same as ``call()`` already
   does. Without this the new ``hindsight.llm.tokens.cached_input`` and
   ``hindsight.llm.tokens.thoughts`` counters would never report the
   reflect-side share of cached/thinking tokens.

Tests (3 new on top of the 12 from earlier on this branch)
----------------------------------------------------------
- ``test_fingerprint_changes_with_tools``: adding a tool changes the
  fingerprint so a loop that adds a tool gets a fresh cache.
- ``test_fingerprint_stable_under_dict_reordering``: dict-key order in
  the OpenAI-style tools list does NOT change the fingerprint.
- ``test_get_or_create_passes_tools_to_create``: the ``caches.create``
  call actually receives the tools in its config — without this the
  cache would silently lack the tool definitions and the first
  ``call_with_tools(cached_content_name=...)`` would 400.

Verification
------------
- ``uv run pytest tests/test_gemini_cache.py tests/test_gemini_safety_settings.py``
  → 28/28 pass (15 cache + 13 safety; the safety-settings suite
  doubles as regression on the ``call_with_tools`` signature change).
- ``uv run ruff check`` on changed files — clean.

Behavioural envelope
--------------------
- Flag still defaults False — no caller is opted in by default.
- When flag is True, both ``retain_extract_facts`` (from the earlier
  commit on this branch) and ``reflect_tool_call`` opt in.
- A cache-side failure (transient SDK error, prefix too small, manager
  uninstantiated) returns None and the caller proceeds uncached. There
  is no path by which caching can break reflect or retain.
2026-06-02 17:13:41 -04:00
Chris Bartholomew 2527c75970 feat(gemini): wire retain fact-extraction to context cache; surface cached + thoughts tokens
Follow-on to the foundation commit on this branch — without this, the
cache manager is unreachable and the metric ignores half the cost
surface. This commit makes the change actually do something when the
flag is flipped on.

What lands
----------
1. Retain fact-extraction (engine/retain/fact_extraction.py) opts into
   the cache. The system prompt and response schema are fingerprinted
   and reused across calls; the user message is the only variable
   part on the wire. A cache lookup failure or "prefix too small"
   response from Gemini transparently falls back to the existing
   uncached path — caching is a soft optimisation, never a blocker.

2. New top-level flag HINDSIGHT_API_LLM_GEMINI_PROMPT_CACHE_ENABLED
   (also exposed as ``llm_gemini_prompt_cache_enabled`` on
   HindsightConfig). Defaults to False so upgrade-and-do-nothing is a
   no-op. Flipping to True opts every Gemini caller (currently only
   retain) into context caching.

3. Two new metrics:
   - hindsight.llm.tokens.cached_input — subset of input tokens billed
     at the cached rate. Lets dashboards split cache-hit vs cache-miss
     volume independently of total throughput.
   - hindsight.llm.tokens.thoughts — reasoning tokens emitted by
     Gemini 2.5+. Billed at the output rate by the provider but
     invisible to candidates_token_count, so absent from output-token
     dashboards today. Surfacing this is required for honest cost
     attribution.

4. Provider plumbing: GeminiLLM gains a ``gemini_prompt_cache_enabled``
   kwarg and a ``get_or_create_cached_prefix(...)`` accessor that lazy-
   builds a GeminiCacheManager on first opt-in. LLMProvider /
   create_llm_provider / ConfiguredLLMProvider pass the flag through
   the standard plumbing alongside the existing safety_settings.

Verification
------------
- ``uv run ruff check`` — clean
- ``uv run pytest tests/test_gemini_cache.py`` — 12 tests including
  two new integration tests that pin (a) flag-off → cache manager
  never built, and (b) flag-on → manager lazy-built, second lookup
  served from in-memory cache, no extra SDK call.
- ``uv run pytest tests/test_gemini_safety_settings.py`` — 13 tests
  still green (no signature drift; the NoOp metrics collector was
  updated alongside the real one).

Rollout
-------
- Land this commit. With the flag default-off, behaviour is identical
  to today: cache code paths exist but are never reached.
- Flip the flag per-env. The metric goes non-zero on cached_input
  within a few calls.
- Watch hindsight.llm.tokens.cached_input vs hindsight.llm.tokens.input
  to confirm cache-hit rate.

What's deliberately NOT in this PR
----------------------------------
- Extending caching to other Gemini callers (reflect tool-call,
  consolidation). Same mechanism applies — copy two lines from the
  retain path. Leave for a follow-up so this lands in one focused PR.
- Cross-pod cache sharing. Each pod warms its own cache. The cost of
  one extra full-price call per pod per fingerprint per TTL window is
  negligible relative to steady-state savings.
2026-06-02 11:27:31 -04:00
Chris Bartholomew d458b715fe feat(gemini): add context-cache foundation (GeminiCacheManager + opt-in call() arg)
Wraps the google-genai SDK's CachedContent API so callers can reuse a
stable (system_instruction + response_schema) prefix across many
requests. Cached input tokens are billed at a fraction of the standard
input rate, which makes workloads with a fixed-prefix / small-user-message
shape — fact extraction, structured tagging, classification — far
cheaper to run.

This PR is foundation-only: no caller is wired up yet. Default
behaviour for every existing path is unchanged because
`cached_content_name` defaults to `None` and the cache manager is
never instantiated until a follow-up wires it in.

What's here
-----------
- `gemini_cache.GeminiCacheManager`: per-process map of prefix
  fingerprint → CachedContent resource name. Thread-safe via a single
  asyncio.Lock. Refreshes proactively at TTL minus a safety margin.
  Stable fingerprint normalisation strips auto-generated Pydantic
  schema titles so dynamically-built schema classes with identical
  shape hash to the same key (relevant for callers that rebuild the
  schema class on every request).
- `gemini_llm.GeminiLLM.call(cached_content_name=...)`: new optional
  arg. When set, the SDK config drops `system_instruction` and
  `response_schema` (those live in the cache) and instead passes
  `cached_content` to GenerateContentConfig. When unset, behaviour is
  byte-identical to before.
- `tests/test_gemini_cache.py`: 10 unit tests covering fingerprint
  stability, dict/list/Pydantic schema cases, get_or_create
  caching/recreate, "minimum token count" soft-fallback, transient
  SDK error soft-fallback, failed-create-doesn't-poison-cache, and
  the TTL refresh boundary.

Failure handling
----------------
- Gemini rejects creates whose prefix is below the model's minimum
  cacheable size with a "minimum"-style error message. The manager
  catches this, logs at DEBUG, and returns None so the caller
  transparently falls back to a non-cached call.
- Any other SDK error is logged at ERROR and also returns None — a
  bad create never crashes a request. Callers are required to treat
  None as "cache unavailable, use the normal path".

Not in this PR
--------------
- Wiring this into the fact-extraction pipeline (or any other caller)
- A metric for cached-token volume
Both will come in a focused follow-up so the foundation can land and
be reviewed independently.
2026-06-02 11:02:45 -04:00
23 changed files with 1958 additions and 139 deletions
@@ -363,6 +363,16 @@ ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOU
# Gemini safety settings
ENV_LLM_GEMINI_SAFETY_SETTINGS = "HINDSIGHT_API_LLM_GEMINI_SAFETY_SETTINGS"
# Gemini prompt caching. When enabled, retain fact-extraction reuses a
# CachedContent prefix for the static system_instruction + response_schema,
# cutting per-call input cost on workloads with many small documents.
# Provider-agnostic prompt-prefix caching. Providers that support it (currently
# Gemini/Vertex via CachedContent) reuse the large, fixed, bank-agnostic system
# prefix at the cached-input rate; providers that don't simply ignore it. On by
# default — the prefix is bank-agnostic so a single cache is shared across all
# banks, and creation soft-fails to an uncached call, so it never breaks a request.
ENV_LLM_PROMPT_CACHE_ENABLED = "HINDSIGHT_API_LLM_PROMPT_CACHE_ENABLED"
# Retain settings
ENV_RETAIN_MAX_COMPLETION_TOKENS = "HINDSIGHT_API_RETAIN_MAX_COMPLETION_TOKENS"
ENV_RETAIN_CHUNK_SIZE = "HINDSIGHT_API_RETAIN_CHUNK_SIZE"
@@ -766,6 +776,7 @@ DEFAULT_RETAIN_BATCH_TOKENS = 10_000 # ~40KB of text # Max chars per sub-batch
DEFAULT_RETAIN_ENTITY_LOOKUP = "trigram" # "full" or "trigram"
DEFAULT_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE = 100 # Unique entity names per pg_trgm candidate lookup query
DEFAULT_RETAIN_BATCH_ENABLED = False # Use LLM Batch API for fact extraction (only when async=True)
DEFAULT_LLM_PROMPT_CACHE_ENABLED = True # Reuse the fixed system prefix via provider prompt caching
DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS = 60 # Batch API polling interval in seconds
# File storage defaults
@@ -1158,6 +1169,10 @@ class HindsightConfig:
# Gemini safety settings (None = use Gemini defaults; list of dicts with category/threshold)
llm_gemini_safety_settings: list | None
# Gemini prompt caching toggle. When True, retain extraction reuses a
# CachedContent prefix for its system prompt + response schema.
llm_prompt_cache_enabled: bool
# Built-in llama.cpp configuration (for provider=llamacpp)
llamacpp_model_path: str | None # Path to GGUF file (None = auto-download default)
llamacpp_gpu_layers: int # -1 = all layers on GPU, 0 = CPU only
@@ -1759,6 +1774,10 @@ class HindsightConfig:
or DEFAULT_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY,
# Gemini safety settings (JSON-encoded list of {category, threshold} dicts)
llm_gemini_safety_settings=json.loads(os.getenv(ENV_LLM_GEMINI_SAFETY_SETTINGS, "null")),
llm_prompt_cache_enabled=os.getenv(
ENV_LLM_PROMPT_CACHE_ENABLED, str(DEFAULT_LLM_PROMPT_CACHE_ENABLED)
).lower()
in ("1", "true", "yes", "on"),
# Built-in llama.cpp configuration
llamacpp_model_path=os.getenv(ENV_LLAMACPP_MODEL_PATH) or None,
llamacpp_gpu_layers=int(os.getenv(ENV_LLAMACPP_GPU_LAYERS, str(DEFAULT_LLAMACPP_GPU_LAYERS))),
@@ -41,7 +41,10 @@ from ..llm_trace import (
from ..llm_wrapper import sanitize_llm_output
from ..memory_engine import Budget, fq_table
from ..retain import embedding_utils
from .prompts import build_batch_consolidation_prompt
from .prompts import (
build_consolidation_input,
build_consolidation_system_prompt,
)
if TYPE_CHECKING:
from asyncpg import Connection
@@ -1648,16 +1651,38 @@ async def _consolidate_batch_with_llm(
f"(out of {max_observations_per_scope}). Prefer UPDATE over CREATE when possible."
)
prompt_template = build_batch_consolidation_prompt(
config.observations_mission,
observation_capacity_note,
# Split the prompt: a bank-agnostic system instruction (rules + input format +
# decision guide + output format) that is byte-identical across batches AND
# across banks, and a per-batch user message (mission + capacity note + facts +
# existing observations). The split lets the system prefix be served from a
# single Gemini context cache shared by every bank — the bank mission, capacity
# note, and response_schema (all bank/batch-variable) are kept OUT of the
# cached prefix so one cache serves all and it never busts within a run.
system_prompt = build_consolidation_system_prompt(
llm_output_language=getattr(config, "llm_output_language", None),
)
prompt = prompt_template.format(
user_content = build_consolidation_input(
facts_text=facts_lines,
observations_text=observations_text,
observations_mission=config.observations_mission,
observation_capacity_note=observation_capacity_note,
)
# Opt into context caching of the stable system prefix when the provider
# supports it (gemini/vertexai with the flag on). response_schema is NOT
# passed to the fingerprint: it varies per batch (max_creates) but is not
# part of the cached prefix, so keying on it would needlessly bust the cache.
cached_prefix_name: str | None = None
provider_impl = getattr(llm_config, "_provider_impl", None)
if provider_impl is not None and provider_impl.supports_prompt_caching():
try:
cached_prefix_name = await provider_impl.get_or_create_cached_prefix(
system_instruction=system_prompt,
)
except Exception:
logger.exception("Consolidation cache prefix lookup failed; falling back to uncached call")
cached_prefix_name = None
# Use a constrained response model when observation limit is active
response_model = _build_response_model(max_creates=remaining_observation_slots)
@@ -1677,12 +1702,17 @@ async def _consolidate_batch_with_llm(
for attempt in range(1, max_attempts + 1):
try:
call_kwargs: dict[str, Any] = {
"messages": [{"role": "user", "content": prompt}],
"messages": [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_content},
],
"response_format": response_model,
"scope": "consolidation",
}
if inner_max_retries is not None:
call_kwargs["max_retries"] = inner_max_retries
if cached_prefix_name is not None:
call_kwargs["cached_prefix"] = cached_prefix_name
response: _ConsolidationBatchResponse = await llm_config.call(**call_kwargs)
# Defensive truncation: some LLM providers may not enforce JSON schema max_length
creates = response.creates
@@ -1699,7 +1729,7 @@ async def _consolidate_batch_with_llm(
updates=updates,
deletes=response.deletes,
obs_count=len(union_observations),
prompt_chars=len(prompt),
prompt_chars=len(system_prompt) + len(user_content),
)
except Exception as exc:
last_exc = exc
@@ -1711,7 +1741,9 @@ async def _consolidate_batch_with_llm(
f"[CONSOLIDATION] LLM batch call failed after {max_attempts} attempts for {batch_label}, "
f"skipping batch. Last error: {last_exc}"
)
return _BatchLLMResult(obs_count=len(union_observations), prompt_chars=len(prompt), failed=True)
return _BatchLLMResult(
obs_count=len(union_observations), prompt_chars=len(system_prompt) + len(user_content), failed=True
)
async def _create_observation_directly(
@@ -37,6 +37,33 @@ _PROCESSING_RULES = """## PROCESSING RULES
9. KEEP DISTINCT TOPICS DISTINCT: do not merge observations about different people, entities, or unrelated topics. Merging is for the same canonical fact recurring — not for related-but-distinct claims."""
# Stable description of the input shape. For the cached split path this lives in
# the system prefix (build_consolidation_system_prompt) so it is not re-sent on
# every batch; the per-batch user message then carries only the actual data.
_INPUT_FORMAT_NOTE = """## INPUT FORMAT
Each request provides new facts and existing observations:
- New facts: one per line, each prefixed with its `[uuid]`, followed by the fact text and optional temporal fields.
- Existing observations: a JSON array pooled from recalls across the new facts. Each entry has:
- `id`: unique identifier — copy this exactly when issuing an UPDATE or DELETE
- `text`: the observation content
- `proof_count`: number of supporting memories
- `occurred_start` / `occurred_end`: temporal range of source facts
- `source_memories`: array of supporting facts with their text and dates"""
# Per-batch data section for the cached split path — the stable format
# explanation above is omitted here (it lives in the cached prefix); only the
# variable facts/observations remain. Placeholders substituted at call time.
_SPLIT_INPUT_SECTION = """## INPUT
### New facts
{facts_text}
### Existing observations
{observations_text}"""
# Data section — format placeholders {facts_text} and {observations_text} are substituted at call time
_INPUT_SECTION = """## INPUT
@@ -146,3 +173,55 @@ def build_batch_consolidation_prompt(
f"{_DECISION_GUIDE}\n\n"
f"{_OUTPUT_SECTION}" + output_language_directive(llm_output_language)
)
def build_consolidation_system_prompt(
llm_output_language: str | None = None,
) -> str:
"""Bank-agnostic, cacheable system instruction for batch consolidation.
Holds only what is constant across banks: processing rules, input format,
decision guide, and output format. The bank's MISSION is deliberately NOT
here — baking it in would make the prefix bank-specific and force a separate
Gemini context cache per mission. The mission, the per-batch INPUT, and any
capacity constraint all ride in the user message (see
:func:`build_consolidation_input`), so this prefix is identical for every
bank and a single CachedContent serves them all. Returns final text
(brace-escaped examples already unescaped) for verbatim use as system message
and cached prefix.
"""
template = (
"You are a memory consolidation system. Synthesize new facts into "
"observations, merging with existing observations when appropriate.\n\n"
f"{_MISSION_PRIORITY_NOTE}\n\n"
f"{_PROCESSING_RULES}\n\n"
f"{_INPUT_FORMAT_NOTE}\n\n"
f"{_DECISION_GUIDE}\n\n"
f"{_OUTPUT_SECTION}" + output_language_directive(llm_output_language)
)
# No {facts_text}/{observations_text} placeholders here — the only braces are
# the doubled {{ }} in the OUTPUT examples, which .format() unescapes.
return template.format()
def build_consolidation_input(
facts_text: str,
observations_text: str,
observations_mission: str | None = None,
observation_capacity_note: str | None = None,
) -> str:
"""Per-batch user message: MISSION + INPUT data + any capacity constraint.
The MISSION lives here (not in the cached system prefix) so the prefix stays
bank-agnostic and one CachedContent serves every bank. The capacity note also
lives here since it varies as observation slots fill.
"""
mission = escape_for_prompt(observations_mission or _DEFAULT_MISSION)
mission_section = f"## MISSION\n\n{mission}\n\n"
capacity_section = ""
if observation_capacity_note:
capacity_section = f"## CAPACITY CONSTRAINT\n\n{escape_for_prompt(observation_capacity_note)}\n\n"
# _SPLIT_INPUT_SECTION omits the stable observation-format explanation (now in
# the cached system prefix) — only the variable facts/observations remain.
template = mission_section + capacity_section + _SPLIT_INPUT_SECTION
return template.format(facts_text=facts_text, observations_text=observations_text)
@@ -69,6 +69,7 @@ class LLMInterface(ABC):
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
cached_prefix: str | None = None,
) -> Any:
"""
Make an LLM API call with retry logic.
@@ -85,6 +86,9 @@ class LLMInterface(ABC):
skip_validation: Return raw JSON without Pydantic validation.
strict_schema: Use strict JSON schema enforcement (OpenAI only).
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
cached_prefix: Opaque handle from ``get_or_create_cached_prefix`` for the
cacheable system prefix, or None. Providers without explicit prompt
caching ignore it (and the wrapper only forwards it when set).
Returns:
If return_usage=False: Parsed response if response_format is provided, otherwise text content.
@@ -108,6 +112,7 @@ class LLMInterface(ABC):
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: str | dict[str, Any] = "auto",
cached_prefix: str | None = None,
) -> LLMToolCallResult:
"""
Make an LLM API call with tool/function calling support.
@@ -137,6 +142,46 @@ class LLMInterface(ABC):
"""
return False
# ── Prompt prefix caching (optional, per-provider) ─────────────────────────
def supports_prompt_caching(self) -> bool:
"""Whether this provider can cache a reusable prompt prefix.
Default False. Providers that return True must implement
``get_or_create_cached_prefix`` and honour the ``cached_prefix`` argument
of ``call`` / ``call_with_tools``.
"""
return False
async def get_or_create_cached_prefix(
self,
*,
system_instruction: str,
response_schema: Any | None = None,
tools: list[dict[str, Any]] | None = None,
) -> str | None:
"""Cache a reusable prompt prefix and return an opaque handle, or None.
The engine has already decided WHAT is cacheable: it puts the stable,
bank-agnostic instructions in ``system_instruction`` (plus ``tools``) and
keeps all per-request / per-bank data (documents, facts, the bank mission)
in the user message. A provider only chooses HOW to cache that prefix:
- Explicit-cache providers (e.g. Gemini ``CachedContent``): create the
cache, return its handle; the engine passes the handle back via
``call(cached_prefix=...)`` and the provider then drops the prefix from
the request, billing it at the cached rate.
- Automatic-cache providers (e.g. OpenAI): no handle needed — caching is
transparent as long as the prefix is a stable leading block, which it
already is. They can keep this default (return None) and still benefit.
- Inline-marker providers (e.g. Anthropic ``cache_control``): mark the
prefix block inside ``call`` instead; may also keep this default.
Returns None when caching is disabled/unsupported or the prefix is too
small; callers MUST fall back to an uncached call in that case.
"""
return None
async def submit_batch(
self,
requests: list[dict[str, Any]],
@@ -255,6 +255,7 @@ def create_llm_provider(
vertexai_region: str | None = None,
vertexai_credentials: Any = None,
gemini_safety_settings: list | None = None,
prompt_cache_enabled: bool = False,
litellmrouter_config: dict[str, Any] | None = None,
) -> Any: # Returns LLMInterface
"""
@@ -343,6 +344,7 @@ def create_llm_provider(
vertexai_region=vertexai_region,
vertexai_credentials=vertexai_credentials,
gemini_safety_settings=gemini_safety_settings,
prompt_cache_enabled=prompt_cache_enabled,
)
elif provider_lower == "anthropic":
@@ -468,6 +470,7 @@ class LLMProvider:
groq_service_tier: str | None = None,
openai_service_tier: str | None = None,
gemini_safety_settings: list | None = None,
prompt_cache_enabled: bool = False,
extra_body: dict[str, Any] | None = None,
default_headers: dict[str, str] | None = None,
litellmrouter_config: dict[str, Any] | None = None,
@@ -506,6 +509,11 @@ class LLMProvider:
self.openai_service_tier = openai_service_tier
# Gemini safety settings (instance default; can be overridden per-request via context var)
self.gemini_safety_settings = gemini_safety_settings
# Gemini prompt caching: when True, retain extraction (and any future
# caller that opts in) will reuse a CachedContent prefix to cut
# input-token cost. Off by default so the change is observable behind
# a flip rather than a silent behaviour change on upgrade.
self.prompt_cache_enabled = prompt_cache_enabled
# Extra body params for OpenAI-compatible providers (e.g. chat_template_kwargs)
self.extra_body = extra_body
# Default headers passed to provider SDK clients (e.g. proxy auth, request tracing).
@@ -624,6 +632,21 @@ class LLMProvider:
except Exception:
pass # Config may not be initialized in test environments
# Prompt-prefix caching is a provider-agnostic toggle (default on): resolve
# it from the static server config for every provider when the caller didn't
# pass an explicit override. Providers that don't support caching ignore the
# value; only those that implement get_or_create_cached_prefix act on it.
if not self.prompt_cache_enabled:
from ..config import DEFAULT_LLM_PROMPT_CACHE_ENABLED, _get_raw_config
try:
raw_config = _get_raw_config()
self.prompt_cache_enabled = bool(
getattr(raw_config, "llm_prompt_cache_enabled", DEFAULT_LLM_PROMPT_CACHE_ENABLED)
)
except Exception:
pass # Config may not be initialized in test environments
# For litellmrouter: prefer an explicit chain from the caller (per-op
# construction in MemoryEngine threads the right chain through). If the caller
# didn't supply one, fall back to the global ``llm_litellmrouter_config`` so
@@ -652,6 +675,7 @@ class LLMProvider:
vertexai_region=vertexai_region,
vertexai_credentials=vertexai_credentials,
gemini_safety_settings=self.gemini_safety_settings,
prompt_cache_enabled=self.prompt_cache_enabled,
litellmrouter_config=router_config,
)
@@ -715,6 +739,7 @@ class LLMProvider:
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
cached_prefix: str | None = None,
) -> Any:
"""
Make an LLM API call with retry logic.
@@ -771,6 +796,11 @@ class LLMProvider:
for sem in _semaphores_for_scope(scope):
await stack.enter_async_context(sem)
# cached_prefix is only set for providers that returned a handle
# from get_or_create_cached_prefix() (e.g. Gemini); it's None for
# the rest. Forward it only when present so providers that don't
# implement caching keep their call() signature untouched.
cache_kwarg = {"cached_prefix": cached_prefix} if cached_prefix is not None else {}
try:
# Delegate to provider implementation
result = await self._provider_impl.call(
@@ -785,6 +815,7 @@ class LLMProvider:
skip_validation=skip_validation,
strict_schema=strict_schema,
return_usage=return_usage,
**cache_kwarg,
)
except Exception as e:
get_span_recorder().record_llm_call(
@@ -824,6 +855,7 @@ class LLMProvider:
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: str | dict[str, Any] = "auto",
cached_prefix: str | None = None,
) -> "LLMToolCallResult":
"""
Make an LLM API call with tool/function calling support.
@@ -864,6 +896,10 @@ class LLMProvider:
for sem in _semaphores_for_scope(scope):
await stack.enter_async_context(sem)
# cached_prefix is only set for providers that returned a handle
# from get_or_create_cached_prefix(); forward it only when present
# so non-caching providers keep their signature (same as call()).
cache_kwarg = {"cached_prefix": cached_prefix} if cached_prefix is not None else {}
try:
# Delegate to provider implementation
result = await self._provider_impl.call_with_tools(
@@ -876,6 +912,7 @@ class LLMProvider:
initial_backoff=initial_backoff,
max_backoff=max_backoff,
tool_choice=tool_choice,
**cache_kwarg,
)
except Exception as e:
get_span_recorder().record_llm_call(
@@ -0,0 +1,316 @@
"""Gemini context-cache manager.
Wraps the ``google-genai`` SDK's CachedContent API to let callers reuse a
stable system_instruction + response_schema prefix across many requests.
Cached input tokens are billed at ~10× lower than fresh input tokens
(check the current Gemini pricing for the exact ratio per model), so for
workloads that repeatedly send a large fixed prefix with a small variable
user message — fact extraction, structured tagging, classification — the
input-cost savings are substantial.
This module owns only the create/refresh/lookup lifecycle. It is up to
the caller to (a) decide that the prefix is stable enough to cache, and
(b) pass the returned cache name to ``GeminiLLM.call()``. When the
returned name is ``None`` (because Gemini rejected the create — most
commonly because the prefix is smaller than the model's minimum), the
caller MUST fall back to a non-cached call.
Cardinality
-----------
The intended cache count per process is small (≲100 entries). Each
entry corresponds to one combination of (model, system_instruction,
response_schema). If a caller sees the cache grow unboundedly it
indicates the system_instruction contains per-request data that should
move into the user message instead.
TTL
---
Gemini's CachedContent has a TTL bounded by the model (currently 1h
for most generally-available models). This manager refreshes proactively
at ``ttl_safety_margin`` before expiry. If a cached entry has expired
between refreshes the next call will recreate it transparently.
"""
from __future__ import annotations
import asyncio
import hashlib
import json
import logging
import time
from dataclasses import dataclass
from typing import Any
logger = logging.getLogger(__name__)
# Default TTL: 55 minutes. Gemini's hard max for CachedContent is 1 hour
# for most models; we refresh 5 minutes early so a request landing right
# at the boundary doesn't race against expiry.
_DEFAULT_TTL_SECONDS = 55 * 60
_DEFAULT_REFRESH_MARGIN_SECONDS = 5 * 60
# Cap on the cache-create network call. It runs while holding the manager lock, so
# a hung create would block every concurrent caller (e.g. all chunks of a 10-chunk
# retain batch waiting on the cold-start create). On timeout the create soft-fails
# to None and callers proceed uncached, rather than stalling the whole batch.
_DEFAULT_CREATE_TIMEOUT_SECONDS = 30.0
@dataclass
class _CacheEntry:
name: str # The CachedContent resource name returned by Gemini.
created_at: float
ttl_seconds: int
class GeminiCacheManager:
"""Per-process map of (prefix fingerprint) → CachedContent name.
Thread-safe across asyncio tasks via a single ``asyncio.Lock``. The
create/refresh calls are serialised; this is fine because cache
creation is a one-shot warm-up per fingerprint (subsequent reads are
pure dict lookups outside the lock).
Not shared across pods — each worker / api replica builds its own
cache. The cost of cold-starting one extra full-price call per pod
per fingerprint per hour is negligible compared to the steady-state
savings.
"""
def __init__(
self,
client: Any,
*,
ttl_seconds: int = _DEFAULT_TTL_SECONDS,
refresh_margin_seconds: int = _DEFAULT_REFRESH_MARGIN_SECONDS,
create_timeout_seconds: float = _DEFAULT_CREATE_TIMEOUT_SECONDS,
) -> None:
self._client = client
self._ttl_seconds = ttl_seconds
self._refresh_margin_seconds = refresh_margin_seconds
self._create_timeout_seconds = create_timeout_seconds
self._entries: dict[str, _CacheEntry] = {}
self._lock = asyncio.Lock()
@staticmethod
def fingerprint(
model: str,
system_instruction: str,
response_schema: Any | None = None,
tools: list[dict[str, Any]] | None = None,
) -> str:
"""Stable hash of the cacheable surface.
``response_schema`` may be a Pydantic class, a dict, or ``None``.
Pydantic schemas are normalised by serialising via
``model_json_schema()`` and stripping the auto-generated
``"title"`` fields so two dynamically-built models with the same
shape but different class names hash identically. This matters
for callers (e.g. fact extraction) that rebuild the schema
class on every request via a builder helper — without the
normalisation the cache would never hit.
``tools`` is the OpenAI-style tools list (each entry has a
``"function"`` dict with name/description/parameters). When
supplied, the tool definitions become part of the cache key so a
loop that adds or renames a tool gets a fresh cache and doesn't
silently use a stale schema. Tools are serialised with
``sort_keys=True`` to neutralise dict-ordering drift.
"""
hasher = hashlib.sha256()
hasher.update(model.encode("utf-8"))
hasher.update(b"\x00")
hasher.update(system_instruction.encode("utf-8"))
hasher.update(b"\x00")
if response_schema is None:
hasher.update(b"none")
elif hasattr(response_schema, "model_json_schema"):
try:
schema = response_schema.model_json_schema()
_strip_titles(schema)
hasher.update(json.dumps(schema, sort_keys=True).encode("utf-8"))
except Exception:
# Fall back to class identity if the schema can't be serialised.
hasher.update(repr(response_schema).encode("utf-8"))
else:
try:
hasher.update(json.dumps(response_schema, sort_keys=True).encode("utf-8"))
except (TypeError, ValueError):
hasher.update(repr(response_schema).encode("utf-8"))
hasher.update(b"\x00")
if tools:
try:
hasher.update(json.dumps(tools, sort_keys=True).encode("utf-8"))
except (TypeError, ValueError):
hasher.update(repr(tools).encode("utf-8"))
else:
hasher.update(b"no-tools")
return hasher.hexdigest()
async def get_or_create(
self,
*,
model: str,
system_instruction: str,
response_schema: Any | None = None,
tools: list[dict[str, Any]] | None = None,
) -> str | None:
"""Return a CachedContent resource name for the given prefix, or
``None`` if Gemini rejects the create (prefix too small, model
does not support caching, etc.).
``tools`` is the OpenAI-style tools list. When supplied, the tool
definitions are baked into the CachedContent so the caller's
``call_with_tools`` doesn't need to resend them on every
iteration. Pass ``None`` for non-tool calls.
``None`` return is a normal, expected value — the caller falls
back to an uncached call and the system continues to work.
"""
key = self.fingerprint(model, system_instruction, response_schema, tools)
async with self._lock:
entry = self._entries.get(key)
if entry is not None and self._is_fresh(entry):
return entry.name
# Need to (re)create. Pop the stale entry first so a failed
# create doesn't leave a name we'd return on the next call.
self._entries.pop(key, None)
try:
cache_name = await self._create_cache(
model=model,
system_instruction=system_instruction,
tools=tools,
)
except _CacheNotEligible as e:
logger.debug(
"GeminiCacheManager: prefix not eligible for caching (model=%s, reason=%s) — caller will fall back",
model,
e,
)
return None
except Exception:
logger.exception(
"GeminiCacheManager: failed to create cached content "
"(model=%s); caller will fall back to uncached call",
model,
)
return None
if cache_name is None:
return None
self._entries[key] = _CacheEntry(
name=cache_name,
created_at=time.monotonic(),
ttl_seconds=self._ttl_seconds,
)
return cache_name
def _is_fresh(self, entry: _CacheEntry) -> bool:
"""An entry is fresh if it's young enough that the next request
won't race against the TTL expiry."""
age = time.monotonic() - entry.created_at
return age < (entry.ttl_seconds - self._refresh_margin_seconds)
def invalidate(self, name: str) -> None:
"""Forget a cache name that the server rejected (expired/deleted/invalid).
Called by the provider when a generate request using this CachedContent
fails, so the next ``get_or_create`` recreates it instead of handing back
the dead name again. Best-effort and sync — drops the matching entry from
the in-process map; the orphaned server-side cache (if any) ages out on
its own TTL.
"""
for key, entry in list(self._entries.items()):
if entry.name == name:
self._entries.pop(key, None)
async def _create_cache(
self,
*,
model: str,
system_instruction: str,
tools: list[dict[str, Any]] | None = None,
) -> str | None:
"""Wrap ``client.aio.caches.create`` with the config we want.
The SDK surface differs slightly across google-genai versions;
this implementation targets the >=1.0.0 line where caches live
under ``client.aio.caches``.
"""
# Lazy import so this module doesn't require the SDK at import time.
from google.genai import types as genai_types
# A CachedContent only holds reusable *input* — system_instruction,
# contents, tools, ttl. ``response_schema``/``response_mime_type`` are
# generation-time output constraints and the SDK rejects them here
# (``CreateCachedContentConfig`` forbids those fields). They are applied
# per-request on the GenerateContentConfig instead — see the call sites,
# which set them alongside ``cached_content``. ``response_schema`` is
# still part of the fingerprint so a schema change keys a fresh cache.
config_kwargs: dict[str, Any] = {
"system_instruction": system_instruction,
"ttl": f"{self._ttl_seconds}s",
}
if tools:
# OpenAI-style {"function": {...}} entries must be converted to
# Gemini's Tool/FunctionDeclaration shape before caching.
gemini_tools = []
for tool in tools:
func = tool.get("function", {})
gemini_tools.append(
genai_types.Tool(
function_declarations=[
genai_types.FunctionDeclaration(
name=func.get("name", ""),
description=func.get("description", ""),
parameters=func.get("parameters"),
)
]
)
)
config_kwargs["tools"] = gemini_tools
try:
cached = await asyncio.wait_for(
self._client.aio.caches.create(
model=model,
config=genai_types.CreateCachedContentConfig(**config_kwargs),
),
timeout=self._create_timeout_seconds,
)
except Exception as e:
# Gemini returns a 400 with a "minimum token count" message
# when the prefix is too small. We treat this as a soft
# "not eligible" signal rather than a real error so callers
# silently fall back to non-cached.
msg = str(e).lower()
if "minimum" in msg or "too small" in msg or "too short" in msg:
raise _CacheNotEligible(str(e)) from e
raise
return getattr(cached, "name", None)
class _CacheNotEligible(Exception):
"""Raised when Gemini rejects the cache create because the prefix
is below the model's minimum cacheable size. Treated as a soft
fallback by the caller, not an error."""
def _strip_titles(node: Any) -> None:
"""Recursively remove auto-generated ``"title"`` keys from a JSON
Schema-like dict tree, in place. Pydantic seeds these from the
Python class name, which means structurally-identical schemas built
from differently-named classes look distinct to a naive hash."""
if isinstance(node, dict):
node.pop("title", None)
for v in node.values():
_strip_titles(v)
elif isinstance(node, list):
for item in node:
_strip_titles(item)
@@ -70,6 +70,14 @@ class GeminiLLM(LLMInterface):
# Safety settings: None means use Gemini's defaults
self._safety_settings: list | None = kwargs.get("gemini_safety_settings")
# Context-cache manager. Lazy-initialized on first cache lookup so
# nothing happens for models/workloads that never reach it. The instance
# default here is off (a directly-constructed GeminiLLM doesn't cache); the
# server-level default is on and flows in via the prompt_cache_enabled kwarg
# resolved from config in LLMProvider.
self._cache_manager: Any | None = None
self._prompt_cache_enabled: bool = bool(kwargs.get("prompt_cache_enabled", False))
if self._is_vertexai:
self._init_vertexai(**kwargs)
else:
@@ -168,6 +176,7 @@ class GeminiLLM(LLMInterface):
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
cached_prefix: str | None = None,
) -> Any:
"""
Make a Gemini/VertexAI API call with retry logic.
@@ -184,6 +193,14 @@ class GeminiLLM(LLMInterface):
skip_validation: Return raw JSON without Pydantic validation.
strict_schema: Use strict JSON schema enforcement (not supported by Gemini).
return_usage: If True, return tuple (result, TokenUsage).
cached_prefix: Optional CachedContent resource name (from
``GeminiCacheManager.get_or_create``). When set, the
system_instruction is assumed to live in the cache; this call
skips resending it and the cached prefix is billed at the
cached-input rate instead of the standard input rate. The
response_schema is still sent per-request (it is not cacheable).
Pass ``None`` to use the
normal uncached path.
Returns:
If return_usage=False: Parsed response if response_format provided, else text.
@@ -191,9 +208,14 @@ class GeminiLLM(LLMInterface):
"""
start_time = time.time()
# Convert OpenAI-style messages to Gemini format
# Convert OpenAI-style messages to Gemini format. We ALWAYS build
# system_instruction (even when a cache is in use): the config builder
# below omits it from the request while the cache carries the prefix, but
# it must be available so the cached-call-failed safety net can re-send it
# inline. Whether it's actually sent is decided in _build_generation_config.
system_instruction = None
gemini_contents = []
using_cache = cached_prefix is not None
for msg in messages:
role = msg.get("role", "user")
@@ -209,7 +231,9 @@ class GeminiLLM(LLMInterface):
else:
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
# Add JSON schema instruction if response_format is provided
# Add the JSON schema as a textual hint in the system_instruction (matching
# the normal uncached path). Structured output is still enforced via
# response_schema regardless; this is just guidance text.
if response_format is not None and hasattr(response_format, "model_json_schema"):
schema = response_format.model_json_schema()
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
@@ -218,32 +242,43 @@ class GeminiLLM(LLMInterface):
else:
system_instruction = schema_msg
# Build generation config
config_kwargs: dict[str, Any] = {}
if system_instruction:
config_kwargs["system_instruction"] = system_instruction
if response_format is not None:
config_kwargs["response_mime_type"] = "application/json"
config_kwargs["response_schema"] = response_format
if temperature is not None:
config_kwargs["temperature"] = temperature
# Gemini's equivalent of OpenAI-style max_completion_tokens is max_output_tokens.
# Without it the model can produce arbitrarily long responses, ignoring the
# caller's intended cap (e.g. mental_models max_tokens during refresh).
if max_completion_tokens is not None:
config_kwargs["max_output_tokens"] = max_completion_tokens
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
effective_safety_settings = _safety_settings_ctx.get()
if effective_safety_settings is None:
effective_safety_settings = self._safety_settings
if effective_safety_settings is not None:
config_kwargs["safety_settings"] = [
genai_types.SafetySetting(category=s["category"], threshold=s["threshold"])
for s in effective_safety_settings
]
generation_config = genai_types.GenerateContentConfig(**config_kwargs) if config_kwargs else None
# Build generation config. ``cached_content`` and ``system_instruction``
# are mutually exclusive (the cache IS the prefix; the SDK rejects
# re-sending it). ``response_schema``/``response_mime_type`` are
# request-level output constraints — NOT cacheable — so they're set on
# every structured call, including cached ones where they ride alongside
# ``cached_content``. Built as a closure so we can rebuild it WITHOUT the
# cache and retry inline if a stale/invalid CachedContent makes the call fail.
def _build_generation_config(use_cache: bool) -> "genai_types.GenerateContentConfig | None":
config_kwargs: dict[str, Any] = {}
if use_cache:
config_kwargs["cached_content"] = cached_prefix
elif system_instruction:
config_kwargs["system_instruction"] = system_instruction
if response_format is not None:
config_kwargs["response_mime_type"] = "application/json"
config_kwargs["response_schema"] = response_format
if temperature is not None:
config_kwargs["temperature"] = temperature
# Gemini's equivalent of OpenAI-style max_completion_tokens is max_output_tokens.
# Without it the model can produce arbitrarily long responses, ignoring the
# caller's intended cap (e.g. mental_models max_tokens during refresh).
if max_completion_tokens is not None:
config_kwargs["max_output_tokens"] = max_completion_tokens
if effective_safety_settings is not None:
config_kwargs["safety_settings"] = [
genai_types.SafetySetting(category=s["category"], threshold=s["threshold"])
for s in effective_safety_settings
]
return genai_types.GenerateContentConfig(**config_kwargs) if config_kwargs else None
cache_active = using_cache
generation_config = _build_generation_config(cache_active)
last_exception = None
@@ -288,15 +323,24 @@ class GeminiLLM(LLMInterface):
else:
result = content
# Extract token usage
# Extract token usage. ``cached_content_token_count`` and
# ``thoughts_token_count`` are populated on the Gemini 2.5+
# family; treat missing fields as 0 so older models still
# record sensible metrics.
input_tokens = 0
output_tokens = 0
cached_input_tokens = 0
thoughts_tokens = 0
cached_tokens = 0
if hasattr(response, "usage_metadata") and response.usage_metadata:
usage = response.usage_metadata
input_tokens = usage.prompt_token_count or 0
output_tokens = usage.candidates_token_count or 0
cached_tokens = getattr(usage, "cached_content_token_count", 0) or 0
cached_input_tokens = getattr(usage, "cached_content_token_count", 0) or 0
thoughts_tokens = getattr(usage, "thoughts_token_count", 0) or 0
# Tracing/TokenUsage consume ``cached_tokens``; metrics consume
# ``cached_input_tokens`` — same value, two downstream names.
cached_tokens = cached_input_tokens
# Record metrics
duration = time.time() - start_time
@@ -309,6 +353,8 @@ class GeminiLLM(LLMInterface):
input_tokens=input_tokens,
output_tokens=output_tokens,
success=True,
cached_input_tokens=cached_input_tokens,
thoughts_tokens=thoughts_tokens,
)
# Record trace span
@@ -370,6 +416,20 @@ class GeminiLLM(LLMInterface):
logger.error(f"Gemini auth error (HTTP {e.code}), not retrying: {str(e)}")
raise
# Cached-request safety net: a stale/invalid/expired CachedContent
# (or an incompatibility like cache + tool_config) surfaces as a 400.
# Retrying the same cached request can't recover, so on the first
# such failure drop the cache, invalidate it so later operations
# recreate it, and retry THIS call inline with the prefix inlined.
# Caching must never break a request.
if cache_active and e.code == 400:
logger.warning(f"Gemini cached call failed (400); retrying uncached. Reason: {str(e)}")
if self._cache_manager is not None and cached_prefix is not None:
self._cache_manager.invalidate(cached_prefix)
cache_active = False
generation_config = _build_generation_config(cache_active)
continue
# Retry on retryable errors (rate limits, server errors, client errors)
if e.code in (400, 429, 500, 502, 503, 504) or (e.code and e.code >= 500):
last_exception = e
@@ -403,6 +463,7 @@ class GeminiLLM(LLMInterface):
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: str | dict[str, Any] = "auto",
cached_prefix: str | None = None,
) -> LLMToolCallResult:
"""
Make a Gemini/VertexAI API call with tool/function calling support.
@@ -417,27 +478,39 @@ class GeminiLLM(LLMInterface):
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
tool_choice: How to choose tools (Gemini uses "auto" only).
cached_prefix: Optional CachedContent resource name (from
``GeminiCacheManager.get_or_create`` with ``tools=...``). When
set, the system_instruction and tool definitions are assumed
to live in the cache; this call will skip resending them and
the cached prefix is billed at the cached-input rate. The
``tools`` argument is still required (the caller may pass
an empty list when the cache holds them) so existing call
sites don't break.
Returns:
LLMToolCallResult with content and/or tool_calls.
"""
start_time = time.time()
using_cache = cached_prefix is not None
# Convert tools to Gemini format
# Convert tools to Gemini format. When the cache is in use, the
# tool definitions are baked into the CachedContent at create time
# and the SDK rejects re-sending them alongside ``cached_content``.
gemini_tools = []
for tool in tools:
func = tool.get("function", {})
gemini_tools.append(
genai_types.Tool(
function_declarations=[
genai_types.FunctionDeclaration(
name=func.get("name", ""),
description=func.get("description", ""),
parameters=func.get("parameters"),
)
]
if not using_cache:
for tool in tools:
func = tool.get("function", {})
gemini_tools.append(
genai_types.Tool(
function_declarations=[
genai_types.FunctionDeclaration(
name=func.get("name", ""),
description=func.get("description", ""),
parameters=func.get("parameters"),
)
]
)
)
)
# Convert messages
system_instruction = None
@@ -450,6 +523,10 @@ class GeminiLLM(LLMInterface):
content = msg.get("content", "")
if role == "system":
# Always capture system_instruction. _build_tools_config omits it
# (and tools) from the request while the cache carries the prefix,
# but it must be available so the cached-call-failed safety net can
# re-send the prefix + tools inline.
system_instruction = (system_instruction + "\n\n" + content) if system_instruction else content
i += 1
elif role == "tool":
@@ -497,49 +574,62 @@ class GeminiLLM(LLMInterface):
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
i += 1
config_kwargs: dict[str, Any] = {"tools": gemini_tools}
if system_instruction:
config_kwargs["system_instruction"] = system_instruction
if temperature is not None:
config_kwargs["temperature"] = temperature
# See note in `call`: Gemini's max_output_tokens is the equivalent of
# OpenAI-style max_completion_tokens.
if max_completion_tokens is not None:
config_kwargs["max_output_tokens"] = max_completion_tokens
# Map OpenAI-style tool_choice to Gemini FunctionCallingConfig
if tool_choice == "required":
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(
mode="ANY",
)
)
elif isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
fn_name = tool_choice.get("function", {}).get("name")
if fn_name:
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(
mode="ANY",
allowed_function_names=[fn_name],
)
)
elif tool_choice == "none":
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(mode="NONE")
)
# "auto" is the default (no tool_config needed)
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
effective_safety_settings = _safety_settings_ctx.get()
if effective_safety_settings is None:
effective_safety_settings = self._safety_settings
if effective_safety_settings is not None:
config_kwargs["safety_settings"] = [
genai_types.SafetySetting(category=s["category"], threshold=s["threshold"])
for s in effective_safety_settings
]
config = genai_types.GenerateContentConfig(**config_kwargs)
# When using a cached prefix, the SDK rejects re-sending system_instruction
# or tools alongside ``cached_content`` — the cache IS the prefix.
# tool_config (mode / allowed_function_names) is a per-request decision and
# stays out of the cache. Built as a closure so we can rebuild it WITHOUT
# the cache and retry inline if a stale/invalid cache makes the call fail.
def _build_tools_config(use_cache: bool) -> "genai_types.GenerateContentConfig":
config_kwargs: dict[str, Any] = {}
if use_cache:
config_kwargs["cached_content"] = cached_prefix
else:
config_kwargs["tools"] = gemini_tools
if system_instruction:
config_kwargs["system_instruction"] = system_instruction
if temperature is not None:
config_kwargs["temperature"] = temperature
# See note in `call`: Gemini's max_output_tokens is the equivalent of
# OpenAI-style max_completion_tokens.
if max_completion_tokens is not None:
config_kwargs["max_output_tokens"] = max_completion_tokens
# Map OpenAI-style tool_choice to Gemini FunctionCallingConfig
if tool_choice == "required":
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(
mode="ANY",
)
)
elif isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
fn_name = tool_choice.get("function", {}).get("name")
if fn_name:
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(
mode="ANY",
allowed_function_names=[fn_name],
)
)
elif tool_choice == "none":
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(mode="NONE")
)
# "auto" is the default (no tool_config needed)
if effective_safety_settings is not None:
config_kwargs["safety_settings"] = [
genai_types.SafetySetting(category=s["category"], threshold=s["threshold"])
for s in effective_safety_settings
]
return genai_types.GenerateContentConfig(**config_kwargs)
cache_active = using_cache
config = _build_tools_config(cache_active)
last_exception = None
for attempt in range(max_retries + 1):
@@ -582,12 +672,18 @@ class GeminiLLM(LLMInterface):
finish_reason = "tool_calls" if tool_calls else "stop"
# Extract token usage
# Extract token usage. ``cached_content_token_count`` and
# ``thoughts_token_count`` are populated on the Gemini 2.5+
# family; absent fields are treated as 0.
input_tokens = 0
output_tokens = 0
cached_input_tokens = 0
thoughts_tokens = 0
if response.usage_metadata:
input_tokens = response.usage_metadata.prompt_token_count or 0
output_tokens = response.usage_metadata.candidates_token_count or 0
cached_input_tokens = getattr(response.usage_metadata, "cached_content_token_count", 0) or 0
thoughts_tokens = getattr(response.usage_metadata, "thoughts_token_count", 0) or 0
# Record metrics
duration = time.time() - start_time
@@ -600,6 +696,8 @@ class GeminiLLM(LLMInterface):
input_tokens=input_tokens,
output_tokens=output_tokens,
success=True,
cached_input_tokens=cached_input_tokens,
thoughts_tokens=thoughts_tokens,
)
# Record OpenTelemetry span
@@ -624,6 +722,7 @@ class GeminiLLM(LLMInterface):
finish_reason=finish_reason,
error=None,
tool_calls=tool_calls_dict,
cached_tokens=cached_input_tokens,
)
return LLMToolCallResult(
@@ -640,6 +739,18 @@ class GeminiLLM(LLMInterface):
logger.error(f"Gemini auth error (HTTP {e.code}), not retrying: {str(e)}")
raise
# Cached-request safety net (see ``call``): a stale/invalid cache or
# a cache+tool_config conflict surfaces as a 400. Drop the cache,
# invalidate it for later operations, and retry THIS call inline
# with the prefix + tools re-sent. Caching must never break a call.
if cache_active and e.code == 400:
logger.warning(f"Gemini cached tool call failed (400); retrying uncached. Reason: {str(e)}")
if self._cache_manager is not None and cached_prefix is not None:
self._cache_manager.invalidate(cached_prefix)
cache_active = False
config = _build_tools_config(cache_active)
continue
# Retry on retryable errors
last_exception = e
if attempt < max_retries:
@@ -656,6 +767,54 @@ class GeminiLLM(LLMInterface):
raise last_exception
raise RuntimeError("Gemini tool call failed")
def supports_prompt_caching(self) -> bool:
"""True when explicit Gemini context caching is enabled for this instance.
Reflects the opt-in flag so callers skip the cache lookup entirely when
it's off; ``get_or_create_cached_prefix`` also returns None in that case.
"""
return self._prompt_cache_enabled
async def get_or_create_cached_prefix(
self,
*,
system_instruction: str,
response_schema: Any | None = None,
tools: list[dict[str, Any]] | None = None,
) -> str | None:
"""Return a CachedContent resource name for the given prefix, or
``None`` if context caching is disabled, the provider doesn't
support it, or Gemini rejects the create (prefix too small, etc.).
``tools`` is the OpenAI-style tools list; pass it when caching a
prefix that will be used by ``call_with_tools()``. The fingerprint
includes the tool definitions so a loop that swaps a tool gets a
fresh cache automatically.
Callers pass the returned name to ``call(cached_prefix=...)``
or ``call_with_tools(cached_prefix=...)`` and treat ``None``
as "cache unavailable — use the normal path". That fallback is
essential: the system must continue to work if caching is disabled,
if Gemini's caching API has an outage, or if the prefix is below
the model's minimum cacheable size.
"""
if not self._prompt_cache_enabled:
return None
if self._client is None:
return None
if self._cache_manager is None:
# Lazy import so the cache module is only loaded when caching
# is actually used.
from hindsight_api.engine.providers.gemini_cache import GeminiCacheManager
self._cache_manager = GeminiCacheManager(self._client)
return await self._cache_manager.get_or_create(
model=self.model,
system_instruction=system_instruction,
response_schema=response_schema,
tools=tools,
)
async def cleanup(self) -> None:
"""Clean up resources (close connections, etc.)."""
# Gemini client doesn't require explicit cleanup
@@ -382,6 +382,28 @@ async def run_reflect_agent(
{"role": "user", "content": query},
]
# Opt into context caching for the agentic tool loop. The system
# prompt and tool definitions are stable for the duration of this
# reflect call (and across reflects against the same bank), so
# caching them once and reusing across every iteration of the loop
# collapses the dominant input cost — the prefix repeated on every
# turn. ``get_or_create_cached_prefix`` returns None when caching is
# disabled, unsupported, or the prefix is too small; the
# ``call_with_tools`` invocation below transparently falls back to
# the uncached path in that case.
cached_prefix_name: str | None = None
provider_impl = getattr(llm_config, "_provider_impl", None)
if provider_impl is not None and provider_impl.supports_prompt_caching():
try:
cached_prefix_name = await provider_impl.get_or_create_cached_prefix(
system_instruction=system_prompt,
tools=tools,
)
except Exception:
# Caching is a soft optimisation; never let a cache-side
# error block a reflect.
cached_prefix_name = None
# Tracking
total_tools_called = 0
tool_trace: list[ToolCall] = []
@@ -576,12 +598,22 @@ async def run_reflect_agent(
iter_tool_choice = "auto"
try:
result = await llm_config.call_with_tools(
ct_kwargs: dict[str, Any] = dict(
messages=messages,
tools=tools,
scope="reflect_tool_call",
tool_choice=iter_tool_choice,
)
# Gemini rejects ``cached_content`` alongside a per-request
# ``tool_config`` (forced tool choice): "CachedContent can not be used
# with GenerateContent request setting system_instruction, tools or
# tool_config." The forced-sequence iterations set tool_config, so only
# the ``auto`` iterations can reference the cache; forced iterations send
# the prefix inline. The cache (tools + system prompt) is identical
# either way, so this just limits *which* iterations are billed cached.
if cached_prefix_name is not None and iter_tool_choice == "auto":
ct_kwargs["cached_prefix"] = cached_prefix_name
result = await llm_config.call_with_tools(**ct_kwargs)
llm_duration = int((time.time() - llm_start) * 1000)
consecutive_errors = 0
total_input_tokens += result.input_tokens
@@ -10,7 +10,7 @@ import json
import logging
import re
from datetime import datetime, timedelta
from typing import Literal, cast
from typing import Any, Literal, cast
from pydantic import BaseModel, ConfigDict, Field, create_model, field_validator
@@ -887,20 +887,15 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
extraction_mode = config.retain_extraction_mode
extract_causal_links = config.retain_extract_causal_links
# Build retain_mission section if set - injected before the mode-specific guidelines
# Escape braces so user-supplied text survives str.format() on the prompt template.
# The per-bank retain mission is NOT baked into this system prompt: it would
# make the prompt bank-specific and force a separate Gemini context cache per
# mission (one per bank). Instead the prompt is bank-agnostic so a single
# CachedContent serves every bank, and the mission rides in the per-request
# user message via _retain_mission_preamble(). The {retain_mission_section}
# placeholder is kept (templates still reference it) but always empty here.
from hindsight_api.engine.prompt_utils import escape_for_prompt
retain_mission = getattr(config, "retain_mission", None)
if retain_mission:
retain_mission_section = (
f"══════════════════════════════════════════════════════════════════════════\n"
f"FOCUS — What to retain for this bank\n"
f"══════════════════════════════════════════════════════════════════════════\n\n"
f"{escape_for_prompt(retain_mission)}\n\n"
)
else:
retain_mission_section = ""
retain_mission_section = ""
# Select base prompt based on extraction mode
if extraction_mode == "custom":
@@ -997,6 +992,26 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
return prompt, response_schema
def _retain_mission_preamble(config) -> str:
"""The bank's retain mission, formatted for the per-request user message.
Kept OUT of the cached system prompt (which must stay bank-agnostic so one
CachedContent serves every bank — otherwise each distinct mission spawns its
own cache) and prepended to the user message instead. Returns "" when unset.
No brace-escaping needed: unlike the system template, the user message is
used verbatim, not passed through str.format().
"""
retain_mission = getattr(config, "retain_mission", None)
if not retain_mission:
return ""
return (
"══════════════════════════════════════════════════════════════════════════\n"
"FOCUS — What to retain for this bank (takes priority over the general guidelines)\n"
"══════════════════════════════════════════════════════════════════════════\n\n"
f"{retain_mission}\n\n"
)
def _build_user_message(
chunk: str,
chunk_index: int,
@@ -1005,8 +1020,14 @@ def _build_user_message(
context: str,
metadata: dict[str, str] | None = None,
agent_name: str | None = None,
mission_preamble: str = "",
) -> str:
"""Build user message for fact extraction."""
"""Build user message for fact extraction.
``mission_preamble`` (the bank's retain mission, possibly empty) is prepended
so the bank-specific focus lives in the variable user turn rather than the
cached, bank-agnostic system prompt.
"""
from .orchestrator import parse_datetime_flexible
sanitized_chunk = _sanitize_text(chunk)
@@ -1039,7 +1060,7 @@ def _build_user_message(
'statements to that speaker and classify them as "world", not "assistant".'
)
return f"""Extract facts from the following text chunk.
return f"""{mission_preamble}Extract facts from the following text chunk.
Chunk: {chunk_index + 1}/{total_chunks}
Event Date: {event_date_str}
@@ -1106,8 +1127,38 @@ async def _extract_facts_from_chunk(
extraction_mode = config.retain_extraction_mode
extract_causal_links = config.retain_extract_causal_links
# Build user message using helper function
user_message = _build_user_message(chunk, chunk_index, total_chunks, event_date, context, metadata, agent_name)
# Build user message — the bank mission rides here (not in the cached prefix).
user_message = _build_user_message(
chunk,
chunk_index,
total_chunks,
event_date,
context,
metadata,
agent_name,
mission_preamble=_retain_mission_preamble(config),
)
# Opt into context caching when the provider supports it. The prompt and
# response_schema are bank-agnostic (the mission lives in the user message),
# so one cached prefix serves every bank; reusing it across many small-payload
# retain calls dramatically lowers per-call input
# cost. ``get_or_create_cached_prefix`` returns None when caching is
# disabled, unsupported, or the prefix is too small; the LLM call
# transparently falls back to the uncached path in that case.
cached_prefix_name: str | None = None
provider_impl = getattr(llm_config, "_provider_impl", None)
if provider_impl is not None and provider_impl.supports_prompt_caching():
try:
cached_prefix_name = await provider_impl.get_or_create_cached_prefix(
system_instruction=prompt,
response_schema=response_schema,
)
except Exception:
# Caching is a soft optimisation — never let a cache-side
# error block a retain operation.
logger.exception("Cache prefix lookup failed; falling back to uncached call")
cached_prefix_name = None
# Retry logic for JSON validation errors
# Use retain-specific overrides if set, otherwise fall back to global LLM config
@@ -1128,7 +1179,7 @@ async def _extract_facts_from_chunk(
config.retain_llm_max_backoff if config.retain_llm_max_backoff is not None else config.llm_max_backoff
)
extraction_response_json, call_usage = await llm_config.call(
call_kwargs: dict[str, Any] = dict(
messages=[{"role": "system", "content": prompt}, {"role": "user", "content": user_message}],
response_format=response_schema,
scope="retain_extract_facts",
@@ -1140,6 +1191,10 @@ async def _extract_facts_from_chunk(
skip_validation=True, # Get raw JSON, we'll validate leniently
return_usage=True,
)
if cached_prefix_name is not None:
call_kwargs["cached_prefix"] = cached_prefix_name
extraction_response_json, call_usage = await llm_config.call(**call_kwargs)
usage = usage + call_usage # Aggregate usage across retries
# Lenient parsing of facts from raw JSON
@@ -1743,6 +1798,7 @@ async def extract_facts_from_contents_batch_api(
item.context,
item.metadata or None,
agent_name,
mission_preamble=_retain_mission_preamble(config),
)
# Build request body using helper function
+51 -4
View File
@@ -184,6 +184,8 @@ class MetricsCollectorBase:
input_tokens: int = 0,
output_tokens: int = 0,
success: bool = True,
cached_input_tokens: int = 0,
thoughts_tokens: int = 0,
):
"""
Record metrics for an LLM call.
@@ -193,9 +195,11 @@ class MetricsCollectorBase:
model: Model name
scope: Scope identifier (e.g., "memory", "reflect", "consolidation")
duration: Call duration in seconds
input_tokens: Number of input/prompt tokens
output_tokens: Number of output/completion tokens
input_tokens: Number of input/prompt tokens (total)
output_tokens: Number of output/completion tokens visible in candidates
success: Whether the call was successful
cached_input_tokens: Subset of input_tokens billed at the cached rate
thoughts_tokens: Reasoning tokens (billed as output, hidden from candidates)
"""
raise NotImplementedError
@@ -233,6 +237,8 @@ class NoOpMetricsCollector(MetricsCollectorBase):
input_tokens: int = 0,
output_tokens: int = 0,
success: bool = True,
cached_input_tokens: int = 0,
thoughts_tokens: int = 0,
):
"""No-op LLM call recording."""
pass
@@ -287,6 +293,27 @@ class MetricsCollector(MetricsCollectorBase):
name="hindsight.llm.calls.total", description="Total number of LLM API calls", unit="calls"
)
# Cached input tokens (subset of input_tokens billed at the cached rate).
# Useful for tracking prompt-cache hit-rate independently of total
# input volume. provider.scope.model labels matche llm_tokens_input.
self.llm_tokens_cached_input = self.meter.create_counter(
name="hindsight.llm.tokens.cached_input",
description="Number of cached input tokens (billed at cached rate) for LLM calls",
unit="tokens",
)
# Thinking / reasoning tokens (Gemini 2.5+ family). Billed at the
# output rate by the provider but invisible to candidates_token_count.
# Surfacing them as a distinct counter is required for honest cost
# attribution: a workload that "looks cheap" by output volume can be
# silently expensive if the model is doing long reasoning chains.
self.llm_tokens_thoughts = self.meter.create_counter(
name="hindsight.llm.tokens.thoughts",
description="Number of reasoning/thinking tokens emitted by the model "
"(billed as output but not surfaced in candidates)",
unit="tokens",
)
# HTTP request metrics
self.http_request_duration = self.meter.create_histogram(
name="hindsight.http.duration", description="Duration of HTTP requests in seconds", unit="s"
@@ -370,6 +397,8 @@ class MetricsCollector(MetricsCollectorBase):
input_tokens: int = 0,
output_tokens: int = 0,
success: bool = True,
cached_input_tokens: int = 0,
thoughts_tokens: int = 0,
):
"""
Record metrics for an LLM call.
@@ -379,9 +408,15 @@ class MetricsCollector(MetricsCollectorBase):
model: Model name
scope: Scope identifier (e.g., "memory", "reflect", "consolidation")
duration: Call duration in seconds
input_tokens: Number of input/prompt tokens
output_tokens: Number of output/completion tokens
input_tokens: Number of input/prompt tokens (total, including cached portion)
output_tokens: Number of output/completion tokens visible in candidates
success: Whether the call was successful
cached_input_tokens: Subset of input_tokens billed at the cached
rate (Gemini context caching). Defaults to 0 when caching is
disabled or the provider doesn't surface this field.
thoughts_tokens: Reasoning/thinking tokens (Gemini 2.5+ family).
Billed at the output rate but not counted in candidates.
Defaults to 0 for providers that don't emit thoughts.
"""
# Base attributes for all metrics
base_attributes = {
@@ -413,6 +448,18 @@ class MetricsCollector(MetricsCollectorBase):
}
self.llm_tokens_output.add(output_tokens, output_attributes)
if cached_input_tokens > 0:
self.llm_tokens_cached_input.add(
cached_input_tokens,
{**base_attributes, "token_bucket": get_token_bucket(cached_input_tokens)},
)
if thoughts_tokens > 0:
self.llm_tokens_thoughts.add(
thoughts_tokens,
{**base_attributes, "token_bucket": get_token_bucket(thoughts_tokens)},
)
@contextmanager
def record_http_request(self, method: str, endpoint: str, status_code_getter: Callable[[], int]):
"""
+57 -3
View File
@@ -3040,10 +3040,13 @@ def _make_mock_llm_one_obs_per_fact():
def callback(messages, scope):
if scope != "consolidation":
return _ConsolidationBatchResponse()
# Parse all fact UUIDs from the prompt — one create per fact
# Parse all fact UUIDs from the prompt — one create per fact. Read only
# the user message(s): consolidation sends the facts there, while the
# stable (cacheable) system message carries example UUIDs in its OUTPUT
# FORMAT samples that must not be mistaken for real facts.
import re
prompt = messages[0]["content"] if messages else ""
prompt = "\n".join(m.get("content", "") for m in messages if m.get("role") == "user")
fact_ids = re.findall(r"\[([0-9a-f-]{36})\]", prompt)
creates = [_CreateAction(text=f"Observation about fact {fid[:8]}", source_fact_ids=[fid]) for fid in fact_ids]
return _ConsolidationBatchResponse(creates=creates)
@@ -3145,7 +3148,9 @@ async def test_max_observations_per_scope_allows_updates_at_capacity(memory: Mem
call_count += 1
import re
prompt = messages[0]["content"] if messages else ""
# Facts live in the user message; the system message (stable, cached)
# carries example UUIDs in its OUTPUT samples — read user only.
prompt = "\n".join(m.get("content", "") for m in messages if m.get("role") == "user")
fact_ids = re.findall(r"\[([0-9a-f-]{36})\]", prompt)
if call_count == 1 and fact_ids:
# First call: create an observation
@@ -3512,3 +3517,52 @@ async def test_enable_auto_consolidation_flag(memory: MemoryEngine, request_cont
finally:
memory._config_resolver._global_config = original_global_config
await memory.delete_bank(bank_id, request_context=request_context)
def test_consolidation_prompt_split_is_cacheable_and_complete():
"""The split consolidation prompt: bank-agnostic system prefix + per-batch user.
The system prefix must be byte-identical across batches AND across banks (the
property that lets a single Gemini context cache serve every bank), carry only
stable instructions, and the per-batch/per-bank data (mission, facts,
observations, capacity note) must live in the user message — never in the
cached prefix.
"""
from hindsight_api.engine.consolidation.prompts import (
build_consolidation_input,
build_consolidation_system_prompt,
)
sys_prompt = build_consolidation_system_prompt()
# Byte-stable across calls and independent of any mission → one cache for all banks.
assert sys_prompt == build_consolidation_system_prompt()
# Instructions only: no per-batch placeholders leaked into the prefix.
assert "{facts_text}" not in sys_prompt
assert "{observations_text}" not in sys_prompt
# JSON examples are unescaped (single braces), i.e. .format() ran.
assert '{"creates"' in sys_prompt
assert "{{" not in sys_prompt
# The stable observation-format boilerplate lives in the cached prefix.
assert "proof_count" in sys_prompt
# Two banks with DIFFERENT missions share the identical cached prefix; the
# mission rides in the per-batch user message instead.
user_a = build_consolidation_input(
facts_text="[id-a] Fact A.", observations_text="[]", observations_mission="Track widgets."
)
user_b = build_consolidation_input(
facts_text="[id-b] Fact B.", observations_text="[]", observations_mission="Track gadgets."
)
assert "Track widgets." in user_a
assert "Track widgets." not in sys_prompt # mission NOT in the cached prefix
assert "Fact A." in user_a
assert user_a != user_b
# The format boilerplate is NOT re-sent per batch (it's in the cached prefix).
assert "proof_count" not in user_a
# The capacity note is per-batch too — kept out of the cached prefix.
capped = build_consolidation_input(
facts_text="[id] F.", observations_text="[]", observation_capacity_note="OBSERVATION LIMIT REACHED"
)
assert "OBSERVATION LIMIT REACHED" in capped
assert "OBSERVATION LIMIT REACHED" not in sys_prompt
@@ -109,7 +109,9 @@ def _mock_llm_one_obs_per_fact():
def callback(messages, scope):
if scope != "consolidation":
return _ConsolidationBatchResponse()
prompt = messages[0]["content"] if messages else ""
# Facts live in the user message; the system message (stable, cached) carries
# example UUIDs in its OUTPUT samples — read user only.
prompt = "\n".join(m.get("content", "") for m in messages if m.get("role") == "user")
fact_ids = re.findall(r"\[([0-9a-f-]{36})\]", prompt)
creates = [
_CreateAction(text=f"Observation about fact {fid[:8]}", source_fact_ids=[fid])
@@ -0,0 +1,469 @@
"""Unit tests for GeminiCacheManager.
The SDK's caches.create call is replaced with a fake throughout — no
network, no real Gemini calls. We assert:
* Identical prefixes return identical fingerprints (cache hits).
* Different prefixes return different fingerprints.
* The first get_or_create for a fingerprint creates; the second within
the TTL window reuses without calling the SDK again.
* "minimum token count" errors from Gemini surface as ``None`` (soft
fallback), not exceptions.
* Other SDK errors also surface as ``None`` so callers don't crash on
transient creation failures.
* The TTL refresh boundary recreates after the safety margin elapses.
"""
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
from hindsight_api.engine.providers.gemini_cache import GeminiCacheManager
def _make_client(create_side_effect=None):
"""Build a fake Gemini client whose ``aio.caches.create`` returns
a SimpleNamespace with ``.name`` (or raises the given exception)."""
create_mock = AsyncMock()
if isinstance(create_side_effect, Exception):
create_mock.side_effect = create_side_effect
elif callable(create_side_effect):
create_mock.side_effect = create_side_effect
else:
create_mock.return_value = SimpleNamespace(
name="cachedContents/test-cache-name-001"
)
client = MagicMock()
client.aio = MagicMock()
client.aio.caches = MagicMock()
client.aio.caches.create = create_mock
return client, create_mock
# ---- Fingerprint properties ----------------------------------------------
def test_fingerprint_is_stable():
fp1 = GeminiCacheManager.fingerprint(
model="gemini-3.1-flash-lite",
system_instruction="Extract facts.",
response_schema=None,
)
fp2 = GeminiCacheManager.fingerprint(
model="gemini-3.1-flash-lite",
system_instruction="Extract facts.",
response_schema=None,
)
assert fp1 == fp2
def test_fingerprint_changes_with_model():
fp1 = GeminiCacheManager.fingerprint("gemini-3.1-flash-lite", "X", None)
fp2 = GeminiCacheManager.fingerprint("gemini-3.1-flash", "X", None)
assert fp1 != fp2
def test_fingerprint_changes_with_system_instruction():
fp1 = GeminiCacheManager.fingerprint("m", "Extract facts.", None)
fp2 = GeminiCacheManager.fingerprint("m", "Extract entities.", None)
assert fp1 != fp2
def test_fingerprint_handles_pydantic_schema():
"""Two equivalent Pydantic schemas should fingerprint identically;
a different shape should not."""
from pydantic import BaseModel
class A(BaseModel):
x: int
y: str
class A_dup(BaseModel):
x: int
y: str
class B(BaseModel):
x: int
y: int # different type
fp_a = GeminiCacheManager.fingerprint("m", "p", A)
fp_dup = GeminiCacheManager.fingerprint("m", "p", A_dup)
fp_b = GeminiCacheManager.fingerprint("m", "p", B)
# A and A_dup have the same JSON schema, even though they're distinct classes.
assert fp_a == fp_dup
assert fp_a != fp_b
# ---- get_or_create lifecycle ---------------------------------------------
@pytest.mark.asyncio
async def test_first_call_creates_subsequent_reuses():
client, create_mock = _make_client()
mgr = GeminiCacheManager(client)
name1 = await mgr.get_or_create(
model="gemini-3.1-flash-lite",
system_instruction="Extract facts.",
response_schema=None,
)
name2 = await mgr.get_or_create(
model="gemini-3.1-flash-lite",
system_instruction="Extract facts.",
response_schema=None,
)
assert name1 == "cachedContents/test-cache-name-001"
assert name2 == name1
# Only ONE underlying create call — second was served from in-memory cache.
assert create_mock.call_count == 1
@pytest.mark.asyncio
async def test_different_prefixes_create_separately():
client, create_mock = _make_client(
create_side_effect=lambda *a, **kw: SimpleNamespace(
name=f"cachedContents/created-{create_mock.call_count}"
)
)
mgr = GeminiCacheManager(client)
name_a = await mgr.get_or_create(
model="m", system_instruction="A", response_schema=None
)
name_b = await mgr.get_or_create(
model="m", system_instruction="B", response_schema=None
)
assert name_a != name_b
assert create_mock.call_count == 2
# ---- Failure / fallback handling -----------------------------------------
@pytest.mark.asyncio
async def test_minimum_token_count_error_returns_none():
"""When the prefix is too short, Gemini rejects with a 'minimum
token count' style message. Manager must surface this as None so
the caller transparently falls back to a non-cached call."""
err = Exception("Cached content must have at least 1024 input tokens (minimum)")
client, _ = _make_client(create_side_effect=err)
mgr = GeminiCacheManager(client)
result = await mgr.get_or_create(
model="m", system_instruction="tiny", response_schema=None
)
assert result is None
@pytest.mark.asyncio
async def test_other_sdk_errors_also_return_none():
"""Transient errors (rate limits, 5xx, etc.) should fail soft so
a single bad create doesn't crash every retain call."""
err = RuntimeError("transient backend error 503")
client, _ = _make_client(create_side_effect=err)
mgr = GeminiCacheManager(client)
result = await mgr.get_or_create(
model="m", system_instruction="ok-sized prefix", response_schema=None
)
assert result is None
@pytest.mark.asyncio
async def test_failed_create_does_not_poison_cache():
"""If create fails on first attempt, a retry should call create
again instead of returning a stale/None entry."""
call_log = []
async def maybe_fail(*args, **kwargs):
call_log.append(1)
if len(call_log) == 1:
raise RuntimeError("first call fails")
return SimpleNamespace(name="cachedContents/recovered")
client = MagicMock()
client.aio = MagicMock()
client.aio.caches = MagicMock()
client.aio.caches.create = maybe_fail
mgr = GeminiCacheManager(client)
first = await mgr.get_or_create(
model="m", system_instruction="prefix", response_schema=None
)
second = await mgr.get_or_create(
model="m", system_instruction="prefix", response_schema=None
)
assert first is None
assert second == "cachedContents/recovered"
assert len(call_log) == 2
# ---- TTL behaviour --------------------------------------------------------
@pytest.mark.asyncio
async def test_refreshes_after_ttl_margin(monkeypatch):
"""An entry created at t=0 with ttl=10 and margin=2 should be
treated as stale at t>=8 and trigger a recreate."""
client, create_mock = _make_client(
create_side_effect=lambda *a, **kw: SimpleNamespace(
name=f"cachedContents/v{create_mock.call_count}"
)
)
mgr = GeminiCacheManager(client, ttl_seconds=10, refresh_margin_seconds=2)
fake_now = {"t": 1000.0}
monkeypatch.setattr(
"hindsight_api.engine.providers.gemini_cache.time.monotonic",
lambda: fake_now["t"],
)
first = await mgr.get_or_create(
model="m", system_instruction="p", response_schema=None
)
assert first == "cachedContents/v1"
# Advance to just before the refresh boundary — should reuse.
fake_now["t"] = 1000.0 + 7.0
again = await mgr.get_or_create(
model="m", system_instruction="p", response_schema=None
)
assert again == "cachedContents/v1"
assert create_mock.call_count == 1
# Advance past the refresh boundary — should recreate.
fake_now["t"] = 1000.0 + 9.0
refreshed = await mgr.get_or_create(
model="m", system_instruction="p", response_schema=None
)
assert refreshed == "cachedContents/v2"
assert create_mock.call_count == 2
# ---- Integration: feature flag + GeminiLLM accessor ----------------------
@pytest.mark.asyncio
async def test_gemini_llm_returns_none_when_cache_disabled():
"""A directly-constructed GeminiLLM (no prompt_cache_enabled kwarg) does not
cache: ``get_or_create_cached_prefix`` returns None without ever building a
cache manager. The server-level default-on flows in via the kwarg (resolved
from config in LLMProvider), not via this constructor default."""
from hindsight_api.engine.providers.gemini_llm import GeminiLLM
llm = GeminiLLM(
provider="gemini",
api_key="not-real-key",
base_url="",
model="gemini-test",
)
# Constructor default is off; even with a stable prefix the cache stays disabled.
result = await llm.get_or_create_cached_prefix(
system_instruction="A reasonably long system prompt " * 50,
response_schema=None,
)
assert result is None
assert llm._cache_manager is None # never built
@pytest.mark.asyncio
async def test_gemini_llm_uses_cache_when_enabled(monkeypatch):
"""When the flag is on, the manager is constructed lazily and its
get_or_create is delegated to. We don't hit the real SDK; we replace
the client's caches.create with a fake."""
from types import SimpleNamespace
from unittest.mock import AsyncMock
from hindsight_api.engine.providers.gemini_llm import GeminiLLM
llm = GeminiLLM(
provider="gemini",
api_key="not-real-key",
base_url="",
model="gemini-test",
prompt_cache_enabled=True,
)
# Replace the SDK-shaped client with a fake whose caches.create returns
# a predictable name. The lazy import inside get_or_create_cached_prefix
# picks up the patched module-level GeminiCacheManager naturally.
fake_create = AsyncMock(
return_value=SimpleNamespace(name="cachedContents/from-llm-test")
)
llm._client = MagicMock()
llm._client.aio = MagicMock()
llm._client.aio.caches = MagicMock()
llm._client.aio.caches.create = fake_create
name = await llm.get_or_create_cached_prefix(
system_instruction="A long enough system prompt for caching",
response_schema=None,
)
assert name == "cachedContents/from-llm-test"
# The manager was lazy-built on first use.
assert llm._cache_manager is not None
# Second call within TTL → no new SDK call.
again = await llm.get_or_create_cached_prefix(
system_instruction="A long enough system prompt for caching",
response_schema=None,
)
assert again == name
assert fake_create.call_count == 1
@pytest.mark.asyncio
async def test_call_falls_back_to_uncached_when_cache_400s():
"""A stale/invalid CachedContent makes the generate call 400. The provider
must drop the cache, invalidate the entry, and retry the SAME call inline
(prefix re-sent) so caching never breaks a request."""
import time
from types import SimpleNamespace
from unittest.mock import AsyncMock
from google.genai import errors as genai_errors
from hindsight_api.engine.providers.gemini_cache import GeminiCacheManager, _CacheEntry
from hindsight_api.engine.providers.gemini_llm import GeminiLLM
llm = GeminiLLM(provider="gemini", api_key="not-real-key", base_url="", model="gemini-test", prompt_cache_enabled=True)
# Seed a cache manager entry that maps to the (now invalid) cache name.
mgr = GeminiCacheManager(client=MagicMock())
mgr._entries["fp"] = _CacheEntry(name="cachedContents/stale", created_at=time.monotonic(), ttl_seconds=3300)
llm._cache_manager = mgr
captured = []
def _gen(*, model, contents, config):
captured.append(config)
if len(captured) == 1:
# First (cached) attempt — Gemini rejects the dead cache.
raise genai_errors.ClientError(
400, {"error": {"code": 400, "status": "INVALID_ARGUMENT", "message": "CachedContent not found"}}
)
# Retry without the cache succeeds.
return SimpleNamespace(
text="extracted",
usage_metadata=SimpleNamespace(
prompt_token_count=10, candidates_token_count=2, cached_content_token_count=0, thoughts_token_count=0
),
candidates=[SimpleNamespace(finish_reason="STOP")],
)
llm._client = MagicMock()
llm._client.aio = MagicMock()
llm._client.aio.models = MagicMock()
llm._client.aio.models.generate_content = AsyncMock(side_effect=_gen)
result = await llm.call(
messages=[{"role": "system", "content": "SYSTEM PREFIX"}, {"role": "user", "content": "doc"}],
cached_prefix="cachedContents/stale",
max_retries=2,
temperature=0.1,
)
# The request succeeded via the uncached retry.
assert result == "extracted"
assert len(captured) == 2
# First attempt referenced the cache; the retry inlined the prefix instead.
assert captured[0].cached_content == "cachedContents/stale"
assert captured[1].cached_content is None
assert captured[1].system_instruction == "SYSTEM PREFIX"
# The dead entry was invalidated so the next operation recreates it.
assert mgr._entries == {}
@pytest.mark.asyncio
async def test_create_cache_times_out_and_falls_back():
"""The create runs under the manager lock, so a hung caches.create would block
every concurrent caller (e.g. all chunks of a retain batch). It must time out
and return None so callers proceed uncached instead of stalling."""
import asyncio
from types import SimpleNamespace
from hindsight_api.engine.providers.gemini_cache import GeminiCacheManager
async def _hang(*args, **kwargs):
await asyncio.sleep(5)
return SimpleNamespace(name="never")
client = MagicMock()
client.aio = MagicMock()
client.aio.caches = MagicMock()
client.aio.caches.create = _hang
mgr = GeminiCacheManager(client, create_timeout_seconds=0.05)
result = await mgr.get_or_create(model="m", system_instruction="long enough prefix " * 20)
assert result is None
assert mgr._entries == {}
# ---- Tools: cache key + create wiring -----------------------------------
def test_fingerprint_changes_with_tools():
"""Two prefixes that differ ONLY in tools must hash differently —
otherwise a loop that adds a tool would silently reuse a stale
cache that doesn't know about it."""
tools_a = [
{"type": "function", "function": {"name": "search", "description": "search", "parameters": {}}}
]
tools_b = [
{"type": "function", "function": {"name": "search", "description": "search", "parameters": {}}},
{"type": "function", "function": {"name": "fetch", "description": "fetch", "parameters": {}}},
]
fp_a = GeminiCacheManager.fingerprint("m", "sys", None, tools=tools_a)
fp_b = GeminiCacheManager.fingerprint("m", "sys", None, tools=tools_b)
assert fp_a != fp_b
def test_fingerprint_stable_under_dict_reordering():
"""The tools list contains dicts; iteration order of dict keys
must not affect the fingerprint (otherwise upstream re-serialisation
would produce phantom cache misses)."""
tools_1 = [{"type": "function", "function": {"description": "d", "name": "n", "parameters": {"a": 1, "b": 2}}}]
tools_2 = [{"function": {"parameters": {"b": 2, "a": 1}, "name": "n", "description": "d"}, "type": "function"}]
fp_1 = GeminiCacheManager.fingerprint("m", "sys", None, tools=tools_1)
fp_2 = GeminiCacheManager.fingerprint("m", "sys", None, tools=tools_2)
assert fp_1 == fp_2
@pytest.mark.asyncio
async def test_get_or_create_passes_tools_to_create():
"""When tools are provided, the underlying caches.create call
must include them so the cached prefix actually contains the tool
definitions."""
captured = {}
async def fake_create(*, model, config):
captured["model"] = model
captured["config_dict"] = config.__dict__ if hasattr(config, "__dict__") else dict(config)
return SimpleNamespace(name="cachedContents/with-tools")
client = MagicMock()
client.aio = MagicMock()
client.aio.caches = MagicMock()
client.aio.caches.create = fake_create
mgr = GeminiCacheManager(client)
tools = [
{"type": "function", "function": {"name": "search", "description": "do a search", "parameters": {"type": "object"}}}
]
name = await mgr.get_or_create(
model="gemini-3.1-flash-lite",
system_instruction="You are a helpful tool-using assistant.",
tools=tools,
)
assert name == "cachedContents/with-tools"
# The Gemini SDK's CreateCachedContentConfig accepted a `tools` list.
cfg = captured["config_dict"]
assert "tools" in cfg, f"tools should be in cache config; got keys: {list(cfg.keys())}"
assert cfg["tools"], "tools list should be non-empty"
@@ -0,0 +1,294 @@
"""Measure the cached/input token ratio per operation against real Gemini.
Each operation re-sends a large constant prefix and a small variable payload:
- ``retain_extract_facts`` — fact-extraction system prompt + schema
- ``reflect_tool_call`` — agent system prompt + tool definitions, reused across
every iteration of the tool loop
- ``consolidation`` — the stable mission/rules/decision/output system prefix,
reused across every consolidation batch
We run real Gemini, route every call through the LLM-request tracer (#1922), and
read back recorded ``cached_tokens`` vs ``input_tokens`` per scope.
Two modes:
- Default (implicit only): Gemini's automatic caching — empirically ~0% for this
low-QPS access pattern. Structural assertions only; the ratio is a measurement.
- ``HINDSIGHT_GEMINI_EXPLICIT_CACHE=1``: enables PR #1936's explicit CachedContent
caching. Then each operation must visibly engage the cache (cached_tokens > 0
and ratio above a conservative floor) — this is the regression guard that the
caching actually works end-to-end through the retain/reflect/consolidation paths.
Gated on ``HINDSIGHT_RUN_GEMINI_EVALS=1`` plus a Gemini API key, since it costs
money and needs network. The default model is ``gemini-2.5-flash`` (override with
``HINDSIGHT_GEMINI_EVAL_MODEL``); explicit caching needs a >=2,048-token prefix.
"""
import asyncio
import os
import uuid
import pytest
from hindsight_api import MemoryEngine, RequestContext
from hindsight_api.engine.consolidation.consolidator import run_consolidation_job
from hindsight_api.engine.llm_trace import LLMRequestEntry
from hindsight_api.engine.llm_wrapper import LLMConfig
_GEMINI_API_KEY = (
os.getenv("HINDSIGHT_GEMINI_API_KEY") or os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
)
_RUN = os.getenv("HINDSIGHT_RUN_GEMINI_EVALS") == "1" and bool(_GEMINI_API_KEY)
pytestmark = pytest.mark.skipif(
not _RUN,
reason=(
"Gemini implicit-cache measurement is gated. Set HINDSIGHT_RUN_GEMINI_EVALS=1 "
"and provide GEMINI_API_KEY/GOOGLE_API_KEY to run."
),
)
# Number of retain "chunks": each is a separate retain call → a separate
# retain_extract_facts LLM call that re-sends the same ~3k-token system prefix.
# That repetition is the precondition for any caching (implicit or explicit) to
# kick in. Override with HINDSIGHT_GEMINI_CACHE_CHUNKS.
_CHUNKS = int(os.getenv("HINDSIGHT_GEMINI_CACHE_CHUNKS", "5"))
# Distinct paragraphs so each retain extracts real, non-duplicate facts. The
# *content* varies per call; the system prompt / schema prefix does not — which
# is exactly the shape caching targets.
_DOCS = [
"Ada Lovelace worked with Charles Babbage on the Analytical Engine in the 1840s. "
"She wrote what is often considered the first algorithm intended for a machine, "
"a method for computing Bernoulli numbers. She lived in London and corresponded "
"extensively with Babbage about the engine's capabilities.",
"Grace Hopper joined the Harvard Mark I team in 1944 and later developed the first "
"compiler, A-0, in 1952. She championed machine-independent programming languages, "
"which led to COBOL. She served in the US Navy and retired as a rear admiral.",
"Katherine Johnson computed orbital mechanics for NASA's first crewed spaceflights. "
"John Glenn personally asked her to verify the electronic computer's calculations "
"before his 1962 Friendship 7 orbit. She worked at Langley Research Center in Virginia.",
"Alan Turing formalized computation with the Turing machine in 1936 and worked at "
"Bletchley Park during World War II breaking the Enigma cipher. He proposed the "
"imitation game, now called the Turing test, in a 1950 paper on machine intelligence.",
"Margaret Hamilton led the software engineering team that wrote the onboard flight "
"software for the Apollo missions at MIT. Her error-detection code prevented an abort "
"during the Apollo 11 landing in 1969. She later coined the term 'software engineering'.",
"Barbara Liskov designed the CLU programming language in the 1970s and introduced data "
"abstraction. The Liskov substitution principle is named after her. She won the Turing "
"Award in 2008 for contributions to programming language and system design.",
"Tim Berners-Lee invented the World Wide Web in 1989 while at CERN, writing the first "
"browser and the HTTP protocol. He founded the World Wide Web Consortium in 1994 to "
"develop open web standards.",
"Radia Perlman invented the spanning-tree protocol while at Digital Equipment Corporation, "
"which made large bridged Ethernet networks possible. She is sometimes called the mother "
"of the internet, a title she has said she dislikes.",
]
# Explicit Gemini prompt caching (PR #1936) — opt-in. Set HINDSIGHT_GEMINI_EXPLICIT_CACHE=1
# to turn it on for this run. On a branch without the feature the flag is simply
# ignored, so the same test measures the implicit baseline there.
_EXPLICIT_CACHE = os.getenv("HINDSIGHT_GEMINI_EXPLICIT_CACHE") == "1"
async def _gemini_engine(memory_no_llm_verify: MemoryEngine) -> MemoryEngine:
"""Point an engine at real Gemini and force-enable the LLM-request tracer.
The fixture builds the engine with tracing disabled (config default). The
recorder reads ``enabled`` once at construction, so we flip the flag directly
rather than rebuilding the engine — equivalent to running with
``HINDSIGHT_API_LLM_TRACE_ENABLED=true``.
When ``HINDSIGHT_GEMINI_EXPLICIT_CACHE=1`` we also enable PR #1936's explicit
CachedContent caching the production way (env var + config-cache clear), so we
can compare its cached/input ratio against the implicit baseline. Otherwise the
only caching observed is Gemini's own implicit caching.
"""
from hindsight_api.config import clear_config_cache
model = os.getenv("HINDSIGHT_GEMINI_EVAL_MODEL", "gemini-2.5-flash")
# Prompt caching is on by default now, so the implicit-baseline run must
# explicitly DISABLE it (not just leave it unset) to measure Gemini's own
# implicit caching. Set the flag in both modes and clear the config cache so
# the per-bank resolver re-reads it.
os.environ["HINDSIGHT_API_LLM_PROMPT_CACHE_ENABLED"] = "true" if _EXPLICIT_CACHE else "false"
clear_config_cache()
cfg = LLMConfig(
provider="gemini",
api_key=_GEMINI_API_KEY or "",
base_url="",
model=model,
prompt_cache_enabled=_EXPLICIT_CACHE,
)
memory_no_llm_verify._llm_config = cfg
memory_no_llm_verify._retain_llm_config = cfg
memory_no_llm_verify._reflect_llm_config = cfg
memory_no_llm_verify._consolidation_llm_config = cfg
memory_no_llm_verify._llm_recorder._enabled = True
mode = "EXPLICIT cache ON" if _EXPLICIT_CACHE else "implicit only"
print(f"\n[gemini-cache] provider=gemini model={model} chunks={_CHUNKS} mode={mode}")
return memory_no_llm_verify
async def _drain_traces(mem: MemoryEngine) -> None:
"""Wait for the recorder's fire-and-forget trace writes to land.
record_llm_call schedules each INSERT as a detached asyncio task tracked in
``_pending`` (bucketed by trace_id). Gather them so the rows are queryable.
Loop a few times because consolidation's attach_memory_ids can spawn a
follow-up write after the first drain.
"""
await mem.wait_for_background_tasks()
rec = mem._llm_recorder
for _ in range(10):
pending = [t for bucket in rec._pending.values() for t in bucket if not t.done()]
if not pending:
break
await asyncio.gather(*pending, return_exceptions=True)
def _report(scope: str, rows: list[LLMRequestEntry]) -> float:
"""Print the cached/input token ratio for a scope and return it.
Gemini's ``prompt_token_count`` (our ``input_tokens``) already includes the
cached prefix, so ``cached_tokens / input_tokens`` is the fraction of prompt
tokens billed at the cheaper cached rate — the number the PR's
``cached_input / input`` dashboard would show.
"""
input_total = sum((r.input_tokens or 0) for r in rows)
cached_total = sum((r.cached_tokens or 0) for r in rows)
output_total = sum((r.output_tokens or 0) for r in rows)
ratio = (cached_total / input_total) if input_total else 0.0
per_call = ", ".join(f"{(r.cached_tokens or 0)}/{(r.input_tokens or 0)}" for r in rows)
mode = "explicit cache ON" if _EXPLICIT_CACHE else "implicit only"
print(
f"\n[gemini-cache] scope={scope!r} calls={len(rows)} ({mode})\n"
f" input_tokens = {input_total}\n"
f" cached_tokens = {cached_total}\n"
f" output_tokens = {output_total}\n"
f" cached/input = {ratio:.1%}\n"
f" per-call cached/input: {per_call}"
)
return ratio
@pytest.mark.hs_llm_core
class TestGeminiCacheRatioPerOperation:
"""Measure cached/input token ratio per operation (retain, reflect, consolidation).
Run with ``HINDSIGHT_GEMINI_EXPLICIT_CACHE=1`` to assert PR #1936's explicit
CachedContent caching actually engages (cached tokens > 0, ratio above a
conservative floor). Without it, the same tests record the implicit-caching
baseline (Gemini gives ~0% for this access pattern) without asserting a floor.
"""
async def _fetch(self, mem: MemoryEngine, bank_id: str, rc: RequestContext, scope: str) -> list[LLMRequestEntry]:
resp = await mem.list_llm_requests(bank_id, request_context=rc, scope=scope, limit=200)
assert resp is not None, "bank should exist"
return [r for r in resp.items if r.status == "success"]
def _assert(self, scope: str, rows: list[LLMRequestEntry], *, min_calls: int, min_ratio: float) -> float:
"""Common per-operation checks; returns the cached/input ratio.
``min_ratio`` is per-operation because the achievable ratio differs by
design: retain re-sends a pure fixed prefix (~90%); consolidation's prefix
is fixed but the facts/observations payload is large (~30%); reflect can
only cache its ``auto`` iterations — Gemini forbids ``cached_content`` with
a per-request ``tool_config`` — and the tool-result context grows, so the
ratio is modest (~10%). The universal guarantee in explicit mode is simply
that caching engaged at all (cached_tokens > 0).
"""
ratio = _report(scope, rows)
cached_total = sum((r.cached_tokens or 0) for r in rows)
assert len(rows) >= min_calls, f"expected >= {min_calls} {scope} calls, got {len(rows)}"
assert all((r.provider == "gemini") for r in rows)
assert sum((r.input_tokens or 0) for r in rows) > 0, "no input tokens recorded"
assert 0.0 <= ratio <= 1.0
if _EXPLICIT_CACHE:
assert cached_total > 0, f"{scope}: explicit cache ON but cached_tokens=0 (caching did not engage)"
assert ratio >= min_ratio, f"{scope}: cached/input {ratio:.1%} below floor {min_ratio:.0%}"
return ratio
async def test_retain_chunks_cached_ratio(self, memory_no_llm_verify, request_context):
"""Retain N distinct chunks → N fact-extraction calls sharing one prefix."""
mem = await _gemini_engine(memory_no_llm_verify)
bank_id = f"gemini-cache-retain-{uuid.uuid4().hex[:8]}"
await mem.get_bank_profile(bank_id, request_context=request_context)
docs = [_DOCS[i % len(_DOCS)] for i in range(_CHUNKS)]
for i, content in enumerate(docs):
await mem.retain_batch_async(
bank_id=bank_id,
contents=[{"content": content}],
request_context=request_context,
document_id=f"doc-{i}",
)
await _drain_traces(mem)
rows = await self._fetch(mem, bank_id, request_context, "retain_extract_facts")
self._assert("retain_extract_facts", rows, min_calls=_CHUNKS, min_ratio=0.5)
await mem.delete_bank(bank_id, request_context=request_context)
async def test_reflect_tool_loop_cached_ratio(self, memory_no_llm_verify, request_context):
"""Reflect runs an agentic tool loop; the system_prompt + tools prefix is
cached once and reused across every iteration (scope ``reflect_tool_call``)."""
mem = await _gemini_engine(memory_no_llm_verify)
bank_id = f"gemini-cache-reflect-{uuid.uuid4().hex[:8]}"
await mem.get_bank_profile(bank_id, request_context=request_context)
await mem.retain_batch_async(
bank_id=bank_id,
contents=[{"content": d} for d in _DOCS],
request_context=request_context,
)
await mem.wait_for_background_tasks()
# A broad question forces the agent to call recall/lookup tools, i.e. to
# iterate the tool loop more than once.
await mem.reflect_async(
bank_id=bank_id,
query="Who were the early pioneers of computing in these memories, and what is each one known for?",
request_context=request_context,
)
await _drain_traces(mem)
rows = await self._fetch(mem, bank_id, request_context, "reflect_tool_call")
if not rows:
# Diagnostic: dump every reflect_tool_call row (incl. errors) so a
# failure in the cached tool-loop path is visible, not silently skipped.
allresp = await mem.list_llm_requests(
bank_id, request_context=request_context, scope="reflect_tool_call", limit=200
)
for r in allresp.items if allresp else []:
print(f"\n[gemini-cache] reflect_tool_call status={r.status} error={r.error}")
pytest.skip("reflect made no SUCCESSFUL reflect_tool_call iterations for this seed")
self._assert("reflect_tool_call", rows, min_calls=1, min_ratio=0.0)
await mem.delete_bank(bank_id, request_context=request_context)
async def test_consolidation_cached_ratio(self, memory_no_llm_verify, request_context):
"""Retain a batch, then consolidate; the stable system prefix is cached and
reused across every consolidation batch (scope ``consolidation``)."""
mem = await _gemini_engine(memory_no_llm_verify)
bank_id = f"gemini-cache-consol-{uuid.uuid4().hex[:8]}"
await mem.get_bank_profile(bank_id, request_context=request_context)
# Seed enough unconsolidated memories that consolidation makes several
# same-prefix LLM calls.
await mem.retain_batch_async(
bank_id=bank_id,
contents=[{"content": d} for d in _DOCS],
request_context=request_context,
)
await mem.wait_for_background_tasks()
await run_consolidation_job(mem, bank_id, request_context)
await _drain_traces(mem)
rows = await self._fetch(mem, bank_id, request_context, "consolidation")
if not rows:
pytest.skip("consolidation made no LLM calls for this seed (nothing to consolidate)")
self._assert("consolidation", rows, min_calls=1, min_ratio=0.15)
await mem.delete_bank(bank_id, request_context=request_context)
+8 -6
View File
@@ -67,9 +67,10 @@ class TestMetricsCollector:
# Create separate mocks for each histogram (operation_duration, llm_duration, http_request_duration)
histogram_mocks = [MagicMock(), MagicMock(), MagicMock()]
meter.create_histogram.side_effect = histogram_mocks
# Create separate mocks for each counter
# (operation_total, llm_tokens_input, llm_tokens_output, llm_calls_total, http_requests_total)
counter_mocks = [MagicMock() for _ in range(5)]
# Create separate mocks for each counter (operation_total, llm_tokens_input,
# llm_tokens_output, llm_calls_total, llm_tokens_cached_input,
# llm_tokens_thoughts, http_requests_total)
counter_mocks = [MagicMock() for _ in range(7)]
meter.create_counter.side_effect = counter_mocks
return meter
@@ -278,9 +279,10 @@ class TestLLMMetrics:
# Create separate mocks for each histogram (operation_duration, llm_duration, http_request_duration)
histogram_mocks = [MagicMock(), MagicMock(), MagicMock()]
meter.create_histogram.side_effect = histogram_mocks
# Create separate mocks for each counter
# (operation_total, llm_tokens_input, llm_tokens_output, llm_calls_total, http_requests_total)
counter_mocks = [MagicMock() for _ in range(5)]
# Create separate mocks for each counter (operation_total, llm_tokens_input,
# llm_tokens_output, llm_calls_total, llm_tokens_cached_input,
# llm_tokens_thoughts, http_requests_total)
counter_mocks = [MagicMock() for _ in range(7)]
meter.create_counter.side_effect = counter_mocks
return meter
@@ -8,7 +8,6 @@ import pytest
from hindsight_api.engine.prompt_utils import escape_for_prompt
# ---------------------------------------------------------------------------
# Unit tests for the shared escape helper
# ---------------------------------------------------------------------------
@@ -117,12 +116,16 @@ class TestRetainBraceSafety:
def test_retain_mission_with_json(self):
from hindsight_api.engine.retain.fact_extraction import (
_build_extraction_prompt_and_schema,
_retain_mission_preamble,
)
config = self._make_config(retain_mission='{"focus": "compliance"}')
prompt, _ = _build_extraction_prompt_and_schema(config)
# prompt is already fully rendered (no remaining placeholders)
assert '{"focus": "compliance"}' in prompt
# The mission no longer lives in the (cached, bank-agnostic) system prompt;
# it rides in the per-request user-message preamble, verbatim and unescaped
# (the preamble is not passed through str.format(), so braces are safe).
assert '{"focus": "compliance"}' not in prompt
assert '{"focus": "compliance"}' in _retain_mission_preamble(config)
def test_custom_instructions_with_braces(self):
from hindsight_api.engine.retain.fact_extraction import (
@@ -139,6 +142,7 @@ class TestRetainBraceSafety:
def test_both_mission_and_custom_with_braces(self):
from hindsight_api.engine.retain.fact_extraction import (
_build_extraction_prompt_and_schema,
_retain_mission_preamble,
)
config = self._make_config(
@@ -147,5 +151,6 @@ class TestRetainBraceSafety:
retain_custom_instructions="Format: {k: v}",
)
prompt, _ = _build_extraction_prompt_and_schema(config)
assert '{"scope": "all"}' in prompt
# Mission → user-message preamble; custom instructions stay in the system prompt.
assert '{"scope": "all"}' in _retain_mission_preamble(config)
assert "{k: v}" in prompt
+72 -9
View File
@@ -2770,30 +2770,93 @@ async def test_retain_batch_with_per_item_tags_on_document(memory, request_conte
print(f"\n=== Cleaned up bank: {bank_id} ===")
def test_retain_mission_injected_into_prompt():
"""Test that retain_mission is injected as a FOCUS section into any extraction mode."""
def test_retain_mission_in_user_preamble_not_cached_prefix():
"""retain_mission rides in the per-request user-message preamble, NOT the
system prompt — so the cached system prefix stays bank-agnostic and a single
Gemini context cache can serve every bank. Independent of extraction mode."""
from unittest.mock import MagicMock
from hindsight_api.engine.retain.fact_extraction import _build_extraction_prompt_and_schema
from hindsight_api.engine.retain.fact_extraction import (
_build_extraction_prompt_and_schema,
_retain_mission_preamble,
)
spec = "Focus on technical decisions and architecture choices only."
# Test with concise mode
config = MagicMock()
config.retain_extraction_mode = "concise"
config.retain_mission = spec
config.retain_custom_instructions = None
config.retain_extract_causal_links = False
# The mission is absent from the (cacheable, bank-agnostic) system prompt...
prompt, _ = _build_extraction_prompt_and_schema(config)
assert spec in prompt
assert "FOCUS" in prompt
assert spec not in prompt
assert "FOCUS" not in prompt
# ...and present in the per-request user-message preamble instead.
preamble = _retain_mission_preamble(config)
assert spec in preamble
assert "FOCUS" in preamble
# retain_mission is injected into verbose mode as well
# Mode-independent: verbose mode → same mission-free prompt, same preamble.
config.retain_extraction_mode = "verbose"
prompt_verbose, _ = _build_extraction_prompt_and_schema(config)
assert spec in prompt_verbose
assert "FOCUS" in prompt_verbose
assert spec not in prompt_verbose
assert spec in _retain_mission_preamble(config)
# The payoff: two banks with DIFFERENT missions produce the IDENTICAL system
# prompt → the same cache fingerprint → one shared CachedContent for both,
# instead of one cache per mission.
config.retain_extraction_mode = "concise"
config.retain_mission = "Track project A architecture decisions."
prompt_a, _ = _build_extraction_prompt_and_schema(config)
config.retain_mission = "Track customer B support incidents."
prompt_b, _ = _build_extraction_prompt_and_schema(config)
assert prompt_a == prompt_b
@pytest.mark.parametrize("mode", ["concise", "verbose"])
def test_retain_cacheable_prefix_invariant_to_per_bank_freetext(mode):
"""The cacheable system prompt must NOT depend on per-bank free-text settings.
For the concise and verbose modes (the ones we cache), the system prefix has
to be byte-identical across banks so a single Gemini CachedContent serves all
of them. The high-cardinality, user-supplied free-text fields — the retain
mission, and custom instructions (which only apply to custom mode anyway) —
must never leak into it; if they did, every bank would fragment the cache.
Structural toggles (causal links, entity labels, output language) MAY change
the prefix — they are low-cardinality and correctly keyed by the cache
fingerprint — so they are deliberately NOT exercised here.
"""
from types import SimpleNamespace
from hindsight_api.engine.retain.fact_extraction import _build_extraction_prompt_and_schema
def make(**overrides):
defaults = {
"retain_extraction_mode": mode,
"retain_extract_causal_links": False,
"retain_mission": None,
"retain_custom_instructions": None,
"retain_taxonomy": None,
"entity_labels": None,
"entities_allow_free_form": True,
"llm_output_language": None,
}
defaults.update(overrides)
return SimpleNamespace(**defaults)
baseline, _ = _build_extraction_prompt_and_schema(make())
# The mission must not change the cacheable prefix, whatever its value.
for mission in ["Track A decisions.", '{"focus": "compliance"}', "Ünïcödé brief", "x" * 600]:
prompt, _ = _build_extraction_prompt_and_schema(make(retain_mission=mission))
assert prompt == baseline, f"retain_mission leaked into the cacheable {mode} prefix"
# Custom instructions are a custom-mode field; they must not touch concise/verbose.
prompt, _ = _build_extraction_prompt_and_schema(make(retain_custom_instructions="Do X with {braces}"))
assert prompt == baseline, f"retain_custom_instructions leaked into the cacheable {mode} prefix"
def test_retain_mission_absent_when_not_set():
+14
View File
@@ -1,5 +1,6 @@
import {LLMProvidersGrid} from '@site/src/components/SupportedGrids';
import {LLMProvidersTable} from '@site/src/components/LLMProvidersTable';
import {LLMProviderCapabilities} from '@site/src/components/LLMProviderCapabilities';
# Models
@@ -57,6 +58,19 @@ Set `HINDSIGHT_API_LLM_PROVIDER=litellmrouter` to run the default LLM through [L
See [Configuration](./configuration#llm-router-litellm-router) for setup.
:::
### Provider Capabilities
Beyond basic generation, some providers support optional features that lower cost or latency. Hindsight uses each feature automatically when the configured provider supports it.
<LLMProviderCapabilities />
- **Batch API** — submits bulk retain extraction through the provider's asynchronous batch endpoint, typically at ~50% lower cost. Used automatically when available; otherwise calls run synchronously.
- **Explicit prompt caching** — reuses the large, fixed system prefix that retain (fact extraction), consolidation, and the reflect tool-loop send on every call, billing it at the provider's cached-input rate. On Gemini/Vertex this uses the `CachedContent` API. **On by default**; disable with `HINDSIGHT_API_LLM_PROMPT_CACHE_ENABLED=false`. Hindsight structures these prompts so the cached prefix is **bank-agnostic** — one cache is shared across all banks rather than one per bank/mission, and creation soft-fails to an uncached call, so it never breaks a request.
:::note
A blank "Explicit prompt caching" cell does not mean a provider has no caching. OpenAI, for example, caches a stable leading prompt prefix **automatically** server-side, so it benefits with no configuration; Anthropic supports caching via `cache_control` breakpoints which can be wired up through the same provider hook. The column tracks only Hindsight's explicit `get_or_create_cached_prefix` hook, which Gemini/Vertex implement today.
:::
### Benchmarks
Not sure which model to use? The **[Model Leaderboard](https://benchmarks.hindsight.vectorize.io/)** benchmarks models across accuracy, speed, cost, and reliability for retain, reflect, and observation consolidation so you can pick the right trade-off for your use case.
@@ -0,0 +1,36 @@
import React from 'react';
import {LLM_PROVIDERS} from '../data/llmProviders';
/**
* Renders the "Provider Capabilities" table from the single-source-of-truth
* provider list in `src/data/llmProviders.tsx`. Adding a provider (or a
* capability flag) is an edit to `llmProviders.json` only — this table, the
* default-models table, and the icon grid all derive from it.
*
* Only real providers are listed (the empty-id "OpenAI Compatible" pseudo-entry
* is skipped). A blank cell means the capability is not supported.
*/
export function LLMProviderCapabilities() {
const rows = LLM_PROVIDERS.filter(p => p.id);
const cell = (on?: boolean) => (on ? '✅' : '—');
return (
<table>
<thead>
<tr>
<th>Provider</th>
<th style={{textAlign: 'center'}}>Batch API</th>
<th style={{textAlign: 'center'}}>Explicit prompt caching</th>
</tr>
</thead>
<tbody>
{rows.map(({id, label, batchApi, promptCaching}) => (
<tr key={id}>
<td>{label} (<code>{id}</code>)</td>
<td style={{textAlign: 'center'}}>{cell(batchApi)}</td>
<td style={{textAlign: 'center'}}>{cell(promptCaching)}</td>
</tr>
))}
</tbody>
</table>
);
}
+4 -4
View File
@@ -1,9 +1,9 @@
[
{"id": "openai", "label": "OpenAI", "iconKey": "openai", "defaultModel": "gpt-4o-mini"},
{"id": "openai", "label": "OpenAI", "iconKey": "openai", "defaultModel": "gpt-4o-mini", "batchApi": true},
{"id": "anthropic", "label": "Anthropic", "iconKey": "anthropic", "defaultModel": "claude-haiku-4-5-20251001"},
{"id": "gemini", "label": "Google Gemini", "iconKey": "gemini", "defaultModel": "gemini-2.5-flash"},
{"id": "vertexai", "label": "Vertex AI", "iconKey": "gemini", "defaultModel": "gemini-2.0-flash-001"},
{"id": "groq", "label": "Groq", "iconKey": "zap", "defaultModel": "openai/gpt-oss-120b"},
{"id": "gemini", "label": "Google Gemini", "iconKey": "gemini", "defaultModel": "gemini-2.5-flash", "promptCaching": true},
{"id": "vertexai", "label": "Vertex AI", "iconKey": "gemini", "defaultModel": "gemini-2.0-flash-001", "promptCaching": true},
{"id": "groq", "label": "Groq", "iconKey": "zap", "defaultModel": "openai/gpt-oss-120b", "batchApi": true},
{"id": "ollama", "label": "Ollama", "iconKey": "ollama", "defaultModel": "gemma3:12b"},
{"id": "ollama-cloud", "label": "Ollama Cloud", "iconKey": "ollama", "defaultModel": "gemma3:12b"},
{"id": "lmstudio", "label": "LM Studio", "iconKey": "brain", "defaultModel": "local-model"},
+13 -4
View File
@@ -42,6 +42,10 @@ export interface LLMProvider {
defaultModel?: string;
/** Optional note rendered in the default-models table. */
defaultModelNote?: string;
/** Supports the asynchronous Batch API (supports_batch_api in the engine). */
batchApi?: boolean;
/** Supports explicit prompt-prefix caching (supports_prompt_caching in the engine). */
promptCaching?: boolean;
}
/**
@@ -51,14 +55,19 @@ export interface LLMProvider {
* consumed by:
* - this module (resolves iconKey -> IconType for the icon grid)
* - LLMProvidersTable React component (renders the default-models table)
* - scripts/generate-docs-skill.sh (renders <LLMProvidersTable /> as
* markdown when copying MDX docs into the agent-facing skill)
* - LLMProviderCapabilities React component (renders the capability table:
* batchApi / promptCaching)
* - scripts/generate-docs-skill.sh (renders <LLMProvidersTable />,
* <LLMProvidersGrid /> and <LLMProviderCapabilities /> as markdown when
* copying MDX docs into the agent-facing skill)
*
* Keep aligned with PROVIDER_DEFAULT_MODELS in
* hindsight-api-slim/hindsight_api/config.py.
* Keep the default model aligned with PROVIDER_DEFAULT_MODELS, and the
* capability flags with supports_batch_api() / supports_prompt_caching() on the
* provider classes, in hindsight-api-slim/hindsight_api/.
*/
export const LLM_PROVIDERS: LLMProvider[] = (providersJson as Array<{
id: string; label: string; iconKey: string; defaultModel?: string; defaultModelNote?: string;
batchApi?: boolean; promptCaching?: boolean;
}>).map(({iconKey, ...rest}) => {
const icon = ICON_REGISTRY[iconKey];
if (!icon) throw new Error(`Unknown iconKey "${iconKey}" for provider "${rest.id || rest.label}"`);
+18
View File
@@ -155,6 +155,24 @@ def render_llm_providers_grid(_match):
content = re.sub(r'<LLMProvidersGrid\s*/>', render_llm_providers_grid, content)
# Render <LLMProviderCapabilities /> as a markdown capability table, sourced from
# the same single-source-of-truth provider list (batchApi / promptCaching flags).
def render_llm_capabilities_table(_match):
providers = json.loads(llm_providers_json.read_text())
rows = [
"| Provider | Batch API | Explicit prompt caching |",
"|----------|:---------:|:-----------------------:|",
]
for p in providers:
if not p.get("id"):
continue
batch = "✅" if p.get("batchApi") else "—"
cache = "✅" if p.get("promptCaching") else "—"
rows.append(f"| {p['label']} (`{p['id']}`) | {batch} | {cache} |")
return "\n".join(rows)
content = re.sub(r'<LLMProviderCapabilities\s*/>', render_llm_capabilities_table, content)
# Convert <Tabs> to markdown sections
# Replace <Tabs> ... </Tabs> with markdown headers
content = re.sub(r'<Tabs>\s*', '', content)
@@ -69,6 +69,37 @@ See [Configuration](./configuration#llm-provider) for setup examples.
Set `HINDSIGHT_API_LLM_PROVIDER=litellmrouter` to run the default LLM through [LiteLLM's Router](https://docs.litellm.ai/docs/routing) — ordered fallback across deployments, load-balanced same-tier routing, weighted picks, per-deployment `rpm`/`tpm` limits, and cooldowns are all available via the [`Router` config](https://docs.litellm.ai/docs/routing#fallbacks). Hindsight passes the JSON config through verbatim.
See [Configuration](./configuration#llm-router-litellm-router) for setup.
### Provider Capabilities
Beyond basic generation, some providers support optional features that lower cost or latency. Hindsight uses each feature automatically when the configured provider supports it.
| Provider | Batch API | Explicit prompt caching |
|----------|:---------:|:-----------------------:|
| OpenAI (`openai`) | ✅ | — |
| Anthropic (`anthropic`) | — | — |
| Google Gemini (`gemini`) | — | ✅ |
| Vertex AI (`vertexai`) | — | ✅ |
| Groq (`groq`) | ✅ | — |
| Ollama (`ollama`) | — | — |
| Ollama Cloud (`ollama-cloud`) | — | — |
| LM Studio (`lmstudio`) | — | — |
| llama.cpp (`llamacpp`) | — | — |
| MiniMax (`minimax`) | — | — |
| DeepSeek (`deepseek`) | — | — |
| z.ai (`zai`) | — | — |
| opencode-go (`opencode-go`) | — | — |
| Volcano Engine (`volcano`) | — | — |
| OpenRouter (`openrouter`) | — | — |
| OpenAI Codex (`openai-codex`) | — | — |
| Claude Code (`claude-code`) | — | — |
| AWS Bedrock (`bedrock`) | — | — |
| LiteLLM (100+) (`litellm`) | — | — |
- **Batch API** — submits bulk retain extraction through the provider's asynchronous batch endpoint, typically at ~50% lower cost. Used automatically when available; otherwise calls run synchronously.
- **Explicit prompt caching** — reuses the large, fixed system prefix that retain (fact extraction), consolidation, and the reflect tool-loop send on every call, billing it at the provider's cached-input rate. On Gemini/Vertex this uses the `CachedContent` API. **On by default**; disable with `HINDSIGHT_API_LLM_PROMPT_CACHE_ENABLED=false`. Hindsight structures these prompts so the cached prefix is **bank-agnostic** — one cache is shared across all banks rather than one per bank/mission, and creation soft-fails to an uncached call, so it never breaks a request.
:::note
A blank "Explicit prompt caching" cell does not mean a provider has no caching. OpenAI, for example, caches a stable leading prompt prefix **automatically** server-side, so it benefits with no configuration; Anthropic supports caching via `cache_control` breakpoints which can be wired up through the same provider hook. The column tracks only Hindsight's explicit `get_or_create_cached_prefix` hook, which Gemini/Vertex implement today.
### Benchmarks
Not sure which model to use? The **[Model Leaderboard](https://benchmarks.hindsight.vectorize.io/)** benchmarks models across accuracy, speed, cost, and reliability for retain, reflect, and observation consolidation so you can pick the right trade-off for your use case.