From d180db31e7c3b9fa959a879273c4376f67d32817 Mon Sep 17 00:00:00 2001 From: Chesars Date: Wed, 17 Dec 2025 00:15:54 -0300 Subject: [PATCH] Use _get_httpx_client for HTTP polling in BFL image edit Replace direct httpx.get() calls with _get_httpx_client() to reuse cached HTTP client, following the pattern used by other providers (RunwayML, Azure AI OCR, Sagemaker, etc.). --- .../image_edit/transformation.py | 5 +-- .../test_bfl_image_edit_transformation.py | 35 +++++++++++++++---- 2 files changed, 31 insertions(+), 9 deletions(-) diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py index ad965263b74..65d83159f75 100644 --- a/litellm/llms/black_forest_labs/image_edit/transformation.py +++ b/litellm/llms/black_forest_labs/image_edit/transformation.py @@ -15,6 +15,7 @@ import httpx from httpx._types import RequestFiles from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from litellm.llms.custom_httpx.http_handler import _get_httpx_client from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.router import GenericLiteLLMParams @@ -247,12 +248,12 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): Raises BlackForestLabsError on failure. """ start_time = time.time() + httpx_client = _get_httpx_client() while time.time() - start_time < max_wait: - response = httpx.get( + response = httpx_client.get( polling_url, headers={"x-key": api_key}, - timeout=30.0, ) if response.status_code != 200: diff --git a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py b/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py index 2c2d207dc72..167aef03140 100644 --- a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py +++ b/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py @@ -242,7 +242,10 @@ class TestBlackForestLabsImageEditTransformation: "result": {"sample": "https://example.com/image.png"}, } - with patch("httpx.get", return_value=mock_response): + mock_client = MagicMock() + mock_client.get.return_value = mock_response + + with patch("litellm.llms.black_forest_labs.image_edit.transformation._get_httpx_client", return_value=mock_client): result = self.config._poll_for_result( polling_url="https://api.bfl.ai/v1/get_result?id=123", api_key="test-key", @@ -266,7 +269,10 @@ class TestBlackForestLabsImageEditTransformation: "result": {"sample": "https://example.com/image.png"}, } - with patch("httpx.get", side_effect=[pending_response, ready_response]): + mock_client = MagicMock() + mock_client.get.side_effect = [pending_response, ready_response] + + with patch("litellm.llms.black_forest_labs.image_edit.transformation._get_httpx_client", return_value=mock_client): result = self.config._poll_for_result( polling_url="https://api.bfl.ai/v1/get_result?id=123", api_key="test-key", @@ -282,7 +288,10 @@ class TestBlackForestLabsImageEditTransformation: mock_response.status_code = 200 mock_response.json.return_value = {"status": "Error"} - with patch("httpx.get", return_value=mock_response): + mock_client = MagicMock() + mock_client.get.return_value = mock_response + + with patch("litellm.llms.black_forest_labs.image_edit.transformation._get_httpx_client", return_value=mock_client): with pytest.raises(BlackForestLabsError) as exc_info: self.config._poll_for_result( polling_url="https://api.bfl.ai/v1/get_result?id=123", @@ -300,7 +309,10 @@ class TestBlackForestLabsImageEditTransformation: mock_response.status_code = 200 mock_response.json.return_value = {"status": "Content Moderated"} - with patch("httpx.get", return_value=mock_response): + mock_client = MagicMock() + mock_client.get.return_value = mock_response + + with patch("litellm.llms.black_forest_labs.image_edit.transformation._get_httpx_client", return_value=mock_client): with pytest.raises(BlackForestLabsError) as exc_info: self.config._poll_for_result( polling_url="https://api.bfl.ai/v1/get_result?id=123", @@ -318,7 +330,10 @@ class TestBlackForestLabsImageEditTransformation: mock_response.status_code = 200 mock_response.json.return_value = {"status": "Pending"} - with patch("httpx.get", return_value=mock_response): + mock_client = MagicMock() + mock_client.get.return_value = mock_response + + with patch("litellm.llms.black_forest_labs.image_edit.transformation._get_httpx_client", return_value=mock_client): with pytest.raises(BlackForestLabsError) as exc_info: self.config._poll_for_result( polling_url="https://api.bfl.ai/v1/get_result?id=123", @@ -336,7 +351,10 @@ class TestBlackForestLabsImageEditTransformation: mock_response.status_code = 500 mock_response.text = "Internal Server Error" - with patch("httpx.get", return_value=mock_response): + mock_client = MagicMock() + mock_client.get.return_value = mock_response + + with patch("litellm.llms.black_forest_labs.image_edit.transformation._get_httpx_client", return_value=mock_client): with pytest.raises(BlackForestLabsError) as exc_info: self.config._poll_for_result( polling_url="https://api.bfl.ai/v1/get_result?id=123", @@ -369,7 +387,10 @@ class TestBlackForestLabsImageEditTransformation: "result": {"sample": "https://example.com/edited-image.png"}, } - with patch("httpx.get", return_value=poll_response): + mock_client = MagicMock() + mock_client.get.return_value = poll_response + + with patch("litellm.llms.black_forest_labs.image_edit.transformation._get_httpx_client", return_value=mock_client): result = self.config.transform_image_edit_response( model=self.model, raw_response=mock_response,