mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
Merge pull request #20930 from BerriAI/litellm_oss_staging_02_11_2026
oss staging 02 / 11/ 2026
This commit is contained in:
commit
a7179797f7
25 changed files with 1240 additions and 145 deletions
|
|
@ -1,52 +1,22 @@
|
|||
# Custom Semgrep Rules
|
||||
# Custom Semgrep rules for LiteLLM
|
||||
|
||||
All `.yml` files under `.semgrep/rules/` run in CI (CircleCI `semgrep` job).
|
||||
Add custom rule YAML files here. Semgrep loads all `.yml`/`.yaml` files under this directory.
|
||||
|
||||
## Add a Rule
|
||||
|
||||
* Add a `.yml` file under `.semgrep/rules/<language>/<domain>/`
|
||||
|
||||
|
||||
[Rule syntax →](https://semgrep.dev/docs/writing-rules/rule-syntax/)
|
||||
|
||||
## Organizing Rules
|
||||
|
||||
### Structure: language → domain
|
||||
|
||||
```
|
||||
.semgrep/rules/<language>/<domain>/<rule-name>.yml
|
||||
```
|
||||
|
||||
Examples:
|
||||
|
||||
- `python/security/unsafe-yaml-load.yml`
|
||||
- `python/reliability/missing-timeout-http.yml`
|
||||
- `python/performance/blocking-io-in-async.yml`
|
||||
|
||||
### Rule metadata
|
||||
|
||||
Match tags to the folder for consistent filtering:
|
||||
|
||||
```yaml
|
||||
metadata:
|
||||
tags: [python, security]
|
||||
```
|
||||
|
||||
### Severity expectations
|
||||
|
||||
All rules must fail CI on findings. No warn-only rules.
|
||||
|
||||
- Use `severity: ERROR` in rule metadata
|
||||
- If a rule is noisy → refine until low false positives before adding
|
||||
|
||||
## Run Locally
|
||||
**Run only custom rules (CI / fail on findings):**
|
||||
|
||||
```bash
|
||||
semgrep scan --config .semgrep/rules . --error
|
||||
```
|
||||
|
||||
With Semgrep registry:
|
||||
**Run with registry + custom rules:**
|
||||
|
||||
```bash
|
||||
semgrep scan --config auto --config .semgrep/rules .
|
||||
```
|
||||
|
||||
**Layout:**
|
||||
|
||||
- `python/` – Python-specific rules (security, patterns)
|
||||
- Add more subdirs as needed (e.g. `generic/` for language-agnostic rules)
|
||||
|
||||
See [Semgrep rule syntax](https://semgrep.dev/docs/writing-rules/rule-syntax/).
|
||||
|
|
|
|||
14
.semgrep/rules/python/unbounded-memory.yml
Normal file
14
.semgrep/rules/python/unbounded-memory.yml
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
# Unbounded memory growth – data structures without a clear max limit
|
||||
# Can lead to OOM under load.
|
||||
|
||||
rules:
|
||||
- id: unbounded-asyncio-queue
|
||||
message: asyncio.Queue() with no maxsize can grow unbounded. Use asyncio.Queue(maxsize=N) for integrations (e.g. log queues).
|
||||
severity: ERROR
|
||||
languages: [python]
|
||||
pattern-either:
|
||||
- pattern: asyncio.Queue()
|
||||
- pattern: asyncio.Queue(maxsize=0)
|
||||
metadata:
|
||||
category: correctness
|
||||
cwe: "CWE-400: Uncontrolled Resource Consumption"
|
||||
|
|
@ -93,6 +93,12 @@ Implement `POST /beta/litellm_basic_guardrail_api`
|
|||
"user_api_key_end_user_id": "end user id associated with the litellm virtual key used",
|
||||
"user_api_key_org_id": "org id associated with the litellm virtual key used"
|
||||
},
|
||||
"request_headers": { // optional: inbound request headers (allowlist). Allowed headers show their value; all others show "[present]" to indicate the header existed.
|
||||
"User-Agent": "OpenAI/Python 2.17.0",
|
||||
"Content-Type": "application/json",
|
||||
"X-Request-Id": "[present]"
|
||||
},
|
||||
"litellm_version": "1.x.y", // optional: LiteLLM library version running this proxy
|
||||
"input_type": "request", // "request" or "response"
|
||||
"litellm_call_id": "unique_call_id", // the call id of the individual LLM call
|
||||
"litellm_trace_id": "trace_id", // the trace id of the LLM call - useful if there are multiple LLM calls for the same conversation
|
||||
|
|
|
|||
|
|
@ -175,6 +175,7 @@ _async_failure_callback: List[Union[str, Callable, "CustomLogger"]] = ( # Custo
|
|||
pre_call_rules: List[Callable] = []
|
||||
post_call_rules: List[Callable] = []
|
||||
turn_off_message_logging: Optional[bool] = False
|
||||
standard_logging_payload_excluded_fields: Optional[List[str]] = None # Fields to exclude from StandardLoggingPayload before callbacks receive it
|
||||
log_raw_request_response: bool = False
|
||||
redact_messages_in_exceptions: Optional[bool] = False
|
||||
redact_user_api_key_info: Optional[bool] = False
|
||||
|
|
|
|||
|
|
@ -774,15 +774,17 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
self, model_call_details: Dict
|
||||
) -> Dict:
|
||||
"""
|
||||
Only redacts messages and responses when self.turn_off_message_logging is True
|
||||
Redacts or excludes fields from StandardLoggingPayload before callbacks receive it.
|
||||
|
||||
This method handles two features:
|
||||
1. turn_off_message_logging: When True, redacts messages and responses
|
||||
2. standard_logging_payload_excluded_fields: Removes specified fields entirely
|
||||
|
||||
By default, self.turn_off_message_logging is False and this does nothing.
|
||||
|
||||
Return a redacted deepcopy of the provided logging payload.
|
||||
Return a modified copy of the provided logging payload.
|
||||
|
||||
This is useful for logging payloads that contain sensitive information.
|
||||
"""
|
||||
import litellm
|
||||
from copy import copy
|
||||
|
||||
from litellm import Choices, Message, ModelResponse
|
||||
|
|
@ -790,14 +792,17 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
turn_off_message_logging: bool = getattr(
|
||||
self, "turn_off_message_logging", False
|
||||
)
|
||||
excluded_fields: Optional[List[str]] = getattr(
|
||||
litellm, "standard_logging_payload_excluded_fields", None
|
||||
)
|
||||
|
||||
if turn_off_message_logging is False:
|
||||
# Early return if no processing needed
|
||||
if turn_off_message_logging is False and not excluded_fields:
|
||||
return model_call_details
|
||||
|
||||
# Only make a shallow copy of the top-level dict to avoid deepcopy issues
|
||||
# with complex objects like AuthenticationError that may be present
|
||||
model_call_details_copy = copy(model_call_details)
|
||||
redacted_str = "redacted-by-litellm"
|
||||
standard_logging_object = model_call_details.get("standard_logging_object")
|
||||
if standard_logging_object is None:
|
||||
return model_call_details_copy
|
||||
|
|
@ -805,39 +810,58 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
# Make a copy of just the standard_logging_object to avoid modifying the original
|
||||
standard_logging_object_copy = copy(standard_logging_object)
|
||||
|
||||
if standard_logging_object_copy.get("messages") is not None:
|
||||
standard_logging_object_copy["messages"] = [
|
||||
Message(content=redacted_str).model_dump()
|
||||
]
|
||||
# Handle excluded fields - remove them entirely from the payload
|
||||
if excluded_fields:
|
||||
for field in excluded_fields:
|
||||
if field in standard_logging_object_copy:
|
||||
del standard_logging_object_copy[field]
|
||||
|
||||
if standard_logging_object_copy.get("response") is not None:
|
||||
response = standard_logging_object_copy["response"]
|
||||
# Check if this is a ResponsesAPIResponse (has "output" field)
|
||||
if isinstance(response, dict) and "output" in response:
|
||||
# Make a copy to avoid modifying the original
|
||||
from copy import deepcopy
|
||||
# Handle turn_off_message_logging - redact messages and responses (if not already excluded)
|
||||
if turn_off_message_logging:
|
||||
redacted_str = "redacted-by-litellm"
|
||||
|
||||
response_copy = deepcopy(response)
|
||||
# Redact content in output array
|
||||
if isinstance(response_copy.get("output"), list):
|
||||
for output_item in response_copy["output"]:
|
||||
if isinstance(output_item, dict) and "content" in output_item:
|
||||
if isinstance(output_item["content"], list):
|
||||
# Redact text in content items
|
||||
for content_item in output_item["content"]:
|
||||
if (
|
||||
isinstance(content_item, dict)
|
||||
and "text" in content_item
|
||||
):
|
||||
content_item["text"] = redacted_str
|
||||
standard_logging_object_copy["response"] = response_copy
|
||||
else:
|
||||
# Standard ModelResponse format
|
||||
model_response = ModelResponse(
|
||||
choices=[Choices(message=Message(content=redacted_str))]
|
||||
)
|
||||
model_response_dict = model_response.model_dump()
|
||||
standard_logging_object_copy["response"] = model_response_dict
|
||||
if (
|
||||
"messages" not in (excluded_fields or [])
|
||||
and standard_logging_object_copy.get("messages") is not None
|
||||
):
|
||||
standard_logging_object_copy["messages"] = [
|
||||
Message(content=redacted_str).model_dump()
|
||||
]
|
||||
|
||||
if (
|
||||
"response" not in (excluded_fields or [])
|
||||
and standard_logging_object_copy.get("response") is not None
|
||||
):
|
||||
response = standard_logging_object_copy["response"]
|
||||
# Check if this is a ResponsesAPIResponse (has "output" field)
|
||||
if isinstance(response, dict) and "output" in response:
|
||||
# Make a copy to avoid modifying the original
|
||||
from copy import deepcopy
|
||||
|
||||
response_copy = deepcopy(response)
|
||||
# Redact content in output array
|
||||
if isinstance(response_copy.get("output"), list):
|
||||
for output_item in response_copy["output"]:
|
||||
if (
|
||||
isinstance(output_item, dict)
|
||||
and "content" in output_item
|
||||
):
|
||||
if isinstance(output_item["content"], list):
|
||||
# Redact text in content items
|
||||
for content_item in output_item["content"]:
|
||||
if (
|
||||
isinstance(content_item, dict)
|
||||
and "text" in content_item
|
||||
):
|
||||
content_item["text"] = redacted_str
|
||||
standard_logging_object_copy["response"] = response_copy
|
||||
else:
|
||||
# Standard ModelResponse format
|
||||
model_response = ModelResponse(
|
||||
choices=[Choices(message=Message(content=redacted_str))]
|
||||
)
|
||||
model_response_dict = model_response.model_dump()
|
||||
standard_logging_object_copy["response"] = model_response_dict
|
||||
|
||||
model_call_details_copy["standard_logging_object"] = (
|
||||
standard_logging_object_copy
|
||||
|
|
|
|||
|
|
@ -98,16 +98,18 @@ class ExceptionCheckers:
|
|||
"""
|
||||
Check if an error string indicates a content policy violation error.
|
||||
"""
|
||||
_lower = error_str.lower()
|
||||
known_exception_substrings = [
|
||||
"invalid_request_error",
|
||||
"content_policy_violation",
|
||||
"responsibleaipolicyviolation",
|
||||
"the response was filtered due to the prompt triggering azure openai's content management",
|
||||
"your task failed as a result of our safety system",
|
||||
"the model produced invalid content",
|
||||
"content_filter_policy",
|
||||
"your request was rejected as a result of our safety system",
|
||||
]
|
||||
for substring in known_exception_substrings:
|
||||
if substring in error_str.lower():
|
||||
if substring in _lower:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
|
@ -2060,6 +2062,19 @@ def exception_type( # type: ignore # noqa: PLR0915
|
|||
if isinstance(body_dict, dict):
|
||||
if isinstance(body_dict.get("error"), dict):
|
||||
azure_error_code = body_dict["error"].get("code") # type: ignore[index]
|
||||
# Also check inner_error for
|
||||
# ResponsibleAIPolicyViolation which indicates a
|
||||
# content policy violation even when the top-level
|
||||
# code is generic (e.g. "invalid_request_error").
|
||||
if azure_error_code != "content_policy_violation":
|
||||
_inner = (
|
||||
body_dict["error"].get("inner_error") # type: ignore[index]
|
||||
or body_dict["error"].get("innererror") # type: ignore[index]
|
||||
)
|
||||
if isinstance(_inner, dict) and _inner.get(
|
||||
"code"
|
||||
) == "ResponsibleAIPolicyViolation":
|
||||
azure_error_code = "content_policy_violation"
|
||||
else:
|
||||
azure_error_code = body_dict.get("code")
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -664,35 +664,34 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
reasoning_effort: Optional[Union[REASONING_EFFORT, str]],
|
||||
model: str,
|
||||
) -> Optional[AnthropicThinkingParam]:
|
||||
if reasoning_effort is None or reasoning_effort == "none":
|
||||
return None
|
||||
if AnthropicConfig._is_claude_opus_4_6(model):
|
||||
return AnthropicThinkingParam(
|
||||
type="adaptive",
|
||||
)
|
||||
elif reasoning_effort == "low":
|
||||
return AnthropicThinkingParam(
|
||||
type="enabled",
|
||||
budget_tokens=DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET,
|
||||
)
|
||||
elif reasoning_effort == "medium":
|
||||
return AnthropicThinkingParam(
|
||||
type="enabled",
|
||||
budget_tokens=DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET,
|
||||
)
|
||||
elif reasoning_effort == "high":
|
||||
return AnthropicThinkingParam(
|
||||
type="enabled",
|
||||
budget_tokens=DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET,
|
||||
)
|
||||
elif reasoning_effort == "minimal":
|
||||
return AnthropicThinkingParam(
|
||||
type="enabled",
|
||||
budget_tokens=DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET,
|
||||
)
|
||||
else:
|
||||
if reasoning_effort is None:
|
||||
return None
|
||||
elif reasoning_effort == "low":
|
||||
return AnthropicThinkingParam(
|
||||
type="enabled",
|
||||
budget_tokens=DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET,
|
||||
)
|
||||
elif reasoning_effort == "medium":
|
||||
return AnthropicThinkingParam(
|
||||
type="enabled",
|
||||
budget_tokens=DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET,
|
||||
)
|
||||
elif reasoning_effort == "high":
|
||||
return AnthropicThinkingParam(
|
||||
type="enabled",
|
||||
budget_tokens=DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET,
|
||||
)
|
||||
elif reasoning_effort == "minimal":
|
||||
return AnthropicThinkingParam(
|
||||
type="enabled",
|
||||
budget_tokens=DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unmapped reasoning effort: {reasoning_effort}")
|
||||
raise ValueError(f"Unmapped reasoning effort: {reasoning_effort}")
|
||||
|
||||
def _extract_json_schema_from_response_format(
|
||||
self, value: Optional[dict]
|
||||
|
|
|
|||
|
|
@ -901,7 +901,20 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
|
||||
if response.json()["status"] == "failed":
|
||||
error_data = response.json()
|
||||
raise AzureOpenAIError(status_code=400, message=json.dumps(error_data))
|
||||
# Preserve Azure error details (e.g. content_policy_violation,
|
||||
# inner_error, content_filter_results) as structured body so
|
||||
# exception_type() can route them correctly.
|
||||
_error_body = error_data.get("error", error_data)
|
||||
_error_msg = (
|
||||
_error_body.get("message", "Image generation failed")
|
||||
if isinstance(_error_body, dict)
|
||||
else json.dumps(error_data)
|
||||
)
|
||||
raise AzureOpenAIError(
|
||||
status_code=400,
|
||||
message=_error_msg,
|
||||
body=error_data,
|
||||
)
|
||||
|
||||
result = response.json()["result"]
|
||||
return httpx.Response(
|
||||
|
|
@ -999,7 +1012,20 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
|
||||
if response.json()["status"] == "failed":
|
||||
error_data = response.json()
|
||||
raise AzureOpenAIError(status_code=400, message=json.dumps(error_data))
|
||||
# Preserve Azure error details (e.g. content_policy_violation,
|
||||
# inner_error, content_filter_results) as structured body so
|
||||
# exception_type() can route them correctly.
|
||||
_error_body = error_data.get("error", error_data)
|
||||
_error_msg = (
|
||||
_error_body.get("message", "Image generation failed")
|
||||
if isinstance(_error_body, dict)
|
||||
else json.dumps(error_data)
|
||||
)
|
||||
raise AzureOpenAIError(
|
||||
status_code=400,
|
||||
message=_error_msg,
|
||||
body=error_data,
|
||||
)
|
||||
|
||||
result = response.json()["result"]
|
||||
return httpx.Response(
|
||||
|
|
|
|||
|
|
@ -14,6 +14,8 @@ import re
|
|||
from typing import Any, Callable, Dict, List, Literal, Optional, Set, Tuple, Union, cast
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import anyio
|
||||
|
||||
from fastapi import HTTPException
|
||||
from httpx import HTTPStatusError
|
||||
from mcp import ReadResourceResult, Resource
|
||||
|
|
@ -1437,6 +1439,9 @@ class MCPServerManager:
|
|||
"""
|
||||
Fetch tools from MCP client with timeout and error handling.
|
||||
|
||||
Uses anyio.fail_after() instead of asyncio.wait_for() to avoid conflicts
|
||||
with the MCP SDK's anyio TaskGroup. See GitHub issue #20715 for details.
|
||||
|
||||
Args:
|
||||
client: MCP client instance
|
||||
server_name: Name of the server for logging
|
||||
|
|
@ -1444,24 +1449,12 @@ class MCPServerManager:
|
|||
Returns:
|
||||
List of tools from the server
|
||||
"""
|
||||
|
||||
async def _list_tools_task():
|
||||
try:
|
||||
try:
|
||||
with anyio.fail_after(30.0):
|
||||
tools = await client.list_tools()
|
||||
verbose_logger.debug(f"Tools from {server_name}: {tools}")
|
||||
return tools
|
||||
except asyncio.CancelledError:
|
||||
verbose_logger.warning(f"Client operation cancelled for {server_name}")
|
||||
return []
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
f"Client operation failed for {server_name}: {str(e)}"
|
||||
)
|
||||
return []
|
||||
|
||||
try:
|
||||
return await asyncio.wait_for(_list_tools_task(), timeout=30.0)
|
||||
except asyncio.TimeoutError:
|
||||
except TimeoutError:
|
||||
verbose_logger.warning(f"Timeout while listing tools from {server_name}")
|
||||
return []
|
||||
except asyncio.CancelledError:
|
||||
|
|
|
|||
|
|
@ -112,7 +112,8 @@ class SpendUpdateQueue(BaseUpdateQueue):
|
|||
for update in updates:
|
||||
_key = f"{update.get('entity_type')}:{update.get('entity_id')}"
|
||||
if _key not in _in_memory_map:
|
||||
_in_memory_map[_key] = update
|
||||
# avoid mutating caller-owned dicts while aggregating queue entries
|
||||
_in_memory_map[_key] = update.copy()
|
||||
else:
|
||||
current_cost = _in_memory_map[_key].get("response_cost", 0) or 0
|
||||
update_cost = update.get("response_cost", 0) or 0
|
||||
|
|
|
|||
|
|
@ -5,10 +5,12 @@
|
|||
# +-------------------------------------------------------------+
|
||||
# Thank you users! We ❤️ you! - Krrish & Ishaan
|
||||
|
||||
import fnmatch
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Any, Dict, Literal, Optional
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._version import version as litellm_version
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
|
|
@ -31,6 +33,110 @@ if TYPE_CHECKING:
|
|||
|
||||
GUARDRAIL_NAME = "generic_guardrail_api"
|
||||
|
||||
# Headers whose values are forwarded as-is (case-insensitive). Glob patterns supported (e.g. x-stainless-*, x-litellm*).
|
||||
_HEADER_VALUE_ALLOWLIST = frozenset({
|
||||
"host",
|
||||
"accept-encoding",
|
||||
"connection",
|
||||
"accept",
|
||||
"content-type",
|
||||
"user-agent",
|
||||
"x-stainless-*",
|
||||
"x-litellm-*",
|
||||
"content-length",
|
||||
})
|
||||
|
||||
# Placeholder for headers that exist but are not on the allowlist (we don't expose their value).
|
||||
_HEADER_PRESENT_PLACEHOLDER = "[present]"
|
||||
|
||||
|
||||
def _header_value_allowed(header_name: str) -> bool:
|
||||
"""Return True if this header's value may be forwarded (allowlist, including globs)."""
|
||||
lower = header_name.lower()
|
||||
if lower in _HEADER_VALUE_ALLOWLIST:
|
||||
return True
|
||||
for pattern in _HEADER_VALUE_ALLOWLIST:
|
||||
if "*" in pattern and fnmatch.fnmatch(lower, pattern):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _sanitize_inbound_headers(headers: Any) -> Optional[Dict[str, str]]:
|
||||
"""
|
||||
Sanitize inbound headers before passing them to a 3rd party guardrail service.
|
||||
|
||||
- Allowlist: only headers in the allowlist have their values forwarded (exact + glob: x-stainless-*, x-litellm-*).
|
||||
- All other headers are included with value "[present]" so the guardrail knows the header existed.
|
||||
- Coerces values to str (for JSON serialization).
|
||||
"""
|
||||
if not headers or not isinstance(headers, dict):
|
||||
return None
|
||||
|
||||
sanitized: Dict[str, str] = {}
|
||||
for k, v in headers.items():
|
||||
if k is None:
|
||||
continue
|
||||
key = str(k)
|
||||
if _header_value_allowed(key):
|
||||
try:
|
||||
sanitized[key] = str(v)
|
||||
except Exception:
|
||||
continue
|
||||
else:
|
||||
sanitized[key] = _HEADER_PRESENT_PLACEHOLDER
|
||||
|
||||
return sanitized or None
|
||||
|
||||
|
||||
def _extract_inbound_headers(
|
||||
request_data: dict, logging_obj: Optional["LiteLLMLoggingObj"]
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""
|
||||
Extract inbound headers from available request context.
|
||||
|
||||
We try multiple locations to support different call paths:
|
||||
- proxy endpoints: request_data["proxy_server_request"]["headers"]
|
||||
- if the guardrail is passed the proxy_server_request object directly
|
||||
- metadata headers captured in litellm_pre_call_utils
|
||||
- response hooks: fallback to logging_obj.model_call_details
|
||||
"""
|
||||
# 1) Most common path (proxy): full request context in proxy_server_request
|
||||
headers = request_data.get("proxy_server_request", {}).get("headers")
|
||||
if headers:
|
||||
return _sanitize_inbound_headers(headers)
|
||||
|
||||
# 2) Some guardrails pass proxy_server_request as request_data itself
|
||||
headers = request_data.get("headers")
|
||||
if headers:
|
||||
return _sanitize_inbound_headers(headers)
|
||||
|
||||
# 3) Pre-call: headers stored in request metadata
|
||||
metadata_headers = (request_data.get("metadata") or {}).get("headers")
|
||||
if metadata_headers:
|
||||
return _sanitize_inbound_headers(metadata_headers)
|
||||
|
||||
litellm_metadata_headers = (request_data.get("litellm_metadata") or {}).get(
|
||||
"headers"
|
||||
)
|
||||
if litellm_metadata_headers:
|
||||
return _sanitize_inbound_headers(litellm_metadata_headers)
|
||||
|
||||
# 4) Post-call: headers not present on response; fallback to logging object
|
||||
if logging_obj and getattr(logging_obj, "model_call_details", None):
|
||||
try:
|
||||
details = logging_obj.model_call_details or {}
|
||||
headers = (
|
||||
details.get("litellm_params", {})
|
||||
.get("metadata", {})
|
||||
.get("headers", None)
|
||||
)
|
||||
if headers:
|
||||
return _sanitize_inbound_headers(headers)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return None
|
||||
|
||||
|
||||
class GenericGuardrailAPI(CustomGuardrail):
|
||||
"""
|
||||
|
|
@ -207,6 +313,7 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
|
||||
# Extract user API key metadata
|
||||
user_metadata = self._extract_user_api_key_metadata(request_data)
|
||||
inbound_headers = _extract_inbound_headers(request_data=request_data, logging_obj=logging_obj)
|
||||
|
||||
# Create request payload
|
||||
guardrail_request = GenericGuardrailAPIRequest(
|
||||
|
|
@ -214,6 +321,8 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
litellm_trace_id=logging_obj.litellm_trace_id if logging_obj else None,
|
||||
texts=texts,
|
||||
request_data=user_metadata,
|
||||
request_headers=inbound_headers,
|
||||
litellm_version=litellm_version,
|
||||
images=images,
|
||||
tools=tools,
|
||||
structured_messages=structured_messages,
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from fastapi import APIRouter, Depends, HTTPException, Query
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import MAX_POLICY_ESTIMATE_IMPACT_ROWS
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
|
|
@ -85,7 +86,6 @@ def _filter_keys_by_tags(keys: list, tag_patterns: list) -> tuple:
|
|||
|
||||
Returns (named_aliases, unnamed_count).
|
||||
"""
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
|
||||
affected: list = []
|
||||
unnamed_count = 0
|
||||
|
|
@ -111,7 +111,6 @@ def _filter_teams_by_tags(teams: list, tag_patterns: list) -> tuple:
|
|||
|
||||
Returns (named_aliases, unnamed_count).
|
||||
"""
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
|
||||
affected: list = []
|
||||
unnamed_count = 0
|
||||
|
|
@ -141,7 +140,6 @@ async def _find_affected_by_team_patterns(
|
|||
|
||||
Returns (new_teams, new_keys, unnamed_keys_count).
|
||||
"""
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
|
||||
new_teams: list = []
|
||||
matched_team_ids: list = []
|
||||
|
|
@ -178,7 +176,6 @@ async def _find_affected_keys_by_alias(
|
|||
prisma_client: object, key_patterns: list, existing_keys: list
|
||||
) -> list:
|
||||
"""Find keys whose alias matches the given patterns."""
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
|
||||
affected: list = []
|
||||
|
||||
|
|
|
|||
|
|
@ -4862,8 +4862,10 @@ async def async_assistants_data_generator(
|
|||
if isinstance(e, HTTPException):
|
||||
raise e
|
||||
else:
|
||||
error_traceback = traceback.format_exc()
|
||||
error_msg = f"{str(e)}\n\n{error_traceback}"
|
||||
# Only include the error message, not the traceback.
|
||||
# The traceback is already logged above via verbose_proxy_logger.exception().
|
||||
# Including it in the SSE response leaks internal details to clients.
|
||||
error_msg = str(e)
|
||||
|
||||
proxy_exception = ProxyException(
|
||||
message=getattr(e, "message", error_msg),
|
||||
|
|
@ -5013,8 +5015,10 @@ async def async_data_generator(
|
|||
elif isinstance(e, StreamingCallbackError):
|
||||
error_msg = str(e)
|
||||
else:
|
||||
error_traceback = traceback.format_exc()
|
||||
error_msg = f"{str(e)}\n\n{error_traceback}"
|
||||
# Only include the error message, not the traceback.
|
||||
# The traceback is already logged above via verbose_proxy_logger.exception().
|
||||
# Including it in the SSE response leaks internal details to clients.
|
||||
error_msg = str(e)
|
||||
|
||||
proxy_exception = ProxyException(
|
||||
message=getattr(e, "message", error_msg),
|
||||
|
|
|
|||
|
|
@ -2283,7 +2283,7 @@ class Router:
|
|||
item = FlowItem(
|
||||
priority=priority, # 👈 SET PRIORITY FOR REQUEST
|
||||
request_id=_request_id, # 👈 SET REQUEST ID
|
||||
model_name="gpt-3.5-turbo", # 👈 SAME as 'Router'
|
||||
model_name=model, # 👈 SAME as 'Router'
|
||||
)
|
||||
### [fin] ###
|
||||
|
||||
|
|
@ -2325,6 +2325,10 @@ class Router:
|
|||
setattr(e, "priority", priority)
|
||||
raise e
|
||||
else:
|
||||
# Clean up the request from the scheduler queue also before raising the timeout exception
|
||||
await self.scheduler.remove_request(
|
||||
request_id=item.request_id, model_name=item.model_name
|
||||
)
|
||||
raise litellm.Timeout(
|
||||
message="Request timed out while polling queue",
|
||||
model=model,
|
||||
|
|
@ -2386,6 +2390,10 @@ class Router:
|
|||
setattr(e, "priority", priority)
|
||||
raise e
|
||||
else:
|
||||
# Clean up the request from the scheduler queue also before raising the timeout exception
|
||||
await self.scheduler.remove_request(
|
||||
request_id=item.request_id, model_name=item.model_name
|
||||
)
|
||||
raise litellm.Timeout(
|
||||
message="Request timed out while polling queue",
|
||||
model=model,
|
||||
|
|
@ -5039,7 +5047,7 @@ class Router:
|
|||
else:
|
||||
_healthy_deployments = []
|
||||
_timeout = self._time_to_sleep_before_retry(
|
||||
e=original_exception,
|
||||
e=e,
|
||||
remaining_retries=remaining_retries,
|
||||
num_retries=num_retries,
|
||||
healthy_deployments=_healthy_deployments,
|
||||
|
|
|
|||
|
|
@ -92,6 +92,17 @@ class Scheduler:
|
|||
|
||||
return True
|
||||
|
||||
async def remove_request(self, request_id: str, model_name: str) -> None:
|
||||
"""
|
||||
Remove a specific request from the priority queue for a model.
|
||||
Used when a request times out while waiting in the queue.
|
||||
"""
|
||||
queue = await self.get_queue(model_name=model_name)
|
||||
filtered_queue = [item for item in queue if item[1] != request_id]
|
||||
heapq.heapify(filtered_queue) # restore heap invariant after filtering
|
||||
await self.save_queue(queue=filtered_queue, model_name=model_name)
|
||||
print_verbose(f"Removed request_id: {request_id} from queue for model: {model_name}")
|
||||
|
||||
async def peek(self, id: str, model_name: str, health_deployments: list) -> bool:
|
||||
"""Return if the id is at the top of the queue. Don't pop the value from heap."""
|
||||
queue = await self.get_queue(model_name=model_name)
|
||||
|
|
|
|||
|
|
@ -60,6 +60,14 @@ class GenericGuardrailAPIRequest(BaseModel):
|
|||
tools: Optional[List[ChatCompletionToolParam]] = None
|
||||
texts: Optional[List[str]] = None
|
||||
request_data: GenericGuardrailAPIMetadata
|
||||
request_headers: Optional[Dict[str, str]] = Field(
|
||||
default=None,
|
||||
description="Sanitized inbound request headers from the original proxy request.",
|
||||
)
|
||||
litellm_version: Optional[str] = Field(
|
||||
default=None,
|
||||
description="LiteLLM library version running this proxy.",
|
||||
)
|
||||
additional_provider_specific_params: Optional[Dict[str, Any]] = None
|
||||
tool_calls: Optional[
|
||||
Union[List[ChatCompletionToolCallChunk], List[ChatCompletionMessageToolCall]]
|
||||
|
|
|
|||
|
|
@ -1406,7 +1406,7 @@ def client(original_function): # noqa: PLR0915
|
|||
# [OPTIONAL] CHECK MAX RETRIES / REQUEST
|
||||
if litellm.num_retries_per_request is not None:
|
||||
# check if previous_models passed in as ['litellm_params']['metadata]['previous_models']
|
||||
previous_models = kwargs.get("metadata", {}).get(
|
||||
previous_models = (kwargs.get("metadata") or {}).get(
|
||||
"previous_models", None
|
||||
)
|
||||
if previous_models is not None:
|
||||
|
|
@ -1483,7 +1483,7 @@ def client(original_function): # noqa: PLR0915
|
|||
# [OPTIONAL] CHECK MAX RETRIES / REQUEST
|
||||
if litellm.num_retries_per_request is not None:
|
||||
# check if previous_models passed in as ['litellm_params']['metadata]['previous_models']
|
||||
previous_models = kwargs.get("metadata", {}).get(
|
||||
previous_models = (kwargs.get("metadata") or {}).get(
|
||||
"previous_models", None
|
||||
)
|
||||
if previous_models is not None:
|
||||
|
|
@ -1678,8 +1678,8 @@ def client(original_function): # noqa: PLR0915
|
|||
"context_window_fallback_dict", {}
|
||||
)
|
||||
|
||||
_is_litellm_router_call = "model_group" in kwargs.get(
|
||||
"metadata", {}
|
||||
_is_litellm_router_call = "model_group" in (
|
||||
kwargs.get("metadata") or {}
|
||||
) # check if call from litellm.router/proxy
|
||||
if (
|
||||
num_retries and not _is_litellm_router_call
|
||||
|
|
@ -1724,8 +1724,8 @@ def client(original_function): # noqa: PLR0915
|
|||
None # set retries to None to prevent infinite loops
|
||||
)
|
||||
|
||||
_is_litellm_router_call = "model_group" in kwargs.get(
|
||||
"metadata", {}
|
||||
_is_litellm_router_call = "model_group" in (
|
||||
kwargs.get("metadata") or {}
|
||||
) # check if call from litellm.router/proxy
|
||||
if (
|
||||
num_retries and not _is_litellm_router_call
|
||||
|
|
@ -1974,8 +1974,8 @@ def client(original_function): # noqa: PLR0915
|
|||
"context_window_fallback_dict", {}
|
||||
)
|
||||
|
||||
_is_litellm_router_call = "model_group" in kwargs.get(
|
||||
"metadata", {}
|
||||
_is_litellm_router_call = "model_group" in (
|
||||
kwargs.get("metadata") or {}
|
||||
) # check if call from litellm.router/proxy
|
||||
|
||||
if (
|
||||
|
|
@ -2008,9 +2008,9 @@ def client(original_function): # noqa: PLR0915
|
|||
kwargs["model"] = context_window_fallback_dict[model]
|
||||
return await original_function(*args, **kwargs)
|
||||
elif call_type == CallTypes.aresponses.value:
|
||||
_is_litellm_router_call = "model_group" in (kwargs.get(
|
||||
"metadata", {}
|
||||
) or {}) # check if call from litellm.router/proxy
|
||||
_is_litellm_router_call = "model_group" in (
|
||||
kwargs.get("metadata") or {}
|
||||
) # check if call from litellm.router/proxy
|
||||
|
||||
if (
|
||||
num_retries and not _is_litellm_router_call
|
||||
|
|
@ -7337,7 +7337,7 @@ def _get_base_model_from_metadata(model_call_details=None):
|
|||
_base_model = litellm_params.get("base_model", None)
|
||||
if _base_model is not None:
|
||||
return _base_model
|
||||
metadata = litellm_params.get("metadata", {})
|
||||
metadata = litellm_params.get("metadata") or {}
|
||||
|
||||
_get_base_model_from_litellm_call_metadata = getattr(
|
||||
sys.modules[__name__], "_get_base_model_from_litellm_call_metadata"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,64 @@
|
|||
"""
|
||||
Tests for _map_reasoning_effort in AnthropicConfig.
|
||||
|
||||
Verifies that reasoning_effort=None returns None for all models,
|
||||
including Claude Opus 4.6.
|
||||
"""
|
||||
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
||||
|
||||
class TestMapReasoningEffort:
|
||||
def test_none_returns_none_for_opus_4_6(self):
|
||||
"""reasoning_effort=None should return None for Opus 4.6, not adaptive."""
|
||||
result = AnthropicConfig._map_reasoning_effort(
|
||||
reasoning_effort=None, model="claude-opus-4-6"
|
||||
)
|
||||
assert result is None
|
||||
|
||||
def test_none_returns_none_for_other_models(self):
|
||||
"""reasoning_effort=None should return None for non-Opus models."""
|
||||
result = AnthropicConfig._map_reasoning_effort(
|
||||
reasoning_effort=None, model="claude-3-7-sonnet-20250219"
|
||||
)
|
||||
assert result is None
|
||||
|
||||
def test_opus_4_6_returns_adaptive_for_low(self):
|
||||
result = AnthropicConfig._map_reasoning_effort(
|
||||
reasoning_effort="low", model="claude-opus-4-6"
|
||||
)
|
||||
assert result["type"] == "adaptive"
|
||||
|
||||
def test_opus_4_6_returns_adaptive_for_high(self):
|
||||
result = AnthropicConfig._map_reasoning_effort(
|
||||
reasoning_effort="high", model="claude-opus-4-6"
|
||||
)
|
||||
assert result["type"] == "adaptive"
|
||||
|
||||
def test_other_model_low_returns_enabled_with_budget(self):
|
||||
result = AnthropicConfig._map_reasoning_effort(
|
||||
reasoning_effort="low", model="claude-3-7-sonnet-20250219"
|
||||
)
|
||||
assert result["type"] == "enabled"
|
||||
assert "budget_tokens" in result
|
||||
|
||||
def test_other_model_high_returns_enabled_with_budget(self):
|
||||
result = AnthropicConfig._map_reasoning_effort(
|
||||
reasoning_effort="high", model="claude-3-7-sonnet-20250219"
|
||||
)
|
||||
assert result["type"] == "enabled"
|
||||
assert "budget_tokens" in result
|
||||
|
||||
def test_none_string_returns_none_for_opus_4_6(self):
|
||||
"""reasoning_effort='none' should return None for Opus 4.6."""
|
||||
result = AnthropicConfig._map_reasoning_effort(
|
||||
reasoning_effort="none", model="claude-opus-4-6"
|
||||
)
|
||||
assert result is None
|
||||
|
||||
def test_none_string_returns_none_for_other_models(self):
|
||||
"""reasoning_effort='none' should return None for non-Opus models."""
|
||||
result = AnthropicConfig._map_reasoning_effort(
|
||||
reasoning_effort="none", model="claude-3-7-sonnet-20250219"
|
||||
)
|
||||
assert result is None
|
||||
88
tests/litellm/test_router_retry_backoff_headers.py
Normal file
88
tests/litellm/test_router_retry_backoff_headers.py
Normal file
|
|
@ -0,0 +1,88 @@
|
|||
"""
|
||||
Tests for router retry backoff behavior.
|
||||
"""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_backoff_uses_current_exception_headers():
|
||||
"""
|
||||
Ensure retry backoff uses the current retry exception, not the initial one.
|
||||
"""
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "sk-test",
|
||||
},
|
||||
}
|
||||
],
|
||||
num_retries=2,
|
||||
)
|
||||
|
||||
first_error = litellm.RateLimitError(
|
||||
message="Rate limited on first attempt",
|
||||
model="gpt-3.5-turbo",
|
||||
llm_provider="openai",
|
||||
)
|
||||
first_error.litellm_response_headers = httpx.Headers({"retry-after": "1"})
|
||||
|
||||
second_error = litellm.RateLimitError(
|
||||
message="Rate limited on second attempt",
|
||||
model="gpt-3.5-turbo",
|
||||
llm_provider="openai",
|
||||
)
|
||||
second_error.litellm_response_headers = httpx.Headers({"retry-after": "15"})
|
||||
|
||||
third_error = litellm.RateLimitError(
|
||||
message="Rate limited on third attempt",
|
||||
model="gpt-3.5-turbo",
|
||||
llm_provider="openai",
|
||||
)
|
||||
third_error.litellm_response_headers = httpx.Headers({"retry-after": "30"})
|
||||
|
||||
raised_errors = [first_error, second_error, third_error]
|
||||
captured_backoff_errors = []
|
||||
|
||||
async def mock_make_call(*args, **kwargs):
|
||||
raise raised_errors.pop(0)
|
||||
|
||||
def mock_time_to_sleep_before_retry(*args, **kwargs):
|
||||
captured_backoff_errors.append(kwargs["e"])
|
||||
return 0.01
|
||||
|
||||
with patch.object(router, "make_call", side_effect=mock_make_call):
|
||||
with patch.object(
|
||||
router,
|
||||
"_async_get_healthy_deployments",
|
||||
return_value=(
|
||||
[{"model_info": {"id": "test-id"}}],
|
||||
[{"model_info": {"id": "test-id"}}],
|
||||
),
|
||||
):
|
||||
with patch.object(
|
||||
router,
|
||||
"_time_to_sleep_before_retry",
|
||||
side_effect=mock_time_to_sleep_before_retry,
|
||||
):
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
await router.acompletion(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
)
|
||||
|
||||
# Router computes backoff once after the initial failure, then once per failed retry.
|
||||
# With num_retries=2 and all attempts failing, that's 1 + 2 = 3 invocations.
|
||||
assert len(captured_backoff_errors) == router.num_retries + 1
|
||||
assert captured_backoff_errors[0] is first_error
|
||||
assert captured_backoff_errors[1] is second_error
|
||||
assert captured_backoff_errors[2] is third_error
|
||||
|
|
@ -117,3 +117,41 @@ async def test_scheduler_prioritized_requests(p0, p1, healthy_deployments):
|
|||
)
|
||||
== False
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scheduler_queue_cleanup_on_timeout():
|
||||
"""
|
||||
Test that a timed-out request is properly removed from the queue.
|
||||
This prevents memory leaks from accumulating timed-out requests.
|
||||
"""
|
||||
scheduler = Scheduler()
|
||||
|
||||
# Add multiple requests with different priorities
|
||||
item1 = FlowItem(priority=0, request_id="req-0", model_name="gpt-3.5-turbo")
|
||||
item2 = FlowItem(priority=1, request_id="req-1", model_name="gpt-3.5-turbo")
|
||||
item3 = FlowItem(priority=2, request_id="req-2", model_name="gpt-3.5-turbo")
|
||||
|
||||
await scheduler.add_request(item1)
|
||||
await scheduler.add_request(item2)
|
||||
await scheduler.add_request(item3)
|
||||
|
||||
# Verify initial queue size
|
||||
queue_before = await scheduler.get_queue(model_name="gpt-3.5-turbo")
|
||||
assert len(queue_before) == 3, f"Expected 3 items in queue, got {len(queue_before)}"
|
||||
|
||||
# Simulate timeout cleanup - remove a non-front request (item2)
|
||||
await scheduler.remove_request(request_id="req-1", model_name="gpt-3.5-turbo")
|
||||
|
||||
# Verify queue was cleaned up
|
||||
queue_after = await scheduler.get_queue(model_name="gpt-3.5-turbo")
|
||||
assert len(queue_after) == 2, f"Expected 2 items after cleanup, got {len(queue_after)}"
|
||||
|
||||
# Verify the correct request was removed
|
||||
remaining_ids = [item[1] for item in queue_after]
|
||||
assert "req-1" not in remaining_ids, "Expected req-1 to be removed"
|
||||
assert "req-0" in remaining_ids, "Expected req-0 to remain"
|
||||
assert "req-2" in remaining_ids, "Expected req-2 to remain"
|
||||
|
||||
# Verify remaining items are in correct priority order (0 should be first)
|
||||
assert queue_after[0][1] == "req-0", "Expected req-0 (priority 0) to be at front"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,415 @@
|
|||
"""
|
||||
Tests for standard_logging_payload_excluded_fields feature.
|
||||
|
||||
This feature allows users to exclude specific fields from StandardLoggingPayload
|
||||
before any callback receives it. This is useful for:
|
||||
- Reducing log sizes (excluding large fields like 'response' or 'messages')
|
||||
- Privacy compliance (excluding sensitive fields)
|
||||
- Cost management (less data stored/transmitted)
|
||||
|
||||
Example config:
|
||||
litellm_settings:
|
||||
success_callback: ["s3"]
|
||||
standard_logging_payload_excluded_fields: ["response", "messages"]
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from copy import deepcopy
|
||||
from typing import Dict, List, Optional
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../.."))
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
|
||||
def create_sample_standard_logging_payload() -> Dict:
|
||||
"""Create a sample StandardLoggingPayload for testing."""
|
||||
return {
|
||||
"id": "test-id-123",
|
||||
"trace_id": "trace-123",
|
||||
"call_type": "completion",
|
||||
"stream": False,
|
||||
"response_cost": 0.001,
|
||||
"cost_breakdown": None,
|
||||
"response_cost_failure_debug_info": None,
|
||||
"status": "success",
|
||||
"status_fields": {},
|
||||
"custom_llm_provider": "openai",
|
||||
"total_tokens": 100,
|
||||
"prompt_tokens": 50,
|
||||
"completion_tokens": 50,
|
||||
"startTime": 1234567890.0,
|
||||
"endTime": 1234567891.0,
|
||||
"completionStartTime": 1234567890.5,
|
||||
"response_time": 1.0,
|
||||
"model_map_information": {},
|
||||
"model": "gpt-4",
|
||||
"model_id": "model-123",
|
||||
"model_group": None,
|
||||
"api_base": "https://api.openai.com/v1",
|
||||
"metadata": {},
|
||||
"cache_hit": False,
|
||||
"cache_key": None,
|
||||
"saved_cache_cost": 0.0,
|
||||
"request_tags": [],
|
||||
"end_user": None,
|
||||
"requester_ip_address": None,
|
||||
"user_agent": None,
|
||||
"messages": [{"role": "user", "content": "Hello, this is sensitive data!"}],
|
||||
"response": {
|
||||
"choices": [
|
||||
{"message": {"content": "This is a sensitive response!"}}
|
||||
]
|
||||
},
|
||||
"error_str": None,
|
||||
"error_information": None,
|
||||
"model_parameters": {},
|
||||
"hidden_params": {},
|
||||
"guardrail_information": None,
|
||||
"standard_built_in_tools_params": None,
|
||||
}
|
||||
|
||||
|
||||
def create_model_call_details(
|
||||
standard_logging_payload: Optional[Dict] = None,
|
||||
) -> Dict:
|
||||
"""Create model_call_details dict with standard_logging_object."""
|
||||
if standard_logging_payload is None:
|
||||
standard_logging_payload = create_sample_standard_logging_payload()
|
||||
return {
|
||||
"standard_logging_object": standard_logging_payload,
|
||||
"other_key": "other_value",
|
||||
}
|
||||
|
||||
|
||||
class TestStandardLoggingPayloadExcludedFields:
|
||||
"""Test suite for standard_logging_payload_excluded_fields feature."""
|
||||
|
||||
def setup_method(self):
|
||||
"""Reset litellm settings before each test."""
|
||||
litellm.standard_logging_payload_excluded_fields = None
|
||||
|
||||
def teardown_method(self):
|
||||
"""Clean up after each test."""
|
||||
litellm.standard_logging_payload_excluded_fields = None
|
||||
|
||||
def test_no_excluded_fields_no_change(self):
|
||||
"""Test that payload is unchanged when no fields are excluded."""
|
||||
logger = CustomLogger()
|
||||
model_call_details = create_model_call_details()
|
||||
original_keys = set(model_call_details["standard_logging_object"].keys())
|
||||
|
||||
result = logger.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details
|
||||
)
|
||||
|
||||
result_keys = set(result["standard_logging_object"].keys())
|
||||
assert result_keys == original_keys
|
||||
|
||||
def test_exclude_single_field(self):
|
||||
"""Test excluding a single field (response)."""
|
||||
litellm.standard_logging_payload_excluded_fields = ["response"]
|
||||
|
||||
logger = CustomLogger()
|
||||
model_call_details = create_model_call_details()
|
||||
|
||||
result = logger.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details
|
||||
)
|
||||
|
||||
assert "response" not in result["standard_logging_object"]
|
||||
assert "messages" in result["standard_logging_object"]
|
||||
assert "model" in result["standard_logging_object"]
|
||||
|
||||
def test_exclude_multiple_fields(self):
|
||||
"""Test excluding multiple fields (response, messages)."""
|
||||
litellm.standard_logging_payload_excluded_fields = ["response", "messages"]
|
||||
|
||||
logger = CustomLogger()
|
||||
model_call_details = create_model_call_details()
|
||||
|
||||
result = logger.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details
|
||||
)
|
||||
|
||||
assert "response" not in result["standard_logging_object"]
|
||||
assert "messages" not in result["standard_logging_object"]
|
||||
assert "model" in result["standard_logging_object"]
|
||||
assert "model_parameters" in result["standard_logging_object"]
|
||||
|
||||
def test_exclude_metadata_field(self):
|
||||
"""Test excluding the metadata field."""
|
||||
litellm.standard_logging_payload_excluded_fields = ["metadata"]
|
||||
|
||||
logger = CustomLogger()
|
||||
payload = create_sample_standard_logging_payload()
|
||||
payload["metadata"] = {"sensitive_key": "sensitive_value"}
|
||||
model_call_details = create_model_call_details(payload)
|
||||
|
||||
result = logger.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details
|
||||
)
|
||||
|
||||
assert "metadata" not in result["standard_logging_object"]
|
||||
|
||||
def test_exclude_hidden_params(self):
|
||||
"""Test excluding hidden_params field."""
|
||||
litellm.standard_logging_payload_excluded_fields = ["hidden_params"]
|
||||
|
||||
logger = CustomLogger()
|
||||
payload = create_sample_standard_logging_payload()
|
||||
payload["hidden_params"] = {"api_key": "sk-secret-key"}
|
||||
model_call_details = create_model_call_details(payload)
|
||||
|
||||
result = logger.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details
|
||||
)
|
||||
|
||||
assert "hidden_params" not in result["standard_logging_object"]
|
||||
|
||||
def test_exclude_nonexistent_field_no_error(self):
|
||||
"""Test that excluding a non-existent field doesn't cause an error."""
|
||||
litellm.standard_logging_payload_excluded_fields = [
|
||||
"nonexistent_field",
|
||||
"response",
|
||||
]
|
||||
|
||||
logger = CustomLogger()
|
||||
model_call_details = create_model_call_details()
|
||||
|
||||
# Should not raise an exception
|
||||
result = logger.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details
|
||||
)
|
||||
|
||||
assert "response" not in result["standard_logging_object"]
|
||||
assert "messages" in result["standard_logging_object"]
|
||||
|
||||
def test_original_payload_not_modified(self):
|
||||
"""Test that the original model_call_details is not modified."""
|
||||
litellm.standard_logging_payload_excluded_fields = ["response", "messages"]
|
||||
|
||||
logger = CustomLogger()
|
||||
model_call_details = create_model_call_details()
|
||||
original_payload = deepcopy(model_call_details)
|
||||
|
||||
logger.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details
|
||||
)
|
||||
|
||||
# Original should still have the fields
|
||||
assert "response" in model_call_details["standard_logging_object"]
|
||||
assert "messages" in model_call_details["standard_logging_object"]
|
||||
assert model_call_details == original_payload
|
||||
|
||||
def test_combined_with_turn_off_message_logging(self):
|
||||
"""Test that excluded_fields works together with turn_off_message_logging."""
|
||||
litellm.standard_logging_payload_excluded_fields = ["metadata", "hidden_params"]
|
||||
|
||||
logger = CustomLogger(turn_off_message_logging=True)
|
||||
model_call_details = create_model_call_details()
|
||||
|
||||
result = logger.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details
|
||||
)
|
||||
|
||||
# excluded_fields should remove these
|
||||
assert "metadata" not in result["standard_logging_object"]
|
||||
assert "hidden_params" not in result["standard_logging_object"]
|
||||
|
||||
# turn_off_message_logging should redact these
|
||||
redacted_str = "redacted-by-litellm"
|
||||
assert (
|
||||
result["standard_logging_object"]["messages"][0]["content"] == redacted_str
|
||||
)
|
||||
assert (
|
||||
result["standard_logging_object"]["response"]["choices"][0]["message"][
|
||||
"content"
|
||||
]
|
||||
== redacted_str
|
||||
)
|
||||
|
||||
def test_excluded_fields_takes_precedence_over_redaction(self):
|
||||
"""Test that if a field is both excluded and would be redacted, it's excluded."""
|
||||
litellm.standard_logging_payload_excluded_fields = ["response"]
|
||||
|
||||
logger = CustomLogger(turn_off_message_logging=True)
|
||||
model_call_details = create_model_call_details()
|
||||
|
||||
result = logger.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details
|
||||
)
|
||||
|
||||
# response should be excluded (not redacted)
|
||||
assert "response" not in result["standard_logging_object"]
|
||||
|
||||
# messages should still be redacted
|
||||
redacted_str = "redacted-by-litellm"
|
||||
assert (
|
||||
result["standard_logging_object"]["messages"][0]["content"] == redacted_str
|
||||
)
|
||||
|
||||
def test_exclude_all_sensitive_fields(self):
|
||||
"""Test excluding all potentially sensitive fields."""
|
||||
litellm.standard_logging_payload_excluded_fields = [
|
||||
"messages",
|
||||
"response",
|
||||
"metadata",
|
||||
"hidden_params",
|
||||
"model_parameters",
|
||||
"error_str",
|
||||
"error_information",
|
||||
]
|
||||
|
||||
logger = CustomLogger()
|
||||
model_call_details = create_model_call_details()
|
||||
|
||||
result = logger.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details
|
||||
)
|
||||
|
||||
standard_obj = result["standard_logging_object"]
|
||||
|
||||
# All sensitive fields should be removed
|
||||
assert "messages" not in standard_obj
|
||||
assert "response" not in standard_obj
|
||||
assert "metadata" not in standard_obj
|
||||
assert "hidden_params" not in standard_obj
|
||||
assert "model_parameters" not in standard_obj
|
||||
assert "error_str" not in standard_obj
|
||||
assert "error_information" not in standard_obj
|
||||
|
||||
# Non-sensitive fields should remain
|
||||
assert "id" in standard_obj
|
||||
assert "model" in standard_obj
|
||||
assert "response_cost" in standard_obj
|
||||
assert "total_tokens" in standard_obj
|
||||
|
||||
def test_empty_excluded_fields_list(self):
|
||||
"""Test that an empty list doesn't affect the payload."""
|
||||
litellm.standard_logging_payload_excluded_fields = []
|
||||
|
||||
logger = CustomLogger()
|
||||
model_call_details = create_model_call_details()
|
||||
original_keys = set(model_call_details["standard_logging_object"].keys())
|
||||
|
||||
result = logger.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details
|
||||
)
|
||||
|
||||
result_keys = set(result["standard_logging_object"].keys())
|
||||
assert result_keys == original_keys
|
||||
|
||||
def test_none_standard_logging_object(self):
|
||||
"""Test handling when standard_logging_object is None."""
|
||||
litellm.standard_logging_payload_excluded_fields = ["response"]
|
||||
|
||||
logger = CustomLogger()
|
||||
model_call_details = {"other_key": "other_value"}
|
||||
|
||||
result = logger.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details
|
||||
)
|
||||
|
||||
# Should return unchanged when no standard_logging_object
|
||||
assert result == model_call_details
|
||||
|
||||
|
||||
class TestExcludedFieldsIntegration:
|
||||
"""Integration tests for excluded fields with actual callbacks."""
|
||||
|
||||
def setup_method(self):
|
||||
"""Reset litellm settings before each test."""
|
||||
litellm.standard_logging_payload_excluded_fields = None
|
||||
litellm.callbacks = []
|
||||
|
||||
def teardown_method(self):
|
||||
"""Clean up after each test."""
|
||||
litellm.standard_logging_payload_excluded_fields = None
|
||||
litellm.callbacks = []
|
||||
|
||||
def test_custom_callback_receives_filtered_payload(self):
|
||||
"""Test that a custom callback receives the filtered payload."""
|
||||
captured_payloads = []
|
||||
|
||||
class TestCallback(CustomLogger):
|
||||
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
captured_payloads.append(kwargs.get("standard_logging_object", {}))
|
||||
|
||||
litellm.standard_logging_payload_excluded_fields = ["response", "messages"]
|
||||
|
||||
callback = TestCallback()
|
||||
model_call_details = create_model_call_details()
|
||||
|
||||
# Simulate what litellm_logging.py does
|
||||
filtered_details = callback.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details
|
||||
)
|
||||
|
||||
callback.log_success_event(
|
||||
kwargs=filtered_details,
|
||||
response_obj=None,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
|
||||
assert len(captured_payloads) == 1
|
||||
assert "response" not in captured_payloads[0]
|
||||
assert "messages" not in captured_payloads[0]
|
||||
assert "model" in captured_payloads[0]
|
||||
|
||||
|
||||
class TestExcludedFieldsConfigLoading:
|
||||
"""Test that the config is properly loaded from litellm_settings."""
|
||||
|
||||
def setup_method(self):
|
||||
"""Reset litellm settings before each test."""
|
||||
litellm.standard_logging_payload_excluded_fields = None
|
||||
|
||||
def teardown_method(self):
|
||||
"""Clean up after each test."""
|
||||
litellm.standard_logging_payload_excluded_fields = None
|
||||
|
||||
def test_config_attribute_exists(self):
|
||||
"""Test that the config attribute exists on litellm module."""
|
||||
assert hasattr(litellm, "standard_logging_payload_excluded_fields")
|
||||
|
||||
def test_config_default_is_none(self):
|
||||
"""Test that the default value is None."""
|
||||
# Reset to ensure we're testing the default
|
||||
litellm.standard_logging_payload_excluded_fields = None
|
||||
assert litellm.standard_logging_payload_excluded_fields is None
|
||||
|
||||
def test_config_can_be_set_to_list(self):
|
||||
"""Test that the config can be set to a list."""
|
||||
litellm.standard_logging_payload_excluded_fields = ["response", "messages"]
|
||||
assert litellm.standard_logging_payload_excluded_fields == [
|
||||
"response",
|
||||
"messages",
|
||||
]
|
||||
|
||||
def test_config_setattr_simulates_proxy_loading(self):
|
||||
"""Test that setattr works as the proxy would use it."""
|
||||
# Simulating how proxy_server.py sets litellm_settings
|
||||
config_value = ["response", "messages", "metadata"]
|
||||
setattr(litellm, "standard_logging_payload_excluded_fields", config_value)
|
||||
|
||||
assert litellm.standard_logging_payload_excluded_fields == config_value
|
||||
|
||||
# Test it actually works in the logger
|
||||
logger = CustomLogger()
|
||||
model_call_details = create_model_call_details()
|
||||
|
||||
result = logger.redact_standard_logging_payload_from_model_call_details(
|
||||
model_call_details
|
||||
)
|
||||
|
||||
assert "response" not in result["standard_logging_object"]
|
||||
assert "messages" not in result["standard_logging_object"]
|
||||
assert "metadata" not in result["standard_logging_object"]
|
||||
|
|
@ -239,4 +239,149 @@ class TestAzureExceptionMapping:
|
|||
assert e.provider_specific_fields is not None
|
||||
assert e.provider_specific_fields["inner_error"]["code"] == "ResponsibleAIPolicyViolation"
|
||||
assert e.provider_specific_fields["inner_error"]["revised_prompt"] == "revised"
|
||||
assert e.provider_specific_fields["inner_error"]["content_filter_results"]["violence"]["filtered"] is True
|
||||
assert e.provider_specific_fields["inner_error"]["content_filter_results"]["violence"]["filtered"] is True
|
||||
|
||||
def test_azure_content_policy_violation_detected_via_inner_error_code(self):
|
||||
"""Regression test for #20811: Azure returns inner_error with
|
||||
ResponsibleAIPolicyViolation but the top-level error message is
|
||||
generic. Previously this fell through to the generic
|
||||
BadRequestError handler and all error details were lost."""
|
||||
|
||||
mock_exception = Exception("Bad request")
|
||||
# This body structure mirrors what Azure OpenAI Images API returns
|
||||
# for DALL-E 3 content policy violations (issue #20811).
|
||||
mock_exception.body = {
|
||||
"error": {
|
||||
"code": "content_policy_violation",
|
||||
"inner_error": {
|
||||
"code": "ResponsibleAIPolicyViolation",
|
||||
"content_filter_results": {
|
||||
"hate": {"filtered": False, "severity": "safe"},
|
||||
"profanity": {"detected": False, "filtered": False},
|
||||
"self_harm": {"filtered": False, "severity": "safe"},
|
||||
"sexual": {"filtered": False, "severity": "safe"},
|
||||
"violence": {"filtered": True, "severity": "low"},
|
||||
},
|
||||
"revised_prompt": (
|
||||
"A dark and intense illustration of a man "
|
||||
"in a dramatic action scene."
|
||||
),
|
||||
},
|
||||
"message": (
|
||||
"Your request was rejected as a result of our safety system."
|
||||
),
|
||||
"type": "invalid_request_error",
|
||||
}
|
||||
}
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
mock_exception.response = mock_response
|
||||
|
||||
with pytest.raises(ContentPolicyViolationError) as exc_info:
|
||||
exception_type(
|
||||
model="azure/dall-e-3",
|
||||
original_exception=mock_exception,
|
||||
custom_llm_provider="azure",
|
||||
)
|
||||
|
||||
e = exc_info.value
|
||||
# Must surface as ContentPolicyViolationError, not generic BadRequestError
|
||||
assert "safety system" in str(e)
|
||||
assert e.provider_specific_fields is not None
|
||||
inner = e.provider_specific_fields["inner_error"]
|
||||
assert inner["code"] == "ResponsibleAIPolicyViolation"
|
||||
assert inner["content_filter_results"]["violence"]["filtered"] is True
|
||||
assert inner["revised_prompt"] is not None
|
||||
|
||||
def test_azure_policy_violation_detected_via_inner_error_without_top_code(self):
|
||||
"""When the top-level code is NOT 'content_policy_violation' but
|
||||
inner_error.code IS 'ResponsibleAIPolicyViolation', the error
|
||||
should still be recognized as a content policy violation."""
|
||||
|
||||
mock_exception = Exception("Some error")
|
||||
mock_exception.body = {
|
||||
"error": {
|
||||
"code": "BadRequest",
|
||||
"inner_error": {
|
||||
"code": "ResponsibleAIPolicyViolation",
|
||||
"content_filter_results": {
|
||||
"violence": {"filtered": True, "severity": "medium"},
|
||||
},
|
||||
},
|
||||
"message": "The request was rejected.",
|
||||
"type": "invalid_request_error",
|
||||
}
|
||||
}
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
mock_exception.response = mock_response
|
||||
|
||||
with pytest.raises(ContentPolicyViolationError) as exc_info:
|
||||
exception_type(
|
||||
model="azure/dall-e-3",
|
||||
original_exception=mock_exception,
|
||||
custom_llm_provider="azure",
|
||||
)
|
||||
|
||||
e = exc_info.value
|
||||
assert e.provider_specific_fields is not None
|
||||
assert e.provider_specific_fields["inner_error"]["code"] == "ResponsibleAIPolicyViolation"
|
||||
|
||||
def test_azure_image_polling_error_preserves_body(self):
|
||||
"""Verify that AzureOpenAIError raised from the DALL-E polling path
|
||||
carries the structured body so exception_type() can inspect it."""
|
||||
from litellm.llms.azure.common_utils import AzureOpenAIError
|
||||
|
||||
error_payload = {
|
||||
"status": "failed",
|
||||
"error": {
|
||||
"code": "content_policy_violation",
|
||||
"message": "Your request was rejected.",
|
||||
"inner_error": {
|
||||
"code": "ResponsibleAIPolicyViolation",
|
||||
"content_filter_results": {
|
||||
"violence": {"filtered": True, "severity": "low"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
# Simulate what the fixed polling path now does
|
||||
_error_body = error_payload.get("error", error_payload)
|
||||
_error_msg = (
|
||||
_error_body.get("message", "Image generation failed")
|
||||
if isinstance(_error_body, dict)
|
||||
else json.dumps(error_payload)
|
||||
)
|
||||
exc = AzureOpenAIError(
|
||||
status_code=400,
|
||||
message=_error_msg,
|
||||
body=error_payload,
|
||||
)
|
||||
|
||||
assert exc.body is not None
|
||||
assert isinstance(exc.body, dict)
|
||||
assert exc.body["error"]["code"] == "content_policy_violation"
|
||||
assert "Your request was rejected" in exc.message
|
||||
|
||||
def test_azure_safety_system_message_detected_as_policy_violation(self):
|
||||
"""Azure's rejection message 'Your request was rejected as a result
|
||||
of our safety system' should be detected by string matching even
|
||||
when the structured body is unavailable."""
|
||||
|
||||
mock_exception = Exception(
|
||||
"Your request was rejected as a result of our safety system. "
|
||||
"The revised prompt may contain text that is not allowed."
|
||||
)
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
mock_exception.response = mock_response
|
||||
|
||||
with pytest.raises(ContentPolicyViolationError):
|
||||
exception_type(
|
||||
model="azure/dall-e-3",
|
||||
original_exception=mock_exception,
|
||||
custom_llm_provider="azure",
|
||||
)
|
||||
|
|
@ -225,6 +225,39 @@ async def test_aggregate_queue_updates_accuracy(spend_queue):
|
|||
assert aggregated["team_list_transactions"]["team1"] == 5.0
|
||||
|
||||
|
||||
def test_get_aggregated_spend_update_queue_item_does_not_mutate_original_updates(
|
||||
spend_queue,
|
||||
):
|
||||
original_update: SpendUpdateQueueItem = {
|
||||
"entity_type": Litellm_EntityType.USER,
|
||||
"entity_id": "user1",
|
||||
"response_cost": 10.0,
|
||||
}
|
||||
duplicate_key_update: SpendUpdateQueueItem = {
|
||||
"entity_type": Litellm_EntityType.USER,
|
||||
"entity_id": "user1",
|
||||
"response_cost": 20.0,
|
||||
}
|
||||
|
||||
aggregated_updates = spend_queue._get_aggregated_spend_update_queue_item(
|
||||
[original_update, duplicate_key_update]
|
||||
)
|
||||
user1_aggregated_update = next(
|
||||
(
|
||||
update
|
||||
for update in aggregated_updates
|
||||
if update.get("entity_type") == Litellm_EntityType.USER
|
||||
and update.get("entity_id") == "user1"
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
assert original_update["response_cost"] == 10.0
|
||||
assert user1_aggregated_update is not None
|
||||
assert user1_aggregated_update["response_cost"] == 30.0
|
||||
assert user1_aggregated_update is not original_update
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_queue_size_reduction_with_large_volume(monkeypatch, spend_queue):
|
||||
"""Test that queue size is actually reduced when dealing with many items"""
|
||||
|
|
|
|||
|
|
@ -14,10 +14,14 @@ import pytest
|
|||
import litellm
|
||||
from litellm import ModelResponse
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
from litellm._version import version as litellm_version
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
GenericGuardrailAPI,
|
||||
)
|
||||
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import (
|
||||
_HEADER_PRESENT_PLACEHOLDER,
|
||||
)
|
||||
from litellm.types.utils import Choices, Message
|
||||
|
||||
|
||||
|
|
@ -351,6 +355,58 @@ class TestMetadataExtraction:
|
|||
# Should be empty dict
|
||||
assert request_metadata == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inbound_headers_and_litellm_version_forwarded_and_sanitized(
|
||||
self, generic_guardrail, mock_request_data_input
|
||||
):
|
||||
"""
|
||||
Ensure inbound proxy request headers are forwarded in JSON payload with allowlist:
|
||||
allowed headers show their value; all other headers show presence only ([present]).
|
||||
"""
|
||||
# Add proxy_server_request headers as they exist in proxy request context
|
||||
request_data = dict(mock_request_data_input)
|
||||
request_data["proxy_server_request"] = {
|
||||
"headers": {
|
||||
"User-Agent": "OpenAI/Python 2.17.0",
|
||||
"Authorization": "Bearer should-not-forward",
|
||||
"Cookie": "session=should-not-forward",
|
||||
"X-Request-Id": "req_123",
|
||||
}
|
||||
}
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"action": "NONE",
|
||||
"texts": ["test"],
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
|
||||
with patch.object(
|
||||
generic_guardrail.async_handler, "post", return_value=mock_response
|
||||
) as mock_post:
|
||||
await generic_guardrail.apply_guardrail(
|
||||
inputs={"texts": ["test"]},
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
call_args = mock_post.call_args
|
||||
json_payload = call_args.kwargs["json"]
|
||||
|
||||
# New fields should exist
|
||||
assert json_payload["litellm_version"] == litellm_version
|
||||
assert "request_headers" in json_payload
|
||||
assert isinstance(json_payload["request_headers"], dict)
|
||||
req_headers = json_payload["request_headers"]
|
||||
|
||||
# Allowed: value forwarded
|
||||
assert req_headers.get("User-Agent") == "OpenAI/Python 2.17.0"
|
||||
|
||||
# Not on allowlist: key present, value is placeholder only
|
||||
assert req_headers.get("Authorization") == _HEADER_PRESENT_PLACEHOLDER
|
||||
assert req_headers.get("Cookie") == _HEADER_PRESENT_PLACEHOLDER
|
||||
assert req_headers.get("X-Request-Id") == _HEADER_PRESENT_PLACEHOLDER
|
||||
|
||||
|
||||
class TestGuardrailActions:
|
||||
"""Test different guardrail action responses"""
|
||||
|
|
|
|||
|
|
@ -3305,3 +3305,73 @@ class TestIsStreamingRequest:
|
|||
|
||||
def test_stream_true_overrides_non_streaming_call_type(self):
|
||||
assert _is_streaming_request(kwargs={"stream": True}, call_type=CallTypes.acompletion) is True
|
||||
|
||||
|
||||
class TestMetadataNoneHandling:
|
||||
"""
|
||||
Test that metadata=None in kwargs doesn't cause TypeError.
|
||||
|
||||
When metadata key exists with value None (e.g., from Azure OpenAI streaming),
|
||||
dict.get("metadata", {}) returns None (key exists, so default is ignored).
|
||||
The fix uses (kwargs.get("metadata") or {}) which handles both missing key
|
||||
and explicit None value.
|
||||
|
||||
Related: #20871
|
||||
"""
|
||||
|
||||
def test_metadata_none_get_previous_models(self):
|
||||
"""kwargs.get("metadata") or {} should return {} when metadata is None."""
|
||||
kwargs = {"metadata": None}
|
||||
previous_models = (kwargs.get("metadata") or {}).get(
|
||||
"previous_models", None
|
||||
)
|
||||
assert previous_models is None
|
||||
|
||||
def test_metadata_none_model_group_check(self):
|
||||
"""'model_group' in (kwargs.get("metadata") or {}) should not raise TypeError."""
|
||||
kwargs = {"metadata": None}
|
||||
_is_litellm_router_call = "model_group" in (
|
||||
kwargs.get("metadata") or {}
|
||||
)
|
||||
assert _is_litellm_router_call is False
|
||||
|
||||
def test_metadata_missing_key(self):
|
||||
"""Should work when metadata key is completely absent."""
|
||||
kwargs = {}
|
||||
previous_models = (kwargs.get("metadata") or {}).get(
|
||||
"previous_models", None
|
||||
)
|
||||
assert previous_models is None
|
||||
|
||||
def test_metadata_present_with_values(self):
|
||||
"""Should work when metadata has actual values."""
|
||||
kwargs = {"metadata": {"previous_models": ["model1"], "model_group": "test"}}
|
||||
previous_models = (kwargs.get("metadata") or {}).get(
|
||||
"previous_models", None
|
||||
)
|
||||
assert previous_models == ["model1"]
|
||||
_is_litellm_router_call = "model_group" in (
|
||||
kwargs.get("metadata") or {}
|
||||
)
|
||||
assert _is_litellm_router_call is True
|
||||
|
||||
def test_metadata_none_causes_error_with_old_pattern(self):
|
||||
"""Demonstrate the bug: dict.get('metadata', {}) returns None when key exists with None value."""
|
||||
kwargs = {"metadata": None}
|
||||
# Old pattern: kwargs.get("metadata", {}) returns None because key exists
|
||||
result = kwargs.get("metadata", {})
|
||||
assert result is None # This is the root cause of the bug
|
||||
|
||||
# Attempting to use .get() on None raises AttributeError or TypeError
|
||||
with pytest.raises((TypeError, AttributeError)):
|
||||
kwargs.get("metadata", {}).get("previous_models", None)
|
||||
|
||||
# Attempting 'in' on None raises TypeError
|
||||
with pytest.raises(TypeError):
|
||||
"model_group" in kwargs.get("metadata", {})
|
||||
|
||||
def test_litellm_params_metadata_none(self):
|
||||
"""litellm_params.get("metadata") or {} should handle None value."""
|
||||
litellm_params = {"metadata": None}
|
||||
metadata = litellm_params.get("metadata") or {}
|
||||
assert metadata == {}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue