scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge branch 'litellm_internal_staging' into litellm_metrics_auth
Some checks failed
Unit Tests: Caching (Redis) / caching-redis (push) Has been cancelled
Details
Unit Tests: Security / security (push) Has been cancelled
Details
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Has been cancelled
Details
Unit Tests: Proxy DB Operations / auth-checks (push) Has been cancelled
Details
Unit Tests: Proxy DB Operations / budgets (push) Has been cancelled
Details
Unit Tests: Proxy DB Operations / custom-logging (push) Has been cancelled
Details
Unit Tests: Proxy DB Operations / db-and-spend (push) Has been cancelled
Details
Unit Tests: Proxy DB Operations / proxy-runtime (push) Has been cancelled
Details
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Has been cancelled
Details
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Has been cancelled
Details
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Has been cancelled
Details
Unit Tests: Proxy DB Operations / key-generation (push) Has been cancelled
Details
Unit Tests: Proxy DB Operations / logging-misc (push) Has been cancelled
Details
Unit Tests: Proxy DB Operations / proxy-server-core (push) Has been cancelled
Details
Unit Tests: Proxy DB Operations / schema-migration (push) Has been cancelled
Details
Unit Tests: Proxy DB Operations / proxy-utils (push) Has been cancelled
Details

Browse source
This commit is contained in:
shivam 2026-04-30 18:41:09 -07:00
commit f4211ff7c4
No known key found for this signature in database
13 changed files with 720 additions and 613 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/caching/dual_cache.py
View file

@ -392,6 +392,7 @@ class DualCache(BaseCache):
value: float,
parent_otel_span: Optional[Span] = None,
local_only: bool = False,
refresh_ttl: bool = False,
**kwargs,
) -> Optional[float]:
"""
@ -399,6 +400,9 @@ class DualCache(BaseCache):
Value - float - the value you want to increment by
Refresh_ttl - bool - if True, resets the Redis TTL on every write.
Default False preserves window-style semantics.
Returns - the incremented value, or None if no cache backend is
available (in_memory_cache is None and Redis failed/is absent).
"""
@ -415,6 +419,7 @@ class DualCache(BaseCache):
value,
parent_otel_span=parent_otel_span,
ttl=kwargs.get("ttl", None),
refresh_ttl=refresh_ttl,
)
return result

10
litellm/caching/redis_cache.py
View file

@ -824,6 +824,7 @@ class RedisCache(BaseCache):
value: float,
ttl: Optional[int] = None,
parent_otel_span: Optional[Span] = None,
refresh_ttl: bool = False,
) -> float:
from redis.asyncio import Redis
@ -834,11 +835,12 @@ class RedisCache(BaseCache):
try:
result = await _redis_client.incrbyfloat(name=key, amount=value)
if _used_ttl is not None:
# check if key already has ttl, if not -> set ttl
current_ttl = await _redis_client.ttl(key)
if current_ttl == -1:
# Key has no expiration
if refresh_ttl:
await _redis_client.expire(key, _used_ttl)
else:
current_ttl = await _redis_client.ttl(key)
if current_ttl == -1:
await _redis_client.expire(key, _used_ttl)
## LOGGING ##
end_time = time.time()

24
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -2535,10 +2535,16 @@ class BaseLLMHTTPHandler:
},
)
delete_kwargs: Dict[str, Any] = {
"url": url,
"headers": headers,
"timeout": timeout,
}
if data:
delete_kwargs["json"] = data
try:
response = await async_httpx_client.delete(
url=url, headers=headers, json=data, timeout=timeout
)
response = await async_httpx_client.delete(**delete_kwargs)
except Exception as e:
raise self._handle_error(
@ -2619,10 +2625,16 @@ class BaseLLMHTTPHandler:
},
)
delete_kwargs: Dict[str, Any] = {
"url": url,
"headers": headers,
"timeout": timeout,
}
if data:
delete_kwargs["json"] = data
try:
response = sync_httpx_client.delete(
url=url, headers=headers, json=data, timeout=timeout
)
response = sync_httpx_client.delete(**delete_kwargs)
except Exception as e:
raise self._handle_error(

120
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -52,6 +52,37 @@ class ResetBudgetJob:
### RESET MULTI-WINDOW BUDGETS ###
await self.reset_budget_windows()
@staticmethod
async def _invalidate_spend_counter(counter_key: str) -> None:
"""Zero a spend counter so a DB-row reset takes effect immediately.
Call AFTER the DB write commits. Clearing Redis before the DB
commit opens a window where get_current_spend reads 0 from Redis
while the DB still holds the pre-reset value, allowing bypass.
"""
try:
from litellm.proxy.proxy_server import spend_counter_cache
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key, value=0.0, ttl=60
)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key, value=0.0, ttl=60
)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to reset spend counter %s in Redis: %s. "
"Budget may be over-enforced until counter expires.",
counter_key,
redis_err,
)
except Exception as e:
verbose_proxy_logger.warning(
"Failed to reset spend counter %s: %s", counter_key, e
)
async def reset_budget_for_litellm_team_members(
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
):
@ -64,46 +95,30 @@ class ResetBudgetJob:
if budget.budget_id is not None
]
# Reset spend counters for affected team members.
# Reset Redis directly so a transient failure doesn't leave stale
# counters that get_current_spend would read as authoritative.
try:
from litellm.proxy.proxy_server import spend_counter_cache
memberships = await self.prisma_client.db.litellm_teammembership.find_many(
where={"budget_id": {"in": budget_ids}}
)
for m in memberships:
counter_key = f"spend:team_member:{m.user_id}:{m.team_id}"
# Always reset in-memory
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key, value=0.0
)
# Explicitly reset Redis with warning on failure
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key, value=0.0
)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to reset team member spend counter in Redis %s: %s. "
"Budget may be over-enforced until counter expires.",
counter_key,
redis_err,
)
except Exception as e:
memberships = []
verbose_proxy_logger.warning(
"Failed to reset team member spend counters: %s", e
"Failed to fetch team memberships for counter invalidation: %s", e
)
return await self.prisma_client.db.litellm_teammembership.update_many(
update_result = await self.prisma_client.db.litellm_teammembership.update_many(
where={"budget_id": {"in": budget_ids}},
data={
"spend": 0,
},
)
for m in memberships:
await self._invalidate_spend_counter(
f"spend:team_member:{m.user_id}:{m.team_id}"
)
return update_result
async def reset_budget_for_keys_linked_to_budgets(
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
):
@ -126,17 +141,36 @@ class ResetBudgetJob:
if not budget_ids:
return
return await self.prisma_client.db.litellm_verificationtoken.update_many(
where={
"budget_id": {"in": budget_ids},
"budget_duration": None, # only keys without their own reset schedule
"spend": {"gt": 0}, # only reset keys that have accumulated spend
},
data={
"spend": 0,
},
where_clause: dict = {
"budget_id": {"in": budget_ids},
"budget_duration": None, # only keys without their own reset schedule
"spend": {"gt": 0}, # only reset keys that have accumulated spend
}
try:
keys = await self.prisma_client.db.litellm_verificationtoken.find_many(
where=where_clause
)
except Exception as e:
keys = []
verbose_proxy_logger.warning(
"Failed to fetch keys for counter invalidation: %s", e
)
update_result = (
await self.prisma_client.db.litellm_verificationtoken.update_many(
where=where_clause,
data={
"spend": 0,
},
)
)
for k in keys:
await self._invalidate_spend_counter(f"spend:key:{k.token}")
return update_result
async def reset_budget_for_litellm_budget_table(self):
"""
Resets the budget for all LiteLLM End-Users (Customers), and Team Members if their budget has expired
@ -365,6 +399,10 @@ class ResetBudgetJob:
data_list=updated_keys,
table_name="key",
)
for k in updated_keys:
token = getattr(k, "token", None)
if token:
await self._invalidate_spend_counter(f"spend:key:{token}")
end_time = time.time()
if len(failed_keys) > 0: # If any keys failed to reset
@ -450,6 +488,12 @@ class ResetBudgetJob:
data_list=updated_users,
table_name="user",
)
for u in updated_users:
user_id = getattr(u, "user_id", None)
if user_id:
await self._invalidate_spend_counter(
f"spend:user:{user_id}"
)
end_time = time.time()
if len(failed_users) > 0: # If any users failed to reset
@ -541,6 +585,12 @@ class ResetBudgetJob:
data_list=updated_teams,
table_name="team",
)
for t in updated_teams:
team_id = getattr(t, "team_id", None)
if team_id:
await self._invalidate_spend_counter(
f"spend:team:{team_id}"
)
end_time = time.time()
if len(failed_teams) > 0: # If any teams failed to reset

14
litellm/proxy/db/spend_counter_reseed.py
View file

@ -129,7 +129,9 @@ class SpendCounterReseed:
"""
lock = await SpendCounterReseed._get_lock(counter_key)
async with lock:
# Re-check after acquiring the lock - another waiter may have warmed it.
# Re-check after acquiring the lock. Skip in-memory on a clean
# Redis miss - in-memory is per-pod-stale.
redis_clean_miss = False
if spend_counter_cache.redis_cache is not None:
try:
val = await spend_counter_cache.redis_cache.async_get_cache(
@ -137,11 +139,13 @@ class SpendCounterReseed:
)
if val is not None:
return float(val)
redis_clean_miss = True
except Exception:
pass
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val)
if not redis_clean_miss:
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val)
db_spend = await SpendCounterReseed.from_db(prisma_client, counter_key)
if db_spend is None:
@ -149,7 +153,7 @@ class SpendCounterReseed:
# Warm even when 0 so subsequent reads hit cache, not DB.
try:
await spend_counter_cache.async_increment_cache(
key=counter_key, value=db_spend
key=counter_key, value=db_spend, refresh_ttl=True
)
except Exception:
verbose_proxy_logger.exception(

171
litellm/proxy/google_endpoints/endpoints.py
View file

@ -1,10 +1,6 @@
from datetime import datetime
from fastapi import APIRouter, Depends, Request, Response
from fastapi.responses import ORJSONResponse
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import ORJSONResponse, StreamingResponse
import litellm
from litellm._uuid import uuid
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
@ -30,12 +26,17 @@ async def google_generate_content(
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
from litellm.proxy.proxy_server import (
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
select_data_generator,
user_api_base,
user_max_tokens,
user_model,
user_request_timeout,
user_temperature,
version,
)
@ -43,48 +44,33 @@ async def google_generate_content(
if "model" not in data:
data["model"] = model_name
# Extract generationConfig and pass it as config parameter
generation_config = data.pop("generationConfig", None)
if generation_config:
data["config"] = generation_config
# Add user authentication metadata for cost tracking
data = await add_litellm_data_to_request(
data=data,
request=request,
user_api_key_dict=user_api_key_dict,
proxy_config=proxy_config,
general_settings=general_settings,
version=version,
)
# Create logging object with full request metadata so callbacks (e.g. S3) get user/trace_id
data["litellm_call_id"] = request.headers.get(
"x-litellm-call-id", str(uuid.uuid4())
)
logging_obj, data = litellm.utils.function_setup(
original_function="agenerate_content",
rules_obj=litellm.utils.Rules(),
start_time=datetime.now(),
**data,
)
data["litellm_logging_obj"] = logging_obj
# call router
if llm_router is None:
raise HTTPException(status_code=500, detail="Router not initialized")
response = await llm_router.agenerate_content(**data)
success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
response=response,
request_data=data,
request=request,
user_api_key_dict=user_api_key_dict,
logging_obj=logging_obj,
version=version,
proxy_logging_obj=proxy_logging_obj,
)
fastapi_response.headers.update(success_headers)
return response
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type="agenerate_content",
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=select_data_generator,
model=model_name,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
version=version,
)
@router.post(
@ -101,73 +87,52 @@ async def google_stream_generate_content(
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
from litellm.proxy.proxy_server import (
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
select_data_generator,
user_api_base,
user_max_tokens,
user_model,
user_request_timeout,
user_temperature,
version,
)
data = await _read_request_body(request=request)
if "model" not in data:
data["model"] = model_name
data["stream"] = True
data["stream"] = True # enforce streaming for this endpoint
# Extract generationConfig and pass it as config parameter
generation_config = data.pop("generationConfig", None)
if generation_config:
data["config"] = generation_config
# Add user authentication metadata for cost tracking
data = await add_litellm_data_to_request(
data=data,
request=request,
user_api_key_dict=user_api_key_dict,
proxy_config=proxy_config,
general_settings=general_settings,
version=version,
)
# Create logging object with full request metadata so streaming END callbacks (e.g. S3) get user/trace_id
data["litellm_call_id"] = request.headers.get(
"x-litellm-call-id", str(uuid.uuid4())
)
logging_obj, data = litellm.utils.function_setup(
original_function="agenerate_content_stream",
rules_obj=litellm.utils.Rules(),
start_time=datetime.now(),
**data,
)
data["litellm_logging_obj"] = logging_obj
# call router
if llm_router is None:
raise HTTPException(status_code=500, detail="Router not initialized")
response = await llm_router.agenerate_content_stream(**data)
success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
response=response,
request_data=data,
request=request,
user_api_key_dict=user_api_key_dict,
logging_obj=logging_obj,
version=version,
proxy_logging_obj=proxy_logging_obj,
)
# Check if response is an async iterator (streaming response)
if response is not None and hasattr(response, "__aiter__"):
return StreamingResponse(
content=response,
media_type="text/event-stream",
headers=success_headers,
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type="agenerate_content_stream",
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=select_data_generator,
model=model_name,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
version=version,
)
fastapi_response.headers.update(success_headers)
return response
@router.post(

21
litellm/proxy/proxy_server.py
View file

@ -1798,12 +1798,16 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float:
3. Reseed from authoritative DB spend (counter expired, cross-pod stale)
4. Caller-supplied fallback (DB unavailable, cold start)
"""
# 1. Try Redis first (cross-pod authoritative)
# 1. Redis first (cross-pod authoritative). On clean miss, skip
# in-memory: per-pod in-memory only has this pod's writes, so it
# would mask cross-pod increments.
redis_clean_miss = False
if spend_counter_cache.redis_cache is not None:
try:
val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
if val is not None:
return float(val)
redis_clean_miss = True
except Exception as e:
verbose_proxy_logger.debug(
"get_current_spend: Redis read failed for %s, falling back to in-memory: %s",
@ -1811,10 +1815,11 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float:
e,
)
# 2. Fall back to in-memory counter (single-instance or Redis failure)
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val)
# 2. In-memory only when Redis is unreachable.
if not redis_clean_miss:
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val)
# 3. Reseed from DB - fallback_spend lags cross-pod, would allow bypass.
db_spend = await SpendCounterReseed.coalesced(
@ -1976,10 +1981,12 @@ async def _init_and_increment_spend_counter(
base_spend = getattr(source, "spend", 0.0) or 0.0
if base_spend > 0:
await spend_counter_cache.async_increment_cache(
key=counter_key, value=base_spend
key=counter_key, value=base_spend, refresh_ttl=True
)
await spend_counter_cache.async_increment_cache(key=counter_key, value=increment)
await spend_counter_cache.async_increment_cache(
key=counter_key, value=increment, refresh_ttl=True
)
async def update_cache( # noqa: PLR0915

44
tests/test_litellm/caching/test_redis_cache.py
View file

@ -50,6 +50,50 @@ async def test_redis_cache_async_increment(namespace, monkeypatch, redis_no_ping
)
@pytest.mark.asyncio
async def test_redis_cache_async_increment_refresh_ttl_true_bumps_existing_ttl(
monkeypatch, redis_no_ping
):
"""With refresh_ttl=True, every increment should call expire() to bump
the TTL, even when the key already has a TTL (counter-style use)."""
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache()
mock_redis_instance = AsyncMock()
mock_redis_instance.__aenter__.return_value = mock_redis_instance
mock_redis_instance.__aexit__.return_value = None
mock_redis_instance.ttl.return_value = 42 # key already has ~42s left
with patch.object(
redis_cache, "init_async_client", return_value=mock_redis_instance
):
await redis_cache.async_increment(
key="spend:team_member:u:t", value=0.05, refresh_ttl=True
)
mock_redis_instance.expire.assert_awaited_once_with("spend:team_member:u:t", 60)
@pytest.mark.asyncio
async def test_redis_cache_async_increment_default_does_not_bump_existing_ttl(
monkeypatch, redis_no_ping
):
"""Default (refresh_ttl=False) preserves window-style semantics: TTL is
set only on first creation, never refreshed (used by rate-limit windows)."""
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache()
mock_redis_instance = AsyncMock()
mock_redis_instance.__aenter__.return_value = mock_redis_instance
mock_redis_instance.__aexit__.return_value = None
mock_redis_instance.ttl.return_value = 42 # key already has ~42s left
with patch.object(
redis_cache, "init_async_client", return_value=mock_redis_instance
):
await redis_cache.async_increment(key="rate_limit:window", value=1)
mock_redis_instance.expire.assert_not_awaited()
@pytest.mark.asyncio
async def test_redis_client_init_with_socket_timeout(monkeypatch, redis_no_ping):
monkeypatch.setenv("REDIS_HOST", "my-fake-host")

74
tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
View file

@ -1,3 +1,4 @@
import asyncio
import os
import sys
from unittest.mock import AsyncMock, Mock, patch
@ -8,6 +9,8 @@ import pytest
sys.path.insert(
0, os.path.abspath("../../../..")
) # Adds the parent directory to the system path
import litellm
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import (
BaseLLMHTTPHandler,
_google_genai_streaming_hidden_params,
@ -103,7 +106,9 @@ def test_fingerprint_agentic_tools_is_deterministic():
tools_a = {"tool_calls": [{"id": "1", "input": {"q": "abc"}, "name": "web_search"}]}
tools_b = {"tool_calls": [{"name": "web_search", "input": {"q": "abc"}, "id": "1"}]}
assert handler._fingerprint_agentic_tools(tools_a) == handler._fingerprint_agentic_tools(tools_b)
assert handler._fingerprint_agentic_tools(
tools_a
) == handler._fingerprint_agentic_tools(tools_b)
@pytest.mark.asyncio
@ -350,3 +355,70 @@ def test_google_genai_streaming_hidden_params_model_info_and_router_fallback():
response_headers=httpx.Headers({}),
)
assert from_router["model_id"] == "router-model-id"
def _build_delete_response_mock(captured: dict):
"""Returns a fake httpx delete that records its kwargs."""
def _response() -> httpx.Response:
return httpx.Response(
status_code=200,
headers={"content-type": "application/json"},
content=b'{"id": "resp_x", "object": "response", "deleted": true}',
request=httpx.Request(method="DELETE", url="https://test.openai.azure.com"),
)
async def fake_async_delete(*args, **kwargs):
captured.update(kwargs)
return _response()
def fake_sync_delete(*args, **kwargs):
captured.update(kwargs)
return _response()
return fake_async_delete, fake_sync_delete
def test_async_delete_responses_omits_body_for_azure():
"""Azure responses DELETE rejects requests with any body. Verify the handler
does not pass `json=` to httpx when the transformer returns an empty dict."""
captured: dict = {}
fake_async_delete, _ = _build_delete_response_mock(captured)
async def run():
with patch.object(AsyncHTTPHandler, "delete", new=fake_async_delete):
await litellm.adelete_responses(
response_id="resp_xyz",
custom_llm_provider="azure",
api_base="https://test.openai.azure.com",
api_key="test-key",
api_version="2025-03-01-preview",
)
asyncio.run(run())
assert "json" not in captured
assert "data" not in captured
assert captured["url"].endswith(
"/openai/responses/resp_xyz?api-version=2025-03-01-preview"
)
def test_sync_delete_responses_omits_body_for_azure():
captured: dict = {}
_, fake_sync_delete = _build_delete_response_mock(captured)
with patch.object(HTTPHandler, "delete", new=fake_sync_delete):
litellm.delete_responses(
response_id="resp_xyz",
custom_llm_provider="azure",
api_base="https://test.openai.azure.com",
api_key="test-key",
api_version="2025-03-01-preview",
)
assert "json" not in captured
assert "data" not in captured
assert captured["url"].endswith(
"/openai/responses/resp_xyz?api-version=2025-03-01-preview"
)

156
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
View file

@ -1049,3 +1049,159 @@ def test_reset_budget_windows_query_error_does_not_break_team_path(monkeypatch):
asyncio.run(job.reset_budget_windows()) # must not raise
prisma_client.db.litellm_teamtable.update.assert_awaited_once()
# ---------------------------------------------------------------------------
# Counter invalidation on budget reset
# ---------------------------------------------------------------------------
def _make_counter_invalidation_job(monkeypatch):
"""Stub spend_counter_cache so we can observe invalidation calls."""
spend_counter_cache = MagicMock()
spend_counter_cache.in_memory_cache.set_cache = MagicMock()
spend_counter_cache.redis_cache = MagicMock()
spend_counter_cache.redis_cache.async_set_cache = AsyncMock()
fake_module = types.ModuleType("litellm.proxy.proxy_server")
fake_module.spend_counter_cache = spend_counter_cache
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
return spend_counter_cache
def test_reset_budget_for_team_members_invalidates_redis_counter(monkeypatch):
"""Team-member budget reset clears the Redis spend counter."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
expired_budget = type("B", (), {"budget_id": "budget-1"})
membership = type(
"Membership",
(),
{"user_id": "alice", "team_id": "team-x", "budget_id": "budget-1"},
)
prisma_client = MagicMock()
prisma_client.db.litellm_teammembership.find_many = AsyncMock(
return_value=[membership]
)
prisma_client.db.litellm_teammembership.update_many = AsyncMock(
return_value={"count": 1}
)
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget]))
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:team_member:alice:team-x", value=0.0, ttl=60
)
counter_cache.redis_cache.async_set_cache.assert_any_await(
key="spend:team_member:alice:team-x", value=0.0, ttl=60
)
def test_reset_budget_for_keys_invalidates_redis_counter(
reset_budget_job, mock_prisma_client, monkeypatch
):
"""Key budget reset must clear the Redis spend counter."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
now = datetime.now(timezone.utc)
mock_prisma_client.data["key"] = [
type(
"Key",
(),
{
"spend": 100.0,
"budget_duration": "30d",
"budget_reset_at": now,
"id": "key-1",
"token": "sk-abc",
},
)
]
asyncio.run(reset_budget_job.reset_budget_for_litellm_keys())
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:key:sk-abc", value=0.0, ttl=60
)
def test_reset_budget_for_users_invalidates_redis_counter(
reset_budget_job, mock_prisma_client, monkeypatch
):
"""User budget reset must clear the Redis spend counter."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
now = datetime.now(timezone.utc)
mock_prisma_client.data["user"] = [
type(
"User",
(),
{
"spend": 50.0,
"budget_duration": "7d",
"budget_reset_at": now,
"id": "user-1",
"user_id": "alice",
},
)
]
asyncio.run(reset_budget_job.reset_budget_for_litellm_users())
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:user:alice", value=0.0, ttl=60
)
def test_reset_budget_for_teams_invalidates_redis_counter(
reset_budget_job, mock_prisma_client, monkeypatch
):
"""Team budget reset must clear the Redis spend counter."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
now = datetime.now(timezone.utc)
mock_prisma_client.data["team"] = [
type(
"Team",
(),
{
"spend": 200.0,
"budget_duration": "1mo",
"budget_reset_at": now,
"id": "team-1",
"team_id": "team-x",
},
)
]
asyncio.run(reset_budget_job.reset_budget_for_litellm_teams())
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:team:team-x", value=0.0, ttl=60
)
def test_reset_budget_for_keys_linked_to_budgets_invalidates_redis_counter(monkeypatch):
"""Resetting keys via budget tier must clear each linked key's counter."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
expired_budget = type("B", (), {"budget_id": "budget-1"})
linked_key = type("Key", (), {"token": "sk-linked"})
prisma_client = MagicMock()
prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[linked_key]
)
prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(
return_value={"count": 1}
)
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget]))
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:key:sk-linked", value=0.0, ttl=60
)

555
tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py
View file

@ -4,7 +4,7 @@ Test to verify the Google GenAI proxy API endpoints
"""
import os
import sys
from unittest.mock import AsyncMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -13,520 +13,171 @@ sys.path.insert(
) # Adds the parent directory to the system path
def test_google_generate_content_endpoint():
"""Test that the google_generate_content endpoint correctly routes requests"""
# Skip this test if we can't import the required modules due to missing dependencies
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
def _build_test_client():
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.google_endpoints.endpoints import router as google_router
from litellm.proxy.google_endpoints.endpoints import router as google_router
app = FastAPI()
app.include_router(google_router)
return TestClient(app)
def _patch_base_process(return_value=None):
"""Patch ProxyBaseLLMRequestProcessing.base_process_llm_request so endpoint
tests don't run the full pipeline. Returns the AsyncMock so callers can
inspect call args."""
if return_value is None:
return_value = {"test": "response"}
return patch(
"litellm.proxy.google_endpoints.endpoints.ProxyBaseLLMRequestProcessing.base_process_llm_request",
new_callable=AsyncMock,
return_value=return_value,
)
def test_google_generate_content_endpoint():
"""generateContent routes through ProxyBaseLLMRequestProcessing with the
agenerate_content route_type — that pipeline runs pre_call_hook +
during_call_hook + post_call_success_hook for every guardrail callback."""
try:
client = _build_test_client()
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
# Create a FastAPI app and include the router (required for FastAPI 0.120+)
app = FastAPI()
app.include_router(google_router)
# Create a test client
client = TestClient(app)
# Mock the router's agenerate_content method
with patch("litellm.proxy.proxy_server.llm_router") as mock_router:
mock_router.agenerate_content = AsyncMock(return_value={"test": "response"})
# Send a request to the endpoint
with _patch_base_process() as mock_base:
response = client.post(
"/v1beta/models/test-model:generateContent",
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
)
# Verify the response
assert response.status_code == 200
assert response.json() == {"test": "response"}
# Verify that agenerate_content was called
mock_router.agenerate_content.assert_called_once()
mock_base.assert_called_once()
kwargs = mock_base.call_args.kwargs
assert kwargs["route_type"] == "agenerate_content"
assert kwargs["model"] == "test-model"
def test_google_stream_generate_content_endpoint():
"""Test that the google_stream_generate_content endpoint correctly routes streaming requests"""
# Skip this test if we can't import the required modules due to missing dependencies
"""streamGenerateContent must route through the same processor with the
streaming route_type so the guardrail pipeline runs."""
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.google_endpoints.endpoints import router as google_router
client = _build_test_client()
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
# Create a FastAPI app and include the router (required for FastAPI 0.120+)
app = FastAPI()
app.include_router(google_router)
# Create a test client
client = TestClient(app)
# Mock the router's agenerate_content_stream method to return a stream
async def mock_stream_generator():
yield 'data: {"test": "stream_chunk_1"}\n\n'
yield 'data: {"test": "stream_chunk_2"}\n\n'
yield "data: [DONE]\n\n"
with patch("litellm.proxy.proxy_server.llm_router") as mock_router:
mock_router.agenerate_content_stream = AsyncMock(
return_value=mock_stream_generator()
)
# Send a request to the endpoint
with (
_patch_base_process() as mock_base,
patch(
"litellm.proxy.google_endpoints.endpoints.ProxyBaseLLMRequestProcessing.__init__",
return_value=None,
) as mock_init,
):
response = client.post(
"/v1beta/models/test-model:streamGenerateContent",
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
)
# Verify the response
assert response.status_code == 200
mock_base.assert_called_once()
kwargs = mock_base.call_args.kwargs
assert kwargs["route_type"] == "agenerate_content_stream"
assert kwargs["model"] == "test-model"
# Verify that agenerate_content_stream was called with correct parameters
mock_router.agenerate_content_stream.assert_called_once()
call_args = mock_router.agenerate_content_stream.call_args
assert call_args[1]["stream"] is True
assert call_args[1]["model"] == "test-model"
assert call_args[1]["contents"] == [
# stream=True must be forced into the data the processor receives.
init_kwargs = mock_init.call_args.kwargs
assert init_kwargs["data"]["stream"] is True
assert init_kwargs["data"]["model"] == "test-model"
assert init_kwargs["data"]["contents"] == [
{"role": "user", "parts": [{"text": "Hello"}]}
]
def test_google_generate_content_with_cost_tracking_metadata():
"""Test that the google_generate_content endpoint includes user metadata for cost tracking"""
def test_google_generate_content_data_flows_through_processor():
"""The body the client sends must reach ProxyBaseLLMRequestProcessing
intact so the pipeline can apply guardrails to it."""
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.google_endpoints.endpoints import router as google_router
client = _build_test_client()
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
# Create a FastAPI app and include the router (required for FastAPI 0.120+)
app = FastAPI()
app.include_router(google_router)
# Create a test client
client = TestClient(app)
# Mock all required proxy server dependencies
with (
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
patch("litellm.proxy.proxy_server.version", "1.0.0"),
_patch_base_process(),
patch(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
) as mock_add_data,
"litellm.proxy.google_endpoints.endpoints.ProxyBaseLLMRequestProcessing.__init__",
return_value=None,
) as mock_init,
):
mock_router.agenerate_content = AsyncMock(return_value={"test": "response"})
# Mock add_litellm_data_to_request to return data with metadata
async def mock_add_litellm_data(
data, request, user_api_key_dict, proxy_config, general_settings, version
):
# Simulate adding user metadata
data["litellm_metadata"] = {
"user_api_key_user_id": "test-user-id",
"user_api_key_team_id": "test-team-id",
"user_api_key": "hashed-key",
}
return data
mock_add_data.side_effect = mock_add_litellm_data
# Send a request to the endpoint
response = client.post(
client.post(
"/v1beta/models/test-model:generateContent",
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
headers={"Authorization": "Bearer sk-test-key"},
)
# Verify the response
assert response.status_code == 200
# Verify that add_litellm_data_to_request was called
mock_add_data.assert_called_once()
# Verify that agenerate_content was called with metadata
mock_router.agenerate_content.assert_called_once()
call_args = mock_router.agenerate_content.call_args
called_data = call_args[1]
# Verify that litellm_metadata exists and contains user information
assert "litellm_metadata" in called_data
assert called_data["litellm_metadata"]["user_api_key_user_id"] == "test-user-id"
assert called_data["litellm_metadata"]["user_api_key_team_id"] == "test-team-id"
def test_google_stream_generate_content_with_cost_tracking_metadata():
"""Test that the google_stream_generate_content endpoint includes user metadata for cost tracking"""
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.google_endpoints.endpoints import router as google_router
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
# Create a FastAPI app and include the router (required for FastAPI 0.120+)
app = FastAPI()
app.include_router(google_router)
# Create a test client
client = TestClient(app)
# Mock the router's agenerate_content_stream method to return a stream
mock_stream = AsyncMock()
mock_stream.__aiter__ = lambda self: mock_stream
mock_stream.__anext__.side_effect = StopAsyncIteration
# Mock all required proxy server dependencies
with (
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
patch("litellm.proxy.proxy_server.version", "1.0.0"),
patch(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
) as mock_add_data,
):
mock_router.agenerate_content_stream = AsyncMock(return_value=mock_stream)
# Mock add_litellm_data_to_request to return data with metadata
async def mock_add_litellm_data(
data, request, user_api_key_dict, proxy_config, general_settings, version
):
# Simulate adding user metadata
data["litellm_metadata"] = {
"user_api_key_user_id": "test-user-id",
"user_api_key_team_id": "test-team-id",
"user_api_key": "hashed-key",
}
return data
mock_add_data.side_effect = mock_add_litellm_data
# Send a request to the endpoint
response = client.post(
"/v1beta/models/test-model:streamGenerateContent",
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
headers={"Authorization": "Bearer sk-test-key"},
)
# Verify the response
assert response.status_code == 200
# Verify that add_litellm_data_to_request was called
mock_add_data.assert_called_once()
# Verify that agenerate_content_stream was called with metadata
mock_router.agenerate_content_stream.assert_called_once()
call_args = mock_router.agenerate_content_stream.call_args
called_data = call_args[1]
# Verify that litellm_metadata exists and contains user information
assert "litellm_metadata" in called_data
assert called_data["litellm_metadata"]["user_api_key_user_id"] == "test-user-id"
assert called_data["litellm_metadata"]["user_api_key_team_id"] == "test-team-id"
# Verify stream is set to True
assert called_data["stream"] is True
def test_google_generate_content_with_system_instruction():
"""
Test that systemInstruction is correctly passed through from the endpoint to the router.
This test verifies the fix for systemInstruction being dropped when forwarding
requests to Vertex AI through the Google GenAI endpoint.
"""
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.google_endpoints.endpoints import router as google_router
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
# Create a FastAPI app and include the router
app = FastAPI()
app.include_router(google_router)
# Create a test client
client = TestClient(app)
# Mock all required proxy server dependencies
with (
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
patch("litellm.proxy.proxy_server.version", "1.0.0"),
patch(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
) as mock_add_data,
):
mock_router.agenerate_content = AsyncMock(return_value={"test": "response"})
# Mock add_litellm_data_to_request to pass through data unchanged
async def mock_add_litellm_data(
data, request, user_api_key_dict, proxy_config, general_settings, version
):
return data
mock_add_data.side_effect = mock_add_litellm_data
# Define the systemInstruction to test
system_instruction = {"parts": [{"text": "Your name is Doodle."}]}
# Send a request with systemInstruction
response = client.post(
"/v1beta/models/gemini-2.5-pro:generateContent",
json={
"systemInstruction": system_instruction,
"contents": [
{"parts": [{"text": "What is your name?"}], "role": "user"}
],
},
headers={"Authorization": "Bearer sk-test-key"},
)
# Verify the response
assert response.status_code == 200
# Verify that agenerate_content was called
mock_router.agenerate_content.assert_called_once()
call_args = mock_router.agenerate_content.call_args
called_data = call_args[1]
# Verify that systemInstruction is present in the call arguments
assert "systemInstruction" in called_data
assert called_data["systemInstruction"] == system_instruction
assert (
called_data["systemInstruction"]["parts"][0]["text"]
== "Your name is Doodle."
)
# Verify contents are also present
assert "contents" in called_data
assert len(called_data["contents"]) == 1
assert called_data["contents"][0]["role"] == "user"
def test_google_generate_content_with_image_config():
"""
Test that imageConfig is correctly passed through from generationConfig to the router.
This test verifies that imageConfig parameters (aspectRatio, imageSize) are preserved
when forwarding requests to Google GenAI through the endpoint.
"""
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.google_endpoints.endpoints import router as google_router
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
# Create a FastAPI app and include the router
app = FastAPI()
app.include_router(google_router)
# Create a test client
client = TestClient(app)
# Mock all required proxy server dependencies
with (
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
patch("litellm.proxy.proxy_server.version", "1.0.0"),
patch(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
) as mock_add_data,
):
mock_router.agenerate_content = AsyncMock(return_value={"test": "response"})
# Mock add_litellm_data_to_request to pass through data unchanged
async def mock_add_litellm_data(
data, request, user_api_key_dict, proxy_config, general_settings, version
):
return data
mock_add_data.side_effect = mock_add_litellm_data
# Send a request with generationConfig containing imageConfig
response = client.post(
"/v1beta/models/gemini-3-pro-image-preview:generateContent",
json={
"contents": [
{
"role": "user",
"parts": [
{
"text": "Create a vibrant infographic about photosynthesis"
}
],
}
],
"contents": [{"role": "user", "parts": [{"text": "Hello"}]}],
"systemInstruction": {"parts": [{"text": "Your name is Doodle."}]},
"generationConfig": {
"responseModalities": ["TEXT", "IMAGE"],
"imageConfig": {"aspectRatio": "9:16", "imageSize": "4K"},
},
},
headers={"Authorization": "Bearer sk-test-key"},
)
# Verify the response
assert response.status_code == 200
# Verify that agenerate_content was called
mock_router.agenerate_content.assert_called_once()
call_args = mock_router.agenerate_content.call_args
called_data = call_args[1]
# Verify that config is present in the call arguments
assert "config" in called_data
# Verify that imageConfig is preserved in the config
assert "imageConfig" in called_data["config"]
assert called_data["config"]["imageConfig"]["aspectRatio"] == "9:16"
assert called_data["config"]["imageConfig"]["imageSize"] == "4K"
# Verify that responseModalities is also preserved
assert "responseModalities" in called_data["config"]
assert called_data["config"]["responseModalities"] == ["TEXT", "IMAGE"]
# Verify contents are also present
assert "contents" in called_data
assert len(called_data["contents"]) == 1
assert called_data["contents"][0]["role"] == "user"
data = mock_init.call_args.kwargs["data"]
assert data["model"] == "test-model"
assert data["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}]
assert data["systemInstruction"] == {
"parts": [{"text": "Your name is Doodle."}]
}
# generationConfig arrives intact here; the rename to `config` is
# done downstream in route_request (see test_route_llm_request).
assert data["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"]
assert data["generationConfig"]["imageConfig"]["aspectRatio"] == "9:16"
def test_google_generate_content_metadata_and_trace_id_callbacks():
"""Test that google_generate_content sets litellm_call_id and logging_obj for callbacks (e.g. S3, Langfuse)"""
def test_google_generate_content_forwards_call_id_header():
"""The endpoint must forward the x-litellm-call-id header to the processor
so the helper can stamp it on the logging object. Trace continuity from
client → callbacks (S3, Langfuse, etc.) depends on this header surviving
the hop through these endpoints."""
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.google_endpoints.endpoints import router as google_router
client = _build_test_client()
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
# Create a FastAPI app and include the router
app = FastAPI()
app.include_router(google_router)
# Create a test client
client = TestClient(app)
# Mock all required proxy server dependencies
with (
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
patch("litellm.proxy.proxy_server.version", "1.0.0"),
patch(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
) as mock_add_data,
):
mock_router.agenerate_content = AsyncMock(return_value={"test": "response"})
# Mock add_litellm_data_to_request to return data with metadata
async def mock_add_litellm_data(
data, request, user_api_key_dict, proxy_config, general_settings, version
):
# Simulate adding user metadata
data["litellm_metadata"] = {
"user_api_key_user_id": "test-user-id",
}
return data
mock_add_data.side_effect = mock_add_litellm_data
# Send a request to the endpoint with x-litellm-call-id header
test_call_id = "test-custom-call-id"
response = client.post(
with _patch_base_process() as mock_base:
client.post(
"/v1beta/models/test-model:generateContent",
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
headers={
"Authorization": "Bearer sk-test-key",
"x-litellm-call-id": test_call_id,
},
headers={"x-litellm-call-id": "trace-abc-123"},
)
assert response.status_code == 200
mock_router.agenerate_content.assert_called_once()
call_args = mock_router.agenerate_content.call_args
called_data = call_args[1]
# Verify that the litellm_logging_obj got assigned in the final called_data to router
assert "litellm_logging_obj" in called_data
assert "litellm_call_id" in called_data
assert called_data["litellm_call_id"] == test_call_id
forwarded_request = mock_base.call_args.kwargs["request"]
assert forwarded_request.headers.get("x-litellm-call-id") == "trace-abc-123"
def test_google_stream_generate_content_metadata_and_trace_id_callbacks():
"""Test that google_stream_generate_content sets litellm_call_id and logging_obj for callbacks"""
def test_google_count_tokens_unchanged():
"""countTokens has its own path and isn't affected by the pipeline change."""
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.google_endpoints.endpoints import router as google_router
client = _build_test_client()
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
app = FastAPI()
app.include_router(google_router)
client = TestClient(app)
fake_response = MagicMock()
fake_response.original_response = {
"totalTokens": 7,
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}],
}
fake_response.total_tokens = 7
mock_stream = AsyncMock()
mock_stream.__aiter__ = lambda self: mock_stream
mock_stream.__anext__.side_effect = StopAsyncIteration
with (
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
patch("litellm.proxy.proxy_server.version", "1.0.0"),
patch(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
) as mock_add_data,
with patch(
"litellm.proxy.proxy_server.token_counter",
new_callable=AsyncMock,
return_value=fake_response,
):
mock_router.agenerate_content_stream = AsyncMock(return_value=mock_stream)
async def mock_add_litellm_data(
data, request, user_api_key_dict, proxy_config, general_settings, version
):
data["litellm_metadata"] = {
"user_api_key_user_id": "test-user-id",
}
return data
mock_add_data.side_effect = mock_add_litellm_data
test_call_id = "test-custom-stream-call-id"
response = client.post(
"/v1beta/models/test-model:streamGenerateContent",
json={"contents": [{"role": "user", "parts": [{"text": "Hello stream"}]}]},
headers={
"Authorization": "Bearer sk-test-key",
"x-litellm-call-id": test_call_id,
},
"/v1beta/models/test-model:countTokens",
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
)
assert response.status_code == 200
mock_router.agenerate_content_stream.assert_called_once()
call_args = mock_router.agenerate_content_stream.call_args
called_data = call_args[1]
assert "litellm_logging_obj" in called_data
assert "litellm_call_id" in called_data
assert called_data["litellm_call_id"] == test_call_id
body = response.json()
assert body["totalTokens"] == 7

87
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -5750,3 +5750,90 @@ class TestLazyFeatureMiddleware:
assert attempts == [
"called"
], f"failing register_fn should be invoked once, not on every request; got {attempts}"
@pytest.mark.asyncio
async def test_get_current_spend_redis_clean_miss_skips_stale_in_memory():
"""When Redis is reachable and cleanly returns None (TTL expired,
counter genuinely absent), the read must reseed from DB - NOT fall
through to per-pod in-memory which only contains this pod's writes.
Pre-fix in multi-pod deployments, in-memory contained a stale local
subset (e.g. $30) while DB had the true cross-pod total ($500). The
fall-through returned $30, enforcement passed, bypass.
"""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.proxy_server import get_current_spend
counter_cache = DualCache()
counter_key = "spend:team_member:user-1:team-1"
# Per-pod stale in-memory: only this pod's writes, not cross-pod truth.
counter_cache.in_memory_cache.set_cache(key=counter_key, value=30.0)
# Redis cleanly returns None (key expired or never written on this pod).
fake_redis = AsyncMock()
fake_redis.async_get_cache = AsyncMock(return_value=None)
fake_redis.async_increment = AsyncMock(return_value=500.0)
counter_cache.redis_cache = fake_redis
# DB has the authoritative cross-pod spend.
db_row = MagicMock()
db_row.spend = 500.0
fake_prisma = MagicMock()
fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=db_row)
import litellm.proxy.proxy_server as ps
orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client
ps.spend_counter_cache = counter_cache
ps.prisma_client = fake_prisma
try:
spend = await get_current_spend(counter_key=counter_key, fallback_spend=0.0)
assert spend == 500.0, (
f"expected DB-authoritative 500.0 on clean Redis miss, got {spend} "
f"(stale per-pod in-memory $30 would have caused multi-pod bypass)"
)
finally:
ps.spend_counter_cache = orig_counter
ps.prisma_client = orig_prisma
@pytest.mark.asyncio
async def test_get_current_spend_redis_error_falls_back_to_in_memory():
"""When Redis raises, the read should still degrade to in-memory rather
than going straight to DB - in-memory is at least same-pod-fresh and
cheaper than a DB query during a Redis outage."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.proxy_server import get_current_spend
counter_cache = DualCache()
counter_key = "spend:team_member:user-1:team-1"
counter_cache.in_memory_cache.set_cache(key=counter_key, value=42.0)
fake_redis = AsyncMock()
fake_redis.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down"))
counter_cache.redis_cache = fake_redis
fake_prisma = MagicMock()
fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(
return_value=MagicMock(spend=999.0)
)
import litellm.proxy.proxy_server as ps
orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client
ps.spend_counter_cache = counter_cache
ps.prisma_client = fake_prisma
try:
spend = await get_current_spend(counter_key=counter_key, fallback_spend=0.0)
assert spend == 42.0, (
f"expected in-memory fallback 42.0 on Redis error, got {spend} "
f"(should not have hit DB when Redis errored)"
)
# DB query should NOT have fired - in-memory short-circuits.
fake_prisma.db.litellm_teammembership.find_unique.assert_not_awaited()
finally:
ps.spend_counter_cache = orig_counter
ps.prisma_client = orig_prisma

52
tests/test_litellm/proxy/test_route_llm_request.py
View file

@ -239,3 +239,55 @@ async def test_route_request_with_router_settings_override_preserves_existing():
assert call_kwargs["num_retries"] == 10
# Key/team timeout should be applied since not in request
assert call_kwargs["timeout"] == 30
@pytest.mark.parametrize(
"route_type", ["agenerate_content", "agenerate_content_stream"]
)
@pytest.mark.asyncio
async def test_route_request_maps_generation_config_for_google_routes(route_type):
"""For Google generate_content routes, route_request must rename
`generationConfig` (Google's wire format) to `config` (the kwarg the
router method expects). Without this mapping the request reaches the
LLM with the field under the wrong name and the config is dropped."""
data = {
"model": "gemini-2.5-flash",
"contents": [{"role": "user", "parts": [{"text": "Hello"}]}],
"generationConfig": {
"responseModalities": ["TEXT", "IMAGE"],
"imageConfig": {"aspectRatio": "9:16", "imageSize": "4K"},
},
}
llm_router = MagicMock()
getattr(llm_router, route_type).return_value = "ok"
await route_request(data, llm_router, None, route_type)
call_kwargs = getattr(llm_router, route_type).call_args[1]
assert "generationConfig" not in call_kwargs
assert "config" in call_kwargs
assert call_kwargs["config"]["responseModalities"] == ["TEXT", "IMAGE"]
assert call_kwargs["config"]["imageConfig"]["aspectRatio"] == "9:16"
assert call_kwargs["config"]["imageConfig"]["imageSize"] == "4K"
@pytest.mark.parametrize(
"route_type", ["agenerate_content", "agenerate_content_stream"]
)
@pytest.mark.asyncio
async def test_route_request_preserves_existing_config_for_google_routes(route_type):
"""If the caller already supplies `config`, route_request must not
overwrite it with `generationConfig`."""
data = {
"model": "gemini-2.5-flash",
"contents": [{"role": "user", "parts": [{"text": "Hello"}]}],
"config": {"existing": True},
"generationConfig": {"shouldNotWin": True},
}
llm_router = MagicMock()
getattr(llm_router, route_type).return_value = "ok"
await route_request(data, llm_router, None, route_type)
call_kwargs = getattr(llm_router, route_type).call_args[1]
assert call_kwargs["config"] == {"existing": True}