diff --git a/src/litserve/loops/base.py b/src/litserve/loops/base.py index 87ac535d..7d85d0dd 100644 --- a/src/litserve/loops/base.py +++ b/src/litserve/loops/base.py @@ -264,6 +264,29 @@ def pre_setup(self, lit_api: LitAPI, spec: Optional[LitSpec]): return original = lit_api.unbatch.__code__ is LitAPI.unbatch.__code__ + if not lit_api.stream and any([ + inspect.isgeneratorfunction(lit_api.predict), + inspect.isgeneratorfunction(lit_api.encode_response), + (original or inspect.isgeneratorfunction(lit_api.unbatch)), + ]): + raise ValueError( + """When `stream=False`, `lit_api.predict`, `lit_api.encode_response` and `lit_api.unbatch` must not be + generator functions. + + Correct example: + + def predict(self, inputs): + ... + return {"output": output} + + Incorrect example: + + def predict(self, inputs): + ... + for i in range(max_token_length): + yield prediction + """ + ) if ( lit_api.stream and lit_api.max_batch_size > 1