scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 139

refactor(mistral): satisfy type-discipline and basedpyright gates for files/batches configs

Browse source
This commit is contained in:
mubashir1osmani 2026-09-09 18:37:01 -04:00
parent 2e5f5a95c8
commit c246f75e3e
10 changed files with 251 additions and 180 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
litellm/cost_calculator.py
View file

@ -2006,13 +2006,11 @@ def ocr_batch_cost(
``ocr_cost``.
"""
has_ocr_pricing: Final = model_info is not None and any(model_info.get(k) is not None for k in _OCR_PRICING_KEYS)
if has_ocr_pricing:
resolved_info: ModelInfo | None = model_info
else:
try:
resolved_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
except Exception:
resolved_info = None
resolved_info: Final = (
model_info
if has_ocr_pricing
else _lookup_model_info_or_none(model=model, custom_llm_provider=custom_llm_provider)
)
if resolved_info is None:
verbose_logger.warning(
"OCR batch cost: model=%s custom_llm_provider=%s has no pricing entry; returning 0.0 cost.",
@ -2022,9 +2020,7 @@ def ocr_batch_cost(
return 0.0, 0.0
page_rate: Final = _first_price(resolved_info, "ocr_cost_per_page_batches", "ocr_cost_per_page")
annotation_rate: Final = _first_price(
resolved_info, "annotation_cost_per_page_batches", "annotation_cost_per_page"
)
annotation_rate: Final = _first_price(resolved_info, "annotation_cost_per_page_batches", "annotation_cost_per_page")
pages_processed: Final = usage_info.pages_processed or 0
annotation_pages: Final = usage_info.pages_processed_annotation or 0
if page_rate is None and pages_processed > 0:
@ -2039,6 +2035,13 @@ def ocr_batch_cost(
return (page_rate or 0.0) * pages_processed + (effective_annotation_rate or 0.0) * annotation_pages, 0.0
def _lookup_model_info_or_none(model: str, custom_llm_provider: str | None) -> ModelInfo | None:
try:
return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
except Exception:
return None
def _first_price(model_info: ModelInfo, *keys: str) -> float | None:
return next((price for price in (model_info.get(k) for k in keys) if isinstance(price, (int, float))), None)

4
litellm/files/main.py
View file

@ -32,9 +32,7 @@ FileCreateProvider = Literal[
FileRetrieveProvider = Literal[
"openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic", "mistral"
]
FileDeleteProvider = Literal[
"openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic", "mistral"
]
FileDeleteProvider = Literal["openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic", "mistral"]
FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic", "mistral"]
import litellm
from litellm import get_secret_str

114
litellm/llms/mistral/batches/transformation.py
View file

@ -7,14 +7,16 @@ Output and error files are OpenAI-shaped JSONL (``{custom_id, response: {status_
so the shared batch cost accounting reads them without a provider branch.
"""
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final, Literal
from typing import Final, Literal, TypeAlias
import httpx
from openai.types.batch import BatchRequestCounts
from openai.types.batch import Errors as BatchErrors
from openai.types.batch_error import BatchError
from pydantic import BaseModel, ConfigDict
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
@ -24,13 +26,14 @@ from litellm.types.utils import LiteLLMBatch, LlmProviders
from ..common_utils import get_mistral_api_base, get_mistral_auth_headers, mistral_error
MistralBatchStatus = Literal[
MistralBatchStatus: TypeAlias = Literal[
"QUEUED", "RUNNING", "SUCCESS", "FAILED", "TIMEOUT_EXCEEDED", "CANCELLATION_REQUESTED", "CANCELLED"
]
OpenAIBatchStatus = Literal[
OpenAIBatchStatus: TypeAlias = Literal[
"validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled"
]
_NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) # mutable-ok: frozen at module scope
_STATUS_MAP: Final[MappingProxyType[MistralBatchStatus, OpenAIBatchStatus]] = MappingProxyType(
{
"QUEUED": "validating",
@ -44,6 +47,23 @@ _STATUS_MAP: Final[MappingProxyType[MistralBatchStatus, OpenAIBatchStatus]] = Ma
)
class MistralCreateBatchJobRequest(TypedDict):
"""Body of ``POST /v1/batch/jobs``."""
input_files: ReadOnly[tuple[str, ...]]
endpoint: ReadOnly[str]
model: ReadOnly[str]
metadata: NotRequired[ReadOnly[Mapping[str, str]]]
class MistralPresignedRequest(TypedDict):
"""A fully-formed request the shared HTTP handler sends as-is (its ``method`` branch)."""
method: ReadOnly[Literal["GET"]]
url: ReadOnly[str]
headers: ReadOnly[Mapping[str, str]]
class MistralBatchError(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
@ -69,7 +89,18 @@ class MistralBatchJob(BaseModel):
output_file: str | None = None
error_file: str | None = None
errors: tuple[MistralBatchError, ...] = ()
metadata: dict[str, str] | None = None
metadata: dict[str, str] | None = None # mutable-ok: LiteLLMBatch.metadata is typed as dict
def _to_batch_errors(errors: Sequence[MistralBatchError]) -> BatchErrors | None:
if not errors:
return None
return BatchErrors(
object="list",
data=[ # mutable-ok: openai Batch.Errors.data is typed as list
BatchError(message=f"{e.message} (x{e.count})" if e.count > 1 else e.message) for e in errors
],
)
def _to_litellm_batch(job: MistralBatchJob) -> LiteLLMBatch:
@ -90,14 +121,7 @@ def _to_litellm_batch(job: MistralBatchJob) -> LiteLLMBatch:
cancelled_at=terminal_at if status == "cancelled" else None,
output_file_id=job.output_file,
error_file_id=job.error_file,
errors=(
BatchErrors(
object="list",
data=[BatchError(message=f"{e.message} (x{e.count})" if e.count > 1 else e.message) for e in job.errors],
)
if job.errors
else None
),
errors=_to_batch_errors(job.errors),
request_counts=BatchRequestCounts(
total=job.total_requests,
completed=job.succeeded_requests,
@ -114,14 +138,14 @@ class MistralBatchesConfig(BaseBatchesConfig):
def validate_environment(
self,
headers: dict,
headers: Mapping[str, str],
model: str,
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
messages: Sequence[AllMessageValues],
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
) -> dict:
) -> dict[str, str]: # mutable-ok: BaseBatchesConfig signature
return get_mistral_auth_headers(headers, api_key)
def get_complete_batch_url(
@ -129,8 +153,8 @@ class MistralBatchesConfig(BaseBatchesConfig):
api_base: str | None,
api_key: str | None,
model: str,
optional_params: dict,
litellm_params: dict,
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
data: CreateBatchRequest,
) -> str:
return f"{get_mistral_api_base(api_base)}/v1/batch/jobs"
@ -139,48 +163,58 @@ class MistralBatchesConfig(BaseBatchesConfig):
self,
model: str,
create_batch_data: CreateBatchRequest,
optional_params: dict,
litellm_params: dict,
) -> dict[str, object]:
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
) -> dict[str, object]: # mutable-ok: BaseBatchesConfig signature
input_file_id: Final = create_batch_data.get("input_file_id")
endpoint: Final = create_batch_data.get("endpoint")
if input_file_id is None or endpoint is None:
raise ValueError("input_file_id and endpoint are required to create a Mistral batch job")
metadata: Final = create_batch_data.get("metadata")
return {
"input_files": [create_batch_data["input_file_id"]],
"endpoint": create_batch_data["endpoint"],
"model": model,
**({"metadata": metadata} if metadata else {}),
**(create_batch_data.get("extra_body") or {}),
}
body: Final = (
MistralCreateBatchJobRequest(
input_files=(input_file_id,), endpoint=endpoint, model=model, metadata=metadata
)
if metadata
else MistralCreateBatchJobRequest(input_files=(input_file_id,), endpoint=endpoint, model=model)
)
return dict(body) # mutable-ok: BaseBatchesConfig signature
def transform_create_batch_response(
self,
model: str | None,
raw_response: httpx.Response,
logging_obj: object,
litellm_params: dict,
litellm_params: Mapping[str, object],
) -> LiteLLMBatch:
return _to_litellm_batch(MistralBatchJob.model_validate(raw_response.json()))
def transform_retrieve_batch_request(
self,
batch_id: str,
optional_params: dict,
litellm_params: dict,
) -> dict[str, object]:
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
) -> dict[str, object]: # mutable-ok: BaseBatchesConfig signature
encoded_batch_id: Final = encode_url_path_segment(batch_id, field_name="batch_id")
return {
"method": "GET",
"url": f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/batch/jobs/{encoded_batch_id}",
"headers": get_mistral_auth_headers({}, litellm_params.get("api_key")),
}
api_base: Final = litellm_params.get("api_base")
api_key: Final = litellm_params.get("api_key")
request: Final = MistralPresignedRequest(
method="GET",
url=f"{get_mistral_api_base(api_base if isinstance(api_base, str) else None)}/v1/batch/jobs/{encoded_batch_id}",
headers=get_mistral_auth_headers(_NO_HEADERS, api_key if isinstance(api_key, str) else None),
)
return dict(request) # mutable-ok: BaseBatchesConfig signature
def transform_retrieve_batch_response(
self,
model: str | None,
raw_response: httpx.Response,
logging_obj: object,
litellm_params: dict,
litellm_params: Mapping[str, object],
) -> LiteLLMBatch:
return _to_litellm_batch(MistralBatchJob.model_validate(raw_response.json()))
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
def get_error_class(
self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers
) -> BaseLLMException:
return mistral_error(error_message, status_code, headers)

13
litellm/llms/mistral/common_utils.py
View file

@ -1,3 +1,4 @@
from collections.abc import Mapping
from typing import Final
import httpx
@ -19,18 +20,22 @@ def get_mistral_api_base(api_base: str | None) -> str:
return resolved.removesuffix("/v1")
def get_mistral_auth_headers(headers: dict, api_key: str | None) -> dict:
def get_mistral_auth_headers(
headers: Mapping[str, str], api_key: str | None
) -> dict[str, str]: # mutable-ok: BaseConfig.validate_environment contract returns dict
resolved_key: Final = api_key or get_secret_str(MISTRAL_API_KEY_ENV_VAR)
if resolved_key is None:
raise ValueError(
"Missing Mistral API Key - A call is being made to Mistral but no key is set either in the environment variables or via params"
)
return {**headers, "Authorization": f"Bearer {resolved_key}"}
return dict(headers, Authorization=f"Bearer {resolved_key}") # mutable-ok: BaseConfig contract returns dict
def mistral_error(error_message: str, status_code: int, headers: dict | httpx.Headers) -> MistralError:
def mistral_error(error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers) -> MistralError:
return MistralError(
status_code=status_code,
message=error_message,
headers=headers if isinstance(headers, httpx.Headers) else httpx.Headers(headers),
headers=headers
if isinstance(headers, httpx.Headers)
else httpx.Headers(dict(headers)), # mutable-ok: httpx.Headers takes a dict
)

129
litellm/llms/mistral/files/transformation.py
View file

@ -7,11 +7,13 @@ Mistral only accepts ``fine-tune``, ``batch`` and ``ocr`` as upload purposes.
"""
import time
from typing import Final, Literal
from collections.abc import Mapping, Sequence
from typing import Final, Literal, TypeAlias
import httpx
from openai.types.file_deleted import FileDeleted
from pydantic import BaseModel, ConfigDict
from typing_extensions import ReadOnly, TypedDict
from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
@ -29,7 +31,16 @@ from litellm.types.utils import LlmProviders
from ..common_utils import get_mistral_api_base, get_mistral_auth_headers, mistral_error
MistralFilePurpose = Literal["fine-tune", "batch", "ocr"]
MistralFilePurpose: TypeAlias = Literal["fine-tune", "batch", "ocr"]
_NO_QUERY_PARAMS: Final[dict[str, str]] = {} # mutable-ok: BaseFilesConfig request transforms return tuple[str, dict]
class MistralMultipartUpload(TypedDict):
"""``files=`` payload for ``POST /v1/files``: each value is an httpx multipart tuple."""
file: ReadOnly[tuple[str, object, str]]
purpose: ReadOnly[tuple[None, MistralFilePurpose]]
class MistralFile(BaseModel):
@ -85,6 +96,11 @@ def _to_mistral_purpose(purpose: str) -> MistralFilePurpose:
return "batch"
def _api_base_from(litellm_params: Mapping[str, object]) -> str:
api_base: Final = litellm_params.get("api_base")
return get_mistral_api_base(api_base if isinstance(api_base, str) else None)
class MistralFilesConfig(BaseFilesConfig):
@property
def custom_llm_provider(self) -> LlmProviders:
@ -95,99 +111,103 @@ class MistralFilesConfig(BaseFilesConfig):
api_base: str | None,
api_key: str | None,
model: str,
optional_params: dict,
litellm_params: dict,
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
stream: bool | None = None,
) -> str:
return f"{get_mistral_api_base(api_base)}/v1/files"
def _file_url(self, file_id: str, litellm_params: dict, suffix: str = "") -> str:
def _file_url(self, file_id: str, litellm_params: Mapping[str, object], suffix: str = "") -> str:
encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id")
return f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/files/{encoded_file_id}{suffix}"
return f"{_api_base_from(litellm_params)}/v1/files/{encoded_file_id}{suffix}"
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
def get_error_class(
self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers
) -> BaseLLMException:
return mistral_error(error_message, status_code, headers)
def validate_environment(
self,
headers: dict,
headers: Mapping[str, str],
model: str,
messages: list,
optional_params: dict,
litellm_params: dict,
messages: Sequence[object],
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
) -> dict:
) -> dict[str, str]: # mutable-ok: BaseFilesConfig signature
return get_mistral_auth_headers(headers, api_key)
def get_supported_openai_params(self, model: str) -> list[OpenAICreateFileRequestOptionalParams]:
return ["purpose"]
def get_supported_openai_params(
self, model: str
) -> list[OpenAICreateFileRequestOptionalParams]: # mutable-ok: BaseFilesConfig signature
return ["purpose"] # mutable-ok: BaseFilesConfig signature
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
non_default_params: Mapping[str, object],
optional_params: dict[str, object], # mutable-ok: BaseConfig signature, returned as-is
model: str,
drop_params: bool,
) -> dict:
) -> dict[str, object]: # mutable-ok: BaseConfig signature
return optional_params
def transform_create_file_request(
self,
model: str,
create_file_data: CreateFileRequest,
optional_params: dict,
litellm_params: dict,
) -> dict:
file_data: Final = create_file_data.get("file")
if file_data is None:
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
) -> dict[str, object]: # mutable-ok: BaseFilesConfig signature
if "file" not in create_file_data:
raise ValueError("File data is required")
extracted: Final = extract_file_data(file_data)
extracted: Final = extract_file_data(create_file_data["file"])
filename: Final = extracted["filename"] or f"file_{int(time.time())}.jsonl"
content_type: Final = extracted.get("content_type") or "application/octet-stream"
return {
"file": (filename, extracted["content"], content_type),
"purpose": (None, _to_mistral_purpose(create_file_data.get("purpose", "batch"))),
}
upload: Final = MistralMultipartUpload(
file=(filename, extracted["content"], content_type),
purpose=(None, _to_mistral_purpose(create_file_data.get("purpose", "batch"))),
)
return dict(upload) # mutable-ok: BaseFilesConfig signature
def transform_create_file_response(
self,
model: str | None,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
litellm_params: Mapping[str, object],
) -> OpenAIFileObject:
return _to_openai_file_object(MistralFile.model_validate(raw_response.json()))
def transform_retrieve_file_request(
self,
file_id: str,
optional_params: dict,
litellm_params: dict,
) -> tuple[str, dict]:
return self._file_url(file_id, litellm_params), {}
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature
return self._file_url(file_id, litellm_params), _NO_QUERY_PARAMS
def transform_retrieve_file_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
litellm_params: Mapping[str, object],
) -> OpenAIFileObject:
return _to_openai_file_object(MistralFile.model_validate(raw_response.json()))
def transform_delete_file_request(
self,
file_id: str,
optional_params: dict,
litellm_params: dict,
) -> tuple[str, dict]:
return self._file_url(file_id, litellm_params), {}
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature
return self._file_url(file_id, litellm_params), _NO_QUERY_PARAMS
def transform_delete_file_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
litellm_params: Mapping[str, object],
) -> FileDeleted:
deleted: Final = MistralFileDeleted.model_validate(raw_response.json())
return FileDeleted(id=deleted.id, deleted=deleted.deleted, object="file")
@ -195,32 +215,39 @@ class MistralFilesConfig(BaseFilesConfig):
def transform_list_files_request(
self,
purpose: str | None,
optional_params: dict,
litellm_params: dict,
) -> tuple[str, dict]:
params: Final = {"purpose": _to_mistral_purpose(purpose)} if purpose else {}
return f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/files", params
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature
url: Final = f"{_api_base_from(litellm_params)}/v1/files"
if not purpose:
return url, _NO_QUERY_PARAMS
return url, {"purpose": _to_mistral_purpose(purpose)} # mutable-ok: BaseFilesConfig signature returns dict
def transform_list_files_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
) -> list[OpenAIFileObject]:
return [_to_openai_file_object(f) for f in MistralFileList.model_validate(raw_response.json()).data]
litellm_params: Mapping[str, object],
) -> list[OpenAIFileObject]: # mutable-ok: BaseFilesConfig signature
return [ # mutable-ok: BaseFilesConfig signature
_to_openai_file_object(f) for f in MistralFileList.model_validate(raw_response.json()).data
]
def transform_file_content_request(
self,
file_content_request: FileContentRequest,
optional_params: dict,
litellm_params: dict,
) -> tuple[str, dict]:
return self._file_url(file_content_request["file_id"], litellm_params, suffix="/content"), {}
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature
file_id: Final = file_content_request.get("file_id")
if file_id is None:
raise ValueError("file_id is required to download file content")
return self._file_url(file_id, litellm_params, suffix="/content"), _NO_QUERY_PARAMS
def transform_file_content_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
litellm_params: Mapping[str, object],
) -> HttpxBinaryResponseContent:
return HttpxBinaryResponseContent(response=raw_response)

2
litellm/proxy/openai_files_endpoints/files_endpoints.py
View file

@ -1130,7 +1130,7 @@ async def get_file(
check_file_id_encoding=True,
)
if should_route:
if should_route and credentials is not None:
# Use model-based routing with credentials from config
prepare_data_with_credentials(
data=data,

46
tests/test_litellm/batches/test_batch_utils.py
View file

@ -645,9 +645,7 @@ async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch):
lambda content, model: pytest.fail("raw vertex path should not run"),
)
result = await bu.calculate_batch_cost_and_usage(
file_content_dictionary=[], custom_llm_provider="vertex_ai"
)
result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=[], custom_llm_provider="vertex_ai")
assert result.cost == 0.0
assert result.usage.total_tokens == 0
assert result.models == []
@ -1250,6 +1248,7 @@ async def test_handle_completed_batch_no_output_file_is_zero(monkeypatch):
result set - zero cost, zero usage, no models - instead of letting the file
fetch raise "Output file id is None" on every aretrieve_batch logging poll.
"""
# The output-file fetch must not even be attempted when there is no output file.
async def _must_not_fetch(*args, **kwargs):
pytest.fail("_fetch_batch_output_file_content should not be called")
@ -1376,7 +1375,10 @@ def test_anthropic_response_body_is_result_message():
def test_anthropic_usage_conversion_includes_cache_tokens():
body = {"model": "claude-sonnet-4-5-20250929", "usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000)}
body = {
"model": "claude-sonnet-4-5-20250929",
"usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000),
}
usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="anthropic")
assert usage.prompt_tokens == 11000
assert usage.completion_tokens == 200
@ -1391,7 +1393,9 @@ def test_bedrock_model_output_line_success_check():
"modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}},
}
assert bu._batch_response_was_successful(row, custom_llm_provider="bedrock") is True
assert bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6"
assert (
bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6"
)
def test_bedrock_cost_uses_deployment_model_name():
@ -1445,7 +1449,13 @@ def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch):
rows = [
{
"custom_id": "req-1",
"response": {"status_code": 200, "body": {"model": "gpt-5.2", "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}},
"response": {
"status_code": 200,
"body": {
"model": "gpt-5.2",
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
},
},
}
]
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
@ -1487,7 +1497,9 @@ def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatc
lambda **kw: pytest.fail("anthropic rows must not go through completion_cost"),
)
result = bu._aggregate_batch_cost_usage_models(entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic")
result = bu._aggregate_batch_cost_usage_models(
entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic"
)
assert result.cost == pytest.approx(0.3)
assert seen[0]["model"] == "claude-sonnet-4-5-20250929"
@ -1522,7 +1534,11 @@ async def test_calculate_batch_cost_and_usage_anthropic_end_to_end():
)
assert result.cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2)
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (11000, 200, 11200)
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (
11000,
200,
11200,
)
assert result.models == ["claude-sonnet-4-5"]
@ -1689,7 +1705,10 @@ async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) ->
def test_bedrock_converse_shaped_batch_usage_is_parsed():
body = {"model": "us.amazon.nova-lite-v1:0", "usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742}}
body = {
"model": "us.amazon.nova-lite-v1:0",
"usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742},
}
usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (2202, 540, 2742)
@ -1739,6 +1758,7 @@ def test_unparsable_bedrock_batch_usage_warns(caplog):
# batch_cost_is_final
# --------------------------------------------------------------------------- #
def _retrieved_batch(
status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None
) -> LiteLLMBatch:
@ -1798,7 +1818,9 @@ def _ocr_row(pages_processed, annotation_pages=None, model="mistral-ocr-latest")
usage_info = {"pages_processed": pages_processed, "doc_size_bytes": 4096}
if annotation_pages is not None:
usage_info["pages_processed_annotation"] = annotation_pages
return _success_row(model=model, pages=[{"index": i, "markdown": "x"} for i in range(pages_processed)], usage_info=usage_info)
return _success_row(
model=model, pages=[{"index": i, "markdown": "x"} for i in range(pages_processed)], usage_info=usage_info
)
def test_ocr_rows_are_priced_per_page_at_batch_rate(monkeypatch):
@ -1835,7 +1857,9 @@ def test_ocr_rows_bill_annotation_pages_separately(monkeypatch):
"annotation_cost_per_page_batches": 0.0025,
},
)
result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(4, annotation_pages=4)], custom_llm_provider="mistral")
result = bu._aggregate_batch_cost_usage_models(
entries=[_ocr_row(4, annotation_pages=4)], custom_llm_provider="mistral"
)
assert result.cost == pytest.approx(4 * 0.002 + 4 * 0.0025)

