mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
Type the batch-retire helpers and rename test helper to avoid shadowing existing _completed_batch
This commit is contained in:
parent
f8b31f493a
commit
67d16a499d
2 changed files with 8 additions and 6 deletions
|
|
@ -1288,7 +1288,7 @@ def batch_cost_poller_is_active() -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _completed_batch_safe_to_retire(response) -> bool:
|
||||
def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool:
|
||||
"""Whether a "completed" batch may be retired from cost recovery.
|
||||
|
||||
``batch_processed=True`` is the sole re-pickup gate for CheckBatchCost's
|
||||
|
|
|
|||
|
|
@ -438,7 +438,9 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
|
|||
)
|
||||
|
||||
|
||||
def _completed_batch(output_file_id, completed=None) -> LiteLLMBatch:
|
||||
def _completed_batch_for_retire(
|
||||
output_file_id: str | None, completed: int | None = None
|
||||
) -> LiteLLMBatch:
|
||||
kwargs = dict(
|
||||
id="batch-1",
|
||||
completion_window="24h",
|
||||
|
|
@ -460,16 +462,16 @@ class TestCompletedBatchSafeToRetire:
|
|||
file has arrived or the provider proves no successful lines (#37713)."""
|
||||
|
||||
def test_output_file_present_is_safe(self):
|
||||
assert _completed_batch_safe_to_retire(_completed_batch("file-out")) is True
|
||||
assert _completed_batch_safe_to_retire(_completed_batch_for_retire("file-out")) is True
|
||||
|
||||
def test_no_output_and_no_successful_lines_is_safe(self):
|
||||
# Every request line errored -> nothing left to recover.
|
||||
assert _completed_batch_safe_to_retire(_completed_batch(None, completed=0)) is True
|
||||
assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, completed=0)) is True
|
||||
|
||||
def test_no_output_but_successful_lines_is_not_safe(self):
|
||||
# The bug: output_file_id is lagging; retiring here loses the spend record.
|
||||
assert _completed_batch_safe_to_retire(_completed_batch(None, completed=5)) is False
|
||||
assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, completed=5)) is False
|
||||
|
||||
def test_no_output_and_unknown_counts_is_not_safe(self):
|
||||
# Counts unknown -> stay eligible so the next poller pass revisits it.
|
||||
assert _completed_batch_safe_to_retire(_completed_batch(None)) is False
|
||||
assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None)) is False
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue