mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
Merge pull request #40767 from BerriAI/litellm_ocr_custom_pricing
fix(cost): honour deployment custom pricing for OCR calls
This commit is contained in:
Mateo Wang
GitHub
commit
566f026c1c
6 changed files with 408 additions and 25 deletions
|
|
@ -58,6 +58,13 @@ impl PythonLogger {
|
|||
params.set_item(name, value)?;
|
||||
}
|
||||
}
|
||||
for name in custom_pricing_fields(py)? {
|
||||
if let Some(value) = kwargs.bind(py).get_item(&name)?
|
||||
&& !value.is_none()
|
||||
{
|
||||
params.set_item(name, value)?;
|
||||
}
|
||||
}
|
||||
update.set_item("litellm_params", params)?;
|
||||
update.set_item("custom_llm_provider", &pre_call.custom_llm_provider)?;
|
||||
self.object(py)
|
||||
|
|
@ -120,6 +127,17 @@ impl PythonLogger {
|
|||
}
|
||||
}
|
||||
|
||||
fn custom_pricing_fields(py: Python<'_>) -> PyResult<Vec<String>> {
|
||||
py.import("litellm.types.utils")?
|
||||
.getattr("CustomPricingLiteLLMParams")?
|
||||
.getattr("model_fields")?
|
||||
.cast_into::<PyDict>()?
|
||||
.keys()
|
||||
.iter()
|
||||
.map(|name| name.extract::<String>())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn redact(
|
||||
py: Python<'_>,
|
||||
params: &Bound<'_, PyDict>,
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, cast
|
|||
|
||||
from httpx import Response
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
import litellm._logging
|
||||
|
|
@ -310,6 +311,15 @@ def _transcription_usage_has_token_details(
|
|||
return (prompt_tokens_val > 0) or (completion_tokens_val > 0)
|
||||
|
||||
|
||||
OCRPricingField = Literal["ocr_cost_per_page", "ocr_cost_per_credit", "annotation_cost_per_page"]
|
||||
|
||||
|
||||
class OCRPricing(TypedDict, total=False):
|
||||
ocr_cost_per_page: ReadOnly[float | None]
|
||||
ocr_cost_per_credit: ReadOnly[float | None]
|
||||
annotation_cost_per_page: ReadOnly[float | None]
|
||||
|
||||
|
||||
def cost_per_token(
|
||||
model: str = "",
|
||||
prompt_tokens: int = 0,
|
||||
|
|
@ -344,6 +354,7 @@ def cost_per_token(
|
|||
response: Any | None = None,
|
||||
### REQUEST MODEL ###
|
||||
request_model: str | None = None, # original request model for router detection
|
||||
custom_model_info: OCRPricing | None = None,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
|
||||
|
|
@ -558,6 +569,7 @@ def cost_per_token(
|
|||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
response=response,
|
||||
model_info=custom_model_info,
|
||||
)
|
||||
elif (
|
||||
call_type == "aretrieve_batch"
|
||||
|
|
@ -1432,20 +1444,9 @@ def completion_cost(
|
|||
)
|
||||
elif call_type in _VIDEO_CALL_TYPES:
|
||||
### VIDEO GENERATION COST CALCULATION ###
|
||||
# Extract custom model_info for deployment-specific pricing
|
||||
_video_model_info: ModelInfo | None = None
|
||||
if custom_pricing and litellm_logging_obj is not None:
|
||||
_litellm_params = getattr(litellm_logging_obj, "litellm_params", None)
|
||||
if _litellm_params is not None:
|
||||
_video_model_info = next(
|
||||
(
|
||||
model_info
|
||||
for _metadata_key in ("metadata", "litellm_metadata")
|
||||
if (model_info := (_litellm_params.get(_metadata_key) or {}).get("model_info"))
|
||||
is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
_video_model_info: ModelInfo | None = _deployment_model_info(
|
||||
litellm_logging_obj, custom_pricing, router_model_id
|
||||
)
|
||||
|
||||
usage_obj = getattr(completion_response, "usage", None)
|
||||
duration_seconds: float | None = None
|
||||
|
|
@ -1665,6 +1666,7 @@ def completion_cost(
|
|||
data_residency=data_residency,
|
||||
vertex_location=vertex_location,
|
||||
response=completion_response,
|
||||
custom_model_info=_ocr_model_info(litellm_logging_obj, custom_pricing, router_model_id),
|
||||
)
|
||||
|
||||
# Get additional costs from provider (e.g., routing fees, infrastructure costs)
|
||||
|
|
@ -1898,16 +1900,82 @@ def response_cost_calculator(
|
|||
raise e
|
||||
|
||||
|
||||
def _deployment_model_info(
|
||||
litellm_logging_obj: LitellmLoggingObject | None,
|
||||
custom_pricing: bool | None,
|
||||
router_model_id: str | None,
|
||||
) -> ModelInfo | None:
|
||||
if not custom_pricing:
|
||||
return None
|
||||
registered_deployment_info: Final = (
|
||||
_cost_map_model_info(router_model_id, None)
|
||||
if router_model_id is not None and router_model_id in litellm.model_cost
|
||||
else None
|
||||
)
|
||||
if registered_deployment_info is not None:
|
||||
return registered_deployment_info
|
||||
if litellm_logging_obj is None:
|
||||
return None
|
||||
litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None)
|
||||
if litellm_params is None:
|
||||
return None
|
||||
return next(
|
||||
(
|
||||
model_info
|
||||
for metadata_key in ("metadata", "litellm_metadata")
|
||||
if (metadata := litellm_params.get(metadata_key)) and (model_info := metadata.get("model_info")) is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def _ocr_model_info(
|
||||
litellm_logging_obj: LitellmLoggingObject | None,
|
||||
custom_pricing: bool | None,
|
||||
router_model_id: str | None,
|
||||
) -> OCRPricing | None:
|
||||
deployment_info: Final = _deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id)
|
||||
litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None) if custom_pricing else None
|
||||
if litellm_params is None:
|
||||
return deployment_info
|
||||
return _layered_ocr_pricing(litellm_params, deployment_info)
|
||||
|
||||
|
||||
def _first_ocr_price(field: OCRPricingField, *sources: Mapping[str, object] | None) -> float | None:
|
||||
return next(
|
||||
(price for source in sources if source is not None and isinstance(price := source.get(field), int | float)),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def _layered_ocr_pricing(*sources: Mapping[str, object] | None) -> OCRPricing:
|
||||
return OCRPricing(
|
||||
ocr_cost_per_page=_first_ocr_price("ocr_cost_per_page", *sources),
|
||||
ocr_cost_per_credit=_first_ocr_price("ocr_cost_per_credit", *sources),
|
||||
annotation_cost_per_page=_first_ocr_price("annotation_cost_per_page", *sources),
|
||||
)
|
||||
|
||||
|
||||
def _cost_map_model_info(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 ocr_cost(
|
||||
model: str,
|
||||
custom_llm_provider: str | None,
|
||||
response: object | None = None,
|
||||
model_info: OCRPricing | None = None,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Args:
|
||||
model: str - model name
|
||||
custom_llm_provider: Optional[str] - custom LLM provider
|
||||
response: Optional[Any] - response object
|
||||
model_info: Optional[OCRPricing] - deployment-specific OCR pricing; each rate it sets
|
||||
overrides the model cost map's, the rest fall back to the map
|
||||
|
||||
Returns:
|
||||
Tuple[float, float]: cost of OCR processing
|
||||
|
|
@ -1925,20 +1993,15 @@ def ocr_cost(
|
|||
if response.usage_info is None:
|
||||
raise ValueError("OCR response usage_info is None")
|
||||
|
||||
try:
|
||||
model_info: ModelInfo | None = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
|
||||
except Exception:
|
||||
model_info = None
|
||||
|
||||
credits: Final = getattr(response.usage_info, "credits", None)
|
||||
cost_per_credit = None
|
||||
if model_info is not None:
|
||||
cost_per_credit = model_info.get("ocr_cost_per_credit")
|
||||
pricing: Final = _layered_ocr_pricing(model_info, _cost_map_model_info(model, custom_llm_provider))
|
||||
|
||||
cost_per_credit: Final = pricing.get("ocr_cost_per_credit")
|
||||
if credits is not None and cost_per_credit is not None:
|
||||
return cost_per_credit * credits, 0.0
|
||||
|
||||
ocr_cost_per_page: Final = model_info.get("ocr_cost_per_page") if model_info is not None else None
|
||||
annotation_cost_per_page: Final = model_info.get("annotation_cost_per_page") if model_info is not None else None
|
||||
ocr_cost_per_page: Final = pricing.get("ocr_cost_per_page")
|
||||
annotation_cost_per_page: Final = pricing.get("annotation_cost_per_page")
|
||||
annotation_rate: Final = annotation_cost_per_page if annotation_cost_per_page is not None else ocr_cost_per_page
|
||||
|
||||
pages_processed: Final = response.usage_info.pages_processed
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ from litellm.llms.base_llm.ocr.transformation import (
|
|||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from litellm.ocr.input import FileReader
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import CustomPricingLiteLLMParams
|
||||
from litellm.utils import ProviderConfigManager, client
|
||||
|
||||
base_llm_http_handler: Final = BaseLLMHTTPHandler()
|
||||
|
|
@ -149,6 +150,7 @@ def _prepare_ocr_request(
|
|||
litellm_params={
|
||||
"litellm_call_id": litellm_call_id,
|
||||
"api_base": resolved_api_base,
|
||||
**litellm_params.model_dump(include=frozenset(CustomPricingLiteLLMParams.model_fields), exclude_none=True),
|
||||
},
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -11,7 +11,8 @@ import pytest
|
|||
|
||||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging, use_custom_pricing_for_model
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo
|
||||
from litellm.llms.custom_httpx import llm_http_handler
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.ocr.legacy import _prepare_ocr_request
|
||||
|
|
@ -198,3 +199,61 @@ def test_generic_azure_connection_still_applies_to_foundry_ocr(monkeypatch: pyte
|
|||
|
||||
assert prepared.api_key == "generic-key"
|
||||
assert prepared.api_base == "https://generic.example.com"
|
||||
|
||||
|
||||
PRICING_OCR_MODEL: Final = "mistral/some-unmapped-ocr-model-for-testing"
|
||||
PRICING_DOCUMENT: Final = {"type": "document_url", "document_url": "https://example.com/doc.pdf"}
|
||||
|
||||
|
||||
def _pricing_logging_obj() -> Logging:
|
||||
return Logging(
|
||||
model=PRICING_OCR_MODEL,
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="ocr",
|
||||
start_time=None,
|
||||
litellm_call_id="test-ocr-request-pricing",
|
||||
function_id="1234",
|
||||
)
|
||||
|
||||
|
||||
def _prepare_with_pricing(kwargs: dict[str, object]) -> Logging:
|
||||
logging_obj: Final = _pricing_logging_obj()
|
||||
_prepare_ocr_request(
|
||||
model=PRICING_OCR_MODEL,
|
||||
document=dict(PRICING_DOCUMENT),
|
||||
api_key="test-key",
|
||||
api_base=None,
|
||||
timeout=None,
|
||||
custom_llm_provider=None,
|
||||
extra_headers=None,
|
||||
kwargs={"litellm_logging_obj": logging_obj, **kwargs},
|
||||
)
|
||||
return logging_obj
|
||||
|
||||
|
||||
def test_prepare_ocr_request_forwards_custom_pricing_to_logging_params() -> None:
|
||||
logging_obj: Final = _prepare_with_pricing({"ocr_cost_per_page": 0.05, "ocr_cost_per_credit": 0.5})
|
||||
|
||||
assert logging_obj.litellm_params["ocr_cost_per_page"] == 0.05
|
||||
assert logging_obj.litellm_params["ocr_cost_per_credit"] == 0.5
|
||||
assert use_custom_pricing_for_model(logging_obj.litellm_params) is True
|
||||
|
||||
|
||||
def test_prepare_ocr_request_without_custom_pricing_leaves_logging_params_unpriced() -> None:
|
||||
logging_obj: Final = _prepare_with_pricing({})
|
||||
|
||||
assert "ocr_cost_per_page" not in logging_obj.litellm_params
|
||||
assert use_custom_pricing_for_model(logging_obj.litellm_params) is False
|
||||
|
||||
|
||||
def test_direct_ocr_call_bills_request_level_per_page_pricing() -> None:
|
||||
assert PRICING_OCR_MODEL not in litellm.model_cost
|
||||
logging_obj: Final = _prepare_with_pricing({"ocr_cost_per_page": 0.05})
|
||||
response: Final = OCRResponse(
|
||||
pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(3)],
|
||||
model=PRICING_OCR_MODEL,
|
||||
usage_info=OCRUsageInfo(pages_processed=3),
|
||||
)
|
||||
|
||||
assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.05 * 3)
|
||||
|
|
|
|||
|
|
@ -17,6 +17,8 @@ from litellm.cost_calculator import (
|
|||
handle_realtime_stream_cost_calculation,
|
||||
response_cost_calculator,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo
|
||||
from litellm.types.llms.openai import OpenAIRealtimeStreamList
|
||||
from litellm.types.rerank import RerankResponse
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -4768,3 +4770,228 @@ def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() ->
|
|||
assert combined.completion_tokens_details.reasoning_tokens == 95
|
||||
assert combined.completion_tokens_details.text_tokens == 38
|
||||
assert combined.completion_tokens_details.audio_tokens == 0
|
||||
|
||||
|
||||
UNMAPPED_OCR_MODEL: Final = "azure_ai/some-unmapped-ocr-model-for-testing"
|
||||
MAPPED_OCR_MODEL: Final = "mistral/mistral-ocr-4-0"
|
||||
|
||||
|
||||
def _ocr_response(model: str, pages_processed: int, credits: float | None = None) -> OCRResponse:
|
||||
return OCRResponse(
|
||||
pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(pages_processed)],
|
||||
model=model,
|
||||
usage_info=OCRUsageInfo(pages_processed=pages_processed, credits=credits),
|
||||
)
|
||||
|
||||
|
||||
def _ocr_logging_obj(litellm_params: dict[str, object]) -> Logging:
|
||||
logging_obj: Final = Logging(
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="ocr",
|
||||
start_time=None,
|
||||
litellm_call_id="test-ocr-custom-pricing",
|
||||
function_id="1234",
|
||||
)
|
||||
logging_obj.update_environment_variables(litellm_params=litellm_params, optional_params={})
|
||||
return logging_obj
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pages_processed", [1, 3, 10])
|
||||
def test_ocr_cost_uses_deployment_per_page_pricing_for_unmapped_model(pages_processed: int):
|
||||
from litellm.cost_calculator import ocr_cost
|
||||
|
||||
assert UNMAPPED_OCR_MODEL not in litellm.model_cost
|
||||
cost, _ = ocr_cost(
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=pages_processed),
|
||||
model_info={"ocr_cost_per_page": 0.004},
|
||||
)
|
||||
assert cost == pytest.approx(0.004 * pages_processed)
|
||||
|
||||
|
||||
def test_ocr_cost_uses_deployment_annotation_only_pricing_for_unmapped_model():
|
||||
from litellm.cost_calculator import ocr_cost
|
||||
|
||||
assert UNMAPPED_OCR_MODEL not in litellm.model_cost
|
||||
response: Final = OCRResponse(
|
||||
pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(3)],
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
usage_info=OCRUsageInfo(pages_processed=3, pages_processed_annotation=2),
|
||||
)
|
||||
cost, _ = ocr_cost(
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
response=response,
|
||||
model_info={"annotation_cost_per_page": 0.01},
|
||||
)
|
||||
assert cost == pytest.approx(0.01 * 2)
|
||||
|
||||
|
||||
def test_ocr_cost_annotation_only_override_keeps_mapped_per_page_rate():
|
||||
from litellm.cost_calculator import ocr_cost
|
||||
|
||||
map_price: Final = litellm.model_cost[MAPPED_OCR_MODEL]["ocr_cost_per_page"]
|
||||
response: Final = OCRResponse(
|
||||
pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(3)],
|
||||
model=MAPPED_OCR_MODEL,
|
||||
usage_info=OCRUsageInfo(pages_processed=3, pages_processed_annotation=2),
|
||||
)
|
||||
cost, _ = ocr_cost(
|
||||
model=MAPPED_OCR_MODEL,
|
||||
custom_llm_provider="mistral",
|
||||
response=response,
|
||||
model_info={"annotation_cost_per_page": 0.01},
|
||||
)
|
||||
assert cost == pytest.approx(map_price * 3 + 0.01 * 2)
|
||||
|
||||
|
||||
def test_ocr_cost_uses_deployment_per_credit_pricing_for_unmapped_model():
|
||||
from litellm.cost_calculator import ocr_cost
|
||||
|
||||
cost, _ = ocr_cost(
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=2, credits=4),
|
||||
model_info={"ocr_cost_per_credit": 0.25},
|
||||
)
|
||||
assert cost == pytest.approx(0.25 * 4)
|
||||
|
||||
|
||||
def test_ocr_cost_unmapped_model_without_deployment_pricing_bills_zero():
|
||||
from litellm.cost_calculator import ocr_cost
|
||||
|
||||
cost, _ = ocr_cost(
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=5),
|
||||
model_info={"id": "some-deployment-id"},
|
||||
)
|
||||
assert cost == 0.0
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_local_model_cost_map")
|
||||
def test_ocr_cost_deployment_pricing_overrides_cost_map_for_mapped_model():
|
||||
from litellm.cost_calculator import ocr_cost
|
||||
|
||||
map_price: Final = litellm.get_model_info(MAPPED_OCR_MODEL)["ocr_cost_per_page"]
|
||||
assert map_price is not None
|
||||
override_price: Final = map_price * 10
|
||||
|
||||
cost, _ = ocr_cost(
|
||||
model=MAPPED_OCR_MODEL,
|
||||
custom_llm_provider="mistral",
|
||||
response=_ocr_response(MAPPED_OCR_MODEL, pages_processed=2),
|
||||
model_info={"ocr_cost_per_page": override_price},
|
||||
)
|
||||
assert cost == pytest.approx(override_price * 2)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_local_model_cost_map")
|
||||
def test_ocr_cost_falls_through_to_cost_map_when_deployment_has_no_ocr_pricing():
|
||||
from litellm.cost_calculator import ocr_cost
|
||||
|
||||
map_price: Final = litellm.get_model_info(MAPPED_OCR_MODEL)["ocr_cost_per_page"]
|
||||
assert map_price is not None
|
||||
|
||||
cost, _ = ocr_cost(
|
||||
model=MAPPED_OCR_MODEL,
|
||||
custom_llm_provider="mistral",
|
||||
response=_ocr_response(MAPPED_OCR_MODEL, pages_processed=2),
|
||||
model_info={"id": "some-deployment-id"},
|
||||
)
|
||||
assert cost == pytest.approx(map_price * 2)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_local_model_cost_map")
|
||||
def test_ocr_cost_ignores_deployment_credit_pricing_when_response_reports_no_credits():
|
||||
from litellm.cost_calculator import ocr_cost
|
||||
|
||||
map_price: Final = litellm.get_model_info(MAPPED_OCR_MODEL)["ocr_cost_per_page"]
|
||||
assert map_price is not None
|
||||
|
||||
cost, _ = ocr_cost(
|
||||
model=MAPPED_OCR_MODEL,
|
||||
custom_llm_provider="mistral",
|
||||
response=_ocr_response(MAPPED_OCR_MODEL, pages_processed=2),
|
||||
model_info={"ocr_cost_per_credit": 0.5},
|
||||
)
|
||||
assert cost == pytest.approx(map_price * 2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
|
||||
def test_completion_cost_ocr_reads_deployment_pricing_from_logging_metadata(metadata_key: str):
|
||||
logging_obj = _ocr_logging_obj({metadata_key: {"model_info": {"ocr_cost_per_page": 0.004}}})
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3),
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
call_type="ocr",
|
||||
custom_pricing=True,
|
||||
litellm_logging_obj=logging_obj,
|
||||
)
|
||||
assert cost == pytest.approx(0.004 * 3)
|
||||
|
||||
|
||||
def test_completion_cost_ocr_prefers_pricing_registered_under_router_model_id(monkeypatch: pytest.MonkeyPatch):
|
||||
deployment_id: Final = "ocr-deployment-priced-through-litellm-params"
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost, deployment_id, {"mode": "ocr", "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.05}
|
||||
)
|
||||
logging_obj = _ocr_logging_obj({"metadata": {"model_info": {"mode": "ocr"}}})
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3),
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
call_type="ocr",
|
||||
custom_pricing=True,
|
||||
router_model_id=deployment_id,
|
||||
litellm_logging_obj=logging_obj,
|
||||
)
|
||||
assert cost == pytest.approx(0.05 * 3)
|
||||
|
||||
|
||||
def test_completion_cost_ocr_bills_request_level_pricing_for_direct_sdk_call():
|
||||
logging_obj = _ocr_logging_obj({"ocr_cost_per_page": 0.05})
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3),
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
call_type="ocr",
|
||||
custom_pricing=True,
|
||||
litellm_logging_obj=logging_obj,
|
||||
)
|
||||
assert cost == pytest.approx(0.05 * 3)
|
||||
|
||||
|
||||
def test_completion_cost_ocr_request_level_pricing_fills_in_deployment_model_info_without_ocr_pricing():
|
||||
logging_obj = _ocr_logging_obj({"ocr_cost_per_page": 0.05, "metadata": {"model_info": {"mode": "ocr"}}})
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3),
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
call_type="ocr",
|
||||
custom_pricing=True,
|
||||
litellm_logging_obj=logging_obj,
|
||||
)
|
||||
assert cost == pytest.approx(0.05 * 3)
|
||||
|
||||
|
||||
def test_completion_cost_ocr_ignores_deployment_pricing_without_custom_pricing_flag():
|
||||
logging_obj = _ocr_logging_obj({"metadata": {"model_info": {"ocr_cost_per_page": 0.004}}})
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3),
|
||||
model=UNMAPPED_OCR_MODEL,
|
||||
custom_llm_provider="azure_ai",
|
||||
call_type="ocr",
|
||||
custom_pricing=False,
|
||||
litellm_logging_obj=logging_obj,
|
||||
)
|
||||
assert cost == 0.0
|
||||
|
|
|
|||
|
|
@ -79,6 +79,20 @@ async def test_proxy_metadata_remains_python_owned(ocr_server: RecordingServer)
|
|||
assert "metadata" not in ocr_server.requests[0].body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_level_custom_pricing_reaches_logging_params_and_bills_the_call(
|
||||
ocr_server: RecordingServer,
|
||||
) -> None:
|
||||
recorder: Final = RecordingLogger()
|
||||
response: Final = await call_aocr(ocr_server, callbacks=[recorder], ocr_cost_per_page=0.05)
|
||||
events: Final = await recorder.wait_for_async("async_log_success_event")
|
||||
|
||||
assert response.usage_info is not None and response.usage_info.pages_processed == 1
|
||||
assert events[0].kwargs["litellm_params"]["ocr_cost_per_page"] == 0.05
|
||||
assert response._hidden_params["response_cost"] == pytest.approx(0.05)
|
||||
assert "ocr_cost_per_page" not in ocr_server.requests[0].body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_replacement_finalized_before_dispatch_in_caller_task(ocr_server: RecordingServer) -> None:
|
||||
caller: Final = asyncio.current_task()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue