18
18
19
19
import io .micrometer .common .KeyValue ;
20
20
import io .micrometer .common .KeyValues ;
21
-
21
+ import org .springframework .ai .chat .metadata .Usage ;
22
+ import org .springframework .ai .embedding .EmbeddingOptions ;
23
+ import org .springframework .ai .embedding .EmbeddingResponse ;
24
+ import org .springframework .ai .embedding .EmbeddingResponseMetadata ;
22
25
import org .springframework .util .StringUtils ;
23
26
27
+ import java .util .Optional ;
28
+
24
29
/**
25
30
* Default conventions to populate observations for embedding model operations.
26
31
*
27
32
* @author Thomas Vitale
28
33
* @author Soby Chacko
34
+ * @author Mengqi Xu
29
35
* @since 1.0.0
30
36
*/
31
37
public class DefaultEmbeddingModelObservationConvention implements EmbeddingModelObservationConvention {
@@ -45,11 +51,11 @@ public String getName() {
45
51
46
52
@ Override
47
53
public String getContextualName (EmbeddingModelObservationContext context ) {
48
- if ( StringUtils . hasText (context .getRequest ().getOptions (). getModel ())) {
49
- return "%s %s" . formatted ( context . getOperationMetadata (). operationType (),
50
- context . getRequest (). getOptions (). getModel ());
51
- }
52
- return context .getOperationMetadata ().operationType ();
54
+ return Optional . ofNullable (context .getRequest ().getOptions ())
55
+ . map ( EmbeddingOptions :: getModel )
56
+ . filter ( StringUtils :: hasText )
57
+ . map ( model -> "%s %s" . formatted ( context . getOperationMetadata (). operationType (), model ))
58
+ . orElseGet (() -> context .getOperationMetadata ().operationType () );
53
59
}
54
60
55
61
@ Override
@@ -69,20 +75,22 @@ protected KeyValue aiProvider(EmbeddingModelObservationContext context) {
69
75
}
70
76
71
77
protected KeyValue requestModel (EmbeddingModelObservationContext context ) {
72
- if (StringUtils .hasText (context .getRequest ().getOptions ().getModel ())) {
73
- return KeyValue .of (EmbeddingModelObservationDocumentation .LowCardinalityKeyNames .REQUEST_MODEL ,
74
- context .getRequest ().getOptions ().getModel ());
75
- }
76
- return REQUEST_MODEL_NONE ;
78
+ return Optional .ofNullable (context .getRequest ().getOptions ())
79
+ .map (EmbeddingOptions ::getModel )
80
+ .filter (StringUtils ::hasText )
81
+ .map (model -> KeyValue .of (EmbeddingModelObservationDocumentation .LowCardinalityKeyNames .REQUEST_MODEL ,
82
+ model ))
83
+ .orElse (REQUEST_MODEL_NONE );
77
84
}
78
85
79
86
protected KeyValue responseModel (EmbeddingModelObservationContext context ) {
80
- if (context .getResponse () != null && context .getResponse ().getMetadata () != null
81
- && StringUtils .hasText (context .getResponse ().getMetadata ().getModel ())) {
82
- return KeyValue .of (EmbeddingModelObservationDocumentation .LowCardinalityKeyNames .RESPONSE_MODEL ,
83
- context .getResponse ().getMetadata ().getModel ());
84
- }
85
- return RESPONSE_MODEL_NONE ;
87
+ return Optional .ofNullable (context .getResponse ())
88
+ .map (EmbeddingResponse ::getMetadata )
89
+ .map (EmbeddingResponseMetadata ::getModel )
90
+ .filter (StringUtils ::hasText )
91
+ .map (model -> KeyValue .of (EmbeddingModelObservationDocumentation .LowCardinalityKeyNames .RESPONSE_MODEL ,
92
+ model ))
93
+ .orElse (RESPONSE_MODEL_NONE );
86
94
}
87
95
88
96
@ Override
@@ -99,36 +107,36 @@ public KeyValues getHighCardinalityKeyValues(EmbeddingModelObservationContext co
99
107
// Request
100
108
101
109
protected KeyValues requestEmbeddingDimension (KeyValues keyValues , EmbeddingModelObservationContext context ) {
102
- if (context .getRequest ().getOptions ().getDimensions () != null ) {
103
- return keyValues
110
+ return Optional .ofNullable (context .getRequest ().getOptions ())
111
+ .map (EmbeddingOptions ::getDimensions )
112
+ .map (dimensions -> keyValues
104
113
.and (EmbeddingModelObservationDocumentation .HighCardinalityKeyNames .REQUEST_EMBEDDING_DIMENSIONS
105
- .asString (), String .valueOf (context .getRequest ().getOptions ().getDimensions ()));
106
- }
107
- return keyValues ;
114
+ .asString (), String .valueOf (dimensions )))
115
+ .orElse (keyValues );
108
116
}
109
117
110
118
// Response
111
119
112
120
protected KeyValues usageInputTokens (KeyValues keyValues , EmbeddingModelObservationContext context ) {
113
- if (context .getResponse () != null && context .getResponse ().getMetadata () != null
114
- && context .getResponse ().getMetadata ().getUsage () != null
115
- && context .getResponse ().getMetadata ().getUsage ().getPromptTokens () != null ) {
116
- return keyValues .and (
121
+ return Optional .ofNullable (context .getResponse ())
122
+ .map (EmbeddingResponse ::getMetadata )
123
+ .map (EmbeddingResponseMetadata ::getUsage )
124
+ .map (Usage ::getPromptTokens )
125
+ .map (promptTokens -> keyValues .and (
117
126
EmbeddingModelObservationDocumentation .HighCardinalityKeyNames .USAGE_INPUT_TOKENS .asString (),
118
- String .valueOf (context .getResponse ().getMetadata ().getUsage ().getPromptTokens ()));
119
- }
120
- return keyValues ;
127
+ String .valueOf (promptTokens )))
128
+ .orElse (keyValues );
121
129
}
122
130
123
131
protected KeyValues usageTotalTokens (KeyValues keyValues , EmbeddingModelObservationContext context ) {
124
- if (context .getResponse () != null && context .getResponse ().getMetadata () != null
125
- && context .getResponse ().getMetadata ().getUsage () != null
126
- && context .getResponse ().getMetadata ().getUsage ().getTotalTokens () != null ) {
127
- return keyValues .and (
132
+ return Optional .ofNullable (context .getResponse ())
133
+ .map (EmbeddingResponse ::getMetadata )
134
+ .map (EmbeddingResponseMetadata ::getUsage )
135
+ .map (Usage ::getTotalTokens )
136
+ .map (totalTokens -> keyValues .and (
128
137
EmbeddingModelObservationDocumentation .HighCardinalityKeyNames .USAGE_TOTAL_TOKENS .asString (),
129
- String .valueOf (context .getResponse ().getMetadata ().getUsage ().getTotalTokens ()));
130
- }
131
- return keyValues ;
138
+ String .valueOf (totalTokens )))
139
+ .orElse (keyValues );
132
140
}
133
141
134
142
}
0 commit comments