aurora

1
import time
2
from enum import Enum
3
from pydantic import BaseModel, Field
4
from typing import List, Optional
5

6

7
class Role(str, Enum):
8
    USER = "user"
9
    ASSISTANT = "assistant"
10
    SYSTEM = "system"
11

12

13
class Finish(str, Enum):
14
    STOP = "stop"
15
    LENGTH = "length"
16

17

18
class ModelCard(BaseModel):
19
    id: str
20
    object: Optional[str] = "model"
21
    created: Optional[int] = Field(default_factory=lambda: int(time.time()))
22
    owned_by: Optional[str] = "owner"
23

24

25
class ModelList(BaseModel):
26
    object: Optional[str] = "list"
27
    data: Optional[List[ModelCard]] = []
28

29

30
class ChatMessage(BaseModel):
31
    role: Role
32
    content: str
33

34

35
class DeltaMessage(BaseModel):
36
    role: Optional[Role] = None
37
    content: Optional[str] = None
38

39

40
class ChatCompletionRequest(BaseModel):
41
    model: str
42
    messages: List[ChatMessage]
43
    do_sample: Optional[bool] = True
44
    temperature: Optional[float] = None
45
    top_p: Optional[float] = None
46
    n: Optional[int] = 1
47
    max_tokens: Optional[int] = None
48
    stream: Optional[bool] = False
49

50

51
class ChatCompletionResponseChoice(BaseModel):
52
    index: int
53
    message: ChatMessage
54
    finish_reason: Finish
55

56

57
class ChatCompletionResponseStreamChoice(BaseModel):
58
    index: int
59
    delta: DeltaMessage
60
    finish_reason: Optional[Finish] = None
61

62

63
class ChatCompletionResponseUsage(BaseModel):
64
    prompt_tokens: int
65
    completion_tokens: int
66
    total_tokens: int
67

68

69
class ChatCompletionResponse(BaseModel):
70
    id: Optional[str] = "chatcmpl-default"
71
    object: Optional[str] = "chat.completion"
72
    created: Optional[int] = Field(default_factory=lambda: int(time.time()))
73
    model: str
74
    choices: List[ChatCompletionResponseChoice]
75
    usage: ChatCompletionResponseUsage
76

77

78
class ChatCompletionStreamResponse(BaseModel):
79
    id: Optional[str] = "chatcmpl-default"
80
    object: Optional[str] = "chat.completion.chunk"
81
    created: Optional[int] = Field(default_factory=lambda: int(time.time()))
82
    model: str
83
    choices: List[ChatCompletionResponseStreamChoice]
84

85

86
class ScoreEvaluationRequest(BaseModel):
87
    model: str
88
    messages: List[str]
89
    max_length: Optional[int] = None
90

91

92
class ScoreEvaluationResponse(BaseModel):
93
    id: Optional[str] = "scoreeval-default"
94
    object: Optional[str] = "score.evaluation"
95
    model: str
96
    scores: List[float]
97