40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
from dataclasses import dataclass
|
|
|
|
from .var import Chat, LLMParam, Who
|
|
|
|
|
|
@dataclass()
|
|
class Payload:
|
|
system_prompt: str
|
|
user_message: str
|
|
temperature: float
|
|
max_tokens: int
|
|
stream: bool
|
|
|
|
def render(self) -> dict:
|
|
return {
|
|
LLMParam.MESSAGES: [
|
|
{Chat.ROLE: Who.SYSTEM, Chat.CONTENT: self.system_prompt},
|
|
{Chat.ROLE: Who.USER, Chat.CONTENT: self.user_message}
|
|
],
|
|
LLMParam.TEMPERATURE: self.temperature,
|
|
LLMParam.MAX_TOKENS: self.max_tokens,
|
|
LLMParam.STREAM: self.stream
|
|
}
|
|
|
|
@dataclass
|
|
class APIEndpoints:
|
|
api_base: str
|
|
chat_completion: str
|
|
|
|
def __post_init__(self):
|
|
self.api_base = self.api_base.rstrip("/")
|
|
|
|
for field_name in self.__dataclass_fields__:
|
|
if field_name == "api_base":
|
|
continue
|
|
|
|
path = getattr(self, field_name)
|
|
clean_path = f"/{path.lstrip('/')}"
|
|
|
|
setattr(self, field_name, f"{self.api_base}{clean_path}") |