scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(roi): persist refreshed identities for reused estimates

Browse source
This commit is contained in:
moe-berri 2026-09-30 13:04:51 -07:00
parent 459d6ff787
commit deb9d2d5e5
2 changed files with 28 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
litellm/proxy/roi_calculator/sync.py
View file

@ -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)

17
tests/unit/proxy/roi_calculator/test_sync.py
View file

@ -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