diff --git a/litellm/proxy/roi_calculator/sync.py b/litellm/proxy/roi_calculator/sync.py index 7bb47ec8e88..6f00bf770ae 100644 --- a/litellm/proxy/roi_calculator/sync.py +++ b/litellm/proxy/roi_calculator/sync.py @@ -313,6 +313,15 @@ def _processed_records(processed: tuple[_ProcessedPull, ...]) -> Mapping[int, RO return MappingProxyType({item.position: item.record for item in processed}) +async def _cache_estimated_pull(repository: _ReportRepository, key: str | None, record: ROIPullRecord) -> None: + if key is None or record["estimate"]["status"] != "estimated": + return + await repository.set_param( + "roi_calculator_pull_" + key, + _JSON_OBJECT_ADAPTER.validate_python(TypeAdapter(ROIPullRecord).dump_python(record, mode="json")), + ) + + class SyncManager: def __init__( self, @@ -497,6 +506,7 @@ class SyncManager: } ) ) + await _cache_estimated_pull(repository, key, cached_record) self._update_estimate_progress(cached_record["estimate"]) return _ProcessedPull(index, cached_record) try: @@ -521,13 +531,7 @@ class SyncManager: ) fetched_key: Final = cache_key(settings, context, repo, evidence_item) record: Final = self._report_record(evidence, estimate, fetched_key) - if fetched_key is not None and estimate["status"] == "estimated": - await repository.set_param( - "roi_calculator_pull_" + fetched_key, - _JSON_OBJECT_ADAPTER.validate_python( - TypeAdapter(ROIPullRecord).dump_python(record, mode="json") - ), - ) + await _cache_estimated_pull(repository, fetched_key, record) self._update_estimate_progress(estimate) return _ProcessedPull(index, record) diff --git a/tests/unit/proxy/roi_calculator/test_sync.py b/tests/unit/proxy/roi_calculator/test_sync.py index 3547b503484..9cbcaca4338 100644 --- a/tests/unit/proxy/roi_calculator/test_sync.py +++ b/tests/unit/proxy/roi_calculator/test_sync.py @@ -572,6 +572,23 @@ async def test_reused_profile_preserves_email_only_when_lookup_fails(profile_sta assert report["pulls"][0]["profile_email"] == expected assert report["pulls"][0]["emails"] == ((expected,) if expected else ()) assert summarize(report, MappingProxyType({}))["metrics"]["cost_per_hour"] == (None if profile_status == 200 else 3) + repository.values = MappingProxyType( + {key: value for key, value in repository.values.items() if key != "roi_calculator_report"} + ) + + def unavailable_profile(request: httpx.Request) -> httpx.Response: + if request.url.path == "/users/alice": + return httpx.Response(503) + return baseline.handle_request(request) + + restarted: Final = SyncManager(clock=_fixed_now) + assert await restarted.start( + _settings(), repository, _spend_reader(), unexpected_completion, httpx.MockTransport(unavailable_profile) + ) + await _wait_until_finished(restarted) + subsequent: Final = TypeAdapter(ROIReport).validate_python(repository.values["roi_calculator_report"]) + assert subsequent["pulls"][0]["profile_email"] == expected + assert subsequent["pulls"][0]["emails"] == ((expected,) if expected else ()) @pytest.mark.asyncio