refactor(types): replace Any with proven types in 39 files

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo 2026-10-18 10:57:03 +00:00
parent e979bd0ecf
commit 638cdc0c82
39 changed files with 457 additions and 187 deletions

View file

@ -2,12 +2,10 @@
Handler for transforming /chat/completions api requests to litellm.responses requests
"""
from collections.abc import AsyncIterable, Coroutine, Iterable, Mapping
from collections.abc import AsyncIterable, Iterable, Mapping
from typing import (
TYPE_CHECKING,
Any,
Final,
Union,
cast, # noqa: TID251 # optional_params reaches the bridge as an untyped dict
)
@ -160,13 +158,7 @@ class ResponsesToCompletionBridgeHandler:
encoding=typed_kwargs.get("encoding"),
)
def completion(
self, *args, **kwargs
) -> Union[
Coroutine[None, None, Union["ModelResponse", "CustomStreamWrapper"]],
"ModelResponse",
"CustomStreamWrapper",
]:
def completion(self, *args, **kwargs) -> object:
if kwargs.get("acompletion") is True:
return self.acompletion(**kwargs)
@ -269,7 +261,7 @@ class ResponsesToCompletionBridgeHandler:
)
return self._apply_post_stream_processing(streamwrapper, model, custom_llm_provider)
async def acompletion(self, *args, **kwargs) -> Union["ModelResponse", "CustomStreamWrapper"]:
async def acompletion(self, *args, **kwargs) -> object:
from litellm import aresponses
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
@ -379,7 +371,7 @@ class ResponsesToCompletionBridgeHandler:
custom_llm_provider: str,
logging_obj: "LiteLLMLoggingObj",
json_mode: bool | None,
) -> "CustomStreamWrapper":
) -> object:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
@ -396,7 +388,7 @@ class ResponsesToCompletionBridgeHandler:
stream: "CustomStreamWrapper",
model: str,
custom_llm_provider: str,
) -> Any:
) -> object:
"""Apply provider-specific post-stream processing if available."""
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager

View file

@ -22,6 +22,8 @@ from typing import Final, TypeVar
from urllib.parse import urlparse
import httpx
from pydantic import ConfigDict, TypeAdapter, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.integrations.batch_utils import (
@ -52,6 +54,15 @@ MONITOR_SCOPE_BY_AUTHORITY_HOST: Final[Mapping[str, str]] = MappingProxyType(
)
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _AzureOAuthTokenResponse(TypedDict, total=False):
access_token: ReadOnly[NotRequired[str | None]]
expires_in: ReadOnly[NotRequired[int | float]]
_AZURE_OAUTH_TOKEN_RESPONSE_ADAPTER: Final = TypeAdapter(_AzureOAuthTokenResponse)
class AzureSentinelLogger(CustomBatchLogger):
"""
Logger that sends LiteLLM logs to Azure Sentinel via Azure Monitor Logs Ingestion API
@ -226,7 +237,7 @@ class AzureSentinelLogger(CustomBatchLogger):
if response.status_code != 200:
raise Exception(f"Failed to get OAuth2 token: {response.status_code} - {response.text}")
token_response: Final = response.json()
token_response: Final = _AZURE_OAUTH_TOKEN_RESPONSE_ADAPTER.validate_python(response.json())
self.oauth_token = token_response.get("access_token")
expires_in: Final = token_response.get("expires_in", 3600)

View file

