From 3b317b0c215df3cffc80e2529b8edf79e50ec75b Mon Sep 17 00:00:00 2001 From: bhimrazy Date: Fri, 10 Jan 2025 02:15:50 +0545 Subject: [PATCH] fix: add missing before and after callback encode response events in streaming loop --- src/litserve/loops/streaming_loops.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/litserve/loops/streaming_loops.py b/src/litserve/loops/streaming_loops.py index 2ed574917..aac0b1fa8 100644 --- a/src/litserve/loops/streaming_loops.py +++ b/src/litserve/loops/streaming_loops.py @@ -69,12 +69,15 @@ def run_streaming_loop( x, ) callback_runner.trigger_event(EventTypes.AFTER_PREDICT, lit_api=lit_api) - + + callback_runner.trigger_event(EventTypes.BEFORE_ENCODE_RESPONSE, lit_api=lit_api) y_enc_gen = _inject_context( context, lit_api.encode_response, y_gen, ) + callback_runner.trigger_event(EventTypes.AFTER_ENCODE_RESPONSE, lit_api=lit_api) + for y_enc in y_enc_gen: y_enc = lit_api.format_encoded_response(y_enc) response_queues[response_queue_id].put((uid, (y_enc, LitAPIStatus.OK)))