Merge pull request #4 from gembancud/main

Add Claude 3.5 Sonnet to ClaudeModel
mochi-neko · Jun 30, 2024 · fc30a7d · fc30a7d
2 parents b54e8d0 + e9744fa
commit fc30a7d
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 7 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -26,8 +26,8 @@ full = ["macros"]
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-bytes = "1.5.*"
-reqwest = { version = "0.11.*", features = ["json", "stream"] }
+bytes = "1.6.*"
+reqwest = { version = "0.12.*", features = ["json", "stream"] }
 serde = { version = "1.0.*", features = ["derive"] }
 serde-json-fmt = "0.1.*"
 serde_json = "1.0.*"
@@ -37,9 +37,9 @@ futures-core = "0.3.*"
 clust_macros = { version = "0.8.0", optional = true }
 
 [dev-dependencies]
-anyhow = "1.0.80"
-clap = { version = "4.5.2", features = ["derive"] }
-tokio = { version = "1.36.0", features = ["macros", "rt-multi-thread", "fs"] }
+anyhow = "1.0.86"
+clap = { version = "4.5.7", features = ["derive"] }
+tokio = { version = "1.38.0", features = ["macros", "rt-multi-thread", "fs"] }
 futures-util = "0.3.30"
 tokio-stream = "0.1.15"
-base64 = "0.22.0"
+base64 = "0.22.1"
diff --git a/src/messages/claude_model.rs b/src/messages/claude_model.rs
@@ -15,6 +15,9 @@ pub enum ClaudeModel {
     // Claude 3 Haiku
     /// Claude 3 Haiku at 2024/03/07.
     Claude3Haiku20240307,
+    // Claude 3.5 Sonnet
+    /// Claude 3.5 Sonnet at 2024/06/20
+    Claude35Sonnet20240620,
 }
 
 impl Default for ClaudeModel {
@@ -38,6 +41,9 @@ impl Display for ClaudeModel {
             | ClaudeModel::Claude3Haiku20240307 => {
                 write!(f, "claude-3-haiku-20240307")
             },
+            | ClaudeModel::Claude35Sonnet20240620 => {
+                write!(f, "claude-3-5-sonnet-20240620")
+            },
         }
     }
 }
@@ -48,6 +54,7 @@ impl ClaudeModel {
             | ClaudeModel::Claude3Opus20240229 => 4096,
             | ClaudeModel::Claude3Sonnet20240229 => 4096,
             | ClaudeModel::Claude3Haiku20240307 => 4096,
+            | ClaudeModel::Claude35Sonnet20240620 => 4096,
         }
     }
 }
@@ -56,7 +63,8 @@ impl_enum_string_serialization!(
     ClaudeModel,
     Claude3Opus20240229 => "claude-3-opus-20240229",
     Claude3Sonnet20240229 => "claude-3-sonnet-20240229",
-    Claude3Haiku20240307 => "claude-3-haiku-20240307"
+    Claude3Haiku20240307 => "claude-3-haiku-20240307",
+    Claude35Sonnet20240620 => "claude-3-5-sonnet-20240620"
 );
 
 #[cfg(test)]
@@ -85,6 +93,10 @@ mod tests {
             ClaudeModel::Claude3Haiku20240307.to_string(),
             "claude-3-haiku-20240307"
         );
+        assert_eq!(
+            ClaudeModel::Claude35Sonnet20240620.to_string(),
+            "claude-3-5-sonnet-20240620"
+        );
     }
 
     #[test]
@@ -101,6 +113,10 @@ mod tests {
             ClaudeModel::Claude3Haiku20240307.max_tokens(),
             4096
         );
+        assert_eq!(
+            ClaudeModel::Claude35Sonnet20240620.max_tokens(),
+            4096
+        );
     }
 
     #[test]
@@ -120,6 +136,13 @@ mod tests {
                 .unwrap(),
             ClaudeModel::Claude3Haiku20240307
         );
+        assert_eq!(
+            serde_json::from_str::<ClaudeModel>(
+                "\"claude-3-5-sonnet-20240620\""
+            )
+            .unwrap(),
+            ClaudeModel::Claude35Sonnet20240620
+        );
     }
 
     #[test]
@@ -136,5 +159,10 @@ mod tests {
             serde_json::to_string(&ClaudeModel::Claude3Haiku20240307).unwrap(),
             "\"claude-3-haiku-20240307\""
         );
+        assert_eq!(
+            serde_json::to_string(&ClaudeModel::Claude35Sonnet20240620)
+                .unwrap(),
+            "\"claude-3-5-sonnet-20240620\""
+        );
     }
 }