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

Make the StreamingChatClient a StreamingModelClient

Browse Source
This commit is contained in:
Christian Tzolov
2024-01-25 16:58:55 +01:00
committed by Mark Pollack
parent 5d55b68380
commit 4217a54c31
23 changed files with 66 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiChatClient.java
View File

@@ -183,7 +183,7 @@ public class AzureOpenAiChatClient implements ChatClient, StreamingChatClient {
}
@Override
public Flux<ChatResponse> generateStream(Prompt prompt) {
public Flux<ChatResponse> streamingCall(Prompt prompt) {
ChatCompletionsOptions options = toAzureChatCompletionsOptions(prompt);
options.setStream(true);

2
models/spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/AzureOpenAiChatClientIT.java
View File

@@ -151,7 +151,7 @@ class AzureOpenAiChatClientIT {
PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("format", format));
Prompt prompt = new Prompt(promptTemplate.createMessage());
String generationTextFromStream = chatClient.generateStream(prompt)
String generationTextFromStream = chatClient.streamingCall(prompt)
.collectList()
.block()
.stream()

2
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/anthropic/BedrockAnthropicChatClient.java
View File

@@ -107,7 +107,7 @@ public class BedrockAnthropicChatClient implements ChatClient, StreamingChatClie
}
@Override
public Flux<ChatResponse> generateStream(Prompt prompt) {
public Flux<ChatResponse> streamingCall(Prompt prompt) {
final String promptValue = MessageToPromptConverter.create().toPrompt(prompt.getInstructions());

2
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereChatClient.java
View File

@@ -122,7 +122,7 @@ public class BedrockCohereChatClient implements ChatClient, StreamingChatClient
}
@Override
public Flux<ChatResponse> generateStream(Prompt prompt) {
public Flux<ChatResponse> streamingCall(Prompt prompt) {
return this.chatApi.chatCompletionStream(this.createRequest(prompt, true)).map(g -> {
if (g.isFinished()) {
String finishReason = g.finishReason().name();

2
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/llama2/BedrockLlama2ChatClient.java
View File

@@ -85,7 +85,7 @@ public class BedrockLlama2ChatClient implements ChatClient, StreamingChatClient
}
@Override
public Flux<ChatResponse> generateStream(Prompt prompt) {
public Flux<ChatResponse> streamingCall(Prompt prompt) {
final String promptValue = MessageToPromptConverter.create().toPrompt(prompt.getInstructions());

2
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanChatClient.java
View File

@@ -84,7 +84,7 @@ public class BedrockTitanChatClient implements ChatClient, StreamingChatClient {
}
@Override
public Flux<ChatResponse> generateStream(Prompt prompt) {
public Flux<ChatResponse> streamingCall(Prompt prompt) {
return this.chatApi.chatCompletionStream(this.createRequest(prompt, true)).map(chunk -> {
Generation generation = new Generation(chunk.outputText());

2
models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/anthropic/BedrockAnthropicChatClientIT.java
View File

@@ -134,7 +134,7 @@ class BedrockAnthropicChatClientIT {
PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("format", format));
Prompt prompt = new Prompt(promptTemplate.createMessage());
String generationTextFromStream = client.generateStream(prompt)
String generationTextFromStream = client.streamingCall(prompt)
.collectList()
.block()
.stream()

2
models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereChatClientIT.java
View File

@@ -134,7 +134,7 @@ class BedrockCohereChatClientIT {
PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("format", format));
Prompt prompt = new Prompt(promptTemplate.createMessage());
String generationTextFromStream = client.generateStream(prompt)
String generationTextFromStream = client.streamingCall(prompt)
.collectList()
.block()
.stream()

2
models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/llama2/BedrockLlama2ChatClientIT.java
View File

@@ -139,7 +139,7 @@ class BedrockLlama2ChatClientIT {
PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("format", format));
Prompt prompt = new Prompt(promptTemplate.createMessage());
String generationTextFromStream = client.generateStream(prompt)
String generationTextFromStream = client.streamingCall(prompt)
.collectList()
.block()
.stream()

2
models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/BedrockTitanChatClientIT.java
View File

@@ -140,7 +140,7 @@ class BedrockTitanChatClientIT {
PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("format", format));
Prompt prompt = new Prompt(promptTemplate.createMessage());
String generationTextFromStream = client.generateStream(prompt)
String generationTextFromStream = client.streamingCall(prompt)
.collectList()
.block()
.stream()

2
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatClient.java
View File

@@ -97,7 +97,7 @@ public class OllamaChatClient implements ChatClient, StreamingChatClient {
}
@Override
public Flux<ChatResponse> generateStream(Prompt prompt) {
public Flux<ChatResponse> streamingCall(Prompt prompt) {
Flux<OllamaApi.ChatResponse> response = this.chatApi.streamingChat(request(prompt, this.model, true));

2
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatClientIT.java
View File

@@ -167,7 +167,7 @@ class OllamaChatClientIT {
PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("format", format));
Prompt prompt = new Prompt(promptTemplate.createMessage());
String generationTextFromStream = client.generateStream(prompt)
String generationTextFromStream = client.streamingCall(prompt)
.collectList()
.block()
.stream()

2
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatClient.java
View File

@@ -127,7 +127,7 @@ public class OpenAiChatClient implements ChatClient, StreamingChatClient {
}
@Override
public Flux<ChatResponse> generateStream(Prompt prompt) {
public Flux<ChatResponse> streamingCall(Prompt prompt) {
return this.retryTemplate.execute(ctx -> {
List<Message> messages = prompt.getInstructions();

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

@@ -140,7 +140,7 @@ class OpenAiChatClientIT extends AbstractIT {
PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("format", format));
Prompt prompt = new Prompt(promptTemplate.createMessage());
String generationTextFromStream = openStreamingChatClient.generateStream(prompt)
String generationTextFromStream = openStreamingChatClient.streamingCall(prompt)
.collectList()
.block()
.stream()

7
spring-ai-core/src/main/java/org/springframework/ai/chat/StreamingChatClient.java
View File

@@ -16,13 +16,10 @@
package org.springframework.ai.chat;
import reactor.core.publisher.Flux;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.StreamingModelClient;
@FunctionalInterface
public interface StreamingChatClient {
Flux<ChatResponse> generateStream(Prompt prompt);
public interface StreamingChatClient extends StreamingModelClient<Prompt, ChatResponse> {
}

43
spring-ai-core/src/main/java/org/springframework/ai/model/StreamingModelClient.java Normal file
View File

@@ -0,0 +1,43 @@
/*
* 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.model;
import reactor.core.publisher.Flux;
/**
* The StreamingModelClient interface provides a generic API for invoking a AI models with
* streaming response. It abstracts the process of sending requests and receiving a
* streaming responses. The interface uses Java generics to accommodate different types of
* requests and responses, enhancing flexibility and adaptability across different AI
* model implementations.
*
* @param <TReq> the generic type of the request to the AI model
* @param <TResChunk> the generic type of a single item in the streaming response from the
* AI model
* @author Christian Tzolov
* @since 0.8.0
*/
public interface StreamingModelClient<TReq extends ModelRequest<?>, TResChunk extends ModelResponse<?>> {
/**
* Executes a method call to the AI model.
* @param request the request object to be sent to the AI model
* @return the streaming response from the AI model
*/
Flux<TResChunk> streamingCall(TReq request);
}

2
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/azure/AzureOpenAiAutoConfigurationIT.java
View File

@@ -90,7 +90,7 @@ public class AzureOpenAiAutoConfigurationIT {
AzureOpenAiChatClient chatClient = context.getBean(AzureOpenAiChatClient.class);
Flux<ChatResponse> response = chatClient.generateStream(new Prompt(List.of(userMessage, systemMessage)));
Flux<ChatResponse> response = chatClient.streamingCall(new Prompt(List.of(userMessage, systemMessage)));
List<ChatResponse> responses = response.collectList().block();
assertThat(responses.size()).isGreaterThan(1);

2
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/anthropic/BedrockAnthropicChatAutoConfigurationIT.java
View File

@@ -83,7 +83,7 @@ public class BedrockAnthropicChatAutoConfigurationIT {
BedrockAnthropicChatClient anthropicChatClient = context.getBean(BedrockAnthropicChatClient.class);
Flux<ChatResponse> response = anthropicChatClient
.generateStream(new Prompt(List.of(userMessage, systemMessage)));
.streamingCall(new Prompt(List.of(userMessage, systemMessage)));
List<ChatResponse> responses = response.collectList().block();
assertThat(responses.size()).isGreaterThan(2);

2
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereChatAutoConfigurationIT.java
View File

@@ -85,7 +85,7 @@ public class BedrockCohereChatAutoConfigurationIT {
BedrockCohereChatClient cohereChatClient = context.getBean(BedrockCohereChatClient.class);
Flux<ChatResponse> response = cohereChatClient
.generateStream(new Prompt(List.of(userMessage, systemMessage)));
.streamingCall(new Prompt(List.of(userMessage, systemMessage)));
List<ChatResponse> responses = response.collectList().block();
assertThat(responses.size()).isGreaterThan(2);

2
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/llama2/BedrockLlama2ChatAutoConfigurationIT.java
View File

@@ -83,7 +83,7 @@ public class BedrockLlama2ChatAutoConfigurationIT {
BedrockLlama2ChatClient llama2ChatClient = context.getBean(BedrockLlama2ChatClient.class);
Flux<ChatResponse> response = llama2ChatClient
.generateStream(new Prompt(List.of(userMessage, systemMessage)));
.streamingCall(new Prompt(List.of(userMessage, systemMessage)));
List<ChatResponse> responses = response.collectList().block();
assertThat(responses.size()).isGreaterThan(2);

2
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/titan/BedrockTitanChatAutoConfigurationIT.java
View File

@@ -82,7 +82,7 @@ public class BedrockTitanChatAutoConfigurationIT {
BedrockTitanChatClient chatClient = context.getBean(BedrockTitanChatClient.class);
Flux<ChatResponse> response = chatClient.generateStream(new Prompt(List.of(userMessage, systemMessage)));
Flux<ChatResponse> response = chatClient.streamingCall(new Prompt(List.of(userMessage, systemMessage)));
List<ChatResponse> responses = response.collectList().block();
assertThat(responses.size()).isGreaterThan(1);

2
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/ollama/OllamaAutoConfigurationIT.java
View File

@@ -100,7 +100,7 @@ public class OllamaAutoConfigurationIT {
OllamaChatClient chatClient = context.getBean(OllamaChatClient.class);
Flux<ChatResponse> response = chatClient.generateStream(new Prompt(List.of(userMessage, systemMessage)));
Flux<ChatResponse> response = chatClient.streamingCall(new Prompt(List.of(userMessage, systemMessage)));
List<ChatResponse> responses = response.collectList().block();
assertThat(responses.size()).isGreaterThan(1);

2
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfigurationIT.java
View File

@@ -59,7 +59,7 @@ public class OpenAiAutoConfigurationIT {
void generateStreaming() {
contextRunner.run(context -> {
OpenAiChatClient client = context.getBean(OpenAiChatClient.class);
Flux<ChatResponse> responseFlux = client.generateStream(new Prompt(new UserMessage("Hello")));
Flux<ChatResponse> responseFlux = client.streamingCall(new Prompt(new UserMessage("Hello")));
String response = responseFlux.collectList().block().stream().map(chatResponse -> {
return chatResponse.getResults().get(0).getOutput().getContent();
}).collect(Collectors.joining());