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
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:
@@ -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
@@ -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()
|
||||
@@ -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
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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__()
|
||||
Reference in New Issue
Block a user