Skip to content

Commit aed37fc

Browse files
ilayaperumalgscionaltera
authored andcommitted
Fix AzureOpenAiChatModelFunctionCallIT#streamFunctionCallUsageTest
- Fix the test to identify the right usage metadata from the stream content Signed-off-by: Ilayaperumal Gopinathan <[email protected]>
1 parent 09b0e7f commit aed37fc

File tree

1 file changed

+10
-4
lines changed

1 file changed

+10
-4
lines changed

models/spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/function/AzureOpenAiChatModelFunctionCallIT.java

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -163,12 +163,18 @@ void streamFunctionCallUsageTest() {
163163
.streamOptions(streamOptions)
164164
.build();
165165

166-
Flux<ChatResponse> response = this.chatModel.stream(new Prompt(messages, promptOptions));
166+
List<ChatResponse> responses = this.chatModel.stream(new Prompt(messages, promptOptions)).collectList().block();
167+
168+
assertThat(responses).isNotEmpty();
169+
170+
ChatResponse finalResponse = responses.get(responses.size() - 2);
171+
172+
logger.info("Final Response: {}", finalResponse);
167173

168-
ChatResponse chatResponse = response.last().block();
169-
logger.info("Response: {}", chatResponse);
174+
assertThat(finalResponse.getMetadata()).isNotNull();
175+
assertThat(finalResponse.getMetadata().getUsage()).isNotNull();
170176

171-
assertThat(chatResponse.getMetadata().getUsage().getTotalTokens()).isGreaterThan(600).isLessThan(800);
177+
assertThat(finalResponse.getMetadata().getUsage().getTotalTokens()).isGreaterThan(600).isLessThan(800);
172178

173179
}
174180

0 commit comments

Comments
 (0)