From 533b3fdb10f6ad3f05f73658528396f11fdb9c5d Mon Sep 17 00:00:00 2001
From: "B.T. Franklin" <brandon.franklin@gmail.com>
Date: Mon, 7 Aug 2023 18:14:39 -0700
Subject: [PATCH] Added finishReason and usage info to the
 ChatCompletionResponse.

---
 .../chat/ChatCompletionResponse.swift         | 32 +++++++++++++++++--
 Sources/CleverBird/chat/ChatMessage.swift     | 18 +++++------
 2 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/Sources/CleverBird/chat/ChatCompletionResponse.swift b/Sources/CleverBird/chat/ChatCompletionResponse.swift
index 2c36338..617a61a 100644
--- a/Sources/CleverBird/chat/ChatCompletionResponse.swift
+++ b/Sources/CleverBird/chat/ChatCompletionResponse.swift
@@ -3,36 +3,62 @@
 struct ChatCompletionResponse: Codable, Identifiable {
 
     struct Choice: Codable {
+
+        enum FinishReason: String, Codable {
+            case stop
+            case maxTokens
+        }
+
         let message: ChatMessage
+        let finishReason: FinishReason
         let functionCall: FunctionCall?
 
         enum CodingKeys: String, CodingKey {
             case message
+            case finishReason
             case functionCall
         }
     }
+
+    struct Usage: Codable {
+        let promptTokens: Int
+        let completionTokens: Int
+        let totalTokens: Int
+    }
+
     let choices: [Choice]
+    let usage: Usage
     let id: String
 
     enum CodingKeys: String, CodingKey {
         case choices
+        case usage
         case id
     }
 
     init(from decoder: Decoder) throws {
         let container = try decoder.container(keyedBy: CodingKeys.self)
-        id = try container.decode(String.self, forKey: .id)
+
+        self.id = try container.decode(String.self, forKey: .id)
+
         var choicesContainer = try container.nestedUnkeyedContainer(forKey: .choices)
         var choices: [Choice] = []
         while !choicesContainer.isAtEnd {
             let choiceContainer = try choicesContainer.nestedContainer(keyedBy: Choice.CodingKeys.self)
+
             var message = try choiceContainer.decode(ChatMessage.self, forKey: .message)
-            let functionCall = try choiceContainer.decodeIfPresent(FunctionCall.self, forKey: .functionCall)
             message.id = id
-            let choice = Choice(message: message, functionCall: functionCall)
+
+            let finishReason = try choiceContainer.decode(Choice.FinishReason.self, forKey: .finishReason)
+
+            let functionCall = try choiceContainer.decodeIfPresent(FunctionCall.self, forKey: .functionCall)
+
+            let choice = Choice(message: message, finishReason: finishReason, functionCall: functionCall)
             choices.append(choice)
         }
         self.choices = choices
+
+        self.usage = try container.decode(Usage.self, forKey: .usage)
     }
 
 }
diff --git a/Sources/CleverBird/chat/ChatMessage.swift b/Sources/CleverBird/chat/ChatMessage.swift
index 5ef4420..58e356e 100644
--- a/Sources/CleverBird/chat/ChatMessage.swift
+++ b/Sources/CleverBird/chat/ChatMessage.swift
@@ -69,19 +69,19 @@ public struct ChatMessage: Codable, Identifiable {
 
     public init(from decoder: Decoder) throws {
         let container = try decoder.container(keyedBy: CodingKeys.self)
-        role = try container.decode(Role.self, forKey: .role)
-        content = try container.decodeIfPresent(String.self, forKey: .content)
-        functionCall = try container.decodeIfPresent(FunctionCall.self, forKey: .functionCall)
-        name = try container.decodeIfPresent(String.self, forKey: .name)
-        id = "pending"
+        self.role = try container.decode(Role.self, forKey: .role)
+        self.content = try container.decodeIfPresent(String.self, forKey: .content)
+        self.functionCall = try container.decodeIfPresent(FunctionCall.self, forKey: .functionCall)
+        self.name = try container.decodeIfPresent(String.self, forKey: .name)
+        self.id = "pending"
     }
 
     public func encode(to encoder: Encoder) throws {
         var container = encoder.container(keyedBy: CodingKeys.self)
-        try container.encode(role, forKey: .role)
-        try container.encodeIfPresent(content, forKey: .content)
-        try container.encodeIfPresent(functionCall, forKey: .functionCall)
-        try container.encodeIfPresent(name, forKey: .name)
+        try container.encode(self.role, forKey: .role)
+        try container.encodeIfPresent(self.content, forKey: .content)
+        try container.encodeIfPresent(self.functionCall, forKey: .functionCall)
+        try container.encodeIfPresent(self.name, forKey: .name)
     }
 }