@ -12,7 +12,7 @@ import os
from collections.abc import Mapping, Sequence
from datetime import datetime
from types import MappingProxyType
from typing import Any, Final, Literal
from typing import Final, Literal
import httpx
from pydantic import ConfigDict, TypeAdapter, ValidationError
@ -54,6 +54,7 @@ from litellm.types.utils import (
PROMPT_CARRYING_GUARDRAIL_FIELDS,
PROMPT_QUOTING_ROUTING_DECISION_FIELDS,
CallTypes,
StandardLoggingMetadata,
StandardLoggingPayload,
StandardLoggingPayloadErrorInformation,
)
@ -93,8 +94,12 @@ _ROUTER_DIMENSIONS: Final[tuple[str, ...]] = ("router_tier", "router_cause", "ro
_COST_DIMENSIONS: Final[tuple[str, ...]] = ("team", "user", "key_alias", "model_group", *_ROUTER_DIMENSIONS)
def _metadata_of(standard_logging_payload: StandardLoggingPayload) -> Mapping[str, Any]:
metadata: Final = standard_logging_payload.get("metadata")
def _logging_metadata_of(standard_logging_payload: StandardLoggingPayload) -> StandardLoggingMetadata | None:
return standard_logging_payload.get("metadata")
def _metadata_of(standard_logging_payload: StandardLoggingPayload) -> Mapping[str, object]:
metadata: Final = _logging_metadata_of(standard_logging_payload)
return metadata or _EMPTY_MAPPING
@ -1057,14 +1062,23 @@ class DataDogLLMObsLogger(CustomBatchLogger):
# Get budget information from metadata
metadata: Final = _metadata_of(standard_logging_payload)
logging_metadata: Final = _logging_metadata_of(standard_logging_payload)
# API key max budget
user_api_key_max_budget: Final = metadata.get("user_api_key_max_budget")
user_api_key_max_budget: Final = (
logging_metadata["user_api_key_max_budget"]
if logging_metadata and "user_api_key_max_budget" in logging_metadata
else None
)
if user_api_key_max_budget is not None:
spend_metrics["user_api_key_max_budget"] = float(user_api_key_max_budget)
# API key spend
user_api_key_spend: Final = metadata.get("user_api_key_spend")
user_api_key_spend: Final = (
logging_metadata["user_api_key_spend"]
if logging_metadata and "user_api_key_spend" in logging_metadata
else None
)
if user_api_key_spend is not None:
try:
spend_metrics["user_api_key_spend"] = float(user_api_key_spend)

View file

@ -5,16 +5,21 @@ Functions for sending Email Alerts
import os
from typing import Final
from pydantic import ConfigDict, TypeAdapter
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm.repositories.team_repository import TeamRepository
from litellm.types.integrations.slack_alerting import WebhookEvent
_QUERY_RAW_RESULT: Final = TypeAdapter(object, config=ConfigDict(strict=True))
_EMAIL_QUERY_RESULT: Final = TypeAdapter(list[object], config=ConfigDict(strict=True))
# we use this for the email header, please send a test email if you change this. verify it looks good on email
LITELLM_LOGO_URL: Final = "https://litellm-listing.s3.amazonaws.com/litellm_logo.png"
LITELLM_SUPPORT_CONTACT: Final = "support@berri.ai"
async def get_all_team_member_emails(team_id: str | None = None) -> list:
async def get_all_team_member_emails(team_id: str | None = None) -> list[str]:
verbose_logger.debug("Email Alerting: Getting all team members for team_id=%s", team_id)
if team_id is None:
return []
@ -51,18 +56,21 @@ async def get_all_team_member_emails(team_id: str | None = None) -> list:
WHERE user_id = ANY($1::TEXT[]);
"""
_result: Final = await prisma_client.db.query_raw(sql_query, _team_member_user_ids)
raw_result: Final = _QUERY_RAW_RESULT.validate_python(
await prisma_client.db.query_raw(sql_query, _team_member_user_ids)
)
result: Final = None if raw_result is None else _EMAIL_QUERY_RESULT.validate_python(raw_result)
verbose_logger.debug("Email Alerting: Got all Emails for team, emails=%s", _result)
verbose_logger.debug("Email Alerting: Got all Emails for team, emails=%s", raw_result)
if _result is None:
if result is None:
return []
emails: Final = []
for user in _result:
if user and isinstance(user, dict) and user.get("user_email", None) is not None:
emails.append(user.get("user_email"))
return emails
return [
user_email
for user in result
if user and isinstance(user, dict) and isinstance((user_email := user.get("user_email")), str)
]
async def send_team_budget_alert(webhook_event: WebhookEvent) -> bool:

View file

@ -8,7 +8,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable
import httpx
from pydantic import TypeAdapter, ValidationError
from pydantic import ConfigDict, TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id
@ -25,6 +25,7 @@ else:
_CODEX_CLIENT_PREFIX_RE: Final = re.compile(r"^codex[-_ /]", re.IGNORECASE)
_MODEL_DUMP: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True))
def is_codex_user_agent(user_agent: str) -> bool:
@ -559,7 +560,8 @@ def preserve_upstream_non_openai_attributes(
"""
# Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings
expected_keys: Final = set(type(model_response).model_fields.keys()).union({"usage"})
for key, value in original_chunk.model_dump().items():
original_attributes: Final = _MODEL_DUMP.validate_python(original_chunk.model_dump())
for key, value in original_attributes.items():
if key not in expected_keys:
setattr(model_response, key, value)

View file

@ -3100,7 +3100,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
## RESPONSE OBJECT
try:
completion_response: Final = raw_response.json()
completion_response: Final = _COMPLETION_RESPONSE.validate_python(raw_response.json())
except Exception as e:
response_headers: Final = getattr(raw_response, "headers", None)
raise AnthropicError(
@ -3186,6 +3186,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
_DECODED_JSON: Final = TypeAdapter(object)
_COMPLETION_RESPONSE: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
_REQUEST_TOOLS: Final = TypeAdapter(tuple[Mapping[str, object], ...])

View file

@ -3,6 +3,7 @@ from typing import Final
from urllib.parse import urlsplit, urlunsplit
from openai import OpenAI
from pydantic import TypeAdapter
import litellm
from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base
@ -18,6 +19,8 @@ from litellm.utils import convert_to_model_response_object
from .cohere_transformation import AzureAICohereConfig
_EMBEDDING_RESPONSE: Final = TypeAdapter[dict[str, object]](dict[str, object])
def _foundry_models_route_base(api_base: str | None) -> str | None:
if api_base is None or not is_foundry_model_inference_base(api_base):
@ -88,7 +91,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
headers={"Authorization": f"Bearer {api_key}"},
)
embedding_response: Final = response.json()
embedding_response: Final = _EMBEDDING_RESPONSE.validate_python(response.json())
embedding_headers: Final = dict(response.headers)
returned_response: Final[EmbeddingResponse] = convert_to_model_response_object(
response_object=embedding_response,
@ -131,7 +134,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
headers={"Authorization": f"Bearer {api_key}"},
)
embedding_response: Final = response.json()
embedding_response: Final = _EMBEDDING_RESPONSE.validate_python(response.json())
embedding_headers: Final = dict(response.headers)
returned_response: Final[EmbeddingResponse] = convert_to_model_response_object(
response_object=embedding_response,

View file

@ -428,7 +428,7 @@ class BaseConfig(ABC):
streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse"],
sync_stream: bool,
json_mode: bool | None = False,
) -> Any:
) -> object:
pass
def get_model_response_iterator_for_request(

View file

@ -14,10 +14,13 @@ from __future__ import annotations
import base64
import os
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Final, cast
import httpx
from httpx._types import RequestFiles
from pydantic import ConfigDict, TypeAdapter, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
@ -39,6 +42,17 @@ else:
LiteLLMLoggingObj = Any
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _NovaCanvasSuccessResponse(TypedDict):
images: NotRequired[ReadOnly[Sequence[str] | None]]
errors: NotRequired[ReadOnly[object]]
message: NotRequired[ReadOnly[object]]
error: NotRequired[ReadOnly[object]]
_NOVA_CANVAS_SUCCESS_RESPONSE: Final = TypeAdapter(_NovaCanvasSuccessResponse)
def _nova_canvas_task_body(
*,
image_b64: str,
@ -401,7 +415,7 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig):
json_mode: bool | None = None,
) -> ImageResponse:
try:
response_data: Final = raw_response.json()
raw_response_data: Final[object] = raw_response.json()
except Exception as e:
raise self.get_error_class(
error_message=f"Error parsing Nova Canvas image edit response: {e}",
@ -411,12 +425,13 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig):
if raw_response.status_code not in (200,):
raise self.get_error_class(
error_message=f"Nova Canvas image edit error: {response_data}",
error_message=f"Nova Canvas image edit error: {raw_response_data}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
images: Final[list[str]] = response_data.get("images") or []
response_data: Final = _NOVA_CANVAS_SUCCESS_RESPONSE.validate_python(raw_response_data)
images: Final[Sequence[str]] = response_data.get("images") or []
if "errors" in response_data and not images:
raise self.get_error_class(

View file

@ -8487,7 +8487,7 @@ class BaseLLMHTTPHandler:
prefetch_resp.raise_for_status()
except Exception as e:
raise self._handle_error(e=e, provider_config=video_provider_config)
prefetched_source_data = prefetch_resp.json()
prefetched_source_data = _JSON_OBJECT.validate_python(prefetch_resp.json())
try:
url, data, files = video_provider_config.transform_video_edit_request(
@ -8585,7 +8585,7 @@ class BaseLLMHTTPHandler:
prefetch_resp.raise_for_status()
except Exception as e:
raise self._handle_error(e=e, provider_config=video_provider_config)
prefetched_source_data = prefetch_resp.json()
prefetched_source_data = _JSON_OBJECT.validate_python(prefetch_resp.json())
try:
url, data, files = video_provider_config.transform_video_edit_request(

View file

@ -1,10 +1,12 @@
import base64
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from io import BufferedReader, BytesIO
from typing import TYPE_CHECKING, Any, Final, cast
import httpx
from httpx._types import RequestFiles
from pydantic import ConfigDict, TypeAdapter, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.images.utils import ImageEditRequestUtils
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
@ -33,6 +35,35 @@ else:
LiteLLMLoggingObj = Any
@with_config(ConfigDict(extra="allow", strict=True))
class _GeminiInlineData(TypedDict):
data: NotRequired[ReadOnly[str | None]]
@with_config(ConfigDict(extra="allow", strict=True))
class _GeminiPart(TypedDict):
inlineData: NotRequired[ReadOnly[_GeminiInlineData | None]]
@with_config(ConfigDict(extra="allow", strict=True))
class _GeminiContent(TypedDict):
parts: NotRequired[ReadOnly[Sequence[_GeminiPart]]]
@with_config(ConfigDict(extra="allow", strict=True))
class _GeminiCandidate(TypedDict):
content: NotRequired[ReadOnly[_GeminiContent]]
@with_config(ConfigDict(extra="allow", strict=True))
class _GeminiImageEditResponse(TypedDict):
candidates: NotRequired[ReadOnly[Sequence[_GeminiCandidate]]]
usageMetadata: NotRequired[ReadOnly[Mapping[str, object]]]
_GEMINI_IMAGE_EDIT_RESPONSE: Final = TypeAdapter(_GeminiImageEditResponse)
class GeminiImageEditConfig(BaseImageEditConfig):
DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta"
SUPPORTED_PARAMS: list[str] = ["n", "size", "imageConfig"]
@ -133,7 +164,8 @@ class GeminiImageEditConfig(BaseImageEditConfig):
headers=raw_response.headers,
)
candidates: Final = response_json.get("candidates", [])
response_body: Final = _GEMINI_IMAGE_EDIT_RESPONSE.validate_python(response_json)
candidates: Final = response_body.get("candidates", [])
data_list: Final[list[ImageObject]] = []
for candidate in candidates:
@ -141,17 +173,18 @@ class GeminiImageEditConfig(BaseImageEditConfig):
parts = content.get("parts", [])
for part in parts:
inline_data = part.get("inlineData")
if inline_data and inline_data.get("data"):
image_data = inline_data.get("data") if inline_data else None
if image_data:
data_list.append(
ImageObject(
b64_json=inline_data["data"],
b64_json=image_data,
url=None,
)
)
model_response.data = cast(list[OpenAIImage], data_list)
if "usageMetadata" in response_json:
model_response.usage = transform_gemini_image_usage(response_json["usageMetadata"])
if "usageMetadata" in response_body:
model_response.usage = transform_gemini_image_usage(dict(response_body["usageMetadata"]))
return model_response
def _prepare_inline_image_parts(self, image: FileTypes | list[FileTypes]) -> list[dict[str, object]]:

View file

@ -12,6 +12,9 @@ import uuid
from collections.abc import Mapping
from typing import Final
from pydantic import ConfigDict, TypeAdapter, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -26,6 +29,14 @@ from .authenticator import get_access_token, get_access_token_async
_file_cache: Final[dict[str, str]] = {}
@with_config(ConfigDict(extra="allow", strict=True))
class _UploadedFile(TypedDict):
id: ReadOnly[NotRequired[str | None]]
_UPLOADED_FILE: Final = TypeAdapter(_UploadedFile)
def _get_url_hash(url: str) -> str:
"""Generate hash for URL to use as cache key."""
return hashlib.sha256(url.encode()).hexdigest()
@ -137,7 +148,7 @@ def upload_file_sync(
timeout=60,
)
response.raise_for_status()
result: Final = response.json()
result: Final = _UPLOADED_FILE.validate_python(response.json())
file_id: Final = result.get("id")
if file_id:
@ -206,7 +217,7 @@ async def upload_file_async(
timeout=60,
)
response.raise_for_status()
result: Final = response.json()
result: Final = _UPLOADED_FILE.validate_python(response.json())
file_id: Final = result.get("id")
if file_id:

View file

@ -2,7 +2,7 @@ import json
import os
import time
from datetime import datetime
from typing import Any, Final
from typing import Final
import httpx
from pydantic import ConfigDict, TypeAdapter, with_config
@ -40,9 +40,15 @@ class _AccessTokenPoll(TypedDict):
error: ReadOnly[NotRequired[object]]
@with_config(ConfigDict(extra="allow", strict=True))
class _ApiKeyInfo(TypedDict):
token: ReadOnly[NotRequired[str | None]]
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
_DEVICE_CODE: Final = TypeAdapter(_DeviceCode)
_ACCESS_TOKEN_POLL: Final = TypeAdapter(_AccessTokenPoll)
_API_KEY_INFO: Final = TypeAdapter(_ApiKeyInfo)
def github_api_headers(
@ -194,7 +200,7 @@ class Authenticator:
verbose_logger.warning("Error reading API endpoint from file: %s", e)
return None
def _refresh_api_key(self) -> dict[str, Any]:
def _refresh_api_key(self) -> _ApiKeyInfo:
"""
Refresh the API key using the access token.
@ -215,7 +221,7 @@ class Authenticator:
response = sync_client.get(api_key_url, headers=headers)
response.raise_for_status()
response_json = response.json()
response_json = _API_KEY_INFO.validate_python(response.json())
if "token" in response_json:
return response_json

View file

@ -17,7 +17,9 @@ from collections.abc import Mapping, Sequence
from enum import Enum
from typing import Any, Final, Protocol, TypedDict
from pydantic import ConfigDict, TypeAdapter, with_config
from typing_extensions import ReadOnly
from typing_extensions import TypedDict as _TypedDict
from litellm._logging import verbose_logger
@ -50,10 +52,14 @@ class _AnthropicToolSpec(TypedDict, total=False):
input_schema: ReadOnly[_ToolArgumentSchema]
class _CodeExecutionArguments(TypedDict, total=False):
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _CodeExecutionArguments(_TypedDict, total=False):
code: ReadOnly[str]
_CODE_EXECUTION_ARGUMENTS: Final = TypeAdapter(_CodeExecutionArguments)
class _GeneratedFile(TypedDict, total=False):
name: ReadOnly[str]
mime_type: ReadOnly[str]
@ -115,7 +121,7 @@ class _CodeExecutionOutcome(TypedDict, total=False):
def _parse_code_execution_arguments(serialized_arguments: str) -> _CodeExecutionArguments:
return json.loads(serialized_arguments)
return _CODE_EXECUTION_ARGUMENTS.validate_python(json.loads(serialized_arguments))
class LiteLLMInternalTools(str, Enum):

View file

@ -14,7 +14,7 @@ from collections.abc import AsyncIterator, Callable, Iterator, Mapping
from typing import TYPE_CHECKING, Any, Final
import httpx
from pydantic import TypeAdapter, ValidationError
from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_logger
@ -188,6 +188,7 @@ def _normalize_tool_choice(selected_params: dict) -> None:
_TOOL_CHOICE_MODES: Final = frozenset({"auto", "none", "required", "any"})
_TOOLS_ADAPTER: Final = TypeAdapter(list[dict[str, object]])
_JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object])
_OCI_JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, JsonValue], config=ConfigDict(strict=True))
def _as_json_object(value: object) -> Mapping[str, object] | None:
@ -714,11 +715,12 @@ class OCIChatConfig(BaseConfig):
status_code=raw_response.status_code,
)
validated_response_json: Final = _OCI_JSON_OBJECT_ADAPTER.validate_python(response_json)
vendor: Final = get_vendor_from_model(model)
if vendor == OCIVendors.COHERE:
model_response = handle_cohere_response(response_json, model, model_response, raw_response)
model_response = handle_cohere_response(validated_response_json, model, model_response, raw_response)
else:
model_response = handle_generic_response(response_json, model, model_response, raw_response)
model_response = handle_generic_response(validated_response_json, model, model_response, raw_response)
model_response.hidden_params["additional_headers"] = raw_response.headers
return model_response
@ -869,9 +871,10 @@ class OCIStreamWrapper(CustomStreamWrapper):
message=f"Chunk cannot be parsed as JSON: {e}",
)
if dict_chunk.get("apiFormat") == "COHERE":
validated_dict_chunk: Final = _OCI_JSON_OBJECT_ADAPTER.validate_python(dict_chunk)
if validated_dict_chunk.get("apiFormat") == "COHERE":
result: Final = handle_cohere_stream_chunk(
dict_chunk,
validated_dict_chunk,
prior_tool_calls_emitted=self._cohere_tool_calls_emitted,
prior_text_emitted=self._cohere_text_emitted,
)
@ -886,7 +889,7 @@ class OCIStreamWrapper(CustomStreamWrapper):
self._cohere_text_emitted = True
break
return self._emit_chunk(result)
return self._emit_chunk(handle_generic_stream_chunk(dict_chunk))
return self._emit_chunk(handle_generic_stream_chunk(validated_dict_chunk))
__all__ = [

View file

@ -12,6 +12,19 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.utils import ModelResponse, ModelResponseStream
_PROVIDER_FIELDS: Final = TypeAdapter(dict[str, object])
_OLLAMA_JSON: Final = TypeAdapter[object](object)
_OLLAMA_FIELDS: Final = TypeAdapter[dict[str, object]](dict[str, object])
_OLLAMA_ENTRIES: Final = TypeAdapter[list[object]](list[object])
def _model_name_from_ollama_tag(entry: object) -> str | None:
if not isinstance(entry, dict):
return None
entry_fields: Final = _OLLAMA_FIELDS.validate_python(entry)
name: Final = entry_fields.get("name") or entry_fields.get("model")
if not isinstance(name, str):
return None
return name if name.startswith("ollama/") else f"ollama/{name}"
class OllamaResponseIdNormalizer:
@ -147,24 +160,22 @@ class OllamaModelInfo(BaseLLMModelInfo):
api_key = self.get_api_key(api_key) if passed_api_base is None or api_key else None
headers: Final = {"Authorization": f"Bearer {api_key}"} if api_key else {}
names: Final[set[str]] = set()
try:
resp: Final = httpx.get(f"{base}/api/tags", headers=headers)
resp.raise_for_status()
data: Final = resp.json()
# Expecting a dict with a 'models' list
models_list = []
if isinstance(data, dict) and "models" in data and isinstance(data["models"], list):
models_list = data["models"]
elif isinstance(data, list):
models_list = data
# Extract model names
for entry in models_list:
if not isinstance(entry, dict):
continue
nm = entry.get("name") or entry.get("model")
if isinstance(nm, str):
names.add(nm if nm.startswith("ollama/") else f"ollama/{nm}")
data: Final = _OLLAMA_JSON.validate_python(resp.json())
data_fields: Final = _OLLAMA_FIELDS.validate_python(data) if isinstance(data, dict) else None
model_field: Final = data_fields.get("models") if data_fields is not None else None
models_list: Final = (
_OLLAMA_ENTRIES.validate_python(model_field)
if isinstance(model_field, list)
else _OLLAMA_ENTRIES.validate_python(data)
if isinstance(data, list)
else []
)
names: Final = frozenset(
name for entry in models_list if (name := _model_name_from_ollama_tag(entry)) is not None
)
except Exception as e:
verbose_logger.warning("Error retrieving ollama tag endpoint: %s", e)
# If tags endpoint fails, fall back to static list
@ -201,13 +212,13 @@ class OllamaModelInfo(BaseLLMModelInfo):
return any(name.lower() in model_cost_keys for name in potential_model_names)
@staticmethod
def _supports_function_calling(ollama_model_info: dict) -> bool:
def _supports_function_calling(ollama_model_info: dict[str, object]) -> bool:
_template: Final[str] = str(ollama_model_info.get("template", "") or "")
return "tools" in _template.lower()
@staticmethod
def _get_max_tokens(ollama_model_info: dict) -> int | None:
_model_info: Final[dict] = ollama_model_info.get("model_info", {})
def _get_max_tokens(ollama_model_info: dict[str, object]) -> object:
_model_info: Final = _OLLAMA_FIELDS.validate_python(ollama_model_info.get("model_info", {}))
for key, value in _model_info.items():
if "context_length" in key:
@ -248,7 +259,7 @@ class OllamaModelInfo(BaseLLMModelInfo):
"max_output_tokens": None,
}
model_info: Final = response.json()
model_info: Final = _OLLAMA_FIELDS.validate_python(response.json())
max_tokens: Final = self._get_max_tokens(model_info)
return {

View file

@ -5,9 +5,9 @@ Ollama /chat/completion calls handled in llm_http_handler.py
"""
from collections.abc import Mapping, Sequence
from typing import Any, Final, Protocol, TypedDict
from typing import Any, Final, Protocol
from typing_extensions import NotRequired, ReadOnly
from pydantic import BaseModel, ConfigDict
import litellm
from litellm.types.utils import EmbeddingResponse
@ -19,11 +19,13 @@ class TokenEncoder(Protocol):
def encode(self, text: str, /) -> Sequence[int]: ...
class OllamaEmbeddingResponse(TypedDict):
class OllamaEmbeddingResponse(BaseModel):
"""Body of an Ollama ``/api/embed`` response."""
embeddings: ReadOnly[list[list[float]]]
prompt_eval_count: ReadOnly[NotRequired[int]]
model_config = ConfigDict(extra="allow", frozen=True, strict=True)
embeddings: list[list[int | float]]
prompt_eval_count: int | None = None
def _ollama_embed_url(api_base: str) -> str:
@ -58,12 +60,12 @@ def _process_ollama_embedding_response(
encoding: TokenEncoder | None,
) -> EmbeddingResponse:
output_data: Final = []
embeddings: Final[list[list[float]]] = response_json["embeddings"]
embeddings: Final[list[list[int | float]]] = response_json.embeddings
for idx, emb in enumerate(embeddings):
output_data.append({"object": "embedding", "index": idx, "embedding": emb})
input_tokens = response_json.get("prompt_eval_count", None)
input_tokens: int | None = response_json.prompt_eval_count
if input_tokens is None:
if encoding is not None:
@ -100,7 +102,7 @@ async def ollama_aembeddings(
data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params)
response: Final = await litellm.module_level_aclient.post(url=_ollama_embed_url(api_base), json=data)
response_json: Final[OllamaEmbeddingResponse] = response.json()
response_json: Final = OllamaEmbeddingResponse.model_validate(response.json())
return _process_ollama_embedding_response(
response_json=response_json,
@ -124,7 +126,7 @@ def ollama_embeddings(
data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params)
response: Final = litellm.module_level_client.post(url=_ollama_embed_url(api_base), json=data)
response_json: Final[OllamaEmbeddingResponse] = response.json()
response_json: Final = OllamaEmbeddingResponse.model_validate(response.json())
return _process_ollama_embedding_response(
response_json=response_json,

View file

@ -6,12 +6,13 @@ import json
from pathlib import Path
from typing import Final
from pydantic import TypeAdapter
from pydantic import ConfigDict, TypeAdapter
from litellm._logging import verbose_logger
from litellm.types.utils import ProviderSpecificModelInfo
_PROVIDER_INFO_ADAPTER: Final = TypeAdapter(ProviderSpecificModelInfo)
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
class SimpleProviderConfig:
@ -53,10 +54,10 @@ class JSONProviderRegistry:
try:
with open(json_path) as f:
data: Final = json.load(f)
data: Final = _JSON_OBJECT.validate_python(json.load(f))
for slug, config in data.items():
cls._providers[slug] = SimpleProviderConfig(slug, config)
cls._providers[slug] = SimpleProviderConfig(slug, _JSON_OBJECT.validate_python(config))
cls._loaded = True
except Exception as e:

View file

@ -41,6 +41,7 @@ MAX_OUTPUT_BYTES: Final = SANDBOX_MAX_OUTPUT_BYTES
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
_JSON_VALUE: Final = TypeAdapter(object)
_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True, strict=True))
class OpenSandboxSandboxConfig(BaseSandboxConfig):
@ -491,21 +492,22 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig):
if not data:
return None
try:
parsed: Final = json.loads(data)
parsed: Final = _JSON_VALUE.validate_python(json.loads(data))
except json.JSONDecodeError:
return None
if not isinstance(parsed, dict):
return None
if "type" not in parsed and "code" in parsed and "message" in parsed:
parsed_object: Final = _JSON_DICT.validate_python(parsed)
if "type" not in parsed_object and "code" in parsed_object and "message" in parsed_object:
return {
"type": "error",
"error": {
"ename": str(parsed["code"]),
"evalue": str(parsed["message"]),
"ename": str(parsed_object["code"]),
"evalue": str(parsed_object["message"]),
"traceback": [],
},
}
return parsed
return parsed_object
@staticmethod
def _normalize_result(message: dict[str, object]) -> dict[str, object]:

View file

@ -6,12 +6,15 @@ from collections.abc import AsyncIterator, Iterator
from typing import Final
import httpx
from pydantic import ConfigDict, TypeAdapter
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.types.llms.openai import OpenAIChatCompletionChunk
from ...custom_httpx.llm_http_handler import BaseLLMHTTPHandler
_STREAM_EVENT_ADAPTER: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
# -------------------------------
# Errors
@ -107,6 +110,14 @@ class _StreamParser:
return None
def _parse_stream_event(payload: str) -> OpenAIChatCompletionChunk | None:
try:
event_object: Final[object] = json.loads(payload)
except Exception:
return None
return _StreamParser.to_openai_chunk(_STREAM_EVENT_ADAPTER.validate_python(event_object))
# -------------------------------
# Iterators
# -------------------------------
@ -146,12 +157,7 @@ class SAPStreamIterator:
raise StopIteration
try:
obj = json.loads(payload)
except Exception:
continue
try:
chunk = _StreamParser.to_openai_chunk(obj)
chunk = _parse_stream_event(payload)
except ValueError as e:
self._safe_close()
raise e
@ -217,12 +223,7 @@ class AsyncSAPStreamIterator:
await self._aclose()
raise StopAsyncIteration
try:
obj = json.loads(payload)
except Exception:
continue
try:
chunk = _StreamParser.to_openai_chunk(obj)
chunk = _parse_stream_event(payload)
except ValueError as e:
await self._aclose()
raise GenAIHubOrchestrationError(502, str(e))

View file

@ -10,6 +10,8 @@ from typing import TYPE_CHECKING, Any, Final, cast
import httpx
from httpx._types import RequestFiles
from pydantic import ConfigDict, TypeAdapter, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.secret_managers.main import get_secret_str
@ -30,6 +32,15 @@ else:
LiteLLMLoggingObj = Any
@with_config(ConfigDict(extra="allow", strict=True))
class _StabilityImageEditResponse(TypedDict):
finish_reason: ReadOnly[NotRequired[object]]
image: ReadOnly[NotRequired[str | None]]
_STABILITY_IMAGE_EDIT_RESPONSE: Final = TypeAdapter(_StabilityImageEditResponse)
class StabilityImageEditConfig(BaseImageEditConfig):
"""
Configuration for Stability AI image edit.
@ -275,8 +286,10 @@ class StabilityImageEditConfig(BaseImageEditConfig):
headers=raw_response.headers,
)
validated_response: Final = _STABILITY_IMAGE_EDIT_RESPONSE.validate_python(response_data)
# Check finish_reason
finish_reason: Final = response_data.get("finish_reason", "")
finish_reason: Final = validated_response.get("finish_reason", "")
if finish_reason == "CONTENT_FILTERED":
raise self.get_error_class(
error_message="Content was filtered by Stability AI safety systems",
@ -289,7 +302,7 @@ class StabilityImageEditConfig(BaseImageEditConfig):
model_response.data = []
# Extract image from response
image_b64: Final = response_data.get("image")
image_b64: Final = validated_response.get("image")
if image_b64:
model_response.data.append(
ImageObject(

View file

@ -1,9 +1,11 @@
import time
from collections.abc import AsyncIterator, Iterator, Mapping
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final
import httpx
from pydantic import ConfigDict, TypeAdapter, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUsageBlock
@ -29,6 +31,25 @@ else:
LiteLLMLoggingObj = Any
@with_config(ConfigDict(extra="allow", strict=True))
class _WatsonXTextGenResponse(TypedDict):
results: ReadOnly[Sequence[object]]
model_id: NotRequired[ReadOnly[str | None]]
created_at: NotRequired[ReadOnly[str | None]]
@with_config(ConfigDict(extra="allow", strict=True))
class _WatsonXTextGenResult(TypedDict):
generated_text: ReadOnly[str]
input_token_count: ReadOnly[int]
generated_token_count: ReadOnly[int]
stop_reason: ReadOnly[str]
_WATSONX_TEXT_GEN_RESPONSE: Final = TypeAdapter(_WatsonXTextGenResponse)
_WATSONX_TEXT_GEN_RESULT: Final = TypeAdapter(_WatsonXTextGenResult)
class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig):
"""
Reference: https://cloud.ibm.com/apidocs/watsonx-ai#text-generation
@ -297,19 +318,22 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig):
status_code=500,
message=f"Error: Invalid response from Watsonx.ai API: {json_resp}",
)
payload: Final = _WATSONX_TEXT_GEN_RESPONSE.validate_python(json_resp)
if model_response is None:
model_response = ModelResponse(model=json_resp.get("model_id", None))
generated_text: Final = json_resp["results"][0]["generated_text"]
prompt_tokens: Final = json_resp["results"][0]["input_token_count"]
completion_tokens: Final = json_resp["results"][0]["generated_token_count"]
model_response = ModelResponse(model=payload.get("model_id", None))
first_result: Final = _WATSONX_TEXT_GEN_RESULT.validate_python(payload["results"][0])
generated_text: Final = first_result["generated_text"]
prompt_tokens: Final = first_result["input_token_count"]
completion_tokens: Final = first_result["generated_token_count"]
model_response.choices[0].message.content = generated_text
model_response.choices[0].finish_reason = map_finish_reason(json_resp["results"][0]["stop_reason"])
if json_resp.get("created_at"):
model_response.choices[0].finish_reason = map_finish_reason(first_result["stop_reason"])
created_at: Final = payload.get("created_at")
if created_at:
try:
created_datetime = datetime.fromisoformat(json_resp["created_at"])
created_datetime = datetime.fromisoformat(created_at)
except ValueError:
# datetime.fromisoformat cannot handle 'Z' in Python 3.10
created_datetime = datetime.fromisoformat(f"{json_resp['created_at'].rstrip('Z')}+00:00")
created_datetime = datetime.fromisoformat(f"{created_at.rstrip('Z')}+00:00")
model_response.created = int(created_datetime.timestamp())
else:
model_response.created = int(time.time())

View file

@ -6,7 +6,6 @@ Canonical definition for ``litellm_organizationmembership``. Re-exported from
"""
from datetime import datetime
from typing import Any
from pydantic import ConfigDict, model_validator
@ -24,7 +23,7 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase):
budget_id: str | None = None
created_at: datetime
updated_at: datetime
user: Any | None = None
user: object | None = None
litellm_budget_table: LiteLLM_BudgetTable | None = None
user_email: str | None = None

View file

@ -2627,7 +2627,7 @@ async def _load_team_membership_on_cache_miss(
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging | None,
) -> LiteLLM_TeamMembership | None:
redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key)
redis_cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=cache_key)
redis_membership: Final = _membership_from_cached_payload(redis_cached)
if not isinstance(redis_membership, _TeamMembershipCacheMiss):
return redis_membership

View file

@ -371,4 +371,4 @@ class IPAddressUtils:
return client_ip
case _HopCountUnset():
pass
return get_request_ip_address(request, use_x_forwarded_for=use_xff)
return get_request_ip_address(request, use_x_forwarded_for=use_xff is True)

View file

@ -7,7 +7,8 @@ from datetime import datetime
from typing import TYPE_CHECKING, Final
import httpx
from pydantic import TypeAdapter
from pydantic import ConfigDict, TypeAdapter, with_config
from typing_extensions import ReadOnly, Required, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.constants import NON_LLM_CONNECTION_TIMEOUT
@ -23,6 +24,14 @@ AUTO_ROUTER_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' fe
_LICENSE_VERDICT: Final = TypeAdapter(object)
@with_config(ConfigDict(extra="allow", strict=True))
class _AirgappedLicensePayload(TypedDict, total=False):
expiration_date: Required[ReadOnly[str]]
_AIRGAPPED_LICENSE_PAYLOAD_ADAPTER: Final = TypeAdapter(_AirgappedLicensePayload)
class LicenseCheck:
"""
- Check if license in env
@ -201,7 +210,7 @@ class LicenseCheck:
)
# Decode and parse the data
license_data: Final = json.loads(message.decode())
license_data: Final = _AIRGAPPED_LICENSE_PAYLOAD_ADAPTER.validate_python(json.loads(message.decode()))
# Check expiration date
expiration_date: Final = datetime.strptime(license_data["expiration_date"], "%Y-%m-%d")

View file

@ -7,6 +7,7 @@ from pydantic import ConfigDict, TypeAdapter
from .exceptions import StreamError, UnauthorizedError
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
_SSE_LINE: Final[TypeAdapter[bytes | bytearray]] = TypeAdapter(
bytes | bytearray, config=ConfigDict(arbitrary_types_allowed=True, strict=True, hide_input_in_errors=True)
)
@ -59,7 +60,7 @@ class ChatClient:
frequency_penalty: float | None = None,
user: str | None = None,
return_request: bool = False,
) -> dict[str, Any] | requests.Request:
) -> dict[str, object] | requests.Request:
"""
Create a chat completion.
@ -114,7 +115,7 @@ class ChatClient:
try:
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
return _JSON_OBJECT.validate_python(response.json())
except requests.exceptions.HTTPError as e:
if e.response.status_code == 401:
raise UnauthorizedError(e)

View file

@ -3,9 +3,12 @@ from typing import Any, Final
from urllib.parse import quote
import requests
from pydantic import ConfigDict, TypeAdapter
from .exceptions import UnauthorizedError
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
class CredentialsManagementClient:
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
@ -75,7 +78,7 @@ class CredentialsManagementClient:
credential_values: Mapping[str, object],
return_request: bool = False,
display_name: str | None = None,
) -> dict[str, Any] | requests.Request:
) -> dict[str, object] | requests.Request:
"""
Create a new credential.
@ -111,7 +114,7 @@ class CredentialsManagementClient:
try:
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
return _JSON_OBJECT.validate_python(response.json())
except requests.exceptions.HTTPError as e:
if e.response.status_code == 401:
raise UnauthorizedError(e)
@ -121,7 +124,7 @@ class CredentialsManagementClient:
self,
credential_name: str,
return_request: bool = False,
) -> dict[str, Any] | requests.Request:
) -> dict[str, object] | requests.Request:
"""
Delete a credential by name.
@ -148,7 +151,7 @@ class CredentialsManagementClient:
try:
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
return _JSON_OBJECT.validate_python(response.json())
except requests.exceptions.HTTPError as e:
if e.response.status_code == 401:
raise UnauthorizedError(e)
@ -172,7 +175,7 @@ class CredentialsManagementClient:
try:
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
return _JSON_OBJECT.validate_python(response.json())
except requests.exceptions.HTTPError as e:
if e.response.status_code == 401:
raise UnauthorizedError(e)
@ -182,7 +185,7 @@ class CredentialsManagementClient:
self,
credential_name: str,
return_request: bool = False,
) -> dict[str, Any] | requests.Request:
) -> dict[str, object] | requests.Request:
"""
Get a credential by name.
@ -209,7 +212,7 @@ class CredentialsManagementClient:
try:
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
return _JSON_OBJECT.validate_python(response.json())
except requests.exceptions.HTTPError as e:
if e.response.status_code == 401:
raise UnauthorizedError(e)

