scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

Merge remote-tracking branch 'origin/main' into litellm_/circleci-specific-sha-0cf414

Browse source
This commit is contained in:
Yuneng Jiang 2026-09-17 18:05:39 -07:00
commit e8c1fe884e
No known key found for this signature in database
95 changed files with 165 additions and 7495 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
.github/workflows/test-unit-proxy-db.yml vendored
View file

@ -94,7 +94,6 @@ jobs:
tests/proxy_unit_tests/test_jwt_key_mapping.py
tests/proxy_unit_tests/test_proxy_custom_auth.py
tests/proxy_unit_tests/test_key_generate_dynamodb.py
tests/proxy_unit_tests/test_deployed_proxy_keygen.py
workers: 4
dist: loadscope
timeout: 15
@ -110,8 +109,6 @@ jobs:
- test-group: proxy-server-core
test-path: >-
tests/proxy_unit_tests/test_proxy_server.py
tests/proxy_unit_tests/test_proxy_server_keys.py
tests/proxy_unit_tests/test_proxy_server_spend.py
tests/proxy_unit_tests/test_aproxy_startup.py
workers: 4
dist: loadscope
@ -120,7 +117,6 @@ jobs:
test-path: >-
tests/proxy_unit_tests/test_proxy_config_unit_test.py
tests/proxy_unit_tests/test_proxy_routes.py
tests/proxy_unit_tests/test_proxy_gunicorn.py
tests/proxy_unit_tests/test_server_root_path.py
tests/proxy_unit_tests/test_proxy_pass_user_config.py
tests/proxy_unit_tests/test_proxy_token_counter.py
@ -198,7 +194,6 @@ jobs:
tests/proxy_unit_tests/test_realtime_cache.py
tests/proxy_unit_tests/test_proxy_exception_mapping.py
tests/proxy_unit_tests/test_custom_tokenizer_bug.py
tests/proxy_unit_tests/test_model_response_typing
workers: 4
dist: loadscope
timeout: 15

95
litellm/integrations/langfuse/langfuse.py
View file

@ -394,35 +394,20 @@ class LangFuseLogger:
status_message=status_message,
)
verbose_logger.debug("OUTPUT IN LANGFUSE: %s; original: %s", output, response_obj)
trace_id = None
generation_id = None
if self._is_langfuse_v2():
trace_id, generation_id = self._log_langfuse_v2(
user_id=user_id,
metadata=metadata,
litellm_params=litellm_params,
output=output,
start_time=start_time,
end_time=end_time,
kwargs=kwargs,
optional_params=optional_params,
input=input,
response_obj=response_obj,
level=level,
litellm_call_id=litellm_call_id,
)
elif response_obj is not None:
self._log_langfuse_v1(
user_id=user_id,
metadata=metadata,
output=output,
start_time=start_time,
end_time=end_time,
kwargs=kwargs,
optional_params=optional_params,
input=input,
response_obj=response_obj,
)
trace_id, generation_id = self._log_langfuse_v2(
user_id=user_id,
metadata=metadata,
litellm_params=litellm_params,
output=output,
start_time=start_time,
end_time=end_time,
kwargs=kwargs,
optional_params=optional_params,
input=input,
response_obj=response_obj,
level=level,
litellm_call_id=litellm_call_id,
)
verbose_logger.debug("Langfuse Layer Logging - final response object: %s", response_obj)
verbose_logger.info("Langfuse Layer Logging - logging success")
@ -518,58 +503,6 @@ class LangFuseLogger:
This approach does not impact latency and runs in the background
"""
def _is_langfuse_v2(self):
import langfuse
return Version(langfuse.version.__version__) >= Version("2.0.0")
def _log_langfuse_v1(
self,
user_id,
metadata,
output,
start_time,
end_time,
kwargs,
optional_params,
input,
response_obj,
):
from langfuse.model import CreateGeneration, CreateTrace
verbose_logger.warning(
"Please upgrade langfuse to v2.0.0 or higher: https://github.com/langfuse/langfuse-python/releases/tag/v2.0.1"
)
trace: Final = self.Langfuse.trace(
CreateTrace(
name=metadata.get("generation_name", "litellm-completion"),
input=input,
output=output,
userId=user_id,
)
)
custom_llm_provider: Final = cast(str | None, kwargs.get("custom_llm_provider"))
model_name: Final = reconstruct_model_name(kwargs.get("model", ""), custom_llm_provider, metadata)
trace.generation(
CreateGeneration(
name=metadata.get("generation_name", "litellm-completion"),
startTime=start_time,
endTime=end_time,
model=model_name,
modelParameters=optional_params,
prompt=input,
completion=output,
usage={
"prompt_tokens": response_obj.usage.prompt_tokens,
"completion_tokens": response_obj.usage.completion_tokens,
},
metadata=metadata,
)
)
def _log_langfuse_v2(
self,
user_id: str | None,

115
litellm/integrations/prometheus.py
View file

@ -995,23 +995,6 @@ class PrometheusLogger(CustomLogger):
return label_filters
def _validate_configured_metric_labels(self, metric_name: str, labels: list[str]):
"""
Ensure that all the configured labels are valid for the metric
Raises ValueError if the metric labels are invalid and pretty prints the error
"""
label_error: Final = self._validate_single_metric_labels(metric_name, labels)
if label_error:
self._pretty_print_invalid_labels_error(
metric_name=label_error.metric_name,
invalid_labels=label_error.invalid_labels,
valid_labels=label_error.valid_labels,
)
raise ValueError(label_error.message)
return True
#########################################################
# Pretty print functions
#########################################################
@ -1090,108 +1073,10 @@ class PrometheusLogger(CustomLogger):
for label_error in validation_results.label_errors:
verbose_logger.error(label_error.message)
def _pretty_print_invalid_labels_error(
self, metric_name: str, invalid_labels: list[str], valid_labels: list[str]
) -> None:
"""Pretty print error message for invalid labels using rich"""
try:
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.text import Text
console: Final = Console()
# Create error panel title
title: Final = Text(
f"🚨🚨 Invalid Labels for Metric: '{metric_name}'\nInvalid labels: {', '.join(invalid_labels)}\nPlease specify only valid labels below",
style="bold red",
)
# Create valid labels table
labels_table: Final = Table(
title="🏷️ Valid Labels for this Metric",
show_header=True,
header_style="bold green",
title_justify="left",
border_style="green",
)
labels_table.add_column("Valid Labels", style="cyan", no_wrap=True)
for label in sorted(valid_labels):
labels_table.add_row(label)
# Print everything in a nice panel
console.print("\n")
console.print(Panel(title, border_style="red"))
console.print(labels_table)
console.print("\n")
except ImportError:
# Fallback to simple logging if rich is not available
verbose_logger.error(
"Invalid labels for metric '%s': %s. Valid labels: %s",
metric_name,
invalid_labels,
sorted(valid_labels),
)
def _pretty_print_invalid_metric_error(self, invalid_metric_name: str, valid_metrics: tuple) -> None:
"""Pretty print error message for invalid metric name using rich"""
try:
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.text import Text
console: Final = Console()
# Create error panel title
title: Final = Text(
f"🚨🚨 Invalid Metric Name: '{invalid_metric_name}'\nPlease specify one of the allowed metrics below",
style="bold red",
)
# Create valid metrics table
metrics_table: Final = Table(
title="📊 Valid Metric Names",
show_header=True,
header_style="bold green",
title_justify="left",
border_style="green",
)
metrics_table.add_column("Available Metrics", style="cyan", no_wrap=True)
for metric in sorted(valid_metrics):
metrics_table.add_row(metric)
# Print everything in a nice panel
console.print("\n")
console.print(Panel(title, border_style="red"))
console.print(metrics_table)
console.print("\n")
except ImportError:
# Fallback to simple logging if rich is not available
verbose_logger.error(
"Invalid metric name: %s. Valid metrics: %s", invalid_metric_name, sorted(valid_metrics)
)
#########################################################
# End of pretty print functions
#########################################################
def _valid_metric_name(self, metric_name: str):
"""
Raises ValueError if the metric name is invalid and pretty prints the error
"""
error: Final = self._validate_single_metric_name(metric_name)
if error:
self._pretty_print_invalid_metric_error(
invalid_metric_name=error.metric_name, valid_metrics=error.valid_metrics
)
raise ValueError(error.message)
def _pretty_print_prometheus_config(self, label_filters: dict[str, list[str]]) -> None:
"""Pretty print the processed prometheus configuration using rich"""
try:

154
litellm/litellm_core_utils/streaming_handler.py
View file

@ -113,14 +113,6 @@ class _PredibaseStreamData(TypedDict):
error: str | None
class _Ai21StreamData(TypedDict):
completions: Sequence[Mapping[str, Mapping[str, str]]]
class _MaritalkStreamData(TypedDict):
answer: str
class _NlpCloudStreamData(TypedDict):
generated_text: str
@ -129,25 +121,6 @@ class _AlephAlphaStreamData(TypedDict):
completions: Sequence[Mapping[str, str]]
class _AzureStreamChoice(TypedDict):
delta: Mapping[str, str] | None
finish_reason: str | None
class _AzureStreamData(TypedDict):
choices: Sequence[_AzureStreamChoice]
class _BasetenModelOutput(TypedDict):
data: NotRequired[Sequence[str]]
class _BasetenStreamData(TypedDict):
token: NotRequired[Mapping[str, str]]
model_output: NotRequired["_BasetenModelOutput | str"]
completion: NotRequired[object]
class _DeltaDumpDict(TypedDict):
role: NotRequired[str | None]
tool_calls: NotRequired[Sequence[Mapping[str, object]]]
@ -572,36 +545,6 @@ class CustomStreamWrapper:
except Exception as e:
raise e
def handle_ai21_chunk(self, chunk): # fake streaming
chunk = chunk.decode("utf-8")
data_json: Final[_Ai21StreamData] = json.loads(chunk)
try:
text: Final = data_json["completions"][0]["data"]["text"]
is_finished: Final = True
finish_reason: Final = "stop"
return {
"text": text,
"is_finished": is_finished,
"finish_reason": finish_reason,
}
except Exception:
raise ValueError(f"Unable to parse response. Original response: {chunk}")
def handle_maritalk_chunk(self, chunk): # fake streaming
chunk = chunk.decode("utf-8")
data_json: Final[_MaritalkStreamData] = json.loads(chunk)
try:
text: Final = data_json["answer"]
is_finished: Final = True
finish_reason: Final = "stop"
return {
"text": text,
"is_finished": is_finished,
"finish_reason": finish_reason,
}
except Exception:
raise ValueError(f"Unable to parse response. Original response: {chunk}")
def handle_nlp_cloud_chunk(self, chunk):
text = ""
is_finished = False
@ -640,46 +583,6 @@ class CustomStreamWrapper:
except Exception:
raise ValueError(f"Unable to parse response. Original response: {chunk}")
def handle_azure_chunk(self, chunk):
is_finished = False
finish_reason = ""
text = ""
print_verbose(f"chunk: {chunk}")
if "data: [DONE]" in chunk:
text = ""
is_finished = True
finish_reason = "stop"
return {
"text": text,
"is_finished": is_finished,
"finish_reason": finish_reason,
}
elif chunk.startswith("data:"):
data_json: Final[_AzureStreamData] = json.loads(chunk[5:]) # chunk.startswith("data:"):
try:
if len(data_json["choices"]) > 0:
delta: Final = data_json["choices"][0]["delta"]
text = "" if delta is None else delta.get("content", "")
if data_json["choices"][0].get("finish_reason", None):
is_finished = True
finish_reason = data_json["choices"][0]["finish_reason"]
print_verbose(f"text: {text}; is_finished: {is_finished}; finish_reason: {finish_reason}")
return {
"text": text,
"is_finished": is_finished,
"finish_reason": finish_reason,
}
except Exception:
raise ValueError(f"Unable to parse response. Original response: {chunk}")
elif "error" in chunk:
raise ValueError(f"Unable to parse response. Original response: {chunk}")
else:
return {
"text": text,
"is_finished": is_finished,
"finish_reason": finish_reason,
}
def handle_replicate_chunk(self, chunk):
try:
text = ""
@ -782,38 +685,6 @@ class CustomStreamWrapper:
except Exception as e:
raise e
def handle_baseten_chunk(self, chunk) -> str:
try:
chunk = chunk.decode("utf-8")
if len(chunk) > 0:
if chunk.startswith("data:"):
data_json: _BasetenStreamData = json.loads(chunk[5:])
if "token" in data_json and "text" in data_json["token"]:
return data_json["token"]["text"]
else:
return ""
data_json = json.loads(chunk)
if "model_output" in data_json:
if (
isinstance(data_json["model_output"], dict)
and "data" in data_json["model_output"]
and isinstance(data_json["model_output"]["data"], list)
):
return data_json["model_output"]["data"][0]
elif isinstance(data_json["model_output"], str):
return data_json["model_output"]
elif "completion" in data_json and isinstance(data_json["completion"], str):
return data_json["completion"]
else:
raise ValueError(f"Unable to parse response. Original response: {chunk}")
else:
return ""
else:
return ""
except Exception as e:
verbose_logger.exception("litellm.CustomStreamWrapper.handle_baseten_chunk(): Exception occured - %s", e)
return ""
def handle_triton_stream(self, chunk):
try:
if isinstance(chunk, dict):
@ -1305,18 +1176,6 @@ class CustomStreamWrapper:
completion_obj["content"] = response_obj["text"]
if response_obj["is_finished"]:
self.received_finish_reason = response_obj["finish_reason"]
elif self.custom_llm_provider and self.custom_llm_provider == "baseten": # baseten doesn't provide streaming
completion_obj["content"] = self.handle_baseten_chunk(chunk)
elif self.custom_llm_provider and self.custom_llm_provider == "ai21": # ai21 doesn't provide streaming
response_obj = self.handle_ai21_chunk(chunk)
completion_obj["content"] = response_obj["text"]
if response_obj["is_finished"]:
self.received_finish_reason = response_obj["finish_reason"]
elif self.custom_llm_provider and self.custom_llm_provider == "maritalk":
response_obj = self.handle_maritalk_chunk(chunk)
completion_obj["content"] = response_obj["text"]
if response_obj["is_finished"]:
self.received_finish_reason = response_obj["finish_reason"]
elif self.custom_llm_provider and self.custom_llm_provider == "vllm":
completion_obj["content"] = chunk[0].outputs[0].text
elif (
@ -1410,19 +1269,6 @@ class CustomStreamWrapper:
new_chunk = stream[:chunk_size]
completion_obj["content"] = new_chunk
self.completion_stream = stream[chunk_size:]
elif self.custom_llm_provider == "palm":
# fake streaming
response_obj = {}
if self.completion_stream is None or len(self.completion_stream) == 0:
if self.received_finish_reason is not None:
raise StopIteration
else:
self.received_finish_reason = "stop"
chunk_size = 30
stream = cast(Any, self.completion_stream)
new_chunk = stream[:chunk_size]
completion_obj["content"] = new_chunk
self.completion_stream = stream[chunk_size:]
elif self.custom_llm_provider == "triton":
response_obj = self.handle_triton_stream(chunk)
completion_obj["content"] = response_obj["text"]

4
litellm/llms/bedrock_mantle/responses/transformation.py
View file

@ -344,6 +344,10 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
kept: Final = [item for item, _ in normalized if item is not None] # mutable-ok: ResponseInputParam is a list
return kept # pyright: ignore[reportReturnType] # Codex passthrough items sit outside the OpenAI input union
@staticmethod
def _model_map_lookup_name(model: str) -> str:
return model.split("/")[-1].removeprefix("openai.")
def map_openai_params(
self,
response_api_optional_params: ResponsesAPIOptionalRequestParams,

129
litellm/llms/deprecated_providers/palm.py
View file

@ -1,27 +1,6 @@
import copy
import time
import traceback
import types
from collections.abc import Callable
from typing import Final
import httpx
import litellm
from litellm.utils import Choices, Message, ModelResponse, Usage
class PalmError(Exception):
def __init__(self, status_code, message):
self.status_code = status_code
self.message = message
self.request = httpx.Request(
method="POST",
url="https://developers.generativeai.google/api/python/google/generativeai/chat",
)
self.response = httpx.Response(status_code=status_code, request=self.request)
super().__init__(self.message) # Call the base class constructor with the parameters it needs
class PalmConfig:
"""
@ -84,111 +63,3 @@ class PalmConfig:
)
and v is not None
}
def completion(
model: str,
messages: list,
model_response: ModelResponse,
print_verbose: Callable,
api_key,
encoding,
logging_obj,
optional_params: dict,
litellm_params=None,
logger_fn=None,
):
try:
import google.generativeai as palm
except Exception:
raise Exception("Importing google.generativeai failed, please run 'pip install -q google-generativeai")
palm.configure(api_key=api_key)
model = model
## Load Config
inference_params: Final = copy.deepcopy(optional_params)
inference_params.pop(
"stream", None
) # palm does not support streaming, so we handle this by fake streaming in main.py
config: Final = litellm.PalmConfig.get_config()
for k, v in config.items():
if (
k not in inference_params
): # completion(top_k=3) > palm_config(top_k=3) <- allows for dynamic variables to be passed in
inference_params[k] = v
prompt = ""
for message in messages:
if "role" in message:
if message["role"] == "user":
prompt += f"{message['content']}"
else:
prompt += f"{message['content']}"
else:
prompt += f"{message['content']}"
## LOGGING
logging_obj.pre_call(
input=prompt,
api_key="",
additional_args={"complete_input_dict": {"inference_params": inference_params}},
)
## COMPLETION CALL
try:
response: Final = palm.generate_text(prompt=prompt, **inference_params)
except Exception as e:
raise PalmError(
message=str(e),
status_code=500,
)
## LOGGING
logging_obj.post_call(
input=prompt,
api_key="",
original_response=response,
additional_args={"complete_input_dict": {}},
)
print_verbose(f"raw model_response: {response}")
## RESPONSE OBJECT
completion_response = response
try:
choices_list: Final = []
for idx, item in enumerate(completion_response.candidates):
if len(item["output"]) > 0:
message_obj = Message(content=item["output"])
else:
message_obj = Message(content=None)
choice_obj = Choices(index=idx + 1, message=message_obj)
choices_list.append(choice_obj)
model_response.choices = choices_list
except Exception:
raise PalmError(message=traceback.format_exc(), status_code=response.status_code)
try:
completion_response = model_response["choices"][0]["message"].get("content")
except Exception:
raise PalmError(
status_code=400,
message=f"No response received. Original response - {response}",
)
## CALCULATING USAGE - baseten charges on time, not tokens - have some mapping of cost here.
prompt_tokens: Final = len(encoding.encode(prompt))
completion_tokens: Final = len(encoding.encode(model_response["choices"][0]["message"].get("content", "")))
model_response.created = int(time.time())
model_response.model = "palm/" + model
usage: Final = Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
)
setattr(model_response, "usage", usage)
return model_response
def embedding():
# logic for parsing in - calling - parsing out model embedding calls
pass

44
litellm/llms/openai/cost_calculation.py
View file

@ -38,7 +38,6 @@ def cost_per_token(
Returns:
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
"""
## CALCULATE INPUT COST
return generic_cost_per_token(
model=model,
usage=usage,
@ -46,49 +45,6 @@ def cost_per_token(
service_tier=service_tier,
data_residency=data_residency,
)
# ### Non-cached text tokens
# non_cached_text_tokens = usage.prompt_tokens
# cached_tokens: Optional[int] = None
# if usage.prompt_tokens_details and usage.prompt_tokens_details.cached_tokens:
# cached_tokens = usage.prompt_tokens_details.cached_tokens
# non_cached_text_tokens = non_cached_text_tokens - cached_tokens
# prompt_cost: float = non_cached_text_tokens * model_info["input_cost_per_token"]
# ## Prompt Caching cost calculation
# if model_info.get("cache_read_input_token_cost") is not None and cached_tokens:
# # Note: We read ._cache_read_input_tokens from the Usage - since cost_calculator.py standardizes the cache read tokens on usage._cache_read_input_tokens
# prompt_cost += cached_tokens * (
# model_info.get("cache_read_input_token_cost", 0) or 0
# )
# _audio_tokens: Optional[int] = (
# usage.prompt_tokens_details.audio_tokens
# if usage.prompt_tokens_details is not None
# else None
# )
# _audio_cost_per_token: Optional[float] = model_info.get(
# "input_cost_per_audio_token"
# )
# if _audio_tokens is not None and _audio_cost_per_token is not None:
# audio_cost: float = _audio_tokens * _audio_cost_per_token
# prompt_cost += audio_cost
# ## CALCULATE OUTPUT COST
# completion_cost: float = (
# usage["completion_tokens"] * model_info["output_cost_per_token"]
# )
# _output_cost_per_audio_token: Optional[float] = model_info.get(
# "output_cost_per_audio_token"
# )
# _output_audio_tokens: Optional[int] = (
# usage.completion_tokens_details.audio_tokens
# if usage.completion_tokens_details is not None
# else None
# )
# if _output_cost_per_audio_token is not None and _output_audio_tokens is not None:
# audio_cost = _output_audio_tokens * _output_cost_per_audio_token
# completion_cost += audio_cost
# return prompt_cost, completion_cost
def cost_per_second(model: str, custom_llm_provider: str | None, duration: float = 0.0) -> tuple[float, float]:

36
litellm/llms/openai/responses/transformation.py
View file

@ -125,6 +125,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
return False
return is_gpt_reasoning_series_name(model)
@staticmethod
def _model_map_lookup_name(model: str) -> str:
return model
@staticmethod
def _supports_reasoning_effort_none(model: str) -> bool:
"""Return True if the model supports reasoning.effort='none'."""
@ -208,8 +212,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
) -> dict:
"""No mapping applied since inputs are in OpenAI spec already.
GPT-5 models have restrictions on temperature (only temperature=1
is accepted unless reasoning_effort='none' on models that support it).
GPT-5 models have restrictions on temperature and top_p (only temperature=1
is accepted, and top_p is rejected, unless reasoning.effort resolves to
'none' on models that support it).
Apply the same validation used by the chat completions path.
"""
params: Final = dict(response_api_optional_params)
@ -234,13 +239,16 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
status_code=400,
)
if self._is_gpt_5_model(model=model):
lookup_name: Final = self._model_map_lookup_name(model)
if self._is_gpt_5_model(model=lookup_name):
reasoning: Final = params.get("reasoning") or {}
effort: Final = reasoning.get("effort") if isinstance(reasoning, dict) else None
supports_none: Final = self._supports_reasoning_effort_none(model=lookup_name)
effort_is_none: Final = supports_none and self._effort_resolves_to_none(lookup_name, effort)
temperature: Final = params.get("temperature")
if temperature is not None and temperature != 1:
reasoning: Final = params.get("reasoning") or {}
effort: Final = reasoning.get("effort") if isinstance(reasoning, dict) else None
supports_none: Final = self._supports_reasoning_effort_none(model=model)
if supports_none and self._effort_resolves_to_none(model, effort):
if effort_is_none:
pass # flexible temperature allowed
elif drop_params or litellm.drop_params:
params.pop("temperature", None)
@ -256,6 +264,20 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
status_code=400,
)
if "top_p" in params and not effort_is_none:
if drop_params or litellm.drop_params:
params.pop("top_p", None)
else:
raise litellm.UnsupportedParamsError(
message=(
f"{model} only supports top_p when reasoning.effort resolves to 'none', "
"either set explicitly on the request or declared as the model's "
"default_reasoning_effort. "
"To drop unsupported params set `litellm.drop_params = True`"
),
status_code=400,
)
return params
def transform_responses_api_request(

2
litellm/llms/vertex_ai/count_tokens/handler.py
View file

@ -20,7 +20,6 @@ class VertexAITokenCounter(GoogleAIStudioTokenCounter, VertexBase):
vertex_credentials: Final = self.get_vertex_ai_credentials(litellm_params=litellm_params)
vertex_project = self.get_vertex_ai_project(litellm_params=litellm_params)
vertex_location: Final = self.get_vertex_ai_location(litellm_params=litellm_params)
should_use_v1beta1_features: Final = self.is_using_v1beta1_features(litellm_params)
_auth_header, vertex_project = await self._ensure_access_token_async(
credentials=vertex_credentials,
project_id=vertex_project,
@ -37,7 +36,6 @@ class VertexAITokenCounter(GoogleAIStudioTokenCounter, VertexBase):
stream=False,
custom_llm_provider="vertex_ai",
api_base=None,
should_use_v1beta1_features=should_use_v1beta1_features,
mode="count_tokens",
)
headers = {

9
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -2701,8 +2701,6 @@ class VertexLLM(VertexBase):
gemini_api_key: str | None = None,
extra_headers: dict | None = None,
) -> CustomStreamWrapper:
should_use_v1beta1_features: Final = self.is_using_v1beta1_features(optional_params=optional_params)
_auth_header, vertex_project = await self._ensure_access_token_async(
credentials=vertex_credentials,
project_id=vertex_project,
@ -2722,7 +2720,6 @@ class VertexLLM(VertexBase):
stream=stream,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
should_use_v1beta1_features=should_use_v1beta1_features,
use_psc_endpoint_format=use_psc_endpoint_format,
)
@ -2797,8 +2794,6 @@ class VertexLLM(VertexBase):
gemini_api_key: str | None = None,
extra_headers: dict | None = None,
) -> ModelResponse | CustomStreamWrapper:
should_use_v1beta1_features: Final = self.is_using_v1beta1_features(optional_params=optional_params)
_auth_header, vertex_project = await self._ensure_access_token_async(
credentials=vertex_credentials,
project_id=vertex_project,
@ -2818,7 +2813,6 @@ class VertexLLM(VertexBase):
stream=stream,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
should_use_v1beta1_features=should_use_v1beta1_features,
use_psc_endpoint_format=use_psc_endpoint_format,
)
@ -2981,8 +2975,6 @@ class VertexLLM(VertexBase):
extra_headers=extra_headers,
)
should_use_v1beta1_features: Final = self.is_using_v1beta1_features(optional_params=optional_params)
_auth_header, vertex_project = self._ensure_access_token(
credentials=vertex_credentials,
project_id=vertex_project,
@ -3002,7 +2994,6 @@ class VertexLLM(VertexBase):
stream=stream,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
should_use_v1beta1_features=should_use_v1beta1_features,
use_psc_endpoint_format=use_psc_endpoint_format,
)
headers: Final = VertexGeminiConfig().validate_environment(

5
litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py
View file

@ -65,8 +65,6 @@ class VertexEmbedding(VertexBase):
litellm_params=litellm_params,
)
should_use_v1beta1_features: Final = self.is_using_v1beta1_features(optional_params=optional_params)
_auth_header, vertex_project = self._ensure_access_token(
credentials=vertex_credentials,
project_id=vertex_project,
@ -85,7 +83,6 @@ class VertexEmbedding(VertexBase):
stream=False,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
should_use_v1beta1_features=should_use_v1beta1_features,
mode="embedding",
use_psc_endpoint_format=use_psc_endpoint_format,
)
@ -160,7 +157,6 @@ class VertexEmbedding(VertexBase):
"""
Async embedding implementation
"""
should_use_v1beta1_features: Final = self.is_using_v1beta1_features(optional_params=optional_params)
_auth_header, vertex_project = await self._ensure_access_token_async(
credentials=vertex_credentials,
project_id=vertex_project,
@ -179,7 +175,6 @@ class VertexEmbedding(VertexBase):
stream=False,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
should_use_v1beta1_features=should_use_v1beta1_features,
mode="embedding",
use_psc_endpoint_format=use_psc_endpoint_format,
)

9
litellm/llms/vertex_ai/vertex_llm_base.py
View file

@ -618,15 +618,6 @@ class VertexBase:
project_id=project_id,
)
def is_using_v1beta1_features(self, optional_params: dict) -> bool:
"""
use this helper to decide if request should be sent to v1 or v1beta1
Returns true if any beta feature is enabled
Returns false in all other cases
"""
return False
def _check_custom_proxy(
self,
api_base: str | None,

2
litellm/main.py
View file

@ -206,7 +206,7 @@ from .llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler
from .llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from .llms.custom_llm import CustomLLM, custom_chat_llm_router
from .llms.databricks.embed.handler import DatabricksEmbeddingHandler
from .llms.deprecated_providers import aleph_alpha, palm
from .llms.deprecated_providers import aleph_alpha
from .llms.gdc.chat.transformation import GDCGeminiConfig
from .llms.gemini.common_utils import get_api_key_from_env
from .llms.groq.chat.handler import GroqChatCompletion

41
litellm/proxy/_logging.py
View file

@ -1,41 +0,0 @@
### DEPRECATED ###
## unused file. initially written for json logging on proxy.
import json
import logging
import os
from logging import Formatter
from typing import Final
from litellm import json_logs
# Set default log level to INFO
log_level: Final = os.getenv("LITELLM_LOG", "INFO")
numeric_level: Final[str] = getattr(logging, log_level.upper())
class JsonFormatter(Formatter):
def __init__(self):
super().__init__()
def format(self, record):
json_record: Final = {
"message": record.getMessage(),
"level": record.levelname,
"timestamp": self.formatTime(record, self.datefmt),
}
return json.dumps(json_record)
logger: Final = logging.root
handler: Final = logging.StreamHandler()
if json_logs:
handler.setFormatter(JsonFormatter())
else:
formatter: Final = logging.Formatter(
"\033[92m%(asctime)s - %(name)s:%(levelname)s\033[0m: %(filename)s:%(lineno)s - %(message)s",
datefmt="%H:%M:%S",
)
handler.setFormatter(formatter)
logger.handlers = [handler]
logger.setLevel(numeric_level)

213
litellm/proxy/common_utils/performance_utils.md
View file

@ -1,213 +0,0 @@
# Performance Utilities Documentation
This module provides performance monitoring and profiling functionality for LiteLLM proxy server using `cProfile` and `line_profiler`.
## Table of Contents
- [Line Profiler Usage](#line-profiler-usage)
- [Example 1: Wrapping a function directly](#example-1-wrapping-a-function-directly)
- [Example 2: Wrapping a module function dynamically](#example-2-wrapping-a-module-function-dynamically)
- [Example 3: Manual stats collection](#example-3-manual-stats-collection)
- [Example 4: Analyzing the profile output](#example-4-analyzing-the-profile-output)
- [Example 5: Using in a decorator pattern](#example-5-using-in-a-decorator-pattern)
- [cProfile Usage](#cprofile-usage)
- [Installation](#installation)
- [Notes](#notes)
## Line Profiler Usage
### Example 1: Wrapping a function directly
This is how it's used in `litellm/utils.py` to profile `wrapper_async`:
```python
from litellm.proxy.common_utils.performance_utils import (
register_shutdown_handler,
wrap_function_directly,
)
def client(original_function):
@wraps(original_function)
async def wrapper_async(*args, **kwargs):
# ... function implementation ...
pass
# Wrap the function with line_profiler
wrapper_async = wrap_function_directly(wrapper_async)
# Register shutdown handler to collect stats on server shutdown
register_shutdown_handler(output_file="wrapper_async_line_profile.lprof")
return wrapper_async
```
### Example 2: Wrapping a module function dynamically
```python
import my_module
from litellm.proxy.common_utils.performance_utils import (
wrap_function_with_line_profiler,
register_shutdown_handler,
)
# Wrap a function in a module
wrap_function_with_line_profiler(my_module, "expensive_function")
# Register shutdown handler
register_shutdown_handler(output_file="my_profile.lprof")
# Now all calls to my_module.expensive_function will be profiled
my_module.expensive_function()
```
### Example 3: Manual stats collection
```python
from litellm.proxy.common_utils.performance_utils import (
wrap_function_directly,
collect_line_profiler_stats,
)
def my_function():
# ... implementation ...
pass
# Wrap the function
my_function = wrap_function_directly(my_function)
# Run your code
my_function()
# Collect stats manually (instead of waiting for shutdown)
collect_line_profiler_stats(output_file="manual_profile.lprof")
```
### Example 4: Analyzing the profile output
After running your code, analyze the `.lprof` file:
```bash
# View the profile
python -m line_profiler wrapper_async_line_profile.lprof
# Save to text file
python -m line_profiler wrapper_async_line_profile.lprof > profile_report.txt
```
The output shows:
- **Line #**: Line number in the source file
- **Hits**: Number of times the line was executed
- **Time**: Total time spent on that line (in microseconds)
- **Per Hit**: Average time per execution
- **% Time**: Percentage of total function time
- **Line Contents**: The actual source code
Example output:
```
Timer unit: 1e-06 s
Total time: 3.73697 s
File: litellm/utils.py
Function: client.<locals>.wrapper_async at line 1657
Line # Hits Time Per Hit % Time Line Contents
==============================================================
1657 @wraps(original_function)
1658 async def wrapper_async(*args, **kwargs):
1659 2005 7577.1 3.8 0.2 print_args_passed_to_litellm(...)
1763 2005 1351909.0 674.3 36.2 result = await original_function(*args, **kwargs)
1846 4010 1543688.1 385.0 41.3 update_response_metadata(...)
```
### Example 5: Using in a decorator pattern
```python
from litellm.proxy.common_utils.performance_utils import (
wrap_function_directly,
register_shutdown_handler,
)
def profile_decorator(func):
# Wrap the function
profiled_func = wrap_function_directly(func)
# Register shutdown handler (only once)
if not hasattr(profile_decorator, '_registered'):
register_shutdown_handler(output_file="decorated_functions.lprof")
profile_decorator._registered = True
return profiled_func
@profile_decorator
async def my_async_function():
# This function will be profiled
pass
```
## cProfile Usage
### Example: Using the profile_endpoint decorator
```python
from litellm.proxy.common_utils.performance_utils import profile_endpoint
@profile_endpoint(sampling_rate=0.1) # Profile 10% of requests
async def my_endpoint():
# ... implementation ...
pass
```
The `sampling_rate` parameter controls what percentage of requests are profiled:
- `1.0`: Profile all requests (100%)
- `0.1`: Profile 1 in 10 requests (10%)
- `0.0`: Profile no requests (0%)
## Installation
`line_profiler` must be installed to use the line profiling functionality:
```bash
uv add --dev line-profiler
```
On Windows with Python 3.14+, you may need to install Microsoft Visual C++ Build Tools to compile `line_profiler` from source.
## Notes
- The profiler aggregates stats by source code location, so multiple instances of the same function (e.g., closures) will be profiled together
- Stats are automatically collected on server shutdown via `atexit` handler when using `register_shutdown_handler()`
- You can also manually collect stats using `collect_line_profiler_stats()`
- The line profiler will fail with an `ImportError` if `line_profiler` is not installed (as configured in `litellm/utils.py`)
## API Reference
### `wrap_function_directly(func: Callable) -> Callable`
Wrap a function directly with line_profiler. This is the recommended way to profile functions, especially closures or functions created dynamically.
**Raises:**
- `ImportError`: If line_profiler is not available
- `RuntimeError`: If line_profiler cannot be enabled or function cannot be wrapped
### `wrap_function_with_line_profiler(module: Any, function_name: str) -> bool`
Dynamically wrap a function in a module with line_profiler.
**Returns:** `True` if wrapping was successful, `False` otherwise
### `collect_line_profiler_stats(output_file: Optional[str] = None) -> None`
Collect and save line_profiler statistics. If `output_file` is provided, saves to file. Otherwise, prints to stdout.
### `register_shutdown_handler(output_file: Optional[str] = None) -> None`
Register an `atexit` handler that will automatically save profiling statistics when the Python process exits. Safe to call multiple times (only registers once).
**Default output file:** `line_profile_stats.lprof` if not specified
### `profile_endpoint(sampling_rate: float = 1.0)`
Decorator to sample endpoint hits and save to a profile file using cProfile.
**Args:**
- `sampling_rate`: Rate of requests to profile (0.0 to 1.0)

299
litellm/proxy/common_utils/performance_utils.py
View file

@ -1,299 +0,0 @@
"""
Performance utilities for LiteLLM proxy server.
This module provides performance monitoring and profiling functionality for endpoint
performance analysis using cProfile with configurable sampling rates, and line_profiler
for line-by-line profiling.
See performance_utils.md for detailed usage examples and documentation.
"""
import atexit
import cProfile
import functools
import inspect
import threading
from collections.abc import Callable
from pathlib import Path as PathLib
from types import ModuleType
from typing import Final, Protocol, TextIO
from litellm._logging import verbose_proxy_logger
class _LineProfiler(Protocol):
"""The line_profiler.LineProfiler surface this module drives."""
def __call__(self, func: Callable[..., object]) -> Callable[..., object]: ...
def add_function(self, func: Callable[..., object]) -> object: ...
def dump_stats(self, filename: str) -> object: ...
def print_stats(self, stream: TextIO) -> object: ...
# Global profiling state
_profile_lock: Final = threading.Lock()
_profiler = None
_last_profile_file_path = None
_sample_counter = 0
_sample_counter_lock: Final = threading.Lock()
# Global line_profiler state
_line_profiler: _LineProfiler | None = None
_line_profiler_lock: Final = threading.Lock()
_wrapped_functions: Final[dict[str, Callable]] = {} # Store original functions
def _should_sample(profile_sampling_rate: float) -> bool:
"""Determine if current request should be sampled based on sampling rate."""
if profile_sampling_rate >= 1.0:
return True # Always sample
elif profile_sampling_rate <= 0.0:
return False # Never sample
# Use deterministic sampling based on counter for consistent rate
global _sample_counter
with _sample_counter_lock:
_sample_counter += 1
# Sample based on rate (e.g., 0.1 means sample every 10th request)
should_sample: Final = (_sample_counter % int(1.0 / profile_sampling_rate)) == 0
return should_sample
def _start_profiling(profile_sampling_rate: float) -> None:
"""Start cProfile profiling once globally."""
global _profiler
with _profile_lock:
if _profiler is None:
_profiler = cProfile.Profile()
_profiler.enable()
verbose_proxy_logger.info("Profiling started with sampling rate: %s", profile_sampling_rate)
def _start_profiling_for_request(profile_sampling_rate: float) -> bool:
"""Start profiling for a specific request (if sampling allows)."""
if _should_sample(profile_sampling_rate):
_start_profiling(profile_sampling_rate)
return True
return False
def _save_stats(profile_file: PathLib) -> None:
"""Save current stats directly to file."""
with _profile_lock:
if _profiler is None:
return
try:
# Disable profiler temporarily to dump stats
_profiler.disable()
_profiler.dump_stats(str(profile_file))
# Re-enable profiler to continue profiling
_profiler.enable()
verbose_proxy_logger.debug("Profiling stats saved to %s", profile_file)
except Exception as e:
verbose_proxy_logger.error("Error saving profiling stats: %s", e)
# Make sure profiler is re-enabled even if there's an error
try:
_profiler.enable()
except Exception:
pass
def profile_endpoint(sampling_rate: float = 1.0):
"""Decorator to sample endpoint hits and save to a profile file.
Args:
sampling_rate: Rate of requests to profile (0.0 to 1.0)
- 1.0: Profile all requests (100%)
- 0.1: Profile 1 in 10 requests (10%)
- 0.0: Profile no requests (0%)
"""
def decorator(func):
def set_last_profile_path(path: PathLib) -> None:
global _last_profile_file_path
_last_profile_file_path = path
if inspect.iscoroutinefunction(func):
@functools.wraps(func)
async def async_wrapper(*args, **kwargs):
is_sampling: Final = _start_profiling_for_request(sampling_rate)
file_path_obj: Final = PathLib("endpoint_profile.pstat")
set_last_profile_path(file_path_obj)
try:
result: Final = await func(*args, **kwargs)
if is_sampling:
_save_stats(file_path_obj)
return result
except Exception:
if is_sampling:
_save_stats(file_path_obj)
raise
return async_wrapper
else:
@functools.wraps(func)
def sync_wrapper(*args, **kwargs):
is_sampling: Final = _start_profiling_for_request(sampling_rate)
file_path_obj: Final = PathLib("endpoint_profile.pstat")
set_last_profile_path(file_path_obj)
try:
result: Final = func(*args, **kwargs)
if is_sampling:
_save_stats(file_path_obj)
return result
except Exception:
if is_sampling:
_save_stats(file_path_obj)
raise
return sync_wrapper
return decorator
def enable_line_profiler() -> None:
"""Enable line_profiler for dynamic function wrapping.
Raises:
ImportError: If line_profiler is not available
"""
global _line_profiler
from line_profiler import LineProfiler # Will raise ImportError if not available
with _line_profiler_lock:
if _line_profiler is None:
_line_profiler = LineProfiler()
verbose_proxy_logger.info("Line profiler enabled")
def wrap_function_with_line_profiler(module: ModuleType, function_name: str) -> bool:
"""Dynamically wrap a function with line_profiler.
Args:
module: The module containing the function
function_name: Name of the function to wrap
Returns:
True if wrapping was successful, False otherwise
"""
try:
enable_line_profiler() # May raise ImportError if not available
except ImportError:
return False
if _line_profiler is None:
return False
try:
original_function: Final = getattr(module, function_name, None)
if original_function is None:
verbose_proxy_logger.warning("Function %s not found in module %s", function_name, module.__name__)
return False
# Store original function if not already wrapped
if function_name not in _wrapped_functions:
_wrapped_functions[function_name] = original_function
# Wrap with line_profiler
profiled_function: Final = _line_profiler(original_function)
setattr(module, function_name, profiled_function)
verbose_proxy_logger.info("Wrapped %s.%s with line_profiler", module.__name__, function_name)
return True
except Exception as e:
verbose_proxy_logger.error("Error wrapping %s with line_profiler: %s", function_name, e)
return False
def wrap_function_directly(func: Callable) -> Callable:
"""Wrap a function directly with line_profiler.
This is the recommended way to profile functions, especially closures or
functions created dynamically (like wrapper_async in litellm/utils.py).
Args:
func: The function to wrap
Returns:
The wrapped function that will be profiled when called
Raises:
ImportError: If line_profiler is not available
RuntimeError: If line_profiler cannot be enabled or function cannot be wrapped
"""
import warnings
enable_line_profiler() # Will raise ImportError if not available
if _line_profiler is None:
raise RuntimeError("Line profiler was not initialized")
# Suppress warnings about __wrapped__ - we intentionally want to profile the wrapper
with warnings.catch_warnings():
warnings.filterwarnings("ignore", message=".*__wrapped__.*", category=UserWarning)
# Add function to line_profiler and wrap it
_line_profiler.add_function(func)
profiled_function: Final = _line_profiler(func)
verbose_proxy_logger.info("Wrapped function %s with line_profiler", func.__name__)
return profiled_function
def collect_line_profiler_stats(output_file: str | None = None) -> None:
"""Collect and save line_profiler statistics.
This can be called manually to collect stats at any time, or it's
automatically called on shutdown if register_shutdown_handler() was used.
Args:
output_file: Optional path to save stats. If None, prints to stdout.
"""
global _line_profiler
with _line_profiler_lock:
if _line_profiler is None:
verbose_proxy_logger.debug("Line profiler not enabled, nothing to collect")
return
try:
if output_file:
# Save to file
output_path: Final = PathLib(output_file)
_line_profiler.dump_stats(str(output_path))
verbose_proxy_logger.info("Line profiler stats saved to %s", output_path)
else:
# Print to stdout
from io import StringIO
stream: Final = StringIO()
_line_profiler.print_stats(stream=stream)
stats_output: Final = stream.getvalue()
verbose_proxy_logger.info("Line profiler stats:\n" + stats_output)
except Exception as e:
verbose_proxy_logger.error("Error collecting line profiler stats: %s", e)
def register_shutdown_handler(output_file: str | None = None) -> None:
"""Register a shutdown handler to collect line_profiler stats.
This registers an atexit handler that will automatically save profiling
statistics when the Python process exits. Safe to call multiple times
(only registers once).
Args:
output_file: Optional path to save stats on shutdown.
Defaults to 'line_profile_stats.lprof'
"""
if output_file is None:
output_file = "line_profile_stats.lprof"
def shutdown_handler():
collect_line_profiler_stats(output_file=output_file)
atexit.register(shutdown_handler)
verbose_proxy_logger.debug("Registered line_profiler shutdown handler for %s", output_file)

5
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -9,6 +9,7 @@ from collections.abc import AsyncGenerator, Callable, Iterable, Mapping, Sequenc
from dataclasses import dataclass
from datetime import datetime
from itertools import groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
from urllib.parse import urlencode, urlparse
@ -991,7 +992,7 @@ async def pass_through_request(
)
upstream_headers: Final = _with_trace_context(headers, parent_span=user_api_key_dict.parent_otel_span)
requested_query_params: dict | None = query_params or dict(request.query_params)
requested_query_params: dict | None = query_params or dict(request.query_params) or None
endpoint_type: Final[EndpointType] = HttpPassThroughEndpointHelpers.get_endpoint_type(str(url))
@ -1193,7 +1194,7 @@ async def pass_through_request(
query=urlencode(
HttpPassThroughEndpointHelpers.get_merged_query_parameters(
existing_url=url,
request_query_params=requested_query_params,
request_query_params=requested_query_params or MappingProxyType({}),
default_query_params=default_query_params,
)
).encode("ascii")

87
tests/image_gen_tests/test_image_variation.py
View file

@ -1,87 +0,0 @@
# What this tests?
## This tests the litellm support for the openai /generations endpoint
import logging
import traceback
from dotenv import load_dotenv
from openai.types.image import Image
from litellm.caching import InMemoryCache
logging.basicConfig(level=logging.DEBUG)
load_dotenv()
import asyncio
import pytest
import litellm
import json
import tempfile
from base_image_generation_test import BaseImageGenTest
import logging
from litellm._logging import verbose_logger
from io import BytesIO
from PIL import Image as PILImage
verbose_logger.setLevel(logging.DEBUG)
@pytest.fixture
def image_url():
# DALL-E 2 image variations require a square PNG (less than 4MB)
# Generate a 1024x1024 square PNG programmatically to avoid network dependency
# and the non-square aspect ratio of the old LiteLLM logo URL
img = PILImage.new("RGBA", (1024, 1024), color=(128, 128, 128, 255))
image_file = BytesIO()
img.save(image_file, format="PNG")
image_file.seek(0)
# openai>=2.24.0 requires BytesIO to have .name for MIME type detection in multipart uploads
image_file.name = "litellm_logo.png"
return image_file
# Commented out: OpenAI /images/variations endpoint deprecated (DALL-E 2 shutdown May 12, 2026)
# def test_openai_image_variation_openai_sdk(image_url):
# from openai import OpenAI
#
# client = OpenAI()
# response = client.images.create_variation(image=image_url, n=2, size="1024x1024")
# print(response)
#
#
# @pytest.mark.parametrize("sync_mode", [True, False])
# @pytest.mark.asyncio
# async def test_openai_image_variation_litellm_sdk(image_url, sync_mode):
# from litellm import image_variation, aimage_variation
#
# if sync_mode:
# image_variation(image=image_url, n=2, size="1024x1024")
# else:
# await aimage_variation(image=image_url, n=2, size="1024x1024")
#
#
# def test_topaz_image_variation(image_url):
# from litellm import image_variation, aimage_variation
# from litellm.llms.custom_httpx.http_handler import HTTPHandler
# from unittest.mock import patch
#
# client = HTTPHandler()
# with patch.object(client, "post") as mock_post:
# try:
# image_variation(
# model="topaz/Standard V2",
# image=image_url,
# n=2,
# size="1024x1024",
# client=client,
# )
# except Exception as e:
# print(e)
# mock_post.assert_called_once()
def test_image_variation_placeholder():
"""Placeholder: variation tests commented out - OpenAI /images/variations deprecated (DALL-E 2 shutdown May 12, 2026)."""
pass

128
tests/local_testing/test_azure_perf.py
View file

@ -1,128 +0,0 @@
# #### What this tests ####
# # This adds perf testing to the router, to ensure it's never > 50ms slower than the azure-openai sdk.
# import sys, os, time, inspect, asyncio, traceback
# from datetime import datetime
# import pytest
# sys.path.insert(0, os.path.abspath("../.."))
# import openai, litellm, uuid
# from openai import AsyncAzureOpenAI
# client = AsyncAzureOpenAI(
# api_key=os.getenv("AZURE_AI_API_KEY"),
# azure_endpoint=os.getenv("AZURE_AI_API_BASE"), # type: ignore
# api_version=os.getenv("AZURE_API_VERSION"),
# )
# model_list = [
# {
# "model_name": "azure-test",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "api_base": os.getenv("AZURE_AI_API_BASE"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# },
# }
# ]
# router = litellm.Router(model_list=model_list) # type: ignore
# async def _openai_completion():
# try:
# start_time = time.time()
# response = await client.chat.completions.create(
# model="chatgpt-v-3",
# messages=[{"role": "user", "content": f"This is a test: {uuid.uuid4()}"}],
# stream=True,
# )
# time_to_first_token = None
# first_token_ts = None
# init_chunk = None
# async for chunk in response:
# if (
# time_to_first_token is None
# and len(chunk.choices) > 0
# and chunk.choices[0].delta.content is not None
# ):
# first_token_ts = time.time()
# time_to_first_token = first_token_ts - start_time
# init_chunk = chunk
# end_time = time.time()
# print(
# "OpenAI Call: ",
# init_chunk,
# start_time,
# first_token_ts,
# time_to_first_token,
# end_time,
# )
# return time_to_first_token
# except Exception as e:
# print(e)
# return None
# async def _router_completion():
# try:
# start_time = time.time()
# response = await router.acompletion(
# model="azure-test",
# messages=[{"role": "user", "content": f"This is a test: {uuid.uuid4()}"}],
# stream=True,
# )
# time_to_first_token = None
# first_token_ts = None
# init_chunk = None
# async for chunk in response:
# if (
# time_to_first_token is None
# and len(chunk.choices) > 0
# and chunk.choices[0].delta.content is not None
# ):
# first_token_ts = time.time()
# time_to_first_token = first_token_ts - start_time
# init_chunk = chunk
# end_time = time.time()
# print(
# "Router Call: ",
# init_chunk,
# start_time,
# first_token_ts,
# time_to_first_token,
# end_time - first_token_ts,
# )
# return time_to_first_token
# except Exception as e:
# print(e)
# return None
# async def test_azure_completion_streaming():
# """
# Test azure streaming call - measure on time to first (non-null) token.
# """
# n = 3 # Number of concurrent tasks
# ## OPENAI AVG. TIME
# tasks = [_openai_completion() for _ in range(n)]
# chat_completions = await asyncio.gather(*tasks)
# successful_completions = [c for c in chat_completions if c is not None]
# total_time = 0
# for item in successful_completions:
# total_time += item
# avg_openai_time = total_time / 3
# ## ROUTER AVG. TIME
# tasks = [_router_completion() for _ in range(n)]
# chat_completions = await asyncio.gather(*tasks)
# successful_completions = [c for c in chat_completions if c is not None]
# total_time = 0
# for item in successful_completions:
# total_time += item
# avg_router_time = total_time / 3
# ## COMPARE
# print(f"avg_router_time: {avg_router_time}; avg_openai_time: {avg_openai_time}")
# assert avg_router_time < avg_openai_time + 0.5
# # asyncio.run(test_azure_completion_streaming())

130
tests/local_testing/test_budget_manager.py
View file

@ -1,130 +0,0 @@
# #### What this tests ####
# # This tests calling batch_completions by running 100 messages together
# import sys, os, json
# import traceback
# import pytest
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import litellm
# litellm.set_verbose = True
# from litellm import completion, BudgetManager
# budget_manager = BudgetManager(project_name="test_project", client_type="hosted")
# ## Scenario 1: User budget enough to make call
# def test_user_budget_enough():
# try:
# user = "1234"
# # create a budget for a user
# budget_manager.create_budget(total_budget=10, user=user, duration="daily")
# # check if a given call can be made
# data = {
# "model": "gpt-3.5-turbo",
# "messages": [{"role": "user", "content": "Hey, how's it going?"}]
# }
# if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user):
# response = completion(**data)
# print(budget_manager.update_cost(completion_obj=response, user=user))
# else:
# response = "Sorry - no budget!"
# print(f"response: {response}")
# except Exception as e:
# pytest.fail(f"An error occurred - {str(e)}")
# ## Scenario 2: User budget not enough to make call
# def test_user_budget_not_enough():
# try:
# user = "12345"
# # create a budget for a user
# budget_manager.create_budget(total_budget=0, user=user, duration="daily")
# # check if a given call can be made
# data = {
# "model": "gpt-3.5-turbo",
# "messages": [{"role": "user", "content": "Hey, how's it going?"}]
# }
# model = data["model"]
# messages = data["messages"]
# if budget_manager.get_current_cost(user=user) < budget_manager.get_total_budget(user=user):
# response = completion(**data)
# print(budget_manager.update_cost(completion_obj=response, user=user))
# else:
# response = "Sorry - no budget!"
# print(f"response: {response}")
# except Exception:
# pytest.fail(f"An error occurred")
# ## Scenario 3: Saving budget to client
# def test_save_user_budget():
# try:
# response = budget_manager.save_data()
# if response["status"] == "error":
# raise Exception(f"An error occurred - {json.dumps(response)}")
# print(response)
# except Exception as e:
# pytest.fail(f"An error occurred: {str(e)}")
# test_save_user_budget()
# ## Scenario 4: Getting list of users
# def test_get_users():
# try:
# response = budget_manager.get_users()
# print(response)
# except Exception:
# pytest.fail(f"An error occurred")
# ## Scenario 5: Reset budget at the end of duration
# def test_reset_on_duration():
# try:
# # First, set a short duration budget for a user
# user = "123456"
# budget_manager.create_budget(total_budget=10, user=user, duration="daily")
# # Use some of the budget
# data = {
# "model": "gpt-3.5-turbo",
# "messages": [{"role": "user", "content": "Hello!"}]
# }
# if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user=user):
# response = litellm.completion(**data)
# print(budget_manager.update_cost(completion_obj=response, user=user))
# assert budget_manager.get_current_cost(user) > 0, f"Test setup failed: Budget did not decrease after completion"
# # Now, we need to simulate the passing of time. Since we don't want our tests to actually take days, we're going
# # to cheat a little -- we'll manually adjust the "created_at" time so it seems like a day has passed.
# # In a real-world testing scenario, we might instead use something like the `freezegun` library to mock the system time.
# one_day_in_seconds = 24 * 60 * 60
# budget_manager.user_dict[user]["last_updated_at"] -= one_day_in_seconds
# # Now the duration should have expired, so our budget should reset
# budget_manager.update_budget_all_users()
# # Make sure the budget was actually reset
# assert budget_manager.get_current_cost(user) == 0, "Budget didn't reset after duration expired"
# except Exception as e:
# pytest.fail(f"An error occurred - {str(e)}")
# ## Scenario 6: passing in text:
# def test_input_text_on_completion():
# try:
# user = "12345"
# budget_manager.create_budget(total_budget=10, user=user, duration="daily")
# input_text = "hello world"
# output_text = "it's a sunny day in san francisco"
# model = "gpt-3.5-turbo"
# budget_manager.update_cost(user=user, model=model, input_text=input_text, output_text=output_text)
# print(budget_manager.get_current_cost(user))
# except Exception as e:
# pytest.fail(f"An error occurred - {str(e)}")
# test_input_text_on_completion()

124
tests/local_testing/test_class.py
View file

@ -1,124 +0,0 @@
# # #### What this tests ####
# # # This tests the LiteLLM Class
# import sys, os
# import traceback
# import pytest
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import litellm
# import asyncio
# # litellm.set_verbose = True
# # from litellm import Router
# import instructor
# from litellm import completion
# from pydantic import BaseModel
# class User(BaseModel):
# name: str
# age: int
# client = instructor.from_litellm(completion)
# litellm.set_verbose = True
# resp = client.chat.completions.create(
# model="gpt-3.5-turbo",
# max_tokens=1024,
# messages=[
# {
# "role": "user",
# "content": "Extract Jason is 25 years old.",
# }
# ],
# response_model=User,
# num_retries=10,
# )
# assert isinstance(resp, User)
# assert resp.name == "Jason"
# assert resp.age == 25
# # from pydantic import BaseModel
# # # This enables response_model keyword
# # # from client.chat.completions.create
# # client = instructor.patch(
# # Router(
# # model_list=[
# # {
# # "model_name": "gpt-3.5-turbo", # openai model name
# # "litellm_params": { # params for litellm completion/embedding call
# # "model": "azure/gpt-4.1-mini",
# # "api_key": os.getenv("AZURE_AI_API_KEY"),
# # "api_version": os.getenv("AZURE_API_VERSION"),
# # "api_base": os.getenv("AZURE_AI_API_BASE"),
# # },
# # }
# # ]
# # )
# # )
# # class UserDetail(BaseModel):
# # name: str
# # age: int
# # user = client.chat.completions.create(
# # model="gpt-3.5-turbo",
# # response_model=UserDetail,
# # messages=[
# # {"role": "user", "content": "Extract Jason is 25 years old"},
# # ],
# # )
# # assert isinstance(user, UserDetail)
# # assert user.name == "Jason"
# # assert user.age == 25
# # print(f"user: {user}")
# # # import instructor
# # # from openai import AsyncOpenAI
# # aclient = instructor.apatch(
# # Router(
# # model_list=[
# # {
# # "model_name": "gpt-3.5-turbo", # openai model name
# # "litellm_params": { # params for litellm completion/embedding call
# # "model": "azure/gpt-4.1-mini",
# # "api_key": os.getenv("AZURE_AI_API_KEY"),
# # "api_version": os.getenv("AZURE_API_VERSION"),
# # "api_base": os.getenv("AZURE_AI_API_BASE"),
# # },
# # }
# # ],
# # default_litellm_params={"acompletion": True},
# # )
# # )
# # class UserExtract(BaseModel):
# # name: str
# # age: int
# # async def main():
# # model = await aclient.chat.completions.create(
# # model="gpt-3.5-turbo",
# # response_model=UserExtract,
# # messages=[
# # {"role": "user", "content": "Extract jason is 25 years old"},
# # ],
# # )
# # print(f"model: {model}")
# # asyncio.run(main())

190
tests/local_testing/test_completion_cost.py
View file

@ -153,23 +153,12 @@ def test_custom_pricing_as_completion_cost_param():
assert round(cost, 5) == round(expected_cost, 5)
def test_get_gpt3_tokens():
max_tokens = get_max_tokens("gpt-3.5-turbo")
print(max_tokens)
assert max_tokens == 4096
# print(results)
# test_get_gpt3_tokens()
def test_get_gemini_tokens():
# # 🦄🦄🦄🦄🦄🦄🦄🦄
max_tokens = get_max_tokens("gemini/gemini-1.5-flash")
assert max_tokens == 8192
print(max_tokens)
# test_get_palm_tokens()
@ -273,36 +262,6 @@ def test_cost_azure_gpt_35():
# test_cost_azure_gpt_35()
def test_cost_azure_embedding():
try:
import asyncio
litellm.set_verbose = True
async def _test():
response = await litellm.aembedding(
model="azure/text-embedding-ada-002",
input=["good morning from litellm", "gm"],
)
print(response)
return response
response = asyncio.run(_test())
cost = litellm.completion_cost(completion_response=response)
print("Cost", cost)
expected_cost = float("7e-07")
assert cost == expected_cost
except Exception as e:
pytest.fail(
f"Cost Calc failed for azure/gpt-3.5-turbo. Expected {expected_cost}, Calculated cost {cost}"
)
# test_cost_azure_embedding()
@ -639,56 +598,6 @@ def test_vertex_ai_medlm_completion_cost():
assert predictive_cost > 0
def test_vertex_ai_claude_completion_cost():
from litellm import Choices, Message, ModelResponse
from litellm.utils import Usage
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
litellm.set_verbose = True
input_tokens = litellm.token_counter(
model="vertex_ai/claude-3-sonnet@20240229",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
)
print(f"input_tokens: {input_tokens}")
output_tokens = litellm.token_counter(
model="vertex_ai/claude-3-sonnet@20240229",
text="It's all going well",
count_response_tokens=True,
)
print(f"output_tokens: {output_tokens}")
response = ModelResponse(
id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac",
choices=[
Choices(
finish_reason=None,
index=0,
message=Message(
content="It's all going well",
role="assistant",
),
)
],
created=1700775391,
model="claude-3-sonnet",
object="chat.completion",
system_fingerprint=None,
usage=Usage(
prompt_tokens=input_tokens,
completion_tokens=output_tokens,
total_tokens=input_tokens + output_tokens,
),
)
cost = litellm.completion_cost(
model="vertex_ai/claude-3-sonnet",
completion_response=response,
messages=[{"role": "user", "content": "Hey, how's it going?"}],
)
predicted_cost = input_tokens * 0.000003 + 0.000015 * output_tokens
assert cost == predicted_cost
def test_vertex_ai_embedding_completion_cost(caplog):
"""
Relevant issue - https://github.com/BerriAI/litellm/issues/4630
@ -1212,105 +1121,6 @@ def test_completion_cost_fireworks_ai(model):
assert cost > 0
def test_cost_azure_openai_prompt_caching():
from litellm.utils import Choices, Message, ModelResponse, Usage
from litellm.types.utils import (
PromptTokensDetailsWrapper,
CompletionTokensDetailsWrapper,
)
from litellm import get_model_info
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "azure/o1-mini"
## LLM API CALL ## (MORE EXPENSIVE)
response_1 = ModelResponse(
id="chatcmpl-3f427194-0840-4d08-b571-56bfe38a5424",
choices=[
Choices(
finish_reason="length",
index=0,
message=Message(
content="Hello! I'm doing well, thank you for",
role="assistant",
tool_calls=None,
function_call=None,
),
)
],
created=1725036547,
model=model,
object="chat.completion",
system_fingerprint=None,
usage=Usage(
completion_tokens=10,
prompt_tokens=14,
total_tokens=24,
completion_tokens_details=CompletionTokensDetailsWrapper(
reasoning_tokens=2
),
),
)
## PROMPT CACHE HIT ## (LESS EXPENSIVE)
response_2 = ModelResponse(
id="chatcmpl-3f427194-0840-4d08-b571-56bfe38a5424",
choices=[
Choices(
finish_reason="length",
index=0,
message=Message(
content="Hello! I'm doing well, thank you for",
role="assistant",
tool_calls=None,
function_call=None,
),
)
],
created=1725036547,
model=model,
object="chat.completion",
system_fingerprint=None,
usage=Usage(
completion_tokens=10,
prompt_tokens=0,
total_tokens=10,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=14,
),
completion_tokens_details=CompletionTokensDetailsWrapper(
reasoning_tokens=2
),
),
)
cost_1 = completion_cost(model=model, completion_response=response_1)
cost_2 = completion_cost(model=model, completion_response=response_2)
assert cost_1 > cost_2
model_info = get_model_info(model=model, custom_llm_provider="azure")
usage = response_2.usage
_expected_cost2 = (
(usage.prompt_tokens - usage.prompt_tokens_details.cached_tokens)
* model_info["input_cost_per_token"]
+ (usage.completion_tokens * model_info["output_cost_per_token"])
+ (
usage.prompt_tokens_details.cached_tokens
* model_info["cache_read_input_token_cost"]
)
)
print("_expected_cost2", _expected_cost2)
print("cost_2", cost_2)
assert (
abs(cost_2 - _expected_cost2) < 1e-5
) # Allow for small floating-point differences
def test_completion_cost_vertex_llama3():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")

90
tests/local_testing/test_langchain_ChatLiteLLM.py
View file

@ -1,90 +0,0 @@
# import os
# import sys, os
# import traceback
# from dotenv import load_dotenv
# load_dotenv()
# import os, io
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import pytest
# import litellm
# from litellm import embedding, completion, text_completion, completion_cost
# from langchain.chat_models import ChatLiteLLM
# from langchain.prompts.chat import (
# ChatPromptTemplate,
# SystemMessagePromptTemplate,
# AIMessagePromptTemplate,
# HumanMessagePromptTemplate,
# )
# from langchain.schema import AIMessage, HumanMessage, SystemMessage
# def test_chat_gpt():
# try:
# chat = ChatLiteLLM(model="gpt-3.5-turbo", max_tokens=10)
# messages = [
# HumanMessage(
# content="what model are you"
# )
# ]
# resp = chat(messages)
# print(resp)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# # test_chat_gpt()
# def test_claude():
# try:
# chat = ChatLiteLLM(model="claude-2", max_tokens=10)
# messages = [
# HumanMessage(
# content="what model are you"
# )
# ]
# resp = chat(messages)
# print(resp)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# # test_claude()
# # def test_openai_with_params():
# # try:
# # api_key = os.environ["OPENAI_API_KEY"]
# # os.environ.pop("OPENAI_API_KEY")
# # print("testing openai with params")
# # llm = ChatLiteLLM(
# # model="gpt-3.5-turbo",
# # openai_api_key=api_key,
# # # Prefer using None which is the default value, endpoint could be empty string
# # openai_api_base= None,
# # max_tokens=20,
# # temperature=0.5,
# # request_timeout=10,
# # model_kwargs={
# # "frequency_penalty": 0,
# # "presence_penalty": 0,
# # },
# # verbose=True,
# # max_retries=0,
# # )
# # messages = [
# # HumanMessage(
# # content="what model are you"
# # )
# # ]
# # resp = llm(messages)
# # print(resp)
# # except Exception as e:
# # pytest.fail(f"Error occurred: {e}")
# # test_openai_with_params()

94
tests/local_testing/test_load_test_router_s3.py
View file

@ -1,94 +0,0 @@
# import sys, os
# import traceback
# from dotenv import load_dotenv
# import copy
# load_dotenv()
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import asyncio
# from litellm import Router, Timeout
# import time
# from litellm.caching.caching import Cache
# import litellm
# litellm.cache = Cache(
# type="s3", s3_bucket_name="litellm-my-test-bucket-2", s3_region_name="us-west-2"
# )
# ### Test calling router with s3 Cache
# async def call_acompletion(semaphore, router: Router, input_data):
# async with semaphore:
# try:
# # Use asyncio.wait_for to set a timeout for the task
# response = await router.acompletion(**input_data)
# # Handle the response as needed
# print(response)
# return response
# except Timeout:
# print(f"Task timed out: {input_data}")
# return None # You may choose to return something else or raise an exception
# async def main():
# # Initialize the Router
# model_list = [
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "gpt-3.5-turbo",
# "api_key": os.getenv("OPENAI_API_KEY"),
# },
# },
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# },
# },
# ]
# router = Router(model_list=model_list, num_retries=3, timeout=10)
# # Create a semaphore with a capacity of 100
# semaphore = asyncio.Semaphore(100)
# # List to hold all task references
# tasks = []
# start_time_all_tasks = time.time()
# # Launch 1000 tasks
# for _ in range(500):
# task = asyncio.create_task(
# call_acompletion(
# semaphore,
# router,
# {
# "model": "gpt-3.5-turbo",
# "messages": [{"role": "user", "content": "Hey, how's it going?"}],
# },
# )
# )
# tasks.append(task)
# # Wait for all tasks to complete
# responses = await asyncio.gather(*tasks)
# # Process responses as needed
# # Record the end time for all tasks
# end_time_all_tasks = time.time()
# # Calculate the total time for all tasks
# total_time_all_tasks = end_time_all_tasks - start_time_all_tasks
# print(f"Total time for all tasks: {total_time_all_tasks} seconds")
# # Calculate the average time per response
# average_time_per_response = total_time_all_tasks / len(responses)
# print(f"Average time per response: {average_time_per_response} seconds")
# print(f"NUMBER OF COMPLETED TASKS: {len(responses)}")
# # Run the main function
# asyncio.run(main())

86
tests/local_testing/test_loadtest_router.py
View file

@ -1,86 +0,0 @@
# import sys, os
# import traceback
# from dotenv import load_dotenv
# import copy
# load_dotenv()
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import asyncio
# from litellm import Router, Timeout
# import time
# async def call_acompletion(semaphore, router: Router, input_data):
# async with semaphore:
# try:
# # Use asyncio.wait_for to set a timeout for the task
# response = await router.acompletion(**input_data)
# # Handle the response as needed
# print(response)
# return response
# except Timeout:
# print(f"Task timed out: {input_data}")
# return None # You may choose to return something else or raise an exception
# async def main():
# # Initialize the Router
# model_list = [
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "gpt-3.5-turbo",
# "api_key": os.getenv("OPENAI_API_KEY"),
# },
# },
# {
# "model_name": "gpt-3.5-turbo",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_AI_API_KEY"),
# "api_base": os.getenv("AZURE_AI_API_BASE"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# },
# },
# ]
# router = Router(model_list=model_list, num_retries=3, timeout=10)
# # Create a semaphore with a capacity of 100
# semaphore = asyncio.Semaphore(100)
# # List to hold all task references
# tasks = []
# start_time_all_tasks = time.time()
# # Launch 1000 tasks
# for _ in range(500):
# task = asyncio.create_task(
# call_acompletion(
# semaphore,
# router,
# {
# "model": "gpt-3.5-turbo",
# "messages": [{"role": "user", "content": "Hey, how's it going?"}],
# },
# )
# )
# tasks.append(task)
# # Wait for all tasks to complete
# responses = await asyncio.gather(*tasks)
# # Process responses as needed
# # Record the end time for all tasks
# end_time_all_tasks = time.time()
# # Calculate the total time for all tasks
# total_time_all_tasks = end_time_all_tasks - start_time_all_tasks
# print(f"Total time for all tasks: {total_time_all_tasks} seconds")
# # Calculate the average time per response
# average_time_per_response = total_time_all_tasks / len(responses)
# print(f"Average time per response: {average_time_per_response} seconds")
# print(f"NUMBER OF COMPLETED TASKS: {len(responses)}")
# # Run the main function
# asyncio.run(main())

382
tests/local_testing/test_logging.py
View file

@ -1,382 +0,0 @@
# #### What this tests ####
# # This tests error logging (with custom user functions) for the raw `completion` + `embedding` endpoints
# # Test Scenarios (test across completion, streaming, embedding)
# ## 1: Pre-API-Call
# ## 2: Post-API-Call
# ## 3: On LiteLLM Call success
# ## 4: On LiteLLM Call failure
# import sys, os, io
# import traceback, logging
# import pytest
# import dotenv
# dotenv.load_dotenv()
# # Create logger
# logger = logging.getLogger(__name__)
# logger.setLevel(logging.DEBUG)
# # Create a stream handler
# stream_handler = logging.StreamHandler(sys.stdout)
# logger.addHandler(stream_handler)
# # Create a function to log information
# def logger_fn(message):
# logger.info(message)
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import litellm
# from litellm import embedding, completion
# from openai.error import AuthenticationError
# litellm.set_verbose = True
# score = 0
# user_message = "Hello, how are you?"
# messages = [{"content": user_message, "role": "user"}]
# # 1. On Call Success
# # normal completion
# # test on openai completion call
# def test_logging_success_completion():
# global score
# try:
# # Redirect stdout
# old_stdout = sys.stdout
# sys.stdout = new_stdout = io.StringIO()
# response = completion(model="gpt-3.5-turbo", messages=messages)
# # Restore stdout
# sys.stdout = old_stdout
# output = new_stdout.getvalue().strip()
# if "Logging Details Pre-API Call" not in output:
# raise Exception("Required log message not found!")
# elif "Logging Details Post-API Call" not in output:
# raise Exception("Required log message not found!")
# elif "Logging Details LiteLLM-Success Call" not in output:
# raise Exception("Required log message not found!")
# score += 1
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# pass
# # ## test on non-openai completion call
# # def test_logging_success_completion_non_openai():
# # global score
# # try:
# # # Redirect stdout
# # old_stdout = sys.stdout
# # sys.stdout = new_stdout = io.StringIO()
# # response = completion(model="claude-3-5-haiku-20241022", messages=messages)
# # # Restore stdout
# # sys.stdout = old_stdout
# # output = new_stdout.getvalue().strip()
# # if "Logging Details Pre-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details Post-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details LiteLLM-Success Call" not in output:
# # raise Exception("Required log message not found!")
# # score += 1
# # except Exception as e:
# # pytest.fail(f"Error occurred: {e}")
# # pass
# # streaming completion
# ## test on openai completion call
# def test_logging_success_streaming_openai():
# global score
# try:
# # litellm.set_verbose = False
# def custom_callback(
# kwargs, # kwargs to completion
# completion_response, # response from completion
# start_time, end_time # start/end time
# ):
# if "complete_streaming_response" in kwargs:
# print(f"Complete Streaming Response: {kwargs['complete_streaming_response']}")
# # Assign the custom callback function
# litellm.success_callback = [custom_callback]
# # Redirect stdout
# old_stdout = sys.stdout
# sys.stdout = new_stdout = io.StringIO()
# response = completion(model="gpt-3.5-turbo", messages=messages, stream=True)
# for chunk in response:
# pass
# # Restore stdout
# sys.stdout = old_stdout
# output = new_stdout.getvalue().strip()
# if "Logging Details Pre-API Call" not in output:
# raise Exception("Required log message not found!")
# elif "Logging Details Post-API Call" not in output:
# raise Exception("Required log message not found!")
# elif "Logging Details LiteLLM-Success Call" not in output:
# raise Exception("Required log message not found!")
# elif "Complete Streaming Response:" not in output:
# raise Exception("Required log message not found!")
# score += 1
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# pass
# # test_logging_success_streaming_openai()
# ## test on non-openai completion call
# def test_logging_success_streaming_non_openai():
# global score
# try:
# # litellm.set_verbose = False
# def custom_callback(
# kwargs, # kwargs to completion
# completion_response, # response from completion
# start_time, end_time # start/end time
# ):
# # print(f"streaming response: {completion_response}")
# if "complete_streaming_response" in kwargs:
# print(f"Complete Streaming Response: {kwargs['complete_streaming_response']}")
# # Assign the custom callback function
# litellm.success_callback = [custom_callback]
# # Redirect stdout
# old_stdout = sys.stdout
# sys.stdout = new_stdout = io.StringIO()
# response = completion(model="claude-3-5-haiku-20241022", messages=messages, stream=True)
# for idx, chunk in enumerate(response):
# pass
# # Restore stdout
# sys.stdout = old_stdout
# output = new_stdout.getvalue().strip()
# if "Logging Details Pre-API Call" not in output:
# raise Exception("Required log message not found!")
# elif "Logging Details Post-API Call" not in output:
# raise Exception("Required log message not found!")
# elif "Logging Details LiteLLM-Success Call" not in output:
# raise Exception("Required log message not found!")
# elif "Complete Streaming Response:" not in output:
# raise Exception(f"Required log message not found! {output}")
# score += 1
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# pass
# # test_logging_success_streaming_non_openai()
# # embedding
# def test_logging_success_embedding_openai():
# try:
# # Redirect stdout
# old_stdout = sys.stdout
# sys.stdout = new_stdout = io.StringIO()
# response = embedding(model="text-embedding-ada-002", input=["good morning from litellm"])
# # Restore stdout
# sys.stdout = old_stdout
# output = new_stdout.getvalue().strip()
# if "Logging Details Pre-API Call" not in output:
# raise Exception("Required log message not found!")
# elif "Logging Details Post-API Call" not in output:
# raise Exception("Required log message not found!")
# elif "Logging Details LiteLLM-Success Call" not in output:
# raise Exception("Required log message not found!")
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# # ## 2. On LiteLLM Call failure
# # ## TEST BAD KEY
# # # normal completion
# # ## test on openai completion call
# # try:
# # temporary_oai_key = os.environ["OPENAI_API_KEY"]
# # os.environ["OPENAI_API_KEY"] = "bad-key"
# # temporary_anthropic_key = os.environ["ANTHROPIC_API_KEY"]
# # os.environ["ANTHROPIC_API_KEY"] = "bad-key"
# # # Redirect stdout
# # old_stdout = sys.stdout
# # sys.stdout = new_stdout = io.StringIO()
# # try:
# # response = completion(model="gpt-3.5-turbo", messages=messages)
# # except AuthenticationError:
# # print(f"raised auth error")
# # pass
# # # Restore stdout
# # sys.stdout = old_stdout
# # output = new_stdout.getvalue().strip()
# # print(output)
# # if "Logging Details Pre-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details Post-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details LiteLLM-Failure Call" not in output:
# # raise Exception("Required log message not found!")
# # os.environ["OPENAI_API_KEY"] = temporary_oai_key
# # os.environ["ANTHROPIC_API_KEY"] = temporary_anthropic_key
# # score += 1
# # except Exception as e:
# # print(f"exception type: {type(e).__name__}")
# # pytest.fail(f"Error occurred: {e}")
# # pass
# # ## test on non-openai completion call
# # try:
# # temporary_oai_key = os.environ["OPENAI_API_KEY"]
# # os.environ["OPENAI_API_KEY"] = "bad-key"
# # temporary_anthropic_key = os.environ["ANTHROPIC_API_KEY"]
# # os.environ["ANTHROPIC_API_KEY"] = "bad-key"
# # # Redirect stdout
# # old_stdout = sys.stdout
# # sys.stdout = new_stdout = io.StringIO()
# # try:
# # response = completion(model="claude-3-5-haiku-20241022", messages=messages)
# # except AuthenticationError:
# # pass
# # if "Logging Details Pre-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details Post-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details LiteLLM-Failure Call" not in output:
# # raise Exception("Required log message not found!")
# # os.environ["OPENAI_API_KEY"] = temporary_oai_key
# # os.environ["ANTHROPIC_API_KEY"] = temporary_anthropic_key
# # score += 1
# # except Exception as e:
# # print(f"exception type: {type(e).__name__}")
# # # Restore stdout
# # sys.stdout = old_stdout
# # output = new_stdout.getvalue().strip()
# # print(output)
# # pytest.fail(f"Error occurred: {e}")
# # # streaming completion
# # ## test on openai completion call
# # try:
# # temporary_oai_key = os.environ["OPENAI_API_KEY"]
# # os.environ["OPENAI_API_KEY"] = "bad-key"
# # temporary_anthropic_key = os.environ["ANTHROPIC_API_KEY"]
# # os.environ["ANTHROPIC_API_KEY"] = "bad-key"
# # # Redirect stdout
# # old_stdout = sys.stdout
# # sys.stdout = new_stdout = io.StringIO()
# # try:
# # response = completion(model="gpt-3.5-turbo", messages=messages)
# # except AuthenticationError:
# # pass
# # # Restore stdout
# # sys.stdout = old_stdout
# # output = new_stdout.getvalue().strip()
# # print(output)
# # if "Logging Details Pre-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details Post-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details LiteLLM-Failure Call" not in output:
# # raise Exception("Required log message not found!")
# # os.environ["OPENAI_API_KEY"] = temporary_oai_key
# # os.environ["ANTHROPIC_API_KEY"] = temporary_anthropic_key
# # score += 1
# # except Exception as e:
# # print(f"exception type: {type(e).__name__}")
# # pytest.fail(f"Error occurred: {e}")
# # ## test on non-openai completion call
# # try:
# # temporary_oai_key = os.environ["OPENAI_API_KEY"]
# # os.environ["OPENAI_API_KEY"] = "bad-key"
# # temporary_anthropic_key = os.environ["ANTHROPIC_API_KEY"]
# # os.environ["ANTHROPIC_API_KEY"] = "bad-key"
# # # Redirect stdout
# # old_stdout = sys.stdout
# # sys.stdout = new_stdout = io.StringIO()
# # try:
# # response = completion(model="claude-3-5-haiku-20241022", messages=messages)
# # except AuthenticationError:
# # pass
# # # Restore stdout
# # sys.stdout = old_stdout
# # output = new_stdout.getvalue().strip()
# # print(output)
# # if "Logging Details Pre-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details Post-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details LiteLLM-Failure Call" not in output:
# # raise Exception("Required log message not found!")
# # score += 1
# # except Exception as e:
# # print(f"exception type: {type(e).__name__}")
# # pytest.fail(f"Error occurred: {e}")
# # # embedding
# # try:
# # temporary_oai_key = os.environ["OPENAI_API_KEY"]
# # os.environ["OPENAI_API_KEY"] = "bad-key"
# # temporary_anthropic_key = os.environ["ANTHROPIC_API_KEY"]
# # os.environ["ANTHROPIC_API_KEY"] = "bad-key"
# # # Redirect stdout
# # old_stdout = sys.stdout
# # sys.stdout = new_stdout = io.StringIO()
# # try:
# # response = embedding(model="text-embedding-ada-002", input=["good morning from litellm"])
# # except AuthenticationError:
# # pass
# # # Restore stdout
# # sys.stdout = old_stdout
# # output = new_stdout.getvalue().strip()
# # print(output)
# # if "Logging Details Pre-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details Post-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details LiteLLM-Failure Call" not in output:
# # raise Exception("Required log message not found!")
# # except Exception as e:
# # print(f"exception type: {type(e).__name__}")
# # pytest.fail(f"Error occurred: {e}")

163
tests/local_testing/test_max_tpm_rpm_limiter.py
View file

@ -1,163 +0,0 @@
### REPLACED BY 'test_parallel_request_limiter.py' ###
# What is this?
## Unit tests for the max tpm / rpm limiter hook for proxy
# import sys, os, asyncio, time, random
# from datetime import datetime
# import traceback
# from dotenv import load_dotenv
# from typing import Optional
# load_dotenv()
# import os
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import pytest
# import litellm
# from litellm import Router
# from litellm.proxy.utils import ProxyLogging, hash_token
# from litellm.proxy._types import UserAPIKeyAuth
# from litellm.caching.caching import DualCache, RedisCache
# from litellm.proxy.hooks.tpm_rpm_limiter import _PROXY_MaxTPMRPMLimiter
# from datetime import datetime
# @pytest.mark.asyncio
# async def test_pre_call_hook_rpm_limits():
# """
# Test if error raised on hitting rpm limits
# """
# litellm.set_verbose = True
# _api_key = hash_token("sk-12345")
# user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=9, rpm_limit=1)
# local_cache = DualCache()
# # redis_usage_cache = RedisCache()
# local_cache.set_cache(
# key=_api_key, value={"api_key": _api_key, "tpm_limit": 9, "rpm_limit": 1}
# )
# tpm_rpm_limiter = _PROXY_MaxTPMRPMLimiter(internal_cache=DualCache())
# await tpm_rpm_limiter.async_pre_call_hook(
# user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
# )
# kwargs = {"litellm_params": {"metadata": {"user_api_key": _api_key}}}
# await tpm_rpm_limiter.async_log_success_event(
# kwargs=kwargs,
# response_obj="",
# start_time="",
# end_time="",
# )
# ## Expected cache val: {"current_requests": 0, "current_tpm": 0, "current_rpm": 1}
# try:
# await tpm_rpm_limiter.async_pre_call_hook(
# user_api_key_dict=user_api_key_dict,
# cache=local_cache,
# data={},
# call_type="",
# )
# pytest.fail(f"Expected call to fail")
# except Exception as e:
# assert e.status_code == 429
# @pytest.mark.asyncio
# async def test_pre_call_hook_team_rpm_limits(
# _redis_usage_cache: Optional[RedisCache] = None,
# ):
# """
# Test if error raised on hitting team rpm limits
# """
# litellm.set_verbose = True
# _api_key = "sk-12345"
# _team_id = "unique-team-id"
# _user_api_key_dict = {
# "api_key": _api_key,
# "max_parallel_requests": 1,
# "tpm_limit": 9,
# "rpm_limit": 10,
# "team_rpm_limit": 1,
# "team_id": _team_id,
# }
# user_api_key_dict = UserAPIKeyAuth(**_user_api_key_dict) # type: ignore
# _api_key = hash_token(_api_key)
# local_cache = DualCache()
# local_cache.set_cache(key=_api_key, value=_user_api_key_dict)
# internal_cache = DualCache(redis_cache=_redis_usage_cache)
# tpm_rpm_limiter = _PROXY_MaxTPMRPMLimiter(internal_cache=internal_cache)
# await tpm_rpm_limiter.async_pre_call_hook(
# user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
# )
# kwargs = {
# "litellm_params": {
# "metadata": {"user_api_key": _api_key, "user_api_key_team_id": _team_id}
# }
# }
# await tpm_rpm_limiter.async_log_success_event(
# kwargs=kwargs,
# response_obj="",
# start_time="",
# end_time="",
# )
# print(f"local_cache: {local_cache}")
# ## Expected cache val: {"current_requests": 0, "current_tpm": 0, "current_rpm": 1}
# try:
# await tpm_rpm_limiter.async_pre_call_hook(
# user_api_key_dict=user_api_key_dict,
# cache=local_cache,
# data={},
# call_type="",
# )
# pytest.fail(f"Expected call to fail")
# except Exception as e:
# assert e.status_code == 429 # type: ignore
# @pytest.mark.asyncio
# async def test_namespace():
# """
# - test if default namespace set via `proxyconfig._init_cache`
# - respected for tpm/rpm caching
# """
# from litellm.proxy.proxy_server import ProxyConfig
# redis_usage_cache: Optional[RedisCache] = None
# cache_params = {"type": "redis", "namespace": "litellm_default"}
# ## INIT CACHE ##
# proxy_config = ProxyConfig()
# setattr(litellm.proxy.proxy_server, "proxy_config", proxy_config)
# proxy_config._init_cache(cache_params=cache_params)
# redis_cache: Optional[RedisCache] = getattr(
# litellm.proxy.proxy_server, "redis_usage_cache"
# )
# ## CHECK IF NAMESPACE SET ##
# assert redis_cache.namespace == "litellm_default"
# ## CHECK IF TPM/RPM RATE LIMITING WORKS ##
# await test_pre_call_hook_team_rpm_limits(_redis_usage_cache=redis_cache)
# current_date = datetime.now().strftime("%Y-%m-%d")
# current_hour = datetime.now().strftime("%H")
# current_minute = datetime.now().strftime("%M")
# precise_minute = f"{current_date}-{current_hour}-{current_minute}"
# cache_key = "litellm_default:usage:{}".format(precise_minute)
# value = await redis_cache.async_get_cache(key=cache_key)
# assert value is not None

243
tests/local_testing/test_mem_leak.py
View file

@ -1,243 +0,0 @@
# import io
# import os
# import sys
# sys.path.insert(0, os.path.abspath("../.."))
# import litellm
# from memory_profiler import profile
# from litellm.utils import (
# ModelResponseIterator,
# ModelResponseListIterator,
# CustomStreamWrapper,
# )
# from litellm.types.utils import ModelResponse, Choices, Message
# import time
# import pytest
# # @app.post("/debug")
# # async def debug(body: ExampleRequest) -> str:
# # return await main_logic(body.query)
# def model_response_list_factory():
# chunks = [
# {
# "id": "chatcmpl-9SQxdH5hODqkWyJopWlaVOOUnFwlj",
# "choices": [
# {
# "delta": {"content": "", "role": "assistant"},
# "finish_reason": None,
# "index": 0,
# }
# ],
# "created": 1716563849,
# "model": "gpt-4o-2024-05-13",
# "object": "chat.completion.chunk",
# "system_fingerprint": "fp_5f4bad809a",
# },
# {
# "id": "chatcmpl-9SQxdH5hODqkWyJopWlaVOOUnFwlj",
# "choices": [
# {"delta": {"content": "This"}, "finish_reason": None, "index": 0}
# ],
# "created": 1716563849,
# "model": "gpt-4o-2024-05-13",
# "object": "chat.completion.chunk",
# "system_fingerprint": "fp_5f4bad809a",
# },
# {
# "id": "chatcmpl-9SQxdH5hODqkWyJopWlaVOOUnFwlj",
# "choices": [
# {"delta": {"content": " is"}, "finish_reason": None, "index": 0}
# ],
# "created": 1716563849,
# "model": "gpt-4o-2024-05-13",
# "object": "chat.completion.chunk",
# "system_fingerprint": "fp_5f4bad809a",
# },
# {
# "id": "chatcmpl-9SQxdH5hODqkWyJopWlaVOOUnFwlj",
# "choices": [
# {"delta": {"content": " a"}, "finish_reason": None, "index": 0}
# ],
# "created": 1716563849,
# "model": "gpt-4o-2024-05-13",
# "object": "chat.completion.chunk",
# "system_fingerprint": "fp_5f4bad809a",
# },
# {
# "id": "chatcmpl-9SQxdH5hODqkWyJopWlaVOOUnFwlj",
# "choices": [
# {"delta": {"content": " dummy"}, "finish_reason": None, "index": 0}
# ],
# "created": 1716563849,
# "model": "gpt-4o-2024-05-13",
# "object": "chat.completion.chunk",
# "system_fingerprint": "fp_5f4bad809a",
# },
# {
# "id": "chatcmpl-9SQxdH5hODqkWyJopWlaVOOUnFwlj",
# "choices": [
# {
# "delta": {"content": " response"},
# "finish_reason": None,
# "index": 0,
# }
# ],
# "created": 1716563849,
# "model": "gpt-4o-2024-05-13",
# "object": "chat.completion.chunk",
# "system_fingerprint": "fp_5f4bad809a",
# },
# {
# "id": "",
# "choices": [
# {
# "finish_reason": None,
# "index": 0,
# "content_filter_offsets": {
# "check_offset": 35159,
# "start_offset": 35159,
# "end_offset": 36150,
# },
# "content_filter_results": {
# "hate": {"filtered": False, "severity": "safe"},
# "self_harm": {"filtered": False, "severity": "safe"},
# "sexual": {"filtered": False, "severity": "safe"},
# "violence": {"filtered": False, "severity": "safe"},
# },
# }
# ],
# "created": 0,
# "model": "",
# "object": "",
# },
# {
# "id": "chatcmpl-9SQxdH5hODqkWyJopWlaVOOUnFwlj",
# "choices": [{"delta": {"content": "."}, "finish_reason": None, "index": 0}],
# "created": 1716563849,
# "model": "gpt-4o-2024-05-13",
# "object": "chat.completion.chunk",
# "system_fingerprint": "fp_5f4bad809a",
# },
# {
# "id": "chatcmpl-9SQxdH5hODqkWyJopWlaVOOUnFwlj",
# "choices": [{"delta": {}, "finish_reason": "stop", "index": 0}],
# "created": 1716563849,
# "model": "gpt-4o-2024-05-13",
# "object": "chat.completion.chunk",
# "system_fingerprint": "fp_5f4bad809a",
# },
# {
# "id": "",
# "choices": [
# {
# "finish_reason": None,
# "index": 0,
# "content_filter_offsets": {
# "check_offset": 36150,
# "start_offset": 36060,
# "end_offset": 37029,
# },
# "content_filter_results": {
# "hate": {"filtered": False, "severity": "safe"},
# "self_harm": {"filtered": False, "severity": "safe"},
# "sexual": {"filtered": False, "severity": "safe"},
# "violence": {"filtered": False, "severity": "safe"},
# },
# }
# ],
# "created": 0,
# "model": "",
# "object": "",
# },
# ]
# chunk_list = []
# for chunk in chunks:
# new_chunk = litellm.ModelResponse(stream=True, id=chunk["id"])
# if "choices" in chunk and isinstance(chunk["choices"], list):
# new_choices = []
# for choice in chunk["choices"]:
# if isinstance(choice, litellm.utils.StreamingChoices):
# _new_choice = choice
# elif isinstance(choice, dict):
# _new_choice = litellm.utils.StreamingChoices(**choice)
# new_choices.append(_new_choice)
# new_chunk.choices = new_choices
# chunk_list.append(new_chunk)
# return ModelResponseListIterator(model_responses=chunk_list)
# async def mock_completion(*args, **kwargs):
# completion_stream = model_response_list_factory()
# return litellm.CustomStreamWrapper(
# completion_stream=completion_stream,
# model="gpt-4-0613",
# custom_llm_provider="cached_response",
# logging_obj=litellm.Logging(
# model="gpt-4-0613",
# messages=[{"role": "user", "content": "Hey"}],
# stream=True,
# call_type="completion",
# start_time=time.time(),
# litellm_call_id="12345",
# function_id="1245",
# ),
# )
# @profile
# async def main_logic() -> str:
# stream = await mock_completion()
# result = ""
# async for chunk in stream:
# result += chunk.choices[0].delta.content or ""
# return result
# import asyncio
# for _ in range(100):
# asyncio.run(main_logic())
# # @pytest.mark.asyncio
# # def test_memory_profile(capsys):
# # # Run the async function
# # result = asyncio.run(main_logic())
# # # Verify the result
# # assert result == "This is a dummy response."
# # # Capture the output
# # captured = capsys.readouterr()
# # # Print memory output for debugging
# # print("Memory Profiler Output:")
# # print(f"captured out: {captured.out}")
# # # Basic memory leak checks
# # for idx, line in enumerate(captured.out.split("\n")):
# # if idx % 2 == 0 and "MiB" in line:
# # print(f"line: {line}")
# # # mem_lines = [line for line in captured.out.split("\n") if "MiB" in line]
# # print(mem_lines)
# # # Ensure we have some memory lines
# # assert len(mem_lines) > 0, "No memory profiler output found"
# # # Optional: Add more specific memory leak detection
# # for line in mem_lines:
# # # Extract memory increment
# # parts = line.split()
# # if len(parts) >= 3:
# # try:
# # mem_increment = float(parts[2].replace("MiB", ""))
# # # Assert that memory increment is below a reasonable threshold
# # assert mem_increment < 1.0, f"Potential memory leak detected: {line}"
# # except (ValueError, IndexError):
# # pass # Skip lines that don't match expected format

153
tests/local_testing/test_mem_usage.py
View file

@ -1,153 +0,0 @@
# #### What this tests ####
# from memory_profiler import profile, memory_usage
# import sys, os, time
# import traceback, asyncio
# import pytest
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import litellm
# from litellm import Router
# from concurrent.futures import ThreadPoolExecutor
# from collections import defaultdict
# from dotenv import load_dotenv
# from litellm._uuid import uuid
# import tracemalloc
# import objgraph
# objgraph.growth(shortnames=True)
# objgraph.show_most_common_types(limit=10)
# from mem_top import mem_top
# load_dotenv()
# model_list = [
# {
# "model_name": "gpt-3.5-turbo", # openai model name
# "litellm_params": { # params for litellm completion/embedding call
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_version": os.getenv("AZURE_API_VERSION"),
# "api_base": os.getenv("AZURE_API_BASE"),
# },
# "tpm": 240000,
# "rpm": 1800,
# },
# {
# "model_name": "bad-model", # openai model name
# "litellm_params": { # params for litellm completion/embedding call
# "model": "azure/gpt-4.1-mini",
# "api_key": "bad-key",
# "api_version": os.getenv("AZURE_API_VERSION"),
# "api_base": os.getenv("AZURE_API_BASE"),
# },
# "tpm": 240000,
# "rpm": 1800,
# },
# {
# "model_name": "text-embedding-ada-002",
# "litellm_params": {
# "model": "azure/text-embedding-ada-002",
# "api_key": os.environ["AZURE_API_KEY"],
# "api_base": os.environ["AZURE_API_BASE"],
# },
# "tpm": 100000,
# "rpm": 10000,
# },
# ]
# litellm.set_verbose = True
# litellm.cache = litellm.Cache(
# type="s3", s3_bucket_name="litellm-my-test-bucket-2", s3_region_name="us-east-1"
# )
# router = Router(
# model_list=model_list,
# fallbacks=[
# {"bad-model": ["gpt-3.5-turbo"]},
# ],
# ) # type: ignore
# async def router_acompletion():
# # embedding call
# question = f"This is a test: {uuid.uuid4()}" * 1
# response = await router.acompletion(
# model="bad-model", messages=[{"role": "user", "content": question}]
# )
# print("completion-resp", response)
# return response
# async def main():
# for i in range(1):
# start = time.time()
# n = 15 # Number of concurrent tasks
# tasks = [router_acompletion() for _ in range(n)]
# chat_completions = await asyncio.gather(*tasks)
# successful_completions = [c for c in chat_completions if c is not None]
# # Write errors to error_log.txt
# with open("error_log.txt", "a") as error_log:
# for completion in chat_completions:
# if isinstance(completion, str):
# error_log.write(completion + "\n")
# print(n, time.time() - start, len(successful_completions))
# print()
# print(vars(router))
# prev_models = router.previous_models
# print("vars in prev_models")
# print(prev_models[0].keys())
# if __name__ == "__main__":
# # Blank out contents of error_log.txt
# open("error_log.txt", "w").close()
# import tracemalloc
# tracemalloc.start(25)
# # ... run your application ...
# asyncio.run(main())
# print(mem_top())
# snapshot = tracemalloc.take_snapshot()
# # top_stats = snapshot.statistics('lineno')
# # print("[ Top 10 ]")
# # for stat in top_stats[:50]:
# # print(stat)
# top_stats = snapshot.statistics("traceback")
# # pick the biggest memory block
# stat = top_stats[0]
# print("%s memory blocks: %.1f KiB" % (stat.count, stat.size / 1024))
# for line in stat.traceback.format():
# print(line)
# print()
# stat = top_stats[1]
# print("%s memory blocks: %.1f KiB" % (stat.count, stat.size / 1024))
# for line in stat.traceback.format():
# print(line)
# print()
# stat = top_stats[2]
# print("%s memory blocks: %.1f KiB" % (stat.count, stat.size / 1024))
# for line in stat.traceback.format():
# print(line)
# print()
# stat = top_stats[3]
# print("%s memory blocks: %.1f KiB" % (stat.count, stat.size / 1024))
# for line in stat.traceback.format():
# print(line)

23
tests/local_testing/test_model_response_typing/server.py
View file

@ -1,23 +0,0 @@
# #### What this tests ####
# # This tests if the litellm model response type is returnable in a flask app
# import sys, os
# import traceback
# from flask import Flask, request, jsonify, abort, Response
# sys.path.insert(0, os.path.abspath('../../..')) # Adds the parent directory to the system path
# import litellm
# from litellm import completion
# litellm.set_verbose = False
# app = Flask(__name__)
# @app.route('/')
# def hello():
# data = request.json
# return completion(**data)
# if __name__ == '__main__':
# from waitress import serve
# serve(app, host='localhost', port=8080, threads=10)

14
tests/local_testing/test_model_response_typing/test.py
View file

@ -1,14 +0,0 @@
# import requests, json
# BASE_URL = 'http://localhost:8080'
# def test_hello_route():
# data = {"model": "claude-3-5-haiku-20241022", "messages": [{"role": "user", "content": "hey, how's it going?"}]}
# headers = {'Content-Type': 'application/json'}
# response = requests.get(BASE_URL, headers=headers, data=json.dumps(data))
# print(response.text)
# assert response.status_code == 200
# print("Hello route test passed!")
# if __name__ == '__main__':
# test_hello_route()

336
tests/local_testing/test_ollama_local.py
View file

@ -1,336 +0,0 @@
# ##### THESE TESTS CAN ONLY RUN LOCALLY WITH THE OLLAMA SERVER RUNNING ######
# # https://ollama.ai/
# import sys, os
# import traceback
# from dotenv import load_dotenv
# load_dotenv()
# import os
# sys.path.insert(0, os.path.abspath('../..')) # Adds the parent directory to the system path
# import pytest
# import litellm
# from litellm import embedding, completion
# import asyncio
# user_message = "respond in 20 words. who are you?"
# messages = [{ "content": user_message,"role": "user"}]
# async def test_ollama_aembeddings():
# litellm.set_verbose = True
# input = "The food was delicious and the waiter..."
# response = await litellm.aembedding(model="ollama/mistral", input=input)
# print(response)
# asyncio.run(test_ollama_aembeddings())
# def test_ollama_embeddings():
# litellm.set_verbose = True
# input = "The food was delicious and the waiter..."
# response = litellm.embedding(model="ollama/mistral", input=input)
# print(response)
# test_ollama_embeddings()
# def test_ollama_streaming():
# try:
# litellm.set_verbose = False
# messages = [
# {"role": "user", "content": "What is the weather like in Boston?"}
# ]
# functions = [
# {
# "name": "get_current_weather",
# "description": "Get the current weather in a given location",
# "parameters": {
# "type": "object",
# "properties": {
# "location": {
# "type": "string",
# "description": "The city and state, e.g. San Francisco, CA"
# },
# "unit": {
# "type": "string",
# "enum": ["celsius", "fahrenheit"]
# }
# },
# "required": ["location"]
# }
# }
# ]
# response = litellm.completion(model="ollama/mistral",
# messages=messages,
# functions=functions,
# stream=True)
# for chunk in response:
# print(f"CHUNK: {chunk}")
# except Exception as e:
# print(e)
# # test_ollama_streaming()
# async def test_async_ollama_streaming():
# try:
# litellm.set_verbose = False
# response = await litellm.acompletion(model="ollama/mistral-openorca",
# messages=[{"role": "user", "content": "Hey, how's it going?"}],
# stream=True)
# async for chunk in response:
# print(f"CHUNK: {chunk}")
# except Exception as e:
# print(e)
# # asyncio.run(test_async_ollama_streaming())
# def test_completion_ollama():
# try:
# litellm.set_verbose = True
# response = completion(
# model="ollama/mistral",
# messages=[{"role": "user", "content": "Hey, how's it going?"}],
# max_tokens=200,
# request_timeout = 10,
# stream=True
# )
# for chunk in response:
# print(chunk)
# print(response)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# # test_completion_ollama()
# def test_completion_ollama_function_calling():
# try:
# litellm.set_verbose = True
# messages = [
# {"role": "user", "content": "What is the weather like in Boston?"}
# ]
# functions = [
# {
# "name": "get_current_weather",
# "description": "Get the current weather in a given location",
# "parameters": {
# "type": "object",
# "properties": {
# "location": {
# "type": "string",
# "description": "The city and state, e.g. San Francisco, CA"
# },
# "unit": {
# "type": "string",
# "enum": ["celsius", "fahrenheit"]
# }
# },
# "required": ["location"]
# }
# }
# ]
# response = completion(
# model="ollama/mistral",
# messages=messages,
# functions=functions,
# max_tokens=200,
# request_timeout = 10,
# )
# for chunk in response:
# print(chunk)
# print(response)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# # test_completion_ollama_function_calling()
# async def async_test_completion_ollama_function_calling():
# try:
# litellm.set_verbose = True
# messages = [
# {"role": "user", "content": "What is the weather like in Boston?"}
# ]
# functions = [
# {
# "name": "get_current_weather",
# "description": "Get the current weather in a given location",
# "parameters": {
# "type": "object",
# "properties": {
# "location": {
# "type": "string",
# "description": "The city and state, e.g. San Francisco, CA"
# },
# "unit": {
# "type": "string",
# "enum": ["celsius", "fahrenheit"]
# }
# },
# "required": ["location"]
# }
# }
# ]
# response = await litellm.acompletion(
# model="ollama/mistral",
# messages=messages,
# functions=functions,
# max_tokens=200,
# request_timeout = 10,
# )
# print(response)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# # asyncio.run(async_test_completion_ollama_function_calling())
# def test_completion_ollama_with_api_base():
# try:
# response = completion(
# model="ollama/llama2",
# messages=messages,
# api_base="http://localhost:11434"
# )
# print(response)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# # test_completion_ollama_with_api_base()
# def test_completion_ollama_custom_prompt_template():
# user_message = "what is litellm?"
# litellm.register_prompt_template(
# model="ollama/llama2",
# roles={
# "system": {"pre_message": "System: "},
# "user": {"pre_message": "User: "},
# "assistant": {"pre_message": "Assistant: "}
# }
# )
# messages = [{ "content": user_message,"role": "user"}]
# litellm.set_verbose = True
# try:
# response = completion(
# model="ollama/llama2",
# messages=messages,
# stream=True
# )
# print(response)
# for chunk in response:
# print(chunk)
# # print(chunk['choices'][0]['delta'])
# except Exception as e:
# traceback.print_exc()
# pytest.fail(f"Error occurred: {e}")
# # test_completion_ollama_custom_prompt_template()
# async def test_completion_ollama_async_stream():
# user_message = "what is the weather"
# messages = [{ "content": user_message,"role": "user"}]
# try:
# response = await litellm.acompletion(
# model="ollama/llama2",
# messages=messages,
# api_base="http://localhost:11434",
# stream=True
# )
# async for chunk in response:
# print(chunk['choices'][0]['delta'])
# print("TEST ASYNC NON Stream")
# response = await litellm.acompletion(
# model="ollama/llama2",
# messages=messages,
# api_base="http://localhost:11434",
# )
# print(response)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# # import asyncio
# # asyncio.run(test_completion_ollama_async_stream())
# def prepare_messages_for_chat(text: str) -> list:
# messages = [
# {"role": "user", "content": text},
# ]
# return messages
# async def ask_question():
# params = {
# "messages": prepare_messages_for_chat("What is litellm? tell me 10 things about it who is sihaan.write an essay"),
# "api_base": "http://localhost:11434",
# "model": "ollama/llama2",
# "stream": True,
# }
# response = await litellm.acompletion(**params)
# return response
# async def main():
# response = await ask_question()
# async for chunk in response:
# print(chunk)
# print("test async completion without streaming")
# response = await litellm.acompletion(
# model="ollama/llama2",
# messages=prepare_messages_for_chat("What is litellm? respond in 2 words"),
# )
# print("response", response)
# def test_completion_expect_error():
# # this tests if we can exception map correctly for ollama
# print("making ollama request")
# # litellm.set_verbose=True
# user_message = "what is litellm?"
# messages = [{ "content": user_message,"role": "user"}]
# try:
# response = completion(
# model="ollama/invalid",
# messages=messages,
# stream=True
# )
# print(response)
# for chunk in response:
# print(chunk)
# # print(chunk['choices'][0]['delta'])
# except Exception as e:
# pass
# pytest.fail(f"Error occurred: {e}")
# # test_completion_expect_error()
# def test_ollama_llava():
# litellm.set_verbose=True
# # same params as gpt-4 vision
# response = completion(
# model = "ollama/llava",
# messages=[
# {
# "role": "user",
# "content": [
# {
# "type": "text",
# "text": "What is in this picture"
# },
# {
# "type": "image_url",
# "image_url": {
# "url": "iVBORw0KGgoAAAANSUhEUgAAAG0AAABmCAYAAADBPx+VAAAACXBIWXMAAAsTAAALEwEAmpwYAAAAAXNSR0IArs4c6QAAAARnQU1BAACxjwv8YQUAAA3VSURBVHgB7Z27r0zdG8fX743i1bi1ikMoFMQloXRpKFFIqI7LH4BEQ+NWIkjQuSWCRIEoULk0gsK1kCBI0IhrQVT7tz/7zZo888yz1r7MnDl7z5xvsjkzs2fP3uu71nNfa7lkAsm7d++Sffv2JbNmzUqcc8m0adOSzZs3Z+/XES4ZckAWJEGWPiCxjsQNLWmQsWjRIpMseaxcuTKpG/7HP27I8P79e7dq1ars/yL4/v27S0ejqwv+cUOGEGGpKHR37tzJCEpHV9tnT58+dXXCJDdECBE2Ojrqjh071hpNECjx4cMHVycM1Uhbv359B2F79+51586daxN/+pyRkRFXKyRDAqxEp4yMlDDzXG1NPnnyJKkThoK0VFd1ELZu3TrzXKxKfW7dMBQ6bcuWLW2v0VlHjx41z717927ba22U9APcw7Nnz1oGEPeL3m3p2mTAYYnFmMOMXybPPXv2bNIPpFZr1NHn4HMw0KRBjg9NuRw95s8PEcz/6DZELQd/09C9QGq5RsmSRybqkwHGjh07OsJSsYYm3ijPpyHzoiacg35MLdDSIS/O1yM778jOTwYUkKNHWUzUWaOsylE00MyI0fcnOwIdjvtNdW/HZwNLGg+sR1kMepSNJXmIwxBZiG8tDTpEZzKg0GItNsosY8USkxDhD0Rinuiko2gfL/RbiD2LZAjU9zKQJj8RDR0vJBR1/Phx9+PHj9Z7REF4nTZkxzX4LCXHrV271qXkBAPGfP/atWvu/PnzHe4C97F48eIsRLZ9+3a3f/9+87dwP1JxaF7/3r17ba+5l4EcaVo0lj3SBq5kGTJSQmLWMjgYNei2GPT1MuMqGTDEFHzeQSP2wi/jGnkmPJ/nhccs44jvDAxpVcxnq0F6eT8h4ni/iIWpR5lPyA6ETkNXoSukvpJAD3AsXLiwpZs49+fPn5ke4j10TqYvegSfn0OnafC+Tv9ooA/JPkgQysqQNBzagXY55nO/oa1F7qvIPWkRL12WRpMWUvpVDYmxAPehxWSe8ZEXL20sadYIozfmNch4QJPAfeJgW3rNsnzphBKNJM2KKODo1rVOMRYik5ETy3ix4qWNI81qAAirizgMIc+yhTytx0JWZuNI03qsrgWlGtwjoS9XwgUhWGyhUaRZZQNNIEwCiXD16tXcAHUs79co0vSD8rrJCIW98pzvxpAWyyo3HYwqS0+H0BjStClcZJT5coMm6D2LOF8TolGJtK9fvyZpyiC5ePFi9nc/oJU4eiEP0jVoAnHa9wyJycITMP78+eMeP37sXrx44d6+fdt6f82aNdkx1pg9e3Zb5W+RSRE+n+VjksQWifvVaTKFhn5O8my63K8Qabdv33b379/PiAP//vuvW7BggZszZ072/+TJk91YgkafPn166zXB1rQHFvouAWHq9z3SEevSUerqCn2/dDCeta2jxYbr69evk4MHDyY7d+7MjhMnTiTPnz9Pfv/+nfQT2ggpO2dMF8cghuoM7Ygj5iWCqRlGFml0QC/ftGmTmzt3rmsaKDsgBSPh0/8yPeLLBihLkOKJc0jp8H8vUzcxIA1k6QJ/c78tWEyj5P3o4u9+jywNPdJi5rAH9x0KHcl4Hg570eQp3+vHXGyrmEeigzQsQsjavXt38ujRo44LQuDDhw+TW7duRS1HGgMxhNXHgflaNTOsHyKvHK5Ijo2jbFjJBQK9YwFd6RVMzfgRBmEfP37suBBm/p49e1qjEP2mwTViNRo0VJWH1deMXcNK08uUjVUu7s/zRaL+oLNxz1bpANco4npUgX4G2eFbpDFyQoQxojBCpEGSytmOH8qrH5Q9vuzD6ofQylkCUmh8DBAr+q8JCyVNtWQIidKQE9wNtLSQnS4jDSsxNHogzFuQBw4cyM61UKVsjfr3ooBkPSqqQHesUPWVtzi9/vQi1T+rJj7WiTz4Pt/l3LxUkr5P2VYZaZ4URpsE+st/dujQoaBBYokbrz/8TJNQYLSonrPS9kUaSkPeZyj1AWSj+d+VBoy1pIWVNed8P0Ll/ee5HdGRhrHhR5GGN0r4LGZBaj8oFDJitBTJzIZgFcmU0Y8ytWMZMzJOaXUSrUs5RxKnrxmbb5YXO9VGUhtpXldhEUogFr3IzIsvlpmdosVcGVGXFWp2oU9kLFL3dEkSz6NHEY1sjSRdIuDFWEhd8KxFqsRi1uM/nz9/zpxnwlESONdg6dKlbsaMGS4EHFHtjFIDHwKOo46l4TxSuxgDzi+rE2jg+BaFruOX4HXa0Nnf1lwAPufZeF8/r6zD97WK2qFnGjBxTw5qNGPxT+5T/r7/7RawFC3j4vTp09koCxkeHjqbHJqArmH5UrFKKksnxrK7FuRIs8STfBZv+luugXZ2pR/pP9Ois4z+TiMzUUkUjD0iEi1fzX8GmXyuxUBRcaUfykV0YZnlJGKQpOiGB76x5GeWkWWJc3mOrK6S7xdND+W5N6XyaRgtWJFe13GkaZnKOsYqGdOVVVbGupsyA/l7emTLHi7vwTdirNEt0qxnzAvBFcnQF16xh/TMpUuXHDowhlA9vQVraQhkudRdzOnK+04ZSP3DUhVSP61YsaLtd/ks7ZgtPcXqPqEafHkdqa84X6aCeL7YWlv6edGFHb+ZFICPlljHhg0bKuk0CSvVznWsotRu433alNdFrqG45ejoaPCaUkWERpLXjzFL2Rpllp7PJU2a/v7Ab8N05/9t27Z16KUqoFGsxnI9EosS2niSYg9SpU6B4JgTrvVW1flt1sT+0ADIJU2maXzcUTraGCRaL1Wp9rUMk16PMom8QhruxzvZIegJjFU7LLCePfS8uaQdPny4jTTL0dbee5mYokQsXTIWNY46kuMbnt8Kmec+LGWtOVIl9cT1rCB0V8WqkjAsRwta93TbwNYoGKsUSChN44lgBNCoHLHzquYKrU6qZ8lolCIN0Rh6cP0Q3U6I6IXILYOQI513hJaSKAorFpuHXJNfVlpRtmYBk1Su1obZr5dnKAO+L10Hrj3WZW+E3qh6IszE37F6EB+68mGpvKm4eb9bFrlzrok7fvr0Kfv727dvWRmdVTJHw0qiiCUSZ6wCK+7XL/AcsgNyL74DQQ730sv78Su7+t/A36MdY0sW5o40ahslXr58aZ5HtZB8GH64m9EmMZ7FpYw4T6QnrZfgenrhFxaSiSGXtPnz57e9TkNZLvTjeqhr734CNtrK41L40sUQckmj1lGKQ0rC37x544r8eNXRpnVE3ZZY7zXo8NomiO0ZUCj2uHz58rbXoZ6gc0uA+F6ZeKS/jhRDUq8MKrTho9fEkihMmhxtBI1DxKFY9XLpVcSkfoi8JGnToZO5sU5aiDQIW716ddt7ZLYtMQlhECdBGXZZMWldY5BHm5xgAroWj4C0hbYkSc/jBmggIrXJWlZM6pSETsEPGqZOndr2uuuR5rF169a2HoHPdurUKZM4CO1WTPqaDaAd+GFGKdIQkxAn9RuEWcTRyN2KSUgiSgF5aWzPTeA/lN5rZubMmR2bE4SIC4nJoltgAV/dVefZm72AtctUCJU2CMJ327hxY9t7EHbkyJFseq+EJSY16RPo3Dkq1kkr7+q0bNmyDuLQcZBEPYmHVdOBiJyIlrRDq41YPWfXOxUysi5fvtyaj+2BpcnsUV/oSoEMOk2CQGlr4ckhBwaetBhjCwH0ZHtJROPJkyc7UjcYLDjmrH7ADTEBXFfOYmB0k9oYBOjJ8b4aOYSe7QkKcYhFlq3QYLQhSidNmtS2RATwy8YOM3EQJsUjKiaWZ+vZToUQgzhkHXudb/PW5YMHD9yZM2faPsMwoc7RciYJXbGuBqJ1UIGKKLv915jsvgtJxCZDubdXr165mzdvtr1Hz5LONA8jrUwKPqsmVesKa49S3Q4WxmRPUEYdTjgiUcfUwLx589ySJUva3oMkP6IYddq6HMS4o55xBJBUeRjzfa4Zdeg56QZ43LhxoyPo7Lf1kNt7oO8wWAbNwaYjIv5lhyS7kRf96dvm5Jah8vfvX3flyhX35cuX6HfzFHOToS1H4BenCaHvO8pr8iDuwoUL7tevX+b5ZdbBair0xkFIlFDlW4ZknEClsp/TzXyAKVOmmHWFVSbDNw1l1+4f90U6IY/q4V27dpnE9bJ+v87QEydjqx/UamVVPRG+mwkNTYN+9tjkwzEx+atCm/X9WvWtDtAb68Wy9LXa1UmvCDDIpPkyOQ5ZwSzJ4jMrvFcr0rSjOUh+GcT4LSg5ugkW1Io0/SCDQBojh0hPlaJdah+tkVYrnTZowP8iq1F1TgMBBauufyB33x1v+NWFYmT5KmppgHC+NkAgbmRkpD3yn9QIseXymoTQFGQmIOKTxiZIWpvAatenVqRVXf2nTrAWMsPnKrMZHz6bJq5jvce6QK8J1cQNgKxlJapMPdZSR64/UivS9NztpkVEdKcrs5alhhWP9NeqlfWopzhZScI6QxseegZRGeg5a8C3Re1Mfl1ScP36ddcUaMuv24iOJtz7sbUjTS4qBvKmstYJoUauiuD3k5qhyr7QdUHMeCgLa1Ear9NquemdXgmum4fvJ6w1lqsuDhNrg1qSpleJK7K3TF0Q2jSd94uSZ60kK1e3qyVpQK6PVWXp2/FC3mp6jBhKKOiY2h3gtUV64TWM6wDETRPLDfSakXmH3w8g9Jlug8ZtTt4kVF0kLUYYmCCtD/DrQ5YhMGbA9L3ucdjh0y8kOHW5gU/VEEmJTcL4Pz/f7mgoAbYkAAAAAElFTkSuQmCC"
# }
# }
# ]
# }
# ],
# )
# print("Response from ollama/llava")
# print(response)
# # test_ollama_llava()
# # PROCESSED CHUNK PRE CHUNK CREATOR

334
tests/local_testing/test_ollama_local_chat.py
View file

@ -1,334 +0,0 @@
# ##### THESE TESTS CAN ONLY RUN LOCALLY WITH THE OLLAMA SERVER RUNNING ######
# # https://ollama.ai/
# import sys, os
# import traceback
# from dotenv import load_dotenv
# load_dotenv()
# import os
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import pytest
# import litellm
# from litellm import embedding, completion
# import asyncio
# user_message = "respond in 20 words. who are you?"
# messages = [{"content": user_message, "role": "user"}]
# def test_ollama_streaming():
# try:
# litellm.set_verbose = False
# messages = [{"role": "user", "content": "What is the weather like in Boston?"}]
# functions = [
# {
# "name": "get_current_weather",
# "description": "Get the current weather in a given location",
# "parameters": {
# "type": "object",
# "properties": {
# "location": {
# "type": "string",
# "description": "The city and state, e.g. San Francisco, CA",
# },
# "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
# },
# "required": ["location"],
# },
# }
# ]
# response = litellm.completion(
# model="ollama_chat/mistral",
# messages=messages,
# functions=functions,
# stream=True,
# )
# for chunk in response:
# print(f"CHUNK: {chunk}")
# except Exception as e:
# print(e)
# # test_ollama_streaming()
# async def test_async_ollama_streaming():
# try:
# litellm.set_verbose = True
# response = await litellm.acompletion(
# model="ollama_chat/llama2",
# messages=[{"role": "user", "content": "Hey, how's it going?"}],
# stream=True,
# )
# async for chunk in response:
# print(f"CHUNK: {chunk}")
# except Exception as e:
# print(e)
# # asyncio.run(test_async_ollama_streaming())
# async def test_async_ollama():
# try:
# litellm.set_verbose = True
# response = await litellm.acompletion(
# model="ollama_chat/llama2",
# messages=[{"role": "user", "content": "Hey, how's it going?"}],
# )
# print("\n response", response)
# except Exception as e:
# print(e)
# # asyncio.run(test_async_ollama())
# def test_completion_ollama():
# try:
# litellm.set_verbose = True
# response = completion(
# model="ollama_chat/mistral",
# messages=[{"role": "user", "content": "Hey, how's it going?"}],
# max_tokens=200,
# request_timeout=10,
# stream=True,
# )
# for chunk in response:
# print(chunk)
# print(response)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# # test_completion_ollama()
# def test_completion_ollama_function_calling():
# try:
# litellm.set_verbose = True
# messages = [{"role": "user", "content": "What is the weather like in Boston?"}]
# functions = [
# {
# "name": "get_current_weather",
# "description": "Get the current weather in a given location",
# "parameters": {
# "type": "object",
# "properties": {
# "location": {
# "type": "string",
# "description": "The city and state, e.g. San Francisco, CA",
# },
# "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
# },
# "required": ["location"],
# },
# }
# ]
# response = completion(
# model="ollama_chat/mistral",
# messages=messages,
# functions=functions,
# max_tokens=200,
# request_timeout=10,
# )
# for chunk in response:
# print(chunk)
# print(response)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# test_completion_ollama_function_calling()
# async def async_test_completion_ollama_function_calling():
# try:
# litellm.set_verbose = True
# messages = [{"role": "user", "content": "What is the weather like in Boston?"}]
# functions = [
# {
# "name": "get_current_weather",
# "description": "Get the current weather in a given location",
# "parameters": {
# "type": "object",
# "properties": {
# "location": {
# "type": "string",
# "description": "The city and state, e.g. San Francisco, CA",
# },
# "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
# },
# "required": ["location"],
# },
# }
# ]
# response = await litellm.acompletion(
# model="ollama/mistral",
# messages=messages,
# functions=functions,
# max_tokens=200,
# request_timeout=10,
# )
# print(response)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# # asyncio.run(async_test_completion_ollama_function_calling())
# def test_completion_ollama_with_api_base():
# try:
# response = completion(
# model="ollama/llama2", messages=messages, api_base="http://localhost:11434"
# )
# print(response)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# # test_completion_ollama_with_api_base()
# def test_completion_ollama_custom_prompt_template():
# user_message = "what is litellm?"
# litellm.register_prompt_template(
# model="ollama/llama2",
# roles={
# "system": {"pre_message": "System: "},
# "user": {"pre_message": "User: "},
# "assistant": {"pre_message": "Assistant: "},
# },
# )
# messages = [{"content": user_message, "role": "user"}]
# litellm.set_verbose = True
# try:
# response = completion(model="ollama/llama2", messages=messages, stream=True)
# print(response)
# for chunk in response:
# print(chunk)
# # print(chunk['choices'][0]['delta'])
# except Exception as e:
# traceback.print_exc()
# pytest.fail(f"Error occurred: {e}")
# # test_completion_ollama_custom_prompt_template()
# async def test_completion_ollama_async_stream():
# user_message = "what is the weather"
# messages = [{"content": user_message, "role": "user"}]
# try:
# response = await litellm.acompletion(
# model="ollama/llama2",
# messages=messages,
# api_base="http://localhost:11434",
# stream=True,
# )
# async for chunk in response:
# print(chunk["choices"][0]["delta"])
# print("TEST ASYNC NON Stream")
# response = await litellm.acompletion(
# model="ollama/llama2",
# messages=messages,
# api_base="http://localhost:11434",
# )
# print(response)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# # import asyncio
# # asyncio.run(test_completion_ollama_async_stream())
# def prepare_messages_for_chat(text: str) -> list:
# messages = [
# {"role": "user", "content": text},
# ]
# return messages
# async def ask_question():
# params = {
# "messages": prepare_messages_for_chat(
# "What is litellm? tell me 10 things about it who is sihaan.write an essay"
# ),
# "api_base": "http://localhost:11434",
# "model": "ollama/llama2",
# "stream": True,
# }
# response = await litellm.acompletion(**params)
# return response
# async def main():
# response = await ask_question()
# async for chunk in response:
# print(chunk)
# print("test async completion without streaming")
# response = await litellm.acompletion(
# model="ollama/llama2",
# messages=prepare_messages_for_chat("What is litellm? respond in 2 words"),
# )
# print("response", response)
# def test_completion_expect_error():
# # this tests if we can exception map correctly for ollama
# print("making ollama request")
# # litellm.set_verbose=True
# user_message = "what is litellm?"
# messages = [{"content": user_message, "role": "user"}]
# try:
# response = completion(model="ollama/invalid", messages=messages, stream=True)
# print(response)
# for chunk in response:
# print(chunk)
# # print(chunk['choices'][0]['delta'])
# except Exception as e:
# pass
# pytest.fail(f"Error occurred: {e}")
# # test_completion_expect_error()
# def test_ollama_llava():
# litellm.set_verbose = True
# # same params as gpt-4 vision
# response = completion(
# model="ollama/llava",
# messages=[
# {
# "role": "user",
# "content": [
# {"type": "text", "text": "What is in this picture"},
# {
# "type": "image_url",
# "image_url": {
# "url": "iVBORw0KGgoAAAANSUhEUgAAAG0AAABmCAYAAADBPx+VAAAACXBIWXMAAAsTAAALEwEAmpwYAAAAAXNSR0IArs4c6QAAAARnQU1BAACxjwv8YQUAAA3VSURBVHgB7Z27r0zdG8fX743i1bi1ikMoFMQloXRpKFFIqI7LH4BEQ+NWIkjQuSWCRIEoULk0gsK1kCBI0IhrQVT7tz/7zZo888yz1r7MnDl7z5xvsjkzs2fP3uu71nNfa7lkAsm7d++Sffv2JbNmzUqcc8m0adOSzZs3Z+/XES4ZckAWJEGWPiCxjsQNLWmQsWjRIpMseaxcuTKpG/7HP27I8P79e7dq1ars/yL4/v27S0ejqwv+cUOGEGGpKHR37tzJCEpHV9tnT58+dXXCJDdECBE2Ojrqjh071hpNECjx4cMHVycM1Uhbv359B2F79+51586daxN/+pyRkRFXKyRDAqxEp4yMlDDzXG1NPnnyJKkThoK0VFd1ELZu3TrzXKxKfW7dMBQ6bcuWLW2v0VlHjx41z717927ba22U9APcw7Nnz1oGEPeL3m3p2mTAYYnFmMOMXybPPXv2bNIPpFZr1NHn4HMw0KRBjg9NuRw95s8PEcz/6DZELQd/09C9QGq5RsmSRybqkwHGjh07OsJSsYYm3ijPpyHzoiacg35MLdDSIS/O1yM778jOTwYUkKNHWUzUWaOsylE00MyI0fcnOwIdjvtNdW/HZwNLGg+sR1kMepSNJXmIwxBZiG8tDTpEZzKg0GItNsosY8USkxDhD0Rinuiko2gfL/RbiD2LZAjU9zKQJj8RDR0vJBR1/Phx9+PHj9Z7REF4nTZkxzX4LCXHrV271qXkBAPGfP/atWvu/PnzHe4C97F48eIsRLZ9+3a3f/9+87dwP1JxaF7/3r17ba+5l4EcaVo0lj3SBq5kGTJSQmLWMjgYNei2GPT1MuMqGTDEFHzeQSP2wi/jGnkmPJ/nhccs44jvDAxpVcxnq0F6eT8h4ni/iIWpR5lPyA6ETkNXoSukvpJAD3AsXLiwpZs49+fPn5ke4j10TqYvegSfn0OnafC+Tv9ooA/JPkgQysqQNBzagXY55nO/oa1F7qvIPWkRL12WRpMWUvpVDYmxAPehxWSe8ZEXL20sadYIozfmNch4QJPAfeJgW3rNsnzphBKNJM2KKODo1rVOMRYik5ETy3ix4qWNI81qAAirizgMIc+yhTytx0JWZuNI03qsrgWlGtwjoS9XwgUhWGyhUaRZZQNNIEwCiXD16tXcAHUs79co0vSD8rrJCIW98pzvxpAWyyo3HYwqS0+H0BjStClcZJT5coMm6D2LOF8TolGJtK9fvyZpyiC5ePFi9nc/oJU4eiEP0jVoAnHa9wyJycITMP78+eMeP37sXrx44d6+fdt6f82aNdkx1pg9e3Zb5W+RSRE+n+VjksQWifvVaTKFhn5O8my63K8Qabdv33b379/PiAP//vuvW7BggZszZ072/+TJk91YgkafPn166zXB1rQHFvouAWHq9z3SEevSUerqCn2/dDCeta2jxYbr69evk4MHDyY7d+7MjhMnTiTPnz9Pfv/+nfQT2ggpO2dMF8cghuoM7Ygj5iWCqRlGFml0QC/ftGmTmzt3rmsaKDsgBSPh0/8yPeLLBihLkOKJc0jp8H8vUzcxIA1k6QJ/c78tWEyj5P3o4u9+jywNPdJi5rAH9x0KHcl4Hg570eQp3+vHXGyrmEeigzQsQsjavXt38ujRo44LQuDDhw+TW7duRS1HGgMxhNXHgflaNTOsHyKvHK5Ijo2jbFjJBQK9YwFd6RVMzfgRBmEfP37suBBm/p49e1qjEP2mwTViNRo0VJWH1deMXcNK08uUjVUu7s/zRaL+oLNxz1bpANco4npUgX4G2eFbpDFyQoQxojBCpEGSytmOH8qrH5Q9vuzD6ofQylkCUmh8DBAr+q8JCyVNtWQIidKQE9wNtLSQnS4jDSsxNHogzFuQBw4cyM61UKVsjfr3ooBkPSqqQHesUPWVtzi9/vQi1T+rJj7WiTz4Pt/l3LxUkr5P2VYZaZ4URpsE+st/dujQoaBBYokbrz/8TJNQYLSonrPS9kUaSkPeZyj1AWSj+d+VBoy1pIWVNed8P0Ll/ee5HdGRhrHhR5GGN0r4LGZBaj8oFDJitBTJzIZgFcmU0Y8ytWMZMzJOaXUSrUs5RxKnrxmbb5YXO9VGUhtpXldhEUogFr3IzIsvlpmdosVcGVGXFWp2oU9kLFL3dEkSz6NHEY1sjSRdIuDFWEhd8KxFqsRi1uM/nz9/zpxnwlESONdg6dKlbsaMGS4EHFHtjFIDHwKOo46l4TxSuxgDzi+rE2jg+BaFruOX4HXa0Nnf1lwAPufZeF8/r6zD97WK2qFnGjBxTw5qNGPxT+5T/r7/7RawFC3j4vTp09koCxkeHjqbHJqArmH5UrFKKksnxrK7FuRIs8STfBZv+luugXZ2pR/pP9Ois4z+TiMzUUkUjD0iEi1fzX8GmXyuxUBRcaUfykV0YZnlJGKQpOiGB76x5GeWkWWJc3mOrK6S7xdND+W5N6XyaRgtWJFe13GkaZnKOsYqGdOVVVbGupsyA/l7emTLHi7vwTdirNEt0qxnzAvBFcnQF16xh/TMpUuXHDowhlA9vQVraQhkudRdzOnK+04ZSP3DUhVSP61YsaLtd/ks7ZgtPcXqPqEafHkdqa84X6aCeL7YWlv6edGFHb+ZFICPlljHhg0bKuk0CSvVznWsotRu433alNdFrqG45ejoaPCaUkWERpLXjzFL2Rpllp7PJU2a/v7Ab8N05/9t27Z16KUqoFGsxnI9EosS2niSYg9SpU6B4JgTrvVW1flt1sT+0ADIJU2maXzcUTraGCRaL1Wp9rUMk16PMom8QhruxzvZIegJjFU7LLCePfS8uaQdPny4jTTL0dbee5mYokQsXTIWNY46kuMbnt8Kmec+LGWtOVIl9cT1rCB0V8WqkjAsRwta93TbwNYoGKsUSChN44lgBNCoHLHzquYKrU6qZ8lolCIN0Rh6cP0Q3U6I6IXILYOQI513hJaSKAorFpuHXJNfVlpRtmYBk1Su1obZr5dnKAO+L10Hrj3WZW+E3qh6IszE37F6EB+68mGpvKm4eb9bFrlzrok7fvr0Kfv727dvWRmdVTJHw0qiiCUSZ6wCK+7XL/AcsgNyL74DQQ730sv78Su7+t/A36MdY0sW5o40ahslXr58aZ5HtZB8GH64m9EmMZ7FpYw4T6QnrZfgenrhFxaSiSGXtPnz57e9TkNZLvTjeqhr734CNtrK41L40sUQckmj1lGKQ0rC37x544r8eNXRpnVE3ZZY7zXo8NomiO0ZUCj2uHz58rbXoZ6gc0uA+F6ZeKS/jhRDUq8MKrTho9fEkihMmhxtBI1DxKFY9XLpVcSkfoi8JGnToZO5sU5aiDQIW716ddt7ZLYtMQlhECdBGXZZMWldY5BHm5xgAroWj4C0hbYkSc/jBmggIrXJWlZM6pSETsEPGqZOndr2uuuR5rF169a2HoHPdurUKZM4CO1WTPqaDaAd+GFGKdIQkxAn9RuEWcTRyN2KSUgiSgF5aWzPTeA/lN5rZubMmR2bE4SIC4nJoltgAV/dVefZm72AtctUCJU2CMJ327hxY9t7EHbkyJFseq+EJSY16RPo3Dkq1kkr7+q0bNmyDuLQcZBEPYmHVdOBiJyIlrRDq41YPWfXOxUysi5fvtyaj+2BpcnsUV/oSoEMOk2CQGlr4ckhBwaetBhjCwH0ZHtJROPJkyc7UjcYLDjmrH7ADTEBXFfOYmB0k9oYBOjJ8b4aOYSe7QkKcYhFlq3QYLQhSidNmtS2RATwy8YOM3EQJsUjKiaWZ+vZToUQgzhkHXudb/PW5YMHD9yZM2faPsMwoc7RciYJXbGuBqJ1UIGKKLv915jsvgtJxCZDubdXr165mzdvtr1Hz5LONA8jrUwKPqsmVesKa49S3Q4WxmRPUEYdTjgiUcfUwLx589ySJUva3oMkP6IYddq6HMS4o55xBJBUeRjzfa4Zdeg56QZ43LhxoyPo7Lf1kNt7oO8wWAbNwaYjIv5lhyS7kRf96dvm5Jah8vfvX3flyhX35cuX6HfzFHOToS1H4BenCaHvO8pr8iDuwoUL7tevX+b5ZdbBair0xkFIlFDlW4ZknEClsp/TzXyAKVOmmHWFVSbDNw1l1+4f90U6IY/q4V27dpnE9bJ+v87QEydjqx/UamVVPRG+mwkNTYN+9tjkwzEx+atCm/X9WvWtDtAb68Wy9LXa1UmvCDDIpPkyOQ5ZwSzJ4jMrvFcr0rSjOUh+GcT4LSg5ugkW1Io0/SCDQBojh0hPlaJdah+tkVYrnTZowP8iq1F1TgMBBauufyB33x1v+NWFYmT5KmppgHC+NkAgbmRkpD3yn9QIseXymoTQFGQmIOKTxiZIWpvAatenVqRVXf2nTrAWMsPnKrMZHz6bJq5jvce6QK8J1cQNgKxlJapMPdZSR64/UivS9NztpkVEdKcrs5alhhWP9NeqlfWopzhZScI6QxseegZRGeg5a8C3Re1Mfl1ScP36ddcUaMuv24iOJtz7sbUjTS4qBvKmstYJoUauiuD3k5qhyr7QdUHMeCgLa1Ear9NquemdXgmum4fvJ6w1lqsuDhNrg1qSpleJK7K3TF0Q2jSd94uSZ60kK1e3qyVpQK6PVWXp2/FC3mp6jBhKKOiY2h3gtUV64TWM6wDETRPLDfSakXmH3w8g9Jlug8ZtTt4kVF0kLUYYmCCtD/DrQ5YhMGbA9L3ucdjh0y8kOHW5gU/VEEmJTcL4Pz/f7mgoAbYkAAAAAElFTkSuQmCC"
# },
# },
# ],
# }
# ],
# )
# print("Response from ollama/llava")
# print(response)
# # test_ollama_llava()
# # PROCESSED CHUNK PRE CHUNK CREATOR

89
tests/local_testing/test_provider_specific_config.py
View file

@ -12,36 +12,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
import litellm
from litellm import RateLimitError, completion
# Huggingface - Expensive to deploy models and keep them running. Maybe we can try doing this via baseten??
# def hf_test_completion_tgi():
# litellm.HuggingfaceConfig(max_new_tokens=200)
# litellm.set_verbose=True
# try:
# # OVERRIDE WITH DYNAMIC MAX TOKENS
# response_1 = litellm.completion(
# model="huggingface/mistralai/Mistral-7B-Instruct-v0.1",
# messages=[{ "content": "Hello, how are you?","role": "user"}],
# api_base="https://n9ox93a8sv5ihsow.us-east-1.aws.endpoints.huggingface.cloud",
# max_tokens=10
# )
# # Add any assertions here to check the response
# print(response_1)
# response_1_text = response_1.choices[0].message.content
# # USE CONFIG TOKENS
# response_2 = litellm.completion(
# model="huggingface/mistralai/Mistral-7B-Instruct-v0.1",
# messages=[{ "content": "Hello, how are you?","role": "user"}],
# api_base="https://n9ox93a8sv5ihsow.us-east-1.aws.endpoints.huggingface.cloud",
# )
# # Add any assertions here to check the response
# print(response_2)
# response_2_text = response_2.choices[0].message.content
# assert len(response_2_text) > len(response_1_text)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# hf_test_completion_tgi()
# Anthropic
@ -322,65 +292,6 @@ def aleph_alpha_test_completion():
# aleph_alpha_test_completion()
# Petals - calls are too slow, will cause circle ci to fail due to delay. Test locally.
# def petals_completion():
# litellm.PetalsConfig(max_new_tokens=10)
# # litellm.set_verbose=True
# try:
# # OVERRIDE WITH DYNAMIC MAX TOKENS
# response_1 = litellm.completion(
# model="petals/petals-team/StableBeluga2",
# messages=[{ "content": "Hello, how are you? Be as verbose as possible","role": "user"}],
# api_base="https://chat.petals.dev/api/v1/generate",
# max_tokens=100
# )
# response_1_text = response_1.choices[0].message.content
# print(f"response_1_text: {response_1_text}")
# # USE CONFIG TOKENS
# response_2 = litellm.completion(
# model="petals/petals-team/StableBeluga2",
# api_base="https://chat.petals.dev/api/v1/generate",
# messages=[{ "content": "Hello, how are you? Be as verbose as possible","role": "user"}],
# )
# response_2_text = response_2.choices[0].message.content
# print(f"response_2_text: {response_2_text}")
# assert len(response_2_text) < len(response_1_text)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# petals_completion()
# VertexAI
# We don't have vertex ai configured for circle ci yet -- need to figure this out.
# def vertex_ai_test_completion():
# litellm.VertexAIConfig(max_output_tokens=10)
# # litellm.set_verbose=True
# try:
# # OVERRIDE WITH DYNAMIC MAX TOKENS
# response_1 = litellm.completion(
# model="chat-bison",
# messages=[{ "content": "Hello, how are you? Be as verbose as possible","role": "user"}],
# max_tokens=100
# )
# response_1_text = response_1.choices[0].message.content
# print(f"response_1_text: {response_1_text}")
# # USE CONFIG TOKENS
# response_2 = litellm.completion(
# model="chat-bison",
# messages=[{ "content": "Hello, how are you? Be as verbose as possible","role": "user"}],
# )
# response_2_text = response_2.choices[0].message.content
# print(f"response_2_text: {response_2_text}")
# assert len(response_2_text) < len(response_1_text)
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# vertex_ai_test_completion()
# Sagemaker

267
tests/local_testing/test_streaming.py
View file

@ -203,38 +203,6 @@ tools_schema = [
}
]
# def test_completion_cohere_stream():
# # this is a flaky test due to the cohere API endpoint being unstable
# try:
# messages = [
# {"role": "system", "content": "You are a helpful assistant."},
# {
# "role": "user",
# "content": "how does a court case get to the Supreme Court?",
# },
# ]
# response = completion(
# model="command-nightly", messages=messages, stream=True, max_tokens=50,
# )
# complete_response = ""
# # Add any assertions here to check the response
# has_finish_reason = False
# for idx, chunk in enumerate(response):
# chunk, finished = streaming_format_tests(idx, chunk)
# has_finish_reason = finished
# if finished:
# break
# complete_response += chunk
# if has_finish_reason is False:
# raise Exception("Finish reason not in final chunk")
# if complete_response.strip() == "":
# raise Exception("Empty response received")
# print(f"completion_response: {complete_response}")
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# test_completion_cohere_stream()
def test_completion_azure_stream_special_char():
litellm.set_verbose = True
@ -466,9 +434,6 @@ def test_completion_azure_stream():
pytest.fail(f"Error occurred: {e}")
# test_completion_azure_stream()
def test_completion_azure_function_calling_stream():
try:
litellm.set_verbose = False
@ -491,9 +456,6 @@ def test_completion_azure_function_calling_stream():
pytest.fail(f"Error occurred: {e}")
# test_completion_azure_function_calling_stream()
@pytest.mark.skip("Flaky ollama test - needs to be fixed")
def test_completion_ollama_hosted_stream():
try:
@ -525,9 +487,6 @@ def test_completion_ollama_hosted_stream():
pytest.fail(f"Error occurred: {e}")
# test_completion_ollama_hosted_stream()
@pytest.mark.parametrize(
"model",
[
@ -658,7 +617,6 @@ async def test_completion_gemini_stream(sync_mode):
pytest.fail(f"Error occurred: {e}")
# asyncio.run(test_acompletion_gemini_stream())
def gemini_mock_post_streaming(url, **kwargs):
# This generator simulates the streaming response with partial JSON content
def stream_response():
@ -856,9 +814,6 @@ def test_completion_mistral_api_mistral_large_function_call_with_streaming():
pytest.fail(f"Error occurred: {e}")
# test_completion_mistral_api_stream()
@pytest.mark.skip()
def test_completion_nlp_cloud_stream():
try:
@ -892,9 +847,6 @@ def test_completion_nlp_cloud_stream():
pytest.fail(f"Error occurred: {e}")
# test_completion_nlp_cloud_stream()
def test_completion_claude_stream_bad_key():
try:
litellm.cache = None
@ -935,10 +887,6 @@ def test_completion_claude_stream_bad_key():
pytest.fail(f"Error occurred: {e}")
# test_completion_claude_stream_bad_key()
# test_completion_replicate_stream()
@pytest.mark.parametrize("provider", ["vertex_ai_beta"]) # ""
def test_vertex_ai_stream(provider):
from test_amazing_vertex_completion import (
@ -997,78 +945,6 @@ def test_vertex_ai_stream(provider):
pytest.fail(f"Error occurred: {e}")
# def test_completion_vertexai_stream():
# try:
# import os
# os.environ["VERTEXAI_PROJECT"] = "pathrise-convert-1606954137718"
# os.environ["VERTEXAI_LOCATION"] = "us-central1"
# messages = [
# {"role": "system", "content": "You are a helpful assistant."},
# {
# "role": "user",
# "content": "how does a court case get to the Supreme Court?",
# },
# ]
# response = completion(
# model="vertex_ai/chat-bison", messages=messages, stream=True, max_tokens=50
# )
# complete_response = ""
# has_finish_reason = False
# # Add any assertions here to check the response
# for idx, chunk in enumerate(response):
# chunk, finished = streaming_format_tests(idx, chunk)
# has_finish_reason = finished
# if finished:
# break
# complete_response += chunk
# if has_finish_reason is False:
# raise Exception("finish reason not set for last chunk")
# if complete_response.strip() == "":
# raise Exception("Empty response received")
# print(f"completion_response: {complete_response}")
# except InvalidRequestError as e:
# pass
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# test_completion_vertexai_stream()
# def test_completion_vertexai_stream_bad_key():
# try:
# import os
# messages = [
# {"role": "system", "content": "You are a helpful assistant."},
# {
# "role": "user",
# "content": "how does a court case get to the Supreme Court?",
# },
# ]
# response = completion(
# model="vertex_ai/chat-bison", messages=messages, stream=True, max_tokens=50
# )
# complete_response = ""
# has_finish_reason = False
# # Add any assertions here to check the response
# for idx, chunk in enumerate(response):
# chunk, finished = streaming_format_tests(idx, chunk)
# has_finish_reason = finished
# if finished:
# break
# complete_response += chunk
# if has_finish_reason is False:
# raise Exception("finish reason not set for last chunk")
# if complete_response.strip() == "":
# raise Exception("Empty response received")
# print(f"completion_response: {complete_response}")
# except InvalidRequestError as e:
# pass
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# test_completion_vertexai_stream_bad_key()
@pytest.mark.skip(reason="Replicate extremely flaky.")
@pytest.mark.parametrize("sync_mode", [False, True])
@pytest.mark.asyncio
@ -1130,39 +1006,6 @@ async def test_completion_replicate_llama3_streaming(sync_mode):
pytest.fail(f"Error occurred: {e}")
# TEMP Commented out - replicate throwing an auth error
# try:
# litellm.set_verbose = True
# messages = [
# {"role": "system", "content": "You are a helpful assistant."},
# {
# "role": "user",
# "content": "how does a court case get to the Supreme Court?",
# },
# ]
# response = completion(
# model="replicate/meta/llama-2-70b-chat:02e509c789964a7ea8736978a43525956ef40397be9033abf9fd2badfe68c9e3", messages=messages, stream=True, max_tokens=50
# )
# complete_response = ""
# has_finish_reason = False
# # Add any assertions here to check the response
# for idx, chunk in enumerate(response):
# chunk, finished = streaming_format_tests(idx, chunk)
# has_finish_reason = finished
# if finished:
# break
# complete_response += chunk
# if has_finish_reason is False:
# raise Exception("finish reason not set for last chunk")
# if complete_response.strip() == "":
# raise Exception("Empty response received")
# print(f"completion_response: {complete_response}")
# except InvalidRequestError as e:
# pass
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
@pytest.mark.parametrize("sync_mode", [True, False]) #
@pytest.mark.parametrize(
"model, region",
@ -1393,11 +1236,6 @@ def test_completion_replicate_stream_bad_key():
pytest.fail(f"Error occurred: {e}")
# test_completion_replicate_stream_bad_key()
# test_completion_bedrock_claude_stream()
@pytest.mark.skip(reason="model end of life")
def test_completion_bedrock_ai21_stream():
try:
@ -1436,9 +1274,6 @@ def test_completion_bedrock_ai21_stream():
pytest.fail(f"Error occurred: {e}")
# test_completion_bedrock_ai21_stream()
def test_completion_bedrock_mistral_stream():
try:
litellm.set_verbose = False
@ -1534,12 +1369,6 @@ def test_sagemaker_weird_response():
pytest.fail(f"An exception occurred - {str(e)}")
# test_sagemaker_weird_response()
# asyncio.run(test_sagemaker_streaming_async())
@pytest.mark.skip(reason="Account deleted by IBM.")
@pytest.mark.asyncio
async def test_completion_watsonx_stream():
@ -1576,32 +1405,6 @@ async def test_completion_watsonx_stream():
pytest.fail(f"Error occurred: {e}")
# test_completion_sagemaker_stream()
# def test_maritalk_streaming():
# messages = [{"role": "user", "content": "Hey"}]
# try:
# response = completion("maritalk", messages=messages, stream=True)
# complete_response = ""
# start_time = time.time()
# for idx, chunk in enumerate(response):
# chunk, finished = streaming_format_tests(idx, chunk)
# complete_response += chunk
# if finished:
# break
# if complete_response.strip() == "":
# raise Exception("Empty response received")
# except Exception:
# pytest.fail(f"error occurred: {traceback.format_exc()}")
# ai21_completion_call()
# ai21_completion_call_bad_key()
@pytest.mark.skip(reason="flaky test")
@pytest.mark.asyncio
async def test_hf_completion_tgi_stream():
@ -1629,60 +1432,6 @@ async def test_hf_completion_tgi_stream():
pytest.fail(f"Error occurred: {e}")
# hf_test_completion_tgi_stream()
# def test_completion_aleph_alpha():
# try:
# response = completion(
# model="luminous-base", messages=messages, stream=True
# )
# # Add any assertions here to check the response
# has_finished = False
# complete_response = ""
# start_time = time.time()
# for idx, chunk in enumerate(response):
# chunk, finished = streaming_format_tests(idx, chunk)
# has_finished = finished
# complete_response += chunk
# if finished:
# break
# if has_finished is False:
# raise Exception("finished reason missing from final chunk")
# if complete_response.strip() == "":
# raise Exception("Empty response received")
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# # test_completion_aleph_alpha()
# def test_completion_aleph_alpha_bad_key():
# try:
# api_key = "bad-key"
# response = completion(
# model="luminous-base", messages=messages, stream=True, api_key=api_key
# )
# # Add any assertions here to check the response
# has_finished = False
# complete_response = ""
# start_time = time.time()
# for idx, chunk in enumerate(response):
# chunk, finished = streaming_format_tests(idx, chunk)
# has_finished = finished
# complete_response += chunk
# if finished:
# break
# if has_finished is False:
# raise Exception("finished reason missing from final chunk")
# if complete_response.strip() == "":
# raise Exception("Empty response received")
# except InvalidRequestError as e:
# pass
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# test_completion_aleph_alpha_bad_key()
# test on openai completion call
def test_openai_chat_completion_call():
litellm.set_verbose = False
@ -1710,9 +1459,6 @@ def test_openai_chat_completion_call():
print(f"complete response: {complete_response}")
# test_openai_chat_completion_call()
def test_openai_chat_completion_complete_response_call():
try:
complete_response = completion(
@ -1727,7 +1473,6 @@ def test_openai_chat_completion_complete_response_call():
pass
# test_openai_chat_completion_complete_response_call()
@pytest.mark.parametrize(
"model",
[
@ -1865,9 +1610,6 @@ def test_openai_text_completion_call():
pass
# test_openai_text_completion_call()
# # test on together ai completion call - starcoder
def test_together_ai_completion_call_mistral():
try:
@ -1931,7 +1673,6 @@ def test_together_ai_completion_call_starcoder_bad_key():
pass
# test_together_ai_completion_call_starcoder_bad_key()
#### Test Function calling + streaming ####
@ -1973,7 +1714,6 @@ def test_completion_openai_with_functions():
pytest.fail(f"Error occurred: {e}")
# test_completion_openai_with_functions()
#### Test Async streaming ####
@ -2005,8 +1745,6 @@ async def completion_call():
pass
# asyncio.run(completion_call())
#### Test Function Calling + Streaming ####
final_openai_function_call_example = {
@ -2310,9 +2048,6 @@ def test_streaming_and_function_calling(model):
raise e
# test_azure_streaming_and_function_calling()
def test_success_callback_streaming():
def success_callback(kwargs, completion_response, start_time, end_time):
print(
@ -2341,8 +2076,6 @@ def test_success_callback_streaming():
print(chunk["choices"][0])
# test_success_callback_streaming()
from typing import List, Optional
#### STREAMING + FUNCTION CALLING ###

63
tests/proxy_unit_tests/test_deployed_proxy_keygen.py
View file

@ -1,63 +0,0 @@
# import sys, os, time
# import traceback
# from dotenv import load_dotenv
# load_dotenv()
# import os, io
# # this file is to test litellm/proxy
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import pytest, logging, requests
# import litellm
# from litellm import embedding, completion, completion_cost, Timeout
# from litellm import RateLimitError
# def test_add_new_key():
# max_retries = 3
# retry_delay = 1 # seconds
# for retry in range(max_retries + 1):
# try:
# # Your test data
# test_data = {
# "models": ["gpt-3.5-turbo", "gpt-4", "claude-2", "azure-model"],
# "aliases": {"mistral-7b": "gpt-3.5-turbo"},
# "duration": "20m",
# }
# print("testing proxy server")
# # Your bearer token
# token = os.getenv("PROXY_MASTER_KEY")
# headers = {"Authorization": f"Bearer {token}"}
# staging_endpoint = "https://litellm-litellm-pr-1376.up.railway.app"
# main_endpoint = "https://litellm-staging.up.railway.app"
# # Make a request to the staging endpoint
# response = requests.post(
# main_endpoint + "/key/generate", json=test_data, headers=headers
# )
# print(f"response: {response.text}")
# if response.status_code == 200:
# result = response.json()
# break # Successful response, exit the loop
# elif response.status_code == 503 and retry < max_retries:
# print(
# f"Retrying in {retry_delay} seconds... (Retry {retry + 1}/{max_retries})"
# )
# time.sleep(retry_delay)
# else:
# assert False, f"Unexpected response status code: {response.status_code}"
# except Exception as e:
# print(traceback.format_exc())
# pytest.fail(f"An error occurred {e}")
# test_add_new_key()

23
tests/proxy_unit_tests/test_model_response_typing/server.py
View file

@ -1,23 +0,0 @@
# #### What this tests ####
# # This tests if the litellm model response type is returnable in a flask app
# import sys, os
# import traceback
# from flask import Flask, request, jsonify, abort, Response
# sys.path.insert(0, os.path.abspath('../../..')) # Adds the parent directory to the system path
# import litellm
# from litellm import completion
# litellm.set_verbose = False
# app = Flask(__name__)
# @app.route('/')
# def hello():
# data = request.json
# return completion(**data)
# if __name__ == '__main__':
# from waitress import serve
# serve(app, host='localhost', port=8080, threads=10)

14
tests/proxy_unit_tests/test_model_response_typing/test.py
View file

@ -1,14 +0,0 @@
# import requests, json
# BASE_URL = 'http://localhost:8080'
# def test_hello_route():
# data = {"model": "claude-3-5-haiku-20241022", "messages": [{"role": "user", "content": "hey, how's it going?"}]}
# headers = {'Content-Type': 'application/json'}
# response = requests.get(BASE_URL, headers=headers, data=json.dumps(data))
# print(response.text)
# assert response.status_code == 200
# print("Hello route test passed!")
# if __name__ == '__main__':
# test_hello_route()

61
tests/proxy_unit_tests/test_proxy_gunicorn.py
View file

@ -1,61 +0,0 @@
# #### What this tests ####
# # Allow the user to easily run the local proxy server with Gunicorn
# # LOCAL TESTING ONLY
# import sys, os, subprocess
# import traceback
# from dotenv import load_dotenv
# load_dotenv()
# import os, io
# # this file is to test litellm/proxy
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import pytest
# import litellm
# ### LOCAL Proxy Server INIT ###
# from litellm.proxy.proxy_server import save_worker_config # Replace with the actual module where your FastAPI router is defined
# filepath = os.path.dirname(os.path.abspath(__file__))
# config_fp = f"{filepath}/test_configs/test_config_custom_auth.yaml"
# def get_openai_info():
# return {
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_base": os.getenv("AZURE_API_BASE"),
# }
# def run_server(host="0.0.0.0",port=8008,num_workers=None):
# if num_workers is None:
# # Set it to min(8,cpu_count())
# import multiprocessing
# num_workers = min(4,multiprocessing.cpu_count())
# ### LOAD KEYS ###
# # Load the Azure keys. For now get them from openai-usage
# azure_info = get_openai_info()
# print(f"Azure info:{azure_info}")
# os.environ["AZURE_API_KEY"] = azure_info['api_key']
# os.environ["AZURE_API_BASE"] = azure_info['api_base']
# os.environ["AZURE_API_VERSION"] = "2023-09-01-preview"
# ### SAVE CONFIG ###
# os.environ["WORKER_CONFIG"] = config_fp
# # In order for the app to behave well with signals, run it with gunicorn
# # The first argument must be the "name of the command run"
# cmd = f"gunicorn litellm.proxy.proxy_server:app --workers {num_workers} --worker-class uvicorn.workers.UvicornWorker --bind {host}:{port}"
# cmd = cmd.split()
# print(f"Running command: {cmd}")
# import sys
# sys.stdout.flush()
# sys.stderr.flush()
# # Make sure to propage env variables
# subprocess.run(cmd) # This line actually starts Gunicorn
# if __name__ == "__main__":
# run_server()

269
tests/proxy_unit_tests/test_proxy_server_keys.py
View file

@ -1,269 +0,0 @@
# import sys, os, time
# import traceback
# from dotenv import load_dotenv
# load_dotenv()
# import os, io
# # this file is to test litellm/proxy
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import pytest, logging
# import litellm
# from litellm import embedding, completion, completion_cost, Timeout
# from litellm import RateLimitError
# import sys, os, time
# import traceback
# from dotenv import load_dotenv
# load_dotenv()
# import os, io
# # this file is to test litellm/proxy
# from concurrent.futures import ThreadPoolExecutor
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import pytest, logging, requests
# import litellm
# from litellm import embedding, completion, completion_cost, Timeout
# from litellm import RateLimitError
# from github import Github
# import subprocess
# # Function to execute a command and return the output
# def run_command(command):
# process = subprocess.Popen(command, stdout=subprocess.PIPE, shell=True)
# output, _ = process.communicate()
# return output.decode().strip()
# # Retrieve the current branch name
# branch_name = run_command("git rev-parse --abbrev-ref HEAD")
# # GitHub personal access token (with repo scope) or use username and password
# access_token = os.getenv("GITHUB_ACCESS_TOKEN")
# # Instantiate the PyGithub library's Github object
# g = Github(access_token)
# # Provide the owner and name of the repository where the pull request is located
# repository_owner = "BerriAI"
# repository_name = "litellm"
# # Get the repository object
# repo = g.get_repo(f"{repository_owner}/{repository_name}")
# # Iterate through the pull requests to find the one related to your branch
# for pr in repo.get_pulls():
# print(f"in here! {pr.head.ref}")
# if pr.head.ref == branch_name:
# pr_number = pr.number
# break
# print(f"The pull request number for branch {branch_name} is: {pr_number}")
# def test_add_new_key():
# max_retries = 3
# retry_delay = 10 # seconds
# for retry in range(max_retries + 1):
# try:
# # Your test data
# test_data = {
# "models": ["gpt-3.5-turbo", "gpt-4", "claude-2", "azure-model"],
# "aliases": {"mistral-7b": "gpt-3.5-turbo"},
# "duration": "20m",
# }
# print("testing proxy server")
# # Your bearer token
# token = os.getenv("PROXY_MASTER_KEY")
# headers = {"Authorization": f"Bearer {token}"}
# endpoint = f"https://litellm-litellm-pr-{pr_number}.up.railway.app"
# # Make a request to the staging endpoint
# response = requests.post(
# endpoint + "/key/generate", json=test_data, headers=headers
# )
# print(f"response: {response.text}")
# if response.status_code == 200:
# result = response.json()
# break # Successful response, exit the loop
# elif response.status_code == 503 and retry < max_retries:
# print(
# f"Retrying in {retry_delay} seconds... (Retry {retry + 1}/{max_retries})"
# )
# time.sleep(retry_delay)
# else:
# assert False, f"Unexpected response status code: {response.status_code}"
# except Exception as e:
# print(traceback.format_exc())
# pytest.fail(f"An error occurred {e}")
# def test_update_new_key():
# try:
# # Your test data
# test_data = {
# "models": ["gpt-3.5-turbo", "gpt-4", "claude-2", "azure-model"],
# "aliases": {"mistral-7b": "gpt-3.5-turbo"},
# "duration": "20m",
# }
# print("testing proxy server")
# # Your bearer token
# token = os.getenv("PROXY_MASTER_KEY")
# headers = {"Authorization": f"Bearer {token}"}
# endpoint = f"https://litellm-litellm-pr-{pr_number}.up.railway.app"
# # Make a request to the staging endpoint
# response = requests.post(
# endpoint + "/key/generate", json=test_data, headers=headers
# )
# assert response.status_code == 200
# result = response.json()
# assert result["key"].startswith("sk-")
# def _post_data():
# json_data = {"models": ["bedrock-models"], "key": result["key"]}
# response = requests.post(
# endpoint + "/key/generate", json=json_data, headers=headers
# )
# print(f"response text: {response.text}")
# assert response.status_code == 200
# return response
# _post_data()
# print(f"Received response: {result}")
# except Exception as e:
# pytest.fail(f"LiteLLM Proxy test failed. Exception: {str(e)}")
# def test_add_new_key_max_parallel_limit():
# try:
# # Your test data
# test_data = {"duration": "20m", "max_parallel_requests": 1}
# # Your bearer token
# token = os.getenv("PROXY_MASTER_KEY")
# headers = {"Authorization": f"Bearer {token}"}
# endpoint = f"https://litellm-litellm-pr-{pr_number}.up.railway.app"
# print(f"endpoint: {endpoint}")
# # Make a request to the staging endpoint
# response = requests.post(
# endpoint + "/key/generate", json=test_data, headers=headers
# )
# assert response.status_code == 200
# result = response.json()
# # load endpoint with model
# model_data = {
# "model_name": "azure-model",
# "litellm_params": {
# "model": "azure/gpt-4.1-mini",
# "api_key": os.getenv("AZURE_API_KEY"),
# "api_base": os.getenv("AZURE_API_BASE"),
# "api_version": os.getenv("AZURE_API_VERSION")
# }
# }
# response = requests.post(endpoint + "/model/new", json=model_data, headers=headers)
# assert response.status_code == 200
# print(f"response text: {response.text}")
# def _post_data():
# json_data = {
# "model": "azure-model",
# "messages": [
# {
# "role": "user",
# "content": f"this is a test request, write a short poem {time.time()}",
# }
# ],
# }
# # Your bearer token
# response = requests.post(
# endpoint + "/chat/completions", json=json_data, headers={"Authorization": f"Bearer {result['key']}"}
# )
# return response
# def _run_in_parallel():
# with ThreadPoolExecutor(max_workers=2) as executor:
# future1 = executor.submit(_post_data)
# future2 = executor.submit(_post_data)
# # Obtain the results from the futures
# response1 = future1.result()
# print(f"response1 text: {response1.text}")
# response2 = future2.result()
# print(f"response2 text: {response2.text}")
# if response1.status_code == 429 or response2.status_code == 429:
# pass
# else:
# raise Exception()
# _run_in_parallel()
# except Exception as e:
# pytest.fail(f"LiteLLM Proxy test failed. Exception: {str(e)}")
# def test_add_new_key_max_parallel_limit_streaming():
# try:
# # Your test data
# test_data = {"duration": "20m", "max_parallel_requests": 1}
# # Your bearer token
# token = os.getenv("PROXY_MASTER_KEY")
# headers = {"Authorization": f"Bearer {token}"}
# endpoint = f"https://litellm-litellm-pr-{pr_number}.up.railway.app"
# # Make a request to the staging endpoint
# response = requests.post(
# endpoint + "/key/generate", json=test_data, headers=headers
# )
# print(f"response: {response.text}")
# assert response.status_code == 200
# result = response.json()
# def _post_data():
# json_data = {
# "model": "azure-model",
# "messages": [
# {
# "role": "user",
# "content": f"this is a test request, write a short poem {time.time()}",
# }
# ],
# "stream": True,
# }
# response = requests.post(
# endpoint + "/chat/completions", json=json_data, headers={"Authorization": f"Bearer {result['key']}"}
# )
# return response
# def _run_in_parallel():
# with ThreadPoolExecutor(max_workers=2) as executor:
# future1 = executor.submit(_post_data)
# future2 = executor.submit(_post_data)
# # Obtain the results from the futures
# response1 = future1.result()
# response2 = future2.result()
# if response1.status_code == 429 or response2.status_code == 429:
# pass
# else:
# raise Exception()
# _run_in_parallel()
# except Exception as e:
# pytest.fail(f"LiteLLM Proxy test failed. Exception: {str(e)}")

82
tests/proxy_unit_tests/test_proxy_server_spend.py
View file

@ -1,82 +0,0 @@
# import openai, json, time, asyncio
# client = openai.AsyncOpenAI(
# api_key="sk-1234",
# base_url="http://0.0.0.0:8000"
# )
# super_fake_messages = [
# {
# "role": "user",
# "content": f"What's the weather like in San Francisco, Tokyo, and Paris? {time.time()}"
# },
# {
# "content": None,
# "role": "assistant",
# "tool_calls": [
# {
# "id": "1",
# "function": {
# "arguments": "{\"location\": \"San Francisco\", \"unit\": \"celsius\"}",
# "name": "get_current_weather"
# },
# "type": "function"
# },
# {
# "id": "2",
# "function": {
# "arguments": "{\"location\": \"Tokyo\", \"unit\": \"celsius\"}",
# "name": "get_current_weather"
# },
# "type": "function"
# },
# {
# "id": "3",
# "function": {
# "arguments": "{\"location\": \"Paris\", \"unit\": \"celsius\"}",
# "name": "get_current_weather"
# },
# "type": "function"
# }
# ]
# },
# {
# "tool_call_id": "1",
# "role": "tool",
# "name": "get_current_weather",
# "content": "{\"location\": \"San Francisco\", \"temperature\": \"90\", \"unit\": \"celsius\"}"
# },
# {
# "tool_call_id": "2",
# "role": "tool",
# "name": "get_current_weather",
# "content": "{\"location\": \"Tokyo\", \"temperature\": \"30\", \"unit\": \"celsius\"}"
# },
# {
# "tool_call_id": "3",
# "role": "tool",
# "name": "get_current_weather",
# "content": "{\"location\": \"Paris\", \"temperature\": \"50\", \"unit\": \"celsius\"}"
# }
# ]
# async def chat_completions():
# super_fake_response = await client.chat.completions.create(
# model="gpt-3.5-turbo",
# messages=super_fake_messages,
# seed=1337,
# stream=False
# ) # get a new response from the model where it can see the function response
# await asyncio.sleep(1)
# return super_fake_response
# async def loadtest_fn(n = 1):
# global num_task_cancelled_errors, exception_counts, chat_completions
# start = time.time()
# tasks = [chat_completions() for _ in range(n)]
# chat_completions = await asyncio.gather(*tasks)
# successful_completions = [c for c in chat_completions if c is not None]
# print(n, time.time() - start, len(successful_completions))
# # print(json.dumps(super_fake_response.model_dump(), indent=4))
# asyncio.run(loadtest_fn())

20
tests/search_tests/test_google_pse_search.py
View file

@ -1,20 +0,0 @@
"""
Tests for Google Programmable Search Engine (PSE) API integration.
"""
import pytest
from tests.search_tests.base_search_unit_tests import BaseSearchTest
# class TestGooglePSESearch(BaseSearchTest):
# """
# Tests for Google PSE Search functionality.
# """
# def get_search_provider(self) -> str:
# """
# Return search_provider for Google PSE Search.
# """
# return "google_pse"

2
tests/test_litellm/batches/test_batch_utils.py
View file

@ -1670,8 +1670,6 @@ async def test_handle_completed_bedrock_batch_prices_from_deployment_model(monke
)
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (1800, 1000, 2800)
# 3e-06 / 1.5e-05 on-demand, halved for batch.
assert result.cost == pytest.approx(1800 * 3e-06 / 2 + 1000 * 1.5e-05 / 2)
# The response model alone cannot price a bedrock batch: this is the $0 bug.
zero_result = await bu._handle_completed_batch(

3
tests/test_litellm/containers/test_container_transformation.py
View file

@ -377,8 +377,6 @@ class TestOpenAIContainerTransformation:
in container._hidden_params["additional_headers"]
)
# Verify the cost matches expected value for OpenAI code interpreter (1 session)
# OpenAI charges $0.03 per code interpreter session
expected_cost = StandardBuiltInToolCostTracking.get_cost_for_code_interpreter(
sessions=1, provider="openai"
)
@ -387,4 +385,3 @@ class TestOpenAIContainerTransformation:
]
assert actual_cost == expected_cost
assert actual_cost == 0.03 # OpenAI code interpreter costs $0.03 per session

6
tests/test_litellm/integrations/test_langfuse.py
View file

@ -117,12 +117,6 @@ class TestLangfuseUsageDetails(unittest.TestCase):
log_event_on_langfuse, self.logger
)
# Make sure _is_langfuse_v2 returns True
def mock_is_langfuse_v2(self):
return True
self.logger._is_langfuse_v2 = types.MethodType(mock_is_langfuse_v2, self.logger)
def tearDown(self):
# Clean up logger instance to prevent state leakage
if hasattr(self, "logger"):

20
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_azure_assistant_cost_tracking.py
View file

@ -90,15 +90,6 @@ class TestAzureAssistantCostTracking:
)
assert cost == 0.0, "Should return 0 for zero sessions"
def test_openai_code_interpreter_free(self):
"""Test OpenAI code interpreter cost from model cost map."""
cost = StandardBuiltInToolCostTracking.get_cost_for_code_interpreter(
sessions=5,
provider="openai",
)
assert (
cost == 0.15
), "OpenAI code interpreter should return 0.15 based on current implementation"
@pytest.mark.parametrize(
"input_tokens,output_tokens,expected_cost",
@ -222,14 +213,3 @@ class TestAzureAssistantCostTracking:
)
assert StandardBuiltInToolCostTracking.get_cost_for_vector_store(None) == 0.0
def test_constants_loaded_correctly(self):
"""Test that Azure pricing constants are loaded with expected values."""
assert AZURE_FILE_SEARCH_COST_PER_GB_PER_DAY == 0.1
# Code interpreter cost is now in model cost map
azure_container_info = litellm.model_cost.get("azure/container", {})
assert azure_container_info.get("code_interpreter_cost_per_session") == 0.03
assert AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS == 3.0
assert AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS == 12.0
assert AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY == 0.1

237
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
View file

@ -1685,35 +1685,6 @@ def test_azure_gpt55_reasoning_effort_flags_match_live_openai_api(
assert m.get("supports_xhigh_reasoning_effort") is expected_xhigh
def test_generic_cost_per_token_anthropic_prompt_caching_with_cache_creation():
model = "claude-haiku-4-5-20251001"
usage = Usage(
completion_tokens=90,
prompt_tokens=28436,
total_tokens=28526,
completion_tokens_details=CompletionTokensDetailsWrapper(
accepted_prediction_tokens=None,
audio_tokens=None,
reasoning_tokens=0,
rejected_prediction_tokens=None,
text_tokens=None,
),
prompt_tokens_details=None,
cache_creation_input_tokens=2000,
)
custom_llm_provider = "anthropic"
prompt_cost, completion_cost = generic_cost_per_token(
model=model,
usage=usage,
custom_llm_provider=custom_llm_provider,
)
print(f"prompt_cost: {prompt_cost}")
assert round(prompt_cost, 3) == 0.029
def test_string_cost_values():
"""Test that cost values defined as strings are properly converted to floats."""
from unittest.mock import patch
@ -2350,140 +2321,6 @@ def test_gemini_image_generation_cost_falls_back_to_flat_image_pricing(_local_mo
assert round(cost, 10) == round(expected_cost, 10)
def test_bedrock_anthropic_prompt_caching():
"""Test Bedrock Anthropic models with prompt caching return correct costs."""
model = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"
usage = Usage(
prompt_tokens=52123,
completion_tokens=497,
total_tokens=52620,
cache_creation_input_tokens=7183,
cache_read_input_tokens=22465,
)
custom_llm_provider = "bedrock"
prompt_cost, completion_cost = generic_cost_per_token(
model=model,
usage=usage,
custom_llm_provider=custom_llm_provider,
)
assert prompt_cost >= 0
assert completion_cost >= 0
assert round(prompt_cost, 3) == 0.111
assert round(completion_cost, 5) == 0.00820
def test_reasoning_tokens_without_text_tokens_gpt5_nano():
"""
Test fix for GitHub issue #18599:
https://github.com/BerriAI/litellm/issues/18599
When OpenAI models (gpt-5-nano, o1, o3) return reasoning_tokens but don't provide
text_tokens, LiteLLM should calculate text_tokens as:
text_tokens = completion_tokens - reasoning_tokens - audio_tokens - image_tokens
This ensures ALL completion tokens are billed, not just reasoning tokens.
"""
model = "gpt-5-nano"
custom_llm_provider = "openai"
# Simulate OpenAI gpt-5-nano response where text_tokens is NOT provided
# completion_tokens: 977 total
# reasoning_tokens: 768
# text_tokens: should be calculated as 977 - 768 = 209
usage = Usage(
prompt_tokens=17,
completion_tokens=977,
total_tokens=994,
completion_tokens_details=CompletionTokensDetailsWrapper(
reasoning_tokens=768,
audio_tokens=0,
# text_tokens NOT provided - this is the key part of the bug
),
)
prompt_cost, completion_cost = generic_cost_per_token(
model=model,
usage=usage,
custom_llm_provider=custom_llm_provider,
)
# gpt-5-nano pricing: $0.05/1M input, $0.40/1M output
expected_prompt_cost = 17 * 0.05 / 1_000_000
expected_completion_cost = 977 * 0.40 / 1_000_000 # ALL tokens, not just reasoning
assert abs(prompt_cost - expected_prompt_cost) < 1e-10, (
f"Prompt cost incorrect: {prompt_cost} vs {expected_prompt_cost}"
)
assert abs(completion_cost - expected_completion_cost) < 1e-10, (
f"Completion cost incorrect: {completion_cost} vs {expected_completion_cost}"
)
# Verify it's NOT using only reasoning_tokens (the bug)
wrong_cost = 768 * 0.40 / 1_000_000 # Only reasoning tokens
assert abs(completion_cost - wrong_cost) > 1e-6, (
"Bug detected: Cost calculation is using only reasoning_tokens instead of all completion_tokens!"
)
def test_image_count_prevents_text_tokens_fallback(_local_model_cost_map):
"""
Test that the text_tokens fallback in generic_cost_per_token does not
override text_tokens=0 when image_count > 0.
Regression test for: Bedrock image embedding double-charging bug.
When image_count > 0, text_tokens=0 is intentional (image-only request),
not "text_tokens not set by provider."
"""
# Simulate Nova image-only embedding: prompt_tokens estimated from
# embedding dimensions (768 for 3072-dim), image_count=1
usage = Usage(
prompt_tokens=768,
completion_tokens=0,
total_tokens=768,
prompt_tokens_details=PromptTokensDetailsWrapper(
image_count=1,
),
)
prompt_cost, completion_cost = generic_cost_per_token(
model="amazon.nova-2-multimodal-embeddings-v1:0",
usage=usage,
custom_llm_provider="bedrock",
)
# Cost should be 1 * input_cost_per_image ($6e-05) = $0.00006
# NOT 768 * input_cost_per_token ($1.35e-07) + $0.00006 = $0.000164
expected_image_cost = 1 * 6e-05
assert prompt_cost == expected_image_cost, (
f"Expected prompt_cost={expected_image_cost} (image-only), "
f"got {prompt_cost}. text_tokens fallback may be double-charging."
)
assert completion_cost == 0.0
def test_query_count_bills_input_cost_per_query(_local_model_cost_map):
usage = Usage(
prompt_tokens=0,
completion_tokens=0,
total_tokens=0,
prompt_tokens_details=PromptTokensDetailsWrapper(query_count=3, image_count=1),
)
prompt_cost, completion_cost = generic_cost_per_token(
model="us.twelvelabs.marengo-embed-3-0-v1:0",
usage=usage,
custom_llm_provider="bedrock",
)
assert prompt_cost == pytest.approx(3 * 7e-05 + 1e-04)
assert completion_cost == 0.0
def test_query_count_is_free_without_a_per_query_price(_local_model_cost_map):
usage = Usage(
prompt_tokens=0,
@ -2692,36 +2529,6 @@ def test_vertex_uplift_invalid_multiplier_defaults_to_one():
)
def test_priority_service_tier_above_threshold_uses_priority_tier_rates_for_cached_tokens(
_local_model_cost_map,
):
"""Regression: for a model that publishes both service_tier and above_threshold rate
variants, a priority request over the threshold must bill cached tokens at
cache_read_input_token_cost_above_200k_tokens_priority (and analogously for
input/output above-threshold), not the standard above-threshold rate."""
usage = Usage(
prompt_tokens=250_000,
completion_tokens=1_000,
total_tokens=251_000,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=200_000, text_tokens=50_000),
completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=1_000),
)
prompt_cost, completion_cost = generic_cost_per_token(
model="gemini-3-pro-preview",
usage=usage,
custom_llm_provider="gemini",
service_tier="priority",
)
# gemini-3-pro-preview priority + above_200k rates from the pricing JSON:
# input 7.2e-6, output 3.24e-5, cache_read 7.2e-7
expected_prompt = 50_000 * 7.2e-6 + 200_000 * 7.2e-7
expected_completion = 1_000 * 3.24e-5
assert prompt_cost == pytest.approx(expected_prompt, rel=1e-9)
assert completion_cost == pytest.approx(expected_completion, rel=1e-9)
def test_service_tier_suffixes_constant_in_sync_with_enum():
from litellm.litellm_core_utils.llm_cost_calc.utils import _SERVICE_TIER_SUFFIXES
from litellm.types.utils import ServiceTier
@ -3614,28 +3421,6 @@ def test_gemini_38_flash_matches_37_flash_promotional_pricing(prefix, _local_mod
assert new_model[field] == old_model[field], field
@pytest.mark.parametrize(
("model", "provider", "image_token_rate"),
[
("gpt-realtime-2.1", "openai", 5e-06),
("gpt-realtime-2.1-mini", "openai", 8e-07),
("azure/gpt-realtime-2.1", "azure", 5e-06),
("azure/gpt-realtime-2.1-mini", "azure", 8e-07),
],
)
def test_realtime_image_tokens_priced_per_token(model, provider, image_token_rate, _local_model_cost_map):
"""Realtime image input is billed per 1M image tokens, not per image."""
usage = Usage(
prompt_tokens=1_100,
completion_tokens=0,
total_tokens=1_100,
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=100, image_tokens=1_000),
)
prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider=provider)
text_rate = litellm.model_cost[model]["input_cost_per_token"]
assert prompt_cost == pytest.approx(100 * text_rate + 1_000 * image_token_rate)
@pytest.mark.parametrize(
("response_quality", "requested_quality", "expected_cost"),
[
@ -3830,28 +3615,6 @@ def test_cached_audio_tokens_fall_back_to_cache_read_input_token_cost() -> None:
assert prompt_cost == pytest.approx(expected)
def test_cache_read_breakdown_splits_cached_audio_at_the_audio_cache_rate(_local_model_cost_map: None) -> None:
usage = Usage(
prompt_tokens=4863,
completion_tokens=1087,
total_tokens=5950,
prompt_tokens_details=PromptTokensDetailsWrapper(
text_tokens=1693,
audio_tokens=3170,
cached_tokens=2816,
cached_tokens_details={"text_tokens": 896, "audio_tokens": 1920},
),
)
breakdown = get_token_type_cost_breakdown(model="gpt-realtime-2.1-mini", custom_llm_provider="openai", usage=usage)
prompt_cost, _ = generic_cost_per_token(model="gpt-realtime-2.1-mini", usage=usage, custom_llm_provider="openai")
assert breakdown.cache_read_cost == pytest.approx(896 * 6e-8 + 1920 * 3e-7)
assert breakdown.rates is not None
assert breakdown.rates.cache_read_input_audio_token_cost == pytest.approx(3e-7)
assert prompt_cost == pytest.approx((1693 - 896) * 6e-7 + (3170 - 1920) * 1e-5 + breakdown.cache_read_cost)
def test_generic_cost_per_token_bills_cache_creation_at_the_input_rate_without_a_write_price():
"""Azure and OpenAI publish no cache-write price and bill cache writes as ordinary input.
A deployment priced with only input, output, and cache-read rates must bill the creation

347
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py
View file

@ -309,102 +309,6 @@ def test_get_cost_for_gemini_web_search(model):
assert cost > 0.0
@pytest.mark.parametrize(
"model,custom_llm_provider",
[
("vertex_ai/gemini-2.5-flash", "vertex_ai"),
("gemini-2.5-flash", "vertex_ai"),
],
)
def test_get_cost_for_vertex_ai_gemini_web_search(model, custom_llm_provider):
"""
Test that Vertex AI Gemini web search costs are tracked when passing
a ModelResponse with usage.prompt_tokens_details.web_search_requests.
This tests the fix for: https://github.com/BerriAI/litellm/issues/XXXXX
The issue: When a ModelResponse is passed, the detection logic only checks
for url_citation annotations, not usage.prompt_tokens_details.web_search_requests.
This causes Vertex AI grounding costs to not be tracked.
"""
from litellm.types.utils import Choices, Message, PromptTokensDetailsWrapper, Usage
# Create a realistic ModelResponse like what Vertex AI returns
response = ModelResponse(
id="test-id",
choices=[
Choices(
finish_reason="stop",
index=0,
message=Message(
content="Test response with grounding", role="assistant"
),
)
],
created=1234567890,
model=model,
object="chat.completion",
system_fingerprint=None,
)
# Add usage with web_search_requests (how Vertex AI indicates grounding was used)
usage = Usage(
prompt_tokens=11,
completion_tokens=100,
total_tokens=111,
prompt_tokens_details=PromptTokensDetailsWrapper(
text_tokens=11, web_search_requests=1 # This should trigger grounding cost
),
)
response.usage = usage
# Calculate cost - should include grounding cost
cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools(
model=model,
usage=usage,
response_object=response, # Pass the ModelResponse
custom_llm_provider=custom_llm_provider,
standard_built_in_tools_params=None,
)
# Vertex AI charges $0.035 per grounded request
assert cost == 0.035, f"Expected $0.035 grounding cost, got ${cost}"
def test_azure_assistant_features_integrated_cost_tracking(monkeypatch):
"""
Test integrated cost tracking for Azure assistant features.
"""
# Force use of local model cost map for CI/CD consistency
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "azure/gpt-4o"
# Test with multiple Azure assistant features
standard_built_in_tools_params = StandardBuiltInToolsParams(
vector_store_usage={"storage_gb": 1.0, "days": 10},
computer_use_usage={"input_tokens": 1000, "output_tokens": 500},
code_interpreter_sessions=2,
)
cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools(
model=model,
response_object=None,
usage=None,
custom_llm_provider="azure",
standard_built_in_tools_params=standard_built_in_tools_params,
)
# Should calculate costs for:
# - Vector store: 1.0 * 10 * 0.1 = $1.00
# - Computer use: (1000/1000 * 3.0) + (500/1000 * 12.0) = $9.00
# - Code interpreter: 2 * 0.03 = $0.06
# Total: $10.06
expected_cost = 1.0 + 9.0 + 0.06
assert abs(cost - expected_cost) < 0.01, f"Expected ~{expected_cost}, got {cost}"
def test_completion_cost_includes_web_search_without_standard_built_in_tools_params():
"""
Test that completion_cost includes web search cost even when
@ -510,68 +414,6 @@ def test_gemini_3x_web_search_billed_per_query(model, local_model_cost_map):
)
@pytest.mark.parametrize(
"model,custom_llm_provider",
[
("gemini/gemini-2.5-flash", "gemini"),
("vertex_ai/gemini-2.5-flash", "vertex_ai"),
],
)
def test_gemini_2x_maps_grounding_billed_at_maps_rate(model, custom_llm_provider, local_model_cost_map):
"""
Grounding with Google Maps is its own SKU: a Maps-only grounded prompt on Gemini 2.x bills the
$0.025 Maps per-prompt fee, not the $0.035 Google Search fee it was previously conflated with,
and not $0 as on Vertex AI where webSearchQueries is never populated for Maps.
Regression for https://github.com/BerriAI/litellm/issues/35906
"""
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
model_info = litellm.get_model_info(model)
expected_cost = model_info["google_maps_grounding_cost_per_query"]
assert expected_cost == pytest.approx(0.025)
usage = Usage(
prompt_tokens=15,
completion_tokens=100,
total_tokens=115,
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=15, google_maps_grounding_requests=1),
)
cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools(
model=model,
usage=usage,
response_object=None,
custom_llm_provider=custom_llm_provider,
standard_built_in_tools_params=None,
)
assert cost == pytest.approx(expected_cost)
def test_gemini_3x_maps_grounding_billed_per_query(local_model_cost_map):
"""Gemini 3.x bills Maps grounding per executed query: N queries cost N * $0.014."""
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
model = "vertex_ai/gemini-3.5-flash"
model_info = litellm.get_model_info(model)
assert model_info["web_search_billing_unit"] == "per_query"
expected_cost = model_info["google_maps_grounding_cost_per_query"] * 2
usage = Usage(
prompt_tokens=15,
completion_tokens=100,
total_tokens=115,
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=15, google_maps_grounding_requests=2),
)
cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools(
model=model,
usage=usage,
response_object=None,
custom_llm_provider="vertex_ai",
standard_built_in_tools_params=None,
)
assert cost == pytest.approx(expected_cost)
assert cost == pytest.approx(0.028)
def test_gemini_combined_search_and_maps_costs_are_additive(local_model_cost_map):
"""A prompt grounded with both Google Search and Google Maps pays both fees."""
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
@ -708,35 +550,6 @@ def _openai_responses_with_web_search_calls(model, num_calls):
)
def test_openai_responses_web_search_priced_per_call(local_model_cost_map):
"""
Regression for LIT-5013 bug 1: OpenAI reasoning models (gpt-5 family, o-series, deep-research)
carry supports_web_search but had no search_context_cost_per_query, so get_cost_for_web_search_request
(no openai branch) returned None and the default fallback billed web search as $0. gpt-5-nano now
prices at $0.01 per call, and two web_search_call items in the Responses output must bill 2 x $0.01.
"""
from litellm.types.utils import Usage
model = "gpt-5-nano"
per_call = litellm.get_model_info(model)["search_context_cost_per_query"][
"search_context_size_medium"
]
assert per_call == 0.01
response = _openai_responses_with_web_search_calls(model, num_calls=2)
cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools(
model=model,
response_object=response,
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
custom_llm_provider="openai",
standard_built_in_tools_params=None,
)
assert cost == pytest.approx(2 * per_call), (
f"gpt-5-nano web search must bill 2 x ${per_call}, got ${cost}"
)
def test_openai_responses_web_search_multiplied_by_call_count(local_model_cost_map):
"""
Regression for LIT-5013 bug 2: web_search_call detection was binary, so a Responses output with
@ -808,88 +621,6 @@ def test_web_search_call_count_reads_dict_output_items(local_model_cost_map):
)
def test_dated_search_preview_entries_carry_search_pricing(local_model_cost_map):
"""
Regression for the live QA finding: OpenAI resolves gpt-4o-search-preview requests to the
dated id gpt-4o-search-preview-2025-03-11, whose cost map entry lacked
search_context_cost_per_query, so the default chat path silently billed the $0.035 search
fee as $0. Dated entries must price identically to their undated siblings.
"""
from litellm.types.utils import Usage
for dated, undated in (
("gpt-4o-search-preview-2025-03-11", "gpt-4o-search-preview"),
("gpt-4o-mini-search-preview-2025-03-11", "gpt-4o-mini-search-preview"),
):
assert (
litellm.get_model_info(dated)["search_context_cost_per_query"]
== litellm.get_model_info(undated)["search_context_cost_per_query"]
)
response = ModelResponse(
model="gpt-4o-search-preview-2025-03-11",
choices=[
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": "headlines",
"annotations": [
{
"type": "url_citation",
"url_citation": {
"url": "https://example.com",
"title": "t",
"start_index": 0,
"end_index": 1,
},
}
],
},
}
],
)
cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools(
model="gpt-4o-search-preview-2025-03-11",
response_object=response,
usage=Usage(prompt_tokens=14, completion_tokens=825, total_tokens=839),
custom_llm_provider="openai",
standard_built_in_tools_params=None,
)
assert cost == pytest.approx(0.025), (
f"dated search-preview id must bill the $0.025 search fee, got ${cost}"
)
@pytest.mark.parametrize(
"web_search_options",
[
None,
WebSearchOptions(search_context_size="low"),
WebSearchOptions(search_context_size="medium"),
WebSearchOptions(search_context_size="high"),
],
)
def test_gpt_4o_mini_snapshot_bills_web_search_like_its_alias(
web_search_options: WebSearchOptions | None, local_model_cost_map: None
) -> None:
alias_info = litellm.get_model_info("gpt-4o-mini")
snapshot_info = litellm.get_model_info("gpt-4o-mini-2024-07-18")
assert not snapshot_info["supports_web_search"]
assert not alias_info["supports_web_search"]
snapshot_cost = StandardBuiltInToolCostTracking.get_cost_for_web_search(
web_search_options=web_search_options, model_info=snapshot_info
)
alias_cost = StandardBuiltInToolCostTracking.get_cost_for_web_search(
web_search_options=web_search_options, model_info=alias_info
)
assert snapshot_cost == alias_cost == 0.025
# Note: File search integration test removed due to complex annotation detection logic
# The unit tests in test_azure_assistant_cost_tracking.py provide comprehensive coverage
@ -999,81 +730,3 @@ def _web_search_cost(model: str, response: ResponsesAPIResponse, custom_llm_prov
)
@pytest.mark.parametrize("model", _BEDROCK_MANTLE_WEB_SEARCH_MODELS)
def test_bedrock_mantle_web_search_billed_per_query(local_model_cost_map, model):
"""Two Bedrock-reported web searches bill 2 x $0.012 under the prefixed and the bare model id alike."""
pricing = litellm.get_model_info(model)["search_context_cost_per_query"]
assert pricing == {
"search_context_size_low": _BEDROCK_MANTLE_WEB_SEARCH_RATE,
"search_context_size_medium": _BEDROCK_MANTLE_WEB_SEARCH_RATE,
"search_context_size_high": _BEDROCK_MANTLE_WEB_SEARCH_RATE,
}
response = _responses_with_web_search(
model,
actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}],
tool_usage={"web_search": {"num_requests": 2}},
)
for cost_model in (model, model.split("/", 1)[1]):
cost = _web_search_cost(cost_model, response, "bedrock_mantle")
assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), (
f"{cost_model} must bill 2 x ${_BEDROCK_MANTLE_WEB_SEARCH_RATE} for 2 web searches, got ${cost}"
)
@pytest.mark.parametrize("num_requests", [1, 0])
def test_web_search_call_count_prefers_provider_reported_num_requests(local_model_cost_map, num_requests):
"""A search plus an open_page fetch bills tool_usage.web_search.num_requests, never the two items."""
model = "bedrock_mantle/openai.gpt-5.6-sol"
response = _responses_with_web_search(
model,
actions=[
{"type": "search", "query": "litellm"},
{"type": "open_page", "url": "https://docs.litellm.ai/"},
],
tool_usage={"web_search": {"num_requests": num_requests}},
)
cost = _web_search_cost(model, response, "bedrock_mantle")
assert cost == pytest.approx(num_requests * _BEDROCK_MANTLE_WEB_SEARCH_RATE), (
f"{num_requests} reported web search requests must bill {num_requests} x "
f"${_BEDROCK_MANTLE_WEB_SEARCH_RATE}, got ${cost}"
)
@pytest.mark.parametrize(
"tool_usage",
[None, {}, {"web_search": None}, {"web_search": {"num_requests": "many"}}, {"web_search": {"num_requests": -1}}],
)
def test_web_search_call_count_falls_back_to_items_without_reported_count(local_model_cost_map, tool_usage):
"""Without a usable reported count the per-call path keeps counting web_search_call items."""
model = "bedrock_mantle/openai.gpt-5.6-sol"
response = _responses_with_web_search(
model,
actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}],
tool_usage=tool_usage,
)
cost = _web_search_cost(model, response, "bedrock_mantle")
assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), (
f"2 web_search_call items with tool_usage={tool_usage!r} must bill 2 x "
f"${_BEDROCK_MANTLE_WEB_SEARCH_RATE}, got ${cost}"
)
def test_web_search_call_count_reads_reported_count_beside_other_tool_usage_entries(local_model_cost_map):
"""OpenAI reports web_search.num_requests next to other tool entries, which must not disable the reported count."""
response = _responses_with_web_search(
"gpt-5.6",
actions=[{"type": "search", "query": "S&P 500 close"}, {"type": "open_page", "url": "https://example.com/"}],
tool_usage={
"image_gen": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0},
"web_search": {"num_requests": 1},
},
)
cost = _web_search_cost("gpt-5.6", response, "openai")
assert cost == pytest.approx(0.01), f"1 reported OpenAI web search must bill 1 x $0.01, not the 2 items, got ${cost}"

48
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -395,53 +395,6 @@ class TestGetRouterDeploymentModelInfo:
logging_obj.litellm_params = {"api_base": ""}
assert logging_obj.get_router_deployment_model_info() is None
@pytest.mark.parametrize(
"declared,expected_input,expected_output",
[
({"input_cost_per_token": 1e-06}, 1e-06, 1.5e-05),
({"output_cost_per_token": 5e-06}, 3e-06, 5e-06),
({"input_cost_per_token": 0.0, "output_cost_per_token": 0.0}, 0.0, 0.0),
],
ids=["input-only", "output-only", "both-zero"],
)
def test_one_sided_override_keeps_the_published_rate_for_the_other_side(
self,
declared: dict[str, float],
expected_input: float,
expected_output: float,
) -> None:
"""A deployment may configure one direction only.
Substituting its pricing wholesale billed the direction it left unset at
zero, because get_model_info fills an absent cost with 0 and that
suppressed the global fallback.
"""
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
model = "bedrock/global.anthropic.claude-sonnet-4-6"
published = litellm.get_model_info(model=model)
assert (published["input_cost_per_token"], published["output_cost_per_token"]) == (3e-06, 1.5e-05)
deployment_id = f"deploy-one-sided-{'-'.join(sorted(declared))}"
litellm.model_cost[deployment_id] = {"id": deployment_id, **declared}
obj = LiteLLMLoggingObj(
model=model,
messages=[],
stream=False,
call_type="aretrieve_batch",
start_time=time.time(),
litellm_call_id="one-sided",
function_id="f",
)
obj.litellm_params = {"litellm_metadata": {"model_info": {"id": deployment_id}}, "model": model}
obj.model_call_details["model"] = model
try:
info = obj.get_router_deployment_model_info()
assert info is not None
assert info["input_cost_per_token"] == expected_input
assert info["output_cost_per_token"] == expected_output
finally:
litellm.model_cost.pop(deployment_id, None)
def test_a_published_batch_rate_never_displaces_a_declared_standard_rate(self) -> None:
"""Ownership is per token direction, not per field.
@ -511,7 +464,6 @@ class TestGetRouterDeploymentModelInfo:
cached_before = dict(litellm.get_model_info(model=deployment_id))
info = obj.get_router_deployment_model_info()
assert info is not None
assert info["output_cost_per_token"] == 1.5e-05
assert dict(litellm.get_model_info(model=deployment_id)) == cached_before
finally:
litellm.model_cost.pop(deployment_id, None)

8
tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py
View file

@ -336,7 +336,6 @@ def test_streaming_preserves_anthropic_1hr_cache_creation_breakdown():
Correct cache-write cost is 50 * 6e-06 (1h) = 0.0003, not 50 * 3.75e-06 = 0.0001875.
"""
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.cost_calculation import cost_per_token
config = AnthropicConfig()
message_start_usage = config.calculate_usage(
@ -400,13 +399,6 @@ def test_streaming_preserves_anthropic_1hr_cache_creation_breakdown():
assert usage.cache_creation_input_tokens == 50
assert usage.cache_read_input_tokens == 8728
prompt_cost, _ = cost_per_token(model="claude-sonnet-4-6", usage=usage)
# text 3*3e-06 + cache_read 8728*3e-07 + cache_write 50*6e-06 (1h rate)
expected = 3 * 3e-06 + 8728 * 3e-07 + 50 * 6e-06
assert prompt_cost == pytest.approx(expected)
# Guard against the regression: 5m-rate fallback would shave the write cost.
buggy = 3 * 3e-06 + 8728 * 3e-07 + 50 * 3.75e-06
assert prompt_cost != pytest.approx(buggy)
def test_streaming_keeps_cache_creation_breakdown_from_final_chunk():

32
tests/test_litellm/litellm_core_utils/test_streaming_handler.py
View file

@ -2589,22 +2589,6 @@ def test_dispatch_petals_empty_stream_after_finish_raises(
_run_dispatch(initialized_custom_stream_wrapper, chunk=None)
def test_dispatch_palm_slices_completion_stream(
initialized_custom_stream_wrapper: CustomStreamWrapper,
):
"""palm uses the same fake-streaming slice strategy as petals."""
initialized_custom_stream_wrapper.custom_llm_provider = "palm"
initialized_custom_stream_wrapper.completion_stream = "B" * 40
result, _, completion_obj = _run_dispatch(
initialized_custom_stream_wrapper, chunk=None
)
assert isinstance(result, _ProviderChunkParsed)
assert completion_obj["content"] == "B" * 30
assert initialized_custom_stream_wrapper.completion_stream == "B" * 10
def test_dispatch_cached_response_extracts_delta(
initialized_custom_stream_wrapper: CustomStreamWrapper,
):
@ -2844,22 +2828,6 @@ def test_dispatch_triton_stream(
assert initialized_custom_stream_wrapper.received_finish_reason == "stop"
def test_dispatch_ai21_decodes_completion(
initialized_custom_stream_wrapper: CustomStreamWrapper,
):
"""ai21 does fake streaming over a single byte-encoded JSON completion."""
initialized_custom_stream_wrapper.custom_llm_provider = "ai21"
chunk = json.dumps({"completions": [{"data": {"text": "ai21 text"}}]}).encode(
"utf-8"
)
result, _, completion_obj = _run_dispatch(initialized_custom_stream_wrapper, chunk)
assert isinstance(result, _ProviderChunkParsed)
assert completion_obj["content"] == "ai21 text"
assert initialized_custom_stream_wrapper.received_finish_reason == "stop"
def test_dispatch_text_completion_openai_with_usage(
initialized_custom_stream_wrapper: CustomStreamWrapper,
):

13
tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py
View file

@ -130,16 +130,3 @@ def test_openai_style_unsupported_param_dropped_with_drop_params():
assert mapped == {}
def test_cost_calculator_uses_aiml_pricing_for_gpt_image_2():
"""Regression: pricing must come from the ``aiml/openai/gpt-image-2`` entry,
not the upstream OpenAI token-based entry.
"""
response = ImageResponse(
data=[
ImageObject(b64_json=None, url="https://example.com/1.png"),
ImageObject(b64_json=None, url="https://example.com/2.png"),
]
)
assert aiml_cost_calculator(
model="openai/gpt-image-2", image_response=response
) == pytest.approx(0.054 * 2)

39
tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py
View file

@ -2442,21 +2442,6 @@ def test_get_max_tokens_for_model_claude_35():
assert max_tokens == 8192
def test_get_max_tokens_for_model_claude_37():
"""
Test that get_max_tokens_for_model returns correct value for Claude 3.7 models.
Claude 3.7 Sonnet has max_output_tokens of 64000 by default.
128K output requires the beta header 'output-128k-2025-02-19'.
Fixes: https://github.com/BerriAI/litellm/issues/8835
"""
config = AnthropicConfig()
# Claude 3.7 Sonnet should return 64000 (64K default, 128K requires beta header)
max_tokens = config.get_max_tokens_for_model("claude-3-7-sonnet-20250219")
assert max_tokens == 64000
def test_get_max_tokens_for_model_unknown():
"""
Test that get_max_tokens_for_model returns 4096 fallback for unknown models.
@ -2631,29 +2616,6 @@ def test_transform_request_injects_dummy_tool_without_tools_param():
assert "dummy_tool" in names
def test_transform_request_uses_dynamic_max_tokens():
"""
Test that transform_request uses dynamic max_tokens based on model
when max_tokens is not explicitly provided.
Fixes: https://github.com/BerriAI/litellm/issues/8835
"""
config = AnthropicConfig()
messages = [{"role": "user", "content": "Hello"}]
# Claude 3.7 model should get 64000 as default max_tokens (from model_prices_and_context_window.json)
result = config.transform_request(
model="claude-3-7-sonnet-20250219",
messages=messages,
optional_params={}, # No max_tokens provided
litellm_params={},
headers={},
)
assert result["max_tokens"] == 64000
def test_transform_request_respects_user_max_tokens():
"""
Test that transform_request respects user-provided max_tokens
@ -2851,7 +2813,6 @@ def test_raw_adaptive_thinking_untouched_for_46_plus_model():
assert result["thinking"] == {"type": "adaptive"}
@pytest.mark.parametrize(
"model, expected",
[

23
tests/test_litellm/llms/anthropic/test_azure_ai_cache_pricing.py
View file

@ -4,7 +4,6 @@ Verifies the fix for issue #19532.
"""
import litellm
from litellm import get_model_info
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
@ -18,25 +17,3 @@ def reload_model_costs():
yield
@pytest.mark.parametrize(
"model,expected_cache_creation_cost,expected_cache_read_cost",
[
("claude-haiku-4-5", 1.25e-06, 1e-07),
("claude-opus-4-5", 6.25e-06, 5e-07),
("claude-opus-4-1", 1.875e-05, 1.5e-06),
("claude-sonnet-4-5", 3.75e-06, 3e-07),
],
)
def test_azure_ai_claude_cache_pricing(
model, expected_cache_creation_cost, expected_cache_read_cost
):
"""Test that Azure AI Claude models have correct cache pricing."""
model_info = get_model_info(model=model, custom_llm_provider="azure_ai")
assert model_info.get("cache_creation_input_token_cost") is not None
assert model_info.get("cache_read_input_token_cost") is not None
assert (
model_info.get("cache_creation_input_token_cost")
== expected_cache_creation_cost
)
assert model_info.get("cache_read_input_token_cost") == expected_cache_read_cost

20
tests/test_litellm/llms/azure/test_audio_transcriptions.py
View file

@ -26,26 +26,6 @@ def _transcription_client() -> AzureOpenAI:
)
def test_azure_ai_transcription_is_priced_at_the_azure_ai_entry():
with AUDIO_FILE.open("rb") as audio:
response = litellm.transcription(
model="azure_ai/whisper",
file=audio,
api_base="https://example.cognitiveservices.azure.com",
api_key="test-key",
api_version="2024-06-01",
client=_transcription_client(),
)
with AUDIO_FILE.open("rb") as audio:
duration = calculate_request_duration(audio)
assert duration is not None and duration > 0
assert response._hidden_params["custom_llm_provider"] == "azure_ai"
assert completion_cost(completion_response=response, call_type="transcription") == pytest.approx(
WHISPER_COST_PER_SECOND * duration
)
def test_azure_transcription_keeps_the_azure_provider():
with AUDIO_FILE.open("rb") as audio:
response = litellm.transcription(

54
tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py
View file

@ -158,13 +158,6 @@ class TestAzureModelRouterFlatCost:
assert prompt_cost == pytest.approx(1000 * ROUTER_FEE_PER_TOKEN, rel=1e-9)
assert completion_cost_usd == 0.0
@pytest.mark.parametrize("router_entry_name", ["model_router", "model-router"])
def test_router_entry_prices_its_own_fee(self, router_entry_name: str) -> None:
usage = Usage(prompt_tokens=1_000_000, completion_tokens=0, total_tokens=1_000_000)
prompt_cost, completion_cost_usd = cost_per_token(model=router_entry_name, usage=usage)
assert prompt_cost == pytest.approx(0.14, rel=1e-9)
assert completion_cost_usd == 0.0
def test_routed_model_is_priced_as_itself(self) -> None:
routed_prompt_cost, routed_completion_cost = _routed_model_cost()
prompt_cost, completion_cost_usd = cost_per_token(model=ROUTED_MODEL, usage=ROUTED_USAGE)
@ -210,24 +203,6 @@ class TestAzureModelRouterFlatCost:
assert prompt_cost == pytest.approx(routed_prompt_cost + ROUTED_FEE, rel=1e-9)
assert completion_cost_usd == pytest.approx(routed_completion_cost, rel=1e-9)
def test_flat_cost_helper(self) -> None:
assert calculate_azure_model_router_flat_cost(
model="azure-model-router", prompt_tokens=10_000
) == pytest.approx(0.0014, rel=1e-9)
assert calculate_azure_model_router_flat_cost(model="gpt-5-nano", prompt_tokens=10_000) == 0.0
def test_flat_cost_reads_the_fee_from_the_deployment_named_entry(self) -> None:
litellm.register_model(
{"azure_ai/model-router": {"input_cost_per_token": 2e-07, "litellm_provider": "azure_ai", "mode": "chat"}}
)
litellm.get_model_info.cache_clear()
assert calculate_azure_model_router_flat_cost(model="model-router", prompt_tokens=1_000_000) == pytest.approx(
0.2, rel=1e-9
)
assert calculate_azure_model_router_flat_cost(
model="azure-model-router", prompt_tokens=1_000_000
) == pytest.approx(0.14, rel=1e-9)
@pytest.mark.usefixtures("local_model_cost_map")
class TestAzureModelRouterCostBreakdown:
@ -350,32 +325,3 @@ class TestAzureAIServiceTierCostCalculation:
assert flex_prompt < standard_prompt
assert flex_completion < standard_completion
def test_codestral_2501_model_info_and_cost(local_model_cost_map):
model_info = get_model_info(model="Codestral-2501", custom_llm_provider="azure_ai")
usage = Usage(prompt_tokens=1_000_000, completion_tokens=1_000_000, total_tokens=2_000_000)
prompt_cost, completion_cost = cost_per_token(model="Codestral-2501", usage=usage)
assert model_info["mode"] == "chat"
assert model_info["max_input_tokens"] == 256000
assert model_info["max_output_tokens"] == 4096
assert prompt_cost == pytest.approx(0.3)
assert completion_cost == pytest.approx(0.9)
def test_mai_thinking_1_model_info_and_cost(local_model_cost_map):
model_info = get_model_info(model="MAI-Thinking-1", custom_llm_provider="azure_ai")
usage = Usage(prompt_tokens=1_000_000, completion_tokens=1_000_000, total_tokens=2_000_000)
prompt_cost, completion_cost = cost_per_token(model="MAI-Thinking-1", usage=usage)
assert model_info["mode"] == "chat"
assert model_info["max_input_tokens"] == 256000
assert model_info["max_output_tokens"] == 64000
assert model_info["cache_read_input_token_cost"] == pytest.approx(2e-07)
assert model_info["supports_reasoning"] is True
assert model_info["supports_function_calling"] is True
assert prompt_cost == pytest.approx(2.0)
assert completion_cost == pytest.approx(8.0)

14
tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py
View file

@ -33,17 +33,3 @@ def use_local_model_cost_map():
monkeypatch.undo()
def test_azure_ai_kimi_k26_cost_per_token(use_local_model_cost_map):
from litellm.llms.azure_ai.cost_calculator import cost_per_token
from litellm.types.utils import Usage
usage = Usage(
prompt_tokens=1_000_000,
completion_tokens=1_000_000,
total_tokens=2_000_000,
)
prompt_cost, completion_cost = cost_per_token(model="kimi-k2.6", usage=usage)
assert prompt_cost == pytest.approx(0.95)
assert completion_cost == pytest.approx(4.0)

8
tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py
View file

@ -1903,7 +1903,6 @@ async def test_unified_bedrock_messages_cache_on_start_only_never_negative_cost(
custom_llm_provider="bedrock",
)
assert cost > 0
assert cost == pytest.approx(0.0093951, rel=0, abs=1e-9)
@pytest.mark.asyncio
@ -1967,13 +1966,6 @@ async def test_unified_bedrock_messages_sse_usage_and_cost_claude_sonnet_46():
assert built.usage.cache_creation_input_tokens == 10553
assert built.usage.cache_read_input_tokens == 25490
cost = completion_cost(
completion_response=built,
model="bedrock/us.anthropic.claude-sonnet-4-6",
custom_llm_provider="bedrock",
)
assert cost == pytest.approx(0.052150725, rel=0, abs=1e-9)
@pytest.mark.parametrize(
"model",

48
tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py
View file

@ -159,51 +159,3 @@ def test_bedrock_gpt_5_6_offers_tools_and_reasoning_effort_but_not_thinking(prof
# Cache-read prices are the `*-cache-read-input-tokens` usagetype rows of the AWS Price List API, us-east-1,
# https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json on 2026-09-15
@pytest.mark.parametrize(
"model,expected_cache_read",
[
("amazon.nova-lite-v1:0", 1.5e-8),
("us.amazon.nova-lite-v1:0", 1.5e-8),
("amazon.nova-micro-v1:0", 8.75e-9),
("us.amazon.nova-micro-v1:0", 8.75e-9),
("amazon.nova-pro-v1:0", 2e-7),
("us.amazon.nova-pro-v1:0", 2e-7),
("us.amazon.nova-premier-v1:0", 6.25e-7),
],
)
def test_bedrock_nova_cache_read_prices(
model, expected_cache_read, local_model_cost_map
):
model_info = litellm.model_cost[model]
assert model_info["cache_read_input_token_cost"] == expected_cache_read
usage = Usage(
prompt_tokens=1_000,
completion_tokens=100,
total_tokens=1_100,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=400),
)
response = _bedrock_response(model, usage)
cost = completion_cost(
completion_response=response,
model=model,
custom_llm_provider="bedrock",
)
expected_cost = (
600 * model_info["input_cost_per_token"]
+ 400 * expected_cache_read
+ 100 * model_info["output_cost_per_token"]
)
assert cost == pytest.approx(expected_cost)
uncached_usage = Usage(
prompt_tokens=1_000,
completion_tokens=100,
total_tokens=1_100,
)
uncached_cost = completion_cost(
completion_response=_bedrock_response(model, uncached_usage),
model=model,
custom_llm_provider="bedrock",
)
assert cost < uncached_cost

78
tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
View file

@ -369,6 +369,52 @@ class TestBedrockMantleResponsesTools:
assert "file_search" in str(mock_warning.call_args)
class TestBedrockMantleSamplingParams:
"""Mantle serves OpenAI's gpt-5 models under their OpenAI sampling rule: top_p and a
non-default temperature are accepted only when reasoning.effort resolves to none, so
the `openai.` catalogue name (region-prefixed on GovCloud) must answer from the OpenAI
model's map entry instead of dropping both params on every request."""
@pytest.mark.parametrize(
"model, effort, survives",
[
("openai.gpt-5.4", None, True),
("openai.gpt-5.5", None, False),
("openai.gpt-5.6-luna", None, False),
("openai.gpt-5.6-luna", "none", True),
("openai.gpt-5.6-luna", "low", False),
("us-gov-west-1/openai.gpt-5.4", None, True),
("us-gov-west-1/openai.gpt-5.6-luna", None, False),
],
)
def test_top_p_and_temperature_follow_the_resolved_effort(self, local_cost_map, model, effort, survives):
params = {"top_p": 0.9, "temperature": 0.2}
if effort is not None:
params["reasoning"] = {"effort": effort}
mapped = BedrockMantleResponsesAPIConfig().map_openai_params(
response_api_optional_params=params,
model=model,
drop_params=True,
)
assert ("top_p" in mapped) is survives
assert ("temperature" in mapped) is survives
def test_top_p_without_drop_params_raises_only_while_reasoning_is_active(self, local_cost_map):
with pytest.raises(litellm.UnsupportedParamsError):
BedrockMantleResponsesAPIConfig().map_openai_params(
response_api_optional_params={"top_p": 0.9},
model="openai.gpt-5.6-luna",
drop_params=False,
)
mapped = BedrockMantleResponsesAPIConfig().map_openai_params(
response_api_optional_params={"top_p": 0.9},
model="openai.gpt-5.4",
drop_params=False,
)
assert mapped["top_p"] == 0.9
class TestBedrockMantleResponsesWebSearch:
"""Web Search on Amazon Bedrock is a server-side built-in tool that Mantle runs
itself when the caller passes {"type": "web_search"} on the Responses path, so
@ -1865,38 +1911,6 @@ class TestBedrockMantleResponsesSigV4:
class TestBedrockMantleResponsesPricing:
@pytest.mark.parametrize(
"model, input_cost, output_cost",
[
("openai.gpt-5.6-sol", 5.5e-06, 3.3e-05),
("openai.gpt-5.6-terra", 2.2e-06, 1.32e-05),
("openai.gpt-5.6-luna", 2.2e-07, 1.32e-06),
],
)
def test_gpt_5_6_responses_call_cost(self, local_cost_map, model, input_cost, output_cost):
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
input_tokens = 100000
output_tokens = 10000
response = ResponsesAPIResponse(
id="resp-1",
created_at=1700000000,
model=model,
output=[],
usage=ResponseAPIUsage(
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=input_tokens + output_tokens,
),
)
cost = litellm.completion_cost(
completion_response=response,
model=f"bedrock_mantle/{model}",
custom_llm_provider="bedrock_mantle",
)
assert cost == pytest.approx(input_tokens * input_cost + output_tokens * output_cost)
def test_models_registered(self, local_cost_map):
assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models

20
tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py
View file

@ -62,23 +62,3 @@ def test_map_openai_params_preserves_max_retries_zero_falsy() -> None:
assert "max_retries" in result and result["max_retries"] == 0, (
f"max_retries=0 (falsy) must not be silently omitted; got: {result!r}"
)
def test_qwen_3_8_27b_cost_and_tokens(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
model = "cerebras/qwen-3.8-27b"
prompt_cost, completion_cost = litellm.cost_per_token(
model=model,
prompt_tokens=1000,
completion_tokens=1000,
)
assert abs(prompt_cost - 0.00099) < 1e-9
assert abs(completion_cost - 0.00149) < 1e-9
model_info = litellm.get_model_info(model)
assert model_info["max_input_tokens"] == 65536
assert model_info["max_output_tokens"] == 32768
assert model_info["supports_vision"] is True
assert model_info["supports_reasoning"] is True
assert model_info["supports_parallel_function_calling"] is True

20
tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py
View file

@ -45,26 +45,6 @@ class TestChatGPTResponsesAPITransformation:
assert isinstance(config, ChatGPTResponsesAPIConfig)
assert config.custom_llm_provider == LlmProviders.CHATGPT
@pytest.mark.parametrize(
"model_name",
[
"chatgpt/gpt-5.5",
"chatgpt/gpt-5.6-luna",
"chatgpt/gpt-5.6-sol",
"chatgpt/gpt-5.6-terra",
],
)
def test_chatgpt_responses_model_metadata(self, model_name: str, local_model_cost_map: None) -> None:
model_info = litellm.get_model_info(model_name)
assert model_info["litellm_provider"] == "chatgpt"
assert model_info["mode"] == "responses"
assert model_info["supported_endpoints"] == [
"/v1/chat/completions",
"/v1/responses",
]
assert model_info["max_input_tokens"] == 1050000
assert model_info["max_output_tokens"] == 128000
@pytest.mark.parametrize(
"model_name",

21
tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py
View file

@ -127,24 +127,3 @@ def test_transform_image_generation_request():
) == {"prompt": "a red bicycle", "quality": "high", "num_images": 2}
@pytest.mark.parametrize(
("model", "expected_cost_for_two_images"),
[
("openai/gpt-image-2", 0.29),
("gpt-image-2", 0.29),
("openai/gpt-image-2/edit", 0.302),
],
)
def test_cost_calculator_uses_registry_price(
model, expected_cost_for_two_images, monkeypatch: pytest.MonkeyPatch
):
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
litellm.get_model_info.cache_clear()
response = ImageResponse(
data=[
ImageObject(url="https://v3b.fal.media/files/b/one.png"),
ImageObject(url="https://v3b.fal.media/files/b/two.png"),
]
)
assert cost_calculator(model=model, image_response=response) == pytest.approx(expected_cost_for_two_images)

17
tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py
View file

@ -145,20 +145,3 @@ def test_transform_request_includes_prompt_and_mapped_params():
}
@pytest.mark.parametrize(
"model", ["fal-ai/nano-banana", "fal-ai/gemini-25-flash-image"]
)
def test_nano_banana_pricing_registered(model):
info = litellm.get_model_info(
model=model, custom_llm_provider=litellm.LlmProviders.FAL_AI.value
)
assert info["output_cost_per_image"] == 0.039
assert info["mode"] == "image_generation"
def test_cost_calculator_scales_with_image_count():
image_response = ImageResponse(
data=[ImageObject(url="https://x/1.png"), ImageObject(url="https://x/2.png")]
)
cost = cost_calculator(model="fal-ai/nano-banana", image_response=image_response)
assert cost == pytest.approx(0.078)

137
tests/test_litellm/llms/fal_ai/test_cost_calculator.py
View file

@ -17,140 +17,3 @@ def _use_local_model_cost_map(monkeypatch):
def _image_response(num_images: int = 1) -> ImageResponse:
return ImageResponse(data=[ImageObject(url="https://example.com/img.png") for _ in range(num_images)])
def test_high_quality_1024x1024_uses_keyed_price():
cost = cost_calculator(
model="openai/gpt-image-2",
image_response=_image_response(),
optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}},
)
assert cost == pytest.approx(0.211)
def test_alias_model_uses_keyed_price():
cost = cost_calculator(
model="gpt-image-2",
image_response=_image_response(),
optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}},
)
assert cost == pytest.approx(0.211)
def test_provider_prefixed_model_uses_keyed_price():
cost = cost_calculator(
model="fal_ai/openai/gpt-image-2",
image_response=_image_response(),
optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}},
)
assert cost == pytest.approx(0.211)
def test_provider_prefixed_edit_model_uses_keyed_edit_price():
cost = cost_calculator(
model="fal_ai/openai/gpt-image-2/edit",
image_response=_image_response(),
optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}},
)
assert cost == pytest.approx(0.219)
def test_default_request_priced_at_default_size_and_quality():
cost = cost_calculator(
model="openai/gpt-image-2",
image_response=_image_response(),
optional_params={},
)
assert cost == pytest.approx(0.145)
def test_auto_quality_priced_as_high():
cost = cost_calculator(
model="openai/gpt-image-2",
image_response=_image_response(),
optional_params={"quality": "auto", "image_size": {"width": 1024, "height": 1024}},
)
assert cost == pytest.approx(0.211)
def test_low_quality_4k_uses_keyed_price():
cost = cost_calculator(
model="openai/gpt-image-2",
image_response=_image_response(),
optional_params={"quality": "low", "image_size": {"width": 3840, "height": 2160}},
)
assert cost == pytest.approx(0.012)
def test_named_fal_size_uses_keyed_price():
cost = cost_calculator(
model="openai/gpt-image-2",
image_response=_image_response(),
optional_params={"quality": "high", "image_size": "square_hd"},
)
assert cost == pytest.approx(0.211)
def test_edit_model_uses_keyed_edit_price():
cost = cost_calculator(
model="openai/gpt-image-2/edit",
image_response=_image_response(),
optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}},
)
assert cost == pytest.approx(0.219)
def test_edit_model_without_size_falls_back_to_flat_price():
cost = cost_calculator(
model="openai/gpt-image-2/edit",
image_response=_image_response(),
optional_params={"quality": "high"},
)
assert cost == pytest.approx(0.151)
def test_missing_optional_params_falls_back_to_flat_price():
cost = cost_calculator(
model="openai/gpt-image-2",
image_response=_image_response(),
optional_params=None,
)
assert cost == pytest.approx(0.145)
def test_unlisted_size_falls_back_to_flat_price():
cost = cost_calculator(
model="openai/gpt-image-2",
image_response=_image_response(),
optional_params={"quality": "high", "image_size": {"width": 999, "height": 999}},
)
assert cost == pytest.approx(0.145)
def test_keyed_price_multiplies_per_image():
cost = cost_calculator(
model="openai/gpt-image-2",
image_response=_image_response(num_images=2),
optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}},
)
assert cost == pytest.approx(0.422)
def test_route_image_generation_passes_optional_params_to_fal():
cost = CostCalculatorUtils.route_image_generation_cost_calculator(
model="openai/gpt-image-2",
completion_response=_image_response(),
custom_llm_provider="fal_ai",
optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}},
)
assert cost == pytest.approx(0.211)
def test_route_image_generation_with_provider_prefixed_model_uses_keyed_price():
cost = CostCalculatorUtils.route_image_generation_cost_calculator(
model="fal_ai/openai/gpt-image-2",
completion_response=_image_response(),
custom_llm_provider="fal_ai",
optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}},
)
assert cost == pytest.approx(0.211)

15
tests/test_litellm/llms/gemini/audio_transcription/test_gemini_audio_transcription_transformation.py
View file

@ -302,18 +302,3 @@ class TestCostRegression:
def local_cost_map(self, monkeypatch):
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
def test_registry_entries(self, local_cost_map):
batch_entry = litellm.model_cost["gemini/gemini-3.5-transcribe"]
assert batch_entry["mode"] == "audio_transcription"
assert batch_entry["input_cost_per_audio_token"] == 2e-06
assert batch_entry["input_cost_per_token"] == 2e-06
assert batch_entry["output_cost_per_token"] == 1.2e-05
assert batch_entry["supported_endpoints"] == ["/v1/audio/transcriptions"]
live_entry = litellm.model_cost["gemini/gemini-3.5-transcribe-live"]
assert live_entry["mode"] == "audio_transcription"
assert live_entry["input_cost_per_audio_token"] == 3.5e-06
assert live_entry["input_cost_per_token"] == 3.5e-06
assert live_entry["output_cost_per_token"] == 2.1e-05
assert live_entry["supported_endpoints"] == ["/v1/realtime"]

48
tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py
View file

@ -1856,54 +1856,6 @@ def test_map_openai_params_drops_stock_voice_case_insensitively():
assert passthrough["generationConfig"]["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore"
def test_gemini_response_done_bills_audio_output_tokens_at_audio_rate(monkeypatch):
"""Regression for the Gemini Live AUDIO output breakdown: responseTokensDetails
must survive into response.done usage and bill at output_cost_per_audio_token,
not the text rate."""
from litellm.cost_calculator import (
RealtimeAPITokenUsageProcessor,
handle_realtime_stream_cost_calculation,
)
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
config = GeminiRealtimeConfig()
done_event = config.transform_response_done_event(
message={
"serverContent": {"turnComplete": True},
"usageMetadata": {
"promptTokenCount": 377,
"responseTokenCount": 51,
"totalTokenCount": 428,
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 377}],
"responseTokensDetails": [{"modality": "AUDIO", "tokenCount": 51}],
"thoughtsTokenCount": 37,
},
},
current_response_id="resp_lit6277",
current_conversation_id="conv_lit6277",
output_items=None,
)
usage = done_event["response"]["usage"]
assert usage["output_tokens_details"]["audio_tokens"] == 51
assert usage["output_token_details"]["audio_tokens"] == 51
results = [done_event]
combined_usage = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(
results=results,
)
assert combined_usage.completion_tokens_details is not None
assert combined_usage.completion_tokens_details.audio_tokens == 51
cost = handle_realtime_stream_cost_calculation(
results=results,
combined_usage_object=combined_usage,
custom_llm_provider="gemini",
litellm_model_name="gemini-2.5-flash-native-audio-preview-12-2025",
)
assert cost == pytest.approx(377 * 5e-07 + 51 * 1.2e-05 + 37 * 2e-06)
@pytest.fixture(autouse=False)
def patch_gemini_transcribe_live_cost_map_entry(monkeypatch):
"""Inject the gemini-3.5-transcribe-live registry entry locally.

33
tests/test_litellm/llms/groq/chat/test_groq_chat_transformation.py
View file

@ -21,7 +21,6 @@ WEB_SEARCH_MODELS = (
COMPOUND_MODELS = ("compound", "compound-mini", "groq/compound", "groq/compound-mini")
class TestGroqWebSearchOptions:
@pytest.mark.parametrize("model", WEB_SEARCH_MODELS + COMPOUND_MODELS)
def test_supported_on_search_capable_models(self, model: str):
@ -204,36 +203,4 @@ class TestGroqWebSearchUsageSignal:
GroqChatConfig()._add_web_search_usage(model_response=model_response)
assert getattr(model_response, "usage", None) is None
@pytest.mark.usefixtures("local_model_cost_map")
@pytest.mark.parametrize(
"executed_tools, expected_cost",
[
(EXECUTED_TOOLS_THREE_SEARCHES_TWO_OPENS, 3 * 0.005 + 2 * 0.001),
(EXECUTED_TOOLS_OPENS_ONLY, 2 * 0.001),
],
)
def test_response_billed_per_action(self, executed_tools: list, expected_cost: float):
response = _groq_completion_with_mocked_response(_searched_groq_response(executed_tools))
assert StandardBuiltInToolCostTracking.response_object_includes_web_search_call(
response_object=response, usage=response.usage
)
cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools(
model="groq/openai/gpt-oss-20b",
response_object=response,
usage=response.usage,
custom_llm_provider="groq",
standard_built_in_tools_params={"web_search_options": {"search_context_size": "high"}},
)
assert cost == pytest.approx(expected_cost)
class TestGroqWebSearchCost:
@pytest.mark.usefixtures("local_model_cost_map")
@pytest.mark.parametrize("model", WEB_SEARCH_MODELS)
@pytest.mark.parametrize("search_context_size", ["low", "medium", "high"])
def test_browser_search_priced_per_search(self, model: str, search_context_size: str):
cost = StandardBuiltInToolCostTracking.get_cost_for_web_search(
web_search_options={"search_context_size": search_context_size},
model_info=litellm.get_model_info(model=model, custom_llm_provider="groq"),
)
assert cost == 0.005

19
tests/test_litellm/llms/inception/test_inception_chat_transformation.py
View file

@ -308,22 +308,3 @@ def test_inception_completion_targets_inception_endpoint():
assert response.choices[0].message.content == "hi"
def test_inception_mercury_2_5_cost_and_tokens(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
model = "inception/mercury-2.5"
prompt_cost, completion_cost = litellm.cost_per_token(
model=model,
prompt_tokens=1000,
completion_tokens=500,
)
assert abs(prompt_cost - 0.0002) < 1e-9
assert abs(completion_cost - 0.000375) < 1e-9
model_info = litellm.get_model_info(model)
assert model_info["max_input_tokens"] == 260000
assert model_info["max_output_tokens"] == 65536
assert model_info["litellm_provider"] == "inception"
assert model_info["mode"] == "chat"
assert model_info["supports_function_calling"] is True
assert model_info["supports_response_schema"] is True

40
tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py
View file

@ -1835,6 +1835,46 @@ class TestResponsesSurfaceSharesTheEffortRule:
)
assert ("temperature" in mapped) is temperature_survives
@pytest.mark.parametrize(
"model, effort, top_p_survives",
[
("gpt-5.1", None, True),
("gpt-5.4", None, True),
("gpt-5.5", None, False),
("gpt-5.6-terra", None, False),
("gpt-5.6-sol", None, False),
("gpt-5.6-terra", "none", True),
("gpt-5.6-terra", "medium", False),
("gpt-6-astra", None, False),
("gpt-6-astra", "low", False),
],
)
def test_top_p_follows_the_resolved_effort(self, local_model_cost_map, model, effort, top_p_survives):
params = {"top_p": 0.9}
if effort is not None:
params["reasoning"] = {"effort": effort}
mapped = OpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params=params,
model=model,
drop_params=True,
)
assert ("top_p" in mapped) is top_p_survives
def test_top_p_raises_without_drop_params(self, local_model_cost_map):
with pytest.raises(litellm.UnsupportedParamsError):
OpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params={"top_p": 0.9},
model="gpt-5.5",
drop_params=False,
)
mapped = OpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params={"top_p": 0.9, "reasoning": {"effort": "none"}},
model="gpt-5.6-terra",
drop_params=False,
)
assert mapped["top_p"] == 0.9
class TestFlattenToolSchemaCombinatorsWiring:
"""Regression tests for MCP tools with a top-level anyOf schema (Codex Desktop).

69
tests/test_litellm/llms/openai_like/test_cognition_provider.py
View file

@ -111,28 +111,6 @@ class TestCognitionProviderIdentity:
class TestCognitionCostTracking:
@pytest.mark.parametrize(
"model, expected_prompt_cost, expected_completion_cost",
[
("cognition/swe-1.7", 0.5, 2.5),
("cognition/swe-1.7-lightning", 2.5, 12.5),
],
)
def test_cost_differs_from_openai_pricing(
self, model: str, expected_prompt_cost: float, expected_completion_cost: float
):
"""A cognition-prefixed model must never be priced off an OpenAI cost entry."""
from litellm.cost_calculator import cost_per_token
prompt_cost, completion_cost = cost_per_token(
model=model,
prompt_tokens=1_000_000,
completion_tokens=1_000_000,
custom_llm_provider="cognition",
)
assert prompt_cost == pytest.approx(expected_prompt_cost)
assert completion_cost == pytest.approx(expected_completion_cost)
def test_lightning_is_five_times_the_standard_tier(self):
standard = litellm.get_model_info(model="cognition/swe-1.7")
@ -151,51 +129,4 @@ class TestCognitionCostTracking:
assert endpoints["embeddings"] is False
class TestCognitionRouting:
@pytest.mark.asyncio
async def test_router_spend_is_attributed_to_cognition_pricing(self):
"""Routed traffic is costed off the cognition entry, not an OpenAI one."""
from litellm import Router
router = Router(
model_list=[
{
"model_name": "swe",
"litellm_params": {"model": "cognition/swe-1.7", "api_key": "sk-test"},
}
]
)
response = await router.acompletion(
model="swe",
messages=[{"role": "user", "content": "hi"}],
mock_response="hello from swe",
)
usage = response.usage
expected = usage.prompt_tokens * 5e-07 + usage.completion_tokens * 2.5e-06
assert response._hidden_params["response_cost"] == pytest.approx(expected)
@pytest.mark.asyncio
async def test_router_spend_uses_the_lightning_entry_for_lightning(self):
"""The Lightning tier is its own model, costed off its own entry."""
from litellm import Router
router = Router(
model_list=[
{
"model_name": "swe-lightning",
"litellm_params": {"model": "cognition/swe-1.7-lightning", "api_key": "sk-test"},
}
]
)
response = await router.acompletion(
model="swe-lightning",
messages=[{"role": "user", "content": "hi"}],
mock_response="hello from swe lightning",
)
usage = response.usage
expected = usage.prompt_tokens * 2.5e-06 + usage.completion_tokens * 1.25e-05
assert response._hidden_params["response_cost"] == pytest.approx(expected)

16
tests/test_litellm/llms/openai_like/test_meta_provider.py
View file

@ -192,20 +192,4 @@ class TestMetaAnthropicMessages:
assert headers["anthropic-version"] == "2023-06-01"
class TestMuseSparkModelInfo:
def test_muse_spark_cost_calculation(self):
from litellm import completion_cost
from litellm.types.utils import ModelResponse, Usage
response = ModelResponse(
model="muse-spark-1.1",
usage=Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500),
)
cost = completion_cost(
completion_response=response,
model="meta/muse-spark-1.1",
custom_llm_provider="meta",
)
expected = 1000 * 1.25e-06 + 500 * 4.25e-06
assert abs(cost - expected) < 1e-12

14
tests/test_litellm/llms/openai_like/test_tensormesh_provider.py
View file

@ -154,17 +154,3 @@ class TestTensormeshCostMap:
for model in TENSORMESH_MODELS:
assert litellm.supports_reasoning(model) is (model in reasoning_models), model
def test_cost_is_wired_and_cache_reads_are_free(self):
prompt_cost, completion_cost = litellm.cost_per_token(
model="tensormesh/openai/gpt-oss-120b",
prompt_tokens=1_000_000,
completion_tokens=1_000_000,
)
assert prompt_cost == pytest.approx(0.15)
assert completion_cost == pytest.approx(0.60)
assert (
litellm.model_cost["tensormesh/openai/gpt-oss-120b"][
"cache_read_input_token_cost"
]
== 0
)

73
tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py
View file

@ -431,76 +431,3 @@ class TestParallelAISearch:
assert result.snippet == ""
assert result.date is None
assert result.model_dump()["excerpts"] == ()
@pytest.mark.parametrize(
"mode,usage,max_results,expected_cost",
[
("turbo", [{"name": "sku_search", "count": 1}], None, 0.001),
("fast", [{"name": "sku_search", "count": 1}], None, 0.001),
("basic", [{"name": "sku_search", "count": 1}], None, 0.005),
("advanced", [{"name": "sku_search", "count": 1}], None, 0.005),
(
"basic",
[
{"name": "sku_search", "count": 1},
{"name": "sku_search_additional_results", "count": 2},
],
20,
0.007,
),
("basic", None, 20, 0.015),
],
)
@pytest.mark.asyncio
async def test_search_cost_uses_mode_and_provider_usage(
self, mode, usage, max_results, expected_cost, bundled_cost_map, respx_mock, httpx_transport
):
response_payload = {**MOCK_V1_RESPONSE, "usage": usage}
respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload)
response = await litellm.asearch(
query="AI developments",
search_provider="parallel_ai",
mode=mode,
max_results=max_results,
)
assert response._hidden_params["response_cost"] == pytest.approx(expected_cost)
@pytest.mark.asyncio
async def test_search_cost_treats_keyword_queries_as_one_request(
self, bundled_cost_map, respx_mock, httpx_transport
):
response_payload = {
**MOCK_V1_RESPONSE,
"usage": [{"name": "sku_search", "count": 1}],
}
respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload)
response = await litellm.asearch(
query=["AI developments", "machine learning trends"],
search_provider="parallel_ai",
mode="basic",
)
assert response._hidden_params["response_cost"] == pytest.approx(0.005)
@pytest.mark.asyncio
async def test_caller_cannot_supply_provider_usage(self, bundled_cost_map, respx_mock, httpx_transport):
"""`_parallel_ai_usage` prices the request, so a caller must not be able to set it.
The provider reports no usage here, which is the case where a caller-supplied
value would otherwise survive into the cost calculation.
"""
response_payload = {k: v for k, v in MOCK_V1_RESPONSE.items() if k != "usage"}
route = respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload)
response = await litellm.asearch(
query="AI developments",
search_provider="parallel_ai",
mode="basic",
_parallel_ai_usage=[{"name": "sku_search", "count": 0}],
)
assert response._hidden_params["response_cost"] == pytest.approx(0.005)
assert "_parallel_ai_usage" not in json.loads(route.calls[0].request.content)

17
tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py
View file

@ -140,23 +140,6 @@ class TestPerplexityCostCalculator:
assert prompt_cost == 0.0
assert completion_cost == 0.008
def test_falls_back_to_manual_calculation_when_no_cost_provided(self):
"""
Test that manual cost calculation is used when Perplexity doesn't
provide the cost object (fallback behavior).
"""
usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)
# No cost object - should use manual calculation
prompt_cost, completion_cost = perplexity_cost_per_token(model="sonar-deep-research", usage=usage)
# Should calculate manually: 100 * 2e-6 + 50 * 8e-6
expected_prompt = 100 * 2e-6
expected_completion = 50 * 8e-6
assert math.isclose(prompt_cost, expected_prompt, rel_tol=1e-6)
assert math.isclose(completion_cost, expected_completion, rel_tol=1e-6)
OFF_PEAK_MODEL = "sonar-off-peak-test"
OFF_PEAK_WINDOW = "14:00-00:00"
INSIDE_WINDOW = datetime(2026, 9, 3, 17, 25, tzinfo=timezone.utc)

21
tests/test_litellm/llms/perplexity/test_perplexity_integration.py
View file

@ -150,24 +150,3 @@ class TestPerplexityIntegration:
assert hasattr(model_response.usage, "prompt_tokens_details")
assert hasattr(model_response.usage, "citation_tokens")
assert model_response.usage.prompt_tokens_details.web_search_requests == 3
@pytest.mark.parametrize("provider_name", ["perplexity", "PERPLEXITY", "Perplexity"])
def test_case_insensitive_provider_matching(self, provider_name):
"""Test that cost calculation works with different case variations of provider name."""
usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)
usage.citation_tokens = 10
usage.prompt_tokens_details = PromptTokensDetailsWrapper(web_search_requests=1)
# Should work regardless of case
prompt_cost, completion_cost_val = cost_per_token(
model="sonar-deep-research",
custom_llm_provider=provider_name.lower(), # Normalize to lowercase
usage_object=usage,
)
# Should calculate costs correctly
expected_prompt_cost = (100 * 2e-6) + (10 * 2e-6)
expected_completion_cost = (50 * 8e-6) + (1 * 0.005)
assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6)
assert math.isclose(completion_cost_val, expected_completion_cost, rel_tol=1e-6)

41
tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py
View file

@ -1056,44 +1056,3 @@ class TestSpendTracking:
litellm.model_cost = original_model_cost
litellm.get_model_info.cache_clear()
def test_should_charge_by_audio_duration(self, monkeypatch):
import litellm
monkeypatch.setattr("time.sleep", lambda *_: None)
responses = {
"POST https://api.soniox.com/v1/transcriptions": [
_make_response({"id": "tx_1", "status": "queued"})
],
"GET https://api.soniox.com/v1/transcriptions/tx_1": [
_make_response(
{"id": "tx_1", "status": "completed", "audio_duration_ms": 600000}
),
],
"GET https://api.soniox.com/v1/transcriptions/tx_1/transcript": [
_make_response({"text": "hello world", "tokens": []}),
],
"DELETE https://api.soniox.com/v1/transcriptions/tx_1": [
_make_response({"deleted": True}),
],
}
resp = SonioxAudioTranscriptionHandler().audio_transcriptions(
audio_file=None,
optional_params={"audio_url": "https://example.com/a.wav"},
litellm_params={},
atranscription=False,
**_common_call_kwargs(_MockSyncClient(responses)),
)
assert resp._hidden_params["audio_transcription_duration"] == pytest.approx(
600.0
)
cost = litellm.completion_cost(
completion_response=resp,
model="soniox/stt-async-v4",
call_type="transcription",
)
# 10 minutes of audio billed at Soniox's ~$0.10/hour async rate.
assert cost > 0
assert cost == pytest.approx((0.10 / 3600) * 600.0, rel=1e-3)

35
tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_audio_transcription_transformation.py
View file

@ -22,16 +22,6 @@ def config():
class TestGetCompleteUrl:
def test_defaults_to_us_regional_host(self, config):
url = config.get_complete_url(
api_base=None,
api_key=None,
model="chirp_3",
optional_params={},
litellm_params={"vertex_project": "test-project"},
)
assert url == "https://us-speech.googleapis.com/v2/projects/test-project/locations/us/recognizers/_:recognize"
def test_uses_vertex_location_for_regional_host(self, config):
url = config.get_complete_url(
api_base=None,
@ -52,16 +42,6 @@ class TestGetCompleteUrl:
)
assert url == "https://speech.googleapis.com/v2/projects/test-project/locations/global/recognizers/_:recognize"
def test_api_base_override(self, config):
url = config.get_complete_url(
api_base="http://localhost:8080/",
api_key=None,
model="chirp_3",
optional_params={},
litellm_params={"vertex_project": "test-project"},
)
assert url == "http://localhost:8080/v2/projects/test-project/locations/us/recognizers/_:recognize"
@pytest.mark.parametrize(
"location,expected_netloc",
[
@ -317,18 +297,3 @@ class TestProviderRouting:
class TestModelCostEntry:
REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))
@pytest.mark.parametrize(
"cost_map_path",
[
"model_prices_and_context_window.json",
"litellm/model_prices_and_context_window_backup.json",
],
)
def test_chirp_3_registered_as_audio_transcription(self, cost_map_path):
with open(os.path.join(self.REPO_ROOT, cost_map_path)) as f:
entry = json.load(f)["vertex_ai/chirp_3"]
assert entry["mode"] == "audio_transcription"
assert entry["litellm_provider"] == "vertex_ai"
assert entry["input_cost_per_second"] == pytest.approx(0.016 / 60, rel=1e-3)
assert entry["supported_endpoints"] == ["/v1/audio/transcriptions"]

34
tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_gemini_transcribe_transformation.py
View file

@ -309,37 +309,3 @@ class TestOptionalParams:
class TestModelCostEntry:
REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))
@pytest.mark.parametrize(
"cost_map_path",
[
"model_prices_and_context_window.json",
"litellm/model_prices_and_context_window_backup.json",
],
)
def test_transcribe_preview_pricing(self, cost_map_path):
with open(os.path.join(self.REPO_ROOT, cost_map_path)) as f:
entry = json.load(f)["vertex_ai/gemini-3.5-transcribe-preview"]
assert entry["mode"] == "audio_transcription"
assert entry["litellm_provider"] == "vertex_ai"
assert entry["input_cost_per_audio_token"] == pytest.approx(2e-06)
assert entry["input_cost_per_token"] == pytest.approx(2e-06)
assert entry["output_cost_per_token"] == pytest.approx(1.2e-05)
assert entry["supported_endpoints"] == ["/v1/audio/transcriptions"]
@pytest.mark.parametrize(
"cost_map_path",
[
"model_prices_and_context_window.json",
"litellm/model_prices_and_context_window_backup.json",
],
)
def test_transcribe_live_preview_pricing(self, cost_map_path):
with open(os.path.join(self.REPO_ROOT, cost_map_path)) as f:
entry = json.load(f)["vertex_ai/gemini-3.5-transcribe-live-preview"]
assert entry["mode"] == "audio_transcription"
assert entry["litellm_provider"] == "vertex_ai"
assert entry["input_cost_per_audio_token"] == pytest.approx(3.5e-06)
assert entry["input_cost_per_token"] == pytest.approx(3.5e-06)
assert entry["output_cost_per_token"] == pytest.approx(2.1e-05)
assert entry["supported_endpoints"] == ["/v1/realtime"]

223
tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py
View file

@ -407,227 +407,4 @@ class TestProcessEmbedContentResponseUsage:
)
assert result.usage.prompt_tokens > 0
def test_file_reference_image_billed_per_image_token_rate(self):
response_json = {
"embedding": {"values": [0.1, 0.2, 0.3]},
"usageMetadata": {
"promptTokenCount": 258,
"totalTokenCount": 258,
"promptTokensDetails": [{"modality": "IMAGE", "tokenCount": 258}],
},
}
result = process_embed_content_response(
input=["files/img123"],
model_response=EmbeddingResponse(),
model=self.MODEL,
response_json=response_json,
resolved_files={
"files/img123": {
"mime_type": "image/png",
"uri": "https://example.com/img123",
}
},
)
assert result.usage.prompt_tokens_details.image_tokens == 258
assert result.usage.prompt_tokens_details.text_tokens == 0
prompt_cost, _ = generic_cost_per_token(
model=self.MODEL,
usage=result.usage,
custom_llm_provider="vertex_ai",
)
assert prompt_cost == pytest.approx(258 * 4.5e-7)
def test_file_reference_non_image_not_counted_as_image(self):
"""A files/... ref resolving to a non-image mime keeps audio token billing."""
response_json = {
"embedding": {"values": [0.1, 0.2]},
"usageMetadata": {
"promptTokenCount": 64,
"totalTokenCount": 64,
"promptTokensDetails": [{"modality": "AUDIO", "tokenCount": 64}],
},
}
result = process_embed_content_response(
input=["files/clip1"],
model_response=EmbeddingResponse(),
model=self.MODEL,
response_json=response_json,
resolved_files={
"files/clip1": {
"mime_type": "audio/mpeg",
"uri": "https://example.com/clip1",
}
},
)
assert result.usage.prompt_tokens_details.audio_tokens == 64
assert result.usage.prompt_tokens_details.image_tokens == 0
prompt_cost, _ = generic_cost_per_token(
model=self.MODEL,
usage=result.usage,
custom_llm_provider="vertex_ai",
)
assert prompt_cost == pytest.approx(64 * 6.5e-6)
def test_video_plus_audio_does_not_double_bill_text(self):
"""Video and audio responses are billed from their respective token counts."""
response_json = {
"embedding": {"values": [0.1]},
"usageMetadata": {
"promptTokenCount": 580,
"totalTokenCount": 580,
"promptTokensDetails": [
{"modality": "VIDEO", "tokenCount": 516},
{"modality": "AUDIO", "tokenCount": 64},
],
},
}
result = process_embed_content_response(
input=["gs://bucket/clip.mp4"],
model_response=EmbeddingResponse(),
model=self.MODEL,
response_json=response_json,
)
assert result.usage.prompt_tokens_details.text_tokens == 0
assert result.usage.prompt_tokens_details.video_tokens == 516
assert result.usage.prompt_tokens_details.audio_tokens == 64
prompt_cost, _ = generic_cost_per_token(
model=self.MODEL,
usage=result.usage,
custom_llm_provider="vertex_ai",
)
assert prompt_cost == pytest.approx(516 * 1.2e-5 + 64 * 6.5e-6)
def test_preview_alias_bills_audio_per_token(self):
response_json = {
"embedding": {"values": [0.1]},
"usageMetadata": {
"promptTokenCount": 64,
"totalTokenCount": 64,
"promptTokensDetails": [{"modality": "AUDIO", "tokenCount": 64}],
},
}
result = process_embed_content_response(
input="audio",
model_response=EmbeddingResponse(),
model="gemini-embedding-2-preview",
response_json=response_json,
)
prompt_cost, _ = generic_cost_per_token(
model="gemini-embedding-2-preview",
usage=result.usage,
custom_llm_provider="vertex_ai",
)
assert prompt_cost == pytest.approx(64 * 6.5e-6)
def test_image_without_modality_details_uses_image_rate(self):
response_json = {
"embedding": {"values": [0.1]},
"usageMetadata": {
"promptTokenCount": 258,
"totalTokenCount": 258,
},
}
result = process_embed_content_response(
input=IMAGE_DATA_URI,
model_response=EmbeddingResponse(),
model=self.MODEL,
response_json=response_json,
)
assert result.usage.prompt_tokens_details.image_tokens == 258
assert result.usage.prompt_tokens_details.text_tokens == 0
prompt_cost, _ = generic_cost_per_token(
model=self.MODEL,
usage=result.usage,
custom_llm_provider="vertex_ai",
)
assert prompt_cost == pytest.approx(258 * 4.5e-7)
@pytest.mark.parametrize(
"input_value,resolved_files,expected_image_tokens",
[
(GCS_URL, {}, 258),
("gs://my-bucket/clip.mp4", {}, 0),
("gs://my-bucket/unknown.bin", {}, 0),
("files/image-123", {"files/image-123": {"mime_type": "image/jpeg"}}, 258),
("files/missing", {}, 0),
("data:application/octet-stream;base64,abc", {}, 0),
([[IMAGE_DATA_URI]], {}, 258),
([], {}, 0),
],
)
def test_missing_modality_details_classifies_image_inputs(self, input_value, resolved_files, expected_image_tokens):
response_json = {
"embedding": {"values": [0.1]},
"usageMetadata": {
"promptTokenCount": 258,
"totalTokenCount": 258,
},
}
result = process_embed_content_response(
input=input_value,
model_response=EmbeddingResponse(),
model=self.MODEL,
response_json=response_json,
resolved_files=resolved_files,
)
assert result.usage.prompt_tokens_details.image_tokens == expected_image_tokens
assert result.usage.prompt_tokens_details.text_tokens == 0
prompt_cost, _ = generic_cost_per_token(
model=self.MODEL,
usage=result.usage,
custom_llm_provider="vertex_ai",
)
expected_rate = 4.5e-7 if expected_image_tokens else 2e-7
assert prompt_cost == pytest.approx(258 * expected_rate)
def test_mixed_text_and_image_without_modality_details_not_billed_as_image(self):
response_json = {
"embedding": {"values": [0.1]},
"usageMetadata": {
"promptTokenCount": 270,
"totalTokenCount": 270,
},
}
result = process_embed_content_response(
input=["a short caption", IMAGE_DATA_URI],
model_response=EmbeddingResponse(),
model=self.MODEL,
response_json=response_json,
)
assert result.usage.prompt_tokens_details.image_tokens == 0
prompt_cost, _ = generic_cost_per_token(
model=self.MODEL,
usage=result.usage,
custom_llm_provider="vertex_ai",
)
assert prompt_cost == pytest.approx(270 * 2e-7)
def test_text_without_modality_details_uses_text_rate(self):
response_json = {
"embedding": {"values": [0.1]},
"usageMetadata": {
"promptTokenCount": 12,
"totalTokenCount": 12,
},
}
result = process_embed_content_response(
input="a short caption",
model_response=EmbeddingResponse(),
model=self.MODEL,
response_json=response_json,
)
assert result.usage.prompt_tokens_details.text_tokens == 0
assert result.usage.prompt_tokens_details.image_tokens == 0
prompt_cost, _ = generic_cost_per_token(
model=self.MODEL,
usage=result.usage,
custom_llm_provider="vertex_ai",
)
assert prompt_cost == pytest.approx(12 * 2e-7)

12
tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py
View file

@ -1150,10 +1150,6 @@ def test_get_token_url():
vertex_ai_location = "us-central1"
vertex_credentials = ""
should_use_v1beta1_features = vertex_llm.is_using_v1beta1_features(
optional_params={"cached_content": "hi"}
)
_, url = vertex_llm._get_token_and_url(
auth_header=None,
vertex_project=vertex_ai_project,
@ -1161,7 +1157,7 @@ def test_get_token_url():
vertex_credentials=vertex_credentials,
gemini_api_key="",
custom_llm_provider="vertex_ai_beta",
should_use_v1beta1_features=should_use_v1beta1_features,
should_use_v1beta1_features=False,
api_base=None,
model="",
stream=False,
@ -1169,10 +1165,6 @@ def test_get_token_url():
print("url=", url)
should_use_v1beta1_features = vertex_llm.is_using_v1beta1_features(
optional_params={"temperature": 0.1}
)
_, url = vertex_llm._get_token_and_url(
auth_header=None,
vertex_project=vertex_ai_project,
@ -1180,7 +1172,7 @@ def test_get_token_url():
vertex_credentials=vertex_credentials,
gemini_api_key="",
custom_llm_provider="vertex_ai_beta",
should_use_v1beta1_features=should_use_v1beta1_features,
should_use_v1beta1_features=False,
api_base=None,
model="",
stream=False,

50
tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py
View file

@ -238,56 +238,6 @@ def test_audio_predict_response_supports_bytes_base64_encoded(
assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.06)
@pytest.mark.parametrize("runtime_entry_is_missing", (True, False))
def test_lyria_predict_cost_falls_back_to_bundled_map_when_runtime_metadata_is_incomplete(
monkeypatch: pytest.MonkeyPatch,
runtime_entry_is_missing: bool,
local_model_cost_map: None,
) -> None:
if runtime_entry_is_missing:
monkeypatch.delitem(litellm.model_cost, "vertex_ai/lyria-002")
else:
monkeypatch.setitem(
litellm.model_cost,
"vertex_ai/lyria-002",
{
key: value
for key, value in litellm.model_cost["vertex_ai/lyria-002"].items()
if key != "output_cost_per_image"
},
)
logging_obj = MagicMock()
logging_obj.model_call_details = {}
response = httpx.Response(
status_code=200,
json={
"predictions": [
{
"audioContent": "clip",
"mimeType": "audio/wav",
}
]
},
)
result = VertexPassthroughLoggingHandler.vertex_passthrough_handler(
httpx_response=response,
logging_obj=logging_obj,
url_route="/v1/projects/test/locations/us-central1/publishers/google/models/lyria-002:predict",
result=response.text,
start_time=datetime.now(),
end_time=datetime.now(),
cache_hit=False,
request_body={"instances": [{"prompt": "ambient piano"}]},
)
if runtime_entry_is_missing:
assert "vertex_ai/lyria-002" not in litellm.model_cost
assert result["kwargs"]["model"] == "lyria-002"
assert result["kwargs"]["response_cost"] == pytest.approx(0.06)
assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.06)
def test_image_predict_response_is_not_billed_as_audio(
local_model_cost_map: None,
) -> None:

30
tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py
View file

@ -123,18 +123,6 @@ class TestVertexAIVideoConfig:
model="veo-002", api_base=None, litellm_params={}
)
def test_get_complete_url_default_location(self):
"""Test URL construction with default location."""
litellm_params = {"vertex_project": "test-project"}
url = self.config.get_complete_url(
model="veo-002", api_base=None, litellm_params=litellm_params
)
# Should default to us-central1
assert "us-central1" in url
# Should NOT include endpoint
assert not url.endswith(":predictLongRunning")
def test_veo_31_lite_provider_routing_from_local_model_map(
self, monkeypatch: pytest.MonkeyPatch
@ -154,24 +142,6 @@ class TestVertexAIVideoConfig:
assert model == "veo-3.1-lite-generate-001"
assert custom_llm_provider == "vertex_ai"
def test_veo_31_lite_cost_uses_resolution_tiers(self):
model_cost = _load_model_cost_map(BACKUP_MODEL_COST_PATH)
model_info = model_cost[VEO_31_LITE_VERTEX_MODEL]
assert video_generation_cost(
model=VEO_31_LITE_VERTEX_MODEL,
duration_seconds=10.0,
custom_llm_provider="vertex_ai",
model_info=dict(model_info),
video_resolution="720p",
) == pytest.approx(0.5)
assert video_generation_cost(
model=VEO_31_LITE_VERTEX_MODEL,
duration_seconds=10.0,
custom_llm_provider="vertex_ai",
model_info=dict(model_info),
video_resolution="1080p",
) == pytest.approx(0.8)
def test_transform_video_create_request(self):
"""Test transformation of video creation request."""

13
tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py
View file

@ -105,16 +105,3 @@ def test_both_cost_maps_agree_on_the_redirected_slugs():
backup = json.loads(BACKUP_PRICES_PATH.read_text(encoding="utf-8"))
for slug in (*REDIRECTED_SLUGS, *CODE_SLUGS, REDIRECT_TARGET, CODE_REDIRECT_TARGET):
assert prices[slug] == backup[slug], slug
def test_every_retired_chat_slug_is_covered(cost_map: dict):
"""The lists above must stay in step with what the registry marks retired."""
marked = {
key
for key, entry in cost_map.items()
if isinstance(entry, dict)
and entry.get("litellm_provider") == "xai"
and "deprecation_date" in entry
and entry.get("mode") == "chat"
}
assert marked == {*REDIRECTED_SLUGS, *CODE_SLUGS}

28
tests/test_litellm/llms/zai/test_zai_provider.py
View file

@ -55,34 +55,6 @@ def test_zai_in_provider_lists():
assert "zai" in litellm.provider_list
def test_zai_glm46_cost_calculation(local_model_cost_map):
"""Test the cost calculation for glm-4.6"""
prompt_cost, completion_cost = cost_per_token(
model="zai/glm-4.6",
prompt_tokens=1000000, # 1M tokens
completion_tokens=1000000,
)
# GLM-4.6: $0.6/M input, $2.2/M output
assert math.isclose(prompt_cost, 0.6, rel_tol=1e-6)
assert math.isclose(completion_cost, 2.2, rel_tol=1e-6)
def test_glm47_cost_calculation(local_model_cost_map):
"""Test cost calculation for GLM-4.7"""
prompt_cost, completion_cost = cost_per_token(
model="zai/glm-4.7",
prompt_tokens=1000000, # 1M tokens
completion_tokens=1000000,
)
# GLM-4.7: $0.6/M input, $2.2/M output (same as GLM-4.6)
assert math.isclose(prompt_cost, 0.6, rel_tol=1e-6)
assert math.isclose(completion_cost, 2.2, rel_tol=1e-6)
@pytest.mark.asyncio
async def test_zai_completion_call(respx_mock, zai_response, monkeypatch):
"""Test completion call with zai provider using mocked response"""

4
tests/test_litellm/passthrough/test_passthrough_main.py
View file

@ -325,7 +325,7 @@ async def test_pass_through_request_stream_param_override(
"POST",
httpx.URL("https://api.anthropic.com/v1/messages"),
json=request_body,
params={},
params=None,
headers={"Authorization": "Bearer test-key"},
)
@ -424,7 +424,7 @@ async def test_pass_through_request_stream_param_no_override(
"POST",
httpx.URL("https://api.anthropic.com/v1/messages"),
headers={"Authorization": "Bearer test-key"},
params={},
params=None,
json=request_body,
)
mock_async_client.send.assert_called_once()

25
tests/test_litellm/proxy/common_utils/test_prompt_cache_pricing.py
View file

@ -7,31 +7,6 @@ from litellm.proxy.common_utils.prompt_cache_pricing import price_cache_tokens
from litellm.types.management_endpoints.prompt_cache_prediction import CacheTokenBuckets
@pytest.mark.parametrize(
("model", "expected"),
[("anthropic/claude-sonnet-4-5", 1.26), ("anthropic/claude-sonnet-4-6", 0.63)],
)
def test_prices_all_cache_buckets_at_total_context_tier(model: str, expected: float) -> None:
tokens: Final = CacheTokenBuckets(
uncached_input_tokens=100_000,
cache_read_input_tokens=50_000,
cache_creation_5m_input_tokens=20_000,
cache_creation_1h_input_tokens=40_000,
)
assert price_cache_tokens(model, "unconfigured-deployment", tokens) == pytest.approx(expected)
@pytest.mark.parametrize(("total", "expected"), [(200_000, 0.387), (200_001, 0.774006)])
def test_long_context_tier_starts_above_threshold(total: int, expected: float) -> None:
tokens: Final = CacheTokenBuckets(
uncached_input_tokens=total - 100_000,
cache_creation_1h_input_tokens=10_000,
cache_read_input_tokens=90_000,
)
actual: Final = price_cache_tokens("anthropic/claude-sonnet-4-5", "unconfigured-deployment", tokens)
assert actual == pytest.approx(expected)
def test_deployment_tariff_wins_without_proxy_discounts_or_margins(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "model_cost", litellm.model_cost.copy())
litellm.Router(

173
tests/test_litellm/proxy/management_endpoints/test_prompt_cache_prediction.py
View file

@ -110,54 +110,6 @@ async def _observe(
await cache.async_set_cache(_cache_key(scope, prefix.fingerprint), observation.model_dump_json(), ttl=3_600)
@pytest.mark.asyncio
@pytest.mark.parametrize(("ttl", "cold_cost"), [("5m", 0.0145), ("1h", 0.022)])
async def test_unobserved_cache_prices_cold_and_warm_bounds(ttl: str, cold_cost: float) -> None:
body: Final = _body(ttl)
arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, DualCache(), Counts())
assert arm.cache_state == "unknown"
assert arm.reason == "no_compatible_observation"
assert arm.evidence is None
assert arm.estimate is not None and arm.cold is not None and arm.warm is not None
assert arm.estimate.input_cost == pytest.approx(cold_cost)
assert arm.cold.input_cost == pytest.approx(cold_cost)
assert arm.warm.input_cost == pytest.approx(0.003)
assert arm.cold.tokens.uncached_input_tokens == 1_000
assert arm.cold.tokens.cache_read_input_tokens == 0
assert arm.cold.tokens.cache_creation_5m_input_tokens == (5_000 if ttl == "5m" else 0)
assert arm.cold.tokens.cache_creation_1h_input_tokens == (5_000 if ttl == "1h" else 0)
assert arm.warm.tokens.cache_read_input_tokens == 5_000
@pytest.mark.asyncio
@pytest.mark.parametrize(
("cached_tokens", "warm_cost", "cold_cost"), [(5_400, 0.00228, 0.0147), (4_600, 0.00372, 0.0143)]
)
@pytest.mark.parametrize("expired", [False, True])
async def test_exact_prefix_conserves_total_with_observed_count_in_all_scenarios(
cached_tokens: int, warm_cost: float, cold_cost: float, expired: bool
) -> None:
cache: Final = DualCache()
body: Final = _body()
await _observe(cache, body, cached_tokens=cached_tokens, expired=expired)
arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, cache, Counts())
assert arm.cache_state == ("stale" if expired else "warm")
assert arm.evidence is not None
assert arm.estimate is not None and arm.warm is not None and arm.cold is not None
assert arm.warm.tokens.cache_read_input_tokens == cached_tokens
assert arm.warm.tokens.cache_creation_5m_input_tokens == 0
assert arm.cold.tokens.cache_creation_5m_input_tokens == cached_tokens
assert arm.cold.tokens.cache_read_input_tokens == 0
for scenario in (arm.estimate, arm.cold, arm.warm):
assert scenario.tokens.total_tokens == 6_000
assert scenario.tokens.uncached_input_tokens == 6_000 - cached_tokens
assert arm.warm.input_cost == pytest.approx(warm_cost)
assert arm.cold.input_cost == pytest.approx(cold_cost)
assert arm.estimate.input_cost == pytest.approx(cold_cost if expired else warm_cost)
@pytest.mark.asyncio
async def test_observed_prefix_larger_than_full_request_returns_unknown() -> None:
cache: Final = DualCache()
@ -170,22 +122,6 @@ async def test_observed_prefix_larger_than_full_request_returns_unknown() -> Non
assert arm.estimate is None and arm.cold is None and arm.warm is None
@pytest.mark.asyncio
@pytest.mark.parametrize(("ttl", "expected"), [("5m", 0.0053), ("1h", 0.0068)])
async def test_append_only_prefix_reads_old_tokens_and_writes_extension(ttl: str, expected: float) -> None:
cache: Final = DualCache()
await _observe(cache, _body(ttl), cached_tokens=4_000)
body: Final = _body(ttl, extended=True)
arm: Final = await endpoint.predict_arm(_deployment(), body, _prefix(body), _CALLER, cache, Counts())
assert arm.cache_state == "partial"
assert arm.estimate is not None
assert arm.estimate.tokens.cache_read_input_tokens == 4_000
assert arm.estimate.tokens.cache_creation_5m_input_tokens == (1_000 if ttl == "5m" else 0)
assert arm.estimate.tokens.cache_creation_1h_input_tokens == (1_000 if ttl == "1h" else 0)
assert arm.estimate.input_cost == pytest.approx(expected)
@pytest.mark.asyncio
async def test_expired_observation_estimates_a_cold_rebuild() -> None:
cache: Final = DualCache()
@ -202,22 +138,6 @@ async def test_expired_observation_estimates_a_cold_rebuild() -> None:
assert arm.estimate.input_cost == arm.cold.input_cost
@pytest.mark.asyncio
async def test_below_model_minimum_prices_all_input_as_uncached() -> None:
body: Final = _body()
arm: Final = await endpoint.predict_arm(
_deployment(), body, _prefix(body), _CALLER, DualCache(), Counts(total=1_500, prefix=1_000)
)
assert arm.cache_state == "disabled"
assert arm.reason == "below_cache_minimum"
assert arm.estimate is not None
assert arm.estimate.tokens.uncached_input_tokens == 1_500
assert arm.estimate.tokens.cache_read_input_tokens == 0
assert arm.estimate.tokens.cache_creation_5m_input_tokens == 0
assert arm.estimate.input_cost == pytest.approx(0.003)
@pytest.mark.asyncio
@pytest.mark.parametrize("counts", [Counts(total=None), Counts(prefix=None), Counts(total=4_000)])
async def test_unavailable_or_inconsistent_token_counts_return_null_estimates(counts: Counts) -> None:
@ -269,20 +189,6 @@ async def test_custom_api_base_from_environment_returns_unknown_before_counting(
assert arm.estimate is None and arm.cold is None and arm.warm is None
@pytest.mark.asyncio
async def test_explicit_official_api_base_overrides_custom_environment(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("ANTHROPIC_API_BASE", "https://custom.invalid")
body: Final = _body()
arm: Final = await endpoint.predict_arm(
_deployment(api_base="https://api.anthropic.com"), body, _prefix(body), _CALLER, DualCache(), Counts()
)
assert arm.cache_state == "unknown"
assert arm.reason == "no_compatible_observation"
assert arm.estimate is not None
assert arm.estimate.input_cost == pytest.approx(0.0145)
@dataclass(frozen=True)
class _ProxyLogging:
internal_usage_cache: InternalUsageCache
@ -343,38 +249,6 @@ async def _post(
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("warm_deployment", "warm_model", "expected_delta", "expected_penalty"),
[("sonnet", "claude-sonnet-5", -0.03325, 0.0), ("opus", "claude-opus-5", 0.007, 0.0115)],
)
async def test_switch_delta_accounts_for_each_deployment_cache(
monkeypatch: pytest.MonkeyPatch,
warm_deployment: str,
warm_model: str,
expected_delta: float,
expected_penalty: float,
) -> None:
cache: Final = DualCache()
body: Final = _body()
await _observe(cache, body, deployment_id=warm_deployment, model=warm_model)
app: Final = _app(monkeypatch, cache, caller=UserAPIKeyAuth(api_key=_CALLER))
response: Final = await _post(app, body)
assert response.status_code == 200, response.text
result: Final = CachePredictionResponse.model_validate(response.json())
assert result.switch_delta == pytest.approx(expected_delta)
assert result.cache_rebuild_penalty == pytest.approx(expected_penalty)
assert result.cache_guarantee is False
assert result.pricing_basis == "input_before_discounts_and_margins"
if warm_deployment == "sonnet":
assert result.switch.cache_state == "warm"
assert result.stay.cache_state == "unknown"
else:
assert result.stay.cache_state == "warm"
assert result.switch.cache_state == "unknown"
@pytest.mark.asyncio
async def test_missing_caller_identity_cannot_reuse_observations(monkeypatch: pytest.MonkeyPatch) -> None:
cache: Final = DualCache()
@ -568,53 +442,6 @@ async def test_each_count_preserves_auth_cached_request_tag_limits(
assert calls.get_nowait() == "claude-opus-5"
@pytest.mark.asyncio
async def test_provider_counter_failure_releases_parallel_capacity(monkeypatch: pytest.MonkeyPatch) -> None:
cache: Final = DualCache()
limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(InternalUsageCache(cache))
caller: Final = UserAPIKeyAuth(api_key=_CALLER, max_parallel_requests=1)
async def fail_count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
raise RuntimeError("provider counter failed")
app: Final = _app(monkeypatch, cache, caller=caller, counts=fail_count, limiter=limiter)
with pytest.raises(RuntimeError, match="provider counter failed"):
await _post(app, _body())
recovered: Final = await _post(_app(monkeypatch, cache, caller=caller, limiter=limiter), _body())
assert recovered.status_code == 200, recovered.text
assert recovered.json()["switch"]["estimate"]["input_cost"] == pytest.approx(0.0145)
@pytest.mark.asyncio
async def test_cancelled_provider_counter_releases_parallel_capacity(monkeypatch: pytest.MonkeyPatch) -> None:
cache: Final = DualCache()
limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(InternalUsageCache(cache))
caller: Final = UserAPIKeyAuth(api_key=_CALLER, max_parallel_requests=1)
started: Final = asyncio.Event()
release: Final = asyncio.Event()
async def wait_count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
started.set()
await release.wait()
return await Counts()(model, api_key, body)
app: Final = _app(monkeypatch, cache, caller=caller, counts=wait_count, limiter=limiter)
pending: Final = asyncio.create_task(_post(app, _body()))
try:
await asyncio.wait_for(started.wait(), timeout=5)
pending.cancel()
with pytest.raises(asyncio.CancelledError):
await pending
release.set()
recovered: Final = await asyncio.wait_for(_post(app, _body()), timeout=5)
assert recovered.status_code == 200, recovered.text
assert recovered.json()["switch"]["estimate"]["input_cost"] == pytest.approx(0.0145)
finally:
pending.cancel()
release.set()
await asyncio.gather(pending, return_exceptions=True)
async def _unexpected_count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
pytest.fail("Unsupported prediction must return before contacting the token counter")

42
tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -7,7 +7,6 @@ from collections.abc import Callable
from contextlib import ExitStack, contextmanager
from io import BytesIO
from types import SimpleNamespace
from typing import Optional
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@ -16,34 +15,32 @@ from fastapi import Request, Response, UploadFile
from starlette.datastructures import FormData, Headers, QueryParams
from starlette.datastructures import UploadFile as StarletteUploadFile
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS,
LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
HttpPassThroughEndpointHelpers,
InitPassThroughEndpointHelpers,
LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
_registered_pass_through_routes,
chat_completion_pass_through_endpoint,
create_pass_through_route,
initialize_pass_through_endpoints,
pass_through_request,
resolve_pass_through_request_timeout,
resolve_llm_passthrough_timeout,
resolve_pass_through_request_timeout,
websocket_passthrough_request,
_with_trace_context,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
)
from litellm.proxy.pass_through_endpoints.success_handler import (
PassThroughEndpointLogging,
)
import litellm
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
)
MESSAGE_START_SSE_FRAME = b'event: message_start\ndata: {"type": "message_start"}\n\n'
@ -2436,10 +2433,10 @@ async def _run_pass_through_and_capture_wire_url(
target: str,
incoming_query: str,
merge_query_params: bool = False,
default_query_params: Optional[dict] = None,
custom_llm_provider: Optional[str] = None,
managed_files_hook: Optional[_FakeManagedFilesHook] = None,
user_api_key_dict: Optional[UserAPIKeyAuth] = None,
default_query_params: dict | None = None,
custom_llm_provider: str | None = None,
managed_files_hook: _FakeManagedFilesHook | None = None,
user_api_key_dict: UserAPIKeyAuth | None = None,
) -> httpx.URL:
import litellm
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
@ -2551,6 +2548,15 @@ async def test_pass_through_request_without_merge_replaces_target_query():
assert dict(wire_url.params) == {"q": "litellm"}
@pytest.mark.asyncio
async def test_pass_through_request_preserves_target_query_without_client_query():
wire_url = await _run_pass_through_and_capture_wire_url(
target="https://example.com/v1/models/gemini:streamGenerateContent?alt=sse",
incoming_query="",
)
assert dict(wire_url.params) == {"alt": "sse"}
@pytest.mark.asyncio
async def test_pass_through_request_merge_query_params_rewrites_managed_ids_on_the_wire():
"""
@ -5361,7 +5367,7 @@ async def test_websocket_passthrough_does_not_close_twice_when_success_logging_f
def _passthrough_kwargs_for_reservation(
user_api_key_dict: UserAPIKeyAuth,
parsed_body: Optional[dict] = None,
parsed_body: dict | None = None,
user_defined_route: bool = False,
) -> dict:
mock_request = MagicMock(spec=Request)

90
tests/test_litellm/proxy/test_proxy_utils.py
View file

@ -2151,96 +2151,6 @@ async def test_proxy_only_error_5xx_keeps_traceback_and_runs_sync_callbacks(monk
assert "test_proxy_utils" in captured["async_traceback"]
def test_create_model_info_response_resolves_alias_to_deployment_model():
"""A public model name that is not itself a cost-map key must not be resolved through
the fallback-generalization rules: `bedrock-claude-opus-5` matches the generic
claude-family baseline (200k/64k) by substring, while the deployment it fronts really
accepts 1M/128k. Regression for the /v1/models alias resolution introduced in v1.94.0."""
from litellm import Router
saved_model_cost = dict(litellm.model_cost)
try:
router = Router(
model_list=[
{
"model_name": "bedrock-claude-opus-5",
"litellm_params": {
"custom_llm_provider": "bedrock",
"model": "bedrock/eu.anthropic.claude-opus-5",
},
"model_info": {"base_model": "eu.anthropic.claude-opus-5"},
}
]
)
response = create_model_info_response(
model_id="bedrock-claude-opus-5", provider="openai", llm_router=router
)
finally:
litellm.model_cost.clear()
litellm.model_cost.update(saved_model_cost)
assert response["max_input_tokens"] == 1000000
assert response["max_output_tokens"] == 128000
def test_create_model_info_response_keeps_exact_alias_over_generalized_deployment_model():
"""Mirror of the alias bug: when the deployment points at a custom backend name that
only matches a generalization rule, the listed name's exact cost-map entry is the
better answer and must win."""
from litellm import Router
saved_model_cost = dict(litellm.model_cost)
try:
router = Router(
model_list=[
{
"model_name": "claude-opus-5",
"litellm_params": {
"custom_llm_provider": "bedrock",
"model": "bedrock/my-claude-opus-5-provisioned",
},
}
]
)
response = create_model_info_response(
model_id="claude-opus-5", provider="openai", llm_router=router
)
finally:
litellm.model_cost.clear()
litellm.model_cost.update(saved_model_cost)
assert response["max_input_tokens"] == 1000000
def test_create_model_info_response_falls_back_to_alias_for_opaque_deployment_name():
"""An Azure deployment named after the resource rather than the model has no cost-map
entry; the listed name still does, and must keep answering."""
from litellm import Router
saved_model_cost = dict(litellm.model_cost)
try:
router = Router(
model_list=[
{
"model_name": "gpt-4o",
"litellm_params": {"model": "azure/my-gpt4o-deployment"},
}
]
)
response = create_model_info_response(
model_id="gpt-4o", provider="openai", llm_router=router
)
finally:
litellm.model_cost.clear()
litellm.model_cost.update(saved_model_cost)
assert response["max_input_tokens"] == 128000
assert response["max_output_tokens"] == 16384
def test_create_model_info_response_resolves_mode_through_deployment_model():
"""`mode` is derived from the same lookup, so an aliased embedding deployment
currently reports no mode at all; it must report `embedding`."""

504
tests/test_litellm/test_cost_calculator.py
View file

@ -203,164 +203,6 @@ def test_cost_calculator_with_usage(_local_model_cost_map, monkeypatch):
assert result == expected_cost, f"Got {result}, Expected {expected_cost}"
def test_transcription_cost_uses_token_pricing(_local_model_cost_map):
from litellm import completion_cost
usage = Usage(
prompt_tokens=14,
completion_tokens=45,
total_tokens=59,
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=0, audio_tokens=14),
)
response = TranscriptionResponse(text="demo text")
response.usage = usage
cost = completion_cost(
completion_response=response,
model="gpt-4o-transcribe",
custom_llm_provider="openai",
call_type="atranscription",
)
expected_cost = (14 * 2.5e-06) + (45 * 1e-05)
assert pytest.approx(cost, rel=1e-6) == expected_cost
def test_transcription_token_pricing_is_provider_aware(_local_model_cost_map):
"""Regression: the token-priced transcription path hardcoded provider openai,
so gemini transcription models raised "This model isn't mapped yet"."""
from litellm import completion_cost
usage = Usage(
prompt_tokens=200,
completion_tokens=10,
total_tokens=210,
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=1, audio_tokens=199),
)
response = TranscriptionResponse(text="demo text")
response.usage = usage
cost = completion_cost(
completion_response=response,
model="gemini/gemini-3.5-transcribe",
custom_llm_provider="gemini",
call_type="atranscription",
)
expected_cost = (199 * 2e-06) + (1 * 2e-06) + (10 * 1.2e-05)
assert pytest.approx(cost, rel=1e-6) == expected_cost
def test_transcription_cost_falls_back_to_duration(_local_model_cost_map):
from litellm import completion_cost
response = TranscriptionResponse(text="demo text")
response.duration = 10.0
cost = completion_cost(
completion_response=response,
model="whisper-1",
custom_llm_provider="openai",
call_type="atranscription",
)
expected_cost = 10.0 * 0.0001
assert pytest.approx(cost, rel=1e-6) == expected_cost
def test_vertex_chirp_3_transcription_cost_from_duration(_local_model_cost_map):
"""Regression: the chirp_3 cost map entry shipped with output_cost_per_second 0.0,
and cost_per_second prefers output_cost_per_second whenever it is not None, so
every transcription priced to $0.00 instead of using input_cost_per_second."""
from litellm import completion_cost
response = TranscriptionResponse(text="demo text")
response.duration = 18.0
cost = completion_cost(
completion_response=response,
model="vertex_ai/chirp_3",
custom_llm_provider="vertex_ai",
call_type="atranscription",
)
expected_cost = 18.0 * 0.00026667
assert cost > 0
assert pytest.approx(cost, rel=1e-6) == expected_cost
def test_handle_realtime_stream_cost_calculation():
from litellm.cost_calculator import RealtimeAPITokenUsageProcessor
# Setup test data
results: OpenAIRealtimeStreamList = [
{"type": "session.created", "session": {"model": "gpt-3.5-turbo"}},
{
"type": "response.done",
"response": {"usage": {"input_tokens": 100, "output_tokens": 50, "total_tokens": 150}},
},
{
"type": "response.done",
"response": {
"usage": {
"input_tokens": 200,
"output_tokens": 100,
"total_tokens": 300,
}
},
},
]
combined_usage_object = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(
results=results,
)
# Test with explicit model name
cost = handle_realtime_stream_cost_calculation(
results=results,
combined_usage_object=combined_usage_object,
custom_llm_provider="openai",
litellm_model_name="gpt-3.5-turbo",
)
# Calculate expected cost
# gpt-3.5-turbo costs: $0.0015/1K tokens input, $0.002/1K tokens output
expected_cost = (300 * 0.0015 / 1000) + ( # input tokens (100 + 200)
150 * 0.002 / 1000
) # output tokens (50 + 100)
assert abs(cost - expected_cost) <= 0.00075 # Allow small floating point differences
# Test with different model name in session
results[0]["session"]["model"] = "gpt-4"
cost = handle_realtime_stream_cost_calculation(
results=results,
combined_usage_object=combined_usage_object,
custom_llm_provider="openai",
litellm_model_name="gpt-3.5-turbo",
)
# Calculate expected cost using gpt-4 rates
# gpt-4 costs: $0.03/1K tokens input, $0.06/1K tokens output
expected_cost = (300 * 0.03 / 1000) + ( # input tokens
150 * 0.06 / 1000
) # output tokens
assert abs(cost - expected_cost) < 0.00076
# Test with no response.done events
results = [{"type": "session.created", "session": {"model": "gpt-3.5-turbo"}}]
combined_usage_object = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(
results=results,
)
cost = handle_realtime_stream_cost_calculation(
results=results,
combined_usage_object=combined_usage_object,
custom_llm_provider="openai",
litellm_model_name="gpt-3.5-turbo",
)
assert cost == 0.0 # No usage, no cost
def test_handle_realtime_stream_cost_calculation_stores_cost_breakdown():
"""Regression: realtime cost must populate logging_obj.cost_breakdown so the
spend logs / UI show input vs output cost (issue: cost_breakdown was None for
@ -557,101 +399,6 @@ def test_realtime_logging_object_does_not_validate_unknown_event_types():
assert len(dumped["results"]) == len(results)
def test_realtime_transcription_duration_cost(monkeypatch):
"""
gpt-realtime-whisper transcription sessions are billed by input audio duration
($0.017/min). The .completed events carry usage {type: duration, seconds: N};
cost must equal total_seconds * input_cost_per_second.
"""
from datetime import datetime
from litellm.litellm_core_utils.litellm_logging import Logging
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
from litellm.cost_calculator import RealtimeAPITokenUsageProcessor
results: OpenAIRealtimeStreamList = [
{
"type": "session.created",
"session": {
"type": "transcription",
"audio": {"input": {"transcription": {"model": "gpt-realtime-whisper"}}},
},
},
{
"type": "conversation.item.input_audio_transcription.completed",
"transcript": "hello",
"usage": {"type": "duration", "seconds": 60.0},
},
{
"type": "conversation.item.input_audio_transcription.completed",
"transcript": "world",
"usage": {"type": "duration", "seconds": 30.0},
},
]
combined = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results=results)
logging_obj = Logging(
model="gpt-realtime-whisper",
messages=[],
stream=False,
call_type="_arealtime",
start_time=datetime.now(),
litellm_call_id="realtime-transcription-cost-breakdown-test",
function_id="realtime-transcription-cost-breakdown-test",
)
cost = handle_realtime_stream_cost_calculation(
results=results,
combined_usage_object=combined,
custom_llm_provider="openai",
litellm_model_name="gpt-realtime-whisper",
litellm_logging_obj=logging_obj,
)
# 90 seconds at $0.017/minute.
expected = 90.0 * (0.017 / 60)
assert abs(cost - expected) < 1e-9
assert cost > 0 # guards against the duration branch being dropped
assert logging_obj.cost_breakdown is not None
assert abs(logging_obj.cost_breakdown["total_cost"] - cost) < 1e-9
# The transcription cost must be attributed in the breakdown, not just folded
# into total_cost, or input_cost + output_cost + additional_costs won't sum to total_cost.
additional_costs = logging_obj.cost_breakdown.get("additional_costs")
assert additional_costs is not None
assert abs(additional_costs["transcription_cost"] - expected) < 1e-9
attributed_total = (
logging_obj.cost_breakdown["input_cost"]
+ logging_obj.cost_breakdown["output_cost"]
+ additional_costs["transcription_cost"]
)
assert abs(attributed_total - logging_obj.cost_breakdown["total_cost"]) < 1e-9
def test_realtime_transcription_duration_cost_resolves_model_from_litellm_name(
monkeypatch,
):
"""When no session event carries the ASR model, the litellm_model_name is used."""
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
results: OpenAIRealtimeStreamList = [
{
"type": "conversation.item.input_audio_transcription.completed",
"usage": {"type": "duration", "seconds": 120.0},
},
]
cost = handle_realtime_stream_cost_calculation(
results=results,
combined_usage_object=Usage(),
custom_llm_provider="azure",
litellm_model_name="azure/gpt-realtime-whisper",
)
assert abs(cost - 120.0 * (0.017 / 60)) < 1e-9
def test_realtime_transcription_no_completed_events_is_zero(monkeypatch):
"""A realtime stream without transcription completed events adds no extra cost."""
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
@ -673,35 +420,6 @@ def test_realtime_transcription_no_completed_events_is_zero(monkeypatch):
)
def test_realtime_transcription_token_billed_fallback(monkeypatch):
"""
Token-billed transcription models price by audio/text tokens. Verify the
fallback path multiplies audio tokens by the model's audio token cost.
"""
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
from litellm.cost_calculator import _transcription_usage_cost
# gpt-4o-transcribe: input_cost_per_audio_token = 2.5e-06, input_cost_per_token = 2.5e-06,
# output_cost_per_token = 1e-05
model_info = litellm.get_model_info(model="gpt-4o-transcribe", custom_llm_provider="openai")
usage = {
"type": "tokens",
"input_tokens": 40,
"output_tokens": 10,
"total_tokens": 50,
"input_token_details": {"audio_tokens": 30, "text_tokens": 10},
}
cost = _transcription_usage_cost(usage, model_info)
expected = (
30 * 2.5e-06 # audio tokens
+ 10 * 2.5e-06 # text tokens
+ 10 * 1e-05 # output tokens
)
assert abs(cost - expected) < 1e-12
def test_transcription_usage_cost_returns_zero_for_unknown_type():
"""An unrecognized usage type yields 0 (safe fallback, no exception)."""
from litellm.cost_calculator import _transcription_usage_cost
@ -1290,78 +1008,6 @@ def test_bedrock_cost_calculator_comparison_with_without_cache():
print(f"Cost with cache: {cost_with_cache}")
def test_gemini_25_implicit_caching_cost():
"""
Test that Gemini 2.5 models correctly calculate costs with implicit caching.
This test reproduces the issue from #11156 where cached tokens should receive
a 75% discount.
"""
from litellm import completion_cost
from litellm.types.utils import (
Choices,
Message,
ModelResponse,
PromptTokensDetailsWrapper,
Usage,
)
# Create a mock response similar to the one in the issue
litellm_model_response = ModelResponse(
id="test-response",
created=1750733889,
model="gemini/gemini-2.5-flash",
object="chat.completion",
system_fingerprint=None,
choices=[
Choices(
finish_reason="stop",
index=0,
message=Message(
content="Understood. This is a test message to check the response from the Gemini model.",
role="assistant",
tool_calls=None,
function_call=None,
),
)
],
usage=Usage(
total_tokens=15050,
prompt_tokens=15033,
completion_tokens=17,
prompt_tokens_details=PromptTokensDetailsWrapper(
audio_tokens=None,
cached_tokens=14316, # This is cachedContentTokenCount from Gemini
),
completion_tokens_details=None,
),
)
# Calculate the cost
result = completion_cost(
completion_response=litellm_model_response,
model="gemini/gemini-2.5-flash",
)
# Current pricing for gemini/gemini-2.5-flash:
# input: $0.30 / 1M tokens (3e-07 per token)
# cache_read: $0.03 / 1M tokens (3e-08 per token)
# output: $2.50 / 1M tokens (2.5e-06 per token)
# Breakdown:
# - Cached tokens: 14316 * 3e-08 = 0.00042948
# - Non-cached tokens: (15033-14316) * 3e-07 = 717 * 3e-07 = 0.00021510
# - Output tokens: 17 * 2.5e-06 = 0.00004250
# Total: 0.00042948 + 0.00021510 + 0.00004250 = 0.00068708
expected_cost = 0.00068708
# Allow for small floating point differences
assert abs(result - expected_cost) < 1e-8, f"Expected cost {expected_cost}, but got {result}"
print(f"✓ Gemini 2.5 implicit caching cost calculation is correct: ${result:.8f}")
def test_log_context_cost_calculation():
"""
Test that log context cost calculation works correctly with tiered pricing.
@ -3730,31 +3376,6 @@ def test_combine_usage_objects_sums_mirrored_cache_write_fields_once():
assert combined_pair.prompt_tokens_details.cache_creation_tokens == 100
def test_completion_cost_prices_anthropic_shaped_cache_read_tokens(_local_model_cost_map):
"""Regression: an Anthropic /v1/messages response reports cache reads as top-level
cache_read_input_tokens with input_tokens excluding them. Reading that usage as
Responses API usage dropped the cache tokens and billed the whole prompt at the
uncached input rate, overstating spend on cache hits."""
response = {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "gpt-5.6-sol",
"stop_reason": "end_turn",
"content": [{"type": "text", "text": "1"}],
"usage": {"input_tokens": 3, "output_tokens": 5, "cache_read_input_tokens": 4014},
}
cost = litellm.completion_cost(
completion_response=response,
model="gpt-5.6-sol",
custom_llm_provider="openai",
)
assert cost == pytest.approx(3 * 4e-6 + 4014 * 4e-7 + 5 * 2e-5, rel=1e-9)
def _together_chat_response(
model: str, prompt_tokens: int, completion_tokens: int, cached_tokens: int
) -> ModelResponse:
@ -3773,60 +3394,6 @@ def _together_chat_response(
)
def test_completion_cost_prices_together_cached_tokens_at_cache_read_rate(_local_model_cost_map):
"""Regression: Together reports prompt_tokens_details.cached_tokens but no together_ai
registry entry carried cache_read_input_token_cost, so cache-hit tokens were priced at
0.0 and spend on cache-heavy workloads was understated."""
cost = completion_cost(
completion_response=_together_chat_response(
model="deepseek-ai/DeepSeek-V4-Flash-0731", prompt_tokens=7864, completion_tokens=16, cached_tokens=7863
),
custom_llm_provider="together_ai",
)
assert cost == pytest.approx(1 * 1.4e-07 + 7863 * 3e-08 + 16 * 2.8e-07, rel=1e-9)
def test_completion_cost_together_mapped_model_skips_size_bucket(_local_model_cost_map):
"""Regression: any together model whose name matches (\\d+b) was rewritten to a
together-ai-* size bucket before the registry lookup, so mapped models like
Muse-Glimmer-30B never used their per-model rates, cache fields included."""
cost = completion_cost(
completion_response=_together_chat_response(
model="meta-models/Muse-Glimmer-30B", prompt_tokens=63, completion_tokens=16, cached_tokens=0
),
custom_llm_provider="together_ai",
)
assert cost == pytest.approx(63 * 3.5e-07 + 16 * 1.5e-06, rel=1e-9)
def test_completion_cost_together_unmapped_model_still_uses_size_bucket(_local_model_cost_map):
cost = completion_cost(
completion_response=_together_chat_response(
model="qwen/Qwen2-72B-Instruct", prompt_tokens=23, completion_tokens=15, cached_tokens=0
),
custom_llm_provider="together_ai",
)
assert cost == pytest.approx((23 + 15) * 9e-07, rel=1e-9)
def test_completion_cost_together_metadata_only_model_still_uses_size_bucket(_local_model_cost_map):
assert "input_cost_per_token" not in litellm.model_cost["together_ai/togethercomputer/CodeLlama-34b-Instruct"]
cost = completion_cost(
completion_response=_together_chat_response(
model="togethercomputer/CodeLlama-34b-Instruct", prompt_tokens=23, completion_tokens=15, cached_tokens=0
),
custom_llm_provider="together_ai",
)
assert cost == pytest.approx((23 + 15) * 8e-07, rel=1e-9)
def test_select_model_name_strips_unregistered_alias_prefix(_local_model_cost_map):
"""A router-facing model_name alias containing "/" whose leading segment is NOT a
registered provider must not be double-prefixed into a non-existent cost key.
@ -4011,31 +3578,6 @@ def test_completion_cost_base_model_ignores_regional_row(_local_model_cost_map):
) == pytest.approx(1000 * flat["input_cost_per_token"])
def test_completion_cost_nonzero_for_slash_alias_model_name(_local_model_cost_map):
"""End-to-end cost through a "/"-containing alias must price above zero (#38069)."""
response = litellm.ModelResponse(
id="x",
choices=[
{
"index": 0,
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
}
],
model="vertex/claude-opus-5",
)
response._hidden_params = {"custom_llm_provider": "vertex_ai"}
response.usage = litellm.Usage(prompt_tokens=100, completion_tokens=50)
cost = litellm.completion_cost(
completion_response=response,
custom_llm_provider="vertex_ai",
)
assert cost == pytest.approx(100 * 5e-6 + 50 * 2.5e-5, rel=1e-9)
def test_select_model_name_unresolvable_alias_unchanged(_local_model_cost_map):
"""An alias that resolves to no known cost key keeps the legacy double-prefixed name."""
@ -4259,52 +3801,6 @@ def test_explicit_pricing_precedes_private_provider_response_model(
assert selected == expected
def test_handle_realtime_stream_cost_calculation_bills_nested_reasoning_tokens_once(
_local_model_cost_map: None,
) -> None:
"""Realtime response.done nests reasoning_tokens inside text_tokens, so they are billed once."""
results: OpenAIRealtimeStreamList = [
{"type": "session.created", "session": {"model": "gpt-realtime-2.1-mini"}},
{
"type": "response.done",
"response": {
"usage": {
"total_tokens": 260,
"input_tokens": 237,
"output_tokens": 23,
"input_token_details": {
"text_tokens": 43,
"audio_tokens": 0,
"image_tokens": 194,
"cached_tokens": 0,
"cached_tokens_details": {"text_tokens": 0, "audio_tokens": 0, "image_tokens": 0},
},
"output_token_details": {"text_tokens": 23, "audio_tokens": 0, "reasoning_tokens": 18},
}
},
},
]
combined_usage_object = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(
results=results,
)
total_cost = handle_realtime_stream_cost_calculation(
results=results,
combined_usage_object=combined_usage_object,
custom_llm_provider="azure",
litellm_model_name="azure/gpt-realtime-2.1-mini",
)
info = litellm.get_model_info(model="azure/gpt-realtime-2.1-mini", custom_llm_provider="azure")
expected = (
43 * info["input_cost_per_token"]
+ 194 * info["input_cost_per_image_token"]
+ 23 * info["output_cost_per_token"]
)
assert total_cost == pytest.approx(expected)
assert total_cost == pytest.approx(0.0002362)
def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() -> None:
"""The combined usage that lands in spend logs keeps reasoning out of text_tokens for every turn."""
results: OpenAIRealtimeStreamList = [

1
tests/test_litellm/test_main.py
View file

@ -3409,7 +3409,6 @@ def test_a_streamed_response_bills_the_usage_the_provider_reported(local_cost_ma
cost = litellm.completion_cost(completion_response=rebuilt, model=STREAM_COST_MODEL)
assert cost == pytest.approx(_priced_at(137, 42))
assert cost == pytest.approx(0.0007625)
def test_streaming_and_not_streaming_bill_the_same_usage_the_same(local_cost_map):

7
tests/test_litellm/test_muse_spark_1_3_model_metadata.py
View file

@ -31,13 +31,6 @@ def test_muse_spark_1_3_routes_to_meta_model_api(model: str):
assert api_base == "https://api.meta.ai/v1"
@pytest.mark.parametrize("model", (MUSE_SPARK_STANDARD, MUSE_SPARK_CONTRIBUTOR))
def test_muse_spark_1_3_web_search_cost_per_query(local_model_cost_map, model: str):
info = litellm.get_model_info(model=model)
assert StandardBuiltInToolCostTracking.get_cost_for_web_search(model_info=info) == WEB_SEARCH_COST_PER_QUERY
@pytest.mark.parametrize("model", (MUSE_SPARK_STANDARD, MUSE_SPARK_CONTRIBUTOR))
def test_muse_spark_1_3_backup_matches_main(model: str):
"""Ensure the bundled model cost map stays in sync with the canonical file."""

15
tests/test_litellm/test_openai_service_tier_long_context_pricing.py
View file

@ -91,18 +91,3 @@ TIERED_COST_CASES = [
("gpt-5.6-luna", "priority", 8e-07, 3.6e-06),
("gpt-6-astra", "priority", 4e-05, 0.00015),
]
@pytest.mark.parametrize("model,tier,input_rate,output_rate", TIERED_COST_CASES)
def test_cost_per_token_bills_long_context_at_the_tier_rate(
model: str, tier: str, input_rate: float, output_rate: float
) -> None:
"""A prompt over 272K on flex or priority must bill at that tier's long-context rate."""
input_cost, output_cost = litellm.cost_per_token(
model=model,
prompt_tokens=LONG_CONTEXT_PROMPT_TOKENS,
completion_tokens=COMPLETION_TOKENS,
service_tier=tier,
)
assert input_cost == pytest.approx(LONG_CONTEXT_PROMPT_TOKENS * input_rate)
assert output_cost == pytest.approx(COMPLETION_TOKENS * output_rate)

121
tests/test_litellm/test_video_generation.py
View file

@ -235,37 +235,6 @@ class TestVideoGeneration:
assert response.status == "completed"
assert response.model == "sora-2"
def test_video_generation_cost_calculation(self):
"""Test video generation cost calculation."""
import json
# Try to load the local model cost map, skip if not found
cost_map_path = "model_prices_and_context_window.json"
if not os.path.exists(cost_map_path):
# Try alternative paths
alt_paths = [
os.path.join(os.path.dirname(__file__), "..", "..", cost_map_path),
os.path.join(
os.path.dirname(__file__), "..", "..", "..", cost_map_path
),
]
for path in alt_paths:
if os.path.exists(path):
cost_map_path = path
break
else:
pytest.skip("model_prices_and_context_window.json not found")
with open(cost_map_path, "r") as f:
litellm.model_cost = json.load(f)
# Test with sora-2 model
cost = default_video_cost_calculator(
model="openai/sora-2", duration_seconds=10.0, custom_llm_provider="openai"
)
# Should calculate cost based on duration (10 seconds * $0.10 per second = $1.00)
assert cost == 1.0
def test_video_generation_cost_calculation_unknown_model(self):
"""Test video generation cost calculation for unknown model."""
@ -502,96 +471,6 @@ class TestVideoGeneration:
)
assert abs(cost - 1.8) < 0.001
def test_completion_cost_video_resolution_tiers_from_cost_map(self, monkeypatch):
"""The 480p/1080p/4k tier keys resolve from the shipped runwayml cost map entries."""
from litellm.cost_calculator import completion_cost
local_map_path = os.path.join(
os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json"
)
with open(local_map_path, "r") as f:
monkeypatch.setattr(litellm, "model_cost", json.load(f))
def cost_for(model: str, resolution: str | None, duration: float) -> float:
mock_response = MagicMock()
mock_response.usage = {
"duration_seconds": duration,
**({"video_resolution": resolution} if resolution else {}),
}
type(mock_response)._hidden_params = {}
return completion_cost(
completion_response=mock_response,
model=model,
call_type="create_video",
custom_llm_provider="runwayml",
)
assert abs(cost_for("runwayml/seedance2", "4k", 8.0) - 12.0) < 0.001
assert abs(cost_for("runwayml/seedance2", "1080p", 8.0) - 3.2) < 0.001
assert abs(cost_for("runwayml/seedance2", "720p", 8.0) - 2.88) < 0.001
assert abs(cost_for("runwayml/seedance2_5", "480p", 8.0) - 1.6) < 0.001
assert abs(cost_for("runwayml/gen4.5", None, 8.0) - 0.96) < 0.001
def test_completion_cost_xai_imagine_video_720p_tier_from_cost_map(self, monkeypatch):
"""720p xAI Imagine Video requests bill the published 720p rate, not the 480p base rate."""
from litellm.cost_calculator import completion_cost
local_map_path = os.path.join(
os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json"
)
with open(local_map_path, "r") as f:
monkeypatch.setattr(litellm, "model_cost", json.load(f))
def cost_for(model: str, resolution: str, duration: float) -> float:
mock_response = MagicMock()
mock_response.usage = {"duration_seconds": duration, "video_resolution": resolution}
type(mock_response)._hidden_params = {}
return completion_cost(
completion_response=mock_response,
model=model,
call_type="create_video",
custom_llm_provider="xai",
)
assert abs(cost_for("xai/grok-imagine-video", "720p", 10.0) - 0.7) < 0.001
assert abs(cost_for("xai/grok-imagine-video-1.5", "720p", 10.0) - 1.4) < 0.001
assert abs(cost_for("xai/grok-imagine-video-1.5", "480p", 10.0) - 0.8) < 0.001
assert abs(cost_for("xai/grok-imagine-video-1.5", "1080p", 10.0) - 2.5) < 0.001
def test_completion_cost_veo_31_tiers_pin_published_rates(self, monkeypatch):
"""The gemini and vertex_ai veo 3.1 entries bill Google's published per-second tier rates."""
from litellm.cost_calculator import completion_cost
local_map_path = os.path.join(
os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json"
)
with open(local_map_path, "r") as f:
monkeypatch.setattr(litellm, "model_cost", json.load(f))
def cost_for(model: str, provider: str, resolution: str | None, duration: float) -> float:
mock_response = MagicMock()
mock_response.usage = {
"duration_seconds": duration,
**({"video_resolution": resolution} if resolution else {}),
}
type(mock_response)._hidden_params = {}
return completion_cost(
completion_response=mock_response,
model=model,
call_type="create_video",
custom_llm_provider=provider,
)
for provider in ("gemini", "vertex_ai"):
for suffix in ("generate-preview", "generate-001"):
standard = f"{provider}/veo-3.1-{suffix}"
fast = f"{provider}/veo-3.1-fast-{suffix}"
assert abs(cost_for(standard, provider, None, 8.0) - 3.2) < 1e-6
assert abs(cost_for(standard, provider, "1080p", 8.0) - 3.2) < 1e-6
assert abs(cost_for(standard, provider, "4k", 8.0) - 4.8) < 1e-6
assert abs(cost_for(fast, provider, "720p", 8.0) - 0.8) < 1e-6
assert abs(cost_for(fast, provider, "1080p", 8.0) - 0.96) < 1e-6
assert abs(cost_for(fast, provider, "4k", 8.0) - 2.4) < 1e-6
def test_video_generation_with_files(self):
"""Test video generation with file uploads."""