diff --git a/litellm/integrations/focus/destinations/vantage_destination.py b/litellm/integrations/focus/destinations/vantage_destination.py index 9e6028900f9..e860af3726b 100644 --- a/litellm/integrations/focus/destinations/vantage_destination.py +++ b/litellm/integrations/focus/destinations/vantage_destination.py @@ -9,6 +9,11 @@ from typing import Any, Optional import httpx from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, + httpxSpecialProvider, +) from .base import FocusDestination, FocusTimeWindow @@ -131,45 +136,38 @@ class FocusVantageDestination(FocusDestination): # rejection (e.g. InvoiceIssuerName, ProviderName, PublisherName). content = _strip_unsupported_columns(content) - # Reuse a single HTTP client for the entire deliver() call - async with httpx.AsyncClient(timeout=60.0) as client: - # Check both size and row-count limits before single-shot upload - lines = content.split(b"\n") - data_line_count = sum(1 for line in lines[1:] if line.strip()) - within_limits = ( - len(content) <= VANTAGE_MAX_BYTES_PER_UPLOAD - and data_line_count <= VANTAGE_MAX_ROWS_PER_UPLOAD - ) - if within_limits: - await self._upload_csv(client, content, filename) - return + client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback, + params={"timeout": 60.0}, + ) - # Otherwise split into batches respecting both limits - await self._upload_batched(client, content, filename) + # Check both size and row-count limits before single-shot upload + lines = content.split(b"\n") + data_line_count = sum(1 for line in lines[1:] if line.strip()) + within_limits = ( + len(content) <= VANTAGE_MAX_BYTES_PER_UPLOAD + and data_line_count <= VANTAGE_MAX_ROWS_PER_UPLOAD + ) + if within_limits: + await self._upload_csv(client, content, filename) + return + + # Otherwise split into batches respecting both limits + await self._upload_batched(client, content, filename) async def _upload_csv( - self, client: httpx.AsyncClient, csv_bytes: bytes, filename: str + self, client: AsyncHTTPHandler, csv_bytes: bytes, filename: str ) -> None: url = f"{self.base_url}/v2/integrations/" f"{self.integration_token}/costs.csv" headers = { "Authorization": f"Bearer {self.api_key}", } - response = await client.post( + await client.post( url, headers=headers, files={"csv": (filename, csv_bytes, "text/csv")}, ) - try: - response.raise_for_status() - except httpx.HTTPStatusError as e: - verbose_logger.error( - "Vantage destination: upload failed for %s — %s — response body: %s", - filename, - e, - response.text, - ) - raise verbose_logger.debug( "Vantage destination: uploaded %d bytes (%s)", @@ -178,7 +176,7 @@ class FocusVantageDestination(FocusDestination): ) async def _upload_batched( - self, client: httpx.AsyncClient, csv_bytes: bytes, filename: str + self, client: AsyncHTTPHandler, csv_bytes: bytes, filename: str ) -> None: """Split the CSV into batches and upload each. @@ -217,7 +215,7 @@ class FocusVantageDestination(FocusDestination): async def _upload_size_limited( self, - client: httpx.AsyncClient, + client: AsyncHTTPHandler, header: bytes, data_lines: list[bytes], filename: str, diff --git a/tests/test_litellm/integrations/focus/test_vantage_destination.py b/tests/test_litellm/integrations/focus/test_vantage_destination.py index 998b447aeb8..10f72399193 100644 --- a/tests/test_litellm/integrations/focus/test_vantage_destination.py +++ b/tests/test_litellm/integrations/focus/test_vantage_destination.py @@ -4,7 +4,7 @@ from __future__ import annotations from datetime import datetime, timedelta, timezone from typing import Any, Dict, List -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -15,6 +15,8 @@ from litellm.integrations.focus.destinations.vantage_destination import ( VANTAGE_MAX_ROWS_PER_UPLOAD, ) +MOCK_TARGET = "litellm.integrations.focus.destinations.vantage_destination.get_async_httpx_client" + def _window(freq: str = "hourly", hour: int = 5) -> FocusTimeWindow: start = datetime(2024, 1, 2, hour, tzinfo=timezone.utc) @@ -72,17 +74,14 @@ async def test_should_skip_empty_content(): @pytest.mark.asyncio async def test_should_upload_csv_to_correct_url(): dest = FocusVantageDestination(prefix="exports", config=_config()) - captured: Dict[str, Any] = {} mock_response = AsyncMock() mock_response.raise_for_status = lambda: None - mock_client = AsyncMock() + mock_client = MagicMock() mock_client.post = AsyncMock(return_value=mock_response) - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=False) - with patch("litellm.integrations.focus.destinations.vantage_destination.httpx.AsyncClient", return_value=mock_client): + with patch(MOCK_TARGET, return_value=mock_client): await dest.deliver( content=b"header\nrow1\n", time_window=_window(), @@ -114,9 +113,7 @@ async def test_should_batch_large_content(): mock_response = AsyncMock() mock_response.raise_for_status = lambda: None - mock_client = AsyncMock() - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=False) + mock_client = MagicMock() async def capture_post(url, **kwargs): files = kwargs.get("files", {}) @@ -126,7 +123,7 @@ async def test_should_batch_large_content(): mock_client.post = capture_post - with patch("litellm.integrations.focus.destinations.vantage_destination.httpx.AsyncClient", return_value=mock_client): + with patch(MOCK_TARGET, return_value=mock_client): await dest.deliver( content=large_content, time_window=_window(), @@ -160,9 +157,7 @@ async def test_should_batch_by_row_count(): mock_response = AsyncMock() mock_response.raise_for_status = lambda: None - mock_client = AsyncMock() - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=False) + mock_client = MagicMock() async def capture_post(url, **kwargs): files = kwargs.get("files", {}) @@ -172,7 +167,7 @@ async def test_should_batch_by_row_count(): mock_client.post = capture_post - with patch("litellm.integrations.focus.destinations.vantage_destination.httpx.AsyncClient", return_value=mock_client): + with patch(MOCK_TARGET, return_value=mock_client): await dest.deliver( content=content, time_window=_window(),