feat: add thinking budget support for Anthropic requests

Author: maan2003

llm_client/src/clients/anthropic.rs

@@ -227,6 +227,13 @@ enum ContentBlockDeltaType {
     },
 }
 
+#[derive(serde::Serialize, Debug, Clone)]
+struct AnthropicThinking {
+    #[serde(rename = "type")]
+    thinking_type: String,
+    budget_tokens: usize,
+}
+
 #[derive(serde::Serialize, Debug, Clone)]
 struct AnthropicRequest {
     system: Vec<AnthropicMessageContent>,
@@ -238,6 +245,8 @@ struct AnthropicRequest {
     stream: bool,
     #[serde(skip_serializing_if = "Option::is_none")]
     max_tokens: Option<usize>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    thinking: Option<AnthropicThinking>,
     model: String,
 }
 
@@ -251,6 +260,12 @@ impl AnthropicRequest {
             Some(tokens) => Some(tokens),
             None => Some(8192),
         };
+        let thinking = completion_request
+            .thinking_budget()
+            .map(|budget| AnthropicThinking {
+                thinking_type: "enabled".to_owned(),
+                budget_tokens: budget,
+            });
         let messages = completion_request.messages();
         // grab the tools over here ONLY from the system message
         let tools = messages
@@ -335,6 +350,7 @@ impl AnthropicRequest {
             tools,
             stream: true,
             max_tokens,
+            thinking,
             model: model_str,
         }
     }
@@ -345,6 +361,12 @@ impl AnthropicRequest {
     ) -> Self {
         let temperature = completion_request.temperature();
         let max_tokens = completion_request.get_max_tokens();
+        let thinking = completion_request
+            .thinking_budget()
+            .map(|budget| AnthropicThinking {
+                thinking_type: "enabled".to_owned(),
+                budget_tokens: budget,
+            });
         let messages = vec![AnthropicMessage::new(
             "user".to_owned(),
             completion_request.prompt().to_owned(),
@@ -356,6 +378,7 @@ impl AnthropicRequest {
             tools: vec![],
             stream: true,
             max_tokens,
+            thinking,
             model: model_str,
         }
     }
@@ -936,4 +959,4 @@ impl LLMClient for AnthropicClient {
 
         Ok(buffered_string)
     }
-}
+}

llm_client/src/clients/types.rs

@@ -686,6 +686,7 @@ pub struct LLMClientCompletionRequest {
     frequency_penalty: Option<f32>,
     stop_words: Option<Vec<String>>,
     max_tokens: Option<usize>,
+    thinking_budget: Option<usize>,
 }
 
 #[derive(Clone)]
@@ -696,6 +697,7 @@ pub struct LLMClientCompletionStringRequest {
     frequency_penalty: Option<f32>,
     stop_words: Option<Vec<String>>,
     max_tokens: Option<usize>,
+    thinking_budget: Option<usize>,
 }
 
 impl LLMClientCompletionStringRequest {
@@ -712,6 +714,7 @@ impl LLMClientCompletionStringRequest {
             frequency_penalty,
             stop_words: None,
             max_tokens: None,
+            thinking_budget: None,
         }
     }
 
@@ -748,6 +751,15 @@ impl LLMClientCompletionStringRequest {
     pub fn get_max_tokens(&self) -> Option<usize> {
         self.max_tokens
     }
+
+    pub fn set_thinking_budget(mut self, thinking_budget: usize) -> Self {
+        self.thinking_budget = Some(thinking_budget);
+        self
+    }
+
+    pub fn thinking_budget(&self) -> Option<usize> {
+        self.thinking_budget
+    }
 }
 
 impl LLMClientCompletionRequest {
@@ -764,6 +776,7 @@ impl LLMClientCompletionRequest {
             frequency_penalty,
             stop_words: None,
             max_tokens: None,
+            thinking_budget: None,
         }
     }
 
@@ -852,6 +865,15 @@ impl LLMClientCompletionRequest {
     pub fn get_max_tokens(&self) -> Option<usize> {
         self.max_tokens
     }
+
+    pub fn set_thinking_budget(mut self, thinking_budget: usize) -> Self {
+        self.thinking_budget = Some(thinking_budget);
+        self
+    }
+
+    pub fn thinking_budget(&self) -> Option<usize> {
+        self.thinking_budget
+    }
 }
 
 #[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]