File size: 4,319 Bytes
79859e3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
from __future__ import annotations
from typing import Optional, List, Dict
from time import time
from .helper import filter_none
try:
from pydantic import BaseModel, Field
except ImportError:
class BaseModel():
@classmethod
def model_construct(cls, **data):
new = cls()
for key, value in data.items():
setattr(new, key, value)
return new
class Field():
def __init__(self, **config):
pass
class ChatCompletionChunk(BaseModel):
id: str
object: str
created: int
model: str
provider: Optional[str]
choices: List[ChatCompletionDeltaChoice]
@classmethod
def model_construct(
cls,
content: str,
finish_reason: str,
completion_id: str = None,
created: int = None
):
return super().model_construct(
id=f"chatcmpl-{completion_id}" if completion_id else None,
object="chat.completion.cunk",
created=created,
model=None,
provider=None,
choices=[ChatCompletionDeltaChoice.model_construct(
ChatCompletionDelta.model_construct(content),
finish_reason
)]
)
class ChatCompletionMessage(BaseModel):
role: str
content: str
@classmethod
def model_construct(cls, content: str):
return super().model_construct(role="assistant", content=content)
class ChatCompletionChoice(BaseModel):
index: int
message: ChatCompletionMessage
finish_reason: str
@classmethod
def model_construct(cls, message: ChatCompletionMessage, finish_reason: str):
return super().model_construct(index=0, message=message, finish_reason=finish_reason)
class ChatCompletion(BaseModel):
id: str
object: str
created: int
model: str
provider: Optional[str]
choices: List[ChatCompletionChoice]
usage: Dict[str, int] = Field(examples=[{
"prompt_tokens": 0, #prompt_tokens,
"completion_tokens": 0, #completion_tokens,
"total_tokens": 0, #prompt_tokens + completion_tokens,
}])
@classmethod
def model_construct(
cls,
content: str,
finish_reason: str,
completion_id: str = None,
created: int = None
):
return super().model_construct(
id=f"chatcmpl-{completion_id}" if completion_id else None,
object="chat.completion",
created=created,
model=None,
provider=None,
choices=[ChatCompletionChoice.model_construct(
ChatCompletionMessage.model_construct(content),
finish_reason
)],
usage={
"prompt_tokens": 0, #prompt_tokens,
"completion_tokens": 0, #completion_tokens,
"total_tokens": 0, #prompt_tokens + completion_tokens,
}
)
class ChatCompletionDelta(BaseModel):
role: str
content: str
@classmethod
def model_construct(cls, content: Optional[str]):
return super().model_construct(role="assistant", content=content)
class ChatCompletionDeltaChoice(BaseModel):
index: int
delta: ChatCompletionDelta
finish_reason: Optional[str]
@classmethod
def model_construct(cls, delta: ChatCompletionDelta, finish_reason: Optional[str]):
return super().model_construct(index=0, delta=delta, finish_reason=finish_reason)
class Image(BaseModel):
url: Optional[str]
b64_json: Optional[str]
revised_prompt: Optional[str]
@classmethod
def model_construct(cls, url: str = None, b64_json: str = None, revised_prompt: str = None):
return super().model_construct(**filter_none(
url=url,
b64_json=b64_json,
revised_prompt=revised_prompt
))
class ImagesResponse(BaseModel):
data: List[Image]
model: str
provider: str
created: int
@classmethod
def model_construct(cls, data: List[Image], created: int = None, model: str = None, provider: str = None):
if created is None:
created = int(time())
return super().model_construct(
data=data,
model=model,
provider=provider,
created=created
)
|