Skip to content

Commit a9843a7

Browse files
peachisaipeachisai
committed
Support tracing RAG retrieval in Spring AI 1.x plugin
1 parent e0e8b3c commit a9843a7

17 files changed

Lines changed: 719 additions & 3 deletions

File tree

‎apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/tag/Tags.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -250,6 +250,21 @@ public static final class HTTP {
250250
*/
251251
public static final StringTag GEN_AI_OUTPUT_MESSAGES = new StringTag(42, "gen_ai.output.messages");
252252

253+
/**
254+
* GEN_AI_DATA_SOURCE_ID represents the data source identifier.
255+
*/
256+
public static final StringTag GEN_AI_DATA_SOURCE_ID = new StringTag(43, "gen_ai.data_source.id");
257+
258+
/**
259+
* GEN_AI_RETRIEVAL_DOCUMENTS represents the documents retrieved.
260+
*/
261+
public static final StringTag GEN_AI_RETRIEVAL_DOCUMENTS = new StringTag(44, "gen_ai.retrieval.documents");
262+
263+
/**
264+
* GEN_AI_RETRIEVAL_QUERY_TEXT represents the query text used for retrieval.
265+
*/
266+
public static final StringTag GEN_AI_RETRIEVAL_QUERY_TEXT = new StringTag(45, "gen_ai.retrieval.query.text");
267+
253268
/**
254269
* Creates a {@code StringTag} with the given key and cache it, if it's created before, simply return it without
255270
* creating a new one.

‎apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/pom.xml‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,13 @@
4646
<scope>provided</scope>
4747
</dependency>
4848

49+
<dependency>
50+
<groupId>org.springframework.ai</groupId>
51+
<artifactId>spring-ai-vector-store</artifactId>
52+
<version>1.1.0</version>
53+
<scope>provided</scope>
54+
</dependency>
55+
4956
</dependencies>
5057

5158
<build>

‎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: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
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.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
22+
import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceConstructorInterceptor;
23+
import org.springframework.ai.embedding.EmbeddingOptions;
24+
import org.springframework.util.StringUtils;
25+
26+
import java.lang.reflect.Field;
27+
import java.lang.reflect.Method;
28+
29+
public class AbstractObservationVectorStoreConstructorInterceptor implements InstanceConstructorInterceptor {
30+
31+
@Override
32+
public void onConstruct(EnhancedInstance objInst, Object[] allArguments) {
33+
if (allArguments != null && allArguments.length > 0) {
34+
String embeddingModelName = resolveModelFromEmbeddingModel(allArguments[0]);
35+
objInst.setSkyWalkingDynamicField(new VectorStoreEnhanceContext(embeddingModelName));
36+
}
37+
}
38+
39+
private String resolveModelFromEmbeddingModel(Object embeddingModel) {
40+
if (embeddingModel == null) {
41+
return null;
42+
}
43+
String model = resolveModelFromOptionsMethod(embeddingModel);
44+
if (StringUtils.hasText(model)) {
45+
return model;
46+
}
47+
model = resolveModelFromOptionsField(embeddingModel, "options");
48+
if (StringUtils.hasText(model)) {
49+
return model;
50+
}
51+
return resolveModelFromOptionsField(embeddingModel, "defaultOptions");
52+
}
53+
54+
private String resolveModelFromOptionsMethod(Object embeddingModel) {
55+
try {
56+
Method method = embeddingModel.getClass().getMethod("getOptions");
57+
return resolveModelFromOptions(method.invoke(embeddingModel));
58+
} catch (Throwable ignored) {
59+
return null;
60+
}
61+
}
62+
63+
private String resolveModelFromOptionsField(Object embeddingModel, String fieldName) {
64+
Class<?> type = embeddingModel.getClass();
65+
while (type != null) {
66+
try {
67+
Field field = type.getDeclaredField(fieldName);
68+
field.setAccessible(true);
69+
return resolveModelFromOptions(field.get(embeddingModel));
70+
} catch (NoSuchFieldException e) {
71+
type = type.getSuperclass();
72+
} catch (Throwable ignored) {
73+
return null;
74+
}
75+
}
76+
return null;
77+
}
78+
79+
private String resolveModelFromOptions(Object options) {
80+
if (options instanceof EmbeddingOptions) {
81+
return ((EmbeddingOptions) options).getModel();
82+
}
83+
return null;
84+
}
85+
}

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

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
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.apache.skywalking.apm.agent.core.context.ContextManager;
22+
import org.apache.skywalking.apm.agent.core.context.tag.Tags;
23+
import org.apache.skywalking.apm.agent.core.context.trace.AbstractSpan;
24+
import org.apache.skywalking.apm.agent.core.context.trace.SpanLayer;
25+
import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
26+
import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceMethodsAroundInterceptor;
27+
import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult;
28+
import org.apache.skywalking.apm.agent.core.util.GsonUtil;
29+
import org.apache.skywalking.apm.network.trace.component.ComponentsDefine;
30+
import org.apache.skywalking.apm.plugin.spring.ai.v1.contant.Constants;
31+
import org.springframework.ai.document.Document;
32+
import org.springframework.ai.vectorstore.SearchRequest;
33+
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
34+
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
35+
import org.springframework.util.StringUtils;
36+
37+
import java.lang.reflect.Method;
38+
import java.util.ArrayList;
39+
import java.util.LinkedHashMap;
40+
import java.util.List;
41+
import java.util.Map;
42+
43+
public class AbstractObservationVectorStoreInterceptor implements InstanceMethodsAroundInterceptor {
44+
45+
@Override
46+
public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
47+
MethodInterceptResult result) throws Throwable {
48+
SearchRequest request = (SearchRequest) allArguments[0];
49+
VectorStoreObservationContext context = createObservationContext(objInst, request);
50+
String dataSourceId = resolveDataSourceId(context, objInst);
51+
52+
AbstractSpan span = ContextManager.createExitSpan(Constants.RETRIEVAL + "/" + dataSourceId, dataSourceId);
53+
SpanLayer.asGenAI(span);
54+
span.setComponent(ComponentsDefine.SPRING_AI);
55+
Tags.GEN_AI_OPERATION_NAME.set(span, Constants.RETRIEVAL);
56+
Tags.GEN_AI_DATA_SOURCE_ID.set(span, dataSourceId);
57+
String model = resolveEmbeddingModelName(objInst);
58+
if (StringUtils.hasText(model)) {
59+
Tags.GEN_AI_REQUEST_MODEL.set(span, model);
60+
}
61+
62+
if (request != null) {
63+
Tags.GEN_AI_TOP_K.set(span, String.valueOf(request.getTopK()));
64+
if (StringUtils.hasText(request.getQuery())) {
65+
Tags.GEN_AI_RETRIEVAL_QUERY_TEXT.set(span, request.getQuery());
66+
}
67+
}
68+
}
69+
70+
@Override
71+
public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
72+
Object ret) throws Throwable {
73+
if (!ContextManager.isActive()) {
74+
return ret;
75+
}
76+
try {
77+
if (ret instanceof List<?>) {
78+
Tags.GEN_AI_RETRIEVAL_DOCUMENTS.set(ContextManager.activeSpan(), toDocumentsJson((List<?>) ret));
79+
}
80+
} finally {
81+
ContextManager.stopSpan();
82+
}
83+
return ret;
84+
}
85+
86+
@Override
87+
public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
88+
Class<?>[] argumentsTypes, Throwable t) {
89+
if (ContextManager.isActive()) {
90+
ContextManager.activeSpan().log(t);
91+
}
92+
}
93+
94+
private VectorStoreObservationContext createObservationContext(EnhancedInstance objInst, SearchRequest request) {
95+
VectorStoreObservationContext.Builder builder = ((AbstractObservationVectorStore) objInst)
96+
.createObservationContextBuilder(VectorStoreObservationContext.Operation.QUERY.value());
97+
if (request != null) {
98+
builder.queryRequest(request);
99+
}
100+
return builder.build();
101+
}
102+
103+
private String resolveEmbeddingModelName(EnhancedInstance objInst) {
104+
Object context = objInst.getSkyWalkingDynamicField();
105+
if (context instanceof VectorStoreEnhanceContext) {
106+
return ((VectorStoreEnhanceContext) context).getEmbeddingModelName();
107+
}
108+
return null;
109+
}
110+
111+
private String resolveDataSourceId(VectorStoreObservationContext context, EnhancedInstance objInst) {
112+
StringBuilder dataSourceId = new StringBuilder();
113+
appendDataSourcePart(dataSourceId, context.getDatabaseSystem());
114+
appendDataSourcePart(dataSourceId, context.getNamespace());
115+
appendDataSourcePart(dataSourceId, context.getCollectionName());
116+
if (dataSourceId.length() > 0) {
117+
return dataSourceId.toString();
118+
}
119+
return objInst.getClass().getSimpleName();
120+
}
121+
122+
private void appendDataSourcePart(StringBuilder dataSourceId, String value) {
123+
if (!StringUtils.hasText(value)) {
124+
return;
125+
}
126+
if (dataSourceId.length() > 0) {
127+
dataSourceId.append('/');
128+
}
129+
dataSourceId.append(value);
130+
}
131+
132+
private String toDocumentsJson(List<?> documents) {
133+
List<Map<String, Object>> retrievalDocuments = new ArrayList<>(documents.size());
134+
for (Object item : documents) {
135+
if (!(item instanceof Document)) {
136+
continue;
137+
}
138+
Document document = (Document) item;
139+
Map<String, Object> documentMap = new LinkedHashMap<>();
140+
documentMap.put("id", document.getId());
141+
if (document.getScore() != null) {
142+
documentMap.put("score", document.getScore());
143+
}
144+
retrievalDocuments.add(documentMap);
145+
}
146+
return GsonUtil.toJson(retrievalDocuments);
147+
}
148+
}

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceMethodsAroundInterceptor;
2727
import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult;
2828
import org.apache.skywalking.apm.plugin.spring.ai.v1.common.ChatModelMetadataResolver;
29+
import org.apache.skywalking.apm.plugin.spring.ai.v1.common.ErrorTypeResolver;
2930
import org.apache.skywalking.apm.plugin.spring.ai.v1.config.SpringAiPluginConfig;
3031
import org.apache.skywalking.apm.plugin.spring.ai.v1.contant.Constants;
3132
import org.apache.skywalking.apm.plugin.spring.ai.v1.messages.InputMessages;
@@ -129,7 +130,9 @@ public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allA
129130
@Override
130131
public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, Throwable t) {
131132
if (ContextManager.isActive()) {
132-
ContextManager.activeSpan().log(t);
133+
AbstractSpan span = ContextManager.activeSpan();
134+
span.log(t);
135+
ErrorTypeResolver.setErrorType(span, t);
133136
}
134137
}
135138

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

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceMethodsAroundInterceptor;
2828
import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult;
2929
import org.apache.skywalking.apm.plugin.spring.ai.v1.common.ChatModelMetadataResolver;
30+
import org.apache.skywalking.apm.plugin.spring.ai.v1.common.ErrorTypeResolver;
3031
import org.apache.skywalking.apm.plugin.spring.ai.v1.config.SpringAiPluginConfig;
3132
import org.apache.skywalking.apm.plugin.spring.ai.v1.contant.Constants;
3233
import org.apache.skywalking.apm.plugin.spring.ai.v1.messages.InputMessages;
@@ -94,11 +95,16 @@ public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allA
9495

9596
return flux
9697
.doOnNext(response -> onStreamNext(span, response, state))
97-
.doOnError(span::log)
98+
.doOnError(t -> recordError(span, t))
9899
.doFinally(signalType -> onStreamFinally(span, allArguments, state))
99100
.contextWrite(c -> c.put(Constants.SKYWALKING_CONTEXT_SNAPSHOT, snapshot));
100101
}
101102

103+
private void recordError(AbstractSpan span, Throwable t) {
104+
span.log(t);
105+
ErrorTypeResolver.setErrorType(span, t);
106+
}
107+
102108
private void onStreamNext(AbstractSpan span, ChatResponse response, StreamState state) {
103109
state.lastResponseRef.set(response);
104110

@@ -248,7 +254,9 @@ private Long readAndClearStartTime() {
248254
@Override
249255
public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes, Throwable t) {
250256
if (ContextManager.isActive()) {
251-
ContextManager.activeSpan().log(t);
257+
AbstractSpan span = ContextManager.activeSpan();
258+
span.log(t);
259+
ErrorTypeResolver.setErrorType(span, t);
252260
}
253261
}
254262

‎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: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
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+
public class VectorStoreEnhanceContext {
22+
23+
private final String embeddingModelName;
24+
25+
public VectorStoreEnhanceContext(String embeddingModelName) {
26+
this.embeddingModelName = embeddingModelName;
27+
}
28+
29+
public String getEmbeddingModelName() {
30+
return embeddingModelName;
31+
}
32+
}

0 commit comments

Comments
 (0)