View file

@ -3,11 +3,14 @@ from collections.abc import Mapping
from typing import Any, Final
import requests
from pydantic import ConfigDict, TypeAdapter
from litellm.litellm_core_utils.secret_redaction import redact_string
from .exceptions import UnauthorizedError
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
class KeysManagementClient:
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
@ -179,7 +182,7 @@ class KeysManagementClient:
try:
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
return _JSON_OBJECT.validate_python(response.json())
except requests.exceptions.HTTPError as e:
if e.response.status_code == 401:
raise UnauthorizedError(e)
@ -223,7 +226,7 @@ class KeysManagementClient:
try:
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
return _JSON_OBJECT.validate_python(response.json())
except requests.exceptions.HTTPError as e:
if e.response.status_code == 401:
raise UnauthorizedError(e)
@ -239,7 +242,7 @@ class KeysManagementClient:
key_alias: str | None = None,
team_id: str | None = None,
user_id: str | None = None,
) -> dict[str, object] | requests.Request:
) -> dict[str, object]:
"""
Update an existing API key's parameters.
@ -285,7 +288,7 @@ class KeysManagementClient:
response: Final = session.send(request.prepare(), timeout=self._timeout)
response_text = response.text
response.raise_for_status()
return response.json()
return _JSON_OBJECT.validate_python(response.json())
except Exception:
raise Exception(f"Error updating key: {response_text}")
@ -314,7 +317,7 @@ class KeysManagementClient:
try:
response: Final = session.send(request.prepare(), timeout=self._timeout)
response.raise_for_status()
return response.json()
return _JSON_OBJECT.validate_python(response.json())
except requests.exceptions.HTTPError as e:
redacted_message: Final = redact_string(str(e))
if e.response.status_code == 401:

View file

@ -4,9 +4,13 @@ import builtins
from typing import Any, Final
import requests
from pydantic import ConfigDict, TypeAdapter
from .exceptions import UnauthorizedError
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
_JSON_OBJECT_LIST: Final = TypeAdapter(builtins.list[dict[str, object]], config=ConfigDict(strict=True))
class TeamsManagementClient:
"""Client for managing teams in LiteLLM proxy."""
@ -125,9 +129,9 @@ class TeamsManagementClient:
raise UnauthorizedError("Authentication failed. Check your API key.")
response.raise_for_status()
return response.json()
return _JSON_OBJECT.validate_python(response.json())
def get_available(self) -> builtins.list[dict[str, Any]]:
def get_available(self) -> builtins.list[dict[str, object]]:
"""
Get list of available teams that the user can join.
@ -146,4 +150,4 @@ class TeamsManagementClient:
raise UnauthorizedError("Authentication failed. Check your API key.")
response.raise_for_status()
return response.json()
return _JSON_OBJECT_LIST.validate_python(response.json())

View file

@ -14,8 +14,8 @@ from ssl import SSLContext
from typing import TYPE_CHECKING, Any, Final
from fastapi import HTTPException
from pydantic import BaseModel
from typing_extensions import NotRequired, TypedDict
from pydantic import BaseModel, ConfigDict, TypeAdapter, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
from websockets.asyncio.client import ClientConnection, connect
from websockets.exceptions import ConnectionClosed
@ -58,38 +58,51 @@ class _WsSslKwargs(TypedDict, total=False):
ssl: bool | str | SSLContext
_CATO_JSON_CONFIG: Final = ConfigDict(extra="allow", strict=True)
@with_config(_CATO_JSON_CONFIG)
class _CatoRequiredAction(TypedDict, total=False):
action_type: str
detection_message: str
action_type: ReadOnly[str]
detection_message: ReadOnly[str]
@with_config(_CATO_JSON_CONFIG)
class _CatoRedactedMessage(TypedDict):
role: NotRequired[str]
content: str | None
content: NotRequired[ReadOnly[str | None]]
@with_config(_CATO_JSON_CONFIG)
class _CatoRedactedChat(TypedDict, total=False):
all_redacted_messages: Sequence[_CatoRedactedMessage]
all_redacted_messages: ReadOnly[Sequence[_CatoRedactedMessage]]
@with_config(_CATO_JSON_CONFIG)
class _CatoAnalysisResult(TypedDict, total=False):
policy_drill_down: Mapping[str, object]
policy_drill_down: ReadOnly[Mapping[str, object]]
@with_config(_CATO_JSON_CONFIG)
class _CatoAnalyzeResponse(TypedDict):
required_action: NotRequired[_CatoRequiredAction | None]
analysis_result: NotRequired[_CatoAnalysisResult]
redacted_chat: NotRequired[_CatoRedactedChat]
required_action: NotRequired[ReadOnly[_CatoRequiredAction | None]]
analysis_result: NotRequired[ReadOnly[_CatoAnalysisResult]]
redacted_chat: NotRequired[ReadOnly[_CatoRedactedChat]]
@with_config(_CATO_JSON_CONFIG)
class _CatoOutputRedaction(TypedDict):
redacted_output: str
redacted_output: ReadOnly[str]
@with_config(_CATO_JSON_CONFIG)
class _CatoStreamMessage(TypedDict, total=False):
verified_chunk: Mapping[str, object]
done: bool
blocking_message: str
verified_chunk: ReadOnly[Mapping[str, object]]
done: ReadOnly[object]
blocking_message: ReadOnly[str]
_CATO_ANALYZE_RESPONSE_ADAPTER: Final = TypeAdapter(_CatoAnalyzeResponse)
_CATO_STREAM_MESSAGE_ADAPTER: Final = TypeAdapter(_CatoStreamMessage)
class CatoNetworksGuardrail(CustomGuardrail):
@ -308,7 +321,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
timeout=self.timeout,
)
response.raise_for_status()
res: Final[_CatoAnalyzeResponse] = response.json()
res: Final = _CATO_ANALYZE_RESPONSE_ADAPTER.validate_python(response.json())
required_action: Final = res.get("required_action")
action_type: Final = required_action and required_action.get("action_type", None)
if action_type is None:
@ -358,7 +371,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
if original_messages:
data["messages"] = [
(
{**original, "content": redacted_messages[idx]["content"]}
{**original, "content": redacted_messages[idx].get("content")}
if idx < len(redacted_messages) and redacted_messages[idx].get("content") is not None
else original
)
@ -449,7 +462,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
timeout=self.timeout,
)
response.raise_for_status()
res: Final[_CatoAnalyzeResponse] = response.json()
res: Final = _CATO_ANALYZE_RESPONSE_ADAPTER.validate_python(response.json())
required_action: Final = res.get("required_action")
action_type: Final = required_action and required_action.get("action_type", None)
if action_type == "block_action" and required_action is not None:
@ -637,7 +650,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
try:
while True:
raw_message = await self._await_cato_message(websocket, sender)
result: _CatoStreamMessage = json.loads(raw_message)
result = _CATO_STREAM_MESSAGE_ADAPTER.validate_python(json.loads(raw_message))
if verified_chunk := result.get("verified_chunk"):
yield ModelResponseStream.model_validate(verified_chunk)
continue

