Skip to content

Commit

Permalink
Add proxy tool calls option to chat models
Browse files Browse the repository at this point in the history
 This commit introduces a new proxyToolCalls option for various chat
 models in the Spring AI project. When enabled, it allows the client to
 handle function calls externally instead of being processed internally
 by Spring AI.

 The change affects multiple chat model implementations, including:
  AnthropicChatModel
  AzureOpenAiChatModel
  MiniMaxChatModel
  MistralAiChatModel
  MoonshotChatModel
  OllamaChatModel
  OpenAiChatModel
  VertexAiGeminiChatModel
  ZhiPuAiChatModel

 The proxyToolCalls option is added to the respective chat options
 classes and integrated into the AbstractToolCallSupport class for
 consistent handling across different implementations.

 The proxyToolCalls option can be set either programmatically via
 the <ModelName>ChatOptions.builder().withProxyToolCalls() method
 or the spring.ai.<model-name>.chat.options.proxy-tool-calls
 application property.

 Documentation for the new option is also updated in the relevant
 Antora pages.

 Resolves spring-projects#1367
  • Loading branch information
tzolov committed Sep 23, 2024
1 parent acb31e7 commit 79b44e4
Show file tree
Hide file tree
Showing 41 changed files with 1,256 additions and 232 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -225,7 +225,8 @@ public ChatResponse call(Prompt prompt) {
return chatResponse;
});

