UN-3810 [FEAT] Prompt caching: SDK master switch + document-context caching (answer_prompt) #2199
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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. | ||
|
|
@@ -162,6 +174,7 @@ | |
| 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 +187,9 @@ | |
| 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 +240,16 @@ | |
|
|
||
| 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 | ||
| 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"], | ||
|
|
@@ -302,32 +328,110 @@ | |
| 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 | ||
| ) | ||
|
|
||
|
Comment on lines
+331
to
+343
Contributor
There was a problem hiding this comment. 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
For 🤖 Prompt for AI Agents |
||
| 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}, | ||
| ] | ||
| # 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": user_content}, | ||
| ] | ||
|
|
||
| @capture_metrics | ||
| 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 | ||
| ID is provided. | ||
|
|
||
| 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: | ||
| dict[str, Any] : A dictionary containing the result of the completion, | ||
| any processed output, and the captured metrics (if applicable). | ||
| """ | ||
| try: | ||
| 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 +555,7 @@ | |
|
|
||
| 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 +606,24 @@ | |
| 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. ``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 = 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 +685,14 @@ | |
| async def acomplete(self, prompt: str, **kwargs: object) -> dict[str, object]: | ||
| """Asynchronous chat completion (wrapper around ``litellm.acompletion``).""" | ||
| try: | ||
| 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,7 +825,49 @@ | |
| 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 | ||
|
|
||
|
Comment on lines
+828
to
+869
Contributor
There was a problem hiding this comment. 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win 🧩 Analysis chain🌐 Web query:
💡 Result: To handle cases where you need to override the model for cost calculation in LiteLLM (for example, when a completion response returns a model name that differs from the one you want to bill against), you should use the base_model parameter [1]. When using the LiteLLM proxy or configuration, you can set the base_model in your model_list configuration. This tells LiteLLM to use the specified base_model for pricing lookups instead of the model name returned in the completion response [1]. If you are calling the completion_cost function directly in your code, you can explicitly pass the desired model name via the model argument to ensure the cost is calculated against that specific pricing tier [2][3]. For more advanced scenarios involving custom pricing, you can: 1. Override specific cost parameters (e.g., input_cost_per_token, output_cost_per_token) within the model_info block of your configuration [4][1]. 2. Use register_model to manually define or override model pricing in the global LiteLLM cost dictionary [2]. 3. For proxy callbacks, it is often more reliable to use the pre-computed response_cost found in the kwargs of the response object rather than manually invoking completion_cost, as this value is typically calculated by LiteLLM using the intended model context before aliasing occurs [5]. When overriding, note that setting custom pricing fields detaches the deployment from the default cost map for that entry [4]. Ensure that you define all necessary fields, as partial overrides can lead to unexpected billing results [4]. Citations:
🏁 Script executed: #!/bin/bash
set -euo pipefail
echo "== locate llm.py and relevant symbols =="
fd -a 'llm\.py$' . | sed 's#^\./##'
echo
rg -n "def _compute_call_cost|_cost_model|cost_model|AWSBedrockLLMParameters|model_id|litellm\.completion_cost|litellm\.cost_per_token" -S .
echo
echo "== inspect llm.py relevant sections =="
LLM_FILES="$(fd 'llm.py$' .)"
for f in $LLM_FILES; do
echo "--- $f ($(wc -l < "$f") lines) ---"
rg -n "_compute_call_cost|_record_usage|_cost_model|cost_model|AWSBedrockLLMParameters|model_id|completion_cost|cost_per_token" -C 4 "$f"
done
echo
echo "== inspect litellm installed versions/docs signatures if available =="
python3 - <<'PY'
import importlib.util
spec = importlib.util.find_spec("litellm")
print("litellm_module_path:", getattr(spec, "origin", None))
if spec:
import litellm
import inspect
print("litellm_version:", getattr(litellm, "__version__", None))
for name in ("completion_cost", "cost_per_token"):
fn = getattr(litellm, name, None)
print(f"litellm.{name}:", fn)
if fn is not None:
try:
print(inspect.signature(fn))
except Exception as e:
print("signature_error:", e)
try:
src = inspect.getsource(fn)
lines = src.splitlines()
for i, line in enumerate(lines[:40], 1):
print(f"{i}:{line}")
except Exception:
pass
PYRepository: Zipstack/unstract Length of output: 11605 🏁 Script executed: #!/bin/bash
set -euo pipefail
echo "== locate llm.py and relevant symbols =="
fd -a 'llm\.py$' . | sed 's#^\./##'
echo
rg -n "def _compute_call_cost|_cost_model|cost_model|AWSBedrockLLMParameters|model_id|litellm\.completion_cost|litellm\.cost_per_token" -S .
echo
echo "== inspect llm.py relevant sections =="
LLM_FILES="$(fd 'llm.py$' .)"
for f in $LLM_FILES; do
echo "--- $f ($(wc -l < "$f") lines) ---"
rg -n "_compute_call_cost|_record_usage|_cost_model|cost_model|AWSBedrockLLMParameters|model_id|completion_cost|cost_per_token" -C 4 "$f"
done
echo
echo "== inspect litellm installed versions/docs signatures if available =="
python3 - <<'PY'
import importlib.util
spec = importlib.util.find_spec("litellm")
print("litellm_module_path:", getattr(spec, "origin", None))
if spec:
import litellm
import inspect
print("litellm_version:", getattr(litellm, "__version__", None))
for name in ("completion_cost", "cost_per_token"):
fn = getattr(litellm, name, None)
print(f"litellm.{name}:", fn)
if fn is not None:
try:
print(inspect.signature(fn))
except Exception as e:
print("signature_error:", e)
try:
src = inspect.getsource(fn)
lines = src.splitlines()
for i, line in enumerate(lines[:80], 1):
print(f"{i}:{line}")
except Exception:
pass
PYRepository: Zipstack/unstract Length of output: 11605 🏁 Script executed: #!/bin/bash
set -euo pipefail
echo "== llm.py cost flow section =="
sed -n '220,260p' unstract/sdk1/src/unstract/sdk1/llm.py
sed -n '420,470p' unstract/sdk1/src/unstract/sdk1/llm.py
sed -n '540,575p' unstract/sdk1/src/unstract/sdk1/llm.py
sed -n '635,655p' unstract/sdk1/src/unstract/sdk1/llm.py
sed -n '700,715p' unstract/sdk1/src/unstract/sdk1/llm.py
sed -n '820,865p' unstract/sdk1/src/unstract/sdk1/llm.py
echo
echo "== base1 cost_model derivation context =="
sed -n '470,510p' unstract/sdk1/src/unstract/sdk1/adapters/base1.py
sed -n '985,1100p' unstract/sdk1/src/unstract/sdk1/adapters/base1.py
sed -n '1525,1545p' unstract/sdk1/src/unstract/sdk1/adapters/base1.py
echo
echo "== prompt caching tests =="
sed -n '1,140p' unstract/sdk1/tests/test_prompt_caching.py
echo
echo "== litellm completion source if available =="
python3 - <<'PY'
import importlib.util
spec = importlib.util.find_spec("litellm")
print("litellm_module_path:", getattr(spec, "origin", None))
if spec:
import litellm
import inspect
print("litellm_version:", getattr(litellm, "__version__", None))
for name in ("completion_cost", "cost_per_token"):
fn = getattr(litellm, name, None)
print(f"litellm.{name}:", fn)
if fn is not None:
try:
print(inspect.signature(fn))
except Exception as e:
print("signature_error:", e)
try:
src = inspect.getsource(fn)
for i, line in enumerate(src.splitlines()[:120], 1):
print(f"{i}:{line}")
print("---")
except Exception as e:
print("source_error:", e)
PYRepository: Zipstack/unstract Length of output: 21711 🌐 Web query:
💡 Result: The cost of an LLM completion request is typically calculated by summing the costs of input (prompt) and output (completion) tokens, often adjusted for factors like prompt caching or specialized modalities (e.g., audio) [1][2][3]. The standard formula for calculating the total cost is: Total Cost = (Input Tokens * Input Cost per Token) + (Output Tokens * Output Cost per Token) [4][3] Technical implementation details often include the following components: 1. Token Counts: The system calculates or extracts the number of prompt tokens and completion tokens generated during the API call [4][5]. 2. Unit Cost Mapping: Costs are determined by referencing a model-specific map that defines the price per unit (e.g., per 1 million tokens or per single token) for input and output [4][6]. 3. Advanced Cost Factors: - Prompt Caching: Costs may be further broken down into distinct rates for cache creation and cache reads [6][3]. - Multi-modal Adjustments: For models processing audio or images, additional costs per audio/image token may be added to the standard input/output calculations [1][2][7]. - Threshold Pricing: Some models apply different rates for tokens once usage exceeds specific thresholds (e.g., beyond 200k tokens) [6][2]. - Alternative Metrics: Some providers or specialized workloads (e.g., certain deployments on AWS Sagemaker or Replicate) may use "cost per second" based on inference runtime rather than token counts [6][5]. In production environments, developers often use middleware or helper utilities to automate these calculations by querying these model-specific price maps based on the provider and model name [4][6][3]. Citations:
🌐 Web query:
💡 Result: The Citations:
🏁 Script executed: #!/bin/bash
set -euo pipefail
if [ -n "$(command -v gh 2>/dev/null)" ]; then
echo "== fetch LiteLLM cost calculator source from main =="
gh api raw https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/cost_calculator.py | sed -n '1,240p'
else
echo "gh unavailable"
fi
echo
echo "== inspect LLM._record_usage =="
sed -n '750,832p' unstract/sdk1/src/unstract/sdk1/llm.py
echo
echo "== inspect cost_per_token implementation from installed package if available by file walk =="
python3 - <<'PY'
import importlib.util, sys, inspect
spec = importlib.util.find_spec("litellm")
if not spec:
print("no literm installed")
sys.exit(0)
print(spec.origin)
from pathlib import Path
base = Path(spec.submodule_search_locations[0]) if spec.submodule_search_locations else Path(spec.origin).parent
print("base", base)
files = list(base.rglob("*.py"))
for f in files[:20]:
print(f)
PYRepository: Zipstack/unstract Length of output: 237 🌐 Web query:
💡 Result: The file litellm/cost_calculator.py within the BerriAI/litellm repository is a core utility module used to calculate the cost of LLM API calls [1]. The mention of completion_response in your query refers to the parameter used by functions in this module—specifically completion_cost—to extract usage information (such as prompt and completion tokens) and the model name to compute the total cost of a response [1]. Key details regarding this file and its usage include: 1. Functionality: The completion_cost function calculates the total cost (in USD) for a given LLM response object [2][1]. If the completion_response is provided, it automatically extracts token counts and the model name to perform the calculation [1]. If not provided, it can calculate costs based on provided model names and text inputs [1]. 2. Dependency: It relies on model pricing information, which LiteLLM typically fetches from a JSON file (model_prices_and_context_window.json) [2][1]. 3. Usage Pattern: Developers typically use it as follows: from litellm import completion_cost cost = completion_cost(completion_response=response_object) 4. Integration: The logic in cost_calculator.py is also used by other LiteLLM components, such as the BudgetManager, to track and update user or model-specific costs [3]. If you are encountering issues with this file, note that it is frequently updated to support new models and providers, and it is standard practice to ensure your LiteLLM version is up to date to maintain accurate cost tracking [4]. Citations:
🌐 Web query:
💡 Result: In the LiteLLM library, the check if completion_response is not None within litellm/cost_calculator.py is used to determine how to proceed with calculating the cost of an LLM API call [1][2]. When completion_response is provided (is not None), the cost calculation logic extracts necessary information—specifically the model name and token usage—directly from the response object [1][2]. This allows LiteLLM to dynamically identify the specific model used for the completion, which is particularly useful in scenarios like dynamic routing (e.g., via Azure Model Router), where the actual model utilized may differ from the initial request model [3]. The code supports both Pydantic BaseModel objects and dictionaries as valid completion_response types, using getattr or dictionary access to retrieve the 'model' field and usage statistics [1]. If completion_response is None, the function falls back to calculating costs based on provided prompt and completion text strings, or other parameters like prompt_tokens and completion_tokens if they were passed separately [2][4][5]. In summary, this conditional check enables the library to handle both post-completion cost tracking (where a response object is available) and pre-completion cost estimation (where only input/output text or token counts are available) [2][5]. Citations:
🌐 Web query:
💡 Result: The completion_cost function in the LiteLLM library is a utility designed to calculate the total cost in USD for a given language model API call [1][2][3]. It is located in the litellm/cost_calculator.py module [1][4]. Key features and usage details include: 1. Functionality: It computes the total cost by combining token counts (for input/prompt and output/completion) with model-specific pricing data [2][3][5]. It utilizes a built-in model cost map [2][5]. 2. Primary Inputs: - completion_response: The response object returned from a LiteLLM completion call [1][3]. - Alternatively, users can provide the model name, the prompt (as a string), and the completion (as a string) to calculate costs manually without a full response object [1][2][3]. 3. Advanced Parameters: - Custom Pricing: Supports overriding default costs using parameters such as custom_cost_per_token or custom_cost_per_second, which is particularly useful for providers like Azure or SageMaker [1][4][6]. - Regional Pricing: Includes parameters like region_name for services with location-based pricing, such as Bedrock [1][4]. - Other: Accepts optional parameters for image generation (size, quality, n), service tiers, and specific routing identifiers [1][4]. 4. Return Value: The function returns a float representing the total USD cost for the specified interaction [1][3]. For implementation examples, users typically pass the response object from a litellm.completion call directly into completion_cost [7][3]. Alternatively, manual calculation can be performed by providing the model, prompt, and completion text directly to the function [2][3][5]. Citations:
🏁 Script executed: #!/bin/bash
set -euo pipefail
URL="https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/cost_calculator.py"
tmp="$(mktemp)"
if command -v curl >/dev/null 2>&1; then
curl -fsSL "$URL" -o "$tmp"
else
python3 - <<'PY' - "$URL" "$tmp"
import urllib.request, sys
urllib.request.urlretrieve(sys.argv[1], sys.argv[2])
PY
fi
echo "== LiteLLM cost calculator snippet =="
sed -n '1,240p' "$tmp"
echo
echo "== model/model_name handling snippets =="
rg -n "model ?=|model_name|completion_response|usage|response_cost|cost_per_token|base_model" "$tmp" -C 3
rm -f "$tmp"
echo
echo "== inspect LLM._record_usage and cost_model caller context =="
sed -n '740,828p' unstract/sdk1/src/unstract/sdk1/llm.pyRepository: Zipstack/unstract Length of output: 50373 🏁 Script executed: #!/bin/bash
set -euo pipefail
URL="https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/cost_calculator.py"
tmp="$(mktemp)"
curl -fsSL "$URL" -o "$tmp"
echo "== LiteLLM completion_price_model selection after model selection =="
sed -n '1100,1240p' "$tmp"
echo
echo "== LiteLLM usage/cost calculation calls for Bedrock route =="
sed -n '1600,1900p' "$tmp"
echo
echo "== LiteLLM generic_cost_per_token signature/body =="
sed -n '290,370p' "$tmp"
rm -f "$tmp"Repository: Zipstack/unstract Length of output: 23019 🌐 Web query:
💡 Result: In LiteLLM, cost calculation for Anthropic models hosted on AWS Bedrock is handled primarily through the cost_per_token function located in litellm/llms/bedrock/cost_calculation.py [1]. This utility calculates costs by utilizing the same logic as standard Anthropic cost calculations, while ensuring the provider is identified as bedrock to facilitate appropriate pricing lookups [1]. To manage or override pricing for these models, LiteLLM provides a custom pricing mechanism [2]. You can override the default cost map by adding a model_info key to your model configuration in the LiteLLM proxy or config file [2]. Example configuration for overriding pricing: model_list: - model_name: "my-bedrock-claude" litellm_params: model: "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" model_info: input_cost_per_token: 0.000003 output_cost_per_token: 0.000015 For providers like Bedrock that support multiple pricing tiers (e.g., PayGo, priority, or flex), LiteLLM automatically selects the appropriate pricing based on the serviceTier returned in the response [2]. Note that there are known edge cases regarding cost calculation: 1. Cross-region inference: Some users have reported issues where region-specific pricing entries (e.g., for us.anthropic.claude-sonnet-4-6) are not correctly matched due to the model name stripping logic, which may default back to the base model's pricing instead [3]. 2. Provider-specific dispatch: In some configurations, if a custom_llm_provider is explicitly set, the cost calculator may dispatch to provider-specific functions (like anthropic_cost_per_token) before checking for custom override values, which can lead to custom zero-cost overrides being ignored [4]. Citations:
Pass the cost override into the response-based cost path.
Proposed fix if has_cache_tokens and response is not None:
try:
- return litellm.completion_cost(completion_response=response)
+ return litellm.completion_cost(completion_response=response, model=model)
except Exception:Add a regression test for prompt-cached calls with 🤖 Prompt for AI Agents |
||
| def _record_usage( | ||
|
Check failure on line 870 in unstract/sdk1/src/unstract/sdk1/llm.py
|
||
| self, | ||
| model: str, | ||
| messages: list[dict[str, str]], | ||
|
|
@@ -733,6 +879,10 @@ | |
| 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 +907,29 @@ | |
| 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%s", | ||
| model, | ||
| llm_api, | ||
| prompt_tokens, | ||
| completion_tokens, | ||
| total_tokens, | ||
| cache_suffix, | ||
| id_suffix, | ||
| ) | ||
|
|
||
| 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 | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The provider-only gate emits Anthropic-specific
cache_controlcontent blocks for every Bedrock model, while the adapter documents support as limited to Anthropic models on Bedrock. Restricting this path to Anthropic Bedrock model identifiers avoids ineffective caching and unsupported message shapes for other model families.Knowledge Base Used: Unstract SDK, Core, and Flags
Prompt To Fix With AI
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!