mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Merge pull request #25117 from stuxf/fix/credential-leak-prevention
security: prevent API key leaks in error tracebacks, logs, and alerts
This commit is contained in:
yuneng-jiang
GitHub
commit
3f3bfdbe33
31 changed files with 890 additions and 245 deletions
|
|
@ -86,6 +86,8 @@ _SECRET_RE = _build_secret_patterns()
|
|||
|
||||
|
||||
def _redact_string(value: str) -> str:
|
||||
if not _ENABLE_SECRET_REDACTION:
|
||||
return value
|
||||
return _SECRET_RE.sub(_REDACTED, value)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from typing import Any, Optional
|
|||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm._logging import _redact_string, verbose_logger
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
from ..exceptions import (
|
||||
|
|
@ -2304,7 +2304,7 @@ def exception_type( # type: ignore # noqa: PLR0915
|
|||
else:
|
||||
# if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors
|
||||
raise APIConnectionError(
|
||||
message=f"{exception_provider} APIConnectionError - {message}\n{traceback.format_exc()}",
|
||||
message=f"{exception_provider} APIConnectionError - {message}\n{_redact_string(traceback.format_exc())}",
|
||||
llm_provider="azure",
|
||||
model=model,
|
||||
litellm_debug_info=extra_information,
|
||||
|
|
@ -2431,7 +2431,7 @@ def exception_type( # type: ignore # noqa: PLR0915
|
|||
else:
|
||||
raise APIConnectionError(
|
||||
message="{}\n{}".format(
|
||||
str(original_exception), traceback.format_exc()
|
||||
str(original_exception), _redact_string(traceback.format_exc())
|
||||
),
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
|
|
@ -2460,7 +2460,7 @@ def exception_type( # type: ignore # noqa: PLR0915
|
|||
setattr(e, "litellm_response_headers", litellm_response_headers)
|
||||
raise e # it's already mapped
|
||||
raised_exc = APIConnectionError(
|
||||
message="{}\n{}".format(original_exception, traceback.format_exc()),
|
||||
message="{}\n{}".format(original_exception, _redact_string(traceback.format_exc())),
|
||||
llm_provider="",
|
||||
model="",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -36,7 +36,7 @@ from litellm import (
|
|||
log_raw_request_response,
|
||||
turn_off_message_logging,
|
||||
)
|
||||
from litellm._logging import _is_debugging_on, verbose_logger
|
||||
from litellm._logging import _is_debugging_on, _redact_string, verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.batches.batch_utils import _handle_completed_batch
|
||||
from litellm.caching.caching import DualCache, InMemoryCache
|
||||
|
|
@ -2854,7 +2854,11 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
self.model_call_details["log_event_type"] = "failed_api_call"
|
||||
self.model_call_details["exception"] = exception
|
||||
self.model_call_details["traceback_exception"] = traceback_exception
|
||||
self.model_call_details["traceback_exception"] = (
|
||||
_redact_string(traceback_exception)
|
||||
if isinstance(traceback_exception, str)
|
||||
else traceback_exception
|
||||
)
|
||||
self.model_call_details["end_time"] = end_time
|
||||
self.model_call_details.setdefault("original_response", None)
|
||||
self.model_call_details["response_cost"] = 0
|
||||
|
|
@ -2877,7 +2881,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
end_time=end_time,
|
||||
logging_obj=self,
|
||||
status="failure",
|
||||
error_str=str(exception),
|
||||
error_str=_redact_string(str(exception)),
|
||||
original_exception=exception,
|
||||
standard_built_in_tools_params=self.standard_built_in_tools_params,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ This requires websockets, and is currently only supported on LiteLLM Proxy.
|
|||
|
||||
from typing import Any, Optional, cast
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._logging import _redact_string, verbose_proxy_logger
|
||||
from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES
|
||||
|
||||
from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
|
|
@ -118,7 +118,7 @@ class AzureOpenAIRealtime(AzureChatCompletion):
|
|||
await realtime_streaming.bidirectional_forward()
|
||||
|
||||
except websockets.exceptions.InvalidStatusCode as e: # type: ignore
|
||||
await websocket.close(code=e.status_code, reason=str(e))
|
||||
await websocket.close(code=e.status_code, reason=_redact_string(str(e)))
|
||||
except Exception:
|
||||
verbose_proxy_logger.exception(
|
||||
"Error in AzureOpenAIRealtime.async_realtime"
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ import asyncio
|
|||
import json
|
||||
from typing import Any, Optional
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._logging import _redact_string, verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
|
||||
from ..base_aws_llm import BaseAWSLLM
|
||||
|
|
@ -152,7 +152,7 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
f"Error in BedrockRealtime.async_realtime: {e}"
|
||||
)
|
||||
try:
|
||||
await websocket.close(code=1011, reason=f"Internal error: {str(e)}")
|
||||
await websocket.close(code=1011, reason=_redact_string(f"Internal error: {str(e)}"))
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -316,16 +316,65 @@ def mask_sensitive_info(error_message):
|
|||
return error_message
|
||||
|
||||
|
||||
def _safe_get_response_text(response: httpx.Response) -> str:
|
||||
"""Safely read response text, falling back to empty string on decoding errors."""
|
||||
try:
|
||||
return response.text
|
||||
except Exception:
|
||||
return ""
|
||||
|
||||
|
||||
async def _safe_aread_response(response: httpx.Response) -> bytes:
|
||||
"""Safely read async response body, falling back to empty bytes on errors."""
|
||||
try:
|
||||
return await response.aread()
|
||||
except Exception:
|
||||
return b""
|
||||
|
||||
|
||||
def _safe_read_response(response: httpx.Response) -> bytes:
|
||||
"""Safely read sync response body, falling back to empty bytes on errors."""
|
||||
try:
|
||||
return response.read()
|
||||
except Exception:
|
||||
return b""
|
||||
|
||||
|
||||
def _raise_masked_sync_error(e: httpx.HTTPStatusError, stream: bool) -> None:
|
||||
"""Raise a MaskedHTTPStatusError for sync HTTP handlers."""
|
||||
if stream:
|
||||
_body = mask_sensitive_info(_safe_read_response(e.response))
|
||||
raise MaskedHTTPStatusError(e, message=_body, text=_body) from None
|
||||
_text = mask_sensitive_info(_safe_get_response_text(e.response))
|
||||
raise MaskedHTTPStatusError(e, message=_text, text=_text) from None
|
||||
|
||||
|
||||
async def _raise_masked_async_error(e: httpx.HTTPStatusError, stream: bool) -> None:
|
||||
"""Raise a MaskedHTTPStatusError for async HTTP handlers."""
|
||||
if stream:
|
||||
_body = mask_sensitive_info(await _safe_aread_response(e.response))
|
||||
raise MaskedHTTPStatusError(e, message=_body, text=_body) from None
|
||||
_text = mask_sensitive_info(_safe_get_response_text(e.response))
|
||||
raise MaskedHTTPStatusError(e, message=_text, text=_text) from None
|
||||
|
||||
|
||||
class MaskedHTTPStatusError(httpx.HTTPStatusError):
|
||||
def __init__(
|
||||
self, original_error, message: Optional[str] = None, text: Optional[str] = None
|
||||
):
|
||||
# Create a new error with the masked URL
|
||||
masked_url = mask_sensitive_info(str(original_error.request.url))
|
||||
# Create a new error that looks like the original, but with a masked URL
|
||||
# Mask the original exception message too (it contains the full URL)
|
||||
masked_original_message = mask_sensitive_info(str(original_error))
|
||||
|
||||
# Safely access response content — decompression can fail (e.g. zlib error)
|
||||
try:
|
||||
response_content = original_error.response.content
|
||||
except Exception:
|
||||
response_content = b""
|
||||
|
||||
super().__init__(
|
||||
message=original_error.message,
|
||||
message=masked_original_message,
|
||||
request=httpx.Request(
|
||||
method=original_error.request.method,
|
||||
url=masked_url,
|
||||
|
|
@ -334,12 +383,13 @@ class MaskedHTTPStatusError(httpx.HTTPStatusError):
|
|||
),
|
||||
response=httpx.Response(
|
||||
status_code=original_error.response.status_code,
|
||||
content=original_error.response.content,
|
||||
content=response_content,
|
||||
headers=original_error.response.headers,
|
||||
),
|
||||
)
|
||||
self.message = message
|
||||
self.text = text
|
||||
self.status_code = original_error.response.status_code
|
||||
|
||||
|
||||
class AsyncHTTPHandler:
|
||||
|
|
@ -501,16 +551,7 @@ class AsyncHTTPHandler:
|
|||
headers=headers,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
if stream is True:
|
||||
setattr(e, "message", await e.response.aread())
|
||||
setattr(e, "text", await e.response.aread())
|
||||
else:
|
||||
setattr(e, "message", mask_sensitive_info(e.response.text))
|
||||
setattr(e, "text", mask_sensitive_info(e.response.text))
|
||||
|
||||
setattr(e, "status_code", e.response.status_code)
|
||||
|
||||
raise e
|
||||
await _raise_masked_async_error(e, stream)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
|
@ -571,12 +612,7 @@ class AsyncHTTPHandler:
|
|||
headers=headers,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
setattr(e, "status_code", e.response.status_code)
|
||||
if stream is True:
|
||||
setattr(e, "message", await e.response.aread())
|
||||
else:
|
||||
setattr(e, "message", e.response.text)
|
||||
raise e
|
||||
await _raise_masked_async_error(e, stream)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
|
@ -637,12 +673,7 @@ class AsyncHTTPHandler:
|
|||
headers=headers,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
setattr(e, "status_code", e.response.status_code)
|
||||
if stream is True:
|
||||
setattr(e, "message", await e.response.aread())
|
||||
else:
|
||||
setattr(e, "message", e.response.text)
|
||||
raise e
|
||||
await _raise_masked_async_error(e, stream)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
|
@ -690,12 +721,7 @@ class AsyncHTTPHandler:
|
|||
finally:
|
||||
await new_client.aclose()
|
||||
except httpx.HTTPStatusError as e:
|
||||
setattr(e, "status_code", e.response.status_code)
|
||||
if stream is True:
|
||||
setattr(e, "message", await e.response.aread())
|
||||
else:
|
||||
setattr(e, "message", e.response.text)
|
||||
raise e
|
||||
await _raise_masked_async_error(e, stream)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
|
@ -1035,16 +1061,7 @@ class HTTPHandler:
|
|||
llm_provider="litellm-httpx-handler",
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
if stream is True:
|
||||
setattr(e, "message", mask_sensitive_info(e.response.read()))
|
||||
setattr(e, "text", mask_sensitive_info(e.response.read()))
|
||||
else:
|
||||
error_text = mask_sensitive_info(e.response.text)
|
||||
setattr(e, "message", error_text)
|
||||
setattr(e, "text", error_text)
|
||||
|
||||
setattr(e, "status_code", e.response.status_code)
|
||||
raise e
|
||||
_raise_masked_sync_error(e, stream)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
|
@ -1083,17 +1100,7 @@ class HTTPHandler:
|
|||
llm_provider="litellm-httpx-handler",
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
if stream is True:
|
||||
setattr(e, "message", mask_sensitive_info(e.response.read()))
|
||||
setattr(e, "text", mask_sensitive_info(e.response.read()))
|
||||
else:
|
||||
error_text = mask_sensitive_info(e.response.text)
|
||||
setattr(e, "message", error_text)
|
||||
setattr(e, "text", error_text)
|
||||
|
||||
setattr(e, "status_code", e.response.status_code)
|
||||
|
||||
raise e
|
||||
_raise_masked_sync_error(e, stream)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
|
@ -1130,6 +1137,8 @@ class HTTPHandler:
|
|||
model="default-model-name",
|
||||
llm_provider="litellm-httpx-handler",
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
_raise_masked_sync_error(e, stream)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
|
@ -1168,17 +1177,7 @@ class HTTPHandler:
|
|||
llm_provider="litellm-httpx-handler",
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
if stream is True:
|
||||
setattr(e, "message", mask_sensitive_info(e.response.read()))
|
||||
setattr(e, "text", mask_sensitive_info(e.response.read()))
|
||||
else:
|
||||
error_text = mask_sensitive_info(e.response.text)
|
||||
setattr(e, "message", error_text)
|
||||
setattr(e, "text", error_text)
|
||||
|
||||
setattr(e, "status_code", e.response.status_code)
|
||||
|
||||
raise e
|
||||
_raise_masked_sync_error(e, stream)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ import litellm
|
|||
import litellm.litellm_core_utils
|
||||
import litellm.types
|
||||
import litellm.types.utils
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm._logging import _redact_string, verbose_logger
|
||||
from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta
|
||||
from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES
|
||||
from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming
|
||||
|
|
@ -4789,12 +4789,12 @@ class BaseLLMHTTPHandler:
|
|||
|
||||
except websockets.exceptions.InvalidStatusCode as e: # type: ignore
|
||||
verbose_logger.exception(f"Error connecting to backend: {e}")
|
||||
await websocket.close(code=e.status_code, reason=str(e))
|
||||
await websocket.close(code=e.status_code, reason=_redact_string(str(e)))
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"Error connecting to backend: {e}")
|
||||
try:
|
||||
await websocket.close(
|
||||
code=1011, reason=f"Internal server error: {str(e)}"
|
||||
code=1011, reason=_redact_string(f"Internal server error: {str(e)}")
|
||||
)
|
||||
except RuntimeError as close_error:
|
||||
if "already completed" in str(close_error) or "websocket.close" in str(
|
||||
|
|
@ -5076,12 +5076,12 @@ class BaseLLMHTTPHandler:
|
|||
|
||||
except websockets.exceptions.InvalidStatusCode as e: # type: ignore
|
||||
verbose_logger.exception(f"Error connecting to responses WS backend: {e}")
|
||||
await websocket.close(code=e.status_code, reason=str(e))
|
||||
await websocket.close(code=e.status_code, reason=_redact_string(str(e)))
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"Error in responses WS: {e}")
|
||||
try:
|
||||
await websocket.close(
|
||||
code=1011, reason=f"Internal server error: {str(e)}"
|
||||
code=1011, reason=_redact_string(f"Internal server error: {str(e)}")
|
||||
)
|
||||
except RuntimeError as close_error:
|
||||
if "already completed" in str(close_error) or "websocket.close" in str(
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ class GeminiModelInfo(BaseLLMModelInfo):
|
|||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""Google AI Studio sends api key in query params"""
|
||||
"""Google AI Studio sends api key via x-goog-api-key header"""
|
||||
return headers
|
||||
|
||||
@property
|
||||
|
|
@ -75,7 +75,8 @@ class GeminiModelInfo(BaseLLMModelInfo):
|
|||
)
|
||||
|
||||
response = litellm.module_level_client.get(
|
||||
url=f"{api_base}{endpoint}?key={api_key}",
|
||||
url=f"{api_base}{endpoint}",
|
||||
headers={"x-goog-api-key": api_key},
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
|
|
|
|||
|
|
@ -86,7 +86,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig):
|
|||
if not final_api_key:
|
||||
raise ValueError("api_key is required")
|
||||
|
||||
url = "{}/{}?key={}".format(api_base, endpoint, final_api_key)
|
||||
url = "{}/{}".format(api_base, endpoint)
|
||||
return url
|
||||
|
||||
def get_supported_openai_params(
|
||||
|
|
@ -231,9 +231,9 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig):
|
|||
)
|
||||
api_base = api_base.rstrip("/")
|
||||
|
||||
url = f"{api_base}/v1beta/{file_part}?key={api_key}"
|
||||
url = f"{api_base}/v1beta/{file_part}"
|
||||
|
||||
# Return empty params dict - API key is already in URL, no query params needed
|
||||
# API key is passed via x-goog-api-key header (set in validate_environment)
|
||||
return url, {}
|
||||
|
||||
def _normalize_gemini_file_id(self, file_id: str) -> str:
|
||||
|
|
|
|||
|
|
@ -75,9 +75,13 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
|
|||
model: str,
|
||||
litellm_params: Optional[GenericLiteLLMParams],
|
||||
) -> dict:
|
||||
"""Google AI Studio uses API key in query params, not headers."""
|
||||
"""Google AI Studio uses x-goog-api-key header for authentication."""
|
||||
headers = headers or {}
|
||||
headers["Content-Type"] = "application/json"
|
||||
if litellm_params:
|
||||
api_key = GeminiModelInfo.get_api_key(litellm_params.get("api_key"))
|
||||
if api_key:
|
||||
headers["x-goog-api-key"] = api_key
|
||||
return headers
|
||||
|
||||
def get_complete_url(
|
||||
|
|
@ -98,11 +102,10 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
|
|||
"Google API key is required. Set GOOGLE_API_KEY or GEMINI_API_KEY environment variable."
|
||||
)
|
||||
|
||||
query_params = f"key={api_key}"
|
||||
if stream:
|
||||
query_params += "&alt=sse"
|
||||
return f"{api_base}/{self.api_version}/interactions?alt=sse"
|
||||
|
||||
return f"{api_base}/{self.api_version}/interactions?{query_params}"
|
||||
return f"{api_base}/{self.api_version}/interactions"
|
||||
|
||||
def transform_request(
|
||||
self,
|
||||
|
|
@ -200,11 +203,10 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
|
|||
) -> Tuple[str, Dict]:
|
||||
"""GET /{api_version}/interactions/{interaction_id}"""
|
||||
resolved_api_base = GeminiModelInfo.get_api_base(api_base)
|
||||
api_key = GeminiModelInfo.get_api_key(litellm_params.api_key)
|
||||
if not api_key:
|
||||
if not GeminiModelInfo.get_api_key(litellm_params.api_key):
|
||||
raise ValueError("Google API key is required")
|
||||
return (
|
||||
f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}?key={api_key}",
|
||||
f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}",
|
||||
{},
|
||||
)
|
||||
|
||||
|
|
@ -234,11 +236,10 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
|
|||
) -> Tuple[str, Dict]:
|
||||
"""DELETE /{api_version}/interactions/{interaction_id}"""
|
||||
resolved_api_base = GeminiModelInfo.get_api_base(api_base)
|
||||
api_key = GeminiModelInfo.get_api_key(litellm_params.api_key)
|
||||
if not api_key:
|
||||
if not GeminiModelInfo.get_api_key(litellm_params.api_key):
|
||||
raise ValueError("Google API key is required")
|
||||
return (
|
||||
f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}?key={api_key}",
|
||||
f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}",
|
||||
{},
|
||||
)
|
||||
|
||||
|
|
@ -265,11 +266,10 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
|
|||
) -> Tuple[str, Dict]:
|
||||
"""POST /{api_version}/interactions/{interaction_id}:cancel (if supported)"""
|
||||
resolved_api_base = GeminiModelInfo.get_api_base(api_base)
|
||||
api_key = GeminiModelInfo.get_api_key(litellm_params.api_key)
|
||||
if not api_key:
|
||||
if not GeminiModelInfo.get_api_key(litellm_params.api_key):
|
||||
raise ValueError("Google API key is required")
|
||||
return (
|
||||
f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}:cancel?key={api_key}",
|
||||
f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}:cancel",
|
||||
{},
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -85,6 +85,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
raise ValueError("api_key is required for Gemini API calls")
|
||||
api_base = api_base.replace("https://", "wss://")
|
||||
api_base = api_base.replace("http://", "ws://")
|
||||
# WebSocket connections do not support custom HTTP headers in all clients,
|
||||
# so the API key must remain as a query parameter here. This is an accepted
|
||||
# limitation; httpx is not used for WebSocket so MaskedHTTPStatusError
|
||||
# already covers the main leak vector.
|
||||
return f"{api_base}/ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent?key={api_key}"
|
||||
|
||||
def map_model_turn_event(
|
||||
|
|
|
|||
|
|
@ -48,7 +48,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
|
|||
def get_auth_credentials(
|
||||
self, litellm_params: dict
|
||||
) -> BaseVectorStoreAuthCredentials:
|
||||
"""Gemini uses API key in query params, not headers."""
|
||||
"""Gemini uses x-goog-api-key header for authentication."""
|
||||
return {}
|
||||
|
||||
def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints:
|
||||
|
|
@ -79,6 +79,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
|
|||
api_key = litellm_params.get("api_key") or get_api_key_from_env()
|
||||
if api_key:
|
||||
self._cached_api_key = api_key
|
||||
headers["x-goog-api-key"] = api_key
|
||||
|
||||
return headers
|
||||
|
||||
|
|
@ -133,13 +134,10 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
|
|||
if model and model.startswith("gemini/"):
|
||||
model = model.replace("gemini/", "")
|
||||
|
||||
# Get API key - Gemini requires it as a query parameter
|
||||
api_key = litellm_params.get("api_key") or GeminiModelInfo.get_api_key()
|
||||
if not api_key:
|
||||
raise ValueError("GEMINI_API_KEY or GOOGLE_API_KEY is required")
|
||||
|
||||
# Build the URL for generateContent with API key
|
||||
url = f"{api_base}/models/{model}:generateContent?key={api_key}"
|
||||
url = f"{api_base}/models/{model}:generateContent"
|
||||
|
||||
# Build file_search tool configuration (using snake_case as per Gemini docs)
|
||||
file_search_config: Dict[str, Any] = {
|
||||
|
|
@ -286,10 +284,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
|
|||
"""
|
||||
url = f"{api_base}/fileSearchStores"
|
||||
|
||||
# Append API key as query parameter (required by Gemini)
|
||||
api_key = self._cached_api_key or get_api_key_from_env()
|
||||
if api_key:
|
||||
url = f"{url}?key={api_key}"
|
||||
# API key is passed via x-goog-api-key header (set in validate_environment)
|
||||
|
||||
request_body: Dict[str, Any] = {}
|
||||
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ This requires websockets, and is currently only supported on LiteLLM Proxy.
|
|||
|
||||
from typing import Any, Optional, cast
|
||||
|
||||
from litellm._logging import _redact_string
|
||||
from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES
|
||||
from litellm.types.realtime import RealtimeQueryParams
|
||||
|
||||
|
|
@ -148,11 +149,11 @@ class OpenAIRealtime(OpenAIChatCompletion):
|
|||
await realtime_streaming.bidirectional_forward()
|
||||
|
||||
except websockets.exceptions.InvalidStatusCode as e: # type: ignore
|
||||
await websocket.close(code=e.status_code, reason=str(e))
|
||||
await websocket.close(code=e.status_code, reason=_redact_string(str(e)))
|
||||
except Exception as e:
|
||||
try:
|
||||
await websocket.close(
|
||||
code=1011, reason=f"Internal server error: {str(e)}"
|
||||
code=1011, reason=_redact_string(f"Internal server error: {str(e)}")
|
||||
)
|
||||
except RuntimeError as close_error:
|
||||
if "already completed" in str(close_error) or "websocket.close" in str(
|
||||
|
|
|
|||
|
|
@ -337,8 +337,13 @@ def _get_gemini_url(
|
|||
mode: all_gemini_url_modes,
|
||||
model: str,
|
||||
stream: Optional[bool],
|
||||
gemini_api_key: Optional[str],
|
||||
) -> Tuple[str, str]:
|
||||
"""Build the Gemini API URL for the given mode.
|
||||
|
||||
The API key is NOT included in the URL. Callers must pass it via the
|
||||
``x-goog-api-key`` header instead to avoid leaking credentials in
|
||||
error tracebacks.
|
||||
"""
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexGeminiConfig,
|
||||
)
|
||||
|
|
@ -352,27 +357,27 @@ def _get_gemini_url(
|
|||
endpoint = "generateContent"
|
||||
if stream is True:
|
||||
endpoint = "streamGenerateContent"
|
||||
url = "https://generativelanguage.googleapis.com/{}/{}:{}?key={}&alt=sse".format(
|
||||
api_version, _gemini_model_name, endpoint, gemini_api_key
|
||||
url = "https://generativelanguage.googleapis.com/{}/{}:{}?alt=sse".format(
|
||||
api_version, _gemini_model_name, endpoint
|
||||
)
|
||||
else:
|
||||
url = "https://generativelanguage.googleapis.com/{}/{}:{}?key={}".format(
|
||||
api_version, _gemini_model_name, endpoint, gemini_api_key
|
||||
url = "https://generativelanguage.googleapis.com/{}/{}:{}".format(
|
||||
api_version, _gemini_model_name, endpoint
|
||||
)
|
||||
elif mode == "embedding":
|
||||
endpoint = "embedContent"
|
||||
url = "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}".format(
|
||||
_gemini_model_name, endpoint, gemini_api_key
|
||||
url = "https://generativelanguage.googleapis.com/v1beta/{}:{}".format(
|
||||
_gemini_model_name, endpoint
|
||||
)
|
||||
elif mode == "batch_embedding":
|
||||
endpoint = "batchEmbedContents"
|
||||
url = "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}".format(
|
||||
_gemini_model_name, endpoint, gemini_api_key
|
||||
url = "https://generativelanguage.googleapis.com/v1beta/{}:{}".format(
|
||||
_gemini_model_name, endpoint
|
||||
)
|
||||
elif mode == "count_tokens":
|
||||
endpoint = "countTokens"
|
||||
url = "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}".format(
|
||||
_gemini_model_name, endpoint, gemini_api_key
|
||||
url = "https://generativelanguage.googleapis.com/v1beta/{}:{}".format(
|
||||
_gemini_model_name, endpoint
|
||||
)
|
||||
elif mode == "image_generation":
|
||||
raise ValueError(
|
||||
|
|
|
|||
|
|
@ -62,10 +62,10 @@ class ContextCachingEndpoints(VertexBase):
|
|||
token, url
|
||||
"""
|
||||
if custom_llm_provider == "gemini":
|
||||
auth_header = None
|
||||
auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment]
|
||||
endpoint = "cachedContents"
|
||||
url = "https://generativelanguage.googleapis.com/v1beta/{}?key={}".format(
|
||||
endpoint, gemini_api_key
|
||||
url = "https://generativelanguage.googleapis.com/v1beta/{}".format(
|
||||
endpoint
|
||||
)
|
||||
elif custom_llm_provider == "vertex_ai":
|
||||
auth_header = vertex_auth_header
|
||||
|
|
@ -353,7 +353,9 @@ class ContextCachingEndpoints(VertexBase):
|
|||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
if token is not None:
|
||||
if isinstance(token, dict):
|
||||
headers.update(token)
|
||||
elif token is not None:
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
if extra_headers is not None:
|
||||
headers.update(extra_headers)
|
||||
|
|
@ -501,7 +503,9 @@ class ContextCachingEndpoints(VertexBase):
|
|||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
if token is not None:
|
||||
if isinstance(token, dict):
|
||||
headers.update(token)
|
||||
elif token is not None:
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
if extra_headers is not None:
|
||||
headers.update(extra_headers)
|
||||
|
|
|
|||
|
|
@ -412,7 +412,7 @@ class VertexBase:
|
|||
url = "{}/models/{}:{}".format(api_base, model, endpoint)
|
||||
if gemini_api_key is None:
|
||||
raise ValueError(
|
||||
"Missing gemini_api_key, please set `GEMINI_API_KEY`"
|
||||
"Missing Gemini API key. Set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable."
|
||||
)
|
||||
if gemini_api_key is not None:
|
||||
auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment]
|
||||
|
|
@ -469,13 +469,16 @@ class VertexBase:
|
|||
"""
|
||||
version: Optional[Literal["v1beta1", "v1"]] = None
|
||||
if custom_llm_provider == "gemini":
|
||||
if not gemini_api_key:
|
||||
raise ValueError(
|
||||
"Missing Gemini API key. Set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable."
|
||||
)
|
||||
url, endpoint = _get_gemini_url(
|
||||
mode=mode,
|
||||
model=model,
|
||||
stream=stream,
|
||||
gemini_api_key=gemini_api_key,
|
||||
)
|
||||
auth_header = None # this field is not used for gemin
|
||||
auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment]
|
||||
else:
|
||||
vertex_location = self.get_vertex_region(
|
||||
vertex_region=vertex_location,
|
||||
|
|
|
|||
|
|
@ -40,6 +40,7 @@ from typing import (
|
|||
get_args,
|
||||
)
|
||||
|
||||
from litellm._logging import _redact_string
|
||||
from litellm._uuid import uuid
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -7244,7 +7245,7 @@ async def ahealth_check(
|
|||
f"Mode {mode} not supported. See modes here: https://docs.litellm.ai/docs/proxy/health"
|
||||
)
|
||||
except Exception as e:
|
||||
stack_trace = traceback.format_exc()
|
||||
stack_trace = _redact_string(traceback.format_exc())
|
||||
if isinstance(stack_trace, str):
|
||||
stack_trace = stack_trace[:1000]
|
||||
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ from fastapi import HTTPException, Request, status
|
|||
from fastapi.responses import JSONResponse, Response, StreamingResponse
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._logging import _redact_string, verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import (
|
||||
DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE,
|
||||
|
|
@ -1785,7 +1785,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
|
||||
if isinstance(e, HTTPException):
|
||||
raise e
|
||||
error_traceback = traceback.format_exc()
|
||||
error_traceback = _redact_string(traceback.format_exc())
|
||||
error_msg = f"{str(e)}\n\n{error_traceback}"
|
||||
proxy_exception = ProxyException(
|
||||
message=getattr(e, "message", error_msg),
|
||||
|
|
|
|||
|
|
@ -77,7 +77,7 @@ from litellm import (
|
|||
ModelResponseStream,
|
||||
Router,
|
||||
)
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._logging import _redact_string, verbose_proxy_logger
|
||||
from litellm._service_logger import ServiceLogging, ServiceTypes
|
||||
from litellm.caching.caching import DualCache, RedisCache
|
||||
from litellm.caching.dual_cache import LimitedSizeOrderedDict
|
||||
|
|
@ -155,7 +155,7 @@ def print_verbose(print_statement):
|
|||
|
||||
verbose_proxy_logger.debug("{}\n{}".format(print_statement, traceback.format_exc()))
|
||||
if litellm.set_verbose:
|
||||
print(f"LiteLLM Proxy: {print_statement}") # noqa
|
||||
print(f"LiteLLM Proxy: {_redact_string(str(print_statement))}") # noqa
|
||||
|
||||
|
||||
def _get_email_logger_class():
|
||||
|
|
@ -1721,6 +1721,7 @@ class ProxyLogging:
|
|||
error_message = str(original_exception)
|
||||
if isinstance(traceback_str, str):
|
||||
error_message += traceback_str[:1000]
|
||||
error_message = _redact_string(error_message)
|
||||
asyncio.create_task(
|
||||
self.alerting_handler(
|
||||
message=f"DB read/write call failed: {error_message}",
|
||||
|
|
@ -1791,7 +1792,7 @@ class ProxyLogging:
|
|||
|
||||
asyncio.create_task(
|
||||
self.alerting_handler(
|
||||
message=f"LLM API call failed: `{exception_str}`",
|
||||
message=_redact_string(f"LLM API call failed: `{exception_str}`"),
|
||||
level="High",
|
||||
alert_type=AlertType.llm_exceptions,
|
||||
request_data=request_data,
|
||||
|
|
|
|||
|
|
@ -143,7 +143,7 @@ class GeminiRAGIngestion(BaseRAGIngestion):
|
|||
Returns:
|
||||
Store name (format: fileSearchStores/xxxxxxx)
|
||||
"""
|
||||
url = f"{base_url}/fileSearchStores?key={api_key}"
|
||||
url = f"{base_url}/fileSearchStores"
|
||||
|
||||
request_body = {"displayName": display_name}
|
||||
|
||||
|
|
@ -154,7 +154,10 @@ class GeminiRAGIngestion(BaseRAGIngestion):
|
|||
response = await client.post(
|
||||
url,
|
||||
json=request_body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"x-goog-api-key": api_key,
|
||||
},
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
|
|
@ -228,7 +231,7 @@ class GeminiRAGIngestion(BaseRAGIngestion):
|
|||
# base_url is like: https://generativelanguage.googleapis.com/v1beta
|
||||
# We need: https://generativelanguage.googleapis.com/upload/v1beta/{store_id}:uploadToFileSearchStore
|
||||
api_base = base_url.replace("/v1beta", "") # Get base without version
|
||||
url = f"{api_base}/upload/v1beta/{vector_store_id}:uploadToFileSearchStore?key={api_key}"
|
||||
url = f"{api_base}/upload/v1beta/{vector_store_id}:uploadToFileSearchStore"
|
||||
|
||||
# Build request body with chunking config and metadata if provided
|
||||
request_body: Dict[str, Any] = {"displayName": filename}
|
||||
|
|
@ -263,6 +266,7 @@ class GeminiRAGIngestion(BaseRAGIngestion):
|
|||
"X-Goog-Upload-Header-Content-Length": str(file_size),
|
||||
"X-Goog-Upload-Header-Content-Type": content_type,
|
||||
"Content-Type": "application/json",
|
||||
"x-goog-api-key": api_key,
|
||||
}
|
||||
|
||||
verbose_logger.debug(f"Initiating resumable upload: {url}")
|
||||
|
|
|
|||
|
|
@ -310,7 +310,7 @@ def test_gemini_multimodal_embedding_e2e():
|
|||
) as mock_get_token:
|
||||
mock_get_token.return_value = (
|
||||
{"x-goog-api-key": "test-key"},
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:embedContent?key=test-key"
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:embedContent"
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
|
|
|
|||
|
|
@ -320,29 +320,30 @@ async def test_generationconfig_to_config_mapping(sample_request_payload):
|
|||
for Google GenAI compatibility in the main functions.
|
||||
"""
|
||||
from litellm.google_genai.main import agenerate_content
|
||||
|
||||
|
||||
# Create a copy of the payload to avoid modifying the fixture
|
||||
test_data = sample_request_payload.copy()
|
||||
|
||||
# Test that agenerate_content can handle generationConfig parameter
|
||||
# This should not raise an error about parameter handling
|
||||
try:
|
||||
# This will fail due to missing API key, but should not fail due to parameter handling
|
||||
|
||||
with patch(
|
||||
"litellm.google_genai.main.base_llm_http_handler.generate_content_handler"
|
||||
) as mock_generate_content_handler:
|
||||
mock_generate_content_handler.return_value = {"text": "mock response"}
|
||||
|
||||
await agenerate_content(
|
||||
model="gemini/gemini-2.5-flash",
|
||||
contents=test_data["contents"],
|
||||
generationConfig=test_data["generationConfig"], # Pass as generationConfig
|
||||
custom_llm_provider="gemini"
|
||||
generationConfig=test_data["generationConfig"],
|
||||
custom_llm_provider="gemini",
|
||||
)
|
||||
except Exception as e:
|
||||
# Should not fail due to parameter handling issues
|
||||
error_msg = str(e).lower()
|
||||
if "generationconfig" in error_msg or "config" in error_msg or "parameter" in error_msg:
|
||||
pytest.fail(f"Parameter handling failed: {e}")
|
||||
# Other errors (like API key missing) are expected
|
||||
print(f"✅ Parameter handling worked (API error expected): {type(e).__name__}")
|
||||
|
||||
print("✅ generationConfig to config mapping test passed")
|
||||
|
||||
mock_generate_content_handler.assert_called_once()
|
||||
generate_content_config_dict = mock_generate_content_handler.call_args.kwargs[
|
||||
"generate_content_config_dict"
|
||||
]
|
||||
assert generate_content_config_dict["temperature"] == 0
|
||||
assert generate_content_config_dict["topP"] == 1
|
||||
assert generate_content_config_dict["responseMimeType"] == "application/json"
|
||||
assert "responseJsonSchema" in generate_content_config_dict
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -405,7 +406,7 @@ async def test_gemini_custom_api_base_proxy_integration():
|
|||
print(f"✅ Custom API base streaming URL test passed: {result_url_streaming}")
|
||||
|
||||
# Test case 3: Error handling - missing API key
|
||||
with pytest.raises(ValueError, match="Missing gemini_api_key"):
|
||||
with pytest.raises(ValueError, match="Missing Gemini API key"):
|
||||
vertex_base._check_custom_proxy(
|
||||
api_base=custom_api_base,
|
||||
custom_llm_provider="gemini",
|
||||
|
|
|
|||
|
|
@ -1127,80 +1127,58 @@ async def test_google_generate_content_with_openai():
|
|||
passed_fields = passed_fields - set(GenericLiteLLMParams.model_fields.keys())
|
||||
# extra_headers is now explicitly passed through for providers that need custom headers
|
||||
assert passed_fields == set(["model", "messages", "extra_headers"]), f"Expected model, messages, and extra_headers to be passed through, got {passed_fields}"
|
||||
@pytest.mark.asyncio
|
||||
async def test_agenerate_content_x_goog_api_key_header():
|
||||
def test_validate_environment_sets_x_goog_api_key():
|
||||
"""
|
||||
Test that agenerate_content passes x-goog-api-key header correctly.
|
||||
|
||||
This test verifies that when calling agenerate_content with a Google GenAI model,
|
||||
the HTTP request includes the x-goog-api-key header with the correct API key value.
|
||||
"""
|
||||
import os
|
||||
import unittest.mock
|
||||
Test that VertexGeminiConfig.validate_environment correctly merges an
|
||||
x-goog-api-key dict into the request headers.
|
||||
|
||||
This is the mechanism by which Google AI Studio (Gemini) requests get
|
||||
authenticated via header instead of a query-string ?key= parameter.
|
||||
"""
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexGeminiConfig,
|
||||
)
|
||||
|
||||
import httpx
|
||||
|
||||
test_api_key = "test-gemini-api-key-123"
|
||||
|
||||
# Mock environment to ensure we use our test API key
|
||||
with unittest.mock.patch.dict(os.environ, {"GEMINI_API_KEY": test_api_key}, clear=False):
|
||||
# Mock the AsyncHTTPHandler's post method to capture headers
|
||||
with unittest.mock.patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", new_callable=unittest.mock.AsyncMock) as mock_post:
|
||||
# Mock a successful response
|
||||
mock_response = unittest.mock.MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {
|
||||
"parts": [{"text": "Hello! How can I help you today?"}],
|
||||
"role": "model"
|
||||
},
|
||||
"finishReason": "STOP",
|
||||
"index": 0
|
||||
}
|
||||
],
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 5,
|
||||
"candidatesTokenCount": 10,
|
||||
"totalTokenCount": 15
|
||||
}
|
||||
}
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {}
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
# Call agenerate_content with Google AI Studio model
|
||||
try:
|
||||
response = await agenerate_content(
|
||||
model="gemini/gemini-1.5-flash",
|
||||
contents=[
|
||||
{"role": "user", "parts": [{"text": "Hello, world!"}]}
|
||||
],
|
||||
api_key=test_api_key
|
||||
)
|
||||
except Exception:
|
||||
# Ignore any response processing errors, we just want to check the headers
|
||||
pass
|
||||
|
||||
# Verify that AsyncHTTPHandler.post was called
|
||||
mock_post.assert_called_once()
|
||||
|
||||
# Get the arguments passed to the post call
|
||||
call_args, call_kwargs = mock_post.call_args
|
||||
|
||||
# Verify that headers contain x-goog-api-key
|
||||
headers = call_kwargs.get("headers", {})
|
||||
assert "x-goog-api-key" in headers, f"x-goog-api-key header not found in headers: {list(headers.keys())}"
|
||||
|
||||
# Verify the API key is set (could be our test key or from api_key parameter)
|
||||
api_key_value = headers["x-goog-api-key"]
|
||||
assert api_key_value == test_api_key, f"Expected x-goog-api-key to be {test_api_key}, got {api_key_value}"
|
||||
|
||||
# Verify other expected headers
|
||||
assert headers.get("Content-Type") == "application/json", f"Expected Content-Type application/json, got {headers.get('Content-Type')}"
|
||||
|
||||
print(f"✓ Test passed: x-goog-api-key header correctly set to {api_key_value}")
|
||||
print(f"✓ All headers: {list(headers.keys())}")
|
||||
# Simulate what _get_token_and_url returns for Gemini: a dict auth_header
|
||||
auth_header_dict = {"x-goog-api-key": test_api_key}
|
||||
|
||||
headers = VertexGeminiConfig().validate_environment(
|
||||
api_key=auth_header_dict,
|
||||
headers=None,
|
||||
model="gemini-2.5-flash",
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
|
||||
assert "x-goog-api-key" in headers, f"x-goog-api-key not in headers: {headers}"
|
||||
assert headers["x-goog-api-key"] == test_api_key
|
||||
assert headers["Content-Type"] == "application/json"
|
||||
|
||||
|
||||
def test_get_gemini_url_excludes_api_key():
|
||||
"""
|
||||
Verify that _get_gemini_url never embeds the API key in the URL.
|
||||
|
||||
API keys in URLs leak through httpx error tracebacks. The key must be
|
||||
sent via the x-goog-api-key header instead.
|
||||
"""
|
||||
from litellm.llms.vertex_ai.common_utils import _get_gemini_url
|
||||
|
||||
for mode in ("chat", "embedding", "batch_embedding", "count_tokens"):
|
||||
url, _ = _get_gemini_url(
|
||||
mode=mode,
|
||||
model="gemini-2.5-flash",
|
||||
stream=False,
|
||||
)
|
||||
assert "key=" not in url, f"API key found in URL for mode={mode}: {url}"
|
||||
|
||||
# Streaming chat should only have ?alt=sse
|
||||
url, _ = _get_gemini_url(mode="chat", model="gemini-2.5-flash", stream=True)
|
||||
assert "key=" not in url, f"API key found in streaming URL: {url}"
|
||||
assert "alt=sse" in url, f"Missing alt=sse in streaming URL: {url}"
|
||||
|
||||
|
||||
def test_inline_data_base64_image_transformation():
|
||||
|
|
|
|||
|
|
@ -0,0 +1,148 @@
|
|||
"""
|
||||
Tests for Gemini Interactions API transformation.
|
||||
|
||||
Covers credential leak prevention changes:
|
||||
- validate_environment sets x-goog-api-key header
|
||||
- get_complete_url excludes API key from URL
|
||||
- get/delete/cancel interaction request URLs exclude API key
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
from litellm.llms.gemini.interactions.transformation import (
|
||||
GoogleAIStudioInteractionsConfig,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
_PATCH_GET_API_KEY = "litellm.llms.gemini.common_utils.GeminiModelInfo.get_api_key"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def config():
|
||||
return GoogleAIStudioInteractionsConfig()
|
||||
|
||||
|
||||
class TestValidateEnvironment:
|
||||
def test_sets_x_goog_api_key_header(self, config):
|
||||
litellm_params = GenericLiteLLMParams(api_key="test-api-key-123")
|
||||
|
||||
headers = config.validate_environment(
|
||||
headers={},
|
||||
model="gemini-2.5-flash",
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
assert headers["x-goog-api-key"] == "test-api-key-123"
|
||||
assert headers["Content-Type"] == "application/json"
|
||||
|
||||
def test_no_api_key_skips_header(self, config):
|
||||
litellm_params = GenericLiteLLMParams(api_key=None)
|
||||
|
||||
with patch(_PATCH_GET_API_KEY, return_value=None):
|
||||
headers = config.validate_environment(
|
||||
headers={},
|
||||
model="gemini-2.5-flash",
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
assert "x-goog-api-key" not in headers
|
||||
assert headers["Content-Type"] == "application/json"
|
||||
|
||||
def test_no_litellm_params_skips_header(self, config):
|
||||
headers = config.validate_environment(
|
||||
headers={},
|
||||
model="gemini-2.5-flash",
|
||||
litellm_params=None,
|
||||
)
|
||||
|
||||
assert "x-goog-api-key" not in headers
|
||||
assert headers["Content-Type"] == "application/json"
|
||||
|
||||
def test_preserves_existing_headers(self, config):
|
||||
litellm_params = GenericLiteLLMParams(api_key="test-key")
|
||||
|
||||
headers = config.validate_environment(
|
||||
headers={"X-Custom": "value"},
|
||||
model="gemini-2.5-flash",
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
assert headers["X-Custom"] == "value"
|
||||
assert headers["x-goog-api-key"] == "test-key"
|
||||
|
||||
|
||||
class TestGetCompleteUrl:
|
||||
def test_url_excludes_api_key(self, config):
|
||||
with patch(_PATCH_GET_API_KEY, return_value="secret-key"):
|
||||
url = config.get_complete_url(
|
||||
api_base=None,
|
||||
model="gemini-2.5-flash",
|
||||
litellm_params={"api_key": "secret-key"},
|
||||
)
|
||||
|
||||
assert "key=" not in url
|
||||
assert "secret-key" not in url
|
||||
assert url.endswith("/interactions")
|
||||
|
||||
def test_stream_url_has_alt_sse_only(self, config):
|
||||
with patch(_PATCH_GET_API_KEY, return_value="secret-key"):
|
||||
url = config.get_complete_url(
|
||||
api_base=None,
|
||||
model="gemini-2.5-flash",
|
||||
litellm_params={"api_key": "secret-key"},
|
||||
stream=True,
|
||||
)
|
||||
|
||||
assert "key=" not in url
|
||||
assert "secret-key" not in url
|
||||
assert "alt=sse" in url
|
||||
|
||||
def test_raises_without_api_key(self, config):
|
||||
with patch(_PATCH_GET_API_KEY, return_value=None):
|
||||
with pytest.raises(ValueError, match="Google API key is required"):
|
||||
config.get_complete_url(
|
||||
api_base=None,
|
||||
model="gemini-2.5-flash",
|
||||
litellm_params={"api_key": None},
|
||||
)
|
||||
|
||||
|
||||
class TestInteractionOperationUrls:
|
||||
"""Test that get/delete/cancel interaction URLs exclude API key."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"method_name,interaction_id,expected_suffix",
|
||||
[
|
||||
("transform_get_interaction_request", "interaction-123", "interaction-123"),
|
||||
("transform_delete_interaction_request", "interaction-456", "interaction-456"),
|
||||
("transform_cancel_interaction_request", "interaction-789", "interaction-789:cancel"),
|
||||
],
|
||||
)
|
||||
def test_url_excludes_key(self, config, method_name, interaction_id, expected_suffix):
|
||||
with patch(_PATCH_GET_API_KEY, return_value="secret-key"):
|
||||
url, params = getattr(config, method_name)(
|
||||
interaction_id=interaction_id,
|
||||
api_base="https://generativelanguage.googleapis.com",
|
||||
litellm_params=GenericLiteLLMParams(api_key="secret-key"),
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert "key=" not in url
|
||||
assert "secret-key" not in url
|
||||
assert expected_suffix in url
|
||||
|
||||
def test_get_interaction_raises_without_key(self, config):
|
||||
with patch(_PATCH_GET_API_KEY, return_value=None):
|
||||
with pytest.raises(ValueError, match="Google API key is required"):
|
||||
config.transform_get_interaction_request(
|
||||
interaction_id="interaction-123",
|
||||
api_base="https://generativelanguage.googleapis.com",
|
||||
litellm_params=GenericLiteLLMParams(api_key=None),
|
||||
headers={},
|
||||
)
|
||||
|
|
@ -0,0 +1,234 @@
|
|||
"""
|
||||
Tests for credential leak prevention in HTTP handlers.
|
||||
|
||||
Covers:
|
||||
- MaskedHTTPStatusError construction and masking behavior
|
||||
- _safe_get_response_text, _safe_aread_response, _safe_read_response helpers
|
||||
- _raise_masked_sync_error and _raise_masked_async_error
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../.."))
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
HTTPHandler,
|
||||
MaskedHTTPStatusError,
|
||||
_raise_masked_async_error,
|
||||
_raise_masked_sync_error,
|
||||
_safe_aread_response,
|
||||
_safe_get_response_text,
|
||||
_safe_read_response,
|
||||
)
|
||||
|
||||
|
||||
def _make_httpx_status_error(
|
||||
status_code: int = 400,
|
||||
url: str = "https://example.com/v1/models?key=SECRET_KEY_123",
|
||||
body: str = "Bad Request",
|
||||
) -> httpx.HTTPStatusError:
|
||||
"""Create a real httpx.HTTPStatusError for testing."""
|
||||
request = httpx.Request("POST", url)
|
||||
response = httpx.Response(status_code, request=request, content=body.encode())
|
||||
return httpx.HTTPStatusError(
|
||||
message=f"Client error '{status_code}' for url '{url}'",
|
||||
request=request,
|
||||
response=response,
|
||||
)
|
||||
|
||||
|
||||
class TestMaskedHTTPStatusError:
|
||||
def test_masks_url_in_request(self):
|
||||
orig = _make_httpx_status_error(url="https://api.example.com?key=MY_SECRET")
|
||||
masked = MaskedHTTPStatusError(orig)
|
||||
|
||||
assert "MY_SECRET" not in str(masked.request.url)
|
||||
assert "[REDACTED_API_KEY]" in str(masked.request.url)
|
||||
|
||||
def test_masks_original_message(self):
|
||||
orig = _make_httpx_status_error(url="https://api.example.com?key=SUPER_SECRET")
|
||||
masked = MaskedHTTPStatusError(orig)
|
||||
|
||||
assert "SUPER_SECRET" not in str(masked)
|
||||
assert "[REDACTED_API_KEY]" in str(masked)
|
||||
|
||||
def test_preserves_status_code(self):
|
||||
orig = _make_httpx_status_error(status_code=403)
|
||||
masked = MaskedHTTPStatusError(orig)
|
||||
|
||||
assert masked.status_code == 403
|
||||
assert masked.response.status_code == 403
|
||||
|
||||
def test_preserves_message_and_text_attrs(self):
|
||||
orig = _make_httpx_status_error()
|
||||
masked = MaskedHTTPStatusError(orig, message="custom msg", text="custom text")
|
||||
|
||||
assert masked.message == "custom msg"
|
||||
assert masked.text == "custom text"
|
||||
|
||||
def test_handles_response_content_decompression_failure(self):
|
||||
"""If response.content raises (e.g. zlib error), should fall back to b''."""
|
||||
orig = _make_httpx_status_error()
|
||||
|
||||
with patch.object(
|
||||
type(orig.response), "content",
|
||||
new_callable=lambda: property(lambda self: (_ for _ in ()).throw(Exception("zlib error"))),
|
||||
):
|
||||
masked = MaskedHTTPStatusError(orig)
|
||||
|
||||
assert masked.response.content == b""
|
||||
assert masked.status_code == 400
|
||||
|
||||
|
||||
class TestSafeResponseHelpers:
|
||||
def test_safe_get_response_text_normal(self):
|
||||
response = httpx.Response(200, content=b"hello world")
|
||||
assert _safe_get_response_text(response) == "hello world"
|
||||
|
||||
def test_safe_get_response_text_error(self):
|
||||
response = MagicMock(spec=httpx.Response)
|
||||
type(response).text = property(lambda self: (_ for _ in ()).throw(UnicodeDecodeError("utf-8", b"", 0, 1, "bad")))
|
||||
assert _safe_get_response_text(response) == ""
|
||||
|
||||
def test_safe_read_response_normal(self):
|
||||
response = httpx.Response(200, content=b"raw bytes")
|
||||
result = _safe_read_response(response)
|
||||
assert result == b"raw bytes"
|
||||
|
||||
def test_safe_read_response_error(self):
|
||||
response = MagicMock(spec=httpx.Response)
|
||||
response.read.side_effect = Exception("read failure")
|
||||
assert _safe_read_response(response) == b""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_safe_aread_response_normal(self):
|
||||
response = MagicMock(spec=httpx.Response)
|
||||
response.aread = AsyncMock(return_value=b"async bytes")
|
||||
result = await _safe_aread_response(response)
|
||||
assert result == b"async bytes"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_safe_aread_response_error(self):
|
||||
response = MagicMock(spec=httpx.Response)
|
||||
response.aread = AsyncMock(side_effect=Exception("async read failure"))
|
||||
result = await _safe_aread_response(response)
|
||||
assert result == b""
|
||||
|
||||
|
||||
class TestRaiseMaskedError:
|
||||
def test_sync_non_stream(self):
|
||||
orig = _make_httpx_status_error(
|
||||
url="https://api.example.com?key=LEAKED_KEY", body="error body"
|
||||
)
|
||||
with pytest.raises(MaskedHTTPStatusError) as exc_info:
|
||||
_raise_masked_sync_error(orig, stream=False)
|
||||
|
||||
err = exc_info.value
|
||||
assert "LEAKED_KEY" not in str(err.request.url)
|
||||
assert err.status_code == 400
|
||||
assert err.text == "error body"
|
||||
|
||||
def test_sync_stream(self):
|
||||
orig = _make_httpx_status_error(
|
||||
url="https://api.example.com?key=LEAKED_KEY", body="stream body"
|
||||
)
|
||||
with pytest.raises(MaskedHTTPStatusError) as exc_info:
|
||||
_raise_masked_sync_error(orig, stream=True)
|
||||
|
||||
err = exc_info.value
|
||||
assert "LEAKED_KEY" not in str(err.request.url)
|
||||
assert err.message is not None
|
||||
|
||||
def test_sync_breaks_exception_chain(self):
|
||||
orig = _make_httpx_status_error()
|
||||
with pytest.raises(MaskedHTTPStatusError) as exc_info:
|
||||
_raise_masked_sync_error(orig, stream=False)
|
||||
|
||||
assert exc_info.value.__cause__ is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_non_stream(self):
|
||||
orig = _make_httpx_status_error(
|
||||
url="https://api.example.com?key=LEAKED_KEY", body="async error"
|
||||
)
|
||||
with pytest.raises(MaskedHTTPStatusError) as exc_info:
|
||||
await _raise_masked_async_error(orig, stream=False)
|
||||
|
||||
err = exc_info.value
|
||||
assert "LEAKED_KEY" not in str(err.request.url)
|
||||
assert err.status_code == 400
|
||||
assert err.text == "async error"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_stream(self):
|
||||
orig = _make_httpx_status_error(
|
||||
url="https://api.example.com?key=LEAKED_KEY", body="async stream"
|
||||
)
|
||||
with pytest.raises(MaskedHTTPStatusError) as exc_info:
|
||||
await _raise_masked_async_error(orig, stream=True)
|
||||
|
||||
err = exc_info.value
|
||||
assert "LEAKED_KEY" not in str(err.request.url)
|
||||
assert err.message is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_breaks_chain(self):
|
||||
orig = _make_httpx_status_error()
|
||||
with pytest.raises(MaskedHTTPStatusError) as exc_info:
|
||||
await _raise_masked_async_error(orig, stream=False)
|
||||
|
||||
assert exc_info.value.__cause__ is None
|
||||
|
||||
|
||||
class TestHTTPHandlerErrorPaths:
|
||||
"""Test that HTTP handler methods raise MaskedHTTPStatusError on HTTPStatusError."""
|
||||
|
||||
@pytest.fixture
|
||||
def sync_handler(self):
|
||||
handler = HTTPHandler()
|
||||
yield handler
|
||||
handler.close()
|
||||
|
||||
@pytest.fixture
|
||||
async def async_handler(self):
|
||||
handler = AsyncHTTPHandler()
|
||||
yield handler
|
||||
await handler.close()
|
||||
|
||||
@pytest.mark.parametrize("method", ["post", "put", "patch", "delete"])
|
||||
def test_sync_raises_masked_error(self, sync_handler, method):
|
||||
with patch.object(
|
||||
sync_handler.client,
|
||||
"send",
|
||||
side_effect=_make_httpx_status_error(url="https://api.test.com?key=SECRET"),
|
||||
):
|
||||
with pytest.raises(MaskedHTTPStatusError) as exc_info:
|
||||
kwargs = {"url": "https://api.test.com?key=SECRET"}
|
||||
if method != "delete":
|
||||
kwargs["data"] = {"test": 1}
|
||||
getattr(sync_handler, method)(**kwargs)
|
||||
|
||||
assert "SECRET" not in str(exc_info.value.request.url)
|
||||
|
||||
@pytest.mark.parametrize("method", ["post", "put", "patch", "delete"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_raises_masked_error(self, async_handler, method):
|
||||
with patch.object(
|
||||
async_handler.client,
|
||||
"send",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=_make_httpx_status_error(url="https://api.test.com?key=SECRET"),
|
||||
):
|
||||
with pytest.raises(MaskedHTTPStatusError) as exc_info:
|
||||
kwargs = {"url": "https://api.test.com?key=SECRET"}
|
||||
if method != "delete":
|
||||
kwargs["data"] = {"test": 1}
|
||||
await getattr(async_handler, method)(**kwargs)
|
||||
|
||||
assert "SECRET" not in str(exc_info.value.request.url)
|
||||
|
|
@ -37,12 +37,12 @@ class TestGoogleAIStudioFilesTransformation:
|
|||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
# Verify URL is constructed exactly as required:
|
||||
# https://generativelanguage.googleapis.com/v1beta/files/{file_id}?key=API_KEY
|
||||
# API key is passed via x-goog-api-key header, not in URL
|
||||
assert (
|
||||
url
|
||||
== "https://generativelanguage.googleapis.com/v1beta/files/test123?key=test-api-key"
|
||||
== "https://generativelanguage.googleapis.com/v1beta/files/test123"
|
||||
)
|
||||
assert "key=" not in url
|
||||
|
||||
# CRITICAL: params should be empty dict, not contain Content-Type or any other params
|
||||
# These would be incorrectly interpreted as query parameters
|
||||
|
|
@ -64,12 +64,12 @@ class TestGoogleAIStudioFilesTransformation:
|
|||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
# Verify URL is constructed exactly as required:
|
||||
# https://generativelanguage.googleapis.com/v1beta/files/{file_id}?key=API_KEY
|
||||
# API key is passed via x-goog-api-key header, not in URL
|
||||
assert (
|
||||
url
|
||||
== "https://generativelanguage.googleapis.com/v1beta/files/test123?key=test-api-key"
|
||||
== "https://generativelanguage.googleapis.com/v1beta/files/test123"
|
||||
)
|
||||
assert "key=" not in url
|
||||
|
||||
# CRITICAL: params should be empty dict
|
||||
assert params == {}, f"Expected empty params dict, got: {params}"
|
||||
|
|
@ -79,11 +79,10 @@ class TestGoogleAIStudioFilesTransformation:
|
|||
|
||||
def test_transform_retrieve_file_request_with_raw_id_only(self):
|
||||
"""
|
||||
Regression guard for the exact retrieval URL format.
|
||||
Regression guard: API key must NOT appear in the URL.
|
||||
|
||||
If someone changes the method and stops producing:
|
||||
https://generativelanguage.googleapis.com/v1beta/files/{file_id}?key=API_KEY
|
||||
this test should fail.
|
||||
The key is sent via x-goog-api-key header to prevent leaking
|
||||
credentials in httpx error tracebacks.
|
||||
"""
|
||||
file_id = "cctqueckiggb"
|
||||
litellm_params = {"api_key": "test-api-key"}
|
||||
|
|
@ -96,8 +95,9 @@ class TestGoogleAIStudioFilesTransformation:
|
|||
|
||||
assert (
|
||||
url
|
||||
== "https://generativelanguage.googleapis.com/v1beta/files/cctqueckiggb?key=test-api-key"
|
||||
== "https://generativelanguage.googleapis.com/v1beta/files/cctqueckiggb"
|
||||
)
|
||||
assert "key=" not in url
|
||||
assert params == {}
|
||||
|
||||
@patch.dict("os.environ", {}, clear=True)
|
||||
|
|
@ -285,10 +285,10 @@ class TestGoogleAIStudioFilesTransformation:
|
|||
litellm_params={},
|
||||
)
|
||||
|
||||
# Verify URL structure
|
||||
# Verify URL structure - API key must NOT be in URL
|
||||
assert api_base in url
|
||||
assert "upload/v1beta/files" in url
|
||||
assert f"key={api_key}" in url
|
||||
assert "key=" not in url
|
||||
|
||||
def test_transform_delete_file_request_with_full_uri(self):
|
||||
"""Test delete file request transformation with full URI"""
|
||||
|
|
|
|||
|
|
@ -93,6 +93,7 @@ def test_completion_pydantic_obj_2():
|
|||
model="gemini/gemini-2.5-flash",
|
||||
messages=messages,
|
||||
response_format=EventsList,
|
||||
api_key="test-api-key",
|
||||
client=client,
|
||||
)
|
||||
# print(response)
|
||||
|
|
@ -285,6 +286,7 @@ def test_function_calling_with_gemini():
|
|||
},
|
||||
},
|
||||
],
|
||||
api_key="test-api-key",
|
||||
client=client,
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -372,7 +374,10 @@ def test_multiple_function_call():
|
|||
|
||||
with patch.object(client, "post", return_value=mock_response) as mock_post:
|
||||
r = litellm.completion(
|
||||
messages=messages, model="gemini/gemini-1.5-flash-002", client=client
|
||||
messages=messages,
|
||||
model="gemini/gemini-1.5-flash-002",
|
||||
api_key="test-api-key",
|
||||
client=client,
|
||||
)
|
||||
assert len(r.choices) > 0
|
||||
|
||||
|
|
@ -478,7 +483,10 @@ def test_multiple_function_call_changed_text_pos():
|
|||
|
||||
with patch.object(client, "post", return_value=mock_response) as mock_post:
|
||||
resp = litellm.completion(
|
||||
messages=messages, model="gemini/gemini-1.5-flash-002", client=client
|
||||
messages=messages,
|
||||
model="gemini/gemini-1.5-flash-002",
|
||||
api_key="test-api-key",
|
||||
client=client,
|
||||
)
|
||||
assert len(resp.choices) > 0
|
||||
mock_post.assert_called_once()
|
||||
|
|
@ -599,6 +607,7 @@ def test_function_calling_with_gemini_multiple_results():
|
|||
messages=messages,
|
||||
tools=tools,
|
||||
tool_choice="required",
|
||||
api_key="test-api-key",
|
||||
client=client,
|
||||
)
|
||||
print("Response\n", response)
|
||||
|
|
@ -1182,6 +1191,7 @@ def test_logprobs():
|
|||
{"role": "user", "content": "What's the weather like in San Francisco?"}
|
||||
],
|
||||
logprobs=True,
|
||||
api_key="test-api-key",
|
||||
client=client,
|
||||
)
|
||||
print(resp)
|
||||
|
|
|
|||
|
|
@ -810,7 +810,7 @@ class TestVertexBase:
|
|||
|
||||
if custom_llm_provider == "gemini" and api_base and gemini_api_key is None:
|
||||
# Test case 5: Should raise ValueError for Gemini without API key
|
||||
with pytest.raises(ValueError, match="Missing gemini_api_key"):
|
||||
with pytest.raises(ValueError, match="Missing Gemini API key"):
|
||||
vertex_base._check_custom_proxy(
|
||||
api_base=api_base,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
|
|
|
|||
|
|
@ -247,8 +247,10 @@ class TestPromptVersionsEndpoint:
|
|||
),
|
||||
}
|
||||
|
||||
# Mock the IN_MEMORY_PROMPT_REGISTRY at the import location
|
||||
with patch("litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY") as mock_registry:
|
||||
# Force the in-memory path so this test is isolated from any leaked prisma mocks.
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", None), patch(
|
||||
"litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY"
|
||||
) as mock_registry:
|
||||
mock_registry.IN_MEMORY_PROMPTS = mock_prompts
|
||||
|
||||
# Test with base prompt ID
|
||||
|
|
@ -293,7 +295,9 @@ class TestPromptVersionsEndpoint:
|
|||
user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY") as mock_registry:
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", None), patch(
|
||||
"litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY"
|
||||
) as mock_registry:
|
||||
mock_registry.IN_MEMORY_PROMPTS = {}
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
|
|
@ -304,4 +308,3 @@ class TestPromptVersionsEndpoint:
|
|||
|
||||
assert exc_info.value.status_code == 404
|
||||
assert "No versions found" in exc_info.value.detail
|
||||
|
||||
|
|
|
|||
|
|
@ -213,6 +213,8 @@ async def test_url_with_format_param(model, sync_mode, monkeypatch):
|
|||
}
|
||||
],
|
||||
}
|
||||
if model.startswith("gemini/"):
|
||||
args["api_key"] = "test-api-key"
|
||||
with patch.object(client, "post", new=MagicMock()) as mock_client:
|
||||
try:
|
||||
if sync_mode:
|
||||
|
|
|
|||
245
tests/test_litellm/test_redact_string_in_error_paths.py
Normal file
245
tests/test_litellm/test_redact_string_in_error_paths.py
Normal file
|
|
@ -0,0 +1,245 @@
|
|||
"""
|
||||
Tests for _redact_string usage in error/logging paths.
|
||||
|
||||
Covers actual execution of redaction in:
|
||||
- WebSocket close reasons in realtime handlers (openai, azure, bedrock)
|
||||
- Gemini RAG ingestion x-goog-api-key header usage
|
||||
- Traceback redaction pattern used in proxy streaming
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../.."))
|
||||
|
||||
from litellm._logging import _ENABLE_SECRET_REDACTION, _redact_string
|
||||
|
||||
|
||||
class TestRedactStringFunction:
|
||||
def test_redacts_bearer_token(self):
|
||||
text = "Authorization: Bearer sk-1234567890abcdefghij"
|
||||
result = _redact_string(text)
|
||||
assert "sk-1234567890abcdefghij" not in result
|
||||
assert "REDACTED" in result
|
||||
|
||||
def test_redacts_api_key_in_url(self):
|
||||
text = "Error at https://example.com?api_key=my-secret-key-value-here"
|
||||
result = _redact_string(text)
|
||||
assert "my-secret-key-value-here" not in result
|
||||
|
||||
def test_redacts_google_api_key(self):
|
||||
text = "key=AIzaSyB1234567890abcdefghijklmnopqrstuvwx"
|
||||
result = _redact_string(text)
|
||||
assert "AIzaSyB1234567890abcdefghijklmnopqrstuvwx" not in result
|
||||
|
||||
def test_passes_clean_text_through(self):
|
||||
text = "This is a normal error message with no secrets"
|
||||
assert _redact_string(text) == text
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _ENABLE_SECRET_REDACTION, reason="redaction disabled via env var"
|
||||
)
|
||||
def test_redaction_enabled_by_default(self):
|
||||
text = "Bearer sk-1234567890abcdefghij"
|
||||
result = _redact_string(text)
|
||||
assert "sk-1234567890abcdefghij" not in result
|
||||
|
||||
|
||||
class TestOpenAIRealtimeRedaction:
|
||||
"""Test that OpenAI realtime handler redacts secrets in websocket close reasons."""
|
||||
|
||||
def _make_patches(self, handler):
|
||||
"""Shared patches for OpenAI realtime handler tests."""
|
||||
return (
|
||||
patch.object(handler, "_construct_url", return_value="wss://api.openai.com/v1/realtime?model=gpt-4"),
|
||||
patch.object(handler, "_get_ssl_config", return_value=None),
|
||||
patch.object(handler, "_get_additional_headers", return_value={}),
|
||||
)
|
||||
|
||||
def _call_kwargs(self):
|
||||
return dict(
|
||||
model="gpt-4",
|
||||
websocket=AsyncMock(),
|
||||
logging_obj=MagicMock(),
|
||||
api_base="https://api.openai.com/",
|
||||
api_key="test-key",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_status_code_redacts_reason(self):
|
||||
import websockets.exceptions
|
||||
|
||||
from litellm.llms.openai.realtime.handler import OpenAIRealtime
|
||||
|
||||
handler = OpenAIRealtime()
|
||||
exc = websockets.exceptions.InvalidStatusCode(403, None)
|
||||
exc.status_code = 403
|
||||
|
||||
kwargs = self._call_kwargs()
|
||||
mock_ws = kwargs["websocket"]
|
||||
p1, p2, p3 = self._make_patches(handler)
|
||||
with p1, p2, p3, patch("websockets.connect", side_effect=exc):
|
||||
await handler.async_realtime(**kwargs)
|
||||
|
||||
mock_ws.close.assert_called_once()
|
||||
assert mock_ws.close.call_args[1]["code"] == 403
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_exception_redacts_reason(self):
|
||||
from litellm.llms.openai.realtime.handler import OpenAIRealtime
|
||||
|
||||
handler = OpenAIRealtime()
|
||||
secret_error = RuntimeError("Connection failed for api_key=sk-1234567890abcdefghij")
|
||||
|
||||
kwargs = self._call_kwargs()
|
||||
mock_ws = kwargs["websocket"]
|
||||
p1, p2, p3 = self._make_patches(handler)
|
||||
with p1, p2, p3, patch("websockets.connect", side_effect=secret_error):
|
||||
await handler.async_realtime(**kwargs)
|
||||
|
||||
mock_ws.close.assert_called_once()
|
||||
assert mock_ws.close.call_args[1]["code"] == 1011
|
||||
assert "sk-1234567890abcdefghij" not in mock_ws.close.call_args[1]["reason"]
|
||||
|
||||
|
||||
class TestAzureRealtimeRedaction:
|
||||
"""Test that Azure realtime handler redacts secrets in websocket close reasons."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_status_code_redacts_reason(self):
|
||||
import websockets.exceptions
|
||||
|
||||
from litellm.llms.azure.realtime.handler import AzureOpenAIRealtime
|
||||
|
||||
handler = AzureOpenAIRealtime()
|
||||
mock_ws = AsyncMock()
|
||||
exc = websockets.exceptions.InvalidStatusCode(403, None)
|
||||
exc.status_code = 403
|
||||
|
||||
with patch.object(handler, "_construct_url", return_value="wss://test.openai.azure.com/openai/realtime"), \
|
||||
patch("websockets.connect", side_effect=exc):
|
||||
await handler.async_realtime(
|
||||
model="gpt-4",
|
||||
websocket=mock_ws,
|
||||
logging_obj=MagicMock(),
|
||||
api_base="https://test.openai.azure.com/",
|
||||
api_key="test-key",
|
||||
api_version="2024-10-01-preview",
|
||||
)
|
||||
|
||||
mock_ws.close.assert_called_once()
|
||||
assert mock_ws.close.call_args[1]["code"] == 403
|
||||
|
||||
|
||||
class TestBedrockRealtimeRedaction:
|
||||
"""Test that _redact_string produces safe close reasons for Bedrock-style errors."""
|
||||
|
||||
def test_internal_error_message_redacted(self):
|
||||
secret_error = RuntimeError("Failed with aws_secret_access_key=AKIAIOSFODNN7EXAMPLE123456")
|
||||
reason = _redact_string(f"Internal error: {str(secret_error)}")
|
||||
assert "AKIAIOSFODNN7EXAMPLE123456" not in reason
|
||||
|
||||
|
||||
class TestLLMHTTPHandlerRealtimeRedaction:
|
||||
"""Test _redact_string on the exact patterns used in llm_http_handler WS close."""
|
||||
|
||||
def test_invalid_status_pattern(self):
|
||||
error_msg = "InvalidStatusCode: 403 for wss://api.example.com?api_key=sk-leaked-key-here"
|
||||
assert "sk-leaked-key-here" not in _redact_string(str(error_msg))
|
||||
|
||||
def test_internal_server_error_pattern(self):
|
||||
error_msg = "Connection failed for api_key=sk-secret-key-12345678"
|
||||
assert "sk-secret-key-12345678" not in _redact_string(f"Internal server error: {error_msg}")
|
||||
|
||||
|
||||
class TestProxyStreamingDataGeneratorRedaction:
|
||||
"""Test _redact_string on traceback.format_exc() — the pattern at common_request_processing.py:1733."""
|
||||
|
||||
def test_redact_traceback_format_exc(self):
|
||||
try:
|
||||
raise RuntimeError(
|
||||
"Failed connecting to api_key=sk-1234567890abcdefghij at https://api.example.com"
|
||||
)
|
||||
except RuntimeError:
|
||||
raw_tb = traceback.format_exc()
|
||||
|
||||
redacted_tb = _redact_string(raw_tb)
|
||||
|
||||
assert "sk-1234567890abcdefghij" not in redacted_tb
|
||||
assert "Traceback" in redacted_tb
|
||||
assert "RuntimeError" in redacted_tb
|
||||
|
||||
|
||||
def _make_mock_ingest_options():
|
||||
mock = MagicMock()
|
||||
mock.vector_store_config = {}
|
||||
mock.ingest_name = "test"
|
||||
mock.chunking_strategy = None
|
||||
mock.embedding_model = None
|
||||
mock.vector_db_type = "gemini"
|
||||
return mock
|
||||
|
||||
|
||||
class TestGeminiIngestionHeaders:
|
||||
"""Test that Gemini RAG ingestion uses x-goog-api-key header."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_file_search_store_sends_header(self):
|
||||
from litellm.rag.ingestion.gemini_ingestion import GeminiRAGIngestion
|
||||
|
||||
ingestion = GeminiRAGIngestion(ingest_options=_make_mock_ingest_options())
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {"name": "fileSearchStores/abc123"}
|
||||
mock_client.post.return_value = mock_response
|
||||
|
||||
with patch(
|
||||
"litellm.rag.ingestion.gemini_ingestion.get_async_httpx_client",
|
||||
return_value=mock_client,
|
||||
):
|
||||
result = await ingestion._create_file_search_store(
|
||||
api_key="test-gemini-key",
|
||||
base_url="https://generativelanguage.googleapis.com/v1beta",
|
||||
display_name="test-store",
|
||||
)
|
||||
|
||||
assert result == "fileSearchStores/abc123"
|
||||
call_kwargs = mock_client.post.call_args
|
||||
assert call_kwargs[1]["headers"]["x-goog-api-key"] == "test-gemini-key"
|
||||
assert "key=" not in call_kwargs[0][0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initiate_resumable_upload_sends_header(self):
|
||||
from litellm.rag.ingestion.gemini_ingestion import GeminiRAGIngestion
|
||||
|
||||
ingestion = GeminiRAGIngestion(ingest_options=_make_mock_ingest_options())
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {"x-goog-upload-url": "https://upload.example.com/upload123"}
|
||||
mock_client.post.return_value = mock_response
|
||||
|
||||
with patch(
|
||||
"litellm.rag.ingestion.gemini_ingestion.get_async_httpx_client",
|
||||
return_value=mock_client,
|
||||
):
|
||||
result = await ingestion._initiate_resumable_upload(
|
||||
api_key="test-gemini-key",
|
||||
base_url="https://generativelanguage.googleapis.com/v1beta",
|
||||
vector_store_id="fileSearchStores/abc123",
|
||||
filename="test.txt",
|
||||
file_size=1024,
|
||||
content_type="text/plain",
|
||||
)
|
||||
|
||||
assert result == "https://upload.example.com/upload123"
|
||||
call_kwargs = mock_client.post.call_args
|
||||
assert call_kwargs[1]["headers"]["x-goog-api-key"] == "test-gemini-key"
|
||||
assert "key=" not in call_kwargs[0][0]
|
||||
Loading…
Add table
Reference in a new issue