Skip to content

Commit c852d8e

Browse files
peachisaipeachisai
committed
Fix suggestions
1 parent e7923be commit c852d8e

6 files changed

Lines changed: 197 additions & 58 deletions

File tree

‎CHANGES.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ Release Notes.
2525
* Only publish `apm-application-toolkit` modules to Maven Central. Agent and plugins are distributed via download package and Docker images.
2626
* Add unified release script (`tools/releasing/release.sh`) with two-step flow: `prepare-vote` and `vote-passed`.
2727
* Fix an issue where `JDBCPluginConfig.Plugin.JDBC.SQL_BODY_MAX_LENGTH` was not honored by clickhouse-0.3.1 and clickhouse-0.3.2.x plugins.
28+
- Add tracing support for vector-store retrieval operations.
2829

2930
All issues and pull requests are [here](https://github.com/apache/skywalking/milestone/249?closed=1)
3031

‎apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/AbstractObservationVectorStoreConstructorInterceptor.java‎

Lines changed: 11 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -20,74 +20,31 @@
2020

2121
import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
2222
import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceConstructorInterceptor;
23-
import org.springframework.ai.embedding.EmbeddingOptions;
2423
import org.springframework.ai.embedding.EmbeddingModel;
25-
import org.springframework.util.StringUtils;
26-
27-
import java.lang.reflect.Field;
28-
import java.lang.reflect.Method;
2924

3025
public class AbstractObservationVectorStoreConstructorInterceptor implements InstanceConstructorInterceptor {
3126

3227
@Override
3328
public void onConstruct(EnhancedInstance objInst, Object[] allArguments) {
3429
if (allArguments != null && allArguments.length > 0) {
35-
String embeddingModelName = resolveModelFromArgument(allArguments[0]);
36-
objInst.setSkyWalkingDynamicField(new VectorStoreEnhanceContext(embeddingModelName));
37-
}
38-
}
39-
40-
private String resolveModelFromArgument(Object argument) {
41-
if (argument instanceof EmbeddingModel) {
42-
return resolveModelFromEmbeddingModel(argument);
43-
}
44-
return null;
45-
}
46-
47-
private String resolveModelFromEmbeddingModel(Object embeddingModel) {
48-
if (embeddingModel == null) {
49-
return null;
50-
}
51-
String model = resolveModelFromOptionsMethod(embeddingModel);
52-
if (StringUtils.hasText(model)) {
53-
return model;
30+
objInst.setSkyWalkingDynamicField(new VectorStoreEnhanceContext(resolveContextFromArgument(allArguments[0])));
5431
}
55-
model = resolveModelFromOptionsField(embeddingModel, "options");
56-
if (StringUtils.hasText(model)) {
57-
return model;
58-
}
59-
return resolveModelFromOptionsField(embeddingModel, "defaultOptions");
6032
}
6133

62-
private String resolveModelFromOptionsMethod(Object embeddingModel) {
63-
try {
64-
Method method = embeddingModel.getClass().getMethod("getOptions");
65-
return resolveModelFromOptions(method.invoke(embeddingModel));
66-
} catch (Throwable ignored) {
67-
return null;
68-
}
69-
}
70-
71-
private String resolveModelFromOptionsField(Object embeddingModel, String fieldName) {
72-
Class<?> type = embeddingModel.getClass();
73-
while (type != null) {
74-
try {
75-
Field field = type.getDeclaredField(fieldName);
76-
field.setAccessible(true);
77-
return resolveModelFromOptions(field.get(embeddingModel));
78-
} catch (NoSuchFieldException e) {
79-
type = type.getSuperclass();
80-
} catch (Throwable ignored) {
81-
return null;
82-
}
34+
private EmbeddingModelEnhanceContext resolveContextFromArgument(Object argument) {
35+
if (argument instanceof EmbeddingModel && argument instanceof EnhancedInstance) {
36+
return getOrCreateContext((EnhancedInstance) argument);
8337
}
8438
return null;
8539
}
8640

87-
private String resolveModelFromOptions(Object options) {
88-
if (options instanceof EmbeddingOptions) {
89-
return ((EmbeddingOptions) options).getModel();
41+
private EmbeddingModelEnhanceContext getOrCreateContext(EnhancedInstance embeddingModel) {
42+
Object context = embeddingModel.getSkyWalkingDynamicField();
43+
if (context instanceof EmbeddingModelEnhanceContext) {
44+
return (EmbeddingModelEnhanceContext) context;
9045
}
91-
return null;
46+
EmbeddingModelEnhanceContext embeddingModelEnhanceContext = new EmbeddingModelEnhanceContext();
47+
embeddingModel.setSkyWalkingDynamicField(embeddingModelEnhanceContext);
48+
return embeddingModelEnhanceContext;
9249
}
9350
}

‎apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/EmbeddingModelEnhanceContext.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*
17+
*/
18+
19+
package org.apache.skywalking.apm.plugin.spring.ai.v1;
20+
21+
import org.springframework.util.StringUtils;
22+
23+
public class EmbeddingModelEnhanceContext {
24+
25+
private volatile String embeddingModelName;
26+
27+
public String getEmbeddingModelName() {
28+
return embeddingModelName;
29+
}
30+
31+
public void setEmbeddingModelNameIfAbsent(String embeddingModelName) {
32+
if (!StringUtils.hasText(this.embeddingModelName) && StringUtils.hasText(embeddingModelName)) {
33+
this.embeddingModelName = embeddingModelName;
34+
}
35+
}
36+
}

‎apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/EmbeddingModelInterceptor.java‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*
17+
*/
18+
19+
package org.apache.skywalking.apm.plugin.spring.ai.v1;
20+
21+
import java.lang.reflect.Method;
22+
23+
import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
24+
import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceMethodsAroundInterceptor;
25+
import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult;
26+
import org.springframework.ai.embedding.EmbeddingOptions;
27+
import org.springframework.util.StringUtils;
28+
29+
public class EmbeddingModelInterceptor implements InstanceMethodsAroundInterceptor {
30+
31+
@Override
32+
public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
33+
MethodInterceptResult result) {
34+
if (allArguments == null || allArguments.length < 2 || !(allArguments[1] instanceof EmbeddingOptions)) {
35+
return;
36+
}
37+
String model = ((EmbeddingOptions) allArguments[1]).getModel();
38+
if (!StringUtils.hasText(model)) {
39+
return;
40+
}
41+
EmbeddingModelEnhanceContext context = getOrCreateContext(objInst);
42+
context.setEmbeddingModelNameIfAbsent(model);
43+
}
44+
45+
@Override
46+
public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
47+
Object ret) {
48+
return ret;
49+
}
50+
51+
@Override
52+
public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
53+
Class<?>[] argumentsTypes, Throwable t) {
54+
}
55+
56+
private EmbeddingModelEnhanceContext getOrCreateContext(EnhancedInstance objInst) {
57+
Object context = objInst.getSkyWalkingDynamicField();
58+
if (context instanceof EmbeddingModelEnhanceContext) {
59+
return (EmbeddingModelEnhanceContext) context;
60+
}
61+
EmbeddingModelEnhanceContext embeddingModelEnhanceContext = new EmbeddingModelEnhanceContext();
62+
objInst.setSkyWalkingDynamicField(embeddingModelEnhanceContext);
63+
return embeddingModelEnhanceContext;
64+
}
65+
}

‎apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/VectorStoreEnhanceContext.java‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,13 +20,16 @@
2020

2121
public class VectorStoreEnhanceContext {
2222

23-
private final String embeddingModelName;
23+
private final EmbeddingModelEnhanceContext embeddingModelEnhanceContext;
2424

25-
public VectorStoreEnhanceContext(String embeddingModelName) {
26-
this.embeddingModelName = embeddingModelName;
25+
public VectorStoreEnhanceContext(EmbeddingModelEnhanceContext embeddingModelEnhanceContext) {
26+
this.embeddingModelEnhanceContext = embeddingModelEnhanceContext;
2727
}
2828

2929
public String getEmbeddingModelName() {
30-
return embeddingModelName;
30+
if (embeddingModelEnhanceContext == null) {
31+
return null;
32+
}
33+
return embeddingModelEnhanceContext.getEmbeddingModelName();
3134
}
3235
}

‎apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/define/EmbeddingModelInstrumentation.java‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*
17+
*/
18+
19+
package org.apache.skywalking.apm.plugin.spring.ai.v1.define;
20+
21+
import net.bytebuddy.description.method.MethodDescription;
22+
import net.bytebuddy.matcher.ElementMatcher;
23+
import org.apache.skywalking.apm.agent.core.plugin.interceptor.ConstructorInterceptPoint;
24+
import org.apache.skywalking.apm.agent.core.plugin.interceptor.InstanceMethodsInterceptPoint;
25+
import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.ClassInstanceMethodsEnhancePluginDefine;
26+
import org.apache.skywalking.apm.agent.core.plugin.match.ClassMatch;
27+
import org.apache.skywalking.apm.agent.core.plugin.match.HierarchyMatch;
28+
29+
import static net.bytebuddy.matcher.ElementMatchers.named;
30+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
31+
import static org.apache.skywalking.apm.agent.core.plugin.bytebuddy.ArgumentTypeNameMatch.takesArgumentWithType;
32+
33+
public class EmbeddingModelInstrumentation extends ClassInstanceMethodsEnhancePluginDefine {
34+
35+
private static final String ENHANCE_CLASS = "org.springframework.ai.embedding.EmbeddingModel";
36+
37+
private static final String INTERCEPT_METHOD = "embed";
38+
39+
private static final String INTERCEPTOR_CLASS =
40+
"org.apache.skywalking.apm.plugin.spring.ai.v1.EmbeddingModelInterceptor";
41+
42+
@Override
43+
protected ClassMatch enhanceClass() {
44+
return HierarchyMatch.byHierarchyMatch(ENHANCE_CLASS);
45+
}
46+
47+
@Override
48+
public ConstructorInterceptPoint[] getConstructorsInterceptPoints() {
49+
return new ConstructorInterceptPoint[0];
50+
}
51+
52+
@Override
53+
public InstanceMethodsInterceptPoint[] getInstanceMethodsInterceptPoints() {
54+
return new InstanceMethodsInterceptPoint[]{
55+
new InstanceMethodsInterceptPoint() {
56+
@Override
57+
public ElementMatcher<MethodDescription> getMethodsMatcher() {
58+
return named(INTERCEPT_METHOD)
59+
.and(takesArguments(3))
60+
.and(takesArgumentWithType(0, "java.util.List"))
61+
.and(takesArgumentWithType(1, "org.springframework.ai.embedding.EmbeddingOptions"))
62+
.and(takesArgumentWithType(2, "org.springframework.ai.embedding.BatchingStrategy"));
63+
}
64+
65+
@Override
66+
public String getMethodsInterceptor() {
67+
return INTERCEPTOR_CLASS;
68+
}
69+
70+
@Override
71+
public boolean isOverrideArgs() {
72+
return false;
73+
}
74+
}
75+
};
76+
}
77+
}

0 commit comments

Comments
 (0)