if (response != null && this.isToolCall(response, Set.of("tool_use"))) {
if (!isProxyToolCalls(prompt, this.defaultOptions) && response != null
&& this.isToolCall(response, Set.of("tool_use"))) {
var toolCallConversation = handleToolCalls(prompt, response);
return this.call(new Prompt(toolCallConversation, prompt.getOptions()));
}
Expand Down Expand Up @@ -256,7 +257,7 @@ public Flux<ChatResponse> stream(Prompt prompt) {
Flux<ChatResponse> chatResponseFlux = response.switchMap(chatCompletionResponse -> {
ChatResponse chatResponse = toChatResponse(chatCompletionResponse);

if (this.isToolCall(chatResponse, Set.of("tool_use"))) {
if (!isProxyToolCalls(prompt, this.defaultOptions) && this.isToolCall(chatResponse, Set.of("tool_use"))) {
var toolCallConversation = handleToolCalls(prompt, chatResponse);
return this.stream(new Prompt(toolCallConversation, prompt.getOptions()));
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,9 @@ public class AnthropicChatOptions implements ChatOptions, FunctionCallingOptions
@NestedConfigurationProperty
@JsonIgnore
private Set<String> functions = new HashSet<>();

@JsonIgnore
private Boolean proxyToolCalls;
// @formatter:on

public static Builder builder() {
Expand Down Expand Up @@ -144,6 +147,11 @@ public Builder withFunction(String functionName) {
return this;
}

public Builder withProxyToolCalls(Boolean proxyToolCalls) {
this.options.proxyToolCalls = proxyToolCalls;
return this;
}

public AnthropicChatOptions build() {
return this.options;
}
Expand Down Expand Up @@ -246,6 +254,15 @@ public Double getPresencePenalty() {
return null;
}

@Override
public Boolean getProxyToolCalls() {
return this.proxyToolCalls;
}

public void setProxyToolCalls(Boolean proxyToolCalls) {
this.proxyToolCalls = proxyToolCalls;
}

@Override
public AnthropicChatOptions copy() {
return fromOptions(this);
Expand All @@ -261,6 +278,7 @@ public static AnthropicChatOptions fromOptions(AnthropicChatOptions fromOptions)
.withTopK(fromOptions.getTopK())
.withFunctionCallbacks(fromOptions.getFunctionCallbacks())
.withFunctions(fromOptions.getFunctions())
.withProxyToolCalls(fromOptions.getProxyToolCalls())
.build();
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,8 @@ public ChatResponse call(Prompt prompt) {

ChatResponse chatResponse = toChatResponse(chatCompletions);

if (isToolCall(chatResponse, Set.of(String.valueOf(CompletionsFinishReason.TOOL_CALLS).toLowerCase()))) {
if (!isProxyToolCalls(prompt, this.defaultOptions)
&& isToolCall(chatResponse, Set.of(String.valueOf(CompletionsFinishReason.TOOL_CALLS).toLowerCase()))) {
var toolCallConversation = handleToolCalls(prompt, chatResponse);
// Recursively call the call method with the tool call message
// conversation that contains the call responses.
Expand Down Expand Up @@ -199,7 +200,8 @@ public Flux<ChatResponse> stream(Prompt prompt) {

ChatResponse chatResponse = toChatResponse(chatCompletions);

if (isToolCall(chatResponse, Set.of(String.valueOf(CompletionsFinishReason.TOOL_CALLS).toLowerCase()))) {
if (!isProxyToolCalls(prompt, this.defaultOptions) && isToolCall(chatResponse,
Set.of(String.valueOf(CompletionsFinishReason.TOOL_CALLS).toLowerCase()))) {
var toolCallConversation = handleToolCalls(prompt, chatResponse);
// Recursively call the call method with the tool call message
// conversation that contains the call responses.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
import org.springframework.ai.model.function.FunctionCallingOptions;
import org.springframework.boot.context.properties.NestedConfigurationProperty;
import org.springframework.util.Assert;
import org.stringtemplate.v4.compiler.CodeGenerator.primary_return;

/**
* The configuration information for a chat completions request. Completions support a
Expand Down Expand Up @@ -161,6 +162,9 @@ public class AzureOpenAiChatOptions implements FunctionCallingOptions, ChatOptio
@JsonIgnore
private Set<String> functions = new HashSet<>();

@JsonIgnore
private Boolean proxyToolCalls;

public static Builder builder() {
return new Builder();
}
Expand Down Expand Up @@ -250,6 +254,11 @@ public Builder withResponseFormat(AzureOpenAiResponseFormat responseFormat) {
return this;
}

public Builder withProxyToolCalls(Boolean proxyToolCalls) {
this.options.proxyToolCalls = proxyToolCalls;
return this;
}

public AzureOpenAiChatOptions build() {
return this.options;
}
Expand Down Expand Up @@ -395,6 +404,15 @@ public Integer getTopK() {
return null;
}

@Override
public Boolean getProxyToolCalls() {
return this.proxyToolCalls;
}

public void setProxyToolCalls(Boolean proxyToolCalls) {
this.proxyToolCalls = proxyToolCalls;
}

@Override
public AzureOpenAiChatOptions copy() {
return fromOptions(this);
Expand All @@ -413,6 +431,7 @@ public static AzureOpenAiChatOptions fromOptions(AzureOpenAiChatOptions fromOpti
.withUser(fromOptions.getUser())
.withFunctionCallbacks(fromOptions.getFunctionCallbacks())
.withFunctions(fromOptions.getFunctions())
.withResponseFormat(fromOptions.getResponseFormat())
.build();
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -190,7 +190,7 @@ public ChatResponse call(Prompt prompt) {

ChatResponse chatResponse = new ChatResponse(generations, from(completionEntity.getBody()));

if (isToolCall(chatResponse,
if (!isProxyToolCalls(prompt, this.defaultOptions) && isToolCall(chatResponse,
Set.of(ChatCompletionFinishReason.TOOL_CALLS.name(), ChatCompletionFinishReason.STOP.name()))) {
var toolCallConversation = handleToolCalls(prompt, chatResponse);
// Recursively call the call method with the tool call message
Expand Down Expand Up @@ -254,7 +254,7 @@ public Flux<ChatResponse> stream(Prompt prompt) {

return chatResponse.flatMap(response -> {

if (isToolCall(response,
if (!isProxyToolCalls(prompt, this.defaultOptions) && isToolCall(response,
Set.of(ChatCompletionFinishReason.TOOL_CALLS.name(), ChatCompletionFinishReason.STOP.name()))) {
var toolCallConversation = handleToolCalls(prompt, response);
// Recursively call the stream method with the tool call message
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -142,6 +142,9 @@ public class MiniMaxChatOptions implements FunctionCallingOptions, ChatOptions {
@NestedConfigurationProperty
@JsonIgnore
private Set<String> functions = new HashSet<>();

@JsonIgnore
private Boolean proxyToolCalls;
// @formatter:on

public static Builder builder() {
Expand Down Expand Up @@ -242,6 +245,11 @@ public Builder withFunction(String functionName) {
return this;
}

public Builder withProxyToolCalls(Boolean proxyToolCalls) {
this.options.proxyToolCalls = proxyToolCalls;
return this;
}

public MiniMaxChatOptions build() {
return this.options;
}
Expand Down Expand Up @@ -394,6 +402,15 @@ public Integer getTopK() {
return null;
}

@Override
public Boolean getProxyToolCalls() {
return this.proxyToolCalls;
}

public void setProxyToolCalls(Boolean proxyToolCalls) {
this.proxyToolCalls = proxyToolCalls;
}

@Override
public int hashCode() {
final int prime = 31;
Expand All @@ -411,6 +428,7 @@ public int hashCode() {
result = prime * result + ((maskSensitiveInfo == null) ? 0 : maskSensitiveInfo.hashCode());
result = prime * result + ((tools == null) ? 0 : tools.hashCode());
result = prime * result + ((toolChoice == null) ? 0 : toolChoice.hashCode());
result = prime * result + ((proxyToolCalls == null) ? 0 : proxyToolCalls.hashCode());
return result;
}

Expand Down Expand Up @@ -501,6 +519,12 @@ else if (!tools.equals(other.tools))
}
else if (!toolChoice.equals(other.toolChoice))
return false;
if (this.proxyToolCalls == null) {
if (other.proxyToolCalls != null)
return false;
}
else if (!proxyToolCalls.equals(other.proxyToolCalls))
return false;
return true;
}

Expand All @@ -525,6 +549,7 @@ public static MiniMaxChatOptions fromOptions(MiniMaxChatOptions fromOptions) {
.withToolChoice(fromOptions.getToolChoice())
.withFunctionCallbacks(fromOptions.getFunctionCallbacks())
.withFunctions(fromOptions.getFunctions())
.withProxyToolCalls(fromOptions.getProxyToolCalls())
.build();
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -183,8 +183,9 @@ public ChatResponse call(Prompt prompt) {
return chatResponse;
});

if (response != null && isToolCall(response, Set.of(MistralAiApi.ChatCompletionFinishReason.TOOL_CALLS.name(),
MistralAiApi.ChatCompletionFinishReason.STOP.name()))) {
if (!isProxyToolCalls(prompt, this.defaultOptions) && response != null
&& isToolCall(response, Set.of(MistralAiApi.ChatCompletionFinishReason.TOOL_CALLS.name(),
MistralAiApi.ChatCompletionFinishReason.STOP.name()))) {
var toolCallConversation = handleToolCalls(prompt, response);
// Recursively call the call method with the tool call message
// conversation that contains the call responses.
Expand Down Expand Up @@ -255,7 +256,7 @@ public Flux<ChatResponse> stream(Prompt prompt) {

// @formatter:off
Flux<ChatResponse> chatResponseFlux = chatResponse.flatMap(response -> {
if (isToolCall(response, Set.of(MistralAiApi.ChatCompletionFinishReason.TOOL_CALLS.name()))) {
if (!isProxyToolCalls(prompt, this.defaultOptions) && isToolCall(response, Set.of(MistralAiApi.ChatCompletionFinishReason.TOOL_CALLS.name()))) {
var toolCallConversation = handleToolCalls(prompt, response);
// Recursively call the stream method with the tool call message
// conversation that contains the call responses.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,9 @@ public class MistralAiChatOptions implements FunctionCallingOptions, ChatOptions
@JsonIgnore
private Set<String> functions = new HashSet<>();

@JsonIgnore
private Boolean proxyToolCalls;

public static Builder builder() {
return new Builder();
}
Expand Down Expand Up @@ -215,6 +218,11 @@ public Builder withFunction(String functionName) {
return this;
}

public Builder withProxyToolCalls(Boolean proxyToolCalls) {
this.options.proxyToolCalls = proxyToolCalls;
return this;
}

public MistralAiChatOptions build() {
return this.options;
}
Expand Down Expand Up @@ -356,6 +364,15 @@ public Integer getTopK() {
return null;
}

@Override
public Boolean getProxyToolCalls() {
return this.proxyToolCalls;
}

public void setProxyToolCalls(Boolean proxyToolCalls) {
this.proxyToolCalls = proxyToolCalls;
}

@Override
public MistralAiChatOptions copy() {
return fromOptions(this);
Expand All @@ -374,7 +391,114 @@ public static MistralAiChatOptions fromOptions(MistralAiChatOptions fromOptions)
.withToolChoice(fromOptions.getToolChoice())
.withFunctionCallbacks(fromOptions.getFunctionCallbacks())
.withFunctions(fromOptions.getFunctions())
.withProxyToolCalls(fromOptions.getProxyToolCalls())
.build();
}

@Override
public int hashCode() {
final int prime = 31;
int result = 1;
result = prime * result + ((model == null) ? 0 : model.hashCode());
result = prime * result + ((temperature == null) ? 0 : temperature.hashCode());
result = prime * result + ((topP == null) ? 0 : topP.hashCode());
result = prime * result + ((maxTokens == null) ? 0 : maxTokens.hashCode());
result = prime * result + ((safePrompt == null) ? 0 : safePrompt.hashCode());
result = prime * result + ((randomSeed == null) ? 0 : randomSeed.hashCode());
result = prime * result + ((responseFormat == null) ? 0 : responseFormat.hashCode());
result = prime * result + ((stop == null) ? 0 : stop.hashCode());
result = prime * result + ((tools == null) ? 0 : tools.hashCode());
result = prime * result + ((toolChoice == null) ? 0 : toolChoice.hashCode());
result = prime * result + ((functionCallbacks == null) ? 0 : functionCallbacks.hashCode());
result = prime * result + ((functions == null) ? 0 : functions.hashCode());
result = prime * result + ((proxyToolCalls == null) ? 0 : proxyToolCalls.hashCode());
return result;
}

@Override
public boolean equals(Object obj) {
if (this == obj)
return true;
if (obj == null)
return false;
if (getClass() != obj.getClass())
return false;
MistralAiChatOptions other = (MistralAiChatOptions) obj;
if (model == null) {
if (other.model != null)
return false;
}
else if (!model.equals(other.model))
return false;
if (temperature == null) {
if (other.temperature != null)
return false;
}
else if (!temperature.equals(other.temperature))
return false;
if (topP == null) {
if (other.topP != null)
return false;
}
else if (!topP.equals(other.topP))
return false;
if (maxTokens == null) {
if (other.maxTokens != null)
return false;
}
else if (!maxTokens.equals(other.maxTokens))
return false;
if (safePrompt == null) {
if (other.safePrompt != null)
return false;
}
else if (!safePrompt.equals(other.safePrompt))
return false;
if (randomSeed == null) {
if (other.randomSeed != null)
return false;
}
else if (!randomSeed.equals(other.randomSeed))
return false;
if (responseFormat == null) {
if (other.responseFormat != null)
return false;
}
else if (!responseFormat.equals(other.responseFormat))
return false;
if (stop == null) {
if (other.stop != null)
return false;
}
else if (!stop.equals(other.stop))
return false;
if (tools == null) {
if (other.tools != null)
return false;
}
else if (!tools.equals(other.tools))
return false;
if (toolChoice != other.toolChoice)
return false;
if (functionCallbacks == null) {
if (other.functionCallbacks != null)
return false;
}
else if (!functionCallbacks.equals(other.functionCallbacks))
return false;
if (functions == null) {
if (other.functions != null)
return false;
}
else if (!functions.equals(other.functions))
return false;
if (proxyToolCalls == null) {
if (other.proxyToolCalls != null)
return false;
}
else if (!proxyToolCalls.equals(other.proxyToolCalls))
return false;
return true;
}

}
Loading

0 comments on commit 79b44e4

Please sign in to comment.