-
Notifications
You must be signed in to change notification settings - Fork 104
Commit
enabling users adding more/custom attributes/data to span enabling us to add pre-built attribute-sets for diffrent APMs Closes #1256
- Loading branch information
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,123 @@ | ||
package io.quarkiverse.langchain4j.test.listeners; | ||
|
||
import static org.assertj.core.api.Assertions.assertThat; | ||
import static org.awaitility.Awaitility.await; | ||
import static org.mockito.Mockito.mock; | ||
import static org.mockito.Mockito.when; | ||
|
||
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 org.mockito.Answers; | ||
|
||
import dev.langchain4j.model.chat.listener.ChatModelErrorContext; | ||
import dev.langchain4j.model.chat.listener.ChatModelRequest; | ||
import dev.langchain4j.model.chat.listener.ChatModelRequestContext; | ||
import dev.langchain4j.model.chat.listener.ChatModelResponse; | ||
import dev.langchain4j.model.chat.listener.ChatModelResponseContext; | ||
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--"); | ||
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--"); | ||
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 = mock(ChatModelRequest.class, Answers.RETURNS_DEEP_STUBS); | ||
when(request.model()).thenReturn("--mock--"); | ||
var response = mock(ChatModelResponse.class, Answers.RETURNS_DEEP_STUBS); | ||
var requestCtx = new ChatModelRequestContext(request, attributes); | ||
var responseContext = new ChatModelResponseContext(response, request, attributes); | ||
var errorCtx = new ChatModelErrorContext( | ||
new RuntimeException("--failed--"), request, response, attributes); | ||
return new MockedContexts(requestCtx, responseContext, errorCtx); | ||
} | ||
} | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,69 @@ | ||
package io.quarkiverse.langchain4j.test.listeners; | ||
|
||
import static org.assertj.core.api.Assertions.assertThat; | ||
|
||
import java.util.List; | ||
import java.util.function.Supplier; | ||
|
||
import jakarta.enterprise.context.control.ActivateRequestContext; | ||
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 dev.langchain4j.data.message.AiMessage; | ||
import dev.langchain4j.model.chat.ChatLanguageModel; | ||
import dev.langchain4j.model.output.Response; | ||
import dev.langchain4j.service.UserMessage; | ||
import io.quarkiverse.langchain4j.RegisterAiService; | ||
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) | ||
.addClasses(AiService.class, EchoChatLanguageModelSupplier.class)); | ||
|
||
@Inject | ||
AiService aiService; | ||
@Inject | ||
@All | ||
List<SpanChatModelListener> spanChatModelListeners; | ||
@Inject | ||
@All | ||
List<ChatModelSpanContributor> chatModelSpanContributors; | ||
|
||
@Test | ||
void shouldNotHaveSpanChatModelListenerWhenNoOtel() { | ||
assertThat(aiService).isNotNull(); | ||
assertThat(spanChatModelListeners).isEmpty(); | ||
assertThat(chatModelSpanContributors).isEmpty(); | ||
} | ||
|
||
@Test | ||
@ActivateRequestContext | ||
void shouldGetResponseFromAiService() { | ||
var actual = aiService.test(); | ||
|
||
assertThat(actual).isEqualTo("test:echo"); | ||
} | ||
|
||
@RegisterAiService(chatLanguageModelSupplier = EchoChatLanguageModelSupplier.class, chatMemoryProviderSupplier = RegisterAiService.NoChatMemoryProviderSupplier.class) | ||
interface AiService { | ||
@UserMessage("test") | ||
String test(); | ||
} | ||
|
||
public static class EchoChatLanguageModelSupplier implements Supplier<ChatLanguageModel> { | ||
@Override | ||
public ChatLanguageModel get() { | ||
return (messages) -> new Response<>(new AiMessage(messages.get(0).text() + ":echo")); | ||
} | ||
} | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,27 @@ | ||
package io.quarkiverse.langchain4j.test.listeners; | ||
Check failure on line 1 in core/deployment/src/test/java/io/quarkiverse/langchain4j/test/listeners/ListenersProcessorOnlySpanChatModelListenerTest.java
|
||
|
||
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(); | ||
} | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,74 @@ | ||
package io.quarkiverse.langchain4j.test.listeners; | ||
Check failure on line 1 in core/deployment/src/test/java/io/quarkiverse/langchain4j/test/listeners/ListenersProcessorSingleChatModelSpanContributorTest.java
|
||
|
||
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--"); | ||
} | ||
} | ||
} |