From b88f64e65f29a34ee7363a3e2cd8189e850e4387 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:17:07 +0000 Subject: [PATCH] refactor(types): replace Any with proven types in 6 files (#47268) * refactor(types): replace Any with proven types in 39 files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(anthropic): use valid response fixtures for strict parsing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(types): revert anthropic chat response seam Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(types): revert validation seams that changed tolerated inputs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(types): revert remaining validation seams to keep behavior unchanged Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(types): drop no-op search router annotation and helper docstring Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(types): keep ip address utils at base Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../integrations/datadog/datadog_llm_obs.py | 24 ++++++++++--- litellm/llms/base_llm/chat/transformation.py | 2 +- litellm/models/organization_membership.py | 3 +- litellm/proxy/auth/auth_checks.py | 2 +- .../workflow_management_endpoints.py | 35 +++++++++++++++---- .../complexity_router/complexity_router.py | 19 +++++++--- 6 files changed, 66 insertions(+), 19 deletions(-) diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 270d799742c..a72023eaedf 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -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) diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 51c8e4b8bef..b2eb4cde5ce 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -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( diff --git a/litellm/models/organization_membership.py b/litellm/models/organization_membership.py index e9697cf57b6..6a0d22cb63b 100644 --- a/litellm/models/organization_membership.py +++ b/litellm/models/organization_membership.py @@ -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 diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 5a09a0efbf6..8887d86570a 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2663,7 +2663,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 diff --git a/litellm/proxy/management_endpoints/workflow_management_endpoints.py b/litellm/proxy/management_endpoints/workflow_management_endpoints.py index 890a344dfe0..649941bc233 100644 --- a/litellm/proxy/management_endpoints/workflow_management_endpoints.py +++ b/litellm/proxy/management_endpoints/workflow_management_endpoints.py @@ -147,6 +147,34 @@ 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: + 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 +415,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 diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 1638f20c8e3..a475358e1f0 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -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