litellm/tests/unit/batches/test_batch_utils.py

"""
Unit tests for litellm/batches/batch_utils.py

batch_utils.py is the batch cost/usage/parsing layer: it turns a batch output
JSONL into spend (cost), token usage, and the list of models seen, and counts
tokens in batch *input* files for rate limiting. A silent bug here mis-bills
real money or lets callers slip past TPM limits, so these tests assert exact
numeric results rather than "ran without error".

Pure functions (parsing, token math, credential extraction, success checks) run
for real with exact-value assertions. The few true external seams - the cost
maps (litellm.completion_cost, batch_cost_calculator), the tokenizer
(token_counter), and remote file fetch (afile_content) - are mocked with
deterministic stand-ins so the arithmetic under test is the only variable.
"""

import json
import logging
from types import MappingProxyType

import httpx
import pytest
import respx
from openai.types.batch import BatchRequestCounts


import litellm
import litellm.batches.batch_utils as bu
from litellm.types.utils import LiteLLMBatch, ModelInfo, Usage

# --------------------------------------------------------------------------- #
# Builders for batch OUTPUT file rows.
# Shape: {"response": {"status_code": 200, "body": {... "usage": {...}}}}
# --------------------------------------------------------------------------- #


def _usage(p, c, t=None):
    return {
        "prompt_tokens": p,
        "completion_tokens": c,
        "total_tokens": t if t is not None else p + c,
    }


def _success_row(model="gpt-4o", usage=None, **body_extra):
    body = {"model": model, **body_extra}
    if usage is not None:
        body["usage"] = usage
    return {"response": {"status_code": 200, "body": body}}


def _failed_row(status_code=500, model="gpt-4o"):
    return {"response": {"status_code": status_code, "body": {"model": model}}}


# =========================================================================== #
# _batch_response_was_successful
# =========================================================================== #


@pytest.mark.parametrize(
    "row,expected",
    [
        ({"response": {"status_code": 200}}, True),
        ({"response": {"status_code": 500}}, False),
        ({"response": {"status_code": 429}}, False),
        ({"response": {}}, False),  # no status_code
        ({}, False),  # no response
        ({"response": None}, False),  # null response
    ],
)
def test_batch_response_was_successful(row, expected):
    assert bu._batch_response_was_successful(row) is expected


# =========================================================================== #
# _get_response_from_batch_job_output_file
# =========================================================================== #


def test_get_response_body_present():
    row = {"response": {"body": {"model": "gpt-4o", "usage": {"x": 1}}}}
    assert bu._get_response_from_batch_job_output_file(row) == {
        "model": "gpt-4o",
        "usage": {"x": 1},
    }


@pytest.mark.parametrize(
    "row",
    [
        {},  # no response
        {"response": {}},  # no body
        {"response": None},  # null response
        {"response": {"body": None}},  # null body
    ],
)
def test_get_response_body_missing_returns_empty(row):
    assert bu._get_response_from_batch_job_output_file(row) == {}


# =========================================================================== #
# _get_batch_job_usage_from_response_body
# =========================================================================== #


