diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py index 875a457708..4e054576e6 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py @@ -1136,6 +1136,24 @@ 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/Claude models on Bedrock support prompt caching, so + # don't advertise the flag for other Bedrock families (Titan, Llama, + # etc.). The LLM layer enforces the same model gate; this just keeps the + # validated metadata honest. + # Check both ``model`` and ``model_id`` so callers routing through a + # Bedrock Application Inference Profile (opaque ARN in ``model``, Claude + # id in ``model_id``) still qualify. + bedrock_model_ids = " ".join( + str(result_metadata.get(field, "")) for field in _MODEL_ID_FIELDS + ).lower() + enable_prompt_caching = bool( + adapter_metadata.get("enable_prompt_caching", False) + ) and ("anthropic" in bedrock_model_ids or "claude" in bedrock_model_ids) + _pack_bedrock_guardrail_config(result_metadata) # Create validation metadata excluding control fields. `auth_type` is @@ -1152,6 +1170,7 @@ def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]: "guardrail_identifier", "guardrail_version", "guardrail_trace", + "enable_prompt_caching", ) } @@ -1170,6 +1189,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 @@ -1237,6 +1257,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", @@ -1244,6 +1270,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 @@ -1259,6 +1286,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 b0a712b49f..e780685180 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -4,7 +4,7 @@ from collections.abc import Callable, Generator, Mapping, Sequence from dataclasses import dataclass, field from enum import Enum -from functools import lru_cache +from functools import cache, lru_cache from typing import Any, NoReturn, cast import litellm @@ -32,6 +32,44 @@ logger = logging.getLogger(__name__) + +# Truthy-looking values that people commonly set expecting a boolean flag to +# turn on — but only "true" enables caching. We warn (once each) so a stray +# ENABLE_PROMPT_CACHING=1 doesn't leave caching silently off. +_PROMPT_CACHING_TRUTHY_LOOKALIKES = frozenset( + {"1", "yes", "y", "on", "t", "enable", "enabled"} +) + + +@cache +def _warn_prompt_caching_lookalike(value: str) -> None: + logger.warning( + "ENABLE_PROMPT_CACHING=%r is not recognized as enabled; only 'true' " + "(case-insensitive) turns prompt caching on — caching stays OFF.", + value, + ) + + +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. + + Only ``"true"`` (case-insensitive) enables it; a truthy-looking value like + ``"1"``/``"yes"``/``"on"`` logs a one-time warning and stays off. + """ + raw = os.environ.get("ENABLE_PROMPT_CACHING", "").strip().lower() + if raw == "true": + return True + if raw in _PROMPT_CACHING_TRUTHY_LOOKALIKES: + _warn_prompt_caching_lookalike(raw) + return False + + # Lets tests force a deterministic completion without a provider or a secret. # Unset in production, where this is a no-op. _MOCK_RESPONSE_ENV = "UNSTRACT_LLM_MOCK_RESPONSE" @@ -187,6 +225,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. @@ -199,6 +238,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 = {} @@ -250,6 +292,16 @@ def __init__( # noqa: C901 self.kwargs = self.adapter.validate(self._adapter_metadata) self._cost_model = self.kwargs.pop("cost_model", None) self.kwargs.pop("context_window", 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 + or is_prompt_caching_enabled() + ) # REF: https://docs.litellm.ai/docs/completion/input#translated-openai-params # supported = get_supported_openai_params(model=self.kwargs["model"], @@ -328,8 +380,129 @@ 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"}) + # Bedrock hosts many model families (Anthropic Claude, Amazon Titan/Nova, + # Meta Llama, Cohere, Mistral, AI21). Only Anthropic/Claude models on + # Bedrock honor ``cache_control``; emitting the blocks for other families + # would be ineffective and could produce unsupported message shapes. These + # substrings identify the cache-capable Bedrock models by their model id. + _BEDROCK_CACHE_MODEL_MARKERS = ("anthropic", "claude") + + def _prompt_caching_active(self) -> bool: + """Whether to emit ``cache_control`` blocks for this call.""" + if not self._enable_prompt_caching: + return False + provider = self.adapter.get_provider() + if provider not in self._PROMPT_CACHE_PROVIDERS: + return False + if provider == "bedrock": + # Gate on the underlying model, not just the provider: only + # Anthropic/Claude models on Bedrock support cache_control. Check + # 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 appears only in ``model_id``. + recognized = any( + marker in str(self.kwargs.get(field, "")).lower() + for field in ("model", "model_id") + for marker in self._BEDROCK_CACHE_MODEL_MARKERS + ) + if not recognized: + # Enabled but the model can't be confirmed as Anthropic/Claude + # (e.g. a fully opaque Application Inference Profile ARN with no + # Claude id in model/model_id). Skipping is safe; leave a + # breadcrumb so operators can diagnose a Claude-on-Bedrock call + # that unexpectedly isn't caching. + logger.debug( + "Prompt caching enabled but skipped for Bedrock: " + "model=%r model_id=%r not recognized as Anthropic/Claude", + self.kwargs.get("model"), + self.kwargs.get("model_id"), + ) + return recognized + return True + + def is_prompt_caching_active(self) -> bool: + """Public: whether ``cache_control`` blocks are emitted for this LLM. + + True only when caching is enabled (adapter flag, constructor arg, or the + ``ENABLE_PROMPT_CACHING`` master switch) *and* the provider/model + supports it. Callers use this to decide whether reordering a prompt into + a cached prefix is worthwhile — reordering for a non-caching provider + changes prompt structure with no benefit. + """ + return self._prompt_caching_active() + + 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. + + - non-empty ``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. + An empty ``cache_prefix`` is treated as absent — Anthropic rejects empty + text content blocks, and an empty prefix carries no caching benefit. + """ + if self._prompt_caching_active() and cache_prefix: + 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}, + ] + # 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 else prompt + return [ + {"role": "system", "content": self._system_prompt}, + {"role": "user", "content": user_content}, + ] + @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 @@ -337,6 +510,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: @@ -344,10 +522,7 @@ 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" ) @@ -355,6 +530,7 @@ def complete(self, prompt: str, **kwargs: object) -> dict[str, object]: completion_kwargs = self.adapter.validate({**self.kwargs, **kwargs}) _inject_mock_response(completion_kwargs) completion_kwargs.pop("cost_model", None) + completion_kwargs.pop("enable_prompt_caching", None) completion_kwargs.pop("context_window", None) # if hasattr(self, "model") and self.model not in O1_MODELS: @@ -479,6 +655,7 @@ def complete_vision( completion_kwargs = self.adapter.validate({**self.kwargs, **kwargs}) _inject_mock_response(completion_kwargs) completion_kwargs.pop("cost_model", None) + completion_kwargs.pop("enable_prompt_caching", None) completion_kwargs.pop("context_window", None) response: dict[str, object] = litellm.completion( @@ -530,17 +707,17 @@ 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" ) @@ -548,6 +725,7 @@ def stream_complete( completion_kwargs = self.adapter.validate({**self.kwargs, **kwargs}) _inject_mock_response(completion_kwargs) completion_kwargs.pop("cost_model", None) + completion_kwargs.pop("enable_prompt_caching", None) completion_kwargs.pop("context_window", None) max_retries = pop_litellm_retry_kwargs( @@ -607,13 +785,20 @@ def stream_complete( message=error_msg, status_code=status_code, actual_err=e ) from e - async def acomplete(self, prompt: str, **kwargs: object) -> dict[str, object]: - """Asynchronous chat completion (wrapper around ``litellm.acompletion``).""" + async def acomplete( + self, + prompt: str, + cache_prefix: str | None = None, + **kwargs: object, + ) -> dict[str, object]: + """Asynchronous chat completion (wrapper around ``litellm.acompletion``). + + ``cache_prefix`` mirrors :meth:`complete` / :meth:`stream_complete`: when + prompt caching is active it is emitted as a cached stable prefix ahead of + ``prompt``; otherwise it is concatenated so the model sees the same text. + """ 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()} async completion API" ) @@ -621,6 +806,7 @@ async def acomplete(self, prompt: str, **kwargs: object) -> dict[str, object]: completion_kwargs = self.adapter.validate({**self.kwargs, **kwargs}) _inject_mock_response(completion_kwargs) completion_kwargs.pop("cost_model", None) + completion_kwargs.pop("enable_prompt_caching", None) completion_kwargs.pop("context_window", None) max_retries = pop_litellm_retry_kwargs( @@ -767,6 +953,57 @@ 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: + # Pass ``model`` so cached calls price against the same model as + # the ``cost_per_token`` fallback below. Without it, + # ``completion_cost`` derives the model from the response and + # ignores any ``cost_model`` override. + return litellm.completion_cost(completion_response=response, model=model) + except Exception: + # Warn (not debug): the cost_per_token fallback prices every + # prompt token at the full input rate, so it over-reports cost + # by up to ~10x on cache hits. Operators watching spend need to + # see that a recorded cost may be inflated. + logger.warning( + "completion_cost() failed for model=%s; falling back to " + "cost_per_token — recorded cost may be OVER-reported for " + "this cached call", + 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, @@ -779,6 +1016,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: @@ -803,30 +1044,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..22912525ef --- /dev/null +++ b/unstract/sdk1/tests/test_prompt_caching.py @@ -0,0 +1,551 @@ +"""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 + is_prompt_caching_active = LLM.is_prompt_caching_active + _PROMPT_CACHE_PROVIDERS = LLM._PROMPT_CACHE_PROVIDERS + _BEDROCK_CACHE_MODEL_MARKERS = LLM._BEDROCK_CACHE_MODEL_MARKERS + + def __init__( + self, + system_prompt: str, + enable_prompt_caching: bool, + provider: str = "anthropic", + model: str = "claude-opus-4-8", + model_id: str | None = None, + ) -> None: + self._system_prompt = system_prompt + self._enable_prompt_caching = enable_prompt_caching + self.adapter = _StubAdapter(provider) + # Only read by the Bedrock model gate; harmless for other providers. + self.kwargs = {"model": model} + if model_id is not None: + self.kwargs["model_id"] = model_id + + +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_empty_cache_prefix_treated_as_absent() -> None: + """An empty cache_prefix must NOT create an empty (Anthropic-rejected) block. + + Caching is active, but ``cache_prefix=""`` should fall back to the plain + single-block user turn rather than emitting ``{"type": "text", "text": ""}``. + """ + llm = _StubLLM("SYSTEM", enable_prompt_caching=True) + assert llm._build_messages("VOLATILE", cache_prefix="") == [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "SYSTEM", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + {"role": "user", "content": "VOLATILE"}, + ] + + +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": "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"}, + ] + + +# ── Bedrock model gate: cache_control only for Anthropic/Claude on Bedrock ─── + +# Bedrock hosts many families; only Anthropic/Claude honor cache_control. +_BEDROCK_ANTHROPIC_MODELS = [ + "bedrock/anthropic.claude-opus-4-8-20260101-v1:0", + "bedrock/us.anthropic.claude-sonnet-4-6-20250101-v1:0", + "anthropic.claude-3-5-sonnet-20241022-v2:0", +] +_BEDROCK_NON_ANTHROPIC_MODELS = [ + "bedrock/amazon.titan-text-premier-v1:0", + "bedrock/meta.llama3-70b-instruct-v1:0", + "bedrock/cohere.command-r-plus-v1:0", + "bedrock/mistral.mistral-large-2407-v1:0", +] + + +@pytest.mark.parametrize("model", _BEDROCK_ANTHROPIC_MODELS) +def test_bedrock_anthropic_model_caches(model: str) -> None: + llm = _StubLLM("SYSTEM", enable_prompt_caching=True, provider="bedrock", model=model) + assert llm._prompt_caching_active() is True + # cache_control block is emitted on the split user turn. + user = llm._build_messages("VOLATILE", cache_prefix="STABLE")[1] + assert user["content"][0]["cache_control"] == {"type": "ephemeral"} + + +@pytest.mark.parametrize("model", _BEDROCK_NON_ANTHROPIC_MODELS) +def test_bedrock_non_anthropic_model_does_not_cache(model: str) -> None: + """Titan/Llama/Cohere/Mistral on Bedrock must not get cache_control blocks.""" + llm = _StubLLM("SYSTEM", enable_prompt_caching=True, provider="bedrock", model=model) + assert llm._prompt_caching_active() is False + # Falls back to the plain string form; prefix still preserved as full text. + assert llm._build_messages("VOLATILE", cache_prefix="STABLE") == [ + {"role": "system", "content": "SYSTEM"}, + {"role": "user", "content": "STABLE" + "VOLATILE"}, + ] + + +def test_bedrock_opaque_inference_profile_uses_model_id() -> None: + """AIP ARN in ``model`` is opaque; the Claude id in ``model_id`` still caches.""" + llm = _StubLLM( + "SYSTEM", + enable_prompt_caching=True, + provider="bedrock", + model="arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcd1234", + model_id="anthropic.claude-opus-4-8-20260101-v1:0", + ) + assert llm._prompt_caching_active() is True + + +def test_bedrock_opaque_profile_without_claude_id_does_not_cache() -> None: + """No Anthropic id in either field -> can't confirm Claude -> no caching.""" + llm = _StubLLM( + "SYSTEM", + enable_prompt_caching=True, + provider="bedrock", + model="arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcd1234", + ) + assert llm._prompt_caching_active() is False + + +def test_is_prompt_caching_active_matches_private() -> None: + """The public probe mirrors the internal gate for callers (e.g. answer_prompt).""" + on = _StubLLM("SYSTEM", enable_prompt_caching=True, provider="anthropic") + assert on.is_prompt_caching_active() is True + off_provider = _StubLLM("SYSTEM", enable_prompt_caching=True, provider="openai") + assert off_provider.is_prompt_caching_active() is False + off_bedrock = _StubLLM( + "SYSTEM", + enable_prompt_caching=True, + provider="bedrock", + model="bedrock/amazon.titan-text-premier-v1:0", + ) + assert off_bedrock.is_prompt_caching_active() is False + + +def test_validate_bedrock_non_anthropic_never_enables_caching() -> None: + """base1 keeps the validated flag honest: non-Anthropic Bedrock -> False.""" + result = AWSBedrockLLMParameters.validate( + { + "model": "amazon.titan-text-premier-v1:0", + "enable_prompt_caching": True, + "aws_region_name": "us-east-1", + } + ) + assert result["enable_prompt_caching"] is False + + +# ── cost accounting: cached calls price against the cost_model override ────── + + +def test_cost_override_passed_to_completion_cost_on_cached_call( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Cached calls price against the ``cost_model`` override. + + ``completion_cost`` must receive ``model=`` so cached and non-cached calls + price against the same model, matching the ``cost_per_token`` fallback path. + """ + import types + + import unstract.sdk1.llm as llm_mod + from unstract.sdk1.llm import LLM + + captured: dict[str, Any] = {} + + def _fake_completion_cost(**kwargs: object) -> float: + captured.update(kwargs) + return 0.42 + + monkeypatch.setattr(llm_mod.litellm, "completion_cost", _fake_completion_cost) + + # Minimal stub instance for ``self`` so a future ``self.`` reference fails + # loudly instead of silently working (unlike a bare ``object()``). + cost = LLM._compute_call_cost( + types.SimpleNamespace(), + model="anthropic/claude-opus-4-8", + prompt_tokens=100, + completion_tokens=10, + has_cache_tokens=True, + response={"id": "resp"}, + ) + assert cost == 0.42 + assert captured.get("model") == "anthropic/claude-opus-4-8" + assert captured.get("completion_response") == {"id": "resp"} + + +def test_compute_cost_falls_back_when_completion_cost_raises( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """completion_cost failure on a cached call falls back to cost_per_token.""" + import types + + import unstract.sdk1.llm as llm_mod + from unstract.sdk1.llm import LLM + + def _boom(**kwargs: object) -> float: + raise RuntimeError("no price map") + + monkeypatch.setattr(llm_mod.litellm, "completion_cost", _boom) + monkeypatch.setattr(llm_mod.litellm, "cost_per_token", lambda **k: (0.1, 0.2)) + + cost = LLM._compute_call_cost( + types.SimpleNamespace(), + model="m", + prompt_tokens=100, + completion_tokens=10, + has_cache_tokens=True, + response={"id": "r"}, + ) + assert cost == pytest.approx(0.3) + + +def test_compute_cost_returns_zero_when_both_paths_fail( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """If both completion_cost and cost_per_token fail, cost is recorded as 0.0.""" + import types + + import unstract.sdk1.llm as llm_mod + from unstract.sdk1.llm import LLM + + def _boom(**kwargs: object) -> float: + raise RuntimeError("nope") + + monkeypatch.setattr(llm_mod.litellm, "completion_cost", _boom) + monkeypatch.setattr(llm_mod.litellm, "cost_per_token", _boom) + + cost = LLM._compute_call_cost( + types.SimpleNamespace(), + model="m", + prompt_tokens=1, + completion_tokens=1, + has_cache_tokens=True, + response={"id": "r"}, + ) + assert cost == 0.0 + + +def test_compute_cost_skips_completion_cost_when_response_is_none( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """has_cache_tokens True but no response -> skip completion_cost, use per-token.""" + import types + + import unstract.sdk1.llm as llm_mod + from unstract.sdk1.llm import LLM + + called = {"completion_cost": False} + + def _cc(**kwargs: object) -> float: + called["completion_cost"] = True + return 9.9 + + monkeypatch.setattr(llm_mod.litellm, "completion_cost", _cc) + monkeypatch.setattr(llm_mod.litellm, "cost_per_token", lambda **k: (0.4, 0.6)) + + cost = LLM._compute_call_cost( + types.SimpleNamespace(), + model="m", + prompt_tokens=1, + completion_tokens=1, + has_cache_tokens=True, + response=None, + ) + assert called["completion_cost"] is False + assert cost == pytest.approx(1.0) + + +# ── control flags never leak into the litellm completion payload ───────────── + + +def test_control_flags_not_forwarded_to_litellm_completion( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """enable_prompt_caching / cost_model / context_window never reach litellm. + + They are applied on our side (cache_control blocks, cost accounting) and must + be popped before ``litellm.completion``; otherwise litellm's Pydantic + validation rejects the unknown params. + """ + import unstract.sdk1.llm as llm_mod + from unstract.sdk1.exceptions import LLMError + from unstract.sdk1.llm import LLM + + captured: dict[str, Any] = {} + + def _capture_completion(**kwargs: object) -> dict[str, object]: + captured.update(kwargs) + raise RuntimeError("stop after capture") + + monkeypatch.setattr(llm_mod.litellm, "completion", _capture_completion) + monkeypatch.delenv("ENABLE_PROMPT_CACHING", raising=False) + + meta = { + "model": "claude-opus-4-8", + "api_key": "sk-test", + "enable_prompt_caching": True, + } + llm = LLM(adapter_id=_ANTHROPIC_ADAPTER_ID, adapter_metadata=meta, system_prompt="S") + with pytest.raises(LLMError): + llm.complete("hi") + + assert captured.get("messages"), "payload should have been built and forwarded" + assert "enable_prompt_caching" not in captured + assert "cost_model" not in captured + assert "context_window" not in captured + + +def test_control_flags_not_forwarded_to_litellm_acompletion( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Same strip guarantee on the async surface (regression for acomplete).""" + import asyncio + + import unstract.sdk1.llm as llm_mod + from unstract.sdk1.exceptions import LLMError + from unstract.sdk1.llm import LLM + + captured: dict[str, Any] = {} + + async def _capture_acompletion(**kwargs: object) -> dict[str, object]: + captured.update(kwargs) + raise RuntimeError("stop after capture") + + monkeypatch.setattr(llm_mod.litellm, "acompletion", _capture_acompletion) + monkeypatch.delenv("ENABLE_PROMPT_CACHING", raising=False) + + meta = { + "model": "claude-opus-4-8", + "api_key": "sk-test", + "enable_prompt_caching": True, + } + llm = LLM(adapter_id=_ANTHROPIC_ADAPTER_ID, adapter_metadata=meta, system_prompt="S") + with pytest.raises(LLMError): + asyncio.run(llm.acomplete("hi", cache_prefix="STABLE")) + + assert captured.get("messages"), "async payload should have been forwarded" + assert "enable_prompt_caching" not in captured + assert "cost_model" not in captured + + +# ── 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( + monkeypatch: pytest.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"} + 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(monkeypatch: pytest.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: pytest.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 + + +def test_env_var_true_is_case_insensitive(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("ENABLE_PROMPT_CACHING", "TRUE") + from unstract.sdk1.llm import is_prompt_caching_enabled + + assert is_prompt_caching_enabled() is True + + +def test_env_var_truthy_lookalike_stays_off_and_warns( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """A value like '1'/'yes'/'on' does NOT enable caching, but warns once.""" + import logging + + from unstract.sdk1.llm import ( + _warn_prompt_caching_lookalike, + is_prompt_caching_enabled, + ) + + _warn_prompt_caching_lookalike.cache_clear() + monkeypatch.setenv("ENABLE_PROMPT_CACHING", "1") + with caplog.at_level(logging.WARNING): + assert is_prompt_caching_enabled() is False + assert any("not recognized as enabled" in r.message for r in caplog.records) diff --git a/workers/executor/executors/answer_prompt.py b/workers/executor/executors/answer_prompt.py index d1eef5b3be..4184b78b71 100644 --- a/workers/executor/executors/answer_prompt.py +++ b/workers/executor/executors/answer_prompt.py @@ -104,6 +104,41 @@ def extract_variable( ) return promptx + @staticmethod + def _llm_caches_prompts(llm: Any) -> bool: + """Whether reordering into a cached prefix pays off for this LLM. + + True only when the LLM reports prompt caching is active for it — caching + enabled (adapter flag / constructor arg / ``ENABLE_PROMPT_CACHING`` + master switch) AND a supported provider/model. LLM objects that don't + expose the capability default to False, so their prompt order is left + unchanged. + """ + checker = getattr(llm, "is_prompt_caching_active", None) + if not callable(checker): + return False + # Fail closed: if the probe raises, keep the original (context-last) + # prompt order rather than risk a reorder for an LLM whose caching + # support is unknown. Hence the deliberately broad excepts below. + try: + return bool(checker()) + except Exception: # noqa: BLE001 - fail closed, see comment above + provider = getattr(getattr(llm, "adapter", None), "get_provider", None) + provider_name = None + if callable(provider): + try: + provider_name = provider() + except Exception: # noqa: BLE001 - provider is best-effort log context + provider_name = None + logger.warning( + "Failed to determine prompt-caching support for LLM " + "(type=%s, provider=%s); using original prompt order", + type(llm).__name__, + provider_name, + exc_info=True, + ) + return False + @staticmethod def construct_and_run_prompt( tool_settings: dict[str, Any], @@ -148,20 +183,36 @@ 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 + # Only reorder into a cached prefix when this LLM actually caches + # (caching enabled AND a supported provider/model). Reordering for an + # unsupported provider would change the prompt structure — moving the + # document context ahead of the instructions — with no caching benefit, + # so those providers keep the original prompt order. + if AnswerPromptService._llm_caches_prompts(llm): + # 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, @@ -189,6 +240,31 @@ def _build_grammar_notes(grammar_list: list[dict[str, Any]]) -> str: ) return notes + @staticmethod + def _prepare_postambles( + postamble: str, + platform_postamble: str, + word_confidence_postamble: str, + prompt_type: str, + ) -> tuple[str, str]: + """Apply JSON + platform/word-confidence formatting to the postambles. + + Shared by :meth:`construct_prompt` and :meth:`construct_cached_prompt` so + the cached and non-cached prompts stay byte-identical apart from the + context/question ordering. Returns the ``(postamble, platform_postamble)`` + pair to interpolate. + """ + 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" + return postamble, platform_postamble + @staticmethod def construct_prompt( preamble: str, @@ -203,25 +279,51 @@ def construct_prompt( """Build the full prompt string with preamble, grammar, postamble, context.""" prompt = f"{preamble}\n\nQuestion or Instruction: {prompt}" prompt += AnswerPromptService._build_grammar_notes(grammar_list) - 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" + postamble, platform_postamble = AnswerPromptService._prepare_postambles( + postamble, platform_postamble, word_confidence_postamble, prompt_type + ) prompt += ( f"\n\n{postamble}\n\nContext:\n---------------\n{context}\n" f"-----------------\n\n{platform_postamble}Answer:" ) 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. + """ + postamble, platform_postamble = AnswerPromptService._prepare_postambles( + postamble, platform_postamble, word_confidence_postamble, prompt_type + ) + 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 +351,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..28fa006620 --- /dev/null +++ b/workers/tests/test_answer_prompt_caching.py @@ -0,0 +1,235 @@ +"""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 reorder only happens when the selected LLM actually caches + (``is_prompt_caching_active()``); unsupported LLMs keep the original order. + +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(): + """Import answer_prompt without triggering the executor package's celery stack. + + The real ``executor`` / ``executor.executors`` package ``__init__``s pull the + full celery worker stack, so we temporarily register lightweight namespace + stubs pointing at the real source dirs, import the module, then remove any + stub we added. The imported module (and its already-resolved ``constants`` / + ``exceptions`` imports) stay cached under their own names, so the stubs are + unneeded afterwards — and removing them keeps ``sys.modules`` clean for other + tests in the same process instead of leaving synthetic packages behind. + """ + injected = [] + 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 + injected.append(pkg) + try: + return importlib.import_module("executor.executors.answer_prompt") + finally: + for pkg in injected: + sys.modules.pop(pkg, None) + + +_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}" + + +# --- shared postamble formatting (cached and uncached must not diverge) ------ + + +def test_prepare_postambles_applies_json_and_platform_formatting(): + json_postamble = os.environ.get( + _mod.PSKeys.JSON_POSTAMBLE, _mod.PSKeys.DEFAULT_JSON_POSTAMBLE + ) + post, plat = A._prepare_postambles("BASE", "PLATFORM", "WORDCONF", _mod.PSKeys.JSON) + assert post == f"BASE\n{json_postamble}" + assert plat == "PLATFORM\n\nWORDCONF\n\n" + + +def test_prepare_postambles_noop_for_text_without_platform(): + post, plat = A._prepare_postambles("BASE", "", "", "text") + assert post == "BASE" + assert plat == "" + + +def test_cached_and_uncached_share_postamble_formatting(): + """Both builders route postambles through the shared helper, so a JSON + + platform postamble is formatted identically — guards against silent + divergence between cached and non-cached prompts (breaks A/B comparison).""" + args = dict(_ARGS) + args.update( + prompt_type=_mod.PSKeys.JSON, + postamble="BASE_POST", + platform_postamble="PLATFORM", + word_confidence_postamble="WORDCONF", + ) + flat = A.construct_prompt(**args) + prefix, volatile = A.construct_cached_prompt(**args) + cached = prefix + volatile + json_postamble = os.environ.get( + _mod.PSKeys.JSON_POSTAMBLE, _mod.PSKeys.DEFAULT_JSON_POSTAMBLE + ) + for piece in ("BASE_POST", "PLATFORM", "WORDCONF", json_postamble): + assert piece in flat, f"missing from construct_prompt: {piece!r}" + assert piece in cached, f"missing from construct_cached_prompt: {piece!r}" + + +# NOTE: the ENABLE_PROMPT_CACHING env-var master switch is owned by the SDK +# (``unstract.sdk1.llm.is_prompt_caching_enabled``) and covered by its tests; +# answer_prompt gates purely on the LLM's ``is_prompt_caching_active()`` probe +# (see ``_llm_caches_prompts`` tests below), so there is no local env flag to +# test here. + + +# --- gate: only reorder into a cached prefix when the LLM actually caches ----- + + +class _FakeLLM: + """Minimal LLM stub exposing the SDK's caching-capability probe.""" + + def __init__(self, active: bool): + self._active = active + + def is_prompt_caching_active(self) -> bool: + return self._active + + +class _LegacyLLM: + """LLM stub without the capability probe (older SDK / mock).""" + + +def test_llm_caches_prompts_probe(): + assert A._llm_caches_prompts(_FakeLLM(True)) is True + assert A._llm_caches_prompts(_FakeLLM(False)) is False + # An LLM that doesn't expose the probe must default to "no caching". + assert A._llm_caches_prompts(_LegacyLLM()) is False + + +def test_llm_caches_prompts_probe_swallows_errors(): + class _BoomLLM: + def is_prompt_caching_active(self): + raise RuntimeError("boom") + + assert A._llm_caches_prompts(_BoomLLM()) is False + + +def _run_and_capture(monkeypatch, llm): + """Run construct_and_run_prompt with run_completion stubbed to capture args.""" + captured: dict = {} + + def _fake_run_completion(**kwargs): + captured.update(kwargs) + return "answer" + + monkeypatch.setattr(A, "run_completion", staticmethod(_fake_run_completion)) + + tool_settings = { + _mod.PSKeys.PREAMBLE: "You are an extractor.", + _mod.PSKeys.POSTAMBLE: "Answer concisely.", + _mod.PSKeys.GRAMMAR: [], + } + output = { + "promptx": "What is the tenant name?", + _mod.PSKeys.NAME: "q1", + _mod.PSKeys.TYPE: "text", + } + A.construct_and_run_prompt( + tool_settings=tool_settings, + output=output, + llm=llm, + context="UNIT 101 John Smith $1,450", + prompt="promptx", + metadata={}, + ) + return captured, output + + +def test_supported_llm_reorders_and_sends_cache_prefix(monkeypatch): + captured, output = _run_and_capture(monkeypatch, _FakeLLM(True)) + # Cached path: context-first prompt + a cache_prefix that is the context. + assert captured["cache_prefix"] is not None + assert captured["cache_prefix"].startswith("Context:") + combined = output[_mod.PSKeys.COMBINED_PROMPT] + assert combined.index("UNIT 101") < combined.index("Question or Instruction") + + +def test_unsupported_llm_keeps_original_order_and_no_cache_prefix(monkeypatch): + # Global flag ON, but the LLM's provider/model does not support caching. + monkeypatch.setenv("ENABLE_PROMPT_CACHING", "true") + captured, output = _run_and_capture(monkeypatch, _FakeLLM(False)) + # No cache prefix, and the original context-last prompt order is preserved. + assert captured["cache_prefix"] is None + combined = output[_mod.PSKeys.COMBINED_PROMPT] + assert combined.index("Question or Instruction") < combined.index("UNIT 101") + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-q"]))