1818
1919import io .micrometer .common .KeyValue ;
2020import 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 ;
2225import org .springframework .util .StringUtils ;
2326
27+ import java .util .Optional ;
28+
2429/**
2530 * Default conventions to populate observations for embedding model operations.
2631 *
2732 * @author Thomas Vitale
2833 * @author Soby Chacko
34+ * @author Mengqi Xu
2935 * @since 1.0.0
3036 */
3137public class DefaultEmbeddingModelObservationConvention implements EmbeddingModelObservationConvention {
@@ -45,11 +51,11 @@ public String getName() {
4551
4652 @ Override
4753 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 () );
5359 }
5460
5561 @ Override
@@ -69,20 +75,22 @@ protected KeyValue aiProvider(EmbeddingModelObservationContext context) {
6975 }
7076
7177 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 );
7784 }
7885
7986 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 );
8694 }
8795
8896 @ Override
@@ -99,36 +107,36 @@ public KeyValues getHighCardinalityKeyValues(EmbeddingModelObservationContext co
99107 // Request
100108
101109 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
104113 .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 );
108116 }
109117
110118 // Response
111119
112120 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 (
117126 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 );
121129 }
122130
123131 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 (
128137 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 );
132140 }
133141
134142}
0 commit comments