Refactor FunctionCallback API to use ToolContext instead of Map<String, Object>

- Replaced `Map<String, Object>` with `ToolContext` in the `FunctionCallback`, `AbstractFunctionCallback`, and related classes.
 - Updated all BiFunction definitions to use `ToolContext` as the second parameter,
   enhancing the clarity and structure of the tool context management.
 - Modified `FunctionCallbackWrapper` and `FunctionCallbackContext` to adapt to the new `ToolContext` parameter.
 - Adjusted the handling of tool context in the documentation and test classes.
 - Updated relevant test cases to reflect the API changes and modified the function handling logic to ensure consistency.
This commit is contained in:
Christian Tzolov
2024-10-07 14:18:40 +02:00
committed by Mark Pollack
parent 662afdf46e
commit 38082de0f8
16 changed files with 127 additions and 71 deletions

View File

@@ -25,6 +25,7 @@ import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.function.FunctionCallbackWrapper;
import org.springframework.ai.openai.OpenAiChatModel;
@@ -71,12 +72,12 @@ class OpenAiChatModelFunctionCallingIT {
@Test
void functionCallWithToolContextTest() {
var biFunction = new BiFunction<MockWeatherService.Request, Map<String, Object>, MockWeatherService.Response>() {
var biFunction = new BiFunction<MockWeatherService.Request, ToolContext, MockWeatherService.Response>() {
@Override
public Response apply(Request request, Map<String, Object> toolContext) {
public Response apply(Request request, ToolContext toolContext) {
assertThat(toolContext).containsEntry("sessionId", "123");
assertThat(toolContext.getContext()).containsEntry("sessionId", "123");
double temperature = 0;
if (request.location().contains("Paris")) {
@@ -133,12 +134,12 @@ class OpenAiChatModelFunctionCallingIT {
@Test
void streamFunctionCallWithToolContextTest() {
var biFunction = new BiFunction<MockWeatherService.Request, Map<String, Object>, MockWeatherService.Response>() {
var biFunction = new BiFunction<MockWeatherService.Request, ToolContext, MockWeatherService.Response>() {
@Override
public Response apply(Request request, Map<String, Object> toolContext) {
public Response apply(Request request, ToolContext toolContext) {
assertThat(toolContext).containsEntry("sessionId", "123");
assertThat(toolContext.getContext()).containsEntry("sessionId", "123");
double temperature = 0;
if (request.location().contains("Paris")) {

View File

@@ -29,7 +29,7 @@ import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.model.function.FunctionCallback;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.openai.OpenAiTestConfiguration;
import org.springframework.ai.openai.api.tool.MockWeatherService;
import org.springframework.ai.openai.api.tool.MockWeatherService.Request;
@@ -115,12 +115,12 @@ class OpenAiChatClientMultipleFunctionCallsIT extends AbstractIT {
@Test
void defaultFunctionCallTestWithToolContext() {
var biFunction = new BiFunction<MockWeatherService.Request, Map<String, Object>, MockWeatherService.Response>() {
var biFunction = new BiFunction<MockWeatherService.Request, ToolContext, MockWeatherService.Response>() {
@Override
public Response apply(Request request, Map<String, Object> toolContext) {
public Response apply(Request request, ToolContext toolContext) {
assertThat(toolContext).containsEntry("sessionId", "123");
assertThat(toolContext.getContext()).containsEntry("sessionId", "123");
double temperature = 0;
if (request.location().contains("Paris")) {
@@ -155,12 +155,12 @@ class OpenAiChatClientMultipleFunctionCallsIT extends AbstractIT {
@Test
void functionCallTestWithToolContext() {
var biFunction = new BiFunction<MockWeatherService.Request, Map<String, Object>, MockWeatherService.Response>() {
var biFunction = new BiFunction<MockWeatherService.Request, ToolContext, MockWeatherService.Response>() {
@Override
public Response apply(Request request, Map<String, Object> toolContext) {
public Response apply(Request request, ToolContext toolContext) {
assertThat(toolContext).containsEntry("sessionId", "123");
assertThat(toolContext.getContext()).containsEntry("sessionId", "123");
double temperature = 0;
if (request.location().contains("Paris")) {

View File

@@ -26,6 +26,7 @@ import org.springframework.ai.chat.client.observation.ChatClientObservationConve
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.converter.StructuredOutputConverter;
@@ -202,7 +203,7 @@ public interface ChatClient {
java.util.function.Function<I, O> function);
<I, O> ChatClientRequestSpec function(String name, String description,
java.util.function.BiFunction<I, Map<String, Object>, O> function);
java.util.function.BiFunction<I, ToolContext, O> function);
<I, O> ChatClientRequestSpec functions(FunctionCallback... functionCallbacks);
@@ -267,7 +268,7 @@ public interface ChatClient {
<I, O> Builder defaultFunction(String name, String description, java.util.function.Function<I, O> function);
<I, O> Builder defaultFunction(String name, String description,
java.util.function.BiFunction<I, Map<String, Object>, O> function);
java.util.function.BiFunction<I, ToolContext, O> function);
Builder defaultFunctions(String... functionNames);

View File

@@ -46,6 +46,7 @@ import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.StreamingChatModel;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.converter.BeanOutputConverter;
@@ -693,7 +694,7 @@ public class DefaultChatClient implements ChatClient {
}
public <I, O> ChatClientRequestSpec function(String name, String description,
java.util.function.BiFunction<I, Map<String, Object>, O> biFunction) {
java.util.function.BiFunction<I, ToolContext, O> biFunction) {
Assert.hasText(name, "the name must be non-null and non-empty");
Assert.hasText(description, "the description must be non-null and non-empty");

View File

@@ -29,6 +29,7 @@ import org.springframework.ai.chat.client.DefaultChatClient.DefaultChatClientReq
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.client.observation.ChatClientObservationConvention;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.model.function.FunctionCallback;
import org.springframework.core.io.Resource;
@@ -142,7 +143,7 @@ public class DefaultChatClientBuilder implements Builder {
}
public <I, O> Builder defaultFunction(String name, String description,
java.util.function.BiFunction<I, Map<String, Object>, O> biFunction) {
java.util.function.BiFunction<I, ToolContext, O> biFunction) {
this.defaultRequest.function(name, description, biFunction);
return this;
}

View File

@@ -56,6 +56,7 @@ import org.springframework.util.StringUtils;
* @param advisors the list of request response advisors
* @param advisorParams the map of advisor parameters
* @param adviseContext the map of advise context
* @param toolContext the tool context
*/
public record AdvisedRequest(ChatModel chatModel, String userText, String systemText, ChatOptions chatOptions,
List<Media> media, List<String> functionNames, List<FunctionCallback> functionCallbacks, List<Message> messages,
@@ -94,8 +95,6 @@ public record AdvisedRequest(ChatModel chatModel, String userText, String system
public static class Builder {
public Map<String, Object> toolContext;
private ChatModel chatModel;
private String userText = "";
@@ -122,6 +121,8 @@ public record AdvisedRequest(ChatModel chatModel, String userText, String system
private Map<String, Object> adviseContext = Map.of();
public Map<String, Object> toolContext = Map.of();
public Builder withChatModel(ChatModel chatModel) {
this.chatModel = chatModel;
return this;

View File

@@ -16,6 +16,7 @@
package org.springframework.ai.chat.model;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
@@ -137,11 +138,13 @@ public abstract class AbstractToolCallSupport {
}
AssistantMessage assistantMessage = toolCallGeneration.get().getOutput();
Map<String, Object> toolContext = null;
if (prompt.getOptions() instanceof FunctionCallingOptions functionCallOptions) {
toolContext = functionCallOptions.getToolContext();
Map<String, Object> toolContextMap = Map.of();
if (prompt.getOptions() instanceof FunctionCallingOptions functionCallOptions
&& !CollectionUtils.isEmpty(functionCallOptions.getToolContext())) {
toolContextMap = functionCallOptions.getToolContext();
}
ToolResponseMessage toolMessageResponse = this.executeFunctions(assistantMessage, toolContext);
ToolResponseMessage toolMessageResponse = this.executeFunctions(assistantMessage,
new ToolContext(toolContextMap));
return this.buildToolCallConversation(prompt.getInstructions(), assistantMessage, toolMessageResponse);
}
@@ -190,7 +193,7 @@ public abstract class AbstractToolCallSupport {
return retrievedFunctionCallbacks;
}
protected ToolResponseMessage executeFunctions(AssistantMessage assistantMessage, Map<String, Object> toolContext) {
protected ToolResponseMessage executeFunctions(AssistantMessage assistantMessage, ToolContext toolContext) {
List<ToolResponseMessage.ToolResponse> toolResponses = new ArrayList<>();

View File

@@ -0,0 +1,59 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.model;
import java.util.Collections;
import java.util.Map;
/**
* Represents the context for tool execution in a function calling scenario.
*
* <p>
* This class encapsulates a map of contextual information that can be passed to tools
* (functions) when they are called. It provides an immutable view of the context to
* ensure thread-safety and prevent modification after creation.
* </p>
*
* <p>
* The context is typically populated from the {@code toolContext} field of
* {@code FunctionCallingOptions} and is used in the function execution process.
* </p>
*
* @author Christian Tzolov
* @since 1.0.0
*/
public class ToolContext {
private final Map<String, Object> context;
/**
* Constructs a new ToolContext with the given context map.
* @param context A map containing the tool context information. This map is wrapped
* in an unmodifiable view to prevent changes.
*/
public ToolContext(Map<String, Object> context) {
this.context = Collections.unmodifiableMap(context);
}
/**
* Returns the immutable context map.
* @return An unmodifiable view of the context map.
*/
public Map<String, Object> getContext() {
return this.context;
}
}

View File

@@ -15,16 +15,16 @@
*/
package org.springframework.ai.model.function;
import java.util.Map;
import java.util.Objects;
import java.util.function.BiFunction;
import java.util.function.Function;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.util.Assert;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.util.Assert;
/**
* Abstract implementation of the {@link FunctionCallback} for interacting with the
* Model's function calling protocol and a {@link Function} wrapping the interaction with
@@ -40,7 +40,7 @@ import org.springframework.util.Assert;
* @param <O> the 3rd party service output type.
* @author Christian Tzolov
*/
abstract class AbstractFunctionCallback<I, O> implements BiFunction<I, Map<String, Object>, O>, FunctionCallback {
abstract class AbstractFunctionCallback<I, O> implements BiFunction<I, ToolContext, O>, FunctionCallback {
private final String name;
@@ -101,7 +101,7 @@ abstract class AbstractFunctionCallback<I, O> implements BiFunction<I, Map<Strin
}
@Override
public String call(String functionInput, Map<String, Object> toolContext) {
public String call(String functionInput, ToolContext toolContext) {
I request = fromJson(functionInput, inputType);
O response = this.apply(request, toolContext);
return this.responseConverter.apply(response);

View File

@@ -17,6 +17,8 @@ package org.springframework.ai.model.function;
import java.util.Map;
import org.springframework.ai.chat.model.ToolContext;
/**
* Represents a model function call handler. Implementations are registered with the
* Models and called on prompts that trigger the function call.
@@ -60,13 +62,13 @@ public interface FunctionCallback {
* @param functionInput JSON string with the function arguments to be passed to the
* function. The arguments are defined as JSON schema usually registered with the the
* model. Arguments are provided by the AI model.
* @param functionContext Map with the function context. The context is used to pass
* @param tooContext Map with the function context. The context is used to pass
* additional user provided state in addition to the arguments provided by the AI
* model.
* @return String containing the function call response.
*/
default String call(String functionInput, Map<String, Object> functionContext) {
if (functionContext != null) {
default String call(String functionInput, ToolContext tooContext) {
if (tooContext != null) {
throw new UnsupportedOperationException("Function context is not supported!");
}
return call(functionInput);

View File

@@ -16,12 +16,10 @@
package org.springframework.ai.model.function;
import java.lang.reflect.Type;
import java.util.Map;
import java.util.function.BiFunction;
import java.util.function.Function;
import com.fasterxml.jackson.annotation.JsonClassDescription;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.beans.BeansException;
import org.springframework.cloud.function.context.catalog.FunctionTypeUtils;
import org.springframework.cloud.function.context.config.FunctionContextUtils;
@@ -33,6 +31,8 @@ import org.springframework.lang.NonNull;
import org.springframework.lang.Nullable;
import org.springframework.util.StringUtils;
import com.fasterxml.jackson.annotation.JsonClassDescription;
/**
* A Spring {@link ApplicationContextAware} implementation that provides a way to retrieve
* a {@link Function} from the Spring context and wrap it into a {@link FunctionCallback}.
@@ -127,7 +127,7 @@ public class FunctionCallbackContext implements ApplicationContextAware {
.build();
}
else if (bean instanceof BiFunction<?, ?, ?> biFunction) {
return FunctionCallbackWrapper.builder((BiFunction<?, Map<String, Object>, ?>) biFunction)
return FunctionCallbackWrapper.builder((BiFunction<?, ToolContext, ?>) biFunction)
.withName(functionName)
.withSchemaType(this.schemaType)
.withDescription(functionDescription)

View File

@@ -15,10 +15,10 @@
*/
package org.springframework.ai.model.function;
import java.util.Map;
import java.util.function.BiFunction;
import java.util.function.Function;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.model.function.FunctionCallbackContext.SchemaType;
import org.springframework.util.Assert;
@@ -37,22 +37,21 @@ import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
*/
public class FunctionCallbackWrapper<I, O> extends AbstractFunctionCallback<I, O> {
private final BiFunction<I, Map<String, Object>, O> biFunction;
private final BiFunction<I, ToolContext, O> biFunction;
private FunctionCallbackWrapper(String name, String description, String inputTypeSchema, Class<I> inputType,
Function<O, String> responseConverter, ObjectMapper objectMapper,
BiFunction<I, Map<String, Object>, O> function) {
Function<O, String> responseConverter, ObjectMapper objectMapper, BiFunction<I, ToolContext, O> function) {
super(name, description, inputTypeSchema, inputType, responseConverter, objectMapper);
Assert.notNull(function, "Function must not be null");
this.biFunction = function;
}
@Override
public O apply(I input, Map<String, Object> context) {
public O apply(I input, ToolContext context) {
return this.biFunction.apply(input, context);
}
public static <I, O> Builder<I, O> builder(BiFunction<I, Map<String, Object>, O> biFunction) {
public static <I, O> Builder<I, O> builder(BiFunction<I, ToolContext, O> biFunction) {
return new Builder<>(biFunction);
}
@@ -68,13 +67,13 @@ public class FunctionCallbackWrapper<I, O> extends AbstractFunctionCallback<I, O
private Class<I> inputType;
private final BiFunction<I, Map<String, Object>, O> biFunction;
private final BiFunction<I, ToolContext, O> biFunction;
private final Function<I, O> function;
private SchemaType schemaType = SchemaType.JSON_SCHEMA;
public Builder(BiFunction<I, Map<String, Object>, O> biFunction) {
public Builder(BiFunction<I, ToolContext, O> biFunction) {
Assert.notNull(biFunction, "Function must not be null");
this.biFunction = biFunction;
this.function = null;
@@ -159,7 +158,7 @@ public class FunctionCallbackWrapper<I, O> extends AbstractFunctionCallback<I, O
this.inputTypeSchema = ModelOptionsUtils.getJsonSchema(this.inputType, upperCaseTypeValues);
}
BiFunction<I, Map<String, Object>, O> finalBiFunction = (this.biFunction != null) ? this.biFunction
BiFunction<I, ToolContext, O> finalBiFunction = (this.biFunction != null) ? this.biFunction
: (request, context) -> this.function.apply(request);
return new FunctionCallbackWrapper<>(this.name, this.description, this.inputTypeSchema, this.inputType,
@@ -167,9 +166,9 @@ public class FunctionCallbackWrapper<I, O> extends AbstractFunctionCallback<I, O
}
@SuppressWarnings("unchecked")
private static <I, O> Class<I> resolveInputType(BiFunction<I, Map<String, Object>, O> biFunction) {
private static <I, O> Class<I> resolveInputType(BiFunction<I, ToolContext, O> biFunction) {
return (Class<I>) TypeResolverHelper
.getBiFunctionInputClass((Class<BiFunction<I, Map<String, Object>, O>>) biFunction.getClass());
.getBiFunctionInputClass((Class<BiFunction<I, ToolContext, O>>) biFunction.getClass());
}
@SuppressWarnings("unchecked")

View File

@@ -77,6 +77,6 @@ public interface FunctionCallingOptions extends ChatOptions {
Map<String, Object> getToolContext();
void setToolContext(Map<String, Object> functionContext);
void setToolContext(Map<String, Object> tooContext);
}

View File

@@ -224,16 +224,18 @@ The https://github.com/spring-projects/spring-ai/blob/main/spring-ai-spring-boot
Spring AI now supports passing additional contextual information to function callbacks through a tool context. This feature allows you to provide extra data that can be used within the function execution, enhancing the flexibility and power of function calling.
The context information that is passed in as the second argument of a `java.util.BiFunction`. The `ToolContext` contains as an immutable `Map<String,Object>` allowing you to access key-value pairs.
==== How to Use Tool Context
You can set the tool context when building your chat options and use a BiFunction for your callback:
[source,java]
----
BiFunction<MockWeatherService.Request, Map<String, Object>, MockWeatherService.Response> weatherFunction =
BiFunction<MockWeatherService.Request, ToolContext, MockWeatherService.Response> weatherFunction =
(request, toolContext) -> {
String sessionId = (String) toolContext.get("sessionId");
String userId = (String) toolContext.get("userId");
String sessionId = (String) toolContext.getContext().get("sessionId");
String userId = (String) toolContext.getContext().get("userId");
// Use sessionId and userId in your function logic
double temperature = 0;

View File

@@ -15,22 +15,22 @@
*/
package org.springframework.ai.autoconfigure.azure.tool;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.ai.autoconfigure.azure.tool.DeploymentNameUtil.getDeploymentName;
import java.util.List;
import java.util.Map;
import java.util.function.BiFunction;
import java.util.function.Function;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.autoconfigure.azure.openai.AzureOpenAiAutoConfiguration;
import org.springframework.ai.azure.openai.AzureOpenAiChatModel;
import org.springframework.ai.azure.openai.AzureOpenAiChatOptions;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.function.FunctionCallingOptionsBuilder.PortableFunctionCallingOptions;
import org.springframework.boot.autoconfigure.AutoConfigurations;
@@ -39,9 +39,6 @@ import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Description;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.ai.autoconfigure.azure.tool.DeploymentNameUtil.getDeploymentName;
@EnabledIfEnvironmentVariable(named = "AZURE_OPENAI_API_KEY", matches = ".+")
@EnabledIfEnvironmentVariable(named = "AZURE_OPENAI_ENDPOINT", matches = ".+")
class FunctionCallWithFunctionBeanIT {
@@ -112,18 +109,6 @@ class FunctionCallWithFunctionBeanIT {
return new MockWeatherService();
}
@Bean
@Description("Get the weather in location")
public Function<MockWeatherService.Request, Function<Map<String, Object>, MockWeatherService.Response>> weatherFunctionWithContext() {
return request -> context -> new MockWeatherService().apply(request);
}
@Bean
@Description("Get the weather in location")
public BiFunction<MockWeatherService.Request, Map<String, Object>, MockWeatherService.Response> weatherFunctionWithContext2() {
return (request, context) -> new MockWeatherService().apply(request);
}
// Relies on the Request's JsonClassDescription annotation to provide the
// function description.
@Bean

View File

@@ -32,6 +32,7 @@ import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.function.FunctionCallingOptions;
import org.springframework.ai.model.function.FunctionCallingOptionsBuilder.PortableFunctionCallingOptions;
@@ -187,7 +188,7 @@ class FunctionCallbackWithPlainFunctionBeanIT {
@Bean
@Description("Get the weather in location")
public BiFunction<MockWeatherService.Request, Map<String, Object>, MockWeatherService.Response> weatherFunctionWithContext() {
public BiFunction<MockWeatherService.Request, ToolContext, MockWeatherService.Response> weatherFunctionWithContext() {
return (request, context) -> {
return new MockWeatherService().apply(request);
};