phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

perf improvement - use tasks for async logging pass through responses

Browse source
This commit is contained in:
Ishaan Jaff 2024-11-23 16:58:25 -08:00
parent e69678a9b3
commit 2c0551f0a8
2 changed files with 32 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

23
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -528,16 +528,19 @@ async def pass_through_request( # noqa: PLR0915
response_body: Optional[dict] = get_response_body(response) response_body: Optional[dict] = get_response_body(response)
passthrough_logging_payload["response_body"] = response_body passthrough_logging_payload["response_body"] = response_body
end_time = datetime.now() end_time = datetime.now()
await pass_through_endpoint_logging.pass_through_async_success_handler(
httpx_response=response, asyncio.create_task(
response_body=response_body, pass_through_endpoint_logging.pass_through_async_success_handler(
url_route=str(url), httpx_response=response,
result="", response_body=response_body,
start_time=start_time, url_route=str(url),
end_time=end_time, result="",
logging_obj=logging_obj, start_time=start_time,
cache_hit=False, end_time=end_time,
**kwargs, logging_obj=logging_obj,
cache_hit=False,
**kwargs,
)
) )
return Response( return Response(

36
litellm/proxy/pass_through_endpoints/streaming_handler.py
View file

@ -58,22 +58,24 @@ class PassThroughStreamingHandler:
# After all chunks are processed, handle post-processing # After all chunks are processed, handle post-processing
end_time = datetime.now() end_time = datetime.now()
await PassThroughStreamingHandler._route_streaming_logging_to_handler( asyncio.create_task(
litellm_logging_obj=litellm_logging_obj, PassThroughStreamingHandler.handle_logging_collected_stream_response(
passthrough_success_handler_obj=passthrough_success_handler_obj, litellm_logging_obj=litellm_logging_obj,
url_route=url_route, passthrough_success_handler_obj=passthrough_success_handler_obj,
request_body=request_body or {}, url_route=url_route,
endpoint_type=endpoint_type, request_body=request_body or {},
start_time=start_time, endpoint_type=endpoint_type,
raw_bytes=raw_bytes, start_time=start_time,
end_time=end_time, raw_bytes=raw_bytes,
end_time=end_time,
)
) )
except Exception as e: except Exception as e:
verbose_proxy_logger.error(f"Error in chunk_processor: {str(e)}") verbose_proxy_logger.error(f"Error in chunk_processor: {str(e)}")
raise raise
@staticmethod @staticmethod
async def _route_streaming_logging_to_handler( async def handle_logging_collected_stream_response(
litellm_logging_obj: LiteLLMLoggingObj, litellm_logging_obj: LiteLLMLoggingObj,
passthrough_success_handler_obj: PassThroughEndpointLogging, passthrough_success_handler_obj: PassThroughEndpointLogging,
url_route: str, url_route: str,
@ -108,9 +110,9 @@ class PassThroughStreamingHandler:
all_chunks=all_chunks, all_chunks=all_chunks,
end_time=end_time, end_time=end_time,
) )
standard_logging_response_object = anthropic_passthrough_logging_handler_result[ standard_logging_response_object = (
"result" anthropic_passthrough_logging_handler_result["result"]
] )
kwargs = anthropic_passthrough_logging_handler_result["kwargs"] kwargs = anthropic_passthrough_logging_handler_result["kwargs"]
elif endpoint_type == EndpointType.VERTEX_AI: elif endpoint_type == EndpointType.VERTEX_AI:
vertex_passthrough_logging_handler_result = ( vertex_passthrough_logging_handler_result = (
@ -125,9 +127,9 @@ class PassThroughStreamingHandler:
end_time=end_time, end_time=end_time,
) )
) )
standard_logging_response_object = vertex_passthrough_logging_handler_result[ standard_logging_response_object = (
"result" vertex_passthrough_logging_handler_result["result"]
] )
kwargs = vertex_passthrough_logging_handler_result["kwargs"] kwargs = vertex_passthrough_logging_handler_result["kwargs"]
if standard_logging_response_object is None: if standard_logging_response_object is None:
@ -168,4 +170,4 @@ class PassThroughStreamingHandler:
# Split by newlines and filter out empty lines # Split by newlines and filter out empty lines
lines = [line.strip() for line in combined_str.split("\n") if line.strip()] lines = [line.strip() for line in combined_str.split("\n") if line.strip()]
return lines return lines