mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
Merge pull request #15240 from BerriAI/litellm_dev_10_06_2025_p1
Azure - passthrough support with router models
This commit is contained in:
commit
077b5e105f
7 changed files with 463 additions and 132 deletions
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
85
litellm/llms/azure/passthrough/transformation.py
Normal file
85
litellm/llms/azure/passthrough/transformation.py
Normal 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)
|
||||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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],
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue