138 lines
3.3 KiB
Python
138 lines
3.3 KiB
Python
from dataclasses import dataclass, field
|
|
from typing import List, Optional, Dict, Any
|
|
|
|
|
|
@dataclass
|
|
class Message:
|
|
role: str
|
|
content: str
|
|
|
|
def to_dict(self) -> Dict[str, str]:
|
|
return {
|
|
'role': self.role,
|
|
'content': self.content,
|
|
}
|
|
|
|
@staticmethod
|
|
def from_dict(data: Dict[str, str]) -> 'Message':
|
|
return Message(role=data['role'], content=data['content'])
|
|
|
|
|
|
@dataclass
|
|
class Delta:
|
|
role: Optional[str] = None
|
|
content: Optional[str] = None
|
|
|
|
def to_dict(self) -> Dict[str, Optional[str]]:
|
|
result = {}
|
|
if self.role is not None:
|
|
result['role'] = self.role
|
|
if self.content is not None:
|
|
result['content'] = self.content
|
|
return result
|
|
|
|
|
|
@dataclass
|
|
class Choice:
|
|
index: int
|
|
message: Message
|
|
finish_reason: str
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return {
|
|
'index': self.index,
|
|
'message': self.message.to_dict(),
|
|
'finish_reason': self.finish_reason,
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class ChatCompletionResponse:
|
|
id: str
|
|
object: str
|
|
created: int
|
|
model: str
|
|
choices: List[Choice] = field(default_factory=list)
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return {
|
|
'id': self.id,
|
|
'object': self.object,
|
|
'created': self.created,
|
|
'model': self.model,
|
|
'choices': [c.to_dict() for c in self.choices],
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class ChunkChoice:
|
|
index: int
|
|
delta: Delta
|
|
finish_reason: Optional[Any] = None
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return {
|
|
'index': self.index,
|
|
'delta': self.delta.to_dict(),
|
|
'finish_reason': self.finish_reason,
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class ChunkResponse:
|
|
id: str
|
|
object: str
|
|
created: int
|
|
model: str
|
|
choices: List[ChunkChoice] = field(default_factory=list)
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return {
|
|
'id': self.id,
|
|
'object': self.object,
|
|
'created': self.created,
|
|
'model': self.model,
|
|
'choices': [c.to_dict() for c in self.choices],
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class ChatCompletionRequest:
|
|
model: str
|
|
messages: List[Message]
|
|
stream: bool = False
|
|
return_images: bool = False
|
|
temperature: float = 0.0
|
|
web_search_options: Optional[Dict[str, str]] = None
|
|
thinking_enabled: bool = False # Enable reasoning by default
|
|
search_enabled: bool = False # Enable search by default
|
|
|
|
@staticmethod
|
|
def from_dict(data: Dict[str, Any]) -> 'ChatCompletionRequest':
|
|
messages = [Message.from_dict(m) for m in data.get('messages', [])]
|
|
web_search_options = data.get('web_search_options', {})
|
|
return ChatCompletionRequest(
|
|
model=data.get('model', 'r1'),
|
|
messages=messages,
|
|
stream=data.get('stream', False),
|
|
return_images=data.get('return_images', False),
|
|
temperature=data.get('temperature', 0.0),
|
|
web_search_options=web_search_options,
|
|
thinking_enabled=False,
|
|
search_enabled=False,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class Model:
|
|
id: str
|
|
object: str
|
|
owned_by: str
|
|
|
|
def to_dict(self) -> Dict[str, str]:
|
|
return {
|
|
'id': self.id,
|
|
'object': self.object,
|
|
'owned_by': self.owned_by,
|
|
}
|