From cb8fcaeee161563aa8938c386153d47fb9b5b081 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 17:53:11 -0700 Subject: [PATCH] fix(accounting): isolate fallback pricing by attempt --- .circleci/scripts/run_integration.sh | 1 - .github/workflows/test-integration.yml | 14 -- codecov.yaml | 2 - litellm/litellm_core_utils/litellm_logging.py | 5 +- litellm/main.py | 54 +++---- .../database/test_harness_lifecycle.py | 12 +- .../test_sagemaker_completion_handler.py | 13 +- .../spend_tracking/test_cost_compiler.py | 41 ++++++ .../test_register_model_custom_pricing.py | 138 ++++++++++++++++-- 9 files changed, 200 insertions(+), 80 deletions(-) diff --git a/.circleci/scripts/run_integration.sh b/.circleci/scripts/run_integration.sh index 211572dd77f..1380b07de53 100644 --- a/.circleci/scripts/run_integration.sh +++ b/.circleci/scripts/run_integration.sh @@ -256,7 +256,6 @@ fi env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \ INTEGRATION_RUN_ID="$integration_identity" \ - PYTEST_ADDOPTS="${PYTEST_ADDOPTS:-}" \ DATABASE_URL="$DATABASE_URL" \ "${shared_services[@]}" "${worker_coverage[@]}" \ INTEGRATION_WORKERS="${INTEGRATION_WORKERS:-1}" \ diff --git a/.github/workflows/test-integration.yml b/.github/workflows/test-integration.yml index 9352482adc7..73159bd64d2 100644 --- a/.github/workflows/test-integration.yml +++ b/.github/workflows/test-integration.yml @@ -43,9 +43,6 @@ jobs: name: integration ${{ matrix.suite }} if: github.event_name != 'pull_request' || github.event.label.name == 'run-ci' environment: integration-license - permissions: - contents: read - id-token: write runs-on: ${{ fromJSON('{"providers":"ubuntu-latest-64-cores","extensions":"ubuntu-latest-64-cores","accounting":"ubuntu-latest-32-cores","management":"ubuntu-latest-32-cores","management-replica":"ubuntu-latest-32-cores"}')[matrix.suite] || 'ubuntu-latest-16-cores' }} timeout-minutes: 60 strategy: @@ -153,7 +150,6 @@ jobs: timeout-minutes: 45 env: SUITE: ${{ matrix.suite }} - PYTEST_ADDOPTS: ${{ matrix.suite == 'database' && '--cov=litellm.proxy.db.accounting_ledger --cov=litellm.proxy.spend_tracking.cost_compiler --cov-branch --cov-report=xml:test-results/accounting-coverage.xml' || '' }} INTEGRATION_POSTGRES_CONTAINER: ${{ job.services.postgres.id }} INTEGRATION_REDIS_CONTAINER: ${{ job.services.redis.id }} LITELLM_LICENSE: ${{ secrets.LITELLM_LICENSE }} @@ -163,16 +159,6 @@ jobs: *) bash .circleci/scripts/run_integration.sh "$SUITE" ;; esac - - name: Upload accounting ledger integration coverage - if: ${{ !cancelled() && matrix.suite == 'database' && hashFiles('test-results/accounting-coverage.xml') != '' }} - uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5.5.5 - with: - use_oidc: true - files: test-results/accounting-coverage.xml - disable_search: true - flags: accounting-ledger - fail_ci_if_error: true - - name: Verify every case ran exactly once if: ${{ !cancelled() && matrix.suite != 'browser' }} timeout-minutes: 2 diff --git a/codecov.yaml b/codecov.yaml index d777602525d..521088868bd 100644 --- a/codecov.yaml +++ b/codecov.yaml @@ -20,8 +20,6 @@ flag_management: # the stale line maps mark comment lines of since-edited files as missed, # sinking patch coverage on unrelated PRs. individual_flags: - - name: accounting-ledger - carryforward: false - name: proxy-mgmt-behavior carryforward: false - name: security diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 30061050f66..41d95903785 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2154,11 +2154,14 @@ class Logging(LiteLLMLoggingBaseClass): self, model: str, provider: str, - overrides: Mapping[str, object], + overrides: Mapping[str, object] | None, deployment_id: str | None = None, ) -> None: from litellm.utils import model_info_with_pricing_overrides + if overrides is None: + self.request_pricing = None + return base: Final = self._request_pricing_model_info(model, provider, overrides, deployment_id) normalized: Final = _REQUEST_MODEL_INFO.validate_python( { diff --git a/litellm/main.py b/litellm/main.py index 170a4cd3551..d6706cdf99d 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1480,10 +1480,16 @@ def _set_request_pricing( ) -> None: if not isinstance(logging_obj, LiteLLMLoggingObj): return + custom_pricing: Final = ( + (kwargs.get("input_cost_per_token") is not None and kwargs.get("output_cost_per_token") is not None) + or kwargs.get("input_cost_per_second") is not None + or kwargs.get("output_cost_per_second") is not None + or kwargs.get("cost_per_second") is not None + ) logging_obj.set_request_pricing( model, custom_llm_provider, - _build_custom_pricing_entry(custom_llm_provider, kwargs, model_info), + _build_custom_pricing_entry(custom_llm_provider, kwargs, model_info) if custom_pricing else None, _get_router_deployment_id(kwargs), ) @@ -5942,19 +5948,13 @@ def completion( "aws_region_name", None ) # support region-based pricing for bedrock - if ( - (input_cost_per_token is not None and output_cost_per_token is not None) - or input_cost_per_second is not None - or output_cost_per_second is not None - or cost_per_second is not None - ): - _set_request_pricing( - logging_obj=litellm_logging_obj, - model=model, - custom_llm_provider=custom_llm_provider, - kwargs=kwargs, - model_info=model_info, - ) + _set_request_pricing( + logging_obj=litellm_logging_obj, + model=model, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + model_info=model_info, + ) ### BUILD CUSTOM PROMPT TEMPLATE -- IF GIVEN ### custom_prompt_dict = {} if initial_prompt_value or roles or final_prompt_value or bos_token or eos_token: @@ -6961,12 +6961,6 @@ def embedding( headers.update(proxy_headers) except Exception as e: verbose_logger.warning("Failed to get proxy auth headers: %s", e) - ### CUSTOM MODEL COST ### - input_cost_per_token: Final = kwargs.get("input_cost_per_token", None) - output_cost_per_token: Final = kwargs.get("output_cost_per_token", None) - cost_per_second: Final = kwargs.get("cost_per_second", None) - input_cost_per_second: Final = kwargs.get("input_cost_per_second", None) - output_cost_per_second: Final = kwargs.get("output_cost_per_second", None) openai_params: Final = [ "user", "dimensions", @@ -7013,19 +7007,13 @@ def embedding( **non_default_params, ) - if ( - (input_cost_per_token is not None and output_cost_per_token is not None) - or input_cost_per_second is not None - or output_cost_per_second is not None - or cost_per_second is not None - ): - _set_request_pricing( - logging_obj=litellm_logging_obj, - model=model, - custom_llm_provider=custom_llm_provider, - kwargs=kwargs, - model_info=kwargs.get("model_info"), - ) + _set_request_pricing( + logging_obj=litellm_logging_obj, + model=model, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + model_info=kwargs.get("model_info"), + ) logging: Final[LiteLLMLoggingObj] = litellm_logging_obj logging.update_environment_variables( diff --git a/tests/integration/database/test_harness_lifecycle.py b/tests/integration/database/test_harness_lifecycle.py index f18cc9cf22c..d600a92c5a2 100644 --- a/tests/integration/database/test_harness_lifecycle.py +++ b/tests/integration/database/test_harness_lifecycle.py @@ -81,7 +81,6 @@ def _nested_pytest_environment(environment: Mapping[str, str]) -> dict[str, str] "INTEGRATION_CASE_SHARD_INDEX", "CIRCLE_NODE_TOTAL", "CIRCLE_NODE_INDEX", - "PYTEST_ADDOPTS", "PYTEST_XDIST_WORKER", "PYTEST_XDIST_WORKER_COUNT", "PYTEST_XDIST_TESTRUNUID", @@ -108,16 +107,8 @@ def test_runner_environment_drops_inherited_provider_api_keys() -> None: @pytest.mark.parametrize("fails", (False, True), ids=("passing-test", "failing-test")) -def test_sdk_teardown_delivers_logging_before_restoring_callbacks( - tmp_path: Path, fails: bool, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_sdk_teardown_delivers_logging_before_restoring_callbacks(tmp_path: Path, fails: bool) -> None: root: Final = Path(__file__).resolve().parents[3] - coverage_report: Final = tmp_path / "outer-coverage.xml" - coverage_report.write_text("parent coverage report") - monkeypatch.setenv( - "PYTEST_ADDOPTS", - f"--cov=litellm.proxy.spend_tracking.cost_compiler --cov-report=xml:{coverage_report}", - ) events: Final = tmp_path / "callbacks.json" probe: Final = tmp_path / "test_callback_teardown.py" probe.write_text( @@ -184,7 +175,6 @@ def test_sdk_teardown_delivers_logging_before_restoring_callbacks( check=False, ) assert result.returncode == (1 if fails else 0), result.stdout + result.stderr - assert coverage_report.read_text() == "parent coverage report", "nested pytest overwrote parent coverage" cases: Final = ElementTree.parse(tmp_path / "junit.xml").findall(".//testcase") assert tuple(case.attrib["name"] for case in cases) == ("test_request", "test_next") assert all(case.find("error") is None for case in cases), result.stdout + result.stderr diff --git a/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py b/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py index 496eb5fe25b..4c253a10116 100644 --- a/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py +++ b/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py @@ -298,13 +298,8 @@ def test_installed_botocore_signs_the_completion_request(): from litellm.llms.sagemaker.completion.handler import SagemakerLLM request = SagemakerLLM()._prepare_request( - credentials=Credentials("test-key", "test-secret"), - model="test-endpoint", - data={"inputs": "ping"}, - messages=[], - litellm_params={}, - optional_params={}, - aws_region_name="us-west-2", + credentials=Credentials("test-key", "test-secret"), model="test-endpoint", data={"inputs": "ping"}, + messages=[], litellm_params={}, optional_params={}, aws_region_name="us-west-2", ) assert request.body == b'{"inputs": "ping"}' assert "/us-west-2/sagemaker/aws4_request" in request.headers["Authorization"] @@ -452,7 +447,9 @@ def test_completion_sagemaker_prompt_template_non_stream( }, ) ) - url: Final = "https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/deepseek_coder_6.7_instruct/invocations" + url: Final = ( + "https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/deepseek_coder_6.7_instruct/invocations" + ) route: Final = respx_mock.post(url).mock( return_value=httpx.Response(200, json={"generated_text": "SageMaker reply"}) ) diff --git a/tests/unit/proxy/spend_tracking/test_cost_compiler.py b/tests/unit/proxy/spend_tracking/test_cost_compiler.py index b898012c57b..077a7b8b379 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_compiler.py +++ b/tests/unit/proxy/spend_tracking/test_cost_compiler.py @@ -1726,3 +1726,44 @@ async def test_image_capture_replays_the_existing_image_calculator( assert float(replay.total_cost) == pytest.approx(source) finally: await meter.stop(1) + + +@pytest.mark.asyncio +async def test_delayed_attempt_capture_keeps_its_prices_after_fallback_reset(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm + from litellm.main import _set_request_pricing + from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage + from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings + from litellm.types.utils import ModelResponse, Usage + from tests.unit.proxy.spend_tracking.test_usage_capture import StoredUsage + from tests.unit.proxy.spend_tracking.test_usage_meter import _manual_flush + + catalog: Final = { + "openai/capture-model": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + } + } + monkeypatch.setattr(litellm, "model_cost", catalog) + monkeypatch.setattr(litellm, "cost_discount_config", {}) + monkeypatch.setattr(litellm, "cost_margin_config", {}) + logger: Final = _capture_logger() + failed_attempt: Final = logger.attempt_scoped_copy() + _set_request_pricing(logger, "capture-model", "openai", {}, None) + response: Final = ModelResponse(model="capture-model", usage=Usage(prompt_tokens=10, completion_tokens=4)) + for attempt, expected in ((logger, 0.018), (failed_attempt, 0.18)): + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + try: + assert attempt.response_cost_calculator(response, ignore_stamped_cost=True) == pytest.approx(expected) + capture_logging_usage(meter, attempt, response, _AT, _AT) + await meter.close() + with monkeypatch.context() as changed_catalog: + changed_catalog.setattr(litellm, "model_cost", {}) + replay: Final = price_usage(database.rows[-1].envelope) + assert isinstance(replay, PricingSnapshot) + assert float(replay.total_cost) == pytest.approx(expected) + finally: + await meter.stop(1) diff --git a/tests/unit/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py index 5d77492ee58..d971f0d7014 100644 --- a/tests/unit/test_register_model_custom_pricing.py +++ b/tests/unit/test_register_model_custom_pricing.py @@ -227,7 +227,10 @@ def test_register_model_strips_none_litellm_provider(): registered = litellm.model_cost.get(model_key) assert registered is not None, f"{model_key} should be in model_cost" # The key may be absent entirely, but if present it must not be None. - assert "litellm_provider" not in registered or registered["litellm_provider"] is not None + assert ( + "litellm_provider" not in registered + or registered["litellm_provider"] is not None + ) # Downstream consumers must accept this entry for any provider, # mirroring what the cost calculator does. assert _check_provider_match(registered, "openai") is True @@ -285,7 +288,9 @@ def test_register_model_strips_none_litellm_provider_from_get_model_info(monkeyp # The strip must have removed the None-valued provider that # ``get_model_info`` returned. The key may be absent entirely, but # it must never be present with value ``None``. - assert "litellm_provider" not in registered or (registered["litellm_provider"] is not None), ( + assert "litellm_provider" not in registered or ( + registered["litellm_provider"] is not None + ), ( "register_model failed to strip litellm_provider=None returned " f"by get_model_info, got {registered.get('litellm_provider')!r}" ) @@ -346,8 +351,14 @@ def test_register_model_inherits_builtin_cache_pricing_for_unmapped_key(monkeypa ) registered = litellm.model_cost[registered_key] - assert registered.get("cache_creation_input_token_cost") == builtin["cache_creation_input_token_cost"] - assert registered.get("cache_read_input_token_cost") == builtin["cache_read_input_token_cost"] + assert ( + registered.get("cache_creation_input_token_cost") + == builtin["cache_creation_input_token_cost"] + ) + assert ( + registered.get("cache_read_input_token_cost") + == builtin["cache_read_input_token_cost"] + ) assert registered["litellm_provider"] == "bedrock" usage = Usage( @@ -411,7 +422,8 @@ def test_register_model_warns_when_no_builtin_match_for_cache_pricing(caplog): ) assert any( - registered_key in record.message and "cache_creation_input_token_cost" in record.message + registered_key in record.message + and "cache_creation_input_token_cost" in record.message for record in caplog.records ), "expected a warning naming the unmapped key and the cache cost fields" finally: @@ -598,14 +610,17 @@ def test_register_model_router_add_deployment_custom_pricing_applies(): # the deployment's provider. from litellm.utils import _check_provider_match - registered_keys = [k for k in (deployment_model, model_key) if k in litellm.model_cost] + registered_keys = [ + k for k in (deployment_model, model_key) if k in litellm.model_cost + ] assert registered_keys, ( - f"Router.add_deployment did not register custom pricing for {model_key} / {deployment_model}" + "Router.add_deployment did not register custom pricing for " + f"{model_key} / {deployment_model}" ) for k in registered_keys: - assert _check_provider_match(litellm.model_cost[k], "openai") is True, ( - f"custom pricing for {k} was dropped by _check_provider_match" - ) + assert ( + _check_provider_match(litellm.model_cost[k], "openai") is True + ), f"custom pricing for {k} was dropped by _check_provider_match" finally: litellm.model_cost.pop(model_key, None) litellm.model_cost.pop(deployment_model, None) @@ -1870,3 +1885,106 @@ def test_sdk_followup_prices_change_only_after_explicit_registration(call_type: finally: _restore_model_cost_entries(original) litellm.open_ai_chat_completion_models.intersection_update(models_before) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ("completion", "acompletion", "embedding", "aembedding")) +@pytest.mark.parametrize("custom_primary", (True, False)) +async def test_router_fallback_prices_belong_to_the_served_attempt( + monkeypatch: pytest.MonkeyPatch, call_type: str, custom_primary: bool +) -> None: + import httpx + import respx + from datetime import datetime, timezone + + model: Final = "openai/fallback-price-regression" + embedding: Final = "embedding" in call_type + keys: Final = (model, model.removeprefix("openai/"), "price-primary", "price-fallback") + original: Final = _snapshot_model_cost_entries(keys) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "cost_discount_config", {}) + monkeypatch.setattr(litellm, "cost_margin_config", {}) + try: + litellm.register_model( + { + model: { + "litellm_provider": "openai", + "mode": "embedding" if embedding else "chat", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + } + }, + persist_across_reloads=False, + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": name, + "litellm_params": { + "model": model, + "api_key": "fake-key", + "api_base": f"https://{name}.example/v1", + **({"input_cost_per_token": rate, "output_cost_per_token": rate} if rate is not None else {}), + }, + "model_info": {"id": f"price-{name}"}, + } + for name, rate in ( + ("primary", 0.0 if custom_primary else None), + ("fallback", None if custom_primary else 0.003), + ) + ], + fallbacks=[{"primary": ["fallback"]}], + num_retries=0, + ) + registered: Final = _snapshot_model_cost_entries(keys) + logger: Final = Logging( + model=model, + messages=[], + stream=False, + call_type=call_type, + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id="fallback-pricing", + function_id="fallback-pricing", + ) + endpoint: Final = "embeddings" if embedding else "chat/completions" + body: Final = ( + { + "object": "list", + "model": model.removeprefix("openai/"), + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "usage": {"prompt_tokens": 10, "total_tokens": 10}, + } + if embedding + else { + "id": "fallback-result", + "object": "chat.completion", + "created": 1, + "model": model.removeprefix("openai/"), + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + ) + with respx.mock(assert_all_called=True) as wire: + primary: Final = wire.post(f"https://primary.example/v1/{endpoint}").mock( + return_value=httpx.Response(429, json={"error": {"message": "rate limit", "type": "rate_limit_error"}}) + ) + fallback: Final = wire.post(f"https://fallback.example/v1/{endpoint}").mock( + return_value=httpx.Response(200, json=body) + ) + request: Final = {"model": "primary", "litellm_logging_obj": logger} + if call_type == "acompletion": + response = await router.acompletion(**request, messages=[{"role": "user", "content": "probe"}]) + elif call_type == "completion": + response = router.completion(**request, messages=[{"role": "user", "content": "probe"}]) + elif call_type == "aembedding": + response = await router.aembedding(**request, input=["probe"]) + else: + response = router.embedding(**request, input=["probe"]) + expected: Final = (0.01 if embedding else 0.02) if custom_primary else (0.03 if embedding else 0.045) + assert logger.response_cost_calculator(response, ignore_stamped_cost=True) == pytest.approx(expected) + assert response._hidden_params["response_cost"] == pytest.approx(expected) + assert primary.called + assert fallback.call_count == 1 + assert _snapshot_model_cost_entries(keys) == registered + finally: + _restore_model_cost_entries(original)