Skip to content

Commit 9168278

Browse files
committed
feat: Enhance Anthropic integration with Thinking
- The `thinking` option is added to `AnthropicChatOptions` and `ChatCompletionRequest`. - The `AnthropicApi` and `AnthropicChatModel` now handle `THINKING` and `REDACTED_THINKING` content blocks in responses. New tests verify parsing of these blocks. - Updated method signatures on ChatCompletionRequestBuilder, deprecating old builders with `with*` prefix in favor of those without. Signed-off-by: Alexandros Pappas <[email protected]>
1 parent 6cb15e4 commit 9168278

File tree

7 files changed

+416
-65
lines changed

7 files changed

+416
-65
lines changed

models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/AnthropicChatModel.java

+38-31
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import java.util.ArrayList;
2020
import java.util.Base64;
21+
import java.util.HashMap;
2122
import java.util.List;
2223
import java.util.Map;
2324
import java.util.Set;
@@ -36,6 +37,7 @@
3637
import org.springframework.ai.tool.definition.ToolDefinition;
3738
import org.springframework.ai.util.json.JsonParser;
3839
import org.springframework.lang.Nullable;
40+
3941
import reactor.core.publisher.Flux;
4042
import reactor.core.publisher.Mono;
4143

@@ -379,46 +381,51 @@ private ChatResponse toChatResponse(ChatCompletionResponse chatCompletion, Usage
379381
return new ChatResponse(List.of());
380382
}
381383

382-
List<Generation> generations = chatCompletion.content()
383-
.stream()
384-
.filter(content -> content.type() != ContentBlock.Type.TOOL_USE)
385-
.map(content -> new Generation(new AssistantMessage(content.text(), Map.of()),
386-
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build()))
387-
.toList();
388-
389-
List<Generation> allGenerations = new ArrayList<>(generations);
384+
List<Generation> generations = new ArrayList<>();
385+
List<AssistantMessage.ToolCall> toolCalls = new ArrayList<>();
386+
for (ContentBlock content : chatCompletion.content()) {
387+
switch (content.type()) {
388+
case TEXT, TEXT_DELTA:
389+
generations.add(new Generation(new AssistantMessage(content.text(), Map.of()),
390+
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build()));
391+
break;
392+
case THINKING, THINKING_DELTA:
393+
System.out.println(
394+
"THINKINGTHINKINGTHINKINGTHINKINGTHINKINGTHINKINGTHINKINGcontent type: " + content.type());
395+
Map<String, Object> thinkingProperties = new HashMap<>();
396+
thinkingProperties.put("signature", content.signature());
397+
generations.add(new Generation(new AssistantMessage(content.thinking(), thinkingProperties),
398+
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build()));
399+
break;
400+
case REDACTED_THINKING:
401+
Map<String, Object> redactedProperties = new HashMap<>();
402+
redactedProperties.put("data", content.data());
403+
generations.add(new Generation(new AssistantMessage(null, redactedProperties),
404+
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build()));
405+
break;
406+
case TOOL_USE:
407+
var functionCallId = content.id();
408+
var functionName = content.name();
409+
var functionArguments = JsonParser.toJson(content.input());
410+
toolCalls.add(
411+
new AssistantMessage.ToolCall(functionCallId, "function", functionName, functionArguments));
412+
break;
413+
}
414+
}
390415

391416
if (chatCompletion.stopReason() != null && generations.isEmpty()) {
392417
Generation generation = new Generation(new AssistantMessage(null, Map.of()),
393418
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build());
394-
allGenerations.add(generation);
419+
generations.add(generation);
395420
}
396421

397-
List<ContentBlock> toolToUseList = chatCompletion.content()
398-
.stream()
399-
.filter(c -> c.type() == ContentBlock.Type.TOOL_USE)
400-
.toList();
401-
402-
if (!CollectionUtils.isEmpty(toolToUseList)) {
403-
List<AssistantMessage.ToolCall> toolCalls = new ArrayList<>();
404-
405-
for (ContentBlock toolToUse : toolToUseList) {
406-
407-
var functionCallId = toolToUse.id();
408-
var functionName = toolToUse.name();
409-
var functionArguments = JsonParser.toJson(toolToUse.input());
410-
411-
toolCalls
412-
.add(new AssistantMessage.ToolCall(functionCallId, "function", functionName, functionArguments));
413-
}
414-
422+
if (!CollectionUtils.isEmpty(toolCalls)) {
415423
AssistantMessage assistantMessage = new AssistantMessage("", Map.of(), toolCalls);
416424
Generation toolCallGeneration = new Generation(assistantMessage,
417425
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build());
418-
allGenerations.add(toolCallGeneration);
426+
generations.add(toolCallGeneration);
419427
}
420-
421-
return new ChatResponse(allGenerations, this.from(chatCompletion, usage));
428+
return new ChatResponse(generations, this.from(chatCompletion, usage));
422429
}
423430

424431
private ChatResponseMetadata from(AnthropicApi.ChatCompletionResponse result) {
@@ -575,7 +582,7 @@ else if (message.getMessageType() == MessageType.TOOL) {
575582
List<ToolDefinition> toolDefinitions = this.toolCallingManager.resolveToolDefinitions(requestOptions);
576583
if (!CollectionUtils.isEmpty(toolDefinitions)) {
577584
request = ModelOptionsUtils.merge(request, this.defaultOptions, ChatCompletionRequest.class);
578-
request = ChatCompletionRequest.from(request).withTools(getFunctionTools(toolDefinitions)).build();
585+
request = ChatCompletionRequest.from(request).tools(getFunctionTools(toolDefinitions)).build();
579586
}
580587

581588
return request;

models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/AnthropicChatOptions.java

+20
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ public class AnthropicChatOptions implements ToolCallingChatOptions {
5656
private @JsonProperty("temperature") Double temperature;
5757
private @JsonProperty("top_p") Double topP;
5858
private @JsonProperty("top_k") Integer topK;
59+
private @JsonProperty("thinking") ChatCompletionRequest.ThinkingConfig thinking;
5960

6061
/**
6162
* Collection of {@link ToolCallback}s to be used for tool calling in the chat
@@ -94,6 +95,7 @@ public static AnthropicChatOptions fromOptions(AnthropicChatOptions fromOptions)
9495
.temperature(fromOptions.getTemperature())
9596
.topP(fromOptions.getTopP())
9697
.topK(fromOptions.getTopK())
98+
.thinking(fromOptions.getThinking())
9799
.toolCallbacks(fromOptions.getToolCallbacks())
98100
.toolNames(fromOptions.getToolNames())
99101
.internalToolExecutionEnabled(fromOptions.isInternalToolExecutionEnabled())
@@ -163,6 +165,14 @@ public void setTopK(Integer topK) {
163165
this.topK = topK;
164166
}
165167

168+
public ChatCompletionRequest.ThinkingConfig getThinking() {
169+
return this.thinking;
170+
}
171+
172+
public void setThinking(ChatCompletionRequest.ThinkingConfig thinking) {
173+
this.thinking = thinking;
174+
}
175+
166176
@Override
167177
@JsonIgnore
168178
public List<FunctionCallback> getToolCallbacks() {
@@ -319,6 +329,16 @@ public Builder topK(Integer topK) {
319329
return this;
320330
}
321331

332+
public Builder thinking(ChatCompletionRequest.ThinkingConfig thinking) {
333+
this.options.thinking = thinking;
334+
return this;
335+
}
336+
337+
public Builder thinking(AnthropicApi.ThinkingType type, Integer budgetTokens) {
338+
this.options.thinking = new ChatCompletionRequest.ThinkingConfig(type, budgetTokens);
339+
return this;
340+
}
341+
322342
public Builder toolCallbacks(List<FunctionCallback> toolCallbacks) {
323343
this.options.setToolCallbacks(toolCallbacks);
324344
return this;

0 commit comments

Comments
 (0)