def test_get_usage_from_response_body():
    usage = bu._get_batch_job_usage_from_response_body(
        {"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}
    )
    assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (
        10,
        5,
        15,
    )


def test_get_usage_from_response_body_missing_is_zero():
    usage = bu._get_batch_job_usage_from_response_body({})
    assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (
        0,
        0,
        0,
    )


# =========================================================================== #
# _get_file_content_as_dictionary  (JSONL parsing)
# =========================================================================== #


def test_parse_jsonl_multiple_lines():
    content = b'{"a": 1}\n{"b": 2}\n{"c": 3}'
    assert bu._get_file_content_as_dictionary(content) == [
        {"a": 1},
        {"b": 2},
        {"c": 3},
    ]


def test_parse_jsonl_trailing_newline_skipped():
    # outer content is stripped; the trailing-newline empty line is dropped.
    content = b'{"a": 1}\n{"b": 2}\n'
    assert bu._get_file_content_as_dictionary(content) == [{"a": 1}, {"b": 2}]


def test_parse_jsonl_empty_content_is_empty_list():
    assert bu._get_file_content_as_dictionary(b"") == []


def test_parse_jsonl_malformed_lines_skipped():
    content = b'{"a": 1}\nnot valid json\n{"b": 2}\n'
    assert bu._get_file_content_as_dictionary(content) == [{"a": 1}, {"b": 2}]


# =========================================================================== #
# _iter_batch_input_lines / _iter_batch_output_entries  (JSONL parsing)
# =========================================================================== #


def test_iter_input_lines_skips_blank_and_strips():
    content = b'{"a":1}\n\n  \n{"b":2}\n'
    assert list(bu._iter_batch_input_lines(content)) == [b'{"a":1}', b'{"b":2}']


def test_iter_input_lines_handles_missing_trailing_newline():
    assert list(bu._iter_batch_input_lines(b'{"a":1}')) == [b'{"a":1}']


def test_iter_input_lines_empty():
    assert list(bu._iter_batch_input_lines(b"")) == []


def test_iter_output_entries_parses_each_row():
    content = b'{"body": {"model": "gpt-4o"}}\n{"body": {"model": "claude-3"}}\n'
    assert list(bu._iter_batch_output_entries(content)) == [
        {"body": {"model": "gpt-4o"}},
        {"body": {"model": "claude-3"}},
    ]


def test_iter_output_entries_skips_malformed_and_non_object_lines():
    content = b'{"ok": 1}\nnot-json\n[1, 2]\n{"ok": 2}\n'
    assert list(bu._iter_batch_output_entries(content)) == [{"ok": 1}, {"ok": 2}]


def test_iter_output_entries_skips_undecodable_line():
    content = b'{"ok": 1}\n{"note": "\xff-bad"}\n{"ok": 2}\n'
    assert list(bu._iter_batch_output_entries(content)) == [{"ok": 1}, {"ok": 2}]


# =========================================================================== #
# _estimate_batch_entry_tokens  (regression: an uncountable/malformed row must
# never contribute zero tokens, or a crafted batch could evade the TPM limit)
# =========================================================================== #


def test_estimate_tokens_scales_with_size():
    # 4 bytes per token, floored, with a minimum of 1.
    assert bu._estimate_batch_entry_tokens(b"a" * 40) == 10


def test_estimate_tokens_never_zero_for_short_rows():
    assert bu._estimate_batch_entry_tokens(b"") == 1
    assert bu._estimate_batch_entry_tokens(b"abc") == 1


# =========================================================================== #
# _aggregate_batch_cost_usage_models: models  (output file)
# =========================================================================== #


def test_output_models_uses_model_name_override(monkeypatch):
    monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
    result = bu._aggregate_batch_cost_usage_models(
        entries=[_success_row(model="ignored")], custom_llm_provider="openai", model_name="forced-model"
    )
    assert result.models == ["forced-model"]


def test_output_models_collects_from_successful_only(monkeypatch):
    monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
    rows = [
        _success_row(model="gpt-4o"),
        _failed_row(model="should-be-skipped"),
        _success_row(model="claude-3"),
    ]
    result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
    assert result.models == ["gpt-4o", "claude-3"]


def test_output_models_skips_successful_without_model(monkeypatch):
    monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
    rows = [{"response": {"status_code": 200, "body": {}}}]
    result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
    assert result.models == []


# =========================================================================== #
# _extract_file_access_credentials
# =========================================================================== #


def test_extract_credentials_only_known_keys():
    params = {
        "api_key": "sk-1",
        "api_base": "https://b",
        "vertex_project": "proj",
        "gcs_bucket_name": "my-bucket",
        "bucket_name": "my-alias-bucket",
        "model": "gpt-4o",  # not a credential key
        "unrelated": "x",
    }
    assert bu._extract_file_access_credentials(params) == {
        "api_key": "sk-1",
        "api_base": "https://b",
        "vertex_project": "proj",
        "gcs_bucket_name": "my-bucket",
        "bucket_name": "my-alias-bucket",
    }


@pytest.mark.parametrize("params", [None, {}])
def test_extract_credentials_empty(params):
    assert bu._extract_file_access_credentials(params) == {}


def test_extract_credentials_all_supported_keys():
    keys = {
        "api_key",
        "api_base",
        "api_version",
        "organization",
        "azure_ad_token",
        "azure_ad_token_provider",
        "vertex_project",
        "vertex_location",
        "vertex_credentials",
        "gcs_bucket_name",
        "bucket_name",
        "s3_endpoint_url",
        "s3_region_name",
        "timeout",
        "max_retries",
    }
    params = {k: f"val-{k}" for k in keys}
    assert bu._extract_file_access_credentials(params) == params


# =========================================================================== #
# _count_prompt_or_input_tokens  (regression-critical: list[list[int]] used to
# count as zero and let callers slip past TPM limits). token_counter stubbed to
# len(text) so every shape has an exact expected value.
# =========================================================================== #


@pytest.fixture
def fake_token_counter(monkeypatch):
    def _tc(model=None, text=None, messages=None, **kw):
        if messages is not None:
            return len(messages)
        if text is not None:
            return len(text)
        return 0

    monkeypatch.setattr(bu, "token_counter", _tc)
    return _tc


def test_count_tokens_str(fake_token_counter):
    assert bu._count_prompt_or_input_tokens("m", "hello") == 5  # len("hello")


def test_count_tokens_list_of_str(fake_token_counter):
    assert bu._count_prompt_or_input_tokens("m", ["ab", "cde"]) == 5  # 2 + 3


def test_count_tokens_list_of_int(fake_token_counter):
    # pre-tokenized prompt: each int counts as one token.
    assert bu._count_prompt_or_input_tokens("m", [1, 2, 3, 4]) == 4


def test_count_tokens_list_of_list_of_int(fake_token_counter):
    # the bug-fix shape: nested pre-tokenized prompts, each int = 1 token.
    assert bu._count_prompt_or_input_tokens("m", [[1, 2, 3], [4, 5]]) == 5


def test_count_tokens_mixed_nested(fake_token_counter):
    # nested list with ints + a string: 2 ints (=2) + len("xyz")=3 -> 5
    assert bu._count_prompt_or_input_tokens("m", [[1, 2, "xyz"]]) == 5


def test_count_tokens_unsupported_shape_is_zero(fake_token_counter):
    assert bu._count_prompt_or_input_tokens("m", 12345) == 0
    assert bu._count_prompt_or_input_tokens("m", {"a": 1}) == 0


# =========================================================================== #
# _count_entry_tokens  (per-entry rate-limit token counting). The individual
# prompt/input/embedding shapes are covered in test_batch_file_validation.py;
# here we pin the body-field precedence and the empty/fallback behavior.
# =========================================================================== #


def test_count_entry_messages_path(fake_token_counter):
    entry = {"body": {"model": "gpt-4o", "messages": [{"role": "user"}, {"role": "x"}]}}
    assert bu._count_entry_tokens(entry) == 2  # len(messages)


def test_count_entry_prompt_path(fake_token_counter):
    assert bu._count_entry_tokens({"body": {"model": "gpt-4o", "prompt": "abcd"}}) == 4


def test_count_entry_input_path(fake_token_counter):
    assert bu._count_entry_tokens({"body": {"model": "gpt-4o", "input": "ab"}}) == 2


def test_count_entry_messages_beats_prompt(fake_token_counter):
    # messages present -> prompt/input are ignored (messages is checked first).
    entry = {
        "body": {
            "model": "gpt-4o",
            "messages": [{"role": "user"}],
            "prompt": "this-should-be-ignored",
        }
    }
    assert bu._count_entry_tokens(entry) == 1


def test_count_entry_prompt_beats_input(fake_token_counter):
    entry = {"body": {"model": "gpt-4o", "prompt": "abc", "input": "this-is-longer"}}
    assert bu._count_entry_tokens(entry) == 3


def test_count_entry_empty_body_is_zero(fake_token_counter):
    assert bu._count_entry_tokens({"body": {}}) == 0
    assert bu._count_entry_tokens({}) == 0


def test_count_entry_uses_model_name_fallback(monkeypatch):
    # No body.model -> the model_name argument is forwarded to the token counter.
    captured = {}

    def _tc(model=None, text=None, messages=None, **kw):
        captured["model"] = model
        return len(text or "")

    monkeypatch.setattr(bu, "token_counter", _tc)
    bu._count_entry_tokens({"body": {"prompt": "ab"}}, model_name="fallback-model")
    assert captured["model"] == "fallback-model"


# =========================================================================== #
# _aggregate_batch_cost_usage_models: usage  (output usage aggregation)
# =========================================================================== #


def test_total_usage_sums_successful_only(monkeypatch):
    monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
    rows = [
        _success_row(usage=_usage(10, 5)),  # 15
        _failed_row(),  # excluded
        _success_row(usage=_usage(20, 10)),  # 30
    ]
    result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (
        30,
        15,
        45,
    )


def test_total_usage_and_cost_normalize_mixed_responses_and_chat():
    responses_row = _success_row(
        usage={
            "input_tokens": 20,
            "output_tokens": 7,
            "total_tokens": 27,
            "input_tokens_details": {"cached_tokens": 3},
        }
    )
    chat_row = _success_row(usage=_usage(10, 5))

    result = bu._aggregate_batch_cost_usage_models(
        entries=[responses_row, chat_row],
        custom_llm_provider="openai",
        model_info={
            "input_cost_per_token_batches": 0.00125,
            "output_cost_per_token_batches": 0.005,
        },
    )

    assert result.usage.prompt_tokens == 30
    assert result.usage.completion_tokens == 12
    assert result.usage.total_tokens == 42
    assert result.usage.cache_read_input_tokens == 3
    assert result.cost == pytest.approx((30 * 0.00125) + (12 * 0.005))


def test_total_cost_applies_the_long_context_batch_tier_per_line():
    long_row = _success_row(usage=_usage(300_000, 10))
    short_row = _success_row(usage=_usage(100, 10))

    result = bu._aggregate_batch_cost_usage_models(
        entries=[long_row, short_row],
        custom_llm_provider="openai",
        model_info=ModelInfo(
            key="lit-batch-tier",
            max_tokens=None,
            max_input_tokens=None,
            max_output_tokens=None,
            input_cost_per_token=2e-6,
            output_cost_per_token=8e-6,
            litellm_provider="openai",
            mode="chat",
            supported_openai_params=None,
            input_cost_per_token_batches=1e-6,
            output_cost_per_token_batches=4e-6,
            input_cost_per_token_above_272k_tokens_batches=2e-6,
            output_cost_per_token_above_272k_tokens_batches=6e-6,
        ),
    )

    assert result.cost == pytest.approx((300_000 * 2e-6) + (10 * 6e-6) + (100 * 1e-6) + (10 * 4e-6))


def test_xai_output_lines_bill_reasoning_tokens_as_completion_tokens():
    row = _success_row(
        model="grok-4.3",
        usage={
            "prompt_tokens": 615,
            "completion_tokens": 3,
            "total_tokens": 993,
            "completion_tokens_details": {"reasoning_tokens": 375},
        },
    )

    result = bu._aggregate_batch_cost_usage_models(
        entries=[row],
        custom_llm_provider="xai",
        model_info=ModelInfo(
            key="xai/grok-4.3",
            max_tokens=None,
            max_input_tokens=None,
            max_output_tokens=None,
            input_cost_per_token=1.25e-6,
            output_cost_per_token=2.5e-6,
            litellm_provider="xai",
            mode="chat",
            supported_openai_params=None,
            input_cost_per_token_batches=1e-6,
            output_cost_per_token_batches=2e-6,
        ),
    )

    assert result.usage.completion_tokens == 378
    assert result.usage.total_tokens == 993
    assert result.cost == pytest.approx((615 * 1e-6) + (378 * 2e-6))


def test_total_usage_empty_is_zero():
    result = bu._aggregate_batch_cost_usage_models(entries=[], custom_llm_provider="openai")
    assert result.cost == 0.0
    assert result.models == []
    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (
        0,
        0,
        0,
    )
    assert result.successful_requests == 0
    assert result.failed_requests == 0


def test_total_usage_includes_reasoning_tokens(monkeypatch):
    monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
    rows = [
        _success_row(
            usage={
                "prompt_tokens": 10,
                "completion_tokens": 50,
                "total_tokens": 60,
                "completion_tokens_details": {"reasoning_tokens": 30},
            }
        ),
        _success_row(
            usage={
                "prompt_tokens": 5,
                "completion_tokens": 20,
                "total_tokens": 25,
                "completion_tokens_details": {"reasoning_tokens": 8},
            }
        ),
        _failed_row(),  # excluded, must not contribute reasoning tokens either
    ]
    result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
    assert result.usage.completion_tokens_details is not None
    assert result.usage.completion_tokens_details.reasoning_tokens == 38


def test_aggregate_counts_successful_and_failed_requests(monkeypatch):
    monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
    rows = [
        _success_row(usage=_usage(10, 5)),
        _failed_row(),
        _success_row(usage=_usage(20, 10)),
        _failed_row(),
        _failed_row(),
    ]
    result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
    assert result.successful_requests == 2
    assert result.failed_requests == 3
    assert result.successful_requests + result.failed_requests == len(rows)


def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch):
    import litellm.cost_calculator as cc

    monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.4, 0.6))
    result = bu._aggregate_batch_cost_usage_models(
        entries=[_success_row(usage=_usage(10, 5))], custom_llm_provider="openai"
    )
    assert isinstance(result, bu.BatchCostUsageResult)
    assert (result.cost, result.models, result.successful_requests, result.failed_requests) == (
        1.0,
        ["gpt-4o"],
        1,
        0,
    )
    assert (result.prompt_cost, result.completion_cost) == (0.4, 0.6)


# =========================================================================== #
# _aggregate_batch_cost_usage_models: cost  (cost maps mocked)
# =========================================================================== #


def test_cost_without_model_info_prices_each_row_by_its_response_model(monkeypatch):
    # model_info is None -> batch_cost_calculator per successful row, model from the response body.
    import litellm.cost_calculator as cc

    calls = []

    def _batch_cost(**kw):
        calls.append(kw)
        return (0.3, 0.2)

    monkeypatch.setattr(cc, "batch_cost_calculator", _batch_cost)
    rows = [
        _success_row(usage=_usage(10, 5)),
        _failed_row(),  # excluded -> not costed
        _success_row(usage=_usage(20, 10)),
    ]

    result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")

    assert result.cost == pytest.approx(1.0)  # 2 successful * (0.3 + 0.2)
    assert (result.prompt_cost, result.completion_cost) == (pytest.approx(0.6), pytest.approx(0.4))
    assert len(calls) == 2  # failed row not costed
    assert all(call["model"] == "gpt-4o" and call["model_info"] is None for call in calls)
    assert result.successful_requests == 2
    assert result.failed_requests == 1


def test_empty_body_line_does_not_zero_whole_batch():
    """A status-200 row with an empty body makes litellm.completion_cost raise;
    that line must be skipped from pricing instead of zeroing the whole batch.

    The provider still reported it as a success, so it stays in
    successful_requests and out of failed_requests - otherwise the counts stop
    reconciling with the provider's own request_counts over a litellm-side
    pricing gap the customer never caused."""
    rows = [
        _success_row(usage=_usage(10, 5)),
        {
            "custom_id": "request-poison-empty",
            "response": {"status_code": 200, "request_id": "inject-empty-body", "body": {}},
        },
        _success_row(usage=_usage(20, 10)),
    ]

    result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")

    assert result.cost > 0.0
    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (30, 15, 45)
    assert result.models == ["gpt-4o", "gpt-4o"]
    assert (result.successful_requests, result.failed_requests) == (3, 0)


def test_cost_from_content_model_info_path(monkeypatch):
    # model_info set -> batch_cost_calculator(prompt_cost, completion_cost).
    import litellm.cost_calculator as cc

    monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.1, 0.2))
    rows = [
        _success_row(usage=_usage(10, 5)),
        _success_row(usage=_usage(20, 10)),
    ]

    result = bu._aggregate_batch_cost_usage_models(
        entries=rows,
        custom_llm_provider="openai",
        model_info={"input_cost_per_token": 0.0},  # type: ignore[arg-type]  # truthy -> model_info path
    )

    assert result.cost == pytest.approx(0.6)  # 2 * (0.1 + 0.2)


def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch):
    """A one-shot generator: any implementation that iterates the entries twice
    (e.g. separate cost and usage passes) sees nothing on the second pass and
    returns wrong totals for at least one of cost/usage/models."""
    import litellm.cost_calculator as cc

    monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.25, 0.25))
    one_shot = (row for row in [_success_row(usage=_usage(10, 5)), _failed_row(), _success_row(usage=_usage(20, 10))])

    result = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai")

    assert result.cost == 1.0
    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (30, 15, 45)
    assert result.models == ["gpt-4o", "gpt-4o"]
    assert result.successful_requests == 2
    assert result.failed_requests == 1


# =========================================================================== #
# calculate_batch_cost_and_usage  (dispatch: vertex-disable-transform vs generic)
# =========================================================================== #


@pytest.mark.asyncio
async def test_calculate_vertex_disable_transform_path(monkeypatch):
    monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False)
    monkeypatch.setattr(
        bu,
        "calculate_vertex_ai_batch_cost_and_usage",
        lambda content, model, model_info=None: bu.BatchCostUsageResult(
            cost=9.9,
            usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3),
            models=["gemini-2.0-flash-001"],
            successful_requests=1,
            failed_requests=0,
        ),
    )
    # generic path must NOT be taken
    monkeypatch.setattr(
        bu,
        "_aggregate_batch_cost_usage_models",
        lambda **kw: pytest.fail("generic path should not run"),
    )

    result = await bu.calculate_batch_cost_and_usage(
        file_content_dictionary=[], custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001"
    )
    assert result.cost == 9.9
    assert result.usage.total_tokens == 3
    assert result.models == ["gemini-2.0-flash-001"]


@pytest.mark.asyncio
async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch):
    """Without a model_name the raw-vertex path cannot price lines; the generic
    aggregation path must run even with the disable flag set."""
    monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False)
    monkeypatch.setattr(
        bu,
        "calculate_vertex_ai_batch_cost_and_usage",
        lambda content, model, model_info=None: pytest.fail("raw vertex path should not run"),
    )

    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 == []


# =========================================================================== #
# calculate_vertex_ai_batch_cost_and_usage  (usageMetadata aggregation)
# =========================================================================== #


def test_vertex_cost_and_usage_aggregation(monkeypatch):
    import litellm.cost_calculator as cc

    monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.1, 0.2))
    responses = [
        {
            "response": {
                "usageMetadata": {
                    "promptTokenCount": 10,
                    "candidatesTokenCount": 5,
                    "totalTokenCount": 15,
                }
            }
        },
        {
            "response": {
                "usageMetadata": {
                    "promptTokenCount": 20,
                    "candidatesTokenCount": 10,
                    "totalTokenCount": 30,
                }
            }
        },
    ]

    result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")

    assert result.cost == pytest.approx(0.6)  # 2 * (0.1 + 0.2)
    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (
        30,
        15,
        45,
    )
    assert result.successful_requests == 2
    assert result.failed_requests == 0


def test_vertex_batch_usage_preserves_modality_token_details(monkeypatch):
    monkeypatch.setitem(
        litellm.model_cost,
        "vertex_ai/gemini-embedding-2",
        {
            "input_cost_per_token_batches": 1e-7,
            "input_cost_per_audio_token_batches": 3.25e-6,
            "input_cost_per_image_token_batches": 2.25e-7,
            "input_cost_per_video_token_batches": 6e-6,
        },
    )
    responses = [
        {
            "key": "id_1",
            "status": "",
            "request": {"content": {"parts": [{"text": "hello"}, {"fileData": {"mimeType": "audio/wav"}}]}},
            "response": {
                "embedding": {"values": [0.1, 0.2]},
                "usageMetadata": {
                    "promptTokenCount": 84,
                    "candidatesTokenCount": 0,
                    "totalTokenCount": 84,
                    "promptTokensDetails": [
                        {"modality": "AUDIO", "tokenCount": 64},
                        {"modality": "TEXT", "tokenCount": 20},
                    ],
                },
            },
        }
    ]

    result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-embedding-2")

    assert (result.successful_requests, result.usage.prompt_tokens) == (1, 84)
    assert result.prompt_cost == pytest.approx(64 * 3.25e-6 + 20 * 1e-7)


def test_vertex_cost_skips_none_response_body(monkeypatch):
    import litellm.cost_calculator as cc

    monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (1.0, 0.0))
    responses = [
        {"response": None},  # skipped
        {
            "response": {
                "usageMetadata": {
                    "promptTokenCount": 7,
                    "candidatesTokenCount": 3,
                    "totalTokenCount": 10,
                }
            }
        },
    ]

    result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")

    assert result.cost == pytest.approx(1.0)  # only one line costed
    assert result.usage.total_tokens == 10
    assert result.successful_requests == 1
    assert result.failed_requests == 1


def test_vertex_usage_total_token_fallback(monkeypatch):
    # no totalTokenCount -> falls back to prompt + completion.
    import litellm.cost_calculator as cc

    monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.0, 0.0))
    responses = [{"response": {"usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 4}}}]

    result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
    assert result.usage.total_tokens == 12


def test_vertex_cost_error_in_line_is_swallowed(monkeypatch):
    # a cost error on one line must not abort aggregation; usage still tallies.
    import litellm.cost_calculator as cc

    def _boom(**kw):
        raise RuntimeError("price map miss")

    monkeypatch.setattr(cc, "batch_cost_calculator", _boom)
    responses = [
        {
            "response": {
                "usageMetadata": {
                    "promptTokenCount": 5,
                    "candidatesTokenCount": 5,
                    "totalTokenCount": 10,
                }
            }
        }
    ]

    result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
    assert result.cost == 0.0
    assert result.usage.total_tokens == 10


# =========================================================================== #
# calculate_batch_cost_and_usage  (async orchestrator)
# =========================================================================== #


@pytest.mark.asyncio
async def test_calculate_batch_cost_and_usage_orchestration(monkeypatch):
    import litellm.cost_calculator as cc

    rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))]
    monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (1.5, 1.0))

    result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=rows, custom_llm_provider="openai")

    assert result.cost == 2.5
    assert (result.prompt_cost, result.completion_cost) == (1.5, 1.0)
    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15)
    assert result.models == ["gpt-4o"]


# =========================================================================== #
# _fetch_batch_output_file_content  (file fetch + credential merge)
# =========================================================================== #


def _batch(output_file_id):
    from litellm.types.llms.openai import Batch

    return Batch(
        id="b",
        completion_window="24h",
        created_at=1,
        endpoint="/v1/chat/completions",
        input_file_id="f",
        object="batch",
        status="completed",
        output_file_id=output_file_id,
    )


def _vertex_openai_row(custom_id, model, prompt_tokens, completion_tokens):
    return {
        "id": f"batch_req_{custom_id}",
        "custom_id": custom_id,
        "response": {
            "status_code": 200,
            "request_id": custom_id,
            "body": {
                "id": f"chatcmpl-{custom_id}",
                "object": "chat.completion",
                "model": model,
                "choices": [
                    {
                        "index": 0,
                        "message": {"role": "assistant", "content": "ok"},
                        "finish_reason": "stop",
                    }
                ],
                "usage": _usage(prompt_tokens, completion_tokens),
            },
        },
        "error": None,
    }


def _vertex_jsonl(rows):
    return "\n".join(json.dumps(row) for row in rows).encode()


@pytest.mark.asyncio
async def test_output_file_content_vertex_fetches_via_afile_content(monkeypatch):
    import litellm.files.main as files_main

    rows = [_vertex_openai_row("request-1", "gemini-3.6-flash", 10, 5)]
    captured: dict = {}

    async def fake_afile_content(**kw):
        captured.update(kw)
        return type("R", (), {"content": _vertex_jsonl(rows)})()

    monkeypatch.setattr(files_main, "afile_content", fake_afile_content)

    result = await bu._fetch_batch_output_file_content(
        _batch("gs://litellm-bucket/output/predictions.jsonl"),
        custom_llm_provider="vertex_ai",
        litellm_params={
            "vertex_project": "proj-1",
            "vertex_location": "us-central1",
            "vertex_credentials": "/path/to/creds.json",
            "gcs_bucket_name": "litellm-bucket",
            "model": "vertex_ai/gemini-3.6-flash",
        },
    )

    assert bu._get_file_content_as_dictionary(result) == rows
    assert captured["file_id"] == "gs://litellm-bucket/output/predictions.jsonl"
    assert captured["custom_llm_provider"] == "vertex_ai"
    assert captured["vertex_project"] == "proj-1"
    assert captured["vertex_location"] == "us-central1"
    assert captured["vertex_credentials"] == "/path/to/creds.json"
    assert captured["gcs_bucket_name"] == "litellm-bucket"
    assert "model" not in captured


