88 lines
2.1 KiB
Python
88 lines
2.1 KiB
Python
import time
|
|
import uuid
|
|
from typing import Any, Dict, List, Literal, Optional
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
class ChatMessage(BaseModel):
|
|
role: Literal["system", "user", "assistant", "tool"]
|
|
content: str
|
|
|
|
|
|
class ChatCompletionRequest(BaseModel):
|
|
model: str = "doubao-pro"
|
|
messages: List[ChatMessage]
|
|
temperature: Optional[float] = 0.7
|
|
max_tokens: Optional[int] = None
|
|
stream: Optional[bool] = False
|
|
top_p: Optional[float] = 1.0
|
|
frequency_penalty: Optional[float] = 0.0
|
|
presence_penalty: Optional[float] = 0.0
|
|
stop: Optional[Any] = None
|
|
user: Optional[str] = None
|
|
new_chat: Optional[bool] = False
|
|
|
|
|
|
class CompletionUsage(BaseModel):
|
|
prompt_tokens: int = 0
|
|
completion_tokens: int = 0
|
|
total_tokens: int = 0
|
|
|
|
|
|
class ChoiceMessage(ChatMessage):
|
|
pass
|
|
|
|
|
|
class Choice(BaseModel):
|
|
index: int = 0
|
|
message: ChoiceMessage
|
|
finish_reason: str = "stop"
|
|
|
|
|
|
class ChunkDelta(BaseModel):
|
|
role: Optional[str] = None
|
|
content: Optional[str] = None
|
|
|
|
|
|
class ChunkChoice(BaseModel):
|
|
index: int = 0
|
|
delta: ChunkDelta = ChunkDelta()
|
|
finish_reason: Optional[str] = None
|
|
|
|
|
|
class ChatCompletionResponse(BaseModel):
|
|
id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:24]}")
|
|
object: str = "chat.completion"
|
|
created: int = Field(default_factory=lambda: int(time.time()))
|
|
model: str = "doubao-pro"
|
|
choices: List[Choice]
|
|
usage: CompletionUsage = CompletionUsage()
|
|
|
|
|
|
class ChatCompletionChunk(BaseModel):
|
|
id: str
|
|
object: str = "chat.completion.chunk"
|
|
created: int = Field(default_factory=lambda: int(time.time()))
|
|
model: str = "doubao-pro"
|
|
choices: List[ChunkChoice]
|
|
|
|
|
|
class ModelInfo(BaseModel):
|
|
id: str
|
|
object: str = "model"
|
|
created: int = Field(default_factory=lambda: int(time.time()))
|
|
owned_by: str = "doubao"
|
|
|
|
|
|
class ModelListResponse(BaseModel):
|
|
object: str = "list"
|
|
data: List[ModelInfo]
|
|
|
|
|
|
SUPPORTED_MODELS = [
|
|
{"id": "doubao-pro", "owned_by": "doubao"},
|
|
{"id": "doubao-lite", "owned_by": "doubao"},
|
|
{"id": "doubao-pro-32k", "owned_by": "doubao"},
|
|
]
|