From dbca136feaaf7f8b1182c4c3c90c32918b1d0bb3 Mon Sep 17 00:00:00 2001
From: Andrei Betlen <abetlen@gmail.com>
Date: Wed, 20 Sep 2023 15:38:26 -0400
Subject: [PATCH] Update llama_types and names to match openai api

---
 llama_cpp/llama_types.py | 99 ++++++++++++++++++++++++++++++++--------
 1 file changed, 80 insertions(+), 19 deletions(-)

diff --git a/llama_cpp/llama_types.py b/llama_cpp/llama_types.py
index 21a0d4f..6ee7ef9 100644
--- a/llama_cpp/llama_types.py
+++ b/llama_cpp/llama_types.py
@@ -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_extensions import TypedDict, NotRequired, Literal
 
@@ -7,16 +13,19 @@ class EmbeddingUsage(TypedDict):
     total_tokens: int
 
 
-class EmbeddingData(TypedDict):
+class Embedding(TypedDict):
     index: int
     object: str
     embedding: List[float]
 
 
-class Embedding(TypedDict):
+EmbeddingData = Embedding
+
+
+class CreateEmbeddingResponse(TypedDict):
     object: Literal["list"]
     model: str
-    data: List[EmbeddingData]
+    data: List[Embedding]
     usage: EmbeddingUsage
 
 
@@ -31,7 +40,7 @@ class CompletionChoice(TypedDict):
     text: str
     index: int
     logprobs: Optional[CompletionLogprobs]
-    finish_reason: Optional[str]
+    finish_reason: Optional[Literal["stop", "length"]]
 
 
 class CompletionUsage(TypedDict):
@@ -40,7 +49,7 @@ class CompletionUsage(TypedDict):
     total_tokens: int
 
 
-class CompletionChunk(TypedDict):
+class CreateCompletionStreamResponse(TypedDict):
     id: str
     object: Literal["text_completion"]
     created: int
@@ -48,7 +57,10 @@ class CompletionChunk(TypedDict):
     choices: List[CompletionChoice]
 
 
-class Completion(TypedDict):
+CompletionChunk = CreateCompletionStreamResponse
+
+
+class CreateCompletionResponse(TypedDict):
     id: str
     object: Literal["text_completion"]
     created: int
@@ -57,29 +69,43 @@ class Completion(TypedDict):
     usage: CompletionUsage
 
 
-class ChatCompletionMessage(TypedDict):
-    role: Literal["assistant", "user", "system"]
-    content: str
+Completion = CreateCompletionResponse
+
+
+class ChatCompletionFunctionCall(TypedDict):
+    name: str
+    arguments: str
+
+
+class ChatCompletionResponseMessage(TypedDict):
+    role: Literal["assistant", "user", "system", "function"]
+    content: Optional[str]
     user: NotRequired[str]
+    function_call: NotRequired[ChatCompletionFunctionCall]
 
 
-class ChatCompletionFunction(TypedDict):
+ChatCompletionMessage = ChatCompletionResponseMessage
+
+
+class ChatCompletionResponseFunction(TypedDict):
     name: str
     description: NotRequired[str]
     parameters: Dict[str, Any]  # TODO: make this more specific
 
 
-class ChatCompletionFunctionCall(TypedDict):
-    name: str
+ChatCompletionFunction = ChatCompletionResponseFunction
 
 
-class ChatCompletionChoice(TypedDict):
+class ChatCompletionResponseChoice(TypedDict):
     index: int
     message: ChatCompletionMessage
     finish_reason: Optional[str]
 
 
-class ChatCompletion(TypedDict):
+ChatCompletionChoice = ChatCompletionResponseChoice
+
+
+class CreateChatCompletionResponse(TypedDict):
     id: str
     object: Literal["chat.completion"]
     created: int
@@ -88,24 +114,59 @@ class ChatCompletion(TypedDict):
     usage: CompletionUsage
 
 
-class ChatCompletionChunkDeltaEmpty(TypedDict):
+ChatCompletion = CreateChatCompletionResponse
+
+
+class ChatCompletionStreamResponseDeltaEmpty(TypedDict):
     pass
 
 
-class ChatCompletionChunkDelta(TypedDict):
+ChatCompletionChunkDeltaEmpty = ChatCompletionStreamResponseDeltaEmpty
+
+
+class ChatCompletionStreamResponseDelta(TypedDict):
     role: NotRequired[Literal["assistant"]]
     content: NotRequired[str]
+    function_call: NotRequired[ChatCompletionFunctionCall]
 
 
-class ChatCompletionChunkChoice(TypedDict):
+ChatCompletionChunkDelta = ChatCompletionStreamResponseDelta
+
+
+class ChatCompletionStreamResponseChoice(TypedDict):
     index: int
     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
     model: str
     object: Literal["chat.completion.chunk"]
     created: int
     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]