Include usage stats

Sources/CleverBird/CleverBirdError.swift

@@ -2,10 +2,60 @@ import Foundation
 
 public enum CleverBirdError: Error, Equatable {
     case requestFailed(message: String)
+    case unauthorized(message: String)
+    case forbidden(message: String)
+    case proxyAuthenticationRequired(message: String)
     case responseParsingFailed(message: String)
     case tokenEncoderCreationFailed(message: String)
     case tokenEncodingError(message: String)
     case invalidMessageContent
     case invalidFunctionMessage
     case invalidEmbeddingRequest(message: String)
+    case tooManyRequests
+}
+
+extension CleverBirdError: LocalizedError {
+    public var errorDescription: String? {
+        switch self {
+        case .requestFailed(message: let message):
+            "Request failed. \(message)"
+        case .unauthorized(message: let message):
+            "Unauthorized. \(message)"
+        case .forbidden(message: let message):
+            "Forbidden. \(message)"
+        case .proxyAuthenticationRequired(message: let message):
+            "Proxy Authentication Required. \(message)"
+        case .responseParsingFailed(message: let message):
+            "Parsing failed. \(message)"
+        case .tokenEncoderCreationFailed(message: let message):
+            "Encoding failed. \(message)"
+        case .tokenEncodingError(message: let message):
+            "Token encoding failed. \(message)"
+        case .invalidMessageContent:
+            "Invalid message content"
+        case .invalidFunctionMessage:
+            "Invalid function message"
+        case .invalidEmbeddingRequest(message: let message):
+            "Invalid embeddings request. \(message)"
+        case .tooManyRequests:
+            "Too many requests"
+        }
+    }
+}
+
+extension CleverBirdError {
+    init(statusCode: Int) {
+        switch statusCode {
+        case 401:
+            self = .unauthorized(message: "")
+        case 403:
+            self = .forbidden(message: "")
+        case 407:
+            self = .proxyAuthenticationRequired(message: "")
+        case 429:
+            self = .tooManyRequests
+        default:
+            self = .requestFailed(message: "Response status code: \(statusCode)")
+        }
+    }
 }

Sources/CleverBird/OpenAIAPIConnection.swift

