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