28
28
import org .springframework .ai .bedrock .titan .api .TitanEmbeddingBedrockApi ;
29
29
import org .springframework .ai .bedrock .titan .api .TitanEmbeddingBedrockApi .TitanEmbeddingRequest ;
30
30
import org .springframework .ai .bedrock .titan .api .TitanEmbeddingBedrockApi .TitanEmbeddingResponse ;
31
+ import org .springframework .ai .chat .metadata .DefaultUsage ;
31
32
import org .springframework .ai .document .Document ;
32
33
import org .springframework .ai .embedding .AbstractEmbeddingModel ;
33
34
import org .springframework .ai .embedding .Embedding ;
34
35
import org .springframework .ai .embedding .EmbeddingOptions ;
35
36
import org .springframework .ai .embedding .EmbeddingRequest ;
36
37
import org .springframework .ai .embedding .EmbeddingResponse ;
38
+ import org .springframework .ai .embedding .EmbeddingResponseMetadata ;
37
39
import org .springframework .util .Assert ;
38
40
39
41
/**
@@ -89,6 +91,7 @@ public EmbeddingResponse call(EmbeddingRequest request) {
89
91
90
92
List <Embedding > embeddings = new ArrayList <>();
91
93
var indexCounter = new AtomicInteger (0 );
94
+ int tokenUsage = 0 ;
92
95
93
96
for (String inputContent : request .getInstructions ()) {
94
97
var apiRequest = createTitanEmbeddingRequest (inputContent , request .getOptions ());
@@ -111,6 +114,10 @@ public EmbeddingResponse call(EmbeddingRequest request) {
111
114
}
112
115
113
116
embeddings .add (new Embedding (response .embedding (), indexCounter .getAndIncrement ()));
117
+
118
+ if (response .inputTextTokenCount () != null ) {
119
+ tokenUsage += response .inputTextTokenCount ();
120
+ }
114
121
}
115
122
catch (Exception ex ) {
116
123
logger .error ("Titan API embedding failed for input at index {}: {}" , indexCounter .get (),
@@ -120,7 +127,10 @@ public EmbeddingResponse call(EmbeddingRequest request) {
120
127
}
121
128
}
122
129
123
- return new EmbeddingResponse (embeddings );
130
+ EmbeddingResponseMetadata embeddingResponseMetadata = new EmbeddingResponseMetadata ("" ,
131
+ getDefaultUsage (tokenUsage ));
132
+
133
+ return new EmbeddingResponse (embeddings , embeddingResponseMetadata );
124
134
}
125
135
126
136
private TitanEmbeddingRequest createTitanEmbeddingRequest (String inputContent , EmbeddingOptions requestOptions ) {
@@ -155,6 +165,10 @@ private String summarizeInput(String input) {
155
165
return input .length () > 100 ? input .substring (0 , 100 ) + "..." : input ;
156
166
}
157
167
168
+ private DefaultUsage getDefaultUsage (int tokens ) {
169
+ return new DefaultUsage (tokens , 0 );
170
+ }
171
+
158
172
public enum InputType {
159
173
160
174
TEXT , IMAGE
0 commit comments