scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 75

fix: Call guardrail during stream processing

Browse source
This commit is contained in:
Yuta Saito 2025-09-04 07:11:44 +09:00
parent d83c420d48
commit 447016817c
1 changed files with 5 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/proxy/common_request_processing.py
View file

@ -108,7 +108,6 @@ async def create_streaming_response(
final_status_code = default_status_code
try:
# Handle coroutine that returns a generator
if asyncio.iscoroutine(generator):
generator = await generator
@ -117,7 +116,6 @@ async def create_streaming_response(
first_chunk_value = await generator.__anext__()
if first_chunk_value is not None:
try:
error_code_from_chunk = await _parse_event_data_for_error(
first_chunk_value
@ -131,7 +129,6 @@ async def create_streaming_response(
verbose_proxy_logger.debug(f"Error parsing first chunk value: {e}")
except StopAsyncIteration:
# Generator was empty. Default status
async def empty_gen() -> AsyncGenerator[str, None]:
if False:
@ -144,7 +141,6 @@ async def create_streaming_response(
status_code=default_status_code,
)
except Exception as e:
# Unexpected error consuming first chunk.
verbose_proxy_logger.exception(
f"Error consuming first chunk from generator: {e}"
@ -167,7 +163,6 @@ async def create_streaming_response(
with tracer.trace(DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE):
yield first_chunk_value
async for chunk in generator:
with tracer.trace(DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE):
yield chunk
@ -460,7 +455,6 @@ class ProxyBaseLLMRequestProcessing:
) or self._is_streaming_response(
response
): # use generate_responses to stream responses
custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=logging_obj.litellm_call_id,
@ -478,7 +472,6 @@ class ProxyBaseLLMRequestProcessing:
if route_type == "allm_passthrough_route":
# Check if response is an async generator
if self._is_streaming_response(response):
if asyncio.iscoroutine(response):
generator = await response
else:
@ -499,7 +492,6 @@ class ProxyBaseLLMRequestProcessing:
headers=custom_headers,
)
else:
selected_data_generator = select_data_generator(
response=response,
user_api_key_dict=user_api_key_dict,
@ -738,7 +730,11 @@ class ProxyBaseLLMRequestProcessing:
verbose_proxy_logger.debug("inside generator")
try:
str_so_far = ""
async for chunk in response:
async for chunk in proxy_logging_obj.async_post_call_streaming_iterator_hook(
user_api_key_dict=user_api_key_dict,
response=response,
request_data=request_data,
):
verbose_proxy_logger.debug(
"async_data_generator: received streaming chunk - {}".format(chunk)
)