Update llama_types and names to match openai api

This commit is contained in:
Andrei Betlen 2023-09-20 15:38:26 -04:00
parent 15000fca69
commit dbca136fea

View file

@ -1,3 +1,9 @@
"""Types and request signatrues for OpenAI compatibility
Based on the OpenAI OpenAPI specification:
https://github.com/openai/openai-openapi/blob/master/openapi.yaml
"""
from typing import Any, List, Optional, Dict, Union from typing import Any, List, Optional, Dict, Union
from typing_extensions import TypedDict, NotRequired, Literal from typing_extensions import TypedDict, NotRequired, Literal
@ -7,16 +13,19 @@ class EmbeddingUsage(TypedDict):
total_tokens: int total_tokens: int
class EmbeddingData(TypedDict): class Embedding(TypedDict):
index: int index: int
object: str object: str
embedding: List[float] embedding: List[float]
class Embedding(TypedDict): EmbeddingData = Embedding
class CreateEmbeddingResponse(TypedDict):
object: Literal["list"] object: Literal["list"]
model: str model: str
data: List[EmbeddingData] data: List[Embedding]
usage: EmbeddingUsage usage: EmbeddingUsage
@ -31,7 +40,7 @@ class CompletionChoice(TypedDict):
text: str text: str
index: int index: int
logprobs: Optional[CompletionLogprobs] logprobs: Optional[CompletionLogprobs]
finish_reason: Optional[str] finish_reason: Optional[Literal["stop", "length"]]
class CompletionUsage(TypedDict): class CompletionUsage(TypedDict):
@ -40,7 +49,7 @@ class CompletionUsage(TypedDict):
total_tokens: int total_tokens: int
class CompletionChunk(TypedDict): class CreateCompletionStreamResponse(TypedDict):
id: str id: str
object: Literal["text_completion"] object: Literal["text_completion"]
created: int created: int
@ -48,7 +57,10 @@ class CompletionChunk(TypedDict):
choices: List[CompletionChoice] choices: List[CompletionChoice]
class Completion(TypedDict): CompletionChunk = CreateCompletionStreamResponse
class CreateCompletionResponse(TypedDict):
id: str id: str
object: Literal["text_completion"] object: Literal["text_completion"]
created: int created: int
@ -57,29 +69,43 @@ class Completion(TypedDict):
usage: CompletionUsage usage: CompletionUsage
class ChatCompletionMessage(TypedDict): Completion = CreateCompletionResponse
role: Literal["assistant", "user", "system"]
content: str
class ChatCompletionFunctionCall(TypedDict):
name: str
arguments: str
class ChatCompletionResponseMessage(TypedDict):
role: Literal["assistant", "user", "system", "function"]
content: Optional[str]
user: NotRequired[str] user: NotRequired[str]
function_call: NotRequired[ChatCompletionFunctionCall]
class ChatCompletionFunction(TypedDict): ChatCompletionMessage = ChatCompletionResponseMessage
class ChatCompletionResponseFunction(TypedDict):
name: str name: str
description: NotRequired[str] description: NotRequired[str]
parameters: Dict[str, Any] # TODO: make this more specific parameters: Dict[str, Any] # TODO: make this more specific
class ChatCompletionFunctionCall(TypedDict): ChatCompletionFunction = ChatCompletionResponseFunction
name: str
class ChatCompletionChoice(TypedDict): class ChatCompletionResponseChoice(TypedDict):
index: int index: int
message: ChatCompletionMessage message: ChatCompletionMessage
finish_reason: Optional[str] finish_reason: Optional[str]
class ChatCompletion(TypedDict): ChatCompletionChoice = ChatCompletionResponseChoice
class CreateChatCompletionResponse(TypedDict):
id: str id: str
object: Literal["chat.completion"] object: Literal["chat.completion"]
created: int created: int
@ -88,24 +114,59 @@ class ChatCompletion(TypedDict):
usage: CompletionUsage usage: CompletionUsage
class ChatCompletionChunkDeltaEmpty(TypedDict): ChatCompletion = CreateChatCompletionResponse
class ChatCompletionStreamResponseDeltaEmpty(TypedDict):
pass pass
class ChatCompletionChunkDelta(TypedDict): ChatCompletionChunkDeltaEmpty = ChatCompletionStreamResponseDeltaEmpty
class ChatCompletionStreamResponseDelta(TypedDict):
role: NotRequired[Literal["assistant"]] role: NotRequired[Literal["assistant"]]
content: NotRequired[str] content: NotRequired[str]
function_call: NotRequired[ChatCompletionFunctionCall]
class ChatCompletionChunkChoice(TypedDict): ChatCompletionChunkDelta = ChatCompletionStreamResponseDelta
class ChatCompletionStreamResponseChoice(TypedDict):
index: int index: int
delta: Union[ChatCompletionChunkDelta, ChatCompletionChunkDeltaEmpty] delta: Union[ChatCompletionChunkDelta, ChatCompletionChunkDeltaEmpty]
finish_reason: Optional[str] finish_reason: Optional[Literal["stop", "length", "function_call"]]
class ChatCompletionChunk(TypedDict): ChatCompletionChunkChoice = ChatCompletionStreamResponseChoice
class ChatCompletionStreamResponse(TypedDict):
id: str id: str
model: str model: str
object: Literal["chat.completion.chunk"] object: Literal["chat.completion.chunk"]
created: int created: int
choices: List[ChatCompletionChunkChoice] choices: List[ChatCompletionChunkChoice]
ChatCompletionChunk = ChatCompletionStreamResponse
JsonType = Union[None, int, str, bool, List["JsonType"], Dict[str, "JsonType"]]
class ChatCompletionFunctions(TypedDict):
name: str
description: NotRequired[str]
parameters: Dict[str, JsonType] # TODO: make this more specific
class ChatCompletionFunctionCallOption(TypedDict):
name: str
class ChatCompletionRequestMessage(TypedDict):
role: Literal["assistant", "user", "system", "function"]
content: Optional[str]
name: NotRequired[str]
funcion_call: NotRequired[ChatCompletionFunctionCall]