mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-05 02:23:42 +00:00
Merge pull request #18833 from BerriAI/litellm_staging_01_08_2026
Litellm staging 01 08 2026
This commit is contained in:
Sameer Kankute
@@ -480,20 +480,21 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
||||
or tool_name == VertexToolName.CODE_EXECUTION.value
|
||||
): # code_execution maintained for backwards compatibility
|
||||
code_execution = self.get_tool_value(tool, "codeExecution")
|
||||
elif tool_name and tool_name == VertexToolName.GOOGLE_SEARCH.value:
|
||||
googleSearch = self.get_tool_value(
|
||||
tool, VertexToolName.GOOGLE_SEARCH.value
|
||||
)
|
||||
elif (
|
||||
tool_name and tool_name == VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value
|
||||
elif tool_name and (
|
||||
tool_name == VertexToolName.GOOGLE_SEARCH.value
|
||||
or tool_name == "google_search"
|
||||
):
|
||||
googleSearchRetrieval = self.get_tool_value(
|
||||
tool, VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value
|
||||
)
|
||||
elif tool_name and tool_name == VertexToolName.ENTERPRISE_WEB_SEARCH.value:
|
||||
enterpriseWebSearch = self.get_tool_value(
|
||||
tool, VertexToolName.ENTERPRISE_WEB_SEARCH.value
|
||||
)
|
||||
googleSearch = self.get_tool_value(tool, tool_name)
|
||||
elif tool_name and (
|
||||
tool_name == VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value
|
||||
or tool_name == "google_search_retrieval"
|
||||
):
|
||||
googleSearchRetrieval = self.get_tool_value(tool, tool_name)
|
||||
elif tool_name and (
|
||||
tool_name == VertexToolName.ENTERPRISE_WEB_SEARCH.value
|
||||
or tool_name == "enterprise_web_search"
|
||||
):
|
||||
enterpriseWebSearch = self.get_tool_value(tool, tool_name)
|
||||
elif tool_name and (
|
||||
tool_name == VertexToolName.URL_CONTEXT.value
|
||||
or tool_name == "urlContext"
|
||||
|
||||
@@ -17,50 +17,141 @@ from litellm.secret_managers.main import str_to_bool
|
||||
|
||||
|
||||
class PrismaWrapper:
|
||||
"""
|
||||
Wrapper around Prisma client that handles RDS IAM token authentication.
|
||||
|
||||
When iam_token_db_auth is enabled, this wrapper:
|
||||
1. Proactively refreshes IAM tokens before they expire (background task)
|
||||
2. Falls back to synchronous refresh if a token is found expired
|
||||
3. Uses proper locking to prevent race conditions during reconnection
|
||||
|
||||
RDS IAM tokens are valid for 15 minutes. This wrapper refreshes them
|
||||
3 minutes before expiration to ensure uninterrupted database connectivity.
|
||||
"""
|
||||
|
||||
# Buffer time in seconds before token expiration to trigger refresh
|
||||
# Refresh 3 minutes (180 seconds) before the token expires
|
||||
TOKEN_REFRESH_BUFFER_SECONDS = 180
|
||||
|
||||
# Fallback refresh interval if token parsing fails (10 minutes)
|
||||
FALLBACK_REFRESH_INTERVAL_SECONDS = 600
|
||||
|
||||
def __init__(self, original_prisma: Any, iam_token_db_auth: bool):
|
||||
self._original_prisma = original_prisma
|
||||
self.iam_token_db_auth = iam_token_db_auth
|
||||
|
||||
# Background token refresh task management
|
||||
self._token_refresh_task: Optional[asyncio.Task] = None
|
||||
self._reconnection_lock = asyncio.Lock()
|
||||
self._last_refresh_time: Optional[datetime] = None
|
||||
|
||||
def _extract_token_from_db_url(self, db_url: Optional[str]) -> Optional[str]:
|
||||
"""
|
||||
Extract the token (password) from the DATABASE_URL.
|
||||
|
||||
The token contains the AWS signature with X-Amz-Date and X-Amz-Expires parameters.
|
||||
|
||||
Important: We must parse the URL while it's still encoded to preserve structure,
|
||||
then decode the password portion. Otherwise the '?' in the token breaks URL parsing.
|
||||
"""
|
||||
if db_url is None:
|
||||
return None
|
||||
try:
|
||||
# Parse URL while still encoded to preserve structure
|
||||
parsed = urllib.parse.urlparse(db_url)
|
||||
if parsed.password:
|
||||
# Now decode just the password/token
|
||||
return urllib.parse.unquote(parsed.password)
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def _parse_token_expiration(self, token: Optional[str]) -> Optional[datetime]:
|
||||
"""
|
||||
Parse the token to extract its expiration time.
|
||||
|
||||
Returns the datetime when the token expires, or None if parsing fails.
|
||||
"""
|
||||
if token is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
# Token format: ...?X-Amz-Date=YYYYMMDDTHHMMSSZ&X-Amz-Expires=900&...
|
||||
if "?" not in token:
|
||||
return None
|
||||
|
||||
query_string = token.split("?", 1)[1]
|
||||
params = urllib.parse.parse_qs(query_string)
|
||||
|
||||
expires_str = params.get("X-Amz-Expires", [None])[0]
|
||||
date_str = params.get("X-Amz-Date", [None])[0]
|
||||
|
||||
if not expires_str or not date_str:
|
||||
return None
|
||||
|
||||
token_created = datetime.strptime(date_str, "%Y%m%dT%H%M%SZ")
|
||||
expires_in = int(expires_str)
|
||||
|
||||
return token_created + timedelta(seconds=expires_in)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(f"Failed to parse token expiration: {e}")
|
||||
return None
|
||||
|
||||
def _calculate_seconds_until_refresh(self) -> float:
|
||||
"""
|
||||
Calculate exactly how many seconds until we need to refresh the token.
|
||||
|
||||
Uses precise timing: sleeps until (token_expiration - buffer_seconds).
|
||||
For a 15-minute (900s) token with 180s buffer, this returns ~720s (12 min).
|
||||
|
||||
Returns:
|
||||
Number of seconds to sleep before the next refresh.
|
||||
Returns 0 if token should be refreshed immediately.
|
||||
Returns FALLBACK_REFRESH_INTERVAL_SECONDS if parsing fails.
|
||||
"""
|
||||
db_url = os.getenv("DATABASE_URL")
|
||||
token = self._extract_token_from_db_url(db_url)
|
||||
expiration_time = self._parse_token_expiration(token)
|
||||
|
||||
if expiration_time is None:
|
||||
# If we can't parse the token, use fallback interval
|
||||
verbose_proxy_logger.debug(
|
||||
f"Could not parse token expiration, using fallback interval of "
|
||||
f"{self.FALLBACK_REFRESH_INTERVAL_SECONDS}s"
|
||||
)
|
||||
return self.FALLBACK_REFRESH_INTERVAL_SECONDS
|
||||
|
||||
# Calculate when we should refresh (expiration - buffer)
|
||||
refresh_at = expiration_time - timedelta(
|
||||
seconds=self.TOKEN_REFRESH_BUFFER_SECONDS
|
||||
)
|
||||
|
||||
# How long until refresh time?
|
||||
now = datetime.utcnow()
|
||||
seconds_until_refresh = (refresh_at - now).total_seconds()
|
||||
|
||||
# If already past refresh time, return 0 (refresh immediately)
|
||||
return max(0, seconds_until_refresh)
|
||||
|
||||
def is_token_expired(self, token_url: Optional[str]) -> bool:
|
||||
"""Check if the token in the given URL is expired."""
|
||||
if token_url is None:
|
||||
return True
|
||||
# Decode the token URL to handle URL-encoded characters
|
||||
decoded_url = urllib.parse.unquote(token_url)
|
||||
|
||||
# Parse the token URL
|
||||
parsed_url = urllib.parse.urlparse(decoded_url)
|
||||
token = self._extract_token_from_db_url(token_url)
|
||||
expiration_time = self._parse_token_expiration(token)
|
||||
|
||||
# Parse the query parameters from the path component (if they exist there)
|
||||
query_params = urllib.parse.parse_qs(parsed_url.query)
|
||||
if expiration_time is None:
|
||||
# If we can't parse the token, assume it's expired to trigger refresh
|
||||
verbose_proxy_logger.debug(
|
||||
"Could not parse token expiration, treating as expired"
|
||||
)
|
||||
return True
|
||||
|
||||
# Get expiration time from the query parameters
|
||||
expires = query_params.get("X-Amz-Expires", [None])[0]
|
||||
if expires is None:
|
||||
raise ValueError("X-Amz-Expires parameter is missing or invalid.")
|
||||
|
||||
expires_int = int(expires)
|
||||
|
||||
# Get the token's creation time from the X-Amz-Date parameter
|
||||
token_time_str = query_params.get("X-Amz-Date", [""])[0]
|
||||
if not token_time_str:
|
||||
raise ValueError("X-Amz-Date parameter is missing or invalid.")
|
||||
|
||||
# Ensure the token time string is parsed correctly
|
||||
try:
|
||||
token_time = datetime.strptime(token_time_str, "%Y%m%dT%H%M%SZ")
|
||||
except ValueError as e:
|
||||
raise ValueError(f"Invalid X-Amz-Date format: {e}")
|
||||
|
||||
# Calculate the expiration time
|
||||
expiration_time = token_time + timedelta(seconds=expires_int)
|
||||
|
||||
# Current time in UTC
|
||||
current_time = datetime.utcnow()
|
||||
|
||||
# Check if the token is expired
|
||||
return current_time > expiration_time
|
||||
return datetime.utcnow() > expiration_time
|
||||
|
||||
def get_rds_iam_token(self) -> Optional[str]:
|
||||
"""Generate a new RDS IAM token and update DATABASE_URL."""
|
||||
if self.iam_token_db_auth:
|
||||
from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token
|
||||
|
||||
@@ -74,7 +165,6 @@ class PrismaWrapper:
|
||||
db_host=db_host, db_port=db_port, db_user=db_user
|
||||
)
|
||||
|
||||
# print(f"token: {token}")
|
||||
_db_url = f"postgresql://{db_user}:{token}@{db_host}:{db_port}/{db_name}"
|
||||
if db_schema:
|
||||
_db_url += f"?schema={db_schema}"
|
||||
@@ -86,6 +176,7 @@ class PrismaWrapper:
|
||||
async def recreate_prisma_client(
|
||||
self, new_db_url: str, http_client: Optional[Any] = None
|
||||
):
|
||||
"""Disconnect and reconnect the Prisma client with a new database URL."""
|
||||
from prisma import Prisma # type: ignore
|
||||
|
||||
try:
|
||||
@@ -100,21 +191,159 @@ class PrismaWrapper:
|
||||
|
||||
await self._original_prisma.connect()
|
||||
|
||||
async def start_token_refresh_task(self) -> None:
|
||||
"""
|
||||
Start the background token refresh task.
|
||||
|
||||
This task proactively refreshes RDS IAM tokens before they expire,
|
||||
preventing connection failures. Should be called after the initial
|
||||
Prisma client connection is established.
|
||||
"""
|
||||
if not self.iam_token_db_auth:
|
||||
verbose_proxy_logger.debug(
|
||||
"IAM token auth not enabled, skipping token refresh task"
|
||||
)
|
||||
return
|
||||
|
||||
if self._token_refresh_task is not None:
|
||||
verbose_proxy_logger.debug("Token refresh task already running")
|
||||
return
|
||||
|
||||
self._token_refresh_task = asyncio.create_task(self._token_refresh_loop())
|
||||
verbose_proxy_logger.info(
|
||||
"Started RDS IAM token proactive refresh background task"
|
||||
)
|
||||
|
||||
async def stop_token_refresh_task(self) -> None:
|
||||
"""
|
||||
Stop the background token refresh task gracefully.
|
||||
|
||||
Should be called during application shutdown to clean up resources.
|
||||
"""
|
||||
if self._token_refresh_task is None:
|
||||
return
|
||||
|
||||
self._token_refresh_task.cancel()
|
||||
try:
|
||||
await self._token_refresh_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._token_refresh_task = None
|
||||
verbose_proxy_logger.info("Stopped RDS IAM token refresh background task")
|
||||
|
||||
async def _token_refresh_loop(self) -> None:
|
||||
"""
|
||||
Background loop that proactively refreshes RDS IAM tokens before expiration.
|
||||
|
||||
Uses precise timing: calculates the exact sleep duration until the token
|
||||
needs to be refreshed (expiration - 3 minute buffer), then refreshes.
|
||||
This is more efficient than polling, requiring only 1 wake-up per token cycle.
|
||||
"""
|
||||
verbose_proxy_logger.info(
|
||||
f"RDS IAM token refresh loop started. "
|
||||
f"Tokens will be refreshed {self.TOKEN_REFRESH_BUFFER_SECONDS}s before expiration."
|
||||
)
|
||||
|
||||
while True:
|
||||
try:
|
||||
# Calculate exactly how long to sleep until next refresh
|
||||
sleep_seconds = self._calculate_seconds_until_refresh()
|
||||
|
||||
if sleep_seconds > 0:
|
||||
verbose_proxy_logger.info(
|
||||
f"RDS IAM token refresh scheduled in {sleep_seconds:.0f} seconds "
|
||||
f"({sleep_seconds / 60:.1f} minutes)"
|
||||
)
|
||||
await asyncio.sleep(sleep_seconds)
|
||||
|
||||
# Refresh the token
|
||||
verbose_proxy_logger.info("Proactively refreshing RDS IAM token...")
|
||||
await self._safe_refresh_token()
|
||||
|
||||
except asyncio.CancelledError:
|
||||
verbose_proxy_logger.info("RDS IAM token refresh loop cancelled")
|
||||
break
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
f"Error in RDS IAM token refresh loop: {e}. "
|
||||
f"Retrying in {self.FALLBACK_REFRESH_INTERVAL_SECONDS}s..."
|
||||
)
|
||||
# On error, wait before retrying to avoid tight error loops
|
||||
try:
|
||||
await asyncio.sleep(self.FALLBACK_REFRESH_INTERVAL_SECONDS)
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
|
||||
async def _safe_refresh_token(self) -> None:
|
||||
"""
|
||||
Refresh the RDS IAM token with proper locking to prevent race conditions.
|
||||
|
||||
Uses an asyncio lock to ensure only one refresh operation happens at a time,
|
||||
preventing multiple concurrent reconnection attempts.
|
||||
"""
|
||||
async with self._reconnection_lock:
|
||||
new_db_url = self.get_rds_iam_token()
|
||||
if new_db_url:
|
||||
await self.recreate_prisma_client(new_db_url)
|
||||
self._last_refresh_time = datetime.utcnow()
|
||||
verbose_proxy_logger.info(
|
||||
"RDS IAM token refreshed successfully. New token valid for ~15 minutes."
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.error(
|
||||
"Failed to generate new RDS IAM token during proactive refresh"
|
||||
)
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
"""
|
||||
Proxy attribute access to the underlying Prisma client.
|
||||
|
||||
If IAM token auth is enabled and the token is expired, this method
|
||||
provides a synchronous fallback to refresh the token. However, this
|
||||
should rarely be needed since the background task proactively refreshes
|
||||
tokens before they expire.
|
||||
|
||||
FIXED: Now properly waits for reconnection to complete before returning,
|
||||
instead of the previous fire-and-forget pattern that caused the bug.
|
||||
"""
|
||||
original_attr = getattr(self._original_prisma, name)
|
||||
|
||||
if self.iam_token_db_auth:
|
||||
db_url = os.getenv("DATABASE_URL")
|
||||
if self.is_token_expired(db_url):
|
||||
db_url = self.get_rds_iam_token()
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
if db_url:
|
||||
# Check if token is expired (should be rare if background task is running)
|
||||
if self.is_token_expired(db_url):
|
||||
verbose_proxy_logger.warning(
|
||||
"RDS IAM token expired in __getattr__ - proactive refresh may have failed. "
|
||||
"Triggering synchronous fallback refresh..."
|
||||
)
|
||||
|
||||
new_db_url = self.get_rds_iam_token()
|
||||
if new_db_url:
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
if loop.is_running():
|
||||
asyncio.run_coroutine_threadsafe(
|
||||
self.recreate_prisma_client(db_url), loop
|
||||
# FIXED: Actually wait for the reconnection to complete!
|
||||
# The previous code used fire-and-forget which caused the bug.
|
||||
future = asyncio.run_coroutine_threadsafe(
|
||||
self.recreate_prisma_client(new_db_url), loop
|
||||
)
|
||||
try:
|
||||
# Wait up to 30 seconds for reconnection
|
||||
future.result(timeout=30)
|
||||
verbose_proxy_logger.info(
|
||||
"Synchronous token refresh completed successfully"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
f"Failed to refresh token synchronously: {e}"
|
||||
)
|
||||
raise
|
||||
else:
|
||||
asyncio.run(self.recreate_prisma_client(db_url))
|
||||
asyncio.run(self.recreate_prisma_client(new_db_url))
|
||||
|
||||
# Get the NEW attribute after reconnection
|
||||
original_attr = getattr(self._original_prisma, name)
|
||||
else:
|
||||
raise ValueError("Failed to get RDS IAM token")
|
||||
|
||||
|
||||
@@ -658,7 +658,7 @@ async def _initialize_shared_aiohttp_session():
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def proxy_startup_event(app: FastAPI):
|
||||
async def proxy_startup_event(app: FastAPI): # noqa: PLR0915
|
||||
global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db, premium_user, _license_check, proxy_batch_polling_interval, shared_aiohttp_session
|
||||
import json
|
||||
|
||||
@@ -788,6 +788,17 @@ async def proxy_startup_event(app: FastAPI):
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"Error closing shared aiohttp session: {e}")
|
||||
|
||||
# Shutdown event - stop RDS IAM token refresh background task
|
||||
if (
|
||||
prisma_client is not None
|
||||
and hasattr(prisma_client, "db")
|
||||
and hasattr(prisma_client.db, "stop_token_refresh_task")
|
||||
):
|
||||
try:
|
||||
await prisma_client.db.stop_token_refresh_task()
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"Error stopping token refresh task: {e}")
|
||||
|
||||
await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues]
|
||||
|
||||
|
||||
@@ -4915,6 +4926,14 @@ class ProxyStartupEvent:
|
||||
|
||||
await prisma_client.connect()
|
||||
|
||||
## Start RDS IAM token refresh background task if enabled ##
|
||||
# This proactively refreshes IAM tokens before they expire,
|
||||
# preventing the 15-minute connection failure bug (#16220)
|
||||
if hasattr(prisma_client, "db") and hasattr(
|
||||
prisma_client.db, "start_token_refresh_task"
|
||||
):
|
||||
await prisma_client.db.start_token_refresh_task()
|
||||
|
||||
## Add necessary views to proxy ##
|
||||
asyncio.create_task(
|
||||
prisma_client.check_view_exists()
|
||||
|
||||
@@ -131,6 +131,7 @@ else:
|
||||
|
||||
unified_guardrail = UnifiedLLMGuardrails()
|
||||
|
||||
_anthropic_async_clients = {}
|
||||
|
||||
def print_verbose(print_statement):
|
||||
"""
|
||||
@@ -4254,11 +4255,16 @@ async def count_tokens_with_anthropic_api(
|
||||
|
||||
if anthropic_api_key and messages:
|
||||
# Call Anthropic API directly for more accurate token counting
|
||||
client = anthropic.Anthropic(api_key=anthropic_api_key)
|
||||
|
||||
# Use cached client if available to avoid socket exhaustion
|
||||
if anthropic_api_key not in _anthropic_async_clients:
|
||||
_anthropic_async_clients[anthropic_api_key] = anthropic.AsyncAnthropic(api_key=anthropic_api_key)
|
||||
|
||||
client = _anthropic_async_clients[anthropic_api_key]
|
||||
|
||||
# Call with explicit parameters to satisfy type checking
|
||||
# Type ignore for now since messages come from generic dict input
|
||||
response = client.beta.messages.count_tokens(
|
||||
response = await client.beta.messages.count_tokens(
|
||||
model=model_to_use,
|
||||
messages=messages, # type: ignore
|
||||
betas=["token-counting-2024-11-01"],
|
||||
|
||||
@@ -7,6 +7,7 @@ sys.path.insert(
|
||||
0, os.path.abspath("../../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
from litellm.llms.vertex_ai.gemini import transformation
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
|
||||
from litellm.types.llms import openai
|
||||
from litellm.types import completion
|
||||
from litellm.types.llms.vertex_ai import RequestBody
|
||||
@@ -225,4 +226,65 @@ async def test__transform_request_body_image_config_with_image_size():
|
||||
assert "generationConfig" in rb
|
||||
assert "imageConfig" in rb["generationConfig"]
|
||||
assert rb["generationConfig"]["imageConfig"]["aspectRatio"] == "16:9"
|
||||
assert rb["generationConfig"]["imageConfig"]["imageSize"] == "4K"
|
||||
assert rb["generationConfig"]["imageConfig"]["imageSize"] == "4K"
|
||||
|
||||
|
||||
def test_map_function_google_search_snake_case():
|
||||
"""
|
||||
Test that google_search tool (snake_case) is properly mapped to googleSearch.
|
||||
Fixes issue where tools=[{"google_search": {}}] was being stripped.
|
||||
"""
|
||||
config = VertexGeminiConfig()
|
||||
optional_params = {}
|
||||
|
||||
# Test snake_case google_search
|
||||
tools = [{"google_search": {}}]
|
||||
result = config._map_function(tools, optional_params)
|
||||
|
||||
assert len(result) == 1
|
||||
assert "googleSearch" in result[0]
|
||||
assert result[0]["googleSearch"] == {}
|
||||
|
||||
|
||||
def test_map_function_google_search_camel_case():
|
||||
"""
|
||||
Test that googleSearch tool (camelCase) still works.
|
||||
"""
|
||||
config = VertexGeminiConfig()
|
||||
optional_params = {}
|
||||
|
||||
# Test camelCase googleSearch
|
||||
tools = [{"googleSearch": {}}]
|
||||
result = config._map_function(tools, optional_params)
|
||||
|
||||
assert len(result) == 1
|
||||
assert "googleSearch" in result[0]
|
||||
assert result[0]["googleSearch"] == {}
|
||||
|
||||
|
||||
def test_map_function_google_search_retrieval_snake_case():
|
||||
"""
|
||||
Test that google_search_retrieval tool (snake_case) is properly mapped.
|
||||
"""
|
||||
config = VertexGeminiConfig()
|
||||
optional_params = {}
|
||||
|
||||
tools = [{"google_search_retrieval": {"dynamic_retrieval_config": {"mode": "MODE_DYNAMIC"}}}]
|
||||
result = config._map_function(tools, optional_params)
|
||||
|
||||
assert len(result) == 1
|
||||
assert "googleSearchRetrieval" in result[0]
|
||||
|
||||
|
||||
def test_map_function_enterprise_web_search_snake_case():
|
||||
"""
|
||||
Test that enterprise_web_search tool (snake_case) is properly mapped.
|
||||
"""
|
||||
config = VertexGeminiConfig()
|
||||
optional_params = {}
|
||||
|
||||
tools = [{"enterprise_web_search": {}}]
|
||||
result = config._map_function(tools, optional_params)
|
||||
|
||||
assert len(result) == 1
|
||||
assert "enterpriseWebSearch" in result[0]
|
||||
@@ -0,0 +1,275 @@
|
||||
"""
|
||||
Tests for the RDS IAM token proactive refresh implementation.
|
||||
|
||||
Tests for GitHub Issue #16220: RDS IAM authentication connection failures after 15 minutes.
|
||||
|
||||
The fix implements:
|
||||
1. Proactive background token refresh (refreshes 3 min before expiration)
|
||||
2. Precise sleep timing (1 wake-up per token cycle instead of polling)
|
||||
3. Proper locking during reconnection
|
||||
4. Fixed __getattr__ fallback that now waits for reconnection
|
||||
|
||||
Run these tests:
|
||||
poetry run pytest tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py -v -s
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import urllib.parse
|
||||
from datetime import datetime, timedelta
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class TestPrismaWrapperTokenRefresh:
|
||||
"""Tests for the PrismaWrapper RDS IAM token refresh implementation."""
|
||||
|
||||
@pytest.fixture
|
||||
def setup_env(self):
|
||||
"""Setup environment variables for testing."""
|
||||
os.environ["DATABASE_HOST"] = "test-host.rds.amazonaws.com"
|
||||
os.environ["DATABASE_PORT"] = "5432"
|
||||
os.environ["DATABASE_USER"] = "test_user"
|
||||
os.environ["DATABASE_NAME"] = "test_db"
|
||||
os.environ["IAM_TOKEN_DB_AUTH"] = "True"
|
||||
yield
|
||||
# Cleanup
|
||||
for key in [
|
||||
"DATABASE_HOST",
|
||||
"DATABASE_PORT",
|
||||
"DATABASE_USER",
|
||||
"DATABASE_NAME",
|
||||
"DATABASE_URL",
|
||||
"IAM_TOKEN_DB_AUTH",
|
||||
"DATABASE_SCHEMA",
|
||||
]:
|
||||
os.environ.pop(key, None)
|
||||
|
||||
def _generate_mock_token(self, expires_in_seconds: int = 900) -> str:
|
||||
"""Generate a mock IAM token with expiration info."""
|
||||
now = datetime.utcnow()
|
||||
date_str = now.strftime("%Y%m%dT%H%M%SZ")
|
||||
# Build the token like AWS does
|
||||
token = f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires={expires_in_seconds}&X-Amz-Signature=abc123"
|
||||
return urllib.parse.quote(token, safe="")
|
||||
|
||||
def _set_database_url_with_token(self, expires_in_seconds: int = 900):
|
||||
"""Set DATABASE_URL with a mock token."""
|
||||
token = self._generate_mock_token(expires_in_seconds)
|
||||
os.environ[
|
||||
"DATABASE_URL"
|
||||
] = f"postgresql://test_user:{token}@test-host:5432/test_db"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_token_expired_fresh(self, setup_env):
|
||||
"""Test that fresh token is not detected as expired."""
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True)
|
||||
|
||||
self._set_database_url_with_token(expires_in_seconds=900)
|
||||
db_url = os.getenv("DATABASE_URL")
|
||||
|
||||
assert wrapper.is_token_expired(db_url) is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_token_expired_old(self, setup_env):
|
||||
"""Test that old token is detected as expired."""
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True)
|
||||
|
||||
# Create an expired token
|
||||
old_date = datetime.utcnow() - timedelta(seconds=901)
|
||||
date_str = old_date.strftime("%Y%m%dT%H%M%SZ")
|
||||
token = (
|
||||
f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires=900&X-Amz-Signature=abc"
|
||||
)
|
||||
encoded_token = urllib.parse.quote(token, safe="")
|
||||
db_url = f"postgresql://test_user:{encoded_token}@test-host:5432/test_db"
|
||||
|
||||
assert wrapper.is_token_expired(db_url) is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_stop_token_refresh_task(self, setup_env):
|
||||
"""Test that token refresh task starts and stops correctly."""
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True)
|
||||
|
||||
# Set a valid token
|
||||
self._set_database_url_with_token(expires_in_seconds=900)
|
||||
|
||||
# Start the task
|
||||
await wrapper.start_token_refresh_task()
|
||||
assert wrapper._token_refresh_task is not None
|
||||
assert not wrapper._token_refresh_task.done()
|
||||
|
||||
# Stop the task
|
||||
await wrapper.stop_token_refresh_task()
|
||||
assert wrapper._token_refresh_task is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_task_not_enabled(self, setup_env):
|
||||
"""Test that task doesn't start when IAM auth is not enabled."""
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
# IAM auth disabled
|
||||
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=False)
|
||||
|
||||
await wrapper.start_token_refresh_task()
|
||||
assert wrapper._token_refresh_task is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_token_expired_null(self, setup_env):
|
||||
"""Test that None token is treated as expired."""
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True)
|
||||
|
||||
assert wrapper.is_token_expired(None) is True
|
||||
|
||||
|
||||
class TestTokenExpirationParsing:
|
||||
"""Tests for token expiration parsing utilities."""
|
||||
|
||||
def test_parse_token_expiration_valid(self):
|
||||
"""Test parsing expiration from a valid token."""
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True)
|
||||
|
||||
# Create a token with known expiration
|
||||
token = "mock-token?X-Amz-Date=20240101T120000Z&X-Amz-Expires=900&X-Amz-Signature=abc"
|
||||
|
||||
expiration = wrapper._parse_token_expiration(token)
|
||||
|
||||
assert expiration is not None
|
||||
expected = datetime(2024, 1, 1, 12, 0, 0) + timedelta(seconds=900)
|
||||
assert expiration == expected
|
||||
|
||||
def test_parse_token_expiration_invalid(self):
|
||||
"""Test that invalid token returns None."""
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True)
|
||||
|
||||
# Invalid tokens
|
||||
assert wrapper._parse_token_expiration(None) is None
|
||||
assert wrapper._parse_token_expiration("no-query-params") is None
|
||||
assert wrapper._parse_token_expiration("?missing=params") is None
|
||||
|
||||
|
||||
class TestBackgroundRefreshLoop:
|
||||
"""Tests for the background refresh loop timing."""
|
||||
|
||||
@pytest.fixture
|
||||
def setup_env(self):
|
||||
"""Setup environment variables for testing."""
|
||||
os.environ["DATABASE_HOST"] = "test-host.rds.amazonaws.com"
|
||||
os.environ["DATABASE_PORT"] = "5432"
|
||||
os.environ["DATABASE_USER"] = "test_user"
|
||||
os.environ["DATABASE_NAME"] = "test_db"
|
||||
yield
|
||||
# Cleanup
|
||||
for key in [
|
||||
"DATABASE_HOST",
|
||||
"DATABASE_PORT",
|
||||
"DATABASE_USER",
|
||||
"DATABASE_NAME",
|
||||
"DATABASE_URL",
|
||||
]:
|
||||
os.environ.pop(key, None)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_calculate_seconds_fallback_when_no_url(self, setup_env):
|
||||
"""Test that fallback is used when DATABASE_URL is not set."""
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True)
|
||||
|
||||
# Don't set DATABASE_URL
|
||||
seconds = wrapper._calculate_seconds_until_refresh()
|
||||
|
||||
# Should return fallback interval
|
||||
assert seconds == wrapper.FALLBACK_REFRESH_INTERVAL_SECONDS
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# DEMONSTRATION SCRIPT
|
||||
# ============================================================================
|
||||
|
||||
|
||||
async def demonstrate_fix():
|
||||
"""
|
||||
Demonstrates the fix for the RDS IAM token expiration bug.
|
||||
|
||||
Shows how the proactive refresh prevents the 15-minute connection failure.
|
||||
"""
|
||||
# Import the actual implementation
|
||||
try:
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
except ImportError:
|
||||
return
|
||||
|
||||
# Setup mock environment
|
||||
os.environ["DATABASE_HOST"] = "mock-rds.region.rds.amazonaws.com"
|
||||
os.environ["DATABASE_PORT"] = "5432"
|
||||
os.environ["DATABASE_USER"] = "iam_user"
|
||||
os.environ["DATABASE_NAME"] = "litellm"
|
||||
|
||||
# Create initial token (expires in 10 seconds for demo)
|
||||
now = datetime.utcnow()
|
||||
date_str = now.strftime("%Y%m%dT%H%M%SZ")
|
||||
token = f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires=10&X-Amz-Signature=abc123"
|
||||
encoded_token = urllib.parse.quote(token, safe="")
|
||||
os.environ[
|
||||
"DATABASE_URL"
|
||||
] = f"postgresql://iam_user:{encoded_token}@mock-rds:5432/litellm"
|
||||
|
||||
# Create mock prisma client
|
||||
mock_prisma = MagicMock()
|
||||
|
||||
wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True)
|
||||
|
||||
# Override buffer for faster demo
|
||||
wrapper.TOKEN_REFRESH_BUFFER_SECONDS = 3
|
||||
wrapper.FALLBACK_REFRESH_INTERVAL_SECONDS = 5
|
||||
_ = wrapper._calculate_seconds_until_refresh() # Verify calculation works
|
||||
db_url = os.getenv("DATABASE_URL")
|
||||
is_expired = wrapper.is_token_expired(db_url)
|
||||
assert is_expired is False, "Fresh token should not be expired!"
|
||||
|
||||
# Mock the _token_refresh_loop to prevent it from actually running
|
||||
async def mock_loop():
|
||||
try:
|
||||
await asyncio.sleep(1000)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
with patch.object(wrapper, "_token_refresh_loop", side_effect=mock_loop):
|
||||
await wrapper.start_token_refresh_task()
|
||||
await wrapper.stop_token_refresh_task()
|
||||
|
||||
# Cleanup
|
||||
for key in [
|
||||
"DATABASE_HOST",
|
||||
"DATABASE_PORT",
|
||||
"DATABASE_USER",
|
||||
"DATABASE_NAME",
|
||||
"DATABASE_URL",
|
||||
]:
|
||||
os.environ.pop(key, None)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(demonstrate_fix())
|
||||
@@ -0,0 +1,42 @@
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch, AsyncMock
|
||||
from litellm.proxy.utils import count_tokens_with_anthropic_api, _anthropic_async_clients
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_tokens_caching():
|
||||
"""
|
||||
Test that count_tokens_with_anthropic_api caches the client.
|
||||
"""
|
||||
# Clear cache
|
||||
_anthropic_async_clients.clear()
|
||||
|
||||
api_key = "sk-ant-test-key"
|
||||
messages = [{"role": "user", "content": "hello"}]
|
||||
model = "claude-3-opus-20240229"
|
||||
|
||||
# Mock anthropic
|
||||
with patch("anthropic.AsyncAnthropic") as mock_cls:
|
||||
mock_client = MagicMock()
|
||||
mock_cls.return_value = mock_client
|
||||
|
||||
# Mock response
|
||||
mock_response = MagicMock()
|
||||
mock_response.input_tokens = 10
|
||||
|
||||
# Setup async return for count_tokens
|
||||
mock_client.beta.messages.count_tokens = AsyncMock(return_value=mock_response)
|
||||
|
||||
# First call
|
||||
with patch.dict("os.environ", {"ANTHROPIC_API_KEY": api_key}):
|
||||
await count_tokens_with_anthropic_api(model, messages)
|
||||
|
||||
assert api_key in _anthropic_async_clients
|
||||
assert _anthropic_async_clients[api_key] == mock_client
|
||||
mock_cls.assert_called_once() # Should be called once
|
||||
|
||||
# Second call
|
||||
with patch.dict("os.environ", {"ANTHROPIC_API_KEY": api_key}):
|
||||
await count_tokens_with_anthropic_api(model, messages)
|
||||
|
||||
# Should still be called once (cached)
|
||||
mock_cls.assert_called_once()
|
||||
Reference in New Issue
Block a user