diff --git a/CHANGES.md b/CHANGES.md index 3a9514fdc6..0bf6869120 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -25,6 +25,7 @@ Release Notes. * Only publish `apm-application-toolkit` modules to Maven Central. Agent and plugins are distributed via download package and Docker images. * Add unified release script (`tools/releasing/release.sh`) with two-step flow: `prepare-vote` and `vote-passed`. * 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. +- Add tracing support for vector-store retrieval operations. All issues and pull requests are [here](https://github.com/apache/skywalking/milestone/249?closed=1) diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/tag/Tags.java b/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/tag/Tags.java index 3d0b9f37cb..8acb2b9879 100644 --- a/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/tag/Tags.java +++ b/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/tag/Tags.java @@ -250,6 +250,21 @@ public static final class HTTP { */ public static final StringTag GEN_AI_OUTPUT_MESSAGES = new StringTag(42, "gen_ai.output.messages"); + /** + * GEN_AI_DATA_SOURCE_ID represents the data source identifier. + */ + public static final StringTag GEN_AI_DATA_SOURCE_ID = new StringTag(43, "gen_ai.data_source.id"); + + /** + * GEN_AI_RETRIEVAL_DOCUMENTS represents the documents retrieved. + */ + public static final StringTag GEN_AI_RETRIEVAL_DOCUMENTS = new StringTag(44, "gen_ai.retrieval.documents"); + + /** + * GEN_AI_RETRIEVAL_QUERY_TEXT represents the query text used for retrieval. + */ + public static final StringTag GEN_AI_RETRIEVAL_QUERY_TEXT = new StringTag(45, "gen_ai.retrieval.query.text"); + /** * Creates a {@code StringTag} with the given key and cache it, if it's created before, simply return it without * creating a new one. diff --git a/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/pom.xml b/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/pom.xml index 295cf9d8ed..583ce9c90c 100644 --- a/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/pom.xml +++ b/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/pom.xml @@ -46,6 +46,13 @@ provided + + org.springframework.ai + spring-ai-vector-store + 1.1.0 + provided + + diff --git a/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 b/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 new file mode 100644 index 0000000000..faefa999bc --- /dev/null +++ b/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 @@ -0,0 +1,48 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package org.apache.skywalking.apm.plugin.spring.ai.v1; + +import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceConstructorInterceptor; +import org.apache.skywalking.apm.plugin.spring.ai.v1.common.EmbeddingModelEnhanceContext; + +public class AbstractObservationVectorStoreConstructorInterceptor implements InstanceConstructorInterceptor { + + @Override + public void onConstruct(EnhancedInstance objInst, Object[] allArguments) { + objInst.setSkyWalkingDynamicField(new VectorStoreEnhanceContext(resolveContextFromArgument(allArguments[0]))); + } + + private EmbeddingModelEnhanceContext resolveContextFromArgument(Object argument) { + if (argument instanceof EnhancedInstance) { + return getOrCreateContext((EnhancedInstance) argument); + } + return null; + } + + private EmbeddingModelEnhanceContext getOrCreateContext(EnhancedInstance embeddingModel) { + Object context = embeddingModel.getSkyWalkingDynamicField(); + if (context instanceof EmbeddingModelEnhanceContext) { + return (EmbeddingModelEnhanceContext) context; + } + EmbeddingModelEnhanceContext embeddingModelEnhanceContext = new EmbeddingModelEnhanceContext(); + embeddingModel.setSkyWalkingDynamicField(embeddingModelEnhanceContext); + return embeddingModelEnhanceContext; + } +} diff --git a/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 b/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 new file mode 100644 index 0000000000..4b278b6c87 --- /dev/null +++ b/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 @@ -0,0 +1,171 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package org.apache.skywalking.apm.plugin.spring.ai.v1; + +import org.apache.skywalking.apm.agent.core.context.ContextManager; +import org.apache.skywalking.apm.agent.core.context.tag.Tags; +import org.apache.skywalking.apm.agent.core.context.trace.AbstractSpan; +import org.apache.skywalking.apm.agent.core.context.trace.SpanLayer; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceMethodsAroundInterceptor; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult; +import org.apache.skywalking.apm.agent.core.util.GsonUtil; +import org.apache.skywalking.apm.network.trace.component.ComponentsDefine; +import org.apache.skywalking.apm.plugin.spring.ai.v1.common.ErrorTypeResolver; +import org.apache.skywalking.apm.plugin.spring.ai.v1.config.SpringAiPluginConfig; +import org.apache.skywalking.apm.plugin.spring.ai.v1.contant.Constants; +import org.springframework.ai.document.Document; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore; +import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext; +import org.springframework.util.StringUtils; + +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +public class AbstractObservationVectorStoreInterceptor implements InstanceMethodsAroundInterceptor { + + @Override + public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class[] argumentsTypes, + MethodInterceptResult result) throws Throwable { + SearchRequest request = (SearchRequest) allArguments[0]; + String dataSourceId = objInst.getClass().getSimpleName(); + + try { + VectorStoreObservationContext context = + createObservationContext(objInst, request); + + String resolved = + resolveDataSourceId(context, objInst); + + if (StringUtils.hasText(resolved)) { + dataSourceId = resolved; + } + } catch (Throwable ignored) { + + } + + AbstractSpan span = ContextManager.createExitSpan(Constants.RETRIEVAL + "/" + dataSourceId, dataSourceId); + + SpanLayer.asGenAI(span); + span.setComponent(ComponentsDefine.SPRING_AI); + Tags.GEN_AI_OPERATION_NAME.set(span, Constants.RETRIEVAL); + Tags.GEN_AI_DATA_SOURCE_ID.set(span, dataSourceId); + String model = resolveEmbeddingModelName(objInst); + if (StringUtils.hasText(model)) { + Tags.GEN_AI_REQUEST_MODEL.set(span, model); + } + + if (request != null) { + Tags.GEN_AI_TOP_K.set(span, String.valueOf(request.getTopK())); + String query = request.getQuery(); + if (StringUtils.hasText(query) && SpringAiPluginConfig.Plugin.SpringAi.COLLECT_RETRIEVAL_QUERY) { + int limit = SpringAiPluginConfig.Plugin.SpringAi.RETRIEVAL_QUERY_LENGTH_LIMIT; + if (limit > 0 && query.length() > limit) { + query = query.substring(0, limit); + } + Tags.GEN_AI_RETRIEVAL_QUERY_TEXT.set(span, query); + } + } + } + + @Override + public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class[] argumentsTypes, + Object ret) throws Throwable { + if (!ContextManager.isActive()) { + return ret; + } + try { + if (ret instanceof List && SpringAiPluginConfig.Plugin.SpringAi.COLLECT_RETRIEVAL_DOCUMENTS) { + Tags.GEN_AI_RETRIEVAL_DOCUMENTS.set(ContextManager.activeSpan(), toDocumentsJson((List) ret)); + } + } finally { + ContextManager.stopSpan(); + } + return ret; + } + + @Override + public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments, + Class[] argumentsTypes, Throwable t) { + if (ContextManager.isActive()) { + AbstractSpan span = ContextManager.activeSpan(); + span.log(t); + ErrorTypeResolver.setErrorType(span, t); + } + } + + private VectorStoreObservationContext createObservationContext(EnhancedInstance objInst, SearchRequest request) { + VectorStoreObservationContext.Builder builder = ((AbstractObservationVectorStore) objInst) + .createObservationContextBuilder(VectorStoreObservationContext.Operation.QUERY.value()); + if (request != null) { + builder.queryRequest(request); + } + return builder.build(); + } + + private String resolveEmbeddingModelName(EnhancedInstance objInst) { + Object context = objInst.getSkyWalkingDynamicField(); + if (context instanceof VectorStoreEnhanceContext) { + return ((VectorStoreEnhanceContext) context).getEmbeddingModelName(); + } + return null; + } + + private String resolveDataSourceId(VectorStoreObservationContext context, EnhancedInstance objInst) { + StringBuilder dataSourceId = new StringBuilder(); + appendDataSourcePart(dataSourceId, context.getDatabaseSystem()); + appendDataSourcePart(dataSourceId, context.getNamespace()); + appendDataSourcePart(dataSourceId, context.getCollectionName()); + if (dataSourceId.length() > 0) { + return dataSourceId.toString(); + } + return objInst.getClass().getSimpleName(); + } + + private void appendDataSourcePart(StringBuilder dataSourceId, String value) { + if (!StringUtils.hasText(value)) { + return; + } + if (dataSourceId.length() > 0) { + dataSourceId.append('/'); + } + dataSourceId.append(value); + } + + private String toDocumentsJson(List documents) { + List> retrievalDocuments = new ArrayList<>(documents.size()); + for (Object item : documents) { + if (!(item instanceof Document)) { + continue; + } + Document document = (Document) item; + Map documentMap = new LinkedHashMap<>(); + documentMap.put("id", document.getId()); + if (document.getScore() != null) { + documentMap.put("score", document.getScore()); + } + retrievalDocuments.add(documentMap); + } + return GsonUtil.toJson(retrievalDocuments); + } +} diff --git a/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 b/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 index 302db3c407..73f62629c1 100644 --- a/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 +++ b/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 @@ -26,6 +26,7 @@ import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceMethodsAroundInterceptor; import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult; import org.apache.skywalking.apm.plugin.spring.ai.v1.common.ChatModelMetadataResolver; +import org.apache.skywalking.apm.plugin.spring.ai.v1.common.ErrorTypeResolver; import org.apache.skywalking.apm.plugin.spring.ai.v1.config.SpringAiPluginConfig; import org.apache.skywalking.apm.plugin.spring.ai.v1.contant.Constants; import org.apache.skywalking.apm.plugin.spring.ai.v1.messages.InputMessages; @@ -129,7 +130,9 @@ public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allA @Override public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments, Class[] argumentsTypes, Throwable t) { if (ContextManager.isActive()) { - ContextManager.activeSpan().log(t); + AbstractSpan span = ContextManager.activeSpan(); + span.log(t); + ErrorTypeResolver.setErrorType(span, t); } } diff --git a/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 b/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 index 1674e82011..7d1b380573 100644 --- a/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 +++ b/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 @@ -27,6 +27,7 @@ import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceMethodsAroundInterceptor; import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult; import org.apache.skywalking.apm.plugin.spring.ai.v1.common.ChatModelMetadataResolver; +import org.apache.skywalking.apm.plugin.spring.ai.v1.common.ErrorTypeResolver; import org.apache.skywalking.apm.plugin.spring.ai.v1.config.SpringAiPluginConfig; import org.apache.skywalking.apm.plugin.spring.ai.v1.contant.Constants; import org.apache.skywalking.apm.plugin.spring.ai.v1.messages.InputMessages; @@ -94,11 +95,16 @@ public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allA return flux .doOnNext(response -> onStreamNext(span, response, state)) - .doOnError(span::log) + .doOnError(t -> recordError(span, t)) .doFinally(signalType -> onStreamFinally(span, allArguments, state)) .contextWrite(c -> c.put(Constants.SKYWALKING_CONTEXT_SNAPSHOT, snapshot)); } + private void recordError(AbstractSpan span, Throwable t) { + span.log(t); + ErrorTypeResolver.setErrorType(span, t); + } + private void onStreamNext(AbstractSpan span, ChatResponse response, StreamState state) { state.lastResponseRef.set(response); @@ -248,7 +254,9 @@ private Long readAndClearStartTime() { @Override public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments, Class[] argumentsTypes, Throwable t) { if (ContextManager.isActive()) { - ContextManager.activeSpan().log(t); + AbstractSpan span = ContextManager.activeSpan(); + span.log(t); + ErrorTypeResolver.setErrorType(span, t); } } diff --git a/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 b/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 new file mode 100644 index 0000000000..8e2b57fae9 --- /dev/null +++ b/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 @@ -0,0 +1,70 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package org.apache.skywalking.apm.plugin.spring.ai.v1; + +import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceMethodsAroundInterceptor; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult; +import org.apache.skywalking.apm.plugin.spring.ai.v1.common.EmbeddingModelEnhanceContext; +import org.springframework.ai.embedding.EmbeddingResponse; +import org.springframework.ai.embedding.EmbeddingResponseMetadata; +import org.springframework.util.StringUtils; + +import java.lang.reflect.Method; + +public class EmbeddingModelInterceptor implements InstanceMethodsAroundInterceptor { + + @Override + public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class[] argumentsTypes, MethodInterceptResult result) { + + } + + @Override + public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class[] argumentsTypes, Object ret) { + if (!(ret instanceof EmbeddingResponse)) { + return ret; + } + + EmbeddingResponseMetadata metadata = ((EmbeddingResponse) ret).getMetadata(); + if (metadata == null) { + return ret; + } + String model = metadata.getModel(); + if (!StringUtils.hasText(model)) { + return ret; + } + EmbeddingModelEnhanceContext context = getOrCreateContext(objInst); + context.setEmbeddingModelNameIfAbsent(model); + return ret; + } + + @Override + public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments, Class[] argumentsTypes, Throwable t) { + } + + private EmbeddingModelEnhanceContext getOrCreateContext(EnhancedInstance objInst) { + Object context = objInst.getSkyWalkingDynamicField(); + if (context instanceof EmbeddingModelEnhanceContext) { + return (EmbeddingModelEnhanceContext) context; + } + EmbeddingModelEnhanceContext embeddingModelEnhanceContext = new EmbeddingModelEnhanceContext(); + objInst.setSkyWalkingDynamicField(embeddingModelEnhanceContext); + return embeddingModelEnhanceContext; + } +} diff --git a/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 b/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 new file mode 100644 index 0000000000..76916a59ea --- /dev/null +++ b/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 @@ -0,0 +1,37 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package org.apache.skywalking.apm.plugin.spring.ai.v1; + +import org.apache.skywalking.apm.plugin.spring.ai.v1.common.EmbeddingModelEnhanceContext; + +public class VectorStoreEnhanceContext { + + private final EmbeddingModelEnhanceContext embeddingModelEnhanceContext; + + public VectorStoreEnhanceContext(EmbeddingModelEnhanceContext embeddingModelEnhanceContext) { + this.embeddingModelEnhanceContext = embeddingModelEnhanceContext; + } + + public String getEmbeddingModelName() { + if (embeddingModelEnhanceContext == null) { + return null; + } + return embeddingModelEnhanceContext.getEmbeddingModelName(); + } +} diff --git a/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/common/EmbeddingModelEnhanceContext.java b/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/common/EmbeddingModelEnhanceContext.java new file mode 100644 index 0000000000..199bd97cd4 --- /dev/null +++ b/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/common/EmbeddingModelEnhanceContext.java @@ -0,0 +1,36 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package org.apache.skywalking.apm.plugin.spring.ai.v1.common; + +import org.springframework.util.StringUtils; + +public class EmbeddingModelEnhanceContext { + + private volatile String embeddingModelName; + + public String getEmbeddingModelName() { + return embeddingModelName; + } + + public void setEmbeddingModelNameIfAbsent(String embeddingModelName) { + if (!StringUtils.hasText(this.embeddingModelName) && StringUtils.hasText(embeddingModelName)) { + this.embeddingModelName = embeddingModelName; + } + } +} diff --git a/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/common/ErrorTypeResolver.java b/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/common/ErrorTypeResolver.java new file mode 100644 index 0000000000..1d3172cf3d --- /dev/null +++ b/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/common/ErrorTypeResolver.java @@ -0,0 +1,80 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package org.apache.skywalking.apm.plugin.spring.ai.v1.common; + +import org.apache.skywalking.apm.agent.core.context.tag.AbstractTag; +import org.apache.skywalking.apm.agent.core.context.tag.Tags; +import org.apache.skywalking.apm.agent.core.context.trace.AbstractSpan; + +import javax.net.ssl.SSLHandshakeException; +import java.net.SocketTimeoutException; +import java.security.cert.CertPathValidatorException; +import java.security.cert.CertificateException; +import java.util.concurrent.TimeoutException; + +public final class ErrorTypeResolver { + + private static final AbstractTag ERROR_TYPE = Tags.ofKey("error.type"); + private static final String TIMEOUT = "timeout"; + private static final String SERVER_CERTIFICATE_INVALID = "server_certificate_invalid"; + + private ErrorTypeResolver() { + } + + public static void setErrorType(AbstractSpan span, Throwable throwable) { + span.tag(ERROR_TYPE, resolve(throwable)); + } + + private static String resolve(Throwable throwable) { + if (matches(throwable, ErrorTypeResolver::isTimeout)) { + return TIMEOUT; + } + if (matches(throwable, ErrorTypeResolver::isCertificateInvalid)) { + return SERVER_CERTIFICATE_INVALID; + } + return throwable.getClass().getName(); + } + + private static boolean isTimeout(Throwable throwable) { + return throwable instanceof SocketTimeoutException + || throwable instanceof TimeoutException + || throwable.getClass().getName().contains("TimeoutException"); + } + + private static boolean isCertificateInvalid(Throwable throwable) { + return throwable instanceof SSLHandshakeException + || throwable instanceof CertificateException + || throwable instanceof CertPathValidatorException; + } + + private static boolean matches(Throwable throwable, Matcher matcher) { + Throwable current = throwable; + while (current != null) { + if (matcher.matches(current)) { + return true; + } + current = current.getCause(); + } + return false; + } + + private interface Matcher { + boolean matches(Throwable throwable); + } +} diff --git a/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/config/SpringAiPluginConfig.java b/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/config/SpringAiPluginConfig.java index d2d5eef6f6..6a348c622b 100644 --- a/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/config/SpringAiPluginConfig.java +++ b/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/config/SpringAiPluginConfig.java @@ -69,6 +69,23 @@ public static class SpringAi { * Whether to collect the execution result (output) of the tool/function call. */ public static boolean COLLECT_TOOL_OUTPUT = false; + + /** + * Whether to collect the query of the rag call. + */ + public static boolean COLLECT_RETRIEVAL_QUERY = false; + + /** + * The maximum characters of the collected rag query content. + * If the content exceeds this limit, it will be truncated. + * Use a negative value to represent no limit, but be aware this could cause OOM. + */ + public static int RETRIEVAL_QUERY_LENGTH_LIMIT = 1024; + + /** + * Whether to collect the documents of the rag call. + */ + public static boolean COLLECT_RETRIEVAL_DOCUMENTS = false; } } } diff --git a/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/contant/Constants.java b/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/contant/Constants.java index 7348a0c11f..688f4c321a 100644 --- a/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/contant/Constants.java +++ b/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/contant/Constants.java @@ -27,5 +27,7 @@ public class Constants { public static final String EXECUTE_TOOL = "execute_tool"; + public static final String RETRIEVAL = "retrieval"; + public static final String DEFAULT_COMPLETIONS_PATH = "/v1/chat/completions"; } diff --git a/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/define/AbstractObservationVectorStoreInstrumentation.java b/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/define/AbstractObservationVectorStoreInstrumentation.java new file mode 100644 index 0000000000..f379c0464c --- /dev/null +++ b/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/spring/ai/v1/define/AbstractObservationVectorStoreInstrumentation.java @@ -0,0 +1,92 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package org.apache.skywalking.apm.plugin.spring.ai.v1.define; + +import net.bytebuddy.description.method.MethodDescription; +import net.bytebuddy.matcher.ElementMatcher; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.ConstructorInterceptPoint; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.InstanceMethodsInterceptPoint; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.ClassInstanceMethodsEnhancePluginDefine; +import org.apache.skywalking.apm.agent.core.plugin.match.ClassMatch; +import org.apache.skywalking.apm.agent.core.plugin.match.HierarchyMatch; +import org.apache.skywalking.apm.agent.core.plugin.match.MultiClassNameMatch; +import org.apache.skywalking.apm.agent.core.plugin.match.logical.LogicalMatchOperation; + +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; +import static org.apache.skywalking.apm.agent.core.plugin.bytebuddy.ArgumentTypeNameMatch.takesArgumentWithType; + +public class AbstractObservationVectorStoreInstrumentation extends ClassInstanceMethodsEnhancePluginDefine { + + private static final String ENHANCE_CLASS = "org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore"; + + private static final String INTERCEPT_METHOD = "doSimilaritySearch"; + + private static final String INTERCEPTOR_CLASS = + "org.apache.skywalking.apm.plugin.spring.ai.v1.AbstractObservationVectorStoreInterceptor"; + + private static final String CONSTRUCTOR_INTERCEPTOR_CLASS = + "org.apache.skywalking.apm.plugin.spring.ai.v1.AbstractObservationVectorStoreConstructorInterceptor"; + + @Override + protected ClassMatch enhanceClass() { + return LogicalMatchOperation.or(HierarchyMatch.byHierarchyMatch(ENHANCE_CLASS), MultiClassNameMatch.byMultiClassMatch(ENHANCE_CLASS)); + } + + @Override + public ConstructorInterceptPoint[] getConstructorsInterceptPoints() { + return new ConstructorInterceptPoint[]{ + new ConstructorInterceptPoint() { + @Override + public ElementMatcher getConstructorMatcher() { + return takesArgumentWithType(0, "org.springframework.ai.embedding.EmbeddingModel"); + } + + @Override + public String getConstructorInterceptor() { + return CONSTRUCTOR_INTERCEPTOR_CLASS; + } + } + }; + } + + @Override + public InstanceMethodsInterceptPoint[] getInstanceMethodsInterceptPoints() { + return new InstanceMethodsInterceptPoint[]{ + new InstanceMethodsInterceptPoint() { + @Override + public ElementMatcher getMethodsMatcher() { + return named(INTERCEPT_METHOD) + .and(takesArguments(1)) + .and(takesArgumentWithType(0, "org.springframework.ai.vectorstore.SearchRequest")); + } + + @Override + public String getMethodsInterceptor() { + return INTERCEPTOR_CLASS; + } + + @Override + public boolean isOverrideArgs() { + return false; + } + } + }; + } +} diff --git a/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 b/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 new file mode 100644 index 0000000000..12a327ae29 --- /dev/null +++ b/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 @@ -0,0 +1,75 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +package org.apache.skywalking.apm.plugin.spring.ai.v1.define; + +import net.bytebuddy.description.method.MethodDescription; +import net.bytebuddy.matcher.ElementMatcher; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.ConstructorInterceptPoint; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.InstanceMethodsInterceptPoint; +import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.ClassInstanceMethodsEnhancePluginDefine; +import org.apache.skywalking.apm.agent.core.plugin.match.ClassMatch; +import org.apache.skywalking.apm.agent.core.plugin.match.HierarchyMatch; + +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; +import static org.apache.skywalking.apm.agent.core.plugin.bytebuddy.ArgumentTypeNameMatch.takesArgumentWithType; + +public class EmbeddingModelInstrumentation extends ClassInstanceMethodsEnhancePluginDefine { + + private static final String ENHANCE_CLASS = "org.springframework.ai.embedding.EmbeddingModel"; + + private static final String INTERCEPT_METHOD = "call"; + + private static final String INTERCEPTOR_CLASS = + "org.apache.skywalking.apm.plugin.spring.ai.v1.EmbeddingModelInterceptor"; + + @Override + protected ClassMatch enhanceClass() { + return HierarchyMatch.byHierarchyMatch(ENHANCE_CLASS); + } + + @Override + public ConstructorInterceptPoint[] getConstructorsInterceptPoints() { + return new ConstructorInterceptPoint[0]; + } + + @Override + public InstanceMethodsInterceptPoint[] getInstanceMethodsInterceptPoints() { + return new InstanceMethodsInterceptPoint[]{ + new InstanceMethodsInterceptPoint() { + @Override + public ElementMatcher getMethodsMatcher() { + return named(INTERCEPT_METHOD) + .and(takesArguments(1)) + .and(takesArgumentWithType(0, "org.springframework.ai.embedding.EmbeddingRequest")); + } + + @Override + public String getMethodsInterceptor() { + return INTERCEPTOR_CLASS; + } + + @Override + public boolean isOverrideArgs() { + return false; + } + } + }; + } +} diff --git a/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/resources/skywalking-plugin.def b/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/resources/skywalking-plugin.def index 5c7eec110c..62e1f6af3e 100644 --- a/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/resources/skywalking-plugin.def +++ b/apm-sniffer/apm-sdk-plugin/spring-plugins/spring-ai-1.x-plugin/src/main/resources/skywalking-plugin.def @@ -15,8 +15,10 @@ # limitations under the License. spring-ai-1.x=org.apache.skywalking.apm.plugin.spring.ai.v1.define.ChatModelInstrumentation +spring-ai-1.x=org.apache.skywalking.apm.plugin.spring.ai.v1.define.EmbeddingModelInstrumentation spring-ai-1.x=org.apache.skywalking.apm.plugin.spring.ai.v1.define.ToolCallbackInstrumentation spring-ai-1.x=org.apache.skywalking.apm.plugin.spring.ai.v1.define.DefaultToolCallingManagerInstrumentation +spring-ai-1.x=org.apache.skywalking.apm.plugin.spring.ai.v1.define.AbstractObservationVectorStoreInstrumentation spring-ai-1.x=org.apache.skywalking.apm.plugin.spring.ai.v1.define.provider.AnthropicApiInstrumentation spring-ai-1.x=org.apache.skywalking.apm.plugin.spring.ai.v1.define.provider.DeepSeekApiInstrumentation spring-ai-1.x=org.apache.skywalking.apm.plugin.spring.ai.v1.define.provider.HuggingfaceChatModelInstrumentation diff --git a/apm-sniffer/config/agent.config b/apm-sniffer/config/agent.config index 38c9bbf8d5..146d0d94fc 100755 --- a/apm-sniffer/config/agent.config +++ b/apm-sniffer/config/agent.config @@ -362,4 +362,12 @@ plugin.springai.content_collect_threshold_tokens=${SW_PLUGIN_SPRINGAI_CONTENT_CO # Whether to collect the arguments (input parameters) of the tool/function call. plugin.springai.collect_tool_input=${SW_PLUGIN_SPRINGAI_COLLECT_TOOL_INPUT:false} # Whether to collect the execution result (output) of the tool/function call. -plugin.springai.collect_tool_output=${SW_PLUGIN_SPRINGAI_COLLECT_TOOL_OUTPUT:false} \ No newline at end of file +plugin.springai.collect_tool_output=${SW_PLUGIN_SPRINGAI_COLLECT_TOOL_OUTPUT:false} +# Whether to collect the query of the rag call. +plugin.springai.collect_retrieval_query=${SW_PLUGIN_SPRINGAI_COLLECT_RETRIEVAL_QUERY:false} +# The maximum characters of the collected rag query. +# If the content exceeds this limit, it will be truncated. +# Use a negative value to represent no limit, but be aware this could cause OOM. +plugin.springai.retrieval_query=${SW_PLUGIN_SPRINGAI_RETRIEVAL_QUERY_LENGTH_LIMIT:1024} +# Whether to collect the documents of the rag call. +plugin.springai.collect_retrieval_documents=${SW_PLUGIN_SPRINGAI_COLLECT_RETRIEVAL_DOCUMENTS:false} \ No newline at end of file diff --git a/test/plugin/scenarios/spring-ai-1.x-scenario/bin/startup.sh b/test/plugin/scenarios/spring-ai-1.x-scenario/bin/startup.sh index 8cb423ece2..9950f6fed6 100644 --- a/test/plugin/scenarios/spring-ai-1.x-scenario/bin/startup.sh +++ b/test/plugin/scenarios/spring-ai-1.x-scenario/bin/startup.sh @@ -18,4 +18,4 @@ home="$(cd "$(dirname $0)"; pwd)" -java -Dskywalking.plugin.springai.collect_input_messages=true -Dskywalking.plugin.springai.collect_output_messages=true -Dskywalking.plugin.springai.collect_tool_input=true -Dskywalking.plugin.springai.collect_tool_output=true -jar ${agent_opts} ${home}/../libs/spring-ai-1.x-scenario.jar & \ No newline at end of file +java -Dskywalking.plugin.springai.collect_input_messages=true -Dskywalking.plugin.springai.collect_output_messages=true -Dskywalking.plugin.springai.collect_tool_input=true -Dskywalking.plugin.springai.collect_tool_output=true -Dskywalking.plugin.springai.collect_retrieval_query=true -Dskywalking.plugin.springai.collect_retrieval_documents=true -jar ${agent_opts} ${home}/../libs/spring-ai-1.x-scenario.jar & \ No newline at end of file diff --git a/test/plugin/scenarios/spring-ai-1.x-scenario/config/expectedData.yaml b/test/plugin/scenarios/spring-ai-1.x-scenario/config/expectedData.yaml index 5f40e79a48..d714ae2542 100644 --- a/test/plugin/scenarios/spring-ai-1.x-scenario/config/expectedData.yaml +++ b/test/plugin/scenarios/spring-ai-1.x-scenario/config/expectedData.yaml @@ -142,6 +142,49 @@ segmentItems: - { key: http.method, value: POST } - { key: http.status_code, value: '200' } + - operationName: retrieval/simple/in-memory-map + parentSpanId: 0 + spanId: 6 + spanLayer: GenAI + startTime: not null + endTime: not null + componentId: 178 + spanType: Exit + peer: simple/in-memory-map + tags: + - { key: gen_ai.operation.name, value: retrieval } + - { key: gen_ai.data_source.id, value: simple/in-memory-map } + - { key: gen_ai.request.model, value: text-embedding-3-small } + - { key: gen_ai.request.top_k, value: '2' } + - { key: gen_ai.retrieval.query.text, value: 'What is Apache SkyWalking?' } + - { key: gen_ai.retrieval.documents, value: not null } + + - operationName: Spring-ai/openai/call + parentSpanId: 0 + spanId: 7 + spanLayer: GenAI + startTime: not null + endTime: not null + componentId: 173 + isError: false + spanType: Exit + peer: http://localhost:8080/spring-ai-1.x-scenario/llm/v1/chat/completions + skipAnalysis: false + tags: + - { key: gen_ai.operation.name, value: chat } + - { key: gen_ai.provider.name, value: openai } + - { key: gen_ai.request.model, value: gpt-4.1-2025-04-14 } + - { key: gen_ai.request.temperature, value: '0.7' } + - { key: gen_ai.request.top_p, value: '0.9' } + - { key: gen_ai.response.id, value: chatcmpl-DknJunZ3tgcSkKiv } + - { key: gen_ai.response.model, value: gpt-4.1-2025-04-14 } + - { key: gen_ai.usage.input_tokens, value: '72' } + - { key: gen_ai.usage.output_tokens, value: '25' } + - { key: gen_ai.client.token.usage, value: '97' } + - { key: gen_ai.response.finish_reasons, value: STOP } + - { key: gen_ai.input.messages, value: not null } + - { key: gen_ai.output.messages, value: not null } + - operationName: GET:/spring-ai-1.x-scenario/case/spring-ai-1.x-scenario-case parentSpanId: -1 spanId: 0 diff --git a/test/plugin/scenarios/spring-ai-1.x-scenario/pom.xml b/test/plugin/scenarios/spring-ai-1.x-scenario/pom.xml index b9d39f538a..0d508db8b0 100644 --- a/test/plugin/scenarios/spring-ai-1.x-scenario/pom.xml +++ b/test/plugin/scenarios/spring-ai-1.x-scenario/pom.xml @@ -55,6 +55,11 @@ spring-ai-starter-model-openai + + org.springframework.ai + spring-ai-vector-store + + org.projectlombok lombok diff --git a/test/plugin/scenarios/spring-ai-1.x-scenario/src/main/java/test/apache/skywalking/apm/testcase/jdk/httpclient/config/ChatClientConfig.java b/test/plugin/scenarios/spring-ai-1.x-scenario/src/main/java/test/apache/skywalking/apm/testcase/jdk/httpclient/config/ChatClientConfig.java index 79fde137ab..02d0f3466b 100644 --- a/test/plugin/scenarios/spring-ai-1.x-scenario/src/main/java/test/apache/skywalking/apm/testcase/jdk/httpclient/config/ChatClientConfig.java +++ b/test/plugin/scenarios/spring-ai-1.x-scenario/src/main/java/test/apache/skywalking/apm/testcase/jdk/httpclient/config/ChatClientConfig.java @@ -17,10 +17,18 @@ package test.apache.skywalking.apm.testcase.jdk.httpclient.config; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.chat.client.ChatClient; import org.springframework.ai.openai.OpenAiChatModel; +import org.springframework.ai.vectorstore.SimpleVectorStore; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; +import org.springframework.context.annotation.Lazy; + +import java.util.ArrayList; +import java.util.List; @Configuration public class ChatClientConfig { @@ -29,4 +37,23 @@ public class ChatClientConfig { public ChatClient openAIChatClient(OpenAiChatModel model) { return ChatClient.create(model); } + + @Bean + @Lazy + public VectorStore vectorStore(EmbeddingModel embeddingModel) { + SimpleVectorStore vectorStore = SimpleVectorStore.builder(embeddingModel).build(); + + List documentList = new ArrayList<>(); + documentList.add(new Document("The 2025 AI Summit is scheduled for October 10-12 in San Francisco. " + + "The event will focus on Generative AI and Autonomous Agents.")); + documentList.add(new Document("Apache SkyWalking is an open-source Application Performance Management system " + + "designed for microservices, cloud native, and container-based architectures.")); + documentList.add(new Document("Spring AI provides a unified interface for interacting with different " + + "AI Models, allowing developers to switch between providers with minimal code changes.")); + documentList.add(new Document("A new distributed tracing protocol, TraceContext v2, was proposed " + + "on August 25, 2025, to improve cross-cloud observability.")); + + vectorStore.add(documentList); + return vectorStore; + } } diff --git a/test/plugin/scenarios/spring-ai-1.x-scenario/src/main/java/test/apache/skywalking/apm/testcase/jdk/httpclient/controller/CaseController.java b/test/plugin/scenarios/spring-ai-1.x-scenario/src/main/java/test/apache/skywalking/apm/testcase/jdk/httpclient/controller/CaseController.java index d6512b726e..73dc1b15bc 100644 --- a/test/plugin/scenarios/spring-ai-1.x-scenario/src/main/java/test/apache/skywalking/apm/testcase/jdk/httpclient/controller/CaseController.java +++ b/test/plugin/scenarios/spring-ai-1.x-scenario/src/main/java/test/apache/skywalking/apm/testcase/jdk/httpclient/controller/CaseController.java @@ -18,12 +18,18 @@ package test.apache.skywalking.apm.testcase.jdk.httpclient.controller; import lombok.RequiredArgsConstructor; +import org.springframework.beans.factory.ObjectProvider; import org.springframework.ai.chat.client.ChatClient; +import org.springframework.ai.document.Document; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.web.bind.annotation.GetMapping; import org.springframework.web.bind.annotation.RequestMapping; import org.springframework.web.bind.annotation.RestController; import test.apache.skywalking.apm.testcase.jdk.httpclient.tool.WeatherTool; +import java.util.stream.Collectors; + @RestController @RequestMapping("/case") @RequiredArgsConstructor @@ -31,6 +37,7 @@ public class CaseController { private final WeatherTool weatherTool; private final ChatClient chatClient; + private final ObjectProvider vectorStoreProvider; @GetMapping("/healthCheck") public String healthCheck() { @@ -63,6 +70,22 @@ public String testCase() throws Exception { .doOnNext(System.out::println) .blockLast(); + String question = "What is Apache SkyWalking?"; + VectorStore vectorStore = vectorStoreProvider.getObject(); + String context = vectorStore.similaritySearch(SearchRequest.builder() + .query(question) + .topK(2) + .build()) + .stream() + .map(Document::getText) + .collect(Collectors.joining("\n")); + + chatClient + .prompt(question) + .system("Answer using only the following context:\n" + context) + .call() + .content(); + return "success"; } } diff --git a/test/plugin/scenarios/spring-ai-1.x-scenario/src/main/java/test/apache/skywalking/apm/testcase/jdk/httpclient/controller/LLMMockController.java b/test/plugin/scenarios/spring-ai-1.x-scenario/src/main/java/test/apache/skywalking/apm/testcase/jdk/httpclient/controller/LLMMockController.java index 5221245280..d63482c5fd 100644 --- a/test/plugin/scenarios/spring-ai-1.x-scenario/src/main/java/test/apache/skywalking/apm/testcase/jdk/httpclient/controller/LLMMockController.java +++ b/test/plugin/scenarios/spring-ai-1.x-scenario/src/main/java/test/apache/skywalking/apm/testcase/jdk/httpclient/controller/LLMMockController.java @@ -21,6 +21,7 @@ import com.alibaba.fastjson.JSONArray; import com.alibaba.fastjson.JSONObject; import jakarta.servlet.http.HttpServletResponse; +import org.springframework.web.bind.annotation.PostMapping; import org.springframework.web.bind.annotation.RequestBody; import org.springframework.web.bind.annotation.RequestMapping; import org.springframework.web.bind.annotation.RestController; @@ -177,6 +178,60 @@ public Object completions(@RequestBody JSONObject request, HttpServletResponse r } """; + String ragLlmResponse = """ + { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "logprobs": null, + "message": { + "annotations": [], + "content": "Apache SkyWalking is an open-source Application Performance Management system designed for microservices, cloud native, and container-based architectures.", + "refusal": null, + "role": "assistant" + } + } + ], + "created": 1780045046, + "id": "chatcmpl-DknJunZ3tgcSkKiv", + "model": "gpt-4.1-2025-04-14", + "object": "chat.completion", + "service_tier": "default", + "system_fingerprint": "fp_a7294185dc", + "usage": { + "completion_tokens": 25, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0 + }, + "latency_checkpoint": { + "engine_tbt_ms": 23, + "engine_ttft_ms": 672, + "engine_ttlt_ms": 1259, + "pre_inference_ms": 157, + "service_tbt_ms": 610, + "service_ttft_ms": 12302, + "service_ttlt_ms": 27518, + "total_duration_ms": 27381, + "user_visible_ttft_ms": 12145 + }, + "prompt_tokens": 72, + "prompt_tokens_details": { + "audio_tokens": 0, + "cached_tokens": 0 + }, + "total_tokens": 97 + } + } + """; + + if (isRagLlmRequest(messages)) { + return JSON.parseObject(ragLlmResponse); + } + if ("tool".equals(lastRole)) { return JSON.parseObject(finalResponse); } @@ -227,4 +282,86 @@ private String escapeJson(String input) { if (input == null) return ""; return input.replace("\"", "\\\"").replace("\n", "\\n").replace("\r", "\\r"); } + + private boolean isRagLlmRequest(JSONArray messages) { + if (messages == null || messages.size() < 2) { + return false; + } + + JSONObject lastMessage = messages.getJSONObject(messages.size() - 1); + if (!"user".equals(lastMessage.getString("role")) + || !"What is Apache SkyWalking?".equals(lastMessage.getString("content"))) { + return false; + } + + for (int i = 0; i < messages.size() - 1; i++) { + JSONObject message = messages.getJSONObject(i); + if (!"system".equals(message.getString("role"))) { + continue; + } + String content = message.getString("content"); + if (content != null + && content.startsWith("Answer using only the following context:") + && content.contains("Apache SkyWalking is an open-source Application Performance Management system")) { + return true; + } + } + return false; + } + + @PostMapping("/v1/embeddings") + public Object embeddings(@RequestBody JSONObject request) { + Object input = request.get("input"); + JSONArray inputs = input instanceof JSONArray ? (JSONArray) input : new JSONArray(); + if (!(input instanceof JSONArray)) { + inputs.add(String.valueOf(input)); + } + + JSONArray data = new JSONArray(); + for (int i = 0; i < inputs.size(); i++) { + JSONObject item = new JSONObject(); + item.put("object", "embedding"); + item.put("index", i); + item.put("embedding", embeddingFor(inputs.getString(i))); + data.add(item); + } + + JSONObject usage = new JSONObject(); + usage.put("prompt_tokens", inputs.size()); + usage.put("total_tokens", inputs.size()); + + JSONObject response = new JSONObject(); + response.put("object", "list"); + response.put("model", "text-embedding-3-small"); + response.put("data", data); + response.put("usage", usage); + return response; + } + + private JSONArray embeddingFor(String input) { + String text = input == null ? "" : input.toLowerCase(); + double[] values = new double[]{ + score(text, "summit", "san francisco", "generative", "autonomous"), + score(text, "skywalking", "apm", "microservices", "cloud native"), + score(text, "spring ai", "models", "providers"), + score(text, "tracecontext", "tracing", "observability"), + 0.1 + }; + + JSONArray embedding = new JSONArray(); + for (double value : values) { + embedding.add(value); + } + return embedding; + } + + private double score(String text, String... keywords) { + double value = 0.0; + for (String keyword : keywords) { + if (text.contains(keyword)) { + value += 1.0; + } + } + return value; + } } diff --git a/test/plugin/scenarios/spring-ai-1.x-scenario/src/main/resources/application.yaml b/test/plugin/scenarios/spring-ai-1.x-scenario/src/main/resources/application.yaml index c4f5c58851..fceabe447e 100644 --- a/test/plugin/scenarios/spring-ai-1.x-scenario/src/main/resources/application.yaml +++ b/test/plugin/scenarios/spring-ai-1.x-scenario/src/main/resources/application.yaml @@ -30,6 +30,9 @@ spring: temperature: 0.7 max-tokens: 1000 top-p: 0.9 + embedding: + options: + model: text-embedding-3-small