This commit is contained in:
Timothy Jaeryang Baek
2024-12-12 23:31:08 -08:00
parent 9a081c8593
commit f9a05dd1e1
6 changed files with 190 additions and 144 deletions

View File

@@ -106,7 +106,7 @@ def openai_chat_message_template(model: str):
def openai_chat_chunk_message_template(
model: str, message: Optional[str] = None
model: str, message: Optional[str] = None, usage: Optional[dict] = None
) -> dict:
template = openai_chat_message_template(model)
template["object"] = "chat.completion.chunk"
@@ -114,17 +114,23 @@ def openai_chat_chunk_message_template(
template["choices"][0]["delta"] = {"content": message}
else:
template["choices"][0]["finish_reason"] = "stop"
if usage:
template["usage"] = usage
return template
def openai_chat_completion_message_template(
model: str, message: Optional[str] = None
model: str, message: Optional[str] = None, usage: Optional[dict] = None
) -> dict:
template = openai_chat_message_template(model)
template["object"] = "chat.completion"
if message is not None:
template["choices"][0]["message"] = {"content": message, "role": "assistant"}
template["choices"][0]["finish_reason"] = "stop"
if usage:
template["usage"] = usage
return template