Merge pull request #1260 from flyinfish/1256

add ChatModelSpanContributor

core/deployment/pom.xml

@@ -112,6 +112,21 @@
             <artifactId>quarkus-test-vertx</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>io.opentelemetry.instrumentation</groupId>
+            <artifactId>opentelemetry-instrumentation-api</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+          <groupId>io.opentelemetry</groupId>
+          <artifactId>opentelemetry-sdk-testing</artifactId>
+          <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>io.quarkus</groupId>
+            <artifactId>quarkus-junit5-mockito</artifactId>
+            <scope>test</scope>
+        </dependency>        
     </dependencies>
     <build>
         <plugins>

core/deployment/src/test/java/io/quarkiverse/langchain4j/test/listeners/ListenersProcessorAbstractSpanChatModelListenerTest.java

@@ -0,0 +1,125 @@
+package io.quarkiverse.langchain4j.test.listeners;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.awaitility.Awaitility.await;
+
+import java.util.HashMap;
+import java.util.List;
+
+import jakarta.enterprise.context.ApplicationScoped;
+import jakarta.inject.Inject;
+
+import org.jboss.shrinkwrap.api.ShrinkWrap;
+import org.jboss.shrinkwrap.api.asset.StringAsset;
+import org.jboss.shrinkwrap.api.spec.JavaArchive;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import dev.langchain4j.data.message.AiMessage;
+import dev.langchain4j.data.message.UserMessage;
+import dev.langchain4j.model.chat.listener.ChatModelErrorContext;
+import dev.langchain4j.model.chat.listener.ChatModelRequestContext;
+import dev.langchain4j.model.chat.listener.ChatModelResponseContext;
+import dev.langchain4j.model.chat.request.ChatRequest;
+import dev.langchain4j.model.chat.request.DefaultChatRequestParameters;
+import dev.langchain4j.model.chat.response.ChatResponse;
+import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter;
+import io.opentelemetry.sdk.trace.data.SpanData;
+import io.opentelemetry.sdk.trace.internal.data.ExceptionEventData;
+import io.quarkiverse.langchain4j.runtime.listeners.ChatModelSpanContributor;
+import io.quarkiverse.langchain4j.runtime.listeners.SpanChatModelListener;
+import io.quarkus.arc.All;
+
+abstract class ListenersProcessorAbstractSpanChatModelListenerTest {
+    @Inject
+    SpanChatModelListener spanChatModelListener;
+    @Inject
+    @All
+    List<ChatModelSpanContributor> contributors;
+    @Inject
+    InMemorySpanExporter exporter;
+
+    static JavaArchive appWithInMemorySpanExporter() {
+        var applicationProperties = """
+                # Since using a InMemorySpanExporter inside our tests,
+                # these properties reduce the export timeout and schedule delay from the BatchSpanProcessor
+                quarkus.otel.bsp.schedule.delay=PT0.001S
+                quarkus.otel.bsp.max.queue.size=1
+                quarkus.otel.bsp.max.export.batch.size=1
+                """;
+        return ShrinkWrap.create(JavaArchive.class)
+                .addAsResource(new StringAsset(applicationProperties), "application.properties")
+                .addClasses(InMemorySpanExporterProducer.class);
+    }
+
+    @BeforeEach
+    void cleanupSpans() {
+        exporter.reset();
+    }
+
+    @Test
+    void shouldCreateSpanOnSuccess() {
+        var ctx = MockedContexts.create();
+
+        spanChatModelListener.onRequest(ctx.requestContext());
+        spanChatModelListener.onResponse(ctx.responseContext());
+
+        await().untilAsserted(() -> assertThat(exporter.getFinishedSpanItems()).hasSize(1));
+        var actualSpan = exporter.getFinishedSpanItems().get(0);
+        assertThat(actualSpan.getName()).isEqualTo("completion --mock-model-name--");
+        verifySuccessfulSpan(actualSpan);
+    }
+
+    protected void verifySuccessfulSpan(SpanData actualSpan) {
+        // nothing here
+    }
+
+    @Test
+    void shouldCreateAndEndSpanOnFailure() {
+        var ctx = MockedContexts.create();
+
+        spanChatModelListener.onRequest(ctx.requestContext());
+        spanChatModelListener.onError(ctx.errorContext());
+
+        await().untilAsserted(() -> assertThat(exporter.getFinishedSpanItems()).hasSize(1));
+        var actualSpan = exporter.getFinishedSpanItems().get(0);
+        assertThat(actualSpan.getName()).isEqualTo("completion --mock-model-name--");
+        assertThat(actualSpan.getEvents())
+                .hasSize(1)
+                .first()
+                .isInstanceOf(ExceptionEventData.class)
+                .extracting(ex -> ((ExceptionEventData) ex).getException().getMessage())
+                .isEqualTo("--failed--");
+        verifyFailedSpan(actualSpan);
+    }
+
+    protected void verifyFailedSpan(SpanData actualSpan) {
+        // nothing here
+    }
+
+    public static class InMemorySpanExporterProducer {
+        @ApplicationScoped
+        InMemorySpanExporter exporter() {
+            return InMemorySpanExporter.create();
+        }
+    }
+
+    record MockedContexts(
+            ChatModelRequestContext requestContext,
+            ChatModelResponseContext responseContext,
+            ChatModelErrorContext errorContext) {
+        static MockedContexts create() {
+            var attributes = new HashMap();
+            var request = ChatRequest.builder().messages(List.of(UserMessage.from("--test-message--")))
+                    .parameters(DefaultChatRequestParameters.builder().modelName("--mock-model-name--").temperature(0.0)
+                            .topP(0.0).build())
+                    .build();
+            var response = ChatResponse.builder().aiMessage(AiMessage.from("--test-response--")).build();
+            var requestCtx = new ChatModelRequestContext(request, attributes);
+            var responseContext = new ChatModelResponseContext(response, request, attributes);
+            var errorCtx = new ChatModelErrorContext(
+                    new RuntimeException("--failed--"), request, attributes);
+            return new MockedContexts(requestCtx, responseContext, errorCtx);
+        }
+    }
+}

core/deployment/src/test/java/io/quarkiverse/langchain4j/test/listeners/ListenersProcessorNoOpentelemetryTest.java

@@ -0,0 +1,37 @@
+package io.quarkiverse.langchain4j.test.listeners;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.List;
+
+import jakarta.inject.Inject;
+
+import org.jboss.shrinkwrap.api.ShrinkWrap;
+import org.jboss.shrinkwrap.api.spec.JavaArchive;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import io.quarkiverse.langchain4j.runtime.listeners.ChatModelSpanContributor;
+import io.quarkiverse.langchain4j.runtime.listeners.SpanChatModelListener;
+import io.quarkus.arc.All;
+import io.quarkus.test.QuarkusUnitTest;
+
+class ListenersProcessorNoOpentelemetryTest {
+    @RegisterExtension
+    static final QuarkusUnitTest unitTest = new QuarkusUnitTest()
+            .setArchiveProducer(
+                    () -> ShrinkWrap.create(JavaArchive.class));
+
+    @Inject
+    @All
+    List<SpanChatModelListener> spanChatModelListeners;
+    @Inject
+    @All
+    List<ChatModelSpanContributor> chatModelSpanContributors;
+
+    @Test
+    void shouldNotHaveSpanChatModelListenerWhenNoOtel() {
+        assertThat(spanChatModelListeners).isEmpty();
+        assertThat(chatModelSpanContributors).isEmpty();
+    }
+}

core/deployment/src/test/java/io/quarkiverse/langchain4j/test/listeners/ListenersProcessorOnlySpanChatModelListenerTest.java

