feat: use dynamic batch size for embedding (#826)

1 files changed, 11 insertions, 6 deletions

diff --git a/src/client/model.rs b/src/client/model.rs
index 50dabeb..08ced10 100644
--- a/src/client/model.rs
+++ b/src/client/model.rs
@@ -157,15 +157,15 @@ impl Model {
             }
             "embedding" => {
                 let ModelData {
-                    max_input_tokens,
                     input_price,
+                    max_tokens_per_chunk,
                     max_batch_size,
                     ..
                 } = &self.data;
-                let max_tokens = format_option_value(max_input_tokens);
+                let max_tokens = format_option_value(max_tokens_per_chunk);
+                let max_batch = format_option_value(max_batch_size);
                 let price = format_option_value(input_price);
-                let batch = format_option_value(max_batch_size);
-                format!("max-tokens:{max_tokens}; price:{price}; batch:{batch}")
+                format!("max-tokens:{max_tokens};max-batch:{max_batch};price:{price}")
             }
             _ => String::new(),
         }
@@ -183,12 +183,16 @@ impl Model {
         self.data.supports_vision
     }
 
+    pub fn max_tokens_per_chunk(&self) -> Option<usize> {
+        self.data.max_tokens_per_chunk
+    }
+
     pub fn default_chunk_size(&self) -> usize {
         self.data.default_chunk_size.unwrap_or(1000)
     }
 
-    pub fn max_batch_size(&self) -> usize {
-        self.data.max_batch_size.unwrap_or(1)
+    pub fn max_batch_size(&self) -> Option<usize> {
+        self.data.max_batch_size
     }
 
     pub fn max_tokens_param(&self) -> Option<isize> {
@@ -266,6 +270,7 @@ pub struct ModelData {
     pub supports_function_calling: bool,
 
     // embedding-only properties
+    pub max_tokens_per_chunk: Option<usize>,
     pub default_chunk_size: Option<usize>,
     pub max_batch_size: Option<usize>,
 }