scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(proxy): only fire post-guardrail log on successful stream completion

Browse source 
Guard task creation with 'if error_message is None' so async_post_guardrail_log_success_event is not called when the stream terminated due to an SSE error.

Made-with: Cursor
This commit is contained in:
devarakondasrikanth 2026-03-13 19:57:57 -07:00
parent 7ac26f23d8
commit 51241870ea
1 changed files with 13 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
litellm/proxy/proxy_server.py
View file

@ -5609,20 +5609,21 @@ async def async_data_generator(
except Exception as e:
yield f"data: {str(e)}\n\n"
# Post-guardrail log: run in background so we don't block yielding [DONE]
def _discard_task(t: asyncio.Task[None]) -> None:
_post_guardrail_log_tasks.discard(t)
# Post-guardrail log: only for successful stream completion, not when we got an SSE error
if error_message is None:
def _discard_task(t: asyncio.Task[None]) -> None:
_post_guardrail_log_tasks.discard(t)
_task = asyncio.create_task(
_async_data_generator_fire_post_guardrail_log(
request_data=request_data,
user_api_key_dict=user_api_key_dict,
chunks_for_log=_streaming_chunks_for_log,
logging_obj=request_data.get("litellm_logging_obj"),
_task = asyncio.create_task(
_async_data_generator_fire_post_guardrail_log(
request_data=request_data,
user_api_key_dict=user_api_key_dict,
chunks_for_log=_streaming_chunks_for_log,
logging_obj=request_data.get("litellm_logging_obj"),
)
)
)
_post_guardrail_log_tasks.add(_task)
_task.add_done_callback(_discard_task)
_post_guardrail_log_tasks.add(_task)
_task.add_done_callback(_discard_task)
# Streaming is done, yield the [DONE] chunk
if error_message is not None: