Compare commits
20
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2ff5d7acb2 | ||
|
|
1910a5fb50 | ||
|
|
7dfebb1147 | ||
|
|
f64a816361 | ||
|
|
2954f0d2a5 | ||
|
|
b5da24a1bf | ||
|
|
477342ce94 | ||
|
|
f1fa1904e0 | ||
|
|
e6a519dc06 | ||
|
|
e1ff77954a | ||
|
|
2fbcf35b54 | ||
|
|
2d56782287 | ||
|
|
9b9769baeb | ||
|
|
fc173a6c5e | ||
|
|
7a3e42ae74 | ||
|
|
ec0920e435 | ||
|
|
08a16f1dde | ||
|
|
9d9ac2b722 | ||
|
|
2527c75970 | ||
|
|
d458b715fe |
@@ -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
|
||||
|
||||
@@ -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]):
|
||||
"""
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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>
|
||||
);
|
||||
}
|
||||
@@ -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"},
|
||||
|
||||
@@ -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}"`);
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user