@pytest.mark.asyncio
async def test_output_file_content_vertex_unified_file_id_extracts_gcs_uri(monkeypatch):
    import base64

    import litellm.files.main as files_main

    captured: dict = {}

    async def fake_afile_content(**kw):
        captured.update(kw)
        return type("R", (), {"content": b'{"a": 1}'})()

    monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
    unified_id = (
        "litellm_proxy:application/jsonl;unified_id,uuid-1;target_model_names,vertex-model;"
        "llm_output_file_id,gs://litellm-bucket/output/predictions.jsonl;llm_output_file_model_id,model-1"
    )
    encoded_id = base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=")

    await bu._fetch_batch_output_file_content(_batch(encoded_id), custom_llm_provider="vertex_ai")

    assert captured["file_id"] == "gs://litellm-bucket/output/predictions.jsonl"
    assert captured["custom_llm_provider"] == "vertex_ai"


@pytest.mark.asyncio
async def test_output_file_content_model_encoded_file_id_decoded_to_provider_id(monkeypatch):
    import litellm.files.main as files_main
    from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model

    captured: dict = {}

    async def fake_afile_content(**kw):
        captured.update(kw)
        return type("R", (), {"content": b'{"a": 1}'})()

    monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
    encoded_id = encode_file_id_with_model("file-Y3FHrMpi7uCkDpY6fgWGeR", "my-batch-model")

    await bu._fetch_batch_output_file_content(_batch(encoded_id), custom_llm_provider="openai")

    assert captured["file_id"] == "file-Y3FHrMpi7uCkDpY6fgWGeR"
    assert captured["custom_llm_provider"] == "openai"


@pytest.mark.asyncio
async def test_output_file_content_raw_openai_file_id_passes_through(monkeypatch):
    import litellm.files.main as files_main

    captured: dict = {}

    async def fake_afile_content(**kw):
        captured.update(kw)
        return type("R", (), {"content": b'{"a": 1}'})()

    monkeypatch.setattr(files_main, "afile_content", fake_afile_content)

    await bu._fetch_batch_output_file_content(_batch("file-abc123"), custom_llm_provider="openai")

    assert captured["file_id"] == "file-abc123"


def _vertex_predictions_row(custom_id, prompt_tokens, completion_tokens):
    return {
        "request": {
            "contents": [{"role": "user", "parts": [{"text": "hi"}]}],
            "labels": {"litellm_custom_id": custom_id},
        },
        "status": "",
        "response": {
            "candidates": [
                {
                    "content": {"role": "model", "parts": [{"text": "ok"}]},
                    "finishReason": "STOP",
                }
            ],
            "usageMetadata": {
                "promptTokenCount": prompt_tokens,
                "candidatesTokenCount": completion_tokens,
                "totalTokenCount": prompt_tokens + completion_tokens,
            },
            "modelVersion": "gemini-3.6-flash",
        },
        "processed_time": "2026-07-30T00:00:00.000000+00:00",
    }


@pytest.fixture
def respx_interceptable_httpx_client(monkeypatch):
    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
    litellm.in_memory_llm_clients_cache.flush_cache()
    yield
    litellm.in_memory_llm_clients_cache.flush_cache()


@pytest.mark.asyncio
@respx.mock
async def test_output_file_content_vertex_managed_uri_accepted_by_real_validation(respx_interceptable_httpx_client):
    managed_output_uri = (
        "gs://litellm-bucket/litellm-vertex-files/publishers/google/models/"
        "gemini-3.6-flash/abc-123/prediction-model/predictions.jsonl"
    )
    rows = [
        _vertex_predictions_row("request-1", 10, 5),
        _vertex_predictions_row("request-2", 20, 10),
    ]
    route = respx.get(url__regex=r"https://storage\.googleapis\.com/storage/v1/b/litellm-bucket/o/.*").mock(
        return_value=httpx.Response(200, content=_vertex_jsonl(rows))
    )

    file_content = await bu._fetch_batch_output_file_content(
        _batch(managed_output_uri),
        custom_llm_provider="vertex_ai",
        litellm_params={
            "api_key": "test-token",
            "vertex_project": "proj-1",
            "vertex_location": "us-central1",
            "gcs_bucket_name": "litellm-bucket",
        },
    )
    result = bu._get_file_content_as_dictionary(file_content)

    assert route.call_count == 1
    request = route.calls.last.request
    assert request.url.raw_path == (
        b"/storage/v1/b/litellm-bucket/o/"
        b"litellm-vertex-files%2Fpublishers%2Fgoogle%2Fmodels%2Fgemini-3.6-flash"
        b"%2Fabc-123%2Fprediction-model%2Fpredictions.jsonl?alt=media"
    )
    assert [row["custom_id"] for row in result] == ["request-1", "request-2"]
    assert all(row["response"]["status_code"] == 200 for row in result)
    assert all(row["response"]["body"]["model"] == "gemini-3.6-flash" for row in result)
    assert [row["response"]["body"]["usage"]["prompt_tokens"] for row in result] == [10, 20]
    assert [row["response"]["body"]["usage"]["completion_tokens"] for row in result] == [5, 10]


@pytest.mark.asyncio
@respx.mock
async def test_output_file_content_vertex_foreign_bucket_rejected_by_real_validation():
    with pytest.raises(Exception, match="does not match the configured storage bucket"):
        await bu._fetch_batch_output_file_content(
            _batch("gs://attacker-bucket/litellm-vertex-files/x/predictions.jsonl"),
            custom_llm_provider="vertex_ai",
            litellm_params={
                "api_key": "test-token",
                "vertex_project": "proj-1",
                "vertex_location": "us-central1",
                "gcs_bucket_name": "litellm-bucket",
            },
        )

    assert respx.mock.calls.call_count == 0


@pytest.mark.asyncio
async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monkeypatch):
    import litellm.files.main as files_main

    rows = [
        _vertex_openai_row("request-1", "gemini-3.6-flash", 10, 5),
        _vertex_openai_row("request-2", "gemini-3.6-flash", 20, 10),
    ]

    async def fake_afile_content(**kw):
        return type("R", (), {"content": _vertex_jsonl(rows)})()

    monkeypatch.setattr(files_main, "afile_content", fake_afile_content)

    result = await bu._handle_completed_batch(
        _batch("gs://litellm-bucket/output/predictions.jsonl"),
        custom_llm_provider="vertex_ai",
        litellm_params={"vertex_project": "proj-1", "vertex_location": "us-central1"},
    )

    pricing = litellm.model_cost["vertex_ai/gemini-3.6-flash"]
    batch_input = pricing["input_cost_per_token_batches"]
    batch_output = pricing["output_cost_per_token_batches"]

    assert batch_input < pricing["input_cost_per_token"]
    assert batch_output < pricing["output_cost_per_token"]
    assert result.cost > 0
    assert result.cost == pytest.approx(30 * batch_input + 15 * batch_output)
    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (30, 15, 45)
    assert result.models == ["gemini-3.6-flash", "gemini-3.6-flash"]
    assert result.successful_requests == 2
    assert result.failed_requests == 0


@pytest.mark.asyncio
async def test_output_file_content_no_output_file_id_raises():
    with pytest.raises(ValueError, match="Output file id is None"):
        await bu._fetch_batch_output_file_content(_batch(None), custom_llm_provider="openai")


@pytest.mark.asyncio
async def test_output_file_content_fetches_and_parses(monkeypatch):
    import litellm.files.main as files_main
    import litellm.proxy.openai_files_endpoints.common_utils as cu

    captured: dict = {}

    async def fake_afile_content(**kw):
        captured.update(kw)
        return type("R", (), {"content": b'{"a": 1}\n{"b": 2}'})()

    monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
    monkeypatch.setattr(cu, "_is_base64_encoded_unified_file_id", lambda fid: False)

    result = await bu._fetch_batch_output_file_content(
        _batch("file-out"),
        custom_llm_provider="azure",
        litellm_params={"api_key": "sk-az", "api_base": "https://az", "model": "x"},
    )

    assert result == b'{"a": 1}\n{"b": 2}'
    # afile_content received the file id + extracted credentials (not "model").
    assert captured["file_id"] == "file-out"
    assert captured["custom_llm_provider"] == "azure"
    assert captured["api_key"] == "sk-az"
    assert captured["api_base"] == "https://az"
    assert "model" not in captured


@pytest.mark.asyncio
async def test_output_file_content_unified_file_id_extraction(monkeypatch):
    # a base64 unified id carries the real provider file id inside
    # "llm_output_file_id,<FID>;" - it must be unwrapped before the fetch.
    import litellm.files.main as files_main
    import litellm.proxy.openai_files_endpoints.common_utils as cu

    captured: dict = {}

    async def fake_afile_content(**kw):
        captured.update(kw)
        return type("R", (), {"content": b'{"a": 1}'})()

    monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
    monkeypatch.setattr(
        cu,
        "_is_base64_encoded_unified_file_id",
        lambda fid: "litellm_proxy;llm_output_file_id,real-file-99;rest",
    )

    await bu._fetch_batch_output_file_content(_batch("encoded-blob"), custom_llm_provider="openai")

    assert captured["file_id"] == "real-file-99"


# =========================================================================== #
# _handle_completed_batch  (async orchestrator: fetch -> single-pass aggregate)
# =========================================================================== #


