53 lines
1.1 KiB
Python
53 lines
1.1 KiB
Python
from enum import Enum
|
|
from typing import List, Optional
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
class Choice(BaseModel):
|
|
index: int
|
|
text: str
|
|
|
|
|
|
class CompletionRequest(BaseModel):
|
|
prompt: str = Field(
|
|
example="def binarySearch(arr, left, right, x):\n mid = (left +",
|
|
description="The context to generate completions for, encoded as a string.",
|
|
)
|
|
|
|
|
|
class CompletionResponse(BaseModel):
|
|
id: str
|
|
created: int
|
|
choices: List[Choice]
|
|
|
|
|
|
class Event(BaseModel):
|
|
type: str
|
|
|
|
|
|
class CompletionEvent(Event):
|
|
id: str
|
|
prompt: str
|
|
created: int
|
|
choices: List[Choice]
|
|
|
|
@classmethod
|
|
def build(cls, request: CompletionRequest, response: CompletionResponse):
|
|
return cls(
|
|
type="completion",
|
|
id=response.id,
|
|
prompt=request.prompt,
|
|
created=response.created,
|
|
choices=response.choices,
|
|
)
|
|
|
|
|
|
class SelectionEvent(Event):
|
|
completion_id: str
|
|
choice_index: int
|
|
|
|
@classmethod
|
|
def build(cls, id, index):
|
|
return cls(type="selection", completion_id=id, choice_index=index)
|