Change max_tokens from u16 to u32 for larger context window in recent…

… models #172 (#174)
Abraxas-365 · Jun 23, 2024 · 14533be · 14533be
1 parent 4472bbf
commit 14533be
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/src/chain/options.rs b/src/chain/options.rs
@@ -4,7 +4,7 @@ use std::pin::Pin;
 use crate::language_models::options::CallOptions;
 
 pub struct ChainCallOptions {
-    pub max_tokens: Option<u16>,
+    pub max_tokens: Option<u32>,
     pub temperature: Option<f32>,
     pub stop_words: Option<Vec<String>>,
     pub streaming_func: Option<
@@ -76,7 +76,7 @@ impl ChainCallOptions {
         llm_option
     }
 
-    pub fn with_max_tokens(mut self, max_tokens: u16) -> Self {
+    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
         self.max_tokens = Some(max_tokens);
         self
     }

diff --git a/src/language_models/options.rs b/src/language_models/options.rs
@@ -7,7 +7,7 @@ use crate::schemas::{FunctionCallBehavior, FunctionDefinition};
 #[derive(Clone)]
 pub struct CallOptions {
     pub candidate_count: Option<usize>,
-    pub max_tokens: Option<u16>,
+    pub max_tokens: Option<u32>,
     pub temperature: Option<f32>,
     pub stop_words: Option<Vec<String>>,
     pub streaming_func: Option<
@@ -56,7 +56,7 @@ impl CallOptions {
     }
 
     // Refactored "with" functions as methods of CallOptions
-    pub fn with_max_tokens(mut self, max_tokens: u16) -> Self {
+    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
         self.max_tokens = Some(max_tokens);
         self
     }

diff --git a/src/llm/claude/models.rs b/src/llm/claude/models.rs
@@ -29,7 +29,7 @@ impl ClaudeMessage {
 pub(crate) struct Payload {
     pub model: String,
     pub messages: Vec<ClaudeMessage>,
-    pub max_tokens: u16,
+    pub max_tokens: u32,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub stream: Option<bool>,
     #[serde(skip_serializing_if = "Option::is_none")]