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

handle vertex pass through separately

Browse source
This commit is contained in:
Ishaan Jaff 2024-11-22 13:18:08 -08:00
parent dcab2d0c6f
commit e829b228b2
1 changed files with 28 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -40,35 +40,39 @@ async def chunk_processor(
- Yields chunks from the response - Yields chunks from the response
- Collect non-empty chunks for post-processing (logging) - Collect non-empty chunks for post-processing (logging)
""" """
collected_chunks: List[str] = [] # List to store all chunks
try: try:
async for chunk in response.aiter_lines(): if endpoint_type == EndpointType.VERTEX_AI:
verbose_proxy_logger.debug(f"Processing chunk: {chunk}") async for chunk in response.aiter_bytes():
if not chunk: yield chunk
continue else:
collected_chunks: List[str] = [] # List to store all chunks
async for chunk in response.aiter_lines():
verbose_proxy_logger.debug(f"Processing chunk: {chunk}")
if not chunk:
continue
# Handle SSE format - pass through the raw SSE format # Handle SSE format - pass through the raw SSE format
if isinstance(chunk, bytes): if isinstance(chunk, bytes):
chunk = chunk.decode("utf-8") chunk = chunk.decode("utf-8")
# Store the chunk for post-processing # Store the chunk for post-processing
if chunk.strip(): # Only store non-empty chunks if chunk.strip(): # Only store non-empty chunks
collected_chunks.append(chunk) collected_chunks.append(chunk)
yield f"{chunk}\n" yield f"{chunk}\n"
# 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 _route_streaming_logging_to_handler( await _route_streaming_logging_to_handler(
litellm_logging_obj=litellm_logging_obj, litellm_logging_obj=litellm_logging_obj,
passthrough_success_handler_obj=passthrough_success_handler_obj, passthrough_success_handler_obj=passthrough_success_handler_obj,
url_route=url_route, url_route=url_route,
request_body=request_body or {}, request_body=request_body or {},
endpoint_type=endpoint_type, endpoint_type=endpoint_type,
start_time=start_time, start_time=start_time,
all_chunks=collected_chunks, all_chunks=collected_chunks,
end_time=end_time, 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)}")