Include usage stats (#19)

* Add Content * Make ChatContent public * Rename ChatContent -> MessageContent * Update ChatMessage * Fix init * Add init for converting image data to base64 * Return emtpy string in case of no text * Get token usage data for streamed chat completion * Add streamOptions to complete method * Update readme * simplify includeUsage * Make Usage properties public * Add completeIncludeUsage method to chatThread * Remove default Swift headers
btfranklin · May 7, 2024 · 09c1477 · 09c1477
1 parent cf2f607
commit 09c1477
Show file tree

Hide file tree

Showing 11 changed files with 73 additions and 27 deletions.
diff --git a/README.md b/README.md
@@ -107,6 +107,19 @@ Just like with the non-streamed completion, the message will be automatically
 appended onto the thread after it has finished streaming, but the stream
 allows you to see it as it's coming through.
 
+To include usage (the number of tokens used in the prompt and completion), add set `streamOptions` in the `complete` method. The usage is available as a property of `StreamableChatThread` after the stream has completed.
+
+```swift
+let chatThread = ChatThread().withStreaming()
+let completionStream = try await chatThread.complete(using: openAIAPIConnection, includeUsage: true)
+for try await messageChunk in completionStream {
+    print("Received message chunk: \(messageChunk)")
+}
+if let usage = completionStream.usage {
+    print("Usage: \(usage)")
+}
+```
+
 Calculate the token count for messages in the chat thread:
 
 ```swift

diff --git a/Sources/CleverBird/chat/ChatCompletionRequestParameters.swift b/Sources/CleverBird/chat/ChatCompletionRequestParameters.swift
@@ -13,6 +13,7 @@ public struct ChatCompletionRequestParameters: Codable {
     public let messages: [ChatMessage]
     public let functions: [Function]?
     public let functionCallMode: FunctionCallMode?
+    public let streamOptions: StreamOptions?
 
     public init(model: ChatModel,
                 temperature: Percentage,
@@ -25,7 +26,8 @@ public struct ChatCompletionRequestParameters: Codable {
                 user: String? = nil,
                 messages: [ChatMessage],
                 functions: [Function]? = nil,
-                functionCallMode: FunctionCallMode? = nil) {
+                functionCallMode: FunctionCallMode? = nil,
+                streamOptions: StreamOptions? = nil) {
         self.model = model
         self.temperature = temperature
         self.topP = topP
@@ -38,5 +40,6 @@ public struct ChatCompletionRequestParameters: Codable {
         self.messages = messages
         self.functions = functions
         self.functionCallMode = functionCallMode
+        self.streamOptions = streamOptions
     }
 }
diff --git a/Sources/CleverBird/chat/ChatCompletionResponse.swift b/Sources/CleverBird/chat/ChatCompletionResponse.swift
@@ -2,6 +2,12 @@
 
 import Foundation
 
+public struct Usage: Codable {
+    public let promptTokens: Int
+    public let completionTokens: Int
+    public let totalTokens: Int
+}
+
 struct ChatCompletionResponse: Codable, Identifiable {
 
     struct Choice: Codable {
@@ -23,12 +29,6 @@ struct ChatCompletionResponse: Codable, Identifiable {
         }
     }
 
-    struct Usage: Codable {
-        let promptTokens: Int
-        let completionTokens: Int
-        let totalTokens: Int
-    }
-
     let choices: [Choice]
     let usage: Usage
     let id: String

diff --git a/Sources/CleverBird/chat/ChatThread+complete.swift b/Sources/CleverBird/chat/ChatThread+complete.swift
@@ -11,7 +11,28 @@ extension ChatThread {
                          frequencyPenalty: Penalty? = nil,
                          functions: [Function]? = nil,
                          functionCallMode: FunctionCallMode? = nil) async throws -> ChatMessage {
-
+        try await completeIncludeUsage(using: connection,
+                                    model: model,
+                                    temperature: temperature,
+                                    topP: topP,
+                                    stop: stop,
+                                    maxTokens: maxTokens,
+                                    presencePenalty: presencePenalty,
+                                    frequencyPenalty: frequencyPenalty,
+                                    functions: functions,
+                                    functionCallMode: functionCallMode).0
+    }
+
+    public func completeIncludeUsage(using connection: OpenAIAPIConnection,
+                         model: ChatModel = .gpt4,
+                         temperature: Percentage = 0.7,
+                         topP: Percentage? = nil,
+                         stop: [String]? = nil,
+                         maxTokens: Int? = nil,
+                         presencePenalty: Penalty? = nil,
+                         frequencyPenalty: Penalty? = nil,
+                         functions: [Function]? = nil,
+                         functionCallMode: FunctionCallMode? = nil) async throws -> (ChatMessage, Usage) {
         let requestBody = ChatCompletionRequestParameters(
             model: model,
             temperature: temperature,
@@ -47,7 +68,7 @@ extension ChatThread {
             // Append the response message to the thread
             addMessage(firstChoiceMessage)
 
-            return firstChoiceMessage
+            return (firstChoiceMessage, completion.usage)
 
         } catch {
             throw CleverBirdError.requestFailed(message: error.localizedDescription)

diff --git a/Sources/CleverBird/chat/MessageContent.swift b/Sources/CleverBird/chat/MessageContent.swift
@@ -1,9 +1,4 @@
-//
-//  ChatContent.swift
-//  
-//
 //  Created by Ronald Mannak on 4/12/24.
-//
 
 import Foundation
 

diff --git a/Sources/CleverBird/chat/streaming/ChatStreamedResponseChunk.swift b/Sources/CleverBird/chat/streaming/ChatStreamedResponseChunk.swift
@@ -19,6 +19,7 @@ struct ChatStreamedResponseChunk: Codable, Identifiable {
     }
     let choices: [Choice]
     let id: String
+    let usage: Usage?
 }
 
 extension ChatStreamedResponseChunk {

diff --git a/Sources/CleverBird/chat/streaming/StreamOptions.swift b/Sources/CleverBird/chat/streaming/StreamOptions.swift
@@ -0,0 +1,12 @@
+//  Created by Ronald Mannak on 5/6/24.
+
+import Foundation
+
+public struct StreamOptions: Codable {
+
+    let includeUsage: Bool
+
+    public init(includeUsage: Bool) {
+        self.includeUsage = includeUsage
+    }
+}
diff --git a/Sources/CleverBird/chat/streaming/StreamableChatThread+complete.swift b/Sources/CleverBird/chat/streaming/StreamableChatThread+complete.swift
@@ -11,7 +11,8 @@ extension StreamableChatThread {
                          stop: [String]? = nil,
                          maxTokens: Int? = nil,
                          presencePenalty: Penalty? = nil,
-                         frequencyPenalty: Penalty? = nil) async throws -> AsyncThrowingStream<String, Swift.Error> {
+                         frequencyPenalty: Penalty? = nil,
+                         includeUsage: Bool = false) async throws -> AsyncThrowingStream<String, Swift.Error> {
 
         let requestBody = ChatCompletionRequestParameters(
             model: model,
@@ -23,7 +24,8 @@ extension StreamableChatThread {
             presencePenalty: presencePenalty,
             frequencyPenalty: frequencyPenalty,
             user: self.chatThread.user,
-            messages: self.chatThread.messages
+            messages: self.chatThread.messages,
+            streamOptions:  includeUsage ? StreamOptions(includeUsage: true) : nil
         )
 
         // Define the callback closure that appends the message to the chat thread
@@ -71,6 +73,10 @@ extension StreamableChatThread {
 
                         responseMessageId = responseChunk.id
 
+                        if let usage = responseChunk.usage {
+                            strongSelf.usage = usage
+                        }
+
                         if let deltaRole = responseChunk.choices.first?.delta.role {
                             responseMessageRole = deltaRole
                             continue
@@ -89,6 +95,9 @@ extension StreamableChatThread {
                         } else {
                             responseMessageContent = deltaContent
                         }
+
+                        strongSelf.usage = responseChunk.usage
+
                         continuation.yield(deltaContent)
                     }
                     // Finished normally

diff --git a/Sources/CleverBird/chat/streaming/StreamableChatThread.swift b/Sources/CleverBird/chat/streaming/StreamableChatThread.swift
@@ -1,9 +1,11 @@
 //  Created by B.T. Franklin on 5/11/23
 
-public class StreamableChatThread {
+public final class StreamableChatThread {
 
     var streamingTask: Task<Void, Error>?
     let chatThread: ChatThread
+
+    public var usage: Usage? = nil
 
     init(chatThread: ChatThread) {
         self.chatThread = chatThread

diff --git a/Tests/CleverBirdTests/MessageContentTests.swift b/Tests/CleverBirdTests/MessageContentTests.swift
@@ -1,9 +1,4 @@
-//
-//  MessageContentTests.swift
-//  
-//
 //  Created by Ronald Mannak on 4/12/24.
-//
 
 import Foundation
 import XCTest

diff --git a/Tests/CleverBirdTests/MessageEncodingTests.swift b/Tests/CleverBirdTests/MessageEncodingTests.swift
@@ -1,9 +1,4 @@
-//
-//  ContentEncodingTests.swift
-//
-//
 //  Created by Ronald Mannak on 4/12/24.
-//
 
 import XCTest
 @testable import CleverBird