@pytest.mark.asyncio
async def test_handle_completed_batch_orchestration(monkeypatch):
    rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))]

    async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
        return _vertex_jsonl(rows)

    import litellm.cost_calculator as cc

    monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
    monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (2.0, 1.3))

    result = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai")

    assert result.cost == 3.3
    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15)
    assert result.models == ["gpt-4o"]


@pytest.mark.asyncio
async def test_handle_completed_batch_counts_error_file_failures(monkeypatch):
    """Regression test: OpenAI writes per-request failures (e.g. a rejected param)
    to a separate error_file_id, never into the output file - so failed_requests
    must include them or it silently undercounts real batch failures."""
    from litellm.types.llms.openai import Batch

    rows = [_success_row(model="gpt-5-mini", usage=_usage(24, 107))]
    error_rows = [
        {
            "id": "batch_req_err1",
            "custom_id": "req-2-bad",
            "response": {"status_code": 400, "body": {"error": {"message": "Invalid 'temperature'"}}},
            "error": None,
        }
    ]

    async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
        return _vertex_jsonl(rows)

    async def fake_afile_content(**kw):
        return type("R", (), {"content": _vertex_jsonl(error_rows)})()

    import litellm.files.main as files_main

    monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
    monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
    monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)

    batch = Batch(
        id="b",
        completion_window="24h",
        created_at=1,
        endpoint="/v1/chat/completions",
        input_file_id="f",
        object="batch",
        status="completed",
        output_file_id="of",
        error_file_id="ef",
    )

    result = await bu._handle_completed_batch(batch, custom_llm_provider="openai")

    assert result.successful_requests == 1
    assert result.failed_requests == 1


@pytest.mark.asyncio
async def test_handle_completed_batch_decodes_model_encoded_error_file_id(monkeypatch):
    """A model-encoded error file id must be decoded to the raw provider id before
    the fetch, exactly like the output file id. Sending the encoded id straight to
    the provider 404s, and the swallowed fetch failure silently reports 0 failures."""
    import base64

    from litellm.types.llms.openai import Batch

    provider_error_file_id = "file-real-error-id"
    encoded_error_file_id = "file-" + base64.urlsafe_b64encode(
        f"litellm:{provider_error_file_id};model,model-abc".encode()
    ).decode().rstrip("=")

    requested_file_ids = []

    async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
        return _vertex_jsonl([_success_row(model="gpt-4o", usage=_usage(10, 5))])

    async def fake_afile_content(**kw):
        requested_file_ids.append(kw["file_id"])
        return type("R", (), {"content": _vertex_jsonl([{"custom_id": "bad-1"}])})()

    import litellm.files.main as files_main

    monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
    monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
    monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)

    batch = Batch(
        id="b",
        completion_window="24h",
        created_at=1,
        endpoint="/v1/chat/completions",
        input_file_id="f",
        object="batch",
        status="completed",
        output_file_id="of",
        error_file_id=encoded_error_file_id,
    )

    result = await bu._handle_completed_batch(batch, custom_llm_provider="openai")

    assert requested_file_ids == [provider_error_file_id]
    assert result.failed_requests == 1


@pytest.mark.asyncio
async def test_handle_completed_batch_no_error_file_id_reports_zero_error_failures(monkeypatch):
    rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))]

    async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
        return _vertex_jsonl(rows)

    monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
    monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)

    result = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai")

    assert result.successful_requests == 1
    assert result.failed_requests == 0


@pytest.mark.asyncio
async def test_handle_completed_batch_no_output_file_is_zero(monkeypatch):
    """
    Regression: an all-error batch completes with output_file_id=None (results go
    to a separate error_file_id). _handle_completed_batch must report an empty
    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")

    monkeypatch.setattr(bu, "_fetch_batch_output_file_content", _must_not_fetch)

    result = await bu._handle_completed_batch(_batch(None), custom_llm_provider="openai")

    assert result.cost == 0.0
    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (0, 0, 0)
    assert result.models == []
    assert result.successful_requests == 0
    assert result.failed_requests == 0


@pytest.mark.asyncio
async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch):
    raw_rows = [{"response": {"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 2}}}]

    async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
        return _vertex_jsonl(raw_rows)

    monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
    monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False)
    seen: dict = {}

    def fake_vertex_calc(content, model, model_info=None):
        seen["content"] = content
        seen["model"] = model
        return bu.BatchCostUsageResult(
            cost=7.7,
            usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3),
            models=["gemini-x"],
            successful_requests=1,
            failed_requests=0,
        )

    monkeypatch.setattr(bu, "calculate_vertex_ai_batch_cost_and_usage", fake_vertex_calc)

    result = await bu._handle_completed_batch(
        _batch("gs://litellm-bucket/output/predictions.jsonl"),
        custom_llm_provider="vertex_ai",
        model_name="gemini-x",
    )

    assert result.cost == 7.7
    assert result.usage.total_tokens == 3
    assert result.models == ["gemini-x"]
    assert list(seen["content"]) == raw_rows
    assert seen["model"] == "gemini-x"


def _anthropic_usage(input_tokens, output_tokens, cache_creation=0, cache_read=0):
    return {
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "cache_creation_input_tokens": cache_creation,
        "cache_read_input_tokens": cache_read,
    }


def _anthropic_succeeded_row(model="claude-sonnet-4-5-20250929", usage=None):
    return {
        "custom_id": "req-1",
        "result": {
            "type": "succeeded",
            "message": {
                "id": "msg_1",
                "type": "message",
                "role": "assistant",
                "model": model,
                "content": [{"type": "text", "text": "ok"}],
                "stop_reason": "end_turn",
                "usage": usage or _anthropic_usage(10, 5),
            },
        },
    }


def _anthropic_errored_row():
    return {
        "custom_id": "req-2",
        "result": {
            "type": "errored",
            "error": {"type": "invalid_request_error", "message": "bad request"},
        },
    }


_ANTHROPIC_MODEL_INFO = {
    "input_cost_per_token": 3e-6,
    "output_cost_per_token": 15e-6,
    "cache_read_input_token_cost": 3e-7,
    "cache_creation_input_token_cost": 3.75e-6,
}


@pytest.mark.parametrize(
    "row,expected",
    [
        (_anthropic_succeeded_row(), True),
        (_anthropic_errored_row(), False),
        ({"custom_id": "x", "result": {"type": "canceled"}}, False),
        ({"custom_id": "x", "result": {"type": "expired"}}, False),
        ({"custom_id": "x"}, False),
        ({"custom_id": "x", "result": None}, False),
    ],
)
def test_anthropic_result_line_success_check(row, expected):
    """
    LIT-4008 regression: anthropic batch results JSONL lines are not
    OpenAI-shaped; success is result.type == "succeeded", not
    response.status_code == 200. Pre-fix every anthropic line parsed as
    unsuccessful, so completed batches were billed $0 forever.
    """
    assert bu._batch_response_was_successful(row, custom_llm_provider="anthropic") is expected


def test_anthropic_response_body_is_result_message():
    row = _anthropic_succeeded_row(model="claude-sonnet-4-5-20250929")
    body = bu._get_response_from_batch_job_output_file(row, custom_llm_provider="anthropic")
    assert body["model"] == "claude-sonnet-4-5-20250929"
    assert body["usage"] == _anthropic_usage(10, 5)


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),
    }
    usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="anthropic")
    assert usage.prompt_tokens == 11000
    assert usage.completion_tokens == 200
    assert usage.total_tokens == 11200
    assert usage.prompt_tokens_details.cached_tokens == 8000
    assert usage.prompt_tokens_details.cache_creation_tokens == 2000


def test_bedrock_model_output_line_success_check():
    row = {
        "recordId": "1",
        "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"
    )


def test_bedrock_cost_uses_deployment_model_name():
    row = {
        "recordId": "1",
        "modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}},
    }
    result = bu._aggregate_batch_cost_usage_models(
        entries=[row],
        custom_llm_provider="bedrock",
        model_name="us.anthropic.claude-sonnet-4-6",
        model_info={},
    )
    assert result.cost > 0
    assert result.models == ["us.anthropic.claude-sonnet-4-6"]


def test_anthropic_total_usage_sums_succeeded_only(monkeypatch):
    import litellm.cost_calculator as cc

    monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.0, 0.0))
    rows = [
        _anthropic_succeeded_row(usage=_anthropic_usage(10, 5)),
        _anthropic_errored_row(),
        _anthropic_succeeded_row(usage=_anthropic_usage(20, 10, cache_read=100)),
    ]
    result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (130, 15, 145)
    assert result.successful_requests == 2
    assert result.failed_requests == 1


def test_anthropic_total_usage_aggregates_cache_token_details(monkeypatch):
    import litellm.cost_calculator as cc

    monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.0, 0.0))
    rows = [
        _anthropic_succeeded_row(usage=_anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000)),
        _anthropic_errored_row(),
        _anthropic_succeeded_row(usage=_anthropic_usage(50, 20, cache_creation=300, cache_read=700)),
    ]
    result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
    assert result.usage.prompt_tokens_details.cached_tokens == 8700
    assert result.usage.prompt_tokens_details.cache_creation_tokens == 2300
    assert result.usage.cache_read_input_tokens == 8700
    assert result.usage.cache_creation_input_tokens == 2300


def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch):
    monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
    rows = [
        {
            "custom_id": "req-1",
            "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")
    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15)
    assert result.usage.prompt_tokens_details is None


def test_anthropic_cost_applies_batch_discount_and_cache_pricing():
    """Anthropic batches bill at 50% of the regular rate for base input,
    cache reads, cache writes, and output tokens alike."""
    rows = [
        _anthropic_succeeded_row(usage=_anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000)),
        _anthropic_errored_row(),
    ]

    result = bu._aggregate_batch_cost_usage_models(
        entries=rows,
        custom_llm_provider="anthropic",
        model_info=_ANTHROPIC_MODEL_INFO,  # type: ignore[arg-type]
    )

    expected_half_price = (1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3.75e-6 + 200 * 15e-6) / 2
    assert result.cost == pytest.approx(expected_half_price)


def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatch):
    import litellm.cost_calculator as cc

    seen = []

    def _fake_batch_cost_calculator(**kw):
        seen.append(kw)
        return (0.1, 0.2)

    monkeypatch.setattr(cc, "batch_cost_calculator", _fake_batch_cost_calculator)
    monkeypatch.setattr(
        litellm,
        "completion_cost",
        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"
    )

    assert result.cost == pytest.approx(0.3)
    assert seen[0]["model"] == "claude-sonnet-4-5-20250929"
    assert seen[0]["custom_llm_provider"] == "anthropic"
    assert seen[0]["usage"].prompt_tokens == 10


def test_anthropic_batch_models_collected_from_succeeded_rows(monkeypatch):
    import litellm.cost_calculator as cc

    monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.0, 0.0))
    rows = [
        _anthropic_succeeded_row(model="claude-sonnet-4-5-20250929"),
        _anthropic_errored_row(),
    ]
    result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
    assert result.models == ["claude-sonnet-4-5-20250929"]


@pytest.mark.asyncio
async def test_calculate_batch_cost_and_usage_anthropic_end_to_end():
    rows = [
        _anthropic_succeeded_row(usage=_anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000)),
        _anthropic_errored_row(),
    ]

    result = await bu.calculate_batch_cost_and_usage(
        file_content_dictionary=rows,
        custom_llm_provider="anthropic",
        model_name="claude-sonnet-4-5",
        model_info=_ANTHROPIC_MODEL_INFO,  # type: ignore[arg-type]
    )

    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.models == ["claude-sonnet-4-5"]


def test_extract_credentials_forwards_the_trusted_model_credential_snapshot():
    """Bedrock resolves a batch's output bucket only from the immutable server-side
    snapshot, never from a request param, so cost accounting on the retrieve path cannot
    read the output file unless this key is forwarded. Without it the accounting raises
    "S3 bucket_name is required" for a bucket the deployment has configured, and the
    batch's cost is never recorded."""
    snapshot = MappingProxyType({"s3_bucket_name": "configured-bucket", "aws_region_name": "us-east-1"})

    credentials = bu._extract_file_access_credentials({"_litellm_internal_model_credentials": snapshot})

    assert credentials["_litellm_internal_model_credentials"] is snapshot


