Repository: gradio-app/gradio
Pull request: Add support for async generators
# Description

Closes: #3618 . @deep-diver would appreciate if you could test this to see if it fixes your issue. Click on the link on the gradio-pr-bot message, Go to Files in the top right, requirements.txt will have a url hosting a wheel built for this PR.

## To test

### Blocks case
```python
import asyncio
import gradio as gr
import time


async def async_iteration(count: int):
    for i in range(count):
        yield i
        await asyncio.sleep(1)


def iteration(count: int):
    for i in range(count):
        yield i
        time.sleep(1)


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            num1 = gr.Number(value=4, precision=0)
            o1 = gr.Number()
            async_iterate = gr.Button(value="Async Iteration")
            async_iterate.click(async_iteration, num1, o1)
        with gr.Column():
            num2 = gr.Number(value=4, precision=0)
            o2 = gr.Number()
            iterate = gr.Button(value="Iterate")
            iterate.click(iteration, num2, o2)

demo.queue(concurrency_count=20).launch()
```

### Interface case
```python
app1 = gr.Interface(async_iteration, gr.Number(precision=0, value=4), gr.Number())
app2 = gr.Interface(iteration, gr.Number(precision=0, value=4), gr.Number())

with gr.Blocks() as demo:
    with gr.Row():
        app1.render()
    with gr.Row():
        app2.render()

demo.queue(concurrency_count=20).launch()
```

# Checklist:

- [ ] I have performed a self-review of my own code
- [ ] I have added a short summary of my change to the CHANGELOG.md
- [ ] My code follows the style guidelines of this project
- [ ] I have commented my code in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] I have added tests that prove my fix is effective or that my feature works
- [ ] New and existing unit tests pass locally with my changes


# A note about the CHANGELOG

Hello 👋 and thank you for contributing to Gradio!

All pull requests must update the change log located in CHANGELOG.md, unless the pull request is labeled with the "no-changelog-update" label.

Please add a brief summary of the change to the Upcoming Release > Full Changelog section of the CHANGELOG.md file and include
a link to the PR (formatted in markdown) and a link to your github profile (if you like). For example, "* Added a cool new feature by `[@myusername](link-to-your-github-profile)` in `[PR 11111](https://github.com/gradio-app/gradio/pull/11111)`".

If you would like to elaborate on your change further, feel free to include a longer explanation in the other sections.
If you would like an image/gif/video showcasing your feature, it may be best to edit the CHANGELOG file using the 
GitHub web UI since that lets you upload files directly via drag-and-drop.

Problem statement:
async generator support 
- [x] I have searched to see if a similar issue already exists.


**Is your feature request related to a problem? Please describe.**  
I'm frustrated when the next user has to wait until the previous request is fully handled when using generator function.

**Describe the solution you'd like**  
Enable async generator in Gradio

**Additional context**  
I am trying to build a Chatbot application with Gradio Chatbot interface. I want to show/response to the connected users in a token by token manner. I can simply `yield` every tokens whenever they are generated. However, this is too time consuming.

Instead, I want to generate a chunk of tokens, and `yield` each token afterwards. In this way, the calculation must be much more efficient. Meanwhile I want to put `asyncio.sleep()` between every `yield` to make end users feel like they are returning back a token by token. 

However, the problem is that I want the other token generations to be processed while each tokens of the previous request are sequentially sent back to the user. It would enable more of concurrent connections, better response time, and eventually better UX overall.

Hints from issue discussion:
I came to this also,  I want to do langchain's streaming in gradio like this:

```python
async def streaming(input, history):

    callback = AsyncIteratorCallbackHandler()

    run = asyncio.create_task(conversation_with_summary.apredict(
        input=input, callbacks=[callback]))

    async for token in callback.aiter():
        print('token:', token)
        history[-1]["content"] += token

        messages = [(history[i]["content"], history[i+1]["content"])
                    for i in range(0, len(history)-1, 2)]
        yield messages

    await run
```

