litellm/litellm/passthrough/main.py

"""
This module is used to pass through requests to the LLM APIs.
"""

from __future__ import annotations

import asyncio
import contextvars
import weakref
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Coroutine, Generator, Iterator
from functools import partial
from types import TracebackType
from typing import Final, cast

import httpx
from httpx._types import CookieTypes, QueryParamTypes, RequestContent, RequestFiles

from litellm._logging import verbose_logger
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.passthrough.utils import CommonUtils
from litellm.utils import client

base_llm_http_handler = BaseLLMHTTPHandler()
from .utils import BasePassthroughUtils

_provider_errors: Final[weakref.WeakSet[httpx.HTTPStatusError]] = weakref.WeakSet()


def raise_for_provider_status(response: httpx.Response) -> None:
    """``response.raise_for_status()`` for the provider's response to a passthrough request, recording the raised
    error as the provider's own."""
    try:
        response.raise_for_status()
    except httpx.HTTPStatusError as error:
        _provider_errors.add(error)
        raise


def is_provider_error(error: httpx.HTTPStatusError) -> bool:
    """Whether ``error`` is the provider's own error response to a passthrough request."""
    return error in _provider_errors


async def _as_async_generator(iterable: AsyncIterator[bytes]) -> AsyncGenerator[bytes, bytes]:
    async for chunk in iterable:
        yield chunk


def _as_generator(iterable: Iterator[bytes]) -> Generator[bytes, bytes, None]:
    yield from iterable


class _SpendCollection:
    """Feeds relayed chunks to the provider's stream collector without letting spend tracking break the relay."""

    def __init__(self, provider_config: BasePassthroughConfig, litellm_logging_obj: LiteLLMLoggingObj) -> None:
        self.collector: Final[PassthroughStreamCollector] = provider_config.create_stream_collector(
            model=litellm_logging_obj.model,
            custom_llm_provider=litellm_logging_obj.model_call_details.get("custom_llm_provider", ""),
            endpoint=litellm_logging_obj.model_call_details.get("endpoint", ""),
        )
        self.chunk_count = 0
        self._failed = False

    def add(self, chunk: bytes) -> None:
        self.chunk_count += 1
        if self._failed:
            return
        try:
            self.collector.add(chunk)
        except Exception as e:  # noqa: BLE001 # Safe catch-all: spend tracking must never break the relayed stream
            self._failed = True
            verbose_logger.exception(
                "Passthrough spend-tracking collector failed; spend dropped for this stream: %s", e
            )

    @property
    def should_flush(self) -> bool:
        return self.chunk_count > 0 and not self._failed


