160 lines
4.2 KiB
Python
160 lines
4.2 KiB
Python
import time
|
|
from typing import Dict, List, Optional, Union
|
|
|
|
from pydantic import BaseModel, Field
|
|
from typing_extensions import Literal
|
|
|
|
|
|
class LogProbs(BaseModel):
|
|
text_offset: List[int] = Field(default_factory=list)
|
|
token_logprobs: List[Optional[float]] = Field(default_factory=list)
|
|
tokens: List[str] = Field(default_factory=list)
|
|
top_logprobs: List[Optional[Dict[str, float]]] = Field(default_factory=list)
|
|
|
|
|
|
class UsageInfo(BaseModel):
|
|
prompt_tokens: int = 0
|
|
total_tokens: int = 0
|
|
completion_tokens: Optional[int] = 0
|
|
|
|
|
|
class CompletionRequest(BaseModel):
|
|
model: str
|
|
prompt: Union[str, List[str]]
|
|
suffix: Optional[str] = None
|
|
max_tokens: Optional[int] = 16
|
|
temperature: Optional[float] = 0.7
|
|
top_p: Optional[float] = 1.0
|
|
n: Optional[int] = 1
|
|
stream: Optional[bool] = False
|
|
logprobs: Optional[int] = None
|
|
echo: Optional[bool] = False
|
|
stop: Optional[Union[str, List[str]]] = Field(default_factory=list)
|
|
presence_penalty: Optional[float] = 0.0
|
|
frequency_penalty: Optional[float] = 0.0
|
|
best_of: Optional[int] = None
|
|
logit_bias: Optional[Dict[str, float]] = None
|
|
user: Optional[str] = None
|
|
|
|
|
|
class CompletionResponseChoice(BaseModel):
|
|
index: int
|
|
text: str
|
|
logprobs: Optional[LogProbs] = None
|
|
finish_reason: Optional[str] = None
|
|
|
|
|
|
class CompletionResponse(BaseModel):
|
|
id: str
|
|
object: str = "text_completion"
|
|
created: int = Field(default_factory=lambda: int(time.time()))
|
|
model: str
|
|
choices: List[CompletionResponseChoice]
|
|
usage: UsageInfo
|
|
|
|
|
|
class CompletionResponseStreamChoice(BaseModel):
|
|
index: int
|
|
text: str
|
|
logprobs: Optional[LogProbs] = None
|
|
finish_reason: Optional[str] = None
|
|
|
|
|
|
class CompletionStreamResponse(BaseModel):
|
|
id: str
|
|
object: str = "text_completion"
|
|
created: int = Field(default_factory=lambda: int(time.time()))
|
|
model: str
|
|
choices: List[CompletionResponseStreamChoice]
|
|
usage: UsageInfo
|
|
|
|
|
|
class ChatCompletionMessageGenericParam(BaseModel):
|
|
role: Literal["system", "assistant"]
|
|
content: str
|
|
|
|
|
|
class ChatCompletionMessageContentTextPart(BaseModel):
|
|
type: Literal["text"]
|
|
text: str
|
|
|
|
|
|
class ChatCompletionMessageContentImageURL(BaseModel):
|
|
url: str
|
|
detail: Optional[Literal["auto", "low", "high"]] = "auto"
|
|
|
|
|
|
class ChatCompletionMessageContentImagePart(BaseModel):
|
|
type: Literal["image_url"]
|
|
image_url: ChatCompletionMessageContentImageURL
|
|
|
|
|
|
ChatCompletionMessageContentPart = Union[
|
|
ChatCompletionMessageContentTextPart, ChatCompletionMessageContentImagePart
|
|
]
|
|
|
|
|
|
class ChatCompletionMessageUserParam(BaseModel):
|
|
role: Literal["user"]
|
|
content: Union[str, List[ChatCompletionMessageContentPart]]
|
|
|
|
|
|
ChatCompletionMessageParam = Union[
|
|
ChatCompletionMessageGenericParam, ChatCompletionMessageUserParam
|
|
]
|
|
|
|
|
|
class ChatCompletionRequest(BaseModel):
|
|
model: str
|
|
messages: Union[str, List[ChatCompletionMessageParam]]
|
|
temperature: Optional[float] = 0.7
|
|
top_p: Optional[float] = 1.0
|
|
n: Optional[int] = 1
|
|
max_tokens: Optional[int] = 16
|
|
stop: Optional[Union[str, List[str]]] = Field(default_factory=list)
|
|
stream: Optional[bool] = False
|
|
presence_penalty: Optional[float] = 0.0
|
|
frequency_penalty: Optional[float] = 0.0
|
|
logit_bias: Optional[Dict[str, float]] = None
|
|
user: Optional[str] = None
|
|
best_of: Optional[int] = None
|
|
|
|
|
|
class ChatMessage(BaseModel):
|
|
role: Optional[str] = None
|
|
content: Optional[str] = None
|
|
|
|
|
|
class ChatCompletionResponseChoice(BaseModel):
|
|
index: int
|
|
message: ChatMessage
|
|
finish_reason: Optional[str] = None
|
|
|
|
|
|
class ChatCompletionResponse(BaseModel):
|
|
id: str
|
|
object: str = "chat.completion"
|
|
created: int = Field(default_factory=lambda: int(time.time()))
|
|
model: str
|
|
choices: List[ChatCompletionResponseChoice]
|
|
usage: UsageInfo
|
|
|
|
|
|
class DeltaMessage(BaseModel):
|
|
role: Optional[str] = None
|
|
content: Optional[str] = None
|
|
|
|
|
|
class ChatCompletionResponseStreamChoice(BaseModel):
|
|
index: int
|
|
delta: DeltaMessage
|
|
finish_reason: Optional[str] = None
|
|
|
|
|
|
class ChatCompletionStreamResponse(BaseModel):
|
|
id: str
|
|
object: str = "chat.completion.chunk"
|
|
created: int = Field(default_factory=lambda: int(time.time()))
|
|
model: str
|
|
choices: List[ChatCompletionResponseStreamChoice]
|