tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 18:21:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Bedrock Passthrough cost tracking (/invoke + /converse routes - streaming + non-streaming) (#12123)

Browse Source 
* refactor(passthrough_endpoints-success-handler): refactor llm passthrough logging logic

isolate the llm translation work to enable cost tracking on sdk

* feat: initial implementation of passthrough SDK cost calculation

enables bedrock passthrough cost tracking to work

* feat(cost_calculator.py): working cost calculation for bedrock passthrough

* feat(litellm_logging.py): consider allm_passthrough in cost tracking

allows async calls (e.g. via proxy) to work

* feat(bedrock/passthrough): working event stream decoding for bedrock passthrough calls + logging instrumentation for passthrough sdk calls (log on stream completion)

Enables bedrock streaming cost calculation

* feat(litellm_logging.py): support streaming passthrough cost tracking

* feat(passthrough/main.py): working async streaming cost calculation

Closes https://github.com/BerriAI/litellm/issues/11359

* feat(proxy_server.py): fix passthrough routing when llm router enabled

* feat: further fixes

* feat(bedrock/): working bedrock passthrough cost tracking (non-streaming)

* feat(litellm_logging.py): working usage tracking for bedrock passthrough calls

ensures tokens are logged

* feat(bedrock/passthrough): add converse passthrough cost tracking support

* feat(base_llm/passthrough): remove redundant function

* refactor(litellm_logging.py): refactor function to be below 50 LOC

* test: update test

* test: remove redundant test
This commit is contained in:
Krish Dholakia
2025-06-27 20:01:12 -07:00
committed by GitHub
parent db2a55a3d2
commit ee6e76e1f9
15 changed files with 841 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/cost_calculator.py
+7 -1
View File
@@ -2,8 +2,9 @@
## File for 'response_cost' calculation in Logging
import time
from functools import lru_cache
from typing import Any, List, Literal, Optional, Tuple, Union, cast
from typing import TYPE_CHECKING, Any, List, Literal, Optional, Tuple, Union, cast
from httpx import Response
from pydantic import BaseModel
import litellm
@@ -93,6 +94,9 @@ from litellm.utils import (
token_counter,
)
if TYPE_CHECKING:
pass
def _cost_per_token_custom_pricing_helper(
prompt_tokens: float = 0,
@@ -655,6 +659,7 @@ def completion_cost( # noqa: PLR0915
potential_model_names = [selected_model]
if model is not None:
potential_model_names.append(model)
for idx, model in enumerate(potential_model_names):
try:
verbose_logger.debug(
@@ -964,6 +969,7 @@ def response_cost_calculator(
ResponsesAPIResponse,
LiteLLMRealtimeStreamLoggingObject,
OpenAIModerationResponse,
Response,
],
model: str,
custom_llm_provider: Optional[str],
litellm/litellm_core_utils/litellm_logging.py
+103 -21
View File
@@ -14,6 +14,7 @@ import uuid
from datetime import datetime as dt_object
from functools import lru_cache
from typing import (
TYPE_CHECKING,
Any,
Callable,
Dict,
@@ -26,6 +27,7 @@ from typing import (
cast,
)
from httpx import Response
from pydantic import BaseModel
import litellm
@@ -81,6 +83,7 @@ from litellm.types.rerank import RerankResponse
from litellm.types.router import CustomPricingLiteLLMParams
from litellm.types.utils import (
CallTypes,
CostResponseTypes,
DynamicPromptManagementParamLiteral,
EmbeddingResponse,
ImageResponse,
@@ -148,6 +151,8 @@ from .initialize_dynamic_callback_params import (
)
from .specialty_caches.dynamic_logging_cache import DynamicLoggingCache
if TYPE_CHECKING:
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
try:
from litellm_enterprise.enterprise_callbacks.callback_controls import (
EnterpriseCallbackControls,
@@ -1254,6 +1259,46 @@ class Logging(LiteLLMLoggingBaseClass):
self.completion_start_time = completion_start_time
self.model_call_details["completion_start_time"] = self.completion_start_time
def normalize_logging_result(self, result: Any) -> Any:
"""
Some endpoints return a different type of result than what is expected by the logging system.
This function is used to normalize the result to the expected type.
"""
logging_result = result
if self.call_type == CallTypes.arealtime.value and isinstance(result, list):
combined_usage_object = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(
results=result
)
logging_result = (
RealtimeAPITokenUsageProcessor.create_logging_realtime_object(
usage=combined_usage_object,
results=result,
)
)
elif (
self.call_type == CallTypes.llm_passthrough_route.value
or self.call_type == CallTypes.allm_passthrough_route.value
) and isinstance(result, Response):
from litellm.utils import ProviderConfigManager
provider_config = ProviderConfigManager.get_provider_passthrough_config(
provider=self.model_call_details.get("custom_llm_provider", ""),
model=self.model,
)
if provider_config is not None:
logging_result = provider_config.logging_non_streaming_response(
model=self.model,
custom_llm_provider=self.model_call_details.get(
"custom_llm_provider", ""
),
httpx_response=result,
request_data=self.model_call_details.get("request_data", {}),
logging_obj=self,
endpoint=self.model_call_details.get("endpoint", ""),
)
return logging_result
def _success_handler_helper_fn(
self,
result=None,
@@ -1287,28 +1332,8 @@ class Logging(LiteLLMLoggingBaseClass):
## if model in model cost map - log the response cost
## else set cost to None
logging_result = result
logging_result = self.normalize_logging_result(result=result)
if self.call_type == CallTypes.arealtime.value and isinstance(result, list):
combined_usage_object = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(
results=result
)
logging_result = (
RealtimeAPITokenUsageProcessor.create_logging_realtime_object(
usage=combined_usage_object,
results=result,
)
)
# self.model_call_details[
# "response_cost"
# ] = handle_realtime_stream_cost_calculation(
# results=result,
# combined_usage_object=combined_usage_object,
# custom_llm_provider=self.custom_llm_provider,
# litellm_model_name=self.model,
# )
# self.model_call_details["combined_usage_object"] = combined_usage_object
if (
standard_logging_object is None
and result is not None
@@ -1426,6 +1451,59 @@ class Logging(LiteLLMLoggingBaseClass):
except Exception as e:
raise Exception(f"[Non-Blocking] LiteLLM.Success_Call Error: {str(e)}")
def _flush_passthrough_collected_chunks_helper(
self,
raw_bytes: List[bytes],
provider_config: "BasePassthroughConfig",
) -> Optional["CostResponseTypes"]:
all_chunks = provider_config._convert_raw_bytes_to_str_lines(raw_bytes)
complete_streaming_response = provider_config.handle_logging_collected_chunks(
all_chunks=all_chunks,
litellm_logging_obj=self,
model=self.model,
custom_llm_provider=self.model_call_details.get("custom_llm_provider", ""),
endpoint=self.model_call_details.get("endpoint", ""),
)
return complete_streaming_response
def flush_passthrough_collected_chunks(
self,
raw_bytes: List[bytes],
provider_config: "BasePassthroughConfig",
):
"""
Flush collected chunks from the logging object
This is used to log the collected chunks once streaming is done on passthrough endpoints
1. Decode the raw bytes to string lines
2. Get the complete streaming response from the provider config
3. Log the complete streaming response (trigger success handler)
This is used for passthrough endpoints
"""
complete_streaming_response = self._flush_passthrough_collected_chunks_helper(
raw_bytes=raw_bytes,
provider_config=provider_config,
)
if complete_streaming_response is not None:
self.success_handler(result=complete_streaming_response)
return
async def async_flush_passthrough_collected_chunks(
self,
raw_bytes: List[bytes],
provider_config: "BasePassthroughConfig",
):
complete_streaming_response = self._flush_passthrough_collected_chunks_helper(
raw_bytes=raw_bytes,
provider_config=provider_config,
)
if complete_streaming_response is not None:
await self.async_success_handler(result=complete_streaming_response)
return
def success_handler( # noqa: PLR0915
self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs
):
@@ -1968,6 +2046,7 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["async_complete_streaming_response"] = (
complete_streaming_response
)
try:
if self.model_call_details.get("cache_hit", False) is True:
self.model_call_details["response_cost"] = 0.0
@@ -2047,6 +2126,7 @@ class Logging(LiteLLMLoggingBaseClass):
)
self.has_run_logging(event_type="async_success")
for callback in callbacks:
# check if callback can run for this request
litellm_params = self.model_call_details.get("litellm_params", {})
@@ -2709,6 +2789,8 @@ class Logging(LiteLLMLoggingBaseClass):
return result
elif isinstance(result, ResponseCompletedEvent):
return result.response
else:
return None
return None
def _handle_anthropic_messages_response_logging(self, result: Any) -> ModelResponse:
litellm/litellm_core_utils/streaming_handler.py
+26
View File
@@ -1891,3 +1891,29 @@ def generic_chunk_has_all_required_fields(chunk: dict) -> bool:
decision = all(key in _all_fields for key in chunk)
return decision
def convert_generic_chunk_to_model_response_stream(
chunk: GChunk,
) -> ModelResponseStream:
from litellm.types.utils import Delta
model_response_stream = ModelResponseStream(
id=str(uuid.uuid4()),
model="",
choices=[
StreamingChoices(
index=chunk.get("index", 0),
delta=Delta(
content=chunk["text"],
tool_calls=chunk.get("tool_use", None),
),
)
],
finish_reason=chunk["finish_reason"] if chunk["is_finished"] else None,
)
if "usage" in chunk and chunk["usage"] is not None:
setattr(model_response_stream, "usage", chunk["usage"])
return model_response_stream
litellm/llms/base_llm/passthrough/transformation.py
+44 -2
View File
@@ -1,10 +1,13 @@
from abc import abstractmethod
from typing import TYPE_CHECKING, Optional, Tuple, Union
from typing import TYPE_CHECKING, List, Optional, Tuple, Union
from ..base_utils import BaseLLMModelInfo
if TYPE_CHECKING:
from httpx import URL, Headers
from httpx import URL, Headers, Response
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.utils import CostResponseTypes
from ..chat.transformation import BaseLLMException
@@ -101,3 +104,42 @@ class BasePassthroughConfig(BaseLLMModelInfo):
return BaseLLMException(
status_code=status_code, message=error_message, headers=headers
)
def logging_non_streaming_response(
self,
model: str,
custom_llm_provider: str,
httpx_response: "Response",
request_data: dict,
logging_obj: "LiteLLMLoggingObj",
endpoint: str,
) -> Optional["CostResponseTypes"]:
pass
def handle_logging_collected_chunks(
self,
all_chunks: List[str],
litellm_logging_obj: "LiteLLMLoggingObj",
model: str,
custom_llm_provider: str,
endpoint: str,
) -> Optional["CostResponseTypes"]:
return None
def _convert_raw_bytes_to_str_lines(self, raw_bytes: List[bytes]) -> List[str]:
"""
Converts a list of raw bytes into a list of string lines, similar to aiter_lines()
Args:
raw_bytes: List of bytes chunks from aiter.bytes()
Returns:
List of string lines, with each line being a complete data: {} chunk
"""
# Combine all bytes and decode to string
combined_str = b"".join(raw_bytes).decode("utf-8")
# Split by newlines and filter out empty lines
lines = [line.strip() for line in combined_str.split("\n") if line.strip()]
return lines
litellm/llms/bedrock/chat/__init__.py
+29 -1
View File
@@ -1,2 +1,30 @@
from typing import Optional
from .converse_handler import BedrockConverseLLM
from .invoke_handler import BedrockLLM
from .invoke_handler import (
AmazonAnthropicClaudeStreamDecoder,
AmazonDeepSeekR1StreamDecoder,
AWSEventStreamDecoder,
BedrockLLM,
)
def get_bedrock_event_stream_decoder(
invoke_provider: Optional[str], model: str, sync_stream: bool, json_mode: bool
):
if invoke_provider and invoke_provider == "anthropic":
decoder: AWSEventStreamDecoder = AmazonAnthropicClaudeStreamDecoder(
model=model,
sync_stream=sync_stream,
json_mode=json_mode,
)
return decoder
elif invoke_provider and invoke_provider == "deepseek_r1":
decoder = AmazonDeepSeekR1StreamDecoder(
model=model,
sync_stream=sync_stream,
)
return decoder
else:
decoder = AWSEventStreamDecoder(model=model)
return decoder
litellm/llms/bedrock/common_utils.py
+66
View File
@@ -2,6 +2,7 @@
Common utilities used across bedrock chat/embedding/image generation
"""
import json
import os
from typing import TYPE_CHECKING, List, Literal, Optional, Union
@@ -458,3 +459,68 @@ class BedrockModelInfo(BaseLLMModelInfo):
):
return "converse"
return "invoke"
class BedrockEventStreamDecoderBase:
"""
Base class for event stream decoding for Bedrock
"""
_response_stream_shape_cache = None
def __init__(self):
from botocore.parsers import EventStreamJSONParser
self.parser = EventStreamJSONParser()
def get_response_stream_shape(self):
if self._response_stream_shape_cache is None:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
bedrock_service_dict = loader.load_service_model(
"bedrock-runtime", "service-2"
)
bedrock_service_model = ServiceModel(bedrock_service_dict)
self._response_stream_shape_cache = bedrock_service_model.shape_for(
"ResponseStream"
)
return self._response_stream_shape_cache
def _parse_message_from_event(self, event) -> Optional[str]:
response_dict = event.to_response_dict()
parsed_response = self.parser.parse(
response_dict, self.get_response_stream_shape()
)
if response_dict["status_code"] != 200:
decoded_body = response_dict["body"].decode()
if isinstance(decoded_body, dict):
error_message = decoded_body.get("message")
elif isinstance(decoded_body, str):
error_message = decoded_body
else:
error_message = ""
exception_status = response_dict["headers"].get(":exception-type")
error_message = exception_status + " " + error_message
raise BedrockError(
status_code=response_dict["status_code"],
message=(
json.dumps(error_message)
if isinstance(error_message, dict)
else error_message
),
)
if "chunk" in parsed_response:
chunk = parsed_response.get("chunk")
if not chunk:
return None
return chunk.get("bytes").decode() # type: ignore[no-any-return]
else:
chunk = response_dict.get("body")
if not chunk:
return None
return chunk.decode() # type: ignore[no-any-return]
litellm/llms/bedrock/passthrough/transformation.py
+143 -3
View File
@@ -1,15 +1,26 @@
from typing import TYPE_CHECKING, Optional, Tuple
import json
from typing import TYPE_CHECKING, List, Optional, Tuple, cast
from httpx import Response
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import BedrockModelInfo
from ..common_utils import BedrockEventStreamDecoderBase, BedrockModelInfo
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.utils import CostResponseTypes
if TYPE_CHECKING:
from httpx import URL
class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BasePassthroughConfig):
class BedrockPassthroughConfig(
BaseAWSLLM, BedrockModelInfo, BedrockEventStreamDecoderBase, BasePassthroughConfig
):
def is_streaming_request(self, endpoint: str, request_data: dict) -> bool:
return "stream" in endpoint
@@ -51,3 +62,132 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BasePassthroughConf
api_base=api_base,
model=model,
)
def logging_non_streaming_response(
self,
model: str,
custom_llm_provider: str,
httpx_response: Response,
request_data: dict,
logging_obj: Logging,
endpoint: str,
) -> Optional["CostResponseTypes"]:
from litellm import encoding
from litellm.types.utils import LlmProviders, ModelResponse
from litellm.utils import ProviderConfigManager
if "invoke" in endpoint:
chat_config_model = "invoke/" + model
elif "converse" in endpoint:
chat_config_model = "converse/" + model
else:
return None
provider_chat_config = ProviderConfigManager.get_provider_chat_config(
provider=LlmProviders(custom_llm_provider),
model=chat_config_model,
)
if provider_chat_config is None:
raise ValueError(f"No provider config found for model: {model}")
litellm_model_response: ModelResponse = provider_chat_config.transform_response(
model=model,
messages=[{"role": "user", "content": "no-message-pass-through-endpoint"}],
raw_response=httpx_response,
model_response=ModelResponse(),
logging_obj=logging_obj,
optional_params={},
litellm_params={},
api_key="",
request_data=request_data,
encoding=encoding,
)
return litellm_model_response
def _convert_raw_bytes_to_str_lines(self, raw_bytes: List[bytes]) -> List[str]:
from botocore.eventstream import EventStreamBuffer
all_chunks = []
event_stream_buffer = EventStreamBuffer()
for chunk in raw_bytes:
event_stream_buffer.add_data(chunk)
for event in event_stream_buffer:
message = self._parse_message_from_event(event)
if message is not None:
all_chunks.append(message)
return all_chunks
def handle_logging_collected_chunks(
self,
all_chunks: List[str],
litellm_logging_obj: "LiteLLMLoggingObj",
model: str,
custom_llm_provider: str,
endpoint: str,
) -> Optional["CostResponseTypes"]:
"""
1. Convert all_chunks to a ModelResponseStream
2. combine model_response_stream to model_response
3. Return the model_response
"""
from litellm.litellm_core_utils.streaming_handler import (
convert_generic_chunk_to_model_response_stream,
generic_chunk_has_all_required_fields,
)
from litellm.llms.bedrock.chat import get_bedrock_event_stream_decoder
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
AmazonInvokeConfig,
)
from litellm.main import stream_chunk_builder
from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
all_translated_chunks = []
if "invoke" in endpoint:
invoke_provider = AmazonInvokeConfig.get_bedrock_invoke_provider(model)
if invoke_provider is None:
raise ValueError(
f"Invalid invoke provider: {invoke_provider}, for model: {model}"
)
obj = get_bedrock_event_stream_decoder(
invoke_provider=invoke_provider,
model=model,
sync_stream=True,
json_mode=False,
)
elif "converse" in endpoint:
obj = get_bedrock_event_stream_decoder(
invoke_provider=None,
model=model,
sync_stream=True,
json_mode=False,
)
else:
return None
for chunk in all_chunks:
message = json.loads(chunk)
translated_chunk = obj._chunk_parser(chunk_data=message)
if isinstance(
translated_chunk, dict
) and generic_chunk_has_all_required_fields(cast(dict, translated_chunk)):
chunk_obj = convert_generic_chunk_to_model_response_stream(
cast(GenericStreamingChunk, translated_chunk)
)
elif isinstance(translated_chunk, ModelResponseStream):
chunk_obj = translated_chunk
else:
continue
all_translated_chunks.append(chunk_obj)
if len(all_translated_chunks) > 0:
model_response = stream_chunk_builder(
chunks=all_translated_chunks,
)
return model_response
return None
litellm/passthrough/main.py
+130 -16
View File
@@ -5,7 +5,17 @@ This module is used to pass through requests to the LLM APIs.
import asyncio
import contextvars
from functools import partial
from typing import TYPE_CHECKING, Any, Coroutine, Optional, Union, cast
from typing import (
TYPE_CHECKING,
Any,
AsyncGenerator,
Coroutine,
Generator,
List,
Optional,
Union,
cast,
)
import httpx
from httpx._types import CookieTypes, QueryParamTypes, RequestFiles
@@ -20,6 +30,7 @@ base_llm_http_handler = BaseLLMHTTPHandler()
from .utils import BasePassthroughUtils
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
@@ -42,7 +53,12 @@ async def allm_passthrough_route(
cookies: Optional[CookieTypes] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
**kwargs,
) -> Union[httpx.Response, Coroutine[Any, Any, httpx.Response]]:
) -> Union[
httpx.Response,
Coroutine[Any, Any, httpx.Response],
Generator[Any, Any, Any],
AsyncGenerator[Any, Any],
]:
"""
Async: Reranks a list of documents based on their relevance to the query
"""
@@ -50,6 +66,26 @@ async def allm_passthrough_route(
loop = 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(
Optional["BasePassthroughConfig"], 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 = partial(
llm_passthrough_route,
method=method,
@@ -76,20 +112,24 @@ async def allm_passthrough_route(
if asyncio.iscoroutine(init_response):
response = await init_response
try:
response.raise_for_status()
except httpx.HTTPStatusError as e:
error_text = await e.response.aread()
error_text_str = error_text.decode("utf-8")
raise Exception(error_text_str)
else:
response = init_response
return response
except Exception as e:
# For passthrough routes, we need to get the provider config to properly handle errors
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,
@@ -97,7 +137,7 @@ async def allm_passthrough_route(
api_base=api_base,
api_key=api_key,
)
# Get provider config if available
provider_config = None
if resolved_custom_llm_provider:
@@ -111,11 +151,11 @@ async def allm_passthrough_route(
except Exception:
# 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,
@@ -142,22 +182,31 @@ def llm_passthrough_route(
cookies: Optional[CookieTypes] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
**kwargs,
) -> Union[httpx.Response, Coroutine[Any, Any, httpx.Response]]:
) -> Union[
httpx.Response,
Coroutine[Any, Any, httpx.Response],
Generator[Any, Any, Any],
AsyncGenerator[Any, Any],
]:
"""
Pass through requests to the LLM APIs.
Step 1. Build the request
Step 2. Send the request
Step 3. Return the response
[TODO] Refactor this into a provider-config pattern, once we expand this to non-vllm providers.
"""
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
if client is None:
if allm_passthrough_route:
client = litellm.module_level_aclient
else:
client = litellm.module_level_client
litellm_logging_obj = 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,
@@ -165,11 +214,15 @@ def llm_passthrough_route(
api_key=api_key,
)
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
litellm_params_dict = get_litellm_params(**kwargs)
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 = cast(
Optional["BasePassthroughConfig"], kwargs.get("provider_config")
@@ -215,7 +268,7 @@ def llm_passthrough_route(
model=model,
)
## SWAP MODEL IN JSON BODY
## 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
@@ -237,12 +290,27 @@ def llm_passthrough_route(
request_data=data or json or {},
)
# Update logging object with streaming status
litellm_logging_obj.stream = is_streaming_request
try:
response = client.client.send(request=request, stream=is_streaming_request)
if asyncio.iscoroutine(response):
return response
if is_streaming_request:
return _async_streaming(response, litellm_logging_obj, provider_config)
else:
return response
response.raise_for_status()
return response
if (
hasattr(response, "iter_bytes") and is_streaming_request
): # yield the chunk, so we can store it in the logging object
return _sync_streaming(response, litellm_logging_obj, provider_config)
else:
# For non-streaming responses, yield the entire response
return response
except Exception as e:
if provider_config is None:
raise e
@@ -250,3 +318,49 @@ def llm_passthrough_route(
e=e,
provider_config=provider_config,
)
def _sync_streaming(
response: httpx.Response,
litellm_logging_obj: "LiteLLMLoggingObj",
provider_config: "BasePassthroughConfig",
):
from litellm.utils import executor
try:
raw_bytes: List[bytes] = []
for chunk in response.iter_bytes(): # type: ignore
raw_bytes.append(chunk)
yield chunk
executor.submit(
litellm_logging_obj.flush_passthrough_collected_chunks,
raw_bytes=raw_bytes,
provider_config=provider_config,
)
except Exception as e:
raise e
async def _async_streaming(
response: Coroutine[Any, Any, httpx.Response],
litellm_logging_obj: "LiteLLMLoggingObj",
provider_config: "BasePassthroughConfig",
):
try:
iter_response = await response
raw_bytes: List[bytes] = []
async for chunk in iter_response.aiter_bytes(): # type: ignore
raw_bytes.append(chunk)
yield chunk
asyncio.create_task(
litellm_logging_obj.async_flush_passthrough_collected_chunks(
raw_bytes=raw_bytes,
provider_config=provider_config,
)
)
except Exception as e:
raise e
litellm/proxy/common_request_processing.py
+42 -25
View File
@@ -108,14 +108,26 @@ async def create_streaming_response(
final_status_code = default_status_code
try:
# Handle coroutine that returns a generator
if asyncio.iscoroutine(generator):
generator = await generator
# Now get the first chunk from the actual generator
first_chunk_value = await generator.__anext__()
if first_chunk_value is not None:
error_code_from_chunk = await _parse_event_data_for_error(first_chunk_value)
if error_code_from_chunk is not None:
final_status_code = error_code_from_chunk
verbose_proxy_logger.debug(
f"Error detected in first stream chunk. Status code set to: {final_status_code}"
try:
error_code_from_chunk = await _parse_event_data_for_error(
first_chunk_value
)
if error_code_from_chunk is not None:
final_status_code = error_code_from_chunk
verbose_proxy_logger.debug(
f"Error detected in first stream chunk. Status code set to: {final_status_code}"
)
except Exception as e:
verbose_proxy_logger.debug(f"Error parsing first chunk value: {e}")
except StopAsyncIteration:
# Generator was empty. Default status
@@ -152,6 +164,7 @@ async def create_streaming_response(
with tracer.trace(DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE):
yield first_chunk_value
async for chunk in generator:
with tracer.trace(DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE):
yield chunk
@@ -432,6 +445,8 @@ class ProxyBaseLLMRequestProcessing:
)
if self._is_streaming_request(
data=self.data, is_streaming_request=is_streaming_request
) or self._is_streaming_response(
response
): # use generate_responses to stream responses
custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
@@ -448,11 +463,28 @@ class ProxyBaseLLMRequestProcessing:
**additional_headers,
)
if route_type == "allm_passthrough_route":
return StreamingResponse(
content=response.aiter_bytes(),
status_code=response.status_code,
headers=custom_headers,
)
# Check if response is an async generator
if self._is_streaming_response(response):
if asyncio.iscoroutine(response):
generator = await response
else:
generator = response
# For passthrough routes, stream directly without error parsing
# since we're dealing with raw binary data (e.g., AWS event streams)
return StreamingResponse(
content=generator,
status_code=status.HTTP_200_OK,
headers=custom_headers,
)
else:
# Traditional HTTP response with aiter_bytes
return StreamingResponse(
content=response.aiter_bytes(),
status_code=response.status_code,
headers=custom_headers,
)
else:
selected_data_generator = select_data_generator(
response=response,
@@ -557,7 +589,6 @@ class ProxyBaseLLMRequestProcessing:
This uses standard Python inspection to detect streaming/async iterator objects
rather than relying on specific wrapper classes.
"""
import asyncio
import inspect
from collections.abc import AsyncGenerator, AsyncIterator
@@ -569,20 +600,6 @@ class ProxyBaseLLMRequestProcessing:
if isinstance(response, (AsyncIterator, AsyncGenerator)):
return True
# Check for __aiter__ method (async iterator protocol)
if hasattr(response, "__aiter__") and callable(getattr(response, "__aiter__")):
return True
# Check if it's a coroutine that might yield an async generator
if asyncio.iscoroutine(response):
return True
# Check for common streaming HTTP response patterns
if hasattr(response, "aiter_bytes") and callable(
getattr(response, "aiter_bytes")
):
return True
return False
def _is_streaming_request(
litellm/proxy/pass_through_endpoints/success_handler.py
+62 -22
View File
@@ -1,6 +1,6 @@
import json
from datetime import datetime
from typing import Optional, Union
from typing import Any, Optional, Union
from urllib.parse import urlparse
import httpx
@@ -88,29 +88,25 @@ class PassThroughEndpointLogging:
cache_hit=False,
**kwargs,
)
async def pass_through_async_success_handler(
def normalize_llm_passthrough_logging_payload(
self,
httpx_response: httpx.Response,
response_body: Optional[dict],
request_body: dict,
logging_obj: LiteLLMLoggingObj,
url_route: str,
result: str,
start_time: datetime,
end_time: datetime,
cache_hit: bool,
request_body: dict,
passthrough_logging_payload: PassthroughStandardLoggingPayload,
**kwargs,
):
standard_logging_response_object: Optional[
PassThroughEndpointLoggingResultValues
] = None
logging_obj.model_call_details["passthrough_logging_payload"] = (
passthrough_logging_payload
)
return_dict = {
"standard_logging_response_object": None,
"kwargs": kwargs,
}
standard_logging_response_object: Optional[Any] = None
if self.is_vertex_route(url_route):
vertex_passthrough_logging_handler_result = (
VertexPassthroughLoggingHandler.vertex_passthrough_handler(
@@ -166,7 +162,33 @@ class PassThroughEndpointLogging:
cohere_passthrough_logging_handler_result["result"]
)
kwargs = cohere_passthrough_logging_handler_result["kwargs"]
elif self.is_assemblyai_route(url_route):
return_dict["standard_logging_response_object"] = (
standard_logging_response_object
)
return_dict["kwargs"] = kwargs
return return_dict
async def pass_through_async_success_handler(
self,
httpx_response: httpx.Response,
response_body: Optional[dict],
logging_obj: LiteLLMLoggingObj,
url_route: str,
result: str,
start_time: datetime,
end_time: datetime,
cache_hit: bool,
request_body: dict,
passthrough_logging_payload: PassthroughStandardLoggingPayload,
**kwargs,
):
standard_logging_response_object: Optional[
PassThroughEndpointLoggingResultValues
] = None
logging_obj.model_call_details["passthrough_logging_payload"] = (
passthrough_logging_payload
)
if self.is_assemblyai_route(url_route):
if (
AssemblyAIPassthroughLoggingHandler._should_log_request(
httpx_response.request.method
@@ -189,18 +211,37 @@ class PassThroughEndpointLogging:
elif self.is_langfuse_route(url_route):
# Don't log langfuse pass-through requests
return
else:
normalized_llm_passthrough_logging_payload = (
self.normalize_llm_passthrough_logging_payload(
httpx_response=httpx_response,
response_body=response_body,
request_body=request_body,
logging_obj=logging_obj,
url_route=url_route,
result=result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
**kwargs,
)
)
standard_logging_response_object = (
normalized_llm_passthrough_logging_payload[
"standard_logging_response_object"
]
)
kwargs = normalized_llm_passthrough_logging_payload["kwargs"]
if standard_logging_response_object is None:
standard_logging_response_object = StandardPassThroughResponseObject(
response=httpx_response.text
)
kwargs = self._set_cost_per_request(
logging_obj=logging_obj,
passthrough_logging_payload=passthrough_logging_payload,
kwargs=kwargs,
)
await self._handle_logging(
logging_obj=logging_obj,
@@ -244,11 +285,10 @@ class PassThroughEndpointLogging:
if route in parsed_url.path:
return True
return False
def _set_cost_per_request(
self,
logging_obj: LiteLLMLoggingObj,
self,
logging_obj: LiteLLMLoggingObj,
passthrough_logging_payload: PassthroughStandardLoggingPayload,
kwargs: dict,
):
@@ -265,8 +305,8 @@ class PassThroughEndpointLogging:
kwargs["response_cost"] = passthrough_logging_payload.get(
"cost_per_request"
)
logging_obj.model_call_details["response_cost"] = passthrough_logging_payload.get(
"cost_per_request"
logging_obj.model_call_details["response_cost"] = (
passthrough_logging_payload.get("cost_per_request")
)
return kwargs
litellm/proxy/route_llm_request.py
+6
View File
@@ -121,20 +121,24 @@ async def route_request(
data["model"] in router_model_names
or data["model"] in llm_router.get_model_ids()
):
return getattr(llm_router, f"{route_type}")(**data)
elif (
llm_router.model_group_alias is not None
and data["model"] in llm_router.model_group_alias
):
return getattr(llm_router, f"{route_type}")(**data)
elif data["model"] in llm_router.deployment_names:
return getattr(llm_router, f"{route_type}")(
**data, specific_deployment=True
)
elif data["model"] not in router_model_names:
if llm_router.router_general_settings.pass_through_all_models:
return getattr(litellm, f"{route_type}")(**data)
elif (
@@ -154,7 +158,9 @@ async def route_request(
elif user_model is not None:
return getattr(litellm, f"{route_type}")(**data)
elif route_type == "allm_passthrough_route":
return getattr(litellm, f"{route_type}")(**data)
# if no route found then it's a bad request
route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type)
raise ProxyModelNotFoundError(
litellm/router.py
+23 -9
View File
@@ -792,6 +792,7 @@ class Router:
generate_content,
generate_content_stream,
)
self.agenerate_content = self.factory_function(
agenerate_content, call_type="agenerate_content"
)
@@ -2479,6 +2480,7 @@ class Router:
"""
handler_name = original_function.__name__
function_name = "_ageneric_api_call_with_fallbacks"
passthrough_on_no_deployment = kwargs.pop("passthrough_on_no_deployment", False)
self._update_kwargs_before_fallbacks(
model=model,
kwargs=kwargs,
@@ -2491,12 +2493,17 @@ class Router:
f"Inside _ageneric_api_call() - handler: {handler_name}, model: {model}; kwargs: {kwargs}"
)
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs)
deployment = await self.async_get_available_deployment(
model=model,
request_kwargs=kwargs,
messages=kwargs.get("messages", None),
specific_deployment=kwargs.pop("specific_deployment", None),
)
try:
deployment = await self.async_get_available_deployment(
model=model,
request_kwargs=kwargs,
messages=kwargs.get("messages", None),
specific_deployment=kwargs.pop("specific_deployment", None),
)
except Exception as e:
if passthrough_on_no_deployment:
return await original_function(model=model, **kwargs)
raise e
self._update_kwargs_with_deployment(
deployment=deployment, kwargs=kwargs, function_name=function_name
@@ -3323,7 +3330,6 @@ class Router:
"aretrieve_fine_tuning_job",
"alist_files",
"aimage_edit",
"allm_passthrough_route",
"agenerate_content",
"agenerate_content_stream",
):
@@ -3331,6 +3337,12 @@ class Router:
original_function=original_function,
**kwargs,
)
elif call_type == "allm_passthrough_route":
return await self._ageneric_api_call_with_fallbacks(
original_function=original_function,
passthrough_on_no_deployment=True,
**kwargs,
)
elif call_type in (
"aget_responses",
"adelete_responses",
@@ -4143,8 +4155,10 @@ class Router:
)
# Determine cooldown time with priority: deployment config > response header > router default
deployment_cooldown = kwargs.get("litellm_params", {}).get("cooldown_time", None)
deployment_cooldown = kwargs.get("litellm_params", {}).get(
"cooldown_time", None
)
header_cooldown = None
if exception_headers is not None:
header_cooldown = litellm.utils._get_retry_after_from_exception_header(
litellm/types/utils.py
+11
View File
@@ -276,6 +276,8 @@ class CallTypes(Enum):
responses = "responses"
aresponses = "aresponses"
alist_input_items = "alist_input_items"
llm_passthrough_route = "llm_passthrough_route"
allm_passthrough_route = "allm_passthrough_route"
#########################################################
# Google GenAI Native Call Types
@@ -2512,3 +2514,12 @@ class CallbacksByType(TypedDict):
success: List[str]
failure: List[str]
success_and_failure: List[str]
CostResponseTypes = Union[
ModelResponse,
TextCompletionResponse,
EmbeddingResponse,
ImageResponse,
TranscriptionResponse,
]
tests/llm_translation/test_bedrock_completion.py
+144 -16
View File
@@ -3216,7 +3216,9 @@ def test_bedrock_meta_llama_function_calling():
print(response)
def test_bedrock_passthrough():
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_bedrock_passthrough(sync_mode: bool):
import litellm
litellm._turn_on_debug()
@@ -3238,20 +3240,86 @@ def test_bedrock_passthrough():
"anthropic_beta": ["claude-code-20250219"],
}
response = litellm.llm_passthrough_route(
model="bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0",
method="POST",
endpoint="/model/us.anthropic.claude-3-5-sonnet-20240620-v1:0/invoke",
data=data,
)
if sync_mode:
response = litellm.llm_passthrough_route(
model="bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0",
method="POST",
endpoint="/model/us.anthropic.claude-3-5-sonnet-20240620-v1:0/invoke",
data=data,
)
else:
response = await litellm.allm_passthrough_route(
model="bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0",
method="POST",
endpoint="/model/us.anthropic.claude-3-5-sonnet-20240620-v1:0/invoke",
data=data,
)
print(response.text)
assert response.status_code == 200
def test_bedrock_streaming_passthrough():
@pytest.mark.asyncio
async def test_bedrock_converse__streaming_passthrough(monkeypatch):
import litellm
from litellm.integrations.custom_logger import CustomLogger
import asyncio
class MockCustomLogger(CustomLogger):
pass
mock_custom_logger = MockCustomLogger()
monkeypatch.setattr(litellm, "callbacks", [mock_custom_logger])
litellm._turn_on_debug()
data = {
"messages": [
{
"role": "user",
"content": [
{
"text": "Write an article about impact of high inflation to GDP of a country"
}
],
}
],
"system": [{"text": "You are an economist with access to lots of data"}],
"inferenceConfig": {"maxTokens": 100, "temperature": 0.5},
}
with patch.object(mock_custom_logger, "async_log_success_event") as mock_callback:
response = await litellm.allm_passthrough_route(
model="bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0",
method="POST",
endpoint="/model/us.anthropic.claude-3-5-sonnet-20240620-v1:0/converse-stream",
data=data,
)
async for chunk in response:
print(chunk)
await asyncio.sleep(1)
mock_callback.assert_called_once()
print(mock_callback.call_args.kwargs.keys())
assert "response_cost" in mock_callback.call_args.kwargs["kwargs"]
assert mock_callback.call_args.kwargs["kwargs"]["response_cost"] > 0
assert "standard_logging_object" in mock_callback.call_args.kwargs["kwargs"]
@pytest.mark.asyncio
async def test_bedrock_streaming_passthrough(monkeypatch):
import litellm
import time
import asyncio
from unittest.mock import MagicMock
from litellm.integrations.custom_logger import CustomLogger
class MockCustomLogger(CustomLogger):
pass
mock_custom_logger = MockCustomLogger()
monkeypatch.setattr(litellm, "callbacks", [mock_custom_logger])
litellm._turn_on_debug()
@@ -3272,12 +3340,72 @@ def test_bedrock_streaming_passthrough():
"anthropic_beta": ["claude-code-20250219"],
}
response = litellm.llm_passthrough_route(
model="bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0",
method="POST",
endpoint="/model/us.anthropic.claude-3-5-sonnet-20240620-v1:0/invoke-with-response-stream",
data=data,
stream=True,
)
with patch.object(mock_custom_logger, "async_log_success_event") as mock_callback:
response = await litellm.allm_passthrough_route(
model="bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0",
method="POST",
endpoint="/model/us.anthropic.claude-3-5-sonnet-20240620-v1:0/invoke-with-response-stream",
data=data,
)
async for chunk in response:
print(chunk)
assert response.status_code == 200
await asyncio.sleep(1)
mock_callback.assert_called_once()
# check standard logging payload created
print(mock_callback.call_args.kwargs.keys())
assert "standard_logging_object" in mock_callback.call_args.kwargs["kwargs"]
assert "response_cost" in mock_callback.call_args.kwargs["kwargs"]
@pytest.mark.asyncio
async def test_bedrock_streaming_passthrough(monkeypatch):
import litellm
import time
import asyncio
from unittest.mock import MagicMock
from litellm.integrations.custom_logger import CustomLogger
class MockCustomLogger(CustomLogger):
pass
mock_custom_logger = MockCustomLogger()
monkeypatch.setattr(litellm, "callbacks", [mock_custom_logger])
litellm._turn_on_debug()
data = {
"max_tokens": 512,
"messages": [{"role": "user", "content": "Hey"}],
"system": [
{
"type": "text",
"text": "Analyze if this message indicates a new conversation topic. If it does, extract a 2-3 word title that captures the new topic. Format your response as a JSON object with two fields: 'isNewTopic' (boolean) and 'title' (string, or null if isNewTopic is false). Only include these fields, no other text.",
}
],
"temperature": 0,
"metadata": {
"user_id": "5dd07c33da27e6d2968d94ea20bf47a7b090b6b158b82328d54da2909a108e84"
},
"anthropic_version": "bedrock-2023-05-31",
"anthropic_beta": ["claude-code-20250219"],
}
with patch.object(mock_custom_logger, "async_log_success_event") as mock_callback:
response = await litellm.allm_passthrough_route(
model="bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0",
method="POST",
endpoint="/model/us.anthropic.claude-3-5-sonnet-20240620-v1:0/invoke-with-response-stream",
data=data,
)
async for chunk in response:
print(chunk)
await asyncio.sleep(1)
mock_callback.assert_called_once()
# check standard logging payload created
print(mock_callback.call_args.kwargs.keys())
assert "standard_logging_object" in mock_callback.call_args.kwargs["kwargs"]
assert "response_cost" in mock_callback.call_args.kwargs["kwargs"]
tests/test_litellm/llms/azure/test_azure_common_utils.py
+5 -3
View File
@@ -27,7 +27,7 @@ def setup_mocks(monkeypatch):
monkeypatch.delenv("AZURE_TENANT_ID", raising=False)
monkeypatch.delenv("AZURE_SCOPE", raising=False)
monkeypatch.delenv("AZURE_AD_TOKEN", raising=False)
with patch(
"litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id"
) as mock_entra_token, patch(
@@ -429,6 +429,8 @@ def test_select_azure_base_url_called(setup_mocks):
"image_edit",
"agenerate_content_stream",
"agenerate_content",
"allm_passthrough_route",
"llm_passthrough_route",
]
],
)
@@ -1033,8 +1035,8 @@ def test_with_existing_azure_ad_token_from_env(setup_mocks):
# mock get_secret_str("AZURE_AD_TOKEN") to "test-token"
with patch("litellm.llms.azure.common_utils.get_secret_str") as mock_get_secret_str:
# Configure the mock to return "test-token" when called with "AZURE_AD_TOKEN"
mock_get_secret_str.side_effect = (
lambda key: "test-token" if key == "AZURE_AD_TOKEN" else None
mock_get_secret_str.side_effect = lambda key: (
"test-token" if key == "AZURE_AD_TOKEN" else None
)
litellm_params = GenericLiteLLMParams()