Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/pass_through_endpoints/llm_provider_handlers/fal_ai_passthrough_logging_handler.py: 39%

1from collections.abc import Mapping, Sequence 

2from typing import Final 

3from urllib.parse import urlparse 

4 

5import httpx 

6 

7import litellm 

8from litellm._logging import verbose_proxy_logger 

9from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj 

10from litellm.llms.fal_ai.cost_calculator import fal_ai_passthrough_cost, fal_ai_queue_base 

11from litellm.proxy._types import PassThroughEndpointLoggingTypedDict 

12from litellm.types.utils import ImageObject, ImageResponse 

13 

14FAL_AI_PROVIDER: Final[str] = litellm.LlmProviders.FAL_AI.value 

15 

16 

17def _url_parts(value: object) -> tuple[Mapping[str, object], ...]: 

18 if isinstance(value, Mapping): 

19 return (value,) if isinstance(value.get("url"), str) else () 

20 if isinstance(value, Sequence) and not isinstance(value, str): 

21 return tuple(item for item in value if isinstance(item, Mapping) and isinstance(item.get("url"), str)) 

22 return () 

23 

24 

25class FalAIPassthroughLoggingHandler: 

26 @staticmethod 

27 def is_fal_ai_route(url_route: str, custom_llm_provider: str | None) -> bool: 

28 return custom_llm_provider == FAL_AI_PROVIDER 

29 

30 def fal_ai_passthrough_handler( 

31 self, 

32 response_body: Mapping[str, object], 

33 request_body: Mapping[str, object], 

34 logging_obj: LiteLLMLoggingObj, 

35 url_route: str, 

36 kwargs: Mapping[str, object], 

37 ) -> PassThroughEndpointLoggingTypedDict: 

38 base_path: Final = httpx.URL(fal_ai_queue_base()).path.strip("/") 

39 raw_path: Final = urlparse(url_route).path.strip("/") 

40 upstream_path: Final = raw_path.removeprefix(f"{base_path}/") if base_path else raw_path 

41 model: Final = upstream_path.partition("/requests/")[0] 

42 is_submit: Final = "/requests/" not in upstream_path 

43 response: Final = ImageResponse( 

44 data=tuple( 

45 ImageObject(url=url) 

46 for value in response_body.values() 

47 for part in _url_parts(value) 

48 if isinstance((url := part.get("url")), str) 

49 ) 

50 ) 

51 response_cost: Final = fal_ai_passthrough_cost(model, request_body) if is_submit else None 

52 response._hidden_params["response_cost"] = response_cost # pyright: ignore[reportPrivateUsage] # the logger reads a precomputed cost off the response's hidden params 

53 logging_obj.model = model # rebind-ok: the spend logger reads model and cost off the shared logging object 

54 logging_obj.model_call_details["model"] = model # rebind-ok: same shared logging object 

55 logging_obj.model_call_details["custom_llm_provider"] = FAL_AI_PROVIDER # rebind-ok: same shared logging object 

56 logging_obj.model_call_details["response_cost"] = response_cost # rebind-ok: same shared logging object 

57 verbose_proxy_logger.debug( 

58 "Fal AI passthrough cost tracking: model %s, cost %s", 

59 model, 

60 response_cost, 

61 ) 

62 logging_result: Final[PassThroughEndpointLoggingTypedDict] = { 

63 "result": response, 

64 "kwargs": { 

65 **kwargs, 

66 "model": model, 

67 "custom_llm_provider": FAL_AI_PROVIDER, 

68 "response_cost": response_cost, 

69 }, 

70 } 

71 return logging_result