mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
Fix data truncation, sub-batch resilience, and env var safety
- Split VantageExportRequest (limit=None) and VantageDryRunRequest (limit=500) so actual exports don't silently truncate large datasets - Add try/except around each sub-batch upload in _upload_size_limited, consistent with _upload_batched's continue-on-failure guarantee - Guard VANTAGE_EXPORT_INTERVAL_SECONDS against non-numeric values with try/except instead of bare int() cast Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
25c8658761
commit
d35abfb55f
4 changed files with 47 additions and 8 deletions
|
|
@ -145,12 +145,15 @@ class FocusVantageDestination(FocusDestination):
|
|||
"""Upload lines in chunks that stay under the 2 MB size limit.
|
||||
|
||||
Individual rows that exceed the limit on their own are skipped with
|
||||
a warning — they cannot be split further.
|
||||
a warning — they cannot be split further. Sub-batch failures are
|
||||
recorded and the first error is re-raised after all sub-batches have
|
||||
been attempted, consistent with ``_upload_batched``.
|
||||
"""
|
||||
current_chunk: list[bytes] = []
|
||||
current_size = len(header) + 1 # header + newline
|
||||
sub_batch = 0
|
||||
header_size = len(header) + 1
|
||||
first_error: Optional[Exception] = None
|
||||
|
||||
for line in data_lines:
|
||||
line_size = len(line) + 1 # line + newline
|
||||
|
|
@ -166,7 +169,15 @@ class FocusVantageDestination(FocusDestination):
|
|||
if current_size + line_size > VANTAGE_MAX_BYTES_PER_UPLOAD and current_chunk:
|
||||
batch_csv = header + b"\n" + b"\n".join(current_chunk) + b"\n"
|
||||
batch_filename = f"{filename}.part{batch_offset}_{sub_batch}"
|
||||
await self._upload_csv(client, batch_csv, batch_filename)
|
||||
try:
|
||||
await self._upload_csv(client, batch_csv, batch_filename)
|
||||
except Exception as e:
|
||||
verbose_logger.error(
|
||||
"Vantage destination: sub-batch %s failed: %s",
|
||||
batch_filename, e,
|
||||
)
|
||||
if first_error is None:
|
||||
first_error = e
|
||||
current_chunk = []
|
||||
current_size = header_size
|
||||
sub_batch += 1
|
||||
|
|
@ -176,4 +187,15 @@ class FocusVantageDestination(FocusDestination):
|
|||
if current_chunk:
|
||||
batch_csv = header + b"\n" + b"\n".join(current_chunk) + b"\n"
|
||||
batch_filename = f"{filename}.part{batch_offset}_{sub_batch}"
|
||||
await self._upload_csv(client, batch_csv, batch_filename)
|
||||
try:
|
||||
await self._upload_csv(client, batch_csv, batch_filename)
|
||||
except Exception as e:
|
||||
verbose_logger.error(
|
||||
"Vantage destination: sub-batch %s failed: %s",
|
||||
batch_filename, e,
|
||||
)
|
||||
if first_error is None:
|
||||
first_error = e
|
||||
|
||||
if first_error is not None:
|
||||
raise first_error
|
||||
|
|
|
|||
|
|
@ -53,7 +53,15 @@ class VantageLogger(FocusLogger):
|
|||
).lower()
|
||||
|
||||
raw_interval = interval_seconds or os.getenv("VANTAGE_EXPORT_INTERVAL_SECONDS")
|
||||
resolved_interval = int(raw_interval) if raw_interval is not None else None
|
||||
resolved_interval: Optional[int] = None
|
||||
if raw_interval is not None:
|
||||
try:
|
||||
resolved_interval = int(raw_interval)
|
||||
except (ValueError, TypeError):
|
||||
verbose_logger.warning(
|
||||
"Invalid VANTAGE_EXPORT_INTERVAL_SECONDS value: %s, ignoring",
|
||||
raw_interval,
|
||||
)
|
||||
|
||||
destination_config: Dict[str, Any] = {}
|
||||
if resolved_api_key:
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
|||
encrypt_value_helper,
|
||||
)
|
||||
from litellm.types.proxy.vantage_endpoints import (
|
||||
VantageDryRunRequest,
|
||||
VantageExportRequest,
|
||||
VantageExportResponse,
|
||||
VantageInitRequest,
|
||||
|
|
@ -356,7 +357,7 @@ async def init_vantage_settings(
|
|||
response_model=VantageExportResponse,
|
||||
)
|
||||
async def vantage_dry_run_export(
|
||||
request: VantageExportRequest,
|
||||
request: VantageDryRunRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""
|
||||
|
|
@ -365,7 +366,7 @@ async def vantage_dry_run_export(
|
|||
Returns the data that would be exported without actually sending it to Vantage.
|
||||
|
||||
Parameters:
|
||||
- limit: Optional limit on number of records to process (default: 500)
|
||||
- limit: Limit on number of records to preview (default: 500)
|
||||
|
||||
Only admin users can perform Vantage exports.
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -29,10 +29,10 @@ class VantageInitResponse(BaseModel):
|
|||
|
||||
|
||||
class VantageExportRequest(BaseModel):
|
||||
"""Request model for Vantage export operations"""
|
||||
"""Request model for Vantage export operations (actual export, no default limit)"""
|
||||
|
||||
limit: Optional[int] = Field(
|
||||
500, description="Limit on number of records to export (default: 500)"
|
||||
None, description="Optional limit on number of records to export (default: no limit)"
|
||||
)
|
||||
start_time_utc: Optional[datetime] = Field(
|
||||
None, description="Start time for data export in UTC"
|
||||
|
|
@ -42,6 +42,14 @@ class VantageExportRequest(BaseModel):
|
|||
)
|
||||
|
||||
|
||||
class VantageDryRunRequest(BaseModel):
|
||||
"""Request model for Vantage dry-run operations (capped for preview)"""
|
||||
|
||||
limit: Optional[int] = Field(
|
||||
500, description="Limit on number of records to preview (default: 500)"
|
||||
)
|
||||
|
||||
|
||||
class VantageExportResponse(BaseModel):
|
||||
"""Response model for Vantage export operations"""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue