From d8e783a76c56c8d871999401d99b17f558f19595 Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Fri, 24 Oct 2025 15:03:44 -0600 Subject: [PATCH 1/2] add prompt fetching to gradle examples --- examples/build.gradle | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/examples/build.gradle b/examples/build.gradle index e6b5d931..932fb303 100644 --- a/examples/build.gradle +++ b/examples/build.gradle @@ -90,3 +90,18 @@ task runAnthropicInstrumentation(type: JavaExec) { suspend = false } } + + +task runPromptFetching(type: JavaExec) { + group = 'Braintrust SDK Examples' + description = 'Run the prompt fetching example' + classpath = sourceSets.main.runtimeClasspath + mainClass = 'dev.braintrust.examples.PromptFetchingExample' + systemProperty 'org.slf4j.simpleLogger.log.dev.braintrust', braintrustLogLevel + debugOptions { + enabled = true + port = 5566 + server = true + suspend = false + } +} From 4434cd146180f655e43d41de7a4117116aa95d31 Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Fri, 24 Oct 2025 00:48:43 -0600 Subject: [PATCH 2/2] unit test harness and refactor out span exporter hack --- src/main/java/dev/braintrust/Braintrust.java | 33 +++-- .../java/dev/braintrust/BraintrustUtils.java | 29 ++++ .../braintrust/api/BraintrustApiClient.java | 102 +++++++++++++- .../braintrust/config/BraintrustConfig.java | 31 +---- src/main/java/dev/braintrust/eval/Eval.java | 19 +-- .../trace/BraintrustSpanExporter.java | 8 +- .../java/dev/braintrust/BraintrustTest.java | 16 ++- .../dev/braintrust/BraintrustUtilsTest.java | 21 +++ src/test/java/dev/braintrust/TestHarness.java | 130 ++++++++++++++++++ .../config/BraintrustConfigTest.java | 41 ------ .../java/dev/braintrust/eval/EvalTest.java | 67 +++------ .../anthropic/BraintrustAnthropicTest.java | 54 ++------ .../openai/BraintrustOpenAITest.java | 72 +++------- .../trace/BraintrustTracingTest.java | 79 +++++------ 14 files changed, 407 insertions(+), 295 deletions(-) create mode 100644 src/main/java/dev/braintrust/BraintrustUtils.java create mode 100644 src/test/java/dev/braintrust/BraintrustUtilsTest.java create mode 100644 src/test/java/dev/braintrust/TestHarness.java diff --git a/src/main/java/dev/braintrust/Braintrust.java b/src/main/java/dev/braintrust/Braintrust.java index 0e5ea61b..d0bc5b33 100644 --- a/src/main/java/dev/braintrust/Braintrust.java +++ b/src/main/java/dev/braintrust/Braintrust.java @@ -37,14 +37,14 @@ @Slf4j public class Braintrust { private static final String SDK_VERSION = SDKMain.loadVersionFromProperties(); - private static final AtomicReference instance = new AtomicReference<>(); + private static final AtomicReference INSTANCE = new AtomicReference<>(); /** * get or create the global braintrust instance. Most users will want to use this method to * access the Braintrust SDK. */ public static Braintrust get() { - var current = instance.get(); + var current = INSTANCE.get(); if (null == current) { return get(BraintrustConfig.fromEnvironment()); } else { @@ -54,18 +54,34 @@ public static Braintrust get() { /** get or create the global braintrust instance from the given config */ public static Braintrust get(BraintrustConfig config) { - var current = instance.get(); + var current = INSTANCE.get(); if (null == current) { - var success = instance.compareAndSet(null, of(config)); + return set(of(config)); + } else { + return current; + } + } + + static Braintrust set(Braintrust braintrust) { + var current = INSTANCE.get(); + if (null == current) { + var success = INSTANCE.compareAndSet(null, braintrust); if (success) { log.info("initialized global Braintrust sdk {}", SDK_VERSION); + } else { + throw new RuntimeException("set must only be called once"); } - return instance.get(); + return braintrust; } else { return current; } } + /** clear global braintrust instance. Only used for testing */ + static void resetForTest() { + INSTANCE.set(null); + } + /** Create a new Braintrust instance from the given config */ public static Braintrust of(BraintrustConfig config) { BraintrustApiClient apiClient = BraintrustApiClient.of(config); @@ -85,7 +101,7 @@ public static Braintrust of(BraintrustConfig config) { @Accessors(fluent = true) private final BraintrustPromptLoader promptLoader; - private Braintrust( + Braintrust( BraintrustConfig config, BraintrustApiClient apiClient, BraintrustPromptLoader promptLoader) { @@ -94,9 +110,10 @@ private Braintrust( this.promptLoader = promptLoader; } + /** the the URI to the configured braintrust org and project */ public URI projectUri() { - // TODO cache? - return config.fetchProjectURI(); + return BraintrustUtils.createProjectURI( + config.appUrl(), apiClient().getOrCreateProjectAndOrgInfo(config())); } /** diff --git a/src/main/java/dev/braintrust/BraintrustUtils.java b/src/main/java/dev/braintrust/BraintrustUtils.java new file mode 100644 index 00000000..1924952a --- /dev/null +++ b/src/main/java/dev/braintrust/BraintrustUtils.java @@ -0,0 +1,29 @@ +package dev.braintrust; + +import dev.braintrust.api.BraintrustApiClient; +import java.net.URI; +import java.net.URISyntaxException; + +public class BraintrustUtils { + /** construct a URI to link to a specific braintrust project within an org */ + public static URI createProjectURI( + String appUrl, BraintrustApiClient.OrganizationAndProjectInfo orgAndProject) { + try { + var baseURI = new URI(appUrl); + var path = + "/app/%s/p/%s" + .formatted( + orgAndProject.orgInfo().name(), orgAndProject.project().name()); + return new URI( + baseURI.getScheme(), + baseURI.getUserInfo(), + baseURI.getHost(), + baseURI.getPort(), + path, + null, + null); + } catch (URISyntaxException e) { + throw new RuntimeException(e); + } + } +} diff --git a/src/main/java/dev/braintrust/api/BraintrustApiClient.java b/src/main/java/dev/braintrust/api/BraintrustApiClient.java index 74d22911..51561d47 100644 --- a/src/main/java/dev/braintrust/api/BraintrustApiClient.java +++ b/src/main/java/dev/braintrust/api/BraintrustApiClient.java @@ -41,6 +41,10 @@ public interface BraintrustApiClient { /** Get project and org info for the given project ID */ Optional getProjectAndOrgInfo(String projectId); + // TODO: cache project+org info? + /** Get project and org info for the given config. Creating them if necessary */ + OrganizationAndProjectInfo getOrCreateProjectAndOrgInfo(BraintrustConfig config); + /** Get a prompt by slug and optional version */ Optional getPrompt( @Nonnull String projectName, @Nonnull String slug, @Nullable String version); @@ -145,6 +149,43 @@ public Optional getProjectAndOrgInfo(String projectI return Optional.of(new OrganizationAndProjectInfo(orgInfo, project)); } + @Override + public OrganizationAndProjectInfo getOrCreateProjectAndOrgInfo(BraintrustConfig config) { + // Get or create project based on config + Project project; + if (config.defaultProjectId().isPresent()) { + var projectId = config.defaultProjectId().get(); + project = + getProject(projectId) + .orElseThrow( + () -> + new ApiException( + "Project with ID '" + + projectId + + "' not found")); + } else if (config.defaultProjectName().isPresent()) { + var projectName = config.defaultProjectName().get(); + project = getOrCreateProject(projectName); + } else { + throw new ApiException( + "Either project ID or project name must be provided in config"); + } + + // Fetch organization info + OrganizationInfo orgInfo = null; + for (var org : login().orgInfo()) { + if (project.orgId().equalsIgnoreCase(org.id())) { + orgInfo = org; + break; + } + } + if (null == orgInfo) { + throw new ApiException("Unable to find organization for project: " + project.id()); + } + + return new OrganizationAndProjectInfo(orgInfo, project); + } + @Override public Optional getPrompt( @Nonnull String projectName, @Nonnull String slug, @Nullable String version) { @@ -293,13 +334,14 @@ class InMemoryImpl implements BraintrustApiClient { private final List prompts = new ArrayList<>(); public InMemoryImpl(OrganizationAndProjectInfo... organizationAndProjectInfos) { - this.organizationAndProjectInfos = List.of(organizationAndProjectInfos); + this.organizationAndProjectInfos = + new ArrayList<>(List.of(organizationAndProjectInfos)); } public InMemoryImpl( List organizationAndProjectInfos, List prompts) { - this.organizationAndProjectInfos = organizationAndProjectInfos; + this.organizationAndProjectInfos = new ArrayList<>(organizationAndProjectInfos); this.prompts.addAll(prompts); } @@ -311,11 +353,24 @@ public Project getOrCreateProject(String projectName) { return orgAndProject.project(); } } - throw new RuntimeException( - "Project '" - + projectName - + "' not found in test data. Please add it to the InMemoryImpl" - + " constructor."); + + // Create new project if not found + var defaultOrgInfo = + organizationAndProjectInfos.isEmpty() + ? new OrganizationInfo("default-org-id", "Default Organization") + : organizationAndProjectInfos.get(0).orgInfo(); + + var newProject = + new Project( + "project-" + UUID.randomUUID().toString(), + projectName, + defaultOrgInfo.id(), + java.time.Instant.now().toString(), + java.time.Instant.now().toString()); + + organizationAndProjectInfos.add( + new OrganizationAndProjectInfo(defaultOrgInfo, newProject)); + return newProject; } @Override @@ -357,6 +412,39 @@ public Optional getProjectAndOrgInfo(String projectI .findFirst(); } + @Override + public OrganizationAndProjectInfo getOrCreateProjectAndOrgInfo(BraintrustConfig config) { + // Get or create project based on config + Project project; + if (config.defaultProjectId().isPresent()) { + var projectId = config.defaultProjectId().get(); + project = + getProject(projectId) + .orElseThrow( + () -> + new ApiException( + "Project with ID '" + + projectId + + "' not found")); + } else if (config.defaultProjectName().isPresent()) { + var projectName = config.defaultProjectName().get(); + project = getOrCreateProject(projectName); + } else { + throw new ApiException( + "Either project ID or project name must be provided in config"); + } + + // Find the organization info for this project + return organizationAndProjectInfos.stream() + .filter(info -> info.project().id().equals(project.id())) + .findFirst() + .orElseThrow( + () -> + new ApiException( + "Unable to find organization for project: " + + project.id())); + } + @Override public Optional getPrompt( @Nonnull String projectName, @Nonnull String slug, @Nullable String version) { diff --git a/src/main/java/dev/braintrust/config/BraintrustConfig.java b/src/main/java/dev/braintrust/config/BraintrustConfig.java index 8eaf666f..2ae84b32 100644 --- a/src/main/java/dev/braintrust/config/BraintrustConfig.java +++ b/src/main/java/dev/braintrust/config/BraintrustConfig.java @@ -1,8 +1,9 @@ package dev.braintrust.config; +import dev.braintrust.Braintrust; +import dev.braintrust.BraintrustUtils; import dev.braintrust.api.BraintrustApiClient; import java.net.URI; -import java.net.URISyntaxException; import java.time.Duration; import java.util.HashMap; import java.util.Map; @@ -87,29 +88,11 @@ public Optional getBraintrustParentValue() { } } - /** fetch all project info and IDs from the braintrust api */ + /** Deprecated. Please use {@link Braintrust#projectUri()} instead */ + @Deprecated public URI fetchProjectURI() { - return fetchProjectURI(BraintrustApiClient.of(this)); - } - - URI fetchProjectURI(BraintrustApiClient client) { - try { - var orgAndProject = client.getProjectAndOrgInfo().orElseThrow(); - var baseURI = new URI(appUrl()); - var path = - "/app/%s/p/%s" - .formatted( - orgAndProject.orgInfo().name(), orgAndProject.project().name()); - return new URI( - baseURI.getScheme(), - baseURI.getUserInfo(), - baseURI.getHost(), - baseURI.getPort(), - path, - null, - null); - } catch (URISyntaxException e) { - throw new RuntimeException(e); - } + var client = BraintrustApiClient.of(this); + var orgAndProject = client.getProjectAndOrgInfo().orElseThrow(); + return BraintrustUtils.createProjectURI(appUrl(), orgAndProject); } } diff --git a/src/main/java/dev/braintrust/eval/Eval.java b/src/main/java/dev/braintrust/eval/Eval.java index a45baa25..11b92a5f 100644 --- a/src/main/java/dev/braintrust/eval/Eval.java +++ b/src/main/java/dev/braintrust/eval/Eval.java @@ -2,6 +2,7 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; +import dev.braintrust.BraintrustUtils; import dev.braintrust.api.BraintrustApiClient; import dev.braintrust.config.BraintrustConfig; import dev.braintrust.trace.BraintrustContext; @@ -9,7 +10,6 @@ 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; @@ -158,19 +158,12 @@ public class Result { @SneakyThrows private Result() { - var baseURI = new URI(config.appUrl()); this.experimentUrl = - new URI( - baseURI.getScheme(), - baseURI.getHost(), - "/app/" - + orgAndProject.orgInfo().name() - + "/p/" - + orgAndProject.project().name() - + "/experiments/" - + experimentName, - null) - .toASCIIString(); + "%s/experiments/%s" + .formatted( + BraintrustUtils.createProjectURI(config.appUrl(), orgAndProject) + .toASCIIString(), + experimentName); } public String createReportString() { diff --git a/src/main/java/dev/braintrust/trace/BraintrustSpanExporter.java b/src/main/java/dev/braintrust/trace/BraintrustSpanExporter.java index 00c4aac0..83f4e775 100644 --- a/src/main/java/dev/braintrust/trace/BraintrustSpanExporter.java +++ b/src/main/java/dev/braintrust/trace/BraintrustSpanExporter.java @@ -9,7 +9,6 @@ import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.CopyOnWriteArrayList; import java.util.stream.Collectors; import lombok.extern.slf4j.Slf4j; @@ -19,9 +18,6 @@ */ @Slf4j class BraintrustSpanExporter implements SpanExporter { - /** Only used in unit tests. */ - static final Map> SPANS_EXPORTED = new ConcurrentHashMap<>(); - private final BraintrustConfig config; private final String tracesEndpoint; private final Map exporterCache = new ConcurrentHashMap<>(); @@ -90,8 +86,8 @@ private CompletableResultCode exportWithParent(String parent, List spa }); if (config.exportSpansInMemoryForUnitTest()) { - SPANS_EXPORTED.putIfAbsent(parent, new CopyOnWriteArrayList<>()); - SPANS_EXPORTED.get(parent).addAll(spans); + // unit test harness hooks up an in-memory exporter so we don't need to do anything + // here return CompletableResultCode.ofSuccess(); } else { var result = exporter.export(spans); diff --git a/src/test/java/dev/braintrust/BraintrustTest.java b/src/test/java/dev/braintrust/BraintrustTest.java index 55e1f28f..1dd7680c 100644 --- a/src/test/java/dev/braintrust/BraintrustTest.java +++ b/src/test/java/dev/braintrust/BraintrustTest.java @@ -4,6 +4,7 @@ import dev.braintrust.config.BraintrustConfig; import io.opentelemetry.api.GlobalOpenTelemetry; +import java.net.URI; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -12,7 +13,6 @@ public class BraintrustTest { @BeforeEach void setUp() { - // Reset global OpenTelemetry before each test GlobalOpenTelemetry.resetForTest(); } @@ -164,4 +164,18 @@ void testGetReturnsSameInstanceAfterInitialization() { assertSame(braintrust1, braintrust2); assertSame(braintrust1, braintrust3); } + + @Test + public void testProjectURI() { + var harness = TestHarness.setup(); + var config = harness.braintrust().config(); + assertEquals( + URI.create( + "%s/app/%s/p/%s" + .formatted( + config.appUrl(), + TestHarness.defaultOrgName().replace(" ", "%20"), + TestHarness.defaultProjectName().replace(" ", "%20"))), + harness.braintrust().projectUri()); + } } diff --git a/src/test/java/dev/braintrust/BraintrustUtilsTest.java b/src/test/java/dev/braintrust/BraintrustUtilsTest.java new file mode 100644 index 00000000..fcf5468a --- /dev/null +++ b/src/test/java/dev/braintrust/BraintrustUtilsTest.java @@ -0,0 +1,21 @@ +package dev.braintrust; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import dev.braintrust.api.BraintrustApiClient; +import java.net.URI; +import org.junit.jupiter.api.Test; + +public class BraintrustUtilsTest { + @Test + public void testBuildProjectUri() { + var orgAndProject = + new BraintrustApiClient.OrganizationAndProjectInfo( + new BraintrustApiClient.OrganizationInfo("123", "some org"), + new BraintrustApiClient.Project( + "456", "some project", "123", "doesntmatter", "doesntmatter")); + assertEquals( + URI.create("http://someserver:3009/app/some%20org/p/some%20project"), + BraintrustUtils.createProjectURI("http://someserver:3009/", orgAndProject)); + } +} diff --git a/src/test/java/dev/braintrust/TestHarness.java b/src/test/java/dev/braintrust/TestHarness.java new file mode 100644 index 00000000..033563db --- /dev/null +++ b/src/test/java/dev/braintrust/TestHarness.java @@ -0,0 +1,130 @@ +package dev.braintrust; + +import static org.junit.jupiter.api.Assertions.assertTrue; + +import dev.braintrust.api.BraintrustApiClient; +import dev.braintrust.config.BraintrustConfig; +import dev.braintrust.prompt.BraintrustPromptLoader; +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.logs.SdkLoggerProvider; +import io.opentelemetry.sdk.metrics.SdkMeterProvider; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import java.util.List; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import javax.annotation.Nonnull; +import lombok.Getter; +import lombok.experimental.Accessors; + +public class TestHarness { + + public static TestHarness setup() { + return setup(createTestConfig()); + } + + public static synchronized TestHarness setup(BraintrustConfig config) { + GlobalOpenTelemetry.resetForTest(); + Braintrust.resetForTest(); + + var apiClient = createApiClient(); + Braintrust braintrust = + Braintrust.set( + new Braintrust( + config, + createApiClient(), + BraintrustPromptLoader.of(config, apiClient))); + var harness = new TestHarness(braintrust); + INSTANCE.set(harness); + GlobalOpenTelemetry.set(harness.openTelemetry()); + return harness; + } + + @Getter + @Accessors(fluent = true) + private static final String defaultProjectId = "01234"; + + @Getter + @Accessors(fluent = true) + private static final String defaultProjectName = "Unit Test"; + + @Getter + @Accessors(fluent = true) + private static final String defaultOrgId = "567890"; + + @Getter + @Accessors(fluent = true) + private static final String defaultOrgName = "Test Org"; + + private static final AtomicReference INSTANCE = new AtomicReference<>(); + + @Getter + @Accessors(fluent = true) + private final OpenTelemetrySdk openTelemetry; + + @Getter + @Accessors(fluent = true) + private final Braintrust braintrust; + + private final @Nonnull InMemorySpanExporter spanExporter; + + private TestHarness(@Nonnull Braintrust braintrust) { + this.braintrust = braintrust; + + var tracerBuilder = SdkTracerProvider.builder(); + this.spanExporter = InMemorySpanExporter.create(); + var loggerBuilder = SdkLoggerProvider.builder(); + var meterBuilder = SdkMeterProvider.builder(); + braintrust.openTelemetryEnable(tracerBuilder, loggerBuilder, meterBuilder); + // Add the in-memory span exporter for testing + tracerBuilder.addSpanProcessor(SimpleSpanProcessor.create(this.spanExporter)); + var openTelemetry = + OpenTelemetrySdk.builder() + .setTracerProvider(tracerBuilder.build()) + .setLoggerProvider(loggerBuilder.build()) + .setMeterProvider(meterBuilder.build()) + .build(); + this.openTelemetry = openTelemetry; + } + + /** flush all pending spans and return all spans which have been exported so far */ + public List awaitExportedSpans() { + assertTrue( + openTelemetry + .getSdkTracerProvider() + .forceFlush() + .join(10, TimeUnit.SECONDS) + .isSuccess()); + return spanExporter.getFinishedSpanItems(); + } + + private static BraintrustApiClient.InMemoryImpl createApiClient() { + var orgInfo = + new dev.braintrust.api.BraintrustApiClient.OrganizationInfo( + defaultOrgId, defaultOrgName); + var project = + new dev.braintrust.api.BraintrustApiClient.Project( + defaultProjectId, + defaultProjectName, + "unit_test_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); + } + + public static BraintrustConfig createTestConfig() { + return BraintrustConfig.of( + "BRAINTRUST_API_KEY", "foobar", + "BRAINTRUST_JAVA_EXPORT_SPANS_IN_MEMORY_FOR_UNIT_TEST", "true", + // NOTE: testhost is not real, just a placeholder value + "BRAINTRUST_API_URL", "https://testhost:8000", + "BRAINTRUST_APP_URL", "https://testhost:3000", + "BRAINTRUST_DEFAULT_PROJECT_NAME", defaultProjectName()); + } +} diff --git a/src/test/java/dev/braintrust/config/BraintrustConfigTest.java b/src/test/java/dev/braintrust/config/BraintrustConfigTest.java index 945c7932..d93ce4e1 100644 --- a/src/test/java/dev/braintrust/config/BraintrustConfigTest.java +++ b/src/test/java/dev/braintrust/config/BraintrustConfigTest.java @@ -28,45 +28,4 @@ void parentUsesProjectId() { "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); - { // standard braintrust.dev - 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()); - } - { // self-hosted with a non-standard port - var config = - BraintrustConfig.of( - "BRAINTRUST_API_KEY", "foobar", - "BRAINTRUST_API_URL", "http://someselfhost:8080", - "BRAINTRUST_APP_URL", "http://someselfhost:3001", - "BRAINTRUST_DEFAULT_PROJECT_ID", projectId); - assertEquals( - "http://someselfhost:3001/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 index 027989fc..873521de 100644 --- a/src/test/java/dev/braintrust/eval/EvalTest.java +++ b/src/test/java/dev/braintrust/eval/EvalTest.java @@ -2,47 +2,31 @@ import static org.junit.jupiter.api.Assertions.*; +import dev.braintrust.TestHarness; import dev.braintrust.api.BraintrustApiClient; -import dev.braintrust.config.BraintrustConfig; 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 { + private TestHarness testHarness; @BeforeEach void beforeEach() { - GlobalOpenTelemetry.resetForTest(); - BraintrustTracingTest.getExportedBraintrustSpans().clear(); + testHarness = TestHarness.setup(); } @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_JAVA_EXPORT_SPANS_IN_MEMORY_FOR_UNIT_TEST", "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() + testHarness + .braintrust() + .evalBuilder() .name(experimentName) - .tracer(BraintrustTracing.getTracer(sdk)) - .config(config) - .apiClient(apiClient) .cases( EvalCase.of("strawberry", "fruit"), EvalCase.of("asparagus", "vegetable")) @@ -57,19 +41,20 @@ public void evalOtelTraceWithProperAttributes() { .build(); var result = eval.run(); assertEquals( - "https://braintrust.dev/app/Test%20Org/p/proj-name/experiments/unit-test-eval", + "%s/experiments/%s" + .formatted(testHarness.braintrust().projectUri(), experimentName), result.getExperimentUrl()); - assertTrue(sdk.getSdkTracerProvider().forceFlush().join(10, TimeUnit.SECONDS).isSuccess()); - assertEquals(1, BraintrustTracingTest.getExportedBraintrustSpans().size()); + var spans = testHarness.awaitExportedSpans(); + assertEquals(6, spans.size(), "each eval case should make three spans"); + // TODO: assert each case makes the expected spans 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( + testHarness + .braintrust() + .apiClient() + .getOrCreateExperiment( + new BraintrustApiClient.CreateExperimentRequest( + TestHarness.defaultProjectId(), experimentName)); + spans.forEach( span -> { var parent = span.getAttributes() @@ -80,20 +65,4 @@ public void evalOtelTraceWithProperAttributes() { "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/instrumentation/anthropic/BraintrustAnthropicTest.java b/src/test/java/dev/braintrust/instrumentation/anthropic/BraintrustAnthropicTest.java index 718c9b41..2ab80b7f 100644 --- a/src/test/java/dev/braintrust/instrumentation/anthropic/BraintrustAnthropicTest.java +++ b/src/test/java/dev/braintrust/instrumentation/anthropic/BraintrustAnthropicTest.java @@ -2,7 +2,6 @@ import static com.github.tomakehurst.wiremock.client.WireMock.*; import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig; -import static dev.braintrust.trace.BraintrustTracingTest.getExportedBraintrustSpans; import static org.junit.jupiter.api.Assertions.*; import com.anthropic.client.AnthropicClient; @@ -11,12 +10,8 @@ import com.anthropic.models.messages.Model; import com.fasterxml.jackson.databind.ObjectMapper; import com.github.tomakehurst.wiremock.junit5.WireMockExtension; -import dev.braintrust.config.BraintrustConfig; -import dev.braintrust.trace.BraintrustTracing; -import io.opentelemetry.api.GlobalOpenTelemetry; +import dev.braintrust.TestHarness; import io.opentelemetry.api.common.AttributeKey; -import io.opentelemetry.sdk.OpenTelemetrySdk; -import java.util.concurrent.TimeUnit; import lombok.SneakyThrows; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -29,16 +24,11 @@ public class BraintrustAnthropicTest { static WireMockExtension wireMock = WireMockExtension.newInstance().options(wireMockConfig().dynamicPort()).build(); - private final BraintrustConfig config = - BraintrustConfig.of( - "BRAINTRUST_API_KEY", "foobar", - "BRAINTRUST_DEFAULT_PROJECT_NAME", "unit-test-project", - "BRAINTRUST_JAVA_EXPORT_SPANS_IN_MEMORY_FOR_UNIT_TEST", "true"); + private TestHarness testHarness; @BeforeEach void beforeEach() { - GlobalOpenTelemetry.resetForTest(); - getExportedBraintrustSpans().clear(); + testHarness = TestHarness.setup(); wireMock.resetAll(); } @@ -74,8 +64,6 @@ void testWrapAnthropic() { } """))); - var openTelemetry = (OpenTelemetrySdk) BraintrustTracing.of(config, true); - // Create Anthropic client pointing to WireMock server AnthropicClient anthropicClient = AnthropicOkHttpClient.builder() @@ -84,7 +72,7 @@ void testWrapAnthropic() { .build(); // Wrap with Braintrust instrumentation - anthropicClient = BraintrustAnthropic.wrap(openTelemetry, anthropicClient); + anthropicClient = BraintrustAnthropic.wrap(testHarness.openTelemetry(), anthropicClient); var request = MessageCreateParams.builder() @@ -106,17 +94,9 @@ void testWrapAnthropic() { assertEquals("The capital of France is Paris.", contentBlock.asText().text()); // Verify spans were exported - assertTrue( - openTelemetry - .getSdkTracerProvider() - .forceFlush() - .join(10, TimeUnit.SECONDS) - .isSuccess()); - var spanData = - getExportedBraintrustSpans().get(config.getBraintrustParentValue().orElseThrow()); - assertNotNull(spanData); - assertEquals(1, spanData.size()); - var span = spanData.get(0); + var spans = testHarness.awaitExportedSpans(); + assertEquals(1, spans.size()); + var span = spans.get(0); // Verify standard GenAI attributes assertEquals( @@ -138,7 +118,7 @@ void testWrapAnthropic() { "msg_test123", span.getAttributes().get(AttributeKey.stringKey("gen_ai.response.id"))); assertEquals( - "project_name:unit-test-project", + "project_name:" + TestHarness.defaultProjectName(), span.getAttributes().get(AttributeKey.stringKey("braintrust.parent"))); assertEquals( 20L, span.getAttributes().get(AttributeKey.longKey("gen_ai.usage.input_tokens"))); @@ -227,8 +207,6 @@ void testWrapAnthropicStreaming() { .withHeader("Content-Type", "text/event-stream") .withBody(streamingResponse))); - var openTelemetry = (OpenTelemetrySdk) BraintrustTracing.of(config, true); - // Create Anthropic client pointing to WireMock server AnthropicClient anthropicClient = AnthropicOkHttpClient.builder() @@ -237,7 +215,7 @@ void testWrapAnthropicStreaming() { .build(); // Wrap with Braintrust instrumentation - anthropicClient = BraintrustAnthropic.wrap(openTelemetry, anthropicClient); + anthropicClient = BraintrustAnthropic.wrap(testHarness.openTelemetry(), anthropicClient); var request = MessageCreateParams.builder() @@ -268,17 +246,9 @@ void testWrapAnthropicStreaming() { wireMock.verify(1, postRequestedFor(urlEqualTo("/v1/messages"))); // Verify spans were exported - assertTrue( - openTelemetry - .getSdkTracerProvider() - .forceFlush() - .join(10, TimeUnit.SECONDS) - .isSuccess()); - var spanData = - getExportedBraintrustSpans().get(config.getBraintrustParentValue().orElseThrow()); - assertNotNull(spanData); - assertEquals(1, spanData.size()); - var span = spanData.get(0); + var spans = testHarness.awaitExportedSpans(); + assertEquals(1, spans.size()); + var span = spans.get(0); // Verify standard GenAI attributes assertEquals( diff --git a/src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java b/src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java index f3b21952..ab80bc65 100644 --- a/src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java +++ b/src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java @@ -2,7 +2,6 @@ import static com.github.tomakehurst.wiremock.client.WireMock.*; import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig; -import static dev.braintrust.trace.BraintrustTracingTest.getExportedBraintrustSpans; import static org.junit.jupiter.api.Assertions.*; import com.fasterxml.jackson.databind.ObjectMapper; @@ -11,19 +10,15 @@ import com.openai.client.okhttp.OpenAIOkHttpClient; import com.openai.models.ChatModel; import com.openai.models.chat.completions.*; +import dev.braintrust.TestHarness; import dev.braintrust.api.BraintrustApiClient; -import dev.braintrust.config.BraintrustConfig; import dev.braintrust.prompt.BraintrustPrompt; import dev.braintrust.trace.Base64Attachment; -import dev.braintrust.trace.BraintrustTracing; -import io.opentelemetry.api.GlobalOpenTelemetry; import io.opentelemetry.api.common.AttributeKey; -import io.opentelemetry.sdk.OpenTelemetrySdk; import java.util.Arrays; import java.util.List; import java.util.Map; import java.util.Optional; -import java.util.concurrent.TimeUnit; import lombok.SneakyThrows; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -37,16 +32,11 @@ public class BraintrustOpenAITest { static WireMockExtension wireMock = WireMockExtension.newInstance().options(wireMockConfig().dynamicPort()).build(); - private final BraintrustConfig config = - BraintrustConfig.of( - "BRAINTRUST_API_KEY", "foobar", - "BRAINTRUST_DEFAULT_PROJECT_NAME", "unit-test-project", - "BRAINTRUST_JAVA_EXPORT_SPANS_IN_MEMORY_FOR_UNIT_TEST", "true"); + private TestHarness testHarness; @BeforeEach void beforeEach() { - GlobalOpenTelemetry.resetForTest(); - getExportedBraintrustSpans().clear(); + testHarness = TestHarness.setup(); wireMock.resetAll(); } @@ -85,8 +75,6 @@ void testWrapOpenAi() { } """))); - var openTelemetry = (OpenTelemetrySdk) BraintrustTracing.of(config, true); - // Create OpenAI client pointing to WireMock server OpenAIClient openAIClient = OpenAIOkHttpClient.builder() @@ -95,7 +83,7 @@ void testWrapOpenAi() { .build(); // Wrap with Braintrust instrumentation - openAIClient = BraintrustOpenAI.wrapOpenAI(openTelemetry, openAIClient); + openAIClient = BraintrustOpenAI.wrapOpenAI(testHarness.openTelemetry(), openAIClient); var request = ChatCompletionCreateParams.builder() @@ -116,17 +104,9 @@ void testWrapOpenAi() { response.choices().get(0).message().content().get()); // Verify spans were exported - assertTrue( - openTelemetry - .getSdkTracerProvider() - .forceFlush() - .join(10, TimeUnit.SECONDS) - .isSuccess()); - var spanData = - getExportedBraintrustSpans().get(config.getBraintrustParentValue().orElseThrow()); - assertNotNull(spanData); - assertEquals(1, spanData.size()); - var span = spanData.get(0); + var spans = testHarness.awaitExportedSpans(); + assertEquals(1, spans.size()); + var span = spans.get(0); assertEquals("openai", span.getAttributes().get(AttributeKey.stringKey("gen_ai.system"))); assertEquals( @@ -152,7 +132,7 @@ void testWrapOpenAi() { + " is the capital of France?\"}]}]", span.getAttributes().get(AttributeKey.stringKey("gen_ai.input.messages"))); assertEquals( - "project_name:unit-test-project", + "project_name:" + TestHarness.defaultProjectName(), span.getAttributes().get(AttributeKey.stringKey("braintrust.parent"))); assertEquals( 20L, span.getAttributes().get(AttributeKey.longKey("gen_ai.usage.input_tokens"))); @@ -212,8 +192,6 @@ void testWrapOpenAiStreaming() { .withHeader("Content-Type", "text/event-stream") .withBody(streamingResponse))); - var openTelemetry = (OpenTelemetrySdk) BraintrustTracing.of(config, true); - // Create OpenAI client pointing to WireMock server OpenAIClient openAIClient = OpenAIOkHttpClient.builder() @@ -222,7 +200,7 @@ void testWrapOpenAiStreaming() { .build(); // Wrap with Braintrust instrumentation - openAIClient = BraintrustOpenAI.wrapOpenAI(openTelemetry, openAIClient); + openAIClient = BraintrustOpenAI.wrapOpenAI(testHarness.openTelemetry(), openAIClient); var request = ChatCompletionCreateParams.builder() @@ -255,17 +233,9 @@ void testWrapOpenAiStreaming() { wireMock.verify(1, postRequestedFor(urlEqualTo("/chat/completions"))); // Verify spans were exported - assertTrue( - openTelemetry - .getSdkTracerProvider() - .forceFlush() - .join(10, TimeUnit.SECONDS) - .isSuccess()); - var spanData = - getExportedBraintrustSpans().get(config.getBraintrustParentValue().orElseThrow()); - assertNotNull(spanData); - assertEquals(1, spanData.size()); - var span = spanData.get(0); + var spans = testHarness.awaitExportedSpans(); + assertEquals(1, spans.size()); + var span = spans.get(0); // Verify span attributes assertEquals("openai", span.getAttributes().get(AttributeKey.stringKey("gen_ai.system"))); @@ -337,8 +307,6 @@ void testWrapOpenAiWithImageAttachment() { } """))); - var openTelemetry = (OpenTelemetrySdk) BraintrustTracing.of(config, true); - // Create OpenAI client pointing to WireMock server OpenAIClient openAIClient = OpenAIOkHttpClient.builder() @@ -347,7 +315,7 @@ void testWrapOpenAiWithImageAttachment() { .build(); // Wrap with Braintrust instrumentation - openAIClient = BraintrustOpenAI.wrapOpenAI(openTelemetry, openAIClient); + openAIClient = BraintrustOpenAI.wrapOpenAI(testHarness.openTelemetry(), openAIClient); String imageDataUrl = Base64Attachment.ofFile( @@ -399,17 +367,9 @@ void testWrapOpenAiWithImageAttachment() { response.choices().get(0).message().content().get()); // Verify spans were exported - assertTrue( - openTelemetry - .getSdkTracerProvider() - .forceFlush() - .join(10, TimeUnit.SECONDS) - .isSuccess()); - var spanData = - getExportedBraintrustSpans().get(config.getBraintrustParentValue().orElseThrow()); - assertNotNull(spanData); - assertEquals(1, spanData.size()); - var span = spanData.get(0); + var spans = testHarness.awaitExportedSpans(); + assertEquals(1, spans.size()); + var span = spans.get(0); // Verify span attributes assertEquals("openai", span.getAttributes().get(AttributeKey.stringKey("gen_ai.system"))); diff --git a/src/test/java/dev/braintrust/trace/BraintrustTracingTest.java b/src/test/java/dev/braintrust/trace/BraintrustTracingTest.java index 3f6eb75c..eff7ccf2 100644 --- a/src/test/java/dev/braintrust/trace/BraintrustTracingTest.java +++ b/src/test/java/dev/braintrust/trace/BraintrustTracingTest.java @@ -2,35 +2,18 @@ import static org.junit.jupiter.api.Assertions.*; -import dev.braintrust.config.BraintrustConfig; -import io.opentelemetry.api.GlobalOpenTelemetry; +import dev.braintrust.TestHarness; 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_JAVA_EXPORT_SPANS_IN_MEMORY_FOR_UNIT_TEST", "true"); + private TestHarness testHarness; @BeforeEach void beforeEach() { - GlobalOpenTelemetry.resetForTest(); - getExportedBraintrustSpans().clear(); + testHarness = TestHarness.setup(); } @Test @@ -45,40 +28,40 @@ void instrumentationInfoIsPresent() { @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"))); + var spans = testHarness.awaitExportedSpans(); + assertEquals(1, spans.size()); + assertEquals(true, spans.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()); + // TestHarness already sets up custom BT tracing with openTelemetryEnable + // We just need to verify it works with a custom tracer + doSimpleOtelTrace(testHarness.openTelemetry().getTracer("some-instrumentation")); + var spans = testHarness.awaitExportedSpans(); + assertEquals(1, spans.size()); + assertEquals(true, spans.get(0).getAttributes().get(AttributeKey.booleanKey("unit-test"))); + } + + @Test + void spanProcessorAddsParentFromConfig() { + // Verify that BraintrustSpanProcessor automatically adds the braintrust.parent + // attribute from the config when no parent is explicitly set on the span + doSimpleOtelTrace(testHarness.openTelemetry().getTracer("my tracer")); + var spans = testHarness.awaitExportedSpans(); + assertEquals(1, spans.size()); + var span = spans.get(0); + + // The span should have the braintrust.parent attribute set from the config + var parentAttribute = span.getAttributes().get(AttributeKey.stringKey("braintrust.parent")); + assertNotNull( + parentAttribute, + "braintrust.parent attribute should be set by BraintrustSpanProcessor"); assertEquals( - true, spanData.get(0).getAttributes().get(AttributeKey.booleanKey("unit-test"))); + "project_name:" + TestHarness.defaultProjectName(), + parentAttribute, + "braintrust.parent should be set from the config's default project name"); } private void doSimpleOtelTrace(Tracer tracer) {