@@ -0,0 +1,27 @@
+package io.quarkiverse.langchain4j.test.listeners;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.List;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import io.quarkus.maven.dependency.Dependency;
+import io.quarkus.test.QuarkusUnitTest;
+
+class ListenersProcessorOnlySpanChatModelListenerTest
+        extends ListenersProcessorAbstractSpanChatModelListenerTest {
+    @RegisterExtension
+    static final QuarkusUnitTest unitTest = new QuarkusUnitTest()
+            .setArchiveProducer(
+                    ListenersProcessorAbstractSpanChatModelListenerTest::appWithInMemorySpanExporter)
+            .setForcedDependencies(
+                    List.of(Dependency.of("io.quarkus", "quarkus-opentelemetry", "3.15.2")));
+
+    @Test
+    void shouldHaveSpanChatModelListenerWithoutContributors() {
+        assertThat(spanChatModelListener).isNotNull();
+        assertThat(contributors).isEmpty();
+    }
+}

core/deployment/src/test/java/io/quarkiverse/langchain4j/test/listeners/ListenersProcessorSingleChatModelSpanContributorTest.java

@@ -0,0 +1,74 @@
+package io.quarkiverse.langchain4j.test.listeners;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.List;
+
+import jakarta.enterprise.context.ApplicationScoped;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import dev.langchain4j.model.chat.listener.ChatModelErrorContext;
+import dev.langchain4j.model.chat.listener.ChatModelRequestContext;
+import dev.langchain4j.model.chat.listener.ChatModelResponseContext;
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.sdk.trace.data.SpanData;
+import io.quarkiverse.langchain4j.runtime.listeners.ChatModelSpanContributor;
+import io.quarkus.maven.dependency.Dependency;
+import io.quarkus.test.QuarkusUnitTest;
+
+class ListenersProcessorSingleChatModelSpanContributorTest
+        extends ListenersProcessorAbstractSpanChatModelListenerTest {
+    @RegisterExtension
+    static final QuarkusUnitTest unitTest = new QuarkusUnitTest()
+            .setArchiveProducer(
+                    () -> appWithInMemorySpanExporter().addClasses(TestChatModelSpanContributor.class))
+            .setForcedDependencies(
+                    List.of(Dependency.of("io.quarkus", "quarkus-opentelemetry", "3.15.2")));
+
+    @Test
+    void shouldHaveSpanChatModelListenerWitContributor() {
+        assertThat(spanChatModelListener).isNotNull();
+        assertThat(contributors).hasSize(1).first().isInstanceOf(TestChatModelSpanContributor.class);
+    }
+
+    @Override
+    protected void verifySuccessfulSpan(SpanData actualSpan) {
+        assertThat(actualSpan.getAttributes().get(AttributeKey.stringKey(("--custom-on-request--"))))
+                .isEqualTo("--value-on-request--");
+        assertThat(actualSpan.getAttributes().get(AttributeKey.stringKey(("--custom-on-response--"))))
+                .isEqualTo("--value-on-response--");
+        assertThat(actualSpan.getAttributes().get(AttributeKey.stringKey(("--custom-on-error--"))))
+                .isNull();
+    }
+
+    @Override
+    protected void verifyFailedSpan(SpanData actualSpan) {
+        assertThat(actualSpan.getAttributes().get(AttributeKey.stringKey(("--custom-on-request--"))))
+                .isEqualTo("--value-on-request--");
+        assertThat(actualSpan.getAttributes().get(AttributeKey.stringKey(("--custom-on-response--"))))
+                .isNull();
+        assertThat(actualSpan.getAttributes().get(AttributeKey.stringKey(("--custom-on-error--"))))
+                .isEqualTo("--value-on-error--");
+    }
+
+    @ApplicationScoped
+    public static class TestChatModelSpanContributor implements ChatModelSpanContributor {
+        @Override
+        public void onRequest(ChatModelRequestContext requestContext, Span currentSpan) {
+            currentSpan.setAttribute("--custom-on-request--", "--value-on-request--");
+        }
+
+        @Override
+        public void onResponse(ChatModelResponseContext responseContext, Span currentSpan) {
+            currentSpan.setAttribute("--custom-on-response--", "--value-on-response--");
+        }
+
+        @Override
+        public void onError(ChatModelErrorContext errorContext, Span currentSpan) {
+            currentSpan.setAttribute("--custom-on-error--", "--value-on-error--");
+        }
+    }
+}