27 lines
917 B
Python
27 lines
917 B
Python
from uuid import uuid4
|
|
from time import time
|
|
from pydantic import BaseModel, Field
|
|
from typing import List, Optional, Union
|
|
from OAI.types.common import LogProbs, UsageStats, CommonCompletionRequest
|
|
|
|
class CompletionRespChoice(BaseModel):
|
|
# Index is 0 since we aren't using multiple choices
|
|
index: int = 0
|
|
finish_reason: str
|
|
logprobs: Optional[LogProbs] = None
|
|
text: str
|
|
|
|
# Inherited from common request
|
|
class CompletionRequest(CommonCompletionRequest):
|
|
# Prompt can also contain token ids, but that's out of scope for this project.
|
|
prompt: Union[str, List[str]]
|
|
|
|
class CompletionResponse(BaseModel):
|
|
id: str = Field(default_factory=lambda: f"cmpl-{uuid4().hex}")
|
|
choices: List[CompletionRespChoice]
|
|
created: int = Field(default_factory=lambda: int(time()))
|
|
model: str
|
|
object: str = "text_completion"
|
|
|
|
# TODO: Add usage stats
|
|
usage: Optional[UsageStats] = None
|