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,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
|
||||
]
|
||||
Reference in New Issue
Block a user