mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_mcp-oauth-db-backed-temp-servers
This commit is contained in:
commit
19c7fe53a8
14 changed files with 887 additions and 32 deletions
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
#########################################################
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
0
tests/test_litellm/rag/__init__.py
Normal file
0
tests/test_litellm/rag/__init__.py
Normal file
266
tests/test_litellm/rag/test_main.py
Normal file
266
tests/test_litellm/rag/test_main.py
Normal 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
|
||||
7
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
7
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue