mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
fix(fal_ai): handle transient result errors and inject clients
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry
parent
eaa6936f13
commit
0b9035b48f
2 changed files with 71 additions and 26 deletions
|
|
@ -1,7 +1,7 @@
|
|||
import math
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeAlias
|
||||
|
|
@ -212,6 +212,16 @@ def _result_error(raw_response: httpx.Response) -> str | None:
|
|||
return response_text or f"fal.ai returned HTTP {raw_response.status_code}"
|
||||
|
||||
|
||||
def _terminal_result_error(raw_response: httpx.Response) -> str | None:
|
||||
if raw_response.status_code == 429 or raw_response.status_code >= 500:
|
||||
return None
|
||||
return _result_error(raw_response)
|
||||
|
||||
|
||||
def _get_fal_ai_async_httpx_client() -> AsyncHTTPHandler:
|
||||
return get_async_httpx_client(llm_provider=LlmProviders.FAL_AI)
|
||||
|
||||
|
||||
def _response_string(response_data: Mapping[str, object], key: str, default: str = "") -> str:
|
||||
value: Final[object] = response_data.get(key)
|
||||
return value if isinstance(value, str) else default
|
||||
|
|
@ -265,6 +275,15 @@ def _status_video_object(
|
|||
|
||||
|
||||
class FalAIVideoConfig(BaseVideoConfig):
|
||||
def __init__(
|
||||
self,
|
||||
sync_client_factory: Callable[[], HTTPHandler] = _get_httpx_client,
|
||||
async_client_factory: Callable[[], AsyncHTTPHandler] = _get_fal_ai_async_httpx_client,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self._sync_client_factory: Final = sync_client_factory
|
||||
self._async_client_factory: Final = async_client_factory
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> _SupportedParams:
|
||||
supported_params: Final[_SupportedParams] = [ # mutable-ok: BaseVideoConfig requires a list
|
||||
"model",
|
||||
|
|
@ -458,11 +477,11 @@ class FalAIVideoConfig(BaseVideoConfig):
|
|||
if result_request is None:
|
||||
return None
|
||||
result_url, result_headers = result_request
|
||||
result_response: Final[httpx.Response] = _get_httpx_client().get(
|
||||
result_response: Final[httpx.Response] = self._sync_client_factory().get(
|
||||
url=result_url,
|
||||
headers=result_headers,
|
||||
)
|
||||
return _result_error(result_response)
|
||||
return _terminal_result_error(result_response)
|
||||
|
||||
async def async_transform_video_status_retrieve_response(
|
||||
self,
|
||||
|
|
@ -488,12 +507,11 @@ class FalAIVideoConfig(BaseVideoConfig):
|
|||
if result_request is None:
|
||||
return None
|
||||
result_url, result_headers = result_request
|
||||
async_httpx_client: Final[AsyncHTTPHandler] = get_async_httpx_client(llm_provider=LlmProviders.FAL_AI)
|
||||
result_response: Final[httpx.Response] = await async_httpx_client.get(
|
||||
result_response: Final[httpx.Response] = await self._async_client_factory().get(
|
||||
url=result_url,
|
||||
headers=result_headers,
|
||||
)
|
||||
return _result_error(result_response)
|
||||
return _terminal_result_error(result_response)
|
||||
|
||||
@staticmethod
|
||||
def _decode_video_id(video_id: str) -> tuple[str, str]:
|
||||
|
|
@ -547,7 +565,7 @@ class FalAIVideoConfig(BaseVideoConfig):
|
|||
response=raw_response,
|
||||
)
|
||||
video_url: Final[str] = self._extract_video_url(_response_data(raw_response))
|
||||
httpx_client: Final[HTTPHandler] = _get_httpx_client()
|
||||
httpx_client: Final[HTTPHandler] = self._sync_client_factory()
|
||||
video_response: Final[httpx.Response] = httpx_client.get( # pyright: ignore[reportUnknownMemberType] # HTTP handler stubs are untyped
|
||||
video_url
|
||||
)
|
||||
|
|
@ -565,7 +583,7 @@ class FalAIVideoConfig(BaseVideoConfig):
|
|||
response=raw_response,
|
||||
)
|
||||
video_url: Final[str] = self._extract_video_url(_response_data(raw_response))
|
||||
async_httpx_client: Final[AsyncHTTPHandler] = get_async_httpx_client(llm_provider=LlmProviders.FAL_AI)
|
||||
async_httpx_client: Final[AsyncHTTPHandler] = self._async_client_factory()
|
||||
video_response: Final[httpx.Response] = await async_httpx_client.get( # pyright: ignore[reportUnknownMemberType] # HTTP handler stubs are untyped
|
||||
video_url
|
||||
)
|
||||
|
|
|
|||
|
|
@ -216,9 +216,10 @@ class TestFalAIVideoTransformation:
|
|||
({"request_id": "abc", "status": "COMPLETED"}, "completed"),
|
||||
],
|
||||
)
|
||||
def test_status_response_mapping(self, response_data, expected_status, monkeypatch):
|
||||
def test_status_response_mapping(self, response_data, expected_status):
|
||||
status_url = "https://queue.fal.run/bytedance/seedance-2.5/requests/abc/status"
|
||||
response = httpx.Response(200, json=response_data, request=httpx.Request("GET", status_url))
|
||||
config = self.config
|
||||
if expected_status == "completed":
|
||||
result_response: Final = httpx.Response(
|
||||
200,
|
||||
|
|
@ -227,9 +228,9 @@ class TestFalAIVideoTransformation:
|
|||
)
|
||||
client: Final = Mock()
|
||||
client.get.return_value = result_response
|
||||
monkeypatch.setattr(fal_video_module, "_get_httpx_client", lambda: client)
|
||||
config = FalAIVideoConfig(sync_client_factory=lambda: client)
|
||||
|
||||
video = self.config.transform_video_status_retrieve_response(
|
||||
video = config.transform_video_status_retrieve_response(
|
||||
raw_response=response,
|
||||
logging_obj=self.logging_obj,
|
||||
custom_llm_provider="fal_ai",
|
||||
|
|
@ -249,7 +250,7 @@ class TestFalAIVideoTransformation:
|
|||
)
|
||||
assert poll_url == status_url
|
||||
|
||||
def test_status_response_error(self, monkeypatch):
|
||||
def test_status_response_error(self):
|
||||
response_data = {
|
||||
"request_id": "abc",
|
||||
"status": "COMPLETED",
|
||||
|
|
@ -268,9 +269,9 @@ class TestFalAIVideoTransformation:
|
|||
)
|
||||
client: Final = Mock()
|
||||
client.get.return_value = result_response
|
||||
monkeypatch.setattr(fal_video_module, "_get_httpx_client", lambda: client)
|
||||
config = FalAIVideoConfig(sync_client_factory=lambda: client)
|
||||
|
||||
video = self.config.transform_video_status_retrieve_response(
|
||||
video = config.transform_video_status_retrieve_response(
|
||||
raw_response=response,
|
||||
logging_obj=self.logging_obj,
|
||||
custom_llm_provider="fal_ai",
|
||||
|
|
@ -279,7 +280,7 @@ class TestFalAIVideoTransformation:
|
|||
assert video.status == "failed"
|
||||
assert video.error == {"code": "fal_error", "message": "generation failed"}
|
||||
|
||||
def test_status_completed_result_error_surfaces_fal_message(self, monkeypatch):
|
||||
def test_status_completed_result_error_surfaces_fal_message(self):
|
||||
status_url = "https://queue.fal.run/minimax/h3/requests/abc/status"
|
||||
auth_headers: Final = {"Authorization": "Key synthetic-fal-key", "Content-Type": "application/json"}
|
||||
response: Final = httpx.Response(
|
||||
|
|
@ -302,9 +303,9 @@ class TestFalAIVideoTransformation:
|
|||
)
|
||||
client: Final = Mock()
|
||||
client.get.return_value = result_response
|
||||
monkeypatch.setattr(fal_video_module, "_get_httpx_client", lambda: client)
|
||||
config = FalAIVideoConfig(sync_client_factory=lambda: client)
|
||||
|
||||
video = self.config.transform_video_status_retrieve_response(
|
||||
video = config.transform_video_status_retrieve_response(
|
||||
raw_response=response,
|
||||
logging_obj=self.logging_obj,
|
||||
custom_llm_provider="fal_ai",
|
||||
|
|
@ -314,8 +315,34 @@ class TestFalAIVideoTransformation:
|
|||
assert "input.reference_image_urls: Failed to download the file" in video.error["message"]
|
||||
client.get.assert_called_once_with(url=result_url, headers=auth_headers)
|
||||
|
||||
@pytest.mark.parametrize("status_code", [429, 503])
|
||||
def test_status_completed_transient_result_error_keeps_completed(self, status_code):
|
||||
status_url = "https://queue.fal.run/minimax/h3/requests/abc/status"
|
||||
response: Final = httpx.Response(
|
||||
200,
|
||||
json={"request_id": "abc", "status": "COMPLETED"},
|
||||
request=httpx.Request("GET", status_url),
|
||||
)
|
||||
result_response: Final = httpx.Response(
|
||||
status_code,
|
||||
json={"detail": "temporary fal failure"},
|
||||
request=httpx.Request("GET", status_url.removesuffix("/status")),
|
||||
)
|
||||
client: Final = Mock()
|
||||
client.get.return_value = result_response
|
||||
config = FalAIVideoConfig(sync_client_factory=lambda: client)
|
||||
|
||||
video = config.transform_video_status_retrieve_response(
|
||||
raw_response=response,
|
||||
logging_obj=self.logging_obj,
|
||||
custom_llm_provider="fal_ai",
|
||||
)
|
||||
|
||||
assert video.status == "completed"
|
||||
assert video.error is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_status_completed_result_error_surfaces_fal_message(self, monkeypatch):
|
||||
async def test_async_status_completed_result_error_surfaces_fal_message(self):
|
||||
status_url = "https://queue.fal.run/minimax/h3/requests/abc/status"
|
||||
auth_headers: Final = {"Authorization": "Key synthetic-fal-key", "Content-Type": "application/json"}
|
||||
response: Final = httpx.Response(
|
||||
|
|
@ -338,9 +365,9 @@ class TestFalAIVideoTransformation:
|
|||
)
|
||||
client: Final = Mock()
|
||||
client.get = AsyncMock(return_value=result_response)
|
||||
monkeypatch.setattr(fal_video_module, "get_async_httpx_client", lambda llm_provider: client)
|
||||
config = FalAIVideoConfig(async_client_factory=lambda: client)
|
||||
|
||||
video = await self.config.async_transform_video_status_retrieve_response(
|
||||
video = await config.async_transform_video_status_retrieve_response(
|
||||
raw_response=response,
|
||||
logging_obj=self.logging_obj,
|
||||
custom_llm_provider="fal_ai",
|
||||
|
|
@ -350,7 +377,7 @@ class TestFalAIVideoTransformation:
|
|||
assert "input.reference_image_urls: Failed to download the file" in video.error["message"]
|
||||
client.get.assert_awaited_once_with(url=result_url, headers=auth_headers)
|
||||
|
||||
def test_status_in_progress_does_not_fetch_result(self, monkeypatch):
|
||||
def test_status_in_progress_does_not_fetch_result(self):
|
||||
status_url = "https://queue.fal.run/minimax/h3/requests/abc/status"
|
||||
response = httpx.Response(
|
||||
200,
|
||||
|
|
@ -359,9 +386,9 @@ class TestFalAIVideoTransformation:
|
|||
)
|
||||
|
||||
client: Final = Mock()
|
||||
monkeypatch.setattr(fal_video_module, "_get_httpx_client", lambda: client)
|
||||
config = FalAIVideoConfig(sync_client_factory=lambda: client)
|
||||
|
||||
video = self.config.transform_video_status_retrieve_response(
|
||||
video = config.transform_video_status_retrieve_response(
|
||||
raw_response=response,
|
||||
logging_obj=self.logging_obj,
|
||||
custom_llm_provider="fal_ai",
|
||||
|
|
@ -391,7 +418,7 @@ class TestFalAIVideoTransformation:
|
|||
assert decoded["video_id"] == "xyz"
|
||||
assert video.model == "workflows/owner/app"
|
||||
|
||||
def test_content_response_downloads_video_url(self, monkeypatch):
|
||||
def test_content_response_downloads_video_url(self):
|
||||
content_response = httpx.Response(
|
||||
200,
|
||||
content=b"video-bytes",
|
||||
|
|
@ -403,11 +430,11 @@ class TestFalAIVideoTransformation:
|
|||
assert url == "https://cdn.example.com/video.mp4"
|
||||
return content_response
|
||||
|
||||
monkeypatch.setattr(fal_video_module, "_get_httpx_client", lambda: FakeHTTPClient())
|
||||
config = FalAIVideoConfig(sync_client_factory=FakeHTTPClient)
|
||||
response = Mock(spec=httpx.Response)
|
||||
response.json.return_value = {"video": {"url": "https://cdn.example.com/video.mp4"}}
|
||||
|
||||
assert self.config.transform_video_content_response(response, self.logging_obj) == b"video-bytes"
|
||||
assert config.transform_video_content_response(response, self.logging_obj) == b"video-bytes"
|
||||
|
||||
def test_content_response_rejects_missing_video(self):
|
||||
response = Mock(spec=httpx.Response)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue