Include metadata when creating ChatResponse in Bedrock Titan ChatModel

Signed-off-by: jitokim <[email protected]>

Author: jitokim

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanChatModel.java

@@ -20,13 +20,15 @@
 
 import reactor.core.publisher.Flux;
 
+import org.springframework.ai.bedrock.BedrockUsage;
 import org.springframework.ai.bedrock.MessageToPromptConverter;
 import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi;
 import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatRequest;
 import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatResponse;
 import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatResponseChunk;
 import org.springframework.ai.chat.messages.AssistantMessage;
 import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
+import org.springframework.ai.chat.metadata.ChatResponseMetadata;
 import org.springframework.ai.chat.metadata.Usage;
 import org.springframework.ai.chat.model.ChatModel;
 import org.springframework.ai.chat.model.ChatResponse;
@@ -42,6 +44,7 @@
  * uses the Titan Chat API.
  *
  * @author Christian Tzolov
+ * @author Jihoon Kim
  * @since 0.8.0
  */
 public class BedrockTitanChatModel implements ChatModel, StreamingChatModel {
@@ -85,8 +88,15 @@ else if (chunk.inputTextTokenCount() != null && chunk.totalOutputTextTokenCount(
 				String completionReason = chunk.completionReason().name();
 				chatGenerationMetadata = ChatGenerationMetadata.from(completionReason, extractUsage(chunk));
 			}
+
+			ChatResponseMetadata chatResponseMetadata = ChatResponseMetadata.builder()
+				.withModel(prompt.getOptions().getModel())
+				.withUsage(BedrockUsage.from(chunk.amazonBedrockInvocationMetrics()))
+				.build();
+
 			return new ChatResponse(
-					List.of(new Generation(new AssistantMessage(chunk.outputText()), chatGenerationMetadata)));
+					List.of(new Generation(new AssistantMessage(chunk.outputText()), chatGenerationMetadata)),
+					chatResponseMetadata);
 		});
 	}