mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
Merge e6bc4e47c7 into 0c98afa780
This commit is contained in:
commit
8b4af15ef5
34 changed files with 1856 additions and 213 deletions
|
|
@ -9,6 +9,7 @@ import litellm
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo
|
||||
from litellm.types.llms.openai import Batch
|
||||
from litellm.types.utils import ModelInfo, Usage
|
||||
from litellm.utils import token_counter
|
||||
|
|
@ -50,7 +51,7 @@ def batch_cost_is_final(batch: Batch) -> bool:
|
|||
|
||||
async def calculate_batch_cost_and_usage(
|
||||
file_content_dictionary: list[dict],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"],
|
||||
model_name: str | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
) -> BatchCostUsageResult:
|
||||
|
|
@ -80,7 +81,7 @@ async def calculate_batch_cost_and_usage(
|
|||
|
||||
async def _handle_completed_batch(
|
||||
batch: Batch,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"],
|
||||
model_name: str | None = None,
|
||||
litellm_params: dict | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
|
|
@ -166,7 +167,7 @@ class _BatchOutputLineStats:
|
|||
|
||||
def _classify_output_line_stats(
|
||||
entries: Iterable[dict],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
|
||||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> Iterator[_BatchOutputLineStats | _LineOutcome]:
|
||||
|
|
@ -185,7 +186,7 @@ def _classify_output_line_stats(
|
|||
|
||||
def _safe_output_line_stats(
|
||||
entry: Mapping[str, object],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
|
||||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> _BatchOutputLineStats | None:
|
||||
|
|
@ -207,7 +208,7 @@ def _safe_output_line_stats(
|
|||
|
||||
def _compute_output_line_stats(
|
||||
entry: Mapping[str, object],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
|
||||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> _BatchOutputLineStats:
|
||||
|
|
@ -218,6 +219,7 @@ def _compute_output_line_stats(
|
|||
response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None
|
||||
completion_details: Final = usage.completion_tokens_details
|
||||
line_prompt_cost, line_completion_cost = _output_line_cost(
|
||||
response_body=response_body,
|
||||
usage=usage,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_name=model_name,
|
||||
|
|
@ -237,19 +239,36 @@ def _compute_output_line_stats(
|
|||
)
|
||||
|
||||
|
||||
def _ocr_usage_info_from_response_body(response_body: Mapping[str, object]) -> OCRUsageInfo | None:
|
||||
"""OCR results report ``usage_info`` (pages) instead of ``usage`` (tokens); None for non-OCR lines."""
|
||||
raw_usage_info: Final = response_body.get("usage_info")
|
||||
if not isinstance(raw_usage_info, Mapping):
|
||||
return None
|
||||
return OCRUsageInfo.model_validate(raw_usage_info)
|
||||
|
||||
|
||||
def _output_line_cost(
|
||||
response_body: Mapping[str, object],
|
||||
usage: Usage,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
|
||||
model_name: str | None,
|
||||
response_model: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> tuple[float, float]:
|
||||
"""(prompt_cost, completion_cost) for one output line, priced at batch rates."""
|
||||
from litellm.cost_calculator import batch_cost_calculator
|
||||
from litellm.cost_calculator import batch_cost_calculator, ocr_batch_cost
|
||||
|
||||
cost_model: Final = (
|
||||
model_name if custom_llm_provider == "bedrock" and model_name else response_model or model_name or ""
|
||||
)
|
||||
ocr_usage: Final = _ocr_usage_info_from_response_body(response_body)
|
||||
if ocr_usage is not None:
|
||||
return ocr_batch_cost(
|
||||
model=cost_model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
usage_info=ocr_usage,
|
||||
model_info=model_info,
|
||||
)
|
||||
return batch_cost_calculator(
|
||||
usage=usage,
|
||||
model=cost_model,
|
||||
|
|
@ -260,7 +279,7 @@ def _output_line_cost(
|
|||
|
||||
def _aggregate_batch_cost_usage_models(
|
||||
entries: Iterable[dict],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
|
||||
model_name: str | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
) -> BatchCostUsageResult:
|
||||
|
|
@ -427,7 +446,7 @@ def _provider_output_file_id(output_file_id: str) -> str:
|
|||
|
||||
async def _fetch_batch_managed_file_content(
|
||||
file_id: str,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai",
|
||||
litellm_params: dict | None = None,
|
||||
) -> bytes:
|
||||
"""
|
||||
|
|
@ -457,7 +476,7 @@ async def _fetch_batch_managed_file_content(
|
|||
|
||||
async def _fetch_batch_output_file_content(
|
||||
batch: Batch,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai",
|
||||
litellm_params: dict | None = None,
|
||||
) -> bytes:
|
||||
"""
|
||||
|
|
@ -479,7 +498,7 @@ async def _fetch_batch_output_file_content(
|
|||
|
||||
async def count_error_file_failed_requests(
|
||||
batch: Batch,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"],
|
||||
litellm_params: dict | None,
|
||||
) -> int:
|
||||
"""Count failed requests reported only in the batch's separate error file.
|
||||
|
|
|
|||
|
|
@ -105,9 +105,11 @@ def _resolve_timeout(
|
|||
@client
|
||||
async def acreate_batch(
|
||||
completion_window: Literal["24h"],
|
||||
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"],
|
||||
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"],
|
||||
input_file_id: str,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy"] = "openai",
|
||||
custom_llm_provider: Literal[
|
||||
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "mistral"
|
||||
] = "openai",
|
||||
metadata: dict[str, str] | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
extra_body: dict[str, str] | None = None,
|
||||
|
|
@ -155,9 +157,11 @@ async def acreate_batch(
|
|||
@client
|
||||
def create_batch(
|
||||
completion_window: Literal["24h"],
|
||||
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"],
|
||||
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"],
|
||||
input_file_id: str,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy"] = "openai",
|
||||
custom_llm_provider: Literal[
|
||||
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "mistral"
|
||||
] = "openai",
|
||||
metadata: dict[str, str] | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
extra_body: dict[str, str] | None = None,
|
||||
|
|
@ -341,7 +345,7 @@ def create_batch(
|
|||
async def aretrieve_batch(
|
||||
batch_id: str,
|
||||
custom_llm_provider: Literal[
|
||||
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic"
|
||||
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral"
|
||||
] = "openai",
|
||||
metadata: dict[str, str] | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
|
|
@ -389,7 +393,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
|
|||
_retrieve_batch_request: RetrieveBatchRequest,
|
||||
_is_async: bool,
|
||||
custom_llm_provider: Literal[
|
||||
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic"
|
||||
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral"
|
||||
] = "openai",
|
||||
logging_obj: LiteLLMLoggingObj | None = None,
|
||||
):
|
||||
|
|
@ -497,7 +501,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
|
|||
message=(
|
||||
f"LiteLLM doesn't support custom_llm_provider={custom_llm_provider} for 'retrieve_batch' without a `model` kwarg. "
|
||||
"Supported via this path: 'openai', 'azure', 'vertex_ai', 'anthropic'. "
|
||||
"'bedrock' is supported but requires `model` to be passed so the provider config can be loaded."
|
||||
"'bedrock' and 'mistral' are supported but require `model` to be passed so the provider config can be loaded."
|
||||
),
|
||||
model="n/a",
|
||||
llm_provider=custom_llm_provider,
|
||||
|
|
@ -514,7 +518,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
|
|||
def retrieve_batch(
|
||||
batch_id: str,
|
||||
custom_llm_provider: Literal[
|
||||
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic"
|
||||
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral"
|
||||
] = "openai",
|
||||
metadata: dict[str, str] | None = None,
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
|
|
|
|||
|
|
@ -139,6 +139,7 @@ if TYPE_CHECKING:
|
|||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
Logging as LitellmLoggingObject,
|
||||
)
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo
|
||||
else:
|
||||
LitellmLoggingObject = Any
|
||||
|
||||
|
|
@ -1982,6 +1983,69 @@ def ocr_cost(
|
|||
return ocr_pages_cost + annotation_pages_cost, 0.0
|
||||
|
||||
|
||||
_OCR_PRICING_KEYS: Final = (
|
||||
"ocr_cost_per_page",
|
||||
"ocr_cost_per_page_batches",
|
||||
"annotation_cost_per_page",
|
||||
"annotation_cost_per_page_batches",
|
||||
)
|
||||
|
||||
|
||||
def ocr_batch_cost(
|
||||
model: str,
|
||||
custom_llm_provider: str | None,
|
||||
usage_info: "OCRUsageInfo",
|
||||
model_info: ModelInfo | None = None,
|
||||
) -> tuple[float, float]:
|
||||
"""Per-page cost of one OCR result inside a batch output file.
|
||||
|
||||
Batch OCR is billed per page at the ``*_batches`` rate, falling back to the
|
||||
synchronous per-page rate when a model has no batch price recorded, the same
|
||||
fallback ``batch_cost_calculator`` applies to per-token batch pricing. Returns
|
||||
``(prompt_cost, completion_cost)`` with the whole cost in the first slot, like
|
||||
``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)
|
||||
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.",
|
||||
model,
|
||||
custom_llm_provider,
|
||||
)
|
||||
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")
|
||||
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:
|
||||
verbose_logger.warning(
|
||||
"OCR batch cost: model=%s custom_llm_provider=%s reported pages_processed=%s but no "
|
||||
"ocr_cost_per_page is configured; returning 0.0 cost for those pages.",
|
||||
model,
|
||||
custom_llm_provider,
|
||||
pages_processed,
|
||||
)
|
||||
effective_annotation_rate: Final = annotation_rate if annotation_rate is not None else page_rate
|
||||
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)
|
||||
|
||||
|
||||
def vector_store_search_cost(
|
||||
model: str | None,
|
||||
custom_llm_provider: str,
|
||||
|
|
|
|||
|
|
@ -27,12 +27,13 @@ FileCreateProvider = Literal[
|
|||
"litellm_proxy",
|
||||
"manus",
|
||||
"anthropic",
|
||||
"mistral",
|
||||
]
|
||||
FileRetrieveProvider = Literal[
|
||||
"openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic"
|
||||
"openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic", "mistral"
|
||||
]
|
||||
FileDeleteProvider = Literal["openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic"]
|
||||
FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic"]
|
||||
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
|
||||
from litellm.files.streaming import FileContentStreamingResponse
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ from collections.abc import AsyncIterator, Iterator
|
|||
from typing import Literal, NamedTuple
|
||||
|
||||
FileContentProvider = Literal[
|
||||
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "manus"
|
||||
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "manus", "mistral"
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1210,8 +1210,8 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
return {"error": f"Unable to parse raw request body. Got - {data}"}
|
||||
return data
|
||||
|
||||
def _get_masked_api_base(self, api_base: str) -> str:
|
||||
return str(mask_api_base_credentials(api_base))
|
||||
def _get_masked_api_base(self, api_base: str | None) -> str:
|
||||
return str(mask_api_base_credentials(api_base or ""))
|
||||
|
||||
def _pre_call(self, input, api_key, model=None, additional_args={}):
|
||||
"""
|
||||
|
|
|
|||
0
litellm/llms/mistral/batches/__init__.py
Normal file
0
litellm/llms/mistral/batches/__init__.py
Normal file
220
litellm/llms/mistral/batches/transformation.py
Normal file
220
litellm/llms/mistral/batches/transformation.py
Normal file
|
|
@ -0,0 +1,220 @@
|
|||
"""
|
||||
Mistral Batch API. Reference: https://docs.mistral.ai/api/#tag/batch
|
||||
|
||||
Mistral runs one model per job (set on the job, not per input line) and accepts
|
||||
``/v1/ocr`` as a batch endpoint, which is how OCR gets its 50% batch discount.
|
||||
Output and error files are OpenAI-shaped JSONL (``{custom_id, response: {status_code, body}}``),
|
||||
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, 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
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.types.llms.openai import AllMessageValues, CreateBatchRequest
|
||||
from litellm.types.utils import LiteLLMBatch, LlmProviders
|
||||
|
||||
from ..common_utils import get_mistral_api_base, get_mistral_auth_headers, mistral_error
|
||||
|
||||
MistralBatchStatus: TypeAlias = Literal[
|
||||
"QUEUED", "RUNNING", "SUCCESS", "FAILED", "TIMEOUT_EXCEEDED", "CANCELLATION_REQUESTED", "CANCELLED"
|
||||
]
|
||||
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",
|
||||
"RUNNING": "in_progress",
|
||||
"SUCCESS": "completed",
|
||||
"FAILED": "failed",
|
||||
"TIMEOUT_EXCEEDED": "expired",
|
||||
"CANCELLATION_REQUESTED": "cancelling",
|
||||
"CANCELLED": "cancelled",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
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")
|
||||
|
||||
message: str
|
||||
count: int = 1
|
||||
|
||||
|
||||
class MistralBatchJob(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
id: str
|
||||
input_files: tuple[str, ...] = ()
|
||||
endpoint: str
|
||||
model: str | None = None
|
||||
status: MistralBatchStatus
|
||||
created_at: int
|
||||
started_at: int | None = None
|
||||
completed_at: int | None = None
|
||||
total_requests: int = 0
|
||||
completed_requests: int = 0
|
||||
succeeded_requests: int = 0
|
||||
failed_requests: int = 0
|
||||
output_file: str | None = None
|
||||
error_file: str | None = None
|
||||
errors: tuple[MistralBatchError, ...] = ()
|
||||
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:
|
||||
status: Final = _STATUS_MAP[job.status]
|
||||
terminal_at: Final = job.completed_at
|
||||
return LiteLLMBatch(
|
||||
id=job.id,
|
||||
object="batch",
|
||||
endpoint=job.endpoint,
|
||||
input_file_id=job.input_files[0] if job.input_files else "",
|
||||
completion_window="24h",
|
||||
status=status,
|
||||
created_at=job.created_at,
|
||||
in_progress_at=job.started_at,
|
||||
completed_at=terminal_at if status == "completed" else None,
|
||||
failed_at=terminal_at if status == "failed" else None,
|
||||
expired_at=terminal_at if status == "expired" else None,
|
||||
cancelled_at=terminal_at if status == "cancelled" else None,
|
||||
output_file_id=job.output_file,
|
||||
error_file_id=job.error_file,
|
||||
errors=_to_batch_errors(job.errors),
|
||||
request_counts=BatchRequestCounts(
|
||||
total=job.total_requests,
|
||||
completed=job.succeeded_requests,
|
||||
failed=job.failed_requests,
|
||||
),
|
||||
metadata=job.metadata,
|
||||
)
|
||||
|
||||
|
||||
class MistralBatchesConfig(BaseBatchesConfig):
|
||||
@property
|
||||
def custom_llm_provider(self) -> LlmProviders:
|
||||
return LlmProviders.MISTRAL
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: Mapping[str, str],
|
||||
model: str,
|
||||
messages: Sequence[AllMessageValues],
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict[str, str]: # mutable-ok: BaseBatchesConfig signature
|
||||
return get_mistral_auth_headers(headers, api_key)
|
||||
|
||||
def get_complete_batch_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
model: str,
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
data: CreateBatchRequest,
|
||||
) -> str:
|
||||
return f"{get_mistral_api_base(api_base)}/v1/batch/jobs"
|
||||
|
||||
def transform_create_batch_request(
|
||||
self,
|
||||
model: str,
|
||||
create_batch_data: CreateBatchRequest,
|
||||
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")
|
||||
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: 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: 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")
|
||||
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: 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: Mapping[str, str] | httpx.Headers
|
||||
) -> BaseLLMException:
|
||||
return mistral_error(error_message, status_code, headers)
|
||||
41
litellm/llms/mistral/common_utils.py
Normal file
41
litellm/llms/mistral/common_utils.py
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
MISTRAL_API_BASE: Final = "https://api.mistral.ai"
|
||||
MISTRAL_API_KEY_ENV_VAR: Final = "MISTRAL_API_KEY"
|
||||
|
||||
|
||||
class MistralError(BaseLLMException):
|
||||
pass
|
||||
|
||||
|
||||
def get_mistral_api_base(api_base: str | None) -> str:
|
||||
"""Return the Mistral origin without a trailing ``/v1``, so callers can append ``/v1/<route>``."""
|
||||
resolved: Final = (api_base or get_secret_str("MISTRAL_API_BASE") or MISTRAL_API_BASE).rstrip("/")
|
||||
return resolved.removesuffix("/v1")
|
||||
|
||||
|
||||
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 dict(headers, Authorization=f"Bearer {resolved_key}") # mutable-ok: BaseConfig contract returns dict
|
||||
|
||||
|
||||
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(dict(headers)), # mutable-ok: httpx.Headers takes a dict
|
||||
)
|
||||
0
litellm/llms/mistral/files/__init__.py
Normal file
0
litellm/llms/mistral/files/__init__.py
Normal file
256
litellm/llms/mistral/files/transformation.py
Normal file
256
litellm/llms/mistral/files/transformation.py
Normal file
|
|
@ -0,0 +1,256 @@
|
|||
"""
|
||||
Mistral Files API. Reference: https://docs.mistral.ai/api/#tag/files
|
||||
|
||||
Mistral's file objects already carry the OpenAI field names (id, bytes, created_at,
|
||||
filename, purpose), so this config is URL routing, auth, and a purpose mapping:
|
||||
Mistral only accepts ``fine-tune``, ``batch`` and ``ocr`` as upload purposes.
|
||||
"""
|
||||
|
||||
import time
|
||||
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
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.files.transformation import BaseFilesConfig, LiteLLMLoggingObj
|
||||
from litellm.types.llms.openai import (
|
||||
CreateFileRequest,
|
||||
FileContentRequest,
|
||||
HttpxBinaryResponseContent,
|
||||
OpenAICreateFileRequestOptionalParams,
|
||||
OpenAIFileObject,
|
||||
OpenAIFilesPurpose,
|
||||
)
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
from ..common_utils import get_mistral_api_base, get_mistral_auth_headers, mistral_error
|
||||
|
||||
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):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
id: str
|
||||
bytes: int = 0
|
||||
created_at: int | None = None
|
||||
filename: str = ""
|
||||
purpose: MistralFilePurpose = "batch"
|
||||
expires_at: int | None = None
|
||||
|
||||
|
||||
class MistralFileList(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
data: tuple[MistralFile, ...] = ()
|
||||
|
||||
|
||||
class MistralFileDeleted(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
id: str
|
||||
deleted: bool = True
|
||||
|
||||
|
||||
def _to_openai_file_object(file: MistralFile) -> OpenAIFileObject:
|
||||
return OpenAIFileObject(
|
||||
id=file.id,
|
||||
bytes=file.bytes,
|
||||
created_at=file.created_at if file.created_at is not None else int(time.time()),
|
||||
filename=file.filename,
|
||||
object="file",
|
||||
purpose=_to_openai_purpose(file.purpose),
|
||||
status="uploaded",
|
||||
expires_at=file.expires_at,
|
||||
)
|
||||
|
||||
|
||||
def _to_openai_purpose(purpose: MistralFilePurpose) -> OpenAIFilesPurpose:
|
||||
match purpose:
|
||||
case "fine-tune" | "batch":
|
||||
return purpose
|
||||
case "ocr":
|
||||
return "user_data"
|
||||
|
||||
|
||||
def _to_mistral_purpose(purpose: str) -> MistralFilePurpose:
|
||||
"""Only Mistral's own purposes pass through. Silently mapping anything else to ``batch``
|
||||
would let an upload skip the proxy's batch-file validation and guardrails, which only
|
||||
run when the caller says ``purpose=batch``."""
|
||||
match purpose:
|
||||
case "batch" | "fine-tune" | "ocr":
|
||||
return purpose
|
||||
case _:
|
||||
raise ValueError(f"Mistral does not support purpose={purpose!r}. Use one of: batch, fine-tune, ocr")
|
||||
|
||||
|
||||
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:
|
||||
return LlmProviders.MISTRAL
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
model: str,
|
||||
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: Mapping[str, object], suffix: str = "") -> str:
|
||||
encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id")
|
||||
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: Mapping[str, str] | httpx.Headers
|
||||
) -> BaseLLMException:
|
||||
return mistral_error(error_message, status_code, headers)
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: Mapping[str, str],
|
||||
model: str,
|
||||
messages: Sequence[object],
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> 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]: # mutable-ok: BaseFilesConfig signature
|
||||
return ["purpose"] # mutable-ok: BaseFilesConfig signature
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: Mapping[str, object],
|
||||
optional_params: dict[str, object], # mutable-ok: BaseConfig signature, returned as-is
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict[str, object]: # mutable-ok: BaseConfig signature
|
||||
return optional_params
|
||||
|
||||
def transform_create_file_request(
|
||||
self,
|
||||
model: str,
|
||||
create_file_data: CreateFileRequest,
|
||||
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(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"
|
||||
upload: Final = MistralMultipartUpload(
|
||||
file=(filename, extracted["content"], content_type),
|
||||
purpose=(None, _to_mistral_purpose(create_file_data.get("purpose") or "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: 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: 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: 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: 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: Mapping[str, object],
|
||||
) -> FileDeleted:
|
||||
deleted: Final = MistralFileDeleted.model_validate(raw_response.json())
|
||||
return FileDeleted(id=deleted.id, deleted=deleted.deleted, object="file")
|
||||
|
||||
def transform_list_files_request(
|
||||
self,
|
||||
purpose: str | None,
|
||||
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: 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: 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: Mapping[str, object],
|
||||
) -> HttpxBinaryResponseContent:
|
||||
return HttpxBinaryResponseContent(response=raw_response)
|
||||
|
|
@ -35500,51 +35500,66 @@
|
|||
"mistral/mistral-ocr-latest": {
|
||||
"litellm_provider": "mistral",
|
||||
"ocr_cost_per_page": 0.004,
|
||||
"ocr_cost_per_page_batches": 0.002,
|
||||
"annotation_cost_per_page": 0.005,
|
||||
"annotation_cost_per_page_batches": 0.0025,
|
||||
"mode": "ocr",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
"/v1/ocr",
|
||||
"/v1/batch"
|
||||
],
|
||||
"source": "https://mistral.ai/pricing#api-pricing"
|
||||
},
|
||||
"mistral/mistral-ocr-4-0": {
|
||||
"litellm_provider": "mistral",
|
||||
"ocr_cost_per_page": 0.004,
|
||||
"ocr_cost_per_page_batches": 0.002,
|
||||
"annotation_cost_per_page": 0.005,
|
||||
"annotation_cost_per_page_batches": 0.0025,
|
||||
"mode": "ocr",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
"/v1/ocr",
|
||||
"/v1/batch"
|
||||
],
|
||||
"source": "https://mistral.ai/pricing#api-pricing"
|
||||
},
|
||||
"mistral/mistral-ocr-4-1": {
|
||||
"annotation_cost_per_page": 0.005,
|
||||
"annotation_cost_per_page_batches": 0.0025,
|
||||
"litellm_provider": "mistral",
|
||||
"mode": "ocr",
|
||||
"ocr_cost_per_page": 0.004,
|
||||
"ocr_cost_per_page_batches": 0.002,
|
||||
"source": "https://docs.mistral.ai/models/model-cards/ocr-4-1",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
"/v1/ocr",
|
||||
"/v1/batch"
|
||||
]
|
||||
},
|
||||
"mistral/mistral-ocr-2505-completion": {
|
||||
"deprecation_date": "2026-05-31",
|
||||
"litellm_provider": "mistral",
|
||||
"ocr_cost_per_page": 0.001,
|
||||
"ocr_cost_per_page_batches": 0.0005,
|
||||
"annotation_cost_per_page": 0.003,
|
||||
"annotation_cost_per_page_batches": 0.0015,
|
||||
"mode": "ocr",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
"/v1/ocr",
|
||||
"/v1/batch"
|
||||
],
|
||||
"source": "https://mistral.ai/pricing#api-pricing"
|
||||
},
|
||||
"mistral/mistral-ocr-2512": {
|
||||
"litellm_provider": "mistral",
|
||||
"ocr_cost_per_page": 0.002,
|
||||
"ocr_cost_per_page_batches": 0.001,
|
||||
"annotation_cost_per_page": 0.003,
|
||||
"annotation_cost_per_page_batches": 0.0015,
|
||||
"mode": "ocr",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
"/v1/ocr",
|
||||
"/v1/batch"
|
||||
],
|
||||
"source": "https://mistral.ai/pricing#api-pricing"
|
||||
},
|
||||
|
|
@ -60665,31 +60680,40 @@
|
|||
"mistral/mistral-ocr-3": {
|
||||
"litellm_provider": "mistral",
|
||||
"ocr_cost_per_page": 0.002,
|
||||
"ocr_cost_per_page_batches": 0.001,
|
||||
"annotation_cost_per_page": 0.003,
|
||||
"annotation_cost_per_page_batches": 0.0015,
|
||||
"mode": "ocr",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
"/v1/ocr",
|
||||
"/v1/batch"
|
||||
],
|
||||
"source": "https://mistral.ai/pricing#api-pricing"
|
||||
},
|
||||
"mistral/mistral-ocr-3-0": {
|
||||
"litellm_provider": "mistral",
|
||||
"ocr_cost_per_page": 0.002,
|
||||
"ocr_cost_per_page_batches": 0.001,
|
||||
"annotation_cost_per_page": 0.003,
|
||||
"annotation_cost_per_page_batches": 0.0015,
|
||||
"mode": "ocr",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
"/v1/ocr",
|
||||
"/v1/batch"
|
||||
],
|
||||
"source": "https://mistral.ai/pricing#api-pricing"
|
||||
},
|
||||
"mistral/mistral-ocr-4": {
|
||||
"annotation_cost_per_page": 0.005,
|
||||
"annotation_cost_per_page_batches": 0.0025,
|
||||
"litellm_provider": "mistral",
|
||||
"mode": "ocr",
|
||||
"ocr_cost_per_page": 0.004,
|
||||
"ocr_cost_per_page_batches": 0.002,
|
||||
"source": "https://docs.mistral.ai/models/model-cards/ocr-4-1",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
"/v1/ocr",
|
||||
"/v1/batch"
|
||||
]
|
||||
},
|
||||
"mistral/voxtral-mini-latest": {
|
||||
|
|
|
|||
|
|
@ -29,14 +29,15 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
|
|||
_is_base64_encoded_unified_file_id,
|
||||
add_internal_model_credentials,
|
||||
apply_team_provider_credentials,
|
||||
authorize_model_for_key,
|
||||
batch_cost_poller_is_active,
|
||||
decode_model_from_file_id,
|
||||
encode_batch_response_ids,
|
||||
encode_file_id_with_model,
|
||||
ensure_batch_response_managed_file_ids,
|
||||
get_authorized_credentials_for_model,
|
||||
get_batch_from_database,
|
||||
get_batch_id_from_unified_batch_id,
|
||||
get_credentials_for_model,
|
||||
get_model_id_from_unified_batch_id,
|
||||
get_models_from_unified_file_id,
|
||||
get_original_file_id,
|
||||
|
|
@ -218,9 +219,10 @@ async def create_batch(
|
|||
|
||||
# SCENARIO 1: File ID is encoded with model info
|
||||
if model_from_file_id is not None and input_file_id:
|
||||
credentials = get_credentials_for_model(
|
||||
credentials = await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=model_from_file_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="batch creation (file created with model)",
|
||||
)
|
||||
|
||||
|
|
@ -285,6 +287,7 @@ async def create_batch(
|
|||
detail={"error": f"Expected 1 model, got {len(target_model_names)}"},
|
||||
)
|
||||
model: Final = target_model_names[0]
|
||||
await authorize_model_for_key(model_id=model, llm_router=llm_router, user_api_key_dict=user_api_key_dict)
|
||||
_create_batch_data["model"] = model
|
||||
|
||||
resolved_storage_url: Final = await _resolve_managed_input_file_storage_url(input_file_id)
|
||||
|
|
@ -310,9 +313,10 @@ async def create_batch(
|
|||
# SCENARIO 2 & 3: Model from header/query OR custom_llm_provider fallback
|
||||
if model_param:
|
||||
# SCENARIO 2: Use model-based routing from header/query/body
|
||||
credentials = get_credentials_for_model(
|
||||
credentials = await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=model_param,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="batch creation",
|
||||
)
|
||||
|
||||
|
|
@ -540,9 +544,10 @@ async def retrieve_batch(
|
|||
# Retrieve from provider (for non-terminal states or if DB lookup failed)
|
||||
# SCENARIO 1: Batch ID is encoded with model info
|
||||
if model_from_id is not None:
|
||||
credentials: Final = get_credentials_for_model(
|
||||
credentials: Final = await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=model_from_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="batch retrieval (batch created with model)",
|
||||
)
|
||||
|
||||
|
|
@ -579,10 +584,17 @@ async def retrieve_batch(
|
|||
)
|
||||
|
||||
if unified_batch_id:
|
||||
unified_model_id: Final = get_model_id_from_unified_batch_id(unified_batch_id)
|
||||
if unified_model_id is not None:
|
||||
await authorize_model_for_key(
|
||||
model_id=llm_router.resolve_model_name_from_model_id(unified_model_id) or unified_model_id,
|
||||
llm_router=llm_router,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
add_internal_model_credentials(
|
||||
data=data,
|
||||
llm_router=llm_router,
|
||||
model_id=get_model_id_from_unified_batch_id(unified_batch_id),
|
||||
model_id=unified_model_id,
|
||||
)
|
||||
|
||||
response = await llm_router.aretrieve_batch(**data)
|
||||
|
|
@ -764,9 +776,10 @@ async def list_batches(
|
|||
data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model")
|
||||
):
|
||||
# SCENARIO 2: Use model-based routing from header/query/body
|
||||
credentials: Final = get_credentials_for_model(
|
||||
credentials: Final = await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=model_param,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="batch listing",
|
||||
)
|
||||
|
||||
|
|
@ -952,9 +965,10 @@ async def cancel_batch(
|
|||
|
||||
# SCENARIO 1: Batch ID is encoded with model info
|
||||
if model_from_id is not None:
|
||||
credentials: Final = get_credentials_for_model(
|
||||
credentials: Final = await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=model_from_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="batch cancellation (batch created with model)",
|
||||
)
|
||||
|
||||
|
|
@ -993,6 +1007,11 @@ async def cancel_batch(
|
|||
status_code=400,
|
||||
detail={"error": "Invalid LiteLLM managed batch ID. Missing model_id."},
|
||||
)
|
||||
await authorize_model_for_key(
|
||||
model_id=llm_router.resolve_model_name_from_model_id(model_id_from_batch) or model_id_from_batch,
|
||||
llm_router=llm_router,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
data["model"] = model_id_from_batch
|
||||
data["batch_id"] = get_batch_id_from_unified_batch_id(unified_batch_id)
|
||||
response = await llm_router.acancel_batch(**data)
|
||||
|
|
|
|||
|
|
@ -351,6 +351,10 @@ def get_credentials_for_model(
|
|||
"""
|
||||
Retrieve API credentials for a model from the LLM Router.
|
||||
|
||||
Does not check whether the caller may use ``model_id``; use
|
||||
``get_authorized_credentials_for_model`` for anything driven by a caller-supplied
|
||||
model name (request body, header, query param, or a model-encoded resource id).
|
||||
|
||||
Args:
|
||||
llm_router: LiteLLM Router instance
|
||||
model_id: Model name or deployment ID
|
||||
|
|
@ -381,6 +385,48 @@ def get_credentials_for_model(
|
|||
return credentials
|
||||
|
||||
|
||||
async def authorize_model_for_key(
|
||||
model_id: str,
|
||||
llm_router: Optional["Router"],
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
) -> None:
|
||||
"""
|
||||
Enforce the caller's model grants on a model name the auth layer never saw.
|
||||
|
||||
The files and batches routes carry their model in a header, query param, or a
|
||||
model-encoded resource id rather than the request body, so ``user_api_key_auth``
|
||||
cannot check it. Run the same key, team (incl. team-member and access-group
|
||||
fallbacks), org and project allowlist checks a chat request would get, so a
|
||||
restricted key cannot borrow another deployment's server-side credentials.
|
||||
|
||||
Raises:
|
||||
ProxyException (403): the caller is not allowed to use ``model_id``
|
||||
"""
|
||||
from litellm.proxy.auth.auth_checks import can_key_call_resolved_model
|
||||
|
||||
await can_key_call_resolved_model(
|
||||
model=model_id,
|
||||
llm_model_list=None,
|
||||
valid_token=user_api_key_dict,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
|
||||
async def get_authorized_credentials_for_model(
|
||||
llm_router: Optional["Router"],
|
||||
model_id: str,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
operation_context: str = "file operation",
|
||||
) -> dict: # mutable-ok: same contract as get_credentials_for_model, callers merge it into request data
|
||||
"""``get_credentials_for_model`` gated by ``authorize_model_for_key``."""
|
||||
await authorize_model_for_key(model_id=model_id, llm_router=llm_router, user_api_key_dict=user_api_key_dict)
|
||||
return get_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=model_id,
|
||||
operation_context=operation_context,
|
||||
)
|
||||
|
||||
|
||||
def get_team_provider_credentials(
|
||||
llm_router: Optional["Router"],
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
|
|
@ -573,21 +619,27 @@ def prepare_data_with_credentials(
|
|||
data["file_id"] = file_id
|
||||
|
||||
|
||||
def handle_model_based_routing(
|
||||
async def handle_model_based_routing(
|
||||
file_id: str,
|
||||
request, # FastAPI Request object
|
||||
llm_router, # Router instance
|
||||
data: dict,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
check_file_id_encoding: bool = True,
|
||||
) -> tuple[bool, str | None, str | None, dict | None]:
|
||||
"""
|
||||
Orchestrate model-based credential routing for file operations.
|
||||
|
||||
The model name comes from the caller (embedded in the file id, or a header, query
|
||||
param or body field), so it is authorized against the caller's key, team, org and
|
||||
project grants before any deployment credentials are resolved.
|
||||
|
||||
Args:
|
||||
file_id: File ID (may contain embedded model info)
|
||||
request: FastAPI request object
|
||||
llm_router: LiteLLM Router instance
|
||||
data: Request data dictionary
|
||||
user_api_key_dict: The authenticated caller
|
||||
check_file_id_encoding: Whether to check for embedded model in file_id
|
||||
|
||||
Returns:
|
||||
|
|
@ -599,6 +651,7 @@ def handle_model_based_routing(
|
|||
|
||||
Raises:
|
||||
HTTPException: If router unavailable or model not found
|
||||
ProxyException: If the caller is not allowed to use the model
|
||||
"""
|
||||
model_from_id, model_from_param = extract_model_from_sources(
|
||||
file_id=file_id,
|
||||
|
|
@ -608,9 +661,10 @@ def handle_model_based_routing(
|
|||
|
||||
# Priority 1: Model embedded in file_id
|
||||
if check_file_id_encoding and model_from_id is not None:
|
||||
credentials = get_credentials_for_model(
|
||||
credentials = await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=model_from_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context=f"file operation (file created with model '{model_from_id}')",
|
||||
)
|
||||
original_file_id: Final = get_original_file_id(file_id)
|
||||
|
|
@ -618,9 +672,10 @@ def handle_model_based_routing(
|
|||
|
||||
# Priority 2: Model from header/query/body
|
||||
elif model_from_param is not None:
|
||||
credentials = get_credentials_for_model(
|
||||
credentials = await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=model_from_param,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="file operation",
|
||||
)
|
||||
return True, model_from_param, None, credentials
|
||||
|
|
|
|||
|
|
@ -68,7 +68,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
|
|||
apply_team_provider_credentials,
|
||||
encode_file_id_with_model,
|
||||
extract_file_creation_params,
|
||||
get_credentials_for_model,
|
||||
get_authorized_credentials_for_model,
|
||||
handle_model_based_routing,
|
||||
prepare_data_with_credentials,
|
||||
validate_file_list_limit,
|
||||
|
|
@ -267,9 +267,10 @@ async def route_create_file(
|
|||
# NEW: Handle model-based routing (no DB required)
|
||||
if model is not None:
|
||||
# Get credentials from model_list via router
|
||||
credentials: Final = get_credentials_for_model(
|
||||
credentials: Final = await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=model,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="file upload",
|
||||
)
|
||||
|
||||
|
|
@ -907,11 +908,12 @@ async def get_file_content(
|
|||
model_used,
|
||||
original_file_id,
|
||||
credentials,
|
||||
) = handle_model_based_routing(
|
||||
) = await handle_model_based_routing(
|
||||
file_id=file_id,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
check_file_id_encoding=True,
|
||||
)
|
||||
|
||||
|
|
@ -1122,15 +1124,16 @@ async def get_file(
|
|||
model_used,
|
||||
original_file_id,
|
||||
credentials,
|
||||
) = handle_model_based_routing(
|
||||
) = await handle_model_based_routing(
|
||||
file_id=file_id,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
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,
|
||||
|
|
@ -1139,7 +1142,10 @@ async def get_file(
|
|||
include_internal_credentials=True,
|
||||
)
|
||||
|
||||
response = await litellm.afile_retrieve(**data)
|
||||
response = await litellm.afile_retrieve(
|
||||
custom_llm_provider=credentials["custom_llm_provider"],
|
||||
**data,
|
||||
)
|
||||
|
||||
# Keep the encoded ID in response if it was originally encoded
|
||||
if original_file_id and response and hasattr(response, "id") and response.id:
|
||||
|
|
@ -1327,11 +1333,12 @@ async def delete_file(
|
|||
model_used,
|
||||
original_file_id,
|
||||
credentials,
|
||||
) = handle_model_based_routing(
|
||||
) = await handle_model_based_routing(
|
||||
file_id=file_id,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
check_file_id_encoding=True,
|
||||
)
|
||||
|
||||
|
|
@ -1514,11 +1521,12 @@ async def list_files(
|
|||
response: Any | None = None
|
||||
|
||||
# Check for model-based credential routing (no file_id encoding check for list)
|
||||
should_route, model_used, _, credentials = handle_model_based_routing(
|
||||
should_route, model_used, _, credentials = await handle_model_based_routing(
|
||||
file_id="", # No file_id for list endpoint
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
check_file_id_encoding=False,
|
||||
)
|
||||
|
||||
|
|
@ -1545,9 +1553,10 @@ async def list_files(
|
|||
status_code=500,
|
||||
detail="LLM Router not initialized. Ensure models added to proxy.",
|
||||
)
|
||||
credentials = get_credentials_for_model(
|
||||
credentials = await get_authorized_credentials_for_model(
|
||||
llm_router=llm_router,
|
||||
model_id=target_model_names_list[0],
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
operation_context="file list",
|
||||
)
|
||||
prepare_data_with_credentials(data=data, credentials=credentials)
|
||||
|
|
|
|||
|
|
@ -144,11 +144,12 @@ async def _update_request_data_with_managed_file_id(
|
|||
model_used,
|
||||
original_file_id,
|
||||
credentials,
|
||||
) = handle_model_based_routing(
|
||||
) = await handle_model_based_routing(
|
||||
file_id=file_id,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
check_file_id_encoding=True,
|
||||
)
|
||||
|
||||
|
|
@ -273,11 +274,12 @@ async def _update_request_data_with_model_routing_hint(
|
|||
_model_used,
|
||||
_original_file_id,
|
||||
credentials,
|
||||
) = handle_model_based_routing(
|
||||
) = await handle_model_based_routing(
|
||||
file_id="",
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
check_file_id_encoding=False,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -499,7 +499,7 @@ class CreateBatchRequest(TypedDict, total=False):
|
|||
"""
|
||||
|
||||
completion_window: Literal["24h"]
|
||||
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"]
|
||||
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"]
|
||||
input_file_id: str
|
||||
metadata: dict[str, str] | None
|
||||
output_expires_after: FileExpiresAfter
|
||||
|
|
|
|||
|
|
@ -322,8 +322,10 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
output_cost_per_second_720p: ReadOnly[float | None]
|
||||
output_cost_per_second_4k: ReadOnly[float | None]
|
||||
ocr_cost_per_page: float | None # for OCR models
|
||||
ocr_cost_per_page_batches: ReadOnly[float | None]
|
||||
ocr_cost_per_credit: float | None # for OCR models priced by credit
|
||||
annotation_cost_per_page: float | None # for OCR models
|
||||
annotation_cost_per_page_batches: ReadOnly[float | None]
|
||||
search_context_cost_per_query: SearchContextCostPerQuery | None # Cost for using web search tool
|
||||
web_search_billing_unit: (
|
||||
Literal["per_query", "per_prompt"] | None
|
||||
|
|
@ -3607,8 +3609,10 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
|
|||
output_cost_per_token_above_512k_tokens: float | None = None
|
||||
output_vector_size: int | None = None
|
||||
ocr_cost_per_page: float | None = None
|
||||
ocr_cost_per_page_batches: float | None = None
|
||||
ocr_cost_per_credit: float | None = None
|
||||
annotation_cost_per_page: float | None = None
|
||||
annotation_cost_per_page_batches: float | None = None
|
||||
regional_processing_uplift_multiplier_eu: float | None = None
|
||||
regional_processing_uplift_multiplier_us: float | None = None
|
||||
regional_endpoint_uplift_multiplier: float | None = None
|
||||
|
|
|
|||
|
|
@ -6012,8 +6012,10 @@ def _get_model_info_helper(
|
|||
tpm=_model_info.get("tpm", None),
|
||||
rpm=_model_info.get("rpm", None),
|
||||
ocr_cost_per_page=_model_info.get("ocr_cost_per_page", None),
|
||||
ocr_cost_per_page_batches=_model_info.get("ocr_cost_per_page_batches", None),
|
||||
ocr_cost_per_credit=_model_info.get("ocr_cost_per_credit", None),
|
||||
annotation_cost_per_page=_model_info.get("annotation_cost_per_page", None),
|
||||
annotation_cost_per_page_batches=_model_info.get("annotation_cost_per_page_batches", None),
|
||||
provider_specific_entry=_model_info.get("provider_specific_entry", None),
|
||||
uses_embed_content=_model_info.get("uses_embed_content", None),
|
||||
supports_image_size=_model_info.get("supports_image_size", None),
|
||||
|
|
@ -8988,6 +8990,10 @@ class ProviderConfigManager:
|
|||
from litellm.llms.anthropic.files.transformation import AnthropicFilesConfig
|
||||
|
||||
return AnthropicFilesConfig()
|
||||
elif LlmProviders.MISTRAL == provider:
|
||||
from litellm.llms.mistral.files.transformation import MistralFilesConfig
|
||||
|
||||
return MistralFilesConfig()
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -8999,6 +9005,10 @@ class ProviderConfigManager:
|
|||
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
|
||||
|
||||
return BedrockBatchesConfig()
|
||||
elif LlmProviders.MISTRAL == provider:
|
||||
from litellm.llms.mistral.batches.transformation import MistralBatchesConfig
|
||||
|
||||
return MistralBatchesConfig()
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -35500,51 +35500,66 @@
|
|||
"mistral/mistral-ocr-latest": {
|
||||
"litellm_provider": "mistral",
|
||||
"ocr_cost_per_page": 0.004,
|
||||
"ocr_cost_per_page_batches": 0.002,
|
||||
"annotation_cost_per_page": 0.005,
|
||||
"annotation_cost_per_page_batches": 0.0025,
|
||||
"mode": "ocr",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
"/v1/ocr",
|
||||
"/v1/batch"
|
||||
],
|
||||
"source": "https://mistral.ai/pricing#api-pricing"
|
||||
},
|
||||
"mistral/mistral-ocr-4-0": {
|
||||
"litellm_provider": "mistral",
|
||||
"ocr_cost_per_page": 0.004,
|
||||
"ocr_cost_per_page_batches": 0.002,
|
||||
"annotation_cost_per_page": 0.005,
|
||||
"annotation_cost_per_page_batches": 0.0025,
|
||||
"mode": "ocr",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
"/v1/ocr",
|
||||
"/v1/batch"
|
||||
],
|
||||
"source": "https://mistral.ai/pricing#api-pricing"
|
||||
},
|
||||
"mistral/mistral-ocr-4-1": {
|
||||
"annotation_cost_per_page": 0.005,
|
||||
"annotation_cost_per_page_batches": 0.0025,
|
||||
"litellm_provider": "mistral",
|
||||
"mode": "ocr",
|
||||
"ocr_cost_per_page": 0.004,
|
||||
"ocr_cost_per_page_batches": 0.002,
|
||||
"source": "https://docs.mistral.ai/models/model-cards/ocr-4-1",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
"/v1/ocr",
|
||||
"/v1/batch"
|
||||
]
|
||||
},
|
||||
"mistral/mistral-ocr-2505-completion": {
|
||||
"deprecation_date": "2026-05-31",
|
||||
"litellm_provider": "mistral",
|
||||
"ocr_cost_per_page": 0.001,
|
||||
"ocr_cost_per_page_batches": 0.0005,
|
||||
"annotation_cost_per_page": 0.003,
|
||||
"annotation_cost_per_page_batches": 0.0015,
|
||||
"mode": "ocr",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
"/v1/ocr",
|
||||
"/v1/batch"
|
||||
],
|
||||
"source": "https://mistral.ai/pricing#api-pricing"
|
||||
},
|
||||
"mistral/mistral-ocr-2512": {
|
||||
"litellm_provider": "mistral",
|
||||
"ocr_cost_per_page": 0.002,
|
||||
"ocr_cost_per_page_batches": 0.001,
|
||||
"annotation_cost_per_page": 0.003,
|
||||
"annotation_cost_per_page_batches": 0.0015,
|
||||
"mode": "ocr",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
"/v1/ocr",
|
||||
"/v1/batch"
|
||||
],
|
||||
"source": "https://mistral.ai/pricing#api-pricing"
|
||||
},
|
||||
|
|
@ -60665,31 +60680,40 @@
|
|||
"mistral/mistral-ocr-3": {
|
||||
"litellm_provider": "mistral",
|
||||
"ocr_cost_per_page": 0.002,
|
||||
"ocr_cost_per_page_batches": 0.001,
|
||||
"annotation_cost_per_page": 0.003,
|
||||
"annotation_cost_per_page_batches": 0.0015,
|
||||
"mode": "ocr",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
"/v1/ocr",
|
||||
"/v1/batch"
|
||||
],
|
||||
"source": "https://mistral.ai/pricing#api-pricing"
|
||||
},
|
||||
"mistral/mistral-ocr-3-0": {
|
||||
"litellm_provider": "mistral",
|
||||
"ocr_cost_per_page": 0.002,
|
||||
"ocr_cost_per_page_batches": 0.001,
|
||||
"annotation_cost_per_page": 0.003,
|
||||
"annotation_cost_per_page_batches": 0.0015,
|
||||
"mode": "ocr",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
"/v1/ocr",
|
||||
"/v1/batch"
|
||||
],
|
||||
"source": "https://mistral.ai/pricing#api-pricing"
|
||||
},
|
||||
"mistral/mistral-ocr-4": {
|
||||
"annotation_cost_per_page": 0.005,
|
||||
"annotation_cost_per_page_batches": 0.0025,
|
||||
"litellm_provider": "mistral",
|
||||
"mode": "ocr",
|
||||
"ocr_cost_per_page": 0.004,
|
||||
"ocr_cost_per_page_batches": 0.002,
|
||||
"source": "https://docs.mistral.ai/models/model-cards/ocr-4-1",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
"/v1/ocr",
|
||||
"/v1/batch"
|
||||
]
|
||||
},
|
||||
"mistral/voxtral-mini-latest": {
|
||||
|
|
|
|||
|
|
@ -53,6 +53,10 @@
|
|||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"annotation_cost_per_page_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"audio_transcription_config": {
|
||||
"type": "string"
|
||||
},
|
||||
|
|
@ -432,6 +436,10 @@
|
|||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"ocr_cost_per_page_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"output_cost_per_audio_token": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
@ -1787,3 +1807,85 @@ class TestBatchCostIsFinal:
|
|||
@pytest.mark.parametrize("status", ["failed", "expired", "cancelled"])
|
||||
def test_other_terminal_statuses_are_final(self, status):
|
||||
assert bu.batch_cost_is_final(_retrieved_batch(status)) is True
|
||||
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
|
||||
def test_ocr_rows_are_priced_per_page_at_batch_rate(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"get_model_info",
|
||||
lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.004, "ocr_cost_per_page_batches": 0.002},
|
||||
)
|
||||
result = bu._aggregate_batch_cost_usage_models(
|
||||
entries=[_ocr_row(3), _ocr_row(5), _failed_row(model="mistral-ocr-latest")],
|
||||
custom_llm_provider="mistral",
|
||||
model_name="mistral/mistral-ocr-latest",
|
||||
)
|
||||
assert result.cost == pytest.approx(8 * 0.002)
|
||||
assert result.prompt_cost == pytest.approx(8 * 0.002)
|
||||
assert result.completion_cost == 0.0
|
||||
assert (result.successful_requests, result.failed_requests) == (2, 1)
|
||||
assert result.usage.total_tokens == 0
|
||||
assert result.models == ["mistral/mistral-ocr-latest"]
|
||||
|
||||
|
||||
def test_ocr_rows_fall_back_to_sync_page_rate_without_batch_price(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "get_model_info", lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.004})
|
||||
result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(2)], custom_llm_provider="mistral")
|
||||
assert result.cost == pytest.approx(2 * 0.004)
|
||||
|
||||
|
||||
def test_ocr_rows_bill_annotation_pages_separately(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"get_model_info",
|
||||
lambda model, custom_llm_provider=None: {
|
||||
"ocr_cost_per_page_batches": 0.002,
|
||||
"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"
|
||||
)
|
||||
assert result.cost == pytest.approx(4 * 0.002 + 4 * 0.0025)
|
||||
|
||||
|
||||
def test_ocr_rows_use_deployment_model_info_pricing_over_cost_map(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
litellm, "get_model_info", lambda model, custom_llm_provider=None: pytest.fail("cost map must not be consulted")
|
||||
)
|
||||
result = bu._aggregate_batch_cost_usage_models(
|
||||
entries=[_ocr_row(10)],
|
||||
custom_llm_provider="mistral",
|
||||
model_info={"ocr_cost_per_page_batches": 0.001},
|
||||
)
|
||||
assert result.cost == pytest.approx(0.01)
|
||||
|
||||
|
||||
def test_ocr_rows_without_pricing_bill_zero_but_count_as_successful(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "get_model_info", lambda model, custom_llm_provider=None: {"mode": "ocr"})
|
||||
result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(3)], custom_llm_provider="mistral")
|
||||
assert result.cost == 0.0
|
||||
assert (result.successful_requests, result.failed_requests) == (1, 0)
|
||||
|
||||
|
||||
def test_chat_rows_from_mistral_still_use_token_pricing(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"get_model_info",
|
||||
lambda model, custom_llm_provider=None: {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002},
|
||||
)
|
||||
result = bu._aggregate_batch_cost_usage_models(
|
||||
entries=[_success_row(model="mistral-small-latest", usage=_usage(10, 5))],
|
||||
custom_llm_provider="mistral",
|
||||
)
|
||||
assert result.cost == pytest.approx((10 * 0.001 + 5 * 0.002) / 2)
|
||||
assert result.usage.total_tokens == 15
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
@ -778,3 +748,43 @@ def test_retrieve__omits_trusted_model_credentials_when_not_supplied(seams):
|
|||
|
||||
litellm_params = logging_obj.update_from_kwargs.call_args.kwargs["litellm_params"]
|
||||
assert "_litellm_internal_model_credentials" not in litellm_params
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
# mistral - a provider-config provider, like bedrock, so it requires `model`
|
||||
# =========================================================================== #
|
||||
|
||||
|
||||
def test_create__mistral_ocr_routes_to_base_http_handler_with_mistral_config(seams):
|
||||
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")
|
||||
forwarded = seams.base_http.create_batch.call_args.kwargs
|
||||
assert type(forwarded["provider_config"]).__name__ == "MistralBatchesConfig"
|
||||
assert forwarded["model"] == "mistral-ocr-latest"
|
||||
assert forwarded["create_batch_data"]["endpoint"] == "/v1/ocr"
|
||||
|
||||
|
||||
def test_create__mistral_without_model_raises_badrequest(seams):
|
||||
with pytest.raises(litellm.exceptions.BadRequestError):
|
||||
bm.create_batch(**CREATE_KW, custom_llm_provider="mistral")
|
||||
|
||||
for m in _all_seam_methods(seams, "create_batch"):
|
||||
m.assert_not_called()
|
||||
|
||||
|
||||
def test_retrieve__mistral_routes_to_base_http_handler_with_mistral_config(seams):
|
||||
result = bm.retrieve_batch(batch_id="job-1", custom_llm_provider="mistral", model="mistral/mistral-ocr-latest")
|
||||
|
||||
assert result is seams.base_http.retrieve_batch.return_value
|
||||
_assert_only(seams.base_http.retrieve_batch, seams, "retrieve_batch")
|
||||
forwarded = seams.base_http.retrieve_batch.call_args.kwargs
|
||||
assert type(forwarded["provider_config"]).__name__ == "MistralBatchesConfig"
|
||||
assert forwarded["batch_id"] == "job-1"
|
||||
|
|
|
|||
|
|
@ -59,6 +59,16 @@ def test_get_masked_api_base(logging_obj):
|
|||
assert type(masked_api_base) == str
|
||||
|
||||
|
||||
def test_pre_call_tolerates_missing_api_base(logging_obj):
|
||||
"""Presigned batch retrieves (Mistral, Bedrock) build their own URL and pass api_base=None
|
||||
to pre_call; masking must not raise or the request's pre-call logging is silently lost."""
|
||||
logging_obj.update_environment_variables(litellm_params={}, optional_params={})
|
||||
|
||||
logging_obj.pre_call(input="", api_key="", additional_args={"api_base": None, "headers": {}})
|
||||
|
||||
assert logging_obj.model_call_details["litellm_params"]["api_base"] == ""
|
||||
|
||||
|
||||
def test_post_call_serializes_dict_with_datetime(logging_obj):
|
||||
import datetime
|
||||
|
||||
|
|
@ -3977,9 +3987,7 @@ def test_get_standard_logging_object_payload_carries_matched_access_groups(loggi
|
|||
"model": "gpt-4o",
|
||||
"messages": [],
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"user_api_key_matched_model_access_groups": ["premium-pool", "shared-pool"]
|
||||
},
|
||||
"metadata": {"user_api_key_matched_model_access_groups": ["premium-pool", "shared-pool"]},
|
||||
"proxy_server_request": {"body": {}},
|
||||
},
|
||||
},
|
||||
|
|
@ -4063,9 +4071,7 @@ def _model_router_response(selected_model: str, stamp: bool):
|
|||
from litellm.types.utils import ModelResponse
|
||||
|
||||
response = ModelResponse(model=selected_model)
|
||||
response._hidden_params = (
|
||||
{AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {}
|
||||
)
|
||||
response._hidden_params = {AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {}
|
||||
return response
|
||||
|
||||
|
||||
|
|
@ -4089,9 +4095,7 @@ def test_standard_logging_payload_uses_stamped_model_router_model(logging_obj):
|
|||
"messages": [],
|
||||
"litellm_params": {"metadata": {}},
|
||||
},
|
||||
init_response_obj=_model_router_response(
|
||||
"azure_ai/grok-4-1-fast-reasoning", stamp=True
|
||||
),
|
||||
init_response_obj=_model_router_response("azure_ai/grok-4-1-fast-reasoning", stamp=True),
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
logging_obj=logging_obj,
|
||||
|
|
@ -4123,9 +4127,7 @@ def test_standard_logging_payload_keeps_requested_model_without_router_stamp(
|
|||
"messages": [],
|
||||
"litellm_params": {"metadata": {}},
|
||||
},
|
||||
init_response_obj=_model_router_response(
|
||||
"azure_ai/grok-4-1-fast-reasoning", stamp=False
|
||||
),
|
||||
init_response_obj=_model_router_response("azure_ai/grok-4-1-fast-reasoning", stamp=False),
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
logging_obj=logging_obj,
|
||||
|
|
@ -5537,9 +5539,7 @@ class TestNonInferenceCallTypesAreNotBilled:
|
|||
init_response_obj=self._retrieved_response(),
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
logging_obj=self._logging_obj(
|
||||
"aget_responses", litellm_metadata=self.BACKGROUND_POLL_METADATA
|
||||
),
|
||||
logging_obj=self._logging_obj("aget_responses", litellm_metadata=self.BACKGROUND_POLL_METADATA),
|
||||
status="success",
|
||||
)
|
||||
|
||||
|
|
@ -5785,9 +5785,7 @@ async def test_streaming_success_callbacks_survive_cost_calculation_failure():
|
|||
releasing.async_log_success_event = AsyncMock()
|
||||
|
||||
patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing])
|
||||
with patcher, patch.object(
|
||||
logging_obj, "_response_cost_calculator", side_effect=ValueError("bad usage block")
|
||||
):
|
||||
with patcher, patch.object(logging_obj, "_response_cost_calculator", side_effect=ValueError("bad usage block")):
|
||||
await logging_obj.async_success_handler(result=_assembled_stream_result())
|
||||
|
||||
assert logging_obj.model_call_details["response_cost"] is None
|
||||
|
|
@ -5800,8 +5798,9 @@ async def test_streaming_success_callbacks_survive_standard_logging_payload_fail
|
|||
releasing.async_log_success_event = AsyncMock()
|
||||
|
||||
patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing])
|
||||
with patcher, patch.object(
|
||||
logging_obj, "_build_standard_logging_payload", side_effect=ValueError("incomplete stream")
|
||||
with (
|
||||
patcher,
|
||||
patch.object(logging_obj, "_build_standard_logging_payload", side_effect=ValueError("incomplete stream")),
|
||||
):
|
||||
await logging_obj.async_success_handler(result=_assembled_stream_result())
|
||||
|
||||
|
|
@ -6074,6 +6073,8 @@ def test_prompt_hooks_skip_prompt_managers_when_no_prompt_id(logging_obj, tmp_pa
|
|||
)
|
||||
for hook in [cb for cb in litellm.callbacks if isinstance(cb, VectorStorePreCallHook)]:
|
||||
litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, hook)
|
||||
|
||||
|
||||
def test_newrelic_dispatch_prefers_otel_v2_when_flag_on(monkeypatch):
|
||||
"""With LITELLM_OTEL_V2 on and operator credentials present, the "newrelic"
|
||||
callback builds the OTel v2 logger (per-team credential routing); with the
|
||||
|
|
@ -6229,7 +6230,9 @@ def test_get_error_information_skips_traceback_for_budget_rejection_with_provide
|
|||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
|
||||
assert litellm.log_client_error_tracebacks is False
|
||||
over_budget = _raise_and_catch(litellm.BudgetExceededError(current_cost=0.01, max_budget=0.0, llm_provider="anthropic"))
|
||||
over_budget = _raise_and_catch(
|
||||
litellm.BudgetExceededError(current_cost=0.01, max_budget=0.0, llm_provider="anthropic")
|
||||
)
|
||||
result = StandardLoggingPayloadSetup.get_error_information(over_budget)
|
||||
assert result["error_code"] == "429"
|
||||
assert result["llm_provider"] == "anthropic"
|
||||
|
|
@ -6746,9 +6749,7 @@ def test_passthrough_embeddings_result_swapped_for_callbacks():
|
|||
],
|
||||
"model": "EmbeddingsGigaR",
|
||||
},
|
||||
request=httpx.Request(
|
||||
"POST", "https://gigachat.devices.sberbank.ru/api/v1/embeddings"
|
||||
),
|
||||
request=httpx.Request("POST", "https://gigachat.devices.sberbank.ru/api/v1/embeddings"),
|
||||
)
|
||||
|
||||
_, _, swapped_result = logging_obj._success_handler_helper_fn(
|
||||
|
|
@ -6767,12 +6768,14 @@ def test_get_status_fields_ranks_guardrail_flagged_between_success_and_intervene
|
|||
request-level guardrail_status but never mask an intervention."""
|
||||
flagged = {"guardrail_status": "guardrail_flagged"}
|
||||
|
||||
assert _get_status_fields(
|
||||
"success", [{"guardrail_status": "success"}, flagged], None
|
||||
)["guardrail_status"] == "guardrail_flagged"
|
||||
assert _get_status_fields(
|
||||
"success", [flagged, {"guardrail_status": "guardrail_intervened"}], None
|
||||
)["guardrail_status"] == "guardrail_intervened"
|
||||
assert (
|
||||
_get_status_fields("success", [{"guardrail_status": "success"}, flagged], None)["guardrail_status"]
|
||||
== "guardrail_flagged"
|
||||
)
|
||||
assert (
|
||||
_get_status_fields("success", [flagged, {"guardrail_status": "guardrail_intervened"}], None)["guardrail_status"]
|
||||
== "guardrail_intervened"
|
||||
)
|
||||
|
||||
|
||||
def test_get_error_information_redacts_provider_key_from_upstream_url():
|
||||
|
|
|
|||
0
tests/test_litellm/llms/mistral/batches/__init__.py
Normal file
0
tests/test_litellm/llms/mistral/batches/__init__.py
Normal file
|
|
@ -0,0 +1,258 @@
|
|||
"""
|
||||
Regression tests for ``MistralBatchesConfig``, the BaseBatchesConfig implementation
|
||||
behind ``custom_llm_provider="mistral"`` on /v1/batches.
|
||||
|
||||
Locks the request shape Mistral's ``POST /v1/batch/jobs`` accepts (input_files list,
|
||||
model set on the job, endpoint passed through untouched so ``/v1/ocr`` batches work),
|
||||
the Mistral -> OpenAI status mapping, request-count and file-id mapping, and auth.
|
||||
Everything runs for real against canned httpx responses; only the API key env var is
|
||||
set.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from litellm.llms.mistral.batches.transformation import MistralBatchesConfig
|
||||
from litellm.llms.mistral.common_utils import MistralError
|
||||
from litellm.types.llms.openai import CreateBatchRequest
|
||||
from litellm.types.utils import LiteLLMBatch, LlmProviders
|
||||
|
||||
STATUS_MAP = {
|
||||
"QUEUED": "validating",
|
||||
"RUNNING": "in_progress",
|
||||
"SUCCESS": "completed",
|
||||
"FAILED": "failed",
|
||||
"TIMEOUT_EXCEEDED": "expired",
|
||||
"CANCELLATION_REQUESTED": "cancelling",
|
||||
"CANCELLED": "cancelled",
|
||||
}
|
||||
|
||||
|
||||
def _job(**overrides):
|
||||
base = {
|
||||
"id": "8ff5e0d1-6bc2-4c3a-9f7d-0d1c2e3f4a5b",
|
||||
"object": "batch",
|
||||
"input_files": ["c1a2b3d4-0000-4000-8000-000000000001"],
|
||||
"endpoint": "/v1/ocr",
|
||||
"model": "mistral-ocr-latest",
|
||||
"status": "SUCCESS",
|
||||
"created_at": 1_757_400_000,
|
||||
"started_at": 1_757_400_010,
|
||||
"completed_at": 1_757_400_500,
|
||||
"total_requests": 3,
|
||||
"completed_requests": 3,
|
||||
"succeeded_requests": 2,
|
||||
"failed_requests": 1,
|
||||
"output_file": "out-0000-4000-8000-000000000002",
|
||||
"error_file": "err-0000-4000-8000-000000000003",
|
||||
"errors": [],
|
||||
"metadata": {"job_type": "testing"},
|
||||
}
|
||||
return {**base, **overrides}
|
||||
|
||||
|
||||
def _response(payload: dict, status_code: int = 200) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
status_code=status_code,
|
||||
content=json.dumps(payload).encode(),
|
||||
request=httpx.Request("GET", "https://api.mistral.ai/v1/batch/jobs/x"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def config() -> MistralBatchesConfig:
|
||||
return MistralBatchesConfig()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def api_key(monkeypatch) -> str:
|
||||
monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test")
|
||||
return "sk-mistral-test"
|
||||
|
||||
|
||||
def test_custom_llm_provider(config):
|
||||
assert config.custom_llm_provider == LlmProviders.MISTRAL
|
||||
|
||||
|
||||
def test_create_request_maps_openai_fields_onto_mistral_job(config):
|
||||
data = CreateBatchRequest(
|
||||
completion_window="24h",
|
||||
endpoint="/v1/ocr",
|
||||
input_file_id="file-123",
|
||||
metadata={"team": "docs"},
|
||||
)
|
||||
body = config.transform_create_batch_request(
|
||||
model="mistral-ocr-latest", create_batch_data=data, optional_params={}, litellm_params={}
|
||||
)
|
||||
assert body == {
|
||||
"input_files": ("file-123",),
|
||||
"endpoint": "/v1/ocr",
|
||||
"model": "mistral-ocr-latest",
|
||||
"metadata": {"team": "docs"},
|
||||
}
|
||||
|
||||
|
||||
def test_create_request_omits_empty_metadata(config):
|
||||
data = CreateBatchRequest(
|
||||
completion_window="24h",
|
||||
endpoint="/v1/chat/completions",
|
||||
input_file_id="file-123",
|
||||
metadata=None,
|
||||
)
|
||||
body = config.transform_create_batch_request(
|
||||
model="mistral-small-latest", create_batch_data=data, optional_params={}, litellm_params={}
|
||||
)
|
||||
assert "metadata" not in body
|
||||
|
||||
|
||||
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(
|
||||
"api_base,expected",
|
||||
[
|
||||
(None, "https://api.mistral.ai/v1/batch/jobs"),
|
||||
("https://api.mistral.ai/v1", "https://api.mistral.ai/v1/batch/jobs"),
|
||||
("https://proxy.example.com/", "https://proxy.example.com/v1/batch/jobs"),
|
||||
],
|
||||
)
|
||||
def test_create_url(config, api_base, expected):
|
||||
url = config.get_complete_batch_url(
|
||||
api_base=api_base, api_key="k", model="m", optional_params={}, litellm_params={}, data={}
|
||||
)
|
||||
assert url == expected
|
||||
|
||||
|
||||
def test_validate_environment_uses_bearer_auth(config, api_key):
|
||||
headers = config.validate_environment(
|
||||
headers={"x-extra": "1"}, model="m", messages=[], optional_params={}, litellm_params={}
|
||||
)
|
||||
assert headers == {"x-extra": "1", "Authorization": f"Bearer {api_key}"}
|
||||
|
||||
|
||||
def test_validate_environment_explicit_key_wins(config, api_key):
|
||||
headers = config.validate_environment(
|
||||
headers={}, model="m", messages=[], optional_params={}, litellm_params={}, api_key="sk-explicit"
|
||||
)
|
||||
assert headers["Authorization"] == "Bearer sk-explicit"
|
||||
|
||||
|
||||
def test_validate_environment_without_key_raises(config, monkeypatch):
|
||||
monkeypatch.delenv("MISTRAL_API_KEY", raising=False)
|
||||
with pytest.raises(ValueError, match="Missing Mistral API Key"):
|
||||
config.validate_environment(headers={}, model="m", messages=[], optional_params={}, litellm_params={})
|
||||
|
||||
|
||||
def test_create_response_maps_job_onto_openai_batch(config):
|
||||
batch = config.transform_create_batch_response(
|
||||
model="mistral-ocr-latest",
|
||||
raw_response=_response(_job(status="QUEUED", started_at=None, completed_at=None)),
|
||||
logging_obj=None,
|
||||
litellm_params={},
|
||||
)
|
||||
assert isinstance(batch, LiteLLMBatch)
|
||||
assert batch.id == "8ff5e0d1-6bc2-4c3a-9f7d-0d1c2e3f4a5b"
|
||||
assert batch.endpoint == "/v1/ocr"
|
||||
assert batch.input_file_id == "c1a2b3d4-0000-4000-8000-000000000001"
|
||||
assert batch.status == "validating"
|
||||
assert batch.created_at == 1_757_400_000
|
||||
assert batch.in_progress_at is None
|
||||
assert batch.completed_at is None
|
||||
assert batch.metadata == {"job_type": "testing"}
|
||||
|
||||
|
||||
def test_retrieve_request_is_presigned_get_with_auth(config, api_key):
|
||||
req = config.transform_retrieve_batch_request(
|
||||
batch_id="job/with slash", optional_params={}, litellm_params={"api_base": "https://api.mistral.ai"}
|
||||
)
|
||||
assert req["method"] == "GET"
|
||||
assert req["url"] == "https://api.mistral.ai/v1/batch/jobs/job%2Fwith%20slash"
|
||||
assert req["headers"] == {"Authorization": f"Bearer {api_key}"}
|
||||
|
||||
|
||||
def test_retrieve_request_prefers_litellm_params_api_key(config, api_key):
|
||||
req = config.transform_retrieve_batch_request(
|
||||
batch_id="job-1", optional_params={}, litellm_params={"api_key": "sk-from-deployment"}
|
||||
)
|
||||
assert req["headers"]["Authorization"] == "Bearer sk-from-deployment"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mistral_status,openai_status", sorted(STATUS_MAP.items()))
|
||||
def test_retrieve_response_status_mapping(config, mistral_status, openai_status):
|
||||
batch = config.transform_retrieve_batch_response(
|
||||
model=None, raw_response=_response(_job(status=mistral_status)), logging_obj=None, litellm_params={}
|
||||
)
|
||||
assert batch.status == openai_status
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mistral_status,populated_field",
|
||||
[
|
||||
("SUCCESS", "completed_at"),
|
||||
("FAILED", "failed_at"),
|
||||
("TIMEOUT_EXCEEDED", "expired_at"),
|
||||
("CANCELLED", "cancelled_at"),
|
||||
],
|
||||
)
|
||||
def test_retrieve_response_terminal_timestamp_lands_on_matching_field(config, mistral_status, populated_field):
|
||||
batch = config.transform_retrieve_batch_response(
|
||||
model=None, raw_response=_response(_job(status=mistral_status)), logging_obj=None, litellm_params={}
|
||||
)
|
||||
terminal_fields = {"completed_at", "failed_at", "expired_at", "cancelled_at"}
|
||||
assert getattr(batch, populated_field) == 1_757_400_500
|
||||
for other in terminal_fields - {populated_field}:
|
||||
assert getattr(batch, other) is None
|
||||
assert batch.in_progress_at == 1_757_400_010
|
||||
|
||||
|
||||
def test_retrieve_response_maps_counts_and_files(config):
|
||||
batch = config.transform_retrieve_batch_response(
|
||||
model=None, raw_response=_response(_job()), logging_obj=None, litellm_params={}
|
||||
)
|
||||
assert batch.request_counts.total == 3
|
||||
assert batch.request_counts.completed == 2
|
||||
assert batch.request_counts.failed == 1
|
||||
assert batch.output_file_id == "out-0000-4000-8000-000000000002"
|
||||
assert batch.error_file_id == "err-0000-4000-8000-000000000003"
|
||||
assert batch.errors is None
|
||||
|
||||
|
||||
def test_retrieve_response_surfaces_job_errors(config):
|
||||
batch = config.transform_retrieve_batch_response(
|
||||
model=None,
|
||||
raw_response=_response(
|
||||
_job(status="FAILED", errors=[{"message": "invalid document", "count": 2}, {"message": "timeout"}])
|
||||
),
|
||||
logging_obj=None,
|
||||
litellm_params={},
|
||||
)
|
||||
assert [e.message for e in batch.errors.data] == ["invalid document (x2)", "timeout"]
|
||||
|
||||
|
||||
def test_retrieve_response_without_files_or_input(config):
|
||||
batch = config.transform_retrieve_batch_response(
|
||||
model=None,
|
||||
raw_response=_response(_job(input_files=[], output_file=None, error_file=None, metadata=None)),
|
||||
logging_obj=None,
|
||||
litellm_params={},
|
||||
)
|
||||
assert batch.input_file_id == ""
|
||||
assert batch.output_file_id is None
|
||||
assert batch.error_file_id is None
|
||||
assert batch.metadata is None
|
||||
|
||||
|
||||
def test_get_error_class(config):
|
||||
err = config.get_error_class("nope", 401, {"x-request-id": "r1"})
|
||||
assert isinstance(err, MistralError)
|
||||
assert err.status_code == 401
|
||||
assert err.message == "nope"
|
||||
0
tests/test_litellm/llms/mistral/files/__init__.py
Normal file
0
tests/test_litellm/llms/mistral/files/__init__.py
Normal file
|
|
@ -0,0 +1,208 @@
|
|||
"""
|
||||
Regression tests for ``MistralFilesConfig``, the BaseFilesConfig implementation behind
|
||||
``custom_llm_provider="mistral"`` on /v1/files.
|
||||
|
||||
Locks the URL routing for each file operation, the multipart upload shape Mistral's
|
||||
``POST /v1/files`` accepts (purpose restricted to fine-tune/batch/ocr), and the
|
||||
Mistral -> OpenAI file object mapping. Runs against canned httpx responses.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from openai.types.file_deleted import FileDeleted
|
||||
|
||||
from litellm.llms.mistral.files.transformation import MistralFilesConfig
|
||||
from litellm.types.llms.openai import CreateFileRequest, FileContentRequest, OpenAIFileObject
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
FILE_ID = "497f6eca-6276-4993-bfeb-53cbbbba6f09"
|
||||
|
||||
|
||||
def _file(**overrides):
|
||||
base = {
|
||||
"id": FILE_ID,
|
||||
"object": "file",
|
||||
"bytes": 13000,
|
||||
"created_at": 1_716_963_433,
|
||||
"filename": "batch_input.jsonl",
|
||||
"purpose": "batch",
|
||||
"sample_type": "batch_request",
|
||||
"num_lines": 3,
|
||||
"source": "upload",
|
||||
}
|
||||
return {**base, **overrides}
|
||||
|
||||
|
||||
def _response(payload) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
status_code=200,
|
||||
content=json.dumps(payload).encode(),
|
||||
request=httpx.Request("GET", "https://api.mistral.ai/v1/files"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def config() -> MistralFilesConfig:
|
||||
return MistralFilesConfig()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def api_key(monkeypatch) -> str:
|
||||
monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test")
|
||||
return "sk-mistral-test"
|
||||
|
||||
|
||||
def test_custom_llm_provider(config):
|
||||
assert config.custom_llm_provider == LlmProviders.MISTRAL
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"api_base,expected",
|
||||
[
|
||||
(None, "https://api.mistral.ai/v1/files"),
|
||||
("https://api.mistral.ai/v1/", "https://api.mistral.ai/v1/files"),
|
||||
("https://proxy.example.com", "https://proxy.example.com/v1/files"),
|
||||
],
|
||||
)
|
||||
def test_upload_url(config, api_base, expected):
|
||||
url = config.get_complete_url(api_base=api_base, api_key="k", model="", optional_params={}, litellm_params={})
|
||||
assert url == expected
|
||||
|
||||
|
||||
def test_validate_environment_uses_bearer_auth(config, api_key):
|
||||
headers = config.validate_environment(headers={}, model="", messages=[], optional_params={}, litellm_params={})
|
||||
assert headers == {"Authorization": f"Bearer {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"
|
||||
),
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
assert body == {
|
||||
"file": ("in.jsonl", b'{"custom_id":"0"}\n', "application/jsonl"),
|
||||
"purpose": (None, "batch"),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("purpose", ["batch", "fine-tune", "ocr"])
|
||||
def test_upload_request_passes_mistral_purposes_through(config, purpose):
|
||||
body = config.transform_create_file_request(
|
||||
model="",
|
||||
create_file_data=CreateFileRequest(file=("f.bin", b"x"), purpose=purpose),
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
assert body["purpose"] == (None, purpose)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("purpose", ["assistants", "user_data", "vision", "evals"])
|
||||
def test_upload_request_rejects_purposes_mistral_lacks(config, purpose):
|
||||
"""Regression: these used to be silently rewritten to ``batch``, so an upload that skipped the
|
||||
proxy's batch-only validation and guardrails still landed on Mistral as a batch input file."""
|
||||
with pytest.raises(ValueError, match=f"purpose={purpose!r}"):
|
||||
config.transform_create_file_request(
|
||||
model="",
|
||||
create_file_data=CreateFileRequest(file=("f.bin", b"x"), purpose=purpose),
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
|
||||
|
||||
def test_upload_request_requires_file(config):
|
||||
with pytest.raises(ValueError, match="File data is required"):
|
||||
config.transform_create_file_request(
|
||||
model="", create_file_data=CreateFileRequest(purpose="batch"), optional_params={}, litellm_params={}
|
||||
)
|
||||
|
||||
|
||||
def test_upload_response_maps_onto_openai_file_object(config):
|
||||
obj = config.transform_create_file_response(
|
||||
model=None, raw_response=_response(_file()), logging_obj=None, litellm_params={}
|
||||
)
|
||||
assert obj == OpenAIFileObject(
|
||||
id=FILE_ID,
|
||||
bytes=13000,
|
||||
created_at=1_716_963_433,
|
||||
filename="batch_input.jsonl",
|
||||
object="file",
|
||||
purpose="batch",
|
||||
status="uploaded",
|
||||
)
|
||||
|
||||
|
||||
def test_file_response_with_ocr_purpose_maps_onto_user_data(config):
|
||||
obj = config.transform_retrieve_file_response(
|
||||
raw_response=_response(_file(purpose="ocr", expires_at=1_800_000_000)), logging_obj=None, litellm_params={}
|
||||
)
|
||||
assert obj.purpose == "user_data"
|
||||
assert obj.expires_at == 1_800_000_000
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"method,suffix",
|
||||
[
|
||||
("transform_retrieve_file_request", ""),
|
||||
("transform_delete_file_request", ""),
|
||||
],
|
||||
)
|
||||
def test_single_file_urls_encode_id_and_honor_api_base(config, method, suffix):
|
||||
url, params = getattr(config, method)(
|
||||
file_id="id/with slash", optional_params={}, litellm_params={"api_base": "https://mistral.internal/v1"}
|
||||
)
|
||||
assert url == f"https://mistral.internal/v1/files/id%2Fwith%20slash{suffix}"
|
||||
assert params == {}
|
||||
|
||||
|
||||
def test_file_content_url(config):
|
||||
url, params = config.transform_file_content_request(
|
||||
file_content_request=FileContentRequest(file_id=FILE_ID), optional_params={}, litellm_params={}
|
||||
)
|
||||
assert url == f"https://api.mistral.ai/v1/files/{FILE_ID}/content"
|
||||
assert params == {}
|
||||
|
||||
|
||||
def test_file_content_response_is_binary_passthrough(config):
|
||||
raw = httpx.Response(
|
||||
200, content=b'{"custom_id":"0","response":{"status_code":200}}\n', request=httpx.Request("GET", "https://x")
|
||||
)
|
||||
out = config.transform_file_content_response(raw_response=raw, logging_obj=None, litellm_params={})
|
||||
assert out.content == b'{"custom_id":"0","response":{"status_code":200}}\n'
|
||||
|
||||
|
||||
def test_delete_response(config):
|
||||
out = config.transform_delete_file_response(
|
||||
raw_response=_response({"id": FILE_ID, "object": "file", "deleted": True}), logging_obj=None, litellm_params={}
|
||||
)
|
||||
assert out == FileDeleted(id=FILE_ID, deleted=True, object="file")
|
||||
|
||||
|
||||
def test_list_request_filters_by_mapped_purpose(config):
|
||||
url, params = config.transform_list_files_request(purpose="batch", optional_params={}, litellm_params={})
|
||||
assert url == "https://api.mistral.ai/v1/files"
|
||||
assert params == {"purpose": "batch"}
|
||||
_, no_params = config.transform_list_files_request(purpose=None, optional_params={}, litellm_params={})
|
||||
assert no_params == {}
|
||||
|
||||
|
||||
def test_list_request_rejects_purposes_mistral_lacks(config):
|
||||
with pytest.raises(ValueError, match="purpose='assistants'"):
|
||||
config.transform_list_files_request(purpose="assistants", optional_params={}, litellm_params={})
|
||||
|
||||
|
||||
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}
|
||||
),
|
||||
logging_obj=None,
|
||||
litellm_params={},
|
||||
)
|
||||
assert [f.id for f in out] == [FILE_ID, "second"]
|
||||
assert out[1].filename == "b.jsonl"
|
||||
|
|
@ -63,7 +63,6 @@ def test_ocr4_cost_scales_with_pages(model: str, pages_processed: int) -> None:
|
|||
assert cost == pytest.approx(OCR4_COST_PER_PAGE * pages_processed)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cost_map_path", [MAIN_COST_MAP, BACKUP_COST_MAP])
|
||||
def test_ocr3_pricing_entry(cost_map_path: Path) -> None:
|
||||
with open(cost_map_path) as f:
|
||||
|
|
@ -72,9 +71,11 @@ def test_ocr3_pricing_entry(cost_map_path: Path) -> None:
|
|||
assert info is not None, f"{OCR3_MODEL} missing from {cost_map_path.name}"
|
||||
assert info["litellm_provider"] == "mistral"
|
||||
assert info["mode"] == "ocr"
|
||||
assert info["supported_endpoints"] == ["/v1/ocr"]
|
||||
assert info["supported_endpoints"] == ["/v1/ocr", "/v1/batch"]
|
||||
assert info["ocr_cost_per_page"] == OCR3_COST_PER_PAGE
|
||||
assert info["annotation_cost_per_page"] == OCR3_ANNOTATION_COST_PER_PAGE
|
||||
assert info["ocr_cost_per_page_batches"] == OCR3_COST_PER_PAGE / 2
|
||||
assert info["annotation_cost_per_page_batches"] == OCR3_ANNOTATION_COST_PER_PAGE / 2
|
||||
|
||||
|
||||
def test_ocr3_model_info_price(local_model_cost_map) -> None:
|
||||
|
|
|
|||
|
|
@ -177,6 +177,10 @@ def harness():
|
|||
logging.get_proxy_hook = MagicMock(return_value=None)
|
||||
|
||||
router = MagicMock(spec=Router)
|
||||
router.model_group_alias = {}
|
||||
router.get_model_access_groups = MagicMock(return_value={})
|
||||
router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id)
|
||||
router.model_list = []
|
||||
router.acreate_batch = AsyncMock(return_value=make_batch())
|
||||
router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup)
|
||||
|
||||
|
|
@ -1161,6 +1165,10 @@ def retrieve_harness():
|
|||
logging.get_proxy_hook = MagicMock(return_value=None)
|
||||
|
||||
router = MagicMock(spec=Router)
|
||||
router.model_group_alias = {}
|
||||
router.get_model_access_groups = MagicMock(return_value={})
|
||||
router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id)
|
||||
router.model_list = []
|
||||
router.aretrieve_batch = AsyncMock(return_value=make_batch())
|
||||
router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup)
|
||||
|
||||
|
|
@ -1616,6 +1624,10 @@ def list_harness():
|
|||
logging.get_proxy_hook = MagicMock(return_value=None)
|
||||
|
||||
router = MagicMock(spec=Router)
|
||||
router.model_group_alias = {}
|
||||
router.get_model_access_groups = MagicMock(return_value={})
|
||||
router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id)
|
||||
router.model_list = []
|
||||
router.alist_batches = AsyncMock(return_value=FakeListPage([]))
|
||||
router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup)
|
||||
|
||||
|
|
@ -2012,6 +2024,10 @@ def cancel_harness():
|
|||
logging.get_proxy_hook = MagicMock(return_value=None)
|
||||
|
||||
router = MagicMock(spec=Router)
|
||||
router.model_group_alias = {}
|
||||
router.get_model_access_groups = MagicMock(return_value={})
|
||||
router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id)
|
||||
router.model_list = []
|
||||
router.acancel_batch = AsyncMock(return_value=make_batch())
|
||||
router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup)
|
||||
|
||||
|
|
@ -2733,8 +2749,6 @@ async def test_cancel__unified_batch_id_allowed_when_managed_files_required(canc
|
|||
assert cancel_harness.router_acancel.call_count == 1
|
||||
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve__managed_batch_defers_cost_to_the_poller_when_it_is_running(retrieve_harness):
|
||||
with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=True)):
|
||||
|
|
@ -2762,3 +2776,101 @@ async def test_retrieve__raw_batch_id_is_untouched_by_the_poller_handoff(retriev
|
|||
|
||||
metadata = retrieve_harness.litellm_aretrieve.await_args.kwargs.get("litellm_metadata") or {}
|
||||
assert metadata.get("batch_ignore_default_logging") is None
|
||||
|
||||
|
||||
def _key_restricted_to(*models: str) -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(api_key="sk-restricted", team_id="team-a", team_models=list(models), models=list(models))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create__header_model_rejects_key_without_model_grant(harness):
|
||||
"""A key not granted the model named in x-litellm-model must not receive that deployment's credentials."""
|
||||
set_body(harness, {"input_file_id": "file-plain", "endpoint": "/v1/chat/completions", "completion_window": "24h"})
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await call_create(harness, user=_key_restricted_to("azure/gpt-4o"), headers={"x-litellm-model": "vertex-model"})
|
||||
|
||||
assert exc_info.value.code == "403"
|
||||
harness.creds_resolver.assert_not_called()
|
||||
harness.litellm_acreate.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create__header_model_allows_key_with_model_grant(harness):
|
||||
set_body(harness, {"input_file_id": "file-plain", "endpoint": "/v1/chat/completions", "completion_window": "24h"})
|
||||
|
||||
await call_create(harness, user=_key_restricted_to("vertex-model"), headers={"x-litellm-model": "vertex-model"})
|
||||
|
||||
harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
|
||||
assert harness.acreate_kwargs()["custom_llm_provider"] == "vertex_ai"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve__model_encoded_id_rejects_key_without_model_grant(retrieve_harness):
|
||||
"""The model embedded in a batch id is caller-controlled, so it is checked against the key's grants too."""
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await call_retrieve(retrieve_harness, AZURE_BATCH_ID, user=_key_restricted_to("vertex-model"))
|
||||
|
||||
assert exc_info.value.code == "403"
|
||||
retrieve_harness.creds_resolver.assert_not_called()
|
||||
retrieve_harness.litellm_aretrieve.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel__model_encoded_id_rejects_key_without_model_grant(cancel_harness):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await call_cancel(cancel_harness, AZURE_BATCH_ID, user=_key_restricted_to("vertex-model"))
|
||||
|
||||
assert exc_info.value.code == "403"
|
||||
cancel_harness.creds_resolver.assert_not_called()
|
||||
cancel_harness.litellm_acancel.assert_not_called()
|
||||
|
||||
|
||||
def _b64_unified_id(decoded: str) -> str:
|
||||
return base64.urlsafe_b64encode(decoded.encode()).decode().rstrip("=")
|
||||
|
||||
|
||||
UNIFIED_FILE_ID_FOR_GPT4O_MINI = _b64_unified_id(
|
||||
"litellm_proxy:application/octet-stream;unified_id,c4843482-b176-4901-8292-7523fd0f2c6e;"
|
||||
"target_model_names,gpt-4o-mini;llm_output_file_id,file-provider;llm_output_file_model_id,dep-1"
|
||||
)
|
||||
UNIFIED_BATCH_ID_FOR_GPT4O_MINI = _b64_unified_id(UNIFIED_BATCH_ID)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create__unified_file_id_rejects_key_without_model_grant(harness):
|
||||
"""The model carried inside a unified file id is caller-controlled too, so it is checked against the key's grants."""
|
||||
set_body(
|
||||
harness,
|
||||
{
|
||||
"input_file_id": UNIFIED_FILE_ID_FOR_GPT4O_MINI,
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
},
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await call_create(harness, user=_key_restricted_to("vertex-model"))
|
||||
|
||||
assert exc_info.value.code == "403"
|
||||
harness.router_acreate.assert_not_called()
|
||||
harness.litellm_acreate.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve__unified_batch_id_rejects_key_without_model_grant(retrieve_harness):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await call_retrieve(retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to("vertex-model"))
|
||||
|
||||
assert exc_info.value.code == "403"
|
||||
retrieve_harness.router_aretrieve.assert_not_called()
|
||||
retrieve_harness.creds_resolver.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel__unified_batch_id_rejects_key_without_model_grant(cancel_harness):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await call_cancel(cancel_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to("vertex-model"))
|
||||
|
||||
assert exc_info.value.code == "403"
|
||||
cancel_harness.router_acancel.assert_not_called()
|
||||
|
|
|
|||
|
|
@ -1877,7 +1877,7 @@ def test_get_file_content_streams_openai_direct_path(
|
|||
monkeypatch.setattr(litellm, "afile_content", _mock_afile_content)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing",
|
||||
lambda **kwargs: (False, None, None, None),
|
||||
AsyncMock(return_value=(False, None, None, None)),
|
||||
)
|
||||
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
|
|
@ -1942,15 +1942,17 @@ def test_get_file_content_routed_provider_skips_streaming_when_resolved_provider
|
|||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing",
|
||||
lambda **kwargs: (
|
||||
True,
|
||||
"azure-gpt-3-5-turbo",
|
||||
"file-original-123",
|
||||
{
|
||||
"custom_llm_provider": "azure",
|
||||
"api_key": "azure-key",
|
||||
"api_base": "https://azure.example.com",
|
||||
},
|
||||
AsyncMock(
|
||||
return_value=(
|
||||
True,
|
||||
"azure-gpt-3-5-turbo",
|
||||
"file-original-123",
|
||||
{
|
||||
"custom_llm_provider": "azure",
|
||||
"api_key": "azure-key",
|
||||
"api_base": "https://azure.example.com",
|
||||
},
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -2015,7 +2017,7 @@ def test_get_file_content_non_openai_provider_skips_streaming_handler(
|
|||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing",
|
||||
lambda **kwargs: (False, None, None, None),
|
||||
AsyncMock(return_value=(False, None, None, None)),
|
||||
)
|
||||
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
|
|
@ -2463,14 +2465,16 @@ def test_list_files_model_routing_does_not_forward_custom_llm_provider_twice(
|
|||
monkeypatch.setattr(litellm, "afile_list", _mock_afile_list)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing",
|
||||
lambda **kwargs: (
|
||||
True,
|
||||
"azure-gpt-4o",
|
||||
None,
|
||||
{
|
||||
"custom_llm_provider": "azure",
|
||||
"api_key": "azure-key",
|
||||
},
|
||||
AsyncMock(
|
||||
return_value=(
|
||||
True,
|
||||
"azure-gpt-4o",
|
||||
None,
|
||||
{
|
||||
"custom_llm_provider": "azure",
|
||||
"api_key": "azure-key",
|
||||
},
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -4819,3 +4823,204 @@ def test_get_file_content_keeps_the_status_of_a_rejection_raised_inside_the_rout
|
|||
error = response.json()["error"]
|
||||
assert error["message"].startswith("Storage backend error")
|
||||
assert (error["type"], error["param"], error["code"]) == ("invalid_request_error", "file_id", "400")
|
||||
|
||||
|
||||
def test_get_file_model_routed_id_forwards_deployment_provider(mocker: MockerFixture, monkeypatch):
|
||||
"""
|
||||
Regression: a file id encoded with a non-OpenAI deployment (here Mistral) must be
|
||||
retrieved from that deployment's provider. Before the fix the retrieve path only
|
||||
forwarded the credentials and let ``custom_llm_provider`` default to openai, so a
|
||||
Mistral file id was sent to api.openai.com with the Mistral key and 401'd.
|
||||
"""
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model
|
||||
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "mistral-ocr",
|
||||
"litellm_params": {"model": "mistral/mistral-ocr-latest", "api_key": "mistral-key"},
|
||||
"model_info": {"id": "mistral-ocr-id"},
|
||||
}
|
||||
]
|
||||
)
|
||||
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
|
||||
proxy_logging_obj.update_request_status = mocker.AsyncMock()
|
||||
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
|
||||
|
||||
captured_kwargs: dict = {}
|
||||
|
||||
async def _mock_afile_retrieve(**kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
return OpenAIFileObject(
|
||||
id="7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df",
|
||||
object="file",
|
||||
bytes=2,
|
||||
created_at=1234567890,
|
||||
filename="batch.jsonl",
|
||||
purpose="batch",
|
||||
status="uploaded",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(litellm, "afile_retrieve", _mock_afile_retrieve)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
api_key="test-key", user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user"
|
||||
)
|
||||
encoded_id = encode_file_id_with_model("7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", "mistral-ocr")
|
||||
|
||||
try:
|
||||
response = client.get(f"/v1/files/{encoded_id}", headers={"Authorization": "Bearer test-key"})
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert captured_kwargs["custom_llm_provider"] == "mistral"
|
||||
assert captured_kwargs["api_key"] == "mistral-key"
|
||||
assert captured_kwargs["file_id"] == "7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df"
|
||||
assert response.json()["id"] == encoded_id
|
||||
|
||||
|
||||
def _mistral_plus_anthropic_router() -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "mistral-ocr",
|
||||
"litellm_params": {"model": "mistral/mistral-ocr-latest", "api_key": "mistral-key"},
|
||||
"model_info": {"id": "mistral-ocr-id"},
|
||||
},
|
||||
{
|
||||
"model_name": "claude-opus-4-6",
|
||||
"litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "anthropic-key"},
|
||||
"model_info": {"id": "claude-id"},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _restricted_key(key_models: list[str]) -> UserAPIKeyAuth:
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
return UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="test-user",
|
||||
team_id="team-a",
|
||||
team_models=["claude-opus-4-6", "mistral-ocr"],
|
||||
models=key_models,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"http_method, path_suffix, litellm_fn",
|
||||
[
|
||||
("get", "", "afile_retrieve"),
|
||||
("get", "/content", "afile_content"),
|
||||
("delete", "", "afile_delete"),
|
||||
],
|
||||
)
|
||||
def test_model_routed_file_ops_reject_key_without_model_grant(
|
||||
mocker: MockerFixture, monkeypatch, http_method: str, path_suffix: str, litellm_fn: str
|
||||
):
|
||||
"""
|
||||
Regression: a key whose allowlist does not include the deployment named in a
|
||||
model-encoded file id must be refused before that deployment's server-side
|
||||
credentials are resolved. Previously any key could name any deployment via the
|
||||
id (or the x-litellm-model header) and act on that provider account's files.
|
||||
"""
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model
|
||||
|
||||
router = _mistral_plus_anthropic_router()
|
||||
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
|
||||
proxy_logging_obj.update_request_status = mocker.AsyncMock()
|
||||
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
|
||||
|
||||
upstream = mocker.AsyncMock(side_effect=AssertionError("provider must not be called"))
|
||||
monkeypatch.setattr(litellm, litellm_fn, upstream)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: _restricted_key(["claude-opus-4-6"])
|
||||
encoded_id = encode_file_id_with_model("7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", "mistral-ocr")
|
||||
|
||||
try:
|
||||
response = getattr(client, http_method)(
|
||||
f"/v1/files/{encoded_id}{path_suffix}", headers={"Authorization": "Bearer test-key"}
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 403, response.text
|
||||
assert "not allowed to access model" in response.text
|
||||
upstream.assert_not_called()
|
||||
|
||||
|
||||
def test_list_files_header_model_rejects_key_without_model_grant(mocker: MockerFixture, monkeypatch):
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
router = _mistral_plus_anthropic_router()
|
||||
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
|
||||
proxy_logging_obj.update_request_status = mocker.AsyncMock()
|
||||
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
|
||||
|
||||
upstream = mocker.AsyncMock(side_effect=AssertionError("provider must not be called"))
|
||||
monkeypatch.setattr(litellm, "afile_list", upstream)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: _restricted_key(["claude-opus-4-6"])
|
||||
|
||||
try:
|
||||
response = client.get(
|
||||
"/v1/files", headers={"Authorization": "Bearer test-key", "x-litellm-model": "mistral-ocr"}
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 403, response.text
|
||||
upstream.assert_not_called()
|
||||
|
||||
|
||||
def test_model_routed_file_retrieve_allows_key_with_model_grant(mocker: MockerFixture, monkeypatch):
|
||||
"""The grant check must not break the happy path: a key allowed the deployment still resolves its credentials."""
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model
|
||||
|
||||
router = _mistral_plus_anthropic_router()
|
||||
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
|
||||
proxy_logging_obj.update_request_status = mocker.AsyncMock()
|
||||
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
|
||||
|
||||
captured_kwargs: dict = {}
|
||||
|
||||
async def _mock_afile_retrieve(**kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
return OpenAIFileObject(
|
||||
id="7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df",
|
||||
object="file",
|
||||
bytes=2,
|
||||
created_at=1234567890,
|
||||
filename="batch.jsonl",
|
||||
purpose="batch",
|
||||
status="uploaded",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(litellm, "afile_retrieve", _mock_afile_retrieve)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: _restricted_key(["mistral-ocr"])
|
||||
encoded_id = encode_file_id_with_model("7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", "mistral-ocr")
|
||||
|
||||
try:
|
||||
response = client.get(f"/v1/files/{encoded_id}", headers={"Authorization": "Bearer test-key"})
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert captured_kwargs["api_key"] == "mistral-key"
|
||||
assert captured_kwargs["custom_llm_provider"] == "mistral"
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
decode_model_from_file_id,
|
||||
get_batch_id_from_unified_batch_id,
|
||||
|
|
@ -58,10 +59,7 @@ def _make_batch_response(
|
|||
|
||||
|
||||
def test_get_batch_id_from_unified_batch_id_handles_appended_fields():
|
||||
decoded_id = (
|
||||
"litellm_proxy;model_id:deployment-123;"
|
||||
"llm_batch_id:batch_openai_123;llm_output_file_id:file-output"
|
||||
)
|
||||
decoded_id = "litellm_proxy;model_id:deployment-123;llm_batch_id:batch_openai_123;llm_output_file_id:file-output"
|
||||
|
||||
assert get_batch_id_from_unified_batch_id(decoded_id) == "batch_openai_123"
|
||||
|
||||
|
|
@ -107,12 +105,10 @@ async def test_create_batch_with_x_litellm_model_encodes_batch_id():
|
|||
}
|
||||
),
|
||||
),
|
||||
patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls,
|
||||
patch(
|
||||
"litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing"
|
||||
) as mock_processor_cls,
|
||||
patch(
|
||||
"litellm.proxy.batches_endpoints.endpoints.get_credentials_for_model",
|
||||
return_value=mock_credentials,
|
||||
"litellm.proxy.batches_endpoints.endpoints.get_authorized_credentials_for_model",
|
||||
new=AsyncMock(return_value=mock_credentials),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.batches_endpoints.endpoints.prepare_data_with_credentials",
|
||||
|
|
@ -165,23 +161,15 @@ async def test_create_batch_with_x_litellm_model_encodes_batch_id():
|
|||
)
|
||||
|
||||
# The batch_id should be encoded with model info
|
||||
assert (
|
||||
response.id != raw_batch_id
|
||||
), f"Expected batch_id to be encoded, but got raw ID: {response.id}"
|
||||
assert response.id.startswith(
|
||||
"batch_"
|
||||
), f"Encoded batch_id should keep batch_ prefix, got: {response.id}"
|
||||
assert response.id != raw_batch_id, f"Expected batch_id to be encoded, but got raw ID: {response.id}"
|
||||
assert response.id.startswith("batch_"), f"Encoded batch_id should keep batch_ prefix, got: {response.id}"
|
||||
|
||||
# Should be decodable back to the original
|
||||
decoded_model = decode_model_from_file_id(response.id)
|
||||
assert (
|
||||
decoded_model == model_name
|
||||
), f"Expected model '{model_name}' from decoded batch_id, got: {decoded_model}"
|
||||
assert decoded_model == model_name, f"Expected model '{model_name}' from decoded batch_id, got: {decoded_model}"
|
||||
|
||||
original_id = get_original_file_id(response.id)
|
||||
assert (
|
||||
original_id == raw_batch_id
|
||||
), f"Expected original ID '{raw_batch_id}', got: {original_id}"
|
||||
assert original_id == raw_batch_id, f"Expected original ID '{raw_batch_id}', got: {original_id}"
|
||||
assert mock_create_batch.call_args.kwargs["metadata"] == {"customer_id": "cust-123"}
|
||||
|
||||
|
||||
|
|
@ -227,12 +215,10 @@ async def test_create_batch_with_x_litellm_model_encodes_output_and_error_file_i
|
|||
}
|
||||
),
|
||||
),
|
||||
patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls,
|
||||
patch(
|
||||
"litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing"
|
||||
) as mock_processor_cls,
|
||||
patch(
|
||||
"litellm.proxy.batches_endpoints.endpoints.get_credentials_for_model",
|
||||
return_value=mock_credentials,
|
||||
"litellm.proxy.batches_endpoints.endpoints.get_authorized_credentials_for_model",
|
||||
new=AsyncMock(return_value=mock_credentials),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.batches_endpoints.endpoints.prepare_data_with_credentials",
|
||||
|
|
@ -316,9 +302,7 @@ async def test_create_batch_without_x_litellm_model_returns_raw_ids(monkeypatch)
|
|||
}
|
||||
),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing"
|
||||
) as mock_processor_cls,
|
||||
patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls,
|
||||
patch(
|
||||
"litellm.acreate_batch",
|
||||
new=AsyncMock(return_value=mock_response),
|
||||
|
|
@ -383,9 +367,7 @@ class TestBatchIdRoundTripWithRetrieve:
|
|||
raw_batch_id = "batch_vllm_12345"
|
||||
|
||||
# What create_batch does:
|
||||
encoded_id = encode_file_id_with_model(
|
||||
file_id=raw_batch_id, model=model_name, id_type="batch"
|
||||
)
|
||||
encoded_id = encode_file_id_with_model(file_id=raw_batch_id, model=model_name, id_type="batch")
|
||||
|
||||
# What retrieve_batch does:
|
||||
decoded_model = decode_model_from_file_id(encoded_id)
|
||||
|
|
@ -410,9 +392,7 @@ class TestBatchIdRoundTripWithRetrieve:
|
|||
]
|
||||
|
||||
for raw_id, model in test_cases:
|
||||
encoded = encode_file_id_with_model(
|
||||
file_id=raw_id, model=model, id_type="batch"
|
||||
)
|
||||
encoded = encode_file_id_with_model(file_id=raw_id, model=model, id_type="batch")
|
||||
assert encoded.startswith("batch_")
|
||||
assert decode_model_from_file_id(encoded) == model
|
||||
assert get_original_file_id(encoded) == raw_id
|
||||
|
|
@ -433,16 +413,10 @@ async def test_cancel_batch_with_unified_id_routes_with_decoded_model_and_batch_
|
|||
mock_request.url.path = f"/v1/batches/{unified_batch_id}/cancel"
|
||||
mock_fastapi_response = MagicMock()
|
||||
mock_fastapi_response.headers = {}
|
||||
mock_user_api_key_dict = MagicMock()
|
||||
mock_user_api_key_dict.parent_otel_span = None
|
||||
mock_user_api_key_dict.user_id = "test_user"
|
||||
mock_user_api_key_dict.allowed_model_region = None
|
||||
mock_user_api_key_dict.team_metadata = {}
|
||||
mock_user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="test_user", team_metadata={})
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing"
|
||||
) as mock_processor_cls,
|
||||
patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls,
|
||||
patch(
|
||||
"litellm.proxy.batches_endpoints.endpoints.update_batch_in_database",
|
||||
new=AsyncMock(),
|
||||
|
|
|
|||
|
|
@ -1017,7 +1017,9 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
|
|||
"input_cost_per_video_per_second_above_128k_tokens": {"type": "number"},
|
||||
"input_dbu_cost_per_token": {"type": "number"},
|
||||
"annotation_cost_per_page": {"type": "number"},
|
||||
"annotation_cost_per_page_batches": {"type": "number"},
|
||||
"ocr_cost_per_page": {"type": "number"},
|
||||
"ocr_cost_per_page_batches": {"type": "number"},
|
||||
"ocr_cost_per_credit": {"type": "number"},
|
||||
"code_interpreter_cost_per_session": {"type": "number"},
|
||||
"inference_geo": {"type": "string"},
|
||||
|
|
|
|||
8
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
8
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -29501,6 +29501,8 @@ export interface components {
|
|||
allow_client_keepalive_override: boolean | null;
|
||||
/** Annotation Cost Per Page */
|
||||
annotation_cost_per_page?: number | null;
|
||||
/** Annotation Cost Per Page Batches */
|
||||
annotation_cost_per_page_batches?: number | null;
|
||||
/** Api Base */
|
||||
api_base?: string | null;
|
||||
/** Api Key */
|
||||
|
|
@ -29714,6 +29716,8 @@ export interface components {
|
|||
ocr_cost_per_credit?: number | null;
|
||||
/** Ocr Cost Per Page */
|
||||
ocr_cost_per_page?: number | null;
|
||||
/** Ocr Cost Per Page Batches */
|
||||
ocr_cost_per_page_batches?: number | null;
|
||||
/** Organization */
|
||||
organization?: string | null;
|
||||
/** Otpm */
|
||||
|
|
@ -39703,6 +39707,8 @@ export interface components {
|
|||
allow_client_keepalive_override: boolean | null;
|
||||
/** Annotation Cost Per Page */
|
||||
annotation_cost_per_page?: number | null;
|
||||
/** Annotation Cost Per Page Batches */
|
||||
annotation_cost_per_page_batches?: number | null;
|
||||
/** Api Base */
|
||||
api_base?: string | null;
|
||||
/** Api Key */
|
||||
|
|
@ -39916,6 +39922,8 @@ export interface components {
|
|||
ocr_cost_per_credit?: number | null;
|
||||
/** Ocr Cost Per Page */
|
||||
ocr_cost_per_page?: number | null;
|
||||
/** Ocr Cost Per Page Batches */
|
||||
ocr_cost_per_page_batches?: number | null;
|
||||
/** Organization */
|
||||
organization?: string | null;
|
||||
/** Otpm */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue