Skip to content

Commit ec1bbb4

Browse files
devonakelleyclaude
andcommitted
feat: capture prompt cache hits in routing outcomes
Track Anthropic cache_read_input_tokens and OpenAI cached_tokens. Report cache_hit and cache_tokens in outcome signals. Add prefer_cached param to Router for future cache-aware routing. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 6cfed3a commit ec1bbb4

2 files changed

Lines changed: 129 additions & 7 deletions

File tree

‎kalibr/intelligence.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,8 @@ def report_outcome(
187187
execution_params: dict | None = None,
188188
model_id: str | None = None,
189189
pipeline_id: str | None = None,
190+
cache_hit: bool | None = None,
191+
cache_tokens: int = 0,
190192
) -> dict[str, Any]:
191193
"""Report execution outcome for a goal.
192194
@@ -251,6 +253,9 @@ def report_outcome(
251253
}
252254
if pipeline_id is not None:
253255
payload["pipeline_id"] = pipeline_id
256+
if cache_hit is not None:
257+
payload["cache_hit"] = cache_hit
258+
payload["cache_tokens"] = int(cache_tokens or 0)
254259

255260
response = self._request(
256261
"POST",

‎kalibr/router.py‎

Lines changed: 124 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,8 @@ def _auto_report_outcome(
5858
score: Optional[float] = None,
5959
failure_category: Optional[str] = None,
6060
failure_reason: Optional[str] = None,
61+
cache_hit: Optional[bool] = None,
62+
cache_tokens: int = 0,
6163
) -> None:
6264
"""Fire-and-forget outcome report. Swallows all errors."""
6365
try:
@@ -71,10 +73,46 @@ def _auto_report_outcome(
7173
failure_reason=failure_reason,
7274
pipeline_id=pipeline_id,
7375
model_id=model_id,
76+
cache_hit=cache_hit,
77+
cache_tokens=cache_tokens,
7478
)
7579
except Exception:
7680
pass
7781

82+
83+
def _extract_cache_info(response: Any) -> tuple[Optional[bool], int]:
84+
"""Extract prompt-cache hit info from a provider response.
85+
86+
Returns ``(cache_hit, cache_tokens)``. ``cache_hit`` is ``None`` when the
87+
provider did not surface any cache field (so callers can distinguish
88+
"unknown" from a confirmed miss). ``cache_tokens`` is the number of tokens
89+
served from cache (0 on miss or unknown).
90+
91+
Supports Anthropic's ``usage.cache_read_input_tokens`` and OpenAI's
92+
``usage.prompt_tokens_details.cached_tokens``. Fails open: any exception
93+
returns ``(None, 0)``.
94+
"""
95+
try:
96+
usage = getattr(response, "usage", None)
97+
if usage is None:
98+
return None, 0
99+
100+
anthropic_cache = getattr(usage, "cache_read_input_tokens", None)
101+
if anthropic_cache is not None:
102+
tokens = int(anthropic_cache or 0)
103+
return tokens > 0, tokens
104+
105+
details = getattr(usage, "prompt_tokens_details", None)
106+
if details is not None:
107+
openai_cache = getattr(details, "cached_tokens", None)
108+
if openai_cache is not None:
109+
tokens = int(openai_cache or 0)
110+
return tokens > 0, tokens
111+
112+
return None, 0
113+
except Exception:
114+
return None, 0
115+
78116
# Type for paths - either string or dict
79117
PathSpec = Union[str, Dict[str, Any]]
80118

@@ -599,6 +637,7 @@ def __init__(
599637
exploration_rate: Optional[float] = None,
600638
auto_register: bool = True,
601639
session_id: Optional[str] = None,
640+
prefer_cached: bool = False,
602641
):
603642
"""
604643
Initialize router.
@@ -639,6 +678,12 @@ def __init__(
639678
momentum and may escalate the model choice if the session is
640679
widening (user frustrated/failing). Falls back to the
641680
KALIBR_SESSION_ID environment variable when not passed.
681+
prefer_cached: When True, future routing decisions will prefer
682+
providers/paths that the intelligence service reports as having
683+
warm prompt caches for this goal. Wired in for forward
684+
compatibility; the intelligence service does not yet return
685+
cache-affinity data, so today the flag is captured on the
686+
Router and reported in outcomes but does not yet bias routing.
642687
"""
643688
self.goal = goal
644689

@@ -664,6 +709,7 @@ def __init__(
664709
self.repair_prompt = repair_prompt
665710
self.exploration_rate = exploration_rate
666711
self.session_id = session_id or os.environ.get("KALIBR_SESSION_ID") or None
712+
self.prefer_cached = prefer_cached
667713
self._last_trace_id: Optional[str] = None
668714
self._last_model_id: Optional[str] = None
669715
self._last_decision: Optional[dict] = None
@@ -873,6 +919,8 @@ def _heal_dispatch(m_id: str, msgs: List[Dict], system_prompt: Optional[str] = N
873919
self._last_model_id = used_model
874920

875921
if heal_result.get("success"):
922+
heal_response = heal_result.get("response")
923+
heal_cache_hit, heal_cache_tokens = _extract_cache_info(heal_response)
876924
heal_count_val = int(heal_result.get("heal_count") or 0)
877925
if heal_count_val > 0:
878926
# Heal telemetry: fire-and-forget signal so the dashboard
@@ -887,6 +935,8 @@ def _heal_dispatch(m_id: str, msgs: List[Dict], system_prompt: Optional[str] = N
887935
"model_id": used_model,
888936
"score": 0.8,
889937
"failure_category": "healed",
938+
"cache_hit": heal_cache_hit,
939+
"cache_tokens": heal_cache_tokens,
890940
},
891941
daemon=True,
892942
).start()
@@ -899,19 +949,23 @@ def _heal_dispatch(m_id: str, msgs: List[Dict], system_prompt: Optional[str] = N
899949
success=True,
900950
model_id=used_model,
901951
pipeline_id=pipeline_id,
952+
cache_hit=heal_cache_hit,
953+
cache_tokens=heal_cache_tokens,
902954
)
903955
self._outcome_reported = True
904956
except Exception:
905957
pass
906958

907-
response = heal_result.get("response")
959+
response = heal_response
908960
if response is None:
909961
response = _make_chat_completion(heal_result.get("result") or "", model=used_model)
910962
response.kalibr_trace_id = trace_id
911963
response.kalibr_healed = bool(heal_result.get("healed"))
912964
response.kalibr_heal_count = heal_result.get("heal_count", 0)
913965
response.kalibr_models_tried = heal_result.get("models_tried") or []
914966
response.kalibr_model_used = used_model
967+
response.kalibr_cache_hit = heal_cache_hit
968+
response.kalibr_cache_tokens = heal_cache_tokens
915969
return response
916970

917971
# Heal loop exhausted all paths — fire failure telemetry.
@@ -989,6 +1043,10 @@ def _heal_dispatch(m_id: str, msgs: List[Dict], system_prompt: Optional[str] = N
9891043
# Success! Update state to reflect which model succeeded
9901044
self._last_model_id = candidate_model
9911045

1046+
# Capture prompt-cache hit info from the provider response so
1047+
# downstream outcome reports can carry it. Fails open.
1048+
cache_hit, cache_tokens = _extract_cache_info(response)
1049+
9921050
# Gate 2: LLM-as-a-judge (optional, runs if judge_model configured)
9931051
if self.judge_model and not self._outcome_reported:
9941052
try:
@@ -1012,7 +1070,15 @@ def _heal_dispatch(m_id: str, msgs: List[Dict], system_prompt: Optional[str] = N
10121070
repair_score = self._run_judge(repair_output, repaired_messages)
10131071
if repair_score >= self.judge_threshold:
10141072
logger.info(f"Prompt repair succeeded (score={repair_score:.2f})")
1015-
self.report(success=True, score=repair_score, pipeline_id=pipeline_id)
1073+
# Re-extract cache info from the repaired response
1074+
cache_hit, cache_tokens = _extract_cache_info(repair_response)
1075+
self.report(
1076+
success=True,
1077+
score=repair_score,
1078+
pipeline_id=pipeline_id,
1079+
cache_hit=cache_hit,
1080+
cache_tokens=cache_tokens,
1081+
)
10161082
self._outcome_reported = True
10171083
response = repair_response
10181084
# Skip to return — fall through to response return below
@@ -1058,7 +1124,13 @@ def _heal_dispatch(m_id: str, msgs: List[Dict], system_prompt: Optional[str] = N
10581124
pass
10591125
continue
10601126
else:
1061-
self.report(success=True, score=judge_score, pipeline_id=pipeline_id)
1127+
self.report(
1128+
success=True,
1129+
score=judge_score,
1130+
pipeline_id=pipeline_id,
1131+
cache_hit=cache_hit,
1132+
cache_tokens=cache_tokens,
1133+
)
10621134
self._outcome_reported = True
10631135
except Exception as e:
10641136
logger.warning(f"Gate 2 judge failed (non-blocking): {e}")
@@ -1072,15 +1144,32 @@ def _heal_dispatch(m_id: str, msgs: List[Dict], system_prompt: Optional[str] = N
10721144
# Priority 1: User-provided continuous scorer
10731145
score = self.score_when(output)
10741146
score = min(1.0, max(0.0, float(score)))
1075-
self.report(success=score >= 0.5, score=score, pipeline_id=pipeline_id)
1147+
self.report(
1148+
success=score >= 0.5,
1149+
score=score,
1150+
pipeline_id=pipeline_id,
1151+
cache_hit=cache_hit,
1152+
cache_tokens=cache_tokens,
1153+
)
10761154
elif self.success_when:
10771155
# Priority 2: User-provided binary scorer
10781156
success = self.success_when(output)
1079-
self.report(success=success, pipeline_id=pipeline_id)
1157+
self.report(
1158+
success=success,
1159+
pipeline_id=pipeline_id,
1160+
cache_hit=cache_hit,
1161+
cache_tokens=cache_tokens,
1162+
)
10801163
else:
10811164
# Priority 3: Default heuristic scoring (zero-config)
10821165
score = self._default_score(response)
1083-
self.report(success=score >= 0.5, score=score, pipeline_id=pipeline_id)
1166+
self.report(
1167+
success=score >= 0.5,
1168+
score=score,
1169+
pipeline_id=pipeline_id,
1170+
cache_hit=cache_hit,
1171+
cache_tokens=cache_tokens,
1172+
)
10841173
except Exception as e:
10851174
logger.warning(f"Auto-outcome evaluation failed: {e}")
10861175

@@ -1098,13 +1187,23 @@ def _heal_dispatch(m_id: str, msgs: List[Dict], system_prompt: Optional[str] = N
10981187
if should_auto_fire:
10991188
threading.Thread(
11001189
target=_auto_report_outcome,
1101-
args=(trace_id, self.goal, True, pipeline_id, candidate_model),
1190+
kwargs={
1191+
"trace_id": trace_id,
1192+
"goal": self.goal,
1193+
"success": True,
1194+
"pipeline_id": pipeline_id,
1195+
"model_id": candidate_model,
1196+
"cache_hit": cache_hit,
1197+
"cache_tokens": cache_tokens,
1198+
},
11021199
daemon=True,
11031200
).start()
11041201

11051202
# Add trace_id to response for explicit linkage
11061203
response.kalibr_trace_id = trace_id
11071204
response.kalibr_model_used = candidate_model
1205+
response.kalibr_cache_hit = cache_hit
1206+
response.kalibr_cache_tokens = cache_tokens
11081207
return response
11091208

11101209
except Exception as e:
@@ -1541,6 +1640,8 @@ def report(
15411640
trace_id: Optional[str] = None,
15421641
failure_category: Optional[str] = None,
15431642
pipeline_id: Optional[str] = None,
1643+
cache_hit: Optional[bool] = None,
1644+
cache_tokens: int = 0,
15441645
):
15451646
"""
15461647
Report outcome for the last completion.
@@ -1556,6 +1657,13 @@ def report(
15561657
trace_id: Optional explicit trace ID (uses last completion's trace_id if not provided)
15571658
failure_category: Optional structured failure category for clustering
15581659
pipeline_id: Optional pipeline identifier for multi-pipeline isolation
1660+
cache_hit: Optional prompt-cache hit flag captured from the provider
1661+
response. ``None`` means the provider did not report a cache
1662+
field. Reported to the intelligence service so future routing
1663+
can prefer cache-warm providers.
1664+
cache_tokens: Number of tokens served from the provider's prompt
1665+
cache on this call. 0 when there was no hit or the field was
1666+
not reported.
15591667
"""
15601668
if self._outcome_reported:
15611669
logger.warning("Outcome already reported for this completion. Each completion() requires a separate report() call.")
@@ -1585,6 +1693,8 @@ def report(
15851693
failure_category=failure_category,
15861694
model_id=self._last_model_id,
15871695
pipeline_id=pipeline_id,
1696+
cache_hit=cache_hit,
1697+
cache_tokens=cache_tokens,
15881698
)
15891699
self._outcome_reported = True
15901700
except Exception as e:
@@ -2130,6 +2240,11 @@ def _anthropic_to_openai_response(self, response: Any, model: str) -> Any:
21302240
if response.content:
21312241
content = response.content[0].text if hasattr(response.content[0], "text") else str(response.content[0])
21322242

2243+
# Preserve Anthropic prompt-cache fields on the normalized usage object so
2244+
# downstream cache-hit extraction works after the conversion.
2245+
cache_read = getattr(response.usage, "cache_read_input_tokens", None)
2246+
cache_creation = getattr(response.usage, "cache_creation_input_tokens", None)
2247+
21332248
return SimpleNamespace(
21342249
id=response.id,
21352250
model=model,
@@ -2147,6 +2262,8 @@ def _anthropic_to_openai_response(self, response: Any, model: str) -> Any:
21472262
prompt_tokens=response.usage.input_tokens,
21482263
completion_tokens=response.usage.output_tokens,
21492264
total_tokens=response.usage.input_tokens + response.usage.output_tokens,
2265+
cache_read_input_tokens=cache_read,
2266+
cache_creation_input_tokens=cache_creation,
21502267
),
21512268
)
21522269

0 commit comments

Comments
 (0)