scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 84

Use _get_httpx_client for HTTP polling in BFL image edit

Browse source 
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.).
This commit is contained in:
Chesars 2025-12-17 00:15:54 -03:00
parent 77ca224a7f
commit d180db31e7
2 changed files with 31 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/llms/black_forest_labs/image_edit/transformation.py
View file

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

35
tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py
View file

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