class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]):
    def __init__(
        self,
        response: Awaitable[httpx.Response],
        litellm_logging_obj: LiteLLMLoggingObj,
        provider_config: BasePassthroughConfig,
    ) -> None:
        self._initialized = False
        self._status_code: int = 0
        self._headers = httpx.Headers()
        self._response_coro = response
        self._response: httpx.Response
        self._iterator: AsyncGenerator[bytes, bytes]
        self._litellm_logging_obj = litellm_logging_obj
        self._spend = _SpendCollection(provider_config, litellm_logging_obj)
        self._flush_scheduled = False
        self._background_tasks: set[asyncio.Task[bool]] = set()  # mutable-ok: instance set for background task tracking
        self._flush: asyncio.Task[bool] | None = None
        self._hidden_params: dict[str, object] = {}  # mutable-ok: router attaches response headers here in place

    @property
    def status_code(self) -> int:
        if not self._initialized:
            raise RuntimeError("AsyncPassthroughStreamingResponse must be awaited before accessing status_code")
        return self._status_code

    @status_code.setter
    def status_code(self, value: int) -> None:
        self._status_code = value

    @property
    def headers(self) -> httpx.Headers:
        if not self._initialized:
            raise RuntimeError("AsyncPassthroughStreamingResponse must be awaited before accessing headers")
        return self._headers

    @headers.setter
    def headers(self, value: httpx.Headers) -> None:
        self._headers = value

    def __await__(self) -> Generator[object, None, AsyncPassthroughStreamingResponse]:
        async def _init():
            if not self._initialized:
                self._response = await self._response_coro
                self.headers = self._response.headers
                self.status_code = self._response.status_code
                self._initialized = True
                try:
                    raise_for_provider_status(self._response)
                    self._iterator = _as_async_generator(self._response.aiter_bytes())
                except Exception:  # noqa: BLE001 # Safe catch-all for cleanup logic
                    try:
                        await self._response.aread()
                    except Exception:  # noqa: BLE001, S110 # Safe catch-all for cleanup logic
                        pass
                    try:
                        await self._response.aclose()
                    except Exception:  # noqa: BLE001, S110 # Safe catch-all for cleanup logic
                        pass
                    raise
            return self

        return _init().__await__()

    def _start_flush(self) -> None:
        if self._flush_scheduled or not self._spend.should_flush:
            return
        self._flush_scheduled = True

        try:
            task: Final = asyncio.create_task(
                self._litellm_logging_obj.async_flush_passthrough_collected_chunks(collector=self._spend.collector)
            )

            self._background_tasks.add(task)
            self._flush = task

            task.add_done_callback(self._background_tasks.discard)
        except Exception as e:  # noqa: BLE001 # Safe catch-all for verbose logging
            verbose_logger.exception(
                "Failed to schedule passthrough spend-tracking flush; %d collected chunks dropped: %s",
                self._spend.chunk_count,
                e,
            )

    async def wait_for_spend_flush(self) -> bool:
        """Wait until the spend-tracking flush scheduled when the stream ended or was closed has finished, and say
        whether it logged the stream, which it does not when the stream returned nothing it could read. Cancelling the
        wait leaves the flush running."""
        flush: Final = self._flush
        if flush is None:
            return False
        outcome: Final = await asyncio.shield(asyncio.gather(flush, return_exceptions=True))
        return outcome[0] is True

    def __aiter__(self) -> AsyncPassthroughStreamingResponse:
        return self

    def aiter_bytes(self) -> AsyncPassthroughStreamingResponse:
        return self

    async def __anext__(self) -> bytes:
        if not self._initialized:
            await self  # pyright: ignore[reportGeneralTypeIssues]  # structural type check misses __await__
        try:
            chunk: Final = await anext(self._iterator)
            self._spend.add(chunk)
        except Exception:  # noqa: BLE001 # Safe catch-all for cleanup logic
            self._start_flush()
            try:
                await self._response.aclose()
            except Exception:  # noqa: BLE001, S110 # Safe catch-all for cleanup logic
                pass
            raise
        else:
            return chunk

    async def asend(self, value: bytes) -> bytes:
        if not self._initialized:
            await self  # pyright: ignore[reportGeneralTypeIssues]  # structural type check misses __await__
        return await self._iterator.asend(value)

    async def athrow(
        self,
        typ: BaseException | type[BaseException],
        val: BaseException | object = None,
        tb: TracebackType | None = None,
    ) -> bytes:
        if not self._initialized:
            await self  # pyright: ignore[reportGeneralTypeIssues]  # structural type check misses __await__
        return await self._iterator.athrow(typ, val, tb)  # pyright: ignore[reportCallIssue, reportArgumentType]  # matches one of the athrow overloads

    async def aclose(self) -> None:
        self._start_flush()
        try:
            if self._initialized:
                await self._iterator.aclose()
                await self._response.aclose()
        except Exception:  # noqa: BLE001, S110 # Safe catch-all for cleanup logic
            pass


