Skip to content

fix: Make sure that XXXResponseSpec instances can only be used once #3886

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: Make sure that XXXResponseSpec instances can only be used once #3886

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions ...ng-ai-client-chat/src/main/java/org/springframework/ai/chat/client/DefaultChatClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;

import io.micrometer.observation.Observation;
Expand Down Expand Up @@ -348,6 +349,11 @@ public static class DefaultCallResponseSpec implements CallResponseSpec {

private final ChatClientObservationConvention observationConvention;

/**
* Used to ensure that the {@link CallResponseSpec} is used only once.
*/
private final AtomicBoolean used = new AtomicBoolean(false);

public DefaultCallResponseSpec(ChatClientRequest chatClientRequest, BaseAdvisorChain advisorChain,
ObservationRegistry observationRegistry, ChatClientObservationConvention observationConvention) {
Assert.notNull(chatClientRequest, "chatClientRequest cannot be null");
Expand Down Expand Up @@ -450,6 +456,10 @@ private ChatClientResponse doGetObservableChatClientResponse(ChatClientRequest c
private ChatClientResponse doGetObservableChatClientResponse(ChatClientRequest chatClientRequest,
@Nullable String outputFormat) {

if (!this.used.compareAndSet(false, true)) {
throw new IllegalStateException("The CallResponseSpec instance can only be used once.");
}

if (outputFormat != null) {
chatClientRequest.context().put(ChatClientAttributes.OUTPUT_FORMAT.getKey(), outputFormat);
}
Expand Down Expand Up @@ -494,6 +504,11 @@ public static class DefaultStreamResponseSpec implements StreamResponseSpec {

private final ChatClientObservationConvention observationConvention;

/**
* Used to ensure that the {@link StreamResponseSpec} is used only once.
*/
private final AtomicBoolean used = new AtomicBoolean(false);

public DefaultStreamResponseSpec(ChatClientRequest chatClientRequest, BaseAdvisorChain advisorChain,
ObservationRegistry observationRegistry, ChatClientObservationConvention observationConvention) {
Assert.notNull(chatClientRequest, "chatClientRequest cannot be null");
Expand All @@ -508,6 +523,11 @@ public DefaultStreamResponseSpec(ChatClientRequest chatClientRequest, BaseAdviso
}

private Flux<ChatClientResponse> doGetObservableFluxChatResponse(ChatClientRequest chatClientRequest) {

if (!this.used.compareAndSet(false, true)) {
throw new IllegalStateException("The StreamResponseSpec instance can only be used once.");
}

return Flux.deferContextual(contextView -> {

ChatClientObservationContext observationContext = ChatClientObservationContext.builder()
Expand Down
27 changes: 27 additions & 0 deletions ...-client-chat/src/test/java/org/springframework/ai/chat/client/DefaultChatClientTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@
import org.springframework.util.MimeTypeUtils;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatNoException;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.BDDMockito.given;
import static org.mockito.Mockito.mock;
Expand Down Expand Up @@ -686,6 +687,19 @@ void buildCallResponseSpecWithNullObservationConvention() {
.hasMessage("observationConvention cannot be null");
}

@Test
void whenUseCallResponseSpecTwiceThenThrow() {
ChatClient chatClient = new DefaultChatClientBuilder(mock(ChatModel.class)).build();
DefaultChatClient.DefaultChatClientRequestSpec chatClientRequestSpec = (DefaultChatClient.DefaultChatClientRequestSpec) chatClient
.prompt("question");
DefaultChatClient.DefaultCallResponseSpec spec = (DefaultChatClient.DefaultCallResponseSpec) chatClientRequestSpec
.call();

assertThatNoException().isThrownBy(() -> spec.content());
assertThatThrownBy(() -> spec.content()).isInstanceOf(IllegalStateException.class)
.hasMessage("The CallResponseSpec instance can only be used once.");
}

@Test
void whenSimplePromptThenChatClientResponse() {
ChatModel chatModel = mock(ChatModel.class);
Expand Down Expand Up @@ -1199,6 +1213,19 @@ void buildStreamResponseSpecWithNullObservationConvention() {
.hasMessage("observationConvention cannot be null");
}

@Test
void whenUseStreamResponseSpecTwiceThenThrow() {
ChatClient chatClient = new DefaultChatClientBuilder(mock(ChatModel.class)).build();
DefaultChatClient.DefaultChatClientRequestSpec chatClientRequestSpec = (DefaultChatClient.DefaultChatClientRequestSpec) chatClient
.prompt("question");
DefaultChatClient.DefaultStreamResponseSpec spec = (DefaultChatClient.DefaultStreamResponseSpec) chatClientRequestSpec
.stream();

assertThatNoException().isThrownBy(() -> spec.content());
assertThatThrownBy(() -> spec.content()).isInstanceOf(IllegalStateException.class)
.hasMessage("The StreamResponseSpec instance can only be used once.");
}

@Test
void whenSimplePromptThenFluxChatClientResponse() {
ChatModel chatModel = mock(ChatModel.class);
Expand Down