diff --git a/examples/src/main/java/dev/braintrust/examples/TriviaEvaluation.java b/examples/src/main/java/dev/braintrust/examples/TriviaEvaluation.java index 73d4524..f6a3da7 100644 --- a/examples/src/main/java/dev/braintrust/examples/TriviaEvaluation.java +++ b/examples/src/main/java/dev/braintrust/examples/TriviaEvaluation.java @@ -61,17 +61,18 @@ public static void main(String[] args) throws Exception { // Show experiment URL var appUrl = config.appUrl().toString().replaceAll("/$", ""); + var orgNameEncoded = + java.net.URLEncoder.encode( + config.orgName().orElse("unknown"), + java.nio.charset.StandardCharsets.UTF_8) + .replace("+", "%20"); var projectNameEncoded = java.net.URLEncoder.encode(project.name(), java.nio.charset.StandardCharsets.UTF_8) .replace("+", "%20"); var experimentUrl = - appUrl.contains("staging") - ? String.format( - "%s/app/%s/p/%s/experiments/%s", - appUrl, project.orgId(), projectNameEncoded, experimentName) - : String.format( - "%s/app/braintrustdata.com/p/%s/experiments/%s", - appUrl, projectNameEncoded, experimentName); + String.format( + "%s/app/%s/p/%s/experiments/%s", + appUrl, orgNameEncoded, projectNameEncoded, experimentName); System.out.println("\nExperiment " + experimentName + " is running at " + experimentUrl); diff --git a/src/main/java/dev/braintrust/api/BraintrustApiClient.java b/src/main/java/dev/braintrust/api/BraintrustApiClient.java index 108a641..9f2293c 100644 --- a/src/main/java/dev/braintrust/api/BraintrustApiClient.java +++ b/src/main/java/dev/braintrust/api/BraintrustApiClient.java @@ -109,6 +109,12 @@ public CompletableFuture> listDatasets(String projectId) { .thenApply(DatasetList::datasets); } + /** Login to get user information including organization details. */ + public CompletableFuture login() { + var request = new LoginRequest(config.apiKey()); + return postAsync("/api/apikey/login", request, LoginResponse.class); + } + // Low-level HTTP methods private CompletableFuture getAsync(String path, Class responseType) { @@ -176,8 +182,8 @@ private T handleResponse(HttpResponse response, Class responseTyp } private boolean isNotFound(Throwable error) { - if (error instanceof ApiException apiException) { - return apiException.getMessage().contains("404"); + if (error instanceof ApiException) { + return ((ApiException) error).getMessage().contains("404"); } return false; } @@ -220,6 +226,7 @@ public record CreateExperimentRequest( String name, Optional description, Optional baseExperimentId) { + public CreateExperimentRequest(String projectId, String name) { this(projectId, name, Optional.empty(), Optional.empty()); } @@ -263,4 +270,11 @@ public DatasetEvent(Object input, Object output) { private record InsertEventsRequest(List events) {} public record InsertEventsResponse(int insertedCount) {} + + // User and Organization models for login functionality + public record OrganizationInfo(String id, String name) {} + + private record LoginRequest(String token) {} + + public record LoginResponse(List orgInfo) {} } diff --git a/src/main/java/dev/braintrust/config/BraintrustConfig.java b/src/main/java/dev/braintrust/config/BraintrustConfig.java index 5589412..6057de6 100644 --- a/src/main/java/dev/braintrust/config/BraintrustConfig.java +++ b/src/main/java/dev/braintrust/config/BraintrustConfig.java @@ -1,5 +1,6 @@ package dev.braintrust.config; +import dev.braintrust.api.BraintrustApiClient; import java.net.URI; import java.time.Duration; import java.util.Optional; @@ -18,6 +19,7 @@ public final class BraintrustConfig { private final URI apiUrl; private final URI appUrl; @Nullable private final String defaultProjectId; + @Nullable private final String orgName; private final boolean enableTraceConsoleLog; private final boolean debug; private final Duration requestTimeout; @@ -27,6 +29,7 @@ private BraintrustConfig(Builder builder) { this.apiUrl = builder.apiUrl; this.appUrl = builder.appUrl; this.defaultProjectId = builder.defaultProjectId; + this.orgName = builder.orgName; this.enableTraceConsoleLog = builder.enableTraceConsoleLog; this.debug = builder.debug; this.requestTimeout = builder.requestTimeout; @@ -48,6 +51,10 @@ public Optional defaultProjectId() { return Optional.ofNullable(defaultProjectId); } + public Optional orgName() { + return Optional.ofNullable(orgName); + } + public boolean enableTraceConsoleLog() { return enableTraceConsoleLog; } @@ -65,12 +72,14 @@ public static Builder builder() { return new Builder(); } - /** Creates a config from environment variables only. */ + /** Creates a config from environment variables and retrieves organization name from API. */ public static BraintrustConfig fromEnvironment() { return builder().build(); } - /** Creates a config with a consumer for customization. */ + /** + * Creates a config with a consumer for customization and retrieves organization name from API. + */ public static BraintrustConfig create(Consumer customizer) { var builder = builder(); customizer.accept(builder); @@ -82,6 +91,7 @@ public static final class Builder { private URI apiUrl; private URI appUrl; private String defaultProjectId; + private String orgName; private boolean enableTraceConsoleLog; private boolean debug; private Duration requestTimeout = Duration.ofSeconds(30); @@ -125,6 +135,11 @@ public Builder defaultProjectId(String projectId) { return this; } + public Builder orgName(String orgName) { + this.orgName = orgName; + return this; + } + public Builder enableTraceConsoleLog(boolean enable) { this.enableTraceConsoleLog = enable; return this; @@ -146,9 +161,40 @@ public BraintrustConfig build() { "API key is required. Set BRAINTRUST_API_KEY environment variable or use" + " apiKey() method."); } + + // If orgName is not already set, try to retrieve it from the API + if (orgName == null || orgName.isBlank()) { + try { + retrieveOrgName(); + } catch (Exception e) { + // Log warning but don't fail the build if org name retrieval fails + System.err.println( + "Warning: Failed to retrieve organization name: " + e.getMessage()); + } + } + return new BraintrustConfig(this); } + private void retrieveOrgName() { + // Create a temporary config without orgName to avoid circular dependency + var tempConfig = new BraintrustConfig(this); + + try (var apiClient = new BraintrustApiClient(tempConfig)) { + var loginResponse = apiClient.login().join(); + var orgInfo = loginResponse.orgInfo(); + + if (orgInfo != null && !orgInfo.isEmpty()) { + // Use the first organization by default, similar to Python SDK + this.orgName = orgInfo.get(0).name(); + } else { + throw new RuntimeException("User is not part of any organizations"); + } + } catch (Exception e) { + throw new RuntimeException("Failed to retrieve organization information", e); + } + } + private static String getEnv(String key, String defaultValue) { String value = System.getenv(key); // Trim any whitespace that might have been accidentally included