From fae63f52b265f46e6030d2cce830e2d1fca6597b Mon Sep 17 00:00:00 2001
From: Ronald Mannak <ronaldmannak@me.com>
Date: Tue, 14 May 2024 21:54:18 -0700
Subject: [PATCH] Add gpt-4o and gpt-4-turbo models (#20)

---
 README.md                                              |  2 +-
 Sources/CleverBird/chat/ChatModel.swift                | 10 ++++++++++
 Sources/CleverBird/chat/ChatThread+complete.swift      |  4 ++--
 Sources/CleverBird/chat/ChatThread+tokenCount.swift    |  4 +++-
 .../chat/streaming/StreamableChatThread+complete.swift |  2 +-
 5 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index dfde867..9871072 100644
--- a/README.md
+++ b/README.md
@@ -59,7 +59,7 @@ The `complete(using:)` method also includes various optional parameters:
 ```swift
 let completion = chatThread.complete(
     using: openAIAPIConnection, 
-    model: .gpt4, 
+    model: .gpt4o, 
     temperature: 0.7, 
     maxTokens: 500
 )
diff --git a/Sources/CleverBird/chat/ChatModel.swift b/Sources/CleverBird/chat/ChatModel.swift
index 7dc69df..a531599 100644
--- a/Sources/CleverBird/chat/ChatModel.swift
+++ b/Sources/CleverBird/chat/ChatModel.swift
@@ -1,6 +1,8 @@
 public enum ChatModel: Codable {
     case gpt35Turbo
     case gpt4
+    case gpt4Turbo
+    case gpt4o
     case specific(String)
 
     public init(from decoder: Decoder) throws {
@@ -10,6 +12,10 @@ public enum ChatModel: Codable {
         switch modelString {
         case _ where modelString.starts(with: "gpt-3.5"):
             self = .gpt35Turbo
+        case _ where modelString.starts(with: "gpt-4o"):
+            self = .gpt4o
+        case _ where modelString.starts(with: "gpt-4-turbo"):
+            self = .gpt4Turbo
         case _ where modelString.starts(with: "gpt-4"):
             self = .gpt4
         default:
@@ -30,6 +36,10 @@ extension ChatModel: CustomStringConvertible {
         switch self {
         case .gpt35Turbo:
             return "gpt-3.5-turbo"
+        case .gpt4o:
+            return "gpt-4o"
+        case .gpt4Turbo:
+            return "gpt-4-turbo"
         case .gpt4:
             return "gpt-4"
         case .specific(let string):
diff --git a/Sources/CleverBird/chat/ChatThread+complete.swift b/Sources/CleverBird/chat/ChatThread+complete.swift
index bc7e600..9b48314 100644
--- a/Sources/CleverBird/chat/ChatThread+complete.swift
+++ b/Sources/CleverBird/chat/ChatThread+complete.swift
@@ -2,7 +2,7 @@
 
 extension ChatThread {
     public func complete(using connection: OpenAIAPIConnection,
-                         model: ChatModel = .gpt4,
+                         model: ChatModel = .gpt4o,
                          temperature: Percentage = 0.7,
                          topP: Percentage? = nil,
                          stop: [String]? = nil,
@@ -24,7 +24,7 @@ extension ChatThread {
     }
     
     public func completeIncludeUsage(using connection: OpenAIAPIConnection,
-                         model: ChatModel = .gpt4,
+                         model: ChatModel = .gpt4o,
                          temperature: Percentage = 0.7,
                          topP: Percentage? = nil,
                          stop: [String]? = nil,
diff --git a/Sources/CleverBird/chat/ChatThread+tokenCount.swift b/Sources/CleverBird/chat/ChatThread+tokenCount.swift
index e4caf88..e29a2e1 100644
--- a/Sources/CleverBird/chat/ChatThread+tokenCount.swift
+++ b/Sources/CleverBird/chat/ChatThread+tokenCount.swift
@@ -17,7 +17,9 @@ extension ChatThread {
         switch model {
         case .gpt35Turbo:
             tokensPerMessage = 4
-        case .gpt4:
+        case .gpt4, .gpt4Turbo:
+            tokensPerMessage = 3
+        case .gpt4o:
             tokensPerMessage = 3
         case .specific(_):
             tokensPerMessage = 3
diff --git a/Sources/CleverBird/chat/streaming/StreamableChatThread+complete.swift b/Sources/CleverBird/chat/streaming/StreamableChatThread+complete.swift
index bd0b2fe..4e035fe 100644
--- a/Sources/CleverBird/chat/streaming/StreamableChatThread+complete.swift
+++ b/Sources/CleverBird/chat/streaming/StreamableChatThread+complete.swift
@@ -5,7 +5,7 @@ import Foundation
 extension StreamableChatThread {
 
     public func complete(using connection: OpenAIAPIConnection,
-                         model: ChatModel = .gpt4,
+                         model: ChatModel = .gpt4o,
                          temperature: Percentage = 0.7,
                          topP: Percentage? = nil,
                          stop: [String]? = nil,