Commit message: Fix tests

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3ec16b9e74..eff2bb2f82 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,7 +2,7 @@
 
 ## New Features:
 
-No changes to highlight.
+- Support for asynchronous iterators by [@freddyaboulton](https://github.com/freddyaboulton) in [PR 3821](https://github.com/gradio-app/gradio/pull/3821)
 
 ## Bug Fixes:
 
diff --git a/client/python/test/test_client.py b/client/python/test/test_client.py
index a24d2e1e7b..a1742f86b9 100644
--- a/client/python/test/test_client.py
+++ b/client/python/test/test_client.py
@@ -588,10 +588,8 @@ def test_numerical_to_label_space(self):
     def test_serializable_in_mapping(self, calculator_demo):
         with connect(calculator_demo) as client:
             assert all(
-                [
-                    isinstance(c, SimpleSerializable)
-                    for c in client.endpoints[0].serializers
-                ]
+                isinstance(c, SimpleSerializable)
+                for c in client.endpoints[0].serializers
             )
 
     @pytest.mark.flaky
diff --git a/gradio/blocks.py b/gradio/blocks.py
index 335d5f2845..a1f3256388 100644
--- a/gradio/blocks.py
+++ b/gradio/blocks.py
@@ -12,7 +12,17 @@
 import webbrowser
 from abc import abstractmethod
 from types import ModuleType
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Set, Tuple, Type
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncIterator,
+    Callable,
+    Dict,
+    List,
+    Set,
+    Tuple,
+    Type,
+)
 
 import anyio
 import requests
@@ -971,7 +981,7 @@ async def call_function(
         self,
         fn_index: int,
         processed_input: List[Any],
-        iterator: Iterator[Any] | None = None,
+        iterator: AsyncIterator[Any] | None = None,
         requests: routes.Request | List[routes.Request] | None = None,
         event_id: str | None = None,
         event_data: EventData | None = None,
@@ -1024,17 +1034,17 @@ async def call_function(
         else:
             prediction = None
 
-        if inspect.isasyncgenfunction(block_fn.fn):
-            raise ValueError("Gradio does not support async generators.")
-        if inspect.isgeneratorfunction(block_fn.fn):
+        if inspect.isgeneratorfunction(block_fn.fn) or inspect.isasyncgenfunction(
+            block_fn.fn
+        ):
             if not self.enable_queue:
                 raise ValueError("Need to enable queue to use generators.")
             try:
                 if iterator is None:
                     iterator = prediction
-                prediction = await anyio.to_thread.run_sync(
-                    utils.async_iteration, iterator, limiter=self.limiter
-                )
+                if inspect.isgenerator(iterator):
+                    iterator = utils.SyncToAsyncIterator(iterator, self.limiter)
+                prediction = await utils.async_iteration(iterator)
                 is_generating = True
             except StopAsyncIteration:
                 n_outputs = len(self.dependencies[fn_index].get("outputs"))
diff --git a/gradio/interface.py b/gradio/interface.py
index c55ac16233..c88fc2eb01 100644
--- a/gradio/interface.py
+++ b/gradio/interface.py
@@ -507,12 +507,17 @@ def render_input_column(
                         # is created. We use whether a generator function is provided
                         # as a proxy of whether the queue will be enabled.
                         # Using a generator function without the queue will raise an error.
-                        if inspect.isgeneratorfunction(self.fn):
+                        if inspect.isgeneratorfunction(
+                            self.fn
+                        ) or inspect.isasyncgenfunction(self.fn):
                             stop_btn = Button("Stop", variant="stop", visible=False)
                 elif self.interface_type == InterfaceTypes.UNIFIED:
                     clear_btn = Button("Clear")
                     submit_btn = Button("Submit", variant="primary")
-                    if inspect.isgeneratorfunction(self.fn) and not self.live:
+                    if (
+                        inspect.isgeneratorfunction(self.fn)
+                        or inspect.isasyncgenfunction(self.fn)
+                    ) and not self.live:
                         stop_btn = Button("Stop", variant="stop")
                     if self.allow_flagging == "manual":
                         flag_btns = self.render_flag_btns()
@@ -543,7 +548,10 @@ def render_output_column(
                 if self.interface_type == InterfaceTypes.OUTPUT_ONLY:
                     clear_btn = Button("Clear")
                     submit_btn = Button("Generate", variant="primary")
-                    if inspect.isgeneratorfunction(self.fn) and not self.live:
+                    if (
+                        inspect.isgeneratorfunction(self.fn)
+                        or inspect.isasyncgenfunction(self.fn)
+                    ) and not self.live:
                         # Stopping jobs only works if the queue is enabled
                         # We don't know if the queue is enabled when the interface
                         # is created. We use whether a generator function is provided
@@ -609,27 +617,25 @@ def attach_submit_events(self, submit_btn: Button | None, stop_btn: Button | Non
             if stop_btn:
 
                 # Wrap the original function to show/hide the "Stop" button
-                def fn(*args):
+                async def fn(*args):
                     # The main idea here is to call the original function
                     # and append some updates to keep the "Submit" button
                     # hidden and the "Stop" button visible
-                    # The 'finally' block hides the "Stop" button and
-                    # shows the "submit" button. Having a 'finally' block
-                    # will make sure the UI is "reset" even if there is an exception
-                    try:
-                        for output in self.fn(*args):
-                            if len(self.output_components) == 1 and not self.batch:
-                                output = [output]
-                            output = list(output)
-                            yield output + [
-                                Button.update(visible=False),
-                                Button.update(visible=True),
-                            ]
-                    finally:
-                        yield [
-                            {"__type__": "generic_update"}
-                            for _ in self.output_components
-                        ] + [Button.update(visible=True), Button.update(visible=False)]
+
+                    if inspect.isasyncgenfunction(self.fn):
+                        iterator = self.fn(*args)
+                    else:
+                        iterator = utils.SyncToAsyncIterator(
+                            self.fn(*args), limiter=self.limiter
+                        )
+                    async for output in iterator:
+                        if len(self.output_components) == 1 and not self.batch:
+                            output = [output]
+                        output = list(output)
+                        yield output + [
+                            Button.update(visible=False),
+                            Button.update(visible=True),
+                        ]
 
                 extra_output = [submit_btn, stop_btn]
             pred = submit_btn.click(
@@ -643,6 +649,19 @@ def fn(*args):
                 batch=self.batch,
                 max_batch_size=self.max_batch_size,
             )
+            if extra_output:
+
+                def cleanup():
+                    return [
+                        {"__type__": "generic_update"} for _ in self.output_components
+                    ] + [Button.update(visible=True), Button.update(visible=False)]
+
+                pred.then(
+                    cleanup,
+                    inputs=None,
+                    outputs=self.output_components + extra_output,
+                    queue=False,
+                )
             if stop_btn:
                 submit_btn.click(
                     lambda: (
diff --git a/gradio/utils.py b/gradio/utils.py
index cc35d1ce35..841c0cae89 100644
--- a/gradio/utils.py
+++ b/gradio/utils.py
@@ -37,16 +37,17 @@
 )
 
 import aiohttp
+import anyio
 import httpx
 import matplotlib
 import requests
 from markdown_it import MarkdownIt
-from pygments import highlight
-from pygments.formatters import HtmlFormatter
-from pygments.lexers import get_lexer_by_name
 from mdit_py_plugins.dollarmath.index import dollarmath_plugin
 from mdit_py_plugins.footnote.index import footnote_plugin
 from pydantic import BaseModel, parse_obj_as
+from pygments import highlight
+from pygments.formatters import HtmlFormatter
+from pygments.lexers import get_lexer_by_name
 
 import gradio
 from gradio.context import Context
@@ -486,7 +487,8 @@ def run_coro_in_background(func: Callable, *args, **kwargs):
     return event_loop.create_task(func(*args, **kwargs))
 
 
-def async_iteration(iterator):
+def run_sync_iterator_async(iterator):
+    """Helper for yielding StopAsyncIteration from sync iterators."""
     try:
         return next(iterator)
     except StopIteration:
@@ -494,6 +496,27 @@ def async_iteration(iterator):
         raise StopAsyncIteration() from None
 
 
+class SyncToAsyncIterator:
+    """Treat a synchronous iterator as async one."""
+
+    def __init__(self, iterator, limiter) -> None:
+        self.iterator = iterator
+        self.limiter = limiter
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        return await anyio.to_thread.run_sync(
+            run_sync_iterator_async, self.iterator, limiter=self.limiter
+        )
+
+
+async def async_iteration(iterator):
+    # anext not introduced until 3.10 :(
+    return await iterator.__anext__()
+
+
 class AsyncRequest:
     """
     The AsyncRequest class is a low-level API that allow you to create asynchronous HTTP requests without a context manager.
diff --git a/test/test_blocks.py b/test/test_blocks.py
index dc1b1f21f4..2dd7cdaf5b 100644
--- a/test/test_blocks.py
+++ b/test/test_blocks.py
@@ -264,6 +264,101 @@ async def test_restart_after_close(self):
                     completed = True
             assert msg["output"]["data"][0] == "Victor"
 
+    @pytest.mark.asyncio
+    async def test_async_generators(self):
+        async def async_iteration(count: int):
+            for i in range(count):
+                yield i
+                await asyncio.sleep(0.2)
+
+        def iteration(count: int):
+            for i in range(count):
+                yield i
+                time.sleep(0.2)
+
+        with gr.Blocks() as demo:
+            with gr.Row():
+                with gr.Column():
+                    num1 = gr.Number(value=4, precision=0)
+                    o1 = gr.Number()
+                    async_iterate = gr.Button(value="Async Iteration")
+                    async_iterate.click(async_iteration, num1, o1)
+                with gr.Column():
+                    num2 = gr.Number(value=4, precision=0)
+                    o2 = gr.Number()
+                    iterate = gr.Button(value="Iterate")
+                    iterate.click(iteration, num2, o2)
+
+        demo.queue(concurrency_count=2).launch(prevent_thread_lock=True)
+
+        def _get_ws_pred(data, fn_index):
+            async def wrapped():
+                async with websockets.connect(
+                    f"{demo.local_url.replace('http', 'ws')}queue/join"
+                ) as ws:
+                    completed = False
+                    while not completed:
+                        msg = json.loads(await ws.recv())
+                        if msg["msg"] == "send_data":
+                            await ws.send(
+                                json.dumps({"data": [data], "fn_index": fn_index})
+                            )
+                        if msg["msg"] == "send_hash":
+                            await ws.send(
+                                json.dumps(
+                                    {"fn_index": fn_index, "session_hash": "shdce"}
+                                )
+                            )
+                        if msg["msg"] == "process_completed":
+                            completed = True
+                    assert msg["output"]["data"][0] == data - 1
+
+            return wrapped
+
+        try:
+            await asyncio.gather(_get_ws_pred(3, 0)(), _get_ws_pred(4, 1)())
+        finally:
+            demo.close()
+    
+    @pytest.mark.asyncio
+    async def test_sync_generators(self):
+
+        def generator(string):
+            for char in string:
+                yield char
+
+        demo = gr.Interface(generator, "text", "text")
+        demo.queue().launch(prevent_thread_lock=True)
+
+        async def _get_ws_pred(data, fn_index):
+            outputs = []
+            async with websockets.connect(
+                f"{demo.local_url.replace('http', 'ws')}queue/join"
+            ) as ws:
+                completed = False
+                while not completed:
+                    msg = json.loads(await ws.recv())
+                    if msg["msg"] == "send_data":
+                        await ws.send(
+                            json.dumps({"data": [data], "fn_index": fn_index})
+                        )
+                    if msg["msg"] == "send_hash":
+                        await ws.send(
+                            json.dumps(
+                                {"fn_index": fn_index, "session_hash": "shdce"}
+                            )
+                        )
+                    if msg['msg'] in ["process_generating"]:
+                        outputs.append(msg['output']['data'])
+                    if msg["msg"] == "process_completed":
+                        completed = True
+            return outputs
+        try:
+            output = await _get_ws_pred(fn_index=0, data="abc")
+            assert [o[0] for o in output] == ['a', 'b', 'c']
+        finally:
+            demo.close()
+
     def test_socket_reuse(self):
 
         try:
@@ -1177,9 +1272,9 @@ def generate(x):
                     ]
                 if msg["msg"] == "process_completed":
                     assert msg["output"]["data"] == [
-                        {"__type__": "update"},
-                        {"visible": True, "__type__": "update"},
+                        '3',
                         {"visible": False, "__type__": "update"},
+                        {"visible": True, "__type__": "update"},
                     ]
                     completed = True
             assert checked_iteration
diff --git a/test/test_routes.py b/test/test_routes.py
index daef41cf38..1b5e155897 100644
--- a/test/test_routes.py
+++ b/test/test_routes.py
@@ -396,69 +396,6 @@ def test_mount_gradio_app_raises_error_if_event_queued_but_queue_disabled(self):
         demo.close()
 
 
-class TestGeneratorRoutes:
-    def test_generator(self):
-        def generator(string):
-            for char in string:
-                yield char
-
-        io = Interface(generator, "text", "text")
-        app, _, _ = io.queue().launch(prevent_thread_lock=True)
-        client = TestClient(app)
-
-        response = client.post(
-            "/api/predict/",
-            json={"data": ["abc"], "fn_index": 0, "session_hash": "11"},
-            headers={"Authorization": f"Bearer {app.queue_token}"},
-        )
-        output = dict(response.json())
-        assert output["data"][0] == "a"
-
-        response = client.post(
-            "/api/predict/",
-            json={"data": ["abc"], "fn_index": 0, "session_hash": "11"},
-            headers={"Authorization": f"Bearer {app.queue_token}"},
-        )
-        output = dict(response.json())
-        assert output["data"][0] == "b"
-
-        response = client.post(
-            "/api/predict/",
-            json={"data": ["abc"], "fn_index": 0, "session_hash": "11"},
-            headers={"Authorization": f"Bearer {app.queue_token}"},
-        )
-        output = dict(response.json())
-        assert output["data"][0] == "c"
-
-        response = client.post(
-            "/api/predict/",
-            json={"data": ["abc"], "fn_index": 0, "session_hash": "11"},
-            headers={"Authorization": f"Bearer {app.queue_token}"},
-        )
-        output = dict(response.json())
-        assert output["data"] == [
-            {"__type__": "update"},
-            {"__type__": "update", "visible": True},
-            {"__type__": "update", "visible": False},
-        ]
-
-        response = client.post(
-            "/api/predict/",
-            json={"data": ["abc"], "fn_index": 0, "session_hash": "11"},
-            headers={"Authorization": f"Bearer {app.queue_token}"},
-        )
-        output = dict(response.json())
-        assert output["data"][0] is None
-
-        response = client.post(
-            "/api/predict/",
-            json={"data": ["abc"], "fn_index": 0, "session_hash": "11"},
-            headers={"Authorization": f"Bearer {app.queue_token}"},
-        )
-        output = dict(response.json())
-        assert output["data"][0] == "a"
-
-
 class TestApp:
     def test_create_app(self):
         app = routes.App.create_app(Interface(lambda x: x, "text", "text"))


Code review comments:
diff_hunk: @@ -643,6 +649,19 @@ def fn(*args):
                 batch=self.batch,
                 max_batch_size=self.max_batch_size,
             )
+            if extra_output:
+
+                def cleanup():
+                    return [
+                        {"__type__": "generic_update"} for _ in self.output_components
+                    ] + [Button.update(visible=True), Button.update(visible=False)]
+
+                pred.then(
path: gradio/interface.py
text: @author:
Refactored the Interface cancel button to use `then` as opposed to a `finally` block. Helps prevent the issues with the generator not being cleaned up