ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Fix] CI/CD – Clean Up Performance PR Changes & others (#17838)

Browse Source
This commit is contained in:
Alexsander Hamir 2025-12-11 12:50:03 -08:00 committed by GitHub
parent 70643a8b9c
commit e9baa83a0f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
22 changed files with 252 additions and 108 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.circleci/config.yml
View File

@ -52,6 +52,7 @@ commands:
pip install "pytest-timeout==2.2.0"
pip install "semantic_router==0.1.10"
pip install "fastapi-offline==1.7.3"
pip install "a2a"
- setup_litellm_enterprise_pip
- save_cache:
paths:

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20251211100212_schema_sync/migration.sql Normal file
View File

@ -0,0 +1,3 @@
-- AlterTable
ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "agent_id" TEXT;

9
litellm/a2a_protocol/main.py
View File

@ -26,7 +26,6 @@ if TYPE_CHECKING:
AgentCard,
SendMessageRequest,
SendStreamingMessageRequest,
SendStreamingMessageResponse,
)
# Runtime imports with availability check
@ -219,6 +218,9 @@ async def asend_message(
raise ValueError("Either a2a_client or api_base is required for standard A2A flow")
a2a_client = await create_a2a_client(base_url=api_base)
# Type assertion: a2a_client is guaranteed to be non-None here
assert a2a_client is not None
agent_name = _get_a2a_model_info(a2a_client, kwargs)
verbose_logger.info(f"A2A send_message request_id={request.id}, agent={agent_name}")
@ -365,11 +367,12 @@ async def asend_message_streaming(
raise ValueError("Either a2a_client or api_base is required for standard A2A flow")
a2a_client = await create_a2a_client(base_url=api_base)
# Type assertion: a2a_client is guaranteed to be non-None here
assert a2a_client is not None
verbose_logger.info(f"A2A send_message_streaming request_id={request.id}")
# Track for logging
import datetime
start_time = datetime.datetime.now()
stream = a2a_client.send_message_streaming(request)

101
litellm/litellm_core_utils/litellm_logging.py
View File

@ -4821,6 +4821,63 @@ def _get_status_fields(
)
def _extract_response_obj_and_hidden_params(
init_response_obj: Union[Any, BaseModel, dict],
original_exception: Optional[Exception],
) -> Tuple[dict, Optional[dict]]:
"""Extract response_obj and hidden_params from init_response_obj."""
hidden_params: Optional[dict] = None
if init_response_obj is None:
response_obj = {}
elif isinstance(init_response_obj, BaseModel):
response_obj = init_response_obj.model_dump()
hidden_params = getattr(init_response_obj, "_hidden_params", None)
elif isinstance(init_response_obj, dict):
response_obj = init_response_obj
else:
response_obj = {}
if original_exception is not None and hidden_params is None:
response_headers = _get_response_headers(original_exception)
if response_headers is not None:
hidden_params = dict(
StandardLoggingHiddenParams(
additional_headers=StandardLoggingPayloadSetup.get_additional_headers(
dict(response_headers)
),
model_id=None,
cache_key=None,
api_base=None,
response_cost=None,
litellm_overhead_time_ms=None,
batch_models=None,
litellm_model_name=None,
usage_object=None,
)
)
return response_obj, hidden_params
def _reconstruct_model_name(
model_name: str,
custom_llm_provider: Optional[str],
metadata: dict,
) -> str:
"""Reconstruct full model name with provider prefix for logging."""
# Check if deployment model name from router metadata is available (has original prefix)
deployment_model_name = metadata.get("deployment")
if deployment_model_name and "/" in deployment_model_name:
# Use the deployment model name which preserves the original provider prefix
return deployment_model_name
elif custom_llm_provider and model_name and "/" not in model_name:
# Only add prefix for Bedrock (not for direct Anthropic API)
# This ensures Bedrock models get the prefix while direct Anthropic models don't
if custom_llm_provider == "bedrock":
return f"{custom_llm_provider}/{model_name}"
return model_name
def get_standard_logging_object_payload(
kwargs: Optional[dict],
init_response_obj: Union[Any, BaseModel, dict],
@ -4835,35 +4892,9 @@ def get_standard_logging_object_payload(
try:
kwargs = kwargs or {}
hidden_params: Optional[dict] = None
if init_response_obj is None:
response_obj = {}
elif isinstance(init_response_obj, BaseModel):
response_obj = init_response_obj.model_dump()
hidden_params = getattr(init_response_obj, "_hidden_params", None)
elif isinstance(init_response_obj, dict):
response_obj = init_response_obj
else:
response_obj = {}
if original_exception is not None and hidden_params is None:
response_headers = _get_response_headers(original_exception)
if response_headers is not None:
hidden_params = dict(
StandardLoggingHiddenParams(
additional_headers=StandardLoggingPayloadSetup.get_additional_headers(
dict(response_headers)
),
model_id=None,
cache_key=None,
api_base=None,
response_cost=None,
litellm_overhead_time_ms=None,
batch_models=None,
litellm_model_name=None,
usage_object=None,
)
)
response_obj, hidden_params = _extract_response_obj_and_hidden_params(
init_response_obj, original_exception
)
# standardize this function to be used across, s3, dynamoDB, langfuse logging
litellm_params = kwargs.get("litellm_params", {}) or {}
@ -4975,6 +5006,14 @@ def get_standard_logging_object_payload(
) and kwargs.get("stream") is True:
stream = True
# Reconstruct full model name with provider prefix for logging
# This ensures Bedrock models like "us.anthropic.claude-3-5-sonnet-20240620-v1:0"
# are logged as "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0"
custom_llm_provider = cast(Optional[str], kwargs.get("custom_llm_provider"))
model_name = _reconstruct_model_name(
kwargs.get("model", "") or "", custom_llm_provider, metadata
)
payload: StandardLoggingPayload = StandardLoggingPayload(
id=str(id),
trace_id=StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id(
@ -4992,13 +5031,13 @@ def get_standard_logging_object_payload(
),
error_str=error_str,
),
custom_llm_provider=cast(Optional[str], kwargs.get("custom_llm_provider")),
custom_llm_provider=custom_llm_provider,
saved_cache_cost=saved_cache_cost,
startTime=start_time_float,
endTime=end_time_float,
completionStartTime=completion_start_time_float,
response_time=response_time,
model=kwargs.get("model", "") or "",
model=model_name,
metadata=clean_metadata,
cache_key=clean_hidden_params["cache_key"],
response_cost=response_cost,

6
litellm/llms/langgraph/chat/transformation.py
View File

@ -148,7 +148,7 @@ class LangGraphConfig(BaseConfig):
OpenAI format: {"role": "user", "content": "..."}
LangGraph format: {"role": "human", "content": "..."}
"""
langgraph_messages = []
langgraph_messages: List[Dict[str, str]] = []
for msg in messages:
role = msg.get("role", "user")
content = msg.get("content", "")
@ -166,6 +166,10 @@ class LangGraphConfig(BaseConfig):
# Handle content that might be a list
if isinstance(content, list):
content = convert_content_list_to_str(msg)
# Ensure content is a string
if not isinstance(content, str):
content = str(content)
langgraph_messages.append({"role": langgraph_role, "content": content})

5
litellm/llms/voyage/rerank/transformation.py
View File

@ -12,7 +12,6 @@ from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.rerank import (
OptionalRerankParams,
RerankBilledUnits,
RerankResponse,
RerankResponseMeta,
@ -48,7 +47,9 @@ class VoyageRerankConfig(BaseRerankConfig):
optional_params["top_k"] = top_n
if return_documents is not None:
optional_params["return_documents"] = return_documents
return dict(OptionalRerankParams(**optional_params))
# Return as dict - OptionalRerankParams is a TypedDict with total=False
# so all fields are optional and we can return the dict directly
return optional_params
def get_complete_url(
self,

2
litellm/proxy/agent_endpoints/a2a_endpoints.py
View File

@ -6,7 +6,7 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM
"""
import json
from typing import Any, Optional
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import JSONResponse, StreamingResponse

5
litellm/proxy/anthropic_endpoints/endpoints.py
View File

@ -9,7 +9,10 @@ from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
create_streaming_response,
)
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.types.utils import TokenCountResponse

6
litellm/proxy/db/db_spend_update_writer.py
View File

@ -418,9 +418,11 @@ class DBSpendUpdateWriter:
)
)
if prisma_client is not None and spend_logs_url is not None:
prisma_client.spend_log_transactions.append(payload)
async with prisma_client._spend_log_transactions_lock:
prisma_client.spend_log_transactions.append(payload)
elif prisma_client is not None:
prisma_client.spend_log_transactions.append(payload)
async with prisma_client._spend_log_transactions_lock:
prisma_client.spend_log_transactions.append(payload)
else:
verbose_proxy_logger.debug(
"prisma_client is None. Skipping writing spend logs to db."

11
litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py
View File

@ -20,7 +20,7 @@ from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import LLMResponseTypes
from litellm.types.utils import Choices, LLMResponseTypes, ModelResponse
class GraySwanGuardrailMissingSecrets(Exception):
@ -256,19 +256,22 @@ class GraySwanGuardrail(CustomGuardrail):
)
# Handle ModelResponse (OpenAI-style chat/text completions)
if hasattr(response, "choices") and response.choices:
# Use isinstance to narrow the type for mypy
if isinstance(response, ModelResponse) and response.choices:
verbose_proxy_logger.debug(
"Gray Swan Guardrail: Replacing response content in ModelResponse format"
)
for choice in response.choices:
# Handle chat completion format (message.content)
if hasattr(choice, "message") and hasattr(
# Choices has message attribute, StreamingChoices has delta
if isinstance(choice, Choices) and hasattr(choice, "message") and hasattr(
choice.message, "content"
):
choice.message.content = violation_message
# Handle text completion format (text)
# Text attribute might be set dynamically, use setattr
elif hasattr(choice, "text"):
choice.text = violation_message
setattr(choice, "text", violation_message)
# Update finish_reason to indicate content filtering
if hasattr(choice, "finish_reason"):

8
litellm/proxy/proxy_server.py
View File

@ -4452,7 +4452,7 @@ class ProxyStartupEvent:
### MONITOR SPEND LOGS QUEUE (queue-size-based job) ###
if general_settings.get("disable_spend_logs", False) is False:
from litellm.proxy.utils import _monitor_spend_logs_queue
# Start background task to monitor spend logs queue size
asyncio.create_task(
_monitor_spend_logs_queue(
@ -5144,14 +5144,16 @@ async def completion( # noqa: PLR0915
if _data.get("stream", None) is not None and _data["stream"] is True:
_text_response = litellm.ModelResponse()
_text_response.choices[0].text = e.message # type: ignore[attr-defined]
# Set text attribute dynamically for text completion format
setattr(_text_response.choices[0], "text", e.message)
_text_response.model = e.model # type: ignore[assignment]
_usage = litellm.Usage(
prompt_tokens=0,
completion_tokens=0,
total_tokens=0,
)
_text_response.usage = _usage # type: ignore[assignment]
# Set usage attribute dynamically (ModelResponse accepts usage in __init__ but it's not in type definition)
setattr(_text_response, "usage", _usage)
_iterator = litellm.utils.ModelResponseIterator(
model_response=_text_response, convert_to_delta=True
)

91
litellm/proxy/utils.py
View File

@ -1717,6 +1717,7 @@ def jsonify_object(data: dict) -> dict:
class PrismaClient:
spend_log_transactions: List = []
_spend_log_transactions_lock = asyncio.Lock()
def __init__(
self,
@ -3356,8 +3357,13 @@ class ProxyUpdateSpend:
MAX_LOGS_PER_INTERVAL = (
10000 # Maximum number of logs to flush in a single interval
)
# Get initial logs to proces
logs_to_process = prisma_client.spend_log_transactions[:MAX_LOGS_PER_INTERVAL]
# Atomically read and remove logs to process (protected by lock)
async with prisma_client._spend_log_transactions_lock:
logs_to_process = prisma_client.spend_log_transactions[:MAX_LOGS_PER_INTERVAL]
# Remove the logs we're about to process
prisma_client.spend_log_transactions = (
prisma_client.spend_log_transactions[len(logs_to_process):]
)
start_time = time.time()
try:
for i in range(n_retry_times + 1):
@ -3379,11 +3385,8 @@ class ProxyUpdateSpend:
)
del json_data
if response.status_code == 200:
prisma_client.spend_log_transactions = (
prisma_client.spend_log_transactions[
len(logs_to_process) :
]
)
# Items already removed from queue at start of function
pass
else:
for j in range(0, len(logs_to_process), BATCH_SIZE):
batch = logs_to_process[j : j + BATCH_SIZE]
@ -3400,10 +3403,9 @@ class ProxyUpdateSpend:
# Explicitly clear batch memory
del batch, batch_with_dates
prisma_client.spend_log_transactions = (
prisma_client.spend_log_transactions[len(logs_to_process) :]
)
remaining_count = len(prisma_client.spend_log_transactions)
# Items already removed from queue at start of function
async with prisma_client._spend_log_transactions_lock:
remaining_count = len(prisma_client.spend_log_transactions)
verbose_proxy_logger.debug(
f"{len(logs_to_process)} logs processed. Remaining in queue: {remaining_count}"
)
@ -3415,9 +3417,8 @@ class ProxyUpdateSpend:
raise
await asyncio.sleep(2**i)
except Exception as e:
prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[
len(logs_to_process) :
]
# Logs already removed from queue at start - don't put them back
# This matches the original behavior where logs are removed even on error
_raise_failed_update_spend_exception(
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
)
@ -3462,12 +3463,24 @@ async def update_spend( # noqa: PLR0915
)
### UPDATE SPEND LOGS ###
# Check queue size with lock protection
async with prisma_client._spend_log_transactions_lock:
queue_size = len(prisma_client.spend_log_transactions)
verbose_proxy_logger.debug(
"Spend Logs transactions: {}".format(len(prisma_client.spend_log_transactions))
"Spend Logs transactions: {}".format(queue_size)
)
# Spend log transactions are now processed by a separate queue-size-based job
# See update_spend_logs_job and _monitor_spend_logs_queue
# Process spend log transactions when called directly.
# This keeps backwards compatibility with the old behavior.
# See update_spend_logs_job and _monitor_spend_logs_queue for the new behavior.
# Safe to keep: under high concurrency this can take up to ~30s to run,
# so it's unlikely to overlap with monitor_spend_logs_queue.
if queue_size > 0:
await update_spend_logs_job(
prisma_client=prisma_client,
db_writer_client=db_writer_client,
proxy_logging_obj=proxy_logging_obj,
)
async def update_spend_logs_job(
@ -3477,17 +3490,19 @@ async def update_spend_logs_job(
):
"""
Job to process spend_log_transactions queue.
This job is triggered based on queue size rather than time.
Processes spend log transactions when the queue reaches a threshold.
"""
n_retry_times = 3
queue_size = len(prisma_client.spend_log_transactions)
# Check queue size with lock protection
async with prisma_client._spend_log_transactions_lock:
queue_size = len(prisma_client.spend_log_transactions)
if queue_size == 0:
return
await ProxyUpdateSpend.update_spend_logs(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
@ -3504,31 +3519,30 @@ async def _monitor_spend_logs_queue(
"""
Background task that monitors the spend_log_transactions queue size
and triggers processing when the threshold is reached.
Args:
prisma_client: Prisma client instance
db_writer_client: Optional HTTP handler for external spend logs endpoint
proxy_logging_obj: Proxy logging object
"""
from litellm.constants import (
SPEND_LOG_QUEUE_POLL_INTERVAL,
SPEND_LOG_QUEUE_SIZE_THRESHOLD,
)
from litellm.constants import SPEND_LOG_QUEUE_SIZE_THRESHOLD, SPEND_LOG_QUEUE_POLL_INTERVAL
threshold = SPEND_LOG_QUEUE_SIZE_THRESHOLD
base_interval = SPEND_LOG_QUEUE_POLL_INTERVAL
max_backoff = 30.0 # Maximum backoff interval in seconds
backoff_multiplier = 1.5 # Exponential backoff multiplier
current_interval = base_interval
verbose_proxy_logger.info(
f"Starting spend logs queue monitor (threshold: {threshold}, poll_interval: {base_interval}s)"
)
while True:
try:
queue_size = len(prisma_client.spend_log_transactions)
# Check queue size with lock protection
async with prisma_client._spend_log_transactions_lock:
queue_size = len(prisma_client.spend_log_transactions)
if queue_size > 0:
if queue_size >= threshold:
verbose_proxy_logger.debug(
@ -3541,10 +3555,8 @@ async def _monitor_spend_logs_queue(
f"Spend logs queue size ({queue_size}) below threshold ({threshold}), processing with backoff"
)
# Exponential backoff when below threshold but still processing
current_interval = min(
current_interval * backoff_multiplier, max_backoff
)
current_interval = min(current_interval * backoff_multiplier, max_backoff)
await update_spend_logs_job(
prisma_client=prisma_client,
db_writer_client=db_writer_client,
@ -3552,10 +3564,8 @@ async def _monitor_spend_logs_queue(
)
else:
# Exponential backoff when no logs to process
current_interval = min(
current_interval * backoff_multiplier, max_backoff
)
current_interval = min(current_interval * backoff_multiplier, max_backoff)
await asyncio.sleep(current_interval)
except Exception as e:
verbose_proxy_logger.error(
@ -3566,6 +3576,7 @@ async def _monitor_spend_logs_queue(
await asyncio.sleep(current_interval)
def _raise_failed_update_spend_exception(
e: Exception, start_time: float, proxy_logging_obj: ProxyLogging
):

6
tests/proxy_unit_tests/test_key_generate_prisma.py
View File

@ -665,7 +665,8 @@ def test_call_with_end_user_over_budget(prisma_client):
asyncio.run(test())
except Exception as e:
print(f"raised error: {e}, traceback: {traceback.format_exc()}")
error_detail = e.message
# Handle DataError and other exceptions that don't have .message attribute
error_detail = getattr(e, 'message', str(e))
assert "ExceededBudget: End User=" in error_detail
assert "over budget" in error_detail
assert isinstance(e, ProxyException)
@ -2081,7 +2082,8 @@ async def test_call_with_key_over_budget_stream(prisma_client):
except Exception as e:
print("Got Exception", e)
error_detail = e.message
# Handle DataError and other exceptions that don't have .message attribute
error_detail = getattr(e, 'message', str(e))
assert "Budget has been exceeded" in error_detail
print(vars(e))

3
tests/proxy_unit_tests/test_proxy_utils.py
View File

@ -1629,12 +1629,15 @@ async def test_end_user_transactions_reset():
@pytest.mark.asyncio
async def test_spend_logs_cleanup_after_error():
# Setup test data
import asyncio
mock_client = MagicMock()
mock_client.spend_log_transactions = [
{"id": 1, "amount": 10.0},
{"id": 2, "amount": 20.0},
{"id": 3, "amount": 30.0},
]
# Add lock for spend_log_transactions (matches real PrismaClient)
mock_client._spend_log_transactions_lock = asyncio.Lock()
# Make the DB operation fail
mock_client.db.litellm_spendlogs.create_many = AsyncMock(
side_effect=Exception("DB Error")

3
tests/proxy_unit_tests/test_unit_test_proxy_hooks.py
View File

@ -17,8 +17,11 @@ async def test_disable_spend_logs():
Test that the spend logs are not written to the database when disable_spend_logs is True
"""
# Mock the necessary components
import asyncio
mock_prisma_client = Mock()
mock_prisma_client.spend_log_transactions = []
# Add lock for spend_log_transactions (matches real PrismaClient)
mock_prisma_client._spend_log_transactions_lock = asyncio.Lock()
with patch("litellm.proxy.proxy_server.disable_spend_logs", True), patch(
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client

28
tests/proxy_unit_tests/test_update_spend.py
View File

@ -28,6 +28,10 @@ class MockPrismaClient:
# Initialize transaction lists
self.spend_log_transactions = []
self.daily_user_spend_transactions = {}
# Add lock for spend_log_transactions (matches real PrismaClient)
import asyncio
self._spend_log_transactions_lock = asyncio.Lock()
def jsonify_object(self, obj):
return obj
@ -207,15 +211,15 @@ async def test_update_spend_logs_multiple_batches_success():
"""
Test successful processing of multiple batches of spend logs
Code sets batch size to 100. This test creates 150 logs, so it should make 2 batches.
Code sets batch size to 1000. This test creates 1500 logs, so it should make 2 batches.
"""
# Setup
prisma_client = MockPrismaClient()
proxy_logging_obj = create_mock_proxy_logging()
# Create 150 test spend logs (1.5x BATCH_SIZE)
# Create 1500 test spend logs (1.5x BATCH_SIZE)
prisma_client.spend_log_transactions = [
{"id": str(i), "spend": 10} for i in range(150)
{"id": str(i), "spend": 10} for i in range(1500)
]
create_many_mock = AsyncMock(return_value=None)
@ -232,12 +236,12 @@ async def test_update_spend_logs_multiple_batches_success():
second_batch = create_many_mock.call_args_list[1][1]["data"]
# Verify batch sizes
assert len(first_batch) == 100
assert len(second_batch) == 50
assert len(first_batch) == 1000
assert len(second_batch) == 500
# Verify exact IDs in each batch
expected_first_batch_ids = {str(i) for i in range(100)}
expected_second_batch_ids = {str(i) for i in range(100, 150)}
expected_first_batch_ids = {str(i) for i in range(1000)}
expected_second_batch_ids = {str(i) for i in range(1000, 1500)}
actual_first_batch_ids = {item["id"] for item in first_batch}
actual_second_batch_ids = {item["id"] for item in second_batch}
@ -253,15 +257,15 @@ async def test_update_spend_logs_multiple_batches_success():
async def test_update_spend_logs_multiple_batches_with_failure():
"""
Test processing of multiple batches where one batch fails.
Creates 400 logs (4 batches) with one batch failing but eventually succeeding after retry.
Creates 4000 logs (4 batches) with one batch failing but eventually succeeding after retry.
"""
# Setup
prisma_client = MockPrismaClient()
proxy_logging_obj = create_mock_proxy_logging()
# Create 400 test spend logs (4x BATCH_SIZE)
# Create 4000 test spend logs (4x BATCH_SIZE)
prisma_client.spend_log_transactions = [
{"id": str(i), "spend": 10} for i in range(400)
{"id": str(i), "spend": 10} for i in range(4000)
]
# Mock to fail on second batch first attempt, then succeed
@ -292,9 +296,9 @@ async def test_update_spend_logs_multiple_batches_with_failure():
# Verify all IDs were processed
processed_ids = {item["id"] for item in all_processed_logs}
# these should have ids 0-399
# these should have ids 0-3999
print("all processed ids", sorted(processed_ids, key=int))
expected_ids = {str(i) for i in range(400)}
expected_ids = {str(i) for i in range(4000)}
assert processed_ids == expected_ids
# Verify all logs were cleared from transactions

1
tests/test_litellm/llms/azure/test_azure_common_utils.py
View File

@ -569,6 +569,7 @@ async def test_ensure_initialize_azure_sdk_client_always_used(call_type):
or call_type == CallTypes.aretrieve_container
or call_type == CallTypes.acreate_container
or call_type == CallTypes.adelete_container
or call_type == CallTypes.alist_container_files
):
# Skip container call types as they're not supported for Azure (only OpenAI)
pytest.skip(f"Skipping {call_type.value} because Azure doesn't support container operations")

7
tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py
View File

@ -3,7 +3,7 @@ Tests for Voyage AI rerank transformation functionality.
"""
import json
from unittest.mock import MagicMock
from unittest.mock import MagicMock, patch
import httpx
import pytest
@ -264,8 +264,11 @@ class TestVoyageRerankTransform:
assert "top_n" in supported_params
assert "return_documents" in supported_params
def test_validate_environment_missing_api_key(self):
@patch("litellm.llms.voyage.rerank.transformation.get_secret_str")
def test_validate_environment_missing_api_key(self, mock_get_secret_str):
"""Test that validate_environment raises error when API key is missing."""
# Mock get_secret_str to return None for both environment variables
mock_get_secret_str.return_value = None
with pytest.raises(ValueError, match="Voyage AI API key is required"):
self.config.validate_environment(
headers={},

6
tests/test_litellm/llms/watsonx/test_watsonx.py
View File

@ -250,6 +250,7 @@ async def test_watsonx_gpt_oss_prompt_transformation(monkeypatch):
"generated_text": "Hello! How can I help you?",
"generated_token_count": 10,
"input_token_count": 5,
"stop_reason": "stop", # Required field for response transformation
}
],
"model_id": "openai/gpt-oss-120b",
@ -282,6 +283,11 @@ async def test_watsonx_gpt_oss_prompt_transformation(monkeypatch):
# Return failure to use tokenizer_config instead
return {"status": "failure"}
# Clear any cached tokenizer config for this model to ensure fresh fetch
hf_model = "openai/gpt-oss-120b"
if hf_model in litellm.known_tokenizer_config:
del litellm.known_tokenizer_config[hf_model]
with patch.object(client, "post") as mock_post, patch.object(
litellm.module_level_client, "post", return_value=mock_token_response
), patch(

47
tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py
View File

@ -4,6 +4,7 @@ Mock tests for A2A endpoints.
Tests that invoke_agent_a2a properly integrates with add_litellm_data_to_request.
"""
import sys
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -67,6 +68,49 @@ async def test_invoke_agent_a2a_adds_litellm_data():
team_id="test-team",
)
# Try to use real a2a.types if available, otherwise create realistic mocks
# This test focuses on LiteLLM integration, not A2A protocol correctness,
# but we want mocks that behave like the real types to catch usage issues
try:
from a2a.types import (
MessageSendParams,
SendMessageRequest,
SendStreamingMessageRequest,
)
# Real types available - use them
use_real_types = True
except ImportError:
# Real types not available - create realistic mocks
use_real_types = False
def make_mock_pydantic_class(name):
"""Create a mock class that behaves like a Pydantic model."""
class MockPydanticClass:
def __init__(self, **kwargs):
self.__dict__.update(kwargs)
# Store kwargs for model_dump() if needed
self._kwargs = kwargs
def model_dump(self, mode="json", exclude_none=False):
"""Mock model_dump method."""
result = dict(self._kwargs)
if exclude_none:
result = {k: v for k, v in result.items() if v is not None}
return result
MockPydanticClass.__name__ = name
return MockPydanticClass
MessageSendParams = make_mock_pydantic_class("MessageSendParams")
SendMessageRequest = make_mock_pydantic_class("SendMessageRequest")
SendStreamingMessageRequest = make_mock_pydantic_class("SendStreamingMessageRequest")
# Create a mock module for a2a.types
mock_a2a_types = MagicMock()
mock_a2a_types.MessageSendParams = MessageSendParams
mock_a2a_types.SendMessageRequest = SendMessageRequest
mock_a2a_types.SendStreamingMessageRequest = SendStreamingMessageRequest
# Patch at the source modules
with patch(
"litellm.proxy.agent_endpoints.a2a_endpoints._get_agent",
@ -90,6 +134,9 @@ async def test_invoke_agent_a2a_adds_litellm_data():
), patch(
"litellm.proxy.proxy_server.version",
"1.0.0",
), patch.dict(
sys.modules,
{"a2a": MagicMock(), "a2a.types": mock_a2a_types},
):
from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a

8
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py
View File

@ -1663,15 +1663,15 @@ async def test_async_log_success_event_with_dict_usage(monkeypatch, token_rate_l
)
# Create a mock response object with usage as a dict (Responses API format)
mock_response = MagicMock()
from litellm.types.utils import BaseLiteLLMOpenAIResponseObject
# Use spec to make isinstance checks work correctly with MagicMock
mock_response = MagicMock(spec=BaseLiteLLMOpenAIResponseObject)
mock_response.usage = {
"prompt_tokens": 25,
"completion_tokens": 35,
"total_tokens": 60
}
# Make isinstance check for BaseLiteLLMOpenAIResponseObject return True
from litellm.types.utils import BaseLiteLLMOpenAIResponseObject
mock_response.__class__ = type('MockResponse', (BaseLiteLLMOpenAIResponseObject,), {})
# Create mock kwargs for the success event
mock_kwargs = {

3
tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
View File

@ -1164,6 +1164,7 @@ class TestSpendLogsPayload:
"proxy_server_request": "{}",
"status": "success",
"mcp_namespaced_tool_name": None,
"agent_id": None,
}
)
@ -1257,6 +1258,7 @@ class TestSpendLogsPayload:
"proxy_server_request": "{}",
"status": "success",
"mcp_namespaced_tool_name": None,
"agent_id": None,
}
)
@ -1348,6 +1350,7 @@ class TestSpendLogsPayload:
"proxy_server_request": "{}",
"status": "success",
"mcp_namespaced_tool_name": None,
"agent_id": None,
}
)