Skip to content

Commit c318a5b

Browse files
temporal-spring-ai: preserve Usage and RateLimit in ChatResponse metadata
ActivityChatModel.toResponse now rehydrates Usage and RateLimit onto the ChatResponseMetadata it returns to workflow code, not just the model name. The activity side (ChatModelActivityImpl) already serialized these into the output record; they were being silently discarded when the workflow side rebuilt the ChatResponse. Usage is rehydrated as a Spring AI DefaultUsage(promptTokens, completionTokens, totalTokens). RateLimit is an interface with no public default impl in spring-ai-model, so we return an anonymous implementation backed by the fields from the activity output record. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 32ae25b commit c318a5b

1 file changed

Lines changed: 63 additions & 1 deletion

File tree

temporal-spring-ai/src/main/java/io/temporal/springai/model/ActivityChatModel.java

Lines changed: 63 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@
1212
import java.util.stream.Collectors;
1313
import org.springframework.ai.chat.messages.*;
1414
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
15+
import org.springframework.ai.chat.metadata.DefaultUsage;
16+
import org.springframework.ai.chat.metadata.RateLimit;
17+
import org.springframework.ai.chat.metadata.Usage;
1518
import org.springframework.ai.chat.model.ChatModel;
1619
import org.springframework.ai.chat.model.ChatResponse;
1720
import org.springframework.ai.chat.model.Generation;
@@ -346,11 +349,70 @@ private ChatResponse toResponse(ChatModelTypes.ChatModelActivityOutput output) {
346349

347350
var builder = ChatResponse.builder().generations(generations);
348351
if (output.metadata() != null) {
349-
builder.metadata(ChatResponseMetadata.builder().model(output.metadata().model()).build());
352+
builder.metadata(toResponseMetadata(output.metadata()));
350353
}
351354
return builder.build();
352355
}
353356

357+
private ChatResponseMetadata toResponseMetadata(
358+
ChatModelTypes.ChatModelActivityOutput.ChatResponseMetadata md) {
359+
ChatResponseMetadata.Builder b = ChatResponseMetadata.builder().model(md.model());
360+
Usage usage = toUsage(md.usage());
361+
if (usage != null) {
362+
b.usage(usage);
363+
}
364+
RateLimit rateLimit = toRateLimit(md.rateLimit());
365+
if (rateLimit != null) {
366+
b.rateLimit(rateLimit);
367+
}
368+
return b.build();
369+
}
370+
371+
private Usage toUsage(ChatModelTypes.ChatModelActivityOutput.ChatResponseMetadata.Usage u) {
372+
if (u == null) {
373+
return null;
374+
}
375+
return new DefaultUsage(u.promptTokens(), u.completionTokens(), u.totalTokens());
376+
}
377+
378+
private RateLimit toRateLimit(
379+
ChatModelTypes.ChatModelActivityOutput.ChatResponseMetadata.RateLimit r) {
380+
if (r == null) {
381+
return null;
382+
}
383+
return new RateLimit() {
384+
@Override
385+
public Long getRequestsLimit() {
386+
return r.requestLimit();
387+
}
388+
389+
@Override
390+
public Long getRequestsRemaining() {
391+
return r.requestRemaining();
392+
}
393+
394+
@Override
395+
public java.time.Duration getRequestsReset() {
396+
return r.requestReset();
397+
}
398+
399+
@Override
400+
public Long getTokensLimit() {
401+
return r.tokenLimit();
402+
}
403+
404+
@Override
405+
public Long getTokensRemaining() {
406+
return r.tokenRemaining();
407+
}
408+
409+
@Override
410+
public java.time.Duration getTokensReset() {
411+
return r.tokenReset();
412+
}
413+
};
414+
}
415+
354416
private AssistantMessage toAssistantMessage(ChatModelTypes.Message message) {
355417
List<AssistantMessage.ToolCall> toolCalls = List.of();
356418
if (!CollectionUtils.isEmpty(message.toolCalls())) {

0 commit comments

Comments
 (0)