UN-3810 [FEAT] Prompt caching: SDK master switch + document-context caching (answer_prompt) - #2199
UN-3810 [FEAT] Prompt caching: SDK master switch + document-context caching (answer_prompt)#2199pk-zipstack wants to merge 8 commits into
Conversation
Gate prompt caching per-adapter (Anthropic / Bedrock-Anthropic) via an enable_prompt_caching flag on adapter metadata or the LLM constructor. Add a reusable cache_prefix on complete()/stream_complete() that caches a stable user-turn prefix (content-block split) so repeated calls reuse it without changing the prompt text the model sees. Record cache token counts and price cache hits via litellm.completion_cost when present. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…r_prompt SDK1: add an ENABLE_PROMPT_CACHING master switch (env var) that auto-enables caching for every LLM on a supported provider, so consumers only pass a cache_prefix. Fix cache_prefix being dropped when caching is inactive (unsupported provider / flag off) so the full prompt is always preserved. answer_prompt: flag-gated context-first restructuring — the reused document context becomes a cached prefix reused across the prompts run on one document, while the per-prompt question is the volatile suffix. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Summary by CodeRabbit
WalkthroughPrompt caching is added to adapter validation, SDK completion flows, usage accounting, and executor prompt construction. Anthropic and Bedrock Claude models can cache stable prompt prefixes when caching is enabled. ChangesPrompt caching
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant AnswerPromptService
participant LLM
participant LiteLLM
participant Provider
AnswerPromptService->>LLM: complete(prompt, cache_prefix)
LLM->>LLM: build provider cache markers
LLM->>LiteLLM: send completion request
LiteLLM->>Provider: submit supported model request
Provider-->>LiteLLM: return completion and cache usage
LiteLLM-->>LLM: return completion result
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
| Filename | Overview |
|---|---|
| unstract/sdk1/src/unstract/sdk1/adapters/base1.py | Carries the prompt-caching control flag through Anthropic and Bedrock validation while restricting Bedrock enablement to identifiers recognized as Claude. |
| unstract/sdk1/src/unstract/sdk1/llm.py | Adds the master switch, provider/model capability checks, cached-prefix message construction across completion APIs, and cache-aware usage pricing. |
| workers/executor/executors/answer_prompt.py | Reorders document context into a cached prefix only when the active LLM confirms caching support, otherwise preserving existing prompt construction. |
| workers/executor/executors/legacy_executor.py | Enables prompt caching when constructing the legacy executor's LLM integration. |
| unstract/sdk1/tests/test_prompt_caching.py | Adds coverage for configuration, provider and Bedrock model gating, inference profiles, message preservation, control-flag stripping, and cached-call cost accounting. |
| workers/tests/test_answer_prompt_caching.py | Adds coverage for cached and uncached answer-prompt construction and completion invocation. |
Sequence Diagram
sequenceDiagram
participant Worker as AnswerPromptService
participant SDK as SDK1 LLM
participant Provider as Anthropic / Bedrock Claude
Worker->>SDK: Check is_prompt_caching_active()
alt Caching supported and enabled
Worker->>Worker: Build context prefix + volatile prompt
Worker->>SDK: "complete(prompt, cache_prefix=context)"
SDK->>Provider: Cached prefix block + volatile block
else Caching inactive or unsupported
Worker->>Worker: Preserve original prompt order
Worker->>SDK: complete(full prompt)
SDK->>Provider: Plain system and user messages
end
Provider-->>SDK: Completion and cache-token usage
SDK-->>Worker: Response with recorded usage
Reviews (5): Last reviewed commit: "UN-3810 [FIX] Document fail-closed inten..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
unstract/sdk1/src/unstract/sdk1/adapters/base1.py (1)
1044-1086: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winBedrock
enable_prompt_cachingisn't gated by model — contradicts its own comment.The comment states "Only Anthropic models on Bedrock support prompt caching," but
enable_prompt_cachingis carried through unconditionally fromadapter_metadataregardless of which Bedrock model is configured. If a non-Anthropic Bedrock model (Titan, Llama, Mistral, etc.) has this flag set,LLM._prompt_caching_active()will still emitcache_controlblocks for it (it only checksprovider == "bedrock", not the model id), which those models don't support.Proposed fix
- enable_prompt_caching = bool(adapter_metadata.get("enable_prompt_caching", False)) + raw_model = adapter_metadata.get("model", "") + is_anthropic_model = "anthropic." in raw_model + enable_prompt_caching = is_anthropic_model and bool( + adapter_metadata.get("enable_prompt_caching", False) + )See consolidated comment below for the paired fix location in
llm.py.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@unstract/sdk1/src/unstract/sdk1/adapters/base1.py` around lines 1044 - 1086, Gate enable_prompt_caching in the Bedrock adapter flow around validation and final assignment so it is true only when the configured Bedrock model is Anthropic; force it false for Titan, Llama, Mistral, and other non-Anthropic models before validated["enable_prompt_caching"] is emitted. Preserve the existing opt-in behavior for supported Anthropic Bedrock models and keep the prompt-caching control field excluded from Pydantic validation.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@unstract/sdk1/src/unstract/sdk1/llm.py`:
- Around line 331-343: The _prompt_caching_active method must also verify that a
Bedrock request uses a supported Anthropic/Claude model before enabling
cache_control emission. Update the provider/model gating in
_prompt_caching_active, preserving the existing Anthropic behavior and disabling
prompt caching for unsupported Bedrock models.
- Around line 828-869: Update _compute_call_cost so the response-based
litellm.completion_cost call passes the cost override via model=model, keeping
cached-call pricing consistent with the fallback cost_per_token path. Add a
regression test covering a prompt-cached call with cost_model set and verifying
the override is used for cost tracking.
In `@workers/executor/executors/answer_prompt.py`:
- Around line 163-187: The prompt construction flow in AnswerPromptService must
only use construct_cached_prompt when caching is globally enabled and supported
by the current LLM; otherwise use construct_prompt, preserve the original prompt
order, and pass no cache prefix. Update
workers/executor/executors/answer_prompt.py lines 163-187 accordingly, and add
the globally enabled unsupported-provider coverage in
workers/tests/test_answer_prompt_caching.py lines 52-91 to verify the original
ordering and absence of a cache prefix.
---
Outside diff comments:
In `@unstract/sdk1/src/unstract/sdk1/adapters/base1.py`:
- Around line 1044-1086: Gate enable_prompt_caching in the Bedrock adapter flow
around validation and final assignment so it is true only when the configured
Bedrock model is Anthropic; force it false for Titan, Llama, Mistral, and other
non-Anthropic models before validated["enable_prompt_caching"] is emitted.
Preserve the existing opt-in behavior for supported Anthropic Bedrock models and
keep the prompt-caching control field excluded from Pydantic validation.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: a725c6fd-d5d8-4ba0-8a50-24fea57e25bf
📒 Files selected for processing (6)
unstract/sdk1/src/unstract/sdk1/adapters/base1.pyunstract/sdk1/src/unstract/sdk1/llm.pyunstract/sdk1/tests/test_prompt_caching.pyworkers/executor/executors/answer_prompt.pyworkers/executor/executors/legacy_executor.pyworkers/tests/test_answer_prompt_caching.py
pk-zipstack
changed the title
…rder for non-caching LLMs, fix cached-call cost Resolves greptile/coderabbit review findings on the prompt-caching PR: 1. Bedrock model gate (llm.py, base1.py): only emit cache_control for Anthropic/Claude models on Bedrock. Titan/Llama/Cohere/Mistral on Bedrock no longer get Anthropic-only cache blocks (ineffective + unsupported shapes). Enforced in LLM._prompt_caching_active() so it also covers the ENABLE_PROMPT_CACHING master-switch path; base1 keeps the validated flag honest for the metadata path. 2. answer_prompt reorder (answer_prompt.py): only reorder into a cached prefix when the LLM actually caches (new public LLM.is_prompt_caching_active()). Unsupported providers keep the original context-last prompt order instead of a no-benefit reorder. 3. Cached-call cost (llm.py): pass model= to litellm.completion_cost so cached calls price against the cost_model override, matching the cost_per_token fallback. Tests: Bedrock anthropic/non-anthropic gating, public probe, base1 non-anthropic flag, cost override regression, and an unsupported-provider answer_prompt case asserting original order + no cache_prefix. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
# Conflicts: # unstract/sdk1/src/unstract/sdk1/llm.py
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
for more information, see https://pre-commit.ci
Review comments addressed + rebased on
|
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
workers/tests/test_answer_prompt_caching.py (1)
28-37: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winScope the
sys.modulesstubbing so it does not leak across the test session.
_load_answer_promptinserts syntheticexecutorandexecutor.executorsmodules intosys.modulesat import time and never removes them. These entries persist for the whole pytest session, and other worker tests import the realexecutorpackage, so later tests can receive stale namespace stubs depending on test order. Move the stubbing into a shared fixture that restoressys.modules, or centralize the same stubbing strategy inworkers/conftest.py.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@workers/tests/test_answer_prompt_caching.py` around lines 28 - 37, Update _load_answer_prompt and _mod initialization so synthetic executor modules are created within a shared pytest fixture rather than at module import time. Ensure the fixture restores the prior sys.modules entries after each test, preserving real executor imports for other worker tests; centralize this setup in workers/conftest.py if appropriate.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@unstract/sdk1/src/unstract/sdk1/llm.py`:
- Around line 413-427: Update the cache-prefix condition in the prompt-building
method to require a truthy cache_prefix rather than only checking for None.
Treat empty strings as absent so they use the normal prompt path and never
create an empty Anthropic text block.
In `@unstract/sdk1/tests/test_prompt_caching.py`:
- Around line 245-267: Update the docstring in
test_cost_override_passed_to_completion_cost_on_cached_call so its summary line
is followed by a blank line before the description, satisfying Ruff D205.
Replace object() when invoking LLM._compute_call_cost with a minimal stub
instance that represents the required self context and can safely support future
self attribute access.
---
Nitpick comments:
In `@workers/tests/test_answer_prompt_caching.py`:
- Around line 28-37: Update _load_answer_prompt and _mod initialization so
synthetic executor modules are created within a shared pytest fixture rather
than at module import time. Ensure the fixture restores the prior sys.modules
entries after each test, preserving real executor imports for other worker
tests; centralize this setup in workers/conftest.py if appropriate.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 7a267d38-4461-4155-a027-57d1ff64aad2
📒 Files selected for processing (6)
unstract/sdk1/src/unstract/sdk1/adapters/base1.pyunstract/sdk1/src/unstract/sdk1/llm.pyunstract/sdk1/tests/test_prompt_caching.pyworkers/executor/executors/answer_prompt.pyworkers/executor/executors/legacy_executor.pyworkers/tests/test_answer_prompt_caching.py
🚧 Files skipped from review as they are similar to previous changes (2)
- workers/executor/executors/legacy_executor.py
- workers/executor/executors/answer_prompt.py
…ache_prefix, test nits Second-pass review fixes (greptile + coderabbit) on the prompt-caching PR: 1. Opaque Bedrock inference profiles (llm.py, base1.py): the Anthropic/ Claude model gate now checks both `model` and `model_id`. When a caller routes through a Bedrock Application Inference Profile, the ARN in `model` is opaque and the Claude id only appears in `model_id`; the previous model-only check silently bypassed caching for those calls. 2. Empty cache_prefix (llm.py): `_build_messages` now treats an empty `cache_prefix` as absent (truthy check, not `is not None`). Anthropic rejects empty text content blocks, and an empty prefix has no caching benefit; the fallback concat path drops it cleanly too. 3. Test quality (test_prompt_caching.py): fixed the D205 docstring, use a SimpleNamespace stub instead of bare object() for `self`, and annotate monkeypatch params. Added tests for the opaque-AIP model_id path (and the no-Claude-id negative) plus the empty-cache_prefix case. ruff + ruff-format clean; 31 SDK + 9 workers tests passing. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
@greptile-apps review The sole 4/5 blocker — "Opaque profiles bypass caching" (
Also addressed in the same commit: empty- |
PR Review Summary: #2199 — UN-3810 [FEAT] Prompt caching: SDK master switch + document-context cachingOverviewThis PR implements opt-in LLM prompt caching for the Unstract SDK, gated by an ENABLE_PROMPT_CACHING env var (default off). It adds cache_control block emission for Overall Assessment: Well-designed feature with strong test coverage and thoughtful gating. The three-layer gating (master switch, adapter flag, provider/model check) is sound, Critical Issues (2 found)
File: unstract/sdk1/src/unstract/sdk1/llm.py complete() and stream_complete() both accept and forward cache_prefix to _build_messages(). However, acomplete() was not updated — it calls self._build_messages(prompt) without a Fix: Add cache_prefix: str | None = None to acomplete() signature and forward it.
File: unstract/sdk1/src/unstract/sdk1/llm.py, _build_messages() When cache_prefix="", the is not None check passes, producing {"type": "text", "text": "", "cache_control": {"type": "ephemeral"}}. Empty content blocks are generally rejected or Fix: Change cache_prefix is not None to cache_prefix (truthy check) in the _build_messages() guard. Similarly in the fallback: cache_prefix + prompt if cache_prefix else prompt. Important Issues (5 found)
Files: llm.py and answer_prompt.py both define identical copies If the env var name or parsing logic changes, both copies must be updated in lockstep. Import from the SDK module instead.
File: workers/tests/test_answer_prompt_caching.py _load_answer_prompt() injects synthetic packages into sys.modules at module load time and never cleans them up. This can contaminate other tests in the same pytest process. Fix: Use a session-scoped fixture with cleanup or unittest.mock.patch.dict(sys.modules, ...).
File: llm.py, _compute_call_cost() When completion_cost() fails for cached calls, the fallback to cost_per_token (which over-reports cost ~10x on cache hits) is logged at debug level — invisible in production. Fix: Promote to logger.warning with a note that cache pricing may be inaccurate.
File: answer_prompt.py Both methods independently compute JSON postamble and platform_postamble formatting. If one is updated without the other, prompt content will silently diverge between cached and Fix: Extract the postamble preparation into a shared private method _prepare_postambles().
File: llm.py, _prompt_caching_active() Opaque ARN-based model IDs (e.g., arn:aws:bedrock:...:application-inference-profile/...) don't contain "anthropic" or "claude" substrings, so Claude models accessed via inference Fix: Document this known limitation. Add a debug-level log when Bedrock caching is skipped so operators can diagnose. Suggestions (4 found)
Only the happy path (where litellm.completion_cost succeeds) is tested. Add tests for: (a) completion_cost raises → fallback to cost_per_token, (b) both fail → returns 0.0, (c)
The explicit completion_kwargs.pop("enable_prompt_caching", None) in complete(), stream_complete(), and acomplete() is untested. If removed, the flag would leak into litellm's
Only "true" enables caching — "1", "yes", "on" don't work. No log message when someone sets a truthy-but-not-"true" value. Add a startup warning.
The warning "Failed to determine prompt-caching support..." is generic with no LLM context to help operators identify which LLM is misbehaving. Strengths
Recommended Action
|
…-log level, test coverage
Resolves harini-venkataraman's PR review (2 critical, 5 important, 4 suggestions):
Critical:
- acomplete() now accepts and forwards cache_prefix, matching complete()/
stream_complete() (was silently un-cached on the async path).
- Empty cache_prefix already treated as absent (truthy guard) — kept + tested.
Important:
- Removed the duplicate is_prompt_caching_enabled() from answer_prompt.py; the
SDK owns the env master switch and production gates on
LLM.is_prompt_caching_active() (DRY).
- Test sys.modules stubs are now removed after import so they can't contaminate
other tests in the same process.
- _compute_call_cost cost-fallback now logs at WARNING (the fallback can
over-report cached-call cost ~10x; debug was invisible in prod).
- Extracted shared _prepare_postambles() so construct_prompt and
construct_cached_prompt can't silently diverge.
- Bedrock gate logs a debug breadcrumb when caching is skipped for an
unrecognized (opaque) model.
Suggestions:
- Tests for _compute_call_cost fallback paths (completion_cost raises, both
fail -> 0.0, response=None skip).
- Tests that enable_prompt_caching/cost_model/context_window never reach
litellm.completion / .acompletion.
- is_prompt_caching_enabled warns once on truthy-lookalikes ("1"/"yes"/"on").
- _llm_caches_prompts warning now includes the LLM type and provider.
ruff + ruff-format clean; 38 SDK + 11 workers tests passing.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
@harini-venkataraman thanks for the thorough review — all 11 points addressed in Critical
Important Suggestions ruff + ruff-format clean; 38 SDK + 11 workers tests passing. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@workers/executor/executors/answer_prompt.py`:
- Around line 122-135: Add a targeted Ruff `noqa` annotation to the intentional
broad exception in the prompt-caching capability probe, documenting that
arbitrary probe failures must fail closed and preserve the original prompt
order. Keep the existing exception handling and logging behavior unchanged,
including the nested provider lookup handling.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 10d97412-d23f-453e-ac57-d534007b1355
📒 Files selected for processing (4)
unstract/sdk1/src/unstract/sdk1/llm.pyunstract/sdk1/tests/test_prompt_caching.pyworkers/executor/executors/answer_prompt.pyworkers/tests/test_answer_prompt_caching.py
…xcepts (BLE001) Add noqa: BLE001 with rationale to the two intentional broad excepts in _llm_caches_prompts: the capability probe must fail closed (keep the original prompt order) and the provider lookup is best-effort log context. Per CodeRabbit. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
Unstract test resultsPer-group results
Critical paths
|
2 participants
Summary
Adds opt-in LLM prompt caching to SDK1 and wires it into the main extraction path (
answer_prompt). Implements the recommendation from the prompt-caching audit. Off by default — a single master switch (ENABLE_PROMPT_CACHING) gates the whole feature, so there is zero behavior change until it is explicitly enabled.Pairs with the cloud-plugins PR (
unstract-cloud#…, table + challenge extractors). This PR should merge first — the cloud plugins depend on the SDK primitives here.Why prompt caching, and how it actually helps
Anthropic/Bedrock prompt caching only pays off when a large, byte-identical prefix repeats within ~5 minutes. It is not a switch you flip and save everywhere — the win comes from identifying, per consumer, the span that actually repeats. This PR builds the mechanism and applies it to the one platform path with a big, genuinely-reused prefix: the document context reused across the many prompts run on one document.
What's in this PR
SDK1 (
unstract/sdk1)ENABLE_PROMPT_CACHINGenv var auto-enables caching for everyLLMon a supported provider (is_prompt_caching_enabled()), so consumers don't each pass a flag; they only decide what to cache viacache_prefix.cache_prefixoncomplete()/stream_complete()— splits the user turn into a cached stable prefix block + the volatile suffix; the model seescache_prefix + promptunchanged.cache_controlis emitted only foranthropic/bedrock(mirrors the existingenable_extended_contextpattern inbase1.py); a no-op on other providers.cache_write=/cache_read=and prices cache hits vialitellm.completion_costwhen present.cache_prefixwas passed, the prefix was silently dropped from the prompt. Now the full prompt is always preserved. This affected the table extractor for non-Claude adapters. Regression-guarded by tests.Workers —
answer_prompt(the ★ consumer)Validation (in a dev environment)
cache_write-ten on the first prompt andcache_readon the rest.Why we're stopping here (scope boundary)
The platform has ~7 LLM call sites. We shipped caching for the three with a real, high-reuse prefix and stopped, after analyzing every remaining consumer:
answer_prompttable_extractor(cloud)challenge(cloud)smart_table_extractorline_item_extractormax_tokens); most runs are single-call → no reuse. Conditional + fiddly.single_pass_extractionagentic_extractionretriever_llmIn short: caching only helps where a large prefix genuinely repeats. Those cases are now covered; the rest would add marginal/conditional benefit at the cost of accuracy-risky rewrites, so we deliberately drew the line here.
agentic_extractionis the one worthwhile follow-up if long agentic extractions become common.Testing
unstract/sdk1:pytest tests/test_prompt_caching.py(18 tests — provider gating,cache_prefixsplit, prefix-preservation fallback, env master switch).workers:pytest tests/test_answer_prompt_caching.py(5 tests — reorder preserves all content, context-first,cache_prefix= context block only, flag defaults off).ENABLE_PROMPT_CACHING=trueon the executor worker, run ≥2 text prompts (chunk size 0) on one multi-page doc with an Anthropic adapter, and confirmcache_readin the executor logs.Follow-ups (not blocking)
cache_creation_input_tokens/cache_read_input_tokens(currently logged only).agentic_extraction.🤖 Generated with Claude Code