wehub-resource-sync
463 lines
16 KiB
Python
463 lines
16 KiB
Python
"""Protocols in MLC LLM for OpenAI API.
|
|
Adapted from FastChat's OpenAI protocol:
|
|
https://github.com/lm-sys/FastChat/blob/main/fastchat/protocol/openai_api_protocol.py
|
|
"""
|
|
|
|
import json
|
|
import time
|
|
from typing import Any, Dict, List, Literal, Optional, Tuple, Union # noqa: UP035
|
|
|
|
import shortuuid
|
|
from pydantic import BaseModel, Field, field_validator, model_validator
|
|
|
|
from .conversation_protocol import Conversation
|
|
from .debug_protocol import DebugConfig
|
|
from .error_protocol import BadRequestError
|
|
|
|
################ Commons ################
|
|
|
|
|
|
# OPenAI API compatible limits
|
|
CHAT_COMPLETION_MAX_TOP_LOGPROBS = 20
|
|
COMPLETION_MAX_TOP_LOGPROBS = 5
|
|
|
|
|
|
class ListResponse(BaseModel):
|
|
object: str = "list"
|
|
data: List[Any] # noqa: UP006
|
|
|
|
|
|
class TopLogProbs(BaseModel):
|
|
token: str
|
|
logprob: float
|
|
bytes: Optional[List[int]] # noqa: UP006
|
|
|
|
|
|
class LogProbsContent(BaseModel):
|
|
token: str
|
|
logprob: float
|
|
bytes: Optional[List[int]] # noqa: UP006
|
|
top_logprobs: List[TopLogProbs] = [] # noqa: UP006
|
|
|
|
|
|
class LogProbs(BaseModel):
|
|
content: List[LogProbsContent] # noqa: UP006
|
|
|
|
|
|
class CompletionLogProbs(BaseModel):
|
|
# The position of the token in the concatenated str: prompt + completion_text
|
|
# TODO(vvchernov): skip optional after support
|
|
text_offset: Optional[List[int]] # noqa: UP006
|
|
token_logprobs: List[float] # noqa: UP006
|
|
tokens: List[str] # noqa: UP006
|
|
top_logprobs: List[Dict[str, float]] # noqa: UP006
|
|
|
|
|
|
class CompletionUsage(BaseModel):
|
|
prompt_tokens: int
|
|
completion_tokens: int
|
|
total_tokens: int
|
|
extra: Optional[Dict[str, Any]] = None # noqa: UP006
|
|
"""Extra metrics and info that may be returned by debug_config
|
|
"""
|
|
|
|
|
|
class StreamOptions(BaseModel):
|
|
include_usage: Optional[bool]
|
|
|
|
|
|
################ v1/embeddings ################
|
|
|
|
|
|
class EmbeddingRequest(BaseModel):
|
|
"""OpenAI "v1/embeddings" request protocol.
|
|
API reference: https://platform.openai.com/docs/api-reference/embeddings/create
|
|
"""
|
|
|
|
input: Union[str, List[str], List[int], List[List[int]]] # noqa: UP006
|
|
model: Optional[str] = None
|
|
encoding_format: Literal["float", "base64"] = "float"
|
|
dimensions: Optional[int] = None
|
|
user: Optional[str] = None
|
|
|
|
@field_validator("input")
|
|
@classmethod
|
|
def validate_input(cls, v):
|
|
"""Check that the input is not an empty list.
|
|
|
|
Note: empty strings are allowed — encoder models produce valid
|
|
embeddings from [CLS]+[SEP] tokens alone.
|
|
"""
|
|
if isinstance(v, list) and len(v) == 0:
|
|
raise ValueError("Input list must not be empty.")
|
|
return v
|
|
|
|
|
|
class EmbeddingObject(BaseModel):
|
|
object: str = "embedding"
|
|
embedding: Union[List[float], str] # noqa: UP006
|
|
index: int
|
|
|
|
|
|
class EmbeddingUsage(BaseModel):
|
|
prompt_tokens: int
|
|
total_tokens: int
|
|
|
|
|
|
class EmbeddingResponse(BaseModel):
|
|
"""OpenAI "v1/embeddings" response protocol.
|
|
API reference: https://platform.openai.com/docs/api-reference/embeddings/object
|
|
"""
|
|
|
|
object: str = "list"
|
|
data: List[EmbeddingObject] # noqa: UP006
|
|
model: Optional[str] = None
|
|
usage: EmbeddingUsage
|
|
|
|
|
|
################ v1/models ################
|
|
|
|
|
|
class ModelResponse(BaseModel):
|
|
"""OpenAI "v1/models" response protocol.
|
|
API reference: https://platform.openai.com/docs/api-reference/models/object
|
|
"""
|
|
|
|
id: str
|
|
created: int = Field(default_factory=lambda: int(time.time()))
|
|
object: str = "model"
|
|
owned_by: str = "MLC-LLM"
|
|
|
|
|
|
################ v1/completions ################
|
|
|
|
|
|
class RequestResponseFormat(BaseModel):
|
|
type: Literal["text", "json_object"] = "text"
|
|
json_schema: Optional[str] = Field(default=None, alias="schema")
|
|
"""This field is named json_schema instead of schema because BaseModel defines a method called
|
|
schema. During construction of RequestResponseFormat, key "schema" still should be used:
|
|
`RequestResponseFormat(type="json_object", schema="{}")`
|
|
"""
|
|
|
|
|
|
class CompletionRequest(BaseModel):
|
|
"""OpenAI completion request protocol.
|
|
API reference: https://platform.openai.com/docs/api-reference/completions/create
|
|
"""
|
|
|
|
model: Optional[str] = None
|
|
prompt: Union[str, List[int]] # noqa: UP006
|
|
best_of: int = 1
|
|
echo: bool = False
|
|
frequency_penalty: Optional[float] = None
|
|
presence_penalty: Optional[float] = None
|
|
logprobs: Optional[int] = None
|
|
logit_bias: Optional[Dict[int, float]] = None # noqa: UP006
|
|
max_tokens: Optional[int] = None
|
|
n: int = 1
|
|
seed: Optional[int] = None
|
|
stop: Optional[Union[str, List[str]]] = None # noqa: UP006
|
|
stream: bool = False
|
|
stream_options: Optional[StreamOptions] = None
|
|
suffix: Optional[str] = None
|
|
temperature: Optional[float] = None
|
|
top_p: Optional[float] = None
|
|
user: Optional[str] = None
|
|
response_format: Optional[RequestResponseFormat] = None
|
|
debug_config: Optional[DebugConfig] = None
|
|
|
|
@field_validator("frequency_penalty", "presence_penalty")
|
|
@classmethod
|
|
def check_penalty_range(cls, penalty_value: Optional[float]) -> Optional[float]:
|
|
"""Check if the penalty value is in range [-2, 2]."""
|
|
if penalty_value and (penalty_value < -2 or penalty_value > 2):
|
|
raise ValueError("Penalty value should be in range [-2, 2].")
|
|
return penalty_value
|
|
|
|
@field_validator("logit_bias")
|
|
@classmethod
|
|
def check_logit_bias(
|
|
cls,
|
|
logit_bias_value: Optional[Dict[int, float]], # noqa: UP006
|
|
) -> Optional[Dict[int, float]]: # noqa: UP006
|
|
"""Check if the logit bias key is given as an integer."""
|
|
if logit_bias_value is None:
|
|
return None
|
|
for token_id, bias in logit_bias_value.items():
|
|
if abs(bias) > 100:
|
|
raise ValueError(
|
|
"Logit bias value should be in range [-100, 100], while value "
|
|
f"{bias} is given for token id {token_id}"
|
|
)
|
|
return logit_bias_value
|
|
|
|
@model_validator(mode="after")
|
|
def check_logprobs(self) -> "CompletionRequest":
|
|
"""Check if the logprobs requirements are valid."""
|
|
if self.logprobs is not None and (
|
|
self.logprobs < 0 or self.logprobs > COMPLETION_MAX_TOP_LOGPROBS
|
|
):
|
|
raise ValueError(f'"logprobs" must be in range [0, {COMPLETION_MAX_TOP_LOGPROBS}]')
|
|
return self
|
|
|
|
|
|
class CompletionResponseChoice(BaseModel):
|
|
finish_reason: Optional[Literal["stop", "length", "preempt"]] = None
|
|
index: int = 0
|
|
logprobs: Optional[CompletionLogProbs] = None
|
|
text: str
|
|
|
|
|
|
class CompletionResponse(BaseModel):
|
|
"""OpenAI completion response protocol.
|
|
API reference: https://platform.openai.com/docs/api-reference/completions/object
|
|
"""
|
|
|
|
id: str
|
|
choices: List[CompletionResponseChoice] # noqa: UP006
|
|
created: int = Field(default_factory=lambda: int(time.time()))
|
|
model: Optional[str] = None
|
|
object: str = "text_completion"
|
|
usage: Optional[CompletionUsage] = None
|
|
|
|
|
|
################ v1/chat/completions ################
|
|
|
|
|
|
class ChatFunction(BaseModel):
|
|
description: Optional[str] = None
|
|
name: str
|
|
parameters: Dict # noqa: UP006
|
|
|
|
|
|
class ChatTool(BaseModel):
|
|
type: Literal["function"]
|
|
function: ChatFunction
|
|
|
|
|
|
class ChatFunctionCall(BaseModel):
|
|
name: str
|
|
arguments: Union[None, Dict[str, Any]] = None # noqa: UP006
|
|
|
|
|
|
class ChatToolCall(BaseModel):
|
|
id: str = Field(default_factory=lambda: f"call_{shortuuid.random()}")
|
|
type: Literal["function"]
|
|
function: ChatFunctionCall
|
|
|
|
|
|
class ChatCompletionMessage(BaseModel):
|
|
content: Optional[Union[str, List[Dict]]] = None # noqa: UP006
|
|
role: Literal["system", "user", "assistant", "tool"]
|
|
name: Optional[str] = None
|
|
tool_calls: Optional[List[ChatToolCall]] = None # noqa: UP006
|
|
tool_call_id: Optional[str] = None
|
|
|
|
|
|
class ChatCompletionRequest(BaseModel):
|
|
"""OpenAI chat completion request protocol.
|
|
API reference: https://platform.openai.com/docs/api-reference/chat/create
|
|
"""
|
|
|
|
messages: List[ChatCompletionMessage] # noqa: UP006
|
|
model: Optional[str] = None
|
|
frequency_penalty: Optional[float] = None
|
|
presence_penalty: Optional[float] = None
|
|
logprobs: bool = False
|
|
top_logprobs: int = 0
|
|
logit_bias: Optional[Dict[int, float]] = None # noqa: UP006
|
|
max_tokens: Optional[int] = None
|
|
n: int = 1
|
|
seed: Optional[int] = None
|
|
stop: Optional[Union[str, List[str]]] = None # noqa: UP006
|
|
stream: bool = False
|
|
stream_options: Optional[StreamOptions] = None
|
|
temperature: Optional[float] = None
|
|
top_p: Optional[float] = None
|
|
tools: Optional[List[ChatTool]] = None # noqa: UP006
|
|
tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None # noqa: UP006
|
|
user: Optional[str] = None
|
|
response_format: Optional[RequestResponseFormat] = None
|
|
# NOTE: debug_config is not part of OpenAI protocol
|
|
# we add it to enable extra debug options
|
|
debug_config: Optional[DebugConfig] = None
|
|
|
|
@field_validator("frequency_penalty", "presence_penalty")
|
|
@classmethod
|
|
def check_penalty_range(cls, penalty_value: Optional[float]) -> Optional[float]:
|
|
"""Check if the penalty value is in range [-2, 2]."""
|
|
if penalty_value and (penalty_value < -2 or penalty_value > 2):
|
|
raise ValueError("Penalty value should be in range [-2, 2].")
|
|
return penalty_value
|
|
|
|
@field_validator("logit_bias")
|
|
@classmethod
|
|
def check_logit_bias(
|
|
cls,
|
|
logit_bias_value: Optional[Dict[int, float]], # noqa: UP006
|
|
) -> Optional[Dict[int, float]]: # noqa: UP006
|
|
"""Check if the logit bias key is given as an integer."""
|
|
if logit_bias_value is None:
|
|
return None
|
|
for token_id, bias in logit_bias_value.items():
|
|
if abs(bias) > 100:
|
|
raise ValueError(
|
|
"Logit bias value should be in range [-100, 100], while value "
|
|
f"{bias} is given for token id {token_id}"
|
|
)
|
|
return logit_bias_value
|
|
|
|
@model_validator(mode="after")
|
|
def check_logprobs(self) -> "ChatCompletionRequest":
|
|
"""Check if the logprobs requirements are valid."""
|
|
if self.top_logprobs < 0 or self.top_logprobs > CHAT_COMPLETION_MAX_TOP_LOGPROBS:
|
|
raise ValueError(
|
|
f'"top_logprobs" must be in range [0, {CHAT_COMPLETION_MAX_TOP_LOGPROBS}]'
|
|
)
|
|
if not self.logprobs and self.top_logprobs > 0:
|
|
raise ValueError('"logprobs" must be True to support "top_logprobs"')
|
|
return self
|
|
|
|
@model_validator(mode="after")
|
|
def check_stream_options(self) -> "ChatCompletionRequest":
|
|
"""Check stream options"""
|
|
if self.stream_options is None:
|
|
return self
|
|
if not self.stream:
|
|
raise ValueError("stream must be set to True when stream_options is present")
|
|
return self
|
|
|
|
@model_validator(mode="after")
|
|
def check_debug_config(self) -> "ChatCompletionRequest":
|
|
"""Check debug config"""
|
|
if self.debug_config is None:
|
|
return self
|
|
|
|
if self.debug_config.special_request is None:
|
|
return self
|
|
|
|
if not self.stream:
|
|
raise ValueError("DebugConfig.special_request requires stream=True")
|
|
|
|
if self.stream_options is None or not self.stream_options.include_usage:
|
|
raise ValueError("DebugConfig.special_request requires include_usage in stream_options")
|
|
|
|
return self
|
|
|
|
def check_message_validity(self) -> None:
|
|
"""Check if the given chat messages are valid. Return error message if invalid."""
|
|
for i, message in enumerate(self.messages):
|
|
if message.role == "system" and i != 0:
|
|
raise BadRequestError(
|
|
f"System prompt at position {i} in the message list is invalid."
|
|
)
|
|
if message.tool_call_id is not None:
|
|
if message.role != "tool":
|
|
raise BadRequestError("Non-tool message having `tool_call_id` is invalid.")
|
|
if isinstance(message.content, list):
|
|
if message.role != "user":
|
|
raise BadRequestError("Non-user message having a list of content is invalid.")
|
|
if message.tool_calls is not None:
|
|
if message.role != "assistant":
|
|
raise BadRequestError("Non-assistant message having `tool_calls` is invalid.")
|
|
raise BadRequestError("Assistant message having `tool_calls` is not supported yet.")
|
|
|
|
def check_function_call_usage(self, conv_template: Conversation) -> None:
|
|
"""Check if function calling is used and update the conversation template.
|
|
Return error message if invalid request format for function calling.
|
|
"""
|
|
|
|
# return if no tools are provided or tool_choice is set to none
|
|
if self.tools is None or (isinstance(self.tool_choice, str) and self.tool_choice == "none"):
|
|
conv_template.use_function_calling = False
|
|
return
|
|
|
|
# select the tool based on the tool_choice if specified
|
|
if isinstance(self.tool_choice, dict):
|
|
if self.tool_choice["type"] != "function":
|
|
raise BadRequestError("Only 'function' tool choice is supported")
|
|
|
|
if len(self.tool_choice["function"]) > 1:
|
|
raise BadRequestError("Only one tool is supported when tool_choice is specified")
|
|
|
|
for tool in self.tools:
|
|
if tool.function.name == self.tool_choice["function"]["name"]:
|
|
conv_template.use_function_calling = True
|
|
conv_template.function_string = tool.function.model_dump_json(by_alias=True)
|
|
return
|
|
|
|
raise BadRequestError(
|
|
f"The tool_choice function {self.tool_choice['function']['name']}"
|
|
" is not found in the tools list"
|
|
)
|
|
|
|
if isinstance(self.tool_choice, str) and self.tool_choice != "auto":
|
|
raise BadRequestError(f"Invalid tool_choice value: {self.tool_choice}")
|
|
|
|
function_list = []
|
|
for tool in self.tools:
|
|
if tool.type != "function":
|
|
raise BadRequestError("Only 'function' tool type is supported")
|
|
function_list.append(tool.function.model_dump(by_alias=True))
|
|
|
|
conv_template.use_function_calling = True
|
|
conv_template.function_string = json.dumps(function_list)
|
|
|
|
|
|
class ChatCompletionResponseChoice(BaseModel):
|
|
finish_reason: Optional[Literal["stop", "length", "tool_calls", "error"]] = None
|
|
index: int = 0
|
|
message: ChatCompletionMessage
|
|
logprobs: Optional[LogProbs] = None
|
|
|
|
|
|
class ChatCompletionStreamResponseChoice(BaseModel):
|
|
finish_reason: Optional[Literal["stop", "length", "tool_calls", "error"]] = None
|
|
index: int = 0
|
|
delta: ChatCompletionMessage
|
|
logprobs: Optional[LogProbs] = None
|
|
|
|
|
|
class ChatCompletionResponse(BaseModel):
|
|
"""OpenAI completion response protocol.
|
|
API reference: https://platform.openai.com/docs/api-reference/chat/object
|
|
"""
|
|
|
|
id: str
|
|
choices: List[ChatCompletionResponseChoice] # noqa: UP006
|
|
created: int = Field(default_factory=lambda: int(time.time()))
|
|
model: Optional[str] = None
|
|
system_fingerprint: str
|
|
object: Literal["chat.completion"] = "chat.completion"
|
|
usage: Optional[CompletionUsage] = None
|
|
|
|
|
|
class ChatCompletionStreamResponse(BaseModel):
|
|
"""OpenAI completion stream response protocol.
|
|
API reference: https://platform.openai.com/docs/api-reference/chat/streaming
|
|
"""
|
|
|
|
id: str
|
|
choices: List[ChatCompletionStreamResponseChoice] # noqa: UP006
|
|
created: int = Field(default_factory=lambda: int(time.time()))
|
|
model: Optional[str] = None
|
|
system_fingerprint: str
|
|
object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
|
|
usage: Optional[CompletionUsage] = None
|
|
|
|
|
|
def openai_api_get_unsupported_fields(
|
|
request: Union[CompletionRequest, ChatCompletionRequest],
|
|
) -> List[str]: # noqa: UP006
|
|
"""Get the unsupported fields in the request."""
|
|
unsupported_field_default_values: List[Tuple[str, Any]] = [ # noqa: UP006
|
|
("best_of", 1),
|
|
]
|
|
|
|
unsupported_fields: List[str] = [] # noqa: UP006
|
|
for field, value in unsupported_field_default_values:
|
|
if hasattr(request, field) and getattr(request, field) != value:
|
|
unsupported_fields.append(field)
|
|
return unsupported_fields
|