diff --git a/.changeset/five-days-join.md b/.changeset/five-days-join.md new file mode 100644 index 0000000000..1be1750c23 --- /dev/null +++ b/.changeset/five-days-join.md @@ -0,0 +1,5 @@ +--- +"gradio": patch +--- + +fix:Fix crash when a streaming `gr.ChatInterface` function yields nothing diff --git a/gradio/chat_interface.py b/gradio/chat_interface.py index f6acd5f4e9..46a7dd8de2 100644 --- a/gradio/chat_interface.py +++ b/gradio/chat_interface.py @@ -44,7 +44,7 @@ from gradio.events import Dependency, EditData, SelectData from gradio.flagging import ChatCSVLogger from gradio.helpers import create_examples as Examples # noqa: N812 -from gradio.helpers import special_args, update +from gradio.helpers import skip, special_args, update from gradio.i18n import I18nData from gradio.layouts import Accordion, Column, Group, Row @@ -972,8 +972,8 @@ async def _stream_fn( yield first_response, history_ else: yield first_response, history_, *additional_outputs - except StopIteration: - yield None, history + except StopAsyncIteration: + yield None, history, *[skip() for _ in self.additional_outputs] async for response in generator: if self.additional_outputs: response, *additional_outputs = response diff --git a/test/test_chat_interface.py b/test/test_chat_interface.py index 527837668b..1be3d8e54c 100644 --- a/test/test_chat_interface.py +++ b/test/test_chat_interface.py @@ -33,6 +33,18 @@ async def async_stream(message, history): yield message[: i + 1] +def stream_or_nothing(message, history): + if message == "skip": + return + yield f"history_len={len(history)}" + + +async def async_stream_or_nothing(message, history): + if message == "skip": + return + yield f"history_len={len(history)}" + + def count(message, history): return str(len(history)) @@ -298,6 +310,40 @@ def test_streaming_api_async(self, connect): wait([job]) assert job.outputs() == ["h", "he", "hel", "hell", "hello"] + def test_streaming_api_empty_generator(self, connect): + chatbot = gr.ChatInterface(stream_or_nothing).queue() + with connect(chatbot) as client: + assert client.predict("hello") == "history_len=0" + assert client.predict("skip") is None + # The skipped turn must leave the history intact and still record + # the user message, so the next turn sees three messages. + assert client.predict("hello") == "history_len=3" + + def test_streaming_api_empty_generator_async(self, connect): + chatbot = gr.ChatInterface(async_stream_or_nothing).queue() + with connect(chatbot) as client: + assert client.predict("hello") == "history_len=0" + assert client.predict("skip") is None + assert client.predict("hello") == "history_len=3" + + def test_streaming_api_empty_generator_with_additional_outputs(self, connect): + def stream_or_nothing_with_extra(message, history): + if message == "skip": + return + yield message, "extra" + + with gr.Blocks() as demo: + code = gr.Code(render=False) + gr.ChatInterface( + stream_or_nothing_with_extra, + additional_outputs=[code], + api_name="chat", + ) + code.render() + with connect(demo) as client: + assert client.predict("hello", api_name="/chat") == ("hello", "extra") + assert client.predict("skip", api_name="/chat") == (None, gr.skip()) + def test_non_streaming_api(self, connect): chatbot = gr.ChatInterface(double) with connect(chatbot) as client: