diff --git a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java index cd4162973..cba0070fd 100644 --- a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java +++ b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java @@ -222,7 +222,12 @@ public class OpenAiChatModel extends return generation; }).toList(); - return new ChatResponse(generations); + if (chatCompletion.usage() != null) { + return new ChatResponse(generations, OpenAiChatResponseMetadata.from(chatCompletion)); + } + else { + return new ChatResponse(generations); + } } catch (Exception e) { logger.error("Error processing chat completion", e); @@ -245,7 +250,7 @@ public class OpenAiChatModel extends .toList(); return new OpenAiApi.ChatCompletion(chunk.id(), choices, chunk.created(), chunk.model(), - chunk.systemFingerprint(), "chat.completion", null); + chunk.systemFingerprint(), "chat.completion", chunk.usage()); } /** @@ -306,6 +311,12 @@ public class OpenAiChatModel extends request, ChatCompletionRequest.class); } + // Remove `streamOptions` from the request if it is not a streaming request + if (request.streamOptions() != null && !stream) { + logger.warn("Removing streamOptions from the request as it is not a streaming request!"); + request = request.withStreamOptions(null); + } + return request; } diff --git a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatOptions.java b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatOptions.java index d6df88ba1..4b7cae9a0 100644 --- a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatOptions.java +++ b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatOptions.java @@ -27,6 +27,7 @@ import org.springframework.ai.model.function.FunctionCallback; import org.springframework.ai.model.function.FunctionCallingOptions; import org.springframework.ai.openai.api.OpenAiApi; import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionRequest.ResponseFormat; +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionRequest.StreamOptions; import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionRequest.ToolChoiceBuilder; import org.springframework.ai.openai.api.OpenAiApi.FunctionTool; import org.springframework.boot.context.properties.NestedConfigurationProperty; @@ -39,6 +40,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; /** * @author Christian Tzolov + * @author Mariusz Bernacki * @since 0.8.0 */ @JsonInclude(Include.NON_NULL) @@ -93,6 +95,10 @@ public class OpenAiChatOptions implements FunctionCallingOptions, ChatOptions { * "json_object" } enables JSON mode, which guarantees the message the model generates is valid JSON. */ private @JsonProperty("response_format") ResponseFormat responseFormat; + /** + * Options for streaming response. Included in the API only if streaming-mode completion is requested. + */ + private @JsonProperty("stream_options") StreamOptions streamOptions; /** * This feature is in Beta. If specified, our system will make a best effort to sample * deterministically, such that repeated requests with the same seed and parameters should return the same result. @@ -226,6 +232,13 @@ public class OpenAiChatOptions implements FunctionCallingOptions, ChatOptions { return this; } + public Builder withStreamUsage(boolean enableStreamUsage) { + if (enableStreamUsage) { + this.options.streamOptions = (enableStreamUsage) ? StreamOptions.INCLUDE_USAGE : null; + } + return this; + } + public Builder withSeed(Integer seed) { this.options.seed = seed; return this; @@ -284,6 +297,14 @@ public class OpenAiChatOptions implements FunctionCallingOptions, ChatOptions { } + public Boolean getStreamUsage() { + return this.streamOptions != null; + } + + public void setStreamUsage(Boolean enableStreamUsage) { + this.streamOptions = (enableStreamUsage) ? StreamOptions.INCLUDE_USAGE : null; + } + public String getModel() { return this.model; } @@ -356,6 +377,14 @@ public class OpenAiChatOptions implements FunctionCallingOptions, ChatOptions { this.responseFormat = responseFormat; } + public StreamOptions getStreamOptions() { + return streamOptions; + } + + public void setStreamOptions(StreamOptions streamOptions) { + this.streamOptions = streamOptions; + } + public Integer getSeed() { return this.seed; } @@ -446,6 +475,7 @@ public class OpenAiChatOptions implements FunctionCallingOptions, ChatOptions { result = prime * result + ((n == null) ? 0 : n.hashCode()); result = prime * result + ((presencePenalty == null) ? 0 : presencePenalty.hashCode()); result = prime * result + ((responseFormat == null) ? 0 : responseFormat.hashCode()); + result = prime * result + ((streamOptions == null) ? 0 : streamOptions.hashCode()); result = prime * result + ((seed == null) ? 0 : seed.hashCode()); result = prime * result + ((stop == null) ? 0 : stop.hashCode()); result = prime * result + ((temperature == null) ? 0 : temperature.hashCode()); @@ -519,6 +549,12 @@ public class OpenAiChatOptions implements FunctionCallingOptions, ChatOptions { } else if (!this.responseFormat.equals(other.responseFormat)) return false; + if (this.streamOptions == null) { + if (other.streamOptions != null) + return false; + } + else if (!this.streamOptions.equals(other.streamOptions)) + return false; if (this.seed == null) { if (other.seed != null) return false; @@ -586,6 +622,7 @@ public class OpenAiChatOptions implements FunctionCallingOptions, ChatOptions { .withN(fromOptions.getN()) .withPresencePenalty(fromOptions.getPresencePenalty()) .withResponseFormat(fromOptions.getResponseFormat()) + .withStreamUsage(fromOptions.getStreamUsage()) .withSeed(fromOptions.getSeed()) .withStop(fromOptions.getStop()) .withTemperature(fromOptions.getTemperature()) diff --git a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java index 64df4e92d..a468cde83 100644 --- a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java +++ b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java @@ -47,6 +47,7 @@ import org.springframework.web.reactive.function.client.WebClient; * * @author Christian Tzolov * @author Michael Lavelle + * @author Mariusz Bernacki */ public class OpenAiApi { @@ -314,6 +315,7 @@ public class OpenAiApi { * @param stop Up to 4 sequences where the API will stop generating further tokens. * @param stream If set, partial message deltas will be sent.Tokens will be sent as data-only server-sent events as * they become available, with the stream terminated by a data: [DONE] message. + * @param streamOptions Options for streaming response. Only set this when you set. * @param temperature What sampling temperature to use, between 0 and 1. Higher values like 0.8 will make the output * more random, while lower values like 0.2 will make it more focused and deterministic. We generally recommend * altering this or top_p but not both. @@ -345,6 +347,7 @@ public class OpenAiApi { @JsonProperty("seed") Integer seed, @JsonProperty("stop") List stop, @JsonProperty("stream") Boolean stream, + @JsonProperty("stream_options") StreamOptions streamOptions, @JsonProperty("temperature") Float temperature, @JsonProperty("top_p") Float topP, @JsonProperty("tools") List tools, @@ -360,7 +363,7 @@ public class OpenAiApi { */ public ChatCompletionRequest(List messages, String model, Float temperature) { this(messages, model, null, null, null, null, null, null, null, - null, null, null, false, temperature, null, + null, null, null, false, null, temperature, null, null, null, null); } @@ -375,7 +378,7 @@ public class OpenAiApi { */ public ChatCompletionRequest(List messages, String model, Float temperature, boolean stream) { this(messages, model, null, null, null, null, null, null, null, - null, null, null, stream, temperature, null, + null, null, null, stream, null, temperature, null, null, null, null); } @@ -391,7 +394,7 @@ public class OpenAiApi { public ChatCompletionRequest(List messages, String model, List tools, Object toolChoice) { this(messages, model, null, null, null, null, null, null, null, - null, null, null, false, 0.8f, null, + null, null, null, false, null, 0.8f, null, tools, toolChoice, null); } @@ -404,10 +407,22 @@ public class OpenAiApi { */ public ChatCompletionRequest(List messages, Boolean stream) { this(messages, null, null, null, null, null, null, null, null, - null, null, null, stream, null, null, + null, null, null, stream, null, null, null, null, null, null); } + /** + * Sets the {@link StreamOptions} for this request. + * + * @param streamOptions The new stream options to use. + * @return A new {@link ChatCompletionRequest} with the specified stream options. + */ + public ChatCompletionRequest withStreamOptions(StreamOptions streamOptions) { + return new ChatCompletionRequest(messages, model, frequencyPenalty, logitBias, logprobs, topLogprobs, maxTokens, n, presencePenalty, + responseFormat, seed, stop, stream, streamOptions, temperature, topP, + tools, toolChoice, user); + } + /** * Helper factory that creates a tool_choice of type 'none', 'auto' or selected function by name. */ @@ -437,6 +452,20 @@ public class OpenAiApi { public record ResponseFormat( @JsonProperty("type") String type) { } + + /** + * @param includeUsage If set, an additional chunk will be streamed + * before the data: [DONE] message. The usage field on this chunk + * shows the token usage statistics for the entire request, and + * the choices field will always be an empty array. All other chunks + * will also include a usage field, but with a null value. + */ + @JsonInclude(Include.NON_NULL) + public record StreamOptions( + @JsonProperty("include_usage") Boolean includeUsage) { + + public static StreamOptions INCLUDE_USAGE = new StreamOptions(true); + } } /** @@ -742,7 +771,8 @@ public class OpenAiApi { @JsonProperty("created") Long created, @JsonProperty("model") String model, @JsonProperty("system_fingerprint") String systemFingerprint, - @JsonProperty("object") String object) { + @JsonProperty("object") String object, + @JsonProperty("usage") Usage usage) { /** * Chat completion choice. @@ -825,7 +855,7 @@ public class OpenAiApi { // Flux> -> Flux> .concatMapIterable(window -> { Mono monoChunk = window.reduce( - new ChatCompletionChunk(null, null, null, null, null, null), + new ChatCompletionChunk(null, null, null, null, null, null, null), (previous, current) -> this.chunkMerger.merge(previous, current)); return List.of(monoChunk); }) diff --git a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiStreamFunctionCallingHelper.java b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiStreamFunctionCallingHelper.java index 6110cc7cf..b0e23ce36 100644 --- a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiStreamFunctionCallingHelper.java +++ b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiStreamFunctionCallingHelper.java @@ -28,6 +28,7 @@ import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionChunk.ChunkChoi import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage.ChatCompletionFunction; import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage.Role; import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage.ToolCall; +import org.springframework.ai.openai.api.OpenAiApi.Usage; import org.springframework.util.CollectionUtils; /** @@ -58,13 +59,14 @@ public class OpenAiStreamFunctionCallingHelper { String systemFingerprint = (current.systemFingerprint() != null ? current.systemFingerprint() : previous.systemFingerprint()); String object = (current.object() != null ? current.object() : previous.object()); + Usage usage = (current.usage() != null ? current.usage() : previous.usage()); ChunkChoice previousChoice0 = (CollectionUtils.isEmpty(previous.choices()) ? null : previous.choices().get(0)); ChunkChoice currentChoice0 = (CollectionUtils.isEmpty(current.choices()) ? null : current.choices().get(0)); ChunkChoice choice = merge(previousChoice0, currentChoice0); List chunkChoices = choice == null ? List.of() : List.of(choice); - return new ChatCompletionChunk(id, chunkChoices, created, model, systemFingerprint, object); + return new ChatCompletionChunk(id, chunkChoices, created, model, systemFingerprint, object, usage); } private ChunkChoice merge(ChunkChoice previous, ChunkChoice current) { diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/MessageTypeContentTests.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/MessageTypeContentTests.java index f6311ca9d..52de86f24 100644 --- a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/MessageTypeContentTests.java +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/MessageTypeContentTests.java @@ -61,7 +61,7 @@ public class MessageTypeContentTests { ArgumentCaptor promptCaptor; Flux fluxResponse = Flux - .generate(() -> new ChatCompletionChunk("id", List.of(), 0l, "model", "fp", "object"), (state, sink) -> { + .generate(() -> new ChatCompletionChunk("id", List.of(), 0l, "model", "fp", "object", null), (state, sink) -> { sink.next(state); sink.complete(); return state; diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatModelIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatModelIT.java index f96450a63..716fe491e 100644 --- a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatModelIT.java +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatModelIT.java @@ -15,6 +15,8 @@ */ package org.springframework.ai.openai.chat; +import static org.assertj.core.api.Assertions.assertThat; + import java.io.IOException; import java.net.URL; import java.util.ArrayList; @@ -29,14 +31,12 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import reactor.core.publisher.Flux; - -import org.springframework.ai.chat.model.ChatResponse; -import org.springframework.ai.chat.model.Generation; import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.Media; import org.springframework.ai.chat.messages.Message; 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.prompt.Prompt; import org.springframework.ai.chat.prompt.PromptTemplate; import org.springframework.ai.chat.prompt.SystemPromptTemplate; @@ -56,7 +56,7 @@ import org.springframework.core.io.ClassPathResource; import org.springframework.core.io.Resource; import org.springframework.util.MimeTypeUtils; -import static org.assertj.core.api.Assertions.assertThat; +import reactor.core.publisher.Flux; @SpringBootTest(classes = OpenAiTestConfiguration.class) @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") @@ -103,6 +103,24 @@ class OpenAiChatModelIT extends AbstractIT { } + @Test + void streamingWithTokenUsage() { + var promptOptions = OpenAiChatOptions.builder().withStreamUsage(true).withSeed(1).build(); + + var prompt = new Prompt("List two colors of the Polish flag. Be brief.", promptOptions); + var streamingTokenUsage = this.chatModel.stream(prompt).blockLast().getMetadata().getUsage(); + var referenceTokenUsage = this.chatModel.call(prompt).getMetadata().getUsage(); + + assertThat(streamingTokenUsage.getPromptTokens()).isGreaterThan(0); + assertThat(streamingTokenUsage.getGenerationTokens()).isGreaterThan(0); + assertThat(streamingTokenUsage.getTotalTokens()).isGreaterThan(0); + + assertThat(streamingTokenUsage.getPromptTokens()).isEqualTo(referenceTokenUsage.getPromptTokens()); + assertThat(streamingTokenUsage.getGenerationTokens()).isEqualTo(referenceTokenUsage.getGenerationTokens()); + assertThat(streamingTokenUsage.getTotalTokens()).isEqualTo(referenceTokenUsage.getTotalTokens()); + + } + @Test void listOutputConverter() { DefaultConversionService conversionService = new DefaultConversionService(); diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiRetryTests.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiRetryTests.java index 8a664e91a..ae17045e9 100644 --- a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiRetryTests.java +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiRetryTests.java @@ -167,7 +167,7 @@ public class OpenAiRetryTests { var choice = new ChatCompletionChunk.ChunkChoice(ChatCompletionFinishReason.STOP, 0, new ChatCompletionMessage("Response", Role.ASSISTANT), null); ChatCompletionChunk expectedChatCompletion = new ChatCompletionChunk("id", List.of(choice), 666l, "model", null, - null); + null, null); when(openAiApi.chatCompletionStream(isA(ChatCompletionRequest.class))) .thenThrow(new TransientAiException("Transient Error 1")) diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat/openai-chat.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat/openai-chat.adoc index de3a50d6f..df16ca67f 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat/openai-chat.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat/openai-chat.adoc @@ -105,6 +105,7 @@ The prefix `spring.ai.openai.chat` is the property prefix that lets you configur | spring.ai.openai.chat.options.toolChoice | Controls which (if any) function is called by the model. none means the model will not call a function and instead generates a message. auto means the model can pick between generating a message or calling a function. Specifying a particular function via {"type: "function", "function": {"name": "my_function"}} forces the model to call that function. none is the default when no functions are present. auto is the default if functions are present. | - | spring.ai.openai.chat.options.user | A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse. | - | spring.ai.openai.chat.options.functions | List of functions, identified by their names, to enable for function calling in a single prompt requests. Functions with those names must exist in the functionCallbacks registry. | - +| spring.ai.openai.chat.options.stream-usage | (For streaming only) Set to add an additional chunk with token usage statistics for the entire request. The `choices` field for this chuk is an empty array and all other chunks will also include a usage field, but with a null value. | false |==== NOTE: You can override the common `spring.ai.openai.base-url` and `spring.ai.openai.api-key` for the `ChatModel` and `EmbeddingModel` implementations. diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfigurationIT.java index fdd3a2abb..2a428ee94 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfigurationIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfigurationIT.java @@ -26,6 +26,7 @@ import org.apache.commons.logging.LogFactory; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.chat.metadata.Usage; import org.springframework.ai.chat.model.ChatResponse; import org.springframework.ai.chat.prompt.Prompt; import org.springframework.ai.embedding.EmbeddingResponse; @@ -113,6 +114,28 @@ public class OpenAiAutoConfigurationIT { }); } + @Test + void streamingWithTokenUsage() { + contextRunner.withPropertyValues("spring.ai.openai.chat.options.stream-usage=true").run(context -> { + OpenAiChatModel chatModel = context.getBean(OpenAiChatModel.class); + + Flux responseFlux = chatModel.stream(new Prompt(new UserMessage("Hello"))); + + Usage[] streamingTokenUsage = new Usage[1]; + String response = responseFlux.collectList().block().stream().map(chatResponse -> { + streamingTokenUsage[0] = chatResponse.getMetadata().getUsage(); + return (chatResponse.getResult() != null) ? chatResponse.getResult().getOutput().getContent() : ""; + }).collect(Collectors.joining()); + + assertThat(streamingTokenUsage[0].getPromptTokens()).isGreaterThan(0); + assertThat(streamingTokenUsage[0].getGenerationTokens()).isGreaterThan(0); + assertThat(streamingTokenUsage[0].getTotalTokens()).isGreaterThan(0); + + assertThat(response).isNotEmpty(); + logger.info("Response: " + response); + }); + } + @Test void embedding() { contextRunner.run(context -> {