refactor(native): move execute_chat and the chat planner out of main.py

litellm/llms/anthropic/pass_through/adapters/handler.py

@@ -40,7 +40,7 @@ from litellm.llms.anthropic.pass_through.utils import (
     proxy_sub_call_logging_kwargs,
     proxy_sub_call_trace_id,
 )
-from litellm.main import aexecute_planned_chat, execute_planned_chat, plan_chat
+from litellm.native.chat_plan import aexecute_planned_chat, execute_planned_chat, plan_chat
 from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS
 from litellm.types.llms.anthropic_messages.anthropic_response import (
     AnthropicMessagesResponse,

litellm/native/chat_mock.py

@@ -0,0 +1,281 @@
+import asyncio
+import time
+from typing import (
+    Final,
+    cast,  # noqa: TID251  # casts moved verbatim from litellm/main.py
+)
+
+import httpx
+import openai
+
+import litellm
+from litellm.constants import DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT
+from litellm.llms.base_llm.base_model_iterator import convert_model_response_to_streaming
+from litellm.native.request_support import admitted_prompt_token_counter
+from litellm.types.utils import ChatCompletionMessageToolCall, ModelResponseStream
+from litellm.utils import (
+    CustomStreamWrapper,
+    ModelResponse,
+    Usage,
+    async_mock_completion_streaming_obj,
+    mock_completion_streaming_obj,
+)
+
+MOCK_RESPONSE_TYPE = str | Exception | dict | ModelResponse | ModelResponseStream
+
+
+def handle_mock_potential_exceptions(
+    mock_response: str | Exception,
+    model: str,
+    custom_llm_provider: str | None = None,
+) -> None:
+    if isinstance(mock_response, Exception):
+        if isinstance(mock_response, openai.APIError):
+            raise mock_response
+        raise litellm.MockException(
+            status_code=getattr(mock_response, "status_code", 500),
+            message=getattr(mock_response, "text", str(mock_response)),
+            llm_provider=getattr(mock_response, "llm_provider", custom_llm_provider or "openai"),
+            model=model,
+            request=httpx.Request(method="POST", url="https://api.openai.com/v1/"),
+        )
+    elif isinstance(mock_response, str) and mock_response == "litellm.RateLimitError":
+        raise litellm.RateLimitError(
+            message="this is a mock rate limit error",
+            llm_provider=getattr(mock_response, "llm_provider", custom_llm_provider or "openai"),
+            model=model,
+        )
+    elif isinstance(mock_response, str) and mock_response == "litellm.ContextWindowExceededError":
+        raise litellm.ContextWindowExceededError(
+            message="this is a mock context window exceeded error",
+            llm_provider=getattr(mock_response, "llm_provider", custom_llm_provider or "openai"),
+            model=model,
+        )
+    elif isinstance(mock_response, str) and mock_response == "litellm.InternalServerError":
+        raise litellm.InternalServerError(
+            message="this is a mock internal server error",
+            llm_provider=getattr(mock_response, "llm_provider", custom_llm_provider or "openai"),
+            model=model,
+        )
+    elif isinstance(mock_response, str) and mock_response.startswith("Exception: content_filter_policy"):
+        raise litellm.MockException(
+            status_code=400,
+            message=mock_response,
+            llm_provider="azure",
+            model=model,
+            request=httpx.Request(method="POST", url="https://api.openai.com/v1/"),
+        )
+
+
+def handle_mock_timeout(
+    mock_timeout: bool | None,
+    timeout: float | str | httpx.Timeout | None,
+    model: str,
+) -> None:
+    if mock_timeout is True and timeout is not None:
+        _sleep_for_timeout(timeout)
+        raise litellm.Timeout(
+            message="This is a mock timeout error",
+            llm_provider="openai",
+            model=model,
+        )
+
+
+async def handle_mock_timeout_async(
+    mock_timeout: bool | None,
+    timeout: float | str | httpx.Timeout | None,
+    model: str,
+) -> None:
+    if mock_timeout is True and timeout is not None:
+        await _sleep_for_timeout_async(timeout)
+        raise litellm.Timeout(
+            message="This is a mock timeout error",
+            llm_provider="openai",
+            model=model,
+        )
+
+
+def _sleep_for_timeout(timeout: float | str | httpx.Timeout):
+    if isinstance(timeout, float):
+        time.sleep(timeout)
+    elif isinstance(timeout, str):
+        time.sleep(float(timeout))
+    elif isinstance(timeout, httpx.Timeout) and timeout.connect is not None:
+        time.sleep(timeout.connect)
+
+
+async def _sleep_for_timeout_async(timeout: float | str | httpx.Timeout):
+    if isinstance(timeout, float):
+        await asyncio.sleep(timeout)
+    elif isinstance(timeout, str):
+        await asyncio.sleep(float(timeout))
+    elif isinstance(timeout, httpx.Timeout) and timeout.connect is not None:
+        await asyncio.sleep(timeout.connect)
+
+
+def mock_completion(
+    model: str,
+    messages: list,
+    stream: bool | None = False,
+    n: int | None = None,
+    mock_response: MOCK_RESPONSE_TYPE | None = "This is a mock request",
+    mock_tool_calls: list | None = None,
+    mock_timeout: bool | None = False,
+    logging=None,
+    custom_llm_provider=None,
+    timeout: float | str | httpx.Timeout | None = None,
+    prompt_tokens: int | None = None,
+    **kwargs,
+):
+    """
+    Generate a mock completion response for testing or debugging purposes.
+
+    This is a helper function that simulates the response structure of the OpenAI completion API.
+
+    Parameters:
+        model (str): The name of the language model for which the mock response is generated.
+        messages (List): A list of message objects representing the conversation context.
+        stream (bool, optional): If True, returns a mock streaming response (default is False).
+        mock_response (str, optional): The content of the mock response (default is "This is a mock request").
+        mock_timeout (bool, optional): If True, the mock response will be a timeout error (default is False).
+        timeout (float, optional): The timeout value to use for the mock response (default is None).
+        **kwargs: Additional keyword arguments that can be used but are not required.
+
+    Returns:
+        litellm.ModelResponse: A ModelResponse simulating a completion response with the specified model, messages, and mock response.
+
+    Raises:
+        Exception: If an error occurs during the generation of the mock completion response.
+    Note:
+        - This function is intended for testing or debugging purposes to generate mock completion responses.
+        - If 'stream' is True, it returns a response that mimics the behavior of a streaming completion.
+    """
+    try:
+        is_acompletion: Final = kwargs.get("acompletion") or False
+        if mock_response is None:
+            mock_response = "This is a mock request"
+
+        handle_mock_timeout(mock_timeout=mock_timeout, timeout=timeout, model=model)
+
+        ## LOGGING
+        if logging is not None:
+            logging.pre_call(
+                input=messages,
+                api_key="mock-key",
+            )
+
+        if isinstance(mock_response, str) or isinstance(mock_response, Exception):
+            handle_mock_potential_exceptions(
+                mock_response=mock_response,
+                model=model,
+                custom_llm_provider=custom_llm_provider,
+            )
+
+        mock_response = cast(
+            str | dict | ModelResponse | ModelResponseStream, mock_response
+        )  # after this point, mock_response is a string, dict, ModelResponse, or ModelResponseStream
+        if isinstance(mock_response, str) and mock_response.startswith("Exception: mock_streaming_error"):
+            mock_response = litellm.MockException(
+                message="This is a mock error raised mid-stream",
+                llm_provider="anthropic",
+                model=model,
+                status_code=529,
+            )
+        time_delay: Final = kwargs.get("mock_delay", None)
+        if time_delay is not None and not is_acompletion:
+            time.sleep(time_delay)
+
+        if isinstance(mock_response, dict):
+            return ModelResponse(**mock_response)
+
+        if isinstance(mock_response, ModelResponse):
+            if not stream:
+                return mock_response
+            # convert to ModelResponseStream
+            mock_response = convert_model_response_to_streaming(mock_response)
+
+        model_response: ModelResponse | ModelResponseStream = ModelResponse()
+
+        if stream is True:
+            model_response = ModelResponseStream()
+            count_prompt_tokens: Final = admitted_prompt_token_counter(prompt_tokens)
+            # don't try to access stream object,
+            if kwargs.get("acompletion", False) is True:
+                return CustomStreamWrapper(
+                    completion_stream=async_mock_completion_streaming_obj(
+                        model_response, mock_response=mock_response, model=model, n=n, prompt_tokens=prompt_tokens
+                    ),
+                    model=model,
+                    custom_llm_provider="openai",
+                    logging_obj=logging,
+                    count_prompt_tokens=count_prompt_tokens,
+                )
+            return CustomStreamWrapper(
+                completion_stream=mock_completion_streaming_obj(
+                    model_response, mock_response=mock_response, model=model, n=n, prompt_tokens=prompt_tokens
+                ),
+                model=model,
+                custom_llm_provider="openai",
+                logging_obj=logging,
+                count_prompt_tokens=count_prompt_tokens,
+            )
+        if isinstance(mock_response, litellm.MockException):
+            raise mock_response
+        # At this point, mock_response must be a string (all other types have been handled or returned early)
+        mock_response = cast(str, mock_response)
+
+        if n is None:
+            model_response.choices[0].message.content = mock_response
+        else:
+            _all_choices: Final = []
+            for i in range(n):
+                _choice = litellm.utils.Choices(
+                    index=i,
+                    message=litellm.utils.Message(content=mock_response, role="assistant"),
+                )
+                _all_choices.append(_choice)
+            model_response.choices = _all_choices
+        model_response.created = int(time.time())
+        model_response.model = model
+
+        if mock_tool_calls:
+            model_response.choices[0].message.tool_calls = [
+                ChatCompletionMessageToolCall(**tool_call) for tool_call in mock_tool_calls
+            ]
+
+        usage_prompt_tokens: Final = (
+            prompt_tokens if prompt_tokens is not None else DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT
+        )
+        setattr(
+            model_response,
+            "usage",
+            Usage(
+                prompt_tokens=usage_prompt_tokens,
+                completion_tokens=DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
+                total_tokens=usage_prompt_tokens + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
+            ),
+        )
+
+        if custom_llm_provider is not None:
+            model_response.hidden_params["custom_llm_provider"] = custom_llm_provider
+        else:
+            try:
+                _, inferred_provider, _, _ = litellm.utils.get_llm_provider(model=model)
+                model_response.hidden_params["custom_llm_provider"] = inferred_provider
+            except Exception:
+                # dont let setting a hidden param block a mock_respose
+                pass
+
+        if logging is not None:
+            logging.post_call(
+                input=messages,
+                api_key="my-secret-key",
+                original_response="my-original-response",
+            )
+
+        return model_response
+
+    except Exception as e:
+        if isinstance(e, openai.APIError):
+            raise e
+        raise Exception(f"Mock completion response failed - {e}")

litellm/native/chat_plan.py

@@ -0,0 +1,277 @@
+import asyncio
+import contextvars
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass
+from functools import partial
+from types import MappingProxyType
+from typing import (
+    Final,
+    cast,  # noqa: TID251  # casts moved verbatim from litellm/main.py
+)
+
+from pydantic import TypeAdapter
+
+import litellm
+from litellm.litellm_core_utils.chat_completion_agentic_loop import maybe_run_chat_completion_agentic_loop
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.litellm_core_utils.prompt_templates.common_utils import get_completion_messages
+from litellm.litellm_core_utils.rules import Rules
+from litellm.native.chat import OPTIONAL_PARAMS_ADAPTER, execute_chat, resolve_dispatched_chat_response
+from litellm.native.request_support import resolve_control_options
+from litellm.types.litellm_params import ControlOptions
+from litellm.types.llms.anthropic import AnthropicThinkingParam
+from litellm.utils import (
+    CustomStreamWrapper,
+    ModelResponse,
+    exception_type,
+    get_llm_provider,
+    get_non_default_completion_params,
+    post_call_processing,
+    validate_and_fix_openai_messages,
+    validate_and_fix_openai_tools,
+    validate_and_fix_thinking_param,
+    validate_chat_completion_tool_choice,
+    validate_openai_optional_params,
+)
+
+_PLANNED_CHAT_ARGS: Final = frozenset(
+    ("kwargs", "non_default_params", "control_options", "prompt_id", "prompt_variables")
+)
+
+
+_EXECUTE_CHAT_CODE: Final = execute_chat.__code__
+
+
+_EXECUTE_CHAT_PARAMS: Final = (
+    frozenset(_EXECUTE_CHAT_CODE.co_varnames[: _EXECUTE_CHAT_CODE.co_kwonlyargcount]) - _PLANNED_CHAT_ARGS
+)
+
+
+_SKIP_RESPONSES_API_BRIDGE_KWARG: Final = "_skip_responses_api_bridge"
+
+
+_CHAT_SURFACE_KWARGS: Final = (
+    "fallbacks",
+    "model_list",
+    "prompt_id",
+    "litellm_system_prompt",
+    "ensure_alternating_roles",
+    "user_continue_message",
+    "assistant_continue_message",
+    "mock_delay",
+    "mock_timeout",
+)
+
+
+_TEXT_COMPLETION_PROVIDERS: Final = frozenset(
+    ("text-completion-openai", "text-completion-codestral", "text-completion-inception")
+)
+
+
+@dataclass(frozen=True, slots=True)
+class PlannedChat:
+    model: str
+    custom_llm_provider: str
+    named: Mapping[str, object]
+    kwargs: Mapping[str, object]
+    non_default_params: Mapping[str, object]
+    control_options: ControlOptions
+
+
+def _optional_str(value: object) -> str | None:
+    return value if isinstance(value, str) else None
+
+
+_OBJECT_LIST: Final = TypeAdapter(list[object])
+
+
+def _object_list(value: object) -> list[object] | None:  # mutable-ok: feeds validate_and_fix_openai_* which take a list
+    return _OBJECT_LIST.validate_python(value) if isinstance(value, list) else None
+
+
+def _tool_choice(value: object) -> str | dict[str, object] | None:  # mutable-ok: its validator takes a dict
+    if isinstance(value, str):
+        return value
+    return OPTIONAL_PARAMS_ADAPTER.validate_python(value) if isinstance(value, dict) else None
+
+
+def _stop(value: object) -> str | list[str] | None:  # mutable-ok: validate_openai_optional_params takes a list
+    if isinstance(value, list):
+        return [item for item in _OBJECT_LIST.validate_python(value) if isinstance(item, str)]
+    return value if isinstance(value, str) else None
+
+
+def _thinking(value: object) -> AnthropicThinkingParam | bool | None:
+    if isinstance(value, dict):
+        return cast(AnthropicThinkingParam, value)  # cast-ok: the caller's thinking dict, normalized by the validator
+    return value if isinstance(value, bool) else None
+
+
+def _chat_entry_points_wrapped() -> bool:
+    """Tracers such as ddtrace wrap `litellm.completion`/`acompletion`, so bridged calls must go through them."""
+    from litellm.chat_completions import dispatch
+
+    return litellm.completion is not dispatch.completion or litellm.acompletion is not dispatch.acompletion
+
+
+def _uses_chat_surface(request: Mapping[str, object]) -> bool:
+    from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler
+    from litellm.types.llms.openai import ToolParam
+
+    tools: Final = request.get("tools")
+    return (
+        litellm.model_fallbacks is not None
+        or _chat_entry_points_wrapped()
+        or any(request.get(key) is not None for key in _CHAT_SURFACE_KWARGS)
+        or (
+            isinstance(tools, list)
+            and LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway(
+                tools=cast(Iterable[ToolParam], tools)  # cast-ok: OpenAI-format tools from the caller's request
+            )
+        )
+    )
+
+
+def plan_chat(request: Mapping[str, object], logging_obj: LiteLLMLoggingObj) -> PlannedChat | None:
+    model: Final = request.get("model")
+    if not isinstance(model, str) or _uses_chat_surface(request):
+        return None
+    from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
+
+    custom_llm_provider: Final = _optional_str(request.get("custom_llm_provider"))
+    api_base: Final = _optional_str(request.get("api_base"))
+    _, resolved_provider, _, _ = get_llm_provider(
+        model=model, custom_llm_provider=custom_llm_provider, api_base=api_base
+    )
+    if resolved_provider in _TEXT_COMPLETION_PROVIDERS:
+        return None
+    kwargs: Final = {
+        key: value
+        for key, value in request.items()
+        if key not in _EXECUTE_CHAT_PARAMS and key != _SKIP_RESPONSES_API_BRIDGE_KWARG
+    }
+    tools: Final = validate_and_fix_openai_tools(tools=_object_list(request.get("tools")))
+    messages: Final = get_completion_messages(
+        messages=validate_and_fix_openai_messages(messages=_object_list(request.get("messages")) or []),
+        ensure_alternating_roles=False,
+        user_continue_message=None,
+        assistant_continue_message=None,
+    )
+    non_default_params: Final = get_non_default_completion_params(kwargs=kwargs)
+    AnthropicCacheControlHook.maybe_seed_default_injection_points(
+        non_default_params=non_default_params,
+        messages=messages,
+        model=model,
+        custom_llm_provider=custom_llm_provider,
+        tools=tools,
+        enable_prompt_caching=request.get("enable_prompt_caching") is True,
+        api_base=api_base,
+    )
+    if logging_obj.should_run_prompt_management_hooks(non_default_params=non_default_params, tools=tools):
+        return None
+    named: Final = {key: value for key, value in request.items() if key in _EXECUTE_CHAT_PARAMS}
+    tool_choice: Final = validate_chat_completion_tool_choice(
+        tool_choice=_tool_choice(request.get("tool_choice")), model=model
+    )
+    return PlannedChat(
+        model=model,
+        custom_llm_provider=resolved_provider,
+        named=MappingProxyType(
+            {
+                **named,
+                "messages": messages,
+                "tools": tools,
+                "tool_choice": tool_choice,
+                "stop": validate_openai_optional_params(stop=_stop(request.get("stop"))),
+                "thinking": validate_and_fix_thinking_param(thinking=_thinking(request.get("thinking"))),
+                "skip_responses_api_bridge": request.get(_SKIP_RESPONSES_API_BRIDGE_KWARG) is True,
+            }
+        ),
+        kwargs=MappingProxyType(kwargs),
+        non_default_params=MappingProxyType(non_default_params),
+        control_options=resolve_control_options(kwargs, model),
+    )
+
+
+def _run_planned_chat(plan: PlannedChat, is_async: bool) -> ModelResponse | CustomStreamWrapper:
+    return execute_chat(
+        **plan.named,  # pyright: ignore[reportArgumentType]  # named holds the caller's validated OpenAI-format params
+        kwargs={**plan.kwargs, "acompletion": is_async},
+        non_default_params=dict(plan.non_default_params),
+        control_options=plan.control_options,
+        prompt_id=None,
+        prompt_variables=None,
+    )
+
+
+def _post_call_checked(
+    plan: PlannedChat, response: ModelResponse | CustomStreamWrapper, is_async: bool
+) -> ModelResponse | CustomStreamWrapper:
+    """Runs the post-call rules and JSON-schema validation the ``completion`` client wrapper runs on its result."""
+    post_call_processing(
+        original_response=response,
+        model=plan.model,
+        optional_params={**plan.kwargs, **plan.named},
+        original_function=litellm.acompletion if is_async else litellm.completion,
+        rules_obj=Rules(),
+    )
+    return response
+
+
+def execute_planned_chat(plan: PlannedChat) -> ModelResponse | CustomStreamWrapper:
+    return _post_call_checked(plan, _run_planned_chat(plan, is_async=False), is_async=False)
+
+
+async def _settled_chat_response(
+    plan: PlannedChat, loop: asyncio.AbstractEventLoop
+) -> ModelResponse | CustomStreamWrapper:
+    pending: Final[object] = await loop.run_in_executor(
+        None, partial(contextvars.copy_context().run, partial(_run_planned_chat, plan, True))
+    )
+    if asyncio.iscoroutine(pending):
+        return await resolve_dispatched_chat_response(pending)
+    if isinstance(pending, (ModelResponse, CustomStreamWrapper)):
+        return pending
+    raise TypeError(f"chat provider returned {type(pending).__name__}, expected a ModelResponse or a stream")
+
+
+async def _with_agentic_loop(plan: PlannedChat, response: ModelResponse) -> ModelResponse | CustomStreamWrapper:
+    looped: Final = await maybe_run_chat_completion_agentic_loop(
+        response=response,
+        model=plan.model,
+        messages=_object_list(plan.named["messages"]) or [],
+        optional_params={
+            key: value
+            for key, value in plan.named.items()
+            if value is not None and key not in ("model", "messages", "stream", "deployment_id")
+        },
+        kwargs=dict(plan.kwargs),
+        logging_obj=plan.kwargs.get("litellm_logging_obj"),
+        custom_llm_provider=plan.custom_llm_provider,
+        stream=plan.named.get("stream") is True,
+    )
+    return response if looped is None else looped
+
+
+async def _aexecute_planned_chat(plan: PlannedChat) -> ModelResponse | CustomStreamWrapper:
+    loop: Final = asyncio.get_event_loop()
+    response: Final = await _settled_chat_response(plan, loop)
+    if isinstance(response, CustomStreamWrapper):
+        response.set_logging_event_loop(loop=loop)
+        return response
+    return await _with_agentic_loop(plan, response)
+
+
+async def aexecute_planned_chat(plan: PlannedChat) -> ModelResponse | CustomStreamWrapper:
+    """Maps a provider failure raised while awaiting the call, as ``acompletion`` does."""
+    try:
+        response: Final = await _aexecute_planned_chat(plan)
+    except Exception as e:  # noqa: BLE001  # every provider failure maps to the litellm exception acompletion raises
+        raise exception_type(
+            model=plan.model,
+            custom_llm_provider=plan.custom_llm_provider or "openai",
+            original_exception=e,
+            completion_kwargs=dict(plan.named),
+            extra_kwargs=dict(plan.kwargs),
+        ) from e
+    return _post_call_checked(plan, response, is_async=True)

litellm/native/request_support.py

@@ -0,0 +1,149 @@
+from collections.abc import Callable, Mapping
+from typing import Final, Protocol
+
+from typing_extensions import assert_never
+
+import litellm
+from litellm.litellm_core_utils.core_helpers import normalize_drop_params
+from litellm.litellm_core_utils.get_litellm_params import InvalidControlOption, parse_control_options
+from litellm.responses.bridge_check import keep_registered_mode
+from litellm.types.litellm_params import ControlOptions
+from litellm.types.llms.base import LiteLLMBaseModel
+from litellm.types.utils import CustomPricingLiteLLMParams
+
+
+def resolve_control_options(kwargs: Mapping[str, object], model: str) -> ControlOptions:
+    control: Final = parse_control_options(kwargs)
+    match control:
+        case ControlOptions():
+            return control
+        case InvalidControlOption(param=param, message=message):
+            if litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True:
+                return ControlOptions()
+            raise litellm.BadRequestError(message=message, model=model, llm_provider=None, body={"param": param})
+        case _:
+            return assert_never(control)
+
+
+class _AdmissionReservation(LiteLLMBaseModel):
+    input_tokens: int | None = None
+
+
+class _AdmissionMetadata(LiteLLMBaseModel):
+    user_api_key_budget_reservation: _AdmissionReservation | None = None
+
+
+def admission_input_tokens(kwargs: Mapping[str, object]) -> int | None:
+    reservations: Final = (
+        _AdmissionMetadata.model_validate(kwargs.get(key) or {}).user_api_key_budget_reservation
+        for key in ("litellm_metadata", "metadata")
+    )
+    return next(
+        (
+            reservation.input_tokens
+            for reservation in reservations
+            if reservation and reservation.input_tokens is not None
+        ),
+        None,
+    )
+
+
+def admitted_prompt_token_counter(prompt_tokens: int | None) -> Callable[[], int] | None:
+    if prompt_tokens is None:
+        return None
+    return lambda: prompt_tokens
+
+
+class _ProxyAuthHeadersProvider(Protocol):
+    def get_auth_headers(self) -> Mapping[str, str]: ...
+
+
+def proxy_auth_headers(proxy_auth: _ProxyAuthHeadersProvider) -> Mapping[str, str]:
+    return proxy_auth.get_auth_headers()
+
+
+def _build_custom_pricing_entry(
+    custom_llm_provider: str,
+    kwargs: dict,
+    model_info: dict | None = None,
+) -> dict:
+    """Build a complete model cost entry from kwargs and model_info.
+
+    Collects all CustomPricingLiteLLMParams fields present in kwargs and
+    merges metadata from model_info (mode, supports_prompt_caching, max_tokens)
+    so that register_model() receives the full pricing configuration.
+    """
+    entry: Final[dict] = {"litellm_provider": custom_llm_provider}
+
+    for field_name in CustomPricingLiteLLMParams.model_fields:
+        value = kwargs.get(field_name)
+        if value is not None:
+            entry[field_name] = value
+
+    if model_info and isinstance(model_info, dict):
+        for key in ("mode", "supports_prompt_caching", "max_tokens"):
+            if key in model_info and model_info[key] is not None:
+                entry.setdefault(key, model_info[key])
+
+    return entry
+
+
+def _get_router_deployment_id(kwargs: dict) -> str | None:
+    for metadata_key in ("litellm_metadata", "metadata"):
+        metadata = kwargs.get(metadata_key) or {}
+        if not isinstance(metadata, dict):
+            continue
+        deployment_model_info = metadata.get("model_info") or {}
+        if not isinstance(deployment_model_info, dict):
+            continue
+        deployment_id = deployment_model_info.get("id")
+        if deployment_id is not None:
+            return str(deployment_id)
+    return None
+
+
+def register_custom_pricing_for_request(
+    model: str,
+    custom_llm_provider: str,
+    kwargs: dict,
+    model_info: dict | None,
+) -> None:
+    """Register per-request custom pricing in litellm.model_cost.
+
+    Router-originated requests (identified by the deployment id the router puts
+    in metadata) get their full pricing registered under that unique id only;
+    the shared ``{provider}/{model}`` key receives the entry with pricing fields
+    stripped, mirroring Router._create_deployment. This keeps one deployment's
+    pricing overrides (e.g. a zero-cost wildcard) from clobbering built-in
+    pricing used by sibling deployments of the same backend model. Direct SDK
+    calls keep the legacy behavior of registering the shared key with pricing.
+    Either way the shared key's ``mode`` follows ``keep_registered_mode``, as on
+    deployment registration, since a call routes by its own ``model_info``.
+    """
+    entry: Final = _build_custom_pricing_entry(
+        custom_llm_provider=custom_llm_provider,
+        kwargs=kwargs,
+        model_info=model_info,
+    )
+    shared_key: Final = f"{custom_llm_provider}/{model}"
+    deployment_id: Final = _get_router_deployment_id(kwargs)
+    if deployment_id is None:
+        litellm.register_model({shared_key: keep_registered_mode(shared_key, entry)}, persist_across_reloads=False)
+        return
+    shared_entry: Final = keep_registered_mode(
+        shared_key, CustomPricingLiteLLMParams.strip_custom_pricing_fields(entry)
+    )
+    litellm.register_model(
+        {deployment_id: entry},
+        persist_across_reloads=False,
+        warning_display_name=shared_key,
+        custom_llm_provider=custom_llm_provider,
+        add_to_provider_model_sets=False,
+    )
+    litellm.register_model(
+        {shared_key: shared_entry},
+        persist_across_reloads=False,
+        warning_display_name=shared_key,
+        custom_llm_provider=custom_llm_provider,
+        add_to_provider_model_sets=False,
+    )

litellm/responses/litellm_completion_transformation/handler.py

@@ -24,14 +24,14 @@ from litellm.types.llms.openai import (
 from litellm.types.utils import ModelResponse
 
 if TYPE_CHECKING:
-    from litellm.main import PlannedChat
+    from litellm.native.chat_plan import PlannedChat
 
 _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(dict[str, object])
 
 
 def _plan_bridged_chat(request: Mapping[str, object]) -> "PlannedChat | None":
     from litellm.litellm_core_utils.litellm_logging import Logging
-    from litellm.main import plan_chat
+    from litellm.native.chat_plan import plan_chat
 
     logging_obj: Final = request.get("litellm_logging_obj")
     return plan_chat(request, logging_obj) if isinstance(logging_obj, Logging) else None
@@ -99,7 +99,7 @@ class LiteLLMCompletionTransformationHandler:
 
         completion_args: Final = {**kwargs, **litellm_completion_request, "_skip_responses_api_bridge": True}
 
-        from litellm.main import execute_planned_chat
+        from litellm.native.chat_plan import execute_planned_chat
 
         planned_chat: Final = _plan_bridged_chat(_REQUEST_KWARGS_ADAPTER.validate_python(completion_args))
         litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = (
@@ -154,7 +154,7 @@ class LiteLLMCompletionTransformationHandler:
 
         acompletion_args: Final = {**kwargs, **litellm_completion_request, "_skip_responses_api_bridge": True}
 
-        from litellm.main import aexecute_planned_chat
+        from litellm.native.chat_plan import aexecute_planned_chat
 
         planned_chat: Final = _plan_bridged_chat(_REQUEST_KWARGS_ADAPTER.validate_python(acompletion_args))
         litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = (

litellm/utils.py

@@ -8843,7 +8843,7 @@ def validate_and_fix_openai_tools(tools: list | None) -> list[dict[str, object]]
     """
     Ensure tools is List[dict] and not List[BaseModel]
     """
-    new_tools: Final = []
+    new_tools: Final[list[dict[str, object]]] = []
     if tools is None:
         return tools
     for tool in tools:
@@ -8964,7 +8964,7 @@ def validate_chat_completion_tool_choice(
     )
 
 
-def validate_openai_optional_params(stop: str | list[str] | None = None, **kwargs) -> str | list[str] | None:
+def validate_openai_optional_params(stop: str | list[str] | None = None, **kwargs: object) -> str | list[str] | None:
     """
     Validates and fixes OpenAI optional parameters.
 
@@ -10901,8 +10901,8 @@ def reasoning_effort_with_summary_alias(
 
 
 def strip_reasoning_summary_aliases_from_optional_params(
-    optional_params: dict,
-) -> tuple[dict, object | None]:
+    optional_params: Mapping[str, object],
+) -> tuple[dict[str, object], object | None]:
     """Copy optional_params; remove reasoningSummary aliases from top-level and extra_body."""
     op: Final = dict(optional_params)
     rs_val = op.pop("reasoningSummary", None)

tests/unit/litellm_core_utils/test_get_litellm_params.py

@@ -30,11 +30,11 @@ def _funnel_kwargs_completion_forwards(monkeypatch, wif_kwargs: dict[str, object
     from unittest.mock import MagicMock
 
     import litellm
-    import litellm.main as litellm_main
+    import litellm.native.chat as native_chat
 
-    spy = MagicMock(wraps=litellm_main.get_litellm_params)
+    spy = MagicMock(wraps=native_chat.get_litellm_params)
     monkeypatch.setattr(  # test-quality-ok: completion() has no injection seam for its kwargs funnel
-        litellm_main, "get_litellm_params", spy
+        native_chat, "get_litellm_params", spy
     )
     litellm.completion(
         model="anthropic/claude-sonnet-5",

tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py

@@ -88,7 +88,7 @@ async def test_aresponses_forwards_timeout_to_acompletion():
     and similar providers, with calls falling back to the provider SDK default
     (~600s for Anthropic).
     """
-    with patch("litellm.main.aexecute_planned_chat", new_callable=AsyncMock) as mock_chat:
+    with patch("litellm.native.chat_plan.aexecute_planned_chat", new_callable=AsyncMock) as mock_chat:
         mock_chat.return_value = ModelResponse(
             id="id",
             created=0,

tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_bridge.py

@@ -29,7 +29,7 @@ async def test_mock_basic_google_ai_studio_responses_api_with_tools():
         ],
     )
 
-    with patch("litellm.main.aexecute_planned_chat", new_callable=AsyncMock) as mock_chat:
+    with patch("litellm.native.chat_plan.aexecute_planned_chat", new_callable=AsyncMock) as mock_chat:
         mock_chat.return_value = mock_response
 
         request_model = "gemini/gemini-2.5-flash"

tests/unit/test_main.py

@@ -46,6 +46,7 @@ from litellm.litellm_core_utils.get_litellm_params import stored_control_options
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
 from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt
 from litellm.llms.custom_httpx.http_handler import HTTPHandler
+from litellm.native import chat as native_chat
 from litellm.types.litellm_params import ControlOptions
 from litellm.types.llms.openai import AllMessageValues, HttpxBinaryResponseContent
 from litellm.types.prompts.init_prompts import PromptSpec
@@ -390,7 +391,7 @@ def test_bedrock_latency_optimized_inference():
     ],
 )
 def test_is_claude_tool_target(custom_llm_provider: str, model: str, expected: bool):
-    assert litellm_main._is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model) is expected
+    assert native_chat._is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model) is expected
 
 
 @pytest.mark.parametrize("key", ["input_examples", "eager_input_streaming"])
@@ -400,7 +401,7 @@ def test_drop_anthropic_only_tool_keys_strips_tool_and_function_levels(key: str)
         "opaque_tool",
     ]
 
-    cleaned = litellm_main._drop_anthropic_only_tool_keys(tools=tools)
+    cleaned = native_chat._drop_anthropic_only_tool_keys(tools=tools)
 
     assert cleaned == [
         {"type": "function", "name": "example_tool", "function": {"name": "example_tool"}},
@@ -2242,7 +2243,7 @@ def test_completion_optional_params_base_model(
 
     Regression for #29618: overwriting ``model`` with a friendly ``base_model``
     label made Bedrock drop ``tools``/``tool_choice`` under ``drop_params``."""
-    with patch("litellm.main.get_optional_params") as mock_get_optional_params:
+    with patch("litellm.native.chat.get_optional_params") as mock_get_optional_params:
         mock_get_optional_params.return_value = MagicMock()
 
         import litellm

tests/unit/test_register_model_custom_pricing.py

@@ -20,7 +20,7 @@ import pytest
 
 import litellm
 from litellm.litellm_core_utils.get_model_cost_map import _expand_model_aliases, adopt_model_cost_map
-from litellm.main import _build_custom_pricing_entry
+from litellm.native.request_support import _build_custom_pricing_entry
 from litellm.utils import _invalidate_model_cost_lowercase_map
 
 
@@ -1677,7 +1677,7 @@ def test_register_model_scopes_builtin_match_to_the_given_provider():
 
 
 def test_per_request_custom_pricing_scopes_a_colliding_deployment_id_to_its_provider():
-    from litellm.main import _register_custom_pricing_for_request
+    from litellm.native.request_support import register_custom_pricing_for_request
 
     deployment_id: Final = "baseten/zai-org/glm-5.2"
     builtin_key: Final = "baseten/zai-org/GLM-5.2"
@@ -1686,7 +1686,7 @@ def test_per_request_custom_pricing_scopes_a_colliding_deployment_id_to_its_prov
     builtin_row_before: Final = copy.deepcopy(litellm.model_cost[builtin_key])
     openai_models_before: Final = frozenset(litellm.open_ai_chat_completion_models)
     try:
-        _register_custom_pricing_for_request(
+        register_custom_pricing_for_request(
             model="zai-org/GLM-5.2",
             custom_llm_provider="openai",
             kwargs={