chore: import upstream snapshot with attribution
PR Test (NPU) / check-changes (push) Has been cancelled
PR Test (NPU) / pr-gate (push) Has been cancelled
PR Test (NPU) / set-image-config (push) Has been cancelled
PR Test (NPU) / stage-b-test-1-npu-a2 (0) (push) Has been cancelled
PR Test (NPU) / stage-b-test-1-npu-a2 (1) (push) Has been cancelled
PR Test (NPU) / stage-b-test-2-npu-a2 (0) (push) Has been cancelled
PR Test (NPU) / stage-b-test-2-npu-a2 (1) (push) Has been cancelled
PR Test (NPU) / stage-b-test-4-npu-a3 (push) Has been cancelled
PR Test (NPU) / stage-b-test-16-npu-a3 (push) Has been cancelled
PR Test (NPU) / multimodal-gen-test-1-npu-a3 (push) Has been cancelled
PR Test (NPU) / multimodal-gen-test-2-npu-a3 (push) Has been cancelled
PR Test (Arm64) / pr-gate (push) Has been cancelled
PR Test (Arm64) / check-changes (push) Has been cancelled
PR Test (Arm64) / build-test (push) Has been cancelled
PR Test (sgl-router) / gate (push) Has been cancelled
PR Test (sgl-router) / tier-1 — lint (push) Has been cancelled
PR Test (sgl-router) / tier-2 — build + test (push) Has been cancelled
PR Test (sgl-router) / tier-3 — docker (placeholder) (push) Has been cancelled
PR Test (sgl-router) / tier-3 — k8s integration (push) Has been cancelled
PR Test (sgl-router) / tier-3 — e2e (push) Has been cancelled
PR Test (sgl-router) / finish (push) Has been cancelled
PR Test (NPU) / single-node-poc (map[name:qwen3_6_27b_w8a8_1p_in64k_out1k_50ms runner:linux-aarch64-a3-2 test_case:test/registered/ascend/performance/qwen3_6_27b/test_npu_qwen3_6_27b_w8a8_1p_in64k_out1k_50ms.py test_type:perf]) (push) Has been cancelled
PR Test (NPU) / pr-test-npu-finish (push) Has been cancelled
PR Test (Xeon) / pr-gate (push) Has been cancelled
PR Test (Xeon) / check-changes (push) Has been cancelled
PR Test (Xeon) / build-test (, xeon-gnr, base-b-test-cpu) (push) Has been cancelled
PR Test (XPU) / check-changes (push) Has been cancelled
PR Test (XPU) / pr-gate (push) Has been cancelled
PR Test (XPU) / stage-a-test-1-gpu-xpu (push) Has been cancelled
PR Test (XPU) / wait-for-stage-a (push) Has been cancelled
PR Test (XPU) / stage-b-test-1-gpu-xpu (push) Has been cancelled
PR Test (XPU) / finish (push) Has been cancelled
CI Model Inventory / build-inventory (push) Has been cancelled
Lint / lint (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark Compilation Check (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark - Manual Policy (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark - Request Processing (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark Summary (push) Has been cancelled
PR Test (SMG) / build-wheel (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on windows (x86_64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on macos (x86_64 - auto) (push) Has been cancelled
PR Test (SMG) / python-unit-tests (push) Has been cancelled
PR Test (SMG) / unit-tests (push) Has been cancelled
PR Test (SMG) / benchmarks (push) Has been cancelled
PR Test (SMG) / chat-completions (push) Has been cancelled
PR Test (SMG) / chat-completions-4gpu (push) Has been cancelled
PR Test (SMG) / e2e (push) Has been cancelled
PR Test (SMG) / docker-build-test (push) Has been cancelled
PR Test (SMG) / k8s-integration (push) Has been cancelled
PR Test (SMG) / finish (push) Has been cancelled
PR Test (SMG) / summarize-benchmarks (push) Has been cancelled
Release SGLang Model Gateway Docker Image / publish (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on macos (aarch64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - musllinux_1_1) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - musllinux_1_1) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / Build SDist (push) Has been cancelled
Release SGLang Model Gateway to PyPI / Upload to PyPI (push) Has been cancelled
Release SGLang Kernels / build-cu129-matrix (aarch64, 12.9, 3.10, arm-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / build-cu129-matrix (x86_64, 12.9, 3.10, x64-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / release-cu129 (push) Has been cancelled
Release SGLang Kernels / build-cu130-matrix (aarch64, 13.0, 3.10, arm-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / build-cu130-matrix (x86_64, 13.0, 3.10, x64-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / release-cu130 (push) Has been cancelled
Release SGLang Kernels / build-rocm-matrix (3.10, 700) (push) Has been cancelled
Release SGLang Kernels / build-rocm-matrix (3.10, 720) (push) Has been cancelled
Release SGLang Kernels / release-rocm700 (push) Has been cancelled
Release SGLang Kernels / release-rocm720 (push) Has been cancelled
Release SGLang Kernels / build-musa43 (43, 3.10) (push) Has been cancelled
Release SGLang Kernels / release-musa43 (push) Has been cancelled

This commit is contained in:
wehub-resource-sync
2026-07-13 12:38:16 +08:00
commit 94057c3d3e
7152 changed files with 2120455 additions and 0 deletions
@@ -0,0 +1,77 @@
from abc import ABC, abstractmethod
from typing import Dict, Iterator, List, Optional, Tuple, Union
import torch
class EngineBase(ABC):
"""
Abstract base class for engine interfaces that support generation, weight updating, and memory control.
This base class provides a unified API for both HTTP-based engines and engines.
"""
@abstractmethod
def generate(
self,
prompt: Optional[Union[List[str], str]] = None,
sampling_params: Optional[Union[List[Dict], Dict]] = None,
input_ids: Optional[Union[List[List[int]], List[int]]] = None,
image_data: Optional[Union[List[str], str]] = None,
return_logprob: Optional[Union[List[bool], bool]] = False,
logprob_start_len: Optional[Union[List[int], int]] = None,
top_logprobs_num: Optional[Union[List[int], int]] = None,
token_ids_logprob: Optional[Union[List[List[int]], List[int]]] = None,
lora_path: Optional[Union[List[Optional[str]], Optional[str]]] = None,
custom_logit_processor: Optional[Union[List[str], str]] = None,
return_hidden_states: Optional[bool] = None,
stream: Optional[bool] = None,
bootstrap_host: Optional[Union[List[str], str]] = None,
bootstrap_port: Optional[Union[List[int], int]] = None,
bootstrap_room: Optional[Union[List[int], int]] = None,
routed_dp_rank: Optional[int] = None,
disagg_prefill_dp_rank: Optional[int] = None,
data_parallel_rank: Optional[int] = None,
rid: Optional[Union[List[str], str]] = None,
priority: Optional[int] = None,
session_id: Optional[str] = None,
) -> Union[Dict, Iterator[Dict]]:
"""Generate outputs based on given inputs."""
pass
@abstractmethod
def flush_cache(self):
"""Flush the cache of the engine."""
pass
@abstractmethod
def update_weights_from_tensor(
self,
named_tensors: List[Tuple[str, torch.Tensor]],
load_format: Optional[str] = None,
flush_cache: bool = True,
):
"""Update model weights with in-memory tensor data."""
pass
def load_lora_adapter(self, lora_name: str, lora_path: str):
"""Load a new LoRA adapter without re-launching the engine."""
pass
def unload_lora_adapter(self, lora_name: str):
"""Unload a LoRA adapter without re-launching the engine."""
pass
@abstractmethod
def release_memory_occupation(self):
"""Release GPU memory occupation temporarily."""
pass
@abstractmethod
def resume_memory_occupation(self):
"""Resume GPU memory occupation which is previously released."""
pass
@abstractmethod
def shutdown(self):
"""Shutdown the engine and clean up resources."""
pass
@@ -0,0 +1,518 @@
"""Pydantic models for Anthropic Messages API protocol.
Mirrors the shape of the official Anthropic Python SDK
(``anthropic-sdk-python``): ``ContentBlock``, ``Tool``, ``MessageStreamEvent``
and ``ContentBlockDelta`` are discriminated unions over the ``type`` field,
so each variant carries only the fields it actually uses.
"""
import uuid
from typing import Annotated, Any, Literal, Optional, Union
from pydantic import (
BaseModel,
Discriminator,
Field,
NonNegativeInt,
Tag,
field_validator,
model_validator,
)
class AnthropicError(BaseModel):
"""Error structure for Anthropic API."""
type: str
message: str
class AnthropicErrorResponse(BaseModel):
"""Error response structure for Anthropic API."""
type: Literal["error"] = "error"
error: AnthropicError
class AnthropicUsage(BaseModel):
"""Token usage information.
``input_tokens``/``output_tokens`` are ``Optional`` because Anthropic's
streaming ``message_delta`` event omits ``input_tokens`` (the spec
requires it only on ``message_start``). Non-streaming responses set both.
"""
input_tokens: Optional[NonNegativeInt] = None
output_tokens: Optional[NonNegativeInt] = None
cache_creation_input_tokens: Optional[NonNegativeInt] = None
cache_read_input_tokens: Optional[NonNegativeInt] = None
# ---------- Content blocks (discriminated by ``type``) ----------
class TextBlock(BaseModel):
type: Literal["text"] = "text"
text: str
class ImageBlock(BaseModel):
type: Literal["image"] = "image"
# Kept loosely typed for compat with both base64 and URL sources; the
# serving layer normalises to OpenAI ``image_url`` parts.
source: Optional[Union[dict[str, Any], str]] = None
class ToolUseBlock(BaseModel):
type: Literal["tool_use"] = "tool_use"
id: str
name: str
input: dict[str, Any] = Field(default_factory=dict)
class ToolResultBlock(BaseModel):
type: Literal["tool_result"] = "tool_result"
tool_use_id: Optional[str] = None
# Some legacy payloads use ``id`` instead of ``tool_use_id``.
id: Optional[str] = None
content: Optional[Union[str, list["AnthropicContentBlock"]]] = None
is_error: Optional[bool] = None
class ToolReferenceBlock(BaseModel):
"""sglang extension: references a deferred-loaded tool by name."""
type: Literal["tool_reference"] = "tool_reference"
name: Optional[str] = None
# Anthropic-style payloads sometimes use ``tool_name``; accept both.
tool_name: Optional[str] = None
id: Optional[str] = None
class SearchResultBlock(BaseModel):
type: Literal["search_result"] = "search_result"
# ``source`` here is a URL/identifier string (unlike ImageBlock.source).
source: Optional[Union[str, dict[str, Any]]] = None
title: Optional[str] = None
content: Optional[list[dict[str, Any]]] = None
class ThinkingBlock(BaseModel):
type: Literal["thinking"] = "thinking"
thinking: str
signature: Optional[str] = None
class RedactedThinkingBlock(BaseModel):
type: Literal["redacted_thinking"] = "redacted_thinking"
data: Optional[str] = None
AnthropicContentBlock = Annotated[
Union[
TextBlock,
ImageBlock,
ToolUseBlock,
ToolResultBlock,
ToolReferenceBlock,
SearchResultBlock,
ThinkingBlock,
RedactedThinkingBlock,
],
Field(discriminator="type"),
]
class AnthropicMessage(BaseModel):
role: Literal["user", "assistant", "system"]
content: Union[str, list[AnthropicContentBlock]]
# ---------- Tools (discriminated by ``type`` family) ----------
class AnthropicCustomTool(BaseModel):
"""Custom tool defined by the API user — requires ``input_schema``."""
type: Optional[Literal["custom"]] = None # absent or explicit "custom"
name: str
description: Optional[str] = None
input_schema: dict[str, Any]
defer_loading: Optional[bool] = None
@field_validator("input_schema")
@classmethod
def _ensure_object_schema(cls, v):
if not isinstance(v, dict):
raise ValueError("input_schema must be a dictionary")
if "type" not in v:
v["type"] = "object"
return v
class AnthropicWebSearchTool(BaseModel):
"""Anthropic ``web_search_*`` server tool family.
No client-side ``input_schema`` — Anthropic provides the backing
search implementation. Tag format is ``web_search_YYYYMMDD``.
"""
type: str = Field(pattern=r"^web_search_\d{8}$")
name: Literal["web_search"] = "web_search"
description: Optional[str] = None
defer_loading: Optional[bool] = None
max_uses: Optional[int] = None
allowed_domains: Optional[list[str]] = None
blocked_domains: Optional[list[str]] = None
class AnthropicComputerTool(BaseModel):
"""Anthropic ``computer_*`` server tool family."""
type: str = Field(pattern=r"^computer_\d{8}$")
name: Literal["computer"] = "computer"
description: Optional[str] = None
defer_loading: Optional[bool] = None
display_width_px: Optional[int] = None
display_height_px: Optional[int] = None
display_number: Optional[int] = None
class AnthropicBashTool(BaseModel):
"""Anthropic ``bash_*`` server tool family."""
type: str = Field(pattern=r"^bash_\d{8}$")
name: Literal["bash"] = "bash"
description: Optional[str] = None
defer_loading: Optional[bool] = None
class AnthropicTextEditorTool(BaseModel):
"""Anthropic ``text_editor_*`` server tool family."""
type: str = Field(pattern=r"^text_editor_\d{8}$")
name: Literal["str_replace_editor", "str_replace_based_edit_tool"]
description: Optional[str] = None
defer_loading: Optional[bool] = None
def _tool_discriminator(v) -> str:
"""Pick the right tool variant from a dict or model instance.
Pydantic discriminators don't accept ``None`` as a tag, and custom
tools allow ``type`` to be absent. Map missing/``custom`` to
``"custom"`` and prefix-match server-tool families.
"""
if isinstance(v, dict):
t = v.get("type")
else:
t = getattr(v, "type", None)
if not t or t == "custom":
return "custom"
if t.startswith("web_search_"):
return "web_search"
if t.startswith("computer_"):
return "computer"
if t.startswith("bash_"):
return "bash"
if t.startswith("text_editor_"):
return "text_editor"
return "custom"
AnthropicTool = Annotated[
Union[
Annotated[AnthropicCustomTool, Tag("custom")],
Annotated[AnthropicWebSearchTool, Tag("web_search")],
Annotated[AnthropicComputerTool, Tag("computer")],
Annotated[AnthropicBashTool, Tag("bash")],
Annotated[AnthropicTextEditorTool, Tag("text_editor")],
],
Discriminator(_tool_discriminator),
]
def is_server_tool(tool) -> bool:
"""Return True for Anthropic built-in server-side tools."""
return isinstance(
tool,
(
AnthropicWebSearchTool,
AnthropicComputerTool,
AnthropicBashTool,
AnthropicTextEditorTool,
),
)
class AnthropicToolChoice(BaseModel):
"""Tool choice strategy."""
type: Literal["auto", "any", "tool", "none"]
name: Optional[str] = None
class AnthropicThinkingParam(BaseModel):
"""Anthropic extended-thinking control on the request.
Mirrors the Anthropic SDK's ``ThinkingConfigParam`` discriminated
union of three variants — see ``anthropic-sdk-python``'s
``thinking_config_{enabled,disabled,adaptive}_param.py``:
* ``enabled`` requires ``budget_tokens`` (≥1024) and accepts
``display``.
* ``disabled`` accepts no other fields.
* ``adaptive`` (Claude 4.7) accepts ``display`` but not
``budget_tokens``.
The serving layer treats ``adaptive`` identically to ``enabled``
because the local OpenAI-compatible backend has no auto-throttle
equivalent. ``budget_tokens`` is accepted on ``enabled`` for SDK
compatibility but the backend has no hard-cap knob to honor it; the
serving layer logs a WARNING so operators see that the requested
budget is not enforced. ``display="omitted"`` is accepted but
similarly cannot suppress reasoning mid-stream and is logged.
"""
type: Literal["enabled", "disabled", "adaptive"]
budget_tokens: Optional[int] = None
display: Optional[Literal["summarized", "omitted"]] = None
@model_validator(mode="after")
def _validate_thinking_shape(self):
# Cross-field rules mirror the SDK's three discriminated variants.
if self.type == "enabled":
if self.budget_tokens is None:
raise ValueError(
"thinking.budget_tokens is required when "
"thinking.type is 'enabled'"
)
if self.budget_tokens < 1024:
raise ValueError(
"thinking.budget_tokens must be >= 1024 "
"(got {})".format(self.budget_tokens)
)
elif self.type == "disabled":
if self.budget_tokens is not None:
raise ValueError(
"thinking.budget_tokens is not allowed when "
"thinking.type is 'disabled'"
)
if self.display is not None:
raise ValueError(
"thinking.display is not allowed when "
"thinking.type is 'disabled'"
)
elif self.type == "adaptive":
if self.budget_tokens is not None:
raise ValueError(
"thinking.budget_tokens is not allowed when "
"thinking.type is 'adaptive'"
)
return self
class AnthropicTaskBudget(BaseModel):
"""Claude 4.7 ``output_config.task_budget`` — soft hint, not a hard cap.
Mirrors ``BetaTokenTaskBudgetParam`` in the Anthropic SDK: ``total``
and ``type`` are required; ``remaining`` is the client-tracked
countdown used for compaction. The hard cap on generation is still
``max_tokens``; we never enforce ``task_budget`` ourselves.
"""
type: Literal["tokens"]
total: int = Field(gt=0)
remaining: Optional[int] = Field(default=None, ge=0)
class AnthropicOutputConfig(BaseModel):
"""Claude 4.7 ``output_config`` block.
``effort`` maps to the OpenAI ``reasoning_effort`` knob (``xhigh`` →
``max`` because the OpenAI Literal does not include ``xhigh``).
``task_budget`` is propagated as a custom-param hint.
"""
effort: Optional[Literal["low", "medium", "high", "xhigh", "max"]] = None
task_budget: Optional[AnthropicTaskBudget] = None
class AnthropicCountTokensRequest(BaseModel):
"""Anthropic count_tokens API request."""
model: str
messages: list[AnthropicMessage]
system: Optional[Union[str, list[AnthropicContentBlock]]] = None
thinking: Optional[AnthropicThinkingParam] = None
tool_choice: Optional[AnthropicToolChoice] = None
tools: Optional[list[AnthropicTool]] = None
# Claude 4.7 / SDK-compatibility fields. Accepted but no-op on count.
output_config: Optional[AnthropicOutputConfig] = None
betas: Optional[list[str]] = None
class AnthropicCountTokensResponse(BaseModel):
"""Anthropic count_tokens API response."""
input_tokens: int
class AnthropicMessagesRequest(BaseModel):
"""Anthropic Messages API request."""
model: str
messages: list[AnthropicMessage]
max_tokens: int
metadata: Optional[dict[str, Any]] = None
stop_sequences: Optional[list[str]] = None
stream: Optional[bool] = False
system: Optional[Union[str, list[AnthropicContentBlock]]] = None
temperature: Optional[float] = None
thinking: Optional[AnthropicThinkingParam] = None
tool_choice: Optional[AnthropicToolChoice] = None
tools: Optional[list[AnthropicTool]] = None
top_k: Optional[int] = None
top_p: Optional[float] = None
# Claude 4.7 fields. The Anthropic SDK / Claude Code attach these even
# when targeting non-Anthropic backends, so the schema must accept them.
output_config: Optional[AnthropicOutputConfig] = None
betas: Optional[list[str]] = None
@field_validator("model")
@classmethod
def _validate_model(cls, v):
if not v:
raise ValueError("Model is required")
return v
@field_validator("max_tokens")
@classmethod
def _validate_max_tokens(cls, v):
if v <= 0:
raise ValueError("max_tokens must be positive")
return v
# ---------- Stream deltas ----------
# Content-block deltas (discriminated by ``type``) vs message-end delta
# (separate model; the wire format does not put ``type`` inside its payload).
class TextDelta(BaseModel):
type: Literal["text_delta"] = "text_delta"
text: str
class InputJsonDelta(BaseModel):
type: Literal["input_json_delta"] = "input_json_delta"
partial_json: str
class ThinkingDelta(BaseModel):
type: Literal["thinking_delta"] = "thinking_delta"
thinking: str
class SignatureDelta(BaseModel):
type: Literal["signature_delta"] = "signature_delta"
signature: str
AnthropicContentDelta = Annotated[
Union[TextDelta, InputJsonDelta, ThinkingDelta, SignatureDelta],
Field(discriminator="type"),
]
class AnthropicMessageEndDelta(BaseModel):
"""Delta carried on ``message_delta`` events.
Anthropic's protocol does NOT put a ``type`` field inside this delta
payload — the SSE ``event:`` header already says ``message_delta``.
Stop reason and stop sequence are the only fields.
"""
stop_reason: Optional[
Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"]
] = None
stop_sequence: Optional[str] = None
# ---------- Stream events (discriminated by ``type``) ----------
class MessageStartEvent(BaseModel):
type: Literal["message_start"] = "message_start"
message: "AnthropicMessagesResponse"
class MessageDeltaEvent(BaseModel):
type: Literal["message_delta"] = "message_delta"
delta: AnthropicMessageEndDelta
usage: AnthropicUsage
class MessageStopEvent(BaseModel):
type: Literal["message_stop"] = "message_stop"
class ContentBlockStartEvent(BaseModel):
type: Literal["content_block_start"] = "content_block_start"
index: int
content_block: AnthropicContentBlock
class ContentBlockDeltaEvent(BaseModel):
type: Literal["content_block_delta"] = "content_block_delta"
index: int
delta: AnthropicContentDelta
class ContentBlockStopEvent(BaseModel):
type: Literal["content_block_stop"] = "content_block_stop"
index: int
class PingEvent(BaseModel):
type: Literal["ping"] = "ping"
class ErrorEvent(BaseModel):
type: Literal["error"] = "error"
error: AnthropicError
AnthropicStreamEvent = Annotated[
Union[
MessageStartEvent,
MessageDeltaEvent,
MessageStopEvent,
ContentBlockStartEvent,
ContentBlockDeltaEvent,
ContentBlockStopEvent,
PingEvent,
ErrorEvent,
],
Field(discriminator="type"),
]
class AnthropicMessagesResponse(BaseModel):
"""Anthropic Messages API response."""
id: str = Field(default_factory=lambda: f"msg_{uuid.uuid4().hex}")
type: Literal["message"] = "message"
role: Literal["assistant"] = "assistant"
content: list[AnthropicContentBlock]
model: str
stop_reason: Optional[
Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"]
] = None
stop_sequence: Optional[str] = None
usage: Optional[AnthropicUsage] = None
# Resolve forward references for nested types.
ToolResultBlock.model_rebuild()
MessageStartEvent.model_rebuild()
File diff suppressed because it is too large Load Diff
+249
View File
@@ -0,0 +1,249 @@
# SPDX-License-Identifier: Apache-2.0
# Copied from vLLM
import logging
from abc import ABC, abstractmethod
from typing import Union
import orjson
logger = logging.getLogger(__name__)
try:
from mcp import ClientSession
except ImportError as e:
mcp = e
from openai_harmony import Author, Message, Role, StreamState, TextContent
from sglang.srt.entrypoints.harmony_utils import (
get_encoding,
get_streamable_parser_for_assistant,
render_for_completion,
)
from sglang.srt.entrypoints.tool import Tool
class ConversationContext(ABC):
@abstractmethod
def append_output(self, output) -> None:
pass
@abstractmethod
async def call_tool(self) -> list[Message]:
pass
@abstractmethod
def need_builtin_tool_call(self) -> bool:
pass
@abstractmethod
def render_for_completion(self) -> list[int]:
pass
class SimpleContext(ConversationContext):
def __init__(self):
self.last_output = None
def append_output(self, output) -> None:
self.last_output = output
def need_builtin_tool_call(self) -> bool:
return False
async def call_tool(self) -> list[Message]:
raise NotImplementedError("Should not be called.")
def render_for_completion(self) -> list[int]:
raise NotImplementedError("Should not be called.")
class HarmonyContext(ConversationContext):
def __init__(
self,
messages: list,
tool_sessions: dict[str, Union["ClientSession", Tool]],
):
# TODO: Remove the hack of Union[ClientSession, Tool] by using MCP
# when demo.
self._messages = messages
self.tool_sessions = tool_sessions
self.parser = get_streamable_parser_for_assistant()
self.num_init_messages = len(messages)
# TODO
self.num_prompt_tokens = 0
self.num_cached_tokens = 0
self.num_output_tokens = 0
self.num_reasoning_tokens = 0
def append_output(self, output) -> None:
if isinstance(output, dict) and "output_ids" in output:
output_token_ids = output["output_ids"]
for token_id in output_token_ids:
self.parser.process(token_id)
output_msgs = self.parser.messages
meta_info = output["meta_info"]
if isinstance(meta_info, dict):
if "prompt_token_ids" in meta_info:
self.num_prompt_tokens = meta_info["prompt_tokens"]
if "cached_tokens" in meta_info:
self.num_cached_tokens = meta_info["cached_tokens"]
if "completion_tokens" in meta_info:
self.num_output_tokens += meta_info["completion_tokens"]
else:
output_msgs = output
self._messages.extend(output_msgs)
@property
def messages(self) -> list:
return self._messages
def need_builtin_tool_call(self) -> bool:
if not self.messages:
return False
last_msg = self.messages[-1]
recipient = last_msg.recipient
return recipient is not None and (
recipient.startswith("browser.") or recipient.startswith("python")
)
async def call_tool(self) -> list[Message]:
if not self.messages:
return []
last_msg = self.messages[-1]
recipient = last_msg.recipient
if recipient is not None:
if recipient.startswith("browser."):
return await self.call_search_tool(
self.tool_sessions["browser"], last_msg
)
elif recipient.startswith("python"):
return await self.call_python_tool(
self.tool_sessions["python"], last_msg
)
raise ValueError("No tool call found")
def render_for_completion(self) -> list[int]:
return render_for_completion(self.messages)
async def call_search_tool(
self, tool_session: Union["ClientSession", Tool], last_msg: Message
) -> list[Message]:
if isinstance(tool_session, Tool):
return await tool_session.get_result(self)
tool_name = last_msg.recipient.split(".")[1]
args = orjson.loads(last_msg.content[0].text)
result = await tool_session.call_tool(tool_name, args)
result_str = result.content[0].text
content = TextContent(text=result_str)
author = Author(role=Role.TOOL, name=last_msg.recipient)
return [Message(author=author, content=[content], recipient=Role.ASSISTANT)]
async def call_python_tool(
self, tool_session: Union["ClientSession", Tool], last_msg: Message
) -> list[Message]:
if isinstance(tool_session, Tool):
return await tool_session.get_result(self)
param = {
"code": last_msg.content[0].text,
}
result = await tool_session.call_tool("python", param)
result_str = result.content[0].text
content = TextContent(text=result_str)
author = Author(role=Role.TOOL, name="python")
return [
Message(
author=author,
content=[content],
channel=last_msg.channel,
recipient=Role.ASSISTANT,
)
]
class StreamingHarmonyContext(HarmonyContext):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.last_output = None
self.parser = get_streamable_parser_for_assistant()
self.encoding = get_encoding()
self.last_tok = None
self.num_processed_tokens = 0
@property
def messages(self) -> list:
return self.parser.messages
def append_output(self, output) -> None:
if isinstance(output, dict) and "output_ids" in output:
# RequestOutput from SGLang with outputs
output_token_ids = output["output_ids"]
# Check if we need to handle cumulative tokens
meta_info = output.get("meta_info", {})
completion_tokens = meta_info.get("completion_tokens")
if (
completion_tokens is not None
and len(output_token_ids) == completion_tokens
):
# Case 1: When --incremental-streaming-output is not set.
# The output_ids contains all tokens generated so far.
# We only need to process the new tokens.
new_token_ids = output_token_ids[self.num_processed_tokens :]
self.num_processed_tokens = len(output_token_ids)
else:
# Case 2: When --incremental-streaming-output is set.
# The output_ids contains only the new tokens.
new_token_ids = output_token_ids
self.num_processed_tokens += len(output_token_ids)
for token_id in new_token_ids:
self.parser.process(token_id)
else:
# Handle the case of tool output in direct message format
assert len(output) == 1, "Tool output should be a single message"
msg = output[0]
# Sometimes the recipient is not set for tool messages,
# so we set it to "assistant"
if msg.author.role == Role.TOOL and msg.recipient is None:
msg.recipient = "assistant"
toks = self.encoding.render(msg)
for tok in toks:
self.parser.process(tok)
self.last_tok = toks[-1]
def is_expecting_start(self) -> bool:
return self.parser.state == StreamState.EXPECT_START
def is_assistant_action_turn(self) -> bool:
return self.last_tok in self.encoding.stop_tokens_for_assistant_actions()
def render_for_completion(self) -> list[int]:
# now this list of tokens as next turn's starting tokens
# `<|start|>assistant``,
# we need to process them in parser.
rendered_tokens = super().render_for_completion()
last_n = -1
to_process = []
while rendered_tokens[last_n] != self.last_tok:
to_process.append(rendered_tokens[last_n])
last_n -= 1
for tok in reversed(to_process):
self.parser.process(tok)
return rendered_tokens
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,105 @@
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import logging
import threading
from typing import Dict, Optional, Tuple
import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.responses import PlainTextResponse
logger = logging.getLogger(__name__)
class EngineInfoBootstrapServer:
"""Lightweight HTTP server for per-rank model info registration.
Runs in a daemon thread on node_rank==0. Each ModelRunner registers its
info via HTTP PUT after model initialization. The Engine
accesses the collected info directly in-process; external consumers can
query via HTTP GET.
Currently supports transfer engine memory registration info.
"""
def __init__(self, host: str, port: int):
self.host = host
self.port = port
# Storage: {tp_rank: (session_id, weights_info_dict)}
self.transfer_engine_info: Dict[int, Tuple] = {}
self.lock = threading.Lock()
app = FastAPI()
@app.get("/health")
def health():
return PlainTextResponse("OK")
@app.put("/register_transfer_engine_info")
def register_transfer_engine_info(data: dict):
try:
tp_rank = data["tp_rank"]
info = data["transfer_engine_info"]
session_id = info["session_id"]
weights_info_dict = info["weights_info_dict"]
with self.lock:
self.transfer_engine_info[tp_rank] = (
session_id,
weights_info_dict,
)
logger.info(
f"Registered transfer engine info for tp_rank={tp_rank}, "
f"session_id={session_id}"
)
return PlainTextResponse("OK")
except Exception as e:
logger.error(f"Failed to register engine info: {e}")
raise HTTPException(status_code=400, detail=str(e))
@app.get("/get_transfer_engine_info")
def get_transfer_engine_info(rank: int):
if rank < 0:
raise HTTPException(status_code=400, detail="Invalid rank parameter")
with self.lock:
info = self.transfer_engine_info.get(rank)
if info is None:
raise HTTPException(
status_code=404,
detail=f"No transfer engine info for rank {rank}",
)
return {"rank": rank, "remote_instance_transfer_engine_info": list(info)}
config = uvicorn.Config(app, host=host, port=port, log_level="warning")
self._server = uvicorn.Server(config)
self._thread = threading.Thread(
target=self._server.run,
daemon=True,
)
self._thread.start()
logger.info(f"EngineInfoBootstrapServer started on {host}:{port}")
def close(self):
self._server.should_exit = True
self._thread.join(timeout=5)
def get_transfer_engine_info(self, rank: int) -> Optional[Tuple]:
"""Direct in-process access for co-located HTTP server (no HTTP round-trip)."""
return self.transfer_engine_info.get(rank)
@@ -0,0 +1,111 @@
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""
Engine mixin that exposes score() and async_score() on the Engine class.
These methods delegate to TokenizerManager.score_request() which is provided
by TokenizerManagerScoreMixin.
"""
from typing import List, Optional, Union
import torch
from sglang.srt.managers.tokenizer_manager_score_mixin import ScoreResult
class EngineScoreMixin:
def score(
self,
query: Optional[Union[str, List[int]]] = None,
items: Optional[Union[str, List[str], List[List[int]]]] = None,
label_token_ids: Optional[List[int]] = None,
apply_softmax: bool = False,
item_first: bool = False,
embed_override_token_id: Optional[int] = None,
query_embed_overrides: Optional[List[torch.Tensor]] = None,
item_embed_overrides: Optional[List[Optional[List[torch.Tensor]]]] = None,
return_pooled_hidden_states: bool = False,
) -> ScoreResult:
"""
Score items against a query using the loaded model.
For generation (CausalLM) models, returns the probability of each label_token_id
being generated after the query+item prompt. Example:
query = "<|user|>Is the following city the capital of France? "
items = ["Paris <|assistant|>", "London <|assistant|>"]
label_token_ids = [2332, 1223] # "Yes" / "No"
# -> [[0.9, 0.1], [0.2, 0.8]]
For SequenceClassification models, returns the pooled class logits directly from
the classification head. label_token_ids is optional and ignored.
Args:
query: The query text or pre-tokenized token IDs.
items: The item text(s) or pre-tokenized token IDs.
label_token_ids: Token IDs to score (required for CausalLM; ignored for
SequenceClassification).
apply_softmax: Whether to normalize scores using softmax.
item_first: If True, prepend items before query (single-item mode only).
embed_override_token_id: Placeholder token ID used to locate override positions.
query_embed_overrides: Embedding vectors replacing placeholder tokens in query.
item_embed_overrides: Per-item embedding vectors replacing placeholder tokens in items.
return_pooled_hidden_states: Whether to include raw pooled transformer
hidden states (before the task head) in the result. Only supported
for non-generation models (SequenceClassification, RewardModel).
Returns:
ScoreResult with scores (one list per item), prompt token count, and
optional pooled_hidden_states tensors.
"""
return self.loop.run_until_complete(
self.tokenizer_manager.score_request(
query=query,
items=items,
label_token_ids=label_token_ids,
apply_softmax=apply_softmax,
item_first=item_first,
embed_override_token_id=embed_override_token_id,
query_embed_overrides=query_embed_overrides,
item_embed_overrides=item_embed_overrides,
request=None,
return_pooled_hidden_states=return_pooled_hidden_states,
)
)
async def async_score(
self,
query: Optional[Union[str, List[int]]] = None,
items: Optional[Union[str, List[str], List[List[int]]]] = None,
label_token_ids: Optional[List[int]] = None,
apply_softmax: bool = False,
item_first: bool = False,
embed_override_token_id: Optional[int] = None,
query_embed_overrides: Optional[List[torch.Tensor]] = None,
item_embed_overrides: Optional[List[Optional[List[torch.Tensor]]]] = None,
return_pooled_hidden_states: bool = False,
) -> ScoreResult:
"""Asynchronous version of score(). See score() for full documentation."""
return await self.tokenizer_manager.score_request(
query=query,
items=items,
label_token_ids=label_token_ids,
apply_softmax=apply_softmax,
item_first=item_first,
embed_override_token_id=embed_override_token_id,
query_embed_overrides=query_embed_overrides,
item_embed_overrides=item_embed_overrides,
request=None,
return_pooled_hidden_states=return_pooled_hidden_states,
)
@@ -0,0 +1,777 @@
"""Python-side bridge between the Rust gRPC server and TokenizerManager.
The RuntimeHandle exposes synchronous methods that Rust can call via PyO3
(with a brief GIL acquisition). Response chunks are pushed into Rust-side
channels via callback objects while all async work stays on the
TokenizerManager's event loop.
"""
import asyncio
import dataclasses
import json
import logging
from types import SimpleNamespace
from typing import Any, Awaitable, Callable, Dict, List, Optional
from pydantic import ValidationError
from sglang.srt.utils.msgspec_utils import msgspec_to_builtins
logger = logging.getLogger(__name__)
class _BadOpenAIRequest(ValueError):
pass
class _CaseInsensitiveHeaders:
__slots__ = ("_data",)
def __init__(self, headers: Optional[Dict[str, str]] = None):
self._data = {k.lower(): v for k, v in (headers or {}).items()}
def get(self, name: str, default: Optional[str] = None) -> Optional[str]:
return self._data.get(name.lower(), default)
class _GrpcRequest:
"""Small FastAPI Request shim used by OpenAIServing* and TokenizerManager."""
def __init__(
self,
headers: Optional[Dict[str, str]] = None,
is_disconnected_fn: Optional[Callable[[], bool]] = None,
):
self.headers = _CaseInsensitiveHeaders(headers)
self.state = SimpleNamespace()
self._is_disconnected_fn = is_disconnected_fn
async def is_disconnected(self) -> bool:
if self._is_disconnected_fn is None:
return False
return bool(self._is_disconnected_fn())
class RuntimeHandle:
"""Thin Python handle that the Rust gRPC server calls into.
Provides synchronous ``submit_*``, ``abort``, and info methods.
Each submit method receives a ``chunk_callback`` (a Rust-side PyO3 object)
that it invokes with ``(chunk_dict, finished, error)`` for each response
chunk produced by TokenizerManager.
"""
def __init__(
self,
tokenizer_manager,
template_manager,
server_args,
scheduler_info: Optional[Dict] = None,
):
self.tokenizer_manager = tokenizer_manager
self.template_manager = template_manager
self.server_args = server_args
self.scheduler_info = scheduler_info or {}
self._openai_serving_classes = None
self.tokenizer_manager.auto_create_handle_loop()
self._event_loop = self.tokenizer_manager.event_loop
@property
def _tm_loop(self):
"""Return the TokenizerManager loop used by communicator RPCs."""
return self._event_loop
def _safe_callback(self, chunk_callback, payload, **kwargs):
"""Invoke a Rust callback and return its ChunkSendStatus, if any."""
try:
return chunk_callback(payload, **kwargs)
except Exception as e:
logger.warning("gRPC chunk_callback failed: %s", e)
return None
def _send_native_error(self, chunk_callback, message: str):
# ChunkCallback extracts the PyDict arg before reading error=.
return self._safe_callback(chunk_callback, {}, finished=True, error=message)
_BACKPRESSURE_TIMEOUT_S = 300.0
@staticmethod
def _is_pending_status(status) -> bool:
return status is not None and status == type(status).Pending
@staticmethod
def _is_closed_status(status) -> bool:
return status is not None and status == type(status).Closed
def _abort_request_id(self, rid) -> None:
if isinstance(rid, list):
for single_rid in rid:
self.tokenizer_manager.abort_request(rid=single_rid)
else:
self.tokenizer_manager.abort_request(rid=rid)
async def _send_with_backpressure(
self,
chunk_callback,
ready_event: Optional[asyncio.Event],
payload,
*,
timeout_abort_rid=None,
**kwargs,
) -> bool:
status = self._safe_callback(chunk_callback, payload, **kwargs)
if status is None or self._is_closed_status(status):
return False
if not self._is_pending_status(status):
return True
if kwargs.get("finished"):
return True
if ready_event is None:
return True
try:
await asyncio.wait_for(
ready_event.wait(), timeout=self._BACKPRESSURE_TIMEOUT_S
)
except asyncio.TimeoutError:
if timeout_abort_rid is not None:
self._abort_request_id(timeout_abort_rid)
logger.warning(
"gRPC chunk backpressure wait timed out after %ss; aborted request",
self._BACKPRESSURE_TIMEOUT_S,
)
else:
logger.warning(
"gRPC chunk backpressure wait timed out after %ss; closing stream",
self._BACKPRESSURE_TIMEOUT_S,
)
return False
ready_event.clear()
return True
def _install_on_ready(self, chunk_callback) -> Optional[asyncio.Event]:
set_on_ready = getattr(chunk_callback, "set_on_ready", None)
if set_on_ready is None:
return None
ready_event = asyncio.Event()
loop = self._tm_loop
def _on_ready() -> None:
loop.call_soon_threadsafe(ready_event.set)
try:
set_on_ready(_on_ready)
except Exception as e:
logger.warning("gRPC set_on_ready failed: %s", e)
raise
return ready_event
@staticmethod
def _uninstall_on_ready(chunk_callback) -> None:
clear = getattr(chunk_callback, "clear_on_ready", None)
if clear is None:
return
try:
clear()
except Exception as e:
logger.warning("gRPC clear_on_ready failed: %s", e)
def _submit_on_tm_loop(self, coro: Awaitable) -> None:
future = asyncio.run_coroutine_threadsafe(coro, self._tm_loop)
future.add_done_callback(self._log_unhandled_future_exception)
@staticmethod
def _log_unhandled_future_exception(future) -> None:
try:
future.result()
except Exception as e:
logger.error(
"gRPC scheduled coroutine raised unhandled exception: %s",
e,
exc_info=True,
)
def _submit_json_unary(
self,
op_name: str,
payload_coro_factory: Callable[[], Awaitable[Any]],
chunk_callback,
*,
error_payload_fn: Optional[Callable[[Exception], Any]] = None,
) -> None:
error_fn = error_payload_fn or (lambda e: {"error": {"message": str(e)}})
async def _run() -> None:
try:
payload = await payload_coro_factory()
self._safe_callback(
chunk_callback,
json.dumps(payload, default=str).encode("utf-8"),
finished=True,
)
except Exception as e:
logger.error("gRPC %s error: %s", op_name, e)
self._safe_callback(
chunk_callback,
json.dumps(error_fn(e), default=str).encode("utf-8"),
finished=True,
error=str(e),
)
self._submit_on_tm_loop(_run())
def _get_openai_serving(self):
"""Lazily initialize OpenAI serving classes."""
if self._openai_serving_classes is not None:
return self._openai_serving_classes
from sglang.srt.entrypoints.openai.serving_chat import OpenAIServingChat
from sglang.srt.entrypoints.openai.serving_classify import (
OpenAIServingClassify,
)
from sglang.srt.entrypoints.openai.serving_completions import (
OpenAIServingCompletion,
)
from sglang.srt.entrypoints.openai.serving_embedding import (
OpenAIServingEmbedding,
)
from sglang.srt.entrypoints.openai.serving_rerank import OpenAIServingRerank
from sglang.srt.entrypoints.openai.serving_score import OpenAIServingScore
self._openai_serving_classes = {
"chat": OpenAIServingChat(self.tokenizer_manager, self.template_manager),
"completion": OpenAIServingCompletion(
self.tokenizer_manager, self.template_manager
),
"embedding": OpenAIServingEmbedding(
self.tokenizer_manager, self.template_manager
),
"classify": OpenAIServingClassify(
self.tokenizer_manager, self.template_manager
),
"score": OpenAIServingScore(self.tokenizer_manager),
"rerank": OpenAIServingRerank(
self.tokenizer_manager, self.template_manager
),
}
return self._openai_serving_classes
def submit_request(
self,
*,
req_type: str,
req_dict: dict,
chunk_callback,
is_disconnected_fn: Optional[Callable[[], bool]] = None,
):
mock_request = (
_GrpcRequest(is_disconnected_fn=is_disconnected_fn)
if is_disconnected_fn is not None
else None
)
if req_type == "generate":
from sglang.srt.managers.io_struct import GenerateReqInput
obj = GenerateReqInput(**req_dict)
stream = req_dict.get("stream", False)
self._submit_on_tm_loop(
self._run_generate(obj, chunk_callback, stream, mock_request)
)
elif req_type == "embed":
from sglang.srt.managers.io_struct import EmbeddingReqInput
obj = EmbeddingReqInput(**req_dict)
self._submit_on_tm_loop(self._run_embed(obj, chunk_callback, mock_request))
else:
raise ValueError(
f"Unknown req_type: {req_type!r} (expected 'generate' or 'embed')"
)
async def _run_generate(self, obj, chunk_callback, stream: bool, request):
ready_event = None
try:
ready_event = self._install_on_ready(chunk_callback) if stream else None
gen = self.tokenizer_manager.generate_request(obj, request=request)
if stream:
async for chunk in gen:
finished = (
chunk.get("meta_info", {}).get("finish_reason") is not None
)
keep_going = await self._send_with_backpressure(
chunk_callback,
ready_event,
chunk,
finished=finished,
timeout_abort_rid=obj.rid,
)
if finished or not keep_going:
return
# Defensive: generator exited without a finish_reason chunk.
self._safe_callback(chunk_callback, {}, finished=True)
else:
result = await gen.__anext__()
self._safe_callback(chunk_callback, result, finished=True)
except StopAsyncIteration:
self._safe_callback(chunk_callback, {}, finished=True)
except Exception as e:
logger.error("gRPC generate error for rid=%s: %s", obj.rid, e)
self._send_native_error(chunk_callback, str(e))
finally:
if stream:
self._uninstall_on_ready(chunk_callback)
async def _run_embed(self, obj, chunk_callback, request):
try:
gen = self.tokenizer_manager.generate_request(obj, request=request)
result = await gen.__anext__()
self._safe_callback(chunk_callback, result, finished=True)
except StopAsyncIteration:
self._safe_callback(chunk_callback, {}, finished=True)
except Exception as e:
logger.error("gRPC embed error for rid=%s: %s", obj.rid, e)
self._send_native_error(chunk_callback, str(e))
# Bounded so a stuck TM loop can't deadlock the gRPC handler thread that
# called abort. abort_request only enqueues a message on the ZMQ socket,
# so a few seconds is generous; if we time out, log and drop — the client
# will retry or give up.
_ABORT_TIMEOUT_S = 5.0
def abort(self, rid: str = "", abort_all: bool = False):
"""Abort a request by request ID or abort all active requests."""
loop = self._tm_loop
try:
running_loop = asyncio.get_running_loop()
except RuntimeError:
running_loop = None
if running_loop is loop:
self.tokenizer_manager.abort_request(rid=rid, abort_all=abort_all)
return
future = asyncio.run_coroutine_threadsafe(
self._abort_async(rid, abort_all),
loop,
)
try:
future.result(timeout=self._ABORT_TIMEOUT_S)
except TimeoutError:
future.cancel()
logger.error(
"gRPC abort timed out after %ss (rid=%r, abort_all=%s); "
"tokenizer_manager loop appears stuck",
self._ABORT_TIMEOUT_S,
rid,
abort_all,
)
async def _abort_async(self, rid: str, abort_all: bool) -> None:
self.tokenizer_manager.abort_request(rid=rid, abort_all=abort_all)
def get_model_info(self) -> str:
model_config = self.tokenizer_manager.model_config
result = {
"model_path": self.tokenizer_manager.model_path,
"tokenizer_path": self.server_args.tokenizer_path,
"is_generation": self.tokenizer_manager.is_generation,
"weight_version": self.server_args.weight_version,
"model_type": getattr(model_config.hf_config, "model_type", None),
"architectures": getattr(model_config.hf_config, "architectures", None),
}
return json.dumps(result, default=str)
def get_server_info(self) -> str:
result: Dict[str, Any] = dataclasses.asdict(self.server_args)
result.update(self.scheduler_info)
return json.dumps(msgspec_to_builtins(result), default=str)
def health_check(self) -> bool:
from sglang.srt.managers.tokenizer_manager import ServerStatus
if self.tokenizer_manager.gracefully_exit:
return False
return self.tokenizer_manager.server_status not in (
ServerStatus.Starting,
ServerStatus.UnHealthy,
)
def tokenize(self, text: str, add_special_tokens: bool = True) -> str:
tokenizer = self.tokenizer_manager.tokenizer
tokens = tokenizer.encode(text, add_special_tokens=add_special_tokens)
result = {
"tokens": tokens,
"count": len(tokens),
"max_model_len": self.tokenizer_manager.model_config.context_len,
"input_text": text,
}
return json.dumps(result)
def detokenize(self, tokens: List[int]) -> str:
tokenizer = self.tokenizer_manager.tokenizer
text = tokenizer.decode(tokens)
return json.dumps({"text": text})
def list_models(self) -> str:
served_model_name = self.tokenizer_manager.served_model_name
models = [
{
"id": served_model_name,
"root": served_model_name,
"max_model_len": self.tokenizer_manager.model_config.context_len,
}
]
if self.server_args.enable_lora and hasattr(
self.tokenizer_manager, "lora_registry"
):
lora_registry = self.tokenizer_manager.lora_registry
for _, lora_ref in lora_registry.get_all_adapters().items():
models.append(
{
"id": lora_ref.lora_name,
"root": lora_ref.lora_path,
"parent": served_model_name,
}
)
return json.dumps(models)
def get_load(self, chunk_callback, dp_rank: Optional[int] = None) -> None:
async def _payload():
result = await self.tokenizer_manager.get_loads(dp_rank=dp_rank)
return [r.to_dict() for r in result]
self._submit_json_unary("get_load", _payload, chunk_callback)
def flush_cache(self, chunk_callback) -> None:
async def _payload():
ret = await self.tokenizer_manager.flush_cache()
return {"success": ret.success, "message": "Cache flushed."}
self._submit_json_unary(
"flush_cache",
_payload,
chunk_callback,
error_payload_fn=lambda e: {"success": False, "message": str(e)},
)
def pause_generation(self, mode: str, chunk_callback) -> None:
async def _payload():
from sglang.srt.managers.io_struct import PauseGenerationReqInput
await self.tokenizer_manager.pause_generation(
PauseGenerationReqInput(mode=mode)
)
return {"message": f"Generation paused (mode={mode})."}
self._submit_json_unary("pause_generation", _payload, chunk_callback)
def continue_generation(self, chunk_callback) -> None:
async def _payload():
from sglang.srt.managers.io_struct import ContinueGenerationReqInput
await self.tokenizer_manager.continue_generation(
ContinueGenerationReqInput()
)
return {"message": "Generation continued."}
self._submit_json_unary("continue_generation", _payload, chunk_callback)
def start_profile(self, output_dir: Optional[str], chunk_callback) -> None:
async def _payload():
from sglang.srt.managers.io_struct import ProfileReq
req = ProfileReq(output_dir=output_dir) if output_dir else ProfileReq()
await self.tokenizer_manager.start_profile(req)
return {"message": "Profiling started."}
self._submit_json_unary("start_profile", _payload, chunk_callback)
def stop_profile(self, chunk_callback) -> None:
async def _payload():
await self.tokenizer_manager.stop_profile()
return {"message": "Profiling stopped."}
self._submit_json_unary("stop_profile", _payload, chunk_callback)
def update_weights_from_disk(
self, model_path: str, load_format: Optional[str], chunk_callback
) -> None:
async def _payload():
from sglang.srt.managers.io_struct import UpdateWeightFromDiskReqInput
obj = UpdateWeightFromDiskReqInput(
model_path=model_path, load_format=load_format
)
success, message, num_paused = (
await self.tokenizer_manager.update_weights_from_disk(obj, request=None)
)
return {
"success": success,
"message": message,
"num_paused_requests": num_paused,
}
self._submit_json_unary(
"update_weights",
_payload,
chunk_callback,
error_payload_fn=lambda e: {"success": False, "message": str(e)},
)
def _submit_openai(
self,
serving_key: str,
streaming: bool,
json_body: bytes,
chunk_callback,
trace_headers: Optional[Dict[str, str]],
is_disconnected_fn: Optional[Callable[[], bool]],
) -> None:
self._submit_on_tm_loop(
self._run_openai_request(
serving_key,
json_body,
chunk_callback,
streaming=streaming,
trace_headers=trace_headers,
is_disconnected_fn=is_disconnected_fn,
)
)
def submit_openai_chat(
self,
*,
json_body: bytes,
chunk_callback,
trace_headers: Optional[Dict[str, str]] = None,
is_disconnected_fn: Optional[Callable[[], bool]] = None,
) -> None:
self._submit_openai(
"chat", True, json_body, chunk_callback, trace_headers, is_disconnected_fn
)
def submit_openai_complete(
self,
*,
json_body: bytes,
chunk_callback,
trace_headers: Optional[Dict[str, str]] = None,
is_disconnected_fn: Optional[Callable[[], bool]] = None,
) -> None:
self._submit_openai(
"completion",
True,
json_body,
chunk_callback,
trace_headers,
is_disconnected_fn,
)
def submit_openai_embed(
self,
*,
json_body: bytes,
chunk_callback,
trace_headers: Optional[Dict[str, str]] = None,
is_disconnected_fn: Optional[Callable[[], bool]] = None,
) -> None:
self._submit_openai(
"embedding",
False,
json_body,
chunk_callback,
trace_headers,
is_disconnected_fn,
)
def submit_openai_classify(
self,
*,
json_body: bytes,
chunk_callback,
trace_headers: Optional[Dict[str, str]] = None,
is_disconnected_fn: Optional[Callable[[], bool]] = None,
) -> None:
self._submit_openai(
"classify",
False,
json_body,
chunk_callback,
trace_headers,
is_disconnected_fn,
)
def submit_openai_score(
self,
*,
json_body: bytes,
chunk_callback,
trace_headers: Optional[Dict[str, str]] = None,
is_disconnected_fn: Optional[Callable[[], bool]] = None,
) -> None:
self._submit_openai(
"score", False, json_body, chunk_callback, trace_headers, is_disconnected_fn
)
def submit_openai_rerank(
self,
*,
json_body: bytes,
chunk_callback,
trace_headers: Optional[Dict[str, str]] = None,
is_disconnected_fn: Optional[Callable[[], bool]] = None,
) -> None:
self._submit_openai(
"rerank",
False,
json_body,
chunk_callback,
trace_headers,
is_disconnected_fn,
)
def _get_openai_request_class(self, serving_key: str):
"""Return the Pydantic request class for a given serving key."""
from sglang.srt.entrypoints.openai.protocol import (
ChatCompletionRequest,
ClassifyRequest,
CompletionRequest,
EmbeddingRequest,
ScoringRequest,
V1RerankReqInput,
)
return {
"chat": ChatCompletionRequest,
"completion": CompletionRequest,
"embedding": EmbeddingRequest,
"classify": ClassifyRequest,
"score": ScoringRequest,
"rerank": V1RerankReqInput,
}[serving_key]
async def _run_openai_request(
self,
serving_key: str,
json_body: bytes,
chunk_callback,
streaming: bool,
trace_headers: Optional[Dict[str, str]] = None,
is_disconnected_fn: Optional[Callable[[], bool]] = None,
):
try:
serving = self._get_openai_serving()[serving_key]
try:
request_dict = json.loads(json_body)
if not isinstance(request_dict, dict):
raise _BadOpenAIRequest(
f"Request body must be a JSON object, got {type(request_dict).__name__}"
)
request_cls = self._get_openai_request_class(serving_key)
request_obj = request_cls(**request_dict)
except (json.JSONDecodeError, ValidationError, _BadOpenAIRequest) as e:
error_body = json.dumps(
{"error": {"message": str(e), "type": "BadRequest"}}
).encode("utf-8")
if streaming:
self._safe_callback(
chunk_callback, error_body, finished=True, error=str(e)
)
else:
self._safe_callback(
chunk_callback, error_body, finished=True, status_code=400
)
return
mock_request = _GrpcRequest(
headers=trace_headers,
is_disconnected_fn=is_disconnected_fn,
)
result = await serving.handle_request(request_obj, mock_request)
if hasattr(result, "body_iterator"):
ready_event = self._install_on_ready(chunk_callback)
data_buf: List[str] = []
stream_closed = False
async def _flush_event() -> bool:
"""Flush buffered SSE data lines as one chunk. Returns False if Rust closed."""
if not data_buf:
return True
body = "\n".join(data_buf)
data_buf.clear()
if body == "[DONE]" or not body:
return True
return await self._send_with_backpressure(
chunk_callback,
ready_event,
body.encode("utf-8"),
finished=False,
)
try:
async for raw_chunk in result.body_iterator:
if isinstance(raw_chunk, bytes):
raw_chunk = raw_chunk.decode("utf-8", errors="replace")
for line in raw_chunk.split("\n"):
line = line.rstrip("\r")
if not line:
if not await _flush_event():
stream_closed = True
break
elif line.startswith(":"):
continue # SSE comment / heartbeat
elif line.startswith("data:"):
value = line[5:]
if value.startswith(" "):
value = value[1:]
data_buf.append(value)
# event:, id:, retry:, unknown fields: ignored
if stream_closed:
break
if not stream_closed:
await _flush_event()
self._safe_callback(chunk_callback, b"", finished=True)
finally:
self._uninstall_on_ready(chunk_callback)
else:
if hasattr(result, "model_dump"):
resp_bytes = json.dumps(result.model_dump()).encode("utf-8")
elif hasattr(result, "body"):
resp_bytes = result.body
elif isinstance(result, (dict, list)):
resp_bytes = json.dumps(result).encode("utf-8")
else:
resp_bytes = str(result).encode("utf-8")
status_code = int(
getattr(result, "status_code", None)
or getattr(result, "code", None)
or 200
)
self._safe_callback(
chunk_callback,
resp_bytes,
finished=True,
status_code=status_code,
)
except Exception as e:
logger.error("gRPC OpenAI %s error: %s", serving_key, e)
error_body = json.dumps({"error": {"message": str(e)}}).encode("utf-8")
if streaming:
self._safe_callback(
chunk_callback, error_body, finished=True, error=str(e)
)
else:
self._safe_callback(
chunk_callback,
error_body,
finished=True,
status_code=int(getattr(e, "status_code", 500)),
)
@@ -0,0 +1,263 @@
"""
Thin gRPC server wrapper — delegates to smg-grpc-servicer package.
A lightweight HTTP sidecar is started alongside the gRPC server to expose:
- /metrics (Prometheus, when --enable-metrics is set)
- /start_profile, /stop_profile (profiling control)
The sidecar is started on --smg-http-sidecar-port (default: --port + 1)
once the gRPC request manager is ready, regardless of whether --enable-metrics
is set.
"""
import inspect
import json
import logging
import time
from aiohttp import web
from sglang.srt.managers.io_struct import ProfileReq, ProfileReqType
from sglang.srt.utils.common import get_bool_env_var
logger = logging.getLogger(__name__)
async def _start_sidecar_server(host: str, port: int, app):
"""Start the aiohttp sidecar and return the runner for cleanup."""
runner = web.AppRunner(app)
await runner.setup()
try:
site = web.TCPSite(runner, host, port)
await site.start()
except BaseException:
await runner.cleanup()
raise
logger.info("HTTP sidecar server started on http://%s:%d", host, port)
return runner
def _add_metrics_routes(app):
"""Add Prometheus /metrics endpoint to the aiohttp app."""
from prometheus_client import (
CollectorRegistry,
multiprocess,
)
from prometheus_client.openmetrics.exposition import (
CONTENT_TYPE_LATEST,
generate_latest,
)
async def metrics_handler(request):
try:
registry = CollectorRegistry()
multiprocess.MultiProcessCollector(registry)
data = generate_latest(registry)
return web.Response(
body=data,
headers={"Content-Type": CONTENT_TYPE_LATEST},
)
except Exception:
logger.exception("Failed to generate Prometheus metrics")
return web.Response(status=500, text="Failed to generate metrics")
app.router.add_get("/metrics", metrics_handler)
def _check_communicator_results(results, action):
"""Return a web.Response error if results indicate failure, else None."""
if not results:
return web.Response(status=500, text="No response from scheduler\n")
failures = [r for r in results if not r.success]
if failures:
msgs = " | ".join(r.message for r in failures)
return web.Response(status=500, text=f"{action} failed: {msgs}\n")
return None
def _add_admin_routes(app, request_manager):
"""Add admin endpoints to the aiohttp app.
Endpoints: /start_profile, /stop_profile.
Business logic (request construction, env var handling, response interpretation)
lives here; request_manager only provides the transport to the scheduler.
"""
async def start_profile_handler(request):
try:
if request.content_length and request.content_length > 0:
try:
body = await request.json()
except json.JSONDecodeError as e:
return web.Response(
status=400,
text=f"Invalid JSON in request body: {e}",
)
else:
body = {}
# Build ProfileReq with env var overrides (same as tokenizer_communicator_mixin)
with_stack = body.get("with_stack")
env_with_stack = get_bool_env_var("SGLANG_PROFILE_WITH_STACK", "true")
with_stack = (with_stack is not False) and env_with_stack
record_shapes = body.get("record_shapes")
env_record_shapes = get_bool_env_var("SGLANG_PROFILE_RECORD_SHAPES", "true")
record_shapes = (record_shapes is not False) and env_record_shapes
req = ProfileReq(
req_type=ProfileReqType.START_PROFILE,
output_dir=body.get("output_dir"),
start_step=body.get("start_step"),
num_steps=body.get("num_steps"),
activities=body.get("activities"),
with_stack=with_stack,
record_shapes=record_shapes,
profile_by_stage=body.get("profile_by_stage", False),
profile_id=str(time.time()),
merge_profiles=body.get("merge_profiles", False),
profile_prefix=body.get("profile_prefix"),
profile_stages=body.get("profile_stages"),
)
results = await request_manager.send_communicator_req(
req, "profile_communicator", timeout=600.0
)
err = _check_communicator_results(results, "Start Profile")
if err:
return err
return web.Response(text="Start profiling.\n")
except Exception as e:
logger.exception("Failed to start profile")
return web.Response(
status=500,
text=f"Internal error: {type(e).__name__}. Check server logs.\n",
)
async def stop_profile_handler(request):
try:
req = ProfileReq(req_type=ProfileReqType.STOP_PROFILE)
results = await request_manager.send_communicator_req(
req, "profile_communicator", timeout=600.0
)
err = _check_communicator_results(results, "Stop profile")
if err:
return err
return web.Response(text="Stop profiling. This will take some time.\n")
except Exception as e:
logger.exception("Failed to stop profile")
return web.Response(
status=500,
text=f"Internal error: {type(e).__name__}. Check server logs.\n",
)
app.router.add_post("/start_profile", start_profile_handler)
app.router.add_post("/stop_profile", stop_profile_handler)
async def serve_grpc(server_args, model_info=None):
"""Start the standalone gRPC server with integrated scheduler."""
try:
from smg_grpc_servicer.sglang.server import serve_grpc as _serve_grpc
except ImportError as e:
raise ImportError(
"gRPC mode requires the smg-grpc-servicer package. "
"If not installed, run: pip install smg-grpc-servicer[sglang]. "
"If already installed, there may be a broken import due to a "
"version mismatch — see the chained exception above for details."
) from e
sidecar_app = web.Application()
sidecar_runner = None
sidecar_port = (
server_args.smg_http_sidecar_port
if server_args.smg_http_sidecar_port is not None
else server_args.port + 1
)
# Metrics setup: must set PROMETHEUS_MULTIPROC_DIR before scheduler
# processes import prometheus_client, since the env var is inherited
# at fork time.
if server_args.enable_metrics:
try:
from sglang.srt.observability.func_timer import enable_func_timer
from sglang.srt.utils import set_prometheus_multiproc_dir
set_prometheus_multiproc_dir()
enable_func_timer()
_add_metrics_routes(sidecar_app)
except Exception as e:
logger.error(
"Failed to set up metrics: %s. Continuing without metrics.",
e,
exc_info=True,
)
async def _on_request_manager_ready(request_manager, srv_args, sched_info):
nonlocal sidecar_runner
try:
_add_admin_routes(sidecar_app, request_manager)
except Exception as e:
logger.error(
"Failed to set up admin routes: %s. "
"Continuing without admin endpoints.",
e,
exc_info=True,
)
try:
sidecar_runner = await _start_sidecar_server(
server_args.host, sidecar_port, sidecar_app
)
except OSError as e:
logger.error(
"Failed to start HTTP sidecar server: %s. "
"Continuing without metrics/profile endpoints.",
e,
exc_info=True,
)
except Exception as e:
logger.error(
"Unexpected error starting HTTP sidecar server: %s. "
"Continuing without metrics/profile endpoints.",
e,
exc_info=True,
)
# Older smg-grpc-servicer releases (≤ 0.5.2) accept only (server_args,
# model_info) and reject the on_request_manager_ready hook. The hook is
# what calls _start_sidecar_server, so dropping the kwarg disables the
# entire HTTP sidecar (Prometheus /metrics and /start_profile +
# /stop_profile). Core gRPC serving still works without it.
serve_kwargs: dict = {}
sidecar_supported = (
"on_request_manager_ready" in inspect.signature(_serve_grpc).parameters
)
if sidecar_supported:
serve_kwargs["on_request_manager_ready"] = _on_request_manager_ready
elif server_args.enable_metrics:
# User explicitly asked for metrics but the installed servicer can't
# start the sidecar that serves them — fail loud rather than silently
# produce a server with no /metrics endpoint.
raise RuntimeError(
"--enable-metrics requires smg-grpc-servicer ≥ 0.5.3 (the version "
"that accepts 'on_request_manager_ready'); installed version "
"lacks the hook so the HTTP sidecar would never start. Upgrade "
"smg-grpc-servicer or remove --enable-metrics."
)
else:
logger.warning(
"Installed smg-grpc-servicer does not accept "
"'on_request_manager_ready'; HTTP sidecar disabled "
"(no /metrics, /start_profile, /stop_profile). "
"Upgrade smg-grpc-servicer to ≥ 0.5.3 to enable it."
)
try:
await _serve_grpc(server_args, model_info, **serve_kwargs)
finally:
if sidecar_runner is not None:
try:
await sidecar_runner.cleanup()
except Exception as e:
logger.exception(
"Failed to cleanly shut down HTTP sidecar server: %s",
e,
)
@@ -0,0 +1,389 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from vLLM: https://github.com/vllm-project/vllm/blob/1b9902806915040ac9b3029f2ab7522ec505afc3/vllm/entrypoints/harmony_utils.py
# Slight differences in processing chat messages
import datetime
import logging
from collections.abc import Iterable
from typing import Literal, Optional, Union
import orjson
from openai.types.responses import (
ResponseOutputItem,
ResponseOutputMessage,
ResponseOutputText,
ResponseReasoningItem,
)
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
from openai.types.responses.response_function_web_search import (
ActionFind,
ActionOpenPage,
ActionSearch,
ResponseFunctionWebSearch,
)
from openai.types.responses.response_reasoning_item import (
Content as ResponseReasoningTextContent,
)
from openai.types.responses.tool import Tool
from openai_harmony import (
Author,
Conversation,
DeveloperContent,
HarmonyEncodingName,
Message,
ReasoningEffort,
Role,
StreamableParser,
SystemContent,
TextContent,
ToolDescription,
load_harmony_encoding,
)
from sglang.srt.entrypoints.openai.protocol import ResponseInputOutputItem
from sglang.srt.utils import random_uuid
logger = logging.getLogger(__name__)
REASONING_EFFORT = {
"high": ReasoningEffort.HIGH,
"medium": ReasoningEffort.MEDIUM,
"low": ReasoningEffort.LOW,
}
_harmony_encoding = None
def get_encoding():
global _harmony_encoding
if _harmony_encoding is None:
_harmony_encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
return _harmony_encoding
def get_system_message(
model_identity: Optional[str] = None,
reasoning_effort: Optional[Literal["high", "medium", "low"]] = None,
start_date: Optional[str] = None,
browser_description: Optional[str] = None,
python_description: Optional[str] = None,
) -> Message:
sys_msg_content = SystemContent.new()
if model_identity is not None:
sys_msg_content = sys_msg_content.with_model_identity(model_identity)
if reasoning_effort is not None:
sys_msg_content = sys_msg_content.with_reasoning_effort(
REASONING_EFFORT[reasoning_effort]
)
if start_date is None:
start_date = datetime.datetime.now().strftime("%Y-%m-%d")
sys_msg_content = sys_msg_content.with_conversation_start_date(start_date)
if browser_description is not None:
sys_msg_content = sys_msg_content.with_tools(browser_description)
if python_description is not None:
sys_msg_content = sys_msg_content.with_tools(python_description)
sys_msg = Message.from_role_and_content(Role.SYSTEM, sys_msg_content)
return sys_msg
def get_developer_message(
instructions: Optional[str] = None, tools: Optional[list[Tool]] = None
) -> Message:
dev_msg_content = DeveloperContent.new()
if instructions is not None:
dev_msg_content = dev_msg_content.with_instructions(instructions)
if tools is not None:
function_tools = []
for tool in tools:
if tool.type in (
"web_search",
"web_search_preview",
"code_interpreter",
):
# These are built-in tools that are added to the system message.
pass
elif tool.type == "function":
function_tools.append(tool)
else:
# No harmony prompt template for the remaining built-ins;
# drop them so the request still runs.
logger.debug(
"harmony: ignoring unsupported response tool type %r",
tool.type,
)
if function_tools:
function_tool_descriptions = [
ToolDescription.new(
name=tool.name,
description=tool.description,
parameters=tool.parameters,
)
for tool in function_tools
]
dev_msg_content = dev_msg_content.with_function_tools(
function_tool_descriptions
)
dev_msg = Message.from_role_and_content(Role.DEVELOPER, dev_msg_content)
return dev_msg
def get_user_message(content: str) -> Message:
return Message.from_role_and_content(Role.USER, content)
def parse_response_input(
response_msg: ResponseInputOutputItem,
prev_responses: list[Union[ResponseOutputItem, ResponseReasoningItem]],
) -> Message:
if not isinstance(response_msg, dict):
response_msg = response_msg.model_dump()
if "type" not in response_msg or response_msg["type"] == "message":
role = response_msg["role"]
content = response_msg["content"]
if role == "system":
# User is trying to set a system message. Change it to:
# <|start|>developer<|message|># Instructions
# {instructions}<|end|>
role = "developer"
text_prefix = "Instructions:\n"
else:
text_prefix = ""
if isinstance(content, str):
msg = Message.from_role_and_content(role, text_prefix + content)
else:
# Filter to text parts first, then enumerate, so the surviving first
# text chunk always carries the system→developer text_prefix even if
# earlier parts were non-text (image/audio) and got dropped.
text_chunks = [
c for c in content if c.get("type") in ("text", "input_text")
]
contents = [
TextContent(text=(text_prefix if i == 0 else "") + c.get("text", ""))
for i, c in enumerate(text_chunks)
]
msg = Message.from_role_and_contents(role, contents)
elif response_msg["type"] == "function_call_output":
call_id = response_msg["call_id"]
call_response: Optional[ResponseFunctionToolCall] = None
for prev_response in reversed(prev_responses):
if (
isinstance(prev_response, ResponseFunctionToolCall)
and prev_response.call_id == call_id
):
call_response = prev_response
break
if call_response is None:
raise ValueError(f"No call message found for {call_id}")
msg = Message.from_author_and_content(
Author.new(Role.TOOL, f"functions.{call_response.name}"),
response_msg["output"],
)
elif response_msg["type"] == "reasoning":
content = response_msg["content"]
assert len(content) == 1
msg = Message.from_role_and_content(Role.ASSISTANT, content[0]["text"])
elif response_msg["type"] == "function_call":
msg = Message.from_role_and_content(Role.ASSISTANT, response_msg["arguments"])
msg = msg.with_channel("commentary")
msg = msg.with_recipient(f"functions.{response_msg['name']}")
msg = msg.with_content_type("json")
else:
raise ValueError(f"Unknown input type: {response_msg['type']}")
return msg
def parse_response_output(output: ResponseOutputItem) -> Message:
if isinstance(output, ResponseOutputMessage):
role = output.role
contents = [TextContent(text=c.text) for c in output.content]
msg = Message.from_role_and_contents(role, contents)
return msg
elif isinstance(output, ResponseFunctionToolCall):
msg = Message.from_role_and_content(Role.ASSISTANT, output.arguments)
msg = msg.with_channel("commentary")
msg = msg.with_recipient(output.name)
msg = msg.with_content_type("json")
return msg
else:
raise ValueError(f"Unknown output type: {type(output)}")
def parse_chat_input(chat_msg) -> Message:
role = chat_msg.role
content = chat_msg.content
if isinstance(content, str):
contents = [TextContent(text=content)]
else:
# TODO: Support refusal.
contents = [TextContent(text=c.text) for c in content]
msg = Message.from_role_and_contents(role, contents)
return msg
def render_for_completion(messages: list[Message]) -> list[int]:
conversation = Conversation.from_messages(messages)
token_ids = get_encoding().render_conversation_for_completion(
conversation, Role.ASSISTANT
)
return token_ids
def get_stop_tokens_for_assistant_actions() -> list[int]:
return get_encoding().stop_tokens_for_assistant_actions()
def get_streamable_parser_for_assistant() -> StreamableParser:
return StreamableParser(get_encoding(), role=Role.ASSISTANT)
def parse_output_message(message: Message):
if message.author.role != "assistant":
# This is a message from a tool to the assistant (e.g., search result).
# Don't include it in the final output for now. This aligns with
# OpenAI's behavior on models like o4-mini.
return []
output_items = []
recipient = message.recipient
if recipient is not None and recipient.startswith("browser."):
if len(message.content) != 1:
raise ValueError("Invalid number of contents in browser message")
content = message.content[0]
browser_call = orjson.loads(content.text)
# TODO: translate to url properly!
if recipient == "browser.search":
action = ActionSearch(
query=f"cursor:{browser_call.get('query', '')}", type="search"
)
elif recipient == "browser.open":
action = ActionOpenPage(
url=f"cursor:{browser_call.get('url', '')}", type="open_page"
)
elif recipient == "browser.find":
action = ActionFind(
pattern=browser_call["pattern"],
url=f"cursor:{browser_call.get('url', '')}",
type="find",
)
else:
raise ValueError(f"Unknown browser action: {recipient}")
web_search_item = ResponseFunctionWebSearch(
id=f"ws_{random_uuid()}",
action=action,
status="completed",
type="web_search_call",
)
output_items.append(web_search_item)
elif message.channel == "analysis":
for content in message.content:
reasoning_item = ResponseReasoningItem(
id=f"rs_{random_uuid()}",
type="reasoning",
summary=[],
content=[
ResponseReasoningTextContent(
text=content.text, type="reasoning_text"
)
],
status=None,
)
output_items.append(reasoning_item)
elif message.channel == "commentary":
if message.recipient.startswith("functions."):
function_name = message.recipient.split(".")[-1]
for content in message.content:
random_id = random_uuid()
response_item = ResponseFunctionToolCall(
arguments=content.text,
call_id=f"call_{random_id}",
type="function_call",
name=function_name,
id=f"ft_{random_id}",
)
output_items.append(response_item)
elif message.recipient.startswith("python") or message.recipient.startswith(
"browser"
):
for content in message.content:
reasoning_item = ResponseReasoningItem(
id=f"rs_{random_uuid()}",
type="reasoning",
summary=[],
content=[
ResponseReasoningTextContent(
text=content.text, type="reasoning_text"
)
],
status=None,
)
output_items.append(reasoning_item)
else:
raise ValueError(f"Unknown recipient: {message.recipient}")
elif message.channel == "final":
contents = []
for content in message.content:
output_text = ResponseOutputText(
text=content.text,
annotations=[], # TODO
type="output_text",
logprobs=None, # TODO
)
contents.append(output_text)
text_item = ResponseOutputMessage(
id=f"msg_{random_uuid()}",
content=contents,
role=message.author.role,
status="completed",
type="message",
)
output_items.append(text_item)
else:
raise ValueError(f"Unknown channel: {message.channel}")
return output_items
def parse_remaining_state(parser: StreamableParser):
if not parser.current_content:
return []
if parser.current_role != Role.ASSISTANT:
return []
current_recipient = parser.current_recipient
if current_recipient is not None and current_recipient.startswith("browser."):
return []
if parser.current_channel == "analysis":
reasoning_item = ResponseReasoningItem(
id=f"rs_{random_uuid()}",
type="reasoning",
summary=[],
content=[
ResponseReasoningTextContent(
text=parser.current_content, type="reasoning_text"
)
],
status=None,
)
return [reasoning_item]
elif parser.current_channel == "final":
output_text = ResponseOutputText(
text=parser.current_content,
annotations=[], # TODO
type="output_text",
logprobs=None, # TODO
)
text_item = ResponseOutputMessage(
id=f"msg_{random_uuid()}",
content=[output_text],
role="assistant",
status="completed",
type="message",
)
return [text_item]
return []
def parse_output_into_messages(token_ids: Iterable[int]):
parser = get_streamable_parser_for_assistant()
for token_id in token_ids:
parser.process(token_id)
return parser
@@ -0,0 +1,95 @@
"""Pure-ASGI middleware that decompresses compressed request bodies.
Gated on `SGLANG_ENABLE_REQUEST_DECOMPRESSION` and request header
`x-body-compressed`, whose value names the method. For example, a caller that
compressed the body with zstd sets the `x-body-compressed: zstd` header.
"""
import asyncio
import io
import logging
import zstandard
from fastapi.responses import Response
from starlette.datastructures import Headers
logger = logging.getLogger(__name__)
def _zstd_decompress(raw: bytes) -> bytes:
return zstandard.ZstdDecompressor().stream_reader(io.BytesIO(raw)).read()
_DECOMPRESSORS = {"zstd": _zstd_decompress}
def _rewrite_headers(headers, new_len):
"""Update headers to reflect body status after decompression."""
out = [
(k, v)
for (k, v) in headers
if k not in (b"content-length", b"x-body-compressed")
]
out.append((b"content-length", str(new_len).encode()))
return out
class RequestDecompressionMiddleware:
"""Decompress request body per request header `x-body-compressed`."""
def __init__(self, app):
self.app = app
async def __call__(self, scope, receive, send):
# No-op passthrough for any request without the compression header.
if scope["type"] != "http":
return await self.app(scope, receive, send)
method = Headers(scope=scope).get("x-body-compressed")
if method is None:
return await self.app(scope, receive, send)
# Fail loud on an unsupported compression method.
decompress = _DECOMPRESSORS.get(method)
if decompress is None:
return await Response(
f"unsupported x-body-compressed {method!r}; "
f"supported: {sorted(_DECOMPRESSORS)}",
status_code=400,
)(scope, receive, send)
# Collect request body.
body = b""
more_body = True
while more_body:
message = await receive()
# Incomplete body (e.g. client disconnect); hand off to later stages.
if message["type"] != "http.request":
return await self.app(scope, receive, send)
body += message.get("body", b"")
more_body = message.get("more_body", False)
# Decompress off the event loop by releasing the GIL around the C decompress.
try:
loop = asyncio.get_running_loop()
body = await loop.run_in_executor(None, decompress, body)
except Exception as e:
logger.warning("request body decompress failed: %s", e)
return await Response("decompress failed", status_code=400)(
scope, receive, send
)
# Update the headers after decompression
scope = dict(scope)
scope["headers"] = _rewrite_headers(scope["headers"], len(body))
# Fake receiver to let later stages see the decompressed body.
body_sent = False
async def wrapped_receive():
nonlocal body_sent
if not body_sent:
body_sent = True
return {"type": "http.request", "body": body, "more_body": False}
return await receive()
await self.app(scope, wrapped_receive, send)
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,146 @@
import multiprocessing
import time
from typing import List, Optional, Tuple
import requests
import torch
from sglang.srt.entrypoints.EngineBase import EngineBase
from sglang.srt.entrypoints.http_server import launch_server
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import MultiprocessingSerializer, kill_process_tree
def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process:
p = multiprocessing.Process(target=launch_server, args=(server_args,))
p.start()
base_url = server_args.url()
timeout = 300.0 # Increased timeout to 5 minutes for downloading large models
start_time = time.perf_counter()
ssl_verify = server_args.ssl_verify()
with requests.Session() as session:
while time.perf_counter() - start_time < timeout:
try:
headers = {
"Content-Type": "application/json; charset=utf-8",
"Authorization": f"Bearer {server_args.api_key}",
}
response = session.get(
f"{base_url}/health_generate", headers=headers, verify=ssl_verify
)
if response.status_code == 200:
return p
except requests.RequestException:
pass
if not p.is_alive():
raise Exception("Server process terminated unexpectedly.")
time.sleep(2)
p.terminate()
raise TimeoutError("Server failed to start within the timeout period.")
class HttpServerEngineAdapter(EngineBase):
"""
You can use this class to launch a server from a VerlEngine instance.
We recommend using this class only you need to use http server.
Otherwise, you can use Engine directly.
"""
def __init__(self, **kwargs):
self.server_args = ServerArgs(**kwargs)
print(
f"Launch HttpServerEngineAdapter at: {self.server_args.host}:{self.server_args.port}"
)
self.process = launch_server_process(self.server_args)
def _make_request(self, endpoint: str, payload: Optional[dict] = None):
"""Make a POST request to the specified endpoint with the given payload.
Args:
endpoint: The API endpoint to call
payload: The JSON payload to send (default: empty dict)
Returns:
The JSON response from the server
"""
url = f"{self.server_args.url()}/{endpoint}"
response = requests.post(
url, json=payload or {}, verify=self.server_args.ssl_verify()
)
response.raise_for_status()
return response.json()
def update_weights_from_tensor(
self,
named_tensors: List[Tuple[str, torch.Tensor]],
load_format: Optional[str] = None,
flush_cache: bool = False,
):
"""
Update model weights from tensor data. The HTTP server will only post meta data, and the real weights will be copied directly from GPUs.
Note: The model should be on GPUs rather than CPU for this functionality to work properly.
If you encounter issues, ensure your model is loaded on GPU devices rather than CPU.
"""
return self._make_request(
"update_weights_from_tensor",
{
"serialized_named_tensors": [
MultiprocessingSerializer.serialize(named_tensors, output_str=True)
for _ in range(self.server_args.tp_size)
],
"load_format": load_format,
"flush_cache": flush_cache,
},
)
def shutdown(self):
kill_process_tree(self.process.pid, wait_timeout=60)
def generate(
self,
prompt=None,
sampling_params=None,
input_ids=None,
image_data=None,
return_logprob=False,
logprob_start_len=None,
top_logprobs_num=None,
token_ids_logprob=None,
lora_path=None,
custom_logit_processor=None,
priority=None,
session_id=None,
):
payload = {
"text": prompt,
"sampling_params": sampling_params,
"input_ids": input_ids,
"image_data": image_data,
"return_logprob": return_logprob,
"logprob_start_len": logprob_start_len,
"top_logprobs_num": top_logprobs_num,
"token_ids_logprob": token_ids_logprob,
"lora_path": lora_path,
"custom_logit_processor": custom_logit_processor,
"priority": priority,
"session_id": session_id,
}
# Filter out None values
payload = {k: v for k, v in payload.items() if v is not None}
return self._make_request("generate", payload)
def release_memory_occupation(self):
return self._make_request("release_memory_occupation")
def resume_memory_occupation(self):
return self._make_request("resume_memory_occupation")
def flush_cache(self):
return self._make_request("flush_cache")
@@ -0,0 +1,112 @@
# SGLang Ollama Integration
Ollama API compatibility for SGLang, plus a Smart Router for intelligent routing between local and remote models.
## Features
1. **Ollama-compatible API** - Use Ollama CLI/library with SGLang backend
2. **Smart Router** - LLM-based routing between local and remote models
## Ollama API
For basic Ollama API usage with SGLang (CLI and Python examples), see the [Ollama API documentation](https://sgl-project.github.io/basic_usage/ollama_api.html).
## Smart Router
### Prerequisites
```bash
pip install ollama
```
Intelligently routes requests between local Ollama and remote SGLang using an LLM judge.
### How It Works
```
User Request
┌─────────────────────┐
│ LLM Judge │ Classifies as SIMPLE or COMPLEX
│ (local model) │
└─────────────────────┘
┌─────────────────────┐
│ SIMPLE → Local │ Fast response from local Ollama
│ COMPLEX → Remote │ Powerful response from SGLang
└─────────────────────┘
```
The LLM judge (running on local Ollama) analyzes each request and decides:
- **SIMPLE**: Quick responses, greetings, factual questions, definitions, basic Q&A
- **COMPLEX**: Deep reasoning, multi-step analysis, long explanations, creative writing
### Setup
**Terminal 1: Local Ollama**
```bash
ollama pull <LOCAL_MODEL> # e.g., llama3.2, mistral, phi3
ollama serve # This will block the terminal
```
**Terminal 2: Remote SGLang (GPU)**
```bash
ssh user@gpu-server
python -m sglang.launch_server --model <REMOTE_MODEL> --port 30001 --host 0.0.0.0
```
**Terminal 3: Smart Router**
```bash
ssh -L 30001:localhost:30001 user@gpu-server -N &
python python/sglang/srt/entrypoints/ollama/smart_router.py
```
### Configuration
```python
from sglang.srt.entrypoints.ollama.smart_router import SmartRouter
router = SmartRouter(
# Local Ollama
local_host="http://localhost:11434",
local_model="llama3.2", # or any Ollama model
# Remote SGLang
remote_host="http://localhost:30001",
remote_model="Qwen/Qwen2.5-1.5B-Instruct", # or any HuggingFace model
# LLM Judge (optional, defaults to local_model)
judge_model="llama3.2",
)
```
### Usage
```python
# Auto-routing via LLM judge
response = router.chat("Hello!", verbose=True)
# [Router] LLM Judge: SIMPLE
# [Router] -> Local Ollama | Model: llama3.2
response = router.chat("Explain quantum computing in detail", verbose=True)
# [Router] LLM Judge: COMPLEX
# [Router] -> Remote SGLang | Model: Qwen/Qwen2.5-1.5B-Instruct
# Force routing (skip LLM judge)
response = router.chat("question", force_local=True)
response = router.chat("question", force_remote=True)
# Streaming
for chunk in router.chat_stream("Tell me a story"):
print(chunk['message']['content'], end='')
```
---
## Value
- **Ollama**: Simple CLI/API developers already know
- **SGLang**: High-performance inference
- **Smart Router**: Intelligent routing - fast local for simple tasks, powerful remote for complex tasks
@@ -0,0 +1 @@
# Ollama-compatible API for SGLang
@@ -0,0 +1,137 @@
"""
Ollama-compatible API protocol definitions.
These models match the Ollama API format:
https://github.com/ollama/ollama/blob/main/docs/api.md
"""
from typing import Any, Dict, List, Literal, Optional, Union
from pydantic import BaseModel, Field
class OllamaMessage(BaseModel):
"""Ollama message format."""
role: str
content: str
images: Optional[List[str]] = None
class OllamaChatRequest(BaseModel):
"""Ollama /api/chat request format."""
model: str
messages: List[OllamaMessage]
stream: bool = True
format: Optional[Union[Literal["json"], Dict[str, Any]]] = None
options: Optional[Dict[str, Any]] = None
keep_alive: Optional[Union[float, str]] = None
think: Optional[Union[bool, Literal["low", "medium", "high"]]] = None
class OllamaChatResponse(BaseModel):
"""Ollama /api/chat response format (non-streaming)."""
model: str
created_at: str
message: OllamaMessage
done: bool = True
done_reason: Optional[str] = "stop"
total_duration: Optional[int] = None
load_duration: Optional[int] = None
prompt_eval_count: Optional[int] = None
prompt_eval_duration: Optional[int] = None
eval_count: Optional[int] = None
eval_duration: Optional[int] = None
class OllamaChatStreamResponse(BaseModel):
"""Ollama /api/chat streaming response chunk."""
model: str
created_at: str
message: OllamaMessage
done: bool = False
done_reason: Optional[str] = None
class OllamaGenerateRequest(BaseModel):
"""Ollama /api/generate request format."""
model: str
prompt: str
suffix: Optional[str] = None
system: Optional[str] = None
template: Optional[str] = None
context: Optional[List[int]] = None
stream: bool = True
raw: bool = False
format: Optional[Union[Literal["json"], Dict[str, Any]]] = None
options: Optional[Dict[str, Any]] = None
keep_alive: Optional[Union[float, str]] = None
images: Optional[List[str]] = None
think: Optional[bool] = None
class OllamaGenerateResponse(BaseModel):
"""Ollama /api/generate response format (non-streaming)."""
model: str
created_at: str
response: str
done: bool = True
done_reason: Optional[str] = "stop"
context: Optional[List[int]] = None
total_duration: Optional[int] = None
load_duration: Optional[int] = None
prompt_eval_count: Optional[int] = None
prompt_eval_duration: Optional[int] = None
eval_count: Optional[int] = None
eval_duration: Optional[int] = None
class OllamaGenerateStreamResponse(BaseModel):
"""Ollama /api/generate streaming response chunk."""
model: str
created_at: str
response: str
done: bool = False
done_reason: Optional[str] = None
class OllamaModelInfo(BaseModel):
"""Model information for /api/tags response."""
name: str
model: str
modified_at: str
size: int
digest: str
details: Optional[Dict[str, Any]] = None
class OllamaTagsResponse(BaseModel):
"""Ollama /api/tags response format."""
models: List[OllamaModelInfo]
class OllamaShowRequest(BaseModel):
"""Ollama /api/show request format."""
model: str
class OllamaShowResponse(BaseModel):
"""Ollama /api/show response format."""
license: str = ""
modelfile: str = ""
parameters: str = ""
template: str = ""
modified_at: str = ""
details: Dict[str, Any] = Field(default_factory=dict)
model_info: Dict[str, Any] = Field(default_factory=dict)
capabilities: List[str] = Field(default_factory=list)
@@ -0,0 +1,349 @@
"""
Ollama-compatible API serving handlers.
This module provides handlers that convert Ollama API requests to SGLang's
internal format and return Ollama-compatible responses.
"""
import time
from datetime import datetime, timezone
from typing import AsyncIterator, Union
import orjson
from fastapi import Request
from fastapi.responses import StreamingResponse
from sglang.srt.entrypoints.ollama.protocol import (
OllamaChatRequest,
OllamaChatResponse,
OllamaChatStreamResponse,
OllamaGenerateRequest,
OllamaGenerateResponse,
OllamaGenerateStreamResponse,
OllamaMessage,
OllamaModelInfo,
OllamaShowResponse,
OllamaTagsResponse,
)
from sglang.srt.managers.io_struct import GenerateReqInput
class OllamaServing:
"""Handler for Ollama-compatible API endpoints."""
def __init__(self, tokenizer_manager):
self.tokenizer_manager = tokenizer_manager
def _get_timestamp(self) -> str:
"""Get current timestamp in Ollama format."""
return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
def _convert_options_to_sampling_params(self, options: dict = None) -> dict:
"""Convert Ollama options to SGLang sampling params."""
sampling_params = {}
if options:
# Map Ollama options to SGLang params
param_mapping = {
"temperature": "temperature",
"top_p": "top_p",
"top_k": "top_k",
"num_predict": "max_new_tokens",
"stop": "stop",
"presence_penalty": "presence_penalty",
"frequency_penalty": "frequency_penalty",
"seed": "seed",
}
for ollama_param, sglang_param in param_mapping.items():
if ollama_param in options:
sampling_params[sglang_param] = options[ollama_param]
# Set a reasonable default for max_new_tokens if not specified
# Ollama users typically expect longer responses than SGLang's default (128)
if "max_new_tokens" not in sampling_params:
sampling_params["max_new_tokens"] = 2048
return sampling_params
async def handle_chat(
self, request: OllamaChatRequest, raw_request: Request
) -> Union[OllamaChatResponse, StreamingResponse]:
"""Handle /api/chat endpoint."""
model_name = self.tokenizer_manager.served_model_name
# Convert messages to SGLang format
messages = [
{"role": msg.role, "content": msg.content} for msg in request.messages
]
# Apply chat template using tokenizer
prompt_ids = self.tokenizer_manager.tokenizer.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
)
# Convert options to sampling params
sampling_params = self._convert_options_to_sampling_params(request.options)
# Create SGLang request with input_ids
gen_request = GenerateReqInput(
input_ids=prompt_ids,
sampling_params=sampling_params,
stream=request.stream,
)
if request.stream:
return await self._stream_chat_response(
gen_request, raw_request, model_name
)
else:
return await self._generate_chat_response(
gen_request, raw_request, model_name
)
async def _generate_chat_response(
self, gen_request: GenerateReqInput, raw_request: Request, model_name: str
) -> OllamaChatResponse:
"""Generate non-streaming chat response."""
start_time = time.time_ns()
# Get response from tokenizer manager
response = await self.tokenizer_manager.generate_request(
gen_request, raw_request
).__anext__()
end_time = time.time_ns()
total_duration = end_time - start_time
output_text = response.get("text", "")
return OllamaChatResponse(
model=model_name,
created_at=self._get_timestamp(),
message=OllamaMessage(role="assistant", content=output_text),
done=True,
done_reason="stop",
total_duration=total_duration,
prompt_eval_count=response.get("meta_info", {}).get("prompt_tokens", None),
eval_count=response.get("meta_info", {}).get("completion_tokens", None),
)
async def _stream_chat_response(
self, gen_request: GenerateReqInput, raw_request: Request, model_name: str
) -> StreamingResponse:
"""Generate streaming chat response."""
async def generate_stream() -> AsyncIterator[bytes]:
previous_text = ""
async for chunk in self.tokenizer_manager.generate_request(
gen_request, raw_request
):
text = chunk.get("text", "")
is_done = chunk.get("meta_info", {}).get("finish_reason") is not None
# Calculate delta (new text since last chunk)
delta = text[len(previous_text) :]
previous_text = text
if is_done:
# Final chunk
response = OllamaChatStreamResponse(
model=model_name,
created_at=self._get_timestamp(),
message=OllamaMessage(role="assistant", content=""),
done=True,
done_reason="stop",
)
else:
response = OllamaChatStreamResponse(
model=model_name,
created_at=self._get_timestamp(),
message=OllamaMessage(role="assistant", content=delta),
done=False,
)
yield orjson.dumps(response.model_dump()) + b"\n"
return StreamingResponse(
generate_stream(),
media_type="application/x-ndjson",
)
async def handle_generate(
self, request: OllamaGenerateRequest, raw_request: Request
) -> Union[OllamaGenerateResponse, StreamingResponse]:
"""Handle /api/generate endpoint."""
model_name = self.tokenizer_manager.served_model_name
# Build prompt
prompt = request.prompt
if request.system:
prompt = f"{request.system}\n\n{prompt}"
# Handle empty prompt - Ollama CLI sends empty requests on initialization
if not prompt or not prompt.strip():
empty_response = OllamaGenerateResponse(
model=model_name,
created_at=self._get_timestamp(),
response="",
done=True,
done_reason="stop",
)
if request.stream:
# Return streaming response with done=True
async def empty_stream() -> AsyncIterator[bytes]:
yield orjson.dumps(empty_response.model_dump()) + b"\n"
return StreamingResponse(
empty_stream(),
media_type="application/x-ndjson",
)
return empty_response
# Convert options to sampling params
sampling_params = self._convert_options_to_sampling_params(request.options)
# Create SGLang request
gen_request = GenerateReqInput(
text=prompt,
sampling_params=sampling_params,
stream=request.stream,
)
if request.stream:
return await self._stream_generate_response(
gen_request, raw_request, model_name
)
else:
return await self._generate_generate_response(
gen_request, raw_request, model_name
)
async def _generate_generate_response(
self, gen_request: GenerateReqInput, raw_request: Request, model_name: str
) -> OllamaGenerateResponse:
"""Generate non-streaming generate response."""
start_time = time.time_ns()
response = await self.tokenizer_manager.generate_request(
gen_request, raw_request
).__anext__()
end_time = time.time_ns()
total_duration = end_time - start_time
output_text = response.get("text", "")
return OllamaGenerateResponse(
model=model_name,
created_at=self._get_timestamp(),
response=output_text,
done=True,
done_reason="stop",
total_duration=total_duration,
prompt_eval_count=response.get("meta_info", {}).get("prompt_tokens", None),
eval_count=response.get("meta_info", {}).get("completion_tokens", None),
)
async def _stream_generate_response(
self, gen_request: GenerateReqInput, raw_request: Request, model_name: str
) -> StreamingResponse:
"""Generate streaming generate response."""
async def generate_stream() -> AsyncIterator[bytes]:
previous_text = ""
async for chunk in self.tokenizer_manager.generate_request(
gen_request, raw_request
):
text = chunk.get("text", "")
is_done = chunk.get("meta_info", {}).get("finish_reason") is not None
# Calculate delta (new text since last chunk)
delta = text[len(previous_text) :]
previous_text = text
if is_done:
response = OllamaGenerateStreamResponse(
model=model_name,
created_at=self._get_timestamp(),
response="",
done=True,
done_reason="stop",
)
else:
response = OllamaGenerateStreamResponse(
model=model_name,
created_at=self._get_timestamp(),
response=delta,
done=False,
)
yield orjson.dumps(response.model_dump()) + b"\n"
return StreamingResponse(
generate_stream(),
media_type="application/x-ndjson",
)
def get_tags(self) -> OllamaTagsResponse:
"""Handle /api/tags endpoint - list available models."""
model_name = self.tokenizer_manager.served_model_name
model_info = OllamaModelInfo(
name=model_name,
model=model_name,
modified_at=self._get_timestamp(),
size=0, # We don't track model size
digest="sha256:sglang0000000000000000000000000000000000000000000000000000000000",
details={
"format": "sglang",
"family": (
model_name.split("/")[-1] if "/" in model_name else model_name
),
"parameter_size": "unknown",
},
)
return OllamaTagsResponse(models=[model_info])
def get_show(self, model: str) -> OllamaShowResponse:
"""Handle /api/show endpoint - show model information."""
model_config = self.tokenizer_manager.model_config
# Extract model family from model name
model_family = model.split("/")[-1] if "/" in model else model
# Remove common suffixes to get base family
for suffix in ["-Instruct", "-Chat", "-Base"]:
if model_family.endswith(suffix):
model_family = model_family[: -len(suffix)]
break
# Build context length info
context_len = model_config.context_len if model_config else 4096
return OllamaShowResponse(
license="", # License info not available from SGLang
modelfile=f"FROM {model}\nPARAMETER num_ctx {context_len}\n",
parameters=f"num_ctx {context_len}",
template="", # Template info not easily accessible
modified_at=self._get_timestamp(),
details={
"parent_model": "",
"format": "sglang",
"family": model_family,
"families": [model_family],
"parameter_size": "unknown",
"quantization_level": "",
},
model_info={
"general.architecture": model_family,
"general.name": model,
"general.parameter_count": 0,
f"{model_family}.context_length": context_len,
f"{model_family}.block_count": 0,
f"{model_family}.embedding_length": 0,
f"{model_family}.attention.head_count": 0,
},
capabilities=["completion"],
)
@@ -0,0 +1,296 @@
"""
Smart Router: Automatically routes requests between local Ollama and remote SGLang.
Uses an LLM judge to classify tasks as simple or complex, then routes accordingly:
- Simple tasks → Local Ollama (fast response)
- Complex tasks → Remote SGLang (powerful model)
Usage:
from sglang.srt.entrypoints.ollama.smart_router import SmartRouter
router = SmartRouter(
local_host="http://localhost:11434",
remote_host="http://sglang-server:30001",
)
response = router.chat("Hello!")
"""
from typing import Optional
import ollama
class SmartRouter:
"""Routes requests between local Ollama and remote SGLang using LLM-based classification."""
# Classification prompt for LLM judge
CLASSIFICATION_PROMPT = """You are a task classifier. Classify the following user request into one of two categories.
Categories:
- SIMPLE: Quick responses, greetings, factual questions, definitions, translations, basic Q&A
- COMPLEX: Tasks requiring deep reasoning, multi-step analysis, long explanations, creative writing, detailed research
Reply with ONLY one word: either SIMPLE or COMPLEX.
User request: "{prompt}"
Category:"""
def __init__(
self,
local_host: str = "http://localhost:11434",
remote_host: str = "http://localhost:30001",
local_model: str = "llama3.2",
remote_model: str = "Qwen/Qwen2.5-1.5B-Instruct",
judge_model: Optional[str] = None,
judge_host: Optional[str] = None,
):
"""
Initialize the smart router.
Args:
local_host: URL of local Ollama server
remote_host: URL of remote SGLang server
local_model: Model name for local Ollama
remote_model: Model name for remote SGLang
judge_model: Model for LLM-based classification (default: same as local_model)
judge_host: Host for judge model (default: same as local_host)
"""
self.local_client = ollama.Client(host=local_host)
self.remote_client = ollama.Client(host=remote_host)
self.local_model = local_model
self.remote_model = remote_model
# Judge model configuration
self.judge_model = judge_model or local_model
self.judge_host = judge_host or local_host
self.judge_client = ollama.Client(host=self.judge_host)
def _classify_with_llm(
self, prompt: str, verbose: bool = False
) -> tuple[bool, str]:
"""
Use LLM to classify the prompt.
Returns:
Tuple of (use_remote, reason)
"""
try:
classification_prompt = self.CLASSIFICATION_PROMPT.format(
prompt=prompt[:500] # Limit prompt length for classification
)
response = self.judge_client.chat(
model=self.judge_model,
messages=[{"role": "user", "content": classification_prompt}],
options={"temperature": 0, "num_predict": 10},
)
result = response["message"]["content"].strip().upper()
if verbose:
print(f"[Router] LLM Judge: {result}")
if "COMPLEX" in result:
return True, "Complex task"
else:
return False, "Simple task"
except Exception as e:
if verbose:
print(f"[Router] LLM Judge failed: {e}, defaulting to local")
return False, "Judge failed, defaulting to local"
def should_use_remote(self, prompt: str, verbose: bool = False) -> tuple[bool, str]:
"""
Determine if the prompt should be routed to remote SGLang.
Args:
prompt: User's input prompt
verbose: Print debug information
Returns:
Tuple of (should_use_remote, reason)
"""
return self._classify_with_llm(prompt, verbose)
def chat(
self,
prompt: str,
messages: Optional[list] = None,
verbose: bool = False,
force_local: bool = False,
force_remote: bool = False,
) -> dict:
"""
Route the request and get response.
Args:
prompt: User's input (used if messages is None)
messages: Full message history (overrides prompt if provided)
verbose: Print routing decision
force_local: Force use of local model
force_remote: Force use of remote model
Returns:
Response dict with 'content', 'model', 'location', 'reason' keys
"""
# Build messages
if messages is None:
messages = [{"role": "user", "content": prompt}]
check_prompt = prompt
else:
# Use the last user message for routing decision
check_prompt = ""
for msg in reversed(messages):
if msg.get("role") == "user":
check_prompt = msg.get("content", "")
break
# Determine routing
if force_remote:
use_remote, reason = True, "Forced remote"
elif force_local:
use_remote, reason = False, "Forced local"
else:
use_remote, reason = self.should_use_remote(check_prompt, verbose)
if use_remote:
client = self.remote_client
model = self.remote_model
location = "Remote SGLang"
else:
client = self.local_client
model = self.local_model
location = "Local Ollama"
if verbose:
print(f"[Router] -> {location} | Model: {model}")
try:
response = client.chat(model=model, messages=messages)
return {
"content": response["message"]["content"],
"model": model,
"location": location,
"reason": reason,
}
except Exception as e:
# Fallback to the other option
if verbose:
print(f"[Router] {location} failed: {e}, falling back...")
fallback_client = (
self.remote_client if not use_remote else self.local_client
)
fallback_model = self.remote_model if not use_remote else self.local_model
fallback_location = "Remote SGLang" if not use_remote else "Local Ollama"
response = fallback_client.chat(model=fallback_model, messages=messages)
return {
"content": response["message"]["content"],
"model": fallback_model,
"location": fallback_location,
"reason": f"Fallback from {location}",
}
def chat_stream(
self,
prompt: str,
messages: Optional[list] = None,
verbose: bool = False,
force_local: bool = False,
force_remote: bool = False,
):
"""
Route the request and stream response.
Yields:
Response chunks
"""
if messages is None:
messages = [{"role": "user", "content": prompt}]
check_prompt = prompt
else:
check_prompt = ""
for msg in reversed(messages):
if msg.get("role") == "user":
check_prompt = msg.get("content", "")
break
if force_remote:
use_remote, reason = True, "Forced remote"
elif force_local:
use_remote, reason = False, "Forced local"
else:
use_remote, reason = self.should_use_remote(check_prompt, verbose)
if use_remote:
client = self.remote_client
model = self.remote_model
location = "Remote SGLang"
else:
client = self.local_client
model = self.local_model
location = "Local Ollama"
if verbose:
print(f"[Router] -> {location} | Model: {model}")
for chunk in client.chat(model=model, messages=messages, stream=True):
yield chunk
def main():
"""Interactive demo of the smart router."""
print("=" * 60)
print("Smart Router: Local Ollama <-> Remote SGLang")
print("=" * 60)
print("\nRouting strategy:")
print(" LLM Judge classifies each request as SIMPLE or COMPLEX")
print(" - SIMPLE tasks -> Local Ollama (fast)")
print(" - COMPLEX tasks -> Remote SGLang (powerful)")
print("\nType 'quit' to exit\n")
router = SmartRouter(
local_host="http://localhost:11434",
remote_host="http://localhost:30001",
local_model="llama3.2",
remote_model="Qwen/Qwen2.5-1.5B-Instruct",
)
messages = []
while True:
try:
user_input = input("You: ").strip()
if user_input.lower() in ["quit", "exit", "q"]:
print("Goodbye!")
break
if not user_input:
continue
messages.append({"role": "user", "content": user_input})
# Use streaming for real-time output
print("\nAssistant: ", end="", flush=True)
full_response = ""
for chunk in router.chat_stream(
prompt=user_input, messages=messages, verbose=True
):
content = chunk.get("message", {}).get("content", "")
if content:
print(content, end="", flush=True)
full_response += content
print("\n")
messages.append({"role": "assistant", "content": full_response})
except KeyboardInterrupt:
print("\nGoodbye!")
break
except Exception as e:
print(f"Error: {e}\n")
if __name__ == "__main__":
main()
@@ -0,0 +1,81 @@
"""Single home for the chat-encoding dispatch.
Which encoder turns chat messages into prompt tokens is a property of the
model, so the serving path and offline tools (benchmarks, evals) must resolve
it here instead of re-deriving it from model architectures themselves.
"""
from __future__ import annotations
from typing import Any, Dict, List, Optional
def resolve_chat_encoding_spec(
*,
hf_config: Any,
tokenizer: Any,
tool_call_parser: Optional[str] = None,
) -> Optional[str]:
"""Return the chat encoding spec for a model: "dsv4", "dsv32", or None.
None means the default path (HF chat template).
"""
if tool_call_parser == "deepseekv4":
return "dsv4"
if tool_call_parser == "deepseekv32":
return "dsv32"
architectures = hf_config.architectures
arch = architectures[0] if architectures else ""
if "DeepseekV4" in arch:
return "dsv4"
has_chat_template = tokenizer is not None and tokenizer.chat_template is not None
if "DeepseekV3" in arch and not has_chat_template:
return "dsv32"
return None
def encode_simple_chat(
*,
tokenizer: Any,
spec: Optional[str],
messages: List[Dict[str, Any]],
thinking_mode: str = "chat",
) -> List[int]:
"""Encode a plain-text chat conversation into prompt token ids.
Minimal encode for offline tools: no tools, no multimodal content, no
continue_final_message; the serving path keeps its full request-level
pipeline in ``serving_chat``. Like
``serving_chat``, an empty system message is prepended when the
conversation does not start with one (for the dsv4/dsv32 encoders this
currently renders to zero tokens, but keeping the insertion explicit ties
this helper to the serving semantics rather than to that coincidence).
"""
if spec in ("dsv4", "dsv32"):
if messages and messages[0]["role"] != "system":
messages = [{"role": "system", "content": ""}] + list(messages)
if spec == "dsv4":
from sglang.srt.entrypoints.openai import encoding_dsv4
real_input = encoding_dsv4.encode_messages(
messages, thinking_mode=thinking_mode
)
else:
from sglang.srt.entrypoints.openai import encoding_dsv32
real_input = encoding_dsv32.encode_messages(
messages, thinking_mode=thinking_mode
)
return tokenizer.encode(real_input)
if getattr(tokenizer, "chat_template", None) is None:
raise ValueError(
"This model has no HF chat template and no custom chat encoder; "
f"cannot encode chat messages with {getattr(tokenizer, 'name_or_path', tokenizer)!r}."
)
return tokenizer.apply_chat_template(
messages, add_generation_prompt=True, tokenize=True
)
@@ -0,0 +1,471 @@
# Adapted from https://huggingface.co/deepseek-ai/DeepSeek-V3.2/blob/main/encoding/encoding_dsv32.py
import copy
import json
import re
from typing import Any, Dict, List, Optional, Tuple, Union
class DS32EncodingError(Exception):
pass
TOOLS_SYSTEM_TEMPLATE = """## Tools
You have access to a set of tools you can use to answer the user's question.
You can invoke functions by writing a "<{dsml_token}function_calls>" block like the following as part of your reply to the user:
<{dsml_token}function_calls>
<{dsml_token}invoke name="$FUNCTION_NAME">
<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}parameter>
...
</{dsml_token}invoke>
<{dsml_token}invoke name="$FUNCTION_NAME2">
...
</{dsml_token}invoke>
</{dsml_token}function_calls>
String and scalar parameters should be specified as is without any escaping or quotes, while lists and objects should use JSON format. The "string" attribute should be set to "true" for string type parameters and "false" for other types (numbers, booleans, arrays, objects).
If the thinking_mode is enabled, then after function results you should strongly consider outputting a thinking block. Here is an example:
<{dsml_token}function_calls>
...
</{dsml_token}function_calls>
<function_results>
...
</function_results>
{thinking_start_token}...thinking about results{thinking_end_token}
Here are the functions available in JSONSchema format:
<functions>
{tool_schemas}
</functions>
"""
bos_token: str = "<begin▁of▁sentence>"
eos_token: str = "<end▁of▁sentence>"
thinking_start_token: str = "<think>"
thinking_end_token: str = "</think>"
dsml_token: str = "DSML"
system_msg_template: str = "{content}"
user_msg_template: str = "<User>{content}<Assistant>"
assistant_msg_template: str = "{reasoning}{content}{tool_calls}<end▁of▁sentence>"
thinking_template = "{reasoning_content}"
response_format_template: str = (
"## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}"
)
tool_call_template: str = (
'<{dsml_token}invoke name="{name}">\n{arguments}\n</{dsml_token}invoke>'
)
tool_calls_template = (
"<{dsml_token}function_calls>\n{tool_calls}\n</{dsml_token}function_calls>"
)
tool_output_template: str = "\n<result>{content}</result>"
def to_json(value: Any) -> str:
try:
return json.dumps(value, ensure_ascii=False)
except:
return json.dumps(value, ensure_ascii=True)
def tools_from_openai_format(tools):
return [tool["function"] for tool in tools]
def tool_calls_from_openai_format(tool_calls):
return [
{
"name": tool_call["function"]["name"],
"arguments": tool_call["function"]["arguments"],
}
for tool_call in tool_calls
]
def tool_calls_to_openai_format(tool_calls):
return [
{
"type": "function",
"function": {
"name": tool_call["name"],
"arguments": tool_call["arguments"],
},
}
for tool_call in tool_calls
]
def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str:
p_dsml_template = """<{dsml_token}parameter name="{key}" string="{is_str}">{value}</{dsml_token}parameter>"""
P_dsml_strs = []
raw_arguments = tool_call["arguments"]
arguments = (
json.loads(raw_arguments) if isinstance(raw_arguments, str) else raw_arguments
)
if not isinstance(arguments, dict):
raise ValueError(
"Assistant tool call function.arguments must be a JSON object."
)
for k, v in arguments.items():
p_dsml_str = p_dsml_template.format(
dsml_token=dsml_token,
key=k,
is_str="true" if isinstance(v, str) else "false",
value=v if isinstance(v, str) else to_json(v),
)
P_dsml_strs.append(p_dsml_str)
return "\n".join(P_dsml_strs)
def decode_dsml_to_arguments(
tool_name: str, tool_args: Dict[str, Tuple[str, str]]
) -> Dict[str, str]:
def _decode_value(key: str, value: str, string: str):
if string == "true":
value = to_json(value)
return f"{to_json(key)}: {value}"
tool_args_json = (
"{"
+ ", ".join(
[_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]
)
+ "}"
)
return dict(name=tool_name, arguments=tool_args_json)
def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str:
tools_json = [to_json(t) for t in tools]
return TOOLS_SYSTEM_TEMPLATE.format(
tool_schemas="\n".join(tools_json),
dsml_token=dsml_token,
thinking_start_token=thinking_start_token,
thinking_end_token=thinking_end_token,
)
def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
last_user_index = -1
for idx in range(len(messages) - 1, -1, -1):
if messages[idx].get("role") in ["user", "developer"]:
last_user_index = idx
break
return last_user_index
def render_message(
index: int, messages: List[Dict[str, Any]], thinking_mode: str
) -> str:
if not (0 <= index < len(messages)):
raise DS32EncodingError(
f"Index {index} out of range for messages list of length {len(messages)}"
)
if thinking_mode not in ["chat", "thinking"]:
raise DS32EncodingError(f"Invalid thinking_mode `{thinking_mode}`")
prompt = ""
msg = messages[index]
last_user_idx = find_last_user_index(messages)
role = msg.get("role")
content = msg.get("content")
tools = msg.get("tools")
response_format = msg.get("response_format")
tool_calls = msg.get("tool_calls")
reasoning_content = msg.get("reasoning_content")
if tools:
tools = tools_from_openai_format(tools)
if tool_calls:
tool_calls = tool_calls_from_openai_format(tool_calls)
if role == "system":
prompt += system_msg_template.format(content=content or "")
if tools:
prompt += "\n\n" + render_tools(tools)
if response_format:
prompt += "\n\n" + response_format_template.format(
schema=to_json(response_format)
)
elif role == "developer":
if not content:
raise DS32EncodingError(f"Invalid message for role `{role}`: {msg}")
content_developer = ""
if tools:
content_developer += "\n\n" + render_tools(tools)
if response_format:
content_developer += "\n\n" + response_format_template.format(
schema=to_json(response_format)
)
content_developer += "\n\n# The user's message is: {}".format(content)
prompt += user_msg_template.format(content=content_developer)
if index == last_user_idx and thinking_mode == "thinking":
prompt += thinking_start_token
else:
prompt += thinking_end_token
elif role == "user":
prompt += user_msg_template.format(content=content)
if index == last_user_idx and thinking_mode == "thinking":
prompt += thinking_start_token
else:
prompt += thinking_end_token
elif role == "tool":
prev_assistant_idx = index - 1
assistant_msg = messages[prev_assistant_idx]
while prev_assistant_idx >= 0 and assistant_msg.get("role") == "tool":
prev_assistant_idx -= 1
assistant_msg = messages[prev_assistant_idx]
if not (
index == 0
or (prev_assistant_idx >= 0 and assistant_msg.get("role") == "assistant")
):
raise DS32EncodingError(f"Invalid messages at {index}:\n{assistant_msg}")
tool_call_order = index - prev_assistant_idx
assistant_tool_calls = assistant_msg.get("tool_calls")
if not (assistant_tool_calls and len(assistant_tool_calls) >= tool_call_order):
raise DS32EncodingError("No tool calls but found tool output")
if tool_call_order == 1:
prompt += "\n\n<function_results>"
prompt += tool_output_template.format(content=content)
if tool_call_order == len(assistant_tool_calls):
prompt += "\n</function_results>"
if index >= last_user_idx and thinking_mode == "thinking":
prompt += "\n\n" + thinking_start_token
else:
prompt += "\n\n" + thinking_end_token
elif role == "assistant":
prev_assistant_idx = index
thinking_part = ""
tool_calls_content = ""
if tool_calls:
tool_calls = [
tool_call_template.format(
dsml_token=dsml_token,
name=tool_call.get("name"),
arguments=encode_arguments_to_dsml(tool_call),
)
for tool_call in tool_calls
]
tool_calls_content += "\n\n" + tool_calls_template.format(
dsml_token=dsml_token, tool_calls="\n".join(tool_calls)
)
summary_content = content or ""
if thinking_mode == "thinking" and index > last_user_idx:
if not (reasoning_content or tool_calls):
raise DS32EncodingError(
f"ThinkingMode: {thinking_mode}, invalid message without reasoning_content/tool_calls `{msg}` after last user message"
)
thinking_part = (
thinking_template.format(reasoning_content=reasoning_content or "")
+ thinking_end_token
)
prompt += assistant_msg_template.format(
reasoning=thinking_part,
content=summary_content,
tool_calls=tool_calls_content,
)
else:
raise NotImplementedError(f"Unknown role: {role}")
return prompt
def drop_thinking_messages(
messages: List[Dict[str, Any]], last_user_idx: Optional[int] = None
) -> List[Dict[str, Any]]:
messages_wo_thinking: List[Dict[str, Any]] = []
last_user_idx = (
find_last_user_index(messages) if last_user_idx is None else last_user_idx
)
for idx, msg in enumerate(messages):
role = msg.get("role")
if role in ["user", "system", "tool"] or idx >= last_user_idx:
messages_wo_thinking.append(msg)
continue
elif role == "assistant":
msg_wo_thinking = copy.copy(msg)
msg_wo_thinking.pop("reasoning_content", None)
messages_wo_thinking.append(msg_wo_thinking)
return messages_wo_thinking
def encode_messages(
messages: List[Dict[str, Any]],
thinking_mode: str,
context: Optional[List[Dict[str, Any]]] = None,
drop_thinking: bool = True,
add_default_bos_token: bool = True,
) -> str:
context = context if context else []
full_messages = context + messages
prompt = bos_token if add_default_bos_token and len(context) == 0 else ""
if thinking_mode == "thinking" and drop_thinking:
full_messages = drop_thinking_messages(full_messages)
for idx in range(len(messages)):
prompt += render_message(
idx + len(context), full_messages, thinking_mode=thinking_mode
)
return prompt
def _read_until_stop(
index: int, text: str, stop: List[str]
) -> Tuple[int, str, Optional[str]]:
min_pos = len(text)
matched_stop = None
for s in stop:
pos = text.find(s, index)
if pos != -1 and pos < min_pos:
min_pos = pos
matched_stop = s
if matched_stop:
content = text[index:min_pos]
return min_pos + len(matched_stop), content, matched_stop
else:
content = text[index:]
return len(text), content, None
def parse_tool_calls(index: int, text: str):
tool_calls: List[Dict[str, Any]] = []
stop_token = None
tool_calls_end_token = f"</{dsml_token}function_calls>"
while index < len(text):
index, _, stop_token = _read_until_stop(
index, text, [f"<{dsml_token}invoke", tool_calls_end_token]
)
if _ != ">\n":
raise DS32EncodingError("Tool call format error")
if stop_token == tool_calls_end_token:
break
if stop_token is None:
raise DS32EncodingError("Missing special token")
index, tool_name_content, stop_token = _read_until_stop(
index, text, [f"<{dsml_token}parameter", f"</{dsml_token}invoke"]
)
p_tool_name = re.findall(
r'^\s*name="(.*?)">\n$', tool_name_content, flags=re.DOTALL
)
if len(p_tool_name) != 1:
raise DS32EncodingError("Tool name format error")
tool_name = p_tool_name[0]
tool_args: Dict[str, Tuple[str, str]] = {}
while stop_token == f"<{dsml_token}parameter":
index, param_content, stop_token = _read_until_stop(
index, text, [f"/{dsml_token}parameter"]
)
param_kv = re.findall(
r'^ name="(.*?)" string="(true|false)">(.*?)<$',
param_content,
flags=re.DOTALL,
)
if len(param_kv) != 1:
raise DS32EncodingError("Parameter format error")
param_name, string, param_value = param_kv[0]
if param_name in tool_args:
raise DS32EncodingError("Duplicate parameter name")
tool_args[param_name] = (param_value, string)
index, content, stop_token = _read_until_stop(
index, text, [f"<{dsml_token}parameter", f"</{dsml_token}invoke"]
)
if content != ">\n":
raise DS32EncodingError("Parameter format error")
tool_call = decode_dsml_to_arguments(tool_name=tool_name, tool_args=tool_args)
tool_calls.append(tool_call)
return index, stop_token, tool_calls
# NOTE: This function is designed to parse only correctly formatted string and will not attempt to correct malformed output that may be generated by the model.
def parse_message_from_completion_text(text: str, thinking_mode: str):
summary_content, reasoning_content, tool_calls = "", "", []
index, stop_token = 0, None
tool_calls_start_token = f"\n\n<{dsml_token}function_calls"
is_thinking, is_tool_calling = thinking_mode == "thinking", False
if is_thinking:
index, content_delta, stop_token = _read_until_stop(
index, text, [thinking_end_token, tool_calls_start_token]
)
reasoning_content = content_delta
if stop_token != thinking_end_token:
raise DS32EncodingError("Invalid thinking format")
index, content_delta, stop_token = _read_until_stop(
index, text, [eos_token, tool_calls_start_token]
)
summary_content = content_delta
if stop_token == tool_calls_start_token:
is_tool_calling = True
else:
if stop_token != eos_token:
raise DS32EncodingError("Invalid summary format")
if is_tool_calling:
index, stop_token, tool_calls = parse_tool_calls(index, text)
index, tool_ends_text, stop_token = _read_until_stop(index, text, [eos_token])
if tool_ends_text:
raise DS32EncodingError("Unexpected content after tool calls")
if not (len(text) == index and stop_token in [eos_token, None]):
raise DS32EncodingError("Unexpected content at end")
for sp_token in [
bos_token,
eos_token,
thinking_start_token,
thinking_end_token,
dsml_token,
]:
if sp_token in summary_content or sp_token in reasoning_content:
raise DS32EncodingError("Unexpected special token in content")
return {
"role": "assistant",
"content": summary_content,
"reasoning_content": reasoning_content,
"tool_calls": tool_calls_to_openai_format(tool_calls),
}
@@ -0,0 +1,854 @@
# Adapted from the DeepSeek-V4 release reference implementation.
"""
DeepSeek-V4 Encoding
A self-contained implementation for encoding/decoding DeepSeek-V4 chat messages
with tool calling, thinking mode, and quick instruction task support.
"""
import copy
import json
import re
from typing import Any, Dict, List, Optional, Tuple, Union
# ============================================================
# Special Tokens
# ============================================================
bos_token: str = "<begin▁of▁sentence>"
eos_token: str = "<end▁of▁sentence>"
thinking_start_token: str = "<think>"
thinking_end_token: str = "</think>"
dsml_token: str = "DSML"
USER_SP_TOKEN = "<User>"
ASSISTANT_SP_TOKEN = "<Assistant>"
LATEST_REMINDER_SP_TOKEN = "<latest_reminder>"
# Task special tokens for internal classification tasks
DS_TASK_SP_TOKENS = {
"action": "<action>",
"query": "<query>",
"authority": "<authority>",
"domain": "<domain>",
"title": "<title>",
"read_url": "<read_url>",
}
VALID_TASKS = set(DS_TASK_SP_TOKENS.keys())
# ============================================================
# Templates
# ============================================================
system_msg_template: str = "{content}"
user_msg_template: str = "{content}"
latest_reminder_msg_template: str = "{content}"
assistant_msg_template: str = "{reasoning}{content}{tool_calls}" + eos_token
assistant_msg_wo_eos_template: str = "{reasoning}{content}{tool_calls}"
thinking_template: str = "{reasoning_content}"
response_format_template: str = (
"## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}"
)
tool_call_template: str = (
'<{dsml_token}invoke name="{name}">\n{arguments}\n</{dsml_token}invoke>'
)
tool_calls_template = (
"<{dsml_token}{tc_block_name}>\n{tool_calls}\n</{dsml_token}{tc_block_name}>"
)
tool_calls_block_name: str = "tool_calls"
tool_output_template: str = "<tool_result>{content}</tool_result>"
REASONING_EFFORT_MAX = (
"Reasoning Effort: Absolute maximum with no shortcuts permitted.\n"
"You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n"
"Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n"
)
TOOLS_TEMPLATE = """## Tools
You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following:
<{dsml_token}tool_calls>
<{dsml_token}invoke name="$TOOL_NAME">
<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}parameter>
...
</{dsml_token}invoke>
<{dsml_token}invoke name="$TOOL_NAME2">
...
</{dsml_token}invoke>
</{dsml_token}tool_calls>
String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response.
Otherwise, output directly after {thinking_end_token} with tool calls or final response.
### Available Tool Schemas
{tool_schemas}
You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
"""
# ============================================================
# Utility Functions
# ============================================================
def to_json(value: Any) -> str:
"""Serialize a value to JSON string."""
try:
return json.dumps(value, ensure_ascii=False)
except:
return json.dumps(value, ensure_ascii=True)
def tools_from_openai_format(tools):
"""Extract function definitions from OpenAI-format tool list."""
return [tool["function"] for tool in tools]
def tool_calls_from_openai_format(tool_calls):
"""Convert OpenAI-format tool calls to internal format."""
return [
{
"name": tool_call["function"]["name"],
"arguments": tool_call["function"]["arguments"],
}
for tool_call in tool_calls
]
def tool_calls_to_openai_format(tool_calls):
"""Convert internal tool calls to OpenAI format."""
return [
{
"type": "function",
"function": {
"name": tool_call["name"],
"arguments": tool_call["arguments"],
},
}
for tool_call in tool_calls
]
def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str:
"""
Encode tool call arguments into DSML parameter format.
Args:
tool_call: Dict with "name" and "arguments" keys.
Returns:
DSML-formatted parameter string.
"""
p_dsml_template = '<{dsml_token}parameter name="{key}" string="{is_str}">{value}</{dsml_token}parameter>'
P_dsml_strs = []
raw_arguments = tool_call["arguments"]
arguments = (
json.loads(raw_arguments) if isinstance(raw_arguments, str) else raw_arguments
)
if not isinstance(arguments, dict):
raise ValueError(
"Assistant tool call function.arguments must be a JSON object."
)
for k, v in arguments.items():
p_dsml_str = p_dsml_template.format(
dsml_token=dsml_token,
key=k,
is_str="true" if isinstance(v, str) else "false",
value=v if isinstance(v, str) else to_json(v),
)
P_dsml_strs.append(p_dsml_str)
return "\n".join(P_dsml_strs)
def decode_dsml_to_arguments(
tool_name: str, tool_args: Dict[str, Tuple[str, str]]
) -> Dict[str, str]:
"""
Decode DSML parameters back to a tool call dict.
Args:
tool_name: Name of the tool.
tool_args: Dict mapping param_name -> (value, is_string_flag).
Returns:
Dict with "name" and "arguments" (JSON string) keys.
"""
def _decode_value(key: str, value: str, string: str):
if string == "true":
value = to_json(value)
return f"{to_json(key)}: {value}"
tool_args_json = (
"{"
+ ", ".join(
[_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]
)
+ "}"
)
return dict(name=tool_name, arguments=tool_args_json)
def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str:
"""
Render tool schemas into the system prompt format.
Args:
tools: List of tool schema dicts (each with name, description, parameters).
Returns:
Formatted tools section string.
"""
tools_json = [to_json(t) for t in tools]
return TOOLS_TEMPLATE.format(
tool_schemas="\n".join(tools_json),
dsml_token=dsml_token,
thinking_start_token=thinking_start_token,
thinking_end_token=thinking_end_token,
)
def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
"""Find the index of the last user/developer message."""
last_user_index = -1
for idx in range(len(messages) - 1, -1, -1):
if messages[idx].get("role") in ["user", "developer"]:
last_user_index = idx
break
return last_user_index
def attach_task_to_last_user_message(messages: List[Dict[str, Any]], task: str) -> None:
"""Set `task` on the most recent user/developer message; raise if none exists."""
idx = find_last_user_index(messages)
if idx == -1:
raise ValueError(
"`task` requires at least one message with role='user' or 'developer'."
)
messages[idx]["task"] = task
# ============================================================
# Message Rendering
# ============================================================
def render_message(
index: int,
messages: List[Dict[str, Any]],
thinking_mode: str,
drop_thinking: bool = True,
reasoning_effort: Optional[str] = None,
) -> str:
"""
Render a single message at the given index into its encoded string form.
This is the core function that converts each message in the conversation
into the DeepSeek-V4 format.
Args:
index: Index of the message to render.
messages: Full list of messages in the conversation.
thinking_mode: Either "chat" or "thinking".
drop_thinking: Whether to drop reasoning content from earlier turns.
reasoning_effort: Optional reasoning effort level ("max", "high", or None).
Returns:
Encoded string for this message.
"""
assert 0 <= index < len(messages)
assert thinking_mode in [
"chat",
"thinking",
], f"Invalid thinking_mode `{thinking_mode}`"
prompt = ""
msg = messages[index]
last_user_idx = find_last_user_index(messages)
role = msg.get("role")
content = msg.get("content")
tools = msg.get("tools")
response_format = msg.get("response_format")
tool_calls = msg.get("tool_calls")
reasoning_content = msg.get("reasoning_content")
wo_eos = msg.get("wo_eos", False)
if tools:
tools = tools_from_openai_format(tools)
if tool_calls:
tool_calls = tool_calls_from_openai_format(tool_calls)
# Reasoning effort prefix (only at index 0 in thinking mode with max effort)
assert reasoning_effort in [
"max",
None,
"high",
], f"Invalid reasoning effort: {reasoning_effort}"
if index == 0 and thinking_mode == "thinking" and reasoning_effort == "max":
prompt += REASONING_EFFORT_MAX
if role == "system":
prompt += system_msg_template.format(content=content or "")
if tools:
prompt += "\n\n" + render_tools(tools)
if response_format:
prompt += "\n\n" + response_format_template.format(
schema=to_json(response_format)
)
elif role == "developer":
assert content, f"Invalid message for role `{role}`: {msg}"
content_developer = USER_SP_TOKEN
content_developer += content
if tools:
content_developer += "\n\n" + render_tools(tools)
if response_format:
content_developer += "\n\n" + response_format_template.format(
schema=to_json(response_format)
)
prompt += user_msg_template.format(content=content_developer)
elif role == "user":
prompt += USER_SP_TOKEN
# Handle content blocks (tool results mixed with text)
content_blocks = msg.get("content_blocks")
if content_blocks:
parts = []
for block in content_blocks:
block_type = block.get("type")
if block_type == "text":
parts.append(block.get("text", ""))
elif block_type == "tool_result":
tool_content = block.get("content", "")
if isinstance(tool_content, list):
text_parts = []
for b in tool_content:
if b.get("type") == "text":
text_parts.append(b.get("text", ""))
else:
text_parts.append(f"[Unsupported {b.get('type')}]")
tool_content = "\n\n".join(text_parts)
parts.append(tool_output_template.format(content=tool_content))
else:
parts.append(f"[Unsupported {block_type}]")
prompt += "\n\n".join(parts)
else:
prompt += content or ""
elif role == "latest_reminder":
prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format(
content=content
)
elif role == "tool":
raise NotImplementedError(
"deepseek_v4 merges tool messages into user; please preprocess with merge_tool_messages()"
)
elif role == "assistant":
thinking_part = ""
tc_content = ""
if tool_calls:
tc_list = [
tool_call_template.format(
dsml_token=dsml_token,
name=tc.get("name"),
arguments=encode_arguments_to_dsml(tc),
)
for tc in tool_calls
]
tc_content += "\n\n" + tool_calls_template.format(
dsml_token=dsml_token,
tool_calls="\n".join(tc_list),
tc_block_name=tool_calls_block_name,
)
summary_content = content or ""
rc = reasoning_content or ""
# Check if previous message has a task - if so, this is a task output (no thinking)
prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None
if thinking_mode == "thinking" and not prev_has_task:
if not drop_thinking or index > last_user_idx:
thinking_part = (
thinking_template.format(reasoning_content=rc) + thinking_end_token
)
else:
thinking_part = ""
if wo_eos:
prompt += assistant_msg_wo_eos_template.format(
reasoning=thinking_part,
content=summary_content,
tool_calls=tc_content,
)
else:
prompt += assistant_msg_template.format(
reasoning=thinking_part,
content=summary_content,
tool_calls=tc_content,
)
else:
raise NotImplementedError(f"Unknown role: {role}")
# Append transition tokens based on what follows
if index + 1 < len(messages) and messages[index + 1].get("role") not in [
"assistant",
"latest_reminder",
]:
return prompt
task = messages[index].get("task")
if task is not None:
# Task special token for internal classification tasks
assert (
task in VALID_TASKS
), f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}"
task_sp_token = DS_TASK_SP_TOKENS[task]
if task != "action":
# Non-action tasks: append task sp token directly after the message
prompt += task_sp_token
else:
# Action task: append Assistant + thinking token + action sp token
prompt += ASSISTANT_SP_TOKEN
prompt += (
thinking_end_token
if thinking_mode != "thinking"
else thinking_start_token
)
prompt += task_sp_token
elif messages[index].get("role") in ["user", "developer"]:
# Normal generation: append Assistant + thinking token
prompt += ASSISTANT_SP_TOKEN
if not drop_thinking and thinking_mode == "thinking":
prompt += thinking_start_token
elif drop_thinking and thinking_mode == "thinking" and index >= last_user_idx:
prompt += thinking_start_token
else:
prompt += thinking_end_token
return prompt
# ============================================================
# Preprocessing
# ============================================================
def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""
Merge tool messages into the preceding user message using content_blocks format.
DeepSeek-V4 does not have a standalone "tool" role; instead, tool results
are encoded as <tool_result> blocks within user messages.
This function converts a standard OpenAI-format conversation (with separate
"tool" role messages) into V4 format where tool results are merged into
user messages.
Args:
messages: List of message dicts in OpenAI format.
Returns:
Processed message list with tool messages merged into user messages.
"""
merged: List[Dict[str, Any]] = []
for msg in messages:
msg = copy.deepcopy(msg)
role = msg.get("role")
if role == "tool":
# Convert tool message to a user message with tool_result block
tool_block = {
"type": "tool_result",
"tool_use_id": msg.get("tool_call_id", ""),
"content": msg.get("content", ""),
}
# Merge into previous message if it's already a user (merged tool)
if (
merged
and merged[-1].get("role") == "user"
and "content_blocks" in merged[-1]
):
merged[-1]["content_blocks"].append(tool_block)
else:
merged.append(
{
"role": "user",
"content_blocks": [tool_block],
}
)
elif role == "user":
text_block = {"type": "text", "text": msg.get("content", "")}
if (
merged
and merged[-1].get("role") == "user"
and "content_blocks" in merged[-1]
and merged[-1].get("task") is None
):
merged[-1]["content_blocks"].append(text_block)
else:
new_msg = {
"role": "user",
"content": msg.get("content", ""),
"content_blocks": [text_block],
}
# Preserve extra fields (task, wo_eos, mask, etc.)
for key in ("task", "wo_eos", "mask"):
if key in msg:
new_msg[key] = msg[key]
merged.append(new_msg)
else:
merged.append(msg)
return merged
def sort_tool_results_by_call_order(
messages: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
"""
Sort tool_result blocks within user messages by the order of tool_calls
in the preceding assistant message.
Args:
messages: Preprocessed message list (after merge_tool_messages).
Returns:
Message list with sorted tool result blocks.
"""
last_tool_call_order: Dict[str, int] = {}
for msg in messages:
role = msg.get("role")
if role == "assistant" and msg.get("tool_calls"):
last_tool_call_order = {}
for idx, tc in enumerate(msg["tool_calls"]):
tc_id = tc.get("id") or tc.get("function", {}).get("id", "")
if tc_id:
last_tool_call_order[tc_id] = idx
elif role == "user" and msg.get("content_blocks"):
tool_blocks = [
b for b in msg["content_blocks"] if b.get("type") == "tool_result"
]
if len(tool_blocks) > 1 and last_tool_call_order:
sorted_blocks = sorted(
tool_blocks,
key=lambda b: last_tool_call_order.get(b.get("tool_use_id", ""), 0),
)
sorted_idx = 0
new_blocks = []
for block in msg["content_blocks"]:
if block.get("type") == "tool_result":
new_blocks.append(sorted_blocks[sorted_idx])
sorted_idx += 1
else:
new_blocks.append(block)
msg["content_blocks"] = new_blocks
return messages
# ============================================================
# Main Encoding Function
# ============================================================
def encode_messages(
messages: List[Dict[str, Any]],
thinking_mode: str,
context: Optional[List[Dict[str, Any]]] = None,
drop_thinking: bool = True,
add_default_bos_token: bool = True,
reasoning_effort: Optional[str] = None,
) -> str:
"""
Encode a list of messages into the DeepSeek-V4 prompt format.
This is the main entry point for encoding conversations. It handles:
- BOS token insertion
- Thinking mode with optional reasoning content dropping
- Tool message merging into user messages
- Multi-turn conversation context
Args:
messages: List of message dicts to encode.
thinking_mode: Either "chat" or "thinking".
context: Optional preceding context messages (already encoded prefix).
drop_thinking: If True, drop reasoning_content from earlier assistant turns
(only keep reasoning for messages after the last user message).
add_default_bos_token: Whether to prepend BOS token at conversation start.
reasoning_effort: Optional reasoning effort level ("max", "high", or None).
Returns:
The encoded prompt string.
"""
context = context if context else []
# Preprocess: merge tool messages and sort tool results
messages = merge_tool_messages(messages)
messages = sort_tool_results_by_call_order(context + messages)[len(context) :]
if context:
context = merge_tool_messages(context)
context = sort_tool_results_by_call_order(context)
full_messages = context + messages
prompt = bos_token if add_default_bos_token and len(context) == 0 else ""
# Resolve drop_thinking: if any message has tools defined, don't drop thinking
effective_drop_thinking = drop_thinking
if any(m.get("tools") for m in full_messages):
effective_drop_thinking = False
if thinking_mode == "thinking" and effective_drop_thinking:
full_messages = _drop_thinking_messages(full_messages)
# After dropping, recalculate how many messages to render
# (context may have shrunk too)
num_to_render = len(full_messages) - len(_drop_thinking_messages(context))
context_len = len(full_messages) - num_to_render
else:
num_to_render = len(messages)
context_len = len(context)
for idx in range(num_to_render):
prompt += render_message(
idx + context_len,
full_messages,
thinking_mode=thinking_mode,
drop_thinking=effective_drop_thinking,
reasoning_effort=reasoning_effort,
)
return prompt
def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""
Drop reasoning_content and non-essential messages before the last user message.
Behavior:
- Messages with role in ["user", "system", "tool", "latest_reminder"] are always kept.
- Messages at or after the last user index are always kept.
- Assistant messages before the last user get reasoning_content removed.
- Developer messages before the last user are dropped entirely.
"""
last_user_idx = find_last_user_index(messages)
result = []
keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"}
for idx, msg in enumerate(messages):
role = msg.get("role")
if role in keep_roles or idx >= last_user_idx:
result.append(msg)
elif role == "assistant":
msg = copy.copy(msg)
msg.pop("reasoning_content", None)
result.append(msg)
# developer and other roles before last_user_idx are dropped
return result
# ============================================================
# Parsing (Decoding model output)
# ============================================================
def _read_until_stop(
index: int, text: str, stop: List[str]
) -> Tuple[int, str, Optional[str]]:
"""
Read text from index until one of the stop strings is found.
Returns:
Tuple of (new_index, content_before_stop, matched_stop_string_or_None).
"""
min_pos = len(text)
matched_stop = None
for s in stop:
pos = text.find(s, index)
if pos != -1 and pos < min_pos:
min_pos = pos
matched_stop = s
if matched_stop:
content = text[index:min_pos]
return min_pos + len(matched_stop), content, matched_stop
else:
content = text[index:]
return len(text), content, None
def parse_tool_calls(
index: int, text: str
) -> Tuple[int, Optional[str], List[Dict[str, str]]]:
"""
Parse DSML tool calls from text starting at the given index.
Args:
index: Starting position in text.
text: The full text to parse.
Returns:
Tuple of (new_index, last_stop_token, list_of_tool_call_dicts).
Each tool call dict has "name" and "arguments" keys.
"""
tool_calls: List[Dict[str, Any]] = []
stop_token = None
tool_calls_end_token = f"</{dsml_token}{tool_calls_block_name}>"
while index < len(text):
index, _, stop_token = _read_until_stop(
index, text, [f"<{dsml_token}invoke", tool_calls_end_token]
)
if _ != ">\n":
raise ValueError(f"Tool call format error: expected '>\\n' but got '{_}'")
if stop_token == tool_calls_end_token:
break
if stop_token is None:
raise ValueError("Missing special token in tool calls")
index, tool_name_content, stop_token = _read_until_stop(
index, text, [f"<{dsml_token}parameter", f"</{dsml_token}invoke"]
)
p_tool_name = re.findall(
r'^\s*name="(.*?)">\n$', tool_name_content, flags=re.DOTALL
)
if len(p_tool_name) != 1:
raise ValueError(f"Tool name format error: '{tool_name_content}'")
tool_name = p_tool_name[0]
tool_args: Dict[str, Tuple[str, str]] = {}
while stop_token == f"<{dsml_token}parameter":
index, param_content, stop_token = _read_until_stop(
index, text, [f"/{dsml_token}parameter"]
)
param_kv = re.findall(
r'^ name="(.*?)" string="(true|false)">(.*?)<$',
param_content,
flags=re.DOTALL,
)
if len(param_kv) != 1:
raise ValueError(f"Parameter format error: '{param_content}'")
param_name, string, param_value = param_kv[0]
if param_name in tool_args:
raise ValueError(f"Duplicate parameter name: '{param_name}'")
tool_args[param_name] = (param_value, string)
index, content, stop_token = _read_until_stop(
index, text, [f"<{dsml_token}parameter", f"</{dsml_token}invoke"]
)
if content != ">\n":
raise ValueError(
f"Parameter format error: expected '>\\n' but got '{content}'"
)
tool_call = decode_dsml_to_arguments(tool_name=tool_name, tool_args=tool_args)
tool_calls.append(tool_call)
return index, stop_token, tool_calls
def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[str, Any]:
"""
Parse a model completion text into a structured assistant message.
This function takes the raw text output from the model (a single assistant turn)
and extracts:
- reasoning_content (thinking block)
- content (summary/response)
- tool_calls (if any)
NOTE: This function is designed to parse only correctly formatted strings and
will raise ValueError for malformed output.
Args:
text: The raw completion text (including EOS token).
thinking_mode: Either "chat" or "thinking".
Returns:
Dict with keys: "role", "content", "reasoning_content", "tool_calls".
tool_calls are in OpenAI format.
"""
summary_content, reasoning_content, tool_calls = "", "", []
index, stop_token = 0, None
tool_calls_start_token = f"\n\n<{dsml_token}{tool_calls_block_name}"
is_thinking = thinking_mode == "thinking"
is_tool_calling = False
if is_thinking:
index, content_delta, stop_token = _read_until_stop(
index, text, [thinking_end_token, tool_calls_start_token]
)
reasoning_content = content_delta
assert (
stop_token == thinking_end_token
), "Invalid thinking format: missing </think>"
index, content_delta, stop_token = _read_until_stop(
index, text, [eos_token, tool_calls_start_token]
)
summary_content = content_delta
if stop_token == tool_calls_start_token:
is_tool_calling = True
else:
assert stop_token == eos_token, "Invalid format: missing EOS token"
if is_tool_calling:
index, stop_token, tool_calls = parse_tool_calls(index, text)
index, tool_ends_text, stop_token = _read_until_stop(index, text, [eos_token])
assert not tool_ends_text, "Unexpected content after tool calls"
assert len(text) == index and stop_token in [
eos_token,
None,
], "Unexpected content at end"
for sp_token in [
bos_token,
eos_token,
thinking_start_token,
thinking_end_token,
dsml_token,
]:
assert (
sp_token not in summary_content and sp_token not in reasoning_content
), f"Unexpected special token '{sp_token}' in content"
return {
"role": "assistant",
"content": summary_content,
"reasoning_content": reasoning_content,
"tool_calls": tool_calls_to_openai_format(tool_calls),
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,5 @@
"""Realtime transcription WebSocket package. Exposes the FastAPI WS entry."""
from sglang.srt.entrypoints.openai.realtime.handler import (
handle_realtime_transcription as handle_realtime_transcription,
)
@@ -0,0 +1,119 @@
"""WebSocket entry for realtime transcription.
Handles accept, concurrency, cleanup. Event loop is in session.py.
"""
from __future__ import annotations
import asyncio
import logging
from fastapi import WebSocket, WebSocketDisconnect
from openai.types.realtime import RealtimeErrorEvent
from openai.types.realtime.realtime_error import RealtimeError
from sglang.srt.entrypoints.openai.realtime.session import RealtimeConnection
from sglang.srt.entrypoints.openai.transcription_adapters.base import (
TranscriptionAdapter,
)
from sglang.srt.managers.tokenizer_manager import TokenizerManager
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import random_uuid
logger = logging.getLogger(__name__)
async def _safe_send(websocket: WebSocket, text: str) -> None:
try:
await websocket.send_text(text)
except (WebSocketDisconnect, RuntimeError) as e:
logger.debug("[realtime] send failed (peer gone): %s", e)
async def _safe_close(websocket: WebSocket) -> None:
try:
await websocket.close()
except (WebSocketDisconnect, RuntimeError) as e:
logger.debug("[realtime] close failed (already closed): %s", e)
async def _reject_before_session(
websocket: WebSocket,
code: str,
message: str,
*,
error_type: str = "invalid_request_error",
) -> None:
"""Reject path that runs before acquiring the session semaphore, so
unsupported / over-capacity peers don't hold a session slot."""
try:
await websocket.accept()
except (WebSocketDisconnect, RuntimeError) as e:
logger.debug("[realtime] reject: accept failed: %s", e)
return
logger.info("[realtime] rejected (%s)", code)
envelope = RealtimeErrorEvent(
event_id=f"event_{random_uuid()}",
type="error",
error=RealtimeError(type=error_type, code=code, message=message),
)
await _safe_send(websocket, envelope.model_dump_json())
await _safe_close(websocket)
async def handle_realtime_transcription(
websocket: WebSocket,
tokenizer_manager: TokenizerManager,
adapter: TranscriptionAdapter,
server_args: ServerArgs,
session_semaphore: asyncio.Semaphore,
) -> None:
"""WS endpoint for /v1/realtime. Pre-session validation runs before
the semaphore so rejects don't consume a session slot; the
``async with`` then guarantees the slot is released even if
RealtimeConnection raises."""
if not adapter.supports_chunked_streaming:
await _reject_before_session(
websocket,
"not_supported",
"Model does not support streaming ASR",
)
return
if session_semaphore.locked():
await _reject_before_session(
websocket,
"too_many_sessions",
f"Maximum concurrent sessions reached "
f"({server_args.asr_max_concurrent_sessions}).",
error_type="rate_limit_exceeded",
)
return
async with session_semaphore:
try:
try:
await websocket.accept()
except (WebSocketDisconnect, RuntimeError) as e:
logger.debug("[realtime] accept failed: %s", e)
return
connection = RealtimeConnection(
websocket, tokenizer_manager, adapter, server_args
)
await connection.run()
except WebSocketDisconnect:
logger.info("[realtime] client disconnected (normal)")
except Exception:
logger.exception("[realtime] unexpected error in session")
envelope = RealtimeErrorEvent(
event_id=f"event_{random_uuid()}",
type="error",
error=RealtimeError(
type="server_error",
code="inference_failed",
message="Internal server error",
),
)
await _safe_send(websocket, envelope.model_dump_json())
finally:
await _safe_close(websocket)
@@ -0,0 +1,78 @@
"""Wire schema for Realtime WS transcription sessions."""
from __future__ import annotations
from typing import Literal, Optional, Union
from openai.types.realtime import SessionUpdateEvent as _SessionUpdateEvent
from openai.types.realtime.audio_transcription import (
AudioTranscription as _AudioTranscription,
)
from openai.types.realtime.realtime_audio_formats import AudioPCM as _AudioPCM
from openai.types.realtime.realtime_audio_formats import AudioPCMA as _AudioPCMA
from openai.types.realtime.realtime_audio_formats import AudioPCMU as _AudioPCMU
from openai.types.realtime.realtime_transcription_session_audio import (
RealtimeTranscriptionSessionAudio as _AudioCfg,
)
from openai.types.realtime.realtime_transcription_session_audio_input import (
RealtimeTranscriptionSessionAudioInput as _AudioInputCfg,
)
from openai.types.realtime.realtime_transcription_session_create_request import (
RealtimeTranscriptionSessionCreateRequest as _SessionCfg,
)
from pydantic import Field
from typing_extensions import Annotated
# Fallback rate when the client omits `audio.input.format.rate`. SDK pins
# `AudioPCM.rate` to Literal[24000], so this matches the only value the SDK
# accepts when the field is present.
DEFAULT_INPUT_SAMPLE_RATE = 24000
# Wire rates we accept on `audio.input.format.rate` and resample to
# `adapter.model_sample_rate` server-side. 24000 matches the SDK pin;
# 16000 and 48000 widen it to cover common ASR-client and consumer-audio
# rates. Add a value here only after verifying transcription quality.
SUPPORTED_INPUT_SAMPLE_RATES = (16000, 24000, 48000)
class AudioPCM(_AudioPCM):
type: Literal["audio/pcm"] = "audio/pcm"
rate: Optional[int] = None
class AudioPCMU(_AudioPCMU):
type: Literal["audio/pcmu"] = "audio/pcmu"
class AudioPCMA(_AudioPCMA):
type: Literal["audio/pcma"] = "audio/pcma"
AudioInputFormat = Annotated[
Union[AudioPCM, AudioPCMU, AudioPCMA],
Field(discriminator="type"),
]
class AudioTranscription(_AudioTranscription):
# SDK pins model to Literal["whisper-1", "gpt-4o-*-transcribe", ...];
# sglang serves arbitrary ASR models (Qwen3-ASR, etc.) and treats the
# client-supplied name as echo-only.
model: Optional[str] = None
class TranscriptionSessionAudioInput(_AudioInputCfg):
format: Optional[AudioInputFormat] = None
transcription: Optional[AudioTranscription] = None
class TranscriptionSessionAudio(_AudioCfg):
input: Optional[TranscriptionSessionAudioInput] = None
class TranscriptionSessionConfig(_SessionCfg):
audio: Optional[TranscriptionSessionAudio] = None
class SessionUpdateEvent(_SessionUpdateEvent):
session: TranscriptionSessionConfig
@@ -0,0 +1,740 @@
"""WebSocket session for realtime ASR.
Pre-commit deltas reference the reserved current_item_id that the
subsequent input_audio_buffer.committed and conversation.item.created
events will announce — sglang-specific, deviates from OpenAI's
commit-only delta emission.
"""
from __future__ import annotations
import asyncio
import io
import json
import logging
import math
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
import numpy as np
import pybase64
import soundfile as sf
from fastapi import WebSocket, WebSocketDisconnect
from openai.types.realtime import (
ConversationItemCreatedEvent,
InputAudioBufferAppendEvent,
InputAudioBufferClearedEvent,
InputAudioBufferClearEvent,
InputAudioBufferCommitEvent,
InputAudioBufferCommittedEvent,
RealtimeErrorEvent,
SessionCreatedEvent,
SessionUpdatedEvent,
)
from openai.types.realtime.conversation_item_input_audio_transcription_completed_event import (
ConversationItemInputAudioTranscriptionCompletedEvent,
UsageTranscriptTextUsageDuration,
)
from openai.types.realtime.conversation_item_input_audio_transcription_delta_event import (
ConversationItemInputAudioTranscriptionDeltaEvent,
)
from openai.types.realtime.conversation_item_input_audio_transcription_failed_event import (
ConversationItemInputAudioTranscriptionFailedEvent,
)
from openai.types.realtime.conversation_item_input_audio_transcription_failed_event import (
Error as TranscriptionFailedError,
)
from openai.types.realtime.realtime_conversation_item_user_message import (
Content as InputAudioContent,
)
from openai.types.realtime.realtime_conversation_item_user_message import (
RealtimeConversationItemUserMessage,
)
from openai.types.realtime.realtime_error import RealtimeError
from pydantic import BaseModel, ValidationError
from sglang.srt.entrypoints.openai.protocol import TranscriptionRequest
from sglang.srt.entrypoints.openai.realtime.protocol import (
DEFAULT_INPUT_SAMPLE_RATE,
SUPPORTED_INPUT_SAMPLE_RATES,
AudioPCM,
SessionUpdateEvent,
TranscriptionSessionAudioInput,
TranscriptionSessionConfig,
)
from sglang.srt.entrypoints.openai.streaming_asr import (
StreamingASRState,
needs_space,
normalize_whitespace,
process_asr_chunk,
)
from sglang.srt.entrypoints.openai.transcription_adapters.base import (
TranscriptionAdapter,
)
from sglang.srt.managers.tokenizer_manager import TokenizerManager
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import random_uuid
logger = logging.getLogger(__name__)
# PCM16: 16-bit samples → 2 bytes each. Used for frame-length validation
# and bytes/sec arithmetic against `np.frombuffer(..., dtype=np.int16)` below.
_SAMPLE_WIDTH = 2
def _resample_to_target_rate(pcm: bytes, src_rate: int, target_rate: int) -> bytes:
if src_rate == target_rate or not pcm:
return pcm
import torch
import torchaudio
samples = np.frombuffer(pcm, dtype=np.int16).astype(np.float32) / 32768.0
audio = torch.from_numpy(samples).unsqueeze(0)
audio = torchaudio.functional.resample(
audio, orig_freq=src_rate, new_freq=target_rate
)
samples = audio.squeeze(0).numpy()
# Clip to int16 range via 2^15 - 1 so a clipped 1.0 stays representable.
return (np.clip(samples, -1.0, 1.0) * 32767.0).astype(np.int16).tobytes()
def _pcm_to_wav(pcm: bytes, sample_rate: int) -> bytes:
samples = np.frombuffer(pcm, dtype=np.int16)
buf = io.BytesIO()
sf.write(buf, samples, sample_rate, format="WAV")
return buf.getvalue()
_CLIENT_EVENT_TYPES: Dict[str, type] = {
"session.update": SessionUpdateEvent,
"input_audio_buffer.append": InputAudioBufferAppendEvent,
"input_audio_buffer.commit": InputAudioBufferCommitEvent,
"input_audio_buffer.clear": InputAudioBufferClearEvent,
}
def _parse_client_event(raw: Dict[str, Any]) -> Optional[BaseModel]:
"""Parse, returning None if type is unknown. Raises ValidationError on
a malformed payload of a known type."""
cls = _CLIENT_EVENT_TYPES.get(raw.get("type"))
if cls is None:
return None
return cls.model_validate(raw)
@dataclass
class _SessionConfig:
"""Session-level configuration negotiated via session.update: audio
input format, requested language, sampling params. Persists until
the session ends; ``configured`` gates audio-frame handling so the
server doesn't run inference on PCM sent before session.update."""
input_sample_rate: int = DEFAULT_INPUT_SAMPLE_RATE
language: Optional[str] = None
client_model: Optional[str] = None
sampling_params: Optional[Dict[str, Any]] = None
configured: bool = False
@dataclass
class _AudioState:
"""Per-item audio state: PCM buffer accumulated from
input_audio_buffer.append, the chunked ASR rollback state, and the
static buffer-size limits set at __init__. pcm_buffer / state /
last_inference_offset reset on commit-roll and clear; the size limits
stay constant for the session's lifetime."""
max_buffer_bytes: int
chunk_size_bytes: int
state: StreamingASRState
pcm_buffer: bytearray = field(default_factory=bytearray)
last_inference_offset: int = 0
@dataclass
class _ItemState:
"""Per-item conversation-item ids and the wire-formatted deltas
emitted so far for the current item. current_item_id is reserved at
__init__ and only announced to the client by
input_audio_buffer.committed."""
current_item_id: str
previous_item_id: Optional[str] = None
emitted_deltas: List[str] = field(default_factory=list)
class RealtimeConnection:
"""One realtime transcription session. Drives the WS receive loop,
dispatches typed client events to the matching _on_* handler, and
triggers chunked ASR inference at audio buffer thresholds."""
def __init__(
self,
websocket: WebSocket,
tokenizer_manager: TokenizerManager,
adapter: TranscriptionAdapter,
server_args: ServerArgs,
) -> None:
self.websocket = websocket
self.tokenizer_manager = tokenizer_manager
self.adapter = adapter
self.server_args = server_args
self.session_id = f"sess_{random_uuid()}"
self._current_client_event_id: Optional[str] = None
self.model_sample_rate = adapter.model_sample_rate
self.bytes_per_second = self.model_sample_rate * _SAMPLE_WIDTH
self.max_buffer_seconds = server_args.asr_max_buffer_seconds
self.config = _SessionConfig()
state = StreamingASRState(**adapter.chunked_streaming_config)
chunk_size_bytes = int(state.chunk_size_sec * self.bytes_per_second)
if chunk_size_bytes <= 0:
raise RuntimeError(
f"adapter.chunked_streaming_config produced non-positive "
f"chunk_size_sec; got {state.chunk_size_sec!r}"
)
self.audio = _AudioState(
max_buffer_bytes=self.max_buffer_seconds * self.bytes_per_second,
chunk_size_bytes=chunk_size_bytes,
state=state,
)
self.item = _ItemState(current_item_id=f"item_{random_uuid()}")
async def run(self) -> None:
await self._send(
SessionCreatedEvent(
event_id=f"event_{random_uuid()}",
type="session.created",
session=self._build_session_info(),
)
)
try:
await self._run_loop()
except WebSocketDisconnect:
logger.info("[realtime] client disconnected: %s", self.session_id)
except Exception:
logger.exception("[realtime] unexpected error: %s", self.session_id)
try:
await self._send_error(
"inference_failed",
"Internal server error",
error_type="server_error",
)
except (WebSocketDisconnect, RuntimeError) as e:
logger.debug(
"[realtime] failed to notify client of unexpected error: %s",
e,
)
async def _run_loop(self) -> None:
"""Receive-and-dispatch loop. Validation errors emit an error event
and continue; fatal append-path errors (buffer overflow, append-time
inference failure) close the WebSocket and terminate the loop.
"""
while True:
self._current_client_event_id = None
message = await self.websocket.receive()
if message["type"] == "websocket.disconnect":
return
text = message.get("text")
if not text:
if message.get("bytes") is not None:
# OpenAI Realtime is base64 PCM in JSON; binary frames aren't supported.
await self._send_error(
"invalid_payload",
"Binary frames are not supported on /v1/realtime; "
"use input_audio_buffer.append with base64 audio.",
)
continue
try:
raw = json.loads(text)
except json.JSONDecodeError:
await self._send_error("invalid_payload", "Invalid JSON")
continue
if not isinstance(raw, dict):
await self._send_error(
"invalid_payload", "Top-level event must be a JSON object"
)
continue
self._current_client_event_id = raw.get("event_id")
try:
event = _parse_client_event(raw)
except ValidationError as e:
# Report first error only; matches OpenAI server behavior.
err = e.errors()[0]
loc = ".".join(str(x) for x in err["loc"])
await self._send_error(
"invalid_value",
err.get("msg") or "Invalid payload",
param=loc or None,
)
continue
if event is None:
await self._send_error(
"unknown_event",
f"Unknown event type: {raw.get('type')!r}",
)
continue
terminate = await self._dispatch(event)
if terminate:
return
async def _dispatch(self, event: BaseModel) -> bool:
"""Returns True if the session should terminate."""
if isinstance(event, InputAudioBufferAppendEvent):
return await self._on_input_audio_buffer_append(event)
if isinstance(event, SessionUpdateEvent):
await self._on_session_update(event)
elif isinstance(event, InputAudioBufferCommitEvent):
await self._on_input_audio_buffer_commit(event)
elif isinstance(event, InputAudioBufferClearEvent):
await self._on_input_audio_buffer_clear(event)
return False
async def _on_session_update(self, event: SessionUpdateEvent) -> None:
cfg = event.session
# Normalize audio to an empty input cfg if absent so downstream
# `audio.X is not None` reads as a business rule, not an existence check.
# transcription stays nullable so partial-update can detect whether
# the client sent the block.
audio = (
cfg.audio.input if cfg.audio else None
) or TranscriptionSessionAudioInput()
transcription = audio.transcription
# Validate first, then mutate config only after the whole update is accepted.
if audio.turn_detection is not None:
await self._send_error(
"not_supported",
"Server-side VAD is not implemented; "
"set audio.input.turn_detection: null and commit explicitly.",
param="session.audio.input.turn_detection",
)
return
if audio.noise_reduction is not None:
await self._send_error(
"not_supported",
"audio.input.noise_reduction is not supported; set to null.",
param="session.audio.input.noise_reduction",
)
return
if transcription is not None and transcription.prompt is not None:
await self._send_error(
"not_supported",
"audio.input.transcription.prompt is not supported.",
param="session.audio.input.transcription.prompt",
)
return
if (
transcription is not None
and transcription.model
and transcription.model != self.server_args.served_model_name
):
await self._send_error(
"not_supported",
f"Model {transcription.model!r} is not served by this endpoint "
f"(serving {self.server_args.served_model_name!r}); set "
f"transcription.model to null or to the server's model name.",
param="session.audio.input.transcription.model",
)
return
new_rate = self.config.input_sample_rate # default: keep current
fmt = audio.format
if fmt is not None:
if not isinstance(fmt, AudioPCM):
# G.711 (pcmu / pcma): not implemented.
await self._send_error(
"not_supported",
f"audio.input.format.type must be 'audio/pcm'; "
f"{fmt.type!r} is not implemented",
param="session.audio.input.format.type",
)
return
if fmt.rate is not None and fmt.rate not in SUPPORTED_INPUT_SAMPLE_RATES:
await self._send_error(
"invalid_value",
f"audio.input.format.rate must be one of "
f"{SUPPORTED_INPUT_SAMPLE_RATES}, got {fmt.rate}",
param="session.audio.input.format.rate",
)
return
new_rate = fmt.rate or DEFAULT_INPUT_SAMPLE_RATE
# Changing the rate mid-item would leave already-buffered PCM
# at the old rate mixed with new audio at the new rate, so
# require the client to commit or clear before switching.
if new_rate != self.config.input_sample_rate and self.audio.pcm_buffer:
await self._send_error(
"invalid_state",
"Cannot change audio.input.format.rate while audio is "
"buffered; commit or clear the current item first.",
param="session.audio.input.format.rate",
)
return
# Mutation pass — no early returns past this point.
self.config.input_sample_rate = new_rate
if transcription is not None:
self.config.client_model = transcription.model
self.config.language = transcription.language
self.config.sampling_params = self.adapter.build_sampling_params(
TranscriptionRequest(language=self.config.language)
)
self.config.configured = True
# Side effects: log + ack.
if cfg.include:
logger.info(
"[realtime] %s: include[] received but not implemented; ignoring: %s",
self.session_id,
cfg.include,
)
if self.config.input_sample_rate != self.model_sample_rate:
logger.info(
"[realtime] %s configured: resample %d%d (ratio %.2f), language=%s",
self.session_id,
self.config.input_sample_rate,
self.model_sample_rate,
self.config.input_sample_rate / self.model_sample_rate,
self.config.language,
)
await self._send(
SessionUpdatedEvent(
event_id=f"event_{random_uuid()}",
type="session.updated",
session=self._build_session_info(),
)
)
async def _on_input_audio_buffer_append(
self, event: InputAudioBufferAppendEvent
) -> bool:
"""Returns True if the session should terminate (buffer overflow or
append-time inference failure)."""
if not self.config.configured:
await self._send_error(
"invalid_state", "Send session.update before audio frames"
)
return False
# Empty audio is a no-op (heartbeat frames); skip b64decode.
if not event.audio:
return False
try:
data = pybase64.b64decode(event.audio, validate=True)
except (ValueError, TypeError):
await self._send_error(
"invalid_audio", "audio field is not valid base64", param="audio"
)
return False
if len(data) % _SAMPLE_WIDTH != 0:
await self._send_error(
"invalid_audio_format",
f"PCM16 frame length must be a multiple of {_SAMPLE_WIDTH} bytes",
)
return False
# Estimate post-resample size before resampling so oversized frames fail early.
src_samples = len(data) // _SAMPLE_WIDTH
target_samples = math.ceil(
src_samples * self.model_sample_rate / self.config.input_sample_rate
)
if (
len(self.audio.pcm_buffer) + target_samples * _SAMPLE_WIDTH
> self.audio.max_buffer_bytes
):
# Close 1009 ("message too big") so clients can distinguish
# session-resource exhaustion from a normal close.
await self._send_error_and_close(
"buffer_overflow",
f"Accumulated audio exceeded {self.max_buffer_seconds}s; "
f"client is sending faster than inference can keep up",
close_code=1009,
)
return True
if self.config.input_sample_rate != self.model_sample_rate:
data = await asyncio.to_thread(
_resample_to_target_rate,
data,
self.config.input_sample_rate,
self.model_sample_rate,
)
self.audio.pcm_buffer.extend(data)
new_audio_bytes = len(self.audio.pcm_buffer) - self.audio.last_inference_offset
if new_audio_bytes >= self.audio.chunk_size_bytes:
ok = await self._run_inference(is_last=False)
if not ok:
# WS already closed inside _run_inference.
return True
return False
async def _on_input_audio_buffer_commit(
self, event: InputAudioBufferCommitEvent
) -> None:
if not self.config.configured:
await self._send_error("invalid_state", "Send session.update before commit")
return
if not self.audio.pcm_buffer and not self.audio.state.full_transcript:
await self._send_error(
"invalid_state", "Cannot commit an empty audio buffer"
)
return
has_new_audio = len(self.audio.pcm_buffer) > self.audio.last_inference_offset
item_id = self.item.current_item_id
prev_item_id = self.item.previous_item_id
partial_transcript = normalize_whitespace("".join(self.item.emitted_deltas))
await self._send(
InputAudioBufferCommittedEvent(
event_id=f"event_{random_uuid()}",
type="input_audio_buffer.committed",
item_id=item_id,
previous_item_id=prev_item_id,
)
)
await self._send(
ConversationItemCreatedEvent(
event_id=f"event_{random_uuid()}",
type="conversation.item.created",
previous_item_id=prev_item_id,
item=RealtimeConversationItemUserMessage(
id=item_id,
type="message",
role="user",
status="completed",
content=[
InputAudioContent(
type="input_audio", transcript=partial_transcript
)
],
),
)
)
# Capture pcm duration before `_start_next_item()` runs: starting
# the next item clears pcm_buffer, so reading it after gives 0.
pcm_duration_seconds = len(self.audio.pcm_buffer) / self.bytes_per_second
if has_new_audio:
ok = await self._run_inference(is_last=True)
if not ok:
# _run_inference already emitted transcription.failed and
# rolled the item; don't also emit completed.
return
elif self.audio.state.full_transcript:
# Audio length was exactly a chunk_size_bytes multiple. Flush
# the tail tokens update() held back.
tail = self.audio.state.finalize()
await self._emit_transcription_delta(tail)
# Build from emitted_deltas, not state.full_transcript: prefix injection
# means the last chunk's full_transcript is only the continuation tail.
transcript = normalize_whitespace("".join(self.item.emitted_deltas))
await self._send(
ConversationItemInputAudioTranscriptionCompletedEvent(
event_id=f"event_{random_uuid()}",
type="conversation.item.input_audio_transcription.completed",
item_id=item_id,
content_index=0,
transcript=transcript,
usage=UsageTranscriptTextUsageDuration(
type="duration", seconds=pcm_duration_seconds
),
)
)
self._start_next_item()
async def _on_input_audio_buffer_clear(
self, event: InputAudioBufferClearEvent
) -> None:
# Reserve a fresh current_item_id so post-clear pre-commit deltas
# don't share an item_id with deltas the client already received
# for the abandoned audio. previous_item_id is NOT touched — the
# cleared item was never committed, so the prior-commit chain
# shouldn't include it.
self._reset_inference_state()
self.item.current_item_id = f"item_{random_uuid()}"
await self._send(
InputAudioBufferClearedEvent(
event_id=f"event_{random_uuid()}", type="input_audio_buffer.cleared"
)
)
async def _run_inference(self, is_last: bool) -> bool:
"""Run ASR on the current cumulative buffer. Returns False on failure:
commit-time emits transcription.failed and rolls the item; append-time
emits a generic error envelope and closes the WebSocket."""
wav_data = await asyncio.to_thread(
_pcm_to_wav, bytes(self.audio.pcm_buffer), self.model_sample_rate
)
try:
delta = await process_asr_chunk(
tokenizer_manager=self.tokenizer_manager,
adapter=self.adapter,
state=self.audio.state,
audio_data=wav_data,
sampling_params=self.config.sampling_params,
is_last=is_last,
)
except Exception:
logger.exception(
"[realtime] inference failed: session=%s item=%s buffer_bytes=%d",
self.session_id,
self.item.current_item_id,
len(self.audio.pcm_buffer),
)
if is_last:
# Commit-time failure: committed + created already emitted,
# so the item exists client-side and transcription.failed
# can reference it. Wire message is hardcoded "Transcription
# failed" — don't leak backend traces to the client; full
# error is in the logger.exception above.
await self._send(
ConversationItemInputAudioTranscriptionFailedEvent(
event_id=f"event_{random_uuid()}",
type="conversation.item.input_audio_transcription.failed",
item_id=self.item.current_item_id,
content_index=0,
error=TranscriptionFailedError(
type="server_error",
code="inference_failed",
message="Transcription failed",
),
)
)
self._start_next_item()
else:
# Append-time failure: the item isn't visible client-side
# yet (committed/created fire at commit), so
# transcription.failed would reference a ghost id.
await self._send_error_and_close(
"inference_failed",
"Transcription failed",
close_code=1011,
)
return False
self.audio.last_inference_offset = len(self.audio.pcm_buffer)
await self._emit_transcription_delta(delta)
return True
async def _emit_transcription_delta(self, delta: str) -> None:
"""emitted_deltas stores wire-formatted text (with leading
boundary spaces baked in), so "".join(...) reconstructs the
cumulative transcript verbatim."""
if not delta:
return
for word in delta.split(" "):
if not word:
continue
prev = self.item.emitted_deltas[-1] if self.item.emitted_deltas else ""
formatted = f" {word}" if needs_space(prev, word) else word
self.item.emitted_deltas.append(formatted)
await self._send(
ConversationItemInputAudioTranscriptionDeltaEvent(
event_id=f"event_{random_uuid()}",
type="conversation.item.input_audio_transcription.delta",
item_id=self.item.current_item_id,
content_index=0,
delta=formatted,
)
)
def _start_next_item(self) -> None:
self.item.previous_item_id = self.item.current_item_id
self.item.current_item_id = f"item_{random_uuid()}"
self._reset_inference_state()
def _reset_inference_state(self) -> None:
"""Missing any of these resets leaks state across items."""
self.audio.state = StreamingASRState(**self.adapter.chunked_streaming_config)
self.audio.pcm_buffer.clear() # in-place; reuses the buffer's allocation
self.item.emitted_deltas.clear()
self.audio.last_inference_offset = 0
def _build_session_info(self) -> TranscriptionSessionConfig:
# id / object aren't SDK fields; round-trip via extra='allow' so
# dumps emit them like the real server.
return TranscriptionSessionConfig.model_validate(
{
"type": "transcription",
"id": self.session_id,
"object": "realtime.transcription_session",
"audio": {
"input": {
"format": {
"type": "audio/pcm",
"rate": self.config.input_sample_rate,
},
"transcription": {
"model": self.config.client_model,
"language": self.config.language,
},
"noise_reduction": None,
"turn_detection": None,
}
},
}
)
async def _send(self, event: BaseModel) -> None:
await self.websocket.send_text(event.model_dump_json())
async def _send_error(
self,
code: str,
message: str,
*,
error_type: str = "invalid_request_error",
param: Optional[str] = None,
) -> None:
envelope = RealtimeErrorEvent(
event_id=f"event_{random_uuid()}",
type="error",
error=RealtimeError(
type=error_type,
code=code,
message=message,
param=param,
event_id=self._current_client_event_id,
),
)
await self.websocket.send_text(envelope.model_dump_json())
async def _send_error_and_close(
self,
code: str,
message: str,
*,
close_code: int,
error_type: str = "server_error",
) -> None:
# Independent try-blocks: a failed send must not skip the close.
# We still need to release local starlette socket state even when
# the wire send doesn't reach the peer.
try:
await self._send_error(code, message, error_type=error_type)
except (WebSocketDisconnect, RuntimeError) as e:
logger.debug("[realtime] send error %s before close failed: %s", code, e)
try:
await self.websocket.close(code=close_code)
except (WebSocketDisconnect, RuntimeError) as e:
logger.debug("[realtime] close %d after %s failed: %s", close_code, code, e)
@@ -0,0 +1,306 @@
from __future__ import annotations
import json
import logging
import uuid
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
import orjson
from fastapi import HTTPException, Request
from fastapi.responses import ORJSONResponse, StreamingResponse
from sglang.srt.entrypoints.openai.encoding_dsv32 import DS32EncodingError
from sglang.srt.entrypoints.openai.protocol import ErrorResponse, OpenAIServingRequest
from sglang.srt.managers.io_struct import EmbeddingReqInput, GenerateReqInput
from sglang.srt.observability.req_time_stats import monotonic_time
from sglang.srt.server_args import ServerArgs
if TYPE_CHECKING:
from sglang.srt.managers.tokenizer_manager import TokenizerManager
logger = logging.getLogger(__name__)
# Base class for specific endpoint handlers
class OpenAIServingBase(ABC):
"""Abstract base class for OpenAI endpoint handlers"""
def __init__(self, tokenizer_manager: TokenizerManager):
self.tokenizer_manager = tokenizer_manager
self.allowed_custom_labels = (
set(
self.tokenizer_manager.server_args.tokenizer_metrics_allowed_custom_labels
)
if isinstance(self.tokenizer_manager.server_args, ServerArgs)
and self.tokenizer_manager.server_args.tokenizer_metrics_allowed_custom_labels
else None
)
def _parse_model_parameter(self, model: str) -> Tuple[str, Optional[str]]:
"""Parse 'base-model:adapter-name' syntax to extract LoRA adapter.
Returns (base_model, adapter_name) or (model, None) if no colon present.
"""
if ":" not in model:
return model, None
# Split on first colon only to handle model paths with multiple colons
parts = model.split(":", 1)
base_model = parts[0].strip()
adapter_name = parts[1].strip() or None
return base_model, adapter_name
def _resolve_lora_path(
self,
request_model: str,
explicit_lora_path: Optional[Union[str, List[Optional[str]]]],
) -> Optional[Union[str, List[Optional[str]]]]:
"""Resolve LoRA adapter with priority: model parameter > explicit lora_path.
Returns adapter name or None. Supports both single values and lists (batches).
"""
_, adapter_from_model = self._parse_model_parameter(request_model)
# Model parameter adapter takes precedence
if adapter_from_model is not None:
return adapter_from_model
# Fall back to explicit lora_path
return explicit_lora_path
async def handle_request(
self, request: OpenAIServingRequest, raw_request: Request
) -> Union[Any, StreamingResponse, ErrorResponse]:
"""Handle the specific request type with common pattern
If you want to override this method, you should be careful to record the validation time.
"""
received_time = monotonic_time()
try:
# Validate request
error_msg = self._validate_request(request)
if error_msg:
return self.create_error_response(error_msg)
# Log the raw OpenAI request payload before conversion to tokenized form.
request_logger = self.tokenizer_manager.request_logger
if request_logger.log_requests and request_logger.log_requests_level >= 2:
request_logger.log_openai_received_request(request, request=raw_request)
# Convert to internal format
adapted_request, processed_request = self._convert_to_internal_request(
request, raw_request
)
if isinstance(adapted_request, (GenerateReqInput, EmbeddingReqInput)):
# Only set timing fields if adapted_request supports them
adapted_request.received_time = received_time
# Note(Xinyuan): raw_request below is only used for detecting the connection of the client
if hasattr(request, "stream") and request.stream:
return await self._handle_streaming_request(
adapted_request, processed_request, raw_request
)
else:
return await self._handle_non_streaming_request(
adapted_request, processed_request, raw_request
)
except HTTPException as e:
return self.create_error_response(
message=e.detail, err_type=str(e.status_code), status_code=e.status_code
)
except ValueError as e:
return self.create_error_response(
message=str(e),
err_type="BadRequest",
status_code=400,
)
except DS32EncodingError as e:
logger.info(f"DS32EncodingError: {e}")
return self.create_error_response(
message=str(e),
err_type="BadRequest",
status_code=400,
)
except Exception as e:
logger.exception(f"Error in request: {e}")
return self.create_error_response(
message=f"Internal server error: {str(e)}",
err_type="InternalServerError",
status_code=500,
)
@abstractmethod
def _request_id_prefix(self) -> str:
"""Generate request ID based on request type"""
pass
def _generate_request_id_base(self, request: OpenAIServingRequest) -> Optional[str]:
"""Generate request ID based on request type"""
return None
# TODO(chang): the rid is used in io_strcut check and often violates `The rid should be a list` AssertionError
# Temporarily return None in this function until the rid logic is clear.
if rid := getattr(request, "rid", None):
return rid
return f"{self._request_id_prefix()}{uuid.uuid4().hex}"
def _compute_extra_key(self, request: OpenAIServingRequest) -> Optional[str]:
"""Compute the final extra_key by concatenating cache_salt and extra_key if both are provided."""
parts = []
for key in ["cache_salt", "extra_key"]:
value = getattr(request, key, None)
if value:
if not isinstance(value, str):
raise TypeError(
f"Value of {key} must be a string, but got {type(value).__name__}"
)
parts.append(value)
return "".join(parts) if parts else None
@abstractmethod
def _convert_to_internal_request(
self,
request: OpenAIServingRequest,
raw_request: Request = None,
) -> tuple[GenerateReqInput, OpenAIServingRequest]:
"""Convert OpenAI request to internal format"""
pass
async def _handle_streaming_request(
self,
adapted_request: GenerateReqInput,
request: OpenAIServingRequest,
raw_request: Request,
) -> Union[StreamingResponse, ErrorResponse, ORJSONResponse]:
"""Handle streaming request
Override this method in child classes that support streaming requests.
"""
return self.create_error_response(
message=f"{self.__class__.__name__} does not support streaming requests",
err_type="NotImplementedError",
status_code=501,
)
async def _handle_non_streaming_request(
self,
adapted_request: GenerateReqInput,
request: OpenAIServingRequest,
raw_request: Request,
) -> Union[Any, ErrorResponse, ORJSONResponse]:
"""Handle non-streaming request
Override this method in child classes that support non-streaming requests.
"""
return self.create_error_response(
message=f"{self.__class__.__name__} does not support non-streaming requests",
err_type="NotImplementedError",
status_code=501,
)
def _validate_request(self, _: OpenAIServingRequest) -> Optional[str]:
"""Validate request"""
pass
def create_error_response(
self,
message: str,
err_type: str = "BadRequestError",
status_code: int = 400,
param: Optional[str] = None,
) -> ORJSONResponse:
"""Create an error response"""
# TODO: remove fastapi dependency in openai and move response handling to the entrypoint
error = ErrorResponse(
object="error",
message=message,
type=err_type,
param=param,
code=status_code,
)
return ORJSONResponse(content=error.model_dump(), status_code=status_code)
def create_streaming_error_response(
self,
message: str,
err_type: str = "BadRequestError",
status_code: int = 400,
) -> str:
"""Create a streaming error response"""
error = ErrorResponse(
object="error",
message=message,
type=err_type,
param=None,
code=status_code,
)
return json.dumps({"error": error.model_dump()})
def extract_custom_labels(self, raw_request):
if (
not self.allowed_custom_labels
or not self.tokenizer_manager.server_args.tokenizer_metrics_custom_labels_header
):
return None
custom_labels = None
header = (
self.tokenizer_manager.server_args.tokenizer_metrics_custom_labels_header
)
try:
raw_labels = (
orjson.loads(raw_request.headers.get(header))
if raw_request and raw_request.headers.get(header)
else None
)
except json.JSONDecodeError as e:
logger.exception(f"Error in request: {e}")
raw_labels = None
if isinstance(raw_labels, dict):
custom_labels = {
label: value
for label, value in raw_labels.items()
if label in self.allowed_custom_labels
}
return custom_labels
def extract_routing_key(self, raw_request):
if raw_request is None:
return None
return raw_request.headers.get("x-smg-routing-key")
def extract_routed_dp_rank_from_header(
self, raw_request: Request, body_routed_dp_rank: Optional[int] = None
) -> Optional[int]:
"""Extract routed_dp_rank from HTTP header, with higher priority than routed_dp_rank in body.
Header name: X-Data-Parallel-Rank (case-insensitive in HTTP/1.1/2)
"""
if raw_request is None:
return body_routed_dp_rank
header_value = raw_request.headers.get("x-data-parallel-rank")
if header_value is not None:
try:
header_dp_rank = int(header_value)
if (
body_routed_dp_rank is not None
and header_dp_rank != body_routed_dp_rank
):
logger.debug(
f"X-Data-Parallel-Rank header ({header_dp_rank}) overrides "
f"body routed_dp_rank ({body_routed_dp_rank})"
)
return header_dp_rank
except ValueError:
raise HTTPException(
status_code=400,
detail=f"Invalid X-Data-Parallel-Rank header: must be an integer, got '{header_value}'",
)
return body_routed_dp_rank
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,204 @@
from __future__ import annotations
import logging
import time
import uuid
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
import torch
import torch.nn.functional as F
from fastapi import Request
from fastapi.responses import ORJSONResponse
from sglang.srt.entrypoints.openai.protocol import (
ClassifyRequest,
ClassifyResponse,
ErrorResponse,
)
from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
from sglang.srt.managers.io_struct import EmbeddingReqInput
if TYPE_CHECKING:
from sglang.srt.managers.tokenizer_manager import TokenizerManager
from sglang.srt.parser.template_manager import TemplateManager
logger = logging.getLogger(__name__)
class OpenAIServingClassify(OpenAIServingBase):
"""Handler for v1/classify requests"""
def __init__(
self,
tokenizer_manager: TokenizerManager,
template_manager: TemplateManager,
):
super().__init__(tokenizer_manager)
self.template_manager = template_manager
self.id2label = self._get_id2label_mapping()
self.model_name = (
self.tokenizer_manager.served_model_name
if self.tokenizer_manager.served_model_name
else self.tokenizer_manager.server_args.model_path
)
if not self.id2label:
raise ValueError("id2label mapping is missing")
def _request_id_prefix(self) -> str:
return "classify-"
def _convert_to_internal_request(
self,
request: ClassifyRequest,
raw_request: Request = None,
) -> tuple[EmbeddingReqInput, ClassifyRequest]:
"""Convert OpenAI embedding request to internal format"""
prompt = request.input
if isinstance(prompt, str):
# Single string input
prompt_kwargs = {"text": prompt}
elif isinstance(prompt, list):
if len(prompt) > 0 and isinstance(prompt[0], str):
prompt_kwargs = {"text": prompt}
else:
# List of integers (token IDs) or empty list
prompt_kwargs = {"input_ids": prompt}
else:
# Other types (should not happen but handle gracefully)
prompt_kwargs = {"input_ids": prompt}
adapted_request = EmbeddingReqInput(
**prompt_kwargs,
rid=request.rid,
priority=request.priority,
)
return adapted_request, request
def _validate_request(self, request: ClassifyRequest) -> Optional[str]:
"""Validate that the input is not empty or whitespace only."""
if not (input := request.input):
return "Input cannot be empty"
# Handle single string
if isinstance(input, str):
if not input.strip():
return "Input cannot be empty or whitespace only"
return None
# Handle list inputs
if isinstance(input, list):
# Check first element to determine type
first_item = input[0]
if isinstance(first_item, str):
# List of strings
for i, item in enumerate(input):
if not isinstance(item, str):
return f"All items in input list must be strings"
if not item.strip():
return f"Input at index {i} cannot be empty or whitespace only"
elif isinstance(first_item, int):
# List of integers (token IDs)
for i, item in enumerate(input):
if not isinstance(item, int):
return f"All items in input list must be integers"
if item < 0:
return f"Token ID at index {i} must be non-negative"
return None
def _get_id2label_mapping(self) -> Optional[Dict[int, str]]:
"""Get id2label mapping from model config."""
try:
hf_config = self.tokenizer_manager.model_config.hf_config
# Check for id2label in hf_config
if hf_config.id2label:
return hf_config.id2label
# Check for num_labels and create default mapping if needed
if hasattr(hf_config, "num_labels") and hf_config.num_labels:
num_labels = hf_config.num_labels
# Create default mapping: {0: "LABEL_0", 1: "LABEL_1", ...}
return {i: f"LABEL_{i}" for i in range(num_labels)}
except Exception as e:
logger.warning(f"Failed to get id2label mapping: {e}")
return None
async def _handle_non_streaming_request(
self,
adapted_request: EmbeddingReqInput,
request: ClassifyRequest,
raw_request: Request,
) -> Union[ClassifyResponse, ErrorResponse, ORJSONResponse]:
"""Handle non-streaming classification request."""
# Generate request ID
try:
ret = await self.tokenizer_manager.generate_request(
adapted_request, raw_request
).__anext__()
except ValueError as e:
return self.create_error_response(str(e))
if not isinstance(ret, list):
ret = [ret]
response = self._build_classify_response(ret)
return response
def _build_classify_response(self, ret: List[Dict[str, Any]]) -> ClassifyResponse:
request_id = f"{self._request_id_prefix()}{uuid.uuid4().hex}"
created_time = int(time.time())
classify_objects = []
prompt_tokens = 0
total_latency = 0.0
for i, item in enumerate(ret):
embedding = item.get("embedding", [])
meta_info = item.get("meta_info", {})
prompt_tokens += meta_info.get("prompt_tokens", 0)
total_latency += meta_info.get("e2e_latency", 0.0)
if embedding:
try:
embedding_tensor = torch.tensor(embedding, dtype=torch.float32)
probs = F.softmax(embedding_tensor, dim=0).tolist()
predicted_class = torch.argmax(embedding_tensor).item()
label = self.id2label[predicted_class]
except Exception as e:
logger.error(f"Error processing embedding for item {i}: {e}")
probs = [1.0]
label = "Default"
else:
probs = [1.0]
label = "Default"
classify_obj = {
"index": i,
"label": label,
"probs": probs,
"num_classes": len(probs),
}
classify_objects.append(classify_obj)
response = {
"id": request_id,
"object": "list",
"created": created_time,
"model": self.model_name,
"data": classify_objects,
"usage": {
"prompt_tokens": prompt_tokens,
"total_tokens": prompt_tokens,
"completion_tokens": 0,
"prompt_tokens_details": None,
},
}
return ClassifyResponse(**response)
@@ -0,0 +1,616 @@
from __future__ import annotations
import logging
import time
from http import HTTPStatus
from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Optional, Union
from fastapi import Request
from fastapi.responses import ORJSONResponse, StreamingResponse
from sglang.srt.entrypoints.openai.protocol import (
CompletionRequest,
CompletionResponse,
CompletionResponseChoice,
CompletionResponseStreamChoice,
CompletionStreamResponse,
ErrorResponse,
SglExt,
)
from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
from sglang.srt.entrypoints.openai.usage_processor import UsageProcessor
from sglang.srt.entrypoints.openai.utils import (
cached_tokens_details_from_dict,
process_cached_tokens_details_from_ret,
process_hidden_states_from_ret,
process_routed_experts_from_ret,
should_include_usage,
to_openai_style_logprobs,
)
from sglang.srt.managers.io_struct import GenerateReqInput
from sglang.srt.parser.code_completion_parser import (
generate_completion_prompt_from_request,
)
from sglang.utils import convert_json_schema_to_str
if TYPE_CHECKING:
from sglang.srt.managers.tokenizer_manager import TokenizerManager
from sglang.srt.parser.template_manager import TemplateManager
logger = logging.getLogger(__name__)
class OpenAIServingCompletion(OpenAIServingBase):
"""Handler for /v1/completion requests"""
def __init__(
self,
tokenizer_manager: TokenizerManager,
template_manager: TemplateManager,
):
super().__init__(tokenizer_manager)
self.template_manager = template_manager
def _request_id_prefix(self) -> str:
return "cmpl-"
def _validate_request(self, request: CompletionRequest) -> Optional[str]:
"""Validate that the input is valid."""
prompt = request.prompt
if not prompt or (isinstance(prompt, list) and all(not p for p in prompt)):
return "Prompt cannot be empty"
return None
def _convert_to_internal_request(
self,
request: CompletionRequest,
raw_request: Request = None,
) -> tuple[GenerateReqInput, CompletionRequest]:
"""Convert OpenAI completion request to internal format"""
# NOTE: with openai API, the prompt's logprobs are always not computed
if request.echo and request.logprobs:
logger.warning(
"Echo is not compatible with logprobs. "
"To compute logprobs of input prompt, please use the native /generate API."
)
# Process prompt
prompt = request.prompt
if self.template_manager.completion_template_name is not None:
prompt = generate_completion_prompt_from_request(request)
# Set logprob start length based on echo and logprobs
if request.echo and request.logprobs:
logprob_start_len = 0
else:
logprob_start_len = -1
# Build sampling parameters
sampling_params = self._build_sampling_params(request)
# Determine prompt format
if isinstance(prompt, str) or (
isinstance(prompt, list) and isinstance(prompt[0], str)
):
prompt_kwargs = {"text": prompt}
else:
prompt_kwargs = {"input_ids": prompt}
# Extract custom labels from raw request headers
custom_labels = self.extract_custom_labels(raw_request)
# Extract routed_dp_rank from header (has higher priority than body)
effective_routed_dp_rank = self.extract_routed_dp_rank_from_header(
raw_request, request.routed_dp_rank
)
# Resolve LoRA adapter from model parameter or explicit lora_path
lora_path = self._resolve_lora_path(request.model, request.lora_path)
adapted_request = GenerateReqInput(
**prompt_kwargs,
sampling_params=sampling_params,
return_logprob=request.logprobs is not None,
top_logprobs_num=request.logprobs if request.logprobs is not None else 0,
logprob_start_len=logprob_start_len,
return_text_in_logprobs=True,
stream=request.stream,
lora_path=lora_path,
bootstrap_host=request.bootstrap_host,
bootstrap_port=request.bootstrap_port,
bootstrap_room=request.bootstrap_room,
routed_dp_rank=effective_routed_dp_rank,
disagg_prefill_dp_rank=request.disagg_prefill_dp_rank,
return_hidden_states=request.return_hidden_states,
return_routed_experts=request.return_routed_experts,
routed_experts_start_len=request.routed_experts_start_len,
rid=request.rid,
session_id=request.session_id,
extra_key=self._compute_extra_key(request),
priority=request.priority,
routing_key=self.extract_routing_key(raw_request),
custom_labels=custom_labels,
custom_logit_processor=request.custom_logit_processor,
images_config=getattr(request, "images_config", None),
)
return adapted_request, request
def _build_sampling_params(self, request: CompletionRequest) -> Dict[str, Any]:
"""Build sampling parameters for the request"""
# Start with common parameters
sampling_params = {
"temperature": request.temperature,
"max_new_tokens": request.max_tokens,
"min_new_tokens": request.min_tokens,
"stop": request.stop,
"stop_token_ids": request.stop_token_ids,
"stop_regex": request.stop_regex,
"top_p": request.top_p,
"top_k": request.top_k,
"min_p": request.min_p,
"presence_penalty": request.presence_penalty,
"frequency_penalty": request.frequency_penalty,
"repetition_penalty": request.repetition_penalty,
"regex": request.regex,
"json_schema": request.json_schema,
"ebnf": request.ebnf,
"n": request.n,
"no_stop_trim": request.no_stop_trim,
"ignore_eos": request.ignore_eos,
"skip_special_tokens": request.skip_special_tokens,
"logit_bias": request.logit_bias,
"custom_params": request.custom_params,
"sampling_seed": request.seed,
}
# Handle response_format constraints
if request.response_format and request.response_format.type == "json_schema":
json_schema = request.response_format.json_schema
schema = getattr(json_schema, "schema_", None)
if schema is None:
raise ValueError(
"schema_ is required for json_schema response format request."
)
sampling_params["json_schema"] = convert_json_schema_to_str(schema)
elif request.response_format and request.response_format.type == "json_object":
sampling_params["json_schema"] = '{"type": "object"}'
elif (
request.response_format and request.response_format.type == "structural_tag"
):
sampling_params["structural_tag"] = convert_json_schema_to_str(
request.response_format.model_dump(by_alias=True)
)
return sampling_params
async def _handle_streaming_request(
self,
adapted_request: GenerateReqInput,
request: CompletionRequest,
raw_request: Request,
) -> Union[StreamingResponse, ErrorResponse]:
"""Handle streaming completion request"""
generator = self._generate_completion_stream(
adapted_request, request, raw_request
)
# Kick-start the generator to trigger validation before HTTP 200 is sent.
try:
first_chunk = await generator.__anext__()
except ValueError as e:
return self.create_error_response(str(e))
async def prepend_first_chunk():
yield first_chunk
async for chunk in generator:
yield chunk
return StreamingResponse(
prepend_first_chunk(),
media_type="text/event-stream",
background=self.tokenizer_manager.create_abort_task(adapted_request),
)
async def _generate_completion_stream(
self,
adapted_request: GenerateReqInput,
request: CompletionRequest,
raw_request: Request,
) -> AsyncGenerator[str, None]:
"""Generate streaming completion response"""
created = int(time.time())
# State tracking for streaming
stream_offsets = {}
n_prev_tokens = {}
# Usage tracking
prompt_tokens = {}
completion_tokens = {}
reasoning_tokens = {}
cached_tokens = {}
hidden_states = {}
routed_experts = {}
cached_tokens_details = {}
stream_started = False
try:
include_usage, continuous_usage_stats = should_include_usage(
request.stream_options,
self.tokenizer_manager.server_args.stream_response_default_include_usage,
)
async for content in self.tokenizer_manager.generate_request(
adapted_request, raw_request
):
index = content.get("index", 0)
text = content["text"]
prompt_tokens[index] = content["meta_info"].get("prompt_tokens", 0)
completion_tokens[index] = content["meta_info"].get(
"completion_tokens", 0
)
reasoning_tokens[index] = content["meta_info"].get(
"reasoning_tokens", 0
)
cached_tokens[index] = content["meta_info"].get("cached_tokens", 0)
hidden_states[index] = content["meta_info"].get("hidden_states", None)
routed_experts[index] = content["meta_info"].get("routed_experts", None)
cached_tokens_details[index] = content["meta_info"].get(
"cached_tokens_details", None
)
is_first_chunk = index not in stream_offsets
offset = stream_offsets.get(index, 0)
# Handle echo for first chunk
if is_first_chunk: # The first chunk
if request.echo:
echo_text = self._get_echo_text(request, index)
text = echo_text + text
# Handle logprobs
logprobs = None
if request.logprobs is not None:
# The first chunk and echo is enabled.
if is_first_chunk and request.echo:
input_token_logprobs = content["meta_info"][
"input_token_logprobs"
]
input_top_logprobs = content["meta_info"]["input_top_logprobs"]
else:
input_token_logprobs = None
input_top_logprobs = None
n_prev_token = n_prev_tokens.get(index, 0)
total_output_logprobs = content["meta_info"][
"output_token_logprobs_length"
]
if (
n_prev_token < total_output_logprobs
or input_token_logprobs is not None
):
output_token_logprobs = content["meta_info"][
"output_token_logprobs"
]
output_top_logprobs = content["meta_info"].get(
"output_top_logprobs", []
)
if (
not self.tokenizer_manager.server_args.incremental_streaming_output
):
output_token_logprobs = output_token_logprobs[
n_prev_token:total_output_logprobs
]
output_top_logprobs = output_top_logprobs[
n_prev_token:total_output_logprobs
]
logprobs = to_openai_style_logprobs(
input_token_logprobs=input_token_logprobs,
input_top_logprobs=input_top_logprobs,
output_token_logprobs=output_token_logprobs,
output_top_logprobs=output_top_logprobs,
)
n_prev_tokens[index] = total_output_logprobs
# Generate delta
delta = text[offset:]
stream_offsets[index] = len(content["text"])
finish_reason = content["meta_info"].get("finish_reason", None)
finish_reason_type = finish_reason["type"] if finish_reason else None
# Abort with an explicit error status_code is a system error
# (timeout, OOM, validation): emit a streaming error chunk.
# A graceful abort (no status_code, e.g. user-initiated via
# /abort_request or session lifecycle cleanup) falls through
# to the normal chunk path, matching the non-stream behavior
# in tokenizer_manager._handle_abort_finish_reason.
if finish_reason_type == "abort" and isinstance(
finish_reason.get("status_code"), HTTPStatus
):
code = finish_reason["status_code"]
error = self.create_streaming_error_response(
finish_reason.get("message", "Generation aborted."),
code.name,
code.value,
)
yield f"data: {error}\n\n"
break
choice_data = CompletionResponseStreamChoice(
index=index,
text=delta,
logprobs=logprobs,
finish_reason=finish_reason_type,
matched_stop=(
finish_reason["matched"]
if finish_reason and "matched" in finish_reason
else None
),
)
chunk = CompletionStreamResponse(
id=content["meta_info"]["id"],
created=created,
object="text_completion",
choices=[choice_data],
model=request.model,
)
# Add usage stats if continuous_usage_stats is enabled
if continuous_usage_stats:
chunk.usage = UsageProcessor.calculate_token_usage(
prompt_tokens=prompt_tokens.get(index, 0),
completion_tokens=completion_tokens.get(index, 0),
reasoning_tokens=reasoning_tokens.get(index, 0),
)
yield f"data: {chunk.model_dump_json()}\n\n"
stream_started = True
if request.return_hidden_states and hidden_states:
for index, choice_hidden_states in hidden_states.items():
if choice_hidden_states:
last_token_hidden_states = (
choice_hidden_states[-1]
if len(choice_hidden_states) > 1
else []
)
hidden_states_chunk = CompletionStreamResponse(
id=content["meta_info"]["id"],
created=created,
object="text_completion",
choices=[
CompletionResponseStreamChoice(
index=index,
text="",
hidden_states=last_token_hidden_states,
finish_reason=None,
)
],
model=request.model,
)
yield f"data: {hidden_states_chunk.model_dump_json()}\n\n"
sglext_routed = None
if request.return_routed_experts and routed_experts:
sglext_routed = next(
(v for v in routed_experts.values() if v is not None), None
)
sglext_details = None
if request.return_cached_tokens_details and cached_tokens_details:
first_details = next(
(v for v in cached_tokens_details.values() if v is not None), None
)
if first_details is not None:
sglext_details = cached_tokens_details_from_dict(first_details)
if sglext_routed is not None or sglext_details is not None:
sglext_chunk = CompletionStreamResponse(
id=content["meta_info"]["id"],
created=created,
object="text_completion",
choices=[], # sglext is at response level
model=request.model,
sglext=SglExt(
routed_experts=sglext_routed,
cached_tokens_details=sglext_details,
),
)
yield f"data: {sglext_chunk.model_dump_json()}\n\n"
# Handle final usage chunk
if include_usage:
usage = UsageProcessor.calculate_streaming_usage(
prompt_tokens,
reasoning_tokens,
completion_tokens,
cached_tokens=cached_tokens,
n_choices=request.n,
enable_cache_report=self.tokenizer_manager.server_args.enable_cache_report,
)
final_usage_chunk = CompletionStreamResponse(
id=content["meta_info"]["id"],
created=created,
choices=[],
model=request.model,
usage=usage,
)
final_usage_data = final_usage_chunk.model_dump_json(exclude_none=True)
yield f"data: {final_usage_data}\n\n"
except Exception as e:
if not stream_started:
raise
error = self.create_streaming_error_response(str(e))
yield f"data: {error}\n\n"
yield "data: [DONE]\n\n"
async def _handle_non_streaming_request(
self,
adapted_request: GenerateReqInput,
request: CompletionRequest,
raw_request: Request,
) -> Union[CompletionResponse, ErrorResponse, ORJSONResponse]:
"""Handle non-streaming completion request"""
try:
generator = self.tokenizer_manager.generate_request(
adapted_request, raw_request
)
ret = await generator.__anext__()
except ValueError as e:
return self.create_error_response(str(e))
if not isinstance(ret, list):
ret = [ret]
response = self._build_completion_response(
request,
ret,
int(time.time()),
)
return response
def _build_completion_response(
self,
request: CompletionRequest,
ret: List[Dict[str, Any]],
created: int,
) -> CompletionResponse:
"""Build completion response from generation results"""
choices = []
echo = False
# Prepare echo prompts if needed
echo_prompts = []
if request.echo:
echo_prompts = self._prepare_echo_prompts(request)
echo = True
# Build sglext at response level (from first ret_item, as these are per-request)
first_ret = ret[0]
routed_experts = process_routed_experts_from_ret(first_ret, request)
cached_tokens_details = process_cached_tokens_details_from_ret(
first_ret, request
)
response_sglext = None
if routed_experts or cached_tokens_details:
response_sglext = SglExt(
routed_experts=routed_experts,
cached_tokens_details=cached_tokens_details,
)
for idx, ret_item in enumerate(ret):
text = ret_item["text"]
# Handle echo
if echo:
prompt_index = idx // request.n
text = echo_prompts[prompt_index] + text
# Handle logprobs
logprobs = None
if request.logprobs is not None:
if echo:
input_token_logprobs = ret_item["meta_info"]["input_token_logprobs"]
input_top_logprobs = ret_item["meta_info"]["input_top_logprobs"]
else:
input_token_logprobs = None
input_top_logprobs = None
logprobs = to_openai_style_logprobs(
input_token_logprobs=input_token_logprobs,
input_top_logprobs=input_top_logprobs,
output_token_logprobs=ret_item["meta_info"].get(
"output_token_logprobs", []
),
output_top_logprobs=ret_item["meta_info"].get(
"output_top_logprobs", []
),
)
# Handle hidden states
hidden_states = process_hidden_states_from_ret(ret_item, request)
finish_reason = ret_item["meta_info"]["finish_reason"]
choice_data = CompletionResponseChoice(
index=idx,
text=text,
logprobs=logprobs,
finish_reason=finish_reason["type"] if finish_reason else None,
matched_stop=(
finish_reason["matched"]
if finish_reason and "matched" in finish_reason
else None
),
hidden_states=hidden_states,
)
choices.append(choice_data)
# Calculate usage
cache_report = self.tokenizer_manager.server_args.enable_cache_report
usage = UsageProcessor.calculate_response_usage(
ret, n_choices=request.n, enable_cache_report=cache_report
)
return CompletionResponse(
id=ret[0]["meta_info"]["id"],
model=request.model,
created=created,
choices=choices,
usage=usage,
metadata={"weight_version": ret[0]["meta_info"]["weight_version"]},
sglext=response_sglext,
)
def _get_echo_text(self, request: CompletionRequest, index: int) -> str:
"""Get echo text for streaming response"""
if isinstance(request.prompt, str):
# for the case of single str prompts
return request.prompt
elif isinstance(request.prompt, list):
if isinstance(request.prompt[0], str):
# for the case of multiple str prompts
return request.prompt[index // request.n]
elif isinstance(request.prompt[0], int):
# for the case of single token ids prompt
return self.tokenizer_manager.tokenizer.decode(
request.prompt, skip_special_tokens=True
)
elif isinstance(request.prompt[0], list) and isinstance(
request.prompt[0][0], int
):
# for the case of multiple token ids prompts
return self.tokenizer_manager.tokenizer.decode(
request.prompt[index // request.n],
skip_special_tokens=True,
)
return ""
def _prepare_echo_prompts(self, request: CompletionRequest) -> List[str]:
"""Prepare echo prompts for non-streaming response"""
# TODO: handle the case prompt is token ids
if isinstance(request.prompt, list) and isinstance(request.prompt[0], str):
# for the case of multiple str prompts
return request.prompt
elif isinstance(request.prompt, list) and isinstance(request.prompt[0], list):
# for the case of multiple token ids prompts
return [
self.tokenizer_manager.tokenizer.decode(
prompt, skip_special_tokens=True
)
for prompt in request.prompt
]
elif isinstance(request.prompt, list) and isinstance(request.prompt[0], int):
# for the case of single token ids prompt
return [
self.tokenizer_manager.tokenizer.decode(
request.prompt, skip_special_tokens=True
)
]
else:
# for the case of single str prompt
return [request.prompt]
@@ -0,0 +1,283 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
import jinja2
from fastapi import Request
from fastapi.responses import ORJSONResponse
from sglang.srt.entrypoints.openai.protocol import (
EmbeddingObject,
EmbeddingRequest,
EmbeddingResponse,
ErrorResponse,
MultimodalEmbeddingInput,
UsageInfo,
)
from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
from sglang.srt.entrypoints.openai.utils import convert_embeds_to_tensors
from sglang.srt.managers.io_struct import EmbeddingReqInput
from sglang.srt.parser.conversation import generate_embedding_convs
from sglang.srt.parser.jinja_template_utils import process_content_for_template_format
if TYPE_CHECKING:
from sglang.srt.managers.tokenizer_manager import TokenizerManager
from sglang.srt.parser.template_manager import TemplateManager
class OpenAIServingEmbedding(OpenAIServingBase):
"""Handler for v1/embeddings requests"""
def __init__(
self,
tokenizer_manager: TokenizerManager,
template_manager: TemplateManager,
):
super().__init__(tokenizer_manager)
self.template_manager = template_manager
def _request_id_prefix(self) -> str:
return "embd-"
def _validate_request(self, request: EmbeddingRequest) -> Optional[str]:
"""Validate that the input is not empty or whitespace only."""
if not (input := request.input):
return "Input cannot be empty"
# Handle single string
if isinstance(input, str):
if not input.strip():
return "Input cannot be empty or whitespace only"
return None
# Handle list inputs
if isinstance(input, list):
if len(input) == 0:
return "Input cannot be empty"
# Check first element to determine type
first_item = input[0]
if isinstance(first_item, str):
# List of strings
for i, item in enumerate(input):
if not isinstance(item, str):
return "All items in input list must be strings"
if not item.strip():
return f"Input at index {i} cannot be empty or whitespace only"
elif isinstance(first_item, int):
# List of integers (token IDs)
for i, item in enumerate(input):
if not isinstance(item, int):
return "All items in input list must be integers"
if item < 0:
return f"Token ID at index {i} must be non-negative"
return None
def _convert_to_internal_request(
self,
request: EmbeddingRequest,
raw_request: Request = None,
) -> tuple[EmbeddingReqInput, EmbeddingRequest]:
"""Convert OpenAI embedding request to internal format"""
prompt = request.input
if isinstance(prompt, str):
# Single string input
prompt_kwargs = {"text": prompt}
elif isinstance(prompt, list):
if len(prompt) > 0 and isinstance(prompt[0], str):
prompt_kwargs = {"text": prompt}
elif len(prompt) > 0 and isinstance(prompt[0], MultimodalEmbeddingInput):
# Handle multimodal embedding inputs
texts = []
images = []
videos = []
for item in prompt:
texts.append(item.text)
images.append(item.image if item.image is not None else None)
videos.append(item.video if item.video is not None else None)
# Precedence: a SGLang-registered conversation template wins
# over the tokenizer's own HF Jinja template when both exist.
generate_prompts = []
if self.template_manager.chat_template_name is not None:
convs = generate_embedding_convs(
texts, images, videos, self.template_manager.chat_template_name
)
for conv in convs:
generate_prompts.append(conv.get_prompt())
elif (
self.tokenizer_manager.tokenizer is not None
and getattr(self.tokenizer_manager.tokenizer, "chat_template", None)
is not None
):
generate_prompts = self._apply_jinja_template_to_embedding_inputs(
texts, images, videos
)
else:
generate_prompts = [
text if text is not None else "padding" for text in texts
]
if len(generate_prompts) == 1:
prompt_kwargs = {
"text": generate_prompts[0],
"image_data": images[0],
"video_data": videos[0],
}
else:
prompt_kwargs = {
"text": generate_prompts,
"image_data": images,
"video_data": videos,
}
else:
# List of integers (token IDs) or empty list
prompt_kwargs = {"input_ids": prompt}
else:
# Other types (should not happen but handle gracefully)
prompt_kwargs = {"input_ids": prompt}
# Resolve LoRA adapter from model parameter or explicit lora_path
lora_path = self._resolve_lora_path(request.model, request.lora_path)
# Validate pairing: both or neither must be provided
if (
request.embed_overrides is not None
and request.embed_override_token_id is None
):
raise ValueError(
"embed_override_token_id is required when embed_overrides is provided"
)
if (
request.embed_override_token_id is not None
and request.embed_overrides is None
):
raise ValueError(
"embed_override_token_id requires embed_overrides to be provided"
)
# Convert float lists to tensors; position resolution is deferred
# to the tokenizer manager (after tokenization for text inputs).
embed_overrides = convert_embeds_to_tensors(request.embed_overrides)
adapted_request = EmbeddingReqInput(
**prompt_kwargs,
rid=request.rid,
priority=request.priority,
routing_key=self.extract_routing_key(raw_request),
dimensions=request.dimensions,
lora_path=lora_path,
embed_override_token_id=request.embed_override_token_id,
embed_overrides=embed_overrides,
)
return adapted_request, request
def _apply_jinja_template_to_embedding_inputs(
self,
texts: List[Optional[str]],
images: List[Optional[str]],
videos: List[Optional[str]],
) -> List[str]:
"""Render each multimodal embedding input through the tokenizer's Jinja chat template.
Image/video bytes are threaded to the engine separately via
``EmbeddingReqInput.image_data``/``video_data``; this method only produces
the prompt string. ``text=None`` emits no text chunk (no ``"padding"``
literal). Jinja failures are re-raised as ``ValueError`` so the caller
returns HTTP 400 instead of 500.
"""
prompts: List[str] = []
template_content_format = self.template_manager.jinja_template_content_format
for text, image, video in zip(texts, images, videos):
content_parts = []
if image is not None:
content_parts.append({"type": "image_url", "image_url": {"url": image}})
if video is not None:
content_parts.append({"type": "video_url", "video_url": {"url": video}})
if text is not None:
content_parts.append({"type": "text", "text": text})
msg_dict = {
"role": "user",
"content": content_parts if content_parts else "",
}
# Empty list args: this helper is only used to normalize the content
# shape (e.g. image_url -> image); real payloads ride on the outer
# images/videos lists, not EmbeddingReqInput fields derived here.
processed_msg = process_content_for_template_format(
msg_dict,
template_content_format,
image_data=[],
video_data=[],
audio_data=[],
modalities=[],
)
try:
prompt = self.tokenizer_manager.tokenizer.apply_chat_template(
[processed_msg],
tokenize=False,
add_generation_prompt=True,
)
except jinja2.TemplateError as template_error:
location = getattr(template_error, "lineno", None)
name = getattr(template_error, "name", None)
suffix = ""
if name or location:
suffix = f" (template={name or '<unknown>'}, line={location})"
raise ValueError(f"{template_error}{suffix}") from template_error
except (TypeError, KeyError, AttributeError) as template_error:
raise ValueError(
f"Failed to render chat template for embedding input: {template_error}"
) from template_error
prompts.append(prompt)
return prompts
async def _handle_non_streaming_request(
self,
adapted_request: EmbeddingReqInput,
request: EmbeddingRequest,
raw_request: Request,
) -> Union[EmbeddingResponse, ErrorResponse, ORJSONResponse]:
"""Handle the embedding request"""
try:
ret = await self.tokenizer_manager.generate_request(
adapted_request, raw_request
).__anext__()
except ValueError as e:
return self.create_error_response(str(e))
if not isinstance(ret, list):
ret = [ret]
response = self._build_embedding_response(ret)
return response
def _build_embedding_response(self, ret: List[Dict[str, Any]]) -> EmbeddingResponse:
"""Build the embedding response"""
embedding_objects = []
prompt_tokens = 0
for idx, ret_item in enumerate(ret):
embedding_objects.append(
EmbeddingObject(
embedding=ret_item["embedding"],
index=idx,
)
)
# Handle missing prompt_tokens gracefully
meta_info = ret_item.get("meta_info", {})
prompt_tokens += meta_info.get("prompt_tokens", 0)
return EmbeddingResponse(
data=embedding_objects,
model=self.tokenizer_manager.model_path,
usage=UsageInfo(
prompt_tokens=prompt_tokens,
total_tokens=prompt_tokens,
),
)
@@ -0,0 +1,609 @@
import heapq
import logging
import math
from typing import Any, Dict, List, Optional, Union
from fastapi import Request
from fastapi.responses import ORJSONResponse
from sglang.srt.entrypoints.openai.protocol import (
ChatCompletionMessageContentImagePart,
ChatCompletionMessageContentTextPart,
ChatCompletionMessageContentVideoPart,
ErrorResponse,
RerankContent,
RerankResponse,
V1RerankReqInput,
)
from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
from sglang.srt.managers.io_struct import EmbeddingReqInput, GenerateReqInput
logger = logging.getLogger(__name__)
def _get_yes_no_token_ids(tokenizer) -> tuple[int, int]:
"""Get token IDs for 'yes' and 'no' from the tokenizer.
Different model sizes may have different token IDs, so we look them up dynamically.
"""
# Try to encode 'yes' and 'no' to get their token IDs
# The tokenizer should return a single token for these common words
try:
yes_tokens = tokenizer.encode("yes", add_special_tokens=False)
no_tokens = tokenizer.encode("no", add_special_tokens=False)
if len(yes_tokens) == 1 and len(no_tokens) == 1:
return yes_tokens[0], no_tokens[0]
# Fallback: try convert_tokens_to_ids
yes_id = tokenizer.convert_tokens_to_ids("yes")
no_id = tokenizer.convert_tokens_to_ids("no")
if yes_id is not None and no_id is not None:
return yes_id, no_id
except Exception as e:
logger.warning(f"Failed to get yes/no token IDs dynamically: {e}")
# Fallback to known Qwen3 token IDs (may not work for all model sizes)
logger.warning("Using fallback token IDs for yes/no (9693/2152)")
return 9693, 2152
def _is_qwen3_reranker_template(chat_template: str) -> bool:
"""Detect if the chat template is for Qwen3 text-only reranker."""
if not chat_template:
return False
t = chat_template.lower()
return ('answer can only be "yes" or "no"' in t) or (
"answer can only be" in t and '"yes"' in t and '"no"' in t
)
def _is_qwen3_vl_reranker_template(chat_template: str) -> bool:
"""Detect if the chat template is for Qwen3-VL multimodal reranker.
VL reranker templates use `query` and `document` as jinja variables
and include vision token placeholders for image/video support.
"""
if not chat_template:
return False
t = chat_template.lower()
# Check for reranker phrase (yes/no judgment)
has_reranker_phrase = ('answer can only be "yes" or "no"' in t) or (
"answer can only be" in t and '"yes"' in t and '"no"' in t
)
# Check for vision token placeholders (unique to VL templates)
has_vision_tokens = "<|vision_start|>" in t or "<|image_pad|>" in t
return has_reranker_phrase and has_vision_tokens
def _is_qwen3_vl_model(model_path: str) -> bool:
"""Check if the model is a Qwen3-VL model based on model path."""
if not model_path:
return False
model_lower = model_path.lower()
return "qwen3-vl" in model_lower or "qwen3vl" in model_lower
def _detect_rerank_backend(
*,
request: V1RerankReqInput,
chat_template: Optional[str],
model_path: str,
) -> str:
"""
Unify rerank routing decisions used by both `_convert_to_internal_request` and
`_handle_non_streaming_request`.
Returns:
"vl_decoder" | "text_decoder" | "cross_encoder"
"""
is_multimodal = request.is_multimodal()
is_vl_model = _is_qwen3_vl_model(model_path)
is_vl_template = _is_qwen3_vl_reranker_template(chat_template)
is_text_template = _is_qwen3_reranker_template(chat_template)
# Prefer VL when template/model indicates VL, or request is multimodal with reranker template.
if is_vl_template or is_vl_model or (is_multimodal and is_text_template):
return "vl_decoder"
if is_text_template:
return "text_decoder"
return "cross_encoder"
def _qwen3_rerank_score(p_yes: float, p_no: float) -> float:
denom = p_yes + p_no
if denom <= 0.0:
return 0.0
return p_yes / denom
def _get_jinja_env():
try:
import jinja2 # Lazy import: server env should provide this dependency.
from jinja2.sandbox import ImmutableSandboxedEnvironment
except ModuleNotFoundError as e:
raise ValueError(
"Rendering Qwen3 reranker prompts requires `jinja2`. "
"Please install it in your runtime environment (e.g., `pip install jinja2`)."
) from e
# Using a sandboxed environment to stop malicious execution during model loading.
return ImmutableSandboxedEnvironment(
loader=jinja2.BaseLoader(),
autoescape=False,
undefined=jinja2.Undefined,
)
def _render_jinja_chat_template(
chat_template: str,
*,
query: RerankContent,
document: RerankContent,
instruct: Optional[str],
) -> str:
"""Render a loaded Jinja chat template for Qwen3 reranker prompts (text-only)."""
env = _get_jinja_env()
template = env.from_string(chat_template)
# For text-only template, extract text content
query_text = query if isinstance(query, str) else _extract_text_from_content(query)
doc_text = (
document if isinstance(document, str) else _extract_text_from_content(document)
)
render_kwargs = {
"messages": [
{"role": "user", "content": query_text},
{"role": "user", "content": doc_text},
]
}
# Only pass instruct when explicitly provided; template uses `default(...)`
# which works only when the variable is undefined (not None).
if instruct:
render_kwargs["instruct"] = instruct
return template.render(**render_kwargs)
def _render_vl_jinja_template(
chat_template: str,
*,
query: List[Dict[str, Any]],
document: List[Dict[str, Any]],
instruct: Optional[str],
) -> str:
"""Render a loaded Jinja chat template for Qwen3-VL reranker prompts (multimodal).
The template expects `query` and `document` as lists of content parts,
where each part has a `type` field (text, image, video) and corresponding data.
"""
env = _get_jinja_env()
template = env.from_string(chat_template)
render_kwargs = {
"query": query,
"document": document,
}
if instruct:
render_kwargs["instruct"] = instruct
return template.render(**render_kwargs)
def _extract_text_from_content(content: RerankContent) -> str:
"""Extract text from multimodal content."""
if isinstance(content, str):
return content
texts = []
for part in content:
if isinstance(part, ChatCompletionMessageContentTextPart):
texts.append(part.text)
elif isinstance(part, dict) and part.get("type") == "text":
texts.append(part.get("text", ""))
return " ".join(texts)
class OpenAIServingRerank(OpenAIServingBase):
"""Handler for /v1/rerank requests"""
def __init__(self, tokenizer_manager, template_manager=None):
super().__init__(tokenizer_manager)
# TemplateManager is optional; rerank uses tokenizer.chat_template today.
# Keeping this explicit makes the dependency clear and supports future extensions.
self.template_manager = template_manager
# Cache yes/no token IDs for Qwen3 reranker scoring
self._yes_token_id, self._no_token_id = _get_yes_no_token_ids(
tokenizer_manager.tokenizer
)
# NOTE: /v1/rerank is not an official OpenAI endpoint. This module may be moved
# to another module in the future.
def _request_id_prefix(self) -> str:
return "rerank-"
def _validate_request(self, request: V1RerankReqInput) -> Optional[str]:
"""Validate rerank request format and content"""
if not request.query:
return "Query cannot be empty"
if isinstance(request.query, str):
if not request.query.strip():
return "Query cannot be empty or whitespace only"
if not request.documents:
return "Documents cannot be empty"
for doc in request.documents:
if not doc:
return "Each document must be a non-empty string"
if isinstance(doc, str) and not doc.strip():
return "Each document cannot be empty or whitespace only"
return None
def _convert_to_internal_request(
self,
request: V1RerankReqInput,
raw_request: Request = None,
) -> tuple[Union[EmbeddingReqInput, V1RerankReqInput], V1RerankReqInput]:
"""
Convert OpenAI rerank request to internal format.
- For Qwen3-VL reranker (multimodal decoder-only): keep the request.
- For Qwen3 reranker (text-only decoder-only): keep the request and score via
`tokenizer_manager.score_prompts(...)` in the handler.
- For cross-encoder rerank models: adapt into `EmbeddingReqInput` pairs.
"""
chat_template = self.tokenizer_manager.tokenizer.chat_template
model_path = getattr(self.tokenizer_manager.model_config, "model_path", "")
backend = _detect_rerank_backend(
request=request,
chat_template=chat_template if isinstance(chat_template, str) else None,
model_path=model_path,
)
if backend in ("vl_decoder", "text_decoder"):
return request, request
# Cross-encoder rerank: Create pairs of [query, document] for each document.
# Note: Cross-encoder only supports text-only content
if request.is_multimodal():
# Extract text for cross-encoder (multimodal not supported)
query_text = _extract_text_from_content(request.query)
doc_texts = [_extract_text_from_content(doc) for doc in request.documents]
pairs = [[query_text, doc] for doc in doc_texts]
else:
pairs = [[request.query, doc] for doc in request.documents]
adapted_request = EmbeddingReqInput(text=pairs, is_cross_encoder_request=True)
return adapted_request, request
async def _handle_non_streaming_request(
self,
adapted_request: Union[EmbeddingReqInput, V1RerankReqInput],
request: V1RerankReqInput,
raw_request: Request,
) -> Union[List[RerankResponse], ErrorResponse, ORJSONResponse]:
"""Handle the rerank request"""
chat_template = getattr(self.tokenizer_manager.tokenizer, "chat_template", None)
model_path = getattr(self.tokenizer_manager.model_config, "model_path", "")
rerank_ret = await self._handle_rerank_paths(
request=request,
raw_request=raw_request,
chat_template=chat_template,
model_path=model_path,
)
if rerank_ret is not None:
return rerank_ret
# Default cross-encoder rerank path (existing behavior).
try:
if not isinstance(adapted_request, EmbeddingReqInput):
raise ValueError(
"Invalid rerank request adaptation. "
"If you are serving a decoder-only reranker (e.g., Qwen3-Reranker), "
"please provide the corresponding --chat-template and launch without --is-embedding."
)
ret = await self.tokenizer_manager.generate_request(
adapted_request, raw_request
).__anext__()
except ValueError as e:
return self.create_error_response(str(e))
if not isinstance(ret, list):
ret = [ret]
responses = self._build_rerank_response(ret, request)
return responses
async def _handle_rerank_paths(
self,
*,
request: V1RerankReqInput,
raw_request: Request,
chat_template: Optional[str],
model_path: str,
) -> Optional[Union[List[RerankResponse], ErrorResponse, ORJSONResponse]]:
"""
Handle decoder-only rerank paths (VL/text) and return a response if matched.
Returns None if the request should fall back to cross-encoder rerank.
"""
backend = _detect_rerank_backend(
request=request,
chat_template=chat_template,
model_path=model_path,
)
# Qwen3-VL reranker path (decoder-only scoring with query/document template format)
if backend == "vl_decoder":
return await self._handle_vl_reranker_request(
request, raw_request, chat_template or ""
)
# Qwen3 text-only reranker path (decoder-only scoring).
if backend == "text_decoder":
return await self._handle_text_reranker_request(
request=request,
raw_request=raw_request,
chat_template=chat_template or "",
)
return None
async def _handle_text_reranker_request(
self,
*,
request: V1RerankReqInput,
raw_request: Request,
chat_template: str,
) -> Union[List[RerankResponse], ErrorResponse]:
"""Handle text-only decoder reranker request via score_prompts()."""
# Qwen3 reranker relies on decoder-only logprobs. If the server is launched
# with --is-embedding, model_config.is_generation is typically False and
# logprob scoring is not supported.
if not self.tokenizer_manager.model_config.is_generation:
return self.create_error_response(
"Detected Qwen3 reranker chat template, but the server is not in generation mode. "
"Please relaunch without --is-embedding for Qwen3-Reranker models."
)
try:
prompts = [
_render_jinja_chat_template(
chat_template,
query=request.query,
document=doc,
instruct=getattr(request, "instruct", None),
)
for doc in request.documents
]
result = await self.tokenizer_manager.score_prompts(
prompts,
label_token_ids=[self._yes_token_id, self._no_token_id],
apply_softmax=False,
request=raw_request,
)
scores = [_qwen3_rerank_score(s[0], s[1]) for s in result.scores]
except ValueError as e:
return self.create_error_response(str(e))
except Exception as e:
# Includes template rendering errors from jinja2.
return self.create_error_response(str(e))
responses = self._build_rerank_response(scores, request)
return responses
async def _handle_vl_reranker_request(
self,
request: V1RerankReqInput,
raw_request: Request,
_chat_template: str,
) -> Union[List[RerankResponse], ErrorResponse]:
"""Handle multimodal VL reranker request using chat completion with logprobs."""
if not self.tokenizer_manager.model_config.is_generation:
return self.create_error_response(
"Detected Qwen3-VL reranker, but the server is not in generation mode. "
"Please relaunch without --is-embedding for Qwen3-VL-Reranker models."
)
try:
scores = []
instruct = getattr(request, "instruct", None)
for doc in request.documents:
# Build multimodal content lists and render prompt using jinja template
query_content, doc_content, image_data, video_data = (
self._build_vl_reranker_content(
query=request.query,
document=doc,
)
)
# Render the chat template directly with query/document variables
prompt = _render_vl_jinja_template(
chat_template=_chat_template,
query=query_content,
document=doc_content,
instruct=instruct,
)
# Create generate request with logprobs
gen_request = GenerateReqInput(
text=prompt,
image_data=image_data if image_data else None,
video_data=video_data if video_data else None,
sampling_params={
"max_new_tokens": 1,
"temperature": 0,
},
return_logprob=True,
top_logprobs_num=50, # Get enough logprobs to find yes/no tokens
logprob_start_len=0,
)
# Execute generation request
ret = await self.tokenizer_manager.generate_request(
gen_request, raw_request
).__anext__()
# Extract yes/no probabilities from logprobs
score = self._extract_score_from_logprobs(ret)
scores.append(score)
responses = self._build_rerank_response(scores, request)
return responses
except ValueError as e:
return self.create_error_response(str(e))
except Exception as e:
logger.exception("Error handling VL reranker request")
return self.create_error_response(str(e))
def _build_vl_reranker_content(
self,
query: RerankContent,
document: RerankContent,
) -> tuple[List[Dict[str, Any]], List[Dict[str, Any]], List[str], List[str]]:
"""Build content lists for VL reranker request.
Returns:
Tuple of (query_content, document_content, image_data, video_data)
where query_content and document_content are lists suitable for jinja template.
"""
image_data = []
video_data = []
# Build query content list
query_content = self._content_to_template_list(query, image_data, video_data)
# Build document content list
doc_content = self._content_to_template_list(document, image_data, video_data)
return query_content, doc_content, image_data, video_data
def _content_to_template_list(
self,
content: RerankContent,
image_data: List[str],
video_data: List[str],
) -> List[Dict[str, Any]]:
"""Convert RerankContent to a list format suitable for jinja template."""
result = []
if isinstance(content, str):
result.append({"type": "text", "text": content})
return result
for part in content:
if isinstance(part, ChatCompletionMessageContentTextPart):
result.append({"type": "text", "text": part.text})
elif isinstance(part, ChatCompletionMessageContentImagePart):
if part.image_url:
image_data.append(part.image_url.url)
result.append({"type": "image"})
elif isinstance(part, ChatCompletionMessageContentVideoPart):
if part.video_url:
video_data.append(part.video_url.url)
result.append({"type": "video"})
elif isinstance(part, dict):
part_type = part.get("type")
if part_type == "text":
result.append({"type": "text", "text": part.get("text", "")})
elif part_type == "image_url":
image_url = part.get("image_url", {})
if isinstance(image_url, dict):
url = image_url.get("url")
else:
url = image_url
if url:
image_data.append(url)
result.append({"type": "image"})
elif part_type == "video_url":
video_url = part.get("video_url", {})
if isinstance(video_url, dict):
url = video_url.get("url")
else:
url = video_url
if url:
video_data.append(url)
result.append({"type": "video"})
return result
def _extract_score_from_logprobs(self, ret: Dict[str, Any]) -> float:
"""Extract reranking score from generation response with logprobs."""
# Get logprobs from the response
meta_info = ret.get("meta_info", {})
output_top_logprobs = meta_info.get("output_top_logprobs", [])
# Use output_top_logprobs[0] - the model's prediction for the first generated token
top_logprobs = output_top_logprobs[0] if output_top_logprobs else []
# Find yes and no token probabilities
# Format: list of tuples (logprob, token_id, token_text)
p_yes = 0.0
p_no = 0.0
found_yes = False
found_no = False
for item in top_logprobs:
logprob, token_id = item[0], item[1]
if token_id == self._yes_token_id:
p_yes = math.exp(logprob)
found_yes = True
elif token_id == self._no_token_id:
p_no = math.exp(logprob)
found_no = True
if found_yes and found_no:
break
return _qwen3_rerank_score(p_yes, p_no)
def _build_rerank_response(
self, ret: Union[List[Dict[str, Any]], List[float]], request: V1RerankReqInput
) -> List[RerankResponse]:
"""Build the rerank response from generation results"""
responses = []
for idx, item in enumerate(ret):
if isinstance(item, dict):
score_val = item.get("embedding")
# Some rerank/reward models return scalar score as embedding[0].
if isinstance(score_val, list):
if len(score_val) == 0 or not isinstance(
score_val[0], (int, float)
):
raise ValueError(
f"Invalid embedding score for rerank at index {idx}: {score_val!r}"
)
score_val = float(score_val[0])
responses.append(
RerankResponse(
score=float(score_val),
document=(
request.documents[idx] if request.return_documents else None
),
index=idx,
meta_info=item.get("meta_info"),
)
)
else:
responses.append(
RerankResponse(
score=float(item),
document=(
request.documents[idx] if request.return_documents else None
),
index=idx,
)
)
# When top_n is set, nlargest avoids fully sorting the candidate list
# (O(N log top_n) vs O(N log N)) — meaningful for large rerank batches.
# Validator (V1RerankReqInput.validate_top_n) guarantees top_n >= 1.
if request.top_n is not None:
return heapq.nlargest(request.top_n, responses, key=lambda x: x.score)
responses.sort(key=lambda x: x.score, reverse=True)
return responses
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,102 @@
import logging
from typing import Union
import torch
from fastapi import Request
from fastapi.responses import ORJSONResponse
from sglang.srt.entrypoints.openai.protocol import (
ErrorResponse,
ScoringRequest,
ScoringResponse,
UsageInfo,
)
from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
logger = logging.getLogger(__name__)
class OpenAIServingScore(OpenAIServingBase):
"""Handler for /v1/score requests"""
# NOTE: /v1/rerank is not an official OpenAI endpoint. This module may be moved
# to another module in the future.
def _request_id_prefix(self) -> str:
return "score-"
def _convert_to_internal_request(
self,
request: ScoringRequest,
raw_request: Request = None,
) -> tuple[ScoringRequest, ScoringRequest]:
"""Convert OpenAI scoring request to internal format"""
# For scoring, we pass the request directly as the tokenizer_manager
# has a specialized score_request method that doesn't use GenerateReqInput
return request, request
async def _handle_non_streaming_request(
self,
adapted_request: ScoringRequest,
request: ScoringRequest,
raw_request: Request,
) -> Union[ScoringResponse, ErrorResponse]:
"""Handle the scoring request"""
try:
# query_embed_overrides is [num_replacements][hidden_size] -> List[Tensor]
query_embed_overrides = (
[
torch.tensor(v, dtype=torch.float32)
for v in request.query_embed_overrides
]
if request.query_embed_overrides is not None
else None
)
# item_embed_overrides is [num_items][num_replacements][hidden_size] -> List[Optional[List[Tensor]]]
item_embed_overrides = (
[
(
[torch.tensor(v, dtype=torch.float32) for v in per_item]
if per_item is not None
else None
)
for per_item in request.item_embed_overrides
]
if request.item_embed_overrides is not None
else None
)
result = await self.tokenizer_manager.score_request(
query=request.query,
items=request.items,
label_token_ids=request.label_token_ids,
apply_softmax=request.apply_softmax,
item_first=request.item_first,
embed_override_token_id=request.embed_override_token_id,
query_embed_overrides=query_embed_overrides,
item_embed_overrides=item_embed_overrides,
request=raw_request,
return_pooled_hidden_states=request.return_pooled_hidden_states,
)
phs_as_lists = None
if result.pooled_hidden_states is not None:
phs_as_lists = [
t.tolist() if t is not None else None
for t in result.pooled_hidden_states
]
response = ScoringResponse(
scores=result.scores,
pooled_hidden_states=phs_as_lists,
model=request.model,
usage=UsageInfo(
prompt_tokens=result.prompt_tokens,
total_tokens=result.prompt_tokens,
),
)
return ORJSONResponse(content=response.model_dump())
except ValueError as e:
return self.create_error_response(str(e))
@@ -0,0 +1,189 @@
import logging
from http import HTTPStatus
from typing import List, Optional, Union
from fastapi import Request
from sglang.srt.entrypoints.openai.protocol import (
DetokenizeRequest,
DetokenizeResponse,
ErrorResponse,
TokenizeRequest,
TokenizeResponse,
)
from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
from sglang.srt.entrypoints.openai.serving_chat import OpenAIServingChat
logger = logging.getLogger(__name__)
class OpenAIServingTokenize(OpenAIServingBase):
"""Handler for /v1/tokenize requests"""
def __init__(self, tokenizer_manager, template_manager=None):
super().__init__(tokenizer_manager)
self.chat_serving: Optional[OpenAIServingChat] = (
OpenAIServingChat(tokenizer_manager, template_manager)
if template_manager is not None
else None
)
def _request_id_prefix(self) -> str:
return "tok-"
def _convert_to_internal_request(
self, request: TokenizeRequest, raw_request: Request
) -> tuple[TokenizeRequest, TokenizeRequest]:
return request, request
async def _handle_non_streaming_request(
self,
adapted_request: TokenizeRequest,
request: TokenizeRequest,
raw_request: Request,
) -> Union[TokenizeResponse, ErrorResponse]:
try:
tokenizer = self.tokenizer_manager.tokenizer
max_model_len = getattr(tokenizer, "model_max_length", -1)
if request.messages is not None:
token_ids = self._tokenize_chat_request(request)
tokens = token_ids
count = len(token_ids)
elif isinstance(request.prompt, str):
token_ids = tokenizer.encode(
request.prompt,
add_special_tokens=request.add_special_tokens,
)
tokens = token_ids
count = len(token_ids)
elif isinstance(request.prompt, list):
token_ids_list = [
tokenizer.encode(
text, add_special_tokens=request.add_special_tokens
)
for text in request.prompt
]
tokens = token_ids_list
count = [len(ids) for ids in token_ids_list]
else:
return self.create_error_response(
f"Invalid prompt type: {type(request.prompt)}. Expected str or List[str]."
)
return TokenizeResponse(
tokens=tokens, count=count, max_model_len=max_model_len
)
except ValueError as e:
return self.create_error_response(str(e))
except Exception as e:
logger.error("Error during tokenization", exc_info=True)
return self.create_error_response(
f"Internal server error during tokenization: {e}",
err_type="InternalServerError",
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
)
def _tokenize_chat_request(self, request: TokenizeRequest) -> List[int]:
if self.chat_serving is None:
raise ValueError("Chat template tokenization requires a template manager.")
chat_request = request.to_chat_completion_request()
validation_error = self.chat_serving._validate_request(chat_request)
if validation_error:
raise ValueError(validation_error)
is_multimodal = self.tokenizer_manager.model_config.is_multimodal
processed_messages = self.chat_serving._process_messages(
chat_request, is_multimodal
)
prompt_ids = processed_messages.prompt_ids
if isinstance(prompt_ids, list) and (
prompt_ids or not processed_messages.prompt
):
return prompt_ids
if isinstance(prompt_ids, str):
return self.tokenizer_manager.tokenizer.encode(
prompt_ids, add_special_tokens=False
)
if processed_messages.prompt:
return self.tokenizer_manager.tokenizer.encode(
processed_messages.prompt, add_special_tokens=False
)
raise ValueError("Failed to render chat messages into token ids.")
class OpenAIServingDetokenize(OpenAIServingBase):
"""Handler for /v1/detokenize requests"""
def _request_id_prefix(self) -> str:
return "detok-"
def _convert_to_internal_request(
self, request: DetokenizeRequest, raw_request: Request
) -> tuple[DetokenizeRequest, DetokenizeRequest]:
return request, request
async def _handle_non_streaming_request(
self,
adapted_request: DetokenizeRequest,
request: DetokenizeRequest,
raw_request: Request,
) -> Union[DetokenizeResponse, ErrorResponse]:
try:
tokenizer = self.tokenizer_manager.tokenizer
if (
isinstance(request.tokens, list)
and request.tokens
and isinstance(request.tokens[0], int)
):
if not all(isinstance(t, int) for t in request.tokens):
return self.create_error_response(
"Invalid input: 'tokens' must be a list of integers."
)
tokens_to_decode = [int(t) for t in request.tokens]
text = tokenizer.decode(
tokens_to_decode, skip_special_tokens=request.skip_special_tokens
)
text_out: Union[str, List[str]] = text
elif (
isinstance(request.tokens, list)
and request.tokens
and isinstance(request.tokens[0], list)
):
texts: List[str] = []
for token_list in request.tokens:
if not all(isinstance(t, int) for t in token_list):
return self.create_error_response(
f"Invalid input: Sublist in 'tokens' must contain only integers. Found: {token_list}"
)
decoded_text = tokenizer.decode(
[int(t) for t in token_list],
skip_special_tokens=request.skip_special_tokens,
)
texts.append(decoded_text)
text_out = texts
elif isinstance(request.tokens, list) and not request.tokens:
text_out = ""
else:
return self.create_error_response(
f"Invalid tokens type: {type(request.tokens)}. Expected List[int] or List[List[int]]."
)
return DetokenizeResponse(text=text_out)
except Exception as e:
logger.error("Error during detokenization", exc_info=True)
if "decode" in str(e).lower():
return self.create_error_response(
f"Error decoding tokens: {e}. Input tokens might be invalid for the model.",
err_type="DecodeError",
status_code=HTTPStatus.BAD_REQUEST,
)
return self.create_error_response(
f"Internal server error during detokenization: {e}",
err_type="InternalServerError",
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
)
@@ -0,0 +1,466 @@
# Copyright 2025 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""
OpenAI-compatible transcription endpoint handler for audio ASR models.
New ASR models are supported by subclassing ``TranscriptionAdapter`` and
registering via the ``@register_transcription_adapter`` decorator.
See ``transcription_adapters/`` for built-in implementations.
"""
from __future__ import annotations
import asyncio
import io
import logging
import math
import time
import uuid
from typing import TYPE_CHECKING, AsyncGenerator, List, Optional, Union
from fastapi import Request, WebSocket
from fastapi.responses import ORJSONResponse, Response, StreamingResponse
from sglang.srt.entrypoints.openai.protocol import (
DeltaMessage,
ErrorResponse,
TranscriptionRequest,
TranscriptionResponse,
TranscriptionStreamChoice,
TranscriptionStreamResponse,
TranscriptionUsage,
TranscriptionVerboseResponse,
)
from sglang.srt.entrypoints.openai.realtime import (
handle_realtime_transcription,
)
from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
from sglang.srt.entrypoints.openai.streaming_asr import (
StreamingASRState,
needs_space,
process_asr_chunk,
split_audio_chunks,
)
from sglang.srt.entrypoints.openai.transcription_adapters import resolve_adapter
from sglang.srt.managers.io_struct import GenerateReqInput
if TYPE_CHECKING:
from sglang.srt.managers.tokenizer_manager import TokenizerManager
logger = logging.getLogger(__name__)
class OpenAIServingTranscription(OpenAIServingBase):
"""Handler for /v1/audio/transcriptions requests"""
def __init__(self, tokenizer_manager: TokenizerManager):
super().__init__(tokenizer_manager)
model_config = tokenizer_manager.model_config
self._adapter = resolve_adapter(
getattr(model_config.hf_config, "architectures", [])
)
# Cap concurrent /v1/realtime sessions. The Semaphore is bound to the
# event loop on first acquire (uvicorn's loop in normal serving).
self._session_semaphore = asyncio.Semaphore(
tokenizer_manager.server_args.asr_max_concurrent_sessions
)
def _request_id_prefix(self) -> str:
return "trsc-"
def _validate_request(self, request: TranscriptionRequest) -> Optional[str]:
"""Validate transcription request."""
# Validation is done in the route handler for form data
return None
def _convert_to_internal_request(
self,
request: TranscriptionRequest,
raw_request: Request = None,
) -> tuple[GenerateReqInput, TranscriptionRequest]:
"""Convert transcription request to internal format."""
if getattr(request, "_fused_autodetect", False):
sampling_params = self._adapter.build_fused_autodetect_params(request)
else:
sampling_params = self._adapter.build_sampling_params(request)
adapted_request = GenerateReqInput(
text="", # Empty text — the multimodal processor sets proper decoder/prompt tokens
audio_data=request.audio_data,
sampling_params=sampling_params,
stream=request.stream,
modalities=["audio"],
routing_key=self.extract_routing_key(raw_request),
)
return adapted_request, request
@staticmethod
def _get_audio_duration(audio_data: bytes) -> float:
"""Calculate audio duration in seconds."""
try:
import soundfile as sf
info = sf.info(io.BytesIO(audio_data))
return info.duration
except Exception as e:
logger.warning(f"Could not calculate audio duration: {e}")
return 0.0
async def create_transcription(
self,
audio_data: bytes,
model: str,
language: Optional[str],
response_format: str,
temperature: float,
stream: bool,
raw_request: Request,
timestamp_granularities: Optional[List[str]] = None,
) -> Union[
TranscriptionResponse,
TranscriptionVerboseResponse,
StreamingResponse,
Response,
ORJSONResponse,
]:
"""Main entry point for transcription requests."""
# Calculate audio duration for usage reporting
audio_duration_s = self._get_audio_duration(audio_data)
# When language is not specified and the adapter supports detection,
# use a single fused request: SGLang's structured generation (regex)
# constrains the first 3 decode tokens to the forced prefix while
# allowing free transcription afterwards — one encoder pass, no
# extra round-trip. The adapter picks the regex variant based on
# whether timestamps were requested, so fused covers all four
# combinations of (stream, timestamp_granularities):
# * non-streaming: parse_fused_output strips the prefix and
# scrubs trailing/embedded special tokens.
# * streaming: the handler buffers until the sentinel,
# re-anchors, and scrubs each delta via
# adapter.strip_special_tokens.
# verbose_json segment timing still comes from _parse_segments
# over output_ids, which is unaffected by the string-level scrub.
use_fused = language is None and self._adapter.supports_language_detection
# Build request
request = TranscriptionRequest(
audio_data=audio_data,
model=model,
language=language,
response_format=response_format,
temperature=temperature,
timestamp_granularities=timestamp_granularities,
stream=stream,
audio_duration_s=audio_duration_s,
)
if use_fused:
request._fused_autodetect = True
# Stash the variant alongside the flag so the adapter dispatch in
# parse_fused_output and the build_fused_autodetect_params regex
# selection see the same boolean — and we don't recompute it on
# every cumulative-text snapshot in streaming.
request._fused_ts_variant = bool(timestamp_granularities)
# Use the base class handle_request pattern
return await self.handle_request(request, raw_request)
async def _handle_non_streaming_request(
self,
adapted_request: GenerateReqInput,
request: TranscriptionRequest,
raw_request: Request,
) -> Union[
TranscriptionResponse,
TranscriptionVerboseResponse,
ErrorResponse,
ORJSONResponse,
Response,
]:
"""Handle non-streaming transcription request."""
try:
ret = await self.tokenizer_manager.generate_request(
adapted_request, raw_request
).__anext__()
except ValueError as e:
return self.create_error_response(str(e))
text = self._adapter.postprocess_text(ret.get("text", ""))
# For fused auto-detect, parse_fused_output returns the scrubbed
# user-visible text. On parse failure (FSM abort, truncation) it
# returns (None, None) and we fall back to strip_special_tokens —
# the language stays unset rather than reporting a bogus detection.
if getattr(request, "_fused_autodetect", False):
lang, visible = self._adapter.parse_fused_output(
text, ts_variant=getattr(request, "_fused_ts_variant", False)
)
if visible is None:
logger.warning(
"Fused auto-detect parse failed on non-streaming response; "
"falling back to raw-text scrub."
)
text = self._adapter.strip_special_tokens(text)
else:
text = visible
if lang is not None:
request.language = lang
logger.info("Auto-detected language: '%s'", lang)
usage = TranscriptionUsage(seconds=int(math.ceil(request.audio_duration_s)))
# Build response based on format
if request.response_format == "text":
return Response(content=text, media_type="text/plain")
if request.response_format == "verbose_json":
tokenizer = self.tokenizer_manager.tokenizer
return self._adapter.build_verbose_response(
request, text, ret, tokenizer, usage
)
# Default JSON format
return TranscriptionResponse(text=text, usage=usage)
async def _handle_streaming_request(
self,
adapted_request: GenerateReqInput,
request: TranscriptionRequest,
raw_request: Request,
) -> StreamingResponse:
"""Handle streaming transcription request."""
if self._adapter.supports_chunked_streaming:
# No background abort_task: each chunk is a separate request;
# client disconnection is detected via is_disconnected() in the loop.
return StreamingResponse(
self._generate_chunked_asr_stream(
adapted_request, request, raw_request
),
media_type="text/event-stream",
)
return StreamingResponse(
self._generate_transcription_stream(adapted_request, request, raw_request),
media_type="text/event-stream",
background=self.tokenizer_manager.create_abort_task(adapted_request),
)
async def _generate_transcription_stream(
self,
adapted_request: GenerateReqInput,
request: TranscriptionRequest,
raw_request: Request,
) -> AsyncGenerator[str, None]:
"""Generate streaming transcription response.
In fused auto-detect mode, each cumulative-text snapshot is passed
through ``parse_fused_output`` — which returns ``(None, None)``
while the forced prefix is still arriving and ``(lang, visible)``
once it's in. ``visible`` is already stripped of the prefix and
scrubbed of embedded special tokens, and it grows monotonically
across snapshots, so deltas are a plain suffix slice.
"""
created_time = int(time.time())
request_id = f"{self._request_id_prefix()}{uuid.uuid4().hex}"
model = request.model
visible_buffer = ""
fused_mode = getattr(request, "_fused_autodetect", False)
ts_variant = getattr(request, "_fused_ts_variant", False)
# When ``incremental_streaming_output`` is enabled, each chunk's
# ``content["text"]`` is the new delta from the detokenizer, not
# the cumulative text. Always reconstruct cumulative text locally
# so the rest of the loop (prefix parse + visible-buffer slice)
# works uniformly under either mode.
incremental = getattr(
self.tokenizer_manager.server_args,
"incremental_streaming_output",
False,
)
cumulative_text = ""
try:
async for content in self.tokenizer_manager.generate_request(
adapted_request, raw_request
):
finish_reason = content["meta_info"]["finish_reason"]
finish_reason_type = finish_reason["type"] if finish_reason else None
chunk_text = content.get("text", "")
if incremental:
cumulative_text += chunk_text
else:
cumulative_text = chunk_text
if fused_mode:
lang, visible = self._adapter.parse_fused_output(
cumulative_text, ts_variant=ts_variant
)
if visible is None:
# Prefix not yet locatable. Keep buffering until the
# stream ends.
if not finish_reason_type:
continue
# Stream ended before the forced prefix was parseable —
# emit an SSE error frame so the client can distinguish
# this from "silent audio, zero transcription" and raise
# a real error instead of quietly succeeding.
logger.warning(
"Fused auto-detect stream finished before prefix "
"was parseable; returning detection-failed error."
)
error = self.create_streaming_error_response(
"language auto-detect failed: forced-prefix sentinel "
"was not produced before stream end"
)
yield f"data: {error}\n\n"
yield "data: [DONE]\n\n"
return
if lang is not None and request.language is None:
request.language = lang
logger.info("Auto-detected language: '%s'", lang)
else:
visible = cumulative_text
delta = visible[len(visible_buffer) :]
visible_buffer = visible
# Send content delta if there's new text
if delta:
choice_data = TranscriptionStreamChoice(
delta=DeltaMessage(content=delta),
finish_reason=None,
)
chunk = TranscriptionStreamResponse(
id=request_id,
created=created_time,
model=model,
choices=[choice_data],
)
yield f"data: {chunk.model_dump_json()}\n\n"
# Send finish reason when done
if finish_reason_type:
choice_data = TranscriptionStreamChoice(
delta=DeltaMessage(),
finish_reason=finish_reason_type,
)
chunk = TranscriptionStreamResponse(
id=request_id,
created=created_time,
model=model,
choices=[choice_data],
)
yield f"data: {chunk.model_dump_json()}\n\n"
except ValueError as e:
error = self.create_streaming_error_response(str(e))
yield f"data: {error}\n\n"
yield "data: [DONE]\n\n"
async def _generate_chunked_asr_stream(
self,
adapted_request: GenerateReqInput,
request: TranscriptionRequest,
raw_request: Request,
) -> AsyncGenerator[str, None]:
"""Chunk-based streaming for ASR with prefix rollback.
Audio is split into chunks and each chunk is processed as an
independent request. Partial transcripts are emitted via SSE
with prefix rollback to reduce boundary jitter.
TODO:
- Token-level streaming within chunks (stream=True)
- Encoder window caching across chunks
- Cross-chunk KV cache reuse
"""
created_time = int(time.time())
request_id = f"{self._request_id_prefix()}{uuid.uuid4().hex}"
model = request.model
state = StreamingASRState(**self._adapter.chunked_streaming_config)
# Track only the trailing char of the cumulative emit; `needs_space`
# uses prev[-1] / cur[0] so we don't need to keep the full buffer.
last_char = ""
try:
chunks = split_audio_chunks(request.audio_data, state.chunk_size_sec)
for i, chunk_audio in enumerate(chunks):
if await raw_request.is_disconnected():
logger.info("[streaming_asr] client disconnected, stopping")
break
is_last = i == len(chunks) - 1
delta = await process_asr_chunk(
tokenizer_manager=self.tokenizer_manager,
adapter=self._adapter,
state=state,
audio_data=chunk_audio,
sampling_params=adapted_request.sampling_params,
is_last=is_last,
raw_request=raw_request,
routing_key=self.extract_routing_key(raw_request),
)
if delta:
for word in delta.split(" "):
if not word:
continue
content = f" {word}" if needs_space(last_char, word) else word
last_char = content[-1]
chunk_resp = TranscriptionStreamResponse(
id=request_id,
created=created_time,
model=model,
choices=[
TranscriptionStreamChoice(
delta=DeltaMessage(content=content),
finish_reason=None,
)
],
)
yield f"data: {chunk_resp.model_dump_json()}\n\n"
# Send final stop
chunk_resp = TranscriptionStreamResponse(
id=request_id,
created=created_time,
model=model,
choices=[
TranscriptionStreamChoice(
delta=DeltaMessage(),
finish_reason="stop",
)
],
)
yield f"data: {chunk_resp.model_dump_json()}\n\n"
except asyncio.CancelledError:
raise
except Exception as e:
logger.exception("[streaming_asr] unrecoverable error")
error = self.create_streaming_error_response(str(e))
yield f"data: {error}\n\n"
yield "data: [DONE]\n\n"
async def handle_websocket(self, websocket: WebSocket) -> None:
await handle_realtime_transcription(
websocket,
tokenizer_manager=self.tokenizer_manager,
adapter=self._adapter,
server_args=self.tokenizer_manager.server_args,
session_semaphore=self._session_semaphore,
)
@@ -0,0 +1,99 @@
"""SSE chunk building utilities for OpenAI chat completions streaming."""
from __future__ import annotations
from typing import List, Optional, Union
import msgspec
_SSE_DATA_B = b"data: "
_SSE_NL_B = b"\n\n"
class StreamDelta(msgspec.Struct, omit_defaults=True):
"""Delta content for streaming responses.
OpenAI Python SDK's ChoiceDelta does not declare reasoning_content; it is
surfaced via pydantic `extra`. With omit_defaults=True, defaulting to
None would drop the key entirely from the SSE payload, making
`data.reasoning_content` raise AttributeError on the client. Keep it
required (no default) so it is always serialized as null or a string.
"""
reasoning_content: Optional[str]
role: Optional[str] = None
content: Optional[str] = None
class StreamChoice(msgspec.Struct):
"""A single choice in a streaming response."""
index: int
delta: StreamDelta
logprobs: Optional[dict] = None
finish_reason: Optional[str] = None
matched_stop: Union[None, int, str] = None
class StreamChunk(msgspec.Struct, omit_defaults=True):
"""A complete streaming chunk."""
id: str
object: str
created: int
model: str
choices: List[StreamChoice]
usage: Optional[dict] = None
_stream_encoder = msgspec.json.Encoder()
def build_sse_content(
chunk_id: str,
created: int,
model: str,
index: int,
role: Optional[str] = None,
content: Optional[str] = None,
reasoning_content: Optional[str] = None,
finish_reason: Optional[str] = None,
logprobs: Optional[dict] = None,
matched_stop: Union[None, int, str] = None,
usage: Optional[dict] = None,
) -> str:
"""Build an SSE chunk string for content/reasoning updates.
Args:
chunk_id: Request ID for this chunk
created: Unix timestamp
model: Model name
index: Choice index
role: Message role (usually "assistant")
content: Text content delta
reasoning_content: Reasoning/thinking content delta
finish_reason: Finish reason if done
logprobs: Log probabilities if requested
matched_stop: Stop token/string that was matched
usage: Token usage statistics
Returns:
SSE-formatted string "data: {...}\\n\\n"
"""
delta = StreamDelta(role=role, content=content, reasoning_content=reasoning_content)
choice = StreamChoice(
index=index,
delta=delta,
logprobs=logprobs,
finish_reason=finish_reason,
matched_stop=matched_stop,
)
chunk = StreamChunk(
id=chunk_id,
object="chat.completion.chunk",
created=created,
model=model,
choices=[choice],
usage=usage,
)
return (_SSE_DATA_B + _stream_encoder.encode(chunk) + _SSE_NL_B).decode()
@@ -0,0 +1,209 @@
import asyncio
import io
import logging
import re
from dataclasses import dataclass
from typing import Any, Dict, List, Optional
import soundfile as sf
from fastapi import Request
from sglang.srt.entrypoints.openai.transcription_adapters.base import (
TranscriptionAdapter,
)
from sglang.srt.managers.io_struct import GenerateReqInput
from sglang.srt.managers.tokenizer_manager import TokenizerManager
logger = logging.getLogger(__name__)
# Collapse whitespace before punctuation so batched-inference token
# boundary jitter (" ," vs ",") doesn't leak into deltas. Covers both
# ASCII punctuation and the CJK / fullwidth equivalents.
_PUNCT_WS_RE = re.compile(r"\s+([,.;:!?,。!?;:、])")
@dataclass
class StreamingASRState:
"""State for chunk-based streaming ASR with prefix rollback.
Parameters are model-specific and should be provided via the
adapter's ``chunked_streaming_config``.
Known limitation: rollback uses str.split() which is ineffective
for CJK languages (no whitespace between words).
TODO: implement token-level rollback to handle all languages
correctly.
"""
chunk_size_sec: float
unfixed_chunk_num: int
unfixed_token_num: int
confirmed_text: str = ""
# Monotonic accumulator; used as prompt prefix so the model sees a
# natural continuation point, not the rolled-back ``confirmed_text``.
emitted_text: str = ""
full_transcript: str = ""
chunk_index: int = 0
def get_prefix_text(self) -> str:
if self.chunk_index < self.unfixed_chunk_num or not self.emitted_text:
return ""
return self.emitted_text
def _record_emit(self, delta: str) -> str:
if delta:
self.emitted_text = (
f"{self.emitted_text} {delta}".strip() if self.emitted_text else delta
)
return delta
def update(self, new_transcript: str) -> str:
old_confirmed = self.confirmed_text
words = new_transcript.split()
if len(words) > self.unfixed_token_num:
self.confirmed_text = " ".join(words[: -self.unfixed_token_num])
else:
self.confirmed_text = ""
self.full_transcript = new_transcript
self.chunk_index += 1
if self.confirmed_text.startswith(old_confirmed):
return self._record_emit(self.confirmed_text[len(old_confirmed) :].strip())
# Model revised earlier text, use word level common prefix to avoid
# re-emitting already-sent content and cutting mid-word.
old_words = old_confirmed.split()
new_words = self.confirmed_text.split()
common_count = 0
for ow, nw in zip(old_words, new_words):
if ow != nw:
break
common_count += 1
return self._record_emit(" ".join(new_words[common_count:]))
def finalize(self) -> str:
confirmed_words = self.confirmed_text.split()
all_words = self.full_transcript.split()
# Use word level common prefix to handle punctuation differences
# between intermediate chunks and the final full transcription.
common_count = 0
for cw, aw in zip(confirmed_words, all_words):
if cw != aw:
break
common_count += 1
self.confirmed_text = self.full_transcript
if common_count == 0 and confirmed_words and all_words:
return self._record_emit(self.full_transcript)
return self._record_emit(" ".join(all_words[common_count:]))
def split_audio_chunks(audio_data: bytes, chunk_size_sec: float) -> List[bytes]:
if not audio_data:
raise ValueError("audio_data is empty")
if chunk_size_sec <= 0:
raise ValueError(f"chunk_size_sec must be positive, got {chunk_size_sec}")
audio_file = io.BytesIO(audio_data)
try:
data, sample_rate = sf.read(audio_file, dtype="float32")
except sf.LibsndfileError as e:
raise ValueError(f"failed to decode audio: {e}") from e
if len(data.shape) > 1:
data = data.mean(axis=1)
chunk_size_samples = int(chunk_size_sec * sample_rate)
total_samples = len(data)
chunks = []
for end in range(
chunk_size_samples, total_samples + chunk_size_samples, chunk_size_samples
):
end = min(end, total_samples)
buf = io.BytesIO()
sf.write(buf, data[:end], sample_rate, format="WAV")
chunks.append(buf.getvalue())
return chunks
def normalize_whitespace(text: str) -> str:
return _PUNCT_WS_RE.sub(r"\1", text)
_NO_SPACE_BEFORE = frozenset(".,!?;:%)]},。!?;:、)】》」』")
_NO_SPACE_AFTER = frozenset("([{(【《「『")
def _is_cjk(c: str) -> bool:
"""Whether char is a CJK-context glyph that doesn't take inter-word
spaces — ideographs, Japanese kana, CJK punctuation, fullwidth forms.
Excludes Hangul / Devanagari / Arabic etc., which are non-ASCII but
space-separated and need the normal boundary space."""
cp = ord(c)
return (
0x3000 <= cp <= 0x303F # CJK Symbols and Punctuation (,。、《》「」…)
or 0x3040 <= cp <= 0x309F # Hiragana
or 0x30A0 <= cp <= 0x30FF # Katakana
or 0x3400 <= cp <= 0x4DBF # CJK Unified Ideographs Ext A
or 0x4E00 <= cp <= 0x9FFF # CJK Unified Ideographs
or 0xFF00 <= cp <= 0xFFEF # Halfwidth & Fullwidth Forms (fullwidth ASCII)
)
def needs_space(prev: str, cur: str) -> bool:
"""Return whether a boundary space is needed between emitted deltas.
Avoid spaces around punctuation and between adjacent CJK-context glyphs.
Shared by the realtime WS and HTTP SSE chunked streaming paths.
"""
if not prev or not cur:
return False
if prev[-1].isspace() or cur[0].isspace():
return False
if cur[0] in _NO_SPACE_BEFORE or prev[-1] in _NO_SPACE_AFTER:
return False
if _is_cjk(prev[-1]) and _is_cjk(cur[0]):
return False
return True
async def process_asr_chunk(
tokenizer_manager: TokenizerManager,
adapter: TranscriptionAdapter,
state: StreamingASRState,
audio_data: bytes,
sampling_params: Dict[str, Any],
is_last: bool,
raw_request: Optional[Request] = None,
routing_key: Optional[str] = None,
) -> str:
"""Run inference on one audio chunk. Shared by the HTTP and WebSocket paths."""
prompt = adapter.prompt_template + state.get_prefix_text()
chunk_request = GenerateReqInput(
text=prompt,
audio_data=audio_data,
sampling_params=sampling_params,
stream=False,
modalities=["audio"],
)
if routing_key is not None:
chunk_request.routing_key = routing_key
try:
ret = None
async for ret in tokenizer_manager.generate_request(chunk_request, raw_request):
break
except asyncio.CancelledError:
raise
except ValueError:
logger.warning(
"[streaming_asr] chunk %d failed", state.chunk_index, exc_info=True
)
raise
if ret is None:
logger.warning("[streaming_asr] empty response for chunk %d", state.chunk_index)
return ""
text = normalize_whitespace(adapter.postprocess_text(ret.get("text", "")))
if is_last:
state.full_transcript = text
return state.finalize()
return state.update(text)
@@ -0,0 +1,191 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import logging
from abc import ABC, abstractmethod
from contextlib import AbstractAsyncContextManager, asynccontextmanager
from typing import Any
try:
from mcp import ClientSession
from mcp.client.sse import sse_client
from mcp.types import ListToolsResult
except ImportError as e:
ClientSession = sse_client = ListToolsResult = e
from openai_harmony import ToolDescription, ToolNamespaceConfig
logger = logging.getLogger(__name__)
async def list_server_and_tools(server_url: str):
async with (
sse_client(url=server_url) as streams,
ClientSession(*streams) as session,
):
initialize_response = await session.initialize()
list_tools_response = await session.list_tools()
return initialize_response, list_tools_response
def trim_schema(schema: dict) -> dict:
# Turn JSON Schema from MCP generated into Harmony's variant.
if "title" in schema:
del schema["title"]
if "default" in schema and schema["default"] is None:
del schema["default"]
if "anyOf" in schema:
# Turn "anyOf": [{"type": "type-1"}, {"type": "type-2"}]
# into "type": ["type-1", "type-2"]
# if there's more than 1 types, also remove "null" type as Harmony will
# just ignore it
types = [
type_dict["type"]
for type_dict in schema["anyOf"]
if type_dict["type"] != "null"
]
schema["type"] = types
del schema["anyOf"]
if "properties" in schema:
schema["properties"] = {
k: trim_schema(v) for k, v in schema["properties"].items()
}
return schema
def post_process_tools_description(
list_tools_result: "ListToolsResult",
) -> "ListToolsResult":
# Adapt the MCP tool result for Harmony
for tool in list_tools_result.tools:
tool.inputSchema = trim_schema(tool.inputSchema)
# Some tools schema don't need to be part of the prompt (e.g. simple text
# in text out for Python)
list_tools_result.tools = [
tool
for tool in list_tools_result.tools
if getattr(tool.annotations, "include_in_prompt", True)
]
return list_tools_result
class ToolServer(ABC):
@abstractmethod
def has_tool(self, tool_name: str):
pass
@abstractmethod
def get_tool_description(self, tool_name: str):
pass
@abstractmethod
def get_tool_session(self, tool_name: str) -> AbstractAsyncContextManager[Any]: ...
class MCPToolServer(ToolServer):
def __init__(self):
self.harmony_tool_descriptions = {}
async def add_tool_server(self, server_url: str):
tool_urls = server_url.split(",")
self.harmony_tool_descriptions = {}
self.urls: dict[str, str] = {}
for url in tool_urls:
url = f"http://{url}/sse"
initialize_response, list_tools_response = await list_server_and_tools(url)
list_tools_response = post_process_tools_description(list_tools_response)
tool_from_mcp = ToolNamespaceConfig(
name=initialize_response.serverInfo.name,
description=initialize_response.instructions,
tools=[
ToolDescription.new(
name=tool.name,
description=tool.description,
parameters=tool.inputSchema,
)
for tool in list_tools_response.tools
],
)
self.harmony_tool_descriptions[tool_from_mcp.name] = tool_from_mcp
if tool_from_mcp.name not in self.urls:
self.urls[tool_from_mcp.name] = url
else:
logger.warning(
"Tool %s already exists. Ignoring duplicate tool server %s",
tool_from_mcp.name,
url,
)
def has_tool(self, tool_name: str):
return tool_name in self.harmony_tool_descriptions
def get_tool_description(self, tool_name: str):
return self.harmony_tool_descriptions.get(tool_name)
@asynccontextmanager
async def get_tool_session(self, tool_name: str):
url = self.urls.get(tool_name)
if url:
async with (
sse_client(url=url) as streams,
ClientSession(*streams) as session,
):
await session.initialize()
yield session
else:
logger.warning("Tool %s not found", tool_name)
class DemoToolServer(ToolServer):
def __init__(self, *, enable_python: bool = True):
from sglang.srt.entrypoints.tool import (
HarmonyBrowserTool,
HarmonyPythonTool,
Tool,
)
self.tools: dict[str, Tool] = {}
browser_tool = HarmonyBrowserTool()
if browser_tool.enabled:
self.tools["browser"] = browser_tool
if enable_python:
python_tool = HarmonyPythonTool()
if python_tool.enabled:
self.tools["python"] = python_tool
def has_tool(self, tool_name: str):
return tool_name in self.tools
def get_tool_description(self, tool_name: str):
if tool_name not in self.tools:
return None
if tool_name == "browser":
return ToolNamespaceConfig.browser()
elif tool_name == "python":
return ToolNamespaceConfig.python()
else:
raise ValueError(f"Unknown tool {tool_name}")
@asynccontextmanager
async def get_tool_session(self, tool_name: str):
yield self.tools[tool_name]
async def aclose(self):
browser = self.tools.get("browser")
exa_client = getattr(browser, "exa_client", None) if browser else None
if exa_client is not None:
await exa_client.close()
class NativeToolServer(DemoToolServer):
"""Built-in SGLang hosted tools that do not require an external MCP server."""
def __init__(self):
super().__init__(enable_python=False)
@@ -0,0 +1,27 @@
# Re-export the public API from base so callers can do:
# from ...transcription_adapters import TranscriptionAdapter, register_transcription_adapter
from sglang.srt.entrypoints.openai.transcription_adapters.base import ( # noqa: F401
TranscriptionAdapter,
register_transcription_adapter,
resolve_adapter,
)
# Import built-in adapters so they self-register via @register_transcription_adapter.
from sglang.srt.entrypoints.openai.transcription_adapters.mimo_v2_asr import ( # noqa: F401
MiMoV2ASRAdapter,
)
from sglang.srt.entrypoints.openai.transcription_adapters.qwen3_asr import ( # noqa: F401
Qwen3ASRAdapter,
)
from sglang.srt.entrypoints.openai.transcription_adapters.whisper import ( # noqa: F401
WhisperAdapter,
)
__all__ = [
"TranscriptionAdapter",
"register_transcription_adapter",
"resolve_adapter",
"WhisperAdapter",
"Qwen3ASRAdapter",
"MiMoV2ASRAdapter",
]
@@ -0,0 +1,162 @@
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import List, Optional
from sglang.srt.entrypoints.openai.protocol import (
TranscriptionRequest,
TranscriptionUsage,
TranscriptionVerboseResponse,
)
class TranscriptionAdapter(ABC):
"""Abstract base for model-specific transcription logic.
Subclass this and decorate with ``@register_transcription_adapter("Key")``
to add support for a new ASR model. See the sibling modules for
the built-in Whisper and Qwen3-ASR implementations.
"""
@abstractmethod
def build_sampling_params(self, request: TranscriptionRequest) -> dict:
"""Return the ``sampling_params`` dict for ``GenerateReqInput``."""
@property
def supports_language_detection(self) -> bool:
"""Whether this model supports automatic language detection.
When True, the adapter must implement the fused autodetect methods
and the standalone detection methods below.
"""
return False
# -- Fused detect+transcribe (used by the server) ----------------------
def build_fused_autodetect_params(self, request) -> dict:
"""Return ``sampling_params`` dict for a fused detect+transcribe request.
Uses structured generation (``regex``) to constrain the output prefix
to a valid language + task token sequence while allowing free
transcription afterwards — all in a single request.
"""
raise NotImplementedError
@staticmethod
def parse_fused_output(
text: str, *, ts_variant: bool = False
) -> tuple[Optional[str], Optional[str]]:
"""Parse the fused output into ``(language_code, user_visible_text)``.
Called by both streaming and non-streaming handlers with the same
contract. ``ts_variant`` indicates which forced-prefix shape was
requested (the caller knows from ``request.timestamp_granularities``);
adapters use it to disambiguate variants whose detokenized prefix
differs in shape from their token-id prefix.
* ``(None, None)`` — the forced prefix is not yet locatable.
Streaming callers keep buffering; non-streaming / end-of-stream
callers treat this as a parse failure and fall back to
``strip_special_tokens`` on the raw text.
* ``(lang, visible)`` — prefix parsed. ``visible`` is fully
user-visible (prefix removed, embedded special tokens scrubbed).
It must grow monotonically across cumulative streaming snapshots
so callers can compute deltas against it directly.
"""
raise NotImplementedError
@staticmethod
def strip_special_tokens(text: str) -> str:
"""Best-effort scrub of model-specific special-token strings.
Used as a fallback when ``parse_fused_output`` reports a parse
failure (e.g. FSM abort). Default is an identity pass-through;
adapters that request generation with ``skip_special_tokens=False``
should override to strip their special-token syntax.
"""
return text
@property
def supports_chunked_streaming(self) -> bool:
"""Whether this model uses chunk-based streaming instead of token-level streaming."""
return False
@property
def model_sample_rate(self) -> int:
"""Target sample rate in Hz the model expects. Realtime WS path
resamples client PCM to this rate before chunking. Default 16000
matches Whisper / Qwen3-ASR; override for models expecting other rates.
"""
return 16000
@property
def prompt_template(self) -> str:
"""Prompt template for chunked streaming requests.
Only used when ``supports_chunked_streaming`` is True.
The default returns an empty string.
"""
return ""
@property
def chunked_streaming_config(self) -> dict:
"""Parameters for ``StreamingASRState`` when using chunked streaming.
Only used when ``supports_chunked_streaming`` is True.
Keys: ``chunk_size_sec``, ``unfixed_chunk_num``, ``unfixed_token_num``.
"""
return {}
def postprocess_text(self, text: str) -> str:
"""Strip model-specific markers from raw decoded text.
The default implementation is a no-op pass-through.
"""
return text
@abstractmethod
def build_verbose_response(
self,
request: TranscriptionRequest,
text: str,
ret: dict,
tokenizer,
usage: TranscriptionUsage,
) -> TranscriptionVerboseResponse:
"""Build a ``verbose_json`` response with segments / timestamps."""
_ADAPTER_REGISTRY: dict[str, type[TranscriptionAdapter]] = {}
_DEFAULT_ADAPTER_KEY = "Whisper"
def register_transcription_adapter(
key: str,
) -> callable:
"""Class decorator that registers a ``TranscriptionAdapter`` subclass.
*key* is matched as a substring against the model's HF ``architectures``
list at init time (e.g. ``"Whisper"`` matches
``"WhisperForConditionalGeneration"``).
"""
def decorator(cls: type[TranscriptionAdapter]) -> type[TranscriptionAdapter]:
_ADAPTER_REGISTRY[key] = cls
return cls
return decorator
def resolve_adapter(architectures: List[str]) -> TranscriptionAdapter:
"""Pick the right adapter by matching architecture names against the registry."""
for arch in architectures or []:
for key, adapter_cls in _ADAPTER_REGISTRY.items():
if key in arch:
return adapter_cls()
default_cls = _ADAPTER_REGISTRY.get(_DEFAULT_ADAPTER_KEY)
if default_cls is None:
raise RuntimeError(
"No transcription adapters registered. "
"Make sure 'transcription_adapters' package is importable."
)
return default_cls()
@@ -0,0 +1,46 @@
from __future__ import annotations
from sglang.srt.entrypoints.openai.protocol import (
TranscriptionRequest,
TranscriptionUsage,
TranscriptionVerboseResponse,
)
from sglang.srt.entrypoints.openai.transcription_adapters.base import (
TranscriptionAdapter,
register_transcription_adapter,
)
@register_transcription_adapter("MiMoV2ASR")
class MiMoV2ASRAdapter(TranscriptionAdapter):
"""Adapter for MiMo-V2-ASR.
The multimodal processor (``MiMoV2ASRProcessor``) prepends the audio
placeholder ``<|sosp|><|empty|>...<|eosp|>`` when ``input_text`` lacks
one, so the request text can stay empty and the adapter only has to
supply sampling params and the verbose-response shape.
"""
def build_sampling_params(self, request: TranscriptionRequest) -> dict:
return {
"temperature": request.temperature,
"max_new_tokens": 448,
}
def build_verbose_response(
self,
request: TranscriptionRequest,
text: str,
ret: dict,
tokenizer,
usage: TranscriptionUsage,
) -> TranscriptionVerboseResponse:
# MiMo-V2-ASR does not emit timestamp tokens; segments stay empty
# until a forced-aligner path is added.
return TranscriptionVerboseResponse(
language=request.language or "auto",
duration=round(request.audio_duration_s, 2),
text=text,
segments=[],
usage=usage,
)
@@ -0,0 +1,69 @@
from __future__ import annotations
from sglang.srt.entrypoints.openai.protocol import (
TranscriptionRequest,
TranscriptionUsage,
TranscriptionVerboseResponse,
)
from sglang.srt.entrypoints.openai.transcription_adapters.base import (
TranscriptionAdapter,
register_transcription_adapter,
)
from sglang.srt.multimodal.processors.qwen3_asr import DEFAULT_ASR_PROMPT
@register_transcription_adapter("Qwen3ASR")
class Qwen3ASRAdapter(TranscriptionAdapter):
ASR_TEXT_TAG = "<asr_text>"
@property
def supports_chunked_streaming(self) -> bool:
return True
@property
def chunked_streaming_config(self) -> dict:
# Qwen3-ASR paper (arXiv:2601.21337), Table 8 uses 4 unfixed chunks.
# We use 2 here for lower latency; tune based on quality needs.
# TODO: allow users to override these via API request parameters.
return {
"chunk_size_sec": 2.0,
"unfixed_chunk_num": 2,
"unfixed_token_num": 5,
}
@property
def prompt_template(self) -> str:
return DEFAULT_ASR_PROMPT
def build_sampling_params(self, request: TranscriptionRequest) -> dict:
temperature = request.temperature
if temperature == 0.0:
temperature = 0.01 # Qwen3-ASR recommended near-greedy temperature
return {
"temperature": temperature,
"max_new_tokens": 256, # Qwen3-ASR default
}
def postprocess_text(self, text: str) -> str:
# Qwen3-ASR outputs "language <lang><asr_text>transcription" format;
# strip the prefix to return clean transcription text.
if self.ASR_TEXT_TAG in text:
return text.split(self.ASR_TEXT_TAG, 1)[-1]
return text
def build_verbose_response(
self,
request: TranscriptionRequest,
text: str,
ret: dict,
tokenizer,
usage: TranscriptionUsage,
) -> TranscriptionVerboseResponse:
# TODO: Qwen3-ASR needs ForcedAligner to produce timestamps
return TranscriptionVerboseResponse(
language=request.language or "auto",
duration=round(request.audio_duration_s, 2),
text=text,
segments=[],
usage=usage,
)
@@ -0,0 +1,330 @@
from __future__ import annotations
import logging
import re
from typing import List, Optional
from transformers.models.whisper.tokenization_whisper import LANGUAGES
from sglang.srt.entrypoints.openai.protocol import (
TranscriptionRequest,
TranscriptionSegment,
TranscriptionUsage,
TranscriptionVerboseResponse,
)
from sglang.srt.entrypoints.openai.transcription_adapters.base import (
TranscriptionAdapter,
register_transcription_adapter,
)
logger = logging.getLogger(__name__)
# Sampling-params key the adapter plants and the multimodal processor pops
# to flip the decoder prompt from the explicit 4-token forced sequence to
# the bare ``<|startoftranscript|>`` (so the FSM regex drives token 1-3
# instead). Centralized so adapter / processor / warmup all reference the
# same string.
FUSED_AUTODETECT_FLAG = "_detect_language"
# The complete set of Whisper language tokens as they appear in the tokenizer
# vocab (<|xx|> / <|xxx|>). Sourced from the upstream ``LANGUAGES`` dict in
# ``transformers.models.whisper.tokenization_whisper`` so newly-added tokens
# (e.g. ``yue`` in Whisper v3) automatically propagate.
#
# Intentionally wider than ``processors.whisper.ISO639_1_SUPPORTED_LANGS``
# (the narrower input-validation set used by ``normalize_language_to_code``)
# — for the FSM regex we want every language the model was trained on so we
# don't silently force a wrong nearest-match code on audio in languages the
# model *can* detect but the input dict doesn't list (yue/Cantonese,
# jw/Javanese, haw/Hawaiian, ba/Bashkir, su/Sundanese, ...). Codes whose
# ``<|xxx|>`` token isn't in an older checkpoint's vocab are harmless —
# xgrammar simply leaves that regex branch with no admissible tokens.
WHISPER_LANG_TOKEN_CODES: frozenset[str] = frozenset(LANGUAGES.keys())
# Two forced-prefix variants, picked at request build time based on whether
# the client asked for timestamp_granularities:
# * notimestamps variant: <|lang|><|transcribe|><|notimestamps|> text...
# — drops segment/word timing, used when the client doesn't request it.
# * timestamps variant: <|lang|><|transcribe|><|0.00|> text <|X.XX|> ...
# — <|0.00|> anchors the first segment at t=0, and the model naturally
# emits further timestamp tokens between segments. _parse_segments
# reconstructs segments from output_ids afterwards.
# sorted() gives a deterministic regex string so the warmup-compiled FSM is
# reused across server restarts.
_LANG_ALT = "|".join(re.escape(c) for c in sorted(WHISPER_LANG_TOKEN_CODES))
_LANG_PREFIX = r"<\|(" + _LANG_ALT + r")\|>"
WHISPER_AUTODETECT_REGEX = (
_LANG_PREFIX + r"<\|transcribe\|>" + r"<\|notimestamps\|>" + r"[\s\S]*"
)
WHISPER_AUTODETECT_TS_REGEX = (
_LANG_PREFIX + r"<\|transcribe\|>" + r"<\|0\.00\|>" + r"[\s\S]*"
)
# Forced-prefix patterns, one per FSM variant. Each is anchored at start
# and rejects anything missing ``<|transcribe|>`` so a bypassed FSM or a
# mid-stream snapshot can't slip through as a valid detection. The two
# patterns differ in what the third forced token is decoded *as*:
#
# * ``_FUSED_PREFIX_RE_NOTS`` — non-timestamps variant. The third token
# is ``<|notimestamps|>`` (id 50364), which detokenizes to its literal
# string. Mid-stream snapshots stuck at ``<|en|><|transcribe|>`` (the
# third token hasn't fired yet) correctly miss this regex, so the
# streaming handler can detect FSM-abort and surface an error.
#
# * ``_FUSED_PREFIX_RE_TS`` — timestamps variant. The third token is
# ``<|0.00|>`` (id 50365), which Whisper's tokenizer decodes to the
# *empty string* even with ``skip_special_tokens=False`` (only
# ``<|notimestamps|>`` survives detokenization; every ``<|X.XX|>``
# maps to ``""``). So the regex must accept just
# ``<|en|><|transcribe|>`` and rely on the FSM having already
# constrained ``output_ids[2] == 50365``. ``_parse_segments`` reads
# the timestamps from ``output_ids`` directly, so segment timing is
# unaffected.
_FUSED_PREFIX_RE_NOTS = re.compile(
r"^" + _LANG_PREFIX + r"<\|transcribe\|><\|notimestamps\|>"
)
_FUSED_PREFIX_RE_TS = re.compile(r"^" + _LANG_PREFIX + r"<\|transcribe\|>")
# Fixed Whisper control tokens (see transformers.models.whisper vocab).
# <|startoftranscript|> / <|startofprev|> / <|startoflm|> only appear at
# the decoder prompt and never in generated output, but they are cheap to
# include and harmless if they ever leak.
_WHISPER_CONTROL_TOKENS = frozenset(
{
"endoftext",
"startoftranscript",
"startofprev",
"startoflm",
"translate",
"transcribe",
"notimestamps",
"nospeech",
}
)
# Scrubs only actual Whisper special-token literals: language codes
# (WHISPER_LANG_TOKEN_CODES), control tokens (_WHISPER_CONTROL_TOKENS),
# and timestamp tokens (<|X.XX|>, where X.XX matches the
# ``{ts_base + k * 0.02}`` schema the model emits). A broad
# ``<\|[^|]+\|>`` would eat legitimate spoken content on audio that
# pronounces angle-bracket / pipe sequences (e.g. someone reading
# ``<|endoftext|>`` out loud). Used to scrub trailing ``<|endoftext|>``
# and embedded ``<|X.XX|>`` timestamp tokens from the user-visible text
# in fused-autodetect responses, where ``skip_special_tokens=False`` is
# needed to preserve the language prefix for parsing but would otherwise
# leak other special tokens downstream.
_SPECIAL_TOKEN_RE = re.compile(
r"<\|(?:"
+ "|".join(sorted(WHISPER_LANG_TOKEN_CODES | _WHISPER_CONTROL_TOKENS))
+ r"|\d+\.\d{2}"
+ r")\|>"
)
@register_transcription_adapter("Whisper")
class WhisperAdapter(TranscriptionAdapter):
TIMESTAMP_BASE_TOKEN_ID = 50365 # <|0.00|>
TIMESTAMP_BASE_OFFSET = 0.02 # each token step = 0.02 s
def build_sampling_params(self, request: TranscriptionRequest) -> dict:
params: dict = {
"temperature": request.temperature,
"max_new_tokens": 448, # Whisper default max tokens
"language": request.language,
}
if request.timestamp_granularities:
params["timestamp_granularities"] = request.timestamp_granularities
return params
# -- language detection ------------------------------------------------
@property
def supports_language_detection(self) -> bool:
return True
def build_fused_autodetect_params(self, request: TranscriptionRequest) -> dict:
"""Build sampling params for a single fused detect+transcribe request.
Uses SGLang's native structured generation (``regex``) to constrain
the first 3 decode tokens. Picks the regex variant based on whether
the client requested ``timestamp_granularities``:
* no timestamps: ``<|lang|><|transcribe|><|notimestamps|>text``
* with timestamps: ``<|lang|><|transcribe|><|0.00|>text<|X.XX|>...``
— ``<|0.00|>`` anchors segment 0 at t=0; the model naturally
emits further timestamp tokens between segments and
``_parse_segments`` reconstructs them from ``output_ids``.
Either way, detection and transcription run in a single encoder
pass with no extra HTTP round-trip.
"""
ts_variant = bool(request.timestamp_granularities)
params: dict = {
"temperature": request.temperature,
# Fused auto-detect decoder prompt is just <|startoftranscript|>
# (1 token, see processors/whisper.py). Whisper's
# max_target_positions is 448, so max_new_tokens caps at 447:
# 1 prompt + 3 forced prefix + up to 444 free transcription = 448.
"max_new_tokens": 447,
"regex": (
WHISPER_AUTODETECT_TS_REGEX if ts_variant else WHISPER_AUTODETECT_REGEX
),
"skip_special_tokens": False,
# parse_fused_output matches a zero-space forced prefix
# (``<|en|><|transcribe|><|notimestamps|>`` glued together).
# Fast Whisper tokenizers decode adjacent added tokens with no
# space, but slow ones insert a space between them. Force
# spaces_between_special_tokens=False so the parse regex is
# correct regardless of tokenizer variant.
"spaces_between_special_tokens": False,
FUSED_AUTODETECT_FLAG: True,
}
if ts_variant:
params["timestamp_granularities"] = request.timestamp_granularities
return params
@staticmethod
def parse_fused_output(
text: str, *, ts_variant: bool = False
) -> tuple[Optional[str], Optional[str]]:
"""Parse fused output into ``(language_code, user_visible_text)``.
Matches the forced prefix the FSM emits. ``ts_variant`` selects
which shape to expect — the caller knows from
``request.timestamp_granularities`` which regex was sent to the
FSM and so which decoded shape to look for:
* ``ts_variant=False`` — ``<|en|><|transcribe|><|notimestamps|> Hello...``
* ``ts_variant=True`` — ``<|en|><|transcribe|> Hello...`` (``<|0.00|>``
is in ``output_ids`` but Whisper detokenizes it to the empty string).
Return cases:
* ``(None, None)`` — the prefix isn't fully in yet (mid-stream
snapshot before ``<|transcribe|>`` lands, or before
``<|notimestamps|>`` lands in the no-ts variant) or the prefix
is malformed. Streaming callers keep buffering; non-streaming /
end-of-stream callers treat this as a parse failure and fall
back to a best-effort scrub of the raw text.
* ``(lang, visible)`` — prefix fully parsed. ``visible`` is the
transcription with the forced prefix removed, any embedded
special tokens (``<|X.XX|>``, ``<|endoftext|>``) scrubbed, and
surrounding whitespace trimmed. It grows monotonically across
streaming chunks because Whisper's special tokens detokenize
atomically, so callers can compute deltas against it directly.
"""
pattern = _FUSED_PREFIX_RE_TS if ts_variant else _FUSED_PREFIX_RE_NOTS
m = pattern.match(text)
if not m:
return None, None
transcription = text[m.end() :]
# Scrub any remaining special tokens. skip_special_tokens=False is
# set on fused requests so the language prefix survives for
# parsing, but that also preserves trailing <|endoftext|> and, in
# the timestamps variant, embedded <|X.XX|> segment tokens. Those
# are unwanted in the user-visible text (verbose_json gets its
# segments from _parse_segments over output_ids instead).
transcription = _SPECIAL_TOKEN_RE.sub("", transcription)
return m.group(1), transcription.strip()
@staticmethod
def strip_special_tokens(text: str) -> str:
"""Remove any ``<|...|>`` special-token strings from *text*.
Used as the best-effort scrub on FSM abort / parse failure when
the full ``parse_fused_output`` path can't locate the prefix.
"""
return _SPECIAL_TOKEN_RE.sub("", text)
# -- end language detection --------------------------------------------
def build_verbose_response(
self,
request: TranscriptionRequest,
text: str,
ret: dict,
tokenizer,
usage: TranscriptionUsage,
) -> TranscriptionVerboseResponse:
output_ids = ret.get("output_ids", [])
parsed_text, segments = self._parse_segments(output_ids, tokenizer)
return TranscriptionVerboseResponse(
# Pass None through when fused auto-detect failed to parse a
# language — the client should see detection-failed, not a silent
# English default. For explicit-language requests request.language
# is already set by the caller.
language=request.language,
duration=round(request.audio_duration_s, 2),
text=parsed_text or text,
segments=segments,
usage=usage,
)
@staticmethod
def _parse_segments(
output_ids: List[int], tokenizer
) -> tuple[str, List[TranscriptionSegment]]:
"""Parse Whisper timestamp tokens from *output_ids* into segments.
The decoder prompt ends with ``<|0.00|>``, so the first segment starts
at t=0. The model then outputs::
text_tokens <|end_ts|> [<|start_ts|> text_tokens <|end_ts|> ...]
Each timestamp token marks the end of the current segment; its value
also becomes the start of the next segment.
"""
eos_token_id = getattr(tokenizer, "eos_token_id", 50257)
ts_base = WhisperAdapter.TIMESTAMP_BASE_TOKEN_ID
ts_step = WhisperAdapter.TIMESTAMP_BASE_OFFSET
segments: list[TranscriptionSegment] = []
full_text_parts: list[str] = []
current_text_tokens: list[int] = []
current_start = 0.0 # First segment starts at 0.0 (from prompt <|0.00|>)
seg_id = 0
for token_id in output_ids:
if token_id >= ts_base:
timestamp = (token_id - ts_base) * ts_step
if current_text_tokens:
seg_text = tokenizer.decode(
current_text_tokens, skip_special_tokens=True
).strip()
if seg_text:
segments.append(
TranscriptionSegment(
id=seg_id,
start=round(current_start, 2),
end=round(timestamp, 2),
text=seg_text,
)
)
full_text_parts.append(seg_text)
seg_id += 1
current_text_tokens = []
current_start = timestamp
elif token_id == eos_token_id:
continue
else:
current_text_tokens.append(token_id)
if current_text_tokens:
seg_text = tokenizer.decode(
current_text_tokens, skip_special_tokens=True
).strip()
if seg_text:
segments.append(
TranscriptionSegment(
id=seg_id,
start=round(current_start, 2),
end=round(current_start, 2),
text=seg_text,
)
)
full_text_parts.append(seg_text)
return " ".join(full_text_parts), segments
@@ -0,0 +1,126 @@
from __future__ import annotations
from typing import Any, Dict, List, Mapping, Optional, final
from sglang.srt.entrypoints.openai.protocol import PromptTokensDetails, UsageInfo
@final
class UsageProcessor:
"""Stateless helpers that turn raw token counts into a UsageInfo."""
@staticmethod
def _details_if_cached(count: int) -> Optional[PromptTokensDetails]:
"""Return PromptTokensDetails only when count > 0 (keeps JSON slim)."""
return PromptTokensDetails(cached_tokens=count) if count > 0 else None
@staticmethod
def calculate_response_usage(
responses: List[Dict[str, Any]],
n_choices: int = 1,
enable_cache_report: bool = False,
image_tokens: int = 0,
audio_tokens: int = 0,
video_tokens: int = 0,
) -> UsageInfo:
completion_tokens = sum(
r["meta_info"].get("completion_tokens", 0) for r in responses
)
prompt_tokens = sum(
responses[i]["meta_info"].get("prompt_tokens", 0)
for i in range(0, len(responses), n_choices)
)
# some API don't have reasoning_tokens semantics
reasoning_tokens = sum(
r["meta_info"].get("reasoning_tokens", 0) for r in responses
)
cached_details = None
if enable_cache_report:
cached_total = sum(
responses[i]["meta_info"].get("cached_tokens", 0)
for i in range(0, len(responses), n_choices)
)
cached_details = UsageProcessor._details_if_cached(cached_total)
return UsageProcessor.calculate_token_usage(
prompt_tokens=prompt_tokens,
reasoning_tokens=reasoning_tokens,
completion_tokens=completion_tokens,
cached_tokens=cached_details,
image_tokens=image_tokens,
audio_tokens=audio_tokens,
video_tokens=video_tokens,
)
@staticmethod
def calculate_streaming_usage(
prompt_tokens: Mapping[int, int],
reasoning_tokens: Mapping[int, int],
completion_tokens: Mapping[int, int],
cached_tokens: Mapping[int, int],
n_choices: int,
enable_cache_report: bool = False,
image_tokens: int = 0,
audio_tokens: int = 0,
video_tokens: int = 0,
) -> UsageInfo:
# index % n_choices == 0 marks the first choice of a prompt
total_prompt_tokens = sum(
tok for idx, tok in prompt_tokens.items() if idx % n_choices == 0
)
total_reasoning_tokens = sum(reasoning_tokens.values())
total_completion_tokens = sum(completion_tokens.values())
cached_details = (
UsageProcessor._details_if_cached(
sum(tok for idx, tok in cached_tokens.items() if idx % n_choices == 0)
)
if enable_cache_report
else None
)
return UsageProcessor.calculate_token_usage(
prompt_tokens=total_prompt_tokens,
reasoning_tokens=total_reasoning_tokens,
completion_tokens=total_completion_tokens,
cached_tokens=cached_details,
image_tokens=image_tokens,
audio_tokens=audio_tokens,
video_tokens=video_tokens,
)
@staticmethod
def calculate_token_usage(
prompt_tokens: int,
completion_tokens: int,
reasoning_tokens: Optional[int] = 0,
cached_tokens: Optional[PromptTokensDetails] = None,
image_tokens: int = 0,
audio_tokens: int = 0,
video_tokens: int = 0,
) -> UsageInfo:
"""Calculate token usage information"""
# `cached_tokens` is already a PromptTokensDetails (or None) carrying the
# cached count. Attach multimodal counts to the same object, creating one
# only when there is something to report so plain-text requests keep
# prompt_tokens_details=None (backward compatible).
details = cached_tokens
if image_tokens or audio_tokens or video_tokens:
if details is None:
details = PromptTokensDetails()
if image_tokens:
details.image_tokens = image_tokens
if audio_tokens:
details.audio_tokens = audio_tokens
if video_tokens:
details.video_tokens = video_tokens
return UsageInfo(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
prompt_tokens_details=details,
reasoning_tokens=reasoning_tokens,
)
@@ -0,0 +1,180 @@
import logging
from typing import Any, Dict, List, Optional, Union
import torch
from sglang.srt.entrypoints.openai.protocol import (
CachedTokensDetails,
ChatCompletionRequest,
CompletionRequest,
LogProbs,
StreamOptions,
)
logger = logging.getLogger(__name__)
def to_openai_style_logprobs(
input_token_logprobs=None,
output_token_logprobs=None,
input_top_logprobs=None,
output_top_logprobs=None,
):
ret_logprobs = LogProbs()
def append_token_logprobs(token_logprobs):
for logprob, _, token_text in token_logprobs:
ret_logprobs.tokens.append(token_text)
ret_logprobs.token_logprobs.append(logprob)
# Not supported yet
ret_logprobs.text_offset.append(-1)
def append_top_logprobs(top_logprobs):
for tokens in top_logprobs:
if tokens is not None:
ret_logprobs.top_logprobs.append(
{token[2]: token[0] for token in tokens}
)
else:
ret_logprobs.top_logprobs.append(None)
if input_token_logprobs is not None:
append_token_logprobs(input_token_logprobs)
if output_token_logprobs is not None:
append_token_logprobs(output_token_logprobs)
if input_top_logprobs is not None:
append_top_logprobs(input_top_logprobs)
if output_top_logprobs is not None:
append_top_logprobs(output_top_logprobs)
return ret_logprobs
def process_hidden_states_from_ret(
ret_item: Dict[str, Any],
request: Union[
ChatCompletionRequest,
CompletionRequest,
],
) -> Optional[List]:
"""Process hidden states from a ret item in non-streaming response.
Args:
ret_item: Response item containing meta_info
request: The original request object
Returns:
Processed hidden states for the last token, or None
"""
if not request.return_hidden_states:
return None
hidden_states = ret_item["meta_info"].get("hidden_states", None)
if hidden_states is not None:
hidden_states = hidden_states[-1] if len(hidden_states) > 1 else []
return hidden_states
def should_include_usage(
stream_options: StreamOptions | None, stream_response_default_include_usage: bool
) -> tuple[bool, bool]:
# When stream_options are specified in the request
if stream_options:
include_usage = (
stream_options.include_usage or stream_response_default_include_usage
)
continuous_usage_stats = bool(stream_options.continuous_usage_stats)
else:
include_usage, continuous_usage_stats = (
stream_response_default_include_usage,
False,
)
return include_usage, continuous_usage_stats
def process_routed_experts_from_ret(
ret_item: Dict[str, Any],
request: Union[
ChatCompletionRequest,
CompletionRequest,
],
) -> Optional[str]:
"""Process routed experts from a ret item in non-streaming response."""
if not getattr(request, "return_routed_experts", False):
return None
return ret_item["meta_info"].get("routed_experts", None)
def cached_tokens_details_from_dict(
details: Dict[str, Any],
) -> CachedTokensDetails:
"""Convert a raw cached_tokens_details dict to a CachedTokensDetails object."""
if "storage" in details:
return CachedTokensDetails(
device=details.get("device", 0),
host=details.get("host", 0),
storage=details.get("storage", 0),
storage_backend=details.get("storage_backend"),
)
else:
return CachedTokensDetails(
device=details.get("device", 0),
host=details.get("host", 0),
)
def process_cached_tokens_details_from_ret(
ret_item: Dict[str, Any],
request: Union[
ChatCompletionRequest,
CompletionRequest,
],
) -> Optional[CachedTokensDetails]:
"""Process cached tokens details from a ret item in non-streaming response."""
if not request.return_cached_tokens_details:
return None
details = ret_item["meta_info"].get("cached_tokens_details", None)
if details is None:
return None
return cached_tokens_details_from_dict(details)
def convert_embeds_to_tensors(
embeds: Optional[Union[List[Optional[List[List[float]]]], List[List[float]]]],
) -> Optional[List[Optional[List[torch.Tensor]]]]:
"""Convert nested float lists from the HTTP API to lists of tensors.
Accepts either:
- None -> returns None
- List[List[float]] (single input) -> [[tensor, ...]]
- List[Optional[List[List[float]]]] (batch) -> [Optional[List[tensor]], ...]
Each innermost List[float] becomes a 1-D torch.Tensor.
Per-input None entries are preserved (no overrides for that input).
"""
if embeds is None:
return None
if len(embeds) == 0:
return []
# Find first non-None entry to detect nesting depth
first_non_none = next((e for e in embeds if e is not None), None)
if first_non_none is None:
# All entries are None
return [None] * len(embeds)
# Detect nesting depth by checking the first non-None entry:
# - Single input [num_replacements][hidden_size]: first element is List[float]
# - Batch [num_inputs][num_replacements][hidden_size]: first element is List[List[float]]
if not first_non_none or not isinstance(first_non_none[0], list):
# Single input: each entry is a float vector
return [[torch.tensor(vec, dtype=torch.float32) for vec in embeds]]
# Otherwise it's batch: [num_inputs][num_replacements][hidden_size]
return [
(
[torch.tensor(vec, dtype=torch.float32) for vec in per_input]
if per_input is not None
else None
)
for per_input in embeds
]
@@ -0,0 +1,33 @@
"""Override object fields based on _HEADER_OVERRIDES from header values.
This mechanism allows upstream callers to leave the body opaque
(no parse/merge/re-serialize).
"""
from fastapi import HTTPException
# request header -> (target attribute, value type)
_HEADER_OVERRIDES = {
"x-override-rid": ("rid", str),
"x-override-bootstrap-host": ("bootstrap_host", str),
"x-override-bootstrap-port": ("bootstrap_port", int),
"x-override-bootstrap-room": ("bootstrap_room", int),
"x-override-conversation-id": ("conversation_id", str),
"x-override-routed-dp-rank": ("routed_dp_rank", int),
"x-override-disagg-prefill-dp-rank": ("disagg_prefill_dp_rank", int),
"x-override-priority": ("priority", int),
}
def apply_header_overrides(obj, headers) -> None:
"""Override request based on header values. Fail the request when any override has issues."""
for header, (attr, cast) in _HEADER_OVERRIDES.items():
value = headers.get(header)
if value is None:
continue
try:
setattr(obj, attr, cast(value))
except ValueError as e:
raise HTTPException(
status_code=400, detail=f"invalid {header} header {value!r}: {e}"
) from e
@@ -0,0 +1 @@
"""Built-in search backends for SGLang entrypoints."""
@@ -0,0 +1,138 @@
# SPDX-License-Identifier: Apache-2.0
import asyncio
from typing import Any, Optional
import aiohttp
import msgspec
from sglang.srt.environ import envs
from sglang.srt.utils import print_warning_once
EXA_API_BASE_URL = "https://api.exa.ai"
EXA_INTEGRATION_HEADER = "x-exa-integration"
EXA_INTEGRATION_NAME = "sglang"
EXA_DEFAULT_NUM_RESULTS = 10
EXA_DEFAULT_SEARCH_TYPE = "auto"
EXA_DEFAULT_HIGHLIGHTS = True
_SEARCH_TYPES = {"instant", "fast", "auto", "deep-lite", "deep", "deep-reasoning"}
class ExaClientError(RuntimeError):
"""Raised when an Exa API request fails."""
class ExaSearchConfig(msgspec.Struct, frozen=True, kw_only=True):
num_results: int = EXA_DEFAULT_NUM_RESULTS
search_type: str = EXA_DEFAULT_SEARCH_TYPE
include_highlights: bool = EXA_DEFAULT_HIGHLIGHTS
@classmethod
def from_env(cls) -> "ExaSearchConfig":
# EnvField handles type parsing (int/bool) and falls back to its own
# default on a parse error; here we only enforce the semantic bounds
# that the generic descriptors cannot express.
num_results = envs.SGLANG_EXA_NUM_RESULTS.get()
if not 1 <= num_results <= 100:
print_warning_once(
f"Ignoring invalid SGLANG_EXA_NUM_RESULTS={num_results!r}; "
f"expected a value from 1 to 100."
)
num_results = EXA_DEFAULT_NUM_RESULTS
search_type = envs.SGLANG_EXA_SEARCH_TYPE.get()
if search_type not in _SEARCH_TYPES:
print_warning_once(
f"Ignoring invalid SGLANG_EXA_SEARCH_TYPE={search_type!r}; "
f"expected one of {', '.join(sorted(_SEARCH_TYPES))}."
)
search_type = EXA_DEFAULT_SEARCH_TYPE
return cls(
num_results=num_results,
search_type=search_type,
include_highlights=envs.SGLANG_EXA_INCLUDE_HIGHLIGHTS.get(),
)
class ExaClient:
def __init__(
self,
api_key: str,
*,
config: Optional[ExaSearchConfig] = None,
base_url: str = EXA_API_BASE_URL,
timeout: float = 30.0,
):
self.api_key = api_key
self.config = config or ExaSearchConfig()
self.base_url = base_url.rstrip("/")
self.timeout = timeout
self._session: Optional[aiohttp.ClientSession] = None
self._session_lock = asyncio.Lock()
def _headers(self) -> dict[str, str]:
return {
"x-api-key": self.api_key,
"Content-Type": "application/json",
EXA_INTEGRATION_HEADER: EXA_INTEGRATION_NAME,
}
def _search_payload(self, query: str) -> dict[str, Any]:
payload: dict[str, Any] = {
"query": query,
"numResults": self.config.num_results,
"type": self.config.search_type,
}
if self.config.include_highlights:
payload["contents"] = {"highlights": True}
return payload
def _contents_payload(self, urls: list[str]) -> dict[str, Any]:
payload: dict[str, Any] = {"urls": urls, "text": True}
if self.config.include_highlights:
payload["highlights"] = True
return payload
async def search(self, query: str) -> dict[str, Any]:
return await self._post("/search", self._search_payload(query))
async def contents(self, urls: list[str]) -> dict[str, Any]:
return await self._post("/contents", self._contents_payload(urls))
async def _get_session(self) -> aiohttp.ClientSession:
# Reuse a single session across requests so the connection pool and
# DNS cache survive; aiohttp.ClientSession is intended to be long-lived.
if self._session is None:
async with self._session_lock:
if self._session is None:
self._session = aiohttp.ClientSession(
timeout=aiohttp.ClientTimeout(total=self.timeout)
)
return self._session
async def close(self):
if self._session is not None:
await self._session.close()
self._session = None
async def _post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
session = await self._get_session()
url = f"{self.base_url}{path}"
async with session.post(
url,
json=payload,
headers=self._headers(),
) as response:
response_text = await response.text()
if response.status >= 400:
raise ExaClientError(
f"Exa API request failed with status {response.status}: "
f"{response_text}"
)
try:
return await response.json()
except Exception as exc:
raise ExaClientError(
f"Failed to decode Exa API response: {response_text}"
) from exc
@@ -0,0 +1,88 @@
"""Utilities for SSL certificate hot-reloading."""
import asyncio
import logging
import ssl
from typing import Optional
from watchfiles import awatch
logger = logging.getLogger(__name__)
class SSLCertRefresher:
"""Monitors SSL certificate files and reloads them when changed.
Uses ``watchfiles.awatch()`` for efficient inotify/kqueue-based
file monitoring. On change the referenced :class:`ssl.SSLContext`
is updated in-place so that new TLS connections automatically pick
up the fresh certificates.
"""
def __init__(
self,
ssl_context: ssl.SSLContext,
key_path: str,
cert_path: str,
ca_path: Optional[str] = None,
) -> None:
self._ssl_context = ssl_context
self._key_path = key_path
self._cert_path = cert_path
self._ca_path = ca_path
self._tasks: list[asyncio.Task] = []
loop = asyncio.get_running_loop()
self._tasks.append(
loop.create_task(self._watch_cert_key(), name="ssl-cert-key-watcher")
)
if self._ca_path:
self._tasks.append(
loop.create_task(self._watch_ca(), name="ssl-ca-watcher")
)
async def _watch_cert_key(self) -> None:
"""Watch cert and key files and reload on change."""
try:
async for _changes in awatch(self._cert_path, self._key_path):
logger.info(
"SSL cert/key file change detected, reloading: " "cert=%s key=%s",
self._cert_path,
self._key_path,
)
try:
self._ssl_context.load_cert_chain(self._cert_path, self._key_path)
logger.info("SSL cert/key reloaded successfully.")
except Exception:
logger.exception(
"Failed to reload SSL cert/key — continuing with "
"previous certificates."
)
except asyncio.CancelledError:
return
async def _watch_ca(self) -> None:
"""Watch CA file and reload on change."""
assert self._ca_path is not None
try:
async for _changes in awatch(self._ca_path):
logger.info(
"SSL CA file change detected, reloading: ca=%s",
self._ca_path,
)
try:
self._ssl_context.load_verify_locations(self._ca_path)
logger.info("SSL CA certificates reloaded successfully.")
except Exception:
logger.exception(
"Failed to reload SSL CA certificates — continuing "
"with previous CA bundle."
)
except asyncio.CancelledError:
return
def stop(self) -> None:
"""Cancel all watching tasks."""
for task in self._tasks:
task.cancel()
self._tasks.clear()
+286
View File
@@ -0,0 +1,286 @@
# SPDX-License-Identifier: Apache-2.0
import logging
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any
import orjson
from sglang.srt.entrypoints.search.exa_client import (
ExaClient,
ExaSearchConfig,
)
from sglang.srt.environ import envs
from sglang.srt.utils import print_info_once, print_warning_once
if TYPE_CHECKING:
# Avoid circular import.
from sglang.srt.entrypoints.context import ConversationContext
logger = logging.getLogger(__name__)
class Tool(ABC):
@abstractmethod
async def get_result(self, context: "ConversationContext") -> Any:
pass
class HarmonyBrowserTool(Tool):
def __init__(self, client: ExaClient | None = None):
self.enabled = True
if client is not None:
self.exa_client = client
print_info_once("Browser tool initialized")
return
api_key = envs.EXA_API_KEY.get()
if not api_key:
self.enabled = False
print_warning_once("EXA_API_KEY is not set, browsing is disabled")
return
self.exa_client = ExaClient(api_key, config=ExaSearchConfig.from_env())
print_info_once("Browser tool initialized")
async def get_result(self, context: "ConversationContext") -> Any:
from openai_harmony import Author, Message, Role, TextContent
last_msg = context.messages[-1]
recipient = last_msg.recipient
if recipient is None or not recipient.startswith("browser."):
raise ValueError("No browser tool call found")
try:
args = orjson.loads(last_msg.content[0].text)
result_text = await self._dispatch_browser_call(context, recipient, args)
except Exception as exc:
logger.exception("Browser tool call failed")
result_text = f"Browser tool call failed: {exc}"
content = TextContent(text=result_text)
author = Author(role=Role.TOOL, name=recipient)
return [Message(author=author, content=[content], recipient=Role.ASSISTANT)]
async def _dispatch_browser_call(
self, context: "ConversationContext", recipient: str, args: dict[str, Any]
) -> str:
if recipient == "browser.search":
query = args.get("query")
if not query:
raise ValueError("browser.search requires a query")
data = await self.exa_client.search(query)
return self._format_search_results(context, query, data)
if recipient == "browser.open":
url = self._resolve_url(context, args)
data = await self.exa_client.contents([url])
return self._format_page_contents(context, url, data)
if recipient == "browser.find":
pattern = args.get("pattern")
if not pattern:
raise ValueError("browser.find requires a pattern")
return await self._find_pattern(context, args, pattern)
raise ValueError(f"Unknown browser action: {recipient}")
def _browser_state(self, context: "ConversationContext") -> dict[str, Any]:
state = getattr(context, "_sglang_exa_browser_state", None)
if state is None:
state = {"pages": {}, "page_text": {}}
setattr(context, "_sglang_exa_browser_state", state)
return state
def _format_search_results(
self, context: "ConversationContext", query: str, data: dict[str, Any]
) -> str:
state = self._browser_state(context)
state["pages"] = {}
state["page_text"] = {}
results = data.get("results") or []
request_id = data.get("requestId")
if request_id:
logger.debug("Exa search request id: %s", request_id)
lines = [f"Search results for: {query}"]
lines.append("Use browser.open with the cursor number to inspect a result.")
for index, result in enumerate(results, start=1):
cursor = str(index)
state["pages"][cursor] = result
title = result.get("title") or "Untitled"
url = result.get("url") or result.get("id") or ""
snippet = self._best_snippet(result)
lines.append("")
lines.append(f"[{cursor}] {title}")
if url:
lines.append(f"URL: {url}")
if snippet:
lines.append(f"Snippet: {snippet}")
if not results:
lines.append("No results found.")
return "\n".join(lines)
def _format_page_contents(
self, context: "ConversationContext", url: str, data: dict[str, Any]
) -> str:
state = self._browser_state(context)
results = data.get("results") or []
if not results:
return f"No page contents returned for {url}."
page = results[0]
cursor = self._cursor_for_url(state, url)
if cursor:
state["pages"][cursor] = page
title = page.get("title") or "Untitled"
page_url = page.get("url") or url
text = page.get("text") or self._best_snippet(page) or page.get("summary") or ""
if cursor:
state["page_text"][cursor] = text
return "\n".join(
[
f"Opened page: {title}",
f"URL: {page_url}",
"",
self._truncate(text, 12000) if text else "No page text available.",
]
)
async def _find_pattern(
self, context: "ConversationContext", args: dict[str, Any], pattern: str
) -> str:
state = self._browser_state(context)
cursor = self._normalize_cursor(args.get("cursor"))
if cursor:
text = state["page_text"].get(cursor)
if text is None:
url = self._resolve_url(context, args)
data = await self.exa_client.contents([url])
self._format_page_contents(context, url, data)
text = state["page_text"].get(cursor, "")
return self._format_matches(pattern, text)
if args.get("url"):
url = self._resolve_url(context, args)
data = await self.exa_client.contents([url])
results = data.get("results") or []
if not results:
return f"No page contents returned for {url}."
page = results[0]
text = (
page.get("text")
or self._best_snippet(page)
or page.get("summary")
or ""
)
return self._format_matches(pattern, text)
searchable_text = "\n\n".join(
self._best_snippet(page) for page in state["pages"].values()
)
return self._format_matches(pattern, searchable_text)
def _resolve_url(self, context: "ConversationContext", args: dict[str, Any]) -> str:
if args.get("url"):
return str(args["url"])
state = self._browser_state(context)
cursor = self._normalize_cursor(args.get("cursor"))
if not cursor:
raise ValueError("browser.open requires a cursor or url")
page = state["pages"].get(cursor)
if page is None:
raise ValueError(f"Unknown browser cursor: {cursor}")
url = page.get("url") or page.get("id")
if not url:
raise ValueError(f"No URL recorded for browser cursor: {cursor}")
return str(url)
def _normalize_cursor(self, cursor: Any) -> str | None:
if cursor is None:
return None
cursor_str = str(cursor)
# GPT-OSS emits 0 as a 1-based cursor; map it to the first result.
if cursor_str == "0":
return "1"
return cursor_str
def _cursor_for_url(self, state: dict[str, Any], url: str) -> str | None:
for cursor, page in state["pages"].items():
if page.get("url") == url or page.get("id") == url:
return cursor
return None
def _best_snippet(self, result: dict[str, Any]) -> str:
highlights = result.get("highlights") or []
if highlights:
return self._truncate(str(highlights[0]), 1000)
summary = result.get("summary")
if summary:
return self._truncate(str(summary), 1000)
text = result.get("text")
if text:
return self._truncate(str(text), 1000)
return ""
def _format_matches(self, pattern: str, text: str) -> str:
if not text:
return f"No text available to search for {pattern!r}."
pattern_lower = pattern.lower()
matches = []
for line in text.splitlines():
if pattern_lower in line.lower():
matches.append(self._truncate(line.strip(), 1000))
if len(matches) >= 10:
break
if not matches:
return f"No matches found for {pattern!r}."
lines = [f"Matches for {pattern!r}:"]
lines.extend(f"- {match}" for match in matches)
return "\n".join(lines)
def _truncate(self, text: str, max_chars: int) -> str:
if len(text) <= max_chars:
return text
return text[: max_chars - 3] + "..."
class HarmonyPythonTool(Tool):
def __init__(self):
self.enabled = True
try:
from gpt_oss.tools.python_docker.docker_tool import PythonTool
except ImportError:
self.enabled = False
print_warning_once("gpt_oss is not installed, code interpreter is disabled")
return
self.python_tool = PythonTool()
print_info_once("Code interpreter tool initialized")
async def get_result(self, context: "ConversationContext") -> Any:
from sglang.srt.entrypoints.context import HarmonyContext
assert isinstance(context, HarmonyContext)
last_msg = context.messages[-1]
tool_output_msgs = []
async for msg in self.python_tool.process(last_msg):
tool_output_msgs.append(msg)
return tool_output_msgs
@property
def tool_config(self) -> Any:
return self.python_tool.tool_config
+123
View File
@@ -0,0 +1,123 @@
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""
/v1/loads API endpoint for comprehensive load metrics.
This module provides the /v1/loads endpoint which returns detailed scheduler
metrics for load balancing, monitoring, and capacity planning.
"""
import time
from datetime import datetime, timezone
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import Response
from sglang.version import __version__
router = APIRouter()
def _get_tokenizer_manager():
"""Dependency to get tokenizer_manager from global state."""
from sglang.srt.entrypoints.http_server import get_global_state
return get_global_state().tokenizer_manager
def _format_loads_prometheus(load_results, include=None) -> Response:
"""Format load metrics in Prometheus text exposition format."""
section_prefixes = {"speculative": "spec", "disaggregation": "disagg"}
metric_samples = {}
for load in load_results:
load_dict = load.to_dict(include)
dp_rank = load_dict.pop("dp_rank")
for key, value in load_dict.items():
if isinstance(value, dict):
prefix = section_prefixes.get(key, key)
for sub_key, sub_value in value.items():
if isinstance(sub_value, (int, float)):
metric_samples.setdefault(
f"sglang_{prefix}_{sub_key}", []
).append((dp_rank, sub_value))
elif isinstance(value, (int, float)):
metric_samples.setdefault(f"sglang_{key}", []).append((dp_rank, value))
lines = []
for metric_name, samples in metric_samples.items():
lines.append(f"# TYPE {metric_name} gauge")
for dp_rank, value in samples:
lines.append(f'{metric_name}{{dp_rank="{dp_rank}"}} {value}')
return Response(
content="\n".join(lines) + "\n",
media_type="text/plain; version=0.0.4; charset=utf-8",
)
@router.get("/v1/loads")
async def get_loads(
dp_rank: Optional[int] = None,
include: Optional[str] = None,
format: Optional[str] = None,
tokenizer_manager=Depends(_get_tokenizer_manager),
):
"""
Get comprehensive load metrics for all DP ranks.
Query Parameters:
dp_rank: Filter to specific DP rank (optional)
include: Comma-separated sections to include (optional)
Options: core, memory, spec, lora, disagg, queues, all
Default: all
format: Response format - 'json' (default) or 'prometheus'
Returns:
JSON response with timestamp, version, and per-DP-rank loads
"""
include_list = [s.strip() for s in include.split(",")] if include else None
start = time.perf_counter()
try:
load_results = await tokenizer_manager.get_loads(
include=include_list,
dp_rank=dp_rank,
)
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e))
finally:
mc = getattr(tokenizer_manager, "metrics_collector", None)
if mc is not None:
mc.get_loads_duration_seconds.labels(**mc.labels).observe(
time.perf_counter() - start
)
include_set = set(include_list) if include_list else None
if format == "prometheus":
return _format_loads_prometheus(load_results, include_set)
loads = []
for load in load_results:
d = load.to_dict(include_set)
loads.append(d)
return {
"timestamp": datetime.now(timezone.utc).isoformat(),
"version": __version__,
"loads": loads,
}
+161
View File
@@ -0,0 +1,161 @@
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, List
import numpy as np
import tqdm
from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST
from sglang.srt.managers.io_struct import GenerateReqInput
if TYPE_CHECKING:
from sglang.srt.managers.tokenizer_manager import TokenizerManager
logger = logging.getLogger(__file__)
_warmup_registry = {}
def warmup(name: str):
def decorator(fn):
_warmup_registry[name] = fn
return fn
return decorator
async def execute_warmups(
disaggregation_mode: str,
warmup_names: List[str],
tokenizer_manager: TokenizerManager,
):
for warmup_name in warmup_names:
if warmup_name not in _warmup_registry:
logger.warning(f"Could not find custom warmup {warmup_name}")
continue
logger.info(f"Running warmup {warmup_name}")
await _warmup_registry[warmup_name](disaggregation_mode, tokenizer_manager)
@warmup("whisper_autodetect")
async def whisper_autodetect(
disaggregation_mode: str, tokenizer_manager: TokenizerManager
):
"""Pre-compile the xgrammar FSM for both Whisper auto-detect regexes.
The first request that uses each structured-generation regex incurs a
~15-20s compilation cost. xgrammar caches compiled grammars by the
exact regex string, so we warm both the notimestamps and timestamps
variants here — otherwise the first ``language=None +
timestamp_granularities`` request would still pay the full spike.
"""
# A short silent audio encoded as base64 WAV (0.1s, 16kHz, mono) —
# soundfile produces the WAV header + PCM data from a list of floats.
import base64
import io
import soundfile as sf
from sglang.srt.entrypoints.openai.transcription_adapters.whisper import (
FUSED_AUTODETECT_FLAG,
WHISPER_AUTODETECT_REGEX,
WHISPER_AUTODETECT_TS_REGEX,
)
sr, dur = 16000, 0.1
n = int(sr * dur)
buf = io.BytesIO()
sf.write(buf, [0.0] * n, sr, format="WAV")
audio_b64 = base64.b64encode(buf.getvalue()).decode()
audio_data_uri = f"data:audio/wav;base64,{audio_b64}"
for variant_name, regex in (
("notimestamps", WHISPER_AUTODETECT_REGEX),
("timestamps", WHISPER_AUTODETECT_TS_REGEX),
):
logger.info(
"Compiling Whisper auto-detect regex FSM (%s, one-time, ~15-20s)...",
variant_name,
)
req = GenerateReqInput(
text="",
audio_data=audio_data_uri,
sampling_params={
"max_new_tokens": 4,
"temperature": 0,
"regex": regex,
"skip_special_tokens": False,
"spaces_between_special_tokens": False,
FUSED_AUTODETECT_FLAG: True,
},
modalities=["audio"],
)
# PD prefill servers assert req.bootstrap_room is not None in the
# default follow_bootstrap_room scheduler; the fake values match
# what the voice_chat warmup uses for the same reason.
if disaggregation_mode != "null":
req.bootstrap_room = 0
req.bootstrap_host = FAKE_BOOTSTRAP_HOST
# Drain the generator so the FSM is fully installed and any
# downstream exception surfaces instead of being swallowed after
# the first yield.
async for _ in tokenizer_manager.generate_request(req, None):
pass
logger.info("Whisper auto-detect regex FSMs compiled.")
@warmup("voice_chat")
async def voice_chat(disaggregation_mode: str, tokenizer_manager: TokenizerManager):
# this warms up the fused_moe triton kernels and caches them
# if we don't do this we break real time inference for voice chat
for i in tqdm.trange(1, 512):
size = i * 4
generate_req_input = GenerateReqInput(
input_ids=(np.random.randint(2**16, size=[size])).tolist(),
sampling_params={
"max_new_tokens": 30,
"temperature": 0.8,
"stop_token_ids": [1],
"min_p": 0.0,
},
)
if disaggregation_mode != "null":
generate_req_input.bootstrap_room = 0
generate_req_input.bootstrap_host = FAKE_BOOTSTRAP_HOST
await tokenizer_manager.generate_request(generate_req_input, None).__anext__()
@warmup("prefill_shapes")
async def prefill_shapes(disaggregation_mode: str, tokenizer_manager: TokenizerManager):
"""Warmup Triton kernels across a wide range of prefill seq_lens (up to 32K).
Uses power-of-2 sizes plus intermediate points to cover the shape space
that fused_moe, attention extend, and other Triton kernels may encounter.
"""
page_size = 64
sizes = set()
base = 64
while base <= 32768:
sizes.add(base)
mid = base * 3 // 2
mid = (mid + page_size - 1) // page_size * page_size
if mid <= 32768:
sizes.add(mid)
base *= 2
sizes = sorted(sizes)
for size in tqdm.tqdm(sizes, desc="Warmup prefill shapes (up to 32K)"):
generate_req_input = GenerateReqInput(
input_ids=(np.random.randint(2**16, size=[size])).tolist(),
sampling_params={
"max_new_tokens": 1,
"temperature": 0.0,
},
)
if disaggregation_mode != "null":
generate_req_input.bootstrap_room = 0
generate_req_input.bootstrap_host = FAKE_BOOTSTRAP_HOST
await tokenizer_manager.generate_request(generate_req_input, None).__anext__()