View file

@ -11,7 +11,7 @@ from collections.abc import Callable, Mapping, Sequence
from typing import TYPE_CHECKING, Final, Literal, Optional
import httpx
from pydantic import JsonValue, TypeAdapter, ValidationError
from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError
from pydantic_core import PydanticSerializationError
from litellm._logging import verbose_proxy_logger
@ -45,6 +45,7 @@ _CHAT_ROW: Final[TypeAdapter[AllMessageValues]] = TypeAdapter(AllMessageValues)
_SENT_ROW: Final = TypeAdapter(Mapping[str, object])
_SENT_ROWS: Final = TypeAdapter(Sequence[Mapping[str, object]])
_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
# Headers whose values are forwarded as-is (case-insensitive). Glob patterns supported (e.g. x-stainless-*, x-litellm*).
_HEADER_VALUE_ALLOWLIST: Final = frozenset(
@ -600,7 +601,7 @@ class GenericGuardrailAPI(CustomGuardrail):
)
response.raise_for_status()
response_json: Final = response.json()
response_json: Final = _JSON_OBJECT.validate_python(response.json())
verbose_proxy_logger.debug("Generic Guardrail API response: %s", response_json)

View file

@ -6,7 +6,8 @@ from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final
import httpx
from typing_extensions import NotRequired, TypedDict
from pydantic import ConfigDict, TypeAdapter, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
@ -35,6 +36,16 @@ class GraphTokenResponse(TypedDict):
expires_in: NotRequired[int]
@with_config(ConfigDict(extra="allow", strict=True))
class _GraphTokenPayload(TypedDict):
access_token: ReadOnly[str]
expires_in: ReadOnly[NotRequired[int | float | str]]
_GRAPH_TOKEN_RESPONSE_ADAPTER: Final = TypeAdapter(_GraphTokenPayload)
_GRAPH_RESPONSE_ADAPTER: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True))
class PurviewGuardrailBase:
"""
Base class for Microsoft Purview guardrails.
@ -112,7 +123,7 @@ class PurviewGuardrailBase:
timeout=self.timeout,
)
response.raise_for_status()
token_data: Final[GraphTokenResponse] = response.json()
token_data: Final = _GRAPH_TOKEN_RESPONSE_ADAPTER.validate_python(response.json())
access_token: Final = token_data["access_token"]
expires_in: Final = int(token_data.get("expires_in", 3599))
# Recompute ``now`` after the await so the expiry reflects when the
@ -148,7 +159,7 @@ class PurviewGuardrailBase:
verbose_proxy_logger.debug("Purview Graph POST %s", url)
response: Final = await self.async_handler.post(url=url, headers=headers, json=json_body, timeout=self.timeout)
response.raise_for_status()
response_json: Final[dict[str, object]] = response.json()
response_json: Final = _GRAPH_RESPONSE_ADAPTER.validate_python(response.json())
response_headers: Final = dict(response.headers)
verbose_proxy_logger.debug("Purview Graph response: %s", response_json)
return response_json, response_headers

View file

@ -67,6 +67,8 @@ ToolCallLike: TypeAlias = (
| ChatCompletionToolCallChunk
)
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True))
class _ToolCallFunctionSlice(LiteLLMBaseModel):
model_config = ConfigDict(from_attributes=True, extra="ignore")
@ -452,7 +454,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
)
response.raise_for_status()
result: Final[dict[str, object]] = response.json()
result: Final = _JSON_OBJECT.validate_python(response.json())
# Validate response format
if "action" not in result:

View file

@ -46,7 +46,7 @@ from fastapi import (
status,
)
from fastapi.responses import JSONResponse
from pydantic import TypeAdapter
from pydantic import ConfigDict, TypeAdapter, with_config
from typing_extensions import ReadOnly, TypedDict
from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, public_catalog_operation
@ -123,6 +123,22 @@ def does_mcp_server_exist(mcp_server_records: Iterable[_HasServerId], mcp_server
DEFAULT_MCP_REGISTRY_VERSION: Final = "1.0.0"
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _MCPRegistryServer(TypedDict, total=False):
name: ReadOnly[str]
title: ReadOnly[str]
description: ReadOnly[str]
category: ReadOnly[str]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _MCPRegistry(TypedDict, total=False):
servers: ReadOnly[Sequence[_MCPRegistryServer]]
_MCP_REGISTRY_ADAPTER: Final = TypeAdapter(_MCPRegistry)
if TYPE_CHECKING:
from prisma import models as prisma_models
@ -3417,19 +3433,23 @@ if MCP_AVAILABLE:
"mcp_registry.json",
)
_mcp_registry_cache: Mapping[str, Sequence[Mapping[str, str]]] | None = None
_mcp_registry_cache: _MCPRegistry | None = None
def _load_mcp_registry() -> Mapping[str, Sequence[Mapping[str, str]]]:
def _load_mcp_registry() -> _MCPRegistry:
"""Load the curated MCP registry from disk. Cached after first read."""
global _mcp_registry_cache
if _mcp_registry_cache is not None:
return _mcp_registry_cache
try:
with open(_MCP_REGISTRY_PATH, "r") as f:
data: Mapping[str, Sequence[Mapping[str, str]]] = json.load(f)
except Exception as e:
verbose_proxy_logger.warning("Failed to load MCP registry from %s: %s", _MCP_REGISTRY_PATH, e)
data = {"servers": []}
def _read_registry() -> _MCPRegistry:
try:
with open(_MCP_REGISTRY_PATH, "r") as f:
return _MCP_REGISTRY_ADAPTER.validate_python(json.load(f))
except Exception as e:
verbose_proxy_logger.warning("Failed to load MCP registry from %s: %s", _MCP_REGISTRY_PATH, e)
return {"servers": []}
data: Final = _read_registry()
_mcp_registry_cache = data
return data

View file

@ -147,6 +147,35 @@ class _MessageCreateData(TypedDict, total=False):
session_id: str
class _EventTable(Protocol):
async def create(self, *, data: _EventCreateData) -> object: ...
class _RunTable(Protocol):
async def update(self, *, where: Mapping[str, str], data: Mapping[str, str]) -> object: ...
class _WorkflowTx(Protocol):
@property
def litellm_workflowevent(self) -> _EventTable: ...
@property
def litellm_workflowrun(self) -> _RunTable: ...
async def _append_event_in_tx(
tx: _WorkflowTx, run_id: str, event_data: _EventCreateData, new_status: str | None
) -> object:
"""Create the event and, when the event type maps to a status, move the run to it in the same transaction."""
event: Final = await tx.litellm_workflowevent.create(data=event_data)
if new_status:
await tx.litellm_workflowrun.update(
where={"run_id": run_id},
data={"status": new_status},
)
return event
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
@ -387,12 +416,7 @@ async def append_workflow_event(
event_data["data"] = _json(data.data)
async with prisma_client.db.tx() as tx:
event: object = await tx.litellm_workflowevent.create(data=event_data)
if new_status:
await tx.litellm_workflowrun.update(
where={"run_id": run_id},
data={"status": new_status},
)
event: object = await _append_event_in_tx(tx, run_id, event_data, new_status)
return event

View file

@ -22,7 +22,7 @@ import json
import random
import re
import time
from collections.abc import Callable, Iterator, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from itertools import accumulate, chain, islice, takewhile
from threading import Lock
from types import MappingProxyType
@ -30,6 +30,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
from pydantic import BaseModel, TypeAdapter, ValidationError, create_model
from pydantic_core import ErrorDetails
from typing_extensions import ReadOnly, TypedDict
from litellm._internal_context import with_service_target
from litellm._logging import verbose_router_logger
@ -411,9 +412,19 @@ def _append_custom_keywords(base_keywords: list[str], custom_keywords: list[str]
return [*base_keywords, *deduped_custom.values()]
def _parent_session_kwargs(request_kwargs: Mapping[str, object] | None) -> Mapping[str, Any]:
class _ParentSessionKwargs(TypedDict, total=False):
litellm_session_id: ReadOnly[object]
litellm_trace_id: ReadOnly[object]
def _parent_session_kwargs(request_kwargs: Mapping[str, object] | None) -> _ParentSessionKwargs:
kwargs: Final = request_kwargs or {}
return {k: kwargs[k] for k in ("litellm_session_id", "litellm_trace_id") if kwargs.get(k) is not None}
session_id: Final = kwargs.get("litellm_session_id")
trace_id: Final = kwargs.get("litellm_trace_id")
return _ParentSessionKwargs(
**({"litellm_session_id": session_id} if session_id is not None else {}),
**({"litellm_trace_id": trace_id} if trace_id is not None else {}),
)
def _response_cost_or_none(response: object) -> float | None:
@ -2828,7 +2839,7 @@ class ComplexityRouter(CustomLogger):
"originating_request_masked": masked_originating_request(request_kwargs),
"body": {"model": llm_config.model, **payload},
}
classify: Final = (
classify: Final[Callable[..., Awaitable[ModelResponse | ResponsesAPIResponse]]] = (
self.litellm_router_instance.aresponses
if encrypted_task is not None
else self.litellm_router_instance.acompletion

View file

@ -90,7 +90,7 @@ class SearchAPIRouter:
def get_matching_search_tools(
router_instance: _SearchToolsRouter,
search_tool_name: str,
) -> list:
) -> list[SearchToolTypedDict]:
"""
Get all search tools matching the given name.

View file

@ -6,7 +6,8 @@ import traceback
from typing import Final
import httpx
from pydantic import TypeAdapter, ValidationError
from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
@ -30,6 +31,14 @@ class _OidcTokenClaims(LiteLLMBaseModel):
exp: float | None = None
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _GitHubOidcTokenResponse(TypedDict):
value: NotRequired[ReadOnly[str | None]]
_GITHUB_OIDC_TOKEN_RESPONSE: Final = TypeAdapter(_GitHubOidcTokenResponse)
def _oidc_token_cache_ttl(oidc_token: str, max_ttl: int) -> int:
"""Cache TTL for a fetched OIDC token: ``max_ttl``, capped so the cache entry
never outlives the token's own ``exp`` claim (minus a safety margin).
@ -271,7 +280,8 @@ def get_secret(
},
)
if response.status_code == 200:
oidc_token = response.json().get("value", None)
response_data: Final = _GITHUB_OIDC_TOKEN_RESPONSE.validate_python(response.json())
oidc_token = response_data.get("value", None)
oidc_cache.set_cache(key=secret_name, value=oidc_token, ttl=300 - 5)
return oidc_token
else: