From d629c9db301f5fa22b0fc60234a8872acceeec9a Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Wed, 22 Jul 2026 18:25:43 +0530 Subject: [PATCH 1/2] [MISC] Add opt-in provider prompt caching to SDK1 LLM 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) --- .../sdk1/src/unstract/sdk1/adapters/base1.py | 18 ++ unstract/sdk1/src/unstract/sdk1/llm.py | 189 +++++++++++++++--- unstract/sdk1/tests/test_prompt_caching.py | 186 +++++++++++++++++ 3 files changed, 364 insertions(+), 29 deletions(-) create mode 100644 unstract/sdk1/tests/test_prompt_caching.py diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py index 3c9aee1d60..de9d60682f 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py @@ -1041,6 +1041,13 @@ def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]: result_metadata["thinking"] = thinking_config result_metadata["temperature"] = 1 + # Prompt caching is opt-in and applied on the message payload (a + # `cache_control` block on the stable system prompt), not as a LiteLLM + # completion param, so it is excluded from Pydantic validation and + # carried through on the validated dict for the LLM layer to read. + # Only Anthropic models on Bedrock support prompt caching. + enable_prompt_caching = bool(adapter_metadata.get("enable_prompt_caching", False)) + _pack_bedrock_guardrail_config(result_metadata) # Create validation metadata excluding control fields. `auth_type` is @@ -1057,6 +1064,7 @@ def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]: "guardrail_identifier", "guardrail_version", "guardrail_trace", + "enable_prompt_caching", ) } @@ -1075,6 +1083,7 @@ def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]: # lenient. Reads auth_type from result_metadata since validation_ # metadata strips it before Pydantic. validated = _resolve_bedrock_aws_credentials(result_metadata, validated) + validated["enable_prompt_caching"] = enable_prompt_caching return _strip_deprecated_sampling_params(validated) @staticmethod @@ -1142,6 +1151,12 @@ def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]: result_metadata["thinking"] = thinking_config result_metadata["temperature"] = 1 + # Prompt caching is opt-in and applied on the message payload (a + # `cache_control` block on the stable system prompt), not as a LiteLLM + # completion param, so it is excluded from Pydantic validation and + # carried through on the validated dict for the LLM layer to read. + enable_prompt_caching = bool(adapter_metadata.get("enable_prompt_caching", False)) + # Create validation metadata excluding control fields exclude_fields = ( "enable_thinking", @@ -1149,6 +1164,7 @@ def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]: "thinking", "enable_extended_context", "extra_headers", + "enable_prompt_caching", ) validation_metadata = { k: v for k, v in result_metadata.items() if k not in exclude_fields @@ -1164,6 +1180,8 @@ def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]: if enable_extended_context: validated["extra_headers"] = {"anthropic-beta": "context-1m-2025-08-07"} + validated["enable_prompt_caching"] = enable_prompt_caching + return _strip_deprecated_sampling_params(validated) @staticmethod diff --git a/unstract/sdk1/src/unstract/sdk1/llm.py b/unstract/sdk1/src/unstract/sdk1/llm.py index b45b856239..2e1074456c 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -162,6 +162,7 @@ def __init__( # noqa: C901 system_prompt: str = "", kwargs: dict[str, object] | None = None, capture_metrics: bool = False, + enable_prompt_caching: bool = False, ) -> None: """Initialize the LLM interface. @@ -174,6 +175,9 @@ def __init__( # noqa: C901 system_prompt: System prompt for the LLM kwargs: Additional keyword arguments for configuration capture_metrics: Whether to capture performance metrics + enable_prompt_caching: Force provider prompt caching on for + supported providers (Anthropic / Bedrock-Anthropic), regardless + of the stored adapter metadata. Ignored for other providers. """ if adapter_metadata is None: adapter_metadata = {} @@ -224,6 +228,15 @@ def __init__( # noqa: C901 self.kwargs = self.adapter.validate(self._adapter_metadata) self._cost_model = self.kwargs.pop("cost_model", None) + # Opt-in provider prompt caching (Anthropic / Bedrock-Anthropic). + # Enabled either via adapter metadata (from the stored adapter + # config) or the explicit constructor arg (for callers that build + # the LLM by ``adapter_instance_id`` and can't edit stored metadata). + # Popped so it never reaches litellm; applied on the message payload. + self._enable_prompt_caching = ( + bool(self.kwargs.pop("enable_prompt_caching", False)) + or enable_prompt_caching + ) # REF: https://docs.litellm.ai/docs/completion/input#translated-openai-params # supported = get_supported_openai_params(model=self.kwargs["model"], @@ -302,8 +315,78 @@ def test_connection(self) -> bool: actual_err=e, ) from e + # Providers for which we emit explicit ``cache_control`` blocks. Anthropic + # and Bedrock-Anthropic support message-level prompt caching this way; + # OpenAI / Azure auto-cache server-side (no marker needed) and other + # providers don't support it, so we never tag their payloads. + _PROMPT_CACHE_PROVIDERS = frozenset({"anthropic", "bedrock"}) + + def _prompt_caching_active(self) -> bool: + """Whether to emit ``cache_control`` blocks for this call.""" + return ( + self._enable_prompt_caching + and self.adapter.get_provider() in self._PROMPT_CACHE_PROVIDERS + ) + + def _build_messages( + self, prompt: str, cache_prefix: str | None = None + ) -> list[dict[str, object]]: + """Build the system + user message list for a chat completion. + + When prompt caching is active (opt-in flag + a supported provider), a + stable prefix is tagged with ``cache_control`` so providers that support + prefix caching (Anthropic, Bedrock-Anthropic) reuse it across calls. + LiteLLM forwards ``cache_control`` blocks to the provider unchanged. + + - ``cache_prefix`` given: the user turn is split into a cached stable + prefix block followed by the per-request volatile block. The text the + model sees is ``cache_prefix + prompt`` — identical to passing the + concatenation as a single prompt, so no prompt semantics change. + - otherwise: the stable system prompt is cached. + + Only the stable portion is tagged; per-request content is never cached. + """ + if self._prompt_caching_active() and cache_prefix is not None: + return [ + {"role": "system", "content": self._system_prompt}, + { + "role": "user", + "content": [ + { + "type": "text", + "text": cache_prefix, + "cache_control": {"type": "ephemeral"}, + }, + {"type": "text", "text": prompt}, + ], + }, + ] + if self._prompt_caching_active(): + return [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": self._system_prompt, + "cache_control": {"type": "ephemeral"}, + } + ], + }, + {"role": "user", "content": prompt}, + ] + return [ + {"role": "system", "content": self._system_prompt}, + {"role": "user", "content": prompt}, + ] + @capture_metrics - def complete(self, prompt: str, **kwargs: object) -> dict[str, object]: + def complete( + self, + prompt: str, + cache_prefix: str | None = None, + **kwargs: object, + ) -> dict[str, object]: """Return a standard chat completion dict with optional metrics capture. Return a standard chat completion dict and optionally captures metrics if run @@ -311,6 +394,11 @@ def complete(self, prompt: str, **kwargs: object) -> dict[str, object]: Args: prompt (str) The input text prompt for generating the completion. + cache_prefix (str | None) Stable text to cache ahead of ``prompt``. + When prompt caching is active for a supported provider, this is + emitted as a ``cache_control`` block so repeated calls sharing + the same prefix reuse it. The model sees ``cache_prefix + prompt`` + unchanged. Ignored when caching is off or unsupported. **kwargs (Any) Additional arguments passed to the completion function. Returns: @@ -318,16 +406,14 @@ def complete(self, prompt: str, **kwargs: object) -> dict[str, object]: any processed output, and the captured metrics (if applicable). """ try: - messages: list[dict[str, str]] = [ - {"role": "system", "content": self._system_prompt}, - {"role": "user", "content": prompt}, - ] + messages = self._build_messages(prompt, cache_prefix=cache_prefix) logger.debug( f"[sdk1][LLM]Invoking {self.adapter.get_provider()} completion API" ) completion_kwargs = self.adapter.validate({**self.kwargs, **kwargs}) completion_kwargs.pop("cost_model", None) + completion_kwargs.pop("enable_prompt_caching", None) # if hasattr(self, "model") and self.model not in O1_MODELS: # completion_kwargs["temperature"] = 0.003 @@ -451,6 +537,7 @@ def complete_vision( completion_kwargs = self.adapter.validate({**self.kwargs, **kwargs}) completion_kwargs.pop("cost_model", None) + completion_kwargs.pop("enable_prompt_caching", None) response: dict[str, object] = litellm.completion( messages=messages, @@ -501,23 +588,24 @@ def stream_complete( self, prompt: str, callback_manager: object | None = None, + cache_prefix: str | None = None, **kwargs: object, ) -> Generator[LLMResponseCompat, None, None]: """Yield LLMResponseCompat objects with text chunks. - Chunks arrive as they stream from the provider. + Chunks arrive as they stream from the provider. ``cache_prefix`` behaves + as in :meth:`complete` — a stable prefix cached ahead of ``prompt`` when + prompt caching is active for a supported provider. """ try: - messages = [ - {"role": "system", "content": self._system_prompt}, - {"role": "user", "content": prompt}, - ] + messages = self._build_messages(prompt, cache_prefix=cache_prefix) logger.debug( f"[sdk1][LLM]Invoking {self.adapter.get_provider()} stream completion API" ) completion_kwargs = self.adapter.validate({**self.kwargs, **kwargs}) completion_kwargs.pop("cost_model", None) + completion_kwargs.pop("enable_prompt_caching", None) max_retries = pop_litellm_retry_kwargs( completion_kwargs, self._get_adapter_info() @@ -579,16 +667,14 @@ def stream_complete( async def acomplete(self, prompt: str, **kwargs: object) -> dict[str, object]: """Asynchronous chat completion (wrapper around ``litellm.acompletion``).""" try: - messages = [ - {"role": "system", "content": self._system_prompt}, - {"role": "user", "content": prompt}, - ] + messages = self._build_messages(prompt) logger.debug( f"[sdk1][LLM]Invoking {self.adapter.get_provider()} async completion API" ) completion_kwargs = self.adapter.validate({**self.kwargs, **kwargs}) completion_kwargs.pop("cost_model", None) + completion_kwargs.pop("enable_prompt_caching", None) max_retries = pop_litellm_retry_kwargs( completion_kwargs, self._get_adapter_info() @@ -721,6 +807,48 @@ def flush_pending_usage(self) -> list[dict]: self._pending_usage = [] return records + def _compute_call_cost( + self, + model: str, + prompt_tokens: int, + completion_tokens: int, + has_cache_tokens: bool, + response: object | None, + ) -> float: + """Compute the dollar cost of a single call. + + When caching is active, cache-read tokens are billed at ~0.1x and + cache-write at ~1.25x of the base input rate. ``litellm.cost_per_token`` + prices every prompt token at the full input rate, so it over-reports + cost on cache hits. In that case let litellm read the cache token counts + off the response for an accurate figure, falling back to the per-token + path (which is exact when no caching is involved). + """ + if has_cache_tokens and response is not None: + try: + return litellm.completion_cost(completion_response=response) + except Exception: + logger.debug( + "completion_cost() failed for model=%s; " + "falling back to cost_per_token", + model, + exc_info=True, + ) + try: + prompt_cost, compl_cost = litellm.cost_per_token( + model=model, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ) + return prompt_cost + compl_cost + except Exception: + logger.warning( + "Failed to compute cost for model=%s; recording as 0.0", + model, + exc_info=True, + ) + return 0.0 + def _record_usage( self, model: str, @@ -733,6 +861,10 @@ def _record_usage( prompt_tokens = usage_data.get("prompt_tokens", 0) completion_tokens = usage_data.get("completion_tokens", 0) total_tokens = usage_data.get("total_tokens", 0) + # Prompt-caching token counts (populated by Anthropic / Bedrock-Anthropic + # when caching is enabled; 0 for every other provider/call). + cache_creation_tokens = usage_data.get("cache_creation_input_tokens", 0) or 0 + cache_read_tokens = usage_data.get("cache_read_input_tokens", 0) or 0 # Fall back to litellm when providers omit prompt tokens — avoids 0-token billing. if prompt_tokens == 0 and messages: @@ -757,30 +889,29 @@ def _record_usage( id_suffix += f" response_id={response_id}" if request_id is not None: id_suffix += f" request_id={request_id}" + cache_suffix = "" + if cache_creation_tokens or cache_read_tokens: + cache_suffix = ( + f" cache_write={cache_creation_tokens} cache_read={cache_read_tokens}" + ) logger.info( - "[sdk1][LLM][%s][%s] Usage: prompt=%d completion=%d total=%d%s", + "[sdk1][LLM][%s][%s] Usage: prompt=%d completion=%d total=%d%s%s", model, llm_api, prompt_tokens, completion_tokens, total_tokens, + cache_suffix, id_suffix, ) - try: - prompt_cost, compl_cost = litellm.cost_per_token( - model=model, - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - ) - cost = prompt_cost + compl_cost - except Exception: - logger.warning( - "Failed to compute cost for model=%s; recording as 0.0", - model, - exc_info=True, - ) - cost = 0.0 + cost = self._compute_call_cost( + model=model, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + has_cache_tokens=bool(cache_creation_tokens or cache_read_tokens), + response=response, + ) # Trailing segment matches legacy Audit semantics (e.g. bedrock/anthropic/claude). display_model = model.rsplit("/", 1)[-1] if model else model diff --git a/unstract/sdk1/tests/test_prompt_caching.py b/unstract/sdk1/tests/test_prompt_caching.py new file mode 100644 index 0000000000..ba2762a771 --- /dev/null +++ b/unstract/sdk1/tests/test_prompt_caching.py @@ -0,0 +1,186 @@ +"""Tests for opt-in provider prompt caching. + +Covers the two halves of the feature: + +1. Adapter ``validate()`` (Anthropic + Bedrock-Anthropic) carries the + ``enable_prompt_caching`` control flag through on the validated dict, but + never leaks it into the LiteLLM completion kwargs Pydantic validates. +2. ``LLM._build_messages()`` tags only the stable system prompt with + ``cache_control`` when caching is enabled, and leaves the string form + untouched otherwise. +""" + +from typing import Any + +import pytest + +from unstract.sdk1.adapters.base1 import ( + AnthropicLLMParameters, + AWSBedrockLLMParameters, +) + +# ── validate(): flag carried through, absent by default ───────────────────── + +VALIDATE_CASES = [ + ("anthropic", AnthropicLLMParameters, "claude-opus-4-8", {"api_key": "k"}), + ( + "bedrock", + AWSBedrockLLMParameters, + "anthropic.claude-opus-4-8-20260101-v1:0", + {"aws_region_name": "us-east-1"}, + ), +] + + +@pytest.mark.parametrize( + "name,cls,model,extra", VALIDATE_CASES, ids=[c[0] for c in VALIDATE_CASES] +) +def test_validate_defaults_prompt_caching_off( + name: str, cls: type, model: str, extra: dict[str, Any] +) -> None: + result = cls.validate({"model": model, **extra}) + assert result["enable_prompt_caching"] is False + + +@pytest.mark.parametrize( + "name,cls,model,extra", VALIDATE_CASES, ids=[c[0] for c in VALIDATE_CASES] +) +def test_validate_carries_prompt_caching_flag( + name: str, cls: type, model: str, extra: dict[str, Any] +) -> None: + result = cls.validate({"model": model, "enable_prompt_caching": True, **extra}) + assert result["enable_prompt_caching"] is True + + +@pytest.mark.parametrize( + "name,cls,model,extra", VALIDATE_CASES, ids=[c[0] for c in VALIDATE_CASES] +) +def test_validate_is_idempotent_on_prompt_caching( + name: str, cls: type, model: str, extra: dict[str, Any] +) -> None: + """Re-validating a validated dict must preserve the flag (round-trip).""" + once = cls.validate({"model": model, "enable_prompt_caching": True, **extra}) + twice = cls.validate({**once}) + assert twice["enable_prompt_caching"] is True + + +# ── _build_messages(): cache_control only on the system prefix ────────────── + + +class _StubAdapter: + def __init__(self, provider: str) -> None: + self._provider = provider + + def get_provider(self) -> str: + return self._provider + + +class _StubLLM: + """Bind the real caching helpers to a stub carrying just the state they read.""" + + from unstract.sdk1.llm import LLM + + _build_messages = LLM._build_messages + _prompt_caching_active = LLM._prompt_caching_active + _PROMPT_CACHE_PROVIDERS = LLM._PROMPT_CACHE_PROVIDERS + + def __init__( + self, + system_prompt: str, + enable_prompt_caching: bool, + provider: str = "anthropic", + ) -> None: + self._system_prompt = system_prompt + self._enable_prompt_caching = enable_prompt_caching + self.adapter = _StubAdapter(provider) + + +def test_build_messages_plain_string_when_caching_off() -> None: + llm = _StubLLM("SYSTEM", enable_prompt_caching=False) + messages = llm._build_messages("USER") + assert messages == [ + {"role": "system", "content": "SYSTEM"}, + {"role": "user", "content": "USER"}, + ] + + +def test_build_messages_tags_only_system_when_caching_on() -> None: + llm = _StubLLM("SYSTEM", enable_prompt_caching=True) + messages = llm._build_messages("USER") + + system, user = messages + assert system["role"] == "system" + assert system["content"] == [ + { + "type": "text", + "text": "SYSTEM", + "cache_control": {"type": "ephemeral"}, + } + ] + # The per-request user prompt is never tagged for caching. + assert user == {"role": "user", "content": "USER"} + + +@pytest.mark.parametrize("provider", ["openai", "azure", "gemini", "mistral"]) +def test_build_messages_not_tagged_for_unsupported_provider(provider: str) -> None: + """Caching flag on, but a provider we don't emit cache_control for -> plain.""" + llm = _StubLLM("SYSTEM", enable_prompt_caching=True, provider=provider) + assert llm._build_messages("USER") == [ + {"role": "system", "content": "SYSTEM"}, + {"role": "user", "content": "USER"}, + ] + + +def test_build_messages_cache_prefix_splits_user_turn() -> None: + llm = _StubLLM("SYSTEM", enable_prompt_caching=True) + messages = llm._build_messages("VOLATILE", cache_prefix="STABLE") + + system, user = messages + # System stays a plain string; the stable prefix is cached in the user turn. + assert system == {"role": "system", "content": "SYSTEM"} + assert user["role"] == "user" + assert user["content"] == [ + {"type": "text", "text": "STABLE", "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": "VOLATILE"}, + ] + # Text-equivalence invariant: the model sees prefix + prompt, unchanged. + seen = "".join(block["text"] for block in user["content"]) + assert seen == "STABLE" + "VOLATILE" + + +def test_build_messages_cache_prefix_ignored_when_caching_off() -> None: + llm = _StubLLM("SYSTEM", enable_prompt_caching=False) + assert llm._build_messages("VOLATILE", cache_prefix="STABLE") == [ + {"role": "system", "content": "SYSTEM"}, + {"role": "user", "content": "VOLATILE"}, + ] + + +# ── constructor opt-in (real LLM, no flag in adapter metadata) ────────────── + +_ANTHROPIC_ADAPTER_ID = "anthropic|90ebd4cd-2f19-4cef-a884-9eeb6ac0f203" + + +def test_constructor_flag_forces_caching_without_metadata() -> None: + """A caller that builds by adapter without the stored flag can still opt in.""" + from unstract.sdk1.llm import LLM + + meta = {"model": "claude-opus-4-8", "api_key": "sk-test"} + llm = LLM( + adapter_id=_ANTHROPIC_ADAPTER_ID, + adapter_metadata=meta, + system_prompt="S", + enable_prompt_caching=True, + ) + assert llm._enable_prompt_caching is True + # cache_prefix path produces the split user turn end to end. + messages = llm._build_messages("VOLATILE", cache_prefix="STABLE") + assert messages[1]["content"][0]["cache_control"] == {"type": "ephemeral"} + + +def test_constructor_flag_defaults_off() -> None: + from unstract.sdk1.llm import LLM + + meta = {"model": "claude-opus-4-8", "api_key": "sk-test"} + llm = LLM(adapter_id=_ANTHROPIC_ADAPTER_ID, adapter_metadata=meta, system_prompt="S") + assert llm._enable_prompt_caching is False From 395426a30b07a5e5681bd80ce4ff1ed66afea974 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Thu, 23 Jul 2026 10:51:19 +0530 Subject: [PATCH 2/2] [MISC] Master-switch prompt caching + cache document context in answer_prompt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- unstract/sdk1/src/unstract/sdk1/llm.py | 20 +++- unstract/sdk1/tests/test_prompt_caching.py | 38 +++++++- workers/executor/executors/answer_prompt.py | 86 ++++++++++++++--- workers/executor/executors/legacy_executor.py | 2 + workers/tests/test_answer_prompt_caching.py | 95 +++++++++++++++++++ 5 files changed, 223 insertions(+), 18 deletions(-) create mode 100644 workers/tests/test_answer_prompt_caching.py diff --git a/unstract/sdk1/src/unstract/sdk1/llm.py b/unstract/sdk1/src/unstract/sdk1/llm.py index 2e1074456c..ca342cb7b5 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -31,6 +31,18 @@ logger = logging.getLogger(__name__) + +def is_prompt_caching_enabled() -> bool: + """Whether LLM prompt caching is enabled platform-wide (opt-in, default off). + + A single master switch (``ENABLE_PROMPT_CACHING`` env var) that turns the + caching capability on for every ``LLM`` on a supported provider, so callers + don't each have to pass ``enable_prompt_caching``. Consumers still decide + *what* to cache by passing ``cache_prefix``. Exposed for consumers that gate + their own prompt-restructuring on the same flag. + """ + return os.environ.get("ENABLE_PROMPT_CACHING", "").strip().lower() == "true" + # Drop unsupported params rather than raising errors. # Set once at module level instead of per-call to avoid repeated # global mutation in concurrent environments. @@ -236,6 +248,7 @@ def __init__( # noqa: C901 self._enable_prompt_caching = ( bool(self.kwargs.pop("enable_prompt_caching", False)) or enable_prompt_caching + or is_prompt_caching_enabled() ) # REF: https://docs.litellm.ai/docs/completion/input#translated-openai-params @@ -375,9 +388,14 @@ def _build_messages( }, {"role": "user", "content": prompt}, ] + # Caching inactive (opt-in off, or an unsupported provider): emit no + # cache_control, but if the caller split the prompt into + # (cache_prefix, prompt) the model must still see the full text — + # concatenate rather than drop the prefix. + user_content = cache_prefix + prompt if cache_prefix is not None else prompt return [ {"role": "system", "content": self._system_prompt}, - {"role": "user", "content": prompt}, + {"role": "user", "content": user_content}, ] @capture_metrics diff --git a/unstract/sdk1/tests/test_prompt_caching.py b/unstract/sdk1/tests/test_prompt_caching.py index ba2762a771..2fe8eee470 100644 --- a/unstract/sdk1/tests/test_prompt_caching.py +++ b/unstract/sdk1/tests/test_prompt_caching.py @@ -148,21 +148,37 @@ def test_build_messages_cache_prefix_splits_user_turn() -> None: assert seen == "STABLE" + "VOLATILE" -def test_build_messages_cache_prefix_ignored_when_caching_off() -> None: +def test_build_messages_cache_prefix_preserved_when_caching_off() -> None: + """Caching off must NOT drop the prefix — the model still sees prefix+prompt. + + Regression guard: a consumer that splits its prompt into (cache_prefix, + prompt) must produce the full text even on an unsupported provider / with + caching disabled, otherwise the prefix (e.g. instructions/context) is lost. + """ llm = _StubLLM("SYSTEM", enable_prompt_caching=False) assert llm._build_messages("VOLATILE", cache_prefix="STABLE") == [ {"role": "system", "content": "SYSTEM"}, - {"role": "user", "content": "VOLATILE"}, + {"role": "user", "content": "STABLE" + "VOLATILE"}, + ] + + +def test_build_messages_cache_prefix_preserved_for_unsupported_provider() -> None: + """Flag on, but a provider we don't emit cache_control for -> full prompt, no tag.""" + llm = _StubLLM("SYSTEM", enable_prompt_caching=True, provider="openai") + assert llm._build_messages("VOLATILE", cache_prefix="STABLE") == [ + {"role": "system", "content": "SYSTEM"}, + {"role": "user", "content": "STABLE" + "VOLATILE"}, ] -# ── constructor opt-in (real LLM, no flag in adapter metadata) ────────────── +# ── constructor / env opt-in (real LLM, no flag in adapter metadata) ──────── _ANTHROPIC_ADAPTER_ID = "anthropic|90ebd4cd-2f19-4cef-a884-9eeb6ac0f203" -def test_constructor_flag_forces_caching_without_metadata() -> None: +def test_constructor_flag_forces_caching_without_metadata(monkeypatch) -> None: """A caller that builds by adapter without the stored flag can still opt in.""" + monkeypatch.delenv("ENABLE_PROMPT_CACHING", raising=False) from unstract.sdk1.llm import LLM meta = {"model": "claude-opus-4-8", "api_key": "sk-test"} @@ -178,9 +194,21 @@ def test_constructor_flag_forces_caching_without_metadata() -> None: assert messages[1]["content"][0]["cache_control"] == {"type": "ephemeral"} -def test_constructor_flag_defaults_off() -> None: +def test_constructor_flag_defaults_off(monkeypatch) -> None: + monkeypatch.delenv("ENABLE_PROMPT_CACHING", raising=False) from unstract.sdk1.llm import LLM meta = {"model": "claude-opus-4-8", "api_key": "sk-test"} llm = LLM(adapter_id=_ANTHROPIC_ADAPTER_ID, adapter_metadata=meta, system_prompt="S") assert llm._enable_prompt_caching is False + + +def test_env_var_enables_caching_platform_wide(monkeypatch) -> None: + """The ENABLE_PROMPT_CACHING master switch turns caching on with no per-call flag.""" + monkeypatch.setenv("ENABLE_PROMPT_CACHING", "true") + from unstract.sdk1.llm import LLM, is_prompt_caching_enabled + + assert is_prompt_caching_enabled() is True + meta = {"model": "claude-opus-4-8", "api_key": "sk-test"} + llm = LLM(adapter_id=_ANTHROPIC_ADAPTER_ID, adapter_metadata=meta, system_prompt="S") + assert llm._enable_prompt_caching is True diff --git a/workers/executor/executors/answer_prompt.py b/workers/executor/executors/answer_prompt.py index d1eef5b3be..feec7ee408 100644 --- a/workers/executor/executors/answer_prompt.py +++ b/workers/executor/executors/answer_prompt.py @@ -23,6 +23,18 @@ logger = logging.getLogger(__name__) +def is_prompt_caching_enabled() -> bool: + """Whether LLM prompt caching is enabled for extraction (opt-in, default off). + + When enabled, the extraction prompt is reordered so the reused document + context becomes a cached prefix (see ``construct_cached_prompt``) and the + extraction ``LLM`` is built with provider caching on. Gated by the + ``ENABLE_PROMPT_CACHING`` env var so the reorder + caching can be rolled out + and A/B-evaluated before it is on everywhere. + """ + return os.environ.get("ENABLE_PROMPT_CACHING", "").strip().lower() == "true" + + def _resolve_host_addresses(host: str) -> set[str]: """Resolve a hostname or IP string to a set of IP address strings.""" try: @@ -148,20 +160,31 @@ def construct_and_run_prompt( if not enable_word_confidence or summarize_as_source: word_confidence_postamble = "" - prompt = AnswerPromptService.construct_prompt( - preamble=tool_settings.get(PSKeys.PREAMBLE, ""), - prompt=output[prompt], - postamble=tool_settings.get(PSKeys.POSTAMBLE, ""), - grammar_list=tool_settings.get(PSKeys.GRAMMAR, []), - context=context, - platform_postamble=platform_postamble, - word_confidence_postamble=word_confidence_postamble, - prompt_type=prompt_type, - ) - output[PSKeys.COMBINED_PROMPT] = prompt + prompt_args = { + "preamble": tool_settings.get(PSKeys.PREAMBLE, ""), + "prompt": output[prompt], + "postamble": tool_settings.get(PSKeys.POSTAMBLE, ""), + "grammar_list": tool_settings.get(PSKeys.GRAMMAR, []), + "context": context, + "platform_postamble": platform_postamble, + "word_confidence_postamble": word_confidence_postamble, + "prompt_type": prompt_type, + } + cache_prefix: str | None = None + if is_prompt_caching_enabled(): + # Reorder so the reused document context becomes a cached prefix. + # The text the model sees is ``cache_prefix + prompt_str``. + cache_prefix, prompt_str = AnswerPromptService.construct_cached_prompt( + **prompt_args + ) + output[PSKeys.COMBINED_PROMPT] = cache_prefix + prompt_str + else: + prompt_str = AnswerPromptService.construct_prompt(**prompt_args) + output[PSKeys.COMBINED_PROMPT] = prompt_str return AnswerPromptService.run_completion( llm=llm, - prompt=prompt, + prompt=prompt_str, + cache_prefix=cache_prefix, metadata=metadata, prompt_key=output[PSKeys.NAME], prompt_type=prompt_type, @@ -218,10 +241,48 @@ def construct_prompt( ) return prompt + @staticmethod + def construct_cached_prompt( + preamble: str, + prompt: str, + postamble: str, + grammar_list: list[dict[str, Any]], + context: str, + platform_postamble: str, + word_confidence_postamble: str, + prompt_type: str = "text", + ) -> tuple[str, str]: + """Build ``(cache_prefix, volatile)`` with the document context first. + + Same content as :meth:`construct_prompt`, but the reused ``context`` is + moved to the front so it forms a cacheable prefix that repeats across + every prompt run against the same document, while the per-prompt + question is the volatile suffix. The model sees ``cache_prefix + + volatile``. This reorders the prompt (context before question instead of + after), which is why it is flag-gated and A/B-evaluated. + """ + if prompt_type == PSKeys.JSON: + json_postamble = os.environ.get( + PSKeys.JSON_POSTAMBLE, PSKeys.DEFAULT_JSON_POSTAMBLE + ) + postamble += f"\n{json_postamble}" + if platform_postamble: + platform_postamble += "\n\n" + if word_confidence_postamble: + platform_postamble += f"{word_confidence_postamble}\n\n" + cache_prefix = f"Context:\n---------------\n{context}\n-----------------\n\n" + volatile = ( + f"{preamble}\n\nQuestion or Instruction: {prompt}" + + AnswerPromptService._build_grammar_notes(grammar_list) + + f"\n\n{postamble}\n\n{platform_postamble}Answer:" + ) + return cache_prefix, volatile + @staticmethod def run_completion( llm: Any, prompt: str, + cache_prefix: str | None = None, metadata: dict[str, str] | None = None, prompt_key: str | None = None, prompt_type: str | None = "text", @@ -249,6 +310,7 @@ def run_completion( try: completion = llm.complete( prompt=prompt, + cache_prefix=cache_prefix, process_text=process_text, extract_json=prompt_type.lower() != PSKeys.TEXT, ) diff --git a/workers/executor/executors/legacy_executor.py b/workers/executor/executors/legacy_executor.py index ce7fbea0d1..b672d56ab6 100644 --- a/workers/executor/executors/legacy_executor.py +++ b/workers/executor/executors/legacy_executor.py @@ -1928,6 +1928,8 @@ def _init_llm_and_retrieval( from executor.executors.constants import PromptServiceConstants as PSKeys try: + # Caching is enabled globally via the ENABLE_PROMPT_CACHING master + # switch (honored inside the SDK ``LLM``); no per-call flag needed. llm = llm_cls( adapter_instance_id=output[PSKeys.LLM], tool=shim, diff --git a/workers/tests/test_answer_prompt_caching.py b/workers/tests/test_answer_prompt_caching.py new file mode 100644 index 0000000000..7582c6f638 --- /dev/null +++ b/workers/tests/test_answer_prompt_caching.py @@ -0,0 +1,95 @@ +"""Guardrail tests for flag-gated prompt-caching in answer_prompt. + +When ``ENABLE_PROMPT_CACHING`` is on, ``construct_cached_prompt`` reorders the +prompt so the reused document context becomes a cacheable prefix (context +first) instead of a suffix (context last). These tests lock in that: + +- the default (flag off) prompt is unchanged (context last), +- the cached variant is a pure *reorder* — every piece of the original prompt + is preserved, only the context moves to the front, +- ``cache_prefix`` is exactly the context block (no per-prompt question), so + it repeats byte-for-byte across prompts on the same document, +- the flag reads the env var and defaults off. + +The executor package's ``__init__`` pulls the full celery stack, so we load the +module with stubbed parent packages — the methods under test are pure strings. +""" + +import importlib +import os +import sys +import types + +import pytest + +_WORKERS = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +def _load_answer_prompt(): + for pkg, rel in [("executor", "executor"), ("executor.executors", "executor/executors")]: + if pkg not in sys.modules: + mod = types.ModuleType(pkg) + mod.__path__ = [os.path.join(_WORKERS, rel)] + sys.modules[pkg] = mod + return importlib.import_module("executor.executors.answer_prompt") + + +_mod = _load_answer_prompt() +A = _mod.AnswerPromptService + +_ARGS = dict( + preamble="You are an extractor.", + prompt="What is the tenant name?", + postamble="Answer concisely.", + grammar_list=[], + context="UNIT 101 John Smith $1,450\nUNIT 102 Maria Davis $1,525", + platform_postamble="", + word_confidence_postamble="", + prompt_type="text", +) + + +def test_default_prompt_is_context_last(): + off = A.construct_prompt(**_ARGS) + assert off.index("Question or Instruction") < off.index("UNIT 101") + + +def test_cached_prompt_is_context_first(): + prefix, volatile = A.construct_cached_prompt(**_ARGS) + full = prefix + volatile + assert full.index("UNIT 101") < full.index("Question or Instruction") + + +def test_cache_prefix_is_context_block_only(): + prefix, _volatile = A.construct_cached_prompt(**_ARGS) + assert prefix.startswith("Context:") + assert "UNIT 101" in prefix + # The volatile per-prompt question must NOT leak into the cached prefix, + # or the prefix would differ per prompt and never hit the cache. + assert "Question or Instruction" not in prefix + + +def test_cached_variant_is_a_pure_reorder_no_content_lost(): + prefix, volatile = A.construct_cached_prompt(**_ARGS) + full = prefix + volatile + for piece in ( + "You are an extractor.", + "Question or Instruction: What is the tenant name?", + "Answer concisely.", + "UNIT 101 John Smith $1,450", + "Answer:", + ): + assert piece in full, f"missing from cached prompt: {piece!r}" + + +def test_flag_defaults_off_and_reads_env(monkeypatch): + monkeypatch.delenv("ENABLE_PROMPT_CACHING", raising=False) + assert _mod.is_prompt_caching_enabled() is False + monkeypatch.setenv("ENABLE_PROMPT_CACHING", "true") + assert _mod.is_prompt_caching_enabled() is True + monkeypatch.setenv("ENABLE_PROMPT_CACHING", "false") + assert _mod.is_prompt_caching_enabled() is False + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-q"]))