class PassthroughStreamingResponse(Generator[bytes, bytes, None]):
    def __init__(
        self,
        response: httpx.Response,
        litellm_logging_obj: LiteLLMLoggingObj,
        provider_config: BasePassthroughConfig,
    ) -> None:
        self._response = response
        self.headers = response.headers
        self.status_code = response.status_code
        self._litellm_logging_obj = litellm_logging_obj
        self._iterator: Generator[bytes, bytes, None] = _as_generator(response.iter_bytes())
        self._spend = _SpendCollection(provider_config, litellm_logging_obj)
        self._flush_scheduled = False

    def _start_flush(self) -> None:
        if self._flush_scheduled or not self._spend.should_flush:
            return
        self._flush_scheduled = True

        from litellm.utils import executor

        try:
            executor.submit(
                self._litellm_logging_obj.flush_passthrough_collected_chunks, collector=self._spend.collector
            )
        except Exception as e:  # noqa: BLE001 # Safe catch-all for verbose logging
            verbose_logger.exception(
                "Failed to schedule passthrough spend-tracking flush; %d collected chunks dropped: %s",
                self._spend.chunk_count,
                e,
            )

    def __iter__(self) -> PassthroughStreamingResponse:
        return self

    def __next__(self) -> bytes:
        try:
            chunk: Final = next(self._iterator)
            self._spend.add(chunk)
        except Exception:  # noqa: BLE001 # Safe catch-all for cleanup logic
            self._start_flush()
            try:
                self._response.close()
            except Exception:  # noqa: BLE001, S110 # Safe catch-all for cleanup logic
                pass
            raise
        else:
            return chunk

    def send(self, value: bytes) -> bytes:
        return self._iterator.send(value)

    def throw(
        self,
        typ: BaseException | type[BaseException],
        val: BaseException | object = None,
        tb: TracebackType | None = None,
    ) -> bytes:
        return self._iterator.throw(typ, val, tb)  # pyright: ignore[reportCallIssue, reportArgumentType]  # matches one of the throw overloads

    def close(self) -> None:
        self._start_flush()
        try:
            self._response.close()
        except Exception:  # noqa: BLE001, S110 # Safe catch-all for cleanup logic
            pass


@client
async def allm_passthrough_route(
    *,
    method: str,
    endpoint: str,
    model: str,
    custom_llm_provider: str | None = None,
    api_base: str | None = None,
    api_key: str | None = None,
    request_query_params: dict | None = None,
    request_headers: dict | None = None,
    content: RequestContent | None = None,
    data: dict | None = None,
    files: RequestFiles | None = None,
    json: object | None = None,
    params: QueryParamTypes | None = None,
    cookies: CookieTypes | None = None,
    client: HTTPHandler | AsyncHTTPHandler | None = None,
    **kwargs,
) -> httpx.Response | AsyncGenerator[bytes, bytes]:
    """
    Async: Reranks a list of documents based on their relevance to the query
    """
    try:
        loop: Final = asyncio.get_event_loop()
        kwargs["allm_passthrough_route"] = True

        model, custom_llm_provider, api_key, api_base = get_llm_provider(
            model=model,
            custom_llm_provider=custom_llm_provider,
            api_base=api_base,
            api_key=api_key,
        )

        from litellm.types.utils import LlmProviders
        from litellm.utils import ProviderConfigManager

        provider_config = cast(
            BasePassthroughConfig | None, kwargs.get("provider_config")
        ) or ProviderConfigManager.get_provider_passthrough_config(
            provider=LlmProviders(custom_llm_provider),
            model=model,
        )

        if provider_config is None:
            raise Exception(f"Provider {custom_llm_provider} not found")

        func: Final = partial(
            llm_passthrough_route,
            method=method,
            endpoint=endpoint,
            model=model,
            custom_llm_provider=custom_llm_provider,
            api_base=api_base,
            api_key=api_key,
            request_query_params=request_query_params,
            request_headers=request_headers,
            content=content,
            data=data,
            files=files,
            json=json,
            params=params,
            cookies=cookies,
            client=client,
            **kwargs,
        )

        ctx: Final = contextvars.copy_context()
        func_with_context: Final = partial(ctx.run, func)
        init_response: Final = await loop.run_in_executor(None, func_with_context)

        # Since allm_passthrough_route=True, we always get a coroutine from _async_passthrough_request
        if asyncio.iscoroutine(init_response):
            response: Final = await init_response

            # Only call raise_for_status if it's a Response object (not a generator)
            if isinstance(response, httpx.Response):
                response.raise_for_status()

            return response
        else:
            # This shouldn't happen when allm_passthrough_route=True, but handle it for type safety
            raise Exception("Expected coroutine from async passthrough route")

    except httpx.HTTPStatusError as e:
        # For HTTP errors, re-raise as-is to preserve the original error details
        # The caller (e.g., proxy layer) can handle conversion to appropriate response format
        raise e
    except Exception as e:
        # For other exceptions, use provider-specific error handling
        from litellm.types.utils import LlmProviders
        from litellm.utils import ProviderConfigManager

        # Get the provider using the same logic as llm_passthrough_route
        _, resolved_custom_llm_provider, _, _ = get_llm_provider(
            model=model,
            custom_llm_provider=custom_llm_provider,
            api_base=api_base,
            api_key=api_key,
        )

        # Get provider config if available
        provider_config = None
        if resolved_custom_llm_provider:
            try:
                provider_config = cast(
                    BasePassthroughConfig | None, kwargs.get("provider_config")
                ) or ProviderConfigManager.get_provider_passthrough_config(
                    provider=LlmProviders(resolved_custom_llm_provider),
                    model=model,
                )
            except Exception:  # noqa: BLE001, S110  # provider config is optional
                # If we can't get provider config, pass None
                pass

        if provider_config is None:
            # If no provider config available, raise the original exception
            raise e

        raise base_llm_http_handler.handle_error(
            e=e,
            provider_config=provider_config,
        )


@client
def llm_passthrough_route(
    *,
    method: str,
    endpoint: str,
    model: str,
    custom_llm_provider: str | None = None,
    api_base: str | None = None,
    api_key: str | None = None,
    request_query_params: dict | None = None,
    request_headers: dict | None = None,
    content: RequestContent | None = None,
    data: dict | None = None,
    files: RequestFiles | None = None,
    json: object | None = None,
    params: QueryParamTypes | None = None,
    cookies: CookieTypes | None = None,
    client: HTTPHandler | AsyncHTTPHandler | None = None,
    guardrailConfig: object = None,
    **kwargs,
) -> (
    httpx.Response
    | Coroutine[object, object, httpx.Response]
    | Coroutine[object, object, httpx.Response | AsyncGenerator[bytes, bytes]]
    | Generator[bytes, bytes, None]
    | AsyncGenerator[bytes, bytes]
):
    """
    Pass through requests to the LLM APIs.

    Step 1. Build the request
    Step 2. Send the request
    Step 3. Return the response
    """
    from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
    from litellm.types.utils import LlmProviders
    from litellm.utils import ProviderConfigManager

    _is_async: Final = bool(kwargs.get("allm_passthrough_route", False))

    litellm_logging_obj: Final = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj"))

    model, custom_llm_provider, api_key, api_base = get_llm_provider(
        model=model,
        custom_llm_provider=custom_llm_provider,
        api_base=api_base,
        api_key=api_key,
    )

    litellm_params_dict: Final = get_litellm_params(api_key=api_key, api_base=api_base, **kwargs)

    if client is None:
        from litellm.llms.custom_httpx.http_handler import (
            get_async_httpx_client,
            get_httpx_client,
        )
        from litellm.passthrough.timeout_utils import resolve_llm_passthrough_timeout
        from litellm.types.llms.custom_http import httpxSpecialProvider

        resolved_timeout: Final = resolve_llm_passthrough_timeout(
            kwargs=kwargs,
            litellm_params=litellm_params_dict,
        )
        if _is_async:
            client = get_async_httpx_client(
                llm_provider=httpxSpecialProvider.PassThroughEndpoint,
                params={"timeout": resolved_timeout},
            )
        else:
            client = get_httpx_client(params={"timeout": resolved_timeout})

    # Add model_id to litellm_params if present in kwargs (for Bedrock Application Inference Profiles)
    if "model_id" in kwargs:
        litellm_params_dict["model_id"] = kwargs["model_id"]

    litellm_logging_obj.update_environment_variables(
        model=model,
        litellm_params=litellm_params_dict,
        optional_params={},
        endpoint=endpoint,
        custom_llm_provider=custom_llm_provider,
        request_data=data if data else json,
    )

    provider_config: Final = cast(
        BasePassthroughConfig | None, kwargs.get("provider_config")
    ) or ProviderConfigManager.get_provider_passthrough_config(
        provider=LlmProviders(custom_llm_provider),
        model=model,
    )
    if provider_config is None:
        raise Exception(f"Provider {custom_llm_provider} not found")

    updated_url, base_target_url = provider_config.get_complete_url(
        api_base=api_base,
        api_key=api_key,
        model=model,
        endpoint=endpoint,
        request_query_params=request_query_params,
        litellm_params=litellm_params_dict,
    )

    # [TODO: Refactor to bedrockpassthroughconfig] need to encode the id of application-inference-profile for bedrock
    if custom_llm_provider == "bedrock" and "application-inference-profile" in endpoint:
        encoded_url_str: Final = CommonUtils.encode_bedrock_runtime_modelid_arn(str(updated_url))
        updated_url = httpx.URL(encoded_url_str)

    # Add or update query parameters
    provider_api_key: Final = provider_config.get_api_key(api_key)

    auth_headers: Final = provider_config.validate_environment(
        headers={},
        model=model,
        messages=[],
        optional_params={},
        litellm_params=litellm_params_dict,
        api_key=provider_api_key,
        api_base=base_target_url,
    )

    headers = BasePassthroughUtils.forward_headers_from_request(
        request_headers=request_headers or {},
        headers=auth_headers,
        forward_headers=False,
    )

    _request_data: dict | None = data if isinstance(data, dict) else (json if isinstance(json, dict) else None)
    headers, signed_json_body = provider_config.sign_request(
        headers=headers,
        litellm_params={**litellm_params_dict, "guardrailConfig": guardrailConfig},
        request_data=_request_data,
        api_base=str(updated_url),
        model=model,
    )

    ## SWAP MODEL IN JSON BODY [TODO: REFACTOR TO A provider_config.transform_request method]
    if json and isinstance(json, dict) and "model" in json:
        json["model"] = model

    request: Final = client.client.build_request(
        method=method,
        url=updated_url,
        content=signed_json_body if signed_json_body is not None else content,
        data=data if (signed_json_body is None and content is None) else None,
        files=files,
        json=json if (signed_json_body is None and content is None) else None,
        params=params,
        headers=headers,
        cookies=cookies,
    )

    ## IS STREAMING REQUEST
    _streaming_request_data: Final[dict[str, object]] = (
        data if isinstance(data, dict) else (json if isinstance(json, dict) else {})
    )
    is_streaming_request: Final = provider_config.is_streaming_request(
        endpoint=endpoint,
        request_data=_streaming_request_data,
    )

    # Update logging object with streaming status
    litellm_logging_obj.stream = is_streaming_request

    ## LOGGING PRE-CALL
    request_data: Final = data if data else json
    litellm_logging_obj.pre_call(
        input=request_data,
        api_key=provider_api_key,
        additional_args={
            "complete_input_dict": request_data,
            "api_base": str(updated_url),
            "headers": headers,
        },
    )

    try:
        if _is_async:
            # Return the coroutine to be awaited by the caller
            return _async_passthrough_request(
                client=client,
                request=request,
                is_streaming_request=is_streaming_request,
                litellm_logging_obj=litellm_logging_obj,
                provider_config=provider_config,
            )
        else:
            # Sync path - client.client.send returns Response directly
            response: httpx.Response = client.client.send(request=request, stream=is_streaming_request)
            try:
                response.raise_for_status()
            except Exception:  # noqa: BLE001 # Safe catch-all for cleanup logic
                try:
                    response.read()
                except Exception:  # noqa: BLE001, S110 # Safe catch-all for cleanup logic
                    pass
                try:
                    response.close()
                except Exception:  # noqa: BLE001, S110 # Safe catch-all for cleanup logic
                    pass
                raise

            if hasattr(response, "iter_bytes") and is_streaming_request:
                return PassthroughStreamingResponse(response, litellm_logging_obj, provider_config)
            else:
                return response
    except Exception as e:
        assert provider_config is not None
        raise base_llm_http_handler.handle_error(
            e=e,
            provider_config=provider_config,
        )


async def _async_passthrough_request(
    client: HTTPHandler | AsyncHTTPHandler,
    request: httpx.Request,
    is_streaming_request: bool,
    litellm_logging_obj: LiteLLMLoggingObj,
    provider_config: BasePassthroughConfig,
) -> httpx.Response | AsyncGenerator[bytes, bytes]:
    """
    Handle async passthrough requests.
    Uses async client to send request and properly handles streaming.
    """
    # client.client.send returns a coroutine for async clients
    response_result: Final = client.client.send(request=request, stream=is_streaming_request)

    # Check if it's a coroutine and await it
    if asyncio.iscoroutine(response_result):
        if is_streaming_request:
            return await AsyncPassthroughStreamingResponse(  # pyright: ignore[reportGeneralTypeIssues]  # structural type check misses __await__
                response=response_result,
                litellm_logging_obj=litellm_logging_obj,
                provider_config=provider_config,
            )
        else:
            response: Final = await response_result
            await response.aread()
            raise_for_provider_status(response)
            return response
    else:
        # Fallback for sync-like behavior (shouldn't happen in async path)
        raise Exception("Expected coroutine from async client")