refactor(completion): extract provider dispatch into typed helpers so basedpyright can analyze it (#30813)

litellm/types/completion.py

@@ -1,8 +1,28 @@
-from typing import Iterable, List, Optional, Union
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Coroutine,
+    Iterable,
+    List,
+    Optional,
+    Union,
+)
 
 from pydantic import BaseModel, ConfigDict
 from typing_extensions import Literal, Required, TypedDict
 
+if TYPE_CHECKING:
+    import httpx
+    from aiohttp import ClientSession
+
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+    from litellm.llms.base_llm import BaseConfig
+    from litellm.utils import CustomStreamWrapper, ModelResponse
+
 
 class ChatCompletionSystemMessageParam(TypedDict, total=False):
     content: Required[str]
@@ -191,3 +211,44 @@ class CompletionRequest(BaseModel):
     model_list: Optional[List[str]] = None
 
     model_config = ConfigDict(protected_namespaces=(), extra="allow")
+
+
+@dataclass(frozen=True, slots=True)
+class _CompletionDispatchContext:
+    _azure_detection_model: str
+    acompletion: bool
+    api_base: Optional[str]
+    api_key: Optional[str]
+    api_version: Optional[str]
+    client: Any
+    custom_llm_provider: str
+    custom_prompt_dict: dict
+    extra_headers: Optional[dict]
+    headers: dict
+    hf_model_name: Optional[str]
+    kwargs: dict
+    litellm_params: dict
+    logger_fn: Optional[Callable]
+    logging: LiteLLMLoggingObj
+    max_retries: Optional[int]
+    max_tokens: Optional[int]
+    messages: list
+    metadata: Optional[dict]
+    model: str
+    model_response: ModelResponse
+    optional_params: dict
+    organization: Optional[str]
+    provider_config: Optional[BaseConfig]
+    shared_session: Optional[ClientSession]
+    stream: Optional[bool]
+    temperature: Optional[float]
+    text_completion: bool
+    timeout: Optional[Union[float, str, httpx.Timeout]]
+    top_p: Optional[float]
+
+
+_CompletionDispatchResult = Union[
+    Coroutine[Any, Any, Union["ModelResponse", "CustomStreamWrapper"]],
+    "ModelResponse",
+    "CustomStreamWrapper",
+]

ruff-strict-budget.json

@@ -300,7 +300,7 @@
     "slack": 3
   },
   "RET504": {
-    "baseline": 709,
+    "baseline": 702,
     "slack": 20
   },
   "RUF010": {

tests/test_litellm/types/test_completion.py

@@ -8,9 +8,16 @@ Usage:
     pytest tests/test_litellm/types/test_completion.py -v
 """
 
+import dataclasses
 from typing import List
 
-from litellm.types.completion import CompletionRequest, ChatCompletionMessageParam
+import pytest
+
+from litellm.types.completion import (
+    ChatCompletionMessageParam,
+    CompletionRequest,
+    _CompletionDispatchContext,
+)
 
 
 def test_completion_request_messages_type_validation():
@@ -146,3 +153,55 @@ def test_completion_request_with_all_params():
     assert request.presence_penalty == 0.0
     assert request.stream is False
     assert request.n == 1
+
+
+def _build_dispatch_context() -> _CompletionDispatchContext:
+    return _CompletionDispatchContext(
+        _azure_detection_model="gpt-4o",
+        acompletion=False,
+        api_base=None,
+        api_key=None,
+        api_version=None,
+        client=None,
+        custom_llm_provider="openai",
+        custom_prompt_dict={},
+        extra_headers=None,
+        headers={},
+        hf_model_name=None,
+        kwargs={},
+        litellm_params={},
+        logger_fn=None,
+        logging=None,  # type: ignore[arg-type]
+        max_retries=None,
+        max_tokens=None,
+        messages=[],
+        metadata=None,
+        model="gpt-4o",
+        model_response=None,  # type: ignore[arg-type]
+        optional_params={},
+        organization=None,
+        provider_config=None,
+        shared_session=None,
+        stream=None,
+        temperature=None,
+        text_completion=False,
+        timeout=None,
+        top_p=None,
+    )
+
+
+def test_dispatch_context_is_frozen():
+    """A helper must not be able to re-route the call by rebinding a dispatch
+    input mid-flight; this pins the frozen invariant the dispatch shape relies on."""
+    ctx = _build_dispatch_context()
+    with pytest.raises(dataclasses.FrozenInstanceError):
+        ctx.model = "claude-haiku-4-5"  # type: ignore[misc]
+    with pytest.raises(dataclasses.FrozenInstanceError):
+        ctx.custom_llm_provider = "anthropic"  # type: ignore[misc]
+
+
+def test_dispatch_context_uses_slots():
+    """slots=True keeps the per-call context lightweight (no per-instance __dict__)."""
+    ctx = _build_dispatch_context()
+    assert not hasattr(ctx, "__dict__")
+    assert hasattr(type(ctx), "__slots__")