From c60bcb5ecec0ae5da95ed4554947f6860a7c03b2 Mon Sep 17 00:00:00 2001 From: Aniket Maurya Date: Tue, 17 Dec 2024 13:51:50 +0000 Subject: [PATCH 1/2] set stream using method --- src/litserve/server.py | 5 ++--- src/litserve/specs/base.py | 3 +++ src/litserve/specs/openai.py | 3 +++ 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/litserve/server.py b/src/litserve/server.py index 42e4f097..116ccfb9 100644 --- a/src/litserve/server.py +++ b/src/litserve/server.py @@ -43,7 +43,6 @@ from litserve.loops import LitLoop, get_default_loop, inference_worker from litserve.middlewares import MaxSizeMiddleware, RequestCountMiddleware from litserve.python_client import client_template -from litserve.specs import OpenAISpec from litserve.specs.base import LitSpec from litserve.utils import LitAPIStatus, WorkerSetupStatus, call_after_stream @@ -146,8 +145,8 @@ def __init__( raise ValueError("batch_timeout must be less than timeout") if max_batch_size <= 0: raise ValueError("max_batch_size must be greater than 0") - if isinstance(spec, OpenAISpec): - stream = True + if isinstance(spec, LitSpec): + spec.set_stream(self) if loop is None: loop = "auto" diff --git a/src/litserve/specs/base.py b/src/litserve/specs/base.py index 6f5afc7b..03f8a672 100644 --- a/src/litserve/specs/base.py +++ b/src/litserve/specs/base.py @@ -26,6 +26,9 @@ def __init__(self): self._server: LitServer = None + def set_stream(self, lit_server: "LitServer"): + pass + def pre_setup(self, lit_api: "LitAPI"): pass diff --git a/src/litserve/specs/openai.py b/src/litserve/specs/openai.py index f7447d88..c556c174 100644 --- a/src/litserve/specs/openai.py +++ b/src/litserve/specs/openai.py @@ -262,6 +262,9 @@ def __init__( self.add_endpoint("/v1/chat/completions", self.chat_completion, ["POST"]) self.add_endpoint("/v1/chat/completions", self.options_chat_completions, ["OPTIONS"]) + def set_stream(self, lit_server: "LitServer"): + lit_server.stream = True + def pre_setup(self, lit_api: "LitAPI"): from litserve import LitAPI From 373c8b489fbd37ccb327bd762fe06eb8bb359019 Mon Sep 17 00:00:00 2001 From: Aniket Maurya Date: Tue, 17 Dec 2024 13:57:19 +0000 Subject: [PATCH 2/2] fix --- src/litserve/server.py | 2 +- src/litserve/specs/base.py | 5 +++-- src/litserve/specs/openai.py | 5 +++-- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/litserve/server.py b/src/litserve/server.py index 116ccfb9..f676c7f1 100644 --- a/src/litserve/server.py +++ b/src/litserve/server.py @@ -146,7 +146,7 @@ def __init__( if max_batch_size <= 0: raise ValueError("max_batch_size must be greater than 0") if isinstance(spec, LitSpec): - spec.set_stream(self) + stream = spec.stream if loop is None: loop = "auto" diff --git a/src/litserve/specs/base.py b/src/litserve/specs/base.py index 03f8a672..2db26e95 100644 --- a/src/litserve/specs/base.py +++ b/src/litserve/specs/base.py @@ -26,8 +26,9 @@ def __init__(self): self._server: LitServer = None - def set_stream(self, lit_server: "LitServer"): - pass + @property + def stream(self): + return False def pre_setup(self, lit_api: "LitAPI"): pass diff --git a/src/litserve/specs/openai.py b/src/litserve/specs/openai.py index c556c174..90a63618 100644 --- a/src/litserve/specs/openai.py +++ b/src/litserve/specs/openai.py @@ -262,8 +262,9 @@ def __init__( self.add_endpoint("/v1/chat/completions", self.chat_completion, ["POST"]) self.add_endpoint("/v1/chat/completions", self.options_chat_completions, ["OPTIONS"]) - def set_stream(self, lit_server: "LitServer"): - lit_server.stream = True + @property + def stream(self): + return True def pre_setup(self, lit_api: "LitAPI"): from litserve import LitAPI