sglang 0.1.16__py3-none-any.whl → 0.1.18__py3-none-any.whl

sglang/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.16"
+__version__ = "0.1.18"
 
 # SGL API Components
 from sglang.api import (
@@ -24,6 +24,7 @@ from sglang.api import (
 
 # SGL Backends
 from sglang.backend.anthropic import Anthropic
+from sglang.backend.litellm import LiteLLM
 from sglang.backend.openai import OpenAI
 from sglang.backend.runtime_endpoint import RuntimeEndpoint
 from sglang.backend.vertexai import VertexAI
@@ -35,6 +36,7 @@ from sglang.global_config import global_config
 __all__ = [
     "global_config",
     "Anthropic",
+    "LiteLLM",
     "OpenAI",
     "RuntimeEndpoint",
     "VertexAI",

sglang/api.py

@@ -1,4 +1,4 @@
-"""Some Public API Definitions"""
+"""Public APIs of the language."""
 
 import os
 import re
@@ -20,13 +20,13 @@ from sglang.lang.ir import (
 
 
 def function(
-    func: Optional[Callable] = None, api_num_spec_tokens: Optional[int] = None
+    func: Optional[Callable] = None, num_api_spec_tokens: Optional[int] = None
 ):
     if func:
-        return SglFunction(func, api_num_spec_tokens=api_num_spec_tokens)
+        return SglFunction(func, num_api_spec_tokens=num_api_spec_tokens)
 
     def decorator(func):
-        return SglFunction(func, api_num_spec_tokens=api_num_spec_tokens)
+        return SglFunction(func, num_api_spec_tokens=num_api_spec_tokens)
 
     return decorator
 
@@ -43,14 +43,14 @@ def set_default_backend(backend: BaseBackend):
     global_config.default_backend = backend
 
 
-def flush_cache(backend: BaseBackend = None):
+def flush_cache(backend: Optional[BaseBackend] = None):
     backend = backend or global_config.default_backend
     if backend is None:
         return False
     return backend.flush_cache()
 
 
-def get_server_args(backend: BaseBackend = None):
+def get_server_args(backend: Optional[BaseBackend] = None):
     backend = backend or global_config.default_backend
     if backend is None:
         return None
@@ -158,7 +158,7 @@ def video(path: str, num_frames: int):
 
 def select(
     name: Optional[str] = None,
-    choices: List[str] = None,
+    choices: Optional[List[str]] = None,
     temperature: float = 0.0,
 ):
     assert choices is not None

sglang/backend/anthropic.py

@@ -74,4 +74,4 @@ class Anthropic(BaseBackend):
             **sampling_params.to_anthropic_kwargs(),
         ) as stream:
             for text in stream.text_stream:
-                yield text, {}
+                yield text, {}

sglang/backend/litellm.py

@@ -0,0 +1,90 @@
+from typing import Mapping, Optional
+
+from sglang.backend.base_backend import BaseBackend
+from sglang.lang.chat_template import get_chat_template_by_model_path
+from sglang.lang.interpreter import StreamExecutor
+from sglang.lang.ir import SglSamplingParams
+
+try:
+    import litellm
+except ImportError as e:
+    litellm = e
+    litellm.num_retries = 1
+
+
+class LiteLLM(BaseBackend):
+    def __init__(
+        self,
+        model_name,
+        chat_template=None,
+        api_key=None,
+        organization: Optional[str] = None,
+        base_url: Optional[str] = None,
+        timeout: Optional[float] = 600,
+        max_retries: Optional[int] = litellm.num_retries,
+        default_headers: Optional[Mapping[str, str]] = None,
+    ):
+        super().__init__()
+
+        if isinstance(litellm, Exception):
+            raise litellm
+
+        self.model_name = model_name
+
+        self.chat_template = chat_template or get_chat_template_by_model_path(
+            model_name
+        )
+
+        self.client_params = {
+            "api_key": api_key,
+            "organization": organization,
+            "base_url": base_url,
+            "timeout": timeout,
+            "max_retries": max_retries,
+            "default_headers": default_headers,
+        }
+
+    def get_chat_template(self):
+        return self.chat_template
+
+    def generate(
+        self,
+        s: StreamExecutor,
+        sampling_params: SglSamplingParams,
+    ):
+        if s.messages_:
+            messages = s.messages_
+        else:
+            messages = [{"role": "user", "content": s.text_}]
+
+        ret = litellm.completion(
+            model=self.model_name,
+            messages=messages,
+            **self.client_params,
+            **sampling_params.to_anthropic_kwargs(),
+        )
+        comp = ret.choices[0].message.content
+
+        return comp, {}
+
+    def generate_stream(
+        self,
+        s: StreamExecutor,
+        sampling_params: SglSamplingParams,
+    ):
+        if s.messages_:
+            messages = s.messages_
+        else:
+            messages = [{"role": "user", "content": s.text_}]
+
+        ret = litellm.completion(
+            model=self.model_name,
+            messages=messages,
+            stream=True,
+            **self.client_params,
+            **sampling_params.to_litellm_kwargs(),
+        )
+        for chunk in ret:
+            text = chunk.choices[0].delta.content
+            if text is not None:
+                yield text, {}

sglang/backend/openai.py

@@ -1,5 +1,7 @@
+import dataclasses
 import logging
 import time
+import warnings
 from typing import Callable, List, Optional, Union
 
 import numpy as np
@@ -41,6 +43,15 @@ INSTRUCT_MODEL_NAMES = [
 ]
 
 
+@dataclasses.dataclass
+class TokenUsage:
+    prompt_tokens: int
+    completion_tokens: int
+
+    def reset(self):
+        self.prompt_tokens = self.completion_tokens = 0
+
+
 class OpenAI(BaseBackend):
     def __init__(
         self,
@@ -80,40 +91,92 @@ class OpenAI(BaseBackend):
             else:
                 self.is_chat_model = True
 
-        self.chat_begin_str = self.chat_template.role_prefix_and_suffix["assistant"][0]
+        self.chat_prefix = self.chat_template.role_prefix_and_suffix["assistant"][0]
+
+        # Usage
+        self.token_usage = TokenUsage(0, 0)
+
+        # API speculative execution
+        # TODO(ying): This does not support multi-threading (run_batch)
+        self.spec_kwargs = {}
+        self.spec_format = []
+        self.spec_max_num_tries = 3
 
     def get_chat_template(self):
         return self.chat_template
 
+    def _prepare_spec_execution(
+        self,
+        sampling_params: SglSamplingParams,
+        num_api_spec_tokens: int,
+        spec_var_name: str,
+    ):
+        if "max_tokens" not in self.spec_kwargs:
+            self.spec_kwargs["max_tokens"] = num_api_spec_tokens
+        else:
+            assert self.spec_kwargs["max_tokens"] == num_api_spec_tokens
+
+        params = sampling_params.to_openai_kwargs()
+        for key, value in params.items():
+            if key in ["stop"]:
+                continue
+            if key in ["max_tokens"]:
+                warnings.warn(
+                    "The parameter max_tokens will be overwritten by speculated number of tokens."
+                )
+                continue
+            if key not in self.spec_kwargs:
+                self.spec_kwargs[key] = value
+            else:
+                assert (
+                    value == self.spec_kwargs[key]
+                ), "sampling parameters should be consistent if turn on api speculative execution."
+        self.spec_format.append(
+            {"text": "", "stop": params["stop"], "name": spec_var_name}
+        )
+        return "", {}
+
     def generate(
         self,
         s: StreamExecutor,
         sampling_params: SglSamplingParams,
+        spec_var_name: str = None,
     ):
         if sampling_params.dtype is None:
             if self.is_chat_model:
-                if not s.text_.endswith(self.chat_begin_str):
-                    raise RuntimeError(
-                        "This use case is not supported. "
-                        "For OpenAI chat models, sgl.gen must be right after sgl.assistant"
+                if s.num_api_spec_tokens is None:
+                    if not s.text_.endswith(self.chat_prefix):
+                        raise RuntimeError(
+                            "This use case is not supported if api speculative execution is off. "
+                            "For OpenAI chat models, sgl.gen must be right after sgl.assistant. "
+                            "Example of adding api speculative execution: @function(num_api_spec_tokens=128)."
+                        )
+                    prompt = s.messages_
+                else:
+                    return self._prepare_spec_execution(
+                        sampling_params, s.num_api_spec_tokens, spec_var_name
                     )
-                prompt = s.messages_
             else:
                 prompt = s.text_
 
             kwargs = sampling_params.to_openai_kwargs()
             comp = openai_completion(
                 client=self.client,
+                token_usage=self.token_usage,
                 is_chat=self.is_chat_model,
                 model=self.model_name,
                 prompt=prompt,
                 **kwargs,
             )
         elif sampling_params.dtype in [str, "str", "string"]:
+            assert (
+                not self.is_chat_model
+            ), "constrained type not supported on chat model"
             kwargs = sampling_params.to_openai_kwargs()
             kwargs.pop("stop")
             comp = openai_completion(
                 client=self.client,
+                token_usage=self.token_usage,
                 is_chat=self.is_chat_model,
                 model=self.model_name,
                 prompt=s.text_ + '"',
@@ -122,10 +185,14 @@ class OpenAI(BaseBackend):
             )
             comp = '"' + comp + '"'
         elif sampling_params.dtype in [int, "int"]:
+            assert (
+                not self.is_chat_model
+            ), "constrained type not supported on chat model"
             kwargs = sampling_params.to_openai_kwargs()
             kwargs.pop("stop")
             comp = openai_completion(
                 client=self.client,
+                token_usage=self.token_usage,
                 is_chat=self.is_chat_model,
                 model=self.model_name,
                 prompt=s.text_,
@@ -138,6 +205,63 @@ class OpenAI(BaseBackend):
 
         return comp, {}
 
+    def spec_fill(self, value: str):
+        assert self.is_chat_model
+        self.spec_format.append({"text": value, "stop": None, "name": None})
+
+    def spec_pattern_match(self, comp):
+        for i, term in enumerate(self.spec_format):
+            text = term["text"]
+            if text != "":
+                if comp.startswith(text):
+                    comp = comp[len(text) :]
+                else:
+                    return False
+            else:
+                pos = comp.find(term["stop"])
+                if pos != -1:
+                    term["text"] = comp[:pos]
+                    comp = comp[pos:]
+                else:
+                    if i == len(self.spec_format) - 1:
+                        term["text"] = comp
+                    else:
+                        return False
+        return True
+
+    def role_end_generate(
+        self,
+        s: StreamExecutor,
+    ):
+        if s.num_api_spec_tokens is None or not s.text_.endswith(self.chat_prefix):
+            return
+
+        comp = ""
+        if not all(x["name"] is None for x in self.spec_format):
+            # TODO(ying): throw errors or warnings
+            for i in range(self.spec_max_num_tries):
+                comp = openai_completion(
+                    client=self.client,
+                    token_usage=self.token_usage,
+                    is_chat=self.is_chat_model,
+                    model=self.model_name,
+                    prompt=s.messages_,
+                    **self.spec_kwargs,
+                )
+                if self.spec_pattern_match(comp):
+                    break
+
+        for term in self.spec_format:
+            s.text_ += term["text"]
+            name = term["name"]
+            if name is not None:
+                s.variables[name] = term["text"]
+                s.meta_info[name] = {}
+                s.variable_event[name].set()
+
+        self.spec_kwargs = {}
+        self.spec_format = []
+
     def generate_stream(
         self,
         s: StreamExecutor,
@@ -145,7 +269,7 @@ class OpenAI(BaseBackend):
     ):
         if sampling_params.dtype is None:
             if self.is_chat_model:
-                if not s.text_.endswith(self.chat_begin_str):
+                if not s.text_.endswith(self.chat_prefix):
                     raise RuntimeError(
                         "This use case is not supported. "
                         "For OpenAI chat models, sgl.gen must be right after sgl.assistant"
@@ -157,6 +281,7 @@ class OpenAI(BaseBackend):
             kwargs = sampling_params.to_openai_kwargs()
             generator = openai_completion_stream(
                 client=self.client,
+                token_usage=self.token_usage,
                 is_chat=self.is_chat_model,
                 model=self.model_name,
                 prompt=prompt,
@@ -202,6 +327,8 @@ class OpenAI(BaseBackend):
             )
             ret_str = ret.choices[0].text
             ret_token = self.tokenizer.encode(ret_str)[0]
+            self.token_usage.prompt_tokens += ret.usage.prompt_tokens
+            self.token_usage.completion_tokens = ret.usage.completion_tokens
 
             # TODO:
             # 1. return logits as the scores
@@ -231,7 +358,9 @@ class OpenAI(BaseBackend):
         return decision, scores, None, None
 
 
-def openai_completion(client, retries=3, is_chat=None, prompt=None, **kwargs):
+def openai_completion(
+    client, token_usage, is_chat=None, retries=3, prompt=None, **kwargs
+):
     for attempt in range(retries):
         try:
             if is_chat:
@@ -245,6 +374,9 @@ def openai_completion(client, retries=3, is_chat=None, prompt=None, **kwargs):
                     comp = [c.text for c in ret.choices]
                 else:
                     comp = ret.choices[0].text
+
+            token_usage.prompt_tokens += ret.usage.prompt_tokens
+            token_usage.completion_tokens += ret.usage.completion_tokens
             break
         except (openai.APIError, openai.APIConnectionError, openai.RateLimitError) as e:
             logger.error(f"OpenAI Error: {e}. Waiting 5 seconds...")
@@ -258,16 +390,23 @@ def openai_completion(client, retries=3, is_chat=None, prompt=None, **kwargs):
     return comp
 
 
-def openai_completion_stream(client, retries=3, is_chat=None, prompt=None, **kwargs):
+def openai_completion_stream(
+    client, token_usage, is_chat=None, retries=3, prompt=None, **kwargs
+):
     for attempt in range(retries):
         try:
             if is_chat:
                 if "stop" in kwargs and kwargs["stop"] is None:
                     kwargs.pop("stop")
                 generator = client.chat.completions.create(
-                    messages=prompt, stream=True, **kwargs
+                    messages=prompt,
+                    stream=True,
+                    stream_options={"include_usage": True},
+                    **kwargs,
                 )
                 for ret in generator:
+                    if len(ret.choices) == 0:
+                        continue
                     try:
                         content = ret.choices[0].delta.content
                     except IndexError:
@@ -275,11 +414,19 @@ def openai_completion_stream(client, retries=3, is_chat=None, prompt=None, **kwa
                     yield content or "", {}
             else:
                 generator = client.completions.create(
-                    prompt=prompt, stream=True, **kwargs
+                    prompt=prompt,
+                    stream=True,
+                    stream_options={"include_usage": True},
+                    **kwargs,
                 )
                 for ret in generator:
+                    if len(ret.choices) == 0:
+                        continue
                     content = ret.choices[0].text
                     yield content or "", {}
+
+            token_usage.prompt_tokens += ret.usage.prompt_tokens
+            token_usage.completion_tokens += ret.usage.completion_tokens
             break
         except (openai.APIError, openai.APIConnectionError, openai.RateLimitError) as e:
             logger.error(f"OpenAI Error: {e}. Waiting 5 seconds...")

sglang/backend/runtime_endpoint.py

@@ -34,7 +34,7 @@ class RuntimeEndpoint(BaseBackend):
             api_key=self.api_key,
             verify=self.verify,
         )
-        assert res.status_code == 200
+        self._assert_success(res)
         self.model_info = res.json()
 
         self.chat_template = get_chat_template_by_model_path(
@@ -50,7 +50,7 @@ class RuntimeEndpoint(BaseBackend):
             auth_token=self.auth_token,
             verify=self.verify,
         )
-        return res.status_code == 200
+        self._assert_success(res)
 
     def get_server_args(self):
         res = http_request(
@@ -58,6 +58,7 @@ class RuntimeEndpoint(BaseBackend):
             auth_token=self.auth_token,
             verify=self.verify,
         )
+        self._assert_success(res)
         return res.json()
 
     def get_chat_template(self):
@@ -71,7 +72,7 @@ class RuntimeEndpoint(BaseBackend):
             api_key=self.api_key,
             verify=self.verify,
         )
-        assert res.status_code == 200
+        self._assert_success(res)
 
     def commit_lazy_operations(self, s: StreamExecutor):
         data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}}
@@ -83,7 +84,7 @@ class RuntimeEndpoint(BaseBackend):
             api_key=self.api_key,
             verify=self.verify,
         )
-        assert res.status_code == 200
+        self._assert_success(res)
 
     def fill_image(self, s: StreamExecutor):
         data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}}
@@ -95,7 +96,7 @@ class RuntimeEndpoint(BaseBackend):
             api_key=self.api_key,
             verify=self.verify,
         )
-        assert res.status_code == 200
+        self._assert_success(res)
 
     def generate(
         self,
@@ -133,6 +134,8 @@ class RuntimeEndpoint(BaseBackend):
             api_key=self.api_key,
             verify=self.verify,
         )
+        self._assert_success(res)
+
         obj = res.json()
         comp = obj["text"]
         return comp, obj["meta_info"]
@@ -167,7 +170,7 @@ class RuntimeEndpoint(BaseBackend):
         data["stream"] = True
         self._add_images(s, data)
 
-        response = http_request(
+        res = http_request(
             self.base_url + "/generate",
             json=data,
             stream=True,
@@ -175,10 +178,11 @@ class RuntimeEndpoint(BaseBackend):
             api_key=self.api_key,
             verify=self.verify,
         )
+        self._assert_success(res)
         pos = 0
 
         incomplete_text = ""
-        for chunk in response.iter_lines(decode_unicode=False):
+        for chunk in res.iter_lines(decode_unicode=False):
             chunk = chunk.decode("utf-8")
             if chunk and chunk.startswith("data:"):
                 if chunk == "data: [DONE]":
@@ -211,7 +215,7 @@ class RuntimeEndpoint(BaseBackend):
             api_key=self.api_key,
             verify=self.verify,
         )
-        assert res.status_code == 200
+        self._assert_success(res)
         prompt_len = res.json()["meta_info"]["prompt_tokens"]
 
         # Compute logprob
@@ -229,7 +233,7 @@ class RuntimeEndpoint(BaseBackend):
             api_key=self.api_key,
             verify=self.verify,
         )
-        assert res.status_code == 200
+        self._assert_success(res)
         obj = res.json()
         normalized_prompt_logprobs = [
             r["meta_info"]["normalized_prompt_logprob"] for r in obj
@@ -253,9 +257,13 @@ class RuntimeEndpoint(BaseBackend):
             api_key=self.api_key,
             verify=self.verify,
         )
-        assert res.status_code == 200
+        self._assert_success(res)
 
     def _add_images(self, s: StreamExecutor, data):
         if s.images_:
             assert len(s.images_) == 1, "Only support one image."
             data["image_data"] = s.images_[0][1]
+
+    def _assert_success(self, res):
+        if res.status_code != 200:
+            raise RuntimeError(res.json())