mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-31 06:56:59 +00:00
refactor(types): replace Any with proven types in 21 files (#48921)
* refactor(types): replace Any with proven types in 35 files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(types): keep anthropic chat transformation json path untyped Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(types): accept null gemini vector store fields and load encryption TypedDicts on python 3.10 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(types): validate Python client user responses and type delete_user as the count it returns * fix(types): keep loose payloads working under the sweep's validators An off-schema Anthropic stream event is now skipped instead of ending the guardrail's text and stop_reason scan, and MAI image usage accepts the loose numbers it did before. New dict and list adapters hide their input from validation errors, RerankResponse indexing keeps its field types, and the delete_user test mocks the int /user/delete returns. * fix(types): make the changed DynamoAI fields ReadOnly and skip non-object A2A stream items * refactor(types): tighten round 2 parse sites and keep payloads out of errors Parse Anthropic SSE events with validate_json and ignore extra fields, validate admin team and key rows per row with hidden inputs, type rerank get and __getitem__ per key, and move the Phoenix prompt version response to a frozen BaseModel * fix(types): keep MAI image calls alive on unreadable usage counts An unreadable MAI usage count now counts as absent instead of failing a billed image call, and the provider's usage fields pass through as sent. The Anthropic stream scan warnings no longer log raw model output, and the off-schema skip has tests that fail on the old scan. * fix(types): count non-object MAI usage as absent and parse Anthropic SSE with json.loads --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
norkelbyte
GitHub
parent
19e538ab11
commit
e0778bc9dd
27 changed files with 496 additions and 117 deletions
|
|
@ -18,7 +18,7 @@ from contextlib import contextmanager
|
|||
from enum import Enum
|
||||
from typing import Any, Final, Literal, cast
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import current_service_target, service_target
|
||||
|
|
@ -43,6 +43,9 @@ from .s3_cache import S3Cache
|
|||
RESPONSE_CACHE_TARGET: Final = "llm_response"
|
||||
_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token"})
|
||||
_DECODED_LITERAL: Final = TypeAdapter(object)
|
||||
_REDIS_STARTUP_NODES_ADAPTER: Final = TypeAdapter(
|
||||
list[dict[str, object]], config=ConfigDict(strict=True, hide_input_in_errors=True)
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
|
|
@ -207,7 +210,7 @@ class Cache:
|
|||
if not redis_startup_nodes:
|
||||
_env_cluster_nodes: Final = litellm.get_secret("REDIS_CLUSTER_NODES")
|
||||
if _env_cluster_nodes is not None and isinstance(_env_cluster_nodes, str):
|
||||
redis_startup_nodes = json.loads(_env_cluster_nodes)
|
||||
redis_startup_nodes = _REDIS_STARTUP_NODES_ADAPTER.validate_python(json.loads(_env_cluster_nodes))
|
||||
|
||||
if redis_startup_nodes:
|
||||
# Only pass GCP parameters if they are provided
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ class ImageEditRequestUtils:
|
|||
image_edit_optional_params: ImageEditOptionalRequestParams,
|
||||
drop_params: bool | None = None,
|
||||
additional_drop_params: Collection[str] | None = None,
|
||||
) -> dict:
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Get optional parameters for the image edit API.
|
||||
|
||||
|
|
|
|||
|
|
@ -5,9 +5,17 @@ Arize Phoenix API client for fetching prompt versions from Arize Phoenix.
|
|||
import urllib.parse
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
|
||||
class _PromptVersionResponse(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", strict=True, frozen=True, hide_input_in_errors=True)
|
||||
|
||||
data: dict[str, object] | None = None
|
||||
|
||||
|
||||
def _sanitize_id(identifier: str) -> str:
|
||||
"""Reject path traversal characters and URL-encode the identifier."""
|
||||
if any(c in identifier for c in ("/", "\\", "#", "?")):
|
||||
|
|
@ -71,8 +79,7 @@ class ArizePhoenixClient:
|
|||
response: Final = self.http_handler.get(url, headers=self.headers)
|
||||
response.raise_for_status()
|
||||
|
||||
data: Final = response.json()
|
||||
return data.get("data")
|
||||
return _PromptVersionResponse.model_validate(response.json()).data
|
||||
|
||||
except Exception as e:
|
||||
# Check if it's an HTTP error
|
||||
|
|
|
|||
|
|
@ -72,7 +72,7 @@ if TYPE_CHECKING:
|
|||
)
|
||||
from litellm.proxy.proxy_server import UserAPIKeyAuth as _UserAPIKeyAuth
|
||||
|
||||
Span = _Span | Any
|
||||
Span = _Span
|
||||
Tracer = _Tracer
|
||||
Context = _Context
|
||||
SpanExporter = _SpanExporter
|
||||
|
|
@ -1440,7 +1440,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
span = None
|
||||
# Only set attributes if the span is still recording (not closed)
|
||||
# Note: parent_span is guaranteed to be not None here
|
||||
if hasattr(parent_span, "set_status"):
|
||||
if parent_span is not None and hasattr(parent_span, "set_status"):
|
||||
parent_span.set_status(Status(StatusCode.OK))
|
||||
self.set_attributes(parent_span, kwargs, response_obj)
|
||||
# Raw-request as direct child of parent_span
|
||||
|
|
@ -2233,7 +2233,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
# record error on parent span (keeps hierarchy shallow)
|
||||
# Only set attributes if the span is still recording (not closed)
|
||||
# Note: parent_otel_span is guaranteed to be not None here
|
||||
if parent_otel_span.is_recording():
|
||||
if parent_otel_span is not None and parent_otel_span.is_recording():
|
||||
self.set_attributes(parent_otel_span, kwargs, response_obj)
|
||||
parent_error_message: Final = self._record_exception_on_span(span=parent_otel_span, kwargs=kwargs)
|
||||
parent_otel_span.set_status(Status(StatusCode.ERROR, parent_error_message))
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ import json
|
|||
from collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -28,6 +29,8 @@ if TYPE_CHECKING:
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True))
|
||||
|
||||
|
||||
class _A2ATextPart(TypedDict, total=False):
|
||||
"""The subset of an A2A message part this handler reads text from."""
|
||||
|
|
@ -333,8 +336,8 @@ class A2AGuardrailHandler(BaseTranslation):
|
|||
obj = item
|
||||
elif isinstance(item, str):
|
||||
try:
|
||||
obj = json.loads(item.strip())
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
obj = _JSON_OBJECT.validate_python(json.loads(item.strip()))
|
||||
except (json.JSONDecodeError, TypeError, ValidationError):
|
||||
continue
|
||||
else:
|
||||
continue
|
||||
|
|
|
|||
|
|
@ -18,9 +18,9 @@ from copy import deepcopy
|
|||
from dataclasses import dataclass
|
||||
from itertools import chain, repeat
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast, overload, runtime_checkable
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config
|
||||
from typing_extensions import ReadOnly, TypedDict, assert_never
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -349,14 +349,40 @@ def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...
|
|||
)
|
||||
|
||||
|
||||
class _AnthropicSSEDelta(TypedDict, total=False):
|
||||
type: ReadOnly[str]
|
||||
@with_config(ConfigDict(extra="ignore", strict=True))
|
||||
class _SSETextDelta(TypedDict, total=False):
|
||||
type: ReadOnly[object]
|
||||
text: ReadOnly[str]
|
||||
stop_reason: ReadOnly[str | None]
|
||||
|
||||
|
||||
class _AnthropicSSEEvent(TypedDict, total=False):
|
||||
delta: ReadOnly[_AnthropicSSEDelta]
|
||||
@with_config(ConfigDict(extra="ignore", strict=True))
|
||||
class _SSETextDeltaEvent(TypedDict, total=False):
|
||||
delta: ReadOnly[_SSETextDelta]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="ignore", strict=True))
|
||||
class _SSEStopDelta(TypedDict, total=False):
|
||||
stop_reason: ReadOnly[object]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="ignore", strict=True))
|
||||
class _SSEStopDeltaEvent(TypedDict, total=False):
|
||||
delta: ReadOnly[_SSEStopDelta]
|
||||
|
||||
|
||||
_SSE_TEXT_DELTA_EVENT: Final = TypeAdapter(_SSETextDeltaEvent)
|
||||
_SSE_STOP_DELTA_EVENT: Final = TypeAdapter(_SSEStopDeltaEvent)
|
||||
_SSEEventT = TypeVar("_SSEEventT")
|
||||
|
||||
|
||||
def _parse_sse_event(adapter: TypeAdapter[_SSEEventT], event_type: str, data_line: str) -> _SSEEventT | None:
|
||||
try:
|
||||
return adapter.validate_python(json.loads(data_line))
|
||||
except json.JSONDecodeError:
|
||||
verbose_proxy_logger.warning("Skipping a %s SSE event that is not valid JSON", event_type)
|
||||
except ValidationError:
|
||||
verbose_proxy_logger.warning("Skipping a %s SSE event with an unexpected shape", event_type)
|
||||
return None
|
||||
|
||||
|
||||
class AnthropicMessagesHandler(BaseTranslation):
|
||||
|
|
@ -1806,13 +1832,10 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
# Only process content_block_delta events
|
||||
if event_type == "content_block_delta" and data_line:
|
||||
try:
|
||||
data: _AnthropicSSEEvent = json.loads(data_line)
|
||||
delta: _AnthropicSSEDelta = data.get("delta", {})
|
||||
if delta.get("type") == "text_delta":
|
||||
text += delta.get("text", "")
|
||||
except json.JSONDecodeError:
|
||||
verbose_proxy_logger.warning("Failed to parse JSON from SSE data: %s", data_line)
|
||||
data = _parse_sse_event(_SSE_TEXT_DELTA_EVENT, event_type, data_line)
|
||||
delta: _SSETextDelta = {} if data is None else data.get("delta", {})
|
||||
if delta.get("type") == "text_delta":
|
||||
text += delta.get("text", "")
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("Error extracting text from SSE: %s", e)
|
||||
|
|
@ -1869,14 +1892,11 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
# Check for message_delta event with stop_reason
|
||||
if event_type == "message_delta" and data_line:
|
||||
try:
|
||||
data: _AnthropicSSEEvent = json.loads(data_line)
|
||||
delta: _AnthropicSSEDelta = data.get("delta", {})
|
||||
stop_reason: str | None = delta.get("stop_reason")
|
||||
if stop_reason is not None:
|
||||
return True
|
||||
except json.JSONDecodeError:
|
||||
verbose_proxy_logger.warning("Failed to parse JSON from SSE data: %s", data_line)
|
||||
data = _parse_sse_event(_SSE_STOP_DELTA_EVENT, event_type, data_line)
|
||||
delta: _SSEStopDelta = {} if data is None else data.get("delta", {})
|
||||
stop_reason: object = delta.get("stop_reason")
|
||||
if stop_reason is not None:
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("Error checking streaming end in SSE: %s", e)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,12 @@
|
|||
from typing import TYPE_CHECKING, Any, Final
|
||||
import math
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Annotated, Final, TypeAlias
|
||||
|
||||
import httpx
|
||||
from pydantic import BeforeValidator, ConfigDict, TypeAdapter, ValidationError, with_config
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.exceptions import UnsupportedParamsError
|
||||
from litellm.litellm_core_utils.url_utils import join_api_base
|
||||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
|
|
@ -18,6 +23,67 @@ if TYPE_CHECKING:
|
|||
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
|
||||
|
||||
|
||||
_TOKEN_COUNT_ADAPTER: Final[TypeAdapter[int | float]] = TypeAdapter(int | float)
|
||||
|
||||
|
||||
def _token_count_or_none(value: object) -> int | float | None:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
count: Final = _TOKEN_COUNT_ADAPTER.validate_python(value)
|
||||
except ValidationError:
|
||||
verbose_logger.warning("MAI image usage has a token count that is not a number, counting it as absent")
|
||||
return None
|
||||
if not math.isfinite(count):
|
||||
verbose_logger.warning("MAI image usage has a token count that is not finite, counting it as absent")
|
||||
return None
|
||||
return count
|
||||
|
||||
|
||||
_TokenCount: TypeAlias = Annotated[int | float | None, BeforeValidator(_token_count_or_none)]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow"))
|
||||
class _MAIImageUsageDetails(TypedDict, total=False):
|
||||
text_tokens: ReadOnly[NotRequired[_TokenCount]]
|
||||
image_tokens: ReadOnly[NotRequired[_TokenCount]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow"))
|
||||
class _MAIImageUsage(TypedDict, total=False):
|
||||
input_tokens_details: ReadOnly[NotRequired[object]]
|
||||
num_input_text_tokens: ReadOnly[NotRequired[_TokenCount]]
|
||||
num_input_image_tokens: ReadOnly[NotRequired[_TokenCount]]
|
||||
input_tokens: ReadOnly[NotRequired[_TokenCount]]
|
||||
output_tokens: ReadOnly[NotRequired[_TokenCount]]
|
||||
num_output_tokens: ReadOnly[NotRequired[_TokenCount]]
|
||||
output_image_tokens: ReadOnly[NotRequired[_TokenCount]]
|
||||
total_tokens: ReadOnly[NotRequired[_TokenCount]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow"))
|
||||
class _MAIImageResponse(TypedDict, total=False):
|
||||
usage: ReadOnly[NotRequired[object]]
|
||||
|
||||
|
||||
_MAI_IMAGE_USAGE_DETAILS_ADAPTER: Final[TypeAdapter[_MAIImageUsageDetails]] = TypeAdapter(_MAIImageUsageDetails)
|
||||
_MAI_IMAGE_USAGE_ADAPTER: Final[TypeAdapter[_MAIImageUsage]] = TypeAdapter(_MAIImageUsage)
|
||||
_MAI_RAW_USAGE_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _raw_usage_or_none(usage: object) -> Mapping[str, object] | None:
|
||||
if usage is None:
|
||||
return None
|
||||
try:
|
||||
return _MAI_RAW_USAGE_ADAPTER.validate_python(usage)
|
||||
except ValidationError:
|
||||
verbose_logger.warning("MAI image usage is not an object, counting it as absent")
|
||||
return None
|
||||
|
||||
|
||||
_MAI_IMAGE_RESPONSE_ADAPTER: Final[TypeAdapter[_MAIImageResponse]] = TypeAdapter(_MAIImageResponse)
|
||||
|
||||
|
||||
class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
|
||||
"""Azure AI Foundry MAI image generation (e.g. MAI-Image-2.5)."""
|
||||
|
||||
|
|
@ -78,9 +144,10 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
|
|||
return "maiimage" in model_normalized
|
||||
|
||||
@staticmethod
|
||||
def normalize_mai_image_usage(usage: dict[str, Any] | None) -> dict[str, Any]:
|
||||
def normalize_mai_image_usage(usage: object) -> Mapping[str, object]:
|
||||
"""Map Azure MAI usage fields to OpenAI ImageUsage schema."""
|
||||
if usage is None:
|
||||
raw_usage: Final = _raw_usage_or_none(usage)
|
||||
if raw_usage is None:
|
||||
return {
|
||||
"input_tokens": 0,
|
||||
"input_tokens_details": {"image_tokens": 0, "text_tokens": 0},
|
||||
|
|
@ -88,10 +155,13 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
|
|||
"total_tokens": 0,
|
||||
}
|
||||
|
||||
normalized_usage: Final = dict(usage)
|
||||
input_tokens_details = normalized_usage.get("input_tokens_details")
|
||||
if not isinstance(input_tokens_details, dict):
|
||||
input_tokens_details = {}
|
||||
normalized_usage: Final[_MAIImageUsage] = _MAI_IMAGE_USAGE_ADAPTER.validate_python(raw_usage)
|
||||
input_tokens_details_value: Final = normalized_usage.get("input_tokens_details")
|
||||
input_tokens_details: Final[_MAIImageUsageDetails] = (
|
||||
_MAI_IMAGE_USAGE_DETAILS_ADAPTER.validate_python(input_tokens_details_value)
|
||||
if isinstance(input_tokens_details_value, dict)
|
||||
else {}
|
||||
)
|
||||
|
||||
text_tokens = normalized_usage.get("num_input_text_tokens")
|
||||
if text_tokens is None:
|
||||
|
|
@ -121,18 +191,16 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
|
|||
if total_tokens is None:
|
||||
total_tokens = input_tokens + output_tokens
|
||||
|
||||
normalized_usage.update(
|
||||
{
|
||||
"input_tokens": input_tokens,
|
||||
"input_tokens_details": {
|
||||
"image_tokens": image_tokens,
|
||||
"text_tokens": text_tokens,
|
||||
},
|
||||
"output_tokens": output_tokens,
|
||||
"total_tokens": total_tokens,
|
||||
}
|
||||
)
|
||||
return normalized_usage
|
||||
return {
|
||||
**raw_usage,
|
||||
"input_tokens": input_tokens,
|
||||
"input_tokens_details": {
|
||||
"image_tokens": image_tokens,
|
||||
"text_tokens": text_tokens,
|
||||
},
|
||||
"output_tokens": output_tokens,
|
||||
"total_tokens": total_tokens,
|
||||
}
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]:
|
||||
return ["n", "size"]
|
||||
|
|
@ -252,12 +320,18 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
|
|||
json_mode: bool | None = None,
|
||||
) -> ImageResponse:
|
||||
try:
|
||||
response: Final = raw_response.json()
|
||||
response_json: Final = _MAI_IMAGE_RESPONSE_ADAPTER.validate_python(raw_response.json())
|
||||
except Exception:
|
||||
raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code)
|
||||
|
||||
if "usage" in response:
|
||||
response["usage"] = self.normalize_mai_image_usage(response.get("usage"))
|
||||
response: Final = (
|
||||
{
|
||||
**response_json,
|
||||
"usage": self.normalize_mai_image_usage(response_json["usage"]),
|
||||
}
|
||||
if "usage" in response_json
|
||||
else dict(response_json)
|
||||
)
|
||||
|
||||
logging_obj.post_call(
|
||||
input=request_data.get("prompt", ""),
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from collections.abc import Callable, Mapping
|
|||
from typing import TYPE_CHECKING, Final, get_args, overload
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.constants import BEDROCK_EMBEDDING_PROVIDERS_LITERAL
|
||||
|
|
@ -48,6 +48,8 @@ from .amazon_titan_v2_transformation import AmazonTitanV2Config
|
|||
from .cohere_transformation import BedrockCohereEmbeddingConfig
|
||||
from .twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig, drop_params_enabled
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True))
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
|
|
@ -155,7 +157,7 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
except httpx.TimeoutException:
|
||||
raise BedrockError(status_code=408, message="Timeout error occurred.")
|
||||
|
||||
return response.json()
|
||||
return _JSON_OBJECT.validate_python(response.json())
|
||||
|
||||
async def _make_async_call(
|
||||
self,
|
||||
|
|
@ -189,7 +191,7 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
except httpx.TimeoutException:
|
||||
raise BedrockError(status_code=408, message="Timeout error occurred.")
|
||||
|
||||
return response.json()
|
||||
return _JSON_OBJECT.validate_python(response.json())
|
||||
|
||||
def _transform_response(
|
||||
self,
|
||||
|
|
@ -689,7 +691,7 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
|
||||
# Parse response
|
||||
if response.status_code == 200:
|
||||
return response.json()
|
||||
return _JSON_OBJECT.validate_python(response.json())
|
||||
else:
|
||||
raise Exception(f"Failed to get async invoke status: {response.status_code} - {response.text}")
|
||||
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from collections.abc import Mapping
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from litellm.litellm_core_utils.audio_utils.utils import process_audio_file
|
||||
from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params
|
||||
|
|
@ -21,6 +22,8 @@ from litellm.utils import convert_to_model_response_object
|
|||
|
||||
from ..common_utils import EdenAIException, authorized_headers, endpoint_url, reported_cost
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True))
|
||||
|
||||
|
||||
def _form_fields(model: str, optional_params: Mapping[str, object]) -> dict[str, object]: # mutable-ok: httpx form data
|
||||
"""LiteLLM parks non-OpenAI params, `model` included, under `extra_body` for the OpenAI SDK; a
|
||||
|
|
@ -75,7 +78,7 @@ class EdenAIAudioTranscriptionConfig(OpenAIWhisperAudioTranscriptionConfig):
|
|||
def transform_audio_transcription_response(self, raw_response: httpx.Response) -> TranscriptionResponse:
|
||||
if "application/json" not in raw_response.headers.get("content-type", ""):
|
||||
return TranscriptionResponse(text=raw_response.text)
|
||||
body: Final = raw_response.json()
|
||||
body: Final = _JSON_OBJECT.validate_python(raw_response.json())
|
||||
response: Final[TranscriptionResponse] = convert_to_model_response_object(
|
||||
response_object=body, model_response_object=TranscriptionResponse(), response_type="audio_transcription"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
|
|
@ -9,6 +10,8 @@ from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
|
|||
from litellm.types.utils import ImageResponse
|
||||
from litellm.utils import convert_to_model_response_object
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True))
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
|
|
@ -66,7 +69,7 @@ class GPTImageGenerationConfig(BaseImageGenerationConfig):
|
|||
api_key: str | None = None,
|
||||
json_mode: bool | None = None,
|
||||
) -> ImageResponse:
|
||||
response: Final = raw_response.json()
|
||||
response: Final = _JSON_OBJECT.validate_python(raw_response.json())
|
||||
|
||||
stringified_response: Final = response
|
||||
## LOGGING
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import os
|
|||
from typing import Final, cast
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
|
|
@ -11,6 +12,8 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
)
|
||||
from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
|
||||
|
||||
_TOKEN_INFO_RESPONSE: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True))
|
||||
|
||||
|
||||
class OAuth2TokenValidationError(ValueError):
|
||||
"""The token-info or introspection endpoint did not vouch for the token: it rejected it, reported it
|
||||
|
|
@ -184,7 +187,7 @@ class Oauth2Handler:
|
|||
response.raise_for_status()
|
||||
|
||||
# If we get here, the request was successful
|
||||
data = response.json()
|
||||
data = _TOKEN_INFO_RESPONSE.validate_python(response.json())
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Oauth2 token validation for token=%s, response from endpoint=%s",
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# stdlib imports
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
|
|
@ -61,6 +61,10 @@ def _get_model_info_obj_from_yaml(model: dict[str, Any]) -> ModelYamlInfo:
|
|||
)
|
||||
|
||||
|
||||
def _str_field(model: Mapping[str, object], key: str, default: str) -> str:
|
||||
return str(model.get(key, default))
|
||||
|
||||
|
||||
def format_iso_datetime_str(iso_datetime_str: str | None) -> str:
|
||||
"""Format an ISO format datetime string to human-readable date with minute resolution."""
|
||||
if not iso_datetime_str:
|
||||
|
|
@ -142,10 +146,10 @@ def list_models(ctx: click.Context, output_format: Literal["table", "json"]) ->
|
|||
created = int(created)
|
||||
|
||||
table.add_row(
|
||||
str(model.get("id", "")),
|
||||
str(model.get("object", "model")),
|
||||
_str_field(model, "id", ""),
|
||||
_str_field(model, "object", "model"),
|
||||
(format_timestamp(created) if isinstance(created, int) else format_iso_datetime_str(created)),
|
||||
str(model.get("owned_by", "")),
|
||||
_str_field(model, "owned_by", ""),
|
||||
)
|
||||
|
||||
rich.print(table)
|
||||
|
|
|
|||
|
|
@ -2,9 +2,13 @@ from collections.abc import Iterable, Mapping
|
|||
from typing import Any, Final
|
||||
|
||||
import requests
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from .exceptions import NotFoundError, UnauthorizedError
|
||||
|
||||
_JSON_VALUE: Final = TypeAdapter(object, config=ConfigDict(strict=True, hide_input_in_errors=True))
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True))
|
||||
|
||||
|
||||
class UsersManagementClient:
|
||||
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
|
||||
|
|
@ -39,7 +43,7 @@ class UsersManagementClient:
|
|||
if response.status_code == 404:
|
||||
raise NotFoundError(response.text)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
return _JSON_OBJECT.validate_python(response.json())
|
||||
|
||||
def get_user_v2(self, user_id: str | None = None) -> dict[str, object]:
|
||||
"""Get user info v2 - lightweight, returns only user object (GET /v2/user/info)"""
|
||||
|
|
@ -51,7 +55,7 @@ class UsersManagementClient:
|
|||
if response.status_code == 404:
|
||||
raise NotFoundError(response.text)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
return _JSON_OBJECT.validate_python(response.json())
|
||||
|
||||
def create_user(self, user_data: Mapping[str, object]) -> dict[str, object]:
|
||||
"""Create a new user (POST /user/new)"""
|
||||
|
|
@ -60,9 +64,9 @@ class UsersManagementClient:
|
|||
if response.status_code == 401:
|
||||
raise UnauthorizedError(response.text)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
return _JSON_OBJECT.validate_python(response.json())
|
||||
|
||||
def delete_user(self, user_ids: list[str]) -> dict[str, object]:
|
||||
def delete_user(self, user_ids: list[str]) -> object:
|
||||
"""Delete users (POST /user/delete)"""
|
||||
url: Final = f"{self.base_url}/user/delete"
|
||||
response: Final = requests.post(
|
||||
|
|
@ -71,4 +75,4 @@ class UsersManagementClient:
|
|||
if response.status_code == 401:
|
||||
raise UnauthorizedError(response.text)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
return _JSON_VALUE.validate_python(response.json())
|
||||
|
|
|
|||
|
|
@ -16,6 +16,10 @@ unknown impl value raises, which the proxy converts to the fail-closed 503.
|
|||
"""
|
||||
|
||||
import os
|
||||
from typing import Final
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter, with_config
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.proxy.management_helpers.team_metadata_validation import (
|
||||
|
|
@ -23,12 +27,20 @@ from litellm.proxy.management_helpers.team_metadata_validation import (
|
|||
TeamMetadataValidationResult,
|
||||
)
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
from typing import Final
|
||||
|
||||
ALLOWED_COST_CENTERS: Final = frozenset({"CC-1001", "CC-1002"})
|
||||
CLOSED_PORT_URL: Final = "http://127.0.0.1:9/validate"
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True))
|
||||
class _CostCenterReply(TypedDict):
|
||||
ok: NotRequired[ReadOnly[object]]
|
||||
reason: NotRequired[ReadOnly[str | None]]
|
||||
|
||||
|
||||
_COST_CENTER_REPLY: Final = TypeAdapter(_CostCenterReply)
|
||||
|
||||
|
||||
async def _validate_allowlist(payload: TeamMetadataValidationPayload) -> TeamMetadataValidationResult:
|
||||
cost_center: Final = payload.metadata.get("cost_center")
|
||||
if cost_center is None:
|
||||
|
|
@ -54,7 +66,7 @@ async def _validate_via_http(payload: TeamMetadataValidationPayload, service_url
|
|||
},
|
||||
timeout=2.0,
|
||||
)
|
||||
body: Final = response.json()
|
||||
body: Final = _COST_CENTER_REPLY.validate_python(response.json())
|
||||
if body.get("ok") is True:
|
||||
return TeamMetadataValidationResult(valid=True)
|
||||
return TeamMetadataValidationResult(
|
||||
|
|
|
|||
|
|
@ -300,7 +300,7 @@ class VigilGuardGuardrail(CustomGuardrail):
|
|||
def _build_output(
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
final_texts: list[str],
|
||||
final_tool_calls: Any,
|
||||
final_tool_calls: _ToolCalls | None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
# When nothing was changed, return the input shape verbatim so the guardrail
|
||||
# logs "allow" rather than "mask". When a text or a tool-call argument was
|
||||
|
|
@ -315,7 +315,7 @@ class VigilGuardGuardrail(CustomGuardrail):
|
|||
guardrailed["images"] = inputs["images"]
|
||||
if "tools" in inputs:
|
||||
guardrailed["tools"] = inputs["tools"]
|
||||
if tool_calls_changed:
|
||||
if tool_calls_changed and final_tool_calls is not None:
|
||||
guardrailed["tool_calls"] = final_tool_calls
|
||||
return guardrailed
|
||||
|
||||
|
|
|
|||
|
|
@ -18,6 +18,8 @@ import logging
|
|||
import os
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
from litellm import DualCache
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -36,6 +38,14 @@ else:
|
|||
InternalUsageCache = Any
|
||||
|
||||
|
||||
_CACHE_AMOUNT_ADAPTER: Final[TypeAdapter[str | bytes | int | float]] = TypeAdapter(
|
||||
str | bytes | int | float, config=ConfigDict(strict=True, hide_input_in_errors=True)
|
||||
)
|
||||
_OPTIONAL_CACHE_AMOUNT_ADAPTER: Final[TypeAdapter[str | bytes | int | float | None]] = TypeAdapter(
|
||||
str | bytes | int | float | None, config=ConfigDict(strict=True, hide_input_in_errors=True)
|
||||
)
|
||||
|
||||
|
||||
# Redis Lua script for atomic float increment with TTL.
|
||||
# INCRBYFLOAT returns the new value as a string.
|
||||
# Only sets EXPIRE on first call (when prior value was nil).
|
||||
|
|
@ -216,9 +226,11 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger):
|
|||
"""Read current accumulated spend for a session."""
|
||||
if self.internal_usage_cache.dual_cache.redis_cache is not None:
|
||||
try:
|
||||
result = await self.internal_usage_cache.dual_cache.redis_cache.async_get_cache(key=cache_key)
|
||||
if result is not None:
|
||||
return float(result)
|
||||
redis_result: Final[str | bytes | int | float | None] = _OPTIONAL_CACHE_AMOUNT_ADAPTER.validate_python(
|
||||
await self.internal_usage_cache.dual_cache.redis_cache.async_get_cache(key=cache_key)
|
||||
)
|
||||
if redis_result is not None:
|
||||
return float(redis_result)
|
||||
return 0.0
|
||||
except Exception as e:
|
||||
log_redis_failure(
|
||||
|
|
@ -228,24 +240,28 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger):
|
|||
e,
|
||||
)
|
||||
|
||||
result = await self.internal_usage_cache.async_get_cache(
|
||||
key=cache_key,
|
||||
litellm_parent_otel_span=None,
|
||||
local_only=True,
|
||||
cache_result: Final[str | bytes | int | float | None] = _OPTIONAL_CACHE_AMOUNT_ADAPTER.validate_python(
|
||||
await self.internal_usage_cache.async_get_cache(
|
||||
key=cache_key,
|
||||
litellm_parent_otel_span=None,
|
||||
local_only=True,
|
||||
)
|
||||
)
|
||||
if result is not None:
|
||||
return float(result)
|
||||
if cache_result is not None:
|
||||
return float(cache_result)
|
||||
return 0.0
|
||||
|
||||
async def _increment_spend(self, cache_key: str, amount: float) -> float:
|
||||
"""Atomically increment the session spend and return the new value."""
|
||||
if self.increment_script is not None:
|
||||
try:
|
||||
result: Final = await self.increment_script(
|
||||
keys=[cache_key],
|
||||
args=[str(amount), self.ttl],
|
||||
script_result: Final[str | bytes | int | float] = _CACHE_AMOUNT_ADAPTER.validate_python(
|
||||
await self.increment_script(
|
||||
keys=[cache_key],
|
||||
args=[str(amount), self.ttl],
|
||||
)
|
||||
)
|
||||
return float(result)
|
||||
return float(script_result)
|
||||
except Exception as e:
|
||||
log_redis_failure(
|
||||
verbose_proxy_logger,
|
||||
|
|
@ -257,12 +273,15 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger):
|
|||
return await self._in_memory_increment_spend(cache_key, amount)
|
||||
|
||||
async def _in_memory_increment_spend(self, cache_key: str, amount: float) -> float:
|
||||
current: Final = await self.internal_usage_cache.async_get_cache(
|
||||
key=cache_key,
|
||||
litellm_parent_otel_span=None,
|
||||
local_only=True,
|
||||
current: Final[str | bytes | int | float | None] = _OPTIONAL_CACHE_AMOUNT_ADAPTER.validate_python(
|
||||
await self.internal_usage_cache.async_get_cache(
|
||||
key=cache_key,
|
||||
litellm_parent_otel_span=None,
|
||||
local_only=True,
|
||||
)
|
||||
)
|
||||
new_value: Final = (float(current) if current is not None else 0.0) + amount
|
||||
current_spend: Final = float(current) if current is not None else 0.0
|
||||
new_value: Final = current_spend + amount
|
||||
await self.internal_usage_cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=new_value,
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ from typing import Any, Final, Literal, Protocol, cast, overload
|
|||
|
||||
import fastapi
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
|
|
@ -124,6 +124,20 @@ _USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig])
|
|||
_USER_CACHE_INVALIDATION_BATCH_SIZE: Final = 50
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True))
|
||||
class _AdminTeamAndKeyRows(TypedDict):
|
||||
keys: ReadOnly[Sequence[object] | None]
|
||||
teams: ReadOnly[Sequence[object] | None]
|
||||
|
||||
|
||||
_ADMIN_TEAM_AND_KEY_ROWS_ADAPTER: Final = TypeAdapter(
|
||||
list[_AdminTeamAndKeyRows], config=ConfigDict(hide_input_in_errors=True)
|
||||
)
|
||||
_ADMIN_KEY_ROW_ADAPTER: Final = TypeAdapter(
|
||||
dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)
|
||||
)
|
||||
|
||||
|
||||
def _user_table(
|
||||
prisma_client: "PrismaClient | None",
|
||||
) -> "TableActions[prisma_models.LiteLLM_UserTable]":
|
||||
|
|
@ -1170,8 +1184,8 @@ async def user_info_v2(
|
|||
|
||||
async def _fetch_admin_teams_and_keys_rows(
|
||||
prisma_client: "PrismaClient", sql_query: str
|
||||
) -> Sequence[Mapping[str, Sequence[Mapping[str, object]] | None]]:
|
||||
return await prisma_client.db.query_raw(sql_query)
|
||||
) -> Sequence[_AdminTeamAndKeyRows]:
|
||||
return _ADMIN_TEAM_AND_KEY_ROWS_ADAPTER.validate_python(await prisma_client.db.query_raw(sql_query))
|
||||
|
||||
|
||||
async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth):
|
||||
|
|
@ -1201,17 +1215,17 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth):
|
|||
|
||||
verbose_proxy_logger.debug("results_keys: %s", results)
|
||||
|
||||
_keys_in_db: Final[Sequence[Mapping[str, object]]] = results[0]["keys"] or []
|
||||
_keys_in_db: Final[Sequence[object]] = results[0]["keys"] or []
|
||||
# cast all keys to LiteLLM_VerificationToken
|
||||
keys_in_db: Final = []
|
||||
for key in _keys_in_db:
|
||||
key_payload = dict[str, object](key)
|
||||
key_payload = _ADMIN_KEY_ROW_ADAPTER.validate_python(key)
|
||||
if key_payload.get("models") is None:
|
||||
key_payload["models"] = []
|
||||
keys_in_db.append(LiteLLM_VerificationToken.model_validate(key_payload))
|
||||
|
||||
# cast all teams to LiteLLM_TeamTable
|
||||
_teams_rows: Final[Sequence[Mapping[str, object]]] = results[0]["teams"] or []
|
||||
_teams_rows: Final[Sequence[object]] = results[0]["teams"] or []
|
||||
_teams_in_db: Final = sorted(
|
||||
(LiteLLM_TeamTable.model_validate(team) for team in _teams_rows),
|
||||
key=lambda x: getattr(x, "team_alias", "") or "",
|
||||
|
|
|
|||
|
|
@ -3,11 +3,13 @@
|
|||
import json
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response
|
||||
from fastapi import status as http_status
|
||||
from pydantic import ConfigDict, TypeAdapter, with_config
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
|
|
@ -44,6 +46,20 @@ _DEFAULT_TRANSCRIPTION_MODEL: Final = "gpt-realtime-whisper"
|
|||
_ALLOWED_SESSION_TYPES: Final = ("realtime", "transcription")
|
||||
|
||||
|
||||
@with_config(ConfigDict(strict=True, extra="allow"))
|
||||
class _RealtimeTokenPayload(TypedDict, total=False):
|
||||
v: NotRequired[ReadOnly[object]]
|
||||
ephemeral_key: NotRequired[ReadOnly[str]]
|
||||
model_id: NotRequired[ReadOnly[str]]
|
||||
user_id: NotRequired[ReadOnly[str | None]]
|
||||
team_id: NotRequired[ReadOnly[str | None]]
|
||||
expires_at: NotRequired[ReadOnly[object]]
|
||||
session_type: NotRequired[ReadOnly[str | None]]
|
||||
|
||||
|
||||
_REALTIME_TOKEN_PAYLOAD_ADAPTER: Final = TypeAdapter(_RealtimeTokenPayload)
|
||||
|
||||
|
||||
def _coerce_realtime_session_type(session_type: str | None) -> str:
|
||||
if session_type in _ALLOWED_SESSION_TYPES:
|
||||
return session_type
|
||||
|
|
@ -201,17 +217,14 @@ def _header_value(headers: Mapping[str, str], name: str, default: str) -> str:
|
|||
|
||||
def _decode_realtime_token_payload(
|
||||
decrypted_value: str,
|
||||
) -> dict[str, Any] | None:
|
||||
) -> _RealtimeTokenPayload | None:
|
||||
"""
|
||||
Decode realtime token payload; returns None for legacy/raw ephemeral tokens.
|
||||
"""
|
||||
try:
|
||||
decoded: Final = json.loads(decrypted_value)
|
||||
decoded: Final = _REALTIME_TOKEN_PAYLOAD_ADAPTER.validate_python(json.loads(decrypted_value))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
if not isinstance(decoded, dict):
|
||||
return None
|
||||
if decoded.get("v") != _REALTIME_TOKEN_VERSION:
|
||||
return None
|
||||
if not isinstance(decoded.get("ephemeral_key"), str):
|
||||
|
|
|
|||
|
|
@ -2,9 +2,10 @@
|
|||
|
||||
import enum
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Literal, TypedDict
|
||||
from typing import Literal, TypedDict
|
||||
|
||||
from pydantic import Field
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
from .base import GuardrailConfigModel
|
||||
|
||||
|
|
@ -42,7 +43,7 @@ class PolicyApplicableTo(str, enum.Enum):
|
|||
class DynamoAIRequest(TypedDict, total=False):
|
||||
"""Request structure for DynamoAI /moderation/analyze endpoint"""
|
||||
|
||||
messages: list[dict[str, Any]]
|
||||
messages: ReadOnly[list[dict[str, object]]]
|
||||
textType: DynamoTextType | None
|
||||
policyIds: list[str]
|
||||
modelId: str | None
|
||||
|
|
@ -58,8 +59,8 @@ class PolicyInfo(TypedDict, total=False):
|
|||
description: str
|
||||
method: PolicyMethod
|
||||
action: Literal["BLOCK", "WARN", "REDACT", "SANITIZE", "NONE"]
|
||||
methodParams: dict[str, Any]
|
||||
decisionParams: dict[str, Any]
|
||||
methodParams: ReadOnly[dict[str, object]]
|
||||
decisionParams: ReadOnly[dict[str, object]]
|
||||
applicableTo: PolicyApplicableTo
|
||||
created_at: str
|
||||
creatorId: str
|
||||
|
|
@ -76,7 +77,7 @@ class AppliedPolicyDto(TypedDict, total=False):
|
|||
"""Applied policy details from DynamoAI response"""
|
||||
|
||||
policy: PolicyInfo
|
||||
outputs: dict[str, Any] | None
|
||||
outputs: ReadOnly[dict[str, object] | None]
|
||||
action: str | None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -89,7 +89,7 @@ class NomaV2StreamingParams(BaseModel):
|
|||
cls, value: object, handler: ValidatorFunctionWrapHandler, info: ValidationInfo
|
||||
) -> object:
|
||||
try:
|
||||
validated: Final = handler(value)
|
||||
validated: Final[object] = handler(value)
|
||||
except ValidationError:
|
||||
verbose_proxy_logger.warning(
|
||||
"Noma v2 guardrail %s: ignoring %s=%.100r, which is not a valid value; the default applies",
|
||||
|
|
|
|||
|
|
@ -4,7 +4,8 @@ https://docs.cohere.com/reference/rerank
|
|||
|
||||
"""
|
||||
|
||||
from typing import Literal
|
||||
from collections.abc import Sequence
|
||||
from typing import Literal, overload
|
||||
|
||||
from pydantic import ConfigDict, PrivateAttr
|
||||
from typing_extensions import ReadOnly, Required, TypedDict
|
||||
|
|
@ -95,11 +96,35 @@ class RerankResponse(LiteLLMBaseModel):
|
|||
def hidden_params(self, hidden_params: dict[str, object]) -> None: # mutable-ok: API requires mutation
|
||||
self._hidden_params = hidden_params
|
||||
|
||||
def __getitem__(self, key):
|
||||
@overload
|
||||
def __getitem__(self, key: Literal["id"]) -> str | None: ...
|
||||
|
||||
@overload
|
||||
def __getitem__(self, key: Literal["results"]) -> list[RerankResponseResult] | None: ...
|
||||
|
||||
@overload
|
||||
def __getitem__(self, key: Literal["meta"]) -> RerankResponseMeta | None: ...
|
||||
|
||||
@overload
|
||||
def __getitem__(self, key: str) -> str | Sequence[RerankResponseResult] | RerankResponseMeta | None: ...
|
||||
|
||||
def __getitem__(self, key: str) -> str | Sequence[RerankResponseResult] | RerankResponseMeta | None:
|
||||
return self.__dict__[key]
|
||||
|
||||
def get(self, key, default=None):
|
||||
@overload
|
||||
def get(self, key: Literal["id"], default: None = None) -> str | None: ...
|
||||
|
||||
@overload
|
||||
def get(self, key: Literal["results"], default: None = None) -> list[RerankResponseResult] | None: ...
|
||||
|
||||
@overload
|
||||
def get(self, key: Literal["meta"], default: None = None) -> RerankResponseMeta | None: ...
|
||||
|
||||
@overload
|
||||
def get(self, key: str, default: object = None) -> object: ...
|
||||
|
||||
def get(self, key: str, default: object = None) -> object:
|
||||
return self.__dict__.get(key, default)
|
||||
|
||||
def __contains__(self, key) -> bool:
|
||||
def __contains__(self, key: str) -> bool:
|
||||
return key in self.__dict__
|
||||
|
|
|
|||
|
|
@ -971,6 +971,23 @@ def test_arize_phoenix_client_get_prompt_version_rejects_traversal():
|
|||
client.get_prompt_version("../../projects")
|
||||
|
||||
|
||||
def test_arize_phoenix_client_error_does_not_echo_prompt_content():
|
||||
import httpx
|
||||
import respx
|
||||
|
||||
from litellm.integrations.arize.arize_phoenix_client import ArizePhoenixClient
|
||||
|
||||
client = ArizePhoenixClient(api_key="test-key", api_base="https://phoenix.test")
|
||||
with respx.mock:
|
||||
respx.get("https://phoenix.test/v1/prompt_versions/pv-1").mock(
|
||||
return_value=httpx.Response(200, json={"data": "secret system prompt text"})
|
||||
)
|
||||
with pytest.raises(Exception, match="Error fetching prompt version 'pv-1'") as exc_info:
|
||||
client.get_prompt_version("pv-1")
|
||||
|
||||
assert "secret system prompt text" not in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def _vcr_outcome_gate(request, vcr):
|
||||
install_live_call_probe(request, vcr)
|
||||
|
|
|
|||
|
|
@ -33,3 +33,7 @@ class TestA2AGuardrailHandlerStreamingScanKey:
|
|||
def test_unparseable_items_are_ignored(self):
|
||||
key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hi"), "not json", b"bytes"])
|
||||
assert key.texts == ("hi",)
|
||||
|
||||
def test_json_items_that_are_not_objects_are_ignored(self):
|
||||
key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hi"), "[1, 2]", "5"])
|
||||
assert key.texts == ("hi",)
|
||||
|
|
|
|||
|
|
@ -302,6 +302,57 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
|
|||
texts.append(data["delta"]["text"])
|
||||
return texts
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"off_schema_delta",
|
||||
[None, "stray", {"type": "text_delta", "text": 5}, {"type": "thinking_delta", "text": None}],
|
||||
)
|
||||
def test_off_schema_delta_is_skipped_without_ending_the_text_scan(self, off_schema_delta):
|
||||
events = [
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "abc"}},
|
||||
{"type": "content_block_delta", "index": 1, "delta": off_schema_delta},
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "def"}},
|
||||
]
|
||||
chunk = "".join(f"event: content_block_delta\ndata: {json.dumps(event)}\n\n" for event in events).encode()
|
||||
|
||||
assert AnthropicMessagesHandler().get_streaming_string_so_far([chunk]) == "abcdef"
|
||||
|
||||
def test_malformed_delta_data_is_skipped_without_ending_the_text_scan(self):
|
||||
text_delta = {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "abc"}}
|
||||
chunk = (
|
||||
f"event: content_block_delta\ndata: {json.dumps(text_delta)}\n\n"
|
||||
"event: content_block_delta\ndata: {not json\n\n"
|
||||
f"event: content_block_delta\ndata: {json.dumps(text_delta)}\n\n"
|
||||
).encode()
|
||||
|
||||
assert AnthropicMessagesHandler().get_streaming_string_so_far([chunk]) == "abcabc"
|
||||
|
||||
def test_text_delta_with_a_lone_surrogate_escape_is_scanned(self):
|
||||
event = {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "pineapple \ud83d"}}
|
||||
chunk = f"event: content_block_delta\ndata: {json.dumps(event)}\n\n".encode()
|
||||
|
||||
assert AnthropicMessagesHandler().get_streaming_string_so_far([chunk]) == "pineapple \ud83d"
|
||||
|
||||
def test_message_delta_with_a_lone_surrogate_escape_still_ends_the_stream(self):
|
||||
event = {"type": "message_delta", "delta": {"stop_reason": "stop_sequence", "stop_sequence": "\ud83d"}}
|
||||
chunk = f"event: message_delta\ndata: {json.dumps(event)}\n\n".encode()
|
||||
|
||||
scan_key = AnthropicMessagesHandler().get_streaming_scan_key([chunk])
|
||||
|
||||
assert scan_key is not None
|
||||
assert scan_key.stream_ended is True
|
||||
|
||||
def test_off_schema_message_delta_does_not_hide_a_later_stop_reason(self):
|
||||
events = [
|
||||
{"type": "message_delta", "delta": None},
|
||||
{"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}},
|
||||
]
|
||||
chunk = "".join(f"event: message_delta\ndata: {json.dumps(event)}\n\n" for event in events).encode()
|
||||
|
||||
scan_key = AnthropicMessagesHandler().get_streaming_scan_key([chunk])
|
||||
|
||||
assert scan_key is not None
|
||||
assert scan_key.stream_ended is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_ended_stream_rewrites_writes_text_back_into_sse_chunks(self):
|
||||
handler = AnthropicMessagesHandler()
|
||||
|
|
|
|||
|
|
@ -386,6 +386,76 @@ class TestAzureMAIImageGeneration:
|
|||
assert normalized["input_tokens"] == 22
|
||||
assert normalized["total_tokens"] == 22
|
||||
|
||||
def test_normalize_mai_usage_accepts_a_bool_token_count(self):
|
||||
normalized = AzureFoundryMAIImageGenerationConfig().normalize_mai_image_usage(
|
||||
{"input_tokens": 22, "output_tokens": True}
|
||||
)
|
||||
|
||||
assert normalized["output_tokens"] == 1
|
||||
assert normalized["total_tokens"] == 23
|
||||
|
||||
def test_normalize_mai_usage_reads_string_token_counts_as_numbers(self):
|
||||
normalized = AzureFoundryMAIImageGenerationConfig().normalize_mai_image_usage(
|
||||
{"input_tokens": "8", "output_tokens": "1024"}
|
||||
)
|
||||
|
||||
assert normalized["input_tokens"] == 8
|
||||
assert normalized["total_tokens"] == 1032
|
||||
|
||||
@pytest.mark.parametrize("unreadable", ["n/a", "inf", "nan", float("inf")])
|
||||
def test_normalize_mai_usage_falls_back_past_an_unreadable_count(self, unreadable):
|
||||
normalized = AzureFoundryMAIImageGenerationConfig().normalize_mai_image_usage(
|
||||
{"input_tokens": 8, "output_tokens": unreadable, "num_output_tokens": 1024}
|
||||
)
|
||||
|
||||
assert normalized["output_tokens"] == 1024
|
||||
assert normalized["total_tokens"] == 1032
|
||||
|
||||
@pytest.mark.parametrize("unreadable", ["n/a", {"n": 1}, [1024]])
|
||||
def test_transform_image_generation_response_keeps_the_image_when_a_count_is_unreadable(self, unreadable):
|
||||
raw_response = MagicMock(spec=httpx.Response)
|
||||
raw_response.json.return_value = {
|
||||
"created": 1780897477,
|
||||
"data": [{"b64_json": "abc123"}],
|
||||
"usage": {"input_tokens": 8, "output_tokens": 1024, "total_tokens": 1032, "num_output_tokens": unreadable},
|
||||
}
|
||||
|
||||
image_response = AzureFoundryMAIImageGenerationConfig().transform_image_generation_response(
|
||||
model="MAI-Image-2.5",
|
||||
raw_response=raw_response,
|
||||
model_response=ImageResponse(),
|
||||
logging_obj=MagicMock(),
|
||||
request_data={"prompt": "A red fox"},
|
||||
optional_params={"width": 1024, "height": 1024},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
assert image_response.data[0].b64_json == "abc123"
|
||||
assert image_response.usage.output_tokens == 1024
|
||||
assert image_response.usage.total_tokens == 1032
|
||||
|
||||
@pytest.mark.parametrize("usage", [[], "n/a", 5])
|
||||
def test_transform_image_generation_response_keeps_the_image_when_usage_is_not_an_object(self, usage):
|
||||
raw_response = MagicMock(spec=httpx.Response)
|
||||
raw_response.json.return_value = {"created": 1780897477, "data": [{"b64_json": "abc123"}], "usage": usage}
|
||||
|
||||
image_response = AzureFoundryMAIImageGenerationConfig().transform_image_generation_response(
|
||||
model="MAI-Image-2.5",
|
||||
raw_response=raw_response,
|
||||
model_response=ImageResponse(),
|
||||
logging_obj=MagicMock(),
|
||||
request_data={"prompt": "A red fox"},
|
||||
optional_params={"width": 1024, "height": 1024},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
assert image_response.data[0].b64_json == "abc123"
|
||||
assert image_response.usage.input_tokens == 0
|
||||
assert image_response.usage.output_tokens == 0
|
||||
assert image_response.usage.total_tokens == 0
|
||||
|
||||
def test_azure_sync_image_generation_uses_mai_response_transform(self):
|
||||
raw_response = MagicMock(spec=httpx.Response)
|
||||
raw_response.json.return_value = {
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@ import pytest
|
|||
import requests
|
||||
|
||||
|
||||
|
||||
from litellm.proxy.client.users import (
|
||||
NotFoundError,
|
||||
UnauthorizedError,
|
||||
|
|
@ -72,9 +71,9 @@ def test_create_user_unauthorized(mock_post, client):
|
|||
@patch("requests.post")
|
||||
def test_delete_user_success(mock_post, client):
|
||||
mock_post.return_value.status_code = 200
|
||||
mock_post.return_value.json.return_value = {"deleted": 1}
|
||||
mock_post.return_value.json.return_value = 1
|
||||
result = client.delete_user(["u1"])
|
||||
assert result["deleted"] == 1
|
||||
assert result == 1
|
||||
mock_post.assert_called_once()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -4292,6 +4292,34 @@ async def test_get_user_info_for_proxy_admin_validates_keys_and_teams():
|
|||
assert returned_key["models"] == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"raw_rows",
|
||||
[
|
||||
[{"teams": None, "keys": "sk-row-content-must-not-leak"}],
|
||||
[{"teams": None, "keys": ["sk-row-content-must-not-leak"]}],
|
||||
],
|
||||
)
|
||||
async def test_get_user_info_for_proxy_admin_errors_do_not_echo_row_content(raw_rows):
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import (
|
||||
_get_user_info_for_proxy_admin,
|
||||
)
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.query_raw = AsyncMock(return_value=raw_rows)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
await _get_user_info_for_proxy_admin(user_api_key_dict=UserAPIKeyAuth(user_id=None))
|
||||
|
||||
assert "sk-row-content-must-not-leak" not in str(exc_info.value)
|
||||
|
||||
|
||||
def _object_permission_mocks(mocker, existing_object_permission_id=None):
|
||||
"""Prisma double whose user row optionally already links a permission row."""
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue