diff --git a/.changeset/thirty-meals-invent.md b/.changeset/thirty-meals-invent.md new file mode 100644 index 0000000000..59a2e5a3ac --- /dev/null +++ b/.changeset/thirty-meals-invent.md @@ -0,0 +1,5 @@ +--- +"gradio": minor +--- + +feat:Forward visitor requests when lazily caching examples diff --git a/gradio/helpers.py b/gradio/helpers.py index 1f5885045f..55c7c43297 100644 --- a/gradio/helpers.py +++ b/gradio/helpers.py @@ -366,9 +366,9 @@ def load_example_input(example_tuple): processed_example = self._get_processed_example(example_value) return utils.resolve_singleton(processed_example) - def load_example_output(example_tuple): + def load_example_output(example_tuple, request: routes.Request): example_id, _ = example_tuple - cached_outputs = self.load_from_cache(example_id) + cached_outputs = self.load_from_cache(example_id, request) return utils.resolve_singleton(cached_outputs) self.cache_event = self.load_input_event = self.dataset.click( @@ -509,11 +509,16 @@ async def _start_caching(self): if self.cache_examples is True: await self.cache() - async def cache(self, example_id: int | None = None) -> None: + async def cache( + self, + example_id: int | None = None, + request: routes.Request | None = None, + ) -> None: """ Caches examples so that their predictions can be shown immediately. Parameters: example_id: The id of the example to process (zero-indexed). If None, all examples are cached. + request: The request that triggered lazy caching, if any. """ if self.root_block is None: raise Error("Cannot cache examples if not in a Blocks context.") @@ -560,7 +565,7 @@ async def cache(self, example_id: int | None = None) -> None: prediction = await self.root_block.process_api( block_fn=self.root_block.default_config.fns[fn_index], inputs=processed_input, - request=None, + request=request, in_event_listener=self.cache_examples != "lazy", ) output = prediction["data"] @@ -579,14 +584,19 @@ async def cache(self, example_id: int | None = None) -> None: # Remove the "fake_event" to prevent bugs in loading interfaces from spaces self.root_block.default_config.fns.pop(fn_index) - def load_from_cache(self, example_id: int) -> list[Any]: + def load_from_cache( + self, + example_id: int, + request: routes.Request | None = None, + ) -> list[Any]: """Loads a particular cached example for the interface. Parameters: example_id: The id of the example to process (zero-indexed). + request: The request that triggered lazy caching, if any. """ cached_index = self._get_cached_index_if_cached(example_id) if cached_index is None: - client_utils.synchronize_async(self.cache, example_id) + client_utils.synchronize_async(self.cache, example_id, request) with open(self.cached_indices_file) as f: cached_index = len(f.readlines()) - 1 diff --git a/test/test_helpers.py b/test/test_helpers.py index b7ad2d3aa4..e9e30aac24 100644 --- a/test/test_helpers.py +++ b/test/test_helpers.py @@ -673,6 +673,33 @@ def image_identity(image, string): data = response.json()["data"] assert data[0]["path"].endswith("image.webp") + def test_lazy_cache_examples_preserves_request(self, patched_cache_folder): + def get_ip_token(value, request: gr.Request): + return request.headers.get("x-ip-token") if request else None + + with gr.Blocks() as demo: + text = gr.Textbox() + output = gr.Textbox() + gr.Examples( + examples=["hello"], + inputs=text, + outputs=output, + fn=get_ip_token, + cache_examples=True, + cache_mode="lazy", + api_name="load_example", + ) + + app = routes.App.create_app(demo) + with TestClient(app) as client: + response = client.post( + f"{API_PREFIX}/api/load_example/", + json={"data": [0]}, + headers={"x-ip-token": "visitor-token"}, + ) + + assert response.json()["data"] == ["visitor-token"] + def test_multiple_file_flagging(tmp_path, connect): with patch("gradio.utils.get_cache_folder", return_value=tmp_path):