Merge pull request #5 from braintrustdata/ark/unit-test-harness

Ark/unit test harness

Author: realark

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
+  }
+}

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<Braintrust> instance = new AtomicReference<>();
+    private static final AtomicReference<Braintrust> 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()));
     }
 
     /**

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);
+        }
+    }
+}

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<OrganizationAndProjectInfo> 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<Prompt> getPrompt(
             @Nonnull String projectName, @Nonnull String slug, @Nullable String version);
@@ -145,6 +149,43 @@ public Optional<OrganizationAndProjectInfo> 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<Prompt> getPrompt(
                 @Nonnull String projectName, @Nonnull String slug, @Nullable String version) {
@@ -293,13 +334,14 @@ class InMemoryImpl implements BraintrustApiClient {
         private final List<Prompt> prompts = new ArrayList<>();
 
         public InMemoryImpl(OrganizationAndProjectInfo... organizationAndProjectInfos) {
-            this.organizationAndProjectInfos = List.of(organizationAndProjectInfos);
+            this.organizationAndProjectInfos =
+                    new ArrayList<>(List.of(organizationAndProjectInfos));
         }
 
         public InMemoryImpl(
                 List<OrganizationAndProjectInfo> organizationAndProjectInfos,
                 List<Prompt> 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<OrganizationAndProjectInfo> 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<Prompt> getPrompt(
                 @Nonnull String projectName, @Nonnull String slug, @Nullable String version) {

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<String> 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);
     }
 }

src/main/java/dev/braintrust/eval/Eval.java

@@ -2,14 +2,14 @@
 
 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;
 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;
@@ -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() {

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<String, List<SpanData>> SPANS_EXPORTED = new ConcurrentHashMap<>();
-
     private final BraintrustConfig config;
     private final String tracesEndpoint;
     private final Map<String, OtlpHttpSpanExporter> exporterCache = new ConcurrentHashMap<>();
@@ -90,8 +86,8 @@ private CompletableResultCode exportWithParent(String parent, List<SpanData> 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);