Run before_fn and after_fn for each generator iteration

.changeset/big-bears-cover.md

@@ -0,0 +1,5 @@
+---
+"gradio": patch
+---
+
+fix:Run before_fn and after_fn for each generator iteration

gradio/utils.py

@@ -635,12 +635,19 @@ def function_wrapper(
 
         @functools.wraps(f)
         async def asyncgen_wrapper(*args, **kwargs):
-            if before_fn:
-                before_fn(*before_args)
-            async for response in f(*args, **kwargs):
+            iterator = f(*args, **kwargs)
+            while True:
+                if before_fn:
+                    before_fn(*before_args)
+                try:
+                    response = await iterator.__anext__()
+                except StopAsyncIteration:
+                    if after_fn:
+                        after_fn(*after_args)
+                    break
+                if after_fn:
+                    after_fn(*after_args)
                 yield response
-            if after_fn:
-                after_fn(*after_args)
 
         return asyncgen_wrapper
 
@@ -661,11 +668,19 @@ async def async_wrapper(*args, **kwargs):
 
         @functools.wraps(f)
         def gen_wrapper(*args, **kwargs):
-            if before_fn:
-                before_fn(*before_args)
-            yield from f(*args, **kwargs)
-            if after_fn:
-                after_fn(*after_args)
+            iterator = f(*args, **kwargs)
+            while True:
+                if before_fn:
+                    before_fn(*before_args)
+                try:
+                    response = next(iterator)
+                except StopIteration:
+                    if after_fn:
+                        after_fn(*after_args)
+                    break
+                if after_fn:
+                    after_fn(*after_args)
+                yield response
 
         return gen_wrapper
 
@@ -705,7 +720,10 @@ def after_fn():
         LocalContext.request.set(None)
 
     return function_wrapper(
-        fn, before_fn=before_fn, before_args=(blocks, event_id), after_fn=after_fn
+        fn,
+        before_fn=before_fn,
+        before_args=(blocks, event_id),
+        after_fn=after_fn,
     )