def test_extract_credentials_forwards_the_deployment_aws_credentials():
    """The retrieve path's logging object carries the deployment's AWS keys in its
    litellm_params, and the S3 read of the output file signs with whatever afile_content
    receives. Dropping them here sent the read to the ambient credential chain, so a
    deployment whose only AWS credentials live in its litellm_params never recorded
    batch cost on retrieve even once the bucket resolved."""
    params = {
        "aws_access_key_id": "AKIA-deployment",
        "aws_secret_access_key": "secret-deployment",
        "aws_session_token": "token-deployment",
        "aws_region_name": "us-west-2",
        "aws_role_name": "arn:aws:iam::123456789012:role/batch-reader",
        "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
    }

    credentials = bu._extract_file_access_credentials(params)

    assert credentials == {key: value for key, value in params.items() if key != "model"}


@pytest.mark.asyncio
async def test_output_file_content_bedrock_reads_with_deployment_aws_credentials(monkeypatch):
    import litellm.files.main as files_main

    captured: dict = {}

    async def fake_afile_content(**kw):
        captured.update(kw)
        return type("R", (), {"content": b""})()

    monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
    snapshot = MappingProxyType({"s3_bucket_name": "configured-bucket", "aws_region_name": "us-west-2"})

    await bu._fetch_batch_output_file_content(
        _batch("s3://configured-bucket/litellm-batch-outputs/job-1/out.jsonl.out"),
        custom_llm_provider="bedrock",
        litellm_params={
            "aws_access_key_id": "AKIA-deployment",
            "aws_secret_access_key": "secret-deployment",
            "aws_session_token": "token-deployment",
            "aws_region_name": "us-west-2",
            "_litellm_internal_model_credentials": snapshot,
            "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
        },
    )

    assert captured["file_id"] == "s3://configured-bucket/litellm-batch-outputs/job-1/out.jsonl.out"
    assert captured["custom_llm_provider"] == "bedrock"
    assert captured["aws_access_key_id"] == "AKIA-deployment"
    assert captured["aws_secret_access_key"] == "secret-deployment"
    assert captured["aws_session_token"] == "token-deployment"
    assert captured["aws_region_name"] == "us-west-2"
    assert captured["_litellm_internal_model_credentials"] is snapshot
    assert "model" not in captured


# =========================================================================== #
# _handle_completed_batch threads the deployment's model identity + pricing
# =========================================================================== #


def _bedrock_row(model: str, input_tokens: int, output_tokens: int) -> dict[str, object]:
    return {
        "modelInput": {"messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]},
        "modelOutput": {
            "model": model,
            "id": "msg_1",
            "type": "message",
            "role": "assistant",
            "content": [{"type": "text", "text": "ok"}],
            "stop_reason": "end_turn",
            "usage": {
                "input_tokens": input_tokens,
                "output_tokens": output_tokens,
                "cache_creation_input_tokens": 0,
                "cache_read_input_tokens": 0,
            },
        },
        "recordId": "r",
    }


@pytest.mark.asyncio
async def test_handle_completed_bedrock_batch_prices_from_deployment_model(monkeypatch) -> None:
    """A bedrock batch must price from the deployment model, not the response model."""
    rows = [_bedrock_row("claude-sonnet-4-6", 18, 10)] * 100

    async def fake_fetch(batch: object, custom_llm_provider: str, litellm_params: dict | None = None) -> bytes:
        return _vertex_jsonl(rows)

    monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)

    result = await bu._handle_completed_batch(
        _batch("of"),
        custom_llm_provider="bedrock",
        model_name="bedrock/global.anthropic.claude-sonnet-4-6",
    )

    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (1800, 1000, 2800)

    # The response model alone cannot price a bedrock batch: this is the $0 bug.
    zero_result = await bu._handle_completed_batch(
        _batch("of"),
        custom_llm_provider="bedrock",
        model_name=None,
    )
    assert zero_result.cost == 0.0
    assert zero_result.usage.total_tokens == 2800


@pytest.mark.asyncio
async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) -> None:
    """A deployment's configured rates must win over the global cost map."""
    rows = [_success_row(model="gemini-2.5-flash", usage=_usage(60, 75))]

    async def fake_fetch(batch: object, custom_llm_provider: str, litellm_params: dict | None = None) -> bytes:
        return _vertex_jsonl(rows)

    monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)

    free_result = await bu._handle_completed_batch(
        _batch("of"),
        custom_llm_provider="vertex_ai",
        model_name="vertex_ai/gemini-2.5-flash",
        model_info={
            "input_cost_per_token": 0.0,
            "output_cost_per_token": 0.0,
            "input_cost_per_token_batches": 0.0,
            "output_cost_per_token_batches": 0.0,
        },
    )
    assert free_result.cost == 0.0

    billed_result = await bu._handle_completed_batch(
        _batch("of"),
        custom_llm_provider="vertex_ai",
        model_name="vertex_ai/gemini-2.5-flash",
        model_info=None,
    )
    assert billed_result.cost > 0.0


# =========================================================================== #
# _get_batch_job_usage_from_response_body: bedrock usage shapes
# =========================================================================== #


def test_bedrock_converse_shaped_batch_usage_is_parsed():
    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)


def test_bedrock_converse_batch_usage_totals_default_when_absent():
    body = {"model": "us.amazon.nova-lite-v1:0", "usage": {"inputTokens": 10, "outputTokens": 4}}
    usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
    assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 4, 14)


def test_bedrock_converse_batch_usage_includes_cache_tokens():
    body = {
        "model": "us.amazon.nova-lite-v1:0",
        "usage": {
            "inputTokens": 100,
            "outputTokens": 20,
            "totalTokens": 120,
            "cacheReadInputTokens": 800,
            "cacheWriteInputTokens": 200,
        },
    }
    usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
    assert usage.prompt_tokens == 1100
    assert usage.completion_tokens == 20
    assert usage.prompt_tokens_details.cached_tokens == 800
    assert usage.prompt_tokens_details.cache_creation_tokens == 200


def test_bedrock_anthropic_shaped_batch_usage_still_parsed():
    """Anthropic-shaped bedrock output (what an Anthropic model's batch emits) must not regress."""
    body = {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 18, "output_tokens": 10}}
    usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
    assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (18, 10, 28)


def test_bedrock_titan_embedding_batch_usage_is_parsed():
    """Titan embedding batch lines carry a top-level inputTextTokenCount and no usage block."""
    body = {"embedding": [0.1, 0.2], "embeddingsByType": {"float": [0.1, 0.2]}, "inputTextTokenCount": 17}
    usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
    assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (17, 0, 17)


def test_bedrock_titan_embedding_batch_is_billed():
    """Binary embedding rows carry only embeddingsByType and must bill like float rows."""
    rows = [
        {"recordId": "0", "modelOutput": {"embedding": [0.1], "inputTextTokenCount": 10}},
        {"recordId": "1", "modelOutput": {"embeddingsByType": {"binary": [1, 0]}, "inputTextTokenCount": 7}},
    ]
    result = bu._aggregate_batch_cost_usage_models(
        entries=rows,
        custom_llm_provider="bedrock",
        model_name="amazon.titan-embed-text-v2:0",
        model_info={"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 0.0},
    )
    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (17, 0, 17)
    assert result.cost == pytest.approx(17 * 1e-6)


@pytest.mark.parametrize(
    "body",
    [
        {"embedding": [0.1], "inputTextTokenCount": "17"},
        {"embedding": [0.1], "inputTextTokenCount": True},
        {"embedding": [0.1], "inputTextTokenCount": None},
        {"results": [{"outputText": "hi", "tokenCount": 2}], "inputTextTokenCount": 17},
    ],
)
def test_bedrock_input_text_token_count_outside_embedding_lines_is_not_billed(body):
    """Only embedding lines are parsed here; Titan text generation lines are left as they were."""
    usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
    assert usage.total_tokens == 0


def test_unparsable_bedrock_batch_usage_warns(caplog):
    """An unrecognized usage shape must be visible, not a silent $0."""
    body = {"model": "amazon.titan-text-lite-v1", "usage": {"inputTextTokenCount": 42}}
    with caplog.at_level(logging.WARNING):
        usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
    assert usage.total_tokens == 0
    assert "does not understand" in caplog.text
    assert "inputTextTokenCount" in caplog.text


class TestFileAccessCredentialsCarryFederation:
    """A federated deployment holds no api_key, so the fetch that reads a finished batch's output
    has to inherit the federation fields or it cannot authenticate and the batch is never billed."""

    def test_federation_fields_survive_extraction(self):
        from litellm.batches.batch_utils import _extract_file_access_credentials

        credentials = _extract_file_access_credentials(
            {
                "model": "anthropic/claude-sonnet-4-5",
                "anthropic_federation_rule_id": "fdrl_x",
                "anthropic_organization_id": "org-x",
                "anthropic_identity_token_file": "/var/run/secrets/anthropic.com/token",
                "something_unrelated": "dropped",
            }
        )

        assert credentials["anthropic_federation_rule_id"] == "fdrl_x"
        assert credentials["anthropic_organization_id"] == "org-x"
        assert credentials["anthropic_identity_token_file"] == "/var/run/secrets/anthropic.com/token"
        assert "something_unrelated" not in credentials

    def test_every_federation_field_is_carried(self):
        """Derived from the kwargs set, so a new federation field is carried without an edit here."""
        from litellm.batches.batch_utils import _extract_file_access_credentials
        from litellm.litellm_core_utils.get_litellm_params import ANTHROPIC_WIF_KWARGS_KEYS

        params = {name: f"value-{name}" for name in ANTHROPIC_WIF_KWARGS_KEYS}

        credentials = _extract_file_access_credentials(params)

        assert set(credentials) == set(ANTHROPIC_WIF_KWARGS_KEYS)


def test_total_cost_bills_cached_tokens_per_line_at_the_batch_cached_rate():
    responses_row = _success_row(
        usage={
            "input_tokens": 300_000,
            "output_tokens": 10,
            "total_tokens": 300_010,
            "input_tokens_details": {"cached_tokens": 299_000},
        }
    )
    chat_row = _success_row(usage={**_usage(100, 10), "prompt_tokens_details": {"cached_tokens": 60}})

    result = bu._aggregate_batch_cost_usage_models(
        entries=[responses_row, chat_row],
        custom_llm_provider="openai",
        model_info=ModelInfo(
            key="lit-batch-cached-tier",
            max_tokens=None,
            max_input_tokens=None,
            max_output_tokens=None,
            input_cost_per_token=2e-6,
            output_cost_per_token=8e-6,
            cache_read_input_token_cost=1e-6,
            litellm_provider="openai",
            mode="chat",
            supported_openai_params=None,
            input_cost_per_token_batches=1e-6,
            output_cost_per_token_batches=4e-6,
            cache_read_input_token_cost_batches=5e-7,
            input_cost_per_token_above_272k_tokens_batches=2e-6,
            output_cost_per_token_above_272k_tokens_batches=6e-6,
            cache_read_input_token_cost_above_272k_tokens_batches=1e-6,
        ),
    )

    long_line = 1_000 * 2e-6 + 299_000 * 1e-6 + 10 * 6e-6
    short_line = 40 * 1e-6 + 60 * 5e-7 + 10 * 4e-6
    assert result.cost == pytest.approx(long_line + short_line)


# --------------------------------------------------------------------------- #
# batch_cost_is_final
# --------------------------------------------------------------------------- #


def _retrieved_batch(
    status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None
) -> LiteLLMBatch:
    return LiteLLMBatch(
        id="batch_abc",
        completion_window="24h",
        created_at=1,
        endpoint="/v1/chat/completions",
        input_file_id="file-in",
        object="batch",
        status="validating",
        output_file_id=output_file_id,
        request_counts=counts,
    ).model_copy(update={"status": status})


class TestBatchCostIsFinal:
    """Every retrieve of one batch writes the same spend row, so the first retrieve
    that prices it decides the row for good. A poll before the output exists must
    therefore not count as final: pricing it recorded $0 and pinned it (LIT-7048)."""

    @pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"])
    def test_in_flight_batch_is_not_final(self, status):
        assert bu.batch_cost_is_final(_retrieved_batch(status)) is False

    @pytest.mark.parametrize("status", ["completed", "complete"])
    def test_completed_with_output_is_final(self, status):
        assert bu.batch_cost_is_final(_retrieved_batch(status, output_file_id="file-out")) is True

    def test_completed_without_output_and_unknown_counts_is_not_final(self):
        assert bu.batch_cost_is_final(_retrieved_batch("completed")) is False

    def test_completed_without_output_and_zero_counts_is_not_final(self):
        counts = BatchRequestCounts(total=0, completed=0, failed=0)
        assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False

    def test_completed_without_output_but_successful_lines_is_not_final(self):
        counts = BatchRequestCounts(total=2, completed=2, failed=0)
        assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False

    @pytest.mark.parametrize("status", ["completed", "complete"])
    def test_completed_without_output_and_every_line_failed_is_final(self, status):
        counts = BatchRequestCounts(total=2, completed=0, failed=2)
        assert bu.batch_cost_is_final(_retrieved_batch(status, counts=counts)) is True

    @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_keep_the_published_page_rate_when_the_deployment_prices_only_annotations(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",
        model_info={"annotation_cost_per_page_batches": 0.01},
    )
    assert result.cost == pytest.approx(4 * 0.002 + 4 * 0.01)


def test_ocr_rows_keep_the_deployment_page_rate_when_the_unmapped_model_has_no_annotation_price(monkeypatch):
    def _unmapped(model, custom_llm_provider=None):
        raise Exception(f"This model isn't mapped yet: {model}")

    monkeypatch.setattr(litellm, "get_model_info", _unmapped)
    result = bu._aggregate_batch_cost_usage_models(
        entries=[_ocr_row(4, annotation_pages=4, model="my-private-ocr-model")],
        custom_llm_provider="mistral",
        model_info={"ocr_cost_per_page_batches": 0.001},
    )
    assert result.cost == pytest.approx(4 * 0.001 + 4 * 0.001)


def test_ocr_rows_bill_the_deployment_sync_page_rate_over_the_published_batch_rate(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(3)],
        custom_llm_provider="mistral",
        model_info={"ocr_cost_per_page": 0.0912},
    )
    assert result.cost == pytest.approx(3 * 0.0912)


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


GROUNDED_USAGE_METADATA = {
    "promptTokenCount": 19,
    "candidatesTokenCount": 59,
    "thoughtsTokenCount": 406,
    "toolUsePromptTokenCount": 73,
    "totalTokenCount": 557,
    "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 19}],
    "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 59}],
    "toolUsePromptTokensDetails": [{"modality": "TEXT", "tokenCount": 73}],
    "trafficType": "ON_DEMAND",
}


PASSTHROUGH_OUTPUT_URI = (
    "gs://litellm-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/u/"
    "predictions.jsonl"
)


UNGROUNDED_USAGE_METADATA = {
    "promptTokenCount": 20,
    "candidatesTokenCount": 48,
    "thoughtsTokenCount": 195,
    "toolUsePromptTokenCount": 73,
    "totalTokenCount": 336,
    "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 20}],
    "trafficType": "ON_DEMAND",
}


def _native_vertex_row(usage_metadata: dict, *, grounded: bool, model_version: str | None = "gemini-2.5-flash"):
    candidate = {"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"}
    grounding = {"groundingMetadata": {"webSearchQueries": ["q"]}} if grounded else {}
    response = {"candidates": [{**candidate, **grounding}], "usageMetadata": usage_metadata}
    return {
        "request": {"contents": [{"role": "user", "parts": [{"text": "q"}]}], "tools": [{"googleSearch": {}}]},
        "status": "",
        "response": {**response, **({"modelVersion": model_version} if model_version else {})},
        "processed_time": "2026-09-23T19:02:00.000+00:00",
    }


def _capture_cost_calls(monkeypatch, prompt_cost=0.5, completion_cost=0.25) -> list:
    import litellm.cost_calculator as cc

    calls: list = []

    def _calc(**kw):
        calls.append(kw)
        return (prompt_cost, completion_cost)

    monkeypatch.setattr(cc, "batch_cost_calculator", _calc)
    return calls


