fix: adds aligned hooks back

Lightning-AI · Jan 13, 2025 · 37679c4 · 37679c4
1 parent c33c326
commit 37679c4
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/litserve/loops/streaming_loops.py b/src/litserve/loops/streaming_loops.py
@@ -74,6 +74,7 @@ def run_streaming_loop(
                 )
                 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,
@@ -84,6 +85,9 @@ def run_streaming_loop(
                     self.put_response(response_queues, response_queue_id, uid, y_enc, LitAPIStatus.OK)
                 self.put_response(response_queues, response_queue_id, uid, "", LitAPIStatus.FINISH_STREAMING)
 
+                callback_runner.trigger_event(EventTypes.AFTER_PREDICT, lit_api=lit_api)
+                callback_runner.trigger_event(EventTypes.AFTER_ENCODE_RESPONSE, lit_api=lit_api)
+
             except HTTPException as e:
                 self.put_response(
                     response_queues,