@@ -9,11 +9,13 @@ public class OpenAIAPIConnection {
     let organization: String?
     public let client: APIClient
     let requestHeaders: [String:String]
+    let chatCompletionPath: String
 
     public init(apiKey: String,
                 organization: String? = nil,
                 scheme: String = "https",
                 host: String = "api.openai.com",
+                chatCompletionPath: String? =  nil,
                 port: Int = 443) {
         self.apiKey = apiKey
         self.organization = organization
@@ -23,6 +25,8 @@ public class OpenAIAPIConnection {
         urlComponents.host = host
         urlComponents.port = port
         let openAIAPIURL = urlComponents.url
+
+        self.chatCompletionPath = chatCompletionPath ?? "/v1/chat/completions"
 
         let clientConfiguration = APIClient.Configuration(baseURL: openAIAPIURL)
         clientConfiguration.encoder.keyEncodingStrategy = .convertToSnakeCase

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
     }
 }

Sources/CleverBird/chat/ChatCompletionResponse.swift

@@ -2,9 +2,15 @@
 
 import Foundation
 
-struct ChatCompletionResponse: Codable, Identifiable {
+public struct Usage: Codable {
+    public let promptTokens: Int
+    public let completionTokens: Int
+    public let totalTokens: Int
+}
+
+public struct ChatCompletionResponse: Codable, Identifiable {
 
-    struct Choice: Codable {
+    public struct Choice: Codable {
 
         enum FinishReason: String, Codable {
             case stop
@@ -23,23 +29,17 @@ 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
+    public let id: String
 
     enum CodingKeys: String, CodingKey {
         case choices
         case usage
         case id
     }
 
-    init(from decoder: Decoder) throws {
+    public init(from decoder: Decoder) throws {
         let container = try decoder.container(keyedBy: CodingKeys.self)
 
         self.id = try container.decode(String.self, forKey: .id)

Sources/CleverBird/chat/ChatMessage.swift

@@ -24,7 +24,7 @@ public struct ChatMessage: Codable, Identifiable {
     public let role: Role
 
     /// The contents of the message. `content` is required for all messages except assistant messages with function calls.
-    public let content: String?
+    public let content: Content?
 
     /// The name and arguments of a function that should be called, as generated by the model.
     public let functionCall: FunctionCall?
@@ -36,14 +36,21 @@ public struct ChatMessage: Codable, Identifiable {
                 content: String? = nil,
                 id: String? = nil,
                 functionCall: FunctionCall? = nil) throws {
+        try self.init(role: role, media: content != nil ? .text(content!) : nil, id: id, functionCall: functionCall)
+    }
 
+    public init(role: Role,
+                media: ChatMessage.Content?,
+                id: String? = nil,
+                functionCall: FunctionCall? = nil) throws {
+
         // Validation: Content is required for all messages except assistant messages with function calls.
-        if content == nil && !(role == .assistant && functionCall != nil) {
+        if media == nil && !(role == .assistant && functionCall != nil) {
             throw CleverBirdError.invalidMessageContent
         }
 
         self.role = role
-        self.content = content
+        self.content = media
         self.name = functionCall?.name
         if role == .function {
             // If the role is "function" I need to set functionCall to nil, otherwise this will
@@ -58,7 +65,9 @@ public struct ChatMessage: Codable, Identifiable {
         } else {
             var hasher = Hasher()
             hasher.combine(self.role)
-            hasher.combine(self.content ?? "")
+            if let content {
+                hasher.combine(content)
+            }
             let hashValue = abs(hasher.finalize())
             let timestamp = Int(Date.now.timeIntervalSince1970*10000)
 
@@ -69,7 +78,7 @@ public struct ChatMessage: Codable, Identifiable {
     public init(from decoder: Decoder) throws {
         let container = try decoder.container(keyedBy: CodingKeys.self)
         self.role = try container.decode(Role.self, forKey: .role)
-        self.content = try container.decodeIfPresent(String.self, forKey: .content)
+        self.content = try container.decodeIfPresent(Content.self, forKey: .content)
         self.functionCall = try container.decodeIfPresent(FunctionCall.self, forKey: .functionCall)
         self.name = try container.decodeIfPresent(String.self, forKey: .name)
         self.id = "pending"
@@ -92,3 +101,59 @@ extension ChatMessage: Equatable {
             && lhs.content == rhs.content
     }
 }
+
+extension ChatMessage {
+
+    public enum Content: Codable, Equatable, CustomStringConvertible, Hashable {
+
+        case text(String)
+        case media([MessageContent])
+
+        public init(from decoder: Decoder) throws {
+            let container = try decoder.singleValueContainer()
+            if let textContent = try? container.decode(String.self) {
+                self = .text(textContent)
+            } else if let chatContents = try? container.decode([MessageContent].self) {
+                self = .media(chatContents)
+            } else {
+                throw DecodingError.typeMismatch(MessageContent.self, DecodingError.Context(codingPath: decoder.codingPath, debugDescription: "Unsupported type for Content"))
+            }
+        }
+
+        public func encode(to encoder: Encoder) throws {
+            var container = encoder.singleValueContainer()
+            switch self {
+            case .text(let text):
+                try container.encode(text)
+            case .media(let contents):
+                try container.encode(contents)
+            }
+        }
+
+        public static func == (lhs: Content, rhs: Content) -> Bool {
+            switch (lhs, rhs) {
+            case (.text(let leftText), .text(let rightText)):
+                return leftText == rightText
+            case (.media(let leftMedia), .media(let rightMedia)):
+                return leftMedia == rightMedia
+            default:
+                return false
+            }
+        }
+
+        public var description: String {
+            switch self {
+            case .media(let messageContents):
+                for messageContent in messageContents {
+                    if case .text(let textValue) = messageContent {
+                        return textValue
+                    }
+                }
+                return ""
+            case .text(let text):
+                return text
+            }
+        }
+    }
+}
+

Sources/CleverBird/chat/ChatModel.swift

@@ -1,6 +1,7 @@
 public enum ChatModel: Codable {
     case gpt35Turbo
     case gpt4
+    case gpt4Turbo
     case specific(String)
 
     public init(from decoder: Decoder) throws {
@@ -11,7 +12,7 @@ public enum ChatModel: Codable {
         case _ where modelString.starts(with: "gpt-3.5"):
             self = .gpt35Turbo
         case _ where modelString.starts(with: "gpt-4"):
-            self = .gpt4
+            self = .gpt4Turbo
         default:
             self = .specific(modelString)
         }
@@ -23,18 +24,23 @@ public enum ChatModel: Codable {
 
         try container.encode(modelString)
     }
-}
-
-extension ChatModel: CustomStringConvertible {
-    public var description: String {
+
+    public var modelString: String {
         switch self {
         case .gpt35Turbo:
             return "gpt-3.5-turbo"
         case .gpt4:
             return "gpt-4"
-        case .specific(let string):
-            return string
+        case .gpt4Turbo:
+            return "gpt-4-turbo"
+        case .specific(let specificString):
+            return specificString
         }
-
+    }
+}
+
+extension ChatModel: CustomStringConvertible {
+    public var description: String {
+        return modelString
     }
 }

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)

Sources/CleverBird/chat/ChatThread+tokenCount.swift

@@ -17,7 +17,7 @@ extension ChatThread {
         switch model {
         case .gpt35Turbo:
             tokensPerMessage = 4
-        case .gpt4:
+        case .gpt4, .gpt4Turbo:
             tokensPerMessage = 3
         case .specific(_):
             tokensPerMessage = 3
@@ -29,7 +29,32 @@ extension ChatThread {
                 let roleTokens = try tokenEncoder.encode(text: message.role.rawValue).count
                 let contentTokens: Int
                 if let content = message.content {
-                    contentTokens = try tokenEncoder.encode(text: content).count
+                    switch content {
+                    case .text(let text):
+                        contentTokens = try tokenEncoder.encode(text: text).count
+                    case .media(let media):
+                        var count = 0
+                        for medium in media {
+                            switch medium {
+                            case .text(let text):
+                                count += try tokenEncoder.encode(text: text).count
+                            case .imageUrl(let url):
+                                // See https://platform.openai.com/docs/guides/vision/calculating-costs
+                                switch url.detail {
+                                    // TODO: calculate real values for auto and high
+                                case .auto:
+                                    count += 1105
+                                case .high:
+                                    count += 1105
+                                case .low:
+                                    count += 85
+                                case .none:
+                                    count += 1105
+                                }
+                            }
+                        }
+                        contentTokens = count
+                    }
                 } else if let functionCall = message.functionCall {
                     let jsonEncoder = JSONEncoder()
                     let jsonData = try jsonEncoder.encode(functionCall)