Fix Moonshot Chat model toolcalling token usage #1927
Conversation
- Accumulate the token usage when toolcalling is invoked
- Fix both call() and stream() methods
- Add `usage` field to the Chat completion choice as the usage is returned via Choice
- Add Mootshot chatmodel ITs for functioncalling tests
|
@mxsl-gr Could you check this PR? |
no problem, I'll check and test this PR later |
@mxsl-gr thank you |
hi, @ilayaperumalg |
|
@mxsl-gr I don't seem to find the suggestions on this PR. Could you check? |
|
|
||
| ChatResponse chatResponse = new ChatResponse(generations, from(completionEntity.getBody())); | ||
| MoonshotApi.Usage usage = completionEntity.getBody().usage(); | ||
| Usage currentUsage = (usage != null) ? MoonshotUsage.from(usage) : new EmptyUsage(); |
There was a problem hiding this comment.
I think this line of code seems unnecessary. MoonshotApi.Usage can directly implement the Usage interface, and UsageUtils.getCumulativeUsage already handles null checks for the input. This approach ensures that callers don’t need to repeatedly check for currentUsage in every call, avoiding redundant code and improving cohesion.
If an EmptyUsage is truly needed, it could simply be a static constant like Usage.empty() or EmptyUsage.instance().
| @@ -286,8 +300,11 @@ public Flux<ChatResponse> stream(Prompt prompt) { | |||
| // @formatter:on | |||
| return buildGeneration(choice, metadata); | |||
| }).toList(); | |||
| MoonshotApi.Usage usage = chatCompletion2.usage(); | |||
| Usage currentUsage = (usage != null) ? MoonshotUsage.from(usage) : new EmptyUsage(); | |||
| */ | ||
| @SpringBootTest | ||
| @EnabledIfEnvironmentVariable(named = "MOONSHOT_API_KEY", matches = ".+") | ||
| public class MoonShotChatModelIT { |
There was a problem hiding this comment.
I recall there's a test class called MoonshotChatModelFunctionCallingIT that already contains unit tests for function calls. Instead of creating a new test class, maybe you could consider adding the new test methods to the existing one. And there is already a class org.springframework.ai.moonshot.chat.MoonshotChatModelIT
Also, the correct model name is Moonshot, not MoonShot.
sorry, i forget finish the review... |
usagefield to the Chat completion choice as the usage is returned via Choice