diff --git a/build.gradle b/build.gradle index 4f14205..c57f871 100644 --- a/build.gradle +++ b/build.gradle @@ -37,6 +37,7 @@ dependencies { api "io.opentelemetry:opentelemetry-sdk-logs:${otelVersion}" implementation "io.opentelemetry:opentelemetry-exporter-otlp:${otelVersion}" implementation "io.opentelemetry:opentelemetry-exporter-logging:${otelVersion}" + implementation "io.opentelemetry:opentelemetry-sdk-testing:${otelVersion}" implementation "io.opentelemetry:opentelemetry-semconv:1.25.0-alpha" // HTTP Client (Java 11+) @@ -61,8 +62,6 @@ dependencies { testImplementation "org.junit.jupiter:junit-jupiter:${junitVersion}" testImplementation "org.junit.jupiter:junit-jupiter-params:${junitVersion}" testImplementation "org.assertj:assertj-core:${assertjVersion}" - testImplementation "org.mockito:mockito-junit-jupiter:${mockitoVersion}" - testImplementation "io.opentelemetry:opentelemetry-sdk-testing:${otelVersion}" } test { diff --git a/examples/src/main/java/dev/braintrust/examples/SimpleOpenTelemetryExample.java b/examples/src/main/java/dev/braintrust/examples/SimpleOpenTelemetryExample.java index 50f275e..857c4c0 100644 --- a/examples/src/main/java/dev/braintrust/examples/SimpleOpenTelemetryExample.java +++ b/examples/src/main/java/dev/braintrust/examples/SimpleOpenTelemetryExample.java @@ -6,7 +6,7 @@ public class SimpleOpenTelemetryExample { public static void main(String[] args) throws Exception { var braintrustConfig = BraintrustConfig.fromEnvironment(); - var openTelemetry = BraintrustTracing.of(braintrustConfig, true); + var openTelemetry = BraintrustTracing.quickstart(); var tracer = BraintrustTracing.getTracer(openTelemetry); var span = tracer.spanBuilder("hello-java").startSpan(); diff --git a/src/main/java/dev/braintrust/api/BraintrustApiClient.java b/src/main/java/dev/braintrust/api/BraintrustApiClient.java index 2322bd6..00adb53 100644 --- a/src/main/java/dev/braintrust/api/BraintrustApiClient.java +++ b/src/main/java/dev/braintrust/api/BraintrustApiClient.java @@ -13,10 +13,10 @@ import java.net.http.HttpRequest; import java.net.http.HttpResponse; import java.time.Duration; -import java.util.List; -import java.util.Optional; +import java.util.*; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; public interface BraintrustApiClient { @@ -224,6 +224,71 @@ private static ObjectMapper createObjectMapper() { } } + /** Implementation for test doubling */ + class InMemoryImpl implements BraintrustApiClient { + private final List organizationAndProjectInfos; + private final Set experiments = + Collections.newSetFromMap(new ConcurrentHashMap<>()); + + public InMemoryImpl(OrganizationAndProjectInfo... organizationAndProjectInfos) { + this.organizationAndProjectInfos = List.of(organizationAndProjectInfos); + } + + @Override + public Project getOrCreateProject(String projectName) { + // Find existing project by name + for (var orgAndProject : organizationAndProjectInfos) { + if (orgAndProject.project().name().equals(projectName)) { + return orgAndProject.project(); + } + } + throw new RuntimeException( + "Project '" + + projectName + + "' not found in test data. Please add it to the InMemoryImpl" + + " constructor."); + } + + @Override + public Optional getProject(String projectId) { + return organizationAndProjectInfos.stream() + .map(OrganizationAndProjectInfo::project) + .filter(project -> project.id().equals(projectId)) + .findFirst(); + } + + @Override + public Experiment getOrCreateExperiment(CreateExperimentRequest request) { + var existing = + experiments.stream() + .filter(exp -> exp.name().equals(request.name())) + .findFirst(); + return existing.orElseGet( + () -> + new Experiment( + request.name().hashCode() + "", + request.projectId(), + request.name(), + request.description(), + "notused", + "notused")); + } + + @Override + public Optional getProjectAndOrgInfo() { + return organizationAndProjectInfos.isEmpty() + ? Optional.empty() + : Optional.of(organizationAndProjectInfos.get(0)); + } + + @Override + public Optional getProjectAndOrgInfo(String projectId) { + return organizationAndProjectInfos.stream() + .filter(orgAndProject -> orgAndProject.project().id().equals(projectId)) + .findFirst(); + } + } + // Request/Response DTOs record CreateProjectRequest(String name) {} diff --git a/src/main/java/dev/braintrust/config/BraintrustConfig.java b/src/main/java/dev/braintrust/config/BraintrustConfig.java index beae4e7..4bbfbc3 100644 --- a/src/main/java/dev/braintrust/config/BraintrustConfig.java +++ b/src/main/java/dev/braintrust/config/BraintrustConfig.java @@ -30,6 +30,10 @@ public final class BraintrustConfig extends BaseConfig { private final Duration requestTimeout = Duration.ofSeconds(getConfig("BRAINTRUST_REQUEST_TIMEOUT", 30)); + /** Setting for unit testing. Do not use in production. */ + private final boolean unitTetJavaExportSpansInMemory = + getConfig("BRAINTRUST_TEST_JAVA_EXPORT_SPANS_IN_MEMORY", false); + public static BraintrustConfig fromEnvironment() { return of(); } @@ -79,15 +83,22 @@ public Optional getBraintrustParentValue() { /** fetch all project info and IDs from the braintrust api */ public URI fetchProjectURI() { + return fetchProjectURI(BraintrustApiClient.of(this)); + } + + URI fetchProjectURI(BraintrustApiClient client) { try { - var client = BraintrustApiClient.of(this); var orgAndProject = client.getProjectAndOrgInfo().orElseThrow(); + var baseURI = new URI(appUrl()); return new URI( - appUrl() + baseURI.getScheme(), + baseURI.getHost(), + baseURI.getPath() + "/app/" + orgAndProject.orgInfo().name() + "/p/" - + orgAndProject.project().name()); + + orgAndProject.project().name(), + null); } catch (URISyntaxException e) { throw new RuntimeException(e); } diff --git a/src/main/java/dev/braintrust/eval/Eval.java b/src/main/java/dev/braintrust/eval/Eval.java index 502d3c6..1332c40 100644 --- a/src/main/java/dev/braintrust/eval/Eval.java +++ b/src/main/java/dev/braintrust/eval/Eval.java @@ -4,15 +4,19 @@ import com.fasterxml.jackson.databind.ObjectMapper; import dev.braintrust.api.BraintrustApiClient; import dev.braintrust.config.BraintrustConfig; +import dev.braintrust.spec.SdkSpec; import dev.braintrust.trace.BraintrustContext; -import dev.braintrust.trace.BraintrustSpanProcessor; import dev.braintrust.trace.BraintrustTracing; +import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.SpanKind; import io.opentelemetry.api.trace.Tracer; +import java.net.URI; import java.util.*; import java.util.function.Function; import javax.annotation.Nonnull; import javax.annotation.Nullable; +import lombok.Getter; +import lombok.SneakyThrows; /** * An evaluation framework for testing AI models. @@ -21,6 +25,8 @@ * @param The type of output produced by the task */ public final class Eval { + private static final AttributeKey PARENT = + AttributeKey.stringKey(SdkSpec.Attributes.PARENT); private static final ObjectMapper JSON_MAPPER = new com.fasterxml.jackson.databind.ObjectMapper(); private final @Nonnull String experimentName; @@ -35,7 +41,7 @@ public final class Eval { private Eval(Builder builder) { this.experimentName = builder.experimentName; this.config = Objects.requireNonNull(builder.config); - this.client = BraintrustApiClient.of(config); + this.client = Objects.requireNonNull(builder.apiClient); if (null == builder.projectId) { this.orgAndProject = client.getProjectAndOrgInfo().orElseThrow(); } else { @@ -73,31 +79,24 @@ private EvalCase.Result evalOne( tracer.spanBuilder("eval") // TODO: allow names for eval cases .setNoParent() // each eval case is its own trace .setSpanKind(SpanKind.CLIENT) - .setAttribute( - BraintrustSpanProcessor.PARENT, "experiment_id:" + experimentId) + .setAttribute(PARENT, "experiment_id:" + experimentId) .setAttribute("braintrust.span_attributes", "{\"type\":\"eval\"}") // FIXME: use proper object mapper for json stuff .setAttribute( "braintrust.input_json", "{ \"input\":\"" + evalCase.input() + "\"}") .setAttribute("braintrust.expected", "\"" + evalCase.expected() + "\"") - // TODO: these attributes are deprecated apparently? Do we need to set them? - .setAttribute(BraintrustSpanProcessor.PARENT_EXPERIMENT_ID, experimentId) - .setAttribute(BraintrustSpanProcessor.PARENT_TYPE, "experiment") .startSpan(); - try (var rootScope = - BraintrustContext.forExperiment(experimentId, rootSpan).makeCurrent()) { + try (var rootScope = BraintrustContext.ofExperiment(experimentId, rootSpan).makeCurrent()) { final OUTPUT result; { // run task var taskSpan = tracer.spanBuilder("task") - .setAttribute( - BraintrustSpanProcessor.PARENT, - "experiment_id:" + experimentId) + .setAttribute(PARENT, "experiment_id:" + experimentId) .setAttribute("braintrust.span_attributes", "{\"type\":\"task\"}") .startSpan(); try (var unused = - BraintrustContext.forExperiment(experimentId, taskSpan).makeCurrent()) { + BraintrustContext.ofExperiment(experimentId, taskSpan).makeCurrent()) { result = task.apply(evalCase); } finally { taskSpan.end(); @@ -114,13 +113,11 @@ private EvalCase.Result evalOne( { // run scorers var scoreSpan = tracer.spanBuilder("score") - .setAttribute( - BraintrustSpanProcessor.PARENT, - "experiment_id:" + experimentId) + .setAttribute(PARENT, "experiment_id:" + experimentId) .setAttribute("braintrust.span_attributes", "{\"type\":\"score\"}") .startSpan(); try (var unused = - BraintrustContext.forExperiment(experimentId, scoreSpan).makeCurrent()) { + BraintrustContext.ofExperiment(experimentId, scoreSpan).makeCurrent()) { // NOTE: linked hash map to preserve ordering. Not in the spec but nice user // experience final HashMap nameToScore = new LinkedHashMap<>(); @@ -158,17 +155,23 @@ private EvalCase.Result evalOne( /** Results of all eval cases of an experiment. */ public class Result { - private final String experimentUrl; + @Getter private final String experimentUrl; + @SneakyThrows private Result() { + var baseURI = new URI(config.appUrl()); this.experimentUrl = - config.appUrl() - + "/app/" - + orgAndProject.orgInfo().name() - + "/p/" - + orgAndProject.project().name() - + "/experiments/" - + experimentName; + new URI( + baseURI.getScheme(), + baseURI.getHost(), + "/app/" + + orgAndProject.orgInfo().name() + + "/p/" + + orgAndProject.project().name() + + "/experiments/" + + experimentName, + null) + .toASCIIString(); } public String createReportString() { @@ -189,6 +192,7 @@ public static Builder builder() { public static final class Builder { private @Nonnull String experimentName = "unnamed-java-eval"; private @Nullable BraintrustConfig config; + private @Nullable BraintrustApiClient apiClient; private @Nullable String projectId; private @Nullable Tracer tracer = null; private @Nonnull List> evalCases = List.of(); @@ -211,6 +215,9 @@ public Eval build() { if (scorers.isEmpty()) { throw new RuntimeException("must provide at least one scorer"); } + if (null == apiClient) { + apiClient = BraintrustApiClient.of(config); + } Objects.requireNonNull(task); return new Eval<>(this); } @@ -230,6 +237,11 @@ public Builder config(BraintrustConfig config) { return this; } + public Builder apiClient(BraintrustApiClient apiClient) { + this.apiClient = apiClient; + return this; + } + public Builder tracer(Tracer tracer) { this.tracer = tracer; return this; diff --git a/src/main/java/dev/braintrust/spec/SdkSpec.java b/src/main/java/dev/braintrust/spec/SdkSpec.java new file mode 100644 index 0000000..811956b --- /dev/null +++ b/src/main/java/dev/braintrust/spec/SdkSpec.java @@ -0,0 +1,12 @@ +package dev.braintrust.spec; + +public final class SdkSpec { + + public static final class Attributes { + public static final String PARENT = "braintrust.parent"; + } + + public static final class Context {} + + private SdkSpec() {} +} diff --git a/src/main/java/dev/braintrust/trace/BraintrustContext.java b/src/main/java/dev/braintrust/trace/BraintrustContext.java index 6a8b3b1..a3967e2 100644 --- a/src/main/java/dev/braintrust/trace/BraintrustContext.java +++ b/src/main/java/dev/braintrust/trace/BraintrustContext.java @@ -3,49 +3,32 @@ import io.opentelemetry.api.trace.Span; import io.opentelemetry.context.Context; import io.opentelemetry.context.ContextKey; +import java.util.Objects; import java.util.Optional; +import javax.annotation.Nonnull; import javax.annotation.Nullable; /** - * Context propagation for Braintrust-specific attributes. Uses OpenTelemetry's Context API for - * thread-safe propagation. + * Used to identify the braintrust parent for spans and experiments. SDK users probably don't want + * to use this and instead should use {@link BraintrustTracing} or {@link dev.braintrust.eval.Eval} */ public final class BraintrustContext { private static final ContextKey KEY = ContextKey.named("braintrust-context"); + // NOTE we're actually not using this right now, but leaving in for the future @Nullable private final String projectId; @Nullable private final String experimentId; - @Nullable private final String parentType; - private BraintrustContext( - @Nullable String projectId, - @Nullable String experimentId, - @Nullable String parentType) { + private BraintrustContext(@Nullable String projectId, @Nullable String experimentId) { this.projectId = projectId; this.experimentId = experimentId; - this.parentType = parentType; - } - - /** Creates a context for a project parent. */ - public static BraintrustContext forProject(String projectId) { - return new BraintrustContext(projectId, null, "project"); - } - - /** Creates a context for an experiment parent. */ - public static BraintrustContext forExperiment(String experimentId) { - return new BraintrustContext(null, experimentId, "experiment"); } /** Creates a context for an experiment parent. */ - public static Context forExperiment(String experimentId, Span span) { - return Context.current() - .with(span) - .with(KEY, new BraintrustContext(null, experimentId, "experiment")); - } - - /** Stores this context in the given Context. */ - public Context storeInContext(Context context) { - return context.with(KEY, this); + public static Context ofExperiment(@Nonnull String experimentId, @Nonnull Span span) { + Objects.requireNonNull(experimentId); + Objects.requireNonNull(span); + return Context.current().with(span).with(KEY, new BraintrustContext(null, experimentId)); } /** Retrieves a BraintrustContext from the given Context. */ @@ -54,12 +37,6 @@ public static BraintrustContext fromContext(Context context) { return context.get(KEY); } - /** Returns the current BraintrustContext from the current Context. */ - @Nullable - public static BraintrustContext current() { - return fromContext(Context.current()); - } - public Optional projectId() { return Optional.ofNullable(projectId); } @@ -67,49 +44,4 @@ public Optional projectId() { public Optional experimentId() { return Optional.ofNullable(experimentId); } - - public Optional parentType() { - return Optional.ofNullable(parentType); - } - - /** Builder for creating contexts with multiple attributes. */ - public static final class Builder { - private String projectId; - private String experimentId; - private String parentType; - - public Builder projectId(String projectId) { - this.projectId = projectId; - this.experimentId = null; // Clear experiment when setting project - if (projectId != null) { - this.parentType = "project"; - } - return this; - } - - public Builder experimentId(String experimentId) { - this.experimentId = experimentId; - this.projectId = null; // Clear project when setting experiment - if (experimentId != null) { - this.parentType = "experiment"; - } - return this; - } - - public BraintrustContext build() { - return new BraintrustContext(projectId, experimentId, parentType); - } - - public Context storeInContext(Context context) { - return build().storeInContext(context); - } - - public Context storeInCurrent() { - return storeInContext(Context.current()); - } - } - - public static Builder builder() { - return new Builder(); - } } diff --git a/src/main/java/dev/braintrust/trace/BraintrustLogExporter.java b/src/main/java/dev/braintrust/trace/BraintrustLogExporter.java index 035f0ad..02493bb 100644 --- a/src/main/java/dev/braintrust/trace/BraintrustLogExporter.java +++ b/src/main/java/dev/braintrust/trace/BraintrustLogExporter.java @@ -49,31 +49,21 @@ public CompletableResultCode export(Collection logs) { } private String getParentFromLog(LogRecordData log) { - // Check for the braintrust.parent attribute var parent = log.getAttributes().get(BraintrustSpanProcessor.PARENT); if (parent != null) { return parent; } - - // Check legacy attributes for backward compatibility - var experimentId = log.getAttributes().get(BraintrustSpanProcessor.PARENT_EXPERIMENT_ID); - if (experimentId != null) { - return "experiment_id:" + experimentId; - } - - var projectId = log.getAttributes().get(BraintrustSpanProcessor.PARENT_PROJECT_ID); - if (projectId != null) { - return "project_id:" + projectId; - } - return config.getBraintrustParentValue().orElse(""); } private CompletableResultCode exportWithParent(String parent, List logs) { try { // Get or create exporter for this parent + if (exporterCache.size() >= 1024) { + BraintrustLogger.info("Clearing exporter cache. This should not happen"); + exporterCache.clear(); + } var exporter = - // FIXME: This will grow unbounded exporterCache.computeIfAbsent( parent, p -> { diff --git a/src/main/java/dev/braintrust/trace/BraintrustSpanExporter.java b/src/main/java/dev/braintrust/trace/BraintrustSpanExporter.java index 0d5a220..5c1f8e4 100644 --- a/src/main/java/dev/braintrust/trace/BraintrustSpanExporter.java +++ b/src/main/java/dev/braintrust/trace/BraintrustSpanExporter.java @@ -10,6 +10,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.stream.Collectors; /** @@ -17,6 +18,8 @@ * attributes. */ class BraintrustSpanExporter implements SpanExporter { + /** Only used in unit tests. */ + static final Map> SPANS_EXPORTED = new ConcurrentHashMap<>(); private final BraintrustConfig config; private final String tracesEndpoint; @@ -49,31 +52,21 @@ public CompletableResultCode export(Collection spans) { } private String getParentFromSpan(SpanData span) { - // Check for the braintrust.parent attribute var parent = span.getAttributes().get(BraintrustSpanProcessor.PARENT); if (parent != null) { return parent; } - - // Check legacy attributes for backward compatibility - var experimentId = span.getAttributes().get(BraintrustSpanProcessor.PARENT_EXPERIMENT_ID); - if (experimentId != null) { - return "experiment_id:" + experimentId; - } - - var projectId = span.getAttributes().get(BraintrustSpanProcessor.PARENT_PROJECT_ID); - if (projectId != null) { - return "project_id:" + projectId; - } - return config.getBraintrustParentValue().orElse(""); } private CompletableResultCode exportWithParent(String parent, List spans) { try { // Get or create exporter for this parent + if (exporterCache.size() >= 1024) { + BraintrustLogger.info("Clearing exporter cache. This should not happen"); + exporterCache.clear(); + } var exporter = - // FIXME: This will grow unbounded exporterCache.computeIfAbsent( parent, p -> { @@ -96,8 +89,13 @@ private CompletableResultCode exportWithParent(String parent, List spa }); BraintrustLogger.debug("Exporting {} spans with x-bt-parent: {}", spans.size(), parent); - // Export the spans - return exporter.export(spans); + if (config.unitTetJavaExportSpansInMemory()) { + SPANS_EXPORTED.putIfAbsent(parent, new CopyOnWriteArrayList<>()); + SPANS_EXPORTED.get(parent).addAll(spans); + return CompletableResultCode.ofSuccess(); + } else { + return exporter.export(spans); + } } catch (Exception e) { BraintrustLogger.error("Failed to export spans", e); return CompletableResultCode.ofFailure(); diff --git a/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java b/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java index ad708ea..616058b 100644 --- a/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java +++ b/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java @@ -2,6 +2,7 @@ import dev.braintrust.config.BraintrustConfig; import dev.braintrust.log.BraintrustLogger; +import dev.braintrust.spec.SdkSpec; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.context.Context; import io.opentelemetry.sdk.common.CompletableResultCode; @@ -21,23 +22,11 @@ * Custom span processor that enriches spans with Braintrust-specific attributes. Supports parent * assignment to projects or experiments. */ -public class BraintrustSpanProcessor implements SpanProcessor { +class BraintrustSpanProcessor implements SpanProcessor { // Braintrust-specific attributes - public static final AttributeKey PARENT = AttributeKey.stringKey("braintrust.parent"); - - // Legacy attributes for backward compatibility (deprecated) - @Deprecated - public static final AttributeKey PARENT_PROJECT_ID = - AttributeKey.stringKey("braintrust.parent.project_id"); - - @Deprecated - public static final AttributeKey PARENT_EXPERIMENT_ID = - AttributeKey.stringKey("braintrust.parent.experiment_id"); - - @Deprecated - public static final AttributeKey PARENT_TYPE = - AttributeKey.stringKey("braintrust.parent.type"); + public static final AttributeKey PARENT = + AttributeKey.stringKey(SdkSpec.Attributes.PARENT); // Usage metrics public static final AttributeKey USAGE_PROMPT_TOKENS = diff --git a/src/main/java/dev/braintrust/trace/BraintrustTracing.java b/src/main/java/dev/braintrust/trace/BraintrustTracing.java index 01fdc5c..4b7dfc3 100644 --- a/src/main/java/dev/braintrust/trace/BraintrustTracing.java +++ b/src/main/java/dev/braintrust/trace/BraintrustTracing.java @@ -32,20 +32,6 @@ public final class BraintrustTracing { private static final String INSTRUMENTATION_NAME = "braintrust-java"; private static final String INSTRUMENTATION_VERSION = "0.0.1"; - private BraintrustTracing() { - // Utility class - } - - /** Gets a tracer with Braintrust instrumentation scope. */ - public static Tracer getTracer() { - return getTracer(GlobalOpenTelemetry.get()); - } - - /** Gets a tracer from a specific OpenTelemetry instance. */ - public static Tracer getTracer(OpenTelemetry openTelemetry) { - return openTelemetry.getTracer(INSTRUMENTATION_NAME, INSTRUMENTATION_VERSION); - } - /** * Quick start method that sets up global OpenTelemetry with Braintrust defaults.
*
@@ -163,4 +149,16 @@ public static void enable( result.isDone(), result.isSuccess())); })); } + + /** Gets a tracer with Braintrust instrumentation scope. */ + public static Tracer getTracer() { + return getTracer(GlobalOpenTelemetry.get()); + } + + /** Gets a tracer from a specific OpenTelemetry instance. */ + public static Tracer getTracer(OpenTelemetry openTelemetry) { + return openTelemetry.getTracer(INSTRUMENTATION_NAME, INSTRUMENTATION_VERSION); + } + + private BraintrustTracing() {} } diff --git a/src/test/java/dev/braintrust/config/BraintrustConfigTest.java b/src/test/java/dev/braintrust/config/BraintrustConfigTest.java index 18204bf..52daef7 100644 --- a/src/test/java/dev/braintrust/config/BraintrustConfigTest.java +++ b/src/test/java/dev/braintrust/config/BraintrustConfigTest.java @@ -7,10 +7,14 @@ class BraintrustConfigTest { @Test void parentDefaultsToProjectName() { - var defaultConfig = BraintrustConfig.of("BRAINTRUST_API_KEY", "foobar"); + var defaultConfig = + BraintrustConfig.of( + "BRAINTRUST_API_KEY", + "foobar", + "BRAINTRUST_DEFAULT_PROJECT_NAME", + "proj-name"); assertEquals( - "project_name:" + defaultConfig.defaultProjectName().orElseThrow(), - defaultConfig.getBraintrustParentValue().orElseThrow()); + "project_name:proj-name", defaultConfig.getBraintrustParentValue().orElseThrow()); } @Test @@ -18,9 +22,38 @@ void parentUsesProjectId() { var defaultConfig = BraintrustConfig.of( "BRAINTRUST_API_KEY", "foobar", + "BRAINTRUST_DEFAULT_PROJECT_NAME", "proj-name", "BRAINTRUST_DEFAULT_PROJECT_ID", "12345"); assertEquals( "project_id:" + defaultConfig.defaultProjectId().orElseThrow(), defaultConfig.getBraintrustParentValue().orElseThrow()); } + + @Test + void projectUriFetching() { + var projectName = "some project"; + var projectId = "123456"; + var orgInfo = + new dev.braintrust.api.BraintrustApiClient.OrganizationInfo("org_123", "Test Org"); + var project = + new dev.braintrust.api.BraintrustApiClient.Project( + projectId, + projectName, + "org_123", + "2023-01-01T00:00:00Z", + "2023-01-01T00:00:00Z"); + var orgAndProjectInfo = + new dev.braintrust.api.BraintrustApiClient.OrganizationAndProjectInfo( + orgInfo, project); + var client = new dev.braintrust.api.BraintrustApiClient.InMemoryImpl(orgAndProjectInfo); + var config = + BraintrustConfig.of( + "BRAINTRUST_API_KEY", "foobar", + "BRAINTRUST_API_URL", "https://api.braintrust.dev", + "BRAINTRUST_APP_URL", "https://braintrust.dev", + "BRAINTRUST_DEFAULT_PROJECT_ID", projectId); + assertEquals( + "https://braintrust.dev/app/Test%20Org/p/some%20project", + config.fetchProjectURI(client).toASCIIString()); + } } diff --git a/src/test/java/dev/braintrust/eval/EvalTest.java b/src/test/java/dev/braintrust/eval/EvalTest.java new file mode 100644 index 0000000..2a1f823 --- /dev/null +++ b/src/test/java/dev/braintrust/eval/EvalTest.java @@ -0,0 +1,100 @@ +package dev.braintrust.eval; + +import static org.junit.jupiter.api.Assertions.*; + +import dev.braintrust.api.BraintrustApiClient; +import dev.braintrust.config.BraintrustConfig; +import dev.braintrust.spec.SdkSpec; +import dev.braintrust.trace.BraintrustTracing; +import dev.braintrust.trace.BraintrustTracingTest; +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import java.util.concurrent.TimeUnit; +import java.util.function.Function; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class EvalTest { + + @BeforeEach + void beforeEach() { + GlobalOpenTelemetry.resetForTest(); + BraintrustTracingTest.getExportedBraintrustSpans().clear(); + } + + @Test + public void evalOtelTraceWithProperAttributes() { + var projectId = "1234"; + var projectName = "proj-name"; + var experimentName = "unit-test-eval"; + var config = + BraintrustConfig.of( + "BRAINTRUST_API_KEY", "foobar", + "BRAINTRUST_TEST_JAVA_EXPORT_SPANS_IN_MEMORY", "true", + "BRAINTRUST_API_URL", "https://api.braintrust.dev", + "BRAINTRUST_APP_URL", "https://braintrust.dev", + "BRAINTRUST_DEFAULT_PROJECT_NAME", projectName); + var apiClient = createApiClient(projectId, projectName); + var sdk = (OpenTelemetrySdk) BraintrustTracing.of(config, true); + + var eval = + Eval.builder() + .name(experimentName) + .tracer(BraintrustTracing.getTracer(sdk)) + .config(config) + .apiClient(apiClient) + .cases( + EvalCase.of("strawberry", "fruit"), + EvalCase.of("asparagus", "vegetable")) + .task((Function) food -> "fruit") + .scorers( + Scorer.of( + "fruit_scorer", + result -> "fruit".equals(result) ? 1.0 : 0.0), + Scorer.of( + "vegetable_scorer", + result -> "vegetable".equals(result) ? 1.0 : 0.0)) + .build(); + var result = eval.run(); + assertEquals( + "https://braintrust.dev/app/Test%20Org/p/proj-name/experiments/unit-test-eval", + result.getExperimentUrl()); + assertTrue(sdk.getSdkTracerProvider().forceFlush().join(10, TimeUnit.SECONDS).isSuccess()); + assertEquals(1, BraintrustTracingTest.getExportedBraintrustSpans().size()); + var experiment = + apiClient.getOrCreateExperiment( + new BraintrustApiClient.CreateExperimentRequest(projectId, experimentName)); + var evalSpans = + BraintrustTracingTest.getExportedBraintrustSpans() + .get("experiment_id:" + experiment.id()); + assertNotNull(evalSpans); + assertNotEquals(0, evalSpans.size()); + evalSpans.forEach( + span -> { + var parent = + span.getAttributes() + .get(AttributeKey.stringKey(SdkSpec.Attributes.PARENT)); + assertEquals( + "experiment_id:" + experiment.id(), + parent, + "all eval spans must set the parent to the experiment id"); + }); + } + + private BraintrustApiClient createApiClient(String projectId, String projectName) { + var orgInfo = + new dev.braintrust.api.BraintrustApiClient.OrganizationInfo("org_123", "Test Org"); + var project = + new dev.braintrust.api.BraintrustApiClient.Project( + projectId, + projectName, + "org_123", + "2023-01-01T00:00:00Z", + "2023-01-01T00:00:00Z"); + var orgAndProjectInfo = + new dev.braintrust.api.BraintrustApiClient.OrganizationAndProjectInfo( + orgInfo, project); + return new dev.braintrust.api.BraintrustApiClient.InMemoryImpl(orgAndProjectInfo); + } +} diff --git a/src/test/java/dev/braintrust/trace/BraintrustTracingTest.java b/src/test/java/dev/braintrust/trace/BraintrustTracingTest.java new file mode 100644 index 0000000..d3c3d4f --- /dev/null +++ b/src/test/java/dev/braintrust/trace/BraintrustTracingTest.java @@ -0,0 +1,83 @@ +package dev.braintrust.trace; + +import static org.junit.jupiter.api.Assertions.*; + +import dev.braintrust.config.BraintrustConfig; +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.logs.SdkLoggerProvider; +import io.opentelemetry.sdk.metrics.SdkMeterProvider; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class BraintrustTracingTest { + public static Map> getExportedBraintrustSpans() { + return BraintrustSpanExporter.SPANS_EXPORTED; + } + + private final BraintrustConfig config = + BraintrustConfig.of( + "BRAINTRUST_API_KEY", "foobar", + "BRAINTRUST_TEST_JAVA_EXPORT_SPANS_IN_MEMORY", "true"); + + @BeforeEach + void beforeEach() { + GlobalOpenTelemetry.resetForTest(); + getExportedBraintrustSpans().clear(); + } + + @Test + void globalBTTracing() { + var sdk = (OpenTelemetrySdk) BraintrustTracing.of(config, true); + doSimpleOtelTrace(BraintrustTracing.getTracer()); + assertTrue(sdk.getSdkTracerProvider().forceFlush().join(10, TimeUnit.SECONDS).isSuccess()); + assertEquals(1, getExportedBraintrustSpans().size()); + var spanData = + getExportedBraintrustSpans().get(config.getBraintrustParentValue().orElseThrow()); + assertNotNull(spanData); + assertEquals(1, spanData.size()); + assertEquals( + true, spanData.get(0).getAttributes().get(AttributeKey.booleanKey("unit-test"))); + } + + @Test + void customBTTracing() { + var tracerBuilder = SdkTracerProvider.builder(); + var loggerBuilder = SdkLoggerProvider.builder(); + var meterBuilder = SdkMeterProvider.builder(); + BraintrustTracing.enable(config, tracerBuilder, loggerBuilder, meterBuilder); + var sdk = + OpenTelemetrySdk.builder() + .setTracerProvider(tracerBuilder.build()) + .setLoggerProvider(loggerBuilder.build()) + .setMeterProvider(meterBuilder.build()) + .build(); + GlobalOpenTelemetry.set(sdk); + doSimpleOtelTrace(sdk.getTracer("some-instrumentation")); + assertTrue(sdk.getSdkTracerProvider().forceFlush().join(10, TimeUnit.SECONDS).isSuccess()); + assertEquals(1, getExportedBraintrustSpans().size()); + var spanData = + getExportedBraintrustSpans().get(config.getBraintrustParentValue().orElseThrow()); + assertNotNull(spanData); + assertEquals(1, spanData.size()); + assertEquals( + true, spanData.get(0).getAttributes().get(AttributeKey.booleanKey("unit-test"))); + } + + private void doSimpleOtelTrace(Tracer tracer) { + // use tracer to create a simple trace with a root span and a child span + var span = tracer.spanBuilder("unit-test-root").startSpan(); + try (var ignored = span.makeCurrent()) { + span.setAttribute("unit-test", true); + } finally { + span.end(); + } + } +}