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

Merge pull request #14979 from BerriAI/litellm_stabilize_ci_cd_09_27_2025_p1

Browse source 
stabilize main github tests
This commit is contained in:
Krish Dholakia 2025-09-27 10:34:37 -07:00 • committed by GitHub
commit ea6833b18d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 136 additions and 72 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

171
litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py
View file

@ -47,7 +47,7 @@ class StandardBuiltInToolCostTracking:
- Code Interpreter (Azure)
"""
standard_built_in_tools_params = standard_built_in_tools_params or {}
# Handle web search
if StandardBuiltInToolCostTracking.response_object_includes_web_search_call(
response_object=response_object, usage=usage
@ -58,7 +58,7 @@ class StandardBuiltInToolCostTracking:
usage=usage,
standard_built_in_tools_params=standard_built_in_tools_params,
)
# Handle file search
if StandardBuiltInToolCostTracking.response_object_includes_file_search_call(
response_object=response_object
@ -68,7 +68,7 @@ class StandardBuiltInToolCostTracking:
custom_llm_provider=custom_llm_provider,
standard_built_in_tools_params=standard_built_in_tools_params,
)
# Handle Azure assistant features
return StandardBuiltInToolCostTracking._handle_azure_assistant_costs(
model=model,
@ -85,14 +85,14 @@ class StandardBuiltInToolCostTracking:
) -> float:
"""Handle web search cost calculation."""
from litellm.llms import get_cost_for_web_search_request
model_info = StandardBuiltInToolCostTracking._safe_get_model_info(
model=model, custom_llm_provider=custom_llm_provider
)
if custom_llm_provider is None and model_info is not None:
custom_llm_provider = model_info["litellm_provider"]
if (
model_info is not None
and usage is not None
@ -105,9 +105,11 @@ class StandardBuiltInToolCostTracking:
)
if result is not None:
return result
return StandardBuiltInToolCostTracking.get_cost_for_web_search(
web_search_options=standard_built_in_tools_params.get("web_search_options", None),
web_search_options=standard_built_in_tools_params.get(
"web_search_options", None
),
model_info=model_info,
)
@ -121,12 +123,17 @@ class StandardBuiltInToolCostTracking:
model_info = StandardBuiltInToolCostTracking._safe_get_model_info(
model=model, custom_llm_provider=custom_llm_provider
)
file_search_usage = standard_built_in_tools_params.get("file_search", {})
file_search_raw = standard_built_in_tools_params.get("file_search", {})
file_search_usage: Optional[FileSearchTool] = (
FileSearchTool(**file_search_raw) if file_search_raw else None
)
# Convert model_info to dict and extract usage parameters
model_info_dict = dict(model_info) if model_info is not None else None
storage_gb, days = StandardBuiltInToolCostTracking._extract_file_search_params(file_search_usage)
storage_gb, days = StandardBuiltInToolCostTracking._extract_file_search_params(
file_search_usage
)
return StandardBuiltInToolCostTracking.get_cost_for_file_search(
file_search=file_search_usage,
provider=custom_llm_provider,
@ -144,11 +151,11 @@ class StandardBuiltInToolCostTracking:
"""Handle Azure assistant features cost calculation."""
if custom_llm_provider != "azure":
return 0.0
model_info = StandardBuiltInToolCostTracking._safe_get_model_info(
model=model, custom_llm_provider=custom_llm_provider
)
total_cost = 0.0
total_cost += StandardBuiltInToolCostTracking._get_vector_store_cost(
model_info, custom_llm_provider, standard_built_in_tools_params
@ -159,31 +166,33 @@ class StandardBuiltInToolCostTracking:
total_cost += StandardBuiltInToolCostTracking._get_code_interpreter_cost(
model_info, custom_llm_provider, standard_built_in_tools_params
)
return total_cost
@staticmethod
def _extract_file_search_params(file_search_usage: Any) -> Tuple[Optional[float], Optional[float]]:
def _extract_file_search_params(
file_search_usage: Any,
) -> Tuple[Optional[float], Optional[float]]:
"""Extract and convert file search parameters safely."""
storage_gb = None
days = None
if isinstance(file_search_usage, dict):
storage_gb_val = file_search_usage.get("storage_gb")
days_val = file_search_usage.get("days")
if storage_gb_val is not None:
try:
storage_gb = float(storage_gb_val) # type: ignore
except (TypeError, ValueError):
storage_gb = None
if days_val is not None:
try:
days = float(days_val) # type: ignore
except (TypeError, ValueError):
days = None
return storage_gb, days
@staticmethod
@ -193,13 +202,17 @@ class StandardBuiltInToolCostTracking:
standard_built_in_tools_params: StandardBuiltInToolsParams,
) -> float:
"""Calculate vector store cost."""
vector_store_usage = standard_built_in_tools_params.get("vector_store_usage", None)
vector_store_usage = standard_built_in_tools_params.get(
"vector_store_usage", None
)
if not vector_store_usage:
return 0.0
model_info_dict = dict(model_info) if model_info is not None else None
vector_store_dict = vector_store_usage if isinstance(vector_store_usage, dict) else {}
vector_store_dict = (
vector_store_usage if isinstance(vector_store_usage, dict) else {}
)
return StandardBuiltInToolCostTracking.get_cost_for_vector_store(
vector_store_usage=vector_store_dict,
provider=custom_llm_provider,
@ -213,13 +226,17 @@ class StandardBuiltInToolCostTracking:
standard_built_in_tools_params: StandardBuiltInToolsParams,
) -> float:
"""Calculate computer use cost."""
computer_use_usage = standard_built_in_tools_params.get("computer_use_usage", {})
computer_use_usage = standard_built_in_tools_params.get(
"computer_use_usage", {}
)
if not computer_use_usage:
return 0.0
model_info_dict = dict(model_info) if model_info is not None else None
input_tokens, output_tokens = StandardBuiltInToolCostTracking._extract_token_counts(computer_use_usage)
input_tokens, output_tokens = (
StandardBuiltInToolCostTracking._extract_token_counts(computer_use_usage)
)
return StandardBuiltInToolCostTracking.get_cost_for_computer_use(
input_tokens=input_tokens,
output_tokens=output_tokens,
@ -234,13 +251,17 @@ class StandardBuiltInToolCostTracking:
standard_built_in_tools_params: StandardBuiltInToolsParams,
) -> float:
"""Calculate code interpreter cost."""
code_interpreter_sessions = standard_built_in_tools_params.get("code_interpreter_sessions", None)
code_interpreter_sessions = standard_built_in_tools_params.get(
"code_interpreter_sessions", None
)
if not code_interpreter_sessions:
return 0.0
model_info_dict = dict(model_info) if model_info is not None else None
sessions = StandardBuiltInToolCostTracking._safe_convert_to_int(code_interpreter_sessions)
sessions = StandardBuiltInToolCostTracking._safe_convert_to_int(
code_interpreter_sessions
)
return StandardBuiltInToolCostTracking.get_cost_for_code_interpreter(
sessions=sessions,
provider=custom_llm_provider,
@ -248,18 +269,24 @@ class StandardBuiltInToolCostTracking:
)
@staticmethod
def _extract_token_counts(computer_use_usage: Any) -> Tuple[Optional[int], Optional[int]]:
def _extract_token_counts(
computer_use_usage: Any,
) -> Tuple[Optional[int], Optional[int]]:
"""Extract and convert token counts safely."""
input_tokens = None
output_tokens = None
if isinstance(computer_use_usage, dict):
input_tokens_val = computer_use_usage.get("input_tokens")
output_tokens_val = computer_use_usage.get("output_tokens")
input_tokens = StandardBuiltInToolCostTracking._safe_convert_to_int(input_tokens_val)
output_tokens = StandardBuiltInToolCostTracking._safe_convert_to_int(output_tokens_val)
input_tokens = StandardBuiltInToolCostTracking._safe_convert_to_int(
input_tokens_val
)
output_tokens = StandardBuiltInToolCostTracking._safe_convert_to_int(
output_tokens_val
)
return input_tokens, output_tokens
@staticmethod
@ -400,8 +427,11 @@ class StandardBuiltInToolCostTracking:
if model_info is None:
return 0.0
search_context_raw = model_info.get("search_context_cost_per_query", {})
search_context_pricing: SearchContextCostPerQuery = (
model_info.get("search_context_cost_per_query", {}) or {}
SearchContextCostPerQuery(**search_context_raw)
if search_context_raw
else SearchContextCostPerQuery()
)
if web_search_options.get("search_context_size", None) == "low":
return search_context_pricing.get("search_context_size_low", 0.0)
@ -424,9 +454,12 @@ class StandardBuiltInToolCostTracking:
"""
if model_info is None:
return 0.0
search_context_raw = model_info.get("search_context_cost_per_query", {}) or {}
search_context_pricing: SearchContextCostPerQuery = (
model_info.get("search_context_cost_per_query", {}) or {}
) or {}
SearchContextCostPerQuery(**search_context_raw)
if search_context_raw
else SearchContextCostPerQuery()
)
return search_context_pricing.get("search_context_size_medium", 0.0)
@staticmethod
@ -445,22 +478,27 @@ class StandardBuiltInToolCostTracking:
"""
if file_search is None:
return 0.0
# Check if model-specific pricing is available
if model_info and "file_search_cost_per_gb_per_day" in model_info and provider == "azure":
if (
model_info
and "file_search_cost_per_gb_per_day" in model_info
and provider == "azure"
):
if storage_gb and days:
return storage_gb * days * model_info["file_search_cost_per_gb_per_day"]
elif model_info and "file_search_cost_per_1k_calls" in model_info:
return model_info["file_search_cost_per_1k_calls"]
# Azure has storage-based pricing for file search
if provider == "azure":
from litellm.constants import AZURE_FILE_SEARCH_COST_PER_GB_PER_DAY
if storage_gb and days:
return storage_gb * days * AZURE_FILE_SEARCH_COST_PER_GB_PER_DAY
# Default to 0 if no storage info provided
return 0.0
# Default to OpenAI pricing (per-call based)
return OPENAI_FILE_SEARCH_COST_PER_1K_CALLS
@ -472,24 +510,25 @@ class StandardBuiltInToolCostTracking:
) -> float:
"""
Calculate cost for vector store usage.
Azure charges based on storage size and duration.
"""
if vector_store_usage is None:
return 0.0
storage_gb = vector_store_usage.get("storage_gb", 0.0)
days = vector_store_usage.get("days", 0.0)
# Check if model-specific pricing is available
if model_info and "vector_store_cost_per_gb_per_day" in model_info:
return storage_gb * days * model_info["vector_store_cost_per_gb_per_day"]
# Azure has different pricing structure for vector store
if provider == "azure":
from litellm.constants import AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY
return storage_gb * days * AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY
# OpenAI doesn't charge separately for vector store (included in embeddings)
return 0.0
@ -502,14 +541,18 @@ class StandardBuiltInToolCostTracking:
) -> float:
"""
Calculate cost for computer use feature.
Azure: $0.003 USD per 1K input tokens, $0.012 USD per 1K output tokens
"""
if provider == "azure" and (input_tokens or output_tokens):
# Check if model-specific pricing is available
if model_info:
input_cost = model_info.get("computer_use_input_cost_per_1k_tokens", 0.0)
output_cost = model_info.get("computer_use_output_cost_per_1k_tokens", 0.0)
input_cost = model_info.get(
"computer_use_input_cost_per_1k_tokens", 0.0
)
output_cost = model_info.get(
"computer_use_output_cost_per_1k_tokens", 0.0
)
if input_cost or output_cost:
total_cost = 0.0
if input_tokens:
@ -517,19 +560,24 @@ class StandardBuiltInToolCostTracking:
if output_tokens:
total_cost += (output_tokens / 1000.0) * output_cost
return total_cost
# Azure default pricing
from litellm.constants import (
AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS,
AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS,
)
total_cost = 0.0
if input_tokens:
total_cost += (input_tokens / 1000.0) * AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS
total_cost += (
input_tokens / 1000.0
) * AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS
if output_tokens:
total_cost += (output_tokens / 1000.0) * AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS
total_cost += (
output_tokens / 1000.0
) * AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS
return total_cost
# OpenAI doesn't charge separately for computer use yet
return 0.0
@ -541,21 +589,22 @@ class StandardBuiltInToolCostTracking:
) -> float:
"""
Calculate cost for code interpreter feature.
Azure: $0.03 USD per session
"""
if sessions is None or sessions == 0:
return 0.0
# Check if model-specific pricing is available
if model_info and "code_interpreter_cost_per_session" in model_info:
return sessions * model_info["code_interpreter_cost_per_session"]
# Azure pricing for code interpreter
if provider == "azure":
from litellm.constants import AZURE_CODE_INTERPRETER_COST_PER_SESSION
return sessions * AZURE_CODE_INTERPRETER_COST_PER_SESSION
# OpenAI doesn't charge separately for code interpreter yet
return 0.0

37
tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py
View file

@ -1,8 +1,19 @@
import sys
import os
import sys
import pytest
from litellm.llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig
from litellm.types.rerank import OptionalRerankParams, RerankResponse, RerankResponseResult, RerankResponseMeta, RerankBilledUnits, RerankTokens, RerankResponseDocument
from litellm.types.rerank import (
OptionalRerankParams,
RerankBilledUnits,
RerankResponse,
RerankResponseDocument,
RerankResponseMeta,
RerankResponseResult,
RerankTokens,
)
class TestHostedVLLMRerankTransform:
def setup_method(self):
@ -27,23 +38,27 @@ class TestHostedVLLMRerankTransform:
assert params["return_documents"] is True
def test_map_cohere_rerank_params_raises_on_max_chunks_per_doc(self):
with pytest.raises(ValueError, match="Hosted VLLM does not support max_chunks_per_doc"):
with pytest.raises(
ValueError, match="Hosted VLLM does not support max_chunks_per_doc"
):
self.config.map_cohere_rerank_params(
non_default_params=None,
model=self.model,
drop_params=False,
query="test query",
documents=["doc1"],
max_chunks_per_doc=5
max_chunks_per_doc=5,
)
def test_get_complete_url(self):
base = "https://api.example.com"
url = self.config.get_complete_url(base, self.model)
assert url == "https://api.example.com/v1/rerank"
# Already ends with /v1/rerank
url2 = self.config.get_complete_url("https://api.example.com/v1/rerank", self.model)
assert url2 == "https://api.example.com/v1/rerank"
assert url == "https://api.example.com/rerank"
# Already ends with /rerank
url2 = self.config.get_complete_url(
"https://api.example.com/rerank", self.model
)
assert url2 == "https://api.example.com/rerank"
# Raises if api_base is None
with pytest.raises(ValueError):
self.config.get_complete_url(None, self.model)
@ -55,7 +70,7 @@ class TestHostedVLLMRerankTransform:
{"index": 0, "relevance_score": 0.9, "document": {"text": "doc1 text"}},
{"index": 1, "relevance_score": 0.7, "document": {"text": "doc2 text"}},
],
"usage": {"total_tokens": 42}
"usage": {"total_tokens": 42},
}
result = self.config._transform_response(response_dict)
assert result.id == "abc123"
@ -75,7 +90,7 @@ class TestHostedVLLMRerankTransform:
response_dict = {
"id": "abc123",
"results": [{"relevance_score": 0.5}],
"usage": {"total_tokens": 10}
"usage": {"total_tokens": 10},
}
with pytest.raises(ValueError, match="Missing required fields in the result="):
self.config._transform_response(response_dict)
self.config._transform_response(response_dict)