feat: add model.extra_body support for chat completion requests

Author: dstoc

README.md

@@ -198,6 +198,22 @@ url = "https://generativelanguage.googleapis.com/v1beta/openai"
 api_key_env = "GEMINI_API_KEY"
 ```
 
+Provider-specific request fields can be added under `model.extra_body`. These are
+merged into the top-level chat-completions JSON body after `ut` assembles its
+standard `model`, `messages`, and `temperature` fields. Reserved keys cannot be
+overridden.
+
+For example, with `llama-server`:
+
+```toml
+[model]
+model = "meta-llama/Llama-3.3-70B-Instruct"
+url = "http://127.0.0.1:8080/v1"
+
+[model.extra_body]
+thinking_budget_tokens = 1024
+```
+
 ## Paste Safety
 
 `ut` captures Sway context at:

src/config.rs

@@ -1,11 +1,14 @@
 use crate::context::AppContext;
 use anyhow::{Context, Result};
 use serde::{Deserialize, Serialize};
+use serde_json::{Map, Value};
 use std::collections::BTreeMap;
 use std::env;
 use std::fs;
 use std::path::PathBuf;
 
+pub const RESERVED_CHAT_COMPLETION_BODY_KEYS: &[&str] = &["model", "messages", "temperature"];
+
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
 #[serde(default)]
 pub struct Config {
@@ -48,6 +51,7 @@ impl Config {
         if self.model.timeout_seconds == 0 {
             anyhow::bail!("model.timeout_seconds must be greater than 0");
         }
+        self.model.validate_extra_body()?;
         if self.status_ui.width == 0 {
             anyhow::bail!("status_ui.width must be greater than 0");
         }
@@ -123,6 +127,7 @@ pub struct ModelConfig {
     pub timeout_seconds: u64,
     pub api_key: Option<String>,
     pub api_key_env: Option<String>,
+    pub extra_body: Map<String, Value>,
 }
 
 impl Default for ModelConfig {
@@ -133,6 +138,7 @@ impl Default for ModelConfig {
             timeout_seconds: 60,
             api_key: None,
             api_key_env: None,
+            extra_body: Map::new(),
         }
     }
 }
@@ -150,6 +156,20 @@ impl ModelConfig {
                     .filter(|value| !value.is_empty())
             })
     }
+
+    pub fn validate_extra_body(&self) -> Result<()> {
+        if let Some(key) = self
+            .extra_body
+            .keys()
+            .find(|key| RESERVED_CHAT_COMPLETION_BODY_KEYS.contains(&key.as_str()))
+        {
+            anyhow::bail!(
+                "model.extra_body must not override reserved chat-completion field {key:?}"
+            );
+        }
+
+        Ok(())
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -280,6 +300,7 @@ mod tests {
         assert_eq!(config.model.url, "http://127.0.0.1:11434/v1");
         assert_eq!(config.model.api_key, None);
         assert_eq!(config.model.api_key_env, None);
+        assert!(config.model.extra_body.is_empty());
     }
 
     #[test]
@@ -322,6 +343,39 @@ mod tests {
         assert!(config.validate().is_err());
     }
 
+    #[test]
+    fn validate_rejects_reserved_model_extra_body_key() {
+        let mut config = Config::default();
+        config.model.extra_body.insert(
+            "messages".to_string(),
+            Value::String("not allowed".to_string()),
+        );
+
+        let error = config.validate().expect_err("config should be rejected");
+        assert!(error
+            .to_string()
+            .contains("model.extra_body must not override reserved chat-completion field"));
+    }
+
+    #[test]
+    fn parses_model_extra_body_table() {
+        let config: Config = toml::from_str(
+            r#"
+                [model]
+                model = "llama"
+
+                [model.extra_body]
+                thinking_budget_tokens = 1024
+            "#,
+        )
+        .expect("config should parse");
+
+        assert_eq!(
+            config.model.extra_body.get("thinking_budget_tokens"),
+            Some(&Value::from(1024))
+        );
+    }
+
     #[test]
     fn validate_rejects_zero_status_ui_dimensions() {
         let mut config = Config::default();

src/dictation.rs

@@ -1,12 +1,12 @@
 use crate::audio::{encode_wav_bytes, AudioPayload};
-use crate::config::ModelConfig;
+use crate::config::{ModelConfig, RESERVED_CHAT_COMPLETION_BODY_KEYS};
 use anyhow::{Context, Result};
 use async_trait::async_trait;
 use reqwest::header::{AUTHORIZATION, CONTENT_TYPE};
 use reqwest::Client;
 use serde::Deserialize;
 use serde::Serialize;
-use serde_json::Value;
+use serde_json::{Map, Value};
 use std::time::Duration;
 
 #[derive(Debug, Clone)]
@@ -57,7 +57,7 @@ impl HttpDictationClient {
             temperature: Some(0.0),
         };
 
-        let request_body = serde_json::to_vec(&body)?;
+        let request_body = build_request_body(&body, &self.model_config.extra_body)?;
         let response = send_http_request(
             &self.endpoint,
             &request_body,
@@ -164,6 +164,34 @@ struct ChatMessageResponse {
     content: Option<Value>,
 }
 
+fn build_request_body(
+    body: &ChatCompletionRequest,
+    extra_body: &Map<String, Value>,
+) -> Result<Vec<u8>> {
+    let mut request_body = serde_json::to_value(body)?
+        .as_object()
+        .cloned()
+        .context("chat completion request body must serialize as a JSON object")?;
+    merge_extra_body(&mut request_body, extra_body)?;
+    serde_json::to_vec(&request_body).context("failed to serialize chat completion request body")
+}
+
+fn merge_extra_body(
+    request_body: &mut Map<String, Value>,
+    extra_body: &Map<String, Value>,
+) -> Result<()> {
+    for (key, value) in extra_body {
+        if RESERVED_CHAT_COMPLETION_BODY_KEYS.contains(&key.as_str()) {
+            anyhow::bail!(
+                "model.extra_body must not override reserved chat-completion field {key:?}"
+            );
+        }
+        request_body.insert(key.clone(), value.clone());
+    }
+
+    Ok(())
+}
+
 fn extract_text(response: &ChatCompletionResponse) -> Result<String> {
     let choice = response
         .choices
@@ -315,4 +343,44 @@ mod tests {
             "https://api.openai.com/v1/chat/completions"
         );
     }
+
+    #[test]
+    fn request_body_merges_model_extra_body() {
+        let body = ChatCompletionRequest {
+            model: "llama".to_string(),
+            messages: vec![ChatMessage::system("prompt".to_string())],
+            temperature: Some(0.0),
+        };
+        let mut extra_body = Map::new();
+        extra_body.insert("thinking_budget_tokens".to_string(), Value::from(1024));
+
+        let request_body = build_request_body(&body, &extra_body).expect("request should build");
+        let request_json: Value =
+            serde_json::from_slice(&request_body).expect("request body should be valid json");
+
+        assert_eq!(
+            request_json.get("model"),
+            Some(&Value::String("llama".to_string()))
+        );
+        assert_eq!(
+            request_json.get("thinking_budget_tokens"),
+            Some(&Value::from(1024))
+        );
+    }
+
+    #[test]
+    fn request_body_rejects_reserved_model_extra_body_keys() {
+        let body = ChatCompletionRequest {
+            model: "llama".to_string(),
+            messages: vec![ChatMessage::system("prompt".to_string())],
+            temperature: Some(0.0),
+        };
+        let mut extra_body = Map::new();
+        extra_body.insert("messages".to_string(), Value::Array(Vec::new()));
+
+        let error = build_request_body(&body, &extra_body).expect_err("merge should fail");
+        assert!(error
+            .to_string()
+            .contains("model.extra_body must not override reserved chat-completion field"));
+    }
 }