Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Extend ChatClient API

Browse Source 
- add support for custom StructuredOutputConverters usng entity(outputConverterInstance)
  - add mutate() method that returns ChatClient.Builder to create a new ChatClient whose
    settings are replicated from the ChatClient's default settings.
  - add prompt().mutate() method that returns ChatClient.Builder to create a new ChatClient
    whose settings are replicated from the current default and prompt settings.
This commit is contained in:
Christian Tzolov
2024-05-24 14:33:21 +02:00
parent c49a5c84ff
commit cd3b374dab
4 changed files with 317 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatClientIT.java → models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/client/OpenAiChatClientIT.java
View File

@@ -13,12 +13,11 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.openai.chat;
package org.springframework.ai.openai.chat.client;
import java.io.IOException;
import java.net.URL;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
@@ -34,6 +33,7 @@ import reactor.core.publisher.Flux;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.converter.BeanOutputConverter;
import org.springframework.ai.converter.ListOutputConverter;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.OpenAiTestConfiguration;
import org.springframework.ai.openai.api.OpenAiApi;
@@ -42,6 +42,7 @@ import org.springframework.ai.openai.testutils.AbstractIT;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.core.convert.support.DefaultConversionService;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.Resource;
import org.springframework.util.MimeTypeUtils;
@@ -79,20 +80,21 @@ class OpenAiChatClientIT extends AbstractIT {
}
@Test
void listOutputConverter() {
void listOutputConverterString() {
// @formatter:off
Collection<String> collection = ChatClient.builder(chatModel).build().prompt()
List<String> collection = ChatClient.builder(chatModel).build().prompt()
.user(u -> u.text("List five {subject}")
.param("subject", "ice cream flavors"))
.call()
.entity(new ParameterizedTypeReference<List<String>>() {});
// @formatter:on
logger.info(collection.toString());
assertThat(collection).hasSize(5);
}
@Test
void listOutputConverter2() {
void listOutputConverterBean() {
// @formatter:off
List<ActorsFilms> actorsFilms = ChatClient.builder(chatModel).build().prompt()
@@ -106,6 +108,24 @@ class OpenAiChatClientIT extends AbstractIT {
assertThat(actorsFilms).hasSize(2);
}
@Test
void customOutputConverter() {
var toStringListConverter = new ListOutputConverter(new DefaultConversionService());
// @formatter:off
List<String> flavors = ChatClient.builder(chatModel).build().prompt()
.user(u -> u.text("List five {subject}")
.param("subject", "ice cream flavors"))
.call()
.entity(toStringListConverter);
// @formatter:on
logger.info("ice cream flavors" + flavors);
assertThat(flavors).hasSize(5);
assertThat(flavors).contains("Vanilla");
}
@Test
void mapOutputConverter() {
// @formatter:off
@@ -196,6 +216,24 @@ class OpenAiChatClientIT extends AbstractIT {
assertThat(response).containsAnyOf("15.0", "15");
}
@Test
void defaultFunctionCallTest() {
// @formatter:off
String response = ChatClient.builder(chatModel)
.defaultFunction("getCurrentWeather", "Get the weather in location", new MockWeatherService())
.defaultUser(u -> u.text("What's the weather like in San Francisco, Tokyo, and Paris?"))
.build()
.prompt().call().content();
// @formatter:on
logger.info("Response: {}", response);
assertThat(response).containsAnyOf("30.0", "30");
assertThat(response).containsAnyOf("10.0", "10");
assertThat(response).containsAnyOf("15.0", "15");
}
@Test
void streamFunctionCallTest() {

51
spring-ai-core/src/main/java/org/springframework/ai/chat/client/ChatClient.java
View File

@@ -39,6 +39,7 @@ import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.converter.BeanOutputConverter;
import org.springframework.ai.converter.StructuredOutputConverter;
import org.springframework.ai.model.function.FunctionCallback;
import org.springframework.ai.model.function.FunctionCallbackWrapper;
import org.springframework.ai.model.function.FunctionCallingOptions;
@@ -74,6 +75,12 @@ public interface ChatClient {
ChatClientPromptRequest prompt(Prompt prompt);
/**
* Return a {@link ChatClient.Builder} to create a new {@link ChatClient} whose
* settings are replicated from the default {@link ChatClientRequest} of this client.
*/
Builder mutate();
interface PromptSpec<T> {
T text(String text);
@@ -223,6 +230,26 @@ public interface ChatClient {
private final Map<String, Object> systemParams = new HashMap<>();
/**
* Return a {@code ChatClient.Builder} to create a new {@code ChatClient} whose
* settings are replicated from this {@code ChatClientRequest}.
*/
public Builder mutate() {
Builder builder = ChatClient.builder(chatModel)
.defaultSystem(s -> s.text(this.systemText).params(this.systemParams))
.defaultUser(u -> u.text(this.userText)
.params(this.userParams)
.media(this.media.toArray(new Media[this.media.size()])))
.defaultOptions(this.chatOptions)
.defaultFunctions(StringUtils.toStringArray(this.functionNames));
// workaround to set the missing fields.
builder.defaultRequest.messages.addAll(this.messages);
builder.defaultRequest.functionCallbacks.addAll(this.functionCallbacks);
return builder;
}
/* copy constructor */
ChatClientRequest(ChatClientRequest ccr) {
this(ccr.chatModel, ccr.userText, ccr.userParams, ccr.systemText, ccr.systemParams, ccr.functionCallbacks,
@@ -411,7 +438,11 @@ public interface ChatClient {
return doSingleWithBeanOutputConverter(new BeanOutputConverter<T>(type));
}
private <T> T doSingleWithBeanOutputConverter(BeanOutputConverter<T> boc) {
public <T> T entity(StructuredOutputConverter<T> structuredOutputConverter) {
return doSingleWithBeanOutputConverter(structuredOutputConverter);
}
private <T> T doSingleWithBeanOutputConverter(StructuredOutputConverter<T> boc) {
var processedUserText = this.request.userText + System.lineSeparator() + System.lineSeparator()
+ "{format}";
var chatResponse = doGetChatResponse(processedUserText, boc.getFormat());
@@ -435,13 +466,15 @@ public interface ChatClient {
userParams.put("format", formatParam);
}
var messages = new ArrayList<Message>();
var messages = new ArrayList<Message>(this.request.messages);
var textsAreValid = (StringUtils.hasText(processedUserText)
|| StringUtils.hasText(this.request.systemText));
var messagesAreValid = !this.request.messages.isEmpty();
Assert.state(!(messagesAreValid && textsAreValid), "you must specify either " + Message.class.getName()
+ " instances or user/system texts, but not both");
if (textsAreValid) {
if (StringUtils.hasText(this.request.systemText) || !this.request.systemParams.isEmpty()) {
var systemMessage = new SystemMessage(
new PromptTemplate(this.request.systemText, this.request.systemParams).render());
messages.add(systemMessage);
}
UserMessage userMessage = null;
if (!CollectionUtils.isEmpty(userParams)) {
userMessage = new UserMessage(new PromptTemplate(processedUserText, userParams).render(),
@@ -450,16 +483,8 @@ public interface ChatClient {
else {
userMessage = new UserMessage(processedUserText, this.request.media);
}
if (StringUtils.hasText(this.request.systemText) || !this.request.systemParams.isEmpty()) {
var systemMessage = new SystemMessage(
new PromptTemplate(this.request.systemText, this.request.systemParams).render());
messages.add(systemMessage);
}
messages.add(userMessage);
}
else {
messages.addAll(this.request.messages);
}
if (this.request.chatOptions instanceof FunctionCallingOptions functionCallingOptions) {
// if (this.request.chatOptions instanceof
// FunctionCallingOptionsBuilder.PortableFunctionCallingOptions

9
spring-ai-core/src/main/java/org/springframework/ai/chat/client/DefaultChatClient.java
View File

@@ -35,6 +35,15 @@ class DefaultChatClient implements ChatClient {
return new ChatClientPromptRequest(this.chatModel, prompt);
}
/**
* Return a {@code ChatClient.Builder} to create a new {@code ChatClient} whose
* settings are replicated from this {@code ChatClientRequest}.
*/
@Override
public Builder mutate() {
return this.defaultChatClientRequest.mutate();
}
/**
* use the new fluid DSL starting in {@link #prompt()}
* @param prompt the {@link Prompt prompt} object

227
spring-ai-core/src/test/java/org/springframework/ai/chat/client/ChatClientTest.java
View File

@@ -19,6 +19,7 @@ package org.springframework.ai.chat.client;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.List;
import java.util.function.Function;
import java.util.stream.Collectors;
import org.junit.jupiter.api.Test;
@@ -38,6 +39,9 @@ import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.function.FunctionCallingOptions;
import org.springframework.ai.model.function.FunctionCallingOptionsBuilder;
import org.springframework.ai.model.function.FunctionCallingOptionsBuilder.PortableFunctionCallingOptions;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.util.MimeTypeUtils;
import static org.assertj.core.api.Assertions.assertThat;
@@ -198,6 +202,229 @@ public class ChatClientTest {
assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
}
static Function<String, String> mockFunction = new Function<String, String>() {
@Override
public String apply(String s) {
return s;
}
};
@Test
public void mutateDefaults() {
PortableFunctionCallingOptions options = new FunctionCallingOptionsBuilder().build();
when(chatModel.getDefaultOptions()).thenReturn(options);
when(chatModel.call(promptCaptor.capture())).thenReturn(new ChatResponse(List.of(new Generation("response"))));
when(chatModel.stream(promptCaptor.capture()))
.thenReturn(Flux.generate(() -> new ChatResponse(List.of(new Generation("response"))), (state, sink) -> {
sink.next(state);
sink.complete();
return state;
}));
// @formatter:off
var chatClient = ChatClient.builder(chatModel)
.defaultSystem(s -> s.text("Default system text {param1}, {param2}")
.param("param1", "value1")
.param("param2", "value2"))
.defaultFunctions("fun1", "fun2")
.defaultFunction("fun3", "fun3description", mockFunction)
.defaultUser(u -> u.text("Default user text {uparam1}, {uparam2}")
.param("uparam1", "value1")
.param("uparam2", "value2")
.media(MimeTypeUtils.IMAGE_JPEG,
new DefaultResourceLoader().getResource("classpath:/bikes.json")))
.build();
// @formatter:on
var content = chatClient.prompt().call().content();
assertThat(content).isEqualTo("response");
Prompt prompt = promptCaptor.getValue();
Message systemMessage = prompt.getInstructions().get(0);
assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
assertThat(systemMessage.getContent()).isEqualTo("Default system text value1, value2");
Message userMessage = prompt.getInstructions().get(1);
assertThat(userMessage.getMessageType()).isEqualTo(MessageType.USER);
assertThat(userMessage.getContent()).isEqualTo("Default user text value1, value2");
assertThat(userMessage.getMedia()).hasSize(1);
assertThat(userMessage.getMedia().iterator().next().getMimeType()).isEqualTo(MimeTypeUtils.IMAGE_JPEG);
var fco = (FunctionCallingOptions) prompt.getOptions();
assertThat(fco.getFunctions()).containsExactly("fun1", "fun2");
assertThat(fco.getFunctionCallbacks().iterator().next().getName()).isEqualTo("fun3");
// Streaming
content = join(chatClient.prompt().stream().content());
assertThat(content).isEqualTo("response");
prompt = promptCaptor.getValue();
systemMessage = prompt.getInstructions().get(0);
assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
assertThat(systemMessage.getContent()).isEqualTo("Default system text value1, value2");
userMessage = prompt.getInstructions().get(1);
assertThat(userMessage.getMessageType()).isEqualTo(MessageType.USER);
assertThat(userMessage.getContent()).isEqualTo("Default user text value1, value2");
assertThat(userMessage.getMedia()).hasSize(1);
assertThat(userMessage.getMedia().iterator().next().getMimeType()).isEqualTo(MimeTypeUtils.IMAGE_JPEG);
fco = (FunctionCallingOptions) prompt.getOptions();
assertThat(fco.getFunctions()).containsExactly("fun1", "fun2");
assertThat(fco.getFunctionCallbacks().iterator().next().getName()).isEqualTo("fun3");
// mutate builder
// @formatter:off
chatClient = chatClient.mutate()
.defaultSystem("Mutated default system text {param1}, {param2}")
.defaultFunctions("fun4")
.defaultUser("Mutated default user text {uparam1}, {uparam2}")
.build();
// @formatter:on
content = chatClient.prompt().call().content();
assertThat(content).isEqualTo("response");
prompt = promptCaptor.getValue();
systemMessage = prompt.getInstructions().get(0);
assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
assertThat(systemMessage.getContent()).isEqualTo("Mutated default system text value1, value2");
userMessage = prompt.getInstructions().get(1);
assertThat(userMessage.getMessageType()).isEqualTo(MessageType.USER);
assertThat(userMessage.getContent()).isEqualTo("Mutated default user text value1, value2");
assertThat(userMessage.getMedia()).hasSize(1);
assertThat(userMessage.getMedia().iterator().next().getMimeType()).isEqualTo(MimeTypeUtils.IMAGE_JPEG);
fco = (FunctionCallingOptions) prompt.getOptions();
assertThat(fco.getFunctions()).containsExactly("fun1", "fun2", "fun4");
assertThat(fco.getFunctionCallbacks().iterator().next().getName()).isEqualTo("fun3");
// Streaming
content = join(chatClient.prompt().stream().content());
assertThat(content).isEqualTo("response");
prompt = promptCaptor.getValue();
systemMessage = prompt.getInstructions().get(0);
assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
assertThat(systemMessage.getContent()).isEqualTo("Mutated default system text value1, value2");
userMessage = prompt.getInstructions().get(1);
assertThat(userMessage.getMessageType()).isEqualTo(MessageType.USER);
assertThat(userMessage.getContent()).isEqualTo("Mutated default user text value1, value2");
assertThat(userMessage.getMedia()).hasSize(1);
assertThat(userMessage.getMedia().iterator().next().getMimeType()).isEqualTo(MimeTypeUtils.IMAGE_JPEG);
fco = (FunctionCallingOptions) prompt.getOptions();
assertThat(fco.getFunctions()).containsExactly("fun1", "fun2", "fun4");
assertThat(fco.getFunctionCallbacks().iterator().next().getName()).isEqualTo("fun3");
}
@Test
public void mutatePrompt() {
PortableFunctionCallingOptions options = new FunctionCallingOptionsBuilder().build();
when(chatModel.getDefaultOptions()).thenReturn(options);
when(chatModel.call(promptCaptor.capture())).thenReturn(new ChatResponse(List.of(new Generation("response"))));
when(chatModel.stream(promptCaptor.capture()))
.thenReturn(Flux.generate(() -> new ChatResponse(List.of(new Generation("response"))), (state, sink) -> {
sink.next(state);
sink.complete();
return state;
}));
// @formatter:off
var chatClient = ChatClient.builder(chatModel)
.defaultSystem(s -> s.text("Default system text {param1}, {param2}")
.param("param1", "value1")
.param("param2", "value2"))
.defaultFunctions("fun1", "fun2")
.defaultFunction("fun3", "fun3description", mockFunction)
.defaultUser(u -> u.text("Default user text {uparam1}, {uparam2}")
.param("uparam1", "value1")
.param("uparam2", "value2")
.media(MimeTypeUtils.IMAGE_JPEG,
new DefaultResourceLoader().getResource("classpath:/bikes.json")))
.build();
var content = chatClient
.prompt()
.system("New default system text {param1}, {param2}")
.user(u -> u.param("uparam1", "userValue1")
.param("uparam2", "userValue2"))
.functions("fun5")
.mutate().build() // mutate and build new prompt
.prompt().call().content();
// @formatter:on
assertThat(content).isEqualTo("response");
Prompt prompt = promptCaptor.getValue();
Message systemMessage = prompt.getInstructions().get(0);
assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
assertThat(systemMessage.getContent()).isEqualTo("New default system text value1, value2");
Message userMessage = prompt.getInstructions().get(1);
assertThat(userMessage.getMessageType()).isEqualTo(MessageType.USER);
assertThat(userMessage.getContent()).isEqualTo("Default user text userValue1, userValue2");
assertThat(userMessage.getMedia()).hasSize(1);
assertThat(userMessage.getMedia().iterator().next().getMimeType()).isEqualTo(MimeTypeUtils.IMAGE_JPEG);
var fco = (FunctionCallingOptions) prompt.getOptions();
assertThat(fco.getFunctions()).containsExactly("fun1", "fun2", "fun5");
assertThat(fco.getFunctionCallbacks().iterator().next().getName()).isEqualTo("fun3");
// Streaming
// @formatter:off
content = join(chatClient
.prompt()
.system("New default system text {param1}, {param2}")
.user(u -> u.param("uparam1", "userValue1")
.param("uparam2", "userValue2"))
.functions("fun5")
.mutate().build() // mutate and build new prompt
.prompt().stream().content());
// @formatter:on
assertThat(content).isEqualTo("response");
prompt = promptCaptor.getValue();
systemMessage = prompt.getInstructions().get(0);
assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
assertThat(systemMessage.getContent()).isEqualTo("New default system text value1, value2");
userMessage = prompt.getInstructions().get(1);
assertThat(userMessage.getMessageType()).isEqualTo(MessageType.USER);
assertThat(userMessage.getContent()).isEqualTo("Default user text userValue1, userValue2");
assertThat(userMessage.getMedia()).hasSize(1);
assertThat(userMessage.getMedia().iterator().next().getMimeType()).isEqualTo(MimeTypeUtils.IMAGE_JPEG);
fco = (FunctionCallingOptions) prompt.getOptions();
assertThat(fco.getFunctions()).containsExactly("fun1", "fun2", "fun5");
assertThat(fco.getFunctionCallbacks().iterator().next().getName()).isEqualTo("fun3");
}
@Test
public void defaultUserText() {