Skip to content

Commit 9c66927

Browse files
committed
feat: add thinking budget support for Anthropic requests
1 parent 552407d commit 9c66927

File tree

2 files changed

+46
-1
lines changed

2 files changed

+46
-1
lines changed

llm_client/src/clients/anthropic.rs

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,13 @@ enum ContentBlockDeltaType {
227227
},
228228
}
229229

230+
#[derive(serde::Serialize, Debug, Clone)]
231+
struct AnthropicThinking {
232+
#[serde(rename = "type")]
233+
thinking_type: String,
234+
budget_tokens: usize,
235+
}
236+
230237
#[derive(serde::Serialize, Debug, Clone)]
231238
struct AnthropicRequest {
232239
system: Vec<AnthropicMessageContent>,
@@ -238,6 +245,8 @@ struct AnthropicRequest {
238245
stream: bool,
239246
#[serde(skip_serializing_if = "Option::is_none")]
240247
max_tokens: Option<usize>,
248+
#[serde(skip_serializing_if = "Option::is_none")]
249+
thinking: Option<AnthropicThinking>,
241250
model: String,
242251
}
243252

@@ -251,6 +260,12 @@ impl AnthropicRequest {
251260
Some(tokens) => Some(tokens),
252261
None => Some(8192),
253262
};
263+
let thinking = completion_request
264+
.thinking_budget()
265+
.map(|budget| AnthropicThinking {
266+
thinking_type: "enabled".to_owned(),
267+
budget_tokens: budget,
268+
});
254269
let messages = completion_request.messages();
255270
// grab the tools over here ONLY from the system message
256271
let tools = messages
@@ -335,6 +350,7 @@ impl AnthropicRequest {
335350
tools,
336351
stream: true,
337352
max_tokens,
353+
thinking,
338354
model: model_str,
339355
}
340356
}
@@ -345,6 +361,12 @@ impl AnthropicRequest {
345361
) -> Self {
346362
let temperature = completion_request.temperature();
347363
let max_tokens = completion_request.get_max_tokens();
364+
let thinking = completion_request
365+
.thinking_budget()
366+
.map(|budget| AnthropicThinking {
367+
thinking_type: "enabled".to_owned(),
368+
budget_tokens: budget,
369+
});
348370
let messages = vec![AnthropicMessage::new(
349371
"user".to_owned(),
350372
completion_request.prompt().to_owned(),
@@ -356,6 +378,7 @@ impl AnthropicRequest {
356378
tools: vec![],
357379
stream: true,
358380
max_tokens,
381+
thinking,
359382
model: model_str,
360383
}
361384
}
@@ -936,4 +959,4 @@ impl LLMClient for AnthropicClient {
936959

937960
Ok(buffered_string)
938961
}
939-
}
962+
}

llm_client/src/clients/types.rs

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -686,6 +686,7 @@ pub struct LLMClientCompletionRequest {
686686
frequency_penalty: Option<f32>,
687687
stop_words: Option<Vec<String>>,
688688
max_tokens: Option<usize>,
689+
thinking_budget: Option<usize>,
689690
}
690691

691692
#[derive(Clone)]
@@ -696,6 +697,7 @@ pub struct LLMClientCompletionStringRequest {
696697
frequency_penalty: Option<f32>,
697698
stop_words: Option<Vec<String>>,
698699
max_tokens: Option<usize>,
700+
thinking_budget: Option<usize>,
699701
}
700702

701703
impl LLMClientCompletionStringRequest {
@@ -712,6 +714,7 @@ impl LLMClientCompletionStringRequest {
712714
frequency_penalty,
713715
stop_words: None,
714716
max_tokens: None,
717+
thinking_budget: None,
715718
}
716719
}
717720

@@ -748,6 +751,15 @@ impl LLMClientCompletionStringRequest {
748751
pub fn get_max_tokens(&self) -> Option<usize> {
749752
self.max_tokens
750753
}
754+
755+
pub fn set_thinking_budget(mut self, thinking_budget: usize) -> Self {
756+
self.thinking_budget = Some(thinking_budget);
757+
self
758+
}
759+
760+
pub fn thinking_budget(&self) -> Option<usize> {
761+
self.thinking_budget
762+
}
751763
}
752764

753765
impl LLMClientCompletionRequest {
@@ -764,6 +776,7 @@ impl LLMClientCompletionRequest {
764776
frequency_penalty,
765777
stop_words: None,
766778
max_tokens: None,
779+
thinking_budget: None,
767780
}
768781
}
769782

@@ -852,6 +865,15 @@ impl LLMClientCompletionRequest {
852865
pub fn get_max_tokens(&self) -> Option<usize> {
853866
self.max_tokens
854867
}
868+
869+
pub fn set_thinking_budget(mut self, thinking_budget: usize) -> Self {
870+
self.thinking_budget = Some(thinking_budget);
871+
self
872+
}
873+
874+
pub fn thinking_budget(&self) -> Option<usize> {
875+
self.thinking_budget
876+
}
855877
}
856878

857879
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]

0 commit comments

Comments
 (0)