def test_vertex_native_cost_bills_embedding_rows(monkeypatch):
    monkeypatch.setitem(litellm.model_cost, "vertex_ai/gemini-embedding-2", {"input_cost_per_token_batches": 1e-7})
    rows = [
        {
            "key": "id_1",
            "status": "",
            "request": {"content": {"parts": [{"text": "hello world"}]}},
            "response": {"embedding": {"values": [0.1, 0.2]}, "usageMetadata": {"promptTokenCount": 2}},
        },
        {
            "key": "id_2",
            "status": "",
            "request": {"content": {"parts": [{"text": "hello"}]}},
            "response": {"embedding": {"values": [0.3]}, "tokenCount": "3"},
        },
        {"key": "id_3", "status": "INVALID_ARGUMENT", "request": {"content": {"parts": [{"text": ""}]}}},
    ]

    result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-embedding-2")

    assert (result.successful_requests, result.failed_requests) == (2, 1)
    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (5, 0, 5)
    assert result.cost == pytest.approx(5 * 1e-7)
    assert result.models == ["gemini-embedding-2"]


@pytest.mark.asyncio
async def test_native_vertex_rows_route_to_vertex_cost_path_without_flag(monkeypatch):
    monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False)
    monkeypatch.setattr(
        bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run")
    )
    calls = _capture_cost_calls(monkeypatch)
    rows = [
        _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True),
        _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False),
    ]

    result = await bu.calculate_batch_cost_and_usage(
        file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash"
    )

    assert result.cost == pytest.approx(1.5)
    assert (result.successful_requests, result.failed_requests) == (2, 0)
    assert result.models == ["gemini-2.5-flash"]
    assert {(call["model"], call["custom_llm_provider"]) for call in calls} == {("gemini-2.5-flash", "vertex_ai")}


@pytest.mark.asyncio
async def test_openai_shaped_vertex_rows_keep_the_generic_path_without_flag(monkeypatch):
    monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False)
    monkeypatch.setattr(
        bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run")
    )
    _capture_cost_calls(monkeypatch)
    rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)]

    result = await bu.calculate_batch_cost_and_usage(
        file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash"
    )

    assert result.successful_requests == 1


@pytest.mark.asyncio
async def test_native_vertex_rows_on_another_provider_keep_the_generic_path(monkeypatch):
    monkeypatch.setattr(
        bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run")
    )
    _capture_cost_calls(monkeypatch)

    result = await bu.calculate_batch_cost_and_usage(
        file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)],
        custom_llm_provider="openai",
    )

    assert result.successful_requests == 0


@pytest.mark.asyncio
async def test_handle_completed_batch_routes_native_rows_without_flag(monkeypatch):
    monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False)
    raw_rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)]

    async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
        return _vertex_jsonl(raw_rows)

    monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
    monkeypatch.setattr(
        bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run")
    )
    calls = _capture_cost_calls(monkeypatch, prompt_cost=0.7, completion_cost=0.3)
    deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6}

    result = await bu._handle_completed_batch(
        _batch(PASSTHROUGH_OUTPUT_URI),
        custom_llm_provider="vertex_ai",
        model_name="gemini-2.5-flash",
        model_info=deployment_model_info,
    )

    assert result.cost == pytest.approx(1.0)
    assert result.usage.total_tokens == 557
    assert [call["model_info"] for call in calls] == [deployment_model_info]


def test_native_vertex_usage_is_billed_like_the_online_path(monkeypatch):
    calls = _capture_cost_calls(monkeypatch)
    grounded = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)
    ungrounded = _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False)

    result = bu.calculate_vertex_ai_batch_cost_and_usage([grounded, ungrounded], "gemini-2.5-flash")

    grounded_usage, ungrounded_usage = (call["usage"] for call in calls)
    assert grounded_usage.prompt_tokens == 19
    assert grounded_usage.completion_tokens == 59 + 406
    assert grounded_usage.completion_tokens_details.reasoning_tokens == 406
    assert ungrounded_usage.prompt_tokens == 20 + 73
    assert ungrounded_usage.completion_tokens == 48 + 195
    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (
        19 + 93,
        465 + 243,
        557 + 336,
    )


def test_native_vertex_rows_are_priced_by_model_version_without_a_model_name(monkeypatch):
    calls = _capture_cost_calls(monkeypatch)
    rows = [
        _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"),
        _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-pro"),
        _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None),
    ]

    result = bu.calculate_vertex_ai_batch_cost_and_usage(rows)

    assert [call["model"] for call in calls] == ["gemini-2.5-flash", "gemini-2.5-pro"]
    assert result.models == ["gemini-2.5-flash", "gemini-2.5-pro"]
    assert result.cost == pytest.approx(1.5)
    assert result.successful_requests == 3
    assert result.usage.total_tokens == 557 + 336 + 336


def test_native_vertex_rows_without_usage_metadata_count_as_failed(monkeypatch):
    _capture_cost_calls(monkeypatch)
    rows = [
        {"request": {"contents": []}, "status": "Error: bad request", "processed_time": "t"},
        {"request": {"contents": []}, "response": {"candidates": []}},
        _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True),
    ]

    result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash")

    assert (result.successful_requests, result.failed_requests) == (1, 2)
    assert result.usage.total_tokens == 557


def test_native_vertex_batch_whose_rows_all_failed_still_names_the_deployment_model(monkeypatch):
    calls = _capture_cost_calls(monkeypatch)
    rows = [{"request": {"contents": []}, "status": "Error: quota exceeded", "processed_time": "t"}] * 2

    result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash")

    assert result.models == ["gemini-2.5-flash"]
    assert (result.successful_requests, result.failed_requests, result.cost) == (0, 2, 0.0)
    assert calls == []


def test_native_vertex_rows_are_priced_with_the_deployment_model_info(monkeypatch):
    calls = _capture_cost_calls(monkeypatch)
    deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6}

    bu.calculate_vertex_ai_batch_cost_and_usage(
        [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)],
        "gemini-2.5-flash",
        model_info=deployment_model_info,
    )

    assert [call["model_info"] for call in calls] == [deployment_model_info]


@pytest.mark.asyncio
async def test_native_vertex_rows_keep_the_deployment_model_info_through_the_batch_entrypoint(monkeypatch):
    calls = _capture_cost_calls(monkeypatch)
    deployment_model_info = {"input_cost_per_token_batches": 1e-6}

    await bu.calculate_batch_cost_and_usage(
        file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)],
        custom_llm_provider="vertex_ai",
        model_name="gemini-2.5-flash",
        model_info=deployment_model_info,
    )

    assert [call["model_info"] for call in calls] == [deployment_model_info]


def test_native_vertex_rows_are_priced_by_the_deployment_model_over_model_version(monkeypatch):
    calls = _capture_cost_calls(monkeypatch)
    rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-pro")]

    result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash")

    assert [call["model"] for call in calls] == ["gemini-2.5-flash"]
    assert result.models == ["gemini-2.5-flash"]


def test_native_vertex_rows_that_fail_response_validation_count_as_failed(monkeypatch):
    calls = _capture_cost_calls(monkeypatch)
    rows = [
        {"request": {"contents": []}, "response": {"candidates": "nope", "usageMetadata": GROUNDED_USAGE_METADATA}},
        _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True),
    ]

    result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash")

    assert (result.successful_requests, result.failed_requests) == (1, 1)
    assert result.usage.total_tokens == 557
    assert len(calls) == 1


@pytest.mark.parametrize("wildcard_model", ["*", "vertex_ai/*"])
def test_native_vertex_rows_under_a_wildcard_deployment_are_priced_by_model_version(monkeypatch, wildcard_model):
    calls = _capture_cost_calls(monkeypatch)
    rows = [
        _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"),
        _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None),
    ]

    result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, wildcard_model)

    assert [call["model"] for call in calls] == ["gemini-2.5-flash", wildcard_model]
    assert result.cost == pytest.approx(1.5)
    assert (result.successful_requests, result.failed_requests) == (2, 0)
    assert result.usage.total_tokens == 557 + 336


def test_native_vertex_row_without_model_version_under_a_wildcard_deployment_bills_its_explicit_prices():
    deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6}
    with_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash")
    without_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version=None)

    twin = bu.calculate_vertex_ai_batch_cost_and_usage([with_version], "vertex_ai/*", model_info=deployment_model_info)
    both = bu.calculate_vertex_ai_batch_cost_and_usage(
        [with_version, without_version], "vertex_ai/*", model_info=deployment_model_info
    )

    assert twin.cost > 0
    assert both.cost == pytest.approx(2 * twin.cost)
    assert (both.successful_requests, both.failed_requests) == (2, 0)


def test_native_vertex_row_the_cost_map_cannot_price_is_billed_at_zero_and_the_rest_still_bills(monkeypatch):
    import litellm.cost_calculator as cc

    def _calc(**kw):
        if kw["model"] == "gemini-unpriced":
            raise ValueError("no pricing")
        return (0.5, 0.25)

    monkeypatch.setattr(cc, "batch_cost_calculator", _calc)
    rows = [
        _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-unpriced"),
        _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-flash"),
    ]

    result = bu.calculate_vertex_ai_batch_cost_and_usage(rows)

    assert result.cost == pytest.approx(0.75)
    assert (result.successful_requests, result.failed_requests) == (2, 0)
    assert result.usage.total_tokens == 557 + 336
    assert result.models == ["gemini-unpriced", "gemini-2.5-flash"]


@pytest.mark.asyncio
async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_known(monkeypatch):
    monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False)
    monkeypatch.setattr(
        bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run")
    )
    calls = _capture_cost_calls(monkeypatch)
    rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)]

    result = await bu.calculate_batch_cost_and_usage(
        file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash"
    )

    assert calls == []
    assert (result.successful_requests, result.failed_requests) == (0, 1)