diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index b4a7682202..207dd5667d 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -54,12 +54,7 @@ async def allm_passthrough_route( cookies: Optional[CookieTypes] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, **kwargs, -) -> Union[ - httpx.Response, - Coroutine[Any, Any, httpx.Response], - Generator[Any, Any, Any], - AsyncGenerator[Any, Any], -]: +) -> Union[httpx.Response, AsyncGenerator[Any, Any]]: """ Async: Reranks a list of documents based on their relevance to the query """ @@ -111,20 +106,23 @@ async def allm_passthrough_route( func_with_context = partial(ctx.run, func) init_response = await loop.run_in_executor(None, func_with_context) + # Since allm_passthrough_route=True, we always get a coroutine from _async_passthrough_request if asyncio.iscoroutine(init_response): response = await init_response - try: - response.raise_for_status() - except httpx.HTTPStatusError as e: - error_text = await e.response.aread() - error_text_str = error_text.decode("utf-8") - raise Exception(error_text_str) - + # Only call raise_for_status if it's a Response object (not a generator) + if isinstance(response, httpx.Response): + try: + response.raise_for_status() + except httpx.HTTPStatusError as e: + error_text = await e.response.aread() + error_text_str = error_text.decode("utf-8") + raise Exception(error_text_str) + + return response else: - response = init_response - - return response + # This shouldn't happen when allm_passthrough_route=True, but handle it for type safety + raise Exception("Expected coroutine from async passthrough route") except Exception as e: # For passthrough routes, we need to get the provider config to properly handle errors @@ -186,6 +184,7 @@ def llm_passthrough_route( ) -> Union[ httpx.Response, Coroutine[Any, Any, httpx.Response], + Coroutine[Any, Any, Union[httpx.Response, AsyncGenerator[Any, Any]]], Generator[Any, Any, Any], AsyncGenerator[Any, Any], ]: @@ -200,8 +199,10 @@ def llm_passthrough_route( from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager + _is_async = allm_passthrough_route + if client is None: - if allm_passthrough_route: + if _is_async: client = litellm.module_level_aclient else: client = litellm.module_level_client @@ -302,24 +303,40 @@ def llm_passthrough_route( # Update logging object with streaming status litellm_logging_obj.stream = is_streaming_request + ## LOGGING PRE-CALL + request_data = data if data else json + litellm_logging_obj.pre_call( + input=request_data, + api_key=provider_api_key, + additional_args={ + "complete_input_dict": request_data, + "api_base": str(updated_url), + "headers": headers, + }, + ) + try: - response = client.client.send(request=request, stream=is_streaming_request) - if asyncio.iscoroutine(response): - if is_streaming_request: - return _async_streaming(response, litellm_logging_obj, provider_config) - else: - return response - response.raise_for_status() - - if ( - hasattr(response, "iter_bytes") and is_streaming_request - ): # yield the chunk, so we can store it in the logging object - - return _sync_streaming(response, litellm_logging_obj, provider_config) + if _is_async: + # Return the coroutine to be awaited by the caller + return _async_passthrough_request( + client=client, + request=request, + is_streaming_request=is_streaming_request, + litellm_logging_obj=litellm_logging_obj, + provider_config=provider_config, + ) else: + # Sync path - client.client.send returns Response directly + response: httpx.Response = client.client.send(request=request, stream=is_streaming_request) # type: ignore + response.raise_for_status() - # For non-streaming responses, yield the entire response - return response + if ( + hasattr(response, "iter_bytes") and is_streaming_request + ): # yield the chunk, so we can store it in the logging object + return _sync_streaming(response, litellm_logging_obj, provider_config) + else: + # For non-streaming responses, yield the entire response + return response except Exception as e: if provider_config is None: raise e @@ -329,6 +346,39 @@ def llm_passthrough_route( ) +async def _async_passthrough_request( + client: Union[HTTPHandler, AsyncHTTPHandler], + request: httpx.Request, + is_streaming_request: bool, + litellm_logging_obj: "LiteLLMLoggingObj", + provider_config: "BasePassthroughConfig", +) -> Union[httpx.Response, AsyncGenerator[Any, Any]]: + """ + Handle async passthrough requests. + Uses async client to send request and properly handles streaming. + """ + # client.client.send returns a coroutine for async clients + response_result = client.client.send(request=request, stream=is_streaming_request) + + # Check if it's a coroutine and await it + if asyncio.iscoroutine(response_result): + if is_streaming_request: + # Pass the coroutine to _async_streaming which will await it + return _async_streaming( + response=response_result, + litellm_logging_obj=litellm_logging_obj, + provider_config=provider_config, + ) + else: + response = await response_result + await response.aread() + response.raise_for_status() + return response + else: + # Fallback for sync-like behavior (shouldn't happen in async path) + raise Exception("Expected coroutine from async client") + + def _sync_streaming( response: httpx.Response, litellm_logging_obj: "LiteLLMLoggingObj",