76
tests/test_litellm/batches/test_main.py
View file

@ -66,9 +66,7 @@ def seams():
stack.enter_context(patch.object(bm, "openai_batches_instance", openai_i))
stack.enter_context(patch.object(bm, "azure_batches_instance", azure_i))
stack.enter_context(patch.object(bm, "vertex_ai_batches_instance", vertex_i))
stack.enter_context(
patch.object(bm, "anthropic_batches_instance", anthropic_i)
)
stack.enter_context(patch.object(bm, "anthropic_batches_instance", anthropic_i))
stack.enter_context(patch.object(bm, "base_llm_http_handler", base_http))
stack.enter_context(patch.object(bm, "BedrockBatchesHandler", bedrock_arn))
yield Seams(
@ -174,9 +172,7 @@ def test_create__provider_config_routes_to_base_http_handler(seams):
"get_provider_batches_config",
return_value=MagicMock(name="provider_config"),
):
result = bm.create_batch(
**CREATE_KW, custom_llm_provider="bedrock", model="bedrock/my-batch-model"
)
result = bm.create_batch(**CREATE_KW, custom_llm_provider="bedrock", model="bedrock/my-batch-model")
assert result is seams.base_http.create_batch.return_value
_assert_only(seams.base_http.create_batch, seams, "create_batch")
@ -281,9 +277,7 @@ def test_retrieve__bedrock_model_invocation_job_arn(seams):
result = bm.retrieve_batch(batch_id=arn, custom_llm_provider="bedrock")
seams.bedrock_arn._handle_model_invocation_job_status.assert_called_once()
assert (
result is seams.bedrock_arn._handle_model_invocation_job_status.return_value
)
assert result is seams.bedrock_arn._handle_model_invocation_job_status.return_value
seams.bedrock_arn._handle_async_invoke_status.assert_not_called()
@ -385,9 +379,7 @@ def test_cancel__unsupported_provider_raises_badrequest(seams):
def test_cancel__async_flag_propagates_is_async(seams):
bm.cancel_batch(
batch_id="batch-1", custom_llm_provider="openai", acancel_batch=True
)
bm.cancel_batch(batch_id="batch-1", custom_llm_provider="openai", acancel_batch=True)
assert seams.openai.cancel_batch.call_args.kwargs["_is_async"] is True
@ -415,9 +407,7 @@ async def test_acreate_batch_delegates_to_create_batch():
@pytest.mark.asyncio
async def test_aretrieve_batch_delegates_to_retrieve_batch():
with patch.object(bm, "retrieve_batch", MagicMock(return_value="SENTINEL")) as m:
result = await bm.aretrieve_batch(
batch_id="batch-1", custom_llm_provider="azure"
)
result = await bm.aretrieve_batch(batch_id="batch-1", custom_llm_provider="azure")
assert result == "SENTINEL"
assert m.call_count == 1
@ -429,9 +419,7 @@ async def test_aretrieve_batch_delegates_to_retrieve_batch():
@pytest.mark.asyncio
async def test_alist_batches_delegates_to_list_batches():
with patch.object(bm, "list_batches", MagicMock(return_value="SENTINEL")) as m:
result = await bm.alist_batches(
after="cur", limit=3, custom_llm_provider="vertex_ai"
)
result = await bm.alist_batches(after="cur", limit=3, custom_llm_provider="vertex_ai")
assert result == "SENTINEL"
assert m.call_count == 1
@ -444,9 +432,7 @@ async def test_alist_batches_delegates_to_list_batches():
@pytest.mark.asyncio
async def test_acancel_batch_delegates_to_cancel_batch():
with patch.object(bm, "cancel_batch", MagicMock(return_value="SENTINEL")) as m:
result = await bm.acancel_batch(
batch_id="batch-1", custom_llm_provider="openai"
)
result = await bm.acancel_batch(batch_id="batch-1", custom_llm_provider="openai")
assert result == "SENTINEL"
assert m.call_count == 1
@ -499,9 +485,7 @@ def _sent(mock_method, *keys):
def test_create__openai_credentials_passthrough(seams):
bm.create_batch(**CREATE_KW, custom_llm_provider="openai", **OPENAI_CREDS)
assert _sent(
seams.openai.create_batch, "api_key", "api_base", "organization", "max_retries"
) == {
assert _sent(seams.openai.create_batch, "api_key", "api_base", "organization", "max_retries") == {
"api_key": "sk-user-openai",
"api_base": "https://openai.user.test",
"organization": "org-user-123",
@ -512,9 +496,7 @@ def test_create__openai_credentials_passthrough(seams):
def test_create__azure_credentials_passthrough(seams):
bm.create_batch(**CREATE_KW, custom_llm_provider="azure", **AZURE_CREDS)
assert _sent(
seams.azure.create_batch, "api_key", "api_base", "api_version"
) == {
assert _sent(seams.azure.create_batch, "api_key", "api_base", "api_version") == {
"api_key": "sk-user-azure",
"api_base": "https://azure.user.test",
"api_version": "2024-12-99",
@ -564,9 +546,7 @@ def test_create__provider_config_credentials_passthrough(seams):
def test_retrieve__openai_credentials_passthrough(seams):
bm.retrieve_batch(batch_id="b1", custom_llm_provider="openai", **OPENAI_CREDS)
assert _sent(
seams.openai.retrieve_batch, "api_key", "api_base", "organization"
) == {
assert _sent(seams.openai.retrieve_batch, "api_key", "api_base", "organization") == {
"api_key": "sk-user-openai",
"api_base": "https://openai.user.test",
"organization": "org-user-123",
@ -576,9 +556,7 @@ def test_retrieve__openai_credentials_passthrough(seams):
def test_retrieve__azure_credentials_passthrough(seams):
bm.retrieve_batch(batch_id="b1", custom_llm_provider="azure", **AZURE_CREDS)
assert _sent(
seams.azure.retrieve_batch, "api_key", "api_base", "api_version"
) == {
assert _sent(seams.azure.retrieve_batch, "api_key", "api_base", "api_version") == {
"api_key": "sk-user-azure",
"api_base": "https://azure.user.test",
"api_version": "2024-12-99",
@ -640,9 +618,7 @@ def test_retrieve__provider_config_credentials_passthrough(seams):
def test_list__openai_credentials_passthrough(seams):
bm.list_batches(custom_llm_provider="openai", **OPENAI_CREDS)
assert _sent(
seams.openai.list_batches, "api_key", "api_base", "organization"
) == {
assert _sent(seams.openai.list_batches, "api_key", "api_base", "organization") == {
"api_key": "sk-user-openai",
"api_base": "https://openai.user.test",
"organization": "org-user-123",
@ -652,9 +628,7 @@ def test_list__openai_credentials_passthrough(seams):
def test_list__azure_credentials_passthrough(seams):
bm.list_batches(custom_llm_provider="azure", **AZURE_CREDS)
assert _sent(
seams.azure.list_batches, "api_key", "api_base", "api_version"
) == {
assert _sent(seams.azure.list_batches, "api_key", "api_base", "api_version") == {
"api_key": "sk-user-azure",
"api_base": "https://azure.user.test",
"api_version": "2024-12-99",
@ -682,9 +656,7 @@ def test_list__vertex_credentials_passthrough(seams):
def test_cancel__openai_credentials_passthrough(seams):
bm.cancel_batch(batch_id="b1", custom_llm_provider="openai", **OPENAI_CREDS)
assert _sent(
seams.openai.cancel_batch, "api_key", "api_base", "organization"
) == {
assert _sent(seams.openai.cancel_batch, "api_key", "api_base", "organization") == {
"api_key": "sk-user-openai",
"api_base": "https://openai.user.test",
"organization": "org-user-123",
@ -694,9 +666,7 @@ def test_cancel__openai_credentials_passthrough(seams):
def test_cancel__azure_credentials_passthrough(seams):
bm.cancel_batch(batch_id="b1", custom_llm_provider="azure", **AZURE_CREDS)
assert _sent(
seams.azure.cancel_batch, "api_key", "api_base", "api_version"
) == {
assert _sent(seams.azure.cancel_batch, "api_key", "api_base", "api_version") == {
"api_key": "sk-user-azure",
"api_base": "https://azure.user.test",
"api_version": "2024-12-99",
@ -786,18 +756,16 @@ def test_retrieve__omits_trusted_model_credentials_when_not_supplied(seams):
def test_create__mistral_ocr_routes_to_base_http_handler_with_mistral_config(seams):
with patch.object(bm.ProviderConfigManager, "get_provider_batches_config", wraps=bm.ProviderConfigManager.get_provider_batches_config) as get_cfg:
result = bm.create_batch(
completion_window="24h",
endpoint="/v1/ocr",
input_file_id="file-abc",
custom_llm_provider="mistral",
model="mistral/mistral-ocr-latest",
)
result = bm.create_batch(
completion_window="24h",
endpoint="/v1/ocr",
input_file_id="file-abc",
custom_llm_provider="mistral",
model="mistral/mistral-ocr-latest",
)
assert result is seams.base_http.create_batch.return_value
_assert_only(seams.base_http.create_batch, seams, "create_batch")
get_cfg.assert_called_once()
forwarded = seams.base_http.create_batch.call_args.kwargs
assert type(forwarded["provider_config"]).__name__ == "MistralBatchesConfig"
assert forwarded["model"] == "mistral-ocr-latest"

16
tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py
View file

@ -92,26 +92,34 @@ def test_create_request_maps_openai_fields_onto_mistral_job(config):
model="mistral-ocr-latest", create_batch_data=data, optional_params={}, litellm_params={}
)
assert body == {
"input_files": ["file-123"],
"input_files": ("file-123",),
"endpoint": "/v1/ocr",
"model": "mistral-ocr-latest",
"metadata": {"team": "docs"},
}
def test_create_request_omits_empty_metadata_and_forwards_extra_body(config):
def test_create_request_omits_empty_metadata(config):
data = CreateBatchRequest(
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id="file-123",
metadata=None,
extra_body={"timeout_hours": 48},
)
body = config.transform_create_batch_request(
model="mistral-small-latest", create_batch_data=data, optional_params={}, litellm_params={}
)
assert "metadata" not in body
assert body["timeout_hours"] == 48
def test_create_request_requires_input_file_and_endpoint(config):
with pytest.raises(ValueError, match="input_file_id and endpoint are required"):
config.transform_create_batch_request(
model="m",
create_batch_data=CreateBatchRequest(completion_window="24h"),
optional_params={},
litellm_params={},
)
@pytest.mark.parametrize(

8
tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py
View file

@ -79,7 +79,9 @@ def test_validate_environment_uses_bearer_auth(config, api_key):
def test_upload_request_is_multipart_with_batch_purpose(config):
body = config.transform_create_file_request(
model="",
create_file_data=CreateFileRequest(file=("in.jsonl", b'{"custom_id":"0"}\n', "application/jsonl"), purpose="batch"),
create_file_data=CreateFileRequest(
file=("in.jsonl", b'{"custom_id":"0"}\n', "application/jsonl"), purpose="batch"
),
optional_params={},
litellm_params={},
)
@ -181,7 +183,9 @@ def test_list_request_filters_by_mapped_purpose(config):
def test_list_response(config):
out = config.transform_list_files_response(
raw_response=_response({"data": [_file(), _file(id="second", filename="b.jsonl")], "object": "list", "total": 2}),
raw_response=_response(
{"data": [_file(), _file(id="second", filename="b.jsonl")], "object": "list", "total": 2}
),
logging_obj=None,
litellm_params={},
)