Skip to content

Commit 636c9fd

Browse files
authored
Set LitServer.stream using LitSpec.stream (#398)
* set stream using method * fix
1 parent fcce973 commit 636c9fd

File tree

3 files changed

+10
-3
lines changed

3 files changed

+10
-3
lines changed

src/litserve/server.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,6 @@
4343
from litserve.loops import LitLoop, get_default_loop, inference_worker
4444
from litserve.middlewares import MaxSizeMiddleware, RequestCountMiddleware
4545
from litserve.python_client import client_template
46-
from litserve.specs import OpenAISpec
4746
from litserve.specs.base import LitSpec
4847
from litserve.utils import LitAPIStatus, WorkerSetupStatus, call_after_stream
4948

@@ -146,8 +145,8 @@ def __init__(
146145
raise ValueError("batch_timeout must be less than timeout")
147146
if max_batch_size <= 0:
148147
raise ValueError("max_batch_size must be greater than 0")
149-
if isinstance(spec, OpenAISpec):
150-
stream = True
148+
if isinstance(spec, LitSpec):
149+
stream = spec.stream
151150

152151
if loop is None:
153152
loop = "auto"

src/litserve/specs/base.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,10 @@ def __init__(self):
2626

2727
self._server: LitServer = None
2828

29+
@property
30+
def stream(self):
31+
return False
32+
2933
def pre_setup(self, lit_api: "LitAPI"):
3034
pass
3135

src/litserve/specs/openai.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -262,6 +262,10 @@ def __init__(
262262
self.add_endpoint("/v1/chat/completions", self.chat_completion, ["POST"])
263263
self.add_endpoint("/v1/chat/completions", self.options_chat_completions, ["OPTIONS"])
264264

265+
@property
266+
def stream(self):
267+
return True
268+
265269
def pre_setup(self, lit_api: "LitAPI"):
266270
from litserve import LitAPI
267271

0 commit comments

Comments
 (0)