Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_mcp-oauth-db-backed-temp-servers
Some checks failed
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled

This commit is contained in:
Devin AI 2026-07-17 17:52:51 +00:00
commit 19c7fe53a8
14 changed files with 887 additions and 32 deletions

View file

@ -315,6 +315,11 @@ disable_token_counter: bool = False
disable_add_transform_inline_image_block: bool = False
disable_add_user_agent_to_request_tags: bool = False
disable_anthropic_gemini_context_caching_transform: bool = False
enable_anthropic_prompt_caching: bool = os.getenv("LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING", "false").lower() == "true"
_anthropic_prompt_caching_ttl_env: Optional[str] = os.getenv("LITELLM_ANTHROPIC_PROMPT_CACHING_TTL")
anthropic_prompt_caching_ttl: Optional[Literal["5m", "1h"]] = (
"1h" if _anthropic_prompt_caching_ttl_env == "1h" else "5m" if _anthropic_prompt_caching_ttl_env == "5m" else None
)
disable_vertex_batch_output_transformation: bool = False
extra_spend_tag_headers: Optional[List[str]] = None
in_memory_llm_clients_cache: "LLMClientCache"

View file

@ -296,18 +296,148 @@ class AnthropicCacheControlHook(CustomPromptManagement):
return processed_messages, processed_system, remaining_points
@staticmethod
def _default_control() -> ChatCompletionCachedContent:
"""Build the cache_control block for auto-injected breakpoints.
Defaults to Anthropic's 5-minute ephemeral cache; honors the optional
``litellm.anthropic_prompt_caching_ttl`` override ("5m" or "1h").
"""
import litellm
ttl = litellm.anthropic_prompt_caching_ttl
if ttl == "5m" or ttl == "1h":
return ChatCompletionCachedContent(type="ephemeral", ttl=ttl)
return ChatCompletionCachedContent(type="ephemeral")
@staticmethod
def _request_has_cache_control(
messages: list[AllMessageValues],
system: str | list | None,
tools: list | None = None,
) -> bool:
"""Return True if the request already carries any client-supplied cache_control.
When the client (e.g. Claude Code) already marks its own breakpoints we
stand down entirely rather than add more, per the auto-caching contract.
Tools count: they are a breakpoint the client can mark, they count toward
the provider's four-block limit, and caching only the tool definitions is
a common pattern, so injecting alongside them can exceed the cap.
"""
if any(AnthropicCacheControlHook._count_cache_control_blocks(msg) for msg in messages):
return True
if isinstance(system, list):
if any(isinstance(block, dict) and block.get("cache_control") is not None for block in system):
return True
if tools is not None:
return any(isinstance(tool, dict) and tool.get("cache_control") is not None for tool in tools)
return False
@staticmethod
def get_default_injection_points(
messages: list[AllMessageValues],
system: str | list | None,
model: str,
custom_llm_provider: str | None,
tools: list | None = None,
) -> list[CacheControlInjectionPoint]:
"""Default breakpoints when ``litellm.enable_anthropic_prompt_caching`` is on.
Caches the system prompt and the trailing turn, so the stable prefix
(system + tools + history) is reused while the breakpoint advances with
the conversation. Returns [] (stand down) when the flag is off, the
provider does not consume cache_control breakpoints (only anthropic /
bedrock do), the model lacks prompt-caching support, or the request
already carries client-supplied cache_control.
"""
import litellm
if litellm.enable_anthropic_prompt_caching is not True:
return []
provider = custom_llm_provider
if provider is None:
from litellm.litellm_core_utils.get_llm_provider_logic import (
get_llm_provider,
)
try:
_, provider, _, _ = get_llm_provider(model=model)
except Exception: # noqa: BLE001 # unroutable model must never block the call, just skip auto-caching
return []
if provider not in ("anthropic", "bedrock"):
return []
from litellm.utils import supports_prompt_caching
if not supports_prompt_caching(model=model, custom_llm_provider=provider):
return []
if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools):
return []
control = AnthropicCacheControlHook._default_control()
points: list[CacheControlInjectionPoint] = [
CacheControlMessageInjectionPoint(location="message", role="system", index=None, control=control),
CacheControlMessageInjectionPoint(location="message", role=None, index=-1, control=control),
]
return points
@staticmethod
def maybe_seed_default_injection_points(
non_default_params: dict[str, Any],
messages: list[AllMessageValues],
model: str,
custom_llm_provider: str | None,
tools: list | None = None,
) -> None:
"""For /chat/completions: add default injection points to the request params.
No-op when injection points are already configured (explicit config wins).
Seeding the param lets the existing prompt-management gate and the
AnthropicCacheControlHook run unchanged.
"""
if non_default_params.get("cache_control_injection_points"):
return
points = AnthropicCacheControlHook.get_default_injection_points(
messages=messages,
system=None,
model=model,
custom_llm_provider=custom_llm_provider,
tools=tools,
)
if points:
non_default_params["cache_control_injection_points"] = points
@staticmethod
def maybe_inject_cache_control(
messages: List[Dict],
system: str | list | None,
kwargs: Dict[str, Any],
model: str | None = None,
custom_llm_provider: str | None = None,
tools: list[dict] | None = None,
) -> Tuple[List[Dict], str | list | None]:
"""Extract cache_control_injection_points from kwargs and apply if present.
When none are configured but ``litellm.enable_anthropic_prompt_caching``
is on, synthesize default breakpoints for the native /v1/messages path.
Pops the key from kwargs; if remaining (non-message) points exist they
are written back so downstream transforms can handle them.
"""
injection_points = kwargs.pop("cache_control_injection_points", None)
configured = cast( # cast-ok: kwargs is untyped; this key only holds the documented injection-point list
list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None)
)
injection_points: list[CacheControlInjectionPoint] = configured or []
if not injection_points and model is not None:
injection_points = AnthropicCacheControlHook.get_default_injection_points(
messages=cast(list[AllMessageValues], messages), # cast-ok: Anthropic-shaped dicts from v1/messages
system=system,
tools=tools,
model=model,
custom_llm_provider=custom_llm_provider,
)
if not injection_points:
return messages, system

View file

@ -1453,6 +1453,9 @@ class Logging(LiteLLMLoggingBaseClass):
response_cost = litellm.response_cost_calculator(**response_cost_calculator_kwargs)
verbose_logger.debug(f"response_cost: {response_cost}")
additional_response_cost: object = self.model_call_details.get("additional_response_cost")
if isinstance(additional_response_cost, (int, float)) and additional_response_cost > 0:
return (response_cost or 0.0) + additional_response_cost
return response_cost
except Exception as e: # error calculating cost
debug_info = StandardLoggingModelCostFailureDebugInformation(

View file

@ -237,7 +237,9 @@ async def anthropic_messages(
AnthropicCacheControlHook,
)
messages, system = AnthropicCacheControlHook.maybe_inject_cache_control(messages, system, kwargs)
messages, system = AnthropicCacheControlHook.maybe_inject_cache_control(
messages, system, kwargs, model=model, custom_llm_provider=custom_llm_provider, tools=tools
)
original_stream = stream or kwargs.get("_websearch_interception_converted_stream", False)
@ -426,7 +428,9 @@ def anthropic_messages_handler(
AnthropicCacheControlHook,
)
messages, system = AnthropicCacheControlHook.maybe_inject_cache_control(messages, system, kwargs)
messages, system = AnthropicCacheControlHook.maybe_inject_cache_control(
messages, system, kwargs, model=model, custom_llm_provider=custom_llm_provider, tools=tools
)
metadata = validate_anthropic_api_metadata(metadata)

View file

@ -510,6 +510,20 @@ async def acompletion(
#########################################################
#########################################################
litellm_logging_obj = kwargs.get("litellm_logging_obj", None)
from litellm.integrations.anthropic_cache_control_hook import (
AnthropicCacheControlHook,
)
from litellm.types.llms.openai import AllMessageValues
AnthropicCacheControlHook.maybe_seed_default_injection_points(
non_default_params=kwargs,
messages=cast(list[AllMessageValues], messages), # cast-ok: acompletion types messages as a bare List
model=model,
custom_llm_provider=cast(Optional[str], custom_llm_provider), # cast-ok: read from untyped kwargs
tools=tools,
)
if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and (
litellm_logging_obj.should_run_prompt_management_hooks(
prompt_id=kwargs.get("prompt_id", None),
@ -5055,6 +5069,19 @@ def completion( # type: ignore
litellm_params = {} # used to prevent unbound var errors
## PROMPT MANAGEMENT HOOKS ##
from litellm.integrations.anthropic_cache_control_hook import (
AnthropicCacheControlHook,
)
from litellm.types.llms.openai import AllMessageValues
AnthropicCacheControlHook.maybe_seed_default_injection_points(
non_default_params=non_default_params,
messages=cast(list[AllMessageValues], messages), # cast-ok: completion types messages as a bare List
model=model,
custom_llm_provider=cast(Optional[str], kwargs.get("custom_llm_provider")), # cast-ok: untyped kwargs
tools=tools,
)
if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and (
litellm_logging_obj.should_run_prompt_management_hooks(
prompt_id=prompt_id, non_default_params=non_default_params

View file

@ -11,14 +11,16 @@ from typing import Any, Dict, Optional, Tuple
import orjson
from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
from fastapi.responses import ORJSONResponse
from fastapi.responses import ORJSONResponse, StreamingResponse
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.proxy._types import *
from litellm.proxy.auth.auth_utils import is_request_body_safe
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
_safe_get_request_headers,
@ -604,6 +606,7 @@ async def rag_query(
general_settings,
llm_router,
proxy_config,
select_data_generator,
version,
)
@ -673,6 +676,31 @@ async def rag_query(
**request_data,
)
hidden_params = getattr(response, "_hidden_params", {}) or {}
custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=hidden_params.get("litellm_call_id", None) or "",
model_id=hidden_params.get("model_id", None) or "",
cache_key=hidden_params.get("cache_key", None) or "",
api_base=hidden_params.get("api_base", None) or "",
version=version,
response_cost=hidden_params.get("response_cost", None),
request_data=request_data,
)
if isinstance(response, CustomStreamWrapper):
return StreamingResponse(
select_data_generator(
response=response,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
request=request,
),
media_type="text/event-stream",
headers=custom_headers,
)
fastapi_response.headers.update(custom_headers)
return response
except HTTPException:

View file

@ -11,12 +11,14 @@ __all__ = ["ingest", "aingest", "query", "aquery"]
import asyncio
import contextvars
from contextlib import contextmanager
from functools import partial
from typing import (
TYPE_CHECKING,
Any,
Coroutine,
Dict,
Iterator,
List,
Optional,
Tuple,
@ -27,6 +29,9 @@ from typing import (
import httpx
import litellm
from litellm._internal_context import is_internal_call
from litellm.cost_calculator import vector_store_search_cost
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion
from litellm.rag.ingestion.bedrock_ingestion import BedrockRAGIngestion
from litellm.rag.ingestion.gemini_ingestion import GeminiRAGIngestion
@ -188,6 +193,25 @@ async def aingest(
)
@contextmanager
def _suppressed_sub_call_billing() -> Iterator[None]:
"""
Suppress a sub-call's own billing event so the parent aquery event bills it.
Every suppressed sub-call's cost must be folded into the parent event:
into the response's hidden response_cost on the non-streaming path, or via
the logging object's additional_response_cost on the streaming path (the
streamed cost is computed from assembled chunks after this pipeline
returns, so there is no response object to fold into here).
"""
previous = is_internal_call.get()
is_internal_call.set(True)
try:
yield
finally:
is_internal_call.set(previous)
async def _execute_query_pipeline(
model: str,
messages: List[Any],
@ -209,27 +233,46 @@ async def _execute_query_pipeline(
raise ValueError("No query found in messages for RAG query")
# 2. Search vector store
search_response = await litellm.vector_stores.asearch(
vector_store_id=retrieval_config["vector_store_id"],
query=query_text,
max_num_results=retrieval_config.get("top_k", 10),
custom_llm_provider=retrieval_config.get("custom_llm_provider", "openai"),
**kwargs,
)
with _suppressed_sub_call_billing():
search_response = await litellm.vector_stores.asearch(
vector_store_id=retrieval_config["vector_store_id"],
query=query_text,
max_num_results=retrieval_config.get("top_k", 10),
custom_llm_provider=retrieval_config.get("custom_llm_provider", "openai"),
**kwargs,
)
search_provider = retrieval_config.get("custom_llm_provider", "openai")
try:
search_cost = sum(
vector_store_search_cost(
model=search_provider if "/" in search_provider else None,
custom_llm_provider=search_provider,
response=search_response,
)
)
except Exception: # noqa: BLE001 - cost accounting must never break the query path
search_cost = 0.0
rerank_response = None
rerank_cost = 0.0
context_chunks = search_response.get("data", [])
# 3. Optional rerank
if rerank and rerank.get("enabled"):
documents = RAGQuery.extract_documents_from_search(search_response)
if documents:
rerank_response = await litellm.arerank(
model=rerank["model"],
query=query_text,
documents=documents,
top_n=rerank.get("top_n", 5),
)
with _suppressed_sub_call_billing():
rerank_response = await litellm.arerank(
model=rerank["model"],
query=query_text,
documents=documents,
top_n=rerank.get("top_n", 5),
)
rerank_hidden_params = getattr(rerank_response, "_hidden_params", None)
if isinstance(rerank_hidden_params, dict):
rerank_response_cost: float | None = rerank_hidden_params.get("response_cost")
rerank_cost = rerank_response_cost or 0.0
context_chunks = RAGQuery.get_top_chunks_from_rerank(search_response, rerank_response)
# 4. Build context message and call completion
@ -237,28 +280,40 @@ async def _execute_query_pipeline(
modified_messages = messages[:-1] + [context_message] + [messages[-1]]
# Use router if available to properly resolve virtual model names
if router is not None:
response = await router.acompletion(
model=model,
messages=modified_messages,
stream=stream,
**kwargs,
)
else:
response = await litellm.acompletion(
model=model,
messages=modified_messages,
stream=stream,
**kwargs,
)
with _suppressed_sub_call_billing():
if router is not None:
response = await router.acompletion(
model=model,
messages=modified_messages,
stream=stream,
**kwargs,
)
else:
response = await litellm.acompletion(
model=model,
messages=modified_messages,
stream=stream,
**kwargs,
)
# 5. Attach search results to response
sub_call_cost = search_cost + rerank_cost
if not stream and isinstance(response, ModelResponse):
response = RAGQuery.add_search_results_to_response(
response=response,
search_results=search_response,
rerank_results=rerank_response,
)
if sub_call_cost > 0:
hidden_params = getattr(response, "_hidden_params", None)
if isinstance(hidden_params, dict):
completion_response_cost: float | None = hidden_params.get("response_cost")
if completion_response_cost is not None:
hidden_params["response_cost"] = completion_response_cost + sub_call_cost
elif sub_call_cost > 0:
logging_obj: object = kwargs.get("litellm_logging_obj")
if isinstance(logging_obj, LiteLLMLoggingObj):
logging_obj.model_call_details["additional_response_cost"] = sub_call_cost
return response # type: ignore[return-value]

View file

@ -529,6 +529,7 @@ class ChatCompletionDeltaToolCallChunk(TypedDict, total=False):
class ChatCompletionCachedContent(TypedDict):
type: Literal["ephemeral"]
ttl: NotRequired[Literal["5m", "1h"]]
class ChatCompletionThinkingBlock(TypedDict, total=False):

View file

@ -403,6 +403,11 @@ class CallTypes(str, Enum):
vector_store_search = "vector_store_search"
avector_store_search = "avector_store_search"
ingest = "ingest"
aingest = "aingest"
query = "query"
aquery = "aquery"
#########################################################
# Container Call Types
#########################################################

View file

@ -2,7 +2,9 @@ import copy
import datetime
import json
import os
import subprocess
import sys
import textwrap
import unittest
from typing import List, Optional, Tuple
from unittest.mock import ANY, MagicMock, Mock, patch
@ -1533,3 +1535,242 @@ class TestApplyToAnthropicMessagesRequest:
sys_blocks = sum(1 for b in (result_sys or []) if isinstance(b, dict) and b.get("cache_control") is not None)
total_blocks = sys_blocks + sum(AnthropicCacheControlHook._count_cache_control_blocks(m) for m in result_msgs)
assert total_blocks <= 4
class TestEnableAnthropicPromptCaching:
"""Auto-injected default breakpoints via litellm.enable_anthropic_prompt_caching."""
MESSAGES: List[AllMessageValues] = [
{"role": "system", "content": "a long system prompt"},
{"role": "user", "content": "first turn"},
{"role": "assistant", "content": "a reply"},
{"role": "user", "content": "latest turn"},
]
def _points(self, model="claude-sonnet-4-5", provider="anthropic", messages=None, system=None, tools=None):
return AnthropicCacheControlHook.get_default_injection_points(
messages=copy.deepcopy(self.MESSAGES) if messages is None else messages,
system=system,
model=model,
custom_llm_provider=provider,
tools=tools,
)
def test_disabled_by_default(self):
assert litellm.enable_anthropic_prompt_caching is False
assert self._points() == []
def test_injects_system_and_trailing_turn(self, monkeypatch):
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
assert self._points() == [
{"location": "message", "role": "system", "index": None, "control": {"type": "ephemeral"}},
{"location": "message", "role": None, "index": -1, "control": {"type": "ephemeral"}},
]
def test_bedrock_claude_is_injected(self, monkeypatch):
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
points = self._points(model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", provider="bedrock")
assert [p["index"] for p in points] == [None, -1]
@pytest.mark.parametrize("model, provider", [("gpt-4o", "openai"), ("gemini-2.0-flash", "gemini")])
def test_non_anthropic_providers_never_injected(self, monkeypatch, model, provider):
"""These report supports_prompt_caching=True but never consume cache_control markers."""
from litellm.utils import supports_prompt_caching
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
assert supports_prompt_caching(model=model, custom_llm_provider=provider) is True
assert self._points(model=model, provider=provider) == []
def test_model_without_caching_support_not_injected(self, monkeypatch):
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
assert self._points(model="anthropic.claude-3-5-sonnet-20240620-v1:0", provider="bedrock") == []
def test_stands_down_when_client_sent_cache_control(self, monkeypatch):
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
messages = [
{"role": "system", "content": [{"type": "text", "text": "s", "cache_control": {"type": "ephemeral"}}]},
{"role": "user", "content": "latest turn"},
]
assert self._points(messages=messages) == []
def test_stands_down_when_system_block_has_cache_control(self, monkeypatch):
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
system = [{"type": "text", "text": "s", "cache_control": {"type": "ephemeral"}}]
assert self._points(messages=[{"role": "user", "content": "hi"}], system=system) == []
@staticmethod
def _tools(count: int, cached: bool) -> List[dict]:
tool: dict = {"type": "function", "function": {"name": "t", "description": "d", "parameters": {}}}
if cached:
tool["cache_control"] = {"type": "ephemeral"}
return [{**tool, "function": {**tool["function"], "name": f"t{i}"}} for i in range(count)]
def test_stands_down_when_only_tools_carry_cache_control(self, monkeypatch):
"""Caching just the tool definitions is a normal client pattern, and those
breakpoints count toward the provider's four-block limit. Three of them plus
our two would be five, which Anthropic rejects outright."""
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
assert self._points(tools=self._tools(3, cached=True)) == []
def test_injects_when_tools_carry_no_cache_control(self, monkeypatch):
"""Tools alone must not suppress injection; only client-marked ones do."""
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
assert [p["index"] for p in self._points(tools=self._tools(3, cached=False))] == [None, -1]
@pytest.mark.parametrize("tools", [None, []])
def test_absent_tools_do_not_suppress_injection(self, monkeypatch, tools):
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
assert [p["index"] for p in self._points(tools=tools)] == [None, -1]
def test_seed_stands_down_when_only_tools_carry_cache_control(self, monkeypatch):
"""Same guard on the /chat/completions seeding path."""
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
params: dict = {}
AnthropicCacheControlHook.maybe_seed_default_injection_points(
non_default_params=params,
messages=copy.deepcopy(self.MESSAGES),
model="claude-sonnet-4-5",
custom_llm_provider="anthropic",
tools=self._tools(3, cached=True),
)
assert "cache_control_injection_points" not in params
def test_v1_messages_stands_down_when_only_tools_carry_cache_control(self, monkeypatch):
"""Same guard on the /v1/messages path, where tools reach the hook directly."""
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
messages = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]
result_msgs, result_sys = AnthropicCacheControlHook.maybe_inject_cache_control(
copy.deepcopy(messages),
"sys",
{},
model="claude-sonnet-4-5",
custom_llm_provider="anthropic",
tools=self._tools(3, cached=True),
)
assert result_sys == "sys"
assert result_msgs == messages
def test_default_ttl_is_anthropics_five_minute_cache(self, monkeypatch):
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
assert all(p["control"] == {"type": "ephemeral"} for p in self._points())
@pytest.mark.parametrize("ttl", ["5m", "1h"])
def test_ttl_override_applied(self, monkeypatch, ttl):
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
monkeypatch.setattr(litellm, "anthropic_prompt_caching_ttl", ttl)
assert all(p["control"] == {"type": "ephemeral", "ttl": ttl} for p in self._points())
def test_seed_does_not_override_configured_points(self, monkeypatch):
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
configured = [{"location": "message", "role": "user", "index": 0}]
params = {"cache_control_injection_points": configured}
AnthropicCacheControlHook.maybe_seed_default_injection_points(
non_default_params=params,
messages=copy.deepcopy(self.MESSAGES),
model="claude-sonnet-4-5",
custom_llm_provider="anthropic",
)
assert params["cache_control_injection_points"] is configured
def test_seed_adds_defaults_when_enabled(self, monkeypatch):
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
params: dict = {}
AnthropicCacheControlHook.maybe_seed_default_injection_points(
non_default_params=params,
messages=copy.deepcopy(self.MESSAGES),
model="claude-sonnet-4-5",
custom_llm_provider="anthropic",
)
assert [p["index"] for p in params["cache_control_injection_points"]] == [None, -1]
def test_seed_is_noop_when_disabled(self):
params: dict = {}
AnthropicCacheControlHook.maybe_seed_default_injection_points(
non_default_params=params,
messages=copy.deepcopy(self.MESSAGES),
model="claude-sonnet-4-5",
custom_llm_provider="anthropic",
)
assert params == {}
def test_v1_messages_applies_defaults_end_to_end(self, monkeypatch):
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
messages = [
{"role": "user", "content": [{"type": "text", "text": "first"}]},
{"role": "assistant", "content": [{"type": "text", "text": "reply"}]},
{"role": "user", "content": [{"type": "text", "text": "latest"}]},
]
result_msgs, result_sys = AnthropicCacheControlHook.maybe_inject_cache_control(
messages,
"a system prompt",
{},
model="claude-sonnet-4-5",
custom_llm_provider="anthropic",
)
assert result_sys == [{"type": "text", "text": "a system prompt", "cache_control": {"type": "ephemeral"}}]
assert result_msgs[-1]["content"][-1]["cache_control"] == {"type": "ephemeral"}
assert "cache_control" not in result_msgs[0]["content"][-1]
def test_v1_messages_is_noop_when_disabled(self):
messages = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]
result_msgs, result_sys = AnthropicCacheControlHook.maybe_inject_cache_control(
messages,
"sys",
{},
model="claude-sonnet-4-5",
custom_llm_provider="anthropic",
)
assert result_sys == "sys"
assert result_msgs == messages
class TestAnthropicPromptCachingEnvVars:
"""Both settings are read from the environment at import, so an admin can enable
auto-caching without a config file. Each case re-imports litellm in a subprocess
so the env is read fresh without contaminating this process's module graph.
"""
@staticmethod
def _import_litellm_with_env(env_override: dict) -> Tuple[bool, Optional[str]]:
env = os.environ.copy()
env.pop("LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING", None)
env.pop("LITELLM_ANTHROPIC_PROMPT_CACHING_TTL", None)
env.update(env_override)
script = textwrap.dedent(
"""
import json, litellm
print(json.dumps([litellm.enable_anthropic_prompt_caching, litellm.anthropic_prompt_caching_ttl]))
"""
)
result = subprocess.run(
[sys.executable, "-c", script], capture_output=True, text=True, env=env, timeout=300
)
assert result.returncode == 0, result.stderr
enabled, ttl = json.loads(result.stdout.strip().splitlines()[-1])
return enabled, ttl
def test_unset_env_leaves_auto_caching_off(self):
assert self._import_litellm_with_env({}) == (False, None)
@pytest.mark.parametrize("value", ["true", "True", "TRUE"])
def test_env_enables_auto_caching_case_insensitively(self, value):
enabled, _ = self._import_litellm_with_env({"LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING": value})
assert enabled is True
@pytest.mark.parametrize("value", ["false", "0", "yes", ""])
def test_env_only_enables_on_true(self, value):
enabled, _ = self._import_litellm_with_env({"LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING": value})
assert enabled is False
@pytest.mark.parametrize("value", ["5m", "1h"])
def test_ttl_env_is_applied(self, value):
_, ttl = self._import_litellm_with_env({"LITELLM_ANTHROPIC_PROMPT_CACHING_TTL": value})
assert ttl == value
@pytest.mark.parametrize("value", ["10m", "1H", "3600", "ephemeral"])
def test_unsupported_ttl_env_falls_back_to_provider_default(self, value):
"""An unparseable TTL must fall back to Anthropic's 5m default, never reach the provider verbatim."""
_, ttl = self._import_litellm_with_env({"LITELLM_ANTHROPIC_PROMPT_CACHING_TTL": value})
assert ttl is None

View file

@ -242,3 +242,88 @@ class TestRagIngestSSRFBlocked:
assert response.status_code != 400, (
f"Clean Bedrock ingest_options should not be rejected: {response.json()}"
)
def test_rag_query_returns_response_cost_header(client_internal_user):
"""
/v1/rag/query must surface the completion cost via the
x-litellm-response-cost response header, like /v1/chat/completions does.
"""
from litellm.types.utils import ModelResponse
mock_response = ModelResponse(
id="chatcmpl-test",
choices=[
{
"index": 0,
"message": {"role": "assistant", "content": "The codename is AZURE-FALCON-42."},
"finish_reason": "stop",
}
],
model="gpt-4o-mini",
usage={"prompt_tokens": 35, "completion_tokens": 14, "total_tokens": 49},
)
mock_response._hidden_params["response_cost"] = 3.45e-06
with patch(
"litellm.proxy.rag_endpoints.endpoints.litellm.aquery",
new_callable=AsyncMock,
return_value=mock_response,
), patch("litellm.vector_store_registry", None), patch(
"litellm.proxy.proxy_server.prisma_client", None
):
response = client_internal_user.post(
"/v1/rag/query",
json={
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "What is the codename?"}],
"retrieval_config": {
"vector_store_id": "vs_test_123",
"custom_llm_provider": "openai",
},
},
)
assert response.status_code == 200, response.json()
assert response.headers.get("x-litellm-response-cost") == "3.45e-06"
def test_rag_query_stream_returns_event_stream(client_internal_user):
"""
A stream=true /v1/rag/query must return an SSE response. Returning the raw
stream wrapper makes FastAPI try to serialize it, which raises and turns
every streaming RAG query into a 500; the stream then never drains, so its
single billing event (which carries the folded sub-call costs) never fires.
"""
import litellm as litellm_module
async def fake_aquery(**kwargs):
return await litellm_module.acompletion(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "What is the codename?"}],
mock_response="The codename is AZURE-FALCON-42.",
stream=True,
api_key="test-key",
)
with patch(
"litellm.proxy.rag_endpoints.endpoints.litellm.aquery",
new=AsyncMock(side_effect=fake_aquery),
), patch("litellm.vector_store_registry", None), patch("litellm.proxy.proxy_server.prisma_client", None):
response = client_internal_user.post(
"/v1/rag/query",
json={
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "What is the codename?"}],
"retrieval_config": {
"vector_store_id": "vs_test_123",
"custom_llm_provider": "openai",
},
"stream": True,
},
)
assert response.status_code == 200, response.text
assert response.headers.get("content-type", "").startswith("text/event-stream")
assert '"object":"chat.completion.chunk"' in response.text
assert "data: [DONE]" in response.text

View file

View file

@ -0,0 +1,266 @@
"""
Tests for the RAG query pipeline in litellm/rag/main.py.
The RAG pipeline forwards its kwargs (including the parent litellm_logging_obj)
into @client-decorated sub-calls (vector store search, completion). Each logging
object allows exactly one async_success event, so if sub-calls are not marked as
internal, the vector store search consumes the slot first and the LLM
completion's usage/cost is never logged (spend tracking and budget enforcement
are bypassed). These tests pin the invariant that the single billing event for
aquery carries the completion response with real usage and cost.
"""
import asyncio
from unittest.mock import patch
import pytest
import litellm
from litellm._internal_context import is_internal_call
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import CallTypes, ModelResponse
class RecordingLogger(CustomLogger):
def __init__(self):
super().__init__()
self.success_events = []
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.success_events.append({"kwargs": kwargs, "response_obj": response_obj})
@pytest.mark.asyncio
@pytest.mark.parametrize("use_router", [False, True])
async def test_aquery_single_billing_event_carries_completion_usage_and_cost(use_router):
"""
litellm.aquery must produce exactly one success event, and that event must
carry the LLM completion (a ModelResponse with non-zero usage and cost),
not the vector store search response. The proxy always passes a router, so
both the router and non-router completion branches are pinned.
"""
recording_logger = RecordingLogger()
original_callbacks = litellm.callbacks
litellm.callbacks = [recording_logger]
router_kwargs = {}
if use_router:
router_kwargs["router"] = litellm.Router(
model_list=[
{
"model_name": "gpt-4o-mini",
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"},
}
]
)
try:
response = await litellm.aquery(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "What is the secret project codename?"}],
retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
mock_response="The secret project codename is AZURE-FALCON-42.",
**router_kwargs,
)
assert isinstance(response, ModelResponse)
assert is_internal_call.get() is False
for _ in range(50):
if recording_logger.success_events:
break
await asyncio.sleep(0.1)
await asyncio.sleep(0.5)
finally:
litellm.callbacks = original_callbacks
assert len(recording_logger.success_events) == 1
event = recording_logger.success_events[0]
response_obj = event["response_obj"]
assert isinstance(response_obj, ModelResponse)
assert response_obj.usage.total_tokens > 0
standard_logging_object = event["kwargs"]["standard_logging_object"]
assert standard_logging_object["call_type"] == "aquery"
assert standard_logging_object["total_tokens"] > 0
assert standard_logging_object["prompt_tokens"] > 0
assert standard_logging_object["completion_tokens"] > 0
assert standard_logging_object["response_cost"] > 0
@pytest.mark.asyncio
async def test_aquery_response_hidden_params_carry_completion_cost():
"""
The aquery response must expose the completion's response_cost via hidden
params, so the proxy can return the x-litellm-response-cost header.
"""
response = await litellm.aquery(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hello"}],
retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
mock_response="hi there",
)
assert isinstance(response, ModelResponse)
response_cost = response._hidden_params.get("response_cost")
assert response_cost is not None
assert response_cost > 0
@pytest.mark.asyncio
async def test_aquery_billed_cost_includes_priced_vector_store_search():
"""
When the vector store provider prices search calls (e.g. per-query cost),
that cost must be folded into the aquery billing instead of being dropped
with the suppressed sub-call event.
"""
recording_logger = RecordingLogger()
original_callbacks = litellm.callbacks
litellm.callbacks = [recording_logger]
try:
with patch("litellm.rag.main.vector_store_search_cost", return_value=(0.002, 0.0)):
response = await litellm.aquery(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hello"}],
retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
mock_response="hi there",
)
for _ in range(50):
if recording_logger.success_events:
break
await asyncio.sleep(0.1)
await asyncio.sleep(0.5)
finally:
litellm.callbacks = original_callbacks
assert isinstance(response, ModelResponse)
total_cost = response._hidden_params.get("response_cost")
assert total_cost is not None
assert total_cost > 0.002
assert len(recording_logger.success_events) == 1
standard_logging_object = recording_logger.success_events[0]["kwargs"]["standard_logging_object"]
assert standard_logging_object["response_cost"] == total_cost
@pytest.mark.asyncio
async def test_aquery_with_rerank_bills_once_and_folds_rerank_cost():
"""
When rerank is enabled, its sub-call must run under the internal-call
context (no standalone billing event) and its cost must be folded into
the single aquery billing event.
"""
from litellm.types.rerank import RerankResponse
recording_logger = RecordingLogger()
original_callbacks = litellm.callbacks
litellm.callbacks = [recording_logger]
rerank_seen = {}
async def fake_arerank(**kwargs):
rerank_seen["internal"] = is_internal_call.get()
rerank_result = RerankResponse(id="rr_1", results=[{"index": 0, "relevance_score": 0.9}], meta={})
rerank_result._hidden_params["response_cost"] = 0.001
return rerank_result
try:
with patch("litellm.arerank", side_effect=fake_arerank):
response = await litellm.aquery(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hello"}],
retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
rerank={"enabled": True, "model": "cohere/rerank-english-v3.0", "top_n": 1},
mock_response="hi there",
)
for _ in range(50):
if recording_logger.success_events:
break
await asyncio.sleep(0.1)
await asyncio.sleep(0.5)
finally:
litellm.callbacks = original_callbacks
assert rerank_seen["internal"] is True
assert is_internal_call.get() is False
assert isinstance(response, ModelResponse)
total_cost = response._hidden_params.get("response_cost")
assert total_cost is not None
assert total_cost > 0.001
assert len(recording_logger.success_events) == 1
standard_logging_object = recording_logger.success_events[0]["kwargs"]["standard_logging_object"]
assert standard_logging_object["call_type"] == "aquery"
assert standard_logging_object["response_cost"] == total_cost
@pytest.mark.asyncio
async def test_aquery_streaming_bills_sub_call_costs_into_final_event():
"""
On the streaming path the response cost is computed from the assembled
chunks after the pipeline returns, so there is no response object to fold
sub-call costs into. The pipeline must instead carry the accumulated
search and rerank cost through the logging object so the single streamed
billing event includes it; otherwise a caller passing stream=true incurs
priced vector search and rerank costs that never reach spend tracking.
"""
from litellm.types.rerank import RerankResponse
recording_logger = RecordingLogger()
original_callbacks = litellm.callbacks
litellm.callbacks = [recording_logger]
rerank_seen = {}
async def fake_arerank(**kwargs):
rerank_seen["internal"] = is_internal_call.get()
rerank_result = RerankResponse(id="rr_1", results=[{"index": 0, "relevance_score": 0.9}], meta={})
rerank_result._hidden_params["response_cost"] = 0.001
return rerank_result
try:
with (
patch("litellm.rag.main.vector_store_search_cost", return_value=(0.002, 0.0)),
patch("litellm.arerank", side_effect=fake_arerank),
):
response = await litellm.aquery(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hello"}],
retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"},
rerank={"enabled": True, "model": "cohere/rerank-english-v3.0", "top_n": 1},
mock_response="hi there",
stream=True,
)
async for _ in response:
pass
for _ in range(50):
if recording_logger.success_events:
break
await asyncio.sleep(0.1)
await asyncio.sleep(0.5)
finally:
litellm.callbacks = original_callbacks
assert rerank_seen["internal"] is True
assert is_internal_call.get() is False
assert len(recording_logger.success_events) == 1
standard_logging_object = recording_logger.success_events[0]["kwargs"]["standard_logging_object"]
assert standard_logging_object["call_type"] == "aquery"
assert standard_logging_object["response_cost"] >= 0.003
def test_rag_call_types_are_registered():
"""
query/aquery/ingest/aingest are @client-decorated entry points, so their
function names must resolve to CallTypes members (deployment hooks and
call-type driven logic silently no-op for unregistered call types).
"""
assert CallTypes("query") is CallTypes.query
assert CallTypes("aquery") is CallTypes.aquery
assert CallTypes("ingest") is CallTypes.ingest
assert CallTypes("aingest") is CallTypes.aingest

View file

@ -21690,7 +21690,7 @@ export interface components {
* CallTypes
* @enum {string}
*/
CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill";
CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "ingest" | "aingest" | "query" | "aquery" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill";
/** CallbackDelete */
CallbackDelete: {
/** Callback Name */
@ -21870,6 +21870,11 @@ export interface components {
};
/** ChatCompletionCachedContent */
ChatCompletionCachedContent: {
/**
* Ttl
* @enum {string}
*/
ttl?: "5m" | "1h";
/**
* Type
* @constant