mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
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:
parent
e979bd0ecf
commit
638cdc0c82
39 changed files with 457 additions and 187 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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], ...])
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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]]:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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__ = [
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue