ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(router.py): write file to all deployments (#10708)

Browse Source 
* fix(router.py): write file to all deployments

allows unified file id to work across multiple deployments

* fix(view_logs/index.tsx): show call type in request logs

* fix(router.py): pass a deep copy of kwargs to avoid conflict across multiple runs

* fix(batch_utils.py): broaden check

* fix(router_utils.py): handle null type for function name

* fix(proxy_track_cost_callback.py): fix ruff check error

* fix(router.py): handle healthy_deployments as a dict

* feat(managed_files.py): support encoding / decoding unified batch id … (#10711)

* feat(managed_files.py): support encoding / decoding unified batch id when using managed files

allows routing retrieve batch to the right model id

* fix: fix linting error

* test: add unit tests

* fix: fix ruff check
This commit is contained in:
Krish Dholakia 2025-05-10 00:08:30 -07:00 committed by GitHub
parent 3fafe37eb9
commit 9bfd3e4819
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
17 changed files with 400 additions and 153 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

50
enterprise/enterprise_hooks/managed_files.py
View File

@ -23,7 +23,7 @@ from litellm.types.llms.openai import (
OpenAIFileObject,
OpenAIFilesPurpose,
)
from litellm.types.utils import SpecialEnums
from litellm.types.utils import LiteLLMBatch, LLMResponseTypes, SpecialEnums
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -136,6 +136,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"pass_through_endpoint",
"rerank",
"acreate_batch",
"aretrieve_batch",
],
) -> Union[Exception, str, Dict, None]:
"""
@ -161,6 +162,21 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
data["model_file_id_mapping"] = model_file_id_mapping
elif call_type == CallTypes.aretrieve_batch.value:
retrieve_batch_id = cast(Optional[str], data.get("batch_id"))
potential_batch_id = (
_is_base64_encoded_unified_file_id(retrieve_batch_id)
if retrieve_batch_id
else False
)
if potential_batch_id:
## for managed batch id - get the model id
model_id = self.get_model_id_from_unified_batch_id(potential_batch_id)
data["model"] = model_id
data["batch_id"] = self.get_batch_id_from_unified_batch_id(
potential_batch_id
)
return data
async def async_pre_call_deployment_hook(
@ -340,6 +356,38 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
return response
def get_unified_batch_id(self, batch_id: str, model_id: str) -> str:
unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(
model_id, batch_id
)
return base64.urlsafe_b64encode(unified_batch_id.encode()).decode().rstrip("=")
def get_model_id_from_unified_batch_id(self, file_id: str) -> str:
## use regex to get the model_id from the file_id
return file_id.split("model_id:")[1].split(";")[0]
def get_batch_id_from_unified_batch_id(self, file_id: str) -> str:
## use regex to get the batch_id from the file_id
return file_id.split("llm_batch_id:")[1].split(",")[0]
async def async_post_call_success_hook(
self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
) -> Any:
if isinstance(response, LiteLLMBatch):
## Check if unified_file_id is in the response
unified_batch_id = response._hidden_params.get(
"unified_file_id"
) # managed file id
model_id = response._hidden_params.get("model_id")
if unified_batch_id and model_id:
response.id = self.get_unified_batch_id(
batch_id=response.id, model_id=model_id
)
return await super().async_post_call_success_hook(
data, user_api_key_dict, response
)
async def afile_retrieve(
self, file_id: str, litellm_parent_otel_span: Optional[Span]
) -> OpenAIFileObject:

2
litellm/batches/main.py
View File

@ -51,7 +51,7 @@ async def acreate_batch(
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
) -> Batch:
) -> LiteLLMBatch:
"""
Async: Creates and executes a batch from an uploaded file of request

1
litellm/litellm_core_utils/litellm_logging.py
View File

@ -3674,6 +3674,7 @@ def get_standard_logging_object_payload(
or litellm_params.get("metadata", None)
or {}
)
completion_start_time = kwargs.get("completion_start_time", end_time)
call_type = kwargs.get("call_type")
cache_hit = kwargs.get("cache_hit", False)

1
litellm/proxy/_experimental/out/onboarding.html
View File

File diff suppressed because one or more lines are too long

15
litellm/proxy/_new_secret_config.yaml
View File

@ -56,4 +56,17 @@ model_list:
model: gpt-image-1
api_key: os.environ/OPENAI_API_KEY
# drop_params: true
- model_name: "gpt-4o-batch"
litellm_params:
model: azure/gpt-4o-mini
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
model_info:
id: my-general-azure-deployment
- model_name: "gpt-4o-batch"
litellm_params:
model: azure/gpt-4o-mini
api_base: https://krris-m2f9a9i7-eastus2.openai.azure.com
api_key: 04d22fb7e9ad4d9c8afe7c6abf97a6fc
model_info:
id: my-unique-azure-deployment

35
litellm/proxy/batches_endpoints/endpoints.py
View File

@ -126,7 +126,7 @@ async def create_batch(
response = await llm_router.acreate_batch(**_create_batch_data) # type: ignore
elif (
unified_file_id
unified_file_id and input_file_id
): # litellm_proxy:application/octet-stream;unified_id,c4843482-b176-4901-8292-7523fd0f2c6e;target_model_names,gpt-4o-mini
target_model_names = get_models_from_unified_file_id(unified_file_id)
## EXPECTS 1 MODEL
@ -150,11 +150,18 @@ async def create_batch(
)
response = await llm_router.acreate_batch(**_create_batch_data)
response.input_file_id = input_file_id
response._hidden_params["unified_file_id"] = unified_file_id
else:
response = await litellm.acreate_batch(
custom_llm_provider=custom_llm_provider, **_create_batch_data # type: ignore
)
### CALL HOOKS ### - modify outgoing data
response = await proxy_logging_obj.post_call_success_hook(
data=data, user_api_key_dict=user_api_key_dict, response=response
)
### ALERTING ###
asyncio.create_task(
proxy_logging_obj.update_request_status(
@ -231,7 +238,6 @@ async def retrieve_batch(
```
"""
from litellm.proxy.proxy_server import (
add_litellm_data_to_request,
general_settings,
llm_router,
proxy_config,
@ -248,22 +254,24 @@ async def retrieve_batch(
data = cast(dict, _retrieve_batch_request)
# setup logging
data["litellm_call_id"] = request.headers.get(
"x-litellm-call-id", str(uuid.uuid4())
)
# Include original request and headers in the data
data = await add_litellm_data_to_request(
data=data,
base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data)
(
data,
litellm_logging_obj,
) = await base_llm_response_processor.common_processing_pre_call_logic(
request=request,
general_settings=general_settings,
user_api_key_dict=user_api_key_dict,
version=version,
proxy_logging_obj=proxy_logging_obj,
proxy_config=proxy_config,
route_type="aretrieve_batch",
)
if litellm.enable_loadbalancing_on_batch_endpoints is True:
if (
litellm.enable_loadbalancing_on_batch_endpoints is True
or data.get("model") is not None
):
if llm_router is None:
raise HTTPException(
status_code=500,
@ -283,6 +291,11 @@ async def retrieve_batch(
custom_llm_provider=custom_llm_provider, **data # type: ignore
)
### CALL HOOKS ### - modify outgoing data
response = await proxy_logging_obj.post_call_success_hook(
data=data, user_api_key_dict=user_api_key_dict, response=response
)
### ALERTING ###
asyncio.create_task(
proxy_logging_obj.update_request_status(

1
litellm/proxy/common_request_processing.py
View File

@ -115,6 +115,7 @@ class ProxyBaseLLMRequestProcessing:
"aget_responses",
"adelete_responses",
"acreate_batch",
"aretrieve_batch",
],
version: Optional[str] = None,
user_model: Optional[str] = None,

8
litellm/proxy/hooks/proxy_track_cost_callback.py
View File

@ -96,11 +96,7 @@ class _ProxyDBLogger(CustomLogger):
start_time=None,
end_time=None, # start/end time for completion
):
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
update_cache,
)
from litellm.proxy.proxy_server import proxy_logging_obj, update_cache
verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback")
try:
@ -134,7 +130,7 @@ class _ProxyDBLogger(CustomLogger):
)
verbose_proxy_logger.debug(
f"user_api_key {user_api_key}, prisma_client: {prisma_client}"
f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}"
)
if _should_track_cost_callback(
user_api_key=user_api_key,

2
litellm/proxy/litellm_pre_call_utils.py
View File

@ -67,7 +67,9 @@ def _get_metadata_variable_name(request: Request) -> str:
"batches",
"/v1/messages",
"responses",
"files",
]
if any(
[
litellm_metadata_route in request.url.path

316
litellm/router.py
View File

@ -106,12 +106,7 @@ from litellm.router_utils.router_callbacks.track_deployment_metrics import (
increment_deployment_successes_for_current_minute,
)
from litellm.scheduler import FlowItem, Scheduler
from litellm.types.llms.openai import (
AllMessageValues,
Batch,
FileTypes,
OpenAIFileObject,
)
from litellm.types.llms.openai import AllMessageValues, FileTypes, OpenAIFileObject
from litellm.types.router import (
CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS,
VALID_LITELLM_ENVIRONMENTS,
@ -135,7 +130,7 @@ from litellm.types.router import (
RoutingStrategy,
)
from litellm.types.services import ServiceTypes
from litellm.types.utils import GenericBudgetConfigType
from litellm.types.utils import GenericBudgetConfigType, LiteLLMBatch
from litellm.types.utils import ModelInfo
from litellm.types.utils import ModelInfo as ModelMapInfo
from litellm.types.utils import StandardLoggingPayload
@ -2743,68 +2738,87 @@ class Router:
f"Inside _atext_completion()- model: {model}; kwargs: {kwargs}"
)
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs)
deployment = await self.async_get_available_deployment(
healthy_deployments = await self.async_get_healthy_deployments(
model=model,
messages=[{"role": "user", "content": "files-api-fake-text"}],
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
data = deployment["litellm_params"].copy()
model_name = data["model"]
model_client = self._get_async_openai_model_client(
deployment=deployment,
kwargs=kwargs,
)
self.total_calls[model_name] += 1
## REPLACE MODEL IN FILE WITH SELECTED DEPLOYMENT ##
stripped_model, custom_llm_provider, _, _ = get_llm_provider(
model=data["model"]
parent_otel_span=parent_otel_span,
)
response = litellm.acreate_file(
**{
**data,
"custom_llm_provider": custom_llm_provider,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
async def create_file_for_deployment(deployment: dict) -> OpenAIFileObject:
kwargs_copy = copy.deepcopy(kwargs)
self._update_kwargs_with_deployment(
deployment=deployment,
kwargs=kwargs_copy,
function_name="acreate_file",
)
data = deployment["litellm_params"].copy()
model_name = data["model"]
rpm_semaphore = self._get_client(
deployment=deployment,
kwargs=kwargs,
client_type="max_parallel_requests",
)
model_client = self._get_async_openai_model_client(
deployment=deployment,
kwargs=kwargs_copy,
)
self.total_calls[model_name] += 1
if rpm_semaphore is not None and isinstance(
rpm_semaphore, asyncio.Semaphore
):
async with rpm_semaphore:
"""
- Check rpm limits before making the call
- If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe)
"""
## REPLACE MODEL IN FILE WITH SELECTED DEPLOYMENT ##
stripped_model, custom_llm_provider, _, _ = get_llm_provider(
model=data["model"]
)
response = litellm.acreate_file(
**{
**data,
"custom_llm_provider": custom_llm_provider,
"caching": self.cache_responses,
"client": model_client,
**kwargs_copy,
}
)
rpm_semaphore = self._get_client(
deployment=deployment,
kwargs=kwargs_copy,
client_type="max_parallel_requests",
)
if rpm_semaphore is not None and isinstance(
rpm_semaphore, asyncio.Semaphore
):
async with rpm_semaphore:
"""
- Check rpm limits before making the call
- If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe)
"""
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
)
response = await response # type: ignore
else:
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
)
response = await response # type: ignore
else:
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
self.success_calls[model_name] += 1
verbose_router_logger.info(
f"litellm.acreate_file(model={model_name})\033[32m 200 OK\033[0m"
)
response = await response # type: ignore
self.success_calls[model_name] += 1
verbose_router_logger.info(
f"litellm.acreate_file(model={model_name})\033[32m 200 OK\033[0m"
)
return response
return response # type: ignore
tasks = []
if isinstance(healthy_deployments, dict):
tasks.append(create_file_for_deployment(healthy_deployments))
else:
for deployment in healthy_deployments:
tasks.append(create_file_for_deployment(deployment))
responses = await asyncio.gather(*tasks)
if len(responses) == 0:
raise Exception("No healthy deployments found.")
return responses[0]
except Exception as e:
verbose_router_logger.exception(
f"litellm.acreate_file(model={model}, {kwargs})\033[31m Exception {str(e)}\033[0m"
@ -2817,7 +2831,7 @@ class Router:
self,
model: str,
**kwargs,
) -> Batch:
) -> LiteLLMBatch:
try:
kwargs["model"] = model
kwargs["original_function"] = self._acreate_batch
@ -2848,7 +2862,7 @@ class Router:
self,
model: str,
**kwargs,
) -> Batch:
) -> LiteLLMBatch:
try:
verbose_router_logger.debug(
f"Inside _acreate_batch()- model: {model}; kwargs: {kwargs}"
@ -2861,7 +2875,6 @@ class Router:
request_kwargs=kwargs,
)
kwargs["model_info"] = deployment.get("model_info", {})
data = deployment["litellm_params"].copy()
model_name = data["model"]
self._update_kwargs_with_deployment(
@ -2913,8 +2926,9 @@ class Router:
self.success_calls[model_name] += 1
verbose_router_logger.info(
f"litellm.acreate_file(model={model_name})\033[32m 200 OK\033[0m"
f"litellm.acreate_batch(model={model_name})\033[32m 200 OK\033[0m"
)
return response # type: ignore
except Exception as e:
verbose_router_logger.exception(
@ -2926,26 +2940,44 @@ class Router:
async def aretrieve_batch(
self,
model: Optional[str] = None,
**kwargs,
) -> Batch:
) -> LiteLLMBatch:
"""
Iterate through all models in a model group to check for batch
Future Improvement - cache the result.
"""
try:
filtered_model_list = self.get_model_list()
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs)
if model is not None:
filtered_model_list: Optional[
Union[List[DeploymentTypedDict], List[Dict], Dict]
] = await self.async_get_healthy_deployments(
model=model,
messages=[{"role": "user", "content": "retrieve-api-fake-text"}],
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
parent_otel_span=parent_otel_span,
)
else:
filtered_model_list = self.get_model_list()
if filtered_model_list is None:
raise Exception("Router not yet initialized.")
receieved_exceptions = []
async def try_retrieve_batch(model_name):
async def try_retrieve_batch(model_name: DeploymentTypedDict):
try:
model = model_name["litellm_params"].get("model")
if model is None:
raise Exception(
f"Model not found in litellm_params for deployment: {model_name}"
)
# Update kwargs with the current model name or any other model-specific adjustments
## SET CUSTOM PROVIDER TO SELECTED DEPLOYMENT ##
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model=model_name["litellm_params"]["model"]
model=model
)
new_kwargs = copy.deepcopy(kwargs)
new_kwargs.pop("custom_llm_provider", None)
@ -2957,15 +2989,35 @@ class Router:
return None
# Check all models in parallel
results = await asyncio.gather(
*[try_retrieve_batch(model) for model in filtered_model_list],
return_exceptions=True,
)
if (
filtered_model_list is not None
and isinstance(filtered_model_list, list)
and len(filtered_model_list) > 0
):
results = await asyncio.gather(
*[
try_retrieve_batch(cast(DeploymentTypedDict, model))
for model in filtered_model_list
],
return_exceptions=True,
)
elif filtered_model_list is not None and isinstance(
filtered_model_list, dict
):
results = await try_retrieve_batch(
cast(DeploymentTypedDict, filtered_model_list)
)
else:
raise Exception("No healthy deployments found.")
# Check for successful responses and handle exceptions
for result in results:
if isinstance(result, Batch):
return result
if results is not None:
if isinstance(results, LiteLLMBatch):
return results
elif isinstance(results, list):
for result in results:
if isinstance(result, LiteLLMBatch):
return result
# If no valid Batch response was found, raise the first encountered exception
if receieved_exceptions:
@ -5936,6 +5988,80 @@ class Router:
return model, healthy_deployments
async def async_get_healthy_deployments(
self,
model: str,
request_kwargs: Dict,
messages: Optional[List[Dict[str, str]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
parent_otel_span: Optional[Span] = None,
) -> Union[List[Dict], Dict]:
"""
Get the healthy deployments for a model.
Returns:
- List[Dict], if multiple models chosen
*OR*
- Dict, if specific model chosen
"""
model, healthy_deployments = self._common_checks_available_deployment(
model=model,
messages=messages,
input=input,
specific_deployment=specific_deployment,
) # type: ignore
if isinstance(healthy_deployments, dict):
return healthy_deployments
cooldown_deployments = await _async_get_cooldown_deployments(
litellm_router_instance=self, parent_otel_span=parent_otel_span
)
verbose_router_logger.debug(
f"async cooldown deployments: {cooldown_deployments}"
)
verbose_router_logger.debug(f"cooldown_deployments: {cooldown_deployments}")
healthy_deployments = self._filter_cooldown_deployments(
healthy_deployments=healthy_deployments,
cooldown_deployments=cooldown_deployments,
)
healthy_deployments = await self.async_callback_filter_deployments(
model=model,
healthy_deployments=healthy_deployments,
messages=(
cast(List[AllMessageValues], messages) if messages is not None else None
),
request_kwargs=request_kwargs,
parent_otel_span=parent_otel_span,
)
if self.enable_pre_call_checks and messages is not None:
healthy_deployments = self._pre_call_checks(
model=model,
healthy_deployments=cast(List[Dict], healthy_deployments),
messages=messages,
request_kwargs=request_kwargs,
)
# check if user wants to do tag based routing
healthy_deployments = await get_deployments_for_tag( # type: ignore
llm_router_instance=self,
model=model,
request_kwargs=request_kwargs,
healthy_deployments=healthy_deployments,
)
if len(healthy_deployments) == 0:
exception = await async_raise_no_deployment_exception(
litellm_router_instance=self,
model=model,
parent_otel_span=parent_otel_span,
)
raise exception
return healthy_deployments
async def async_get_available_deployment(
self,
model: str,
@ -5965,61 +6091,17 @@ class Router:
)
try:
parent_otel_span = _get_parent_otel_span_from_kwargs(request_kwargs)
model, healthy_deployments = self._common_checks_available_deployment(
healthy_deployments = await self.async_get_healthy_deployments(
model=model,
request_kwargs=request_kwargs,
messages=messages,
input=input,
specific_deployment=specific_deployment,
) # type: ignore
parent_otel_span=parent_otel_span,
)
if isinstance(healthy_deployments, dict):
return healthy_deployments
cooldown_deployments = await _async_get_cooldown_deployments(
litellm_router_instance=self, parent_otel_span=parent_otel_span
)
verbose_router_logger.debug(
f"async cooldown deployments: {cooldown_deployments}"
)
verbose_router_logger.debug(f"cooldown_deployments: {cooldown_deployments}")
healthy_deployments = self._filter_cooldown_deployments(
healthy_deployments=healthy_deployments,
cooldown_deployments=cooldown_deployments,
)
healthy_deployments = await self.async_callback_filter_deployments(
model=model,
healthy_deployments=healthy_deployments,
messages=(
cast(List[AllMessageValues], messages)
if messages is not None
else None
),
request_kwargs=request_kwargs,
parent_otel_span=parent_otel_span,
)
if self.enable_pre_call_checks and messages is not None:
healthy_deployments = self._pre_call_checks(
model=model,
healthy_deployments=cast(List[Dict], healthy_deployments),
messages=messages,
request_kwargs=request_kwargs,
)
# check if user wants to do tag based routing
healthy_deployments = await get_deployments_for_tag( # type: ignore
llm_router_instance=self,
model=model,
request_kwargs=request_kwargs,
healthy_deployments=healthy_deployments,
)
if len(healthy_deployments) == 0:
exception = await async_raise_no_deployment_exception(
litellm_router_instance=self,
model=model,
parent_otel_span=parent_otel_span,
)
raise exception
start_time = time.time()
if (
self.routing_strategy == "usage-based-routing-v2"

4
litellm/router_strategy/lowest_tpm_rpm.py
View File

@ -106,6 +106,8 @@ class LowestTPMLoggingHandler(CustomLogger):
elif isinstance(id, int):
id = str(id)
if "usage" not in response_obj:
return
total_tokens = response_obj["usage"]["total_tokens"]
# ------------
@ -144,7 +146,7 @@ class LowestTPMLoggingHandler(CustomLogger):
if self.test_flag:
self.logged_success += 1
except Exception as e:
verbose_router_logger.error(
verbose_router_logger.exception(
"litellm.router_strategy.lowest_tpm_rpm.py::async_log_success_event(): Exception occured - {}".format(
str(e)
)

8
litellm/router_utils/batch_utils.py
View File

@ -48,7 +48,7 @@ def replace_model_in_jsonl(
return None
def _get_router_metadata_variable_name(function_name) -> str:
def _get_router_metadata_variable_name(function_name: Optional[str]) -> str:
"""
Helper to return what the "metadata" field should be called in the request data
@ -57,9 +57,11 @@ def _get_router_metadata_variable_name(function_name) -> str:
For ALL other endpoints we call this "metadata
"""
ROUTER_METHODS_USING_LITELLM_METADATA = set(
["batch", "generic_api_call", "_acreate_batch"]
["batch", "generic_api_call", "_acreate_batch", "file"]
)
if function_name in ROUTER_METHODS_USING_LITELLM_METADATA:
if function_name and any(
method in function_name for method in ROUTER_METHODS_USING_LITELLM_METADATA
):
return "litellm_metadata"
else:
return "metadata"

4
litellm/types/utils.py
View File

@ -2316,9 +2316,11 @@ class SpecialEnums(Enum):
"litellm:custom_llm_provider:{};model_id:{};response_id:{}"
)
LITELLM_MANAGED_BATCH_COMPLETE_STR = "litellm_proxy;model_id:{};llm_batch_id:{}"
LLMResponseTypes = Union[
ModelResponse, EmbeddingResponse, ImageResponse, OpenAIFileObject
ModelResponse, EmbeddingResponse, ImageResponse, OpenAIFileObject, LiteLLMBatch
]

18
tests/enterprise/enterprise_hooks/test_managed_files.py
View File

@ -40,6 +40,24 @@ def test_get_file_ids_from_messages():
]
@pytest.mark.asyncio
async def test_async_pre_call_hook_batch_retrieve():
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=MagicMock()
)
data = {
"user_api_key_dict": {"parent_otel_span": MagicMock()},
"data": {
"batch_id": "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1nZW5lcmFsLWF6dXJlLWRlcGxveW1lbnQ7bGxtX2JhdGNoX2lkOmJhdGNoX2EzMjJiNmJhLWFjN2UtNDg4OC05MjljLTFhZDM0NDJmMDZlZA",
},
"call_type": "aretrieve_batch",
"cache": MagicMock(),
}
response = await proxy_managed_files.async_pre_call_hook(**data)
assert response["batch_id"] == "batch_a322b6ba-ac7e-4888-929c-1ad3442f06ed"
assert response["model"] == "my-general-azure-deployment"
# def test_list_managed_files():
# proxy_managed_files = _PROXY_LiteLLMManagedFiles(DualCache())

57
tests/litellm/test_router.py
View File

@ -118,3 +118,60 @@ async def test_router_with_tags_and_fallbacks():
mock_testing_fallbacks=True,
metadata={"tags": ["test"]},
)
@pytest.mark.asyncio
async def test_router_acreate_file():
"""
Write to all deployments of a model
"""
from unittest.mock import MagicMock, call, patch
router = litellm.Router(
model_list=[
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "gpt-3.5-turbo"},
},
{"model_name": "gpt-3.5-turbo", "litellm_params": {"model": "gpt-4o-mini"}},
],
)
with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file:
mock_acreate_file.return_value = MagicMock()
response = await router.acreate_file(
model="gpt-3.5-turbo",
purpose="test",
file=MagicMock(),
)
# assert that the mock_acreate_file was called twice
assert mock_acreate_file.call_count == 2
@pytest.mark.asyncio
async def test_router_async_get_healthy_deployments():
"""
Test that afile_content returns the correct file content
"""
router = litellm.Router(
model_list=[
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "gpt-3.5-turbo"},
},
],
)
result = await router.async_get_healthy_deployments(
model="gpt-3.5-turbo",
request_kwargs={},
messages=None,
input=None,
specific_deployment=False,
parent_otel_span=None,
)
assert len(result) == 1
assert result[0]["model_name"] == "gpt-3.5-turbo"
assert result[0]["litellm_params"]["model"] == "gpt-3.5-turbo"

6
tests/router_unit_tests/test_router_batch_utils.py
View File

@ -82,3 +82,9 @@ def test_router_metadata_variable_name():
assert (
_get_router_metadata_variable_name(function_name="batch") == "litellm_metadata"
)
assert (
_get_router_metadata_variable_name(function_name="acreate_file") == "litellm_metadata"
)
assert (
_get_router_metadata_variable_name(function_name="aget_file") == "litellm_metadata"
)

25
ui/litellm-dashboard/src/components/view_logs/index.tsx
View File

@ -769,6 +769,10 @@ export function RequestViewer({ row }: { row: Row<LogEntry> }) {
<span className="font-medium w-1/3">Model ID:</span>
<span>{row.original.model_id}</span>
</div>
<div className="flex">
<span className="font-medium w-1/3">Call Type:</span>
<span>{row.original.call_type}</span>
</div>
<div className="flex">
<span className="font-medium w-1/3">Provider:</span>
<span>{row.original.custom_llm_provider || "-"}</span>
@ -779,10 +783,12 @@ export function RequestViewer({ row }: { row: Row<LogEntry> }) {
<span className="max-w-[15ch] truncate block">{row.original.api_base || "-"}</span>
</Tooltip>
</div>
<div className="flex">
<span className="font-medium w-1/3">Start Time:</span>
<span>{row.original.startTime}</span>
</div>
{row?.original?.requester_ip_address && (
<div className="flex">
<span className="font-medium w-1/3">IP Address:</span>
<span>{row?.original?.requester_ip_address}</span>
</div>
)}
</div>
<div className="space-y-2">
@ -798,12 +804,7 @@ export function RequestViewer({ row }: { row: Row<LogEntry> }) {
<span className="font-medium w-1/3">Cache Hit:</span>
<span>{row.original.cache_hit}</span>
</div>
{row?.original?.requester_ip_address && (
<div className="flex">
<span className="font-medium w-1/3">IP Address:</span>
<span>{row?.original?.requester_ip_address}</span>
</div>
)}
<div className="flex">
<span className="font-medium w-1/3">Status:</span>
<span className={`px-2 py-1 rounded-md text-xs font-medium inline-block text-center w-16 ${
@ -815,6 +816,10 @@ export function RequestViewer({ row }: { row: Row<LogEntry> }) {
</span>
</div>
<div className="flex">
<span className="font-medium w-1/3">Start Time:</span>
<span>{row.original.startTime}</span>
</div>
<div className="flex">
<span className="font-medium w-1/3">End Time:</span>
<span>{row.original.endTime}</span>