Merge pull request #15240 from BerriAI/litellm_dev_10_06_2025_p1

Azure - passthrough support with router models
This commit is contained in:
Krish Dholakia 2025-10-06 20:06:43 -07:00 committed by GitHub
commit 077b5e105f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 463 additions and 132 deletions

View file

@ -665,10 +665,6 @@ class BaseAzureLLM(BaseOpenAILLM):
) -> dict:
litellm_params = litellm_params or GenericLiteLLMParams()
# If api-key is already in headers, preserve it
if "api-key" in headers:
return headers
api_key = (
litellm_params.api_key
or litellm.api_key
@ -693,7 +689,7 @@ class BaseAzureLLM(BaseOpenAILLM):
def _get_base_azure_url(
api_base: Optional[str],
litellm_params: Optional[Union[GenericLiteLLMParams, Dict[str, Any]]],
route: Literal["/openai/responses", "/openai/vector_stores"],
route: Union[Literal["/openai/responses", "/openai/vector_stores"], str],
default_api_version: Optional[Union[str, Literal["latest", "preview"]]] = None,
) -> str:
"""

View file

@ -0,0 +1,85 @@
from typing import TYPE_CHECKING, List, Optional, Tuple
import httpx
from litellm.llms.azure.common_utils import BaseAzureLLM
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
if TYPE_CHECKING:
from httpx import URL
class AzurePassthroughConfig(BasePassthroughConfig):
def is_streaming_request(self, endpoint: str, request_data: dict) -> bool:
return "stream" in request_data
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
endpoint: str,
request_query_params: Optional[dict],
litellm_params: dict,
) -> Tuple["URL", str]:
base_target_url = self.get_api_base(api_base)
if base_target_url is None:
raise Exception("Azure api base not found")
litellm_metadata = litellm_params.get("litellm_metadata") or {}
model_group = litellm_metadata.get("model_group")
if model_group and model_group in endpoint:
endpoint = endpoint.replace(model_group, model)
complete_url = BaseAzureLLM._get_base_azure_url(
api_base=base_target_url,
litellm_params=litellm_params,
route=endpoint,
default_api_version=litellm_params.get("api_version"),
)
return (
httpx.URL(complete_url),
base_target_url,
)
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
return BaseAzureLLM._base_validate_azure_environment(
headers=headers,
litellm_params=GenericLiteLLMParams(
**{**litellm_params, "api_key": api_key}
),
)
@staticmethod
def get_api_base(
api_base: Optional[str] = None,
) -> Optional[str]:
return api_base or get_secret_str("AZURE_API_BASE")
@staticmethod
def get_api_key(
api_key: Optional[str] = None,
) -> Optional[str]:
return api_key or get_secret_str("AZURE_API_KEY")
@staticmethod
def get_base_model(model: str) -> Optional[str]:
return model
def get_models(
self, api_key: Optional[str] = None, api_base: Optional[str] = None
) -> List[str]:
return super().get_models(api_key, api_base)

View file

@ -242,12 +242,14 @@ def llm_passthrough_route(
request_query_params=request_query_params,
litellm_params=litellm_params_dict,
)
# need to encode the id of application-inference-profile for bedrock
# [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 = CommonUtils.encode_bedrock_runtime_modelid_arn(str(updated_url))
encoded_url_str = CommonUtils.encode_bedrock_runtime_modelid_arn(
str(updated_url)
)
updated_url = httpx.URL(encoded_url_str)
# Add or update query parameters
provider_api_key = provider_config.get_api_key(api_key)

View file

@ -21,9 +21,7 @@ from litellm.constants import BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.proxy._types import *
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.user_api_key_auth import (
user_api_key_auth,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
get_form_data,
@ -31,6 +29,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
)
from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
HttpPassThroughEndpointHelpers,
create_pass_through_route,
create_websocket_passthrough_route,
websocket_passthrough_request,
@ -57,7 +56,9 @@ def create_request_copy(request: Request):
}
def is_passthrough_request_using_router_model(request_body: dict, llm_router: Optional[litellm.Router]) -> bool:
def is_passthrough_request_using_router_model(
request_body: dict, llm_router: Optional[litellm.Router]
) -> bool:
"""
Returns True if the model is in the llm_router model names
"""
@ -93,12 +94,16 @@ async def llm_passthrough_factory_proxy_route(
model=None,
)
if provider_config is None:
raise HTTPException(status_code=404, detail=f"Provider {custom_llm_provider} not found")
raise HTTPException(
status_code=404, detail=f"Provider {custom_llm_provider} not found"
)
base_target_url = provider_config.get_api_base()
if base_target_url is None:
raise HTTPException(status_code=404, detail=f"Provider {custom_llm_provider} api base not found")
raise HTTPException(
status_code=404, detail=f"Provider {custom_llm_provider} api base not found"
)
encoded_endpoint = httpx.URL(endpoint).path
@ -177,11 +182,17 @@ async def gemini_proxy_route(
[Docs](https://docs.litellm.ai/docs/pass_through/google_ai_studio)
"""
## CHECK FOR LITELLM API KEY IN THE QUERY PARAMS - ?..key=LITELLM_API_KEY
google_ai_studio_api_key = request.query_params.get("key") or request.headers.get("x-goog-api-key")
google_ai_studio_api_key = request.query_params.get("key") or request.headers.get(
"x-goog-api-key"
)
user_api_key_dict = await user_api_key_auth(request=request, api_key=f"Bearer {google_ai_studio_api_key}")
user_api_key_dict = await user_api_key_auth(
request=request, api_key=f"Bearer {google_ai_studio_api_key}"
)
base_target_url = os.getenv("GEMINI_API_BASE") or "https://generativelanguage.googleapis.com"
base_target_url = (
os.getenv("GEMINI_API_BASE") or "https://generativelanguage.googleapis.com"
)
encoded_endpoint = httpx.URL(endpoint).path
# Ensure endpoint starts with '/' for proper URL construction
@ -293,13 +304,12 @@ async def vllm_proxy_route(
"""
[Docs](https://docs.litellm.ai/docs/pass_through/vllm)
"""
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
HttpPassThroughEndpointHelpers,
)
from litellm.proxy.proxy_server import llm_router
request_body = await get_request_body(request)
is_router_model = is_passthrough_request_using_router_model(request_body, llm_router)
is_router_model = is_passthrough_request_using_router_model(
request_body, llm_router
)
is_streaming_request = is_passthrough_request_streaming(request_body)
if is_router_model and llm_router:
result = cast(
@ -314,7 +324,11 @@ async def vllm_proxy_route(
content=None,
data=None,
files=None,
json=(request_body if request.headers.get("content-type") == "application/json" else None),
json=(
request_body
if request.headers.get("content-type") == "application/json"
else None
),
params=None,
headers=None,
cookies=None,
@ -492,7 +506,9 @@ async def handle_bedrock_count_tokens(
# Extract model from request body
model = request_body.get("model")
if not model:
raise HTTPException(status_code=400, detail={"error": "Model is required in request body"})
raise HTTPException(
status_code=400, detail={"error": "Model is required in request body"}
)
# Get model parameters from router
litellm_params = {"user_api_key_dict": user_api_key_dict}
@ -531,7 +547,9 @@ async def handle_bedrock_count_tokens(
raise
except Exception as e:
verbose_proxy_logger.error(f"Error in handle_bedrock_count_tokens: {str(e)}")
raise HTTPException(status_code=500, detail={"error": f"CountTokens processing error: {str(e)}"})
raise HTTPException(
status_code=500, detail={"error": f"CountTokens processing error: {str(e)}"}
)
async def bedrock_llm_proxy_route(
@ -583,7 +601,8 @@ async def bedrock_llm_proxy_route(
raise HTTPException(
status_code=400,
detail={
"error": "Model missing from endpoint. Expected format: /model/<Model>/<endpoint>. Got: " + endpoint,
"error": "Model missing from endpoint. Expected format: /model/<Model>/<endpoint>. Got: "
+ endpoint,
},
)
@ -647,7 +666,9 @@ async def bedrock_proxy_route(
aws_region_name = litellm.utils.get_secret(secret_name="AWS_REGION_NAME")
if _is_bedrock_agent_runtime_route(endpoint=endpoint): # handle bedrock agents
base_target_url = f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com"
base_target_url = (
f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com"
)
else:
return await bedrock_llm_proxy_route(
endpoint=endpoint,
@ -677,7 +698,9 @@ async def bedrock_proxy_route(
data = await request.json()
except Exception as e:
raise HTTPException(status_code=400, detail={"error": e})
_request = AWSRequest(method="POST", url=str(updated_url), data=json.dumps(data), headers=headers)
_request = AWSRequest(
method="POST", url=str(updated_url), data=json.dumps(data), headers=headers
)
sigv4.add_auth(_request)
prepped = _request.prepare()
@ -738,8 +761,14 @@ async def assemblyai_proxy_route(
[Docs](https://api.assemblyai.com)
"""
# Set base URL based on the route
assembly_region = AssemblyAIPassthroughLoggingHandler._get_assembly_region_from_url(url=str(request.url))
base_target_url = AssemblyAIPassthroughLoggingHandler._get_assembly_base_url_from_region(region=assembly_region)
assembly_region = AssemblyAIPassthroughLoggingHandler._get_assembly_region_from_url(
url=str(request.url)
)
base_target_url = (
AssemblyAIPassthroughLoggingHandler._get_assembly_base_url_from_region(
region=assembly_region
)
)
encoded_endpoint = httpx.URL(endpoint).path
# Ensure endpoint starts with '/' for proper URL construction
if not encoded_endpoint.startswith("/"):
@ -794,17 +823,91 @@ async def azure_proxy_route(
Call any azure endpoint using the proxy.
Just use `{PROXY_BASE_URL}/azure/{endpoint:path}`
Checks if the deployment id in the url is a litellm model name. If so, it will route using the llm_router.allm_passthrough_route.
"""
from litellm.proxy.proxy_server import llm_router
parts = endpoint.split(
"/"
) # azure model is in the url - e.g. https://{endpoint}/openai/deployments/{deployment-id}/completions?api-version=2024-10-21
if len(parts) > 1 and llm_router:
for part in parts:
is_router_model = is_passthrough_request_using_router_model(
request_body={"model": part}, llm_router=llm_router
)
if is_router_model:
request_body = await get_request_body(request)
is_streaming_request = is_passthrough_request_streaming(request_body)
result = await llm_router.allm_passthrough_route(
model=part,
method=request.method,
endpoint=endpoint,
request_query_params=request.query_params,
request_headers=dict(request.headers),
stream=request_body.get("stream", False),
content=None,
data=None,
files=None,
json=(
request_body
if request.headers.get("content-type") == "application/json"
else None
),
params=None,
headers=None,
cookies=None,
)
if is_streaming_request:
# Check if result is an async generator (from _async_streaming)
import inspect
if inspect.isasyncgen(result):
# Result is already an async generator, use it directly
return StreamingResponse(
content=result,
status_code=200,
headers={"content-type": "text/event-stream"},
)
else:
# Result is an httpx.Response, use aiter_bytes()
result = cast(httpx.Response, result)
return StreamingResponse(
content=result.aiter_bytes(),
status_code=result.status_code,
headers=HttpPassThroughEndpointHelpers.get_response_headers(
headers=result.headers,
custom_headers=None,
),
)
# Non-streaming response
result = cast(httpx.Response, result)
content = await result.aread()
return Response(
content=content,
status_code=result.status_code,
headers=HttpPassThroughEndpointHelpers.get_response_headers(
headers=result.headers,
custom_headers=None,
),
)
base_target_url = get_secret_str(secret_name="AZURE_API_BASE")
if base_target_url is None:
raise Exception("Required 'AZURE_API_BASE' in environment to make pass-through calls to Azure.")
raise Exception(
"Required 'AZURE_API_BASE' in environment to make pass-through calls to Azure."
)
# Add or update query parameters
azure_api_key = passthrough_endpoint_router.get_credentials(
custom_llm_provider=litellm.LlmProviders.AZURE.value,
region_name=None,
)
if azure_api_key is None:
raise Exception("Required 'AZURE_API_KEY' in environment to make pass-through calls to Azure.")
raise Exception(
"Required 'AZURE_API_KEY' in environment to make pass-through calls to Azure."
)
return await BaseOpenAIPassThroughHandler._base_openai_pass_through_handler(
endpoint=endpoint,
@ -828,7 +931,9 @@ class BaseVertexAIPassThroughHandler(ABC):
@staticmethod
@abstractmethod
def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: Optional[str]) -> str:
def update_base_target_url_with_credential_location(
base_target_url: str, vertex_location: Optional[str]
) -> str:
pass
@ -838,7 +943,9 @@ class VertexAIDiscoveryPassThroughHandler(BaseVertexAIPassThroughHandler):
return "https://discoveryengine.googleapis.com/"
@staticmethod
def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: Optional[str]) -> str:
def update_base_target_url_with_credential_location(
base_target_url: str, vertex_location: Optional[str]
) -> str:
return base_target_url
@ -848,7 +955,9 @@ class VertexAIPassThroughHandler(BaseVertexAIPassThroughHandler):
return get_vertex_base_url(vertex_location)
@staticmethod
def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: Optional[str]) -> str:
def update_base_target_url_with_credential_location(
base_target_url: str, vertex_location: Optional[str]
) -> str:
return get_vertex_base_url(vertex_location)
@ -914,14 +1023,18 @@ async def _base_vertex_proxy_route(
location=vertex_location,
)
base_target_url = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location)
base_target_url = get_vertex_pass_through_handler.get_default_base_target_url(
vertex_location
)
headers_passed_through = False
# Use headers from the incoming request if no vertex credentials are found
if vertex_credentials is None or vertex_credentials.vertex_project is None:
headers = dict(request.headers) or {}
headers_passed_through = True
verbose_proxy_logger.debug("default_vertex_config not set, incoming request headers %s", headers)
verbose_proxy_logger.debug(
"default_vertex_config not set, incoming request headers %s", headers
)
headers.pop("content-length", None)
headers.pop("host", None)
else:
@ -1087,7 +1200,9 @@ async def openai_proxy_route(
region_name=None,
)
if openai_api_key is None:
raise Exception("Required 'OPENAI_API_KEY' in environment to make pass-through calls to OpenAI.")
raise Exception(
"Required 'OPENAI_API_KEY' in environment to make pass-through calls to OpenAI."
)
return await BaseOpenAIPassThroughHandler._base_openai_pass_through_handler(
endpoint=endpoint,
@ -1133,7 +1248,9 @@ class BaseOpenAIPassThroughHandler:
endpoint_func = create_pass_through_route(
endpoint=endpoint,
target=str(updated_url),
custom_headers=BaseOpenAIPassThroughHandler._assemble_headers(api_key=api_key, request=request),
custom_headers=BaseOpenAIPassThroughHandler._assemble_headers(
api_key=api_key, request=request
),
) # dynamically construct pass-through endpoint based on incoming path
received_value = await endpoint_func(
request,
@ -1150,7 +1267,10 @@ class BaseOpenAIPassThroughHandler:
"""
Appends the OpenAI-Beta header to the headers if the request is an OpenAI Assistants API request
"""
if RouteChecks._is_assistants_api_request(request) is True and "OpenAI-Beta" not in headers:
if (
RouteChecks._is_assistants_api_request(request) is True
and "OpenAI-Beta" not in headers
):
headers["OpenAI-Beta"] = "assistants=v2"
return headers
@ -1166,7 +1286,9 @@ class BaseOpenAIPassThroughHandler:
)
@staticmethod
def _join_url_paths(base_url: httpx.URL, path: str, custom_llm_provider: litellm.LlmProviders) -> str:
def _join_url_paths(
base_url: httpx.URL, path: str, custom_llm_provider: litellm.LlmProviders
) -> str:
"""
Properly joins a base URL with a path, preserving any existing path in the base URL.
"""
@ -1182,9 +1304,14 @@ class BaseOpenAIPassThroughHandler:
joined_path_str = str(base_url.copy_with(path=full_path))
# Apply OpenAI-specific path handling for both branches
if custom_llm_provider == litellm.LlmProviders.OPENAI and "/v1/" not in joined_path_str:
if (
custom_llm_provider == litellm.LlmProviders.OPENAI
and "/v1/" not in joined_path_str
):
# Insert v1 after api.openai.com for OpenAI requests
joined_path_str = joined_path_str.replace("api.openai.com/", "api.openai.com/v1/")
joined_path_str = joined_path_str.replace(
"api.openai.com/", "api.openai.com/v1/"
)
return joined_path_str
@ -1231,9 +1358,7 @@ async def vertex_ai_live_websocket_passthrough(
if vertex_credentials_config is not None:
resolved_project = resolved_project or vertex_credentials_config.vertex_project
temp_location = (
resolved_location or vertex_credentials_config.vertex_location
)
temp_location = resolved_location or vertex_credentials_config.vertex_location
# Ensure resolved_location is a string
if isinstance(temp_location, dict):
resolved_location = str(temp_location)
@ -1241,7 +1366,11 @@ async def vertex_ai_live_websocket_passthrough(
resolved_location = str(temp_location)
else:
resolved_location = None
credentials_value = str(vertex_credentials_config.vertex_credentials) if vertex_credentials_config.vertex_credentials is not None else None
credentials_value = (
str(vertex_credentials_config.vertex_credentials)
if vertex_credentials_config.vertex_credentials is not None
else None
)
try:
resolved_location = resolved_location or (
@ -1302,7 +1431,7 @@ async def vertex_ai_live_websocket_passthrough(
# Use the new WebSocket passthrough pattern
if user_api_key_dict is None:
raise ValueError("user_api_key_dict is required for WebSocket passthrough")
return await websocket_passthrough_request(
websocket=websocket,
target=service_url,

View file

@ -3601,8 +3601,10 @@ def is_known_model(model: Optional[str], llm_router: Optional[Router]) -> bool:
return False
model_names = llm_router.get_model_names()
model_names_set = set(model_names)
is_in_list = False
if model in model_names:
if model in model_names_set:
is_in_list = True
return is_in_list

View file

@ -532,9 +532,6 @@ def get_dynamic_callbacks(
return returned_callbacks
def function_setup( # noqa: PLR0915
original_function: str, rules_obj, start_time, *args, **kwargs
): # just run once to check if user wants to send their data anywhere - PostHog/Sentry/Slack/etc.
@ -802,7 +799,7 @@ def function_setup( # noqa: PLR0915
call_type=call_type,
):
stream = True
logging_obj = get_litellm_logging_class()( # Victim for object pool
logging_obj = get_litellm_logging_class()( # Victim for object pool
model=model, # type: ignore
messages=messages,
stream=stream,
@ -1429,7 +1426,8 @@ def client(original_function): # noqa: PLR0915
if _caching_handler_response is not None:
if (
_caching_handler_response.cached_result is not None
and _caching_handler_response.final_embedding_cached_response is None
and _caching_handler_response.final_embedding_cached_response
is None
):
return _caching_handler_response.cached_result
@ -1699,7 +1697,6 @@ def _is_streaming_request(
return False
def _select_tokenizer(
model: str, custom_tokenizer: Optional[CustomHuggingfaceTokenizer] = None
):
@ -4883,16 +4880,24 @@ def _get_model_info_helper( # noqa: PLR0915
max_input_tokens=_model_info.get("max_input_tokens", None),
max_output_tokens=_model_info.get("max_output_tokens", None),
input_cost_per_token=_input_cost_per_token,
input_cost_per_token_flex=_model_info.get("input_cost_per_token_flex", None),
input_cost_per_token_priority=_model_info.get("input_cost_per_token_priority", None),
input_cost_per_token_flex=_model_info.get(
"input_cost_per_token_flex", None
),
input_cost_per_token_priority=_model_info.get(
"input_cost_per_token_priority", None
),
cache_creation_input_token_cost=_model_info.get(
"cache_creation_input_token_cost", None
),
cache_read_input_token_cost=_model_info.get(
"cache_read_input_token_cost", None
),
cache_read_input_token_cost_flex=_model_info.get("cache_read_input_token_cost_flex", None),
cache_read_input_token_cost_priority=_model_info.get("cache_read_input_token_cost_priority", None),
cache_read_input_token_cost_flex=_model_info.get(
"cache_read_input_token_cost_flex", None
),
cache_read_input_token_cost_priority=_model_info.get(
"cache_read_input_token_cost_priority", None
),
cache_creation_input_token_cost_above_1hr=_model_info.get(
"cache_creation_input_token_cost_above_1hr", None
),
@ -4917,8 +4922,12 @@ def _get_model_info_helper( # noqa: PLR0915
"output_cost_per_token_batches"
),
output_cost_per_token=_output_cost_per_token,
output_cost_per_token_flex=_model_info.get("output_cost_per_token_flex", None),
output_cost_per_token_priority=_model_info.get("output_cost_per_token_priority", None),
output_cost_per_token_flex=_model_info.get(
"output_cost_per_token_flex", None
),
output_cost_per_token_priority=_model_info.get(
"output_cost_per_token_priority", None
),
output_cost_per_audio_token=_model_info.get(
"output_cost_per_audio_token", None
),
@ -6450,7 +6459,7 @@ def get_valid_models(
try:
################################
# init litellm_params
# init litellm_params
#################################
if litellm_params is None:
litellm_params = LiteLLM_Params(model="")
@ -6459,7 +6468,7 @@ def get_valid_models(
if api_base is not None:
litellm_params.api_base = api_base
#################################
check_provider_endpoint = (
check_provider_endpoint or litellm.check_provider_endpoint
)
@ -6934,7 +6943,10 @@ class ProviderConfigManager:
return litellm.LlamaAPIConfig()
elif litellm.LlmProviders.TEXT_COMPLETION_OPENAI == provider:
return litellm.OpenAITextCompletionConfig()
elif litellm.LlmProviders.COHERE_CHAT == provider or litellm.LlmProviders.COHERE == provider:
elif (
litellm.LlmProviders.COHERE_CHAT == provider
or litellm.LlmProviders.COHERE == provider
):
return litellm.CohereChatConfig()
elif litellm.LlmProviders.SNOWFLAKE == provider:
return litellm.SnowflakeConfig()
@ -7361,7 +7373,12 @@ class ProviderConfigManager:
)
return VLLMPassthroughConfig()
elif LlmProviders.AZURE == provider:
from litellm.llms.azure.passthrough.transformation import (
AzurePassthroughConfig,
)
return AzurePassthroughConfig()
return None
@staticmethod
@ -7548,9 +7565,7 @@ class ProviderConfigManager:
return RecraftImageEditConfig()
elif LlmProviders.AZURE_AI == provider:
from litellm.llms.azure_ai.image_edit import (
get_azure_ai_image_edit_config,
)
from litellm.llms.azure_ai.image_edit import get_azure_ai_image_edit_config
return get_azure_ai_image_edit_config(model)
elif LlmProviders.LITELLM_PROXY == provider:
@ -7605,7 +7620,9 @@ def get_end_user_id_for_cost_tracking(
service_type: "litellm_logging" or "prometheus" - used to allow prometheus only disable cost tracking.
"""
_metadata = cast(dict, get_litellm_metadata_from_kwargs(dict(litellm_params=litellm_params)))
_metadata = cast(
dict, get_litellm_metadata_from_kwargs(dict(litellm_params=litellm_params))
)
end_user_id = cast(
Optional[str],

View file

@ -53,27 +53,38 @@ def test_llm_passthrough_route():
def test_bedrock_application_inference_profile_url_encoding():
client = HTTPHandler()
mock_provider_config = MagicMock()
mock_provider_config.get_complete_url.return_value = (
httpx.URL("https://bedrock-runtime.us-east-1.amazonaws.com/model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse"),
"https://bedrock-runtime.us-east-1.amazonaws.com"
httpx.URL(
"https://bedrock-runtime.us-east-1.amazonaws.com/model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse"
),
"https://bedrock-runtime.us-east-1.amazonaws.com",
)
mock_provider_config.get_api_key.return_value = "test-key"
mock_provider_config.validate_environment.return_value = {}
mock_provider_config.sign_request.return_value = ({}, None)
mock_provider_config.is_streaming_request.return_value = False
with patch("litellm.utils.ProviderConfigManager.get_provider_passthrough_config", return_value=mock_provider_config), \
patch("litellm.litellm_core_utils.get_litellm_params.get_litellm_params", return_value={}), \
patch("litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("test-model", "bedrock", "test-key", "test-base")), \
patch.object(client.client, "send", return_value=MagicMock(status_code=200)) as mock_send, \
patch.object(client.client, "build_request") as mock_build_request:
with patch(
"litellm.utils.ProviderConfigManager.get_provider_passthrough_config",
return_value=mock_provider_config,
), patch(
"litellm.litellm_core_utils.get_litellm_params.get_litellm_params",
return_value={},
), patch(
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
return_value=("test-model", "bedrock", "test-key", "test-base"),
), patch.object(
client.client, "send", return_value=MagicMock(status_code=200)
) as mock_send, patch.object(
client.client, "build_request"
) as mock_build_request:
# Mock logging object
mock_logging_obj = MagicMock()
mock_logging_obj.update_environment_variables = MagicMock()
response = llm_passthrough_route(
model="arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd",
endpoint="model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse",
@ -86,7 +97,7 @@ def test_bedrock_application_inference_profile_url_encoding():
# Verify that build_request was called with the encoded URL
mock_build_request.assert_called_once()
call_args = mock_build_request.call_args
# The URL should have the application-inference-profile ID encoded
actual_url = str(call_args.kwargs["url"])
assert "application-inference-profile%2Fr742sbn2zckd" in actual_url
@ -95,28 +106,39 @@ def test_bedrock_application_inference_profile_url_encoding():
def test_bedrock_non_application_inference_profile_no_encoding():
client = HTTPHandler()
# Mock the provider config and its methods
mock_provider_config = MagicMock()
mock_provider_config.get_complete_url.return_value = (
httpx.URL("https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-sonnet-20240229-v1:0/converse"),
"https://bedrock-runtime.us-east-1.amazonaws.com"
httpx.URL(
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-sonnet-20240229-v1:0/converse"
),
"https://bedrock-runtime.us-east-1.amazonaws.com",
)
mock_provider_config.get_api_key.return_value = "test-key"
mock_provider_config.validate_environment.return_value = {}
mock_provider_config.sign_request.return_value = ({}, None)
mock_provider_config.is_streaming_request.return_value = False
with patch("litellm.utils.ProviderConfigManager.get_provider_passthrough_config", return_value=mock_provider_config), \
patch("litellm.litellm_core_utils.get_litellm_params.get_litellm_params", return_value={}), \
patch("litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("test-model", "bedrock", "test-key", "test-base")), \
patch.object(client.client, "send", return_value=MagicMock(status_code=200)) as mock_send, \
patch.object(client.client, "build_request") as mock_build_request:
with patch(
"litellm.utils.ProviderConfigManager.get_provider_passthrough_config",
return_value=mock_provider_config,
), patch(
"litellm.litellm_core_utils.get_litellm_params.get_litellm_params",
return_value={},
), patch(
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
return_value=("test-model", "bedrock", "test-key", "test-base"),
), patch.object(
client.client, "send", return_value=MagicMock(status_code=200)
) as mock_send, patch.object(
client.client, "build_request"
) as mock_build_request:
# Mock logging object
mock_logging_obj = MagicMock()
mock_logging_obj.update_environment_variables = MagicMock()
response = llm_passthrough_route(
model="anthropic.claude-3-sonnet-20240229-v1:0",
endpoint="model/anthropic.claude-3-sonnet-20240229-v1:0/converse",
@ -129,7 +151,7 @@ def test_bedrock_non_application_inference_profile_no_encoding():
# Verify that build_request was called with the original URL (no encoding)
mock_build_request.assert_called_once()
call_args = mock_build_request.call_args
# The URL should NOT have application-inference-profile encoding
actual_url = str(call_args.kwargs["url"])
assert "application-inference-profile%2F" not in actual_url
@ -151,21 +173,21 @@ def test_update_stream_param_based_on_request_body():
parsed_body=parsed_body, stream=False
)
assert result is True
# Test 2: no stream in request body should return original stream param
parsed_body = {"model": "test-model"}
result = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body(
parsed_body=parsed_body, stream=False
)
assert result is False
# Test 3: stream=False in request body should return False
parsed_body = {"stream": False, "model": "test-model"}
result = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body(
parsed_body=parsed_body, stream=True
)
assert result is False
# Test 4: no stream param provided, no stream in body
parsed_body = {"model": "test-model"}
result = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body(
@ -178,14 +200,14 @@ def test_update_stream_param_based_on_request_body():
def mock_request():
"""Create a mock request with headers"""
from typing import Optional
class QueryParams:
def __init__(self):
self._dict = {}
def __iter__(self):
return iter(self._dict)
def items(self):
return self._dict.items()
@ -210,6 +232,7 @@ def mock_request():
def mock_user_api_key_dict():
"""Create a mock user API key dictionary"""
from litellm.proxy._types import UserAPIKeyAuth
return UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
@ -223,8 +246,8 @@ async def test_pass_through_request_stream_param_override(
mock_request, mock_user_api_key_dict
):
"""
Test that when stream=None is passed as parameter but stream=True
is in request body, the request body value takes precedence and
Test that when stream=None is passed as parameter but stream=True
is in request body, the request body value takes precedence and
the eventual POST request uses streaming.
"""
from unittest.mock import AsyncMock, Mock, patch
@ -238,29 +261,29 @@ async def test_pass_through_request_stream_param_override(
"model": "claude-3-5-sonnet-20241022",
"max_tokens": 256,
"messages": [{"role": "user", "content": "Hello, world"}],
"stream": True # This should override the function parameter
"stream": True, # This should override the function parameter
}
# Create a mock streaming response
mock_response = AsyncMock()
mock_response.status_code = 200
mock_response.headers = {"content-type": "text/event-stream"}
# Mock the streaming response behavior
async def mock_aiter_bytes():
yield b'data: {"content": "Hello"}\n\n'
yield b'data: {"content": "World"}\n\n'
yield b'data: [DONE]\n\n'
yield b"data: [DONE]\n\n"
mock_response.aiter_bytes = mock_aiter_bytes
# Create mocks for the async client
mock_async_client = AsyncMock()
mock_request_obj = AsyncMock()
# Mock build_request to return a request object (it's a sync method)
mock_async_client.build_request = Mock(return_value=mock_request_obj)
# Mock send to return the streaming response
mock_async_client.send.return_value = mock_response
@ -269,9 +292,7 @@ async def test_pass_through_request_stream_param_override(
mock_client_obj.client = mock_async_client
# Create the request
request = mock_request(
headers={}, method="POST", request_body=request_body
)
request = mock_request(headers={}, method="POST", request_body=request_body)
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client",
@ -298,33 +319,32 @@ async def test_pass_through_request_stream_param_override(
httpx.URL("https://api.anthropic.com/v1/messages"),
json=request_body,
params={},
headers={
"Authorization": "Bearer test-key"
},
headers={"Authorization": "Bearer test-key"},
)
# Verify that send was called with stream=True
mock_async_client.send.assert_called_once_with(
mock_request_obj,
stream=True # This proves that stream=True from request body was used
mock_request_obj,
stream=True, # This proves that stream=True from request body was used
)
# Verify that the non-streaming request method was NOT called
mock_async_client.request.assert_not_called()
# Verify response is a StreamingResponse
from fastapi.responses import StreamingResponse
assert isinstance(response, StreamingResponse)
assert response.status_code == 200
@pytest.mark.asyncio
@pytest.mark.asyncio
async def test_pass_through_request_stream_param_no_override(
mock_request, mock_user_api_key_dict
):
"""
Test that when stream=False is passed as parameter and no stream
is in request body, the function parameter is used and
Test that when stream=False is passed as parameter and no stream
is in request body, the function parameter is used and
the eventual request uses non-streaming.
"""
from unittest.mock import AsyncMock, Mock, patch
@ -335,7 +355,7 @@ async def test_pass_through_request_stream_param_no_override(
# Create request body without stream parameter
request_body = {
"model": "claude-3-5-sonnet-20241022",
"model": "claude-3-5-sonnet-20241022",
"max_tokens": 256,
"messages": [{"role": "user", "content": "Hello, world"}],
# No stream parameter - should use function parameter stream=False
@ -346,15 +366,15 @@ async def test_pass_through_request_stream_param_no_override(
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response._content = b'{"response": "Hello world"}'
async def mock_aread():
return mock_response._content
mock_response.aread = mock_aread
# Create mocks for the async client
mock_async_client = AsyncMock()
# Mock request to return the non-streaming response
mock_async_client.request.return_value = mock_response
@ -363,9 +383,7 @@ async def test_pass_through_request_stream_param_no_override(
mock_client_obj.client = mock_async_client
# Create the request
request = mock_request(
headers={}, method="POST", request_body=request_body
)
request = mock_request(headers={}, method="POST", request_body=request_body)
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client",
@ -388,23 +406,105 @@ async def test_pass_through_request_stream_param_no_override(
# Verify that build_request was NOT called (no streaming path)
mock_async_client.build_request.assert_not_called()
# Verify that send was NOT called (no streaming path)
mock_async_client.send.assert_not_called()
# Verify that the non-streaming request method WAS called
mock_async_client.request.assert_called_once_with(
method="POST",
url=httpx.URL("https://api.anthropic.com/v1/messages"),
headers={
"Authorization": "Bearer test-key"
},
headers={"Authorization": "Bearer test-key"},
params={},
json=request_body,
)
# Verify response is a regular Response (not StreamingResponse)
from fastapi.responses import Response, StreamingResponse
assert not isinstance(response, StreamingResponse)
assert isinstance(response, Response)
assert response.status_code == 200
assert response.status_code == 200
def test_azure_with_custom_api_base_and_key():
"""
Test that llm_passthrough_route correctly handles Azure OpenAI
with custom api_base and api_key.
"""
client = HTTPHandler()
# Mock the provider config and its methods
mock_provider_config = MagicMock()
mock_provider_config.get_complete_url.return_value = (
httpx.URL(
"https://my-custom-base/openai/deployments/gpt-4.1/chat/completions?api-version=2024-02-01"
),
"https://my-custom-base",
)
mock_provider_config.get_api_key.return_value = "my-custom-key"
mock_provider_config.validate_environment.return_value = {
"api-key": "my-custom-key"
}
mock_provider_config.sign_request.return_value = (
{"api-key": "my-custom-key"},
None,
)
mock_provider_config.is_streaming_request.return_value = False
with patch(
"litellm.utils.ProviderConfigManager.get_provider_passthrough_config",
return_value=mock_provider_config,
), patch(
"litellm.litellm_core_utils.get_litellm_params.get_litellm_params",
return_value={},
), patch(
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
return_value=("gpt-4.1", "azure", "my-custom-key", "https://my-custom-base"),
), patch.object(
client.client,
"send",
return_value=MagicMock(
status_code=200, json=lambda: {"id": "chatcmpl-123", "choices": []}
),
) as mock_send, patch.object(
client.client, "build_request"
) as mock_build_request:
# Mock logging object
mock_logging_obj = MagicMock()
mock_logging_obj.update_environment_variables = MagicMock()
response = llm_passthrough_route(
model="azure/gpt-4.1",
endpoint="openai/deployments/gpt-4.1/chat/completions",
method="POST",
custom_llm_provider="azure",
api_base="https://my-custom-base",
api_key="my-custom-key",
json={
"model": "gpt-4.1",
"messages": [{"role": "user", "content": "Hello!"}],
},
client=client,
litellm_logging_obj=mock_logging_obj,
)
# Verify that build_request was called with the correct parameters
mock_build_request.assert_called_once()
call_args = mock_build_request.call_args
# Verify the URL contains the custom base
actual_url = str(call_args.kwargs["url"])
assert "my-custom-base" in actual_url
assert "gpt-4.1" in actual_url
# Verify the headers contain the custom API key
headers = call_args.kwargs["headers"]
assert headers["api-key"] == "my-custom-key"
# Verify the model in JSON body is updated
json_body = call_args.kwargs["json"]
assert json_body["model"] == "gpt-4.1"
assert response.status_code == 200