1 files changed, 21 insertions, 0 deletions
diff --git a/src/client/model.rs b/src/client/model.rs
index 88d3294..b29166e 100644
--- a/src/client/model.rs
+++ b/src/client/model.rs
@@ -12,6 +12,7 @@ pub struct Model {
     pub client_name: String,
     pub name: String,
     pub max_tokens: Option<usize>,
+    pub extra_fields: Option<serde_json::Map<String, serde_json::Value>>,
     pub tokens_count_factors: TokensCountFactors,
     pub capabilities: ModelCapabilities,
 }
@@ -27,6 +28,7 @@ impl Model {
         Self {
             client_name: client_name.into(),
             name: name.into(),
+            extra_fields: None,
             max_tokens: None,
             tokens_count_factors: Default::default(),
             capabilities: ModelCapabilities::Text,
@@ -73,6 +75,14 @@ impl Model {
         self
     }
 
+    pub fn set_extra_fields(
+        mut self,
+        extra_fields: Option<serde_json::Map<String, serde_json::Value>>,
+    ) -> Self {
+        self.extra_fields = extra_fields;
+        self
+    }
+
     pub fn set_max_tokens(mut self, max_tokens: Option<usize>) -> Self {
         match max_tokens {
             None | Some(0) => self.max_tokens = None,
@@ -122,12 +132,23 @@ impl Model {
         }
         Ok(())
     }
+
+    pub fn merge_extra_fields(&self, body: &mut serde_json::Value) {
+        if let (Some(body), Some(extra_fields)) = (body.as_object_mut(), &self.extra_fields) {
+            for (k, v) in extra_fields {
+                if !body.contains_key(k) {
+                    body.insert(k.clone(), v.clone());
+                }
+            }
+        }
+    }
 }
 
 #[derive(Debug, Clone, Deserialize)]
 pub struct ModelConfig {
     pub name: String,
     pub max_tokens: Option<usize>,
+    pub extra_fields: Option<serde_json::Map<String, serde_json::Value>>,
     #[serde(deserialize_with = "deserialize_capabilities")]
     #[serde(default = "default_capabilities")]
     pub capabilities: ModelCapabilities,