From 0918db94811eb9d1b85ed74db519909ea0a43742 Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Mon, 22 Sep 2025 12:15:24 -0600 Subject: [PATCH 1/5] rm dead claude code --- .../claude/examples/DatasetExample.java | 329 ------------- .../claude/examples/DebugAuthHeader.java | 91 ---- .../claude/examples/EvaluationExample.java | 177 ------- .../claude/examples/ExperimentExample.java | 223 --------- .../examples/ExperimentWithApiExample.java | 143 ------ .../claude/examples/HttpExporterTest.java | 45 -- .../claude/examples/MinimalTest.java | 66 --- .../claude/examples/OpenTelemetryExample.java | 436 ------------------ .../claude/examples/ProjectExample.java | 151 ------ .../claude/examples/QuickStartExample.java | 114 ----- .../claude/examples/SimpleExample.java | 234 ---------- .../examples/SimpleExperimentExample.java | 148 ------ .../SimpleExperimentWithRegistration.java | 227 --------- .../claude/examples/TestApiEndpoints.java | 76 --- .../claude/examples/TriviaEvaluation.java | 245 ---------- .../braintrust/claude/api/ApiException.java | 13 - .../claude/api/BraintrustApiClient.java | 280 ----------- .../dev/braintrust/claude/api/Experiment.java | 120 ----- .../dev/braintrust/claude/api/Project.java | 52 --- .../claude/config/BraintrustConfig.java | 204 -------- .../braintrust/claude/eval/Evaluation.java | 417 ----------------- .../claude/eval/EvaluationException.java | 13 - .../claude/eval/EvaluationResult.java | 41 -- .../claude/eval/EvaluationResults.java | 80 ---- .../claude/eval/EvaluationSummary.java | 23 - .../claude/eval/ScoreStatistics.java | 24 - .../dev/braintrust/claude/eval/Scorer.java | 214 --------- .../anthropic/AnthropicInterceptor.java | 269 ----------- .../openai/OpenAIInterceptor.java | 203 -------- .../claude/log/BraintrustLogger.java | 257 ----------- .../claude/openai/BraintrustOpenAI.java | 213 --------- .../claude/trace/BraintrustContext.java | 107 ----- .../claude/trace/BraintrustSpanExporter.java | 122 ----- .../claude/trace/BraintrustSpanProcessor.java | 256 ---------- .../claude/trace/BraintrustTracing.java | 200 -------- .../claude/trace/DebugSpanExporter.java | 84 ---- .../claude/api/BraintrustApiClientTest.java | 290 ------------ .../claude/config/BraintrustConfigTest.java | 149 ------ .../eval/EvaluationIntegrationTest.java | 420 ----------------- .../claude/eval/EvaluationTest.java | 231 ---------- .../braintrust/claude/eval/ScorerTest.java | 225 --------- .../anthropic/AnthropicInterceptorTest.java | 430 ----------------- .../openai/OpenAIInterceptorTest.java | 271 ----------- .../claude/log/BraintrustLoggerTest.java | 246 ---------- .../claude/test/otel/SpanAssertions.java | 224 --------- .../claude/trace/BraintrustContextTest.java | 179 ------- .../trace/BraintrustSpanProcessorTest.java | 142 ------ .../claude/trace/BraintrustTracingTest.java | 274 ----------- .../config/BraintrustConfigTest.java | 8 + 49 files changed, 8 insertions(+), 8978 deletions(-) delete mode 100644 examples/src/main/java/dev/braintrust/claude/examples/DatasetExample.java delete mode 100644 examples/src/main/java/dev/braintrust/claude/examples/DebugAuthHeader.java delete mode 100644 examples/src/main/java/dev/braintrust/claude/examples/EvaluationExample.java delete mode 100644 examples/src/main/java/dev/braintrust/claude/examples/ExperimentExample.java delete mode 100644 examples/src/main/java/dev/braintrust/claude/examples/ExperimentWithApiExample.java delete mode 100644 examples/src/main/java/dev/braintrust/claude/examples/HttpExporterTest.java delete mode 100644 examples/src/main/java/dev/braintrust/claude/examples/MinimalTest.java delete mode 100644 examples/src/main/java/dev/braintrust/claude/examples/OpenTelemetryExample.java delete mode 100644 examples/src/main/java/dev/braintrust/claude/examples/ProjectExample.java delete mode 100644 examples/src/main/java/dev/braintrust/claude/examples/QuickStartExample.java delete mode 100644 examples/src/main/java/dev/braintrust/claude/examples/SimpleExample.java delete mode 100644 examples/src/main/java/dev/braintrust/claude/examples/SimpleExperimentExample.java delete mode 100644 examples/src/main/java/dev/braintrust/claude/examples/SimpleExperimentWithRegistration.java delete mode 100644 examples/src/main/java/dev/braintrust/claude/examples/TestApiEndpoints.java delete mode 100644 examples/src/main/java/dev/braintrust/claude/examples/TriviaEvaluation.java delete mode 100644 src/main/java/dev/braintrust/claude/api/ApiException.java delete mode 100644 src/main/java/dev/braintrust/claude/api/BraintrustApiClient.java delete mode 100644 src/main/java/dev/braintrust/claude/api/Experiment.java delete mode 100644 src/main/java/dev/braintrust/claude/api/Project.java delete mode 100644 src/main/java/dev/braintrust/claude/config/BraintrustConfig.java delete mode 100644 src/main/java/dev/braintrust/claude/eval/Evaluation.java delete mode 100644 src/main/java/dev/braintrust/claude/eval/EvaluationException.java delete mode 100644 src/main/java/dev/braintrust/claude/eval/EvaluationResult.java delete mode 100644 src/main/java/dev/braintrust/claude/eval/EvaluationResults.java delete mode 100644 src/main/java/dev/braintrust/claude/eval/EvaluationSummary.java delete mode 100644 src/main/java/dev/braintrust/claude/eval/ScoreStatistics.java delete mode 100644 src/main/java/dev/braintrust/claude/eval/Scorer.java delete mode 100644 src/main/java/dev/braintrust/claude/instrumentation/anthropic/AnthropicInterceptor.java delete mode 100644 src/main/java/dev/braintrust/claude/instrumentation/openai/OpenAIInterceptor.java delete mode 100644 src/main/java/dev/braintrust/claude/log/BraintrustLogger.java delete mode 100644 src/main/java/dev/braintrust/claude/openai/BraintrustOpenAI.java delete mode 100644 src/main/java/dev/braintrust/claude/trace/BraintrustContext.java delete mode 100644 src/main/java/dev/braintrust/claude/trace/BraintrustSpanExporter.java delete mode 100644 src/main/java/dev/braintrust/claude/trace/BraintrustSpanProcessor.java delete mode 100644 src/main/java/dev/braintrust/claude/trace/BraintrustTracing.java delete mode 100644 src/main/java/dev/braintrust/claude/trace/DebugSpanExporter.java delete mode 100644 src/test/java/dev/braintrust/claude/api/BraintrustApiClientTest.java delete mode 100644 src/test/java/dev/braintrust/claude/config/BraintrustConfigTest.java delete mode 100644 src/test/java/dev/braintrust/claude/eval/EvaluationIntegrationTest.java delete mode 100644 src/test/java/dev/braintrust/claude/eval/EvaluationTest.java delete mode 100644 src/test/java/dev/braintrust/claude/eval/ScorerTest.java delete mode 100644 src/test/java/dev/braintrust/claude/instrumentation/anthropic/AnthropicInterceptorTest.java delete mode 100644 src/test/java/dev/braintrust/claude/instrumentation/openai/OpenAIInterceptorTest.java delete mode 100644 src/test/java/dev/braintrust/claude/log/BraintrustLoggerTest.java delete mode 100644 src/test/java/dev/braintrust/claude/test/otel/SpanAssertions.java delete mode 100644 src/test/java/dev/braintrust/claude/trace/BraintrustContextTest.java delete mode 100644 src/test/java/dev/braintrust/claude/trace/BraintrustSpanProcessorTest.java delete mode 100644 src/test/java/dev/braintrust/claude/trace/BraintrustTracingTest.java create mode 100644 src/test/java/dev/braintrust/config/BraintrustConfigTest.java diff --git a/examples/src/main/java/dev/braintrust/claude/examples/DatasetExample.java b/examples/src/main/java/dev/braintrust/claude/examples/DatasetExample.java deleted file mode 100644 index 34c8ae0..0000000 --- a/examples/src/main/java/dev/braintrust/claude/examples/DatasetExample.java +++ /dev/null @@ -1,329 +0,0 @@ -package dev.braintrust.claude.examples; - -import dev.braintrust.claude.api.BraintrustApiClient; -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.eval.EvaluationResult; -import dev.braintrust.claude.eval.EvaluationResults; -import dev.braintrust.claude.eval.Evaluation; -import dev.braintrust.claude.eval.Scorer; -import dev.braintrust.claude.trace.BraintrustContext; -import dev.braintrust.claude.trace.BraintrustTracing; -import io.opentelemetry.api.trace.Span; -import io.opentelemetry.context.Context; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.concurrent.CompletableFuture; - -/** - * Example showing dataset management and API usage with modern Java patterns. Demonstrates - * CompletableFuture composition, Optional handling, and functional transformations. - */ -public class DatasetExample { - - // Domain models using records - record QAExample(String question, String context, String answer) {} - - record ModelOutput(String answer, double confidence) {} - - public static void main(String[] args) throws Exception { - // Initialize configuration - var config = BraintrustConfig.fromEnvironment(); - var openTelemetry = BraintrustTracing.quickstart(config); - var apiClient = new BraintrustApiClient(config); - - // Create or get project - var projectFuture = - apiClient - .createProject("qa-evaluation-demo") - .exceptionally( - error -> { - System.err.println("Failed to create project: " + error); - throw new RuntimeException(error); - }); - - var project = projectFuture.get(); - System.out.println("Using project: " + project.name() + " (ID: " + project.id() + ")"); - - // Set project context for all spans - var context = BraintrustContext.forProject(project.id()).storeInContext(Context.current()); - - try (var scope = context.makeCurrent()) { - // Create dataset - var datasetFuture = createOrUpdateDataset(apiClient, project.id()); - var dataset = datasetFuture.get(); - - // Load evaluation data - var testData = loadTestData(); - - // Run evaluation with experiment tracking - var experimentFuture = - apiClient.createExperiment( - new BraintrustApiClient.CreateExperimentRequest( - project.id(), - "qa-model-v2-eval", - Optional.of("Evaluating QA model v2 with confidence scores"), - Optional.empty())); - - var experiment = experimentFuture.get(); - - // Run evaluation with experiment context - var results = - Evaluation.builder() - .name("QA Model Evaluation") - .data(testData) - .task(DatasetExample::runQAModel) - .scorer(createAnswerScorer()) - .scorer(createConfidenceScorer()) - .scorer(createLengthScorer()) - .experimentId(experiment.id()) - .parallel(true) - .run(); - - // Analyze results - analyzeResults(results); - - // Upload results to dataset for future use - uploadResultsToDataset(apiClient, dataset.id(), results).get(); - - System.out.println("\nEvaluation completed successfully!"); - } - - // Clean up - apiClient.close(); - Thread.sleep(2000); // Allow spans to export - } - - private static CompletableFuture createOrUpdateDataset( - BraintrustApiClient client, String projectId) { - - return client.createDataset( - new BraintrustApiClient.CreateDatasetRequest( - projectId, - "qa-test-dataset", - Optional.of("Question-answering test cases"))) - .thenCompose( - dataset -> { - // Add sample events to dataset - var events = - List.of( - new BraintrustApiClient.DatasetEvent( - Map.of( - "question", - "What is the capital of" - + " France?", - "context", - "France is a country in" - + " Europe."), - "Paris"), - new BraintrustApiClient.DatasetEvent( - Map.of( - "question", - "Who wrote Romeo and Juliet?", - "context", - "Shakespeare was an English" - + " playwright."), - "William Shakespeare")); - - return client.insertDatasetEvents(dataset.id(), events) - .thenApply( - response -> { - System.out.println( - "Inserted " - + response.insertedCount() - + " events"); - return dataset; - }); - }); - } - - private static List loadTestData() { - return List.of( - new QAExample( - "What is the capital of France?", - "France is a country in Western Europe. Its capital city is known for the" - + " Eiffel Tower.", - "Paris"), - new QAExample( - "Who wrote Romeo and Juliet?", - "William Shakespeare was an English playwright and poet of the Renaissance" - + " period.", - "William Shakespeare"), - new QAExample( - "What is the speed of light?", - "Light travels at approximately 299,792,458 meters per second in a vacuum.", - "299,792,458 meters per second"), - new QAExample( - "What is machine learning?", - "Machine learning is a subset of artificial intelligence that enables" - + " systems to learn from data.", - "A subset of AI that enables systems to learn from data")); - } - - private static ModelOutput runQAModel(QAExample example) { - var span = Span.current(); - span.setAttribute("qa.question_length", example.question.length()); - span.setAttribute("qa.context_length", example.context.length()); - - // Simulate model inference - try { - Thread.sleep(50 + (long) (Math.random() * 100)); - - // Simple rule-based "model" for demo - var confidence = 0.5 + Math.random() * 0.5; - var answer = generateAnswer(example); - - span.setAttribute("qa.answer_length", answer.length()); - span.setAttribute("qa.confidence", confidence); - - return new ModelOutput(answer, confidence); - - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException(e); - } - } - - private static String generateAnswer(QAExample example) { - // Simulate different quality answers - var random = Math.random(); - if (random > 0.7) { - return example.answer; // Correct answer - } else if (random > 0.3) { - return example.answer.substring(0, Math.min(10, example.answer.length())) - + "..."; // Partial - } else { - return "I don't know"; // Wrong answer - } - } - - private static Scorer createAnswerScorer() { - return Scorer.of( - "answer_accuracy", - (example, output) -> { - if (output.answer.equals(example.answer)) { - return 1.0; - } else if (example.answer.toLowerCase().contains(output.answer.toLowerCase()) - || output.answer.toLowerCase().contains(example.answer.toLowerCase())) { - return 0.5; - } - return 0.0; - }); - } - - private static Scorer createConfidenceScorer() { - return Scorer.of( - "confidence_calibration", - (example, output) -> { - var isCorrect = output.answer.equals(example.answer); - var confidence = output.confidence; - - // Good calibration: high confidence when correct, low when incorrect - if (isCorrect && confidence > 0.7) { - return 1.0; - } else if (!isCorrect && confidence < 0.3) { - return 1.0; - } else { - return 1.0 - Math.abs(confidence - (isCorrect ? 1.0 : 0.0)); - } - }); - } - - private static Scorer createLengthScorer() { - return Scorer.of( - "answer_completeness", - (example, output) -> { - var expectedLength = example.answer.length(); - var actualLength = output.answer.length(); - - if (actualLength == 0) return 0.0; - - var ratio = (double) actualLength / expectedLength; - if (ratio > 2.0) return 0.5; // Too verbose - if (ratio < 0.5) return ratio; // Too short - return 1.0; - }); - } - - private static void analyzeResults( - EvaluationResults results) { - System.out.println("\n=== Evaluation Analysis ==="); - - // Overall metrics - System.out.println(results.generateReport()); - - // Confidence analysis - var avgConfidence = - results.successful().mapToDouble(r -> r.output().confidence).average().orElse(0.0); - - System.out.printf("\nAverage model confidence: %.3f%n", avgConfidence); - - // Find overconfident wrong answers - System.out.println("\nOverconfident mistakes (confidence > 0.7, score < 0.5):"); - results.successful() - .filter(r -> r.output().confidence > 0.7 && r.averageScore() < 0.5) - .forEach( - r -> - System.out.printf( - " Q: %s%n A: %s (confidence: %.2f)%n", - r.input().question, - r.output().answer, - r.output().confidence)); - - // Score correlation analysis - var scoresByQuestion = - results.successful() - .collect( - java.util.stream.Collectors.groupingBy( - r -> r.input().question, - java.util.stream.Collectors.averagingDouble( - EvaluationResult - ::averageScore))); - - System.out.println("\nHardest questions:"); - scoresByQuestion.entrySet().stream() - .sorted(Map.Entry.comparingByValue()) - .limit(2) - .forEach( - entry -> - System.out.printf( - " %s (avg score: %.2f)%n", - entry.getKey(), entry.getValue())); - } - - private static CompletableFuture uploadResultsToDataset( - BraintrustApiClient client, - String datasetId, - EvaluationResults results) { - - // Convert results to dataset events - var events = - results.successful() - .map( - r -> - new BraintrustApiClient.DatasetEvent( - Map.of( - "question", r.input().question, - "context", r.input().context, - "expected", r.input().answer), - Optional.of( - Map.of( - "answer", r.output().answer, - "confidence", r.output().confidence, - "scores", r.scores())), - Optional.of( - Map.of( - "evaluation_run", results.name(), - "duration_ms", - r.duration().toMillis())))) - .toList(); - - return client.insertDatasetEvents(datasetId, events) - .thenAccept( - response -> - System.out.println( - "\nUploaded " - + response.insertedCount() - + " evaluation results to dataset")); - } -} diff --git a/examples/src/main/java/dev/braintrust/claude/examples/DebugAuthHeader.java b/examples/src/main/java/dev/braintrust/claude/examples/DebugAuthHeader.java deleted file mode 100644 index 2e4c757..0000000 --- a/examples/src/main/java/dev/braintrust/claude/examples/DebugAuthHeader.java +++ /dev/null @@ -1,91 +0,0 @@ -package dev.braintrust.claude.examples; - -import dev.braintrust.claude.api.BraintrustApiClient; -import dev.braintrust.claude.config.BraintrustConfig; -import java.net.URI; -import java.net.http.HttpClient; -import java.net.http.HttpRequest; -import java.net.http.HttpResponse; - -public class DebugAuthHeader { - public static void main(String[] args) throws Exception { - // Get the config - var config = BraintrustConfig.fromEnvironment(); - - System.out.println("=== Debug Authorization Header ==="); - System.out.println("API URL: " + config.apiUrl()); - System.out.println("API Key: " + maskApiKey(config.apiKey())); - System.out.println("API Key length: " + config.apiKey().length()); - System.out.println( - "API Key starts with: " - + config.apiKey().substring(0, Math.min(10, config.apiKey().length()))); - - // Check for any special characters - System.out.println("\nChecking for special characters in API key:"); - for (int i = 0; i < config.apiKey().length(); i++) { - char c = config.apiKey().charAt(i); - if (!Character.isLetterOrDigit(c) && c != '-' && c != '_') { - System.out.println( - " Found special char at position " - + i - + ": '" - + c - + "' (ASCII: " - + (int) c - + ")"); - } - } - - // Try a raw HTTP request to see exactly what we're sending - System.out.println("\n=== Testing Raw HTTP Request ==="); - var httpClient = HttpClient.newHttpClient(); - var authHeader = "Bearer " + config.apiKey(); - System.out.println("Full Authorization header: '" + authHeader + "'"); - System.out.println("Header length: " + authHeader.length()); - - var request = - HttpRequest.newBuilder() - .uri(URI.create(config.apiUrl() + "/projects")) - .header("Authorization", authHeader) - .header("Accept", "application/json") - .GET() - .build(); - - System.out.println("\nRequest headers:"); - request.headers() - .map() - .forEach( - (key, values) -> { - System.out.println(" " + key + ": " + values); - }); - - try { - var response = httpClient.send(request, HttpResponse.BodyHandlers.ofString()); - System.out.println("\nResponse status: " + response.statusCode()); - System.out.println("Response body: " + response.body()); - } catch (Exception e) { - System.err.println("Request failed: " + e.getMessage()); - e.printStackTrace(); - } - - // Also try using the API client - System.out.println("\n=== Testing via API Client ==="); - try { - var apiClient = new BraintrustApiClient(config); - var projects = apiClient.listProjects().get(); - System.out.println("Success! Found " + projects.size() + " projects"); - } catch (Exception e) { - System.err.println("API client failed: " + e.getMessage()); - if (e.getCause() != null) { - System.err.println("Cause: " + e.getCause().getMessage()); - } - } - } - - private static String maskApiKey(String apiKey) { - if (apiKey == null || apiKey.length() < 8) { - return "***"; - } - return apiKey.substring(0, 4) + "..." + apiKey.substring(apiKey.length() - 4); - } -} diff --git a/examples/src/main/java/dev/braintrust/claude/examples/EvaluationExample.java b/examples/src/main/java/dev/braintrust/claude/examples/EvaluationExample.java deleted file mode 100644 index df1df40..0000000 --- a/examples/src/main/java/dev/braintrust/claude/examples/EvaluationExample.java +++ /dev/null @@ -1,177 +0,0 @@ -package dev.braintrust.claude.examples; - -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.eval.Evaluation; -import dev.braintrust.claude.eval.Scorer; -import dev.braintrust.claude.trace.BraintrustTracing; -import java.time.Duration; -import java.util.Map; -import java.util.Random; -import java.util.concurrent.ForkJoinPool; -import java.util.regex.Pattern; -import java.util.stream.IntStream; - -/** - * Example demonstrating the evaluation framework with idiomatic Java. Shows use of records, - * streams, lambdas, and method references. - */ -public class EvaluationExample { - - // Test case using Java record - record MathProblem(String question, int expected) {} - - // AI model simulator - static class SimpleMathModel { - private final Random random = new Random(); - - public String solve(String question) { - // Extract numbers from question (simplified) - var numbers = - Pattern.compile("\\d+") - .matcher(question) - .results() - .map(mr -> Integer.parseInt(mr.group())) - .toList(); - - if (numbers.size() >= 2) { - // Simulate solving with some error - var result = numbers.get(0) + numbers.get(1); - if (random.nextDouble() > 0.9) { - result += random.nextInt(3) - 1; // Add small error - } - return String.valueOf(result); - } - return "Unable to solve"; - } - } - - public static void main(String[] args) { - // Initialize Braintrust - var config = BraintrustConfig.fromEnvironment(); - BraintrustTracing.quickstart(config); - - // Create test dataset using streams - var testData = - IntStream.range(1, 21) - .mapToObj( - i -> - new MathProblem( - String.format("What is %d + %d?", i, i * 2), - i + i * 2)) - .toList(); - - var model = new SimpleMathModel(); - - // Run evaluation with multiple scorers - var results = - Evaluation.builder() - .name("Math Problem Evaluation") - .data(testData) - .task(problem -> model.solve(problem.question)) - .scorer(Scorer.StringScorers.exactMatch(p -> String.valueOf(p.expected))) - .scorer(createNumericScorer()) - .scorer(createPartialCreditScorer()) - .parallel(true) - .timeout(Duration.ofSeconds(5)) - .executor(ForkJoinPool.commonPool()) - .run(); - - // Print comprehensive report - System.out.println(results.generateReport()); - - // Analyze results using streams - System.out.println("\nDetailed Analysis:"); - - // Group by score ranges - var scoreDistribution = - results.results().stream() - .collect( - java.util.stream.Collectors.groupingBy( - r -> (int) (r.averageScore() * 10) / 10.0, - java.util.stream.Collectors.counting())); - - System.out.println("Score Distribution:"); - scoreDistribution.entrySet().stream() - .sorted(Map.Entry.comparingByKey()) - .forEach( - entry -> - System.out.printf( - " %.1f: %d results%n", entry.getKey(), entry.getValue())); - - // Find problematic cases - System.out.println("\nProblematic Cases (score < 0.5):"); - results.results().stream() - .filter(r -> r.averageScore() < 0.5) - .forEach( - r -> - System.out.printf( - " Q: %s, Expected: %d, Got: %s%n", - r.input().question, r.input().expected, r.output())); - - // Performance metrics - var durations = - results.results().stream() - .mapToLong(r -> r.duration().toMillis()) - .summaryStatistics(); - - System.out.printf("\nPerformance Metrics:%n"); - System.out.printf(" Min: %dms%n", durations.getMin()); - System.out.printf(" Max: %dms%n", durations.getMax()); - System.out.printf(" Avg: %.2fms%n", durations.getAverage()); - - // Wait for spans to export - try { - Thread.sleep(2000); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - } - } - - private static Scorer createNumericScorer() { - return Scorer.of( - "numeric_accuracy", - (problem, output) -> { - try { - var actual = Integer.parseInt(output); - var expected = problem.expected; - var error = Math.abs(actual - expected); - - // Give partial credit for close answers - return switch (error) { - case 0 -> 1.0; - case 1 -> 0.8; - case 2 -> 0.5; - default -> 0.0; - }; - } catch (NumberFormatException e) { - return 0.0; - } - }); - } - - private static Scorer createPartialCreditScorer() { - return new Scorer() { - @Override - public double score(MathProblem input, String output) { - // Check if output contains the correct answer - var expectedStr = String.valueOf(input.expected); - if (output.contains(expectedStr)) { - return output.equals(expectedStr) ? 1.0 : 0.7; - } - - // Check if it's a valid number - try { - Integer.parseInt(output); - return 0.3; // Valid number but wrong - } catch (NumberFormatException e) { - return 0.0; // Not a number - } - } - - @Override - public String name() { - return "partial_credit"; - } - }; - } -} diff --git a/examples/src/main/java/dev/braintrust/claude/examples/ExperimentExample.java b/examples/src/main/java/dev/braintrust/claude/examples/ExperimentExample.java deleted file mode 100644 index 19648df..0000000 --- a/examples/src/main/java/dev/braintrust/claude/examples/ExperimentExample.java +++ /dev/null @@ -1,223 +0,0 @@ -package dev.braintrust.claude.examples; - -import com.openai.client.OpenAIClient; -import com.openai.client.okhttp.OpenAIOkHttpClient; -import com.openai.models.ChatModel; -import com.openai.models.chat.completions.ChatCompletionCreateParams; -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.eval.Evaluation; -import dev.braintrust.claude.trace.BraintrustTracing; -import io.opentelemetry.api.trace.SpanKind; -import io.opentelemetry.context.Scope; -import java.util.List; -import java.util.concurrent.CompletableFuture; - -/** - * Example demonstrating how to run an experiment/evaluation in Braintrust. This evaluates different - * prompts for a question-answering task. - */ -public class ExperimentExample { - - // Test case for Q&A evaluation - record QATestCase(String question, String expectedAnswer, String category) {} - - // Output from our Q&A system - record QAOutput(String answer, long promptTokens, long completionTokens, double latencyMs) {} - - public static void main(String[] args) { - // Validate environment - if (System.getenv("BRAINTRUST_API_KEY") == null) { - System.err.println("Please set BRAINTRUST_API_KEY environment variable"); - System.exit(1); - } - if (System.getenv("OPENAI_API_KEY") == null) { - System.err.println("Please set OPENAI_API_KEY environment variable"); - System.exit(1); - } - - // Initialize Braintrust - var config = BraintrustConfig.fromEnvironment(); - var openTelemetry = BraintrustTracing.quickstart(config); - var tracer = BraintrustTracing.getTracer(openTelemetry); - - // Initialize OpenAI client - var openAiClient = OpenAIOkHttpClient.fromEnv(); - - // Define test cases - var testCases = - List.of( - new QATestCase("What is the capital of France?", "Paris", "geography"), - new QATestCase("What is 2 + 2?", "4", "math"), - new QATestCase("Who wrote Romeo and Juliet?", "Shakespeare", "literature"), - new QATestCase( - "What is the largest planet in our solar system?", - "Jupiter", - "science"), - new QATestCase("What year did World War II end?", "1945", "history")); - - // Run experiment with simple prompt - System.out.println("=== Running Experiment: Simple Prompt ==="); - runExperiment( - "QA Simple Prompt v1", - testCases, - "Answer the following question concisely:", - openAiClient, - tracer); - - // Run experiment with detailed prompt - System.out.println("\n=== Running Experiment: Detailed Prompt ==="); - runExperiment( - "QA Detailed Prompt v1", - testCases, - "You are a knowledgeable assistant. Answer the following question accurately and" - + " concisely. If you're not sure, say 'I don't know':", - openAiClient, - tracer); - - // Run experiment with chain-of-thought prompt - System.out.println("\n=== Running Experiment: Chain-of-Thought Prompt ==="); - runExperiment( - "QA Chain-of-Thought v1", - testCases, - "Think step by step and answer the following question. Show your reasoning briefly," - + " then give the final answer:", - openAiClient, - tracer); - - System.out.println("\n=== Experiments Complete ==="); - System.out.println("View results at: https://www.braintrust.dev"); - - // Wait for traces to export - try { - Thread.sleep(5000); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - } - } - - private static void runExperiment( - String experimentName, - List testCases, - String systemPrompt, - OpenAIClient openAiClient, - io.opentelemetry.api.trace.Tracer tracer) { - // Build the evaluation - var evaluation = - Evaluation.builder() - .name(experimentName) - .data(testCases) - .task( - testCase -> { - try { - return evaluateTestCase( - testCase, - systemPrompt, - openAiClient, - tracer) - .get(); - } catch (Exception e) { - throw new RuntimeException( - "Failed to evaluate test case", e); - } - }) - .scorer( - "correctness", - (testCase, output) -> { - // Check if the answer contains the expected answer - // (case-insensitive) - var answerLower = output.answer.toLowerCase(); - var expectedLower = testCase.expectedAnswer.toLowerCase(); - return answerLower.contains(expectedLower) ? 1.0 : 0.0; - }) - .scorer( - "response_time", - (testCase, output) -> { - // Score based on response time (faster is better) - // Perfect score for < 100ms, 0 score for > 2000ms - if (output.latencyMs < 100) return 1.0; - if (output.latencyMs > 2000) return 0.0; - return 1.0 - (output.latencyMs - 100) / 1900.0; - }) - .parallel(true) // Run test cases in parallel - .build(); - - // Run the evaluation - System.out.println("Running evaluation: " + experimentName); - var results = evaluation.run(); - - // Print summary - var summary = results.summary(); - System.out.printf("Completed: %d test cases%n", summary.totalCount()); - System.out.printf("Success rate: %.1f%%%n", summary.successRate() * 100); - - // Print score statistics - summary.scoreStatistics() - .forEach( - (scorer, stats) -> { - System.out.printf( - "%s - Mean: %.2f, Min: %.2f, Max: %.2f%n", - scorer, stats.mean(), stats.min(), stats.max()); - }); - - System.out.printf("Average duration: %.0f ms%n", summary.averageDuration().toMillis()); - } - - private static CompletableFuture evaluateTestCase( - QATestCase testCase, - String systemPrompt, - OpenAIClient client, - io.opentelemetry.api.trace.Tracer tracer) { - return CompletableFuture.supplyAsync( - () -> { - var span = - tracer.spanBuilder("evaluate_qa") - .setSpanKind(SpanKind.INTERNAL) - .setAttribute("test.question", testCase.question) - .setAttribute("test.category", testCase.category) - .startSpan(); - - try (Scope scope = span.makeCurrent()) { - var startTime = System.currentTimeMillis(); - - // Create OpenAI request - var request = - ChatCompletionCreateParams.builder() - .model(ChatModel.GPT_4O) - .addSystemMessage(systemPrompt) - .addUserMessage(testCase.question) - .maxTokens(100L) - .temperature(0.1) // Low temperature for consistent answers - .build(); - - // Call OpenAI - var response = client.chat().completions().create(request); - var answer = response.choices().get(0).message().content().orElse(""); - - var latency = System.currentTimeMillis() - startTime; - - // Extract token usage - var usage = response.usage().orElse(null); - var promptTokens = usage != null ? usage.promptTokens() : 0; - var completionTokens = usage != null ? usage.completionTokens() : 0; - - // Add span attributes - span.setAttribute("llm.model", "gpt-4o"); - span.setAttribute("llm.prompt_tokens", promptTokens); - span.setAttribute("llm.completion_tokens", completionTokens); - span.setAttribute("llm.latency_ms", latency); - span.setAttribute("llm.response", answer); - - span.setStatus(io.opentelemetry.api.trace.StatusCode.OK); - - return new QAOutput(answer, promptTokens, completionTokens, latency); - - } catch (Exception e) { - span.recordException(e); - span.setStatus(io.opentelemetry.api.trace.StatusCode.ERROR, e.getMessage()); - throw new RuntimeException("Failed to evaluate test case", e); - } finally { - span.end(); - } - }); - } -} diff --git a/examples/src/main/java/dev/braintrust/claude/examples/ExperimentWithApiExample.java b/examples/src/main/java/dev/braintrust/claude/examples/ExperimentWithApiExample.java deleted file mode 100644 index 20855af..0000000 --- a/examples/src/main/java/dev/braintrust/claude/examples/ExperimentWithApiExample.java +++ /dev/null @@ -1,143 +0,0 @@ -package dev.braintrust.claude.examples; - -import dev.braintrust.claude.api.BraintrustApiClient; -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.eval.Evaluation; -import dev.braintrust.claude.trace.BraintrustTracing; -import java.util.List; -import java.util.Optional; - -/** - * Example showing how to create an experiment via the API and use it with evaluations. This - * demonstrates the full workflow of experiment management in Braintrust. - */ -public class ExperimentWithApiExample { - - record TestCase(String input, String expected) {} - - record Output(String result, double confidence) {} - - public static void main(String[] args) throws Exception { - // Initialize Braintrust - var config = BraintrustConfig.fromEnvironment(); - var apiClient = new BraintrustApiClient(config); - var openTelemetry = BraintrustTracing.quickstart(config); - - // Step 1: Create or get a project - System.out.println("=== Setting up Project ==="); - var projectName = "Experiment API Example"; - var projects = apiClient.listProjects().get(); - - var project = - projects.stream() - .filter(p -> p.name().equals(projectName)) - .findFirst() - .orElseGet( - () -> { - try { - System.out.println("Creating new project: " + projectName); - return apiClient.createProject(projectName).get(); - } catch (Exception e) { - throw new RuntimeException("Failed to create project", e); - } - }); - - System.out.println("Using project: " + project.name() + " (ID: " + project.id() + ")"); - - // Step 2: Create an experiment - System.out.println("\n=== Creating Experiment ==="); - var experimentName = "Text Classification v" + System.currentTimeMillis(); - var experiment = - apiClient - .createExperiment( - new BraintrustApiClient.CreateExperimentRequest( - project.id(), - experimentName, - Optional.of("Testing text classification accuracy"), - Optional.empty() // No base experiment - )) - .get(); - - System.out.println( - "Created experiment: " + experiment.name() + " (ID: " + experiment.id() + ")"); - - // Step 3: Define test data - var testCases = - List.of( - new TestCase("I love this product!", "positive"), - new TestCase("This is terrible", "negative"), - new TestCase("It's okay, nothing special", "neutral"), - new TestCase("Amazing experience, highly recommend!", "positive"), - new TestCase("Waste of money", "negative"), - new TestCase("Average quality for the price", "neutral")); - - // Step 4: Run evaluation with the experiment - System.out.println("\n=== Running Evaluation ==="); - var evaluation = - Evaluation.builder() - .name("Sentiment Analysis") - .data(testCases) - .task(testCase -> classifyText(testCase.input)) - .scorer( - "accuracy", - (test, output) -> output.result.equals(test.expected) ? 1.0 : 0.0) - .scorer("confidence", (test, output) -> output.confidence) - .experimentId(experiment.id()) // Link to the experiment - .build(); - - var results = evaluation.run(); - - // Step 5: Print results - System.out.println("\n=== Results ==="); - var summary = results.summary(); - System.out.printf("Total test cases: %d%n", summary.totalCount()); - System.out.printf("Success rate: %.1f%%%n", summary.successRate() * 100); - - summary.scoreStatistics() - .forEach( - (scorer, stats) -> { - System.out.printf("\n%s:%n", scorer); - System.out.printf(" Mean: %.3f%n", stats.mean()); - System.out.printf(" Min: %.3f%n", stats.min()); - System.out.printf(" Max: %.3f%n", stats.max()); - }); - - // Step 6: Compare with previous experiments (if any) - System.out.println("\n=== Experiment Comparison ==="); - var allExperiments = apiClient.listExperiments(project.id()).get(); - System.out.printf("Total experiments in project: %d%n", allExperiments.size()); - - // In a real scenario, you would compare metrics across experiments - System.out.println("\nView detailed results at: https://www.braintrust.dev"); - System.out.println("Project: " + project.name()); - System.out.println("Experiment: " + experiment.name()); - - // Wait for traces to export - Thread.sleep(3000); - } - - // Simulated text classifier - private static Output classifyText(String text) { - var lowerText = text.toLowerCase(); - - // Simple keyword-based classification - if (lowerText.contains("love") - || lowerText.contains("amazing") - || lowerText.contains("great") - || lowerText.contains("recommend")) { - return new Output("positive", 0.9); - } else if (lowerText.contains("terrible") - || lowerText.contains("hate") - || lowerText.contains("waste") - || lowerText.contains("bad")) { - return new Output("negative", 0.85); - } else if (lowerText.contains("okay") - || lowerText.contains("average") - || lowerText.contains("nothing special")) { - return new Output("neutral", 0.7); - } else { - // Default to neutral with lower confidence - return new Output("neutral", 0.5); - } - } -} diff --git a/examples/src/main/java/dev/braintrust/claude/examples/HttpExporterTest.java b/examples/src/main/java/dev/braintrust/claude/examples/HttpExporterTest.java deleted file mode 100644 index 29543c2..0000000 --- a/examples/src/main/java/dev/braintrust/claude/examples/HttpExporterTest.java +++ /dev/null @@ -1,45 +0,0 @@ -package dev.braintrust.claude.examples; - -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.trace.BraintrustTracing; - -/** Quick test to verify HTTP exporter configuration. */ -public class HttpExporterTest { - public static void main(String[] args) { - System.out.println("Testing Braintrust HTTP exporter configuration...\n"); - - // Test with dummy credentials - var config = - BraintrustConfig.builder() - .apiKey("test-key") - .apiUrl("https://api.braintrust.dev") // Real endpoint - .build(); - - System.out.println("Configuration:"); - System.out.println("- API URL: " + config.apiUrl()); - System.out.println("- OTLP Endpoint: " + config.apiUrl() + "/otlp/v1/traces"); - System.out.println("- Timeout: " + config.requestTimeout()); - - try { - var openTelemetry = - BraintrustTracing.quickstart( - config, builder -> builder.serviceName("http-exporter-test")); - - var tracer = BraintrustTracing.getTracer(openTelemetry); - - // Create a simple span - var span = tracer.spanBuilder("test.span").startSpan(); - span.setAttribute("test", true); - span.end(); - - System.out.println("\nSpan created and exported via HTTP"); - System.out.println("Expected result: HTTP 401/403 error (invalid API key)"); - - // Wait for export - Thread.sleep(2000); - - } catch (Exception e) { - System.err.println("Error: " + e.getMessage()); - } - } -} diff --git a/examples/src/main/java/dev/braintrust/claude/examples/MinimalTest.java b/examples/src/main/java/dev/braintrust/claude/examples/MinimalTest.java deleted file mode 100644 index 7144d1f..0000000 --- a/examples/src/main/java/dev/braintrust/claude/examples/MinimalTest.java +++ /dev/null @@ -1,66 +0,0 @@ -package dev.braintrust.claude.examples; - -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.log.BraintrustLogger; -import dev.braintrust.claude.trace.BraintrustTracing; -import io.opentelemetry.api.trace.StatusCode; -import io.opentelemetry.context.Scope; -import java.time.Duration; - -/** Minimal test to debug HTTP 400 errors. */ -public class MinimalTest { - - public static void main(String[] args) { - // Enable debug logging - BraintrustLogger.setDebugEnabled(true); - - var apiKey = System.getenv("BRAINTRUST_API_KEY"); - if (apiKey == null || apiKey.isEmpty()) { - System.err.println("Please set BRAINTRUST_API_KEY"); - System.exit(1); - } - - try { - // Minimal configuration - var config = BraintrustConfig.builder().apiKey(apiKey).debug(true).build(); - - System.out.println("Creating minimal test span..."); - System.out.println("Endpoint: " + config.apiUrl() + "/otel/v1/traces"); - - var openTelemetry = - BraintrustTracing.quickstart( - config, - builder -> - builder.serviceName("minimal-test") - .exportInterval(Duration.ofSeconds(1)) - .maxExportBatchSize(1) // Send one span at a time - ); - - var tracer = BraintrustTracing.getTracer(openTelemetry); - - // Create the simplest possible span - var span = tracer.spanBuilder("test").startSpan(); - - try (Scope scope = span.makeCurrent()) { - // Minimal attributes - span.setAttribute("test", true); - span.setStatus(StatusCode.OK); - Thread.sleep(100); // Give it some duration - } finally { - span.end(); - } - - System.out.println("Span created, waiting for export..."); - Thread.sleep(3000); - - System.out.println("\nIf you see HTTP 400 errors, possible causes:"); - System.out.println("1. OTLP protocol version mismatch"); - System.out.println("2. Missing required attributes"); - System.out.println("3. Incorrect protobuf encoding"); - System.out.println("4. Compression issues"); - - } catch (Exception e) { - e.printStackTrace(); - } - } -} diff --git a/examples/src/main/java/dev/braintrust/claude/examples/OpenTelemetryExample.java b/examples/src/main/java/dev/braintrust/claude/examples/OpenTelemetryExample.java deleted file mode 100644 index 6be84b0..0000000 --- a/examples/src/main/java/dev/braintrust/claude/examples/OpenTelemetryExample.java +++ /dev/null @@ -1,436 +0,0 @@ -package dev.braintrust.claude.examples; - -import com.openai.client.OpenAIClient; -import com.openai.client.okhttp.OpenAIOkHttpClient; -import com.openai.models.ChatModel; -import com.openai.models.chat.completions.ChatCompletion; -import com.openai.models.chat.completions.ChatCompletionCreateParams; -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.log.BraintrustLogger; -import dev.braintrust.claude.trace.BraintrustSpanProcessor; -import dev.braintrust.claude.trace.BraintrustTracing; -import io.opentelemetry.api.trace.Span; -import io.opentelemetry.api.trace.SpanKind; -import io.opentelemetry.api.trace.StatusCode; -import io.opentelemetry.context.Scope; -import java.time.Duration; - -/** - * Comprehensive example demonstrating Braintrust with OpenTelemetry and OpenAI. - * - *

This example shows: - Setting up Braintrust with OpenTelemetry - Making real OpenAI API calls - * using the official SDK - Wrapping calls with Braintrust instrumentation - Logging results with - * custom metadata and scores - */ -public class OpenTelemetryExample { - - private static final String OPENAI_API_KEY = System.getenv("OPENAI_API_KEY"); - private static final String BRAINTRUST_API_KEY = System.getenv("BRAINTRUST_API_KEY"); - - public static void main(String[] args) { - // Validate environment variables - if (OPENAI_API_KEY == null || OPENAI_API_KEY.isEmpty()) { - System.err.println("Error: OPENAI_API_KEY environment variable is not set"); - System.exit(1); - } - - if (BRAINTRUST_API_KEY == null || BRAINTRUST_API_KEY.isEmpty()) { - System.err.println("Error: BRAINTRUST_API_KEY environment variable is not set"); - System.exit(1); - } - - try { - - // Initialize Braintrust with OpenTelemetry - // You can set the project ID here or via BRAINTRUST_PROJECT_ID env var - String projectId = System.getenv("BRAINTRUST_PROJECT_ID"); - if (projectId == null || projectId.isEmpty()) { - projectId = "opentelemetry-example"; // Default project ID - } - - var config = - BraintrustConfig.builder() - .apiKey(BRAINTRUST_API_KEY) - .defaultProjectId(projectId) // Set default project - .build(); - - var openTelemetry = - BraintrustTracing.quickstart( - config, - builder -> - builder.serviceName("opentelemetry-example") - .resourceAttribute("environment", "development") - .resourceAttribute("team", "engineering") - .exportInterval(Duration.ofSeconds(2))); - - var tracer = BraintrustTracing.getTracer(openTelemetry); - - // Create OpenAI client - var openAiClient = OpenAIOkHttpClient.fromEnv(); // Uses OPENAI_API_KEY env var - - // Example 1: Simple non-streaming completion - System.out.println("\n=== Example 1: Non-Streaming Completion ==="); - performNonStreamingCompletion(openAiClient, tracer); - - // Example 2: Complex multi-step operation with custom scoring - System.out.println("\n=== Example 2: Multi-Step Operation with Scoring ==="); - performComplexOperation(openAiClient, tracer); - - // Example 3: Error handling demonstration - System.out.println("\n=== Example 3: Error Handling ==="); - demonstrateErrorHandling(openAiClient, tracer); - - // Initialize logger and test it - System.out.println("\n=== Example 4: Braintrust Logger ==="); - demonstrateLogger(); - - // Give time for spans to export - System.out.println("\nWaiting for spans to export..."); - Thread.sleep(5000); - - System.out.println("\nAll examples completed successfully!"); - System.out.println("Check your Braintrust dashboard to see the traces."); - - } catch (Exception e) { - System.err.println("Error running example: " + e.getMessage()); - e.printStackTrace(); - System.exit(1); - } - } - - private static void performNonStreamingCompletion( - OpenAIClient client, io.opentelemetry.api.trace.Tracer tracer) { - // Create a parent span for the operation - var parentSpan = - tracer.spanBuilder("example.non_streaming_completion") - .setSpanKind(SpanKind.INTERNAL) - .setAttribute("example.type", "non_streaming") - .startSpan(); - - try (Scope scope = parentSpan.makeCurrent()) { - // Create chat completion request - var request = - ChatCompletionCreateParams.builder() - .model(ChatModel.GPT_4O) - .addSystemMessage("You are a helpful assistant that writes poetry.") - .addUserMessage("Write a haiku about OpenTelemetry and observability.") - .maxTokens(100L) - .temperature(0.7) - .build(); - - // Wrap the OpenAI call with tracing - var apiSpan = - tracer.spanBuilder("openai.chat.completion") - .setSpanKind(SpanKind.CLIENT) - .setAttribute("openai.model", request.model().toString()) - .setAttribute("openai.max_tokens", request.maxTokens().orElse(0L)) - .setAttribute("openai.temperature", request.temperature().orElse(1.0)) - .startSpan(); - - ChatCompletion response; - try (Scope apiScope = apiSpan.makeCurrent()) { - response = client.chat().completions().create(request); - - // Add usage metrics - if (response.usage().isPresent()) { - var usage = response.usage().get(); - apiSpan.setAttribute( - BraintrustSpanProcessor.USAGE_PROMPT_TOKENS, usage.promptTokens()); - apiSpan.setAttribute( - BraintrustSpanProcessor.USAGE_COMPLETION_TOKENS, - usage.completionTokens()); - apiSpan.setAttribute( - BraintrustSpanProcessor.USAGE_TOTAL_TOKENS, usage.totalTokens()); - - var estimatedCost = - estimateCost("gpt-4o", usage.promptTokens(), usage.completionTokens()); - apiSpan.setAttribute(BraintrustSpanProcessor.USAGE_COST, estimatedCost); - } - - apiSpan.setStatus(StatusCode.OK); - } catch (Exception e) { - apiSpan.recordException(e); - apiSpan.setStatus(StatusCode.ERROR, e.getMessage()); - throw e; - } finally { - apiSpan.end(); - } - - // Log the response - var completion = response.choices().get(0).message().content().orElse(""); - System.out.println("Haiku: " + completion); - - // Add custom metadata and scoring - var currentSpan = Span.current(); - currentSpan.setAttribute("example.haiku", completion); - currentSpan.setAttribute("example.quality_score", evaluateHaikuQuality(completion)); - - // Add Braintrust-specific scoring - BraintrustTracing.SpanUtils.addScore("haiku_quality", evaluateHaikuQuality(completion)); - - parentSpan.setStatus(StatusCode.OK); - - } catch (Exception e) { - parentSpan.recordException(e); - parentSpan.setStatus(StatusCode.ERROR, e.getMessage()); - System.err.println("Error in non-streaming completion: " + e.getMessage()); - } finally { - parentSpan.end(); - } - } - - private static void performComplexOperation( - OpenAIClient client, io.opentelemetry.api.trace.Tracer tracer) { - var rootSpan = - tracer.spanBuilder("example.complex_operation") - .setSpanKind(SpanKind.INTERNAL) - .setAttribute("example.type", "complex") - .startSpan(); - - try (Scope rootScope = rootSpan.makeCurrent()) { - // Step 1: Generate a story outline - var outlineSpan = - tracer.spanBuilder("generate_outline") - .setSpanKind(SpanKind.INTERNAL) - .startSpan(); - - String outline; - try (Scope outlineScope = outlineSpan.makeCurrent()) { - outline = generateStoryOutline(client); - outlineSpan.setAttribute("outline.length", outline.length()); - BraintrustTracing.SpanUtils.addScore("outline_quality", 0.85); - outlineSpan.setStatus(StatusCode.OK); - } finally { - outlineSpan.end(); - } - - // Step 2: Expand the outline into a story - var storySpan = - tracer.spanBuilder("expand_story").setSpanKind(SpanKind.INTERNAL).startSpan(); - - String story; - try (Scope storyScope = storySpan.makeCurrent()) { - story = expandOutlineToStory(client, outline); - storySpan.setAttribute("story.length", story.length()); - storySpan.setAttribute("story.word_count", story.split("\\s+").length); - BraintrustTracing.SpanUtils.addScore("story_coherence", 0.92); - storySpan.setStatus(StatusCode.OK); - } finally { - storySpan.end(); - } - - // Step 3: Generate a title - var titleSpan = - tracer.spanBuilder("generate_title").setSpanKind(SpanKind.INTERNAL).startSpan(); - - String title; - try (Scope titleScope = titleSpan.makeCurrent()) { - title = generateTitle(client, story); - titleSpan.setAttribute("title", title); - BraintrustTracing.SpanUtils.addScore("title_relevance", 0.88); - titleSpan.setStatus(StatusCode.OK); - } finally { - titleSpan.end(); - } - - // Add overall operation metadata - rootSpan.setAttribute("operation.success", true); - rootSpan.setAttribute("operation.total_api_calls", 3); - BraintrustTracing.SpanUtils.addScore("overall_quality", 0.88); - - System.out.println("\nGenerated Story:"); - System.out.println("Title: " + title); - System.out.println( - "\nStory Preview: " - + story.substring(0, Math.min(200, story.length())) - + "..."); - - rootSpan.setStatus(StatusCode.OK); - - } catch (Exception e) { - rootSpan.recordException(e); - rootSpan.setStatus(StatusCode.ERROR, e.getMessage()); - System.err.println("Error in complex operation: " + e.getMessage()); - } finally { - rootSpan.end(); - } - } - - private static void demonstrateErrorHandling( - OpenAIClient client, io.opentelemetry.api.trace.Tracer tracer) { - var span = - tracer.spanBuilder("example.error_handling") - .setSpanKind(SpanKind.INTERNAL) - .setAttribute("example.type", "error_demo") - .startSpan(); - - try (Scope scope = span.makeCurrent()) { - // Intentionally cause an error with invalid parameters - var request = - ChatCompletionCreateParams.builder() - .model(ChatModel.of("invalid-model-name")) - .addUserMessage("This should fail") - .maxTokens(-1L) // Invalid max tokens - .build(); - - try { - var apiSpan = - tracer.spanBuilder("openai.chat.completion.error") - .setSpanKind(SpanKind.CLIENT) - .startSpan(); - - try (Scope apiScope = apiSpan.makeCurrent()) { - client.chat().completions().create(request); - apiSpan.setStatus(StatusCode.OK); - } finally { - apiSpan.end(); - } - - System.out.println("Unexpected success!"); - - } catch (Exception e) { - System.out.println("Expected error occurred: " + e.getMessage()); - span.recordException(e); - span.setAttribute("error.handled", true); - span.setAttribute("error.type", e.getClass().getSimpleName()); - } - - // Demonstrate recovery - span.addEvent("Attempting recovery with valid parameters"); - - var validRequest = - ChatCompletionCreateParams.builder() - .model(ChatModel.GPT_4O) - .addUserMessage("Say 'recovered!'") - .maxTokens(10L) - .build(); - - var recovery = client.chat().completions().create(validRequest); - - System.out.println( - "Recovery successful: " - + recovery.choices().get(0).message().content().orElse("")); - - span.setAttribute("recovery.success", true); - span.setStatus(StatusCode.OK, "Recovered from error"); - - } catch (Exception e) { - span.recordException(e); - span.setStatus(StatusCode.ERROR, e.getMessage()); - System.err.println("Unrecoverable error: " + e.getMessage()); - } finally { - span.end(); - } - } - - private static void demonstrateLogger() { - // Enable debug logging - BraintrustLogger.setDebugEnabled(true); - - // Log various events - BraintrustLogger.info( - "Application started with version={} environment={}", "1.0.0", "development"); - - BraintrustLogger.info( - "API performance metrics: score={} endpoint={} latency_ms={}", - 0.95, - "/api/chat", - 250); - - BraintrustLogger.info( - "Model quality score={} model={} task={}", 0.88, "gpt-4o", "summarization"); - - BraintrustLogger.debug("Debug information: request_id={}", "12345"); - - // Simulate an error - try { - throw new RuntimeException("Simulated error for logging"); - } catch (Exception e) { - BraintrustLogger.error( - "Simulated error occurred: error_type={} handled={}", e, "simulation", true); - } - - System.out.println("Logger demonstration completed"); - } - - // Helper methods - - private static String generateStoryOutline(OpenAIClient client) { - var request = - ChatCompletionCreateParams.builder() - .model(ChatModel.GPT_4O) - .addSystemMessage("You are a creative writing assistant.") - .addUserMessage( - "Create a brief outline for a short story about AI and humanity." - + " Include 3 main plot points.") - .maxTokens(150L) - .temperature(0.8) - .build(); - - var response = client.chat().completions().create(request); - return response.choices().get(0).message().content().orElse(""); - } - - private static String expandOutlineToStory(OpenAIClient client, String outline) { - var request = - ChatCompletionCreateParams.builder() - .model(ChatModel.GPT_4O) - .addSystemMessage("You are a creative writing assistant.") - .addUserMessage("Expand this outline into a 200-word story:\n\n" + outline) - .maxTokens(300L) - .temperature(0.7) - .build(); - - var response = client.chat().completions().create(request); - return response.choices().get(0).message().content().orElse(""); - } - - private static String generateTitle(OpenAIClient client, String story) { - var request = - ChatCompletionCreateParams.builder() - .model(ChatModel.GPT_4O) - .addSystemMessage("You are a creative writing assistant.") - .addUserMessage("Generate a compelling title for this story:\n\n" + story) - .maxTokens(20L) - .temperature(0.9) - .build(); - - var response = client.chat().completions().create(request); - return response.choices().get(0).message().content().orElse("").trim(); - } - - private static double evaluateHaikuQuality(String haiku) { - // Simple heuristic: check if it has approximately 3 lines - var lines = haiku.split("\n"); - if (lines.length >= 3) { - return 0.9; - } else if (lines.length == 2) { - return 0.7; - } else { - return 0.5; - } - } - - private static double estimateCost(String model, long promptTokens, long completionTokens) { - // Simplified pricing based on OpenAI's pricing (as of 2024) - var modelLower = model.toLowerCase(); - double promptCostPer1k, completionCostPer1k; - - if (modelLower.contains("gpt-4o")) { - promptCostPer1k = 0.005; // $5 per 1M input tokens - completionCostPer1k = 0.015; // $15 per 1M output tokens - } else if (modelLower.contains("gpt-4")) { - promptCostPer1k = 0.03; - completionCostPer1k = 0.06; - } else if (modelLower.contains("gpt-3.5-turbo")) { - promptCostPer1k = 0.0005; - completionCostPer1k = 0.0015; - } else { - promptCostPer1k = 0.002; - completionCostPer1k = 0.002; - } - - return (promptTokens / 1000.0 * promptCostPer1k) - + (completionTokens / 1000.0 * completionCostPer1k); - } -} diff --git a/examples/src/main/java/dev/braintrust/claude/examples/ProjectExample.java b/examples/src/main/java/dev/braintrust/claude/examples/ProjectExample.java deleted file mode 100644 index 0caea5e..0000000 --- a/examples/src/main/java/dev/braintrust/claude/examples/ProjectExample.java +++ /dev/null @@ -1,151 +0,0 @@ -package dev.braintrust.claude.examples; - -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.trace.BraintrustContext; -import dev.braintrust.claude.trace.BraintrustTracing; -import io.opentelemetry.api.trace.SpanKind; -import io.opentelemetry.api.trace.StatusCode; -import io.opentelemetry.context.Scope; -import java.time.Duration; - -/** - * Example demonstrating how Braintrust project names work with OpenTelemetry. - * - *

In Braintrust, traces are organized by project. You can set the project in several ways: 1. - * Via span attributes (braintrust.parent.project_id) 2. Via BraintrustContext 3. Via environment - * variable BRAINTRUST_PROJECT_ID - * - *

If no project is specified, Braintrust will create a default project based on the service - * name. - */ -public class ProjectExample { - - public static void main(String[] args) { - try { - var apiKey = System.getenv("BRAINTRUST_API_KEY"); - if (apiKey == null || apiKey.isEmpty()) { - System.out.println("Set BRAINTRUST_API_KEY to see traces in Braintrust dashboard"); - System.exit(1); - } - - var config = BraintrustConfig.builder().apiKey(apiKey).build(); - - // The service name is used as a default if no project is specified - var openTelemetry = - BraintrustTracing.quickstart( - config, - builder -> - builder.serviceName( - "my-ai-application") // This becomes the default - // project name - .resourceAttribute("environment", "production") - .exportInterval(Duration.ofSeconds(2))); - - var tracer = BraintrustTracing.getTracer(openTelemetry); - - System.out.println("=== Braintrust Project Name Examples ===\n"); - - // Example 1: Default project (uses service name) - System.out.println("1. Default project (service name: my-ai-application)"); - var span1 = - tracer.spanBuilder("example.default_project") - .setSpanKind(SpanKind.INTERNAL) - .startSpan(); - - try (Scope scope1 = span1.makeCurrent()) { - span1.setAttribute("info", "This will go to project 'my-ai-application'"); - span1.setStatus(StatusCode.OK); - } finally { - span1.end(); - } - - // Example 2: Explicit project via span attribute - System.out.println("2. Explicit project via span attribute"); - var span2 = - tracer.spanBuilder("example.explicit_project") - .setSpanKind(SpanKind.INTERNAL) - .startSpan(); - - try (Scope scope2 = span2.makeCurrent()) { - // Set project ID explicitly - BraintrustTracing.SpanUtils.setParentProject("customer-support-bot"); - span2.setAttribute("info", "This will go to project 'customer-support-bot'"); - span2.setStatus(StatusCode.OK); - } finally { - span2.end(); - } - - // Example 3: Using BraintrustContext - System.out.println("3. Using BraintrustContext for multiple spans"); - var context = - BraintrustContext.builder().projectId("recommendation-engine").storeInCurrent(); - - try (Scope contextScope = context.makeCurrent()) { - // All spans created within this scope will belong to the recommendation-engine - // project - - var span3a = - tracer.spanBuilder("example.context_span_1") - .setSpanKind(SpanKind.INTERNAL) - .startSpan(); - - try (Scope scope3a = span3a.makeCurrent()) { - span3a.setAttribute("info", "Part of recommendation-engine project"); - span3a.setStatus(StatusCode.OK); - } finally { - span3a.end(); - } - - var span3b = - tracer.spanBuilder("example.context_span_2") - .setSpanKind(SpanKind.INTERNAL) - .startSpan(); - - try (Scope scope3b = span3b.makeCurrent()) { - span3b.setAttribute("info", "Also part of recommendation-engine project"); - span3b.setStatus(StatusCode.OK); - } finally { - span3b.end(); - } - } - - // Example 4: Experiment tracking - System.out.println("4. Experiment tracking"); - var experimentContext = - BraintrustContext.builder() - .experimentId("exp-12345") // Experiments are a special type of project - .storeInCurrent(); - - try (Scope expScope = experimentContext.makeCurrent()) { - var span4 = - tracer.spanBuilder("example.experiment_run") - .setSpanKind(SpanKind.INTERNAL) - .startSpan(); - - try (Scope scope4 = span4.makeCurrent()) { - span4.setAttribute("info", "This is part of experiment exp-12345"); - BraintrustTracing.SpanUtils.addScore("accuracy", 0.95); - span4.setStatus(StatusCode.OK); - } finally { - span4.end(); - } - } - - // Wait for export - System.out.println("\nWaiting for spans to export..."); - Thread.sleep(3000); - - System.out.println("\nExample completed!"); - System.out.println( - "\nIn Braintrust dashboard, you should see traces in these projects:"); - System.out.println("- my-ai-application (default from service name)"); - System.out.println("- customer-support-bot (explicit project)"); - System.out.println("- recommendation-engine (via context)"); - System.out.println("- Experiment exp-12345 (if it exists)"); - - } catch (Exception e) { - System.err.println("Error: " + e.getMessage()); - e.printStackTrace(); - } - } -} diff --git a/examples/src/main/java/dev/braintrust/claude/examples/QuickStartExample.java b/examples/src/main/java/dev/braintrust/claude/examples/QuickStartExample.java deleted file mode 100644 index 3c28e4b..0000000 --- a/examples/src/main/java/dev/braintrust/claude/examples/QuickStartExample.java +++ /dev/null @@ -1,114 +0,0 @@ -package dev.braintrust.claude.examples; - -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.trace.BraintrustTracing; -import io.opentelemetry.api.trace.Span; -import io.opentelemetry.context.Scope; -import java.time.Duration; -import java.util.List; -import java.util.Random; - -/** - * Quick start example showing basic Braintrust tracing setup. Demonstrates idiomatic Java usage - * with try-with-resources and functional style. - */ -public class QuickStartExample { - - public static void main(String[] args) throws Exception { - // Initialize Braintrust with environment configuration - var openTelemetry = BraintrustTracing.quickstart(); - - // Or with custom configuration using builder pattern - var customConfig = - BraintrustConfig.create( - builder -> - builder.apiKey("your-api-key") - .defaultProjectId("my-project") - .enableTraceConsoleLog(true) - .debug(true)); - - var tracer = BraintrustTracing.getTracer(openTelemetry); - - // Example 1: Simple span creation - try (var ignored = tracer.spanBuilder("simple-operation").startSpan().makeCurrent()) { - - System.out.println("Performing simple operation..."); - Thread.sleep(100); - - // Add custom attributes - Span.current().setAttribute("custom.attribute", "value"); - Span.current().setAttribute("custom.count", 42); - } - - // Example 2: Nested spans with functional style - processItems(List.of("item1", "item2", "item3"), tracer); - - // Example 3: Error handling - try (var ignored = tracer.spanBuilder("error-example").startSpan().makeCurrent()) { - - simulateError(); - - } catch (Exception e) { - Span.current().recordException(e); - System.err.println("Caught expected error: " + e.getMessage()); - } - - // Allow time for spans to export - Thread.sleep(2000); - System.out.println("Quick start example completed!"); - } - - private static void processItems(List items, io.opentelemetry.api.trace.Tracer tracer) { - var span = - tracer.spanBuilder("process-items") - .setAttribute("item.count", items.size()) - .startSpan(); - - try (Scope scope = span.makeCurrent()) { - // Use Java streams for functional processing - var results = items.parallelStream().map(item -> processItem(item, tracer)).toList(); - - // Add summary metrics - var successCount = results.stream().filter(r -> r).count(); - span.setAttribute("process.success_count", successCount); - span.setAttribute("process.failure_count", items.size() - successCount); - - } finally { - span.end(); - } - } - - private static boolean processItem(String item, io.opentelemetry.api.trace.Tracer tracer) { - var span = tracer.spanBuilder("process-item").setAttribute("item.name", item).startSpan(); - - try (Scope scope = span.makeCurrent()) { - // Simulate processing with random delay - var random = new Random(); - var duration = Duration.ofMillis(50 + random.nextInt(150)); - Thread.sleep(duration.toMillis()); - - // Simulate random success/failure - var success = random.nextDouble() > 0.2; - span.setAttribute("item.success", success); - - if (!success) { - span.addEvent("Processing failed"); - span.setAttribute("failure.reason", "Random failure for demo"); - span.setAttribute("failure.retry_possible", true); - } - - return success; - - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - span.recordException(e); - return false; - } finally { - span.end(); - } - } - - private static void simulateError() { - throw new RuntimeException("Simulated error for demonstration"); - } -} diff --git a/examples/src/main/java/dev/braintrust/claude/examples/SimpleExample.java b/examples/src/main/java/dev/braintrust/claude/examples/SimpleExample.java deleted file mode 100644 index d15b162..0000000 --- a/examples/src/main/java/dev/braintrust/claude/examples/SimpleExample.java +++ /dev/null @@ -1,234 +0,0 @@ -package dev.braintrust.claude.examples; - -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.log.BraintrustLogger; -import dev.braintrust.claude.trace.BraintrustSpanProcessor; -import dev.braintrust.claude.trace.BraintrustTracing; -import io.opentelemetry.api.trace.SpanKind; -import io.opentelemetry.api.trace.StatusCode; -import io.opentelemetry.context.Scope; -import java.time.Duration; -import java.util.*; -import java.util.concurrent.ThreadLocalRandom; - -/** - * Simple example demonstrating Braintrust with OpenTelemetry. This example doesn't require external - * API keys and shows basic usage. - */ -public class SimpleExample { - - public static void main(String[] args) { - try { - // Initialize Braintrust with OpenTelemetry - // This will use BRAINTRUST_API_KEY if set, or run in demo mode - var configBuilder = BraintrustConfig.builder(); - - // Use API key if available, otherwise use demo mode - var apiKey = System.getenv("BRAINTRUST_API_KEY"); - if (apiKey != null && !apiKey.isEmpty()) { - configBuilder.apiKey(apiKey); - } else { - System.out.println("Running in demo mode (no BRAINTRUST_API_KEY set)"); - configBuilder.apiKey("demo-key"); - configBuilder.apiUrl("http://localhost:8080"); // Won't actually connect - } - - var config = configBuilder.build(); - - var openTelemetry = - BraintrustTracing.quickstart( - config, - builder -> - builder.serviceName("simple-example") - .resourceAttribute("environment", "development") - .resourceAttribute("version", "1.0.0") - .exportInterval(Duration.ofSeconds(2))); - - var tracer = BraintrustTracing.getTracer(openTelemetry); - - // Example 1: Simple operation with tracing - System.out.println("\n=== Example 1: Simple Operation ==="); - performSimpleOperation(tracer); - - // Example 2: Nested operations - System.out.println("\n=== Example 2: Nested Operations ==="); - performNestedOperations(tracer); - - // Example 3: Simulated LLM call - System.out.println("\n=== Example 3: Simulated LLM Call ==="); - performSimulatedLLMCall(tracer); - - // Example 4: Logger demonstration - System.out.println("\n=== Example 4: Logger Demo ==="); - demonstrateLogger(); - - // Give time for spans to export - System.out.println("\nWaiting for spans to export..."); - Thread.sleep(3000); - - System.out.println("\nExample completed successfully!"); - if (apiKey == null || apiKey.isEmpty()) { - System.out.println( - "\n" - + "Note: To see traces in Braintrust, set BRAINTRUST_API_KEY" - + " environment variable"); - } - - } catch (Exception e) { - System.err.println("Error running example: " + e.getMessage()); - e.printStackTrace(); - System.exit(1); - } - } - - private static void performSimpleOperation(io.opentelemetry.api.trace.Tracer tracer) { - var span = - tracer.spanBuilder("simple.operation") - .setSpanKind(SpanKind.INTERNAL) - .setAttribute("operation.type", "calculation") - .startSpan(); - - try (Scope scope = span.makeCurrent()) { - // Simulate some work - var result = calculateSomething(); - - // Add attributes - span.setAttribute("result", result); - span.setAttribute("success", true); - - // Add a custom score - BraintrustTracing.SpanUtils.addScore("accuracy", 0.95); - - System.out.println("Calculation result: " + result); - span.setStatus(StatusCode.OK); - - } catch (Exception e) { - span.recordException(e); - span.setStatus(StatusCode.ERROR, e.getMessage()); - } finally { - span.end(); - } - } - - private static void performNestedOperations(io.opentelemetry.api.trace.Tracer tracer) { - var rootSpan = tracer.spanBuilder("nested.root").setSpanKind(SpanKind.INTERNAL).startSpan(); - - try (Scope rootScope = rootSpan.makeCurrent()) { - // Child operation 1 - var child1Span = - tracer.spanBuilder("nested.child1").setSpanKind(SpanKind.INTERNAL).startSpan(); - - try (Scope child1Scope = child1Span.makeCurrent()) { - Thread.sleep(50); - child1Span.setAttribute("child1.result", "success"); - child1Span.setStatus(StatusCode.OK); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - child1Span.recordException(e); - child1Span.setStatus(StatusCode.ERROR); - } finally { - child1Span.end(); - } - - // Child operation 2 - var child2Span = - tracer.spanBuilder("nested.child2").setSpanKind(SpanKind.INTERNAL).startSpan(); - - try (Scope child2Scope = child2Span.makeCurrent()) { - Thread.sleep(75); - child2Span.setAttribute("child2.result", "success"); - BraintrustTracing.SpanUtils.addScore("performance", 0.88); - child2Span.setStatus(StatusCode.OK); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - child2Span.recordException(e); - child2Span.setStatus(StatusCode.ERROR); - } finally { - child2Span.end(); - } - - System.out.println("Nested operations completed"); - rootSpan.setStatus(StatusCode.OK); - - } finally { - rootSpan.end(); - } - } - - private static void performSimulatedLLMCall(io.opentelemetry.api.trace.Tracer tracer) { - var span = - tracer.spanBuilder("llm.chat.completion") - .setSpanKind(SpanKind.CLIENT) - .setAttribute("llm.model", "gpt-4o") - .setAttribute("llm.temperature", 0.7) - .setAttribute("llm.max_tokens", 100) - .startSpan(); - - try (Scope scope = span.makeCurrent()) { - // Simulate LLM processing time - Thread.sleep(200); - - // Simulate token usage - var promptTokens = 45L; - var completionTokens = 78L; - var totalTokens = promptTokens + completionTokens; - - span.setAttribute(BraintrustSpanProcessor.USAGE_PROMPT_TOKENS, promptTokens); - span.setAttribute(BraintrustSpanProcessor.USAGE_COMPLETION_TOKENS, completionTokens); - span.setAttribute(BraintrustSpanProcessor.USAGE_TOTAL_TOKENS, totalTokens); - - // Simulate cost calculation - var cost = (promptTokens * 0.0005 + completionTokens * 0.0015) / 1000; - span.setAttribute(BraintrustSpanProcessor.USAGE_COST, cost); - - // Add quality scores - BraintrustTracing.SpanUtils.addScore("relevance", 0.92); - BraintrustTracing.SpanUtils.addScore("coherence", 0.89); - - System.out.println("Simulated LLM response generated"); - System.out.println( - "Tokens used: " - + totalTokens - + " (cost: $" - + String.format("%.4f", cost) - + ")"); - - span.setStatus(StatusCode.OK); - - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - span.recordException(e); - span.setStatus(StatusCode.ERROR, e.getMessage()); - } finally { - span.end(); - } - } - - private static void demonstrateLogger() { - // Enable debug logging - BraintrustLogger.setDebugEnabled(true); - - // Various log levels - BraintrustLogger.debug("Debug message: processing_id={}", UUID.randomUUID()); - BraintrustLogger.info("Info message: status={}", "running"); - BraintrustLogger.warn("Warning message: retry_count={}", 3); - - // Log with multiple parameters - BraintrustLogger.info( - "Processing complete: items={} duration_ms={} success_rate={}", 100, 1234, 0.95); - - // Error logging - try { - throw new IllegalArgumentException("Invalid parameter"); - } catch (Exception e) { - BraintrustLogger.error("Error occurred during processing", e); - } - - System.out.println("Logger demonstration completed"); - } - - private static double calculateSomething() { - // Simulate some calculation - return ThreadLocalRandom.current().nextDouble() * 100; - } -} diff --git a/examples/src/main/java/dev/braintrust/claude/examples/SimpleExperimentExample.java b/examples/src/main/java/dev/braintrust/claude/examples/SimpleExperimentExample.java deleted file mode 100644 index 6fee06b..0000000 --- a/examples/src/main/java/dev/braintrust/claude/examples/SimpleExperimentExample.java +++ /dev/null @@ -1,148 +0,0 @@ -package dev.braintrust.claude.examples; - -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.eval.Evaluation; -import dev.braintrust.claude.trace.BraintrustTracing; -import java.util.List; -import java.util.Random; - -/** - * Simple experiment example that doesn't require external APIs. Demonstrates evaluating a basic - * string matching algorithm. - */ -public class SimpleExperimentExample { - - // Test case for string matching - record MatchTestCase(String text, String pattern, boolean shouldMatch) {} - - // Output from our matcher - record MatchOutput(boolean matched, long durationNanos, String algorithm) {} - - public static void main(String[] args) { - // Initialize Braintrust - var config = BraintrustConfig.fromEnvironment(); - var openTelemetry = BraintrustTracing.quickstart(config); - - // Define test cases - var testCases = - List.of( - new MatchTestCase("hello world", "world", true), - new MatchTestCase("hello world", "World", false), // Case sensitive - new MatchTestCase("the quick brown fox", "quick", true), - new MatchTestCase("the quick brown fox", "slow", false), - new MatchTestCase("OpenTelemetry is great", "Telemetry", true), - new MatchTestCase("Braintrust evaluation", "eval", true), - new MatchTestCase("", "test", false), - new MatchTestCase("test", "", true), // Empty pattern matches - new MatchTestCase("abcdefghijklmnop", "mnop", true), - new MatchTestCase("performance test", "perf", true)); - - // Run experiment with simple contains - System.out.println("=== Experiment 1: Simple Contains Matcher ==="); - runExperiment( - "String Matcher - Contains", - testCases, - SimpleExperimentExample::simpleContainsMatcher); - - // Run experiment with regex matcher - System.out.println("\n=== Experiment 2: Regex Matcher ==="); - runExperiment("String Matcher - Regex", testCases, SimpleExperimentExample::regexMatcher); - - // Run experiment with fuzzy matcher (simulated) - System.out.println("\n=== Experiment 3: Fuzzy Matcher ==="); - runExperiment("String Matcher - Fuzzy", testCases, SimpleExperimentExample::fuzzyMatcher); - - System.out.println("\n=== All Experiments Complete ==="); - System.out.println("Check your Braintrust dashboard to compare results!"); - - // Wait for export - try { - Thread.sleep(3000); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - } - } - - private static void runExperiment(String name, List testCases, Matcher matcher) { - var evaluation = - Evaluation.builder() - .name(name) - .data(testCases) - .task(testCase -> matcher.match(testCase)) - .scorer( - "accuracy", - (testCase, output) -> { - var correct = output.matched == testCase.shouldMatch; - return correct ? 1.0 : 0.0; - }) - .build(); - - System.out.println("Running: " + name); - var results = evaluation.run(); - var summary = results.summary(); - - // Calculate average score from the "accuracy" scorer - var accuracyStats = summary.scoreStatistics().get("accuracy"); - var avgScore = accuracyStats != null ? accuracyStats.mean() : 0.0; - - System.out.printf( - "Results: %d/%d correct (%.1f%%)%n", - (int) (avgScore * testCases.size()), testCases.size(), avgScore * 100); - } - - // Different matcher implementations - - interface Matcher { - MatchOutput match(MatchTestCase testCase); - } - - private static MatchOutput simpleContainsMatcher(MatchTestCase testCase) { - var start = System.nanoTime(); - var matched = testCase.text.contains(testCase.pattern); - var duration = System.nanoTime() - start; - - // Simulate some processing time - simulateWork(10); - - return new MatchOutput(matched, duration, "contains"); - } - - private static MatchOutput regexMatcher(MatchTestCase testCase) { - var start = System.nanoTime(); - var matched = testCase.text.matches(".*" + testCase.pattern + ".*"); - var duration = System.nanoTime() - start; - - // Simulate some processing time - simulateWork(20); - - return new MatchOutput(matched, duration, "regex"); - } - - private static MatchOutput fuzzyMatcher(MatchTestCase testCase) { - var start = System.nanoTime(); - - // Simple fuzzy matching: allow case-insensitive and partial matches - var matched = testCase.text.toLowerCase().contains(testCase.pattern.toLowerCase()); - - // Add some randomness to simulate fuzzy matching confidence - if (matched && testCase.pattern.length() > 3) { - // Sometimes "miss" long patterns to simulate fuzzy behavior - matched = new Random().nextDouble() > 0.1; - } - - var duration = System.nanoTime() - start; - - // Simulate more processing time for fuzzy matching - simulateWork(50); - - return new MatchOutput(matched, duration, "fuzzy"); - } - - private static void simulateWork(int microseconds) { - try { - Thread.sleep(0, microseconds * 1000); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - } - } -} diff --git a/examples/src/main/java/dev/braintrust/claude/examples/SimpleExperimentWithRegistration.java b/examples/src/main/java/dev/braintrust/claude/examples/SimpleExperimentWithRegistration.java deleted file mode 100644 index 659d49d..0000000 --- a/examples/src/main/java/dev/braintrust/claude/examples/SimpleExperimentWithRegistration.java +++ /dev/null @@ -1,227 +0,0 @@ -package dev.braintrust.claude.examples; - -import dev.braintrust.claude.api.Experiment; -import dev.braintrust.claude.api.Project; -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.eval.Evaluation; -import dev.braintrust.claude.trace.BraintrustContext; -import dev.braintrust.claude.trace.BraintrustTracing; -import io.opentelemetry.context.Context; -import java.util.List; -import java.util.concurrent.TimeUnit; - -/** - * Example showing the same DX as the Go SDK with automatic experiment registration. This mirrors - * the Go examples where RegisterExperiment is called directly. - */ -public class SimpleExperimentWithRegistration { - - record MathProblem(String question, double expected) {} - - record Answer(double result, long computeTimeMs) {} - - public static void main(String[] args) throws Exception { - // Initialize Braintrust - var config = BraintrustConfig.fromEnvironment(); - var openTelemetry = BraintrustTracing.quickstart(config); - - System.out.println("=== Braintrust Experiment Example (Go SDK Style) ===\n"); - - // Step 1: Register a project (same as Go: api.RegisterProject) - System.out.println("Registering project..."); - var project = Project.registerProjectSync("Math Solver Evaluation"); - System.out.println("Project ID: " + project.id()); - - // Step 2: Register an experiment (same as Go: api.RegisterExperiment) - var experimentName = "accuracy-test-" + System.currentTimeMillis(); - System.out.println("\nRegistering experiment..."); - var experiment = Experiment.registerExperimentSync(experimentName, project.id()); - System.out.println("Experiment ID: " + experiment.id()); - - // Get the app URL from config - var appUrl = config.appUrl().toString().replaceAll("/$", ""); - // URLEncoder.encode uses + for spaces, but we need %20 for URL paths - var projectNameEncoded = - java.net.URLEncoder.encode(project.name(), java.nio.charset.StandardCharsets.UTF_8) - .replace("+", "%20"); - - // For staging, we need to handle the org differently - var baseExperimentUrl = - 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); - - // TODO: Add comparison parameter if there's a previous experiment - var experimentUrl = baseExperimentUrl; - - System.out.println("\nExperiment " + experimentName + " is running at " + experimentUrl); - - // Step 3: Set up tracing context with experiment (similar to Go's trace.SetParent) - var braintrustContext = BraintrustContext.forExperiment(experiment.id()); - var context = braintrustContext.storeInContext(Context.current()).makeCurrent(); - - // Step 4: Define test data - var problems = - List.of( - new MathProblem("What is 2 + 2?", 4.0), - new MathProblem("What is 10 * 5?", 50.0), - new MathProblem("What is 100 / 4?", 25.0), - new MathProblem("What is 7 * 8?", 56.0), - new MathProblem("What is 15 - 9?", 6.0)); - - // Step 5: Run evaluation with the experiment - System.out.println("\nRunning evaluation..."); - System.out.printf("Math Problem Solver (data): %d items%n", problems.size()); - - var evaluation = - Evaluation.builder() - .name("Math Problem Solver") - .data(problems) - .task(SimpleExperimentWithRegistration::solveMathProblem) - .scorer( - "accuracy", - (problem, answer) -> - Math.abs(answer.result - problem.expected) < 0.0001 - ? 1.0 - : 0.0) - .scorer( - "speed", - (problem, answer) -> - answer.computeTimeMs < 10 - ? 1.0 - : answer.computeTimeMs < 50 ? 0.5 : 0.0) - .experimentId(experiment.id()) // Link to the registered experiment - .build(); - - var startTime = System.currentTimeMillis(); - var results = evaluation.run(); - var endTime = System.currentTimeMillis(); - - System.out.printf( - "Math Problem Solver (tasks): %d/%d completed in %.2fs%n", - problems.size(), problems.size(), (endTime - startTime) / 1000.0); - - // Step 6: Print results - System.out.println("\n========================= SUMMARY ========================="); - System.out.println(experimentName + ":"); - - var summary = results.summary(); - - // Score summaries - summary.scoreStatistics() - .forEach( - (scorer, stats) -> { - System.out.printf( - "%.2f%% '%s' score (mean: %.3f, min: %.3f, max: %.3f)%n", - stats.mean() * 100, - scorer, - stats.mean(), - stats.min(), - stats.max()); - }); - - System.out.println(); - System.out.printf( - "%.2fs duration (%.2fs total across %d items)%n", - (endTime - startTime) / 1000.0, - summary.averageDuration().toMillis() * summary.totalCount() / 1000.0, - summary.totalCount()); - - if (summary.errorCount() > 0) { - System.out.printf( - "%d errors (%.1f%% error rate)%n", - summary.errorCount(), (summary.errorCount() * 100.0) / summary.totalCount()); - } - - System.out.println("\n=== View in Braintrust ==="); - System.out.println("See results for " + experimentName + " at " + experimentUrl); - - // Clean up context - context.close(); - - // Force flush all spans and shutdown properly - System.out.println("\nFlushing traces to Braintrust..."); - if (openTelemetry instanceof io.opentelemetry.sdk.OpenTelemetrySdk sdk) { - // First force flush to ensure all spans are sent - var flushResult = sdk.getSdkTracerProvider().forceFlush(); - if (!flushResult.join(10, TimeUnit.SECONDS).isSuccess()) { - System.err.println("Warning: Some spans may not have been flushed"); - } - - // Give a bit more time for any in-flight exports to complete - Thread.sleep(1000); - - // Now shutdown gracefully - var shutdownResult = sdk.getSdkTracerProvider().shutdown(); - if (shutdownResult.join(10, TimeUnit.SECONDS).isSuccess()) { - System.out.println("Traces flushed and exporter shut down successfully!"); - } else { - System.err.println("Warning: Shutdown may not have completed cleanly"); - } - } else { - // Fallback: just wait a bit for traces to be exported - Thread.sleep(5000); - System.out.println("Waited for trace export."); - } - } - - // Even simpler: Use the convenience method that does both project and experiment - public static void demonstrateSimplestApproach() throws Exception { - System.out.println("\n=== Simplest Approach (like Go's GetOrCreateExperiment) ==="); - - var problems = - List.of( - new MathProblem("What is 5 + 5?", 10.0), - new MathProblem("What is 3 * 3?", 9.0)); - - // This single line replaces project + experiment registration! - var evaluation = - Evaluation.builder() - .name("Quick Math Test") - .data(problems) - .task(SimpleExperimentWithRegistration::solveMathProblem) - .scorer( - "accuracy", - (problem, answer) -> - Math.abs(answer.result - problem.expected) < 0.0001 - ? 1.0 - : 0.0) - // This automatically creates project + experiment (like Go's helper) - .experiment("quick-math-experiment", "Quick Math Project") - .build(); - - var results = evaluation.run(); - System.out.println("Evaluation complete! Check Braintrust dashboard."); - } - - private static Answer solveMathProblem(MathProblem problem) { - var start = System.currentTimeMillis(); - - // Parse and solve the math problem (simplified) - var result = - switch (problem.question) { - case "What is 2 + 2?" -> 4.0; - case "What is 10 * 5?" -> 50.0; - case "What is 100 / 4?" -> 25.0; - case "What is 7 * 8?" -> 56.0; - case "What is 15 - 9?" -> 6.0; - case "What is 5 + 5?" -> 10.0; - case "What is 3 * 3?" -> 9.0; - default -> 0.0; // Wrong answer for unknown problems - }; - - // Simulate some computation time - try { - Thread.sleep((long) (Math.random() * 20)); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - } - - var computeTime = System.currentTimeMillis() - start; - return new Answer(result, computeTime); - } -} diff --git a/examples/src/main/java/dev/braintrust/claude/examples/TestApiEndpoints.java b/examples/src/main/java/dev/braintrust/claude/examples/TestApiEndpoints.java deleted file mode 100644 index 0e5fa92..0000000 --- a/examples/src/main/java/dev/braintrust/claude/examples/TestApiEndpoints.java +++ /dev/null @@ -1,76 +0,0 @@ -package dev.braintrust.claude.examples; - -import dev.braintrust.claude.config.BraintrustConfig; -import java.net.URI; -import java.net.http.HttpClient; -import java.net.http.HttpRequest; -import java.net.http.HttpResponse; - -public class TestApiEndpoints { - public static void main(String[] args) throws Exception { - var config = BraintrustConfig.fromEnvironment(); - var httpClient = HttpClient.newHttpClient(); - - System.out.println("=== Testing Braintrust API Endpoints ==="); - System.out.println("Base URL: " + config.apiUrl()); - System.out.println(); - - // Test various endpoints to see which ones exist - String[] endpoints = { - "/projects", - "/v1/projects", - "/v1/project", // Go SDK uses singular - "/experiments", - "/v1/experiments", - "/v1/experiment", // Go SDK uses singular - "/v1/dataset", - "/otel/v1/traces" - }; - - for (String endpoint : endpoints) { - testEndpoint(httpClient, config, endpoint); - } - } - - private static void testEndpoint( - HttpClient httpClient, BraintrustConfig config, String endpoint) { - try { - var request = - HttpRequest.newBuilder() - .uri(URI.create(config.apiUrl() + endpoint)) - .header("Authorization", "Bearer " + config.apiKey()) - .header("Accept", "application/json") - .GET() - .build(); - - var response = httpClient.send(request, HttpResponse.BodyHandlers.ofString()); - - System.out.printf( - "%-25s -> %d %s%n", - endpoint, - response.statusCode(), - response.statusCode() == 200 - ? "✓" - : response.statusCode() == 403 - ? "(403: " + response.body() + ")" - : response.statusCode() == 401 - ? "(unauthorized)" - : "(" + extractError(response.body()) + ")"); - - } catch (Exception e) { - System.out.printf("%-25s -> ERROR: %s%n", endpoint, e.getMessage()); - } - } - - private static String extractError(String body) { - // Try to extract error message from response - if (body.contains("\"error\"")) { - int start = body.indexOf("\"error\""); - int end = body.indexOf("}", start); - if (start > 0 && end > start) { - return body.substring(start, Math.min(end + 1, body.length())); - } - } - return body.length() > 50 ? body.substring(0, 50) + "..." : body; - } -} diff --git a/examples/src/main/java/dev/braintrust/claude/examples/TriviaEvaluation.java b/examples/src/main/java/dev/braintrust/claude/examples/TriviaEvaluation.java deleted file mode 100644 index f44afda..0000000 --- a/examples/src/main/java/dev/braintrust/claude/examples/TriviaEvaluation.java +++ /dev/null @@ -1,245 +0,0 @@ -package dev.braintrust.claude.examples; - -import com.openai.client.okhttp.OpenAIOkHttpClient; -import com.openai.models.ChatModel; -import com.openai.models.chat.completions.ChatCompletionCreateParams; -import dev.braintrust.claude.api.Experiment; -import dev.braintrust.claude.api.Project; -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.eval.Evaluation; -import dev.braintrust.claude.openai.BraintrustOpenAI; -import dev.braintrust.claude.trace.BraintrustContext; -import dev.braintrust.claude.trace.BraintrustTracing; -import io.opentelemetry.api.OpenTelemetry; -import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.context.Context; -import java.util.List; -import java.util.concurrent.TimeUnit; - -/** - * Example showing a real evaluation with OpenAI API calls and scoring. This evaluates an LLM's - * ability to answer trivia questions correctly. - */ -public class TriviaEvaluation { - - record TriviaQuestion(String question, String expectedAnswer) {} - - record LLMResponse(String answer, long latencyMs, int promptTokens, int completionTokens) {} - - private static OpenTelemetry openTelemetry; - private static Tracer tracer; - private static BraintrustOpenAI wrappedOpenAI; - - public static void main(String[] args) throws Exception { - // Initialize Braintrust - var config = BraintrustConfig.fromEnvironment(); - openTelemetry = BraintrustTracing.quickstart(config); - tracer = openTelemetry.getTracer("trivia-evaluation"); - - System.out.println("=== Braintrust Trivia Evaluation ===\n"); - - // Initialize OpenAI client - var openAIKey = System.getenv("OPENAI_API_KEY"); - if (openAIKey == null || openAIKey.isEmpty()) { - System.err.println("Error: OPENAI_API_KEY environment variable not set!"); - System.exit(1); - } - - var openai = OpenAIOkHttpClient.builder().apiKey(openAIKey).build(); - wrappedOpenAI = BraintrustOpenAI.wrap(openai, tracer); - - // Step 1: Register project - System.out.println("Registering project..."); - var project = Project.registerProjectSync("Trivia Bot Evaluation"); - System.out.println("Project ID: " + project.id()); - - // Step 2: Register experiment - var experimentName = "trivia-eval-" + System.currentTimeMillis(); - System.out.println("\nRegistering experiment..."); - var experiment = Experiment.registerExperimentSync(experimentName, project.id()); - System.out.println("Experiment ID: " + experiment.id()); - - // 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 = - String.format( - "%s/app/%s/p/%s/experiments/%s", - appUrl, orgNameEncoded, projectNameEncoded, experimentName); - - System.out.println("\nExperiment " + experimentName + " is running at " + experimentUrl); - - // Step 3: Set up tracing context - var braintrustContext = BraintrustContext.forExperiment(experiment.id()); - var context = braintrustContext.storeInContext(Context.current()).makeCurrent(); - - // Step 4: Define trivia questions dataset - var questions = - List.of( - new TriviaQuestion("Which country has the largest population?", "China"), - new TriviaQuestion("What is the capital of France?", "Paris"), - new TriviaQuestion("Who painted the Mona Lisa?", "Leonardo da Vinci"), - new TriviaQuestion( - "What is the largest planet in our solar system?", "Jupiter"), - new TriviaQuestion("In what year did World War II end?", "1945")); - - // Step 5: Run evaluation - System.out.println("\nRunning evaluation..."); - System.out.printf("Trivia Bot (data): %d questions%n", questions.size()); - - var evaluation = - Evaluation.builder() - .name("Trivia Bot") - .data(questions) - .task(TriviaEvaluation::askLLM) - .scorer( - "accuracy", - (question, response) -> { - if (response == null || response.answer == null) { - return 0.0; - } - // Simple string comparison (case-insensitive) - var normalizedExpected = - question.expectedAnswer.toLowerCase().trim(); - var normalizedActual = response.answer.toLowerCase().trim(); - return normalizedActual.contains(normalizedExpected) - ? 1.0 - : 0.0; - }) - .scorer( - "response_time", - (question, response) -> { - if (response == null) return 0.0; - // Score based on response time (faster is better) - if (response.latencyMs < 1000) return 1.0; - if (response.latencyMs < 2000) return 0.5; - return 0.0; - }) - .experimentId(experiment.id()) - .build(); - - var startTime = System.currentTimeMillis(); - var results = evaluation.run(); - var endTime = System.currentTimeMillis(); - - System.out.printf( - "Trivia Bot (tasks): %d/%d completed in %.2fs%n", - questions.size(), questions.size(), (endTime - startTime) / 1000.0); - - // Step 6: Print results - System.out.println("\n========================= SUMMARY ========================="); - System.out.println(experimentName + ":"); - - var summary = results.summary(); - - // Score summaries - summary.scoreStatistics() - .forEach( - (scorer, stats) -> { - System.out.printf( - "%.2f%% '%s' score (mean: %.3f, min: %.3f, max: %.3f)%n", - stats.mean() * 100, - scorer, - stats.mean(), - stats.min(), - stats.max()); - }); - - System.out.println(); - - // Calculate total tokens - var totalPromptTokens = - results.results().stream() - .mapToInt(r -> r.output() != null ? r.output().promptTokens : 0) - .sum(); - var totalCompletionTokens = - results.results().stream() - .mapToInt(r -> r.output() != null ? r.output().completionTokens : 0) - .sum(); - - System.out.printf( - "%.2fs duration (%.2fs average per question)%n", - (endTime - startTime) / 1000.0, summary.averageDuration().toMillis() / 1000.0); - - System.out.printf( - "%d prompt tokens, %d completion tokens (%d total)%n", - totalPromptTokens, - totalCompletionTokens, - totalPromptTokens + totalCompletionTokens); - - if (summary.errorCount() > 0) { - System.out.printf( - "%d errors (%.1f%% error rate)%n", - summary.errorCount(), (summary.errorCount() * 100.0) / summary.totalCount()); - } - - System.out.println("\n=== View in Braintrust ==="); - System.out.println("See results for " + experimentName + " at " + experimentUrl); - - // Clean up - context.close(); - - // Flush traces - System.out.println("\nFlushing traces to Braintrust..."); - if (openTelemetry instanceof io.opentelemetry.sdk.OpenTelemetrySdk sdk) { - var flushResult = sdk.getSdkTracerProvider().forceFlush(); - if (!flushResult.join(10, TimeUnit.SECONDS).isSuccess()) { - System.err.println("Warning: Some spans may not have been flushed"); - } - Thread.sleep(1000); - var shutdownResult = sdk.getSdkTracerProvider().shutdown(); - if (shutdownResult.join(10, TimeUnit.SECONDS).isSuccess()) { - System.out.println("Traces flushed and exporter shut down successfully!"); - } - } - } - - private static LLMResponse askLLM(TriviaQuestion question) { - var startTime = System.currentTimeMillis(); - - try { - var systemMessage = - "You are a helpful assistant that answers trivia questions. Give" - + " brief, direct answers without explanation."; - - var request = - ChatCompletionCreateParams.builder() - .model(ChatModel.GPT_4O_MINI) // Using GPT-4o-mini for cost efficiency - .addSystemMessage(systemMessage) - .addUserMessage(question.question()) - .maxTokens(50L) - .temperature(0.0) // Deterministic for consistent results - .build(); - - // Create message info for tracing - var messages = - List.of( - BraintrustOpenAI.MessageInfo.system(systemMessage), - BraintrustOpenAI.MessageInfo.user(question.question())); - - // Use the wrapped client which will handle all tracing - var response = wrappedOpenAI.chat().completions().create(request, messages); - var endTime = System.currentTimeMillis(); - - var answer = response.choices().get(0).message().content().orElse(""); - var usage = response.usage().orElse(null); - - return new LLMResponse( - answer, - endTime - startTime, - usage != null ? (int) usage.promptTokens() : 0, - usage != null ? (int) usage.completionTokens() : 0); - - } catch (Exception e) { - System.err.println("Error calling OpenAI: " + e.getMessage()); - return new LLMResponse(null, System.currentTimeMillis() - startTime, 0, 0); - } - } -} diff --git a/src/main/java/dev/braintrust/claude/api/ApiException.java b/src/main/java/dev/braintrust/claude/api/ApiException.java deleted file mode 100644 index 1ec80a0..0000000 --- a/src/main/java/dev/braintrust/claude/api/ApiException.java +++ /dev/null @@ -1,13 +0,0 @@ -package dev.braintrust.claude.api; - -/** Exception thrown when API operations fail. */ -public class ApiException extends RuntimeException { - - public ApiException(String message) { - super(message); - } - - public ApiException(String message, Throwable cause) { - super(message, cause); - } -} diff --git a/src/main/java/dev/braintrust/claude/api/BraintrustApiClient.java b/src/main/java/dev/braintrust/claude/api/BraintrustApiClient.java deleted file mode 100644 index 2f2824b..0000000 --- a/src/main/java/dev/braintrust/claude/api/BraintrustApiClient.java +++ /dev/null @@ -1,280 +0,0 @@ -package dev.braintrust.claude.api; - -import com.fasterxml.jackson.annotation.JsonInclude; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.databind.PropertyNamingStrategies; -import com.fasterxml.jackson.datatype.jdk8.Jdk8Module; -import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.log.BraintrustLogger; -import java.io.IOException; -import java.net.URI; -import java.net.http.HttpClient; -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.concurrent.CompletableFuture; -import java.util.concurrent.CompletionException; - -/** - * Modern API client for Braintrust using Java 11+ HTTP client. Provides both synchronous and - * asynchronous methods. - */ -public class BraintrustApiClient implements AutoCloseable { - - private final BraintrustConfig config; - private final HttpClient httpClient; - private final ObjectMapper objectMapper; - - public BraintrustApiClient(BraintrustConfig config) { - this(config, createDefaultHttpClient(config)); - } - - public BraintrustApiClient(BraintrustConfig config, HttpClient httpClient) { - this.config = config; - this.httpClient = httpClient; - this.objectMapper = createObjectMapper(); - } - - /** Creates or gets a project by name. */ - public CompletableFuture createProject(String name) { - var request = new CreateProjectRequest(name); - return postAsync("/v1/project", request, Project.class); - } - - /** Gets a project by ID. */ - public CompletableFuture> getProject(String projectId) { - return getAsync("/v1/project/" + projectId, Project.class) - .handle( - (project, error) -> { - if (error != null && isNotFound(error)) { - return Optional.empty(); - } - if (error != null) { - throw new CompletionException(error); - } - return Optional.of(project); - }); - } - - /** Lists all projects. */ - public CompletableFuture> listProjects() { - return getAsync("/v1/project", ProjectList.class).thenApply(ProjectList::projects); - } - - /** Creates an experiment. */ - public CompletableFuture createExperiment(CreateExperimentRequest request) { - return postAsync("/v1/experiment", request, Experiment.class); - } - - /** Gets an experiment by ID. */ - public CompletableFuture> getExperiment(String experimentId) { - return getAsync("/v1/experiment/" + experimentId, Experiment.class) - .handle( - (experiment, error) -> { - if (error != null && isNotFound(error)) { - return Optional.empty(); - } - if (error != null) { - throw new CompletionException(error); - } - return Optional.of(experiment); - }); - } - - /** Lists experiments for a project. */ - public CompletableFuture> listExperiments(String projectId) { - return getAsync("/v1/experiment?project_id=" + projectId, ExperimentList.class) - .thenApply(ExperimentList::experiments); - } - - /** Creates a dataset. */ - public CompletableFuture createDataset(CreateDatasetRequest request) { - return postAsync("/v1/dataset", request, Dataset.class); - } - - /** Inserts events into a dataset. */ - public CompletableFuture insertDatasetEvents( - String datasetId, List events) { - var request = new InsertEventsRequest(events); - return postAsync( - "/v1/dataset/" + datasetId + "/insert", request, InsertEventsResponse.class); - } - - /** Lists datasets for a project. */ - public CompletableFuture> listDatasets(String projectId) { - return getAsync("/v1/dataset?project_id=" + projectId, DatasetList.class) - .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) { - var request = - HttpRequest.newBuilder() - .uri(URI.create(config.apiUrl() + path)) - .header("Authorization", "Bearer " + config.apiKey()) - .header("Accept", "application/json") - .timeout(config.requestTimeout()) - .GET() - .build(); - - return sendAsync(request, responseType); - } - - private CompletableFuture postAsync(String path, Object body, Class responseType) { - try { - var jsonBody = objectMapper.writeValueAsString(body); - - var request = - HttpRequest.newBuilder() - .uri(URI.create(config.apiUrl() + path)) - .header("Authorization", "Bearer " + config.apiKey()) - .header("Content-Type", "application/json") - .header("Accept", "application/json") - .timeout(config.requestTimeout()) - .POST(HttpRequest.BodyPublishers.ofString(jsonBody)) - .build(); - - return sendAsync(request, responseType); - } catch (IOException e) { - return CompletableFuture.failedFuture( - new ApiException("Failed to serialize request body", e)); - } - } - - private CompletableFuture sendAsync(HttpRequest request, Class responseType) { - BraintrustLogger.debug("API Request: {} {}", request.method(), request.uri()); - - return httpClient - .sendAsync(request, HttpResponse.BodyHandlers.ofString()) - .thenApply(response -> handleResponse(response, responseType)); - } - - private T handleResponse(HttpResponse response, Class responseType) { - BraintrustLogger.debug("API Response: {} - {}", response.statusCode(), response.body()); - - if (response.statusCode() >= 200 && response.statusCode() < 300) { - try { - return objectMapper.readValue(response.body(), responseType); - } catch (IOException e) { - BraintrustLogger.warn("Failed to parse response body", e); - throw new ApiException("Failed to parse response body", e); - } - } else { - BraintrustLogger.warn( - "API request failed with status {}: {}", - response.statusCode(), - response.body()); - throw new ApiException( - String.format( - "API request failed with status %d: %s", - response.statusCode(), response.body())); - } - } - - private boolean isNotFound(Throwable error) { - if (error instanceof ApiException) { - return ((ApiException) error).getMessage().contains("404"); - } - return false; - } - - private static HttpClient createDefaultHttpClient(BraintrustConfig config) { - return HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build(); - } - - private static ObjectMapper createObjectMapper() { - return new ObjectMapper() - .registerModule(new JavaTimeModule()) - .registerModule(new Jdk8Module()) // For Optional support - .setPropertyNamingStrategy(PropertyNamingStrategies.SNAKE_CASE) - .setSerializationInclusion( - JsonInclude.Include.NON_ABSENT) // Skip null and absent Optional - .configure( - com.fasterxml.jackson.databind.DeserializationFeature - .FAIL_ON_UNKNOWN_PROPERTIES, - false); // Ignore unknown fields from API - } - - @Override - public void close() { - // HttpClient doesn't need explicit closing in Java 11+ - } - - // Request/Response DTOs - - public record CreateProjectRequest(String name) {} - - public record Project( - String id, String name, String orgId, String createdAt, String updatedAt) {} - - private record ProjectList(List projects) {} - - private record ExperimentList(List experiments) {} - - public record CreateExperimentRequest( - String projectId, - String name, - Optional description, - Optional baseExperimentId) { - - public CreateExperimentRequest(String projectId, String name) { - this(projectId, name, Optional.empty(), Optional.empty()); - } - } - - public record Experiment( - String id, - String projectId, - String name, - Optional description, - String createdAt, - String updatedAt) {} - - public record CreateDatasetRequest( - String projectId, String name, Optional description) { - public CreateDatasetRequest(String projectId, String name) { - this(projectId, name, Optional.empty()); - } - } - - public record Dataset( - String id, - String projectId, - String name, - Optional description, - String createdAt, - String updatedAt) {} - - private record DatasetList(List datasets) {} - - public record DatasetEvent(Object input, Optional output, Optional metadata) { - public DatasetEvent(Object input) { - this(input, Optional.empty(), Optional.empty()); - } - - public DatasetEvent(Object input, Object output) { - this(input, Optional.of(output), Optional.empty()); - } - } - - 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/claude/api/Experiment.java b/src/main/java/dev/braintrust/claude/api/Experiment.java deleted file mode 100644 index cc7cc46..0000000 --- a/src/main/java/dev/braintrust/claude/api/Experiment.java +++ /dev/null @@ -1,120 +0,0 @@ -package dev.braintrust.claude.api; - -import dev.braintrust.claude.config.BraintrustConfig; -import java.util.Optional; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.CompletionException; - -/** High-level API for experiment management, matching the Go SDK's DX. */ -public final class Experiment { - - private Experiment() { - // Utility class - } - - /** - * Registers a new experiment in Braintrust. This creates a new experiment every time (matching - * Go's EnsureNew: true behavior). - * - * @param name The name of the experiment - * @param projectId The ID of the project to create the experiment in - * @return The created experiment - */ - public static CompletableFuture registerExperiment( - String name, String projectId) { - var config = BraintrustConfig.fromEnvironment(); - var client = new BraintrustApiClient(config); - - var request = - new BraintrustApiClient.CreateExperimentRequest( - projectId, name, Optional.empty(), Optional.empty()); - - return client.createExperiment(request); - } - - /** - * Registers a new experiment in Braintrust with custom configuration. - * - * @param name The name of the experiment - * @param projectId The ID of the project to create the experiment in - * @param config Custom Braintrust configuration - * @return The created experiment - */ - public static CompletableFuture registerExperiment( - String name, String projectId, BraintrustConfig config) { - var client = new BraintrustApiClient(config); - - var request = - new BraintrustApiClient.CreateExperimentRequest( - projectId, name, Optional.empty(), Optional.empty()); - - return client.createExperiment(request); - } - - /** - * Gets or creates an experiment (convenience method). This first creates/gets the project, then - * creates the experiment. - * - * @param experimentName The name of the experiment - * @param projectName The name of the project - * @return The experiment ID - */ - public static CompletableFuture getOrCreateExperiment( - String experimentName, String projectName) { - var config = BraintrustConfig.fromEnvironment(); - return getOrCreateExperiment(experimentName, projectName, config); - } - - /** - * Gets or creates an experiment with custom configuration. - * - * @param experimentName The name of the experiment - * @param projectName The name of the project - * @param config Custom Braintrust configuration - * @return The experiment ID - */ - public static CompletableFuture getOrCreateExperiment( - String experimentName, String projectName, BraintrustConfig config) { - var client = new BraintrustApiClient(config); - - // First, get or create the project - return Project.registerProject(projectName, config) - .thenCompose( - project -> { - // Then create the experiment - var request = - new BraintrustApiClient.CreateExperimentRequest( - project.id(), - experimentName, - Optional.empty(), - Optional.empty()); - return client.createExperiment(request); - }) - .thenApply(experiment -> experiment.id()); - } - - /** Synchronous version of registerExperiment for simpler usage. */ - public static BraintrustApiClient.Experiment registerExperimentSync( - String name, String projectId) { - try { - return registerExperiment(name, projectId).get(); - } catch (Exception e) { - if (e instanceof CompletionException && e.getCause() != null) { - throw new RuntimeException("Failed to register experiment", e.getCause()); - } - throw new RuntimeException("Failed to register experiment", e); - } - } - - /** Synchronous version of getOrCreateExperiment. */ - public static String getOrCreateExperimentSync(String experimentName, String projectName) { - try { - return getOrCreateExperiment(experimentName, projectName).get(); - } catch (Exception e) { - if (e instanceof CompletionException && e.getCause() != null) { - throw new RuntimeException("Failed to get or create experiment", e.getCause()); - } - throw new RuntimeException("Failed to get or create experiment", e); - } - } -} diff --git a/src/main/java/dev/braintrust/claude/api/Project.java b/src/main/java/dev/braintrust/claude/api/Project.java deleted file mode 100644 index d36f984..0000000 --- a/src/main/java/dev/braintrust/claude/api/Project.java +++ /dev/null @@ -1,52 +0,0 @@ -package dev.braintrust.claude.api; - -import dev.braintrust.claude.config.BraintrustConfig; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.CompletionException; - -/** High-level API for project management, matching the Go SDK's DX. */ -public final class Project { - - private Project() { - // Utility class - } - - /** - * Registers a project in Braintrust. If a project with the same name already exists, returns - * the existing project. - * - * @param name The name of the project - * @return The project (created or existing) - */ - public static CompletableFuture registerProject(String name) { - var config = BraintrustConfig.fromEnvironment(); - return registerProject(name, config); - } - - /** - * Registers a project with custom configuration. - * - * @param name The name of the project - * @param config Custom Braintrust configuration - * @return The project (created or existing) - */ - public static CompletableFuture registerProject( - String name, BraintrustConfig config) { - var client = new BraintrustApiClient(config); - - // The API handles "get or create" semantics - return client.createProject(name); - } - - /** Synchronous version of registerProject for simpler usage. */ - public static BraintrustApiClient.Project registerProjectSync(String name) { - try { - return registerProject(name).get(); - } catch (Exception e) { - if (e instanceof CompletionException && e.getCause() != null) { - throw new RuntimeException("Failed to register project", e.getCause()); - } - throw new RuntimeException("Failed to register project", e); - } - } -} diff --git a/src/main/java/dev/braintrust/claude/config/BraintrustConfig.java b/src/main/java/dev/braintrust/claude/config/BraintrustConfig.java deleted file mode 100644 index 1165cff..0000000 --- a/src/main/java/dev/braintrust/claude/config/BraintrustConfig.java +++ /dev/null @@ -1,204 +0,0 @@ -package dev.braintrust.claude.config; - -import dev.braintrust.claude.api.BraintrustApiClient; -import java.net.URI; -import java.time.Duration; -import java.util.Optional; -import java.util.function.Consumer; -import javax.annotation.Nullable; - -/** - * Configuration for Braintrust SDK. Uses builder pattern with sensible defaults and environment - * variable support. - */ -public final class BraintrustConfig { - private static final String DEFAULT_API_URL = "https://api.braintrust.dev"; - private static final String DEFAULT_APP_URL = "https://www.braintrust.dev"; - - private final String apiKey; - 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; - - private BraintrustConfig(Builder builder) { - this.apiKey = builder.apiKey; - 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; - } - - public String apiKey() { - return apiKey; - } - - public URI apiUrl() { - return apiUrl; - } - - public URI appUrl() { - return appUrl; - } - - public Optional defaultProjectId() { - return Optional.ofNullable(defaultProjectId); - } - - public Optional orgName() { - return Optional.ofNullable(orgName); - } - - public boolean enableTraceConsoleLog() { - return enableTraceConsoleLog; - } - - public boolean debug() { - return debug; - } - - public Duration requestTimeout() { - return requestTimeout; - } - - /** Creates a new builder initialized with environment variables. */ - public static Builder builder() { - return new Builder(); - } - - /** 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 and retrieves organization name from API. - */ - public static BraintrustConfig create(Consumer customizer) { - var builder = builder(); - customizer.accept(builder); - return builder.build(); - } - - public static final class Builder { - private String apiKey; - private URI apiUrl; - private URI appUrl; - private String defaultProjectId; - private String orgName; - private boolean enableTraceConsoleLog; - private boolean debug; - private Duration requestTimeout = Duration.ofSeconds(30); - - private Builder() { - // Initialize from environment - this.apiKey = getEnv("BRAINTRUST_API_KEY", null); - this.apiUrl = URI.create(getEnv("BRAINTRUST_API_URL", DEFAULT_API_URL)); - this.appUrl = URI.create(getEnv("BRAINTRUST_APP_URL", DEFAULT_APP_URL)); - this.defaultProjectId = getEnv("BRAINTRUST_DEFAULT_PROJECT_ID", null); - this.enableTraceConsoleLog = - Boolean.parseBoolean(getEnv("BRAINTRUST_ENABLE_TRACE_CONSOLE_LOG", "false")); - this.debug = Boolean.parseBoolean(getEnv("BRAINTRUST_DEBUG", "false")); - } - - public Builder apiKey(String apiKey) { - this.apiKey = apiKey; - return this; - } - - public Builder apiUrl(String apiUrl) { - return apiUrl(URI.create(apiUrl)); - } - - public Builder apiUrl(URI apiUrl) { - this.apiUrl = apiUrl; - return this; - } - - public Builder appUrl(String appUrl) { - return appUrl(URI.create(appUrl)); - } - - public Builder appUrl(URI appUrl) { - this.appUrl = appUrl; - return this; - } - - public Builder defaultProjectId(String projectId) { - this.defaultProjectId = projectId; - return this; - } - - public Builder orgName(String orgName) { - this.orgName = orgName; - return this; - } - - public Builder enableTraceConsoleLog(boolean enable) { - this.enableTraceConsoleLog = enable; - return this; - } - - public Builder debug(boolean debug) { - this.debug = debug; - return this; - } - - public Builder requestTimeout(Duration timeout) { - this.requestTimeout = timeout; - return this; - } - - public BraintrustConfig build() { - if (apiKey == null || apiKey.isBlank()) { - throw new IllegalStateException( - "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 - return value != null ? value.trim() : defaultValue; - } - } -} diff --git a/src/main/java/dev/braintrust/claude/eval/Evaluation.java b/src/main/java/dev/braintrust/claude/eval/Evaluation.java deleted file mode 100644 index 01cdec6..0000000 --- a/src/main/java/dev/braintrust/claude/eval/Evaluation.java +++ /dev/null @@ -1,417 +0,0 @@ -package dev.braintrust.claude.eval; - -import dev.braintrust.claude.api.Experiment; -import dev.braintrust.claude.trace.BraintrustSpanProcessor; -import dev.braintrust.claude.trace.BraintrustTracing; -import io.opentelemetry.api.trace.SpanKind; -import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.context.Context; -import java.time.Duration; -import java.time.Instant; -import java.util.*; -import java.util.concurrent.*; -import java.util.function.BiFunction; -import java.util.function.Function; -import java.util.stream.Collectors; -import java.util.stream.Stream; -import java.util.stream.StreamSupport; -import javax.annotation.Nullable; - -/** - * Main evaluation framework for testing AI models. Uses Java generics for type safety and streams - * for functional processing. - * - * @param The type of input data for the evaluation - * @param The type of output produced by the task - */ -public final class Evaluation { - private final String name; - private final Iterable data; - private final Function task; - private final List> scorers; - private final EvaluationOptions options; - private final Tracer tracer; - - private Evaluation(Builder builder) { - this.name = builder.name; - this.data = builder.data; - this.task = builder.task; - this.scorers = List.copyOf(builder.scorers); - this.options = builder.options; - this.tracer = BraintrustTracing.getTracer(); - } - - /** Runs the evaluation and returns results. */ - public EvaluationResults run() { - // Each evaluation case should be its own root trace, not a child - var results = streamData().map(this::evaluateOne).collect(Collectors.toList()); - - return new EvaluationResults<>(name, results, calculateSummary(results)); - } - - /** Runs the evaluation asynchronously. */ - public CompletableFuture> runAsync() { - return CompletableFuture.supplyAsync(this::run, options.executor); - } - - private Stream streamData() { - return data instanceof Collection collection - ? options.parallel ? collection.parallelStream() : collection.stream() - : StreamSupport.stream(data.spliterator(), options.parallel); - } - - private EvaluationResult evaluateOne(INPUT input) { - // Create a new root span for each evaluation case - var span = - tracer.spanBuilder("eval") - .setNoParent() // Make it a root span - .setSpanKind(SpanKind.CLIENT) - .startSpan(); - - // Set the experiment parent if configured - if (options.experimentId != null) { - span.setAttribute(BraintrustSpanProcessor.PARENT_EXPERIMENT_ID, options.experimentId); - span.setAttribute(BraintrustSpanProcessor.PARENT_TYPE, "experiment"); - span.setAttribute( - BraintrustSpanProcessor.PARENT, "experiment_id:" + options.experimentId); - } - - // Set span attributes matching Go SDK pattern - span.setAttribute("braintrust.span_attributes", "{\"type\":\"eval\"}"); - - try (var scope = span.makeCurrent()) { - var startTime = Instant.now(); - - // Add input to span - extract the question field for MathProblem - if (input != null) { - try { - // For MathProblem record, get the question field - var questionField = input.getClass().getDeclaredField("question"); - questionField.setAccessible(true); - var question = questionField.get(input); - if (question != null) { - span.setAttribute("input", question.toString()); - } - } catch (NoSuchFieldException | IllegalAccessException e) { - // Fallback to toString - span.setAttribute("input", input.toString()); - } - - // Also set input as JSON matching Go SDK - try { - var objectMapper = new com.fasterxml.jackson.databind.ObjectMapper(); - var inputJson = objectMapper.writeValueAsString(input); - span.setAttribute("braintrust.input_json", inputJson); - } catch (Exception e) { - // Ignore JSON serialization errors - } - } - - // Try to extract expected value if the input has an "expected" field - try { - var expectedField = input.getClass().getDeclaredField("expected"); - expectedField.setAccessible(true); - var expectedValue = expectedField.get(input); - if (expectedValue != null) { - span.setAttribute("expected", expectedValue.toString()); - - // Also set as JSON attribute - try { - var objectMapper = new com.fasterxml.jackson.databind.ObjectMapper(); - var expectedJson = objectMapper.writeValueAsString(expectedValue); - span.setAttribute("braintrust.expected", expectedJson); - } catch (Exception e) { - // Fallback to string value - span.setAttribute("braintrust.expected", expectedValue.toString()); - } - } - } catch (NoSuchFieldException | IllegalAccessException e) { - // No expected field, which is fine - } - - OUTPUT output = null; - @Nullable Exception error = null; - - // Create a child "task" span for the actual execution - var taskSpan = - tracer.spanBuilder("task") - .setParent(Context.current()) - .setSpanKind(SpanKind.INTERNAL) - .startSpan(); - - // Set span attributes matching Go SDK pattern - taskSpan.setAttribute("braintrust.span_attributes", "{\"type\":\"task\"}"); - - try (var taskScope = taskSpan.makeCurrent()) { - // Run the task - try { - output = options.timeout != null ? runWithTimeout(input) : task.apply(input); - } catch (Exception e) { - error = e; - taskSpan.recordException(e); - taskSpan.setStatus(io.opentelemetry.api.trace.StatusCode.ERROR); - } - } finally { - taskSpan.end(); - } - - var duration = Duration.between(startTime, Instant.now()); - - // Add output to span - if (output != null) { - // Extract the result from Answer record - try { - var resultField = output.getClass().getDeclaredField("result"); - resultField.setAccessible(true); - var result = resultField.get(output); - if (result != null) { - span.setAttribute("output", result.toString()); - } - } catch (NoSuchFieldException | IllegalAccessException e) { - span.setAttribute("output", output.toString()); - } - } - - // Calculate scores in a child "score" span - Map scores = Map.of(); - if (error == null && output != null) { - var scoreSpan = - tracer.spanBuilder("score") - .setParent(Context.current()) - .setSpanKind(SpanKind.INTERNAL) - .startSpan(); - - // Set span attributes matching Go SDK pattern - scoreSpan.setAttribute("braintrust.span_attributes", "{\"type\":\"score\"}"); - - try (var scoreScope = scoreSpan.makeCurrent()) { - scores = calculateScores(input, output, error); - - // Add scores to the parent eval span - scores.forEach( - (name, score) -> { - span.setAttribute("score." + name, score); - scoreSpan.setAttribute(name, score); - }); - - // Also set scores as JSON for Braintrust parsing (matching Go SDK) - if (!scores.isEmpty()) { - try { - var objectMapper = new com.fasterxml.jackson.databind.ObjectMapper(); - var scoresJson = objectMapper.writeValueAsString(scores); - scoreSpan.setAttribute("braintrust.scores", scoresJson); - } catch (Exception e) { - // Log error but don't fail the evaluation - span.recordException(e); - } - } - } finally { - scoreSpan.end(); - } - } - - return new EvaluationResult<>( - input, output, scores, duration, Optional.ofNullable(error)); - } finally { - span.end(); - } - } - - private OUTPUT runWithTimeout(INPUT input) throws Exception { - var future = CompletableFuture.supplyAsync(() -> task.apply(input), options.executor); - try { - return future.get(options.timeout.toMillis(), TimeUnit.MILLISECONDS); - } catch (TimeoutException e) { - future.cancel(true); - throw new EvaluationException("Task timed out after " + options.timeout, e); - } catch (ExecutionException e) { - throw new EvaluationException("Task execution failed", e.getCause()); - } - } - - private Map calculateScores( - INPUT input, @Nullable OUTPUT output, @Nullable Exception error) { - if (error != null || scorers.isEmpty()) { - return Map.of(); - } - - return scorers.stream() - .collect( - Collectors.toMap( - Scorer::name, - scorer -> scorer.score(input, output), - (a, b) -> b, // In case of duplicate names, use the last one - LinkedHashMap::new // Preserve order - )); - } - - private EvaluationSummary calculateSummary(List> results) { - var totalCount = results.size(); - var errorCount = (int) results.stream().filter(r -> r.error().isPresent()).count(); - var successCount = totalCount - errorCount; - - var avgDuration = - results.stream().mapToLong(r -> r.duration().toMillis()).average().orElse(0.0); - - var scoreStats = calculateScoreStatistics(results); - - return new EvaluationSummary( - totalCount, - successCount, - errorCount, - Duration.ofMillis((long) avgDuration), - scoreStats); - } - - private Map calculateScoreStatistics( - List> results) { - var scoresByName = new HashMap>(); - - results.stream() - .flatMap(r -> r.scores().entrySet().stream()) - .forEach( - entry -> - scoresByName - .computeIfAbsent(entry.getKey(), k -> new ArrayList<>()) - .add(entry.getValue())); - - return scoresByName.entrySet().stream() - .collect( - Collectors.toMap( - Map.Entry::getKey, entry -> calculateStats(entry.getValue()))); - } - - private ScoreStatistics calculateStats(List values) { - var stats = values.stream().mapToDouble(Double::doubleValue).summaryStatistics(); - - var sorted = values.stream().sorted().toList(); - var median = - sorted.size() % 2 == 0 - ? (sorted.get(sorted.size() / 2 - 1) + sorted.get(sorted.size() / 2)) / 2.0 - : sorted.get(sorted.size() / 2); - - return new ScoreStatistics( - stats.getAverage(), - stats.getMin(), - stats.getMax(), - median, - calculateStdDev(values, stats.getAverage()), - stats.getCount()); - } - - private double calculateStdDev(List values, double mean) { - var variance = - values.stream().mapToDouble(v -> Math.pow(v - mean, 2)).average().orElse(0.0); - return Math.sqrt(variance); - } - - /** Creates a new evaluation builder. */ - public static Builder builder() { - return new Builder<>(); - } - - /** Builder for creating evaluations with fluent API. */ - public static final class Builder { - private String name = "evaluation"; - private Iterable data = List.of(); - private Function task; - private final List> scorers = new ArrayList<>(); - private final EvaluationOptions options = new EvaluationOptions(); - - public Builder name(String name) { - this.name = name; - return this; - } - - public Builder data(Iterable data) { - this.data = data; - return this; - } - - public Builder data(Stream data) { - this.data = data.toList(); - return this; - } - - @SafeVarargs - public final Builder data(INPUT... data) { - this.data = List.of(data); - return this; - } - - public Builder task(Function task) { - this.task = task; - return this; - } - - public Builder scorer(Scorer scorer) { - this.scorers.add(scorer); - return this; - } - - public Builder scorer( - String name, BiFunction scoreFunc) { - return scorer(Scorer.of(name, scoreFunc)); - } - - public Builder experimentId(String experimentId) { - this.options.experimentId = experimentId; - return this; - } - - /** - * Convenience method that creates a new experiment and uses its ID. Matches the Go SDK - * pattern of automatically registering experiments. - */ - public Builder experiment(String experimentName, String projectName) { - try { - // This matches Go's pattern of creating experiment during setup - var experimentId = - Experiment.getOrCreateExperiment( - experimentName, projectName) - .get(); - this.options.experimentId = experimentId; - return this; - } catch (Exception e) { - throw new RuntimeException("Failed to register experiment", e); - } - } - - public Builder parallel(boolean parallel) { - this.options.parallel = parallel; - return this; - } - - public Builder timeout(Duration timeout) { - this.options.timeout = timeout; - return this; - } - - public Builder executor(ExecutorService executor) { - this.options.executor = executor; - return this; - } - - public Evaluation build() { - if (task == null) { - throw new IllegalStateException("Task function is required"); - } - return new Evaluation<>(this); - } - - public EvaluationResults run() { - return build().run(); - } - - public CompletableFuture> runAsync() { - return build().runAsync(); - } - } - - private static final class EvaluationOptions { - @Nullable String experimentId; - boolean parallel = false; - @Nullable Duration timeout; - ExecutorService executor = ForkJoinPool.commonPool(); - } -} diff --git a/src/main/java/dev/braintrust/claude/eval/EvaluationException.java b/src/main/java/dev/braintrust/claude/eval/EvaluationException.java deleted file mode 100644 index 9b5beea..0000000 --- a/src/main/java/dev/braintrust/claude/eval/EvaluationException.java +++ /dev/null @@ -1,13 +0,0 @@ -package dev.braintrust.claude.eval; - -/** Exception thrown during evaluation execution. */ -public class EvaluationException extends RuntimeException { - - public EvaluationException(String message) { - super(message); - } - - public EvaluationException(String message, Throwable cause) { - super(message, cause); - } -} diff --git a/src/main/java/dev/braintrust/claude/eval/EvaluationResult.java b/src/main/java/dev/braintrust/claude/eval/EvaluationResult.java deleted file mode 100644 index 5bcb8ef..0000000 --- a/src/main/java/dev/braintrust/claude/eval/EvaluationResult.java +++ /dev/null @@ -1,41 +0,0 @@ -package dev.braintrust.claude.eval; - -import java.time.Duration; -import java.util.Map; -import java.util.Optional; -import javax.annotation.Nullable; - -/** - * Represents the result of evaluating a single test case. - * - * @param input The input data for this test case - * @param output The output produced by the task (null if error occurred) - * @param scores Map of scorer names to their scores - * @param duration How long the task took to execute - * @param error Any error that occurred during execution - */ -public record EvaluationResult( - INPUT input, - @Nullable OUTPUT output, - Map scores, - Duration duration, - Optional error) { - - /** Returns whether this evaluation was successful (no errors). */ - public boolean isSuccess() { - return error.isEmpty(); - } - - /** Returns the average score across all scorers. */ - public double averageScore() { - if (scores.isEmpty()) { - return 0.0; - } - return scores.values().stream().mapToDouble(Double::doubleValue).average().orElse(0.0); - } - - /** Gets a specific score by name. */ - public Optional getScore(String scorerName) { - return Optional.ofNullable(scores.get(scorerName)); - } -} diff --git a/src/main/java/dev/braintrust/claude/eval/EvaluationResults.java b/src/main/java/dev/braintrust/claude/eval/EvaluationResults.java deleted file mode 100644 index 8bd7d36..0000000 --- a/src/main/java/dev/braintrust/claude/eval/EvaluationResults.java +++ /dev/null @@ -1,80 +0,0 @@ -package dev.braintrust.claude.eval; - -import java.util.List; -import java.util.Map; -import java.util.stream.Collectors; -import java.util.stream.Stream; - -/** Container for all evaluation results with summary statistics. */ -public record EvaluationResults( - String name, List> results, EvaluationSummary summary) { - - /** Returns only successful results. */ - public Stream> successful() { - return results.stream().filter(EvaluationResult::isSuccess); - } - - /** Returns only failed results. */ - public Stream> failed() { - return results.stream().filter(r -> !r.isSuccess()); - } - - /** Groups results by their success status. */ - public Map>> partitionBySuccess() { - return results.stream().collect(Collectors.partitioningBy(EvaluationResult::isSuccess)); - } - - /** Returns results sorted by average score (descending). */ - public List> sortedByScore() { - return results.stream() - .sorted((a, b) -> Double.compare(b.averageScore(), a.averageScore())) - .toList(); - } - - /** Filters results by minimum average score. */ - public List> withMinimumScore(double threshold) { - return results.stream().filter(r -> r.averageScore() >= threshold).toList(); - } - - /** Creates a detailed report of the evaluation. */ - public String generateReport() { - var sb = new StringBuilder(); - sb.append("Evaluation: ").append(name).append("\n"); - sb.append("=====================================\n\n"); - - sb.append("Summary:\n"); - sb.append(" Total: ").append(summary.totalCount()).append("\n"); - sb.append(" Successful: ").append(summary.successCount()).append("\n"); - sb.append(" Failed: ").append(summary.errorCount()).append("\n"); - sb.append(" Success Rate: ") - .append(String.format("%.2f%%", summary.successRate() * 100)) - .append("\n"); - sb.append(" Avg Duration: ").append(summary.averageDuration().toMillis()).append("ms\n\n"); - - if (!summary.scoreStatistics().isEmpty()) { - sb.append("Score Statistics:\n"); - summary.scoreStatistics() - .forEach( - (scorer, stats) -> { - sb.append(" ").append(scorer).append(":\n"); - sb.append(" Mean: ") - .append(String.format("%.4f", stats.mean())) - .append("\n"); - sb.append(" Min: ") - .append(String.format("%.4f", stats.min())) - .append("\n"); - sb.append(" Max: ") - .append(String.format("%.4f", stats.max())) - .append("\n"); - sb.append(" Median: ") - .append(String.format("%.4f", stats.median())) - .append("\n"); - sb.append(" Std Dev: ") - .append(String.format("%.4f", stats.standardDeviation())) - .append("\n"); - }); - } - - return sb.toString(); - } -} diff --git a/src/main/java/dev/braintrust/claude/eval/EvaluationSummary.java b/src/main/java/dev/braintrust/claude/eval/EvaluationSummary.java deleted file mode 100644 index a2d7d1a..0000000 --- a/src/main/java/dev/braintrust/claude/eval/EvaluationSummary.java +++ /dev/null @@ -1,23 +0,0 @@ -package dev.braintrust.claude.eval; - -import java.time.Duration; -import java.util.Map; - -/** Summary statistics for an evaluation run. */ -public record EvaluationSummary( - int totalCount, - int successCount, - int errorCount, - Duration averageDuration, - Map scoreStatistics) { - - /** Calculates the success rate as a percentage. */ - public double successRate() { - return totalCount == 0 ? 0.0 : (double) successCount / totalCount; - } - - /** Calculates the error rate as a percentage. */ - public double errorRate() { - return totalCount == 0 ? 0.0 : (double) errorCount / totalCount; - } -} diff --git a/src/main/java/dev/braintrust/claude/eval/ScoreStatistics.java b/src/main/java/dev/braintrust/claude/eval/ScoreStatistics.java deleted file mode 100644 index 40f2053..0000000 --- a/src/main/java/dev/braintrust/claude/eval/ScoreStatistics.java +++ /dev/null @@ -1,24 +0,0 @@ -package dev.braintrust.claude.eval; - -/** Statistical summary of scores from an evaluation. */ -public record ScoreStatistics( - double mean, double min, double max, double median, double standardDeviation, long count) { - - /** Returns the range (max - min). */ - public double range() { - return max - min; - } - - /** Returns the coefficient of variation (CV). */ - public double coefficientOfVariation() { - return mean == 0 ? 0.0 : standardDeviation / mean; - } - - /** Formats the statistics as a human-readable string. */ - @Override - public String toString() { - return String.format( - "ScoreStats{mean=%.4f, min=%.4f, max=%.4f, median=%.4f, stdDev=%.4f, n=%d}", - mean, min, max, median, standardDeviation, count); - } -} diff --git a/src/main/java/dev/braintrust/claude/eval/Scorer.java b/src/main/java/dev/braintrust/claude/eval/Scorer.java deleted file mode 100644 index c54a421..0000000 --- a/src/main/java/dev/braintrust/claude/eval/Scorer.java +++ /dev/null @@ -1,214 +0,0 @@ -package dev.braintrust.claude.eval; - -import java.util.function.BiFunction; - -/** - * Represents a scoring function for evaluation results. Can be implemented as a functional - * interface or using the factory methods. - * - * @param The type of input data - * @param The type of output data - */ -@FunctionalInterface -public interface Scorer { - - /** - * Calculates a score for the given input and output. - * - * @param input The input to the task - * @param output The output produced by the task - * @return A score between 0.0 and 1.0, where 1.0 is perfect - */ - double score(INPUT input, OUTPUT output); - - /** Returns the name of this scorer. Default implementation returns the class name. */ - default String name() { - return getClass().getSimpleName(); - } - - /** Creates a scorer with a specific name. */ - static Scorer of(String name, BiFunction scoreFunc) { - return new Scorer() { - @Override - public double score(I input, O output) { - return scoreFunc.apply(input, output); - } - - @Override - public String name() { - return name; - } - }; - } - - /** Common scorers for string outputs. */ - final class StringScorers { - private StringScorers() {} - - /** Exact match scorer. */ - public static Scorer exactMatch( - java.util.function.Function expectedFunc) { - return of( - "exact_match", - (input, output) -> { - var expected = expectedFunc.apply(input); - return expected.equals(output) ? 1.0 : 0.0; - }); - } - - /** Levenshtein distance scorer (normalized). */ - public static Scorer levenshtein( - java.util.function.Function expectedFunc) { - return of( - "levenshtein", - (input, output) -> { - var expected = expectedFunc.apply(input); - if (expected == null || output == null) { - return 0.0; - } - - var distance = calculateLevenshteinDistance(expected, output); - var maxLength = Math.max(expected.length(), output.length()); - - return maxLength == 0 ? 1.0 : 1.0 - (double) distance / maxLength; - }); - } - - /** Contains substring scorer. */ - public static Scorer contains( - java.util.function.Function substringFunc) { - return of( - "contains", - (input, output) -> { - var substring = substringFunc.apply(input); - return output != null && output.contains(substring) ? 1.0 : 0.0; - }); - } - - /** Regex match scorer. */ - public static Scorer regex( - java.util.function.Function patternFunc) { - return of( - "regex_match", - (input, output) -> { - if (output == null) return 0.0; - var pattern = patternFunc.apply(input); - return pattern.matcher(output).matches() ? 1.0 : 0.0; - }); - } - - private static int calculateLevenshteinDistance(String s1, String s2) { - int[][] dp = new int[s1.length() + 1][s2.length() + 1]; - - for (int i = 0; i <= s1.length(); i++) { - dp[i][0] = i; - } - - for (int j = 0; j <= s2.length(); j++) { - dp[0][j] = j; - } - - for (int i = 1; i <= s1.length(); i++) { - for (int j = 1; j <= s2.length(); j++) { - if (s1.charAt(i - 1) == s2.charAt(j - 1)) { - dp[i][j] = dp[i - 1][j - 1]; - } else { - dp[i][j] = - 1 - + Math.min( - dp[i - 1][j - 1], - Math.min(dp[i - 1][j], dp[i][j - 1])); - } - } - } - - return dp[s1.length()][s2.length()]; - } - } - - /** Common scorers for numeric outputs. */ - final class NumericScorers { - private NumericScorers() {} - - /** Absolute error scorer (normalized by expected value). */ - public static Scorer absoluteError( - java.util.function.Function expectedFunc, double threshold) { - return of( - "absolute_error", - (input, output) -> { - if (output == null) return 0.0; - - var expected = expectedFunc.apply(input).doubleValue(); - var actual = output.doubleValue(); - var error = Math.abs(expected - actual); - - return error <= threshold - ? 1.0 - : Math.max(0.0, 1.0 - error / Math.abs(expected)); - }); - } - - /** Relative error scorer. */ - public static Scorer relativeError( - java.util.function.Function expectedFunc, double maxError) { - return of( - "relative_error", - (input, output) -> { - if (output == null) return 0.0; - - var expected = expectedFunc.apply(input).doubleValue(); - var actual = output.doubleValue(); - - if (expected == 0) { - return actual == 0 ? 1.0 : 0.0; - } - - var relError = Math.abs((expected - actual) / expected); - return Math.max(0.0, 1.0 - relError / maxError); - }); - } - } - - /** Combines multiple scorers into one. */ - static Scorer combine( - String name, - java.util.List> scorers, - java.util.function.Function, Double> combiner) { - return new Scorer() { - @Override - public double score(I input, O output) { - var scores = scorers.stream().map(scorer -> scorer.score(input, output)).toList(); - return combiner.apply(scores); - } - - @Override - public String name() { - return name; - } - }; - } - - /** Creates a weighted average scorer. */ - static Scorer weightedAverage(java.util.Map, Double> scorerWeights) { - return new Scorer() { - @Override - public double score(I input, O output) { - var totalWeight = - scorerWeights.values().stream().mapToDouble(Double::doubleValue).sum(); - - return scorerWeights.entrySet().stream() - .mapToDouble( - entry -> - entry.getKey().score(input, output) - * entry.getValue()) - .sum() - / totalWeight; - } - - @Override - public String name() { - return "weighted_average"; - } - }; - } -} diff --git a/src/main/java/dev/braintrust/claude/instrumentation/anthropic/AnthropicInterceptor.java b/src/main/java/dev/braintrust/claude/instrumentation/anthropic/AnthropicInterceptor.java deleted file mode 100644 index 9013952..0000000 --- a/src/main/java/dev/braintrust/claude/instrumentation/anthropic/AnthropicInterceptor.java +++ /dev/null @@ -1,269 +0,0 @@ -package dev.braintrust.claude.instrumentation.anthropic; - -import dev.braintrust.claude.trace.BraintrustSpanProcessor; -import dev.braintrust.claude.trace.BraintrustTracing; -import io.opentelemetry.api.trace.SpanKind; -import io.opentelemetry.api.trace.StatusCode; -import io.opentelemetry.context.Scope; -import java.util.concurrent.CompletableFuture; -import java.util.function.Supplier; - -/** - * Interceptor for Anthropic Claude API calls that creates OTEL spans. Uses functional interfaces to - * avoid hard dependency on Anthropic SDK. - */ -public class AnthropicInterceptor { - - /** Wraps a message request with tracing. */ - public static RESP traceMessage( - REQ request, - RequestExtractor requestExtractor, - ResponseExtractor responseExtractor, - Supplier executeRequest) { - // Get tracer lazily to ensure test setup is complete - var tracer = BraintrustTracing.getTracer(); - var requestDetails = requestExtractor.extract(request); - - var span = - tracer.spanBuilder("anthropic.messages") - .setSpanKind(SpanKind.CLIENT) - .setAttribute("anthropic.model", requestDetails.model()) - .setAttribute("anthropic.max_tokens", requestDetails.maxTokens()) - .setAttribute("anthropic.temperature", requestDetails.temperature()) - .startSpan(); - - try (Scope scope = span.makeCurrent()) { - var response = executeRequest.get(); - - // Extract and record usage metrics - var responseDetails = responseExtractor.extract(response); - if (responseDetails.usage() != null) { - span.setAttribute( - BraintrustSpanProcessor.USAGE_PROMPT_TOKENS, - responseDetails.usage().inputTokens()); - span.setAttribute( - BraintrustSpanProcessor.USAGE_COMPLETION_TOKENS, - responseDetails.usage().outputTokens()); - span.setAttribute( - BraintrustSpanProcessor.USAGE_TOTAL_TOKENS, - responseDetails.usage().inputTokens() - + responseDetails.usage().outputTokens()); - - // Estimate cost - var estimatedCost = - estimateCost( - requestDetails.model(), - responseDetails.usage().inputTokens(), - responseDetails.usage().outputTokens()); - span.setAttribute(BraintrustSpanProcessor.USAGE_COST, estimatedCost); - } - - span.setStatus(StatusCode.OK); - return response; - - } catch (Exception e) { - span.recordException(e); - span.setStatus(StatusCode.ERROR, e.getMessage()); - throw e; - } finally { - span.end(); - } - } - - /** Wraps an async message request. */ - public static CompletableFuture traceMessageAsync( - REQ request, - RequestExtractor requestExtractor, - ResponseExtractor responseExtractor, - Supplier> executeRequest) { - // Get tracer lazily to ensure test setup is complete - var tracer = BraintrustTracing.getTracer(); - var requestDetails = requestExtractor.extract(request); - - var span = - tracer.spanBuilder("anthropic.messages") - .setSpanKind(SpanKind.CLIENT) - .setAttribute("anthropic.model", requestDetails.model()) - .setAttribute("anthropic.max_tokens", requestDetails.maxTokens()) - .setAttribute("anthropic.temperature", requestDetails.temperature()) - .startSpan(); - - try (Scope scope = span.makeCurrent()) { - return executeRequest - .get() - .whenComplete( - (response, error) -> { - if (error != null) { - span.recordException(error); - span.setStatus(StatusCode.ERROR, error.getMessage()); - } else { - var responseDetails = responseExtractor.extract(response); - if (responseDetails.usage() != null) { - span.setAttribute( - BraintrustSpanProcessor.USAGE_PROMPT_TOKENS, - responseDetails.usage().inputTokens()); - span.setAttribute( - BraintrustSpanProcessor.USAGE_COMPLETION_TOKENS, - responseDetails.usage().outputTokens()); - span.setAttribute( - BraintrustSpanProcessor.USAGE_TOTAL_TOKENS, - responseDetails.usage().inputTokens() - + responseDetails.usage().outputTokens()); - - var estimatedCost = - estimateCost( - requestDetails.model(), - responseDetails.usage().inputTokens(), - responseDetails.usage().outputTokens()); - span.setAttribute( - BraintrustSpanProcessor.USAGE_COST, estimatedCost); - } - span.setStatus(StatusCode.OK); - } - span.end(); - }); - } catch (Exception e) { - span.recordException(e); - span.setStatus(StatusCode.ERROR, e.getMessage()); - span.end(); - throw e; - } - } - - /** Wraps a streaming message request. */ - public static STREAM traceStreamingMessage( - REQ request, - RequestExtractor requestExtractor, - Supplier executeRequest, - StreamWrapper streamWrapper) { - // Get tracer lazily to ensure test setup is complete - var tracer = BraintrustTracing.getTracer(); - var requestDetails = requestExtractor.extract(request); - - var span = - tracer.spanBuilder("anthropic.messages.stream") - .setSpanKind(SpanKind.CLIENT) - .setAttribute("anthropic.model", requestDetails.model()) - .setAttribute("anthropic.max_tokens", requestDetails.maxTokens()) - .setAttribute("anthropic.temperature", requestDetails.temperature()) - .setAttribute("anthropic.stream", true) - .startSpan(); - - try (Scope scope = span.makeCurrent()) { - var stream = executeRequest.get(); - - // Track token accumulator - var tokenAccumulator = new TokenAccumulator(); - - // Wrap the stream to capture metrics - return streamWrapper.wrap( - stream, - chunk -> tokenAccumulator.accumulate(chunk), - () -> { - // Stream completed - span.setAttribute( - BraintrustSpanProcessor.USAGE_PROMPT_TOKENS, - tokenAccumulator.getInputTokens()); - span.setAttribute( - BraintrustSpanProcessor.USAGE_COMPLETION_TOKENS, - tokenAccumulator.getOutputTokens()); - span.setAttribute( - BraintrustSpanProcessor.USAGE_TOTAL_TOKENS, - tokenAccumulator.getTotalTokens()); - - var estimatedCost = - estimateCost( - requestDetails.model(), - tokenAccumulator.getInputTokens(), - tokenAccumulator.getOutputTokens()); - span.setAttribute(BraintrustSpanProcessor.USAGE_COST, estimatedCost); - - span.setStatus(StatusCode.OK); - span.end(); - }, - error -> { - span.recordException(error); - span.setStatus(StatusCode.ERROR, error.getMessage()); - span.end(); - }); - - } catch (Exception e) { - span.recordException(e); - span.setStatus(StatusCode.ERROR, e.getMessage()); - span.end(); - throw e; - } - } - - private static double estimateCost(String model, long inputTokens, long outputTokens) { - // Simplified cost calculation - should use actual pricing - var costPer1kInput = getCostPer1kTokens(model, true); - var costPer1kOutput = getCostPer1kTokens(model, false); - - return (inputTokens / 1000.0 * costPer1kInput) + (outputTokens / 1000.0 * costPer1kOutput); - } - - private static double getCostPer1kTokens(String model, boolean isInput) { - // Simplified pricing - should be configurable - var modelLower = model.toLowerCase(); - if (modelLower.contains("claude-3-opus")) { - return isInput ? 0.015 : 0.075; - } else if (modelLower.contains("claude-3-sonnet")) { - return isInput ? 0.003 : 0.015; - } else if (modelLower.contains("claude-3-haiku")) { - return isInput ? 0.00025 : 0.00125; - } else if (modelLower.contains("claude-2")) { - return isInput ? 0.008 : 0.024; - } else { - return 0.0; - } - } - - /** Extracts request details. */ - public interface RequestExtractor { - RequestDetails extract(REQ request); - } - - /** Extracts response details. */ - public interface ResponseExtractor { - ResponseDetails extract(RESP response); - } - - /** Wraps a stream to capture metrics. */ - public interface StreamWrapper { - STREAM wrap( - STREAM stream, - java.util.function.Consumer onChunk, - Runnable onComplete, - java.util.function.Consumer onError); - } - - public record RequestDetails(String model, int maxTokens, double temperature) {} - - public record ResponseDetails(Usage usage) {} - - public record Usage(long inputTokens, long outputTokens) {} - - /** Helper class to accumulate token counts from streaming responses. */ - private static class TokenAccumulator { - private long inputTokens = 0; - private long outputTokens = 0; - - public void accumulate(Object chunk) { - // This would need to be implemented based on the actual chunk type - // For now, this is a placeholder - } - - public long getInputTokens() { - return inputTokens; - } - - public long getOutputTokens() { - return outputTokens; - } - - public long getTotalTokens() { - return inputTokens + outputTokens; - } - } -} diff --git a/src/main/java/dev/braintrust/claude/instrumentation/openai/OpenAIInterceptor.java b/src/main/java/dev/braintrust/claude/instrumentation/openai/OpenAIInterceptor.java deleted file mode 100644 index 673bcb8..0000000 --- a/src/main/java/dev/braintrust/claude/instrumentation/openai/OpenAIInterceptor.java +++ /dev/null @@ -1,203 +0,0 @@ -package dev.braintrust.claude.instrumentation.openai; - -import dev.braintrust.claude.log.BraintrustLogger; -import dev.braintrust.claude.trace.BraintrustSpanProcessor; -import dev.braintrust.claude.trace.BraintrustTracing; -import io.opentelemetry.api.trace.SpanKind; -import io.opentelemetry.api.trace.StatusCode; -import io.opentelemetry.context.Scope; -import java.util.function.Supplier; - -/** - * Interceptor for OpenAI API calls that creates OTEL spans. Uses functional interfaces to avoid - * hard dependency on OpenAI SDK. - */ -public class OpenAIInterceptor { - - /** - * Wraps a completion request with tracing. - * - * @param operation The operation name (e.g., "chat.completion", "completion") - * @param request The request object - * @param requestExtractor Function to extract request details - * @param responseExtractor Function to extract response details - * @param executeRequest The actual request execution - * @return The response - */ - public static RESP traceCompletion( - String operation, - REQ request, - RequestExtractor requestExtractor, - ResponseExtractor responseExtractor, - Supplier executeRequest) { - // Get tracer lazily to ensure test setup is complete - var tracer = BraintrustTracing.getTracer(); - var requestDetails = requestExtractor.extract(request); - - BraintrustLogger.debug( - "OpenAI {} request: model={}, maxTokens={}, temperature={}", - operation, - requestDetails.model(), - requestDetails.maxTokens(), - requestDetails.temperature()); - - var span = - tracer.spanBuilder("openai." + operation) - .setSpanKind(SpanKind.CLIENT) - .setAttribute("openai.model", requestDetails.model()) - .setAttribute("openai.max_tokens", requestDetails.maxTokens()) - .setAttribute("openai.temperature", requestDetails.temperature()) - .startSpan(); - - try (Scope scope = span.makeCurrent()) { - var response = executeRequest.get(); - - // Extract and record usage metrics - var responseDetails = responseExtractor.extract(response); - if (responseDetails.usage() != null) { - BraintrustLogger.debug( - "OpenAI {} response: promptTokens={}, completionTokens={}, totalTokens={}", - operation, - responseDetails.usage().promptTokens(), - responseDetails.usage().completionTokens(), - responseDetails.usage().totalTokens()); - - span.setAttribute( - BraintrustSpanProcessor.USAGE_PROMPT_TOKENS, - responseDetails.usage().promptTokens()); - span.setAttribute( - BraintrustSpanProcessor.USAGE_COMPLETION_TOKENS, - responseDetails.usage().completionTokens()); - span.setAttribute( - BraintrustSpanProcessor.USAGE_TOTAL_TOKENS, - responseDetails.usage().totalTokens()); - - // Estimate cost (simplified - should use proper pricing) - var estimatedCost = - estimateCost( - requestDetails.model(), - responseDetails.usage().promptTokens(), - responseDetails.usage().completionTokens()); - span.setAttribute(BraintrustSpanProcessor.USAGE_COST, estimatedCost); - } - - span.setStatus(StatusCode.OK); - return response; - - } catch (Exception e) { - BraintrustLogger.warn("OpenAI {} request failed: {}", operation, e.getMessage(), e); - span.recordException(e); - span.setStatus(StatusCode.ERROR, e.getMessage()); - throw e; - } finally { - span.end(); - } - } - - /** Wraps a streaming completion request. */ - public static STREAM traceStreamingCompletion( - String operation, - REQ request, - RequestExtractor requestExtractor, - Supplier executeRequest, - StreamWrapper streamWrapper) { - // Get tracer lazily to ensure test setup is complete - var tracer = BraintrustTracing.getTracer(); - var requestDetails = requestExtractor.extract(request); - - var span = - tracer.spanBuilder("openai." + operation + ".stream") - .setSpanKind(SpanKind.CLIENT) - .setAttribute("openai.model", requestDetails.model()) - .setAttribute("openai.max_tokens", requestDetails.maxTokens()) - .setAttribute("openai.temperature", requestDetails.temperature()) - .setAttribute("openai.stream", true) - .startSpan(); - - try (Scope scope = span.makeCurrent()) { - var stream = executeRequest.get(); - - // Wrap the stream to capture metrics when it completes - return streamWrapper.wrap( - stream, - usage -> { - if (usage != null) { - span.setAttribute( - BraintrustSpanProcessor.USAGE_PROMPT_TOKENS, - usage.promptTokens()); - span.setAttribute( - BraintrustSpanProcessor.USAGE_COMPLETION_TOKENS, - usage.completionTokens()); - span.setAttribute( - BraintrustSpanProcessor.USAGE_TOTAL_TOKENS, - usage.totalTokens()); - - var estimatedCost = - estimateCost( - requestDetails.model(), - usage.promptTokens(), - usage.completionTokens()); - span.setAttribute(BraintrustSpanProcessor.USAGE_COST, estimatedCost); - } - span.setStatus(StatusCode.OK); - span.end(); - }, - error -> { - span.recordException(error); - span.setStatus(StatusCode.ERROR, error.getMessage()); - span.end(); - }); - - } catch (Exception e) { - span.recordException(e); - span.setStatus(StatusCode.ERROR, e.getMessage()); - span.end(); - throw e; - } - } - - private static double estimateCost(String model, long promptTokens, long completionTokens) { - // Simplified cost calculation - should use actual pricing - var costPer1kPrompt = getCostPer1kTokens(model, true); - var costPer1kCompletion = getCostPer1kTokens(model, false); - - return (promptTokens / 1000.0 * costPer1kPrompt) - + (completionTokens / 1000.0 * costPer1kCompletion); - } - - private static double getCostPer1kTokens(String model, boolean isPrompt) { - // Simplified pricing - should be configurable - var modelLower = model.toLowerCase(); - if (modelLower.contains("gpt-4")) { - return isPrompt ? 0.03 : 0.06; - } else if (modelLower.contains("gpt-3.5")) { - return isPrompt ? 0.0015 : 0.002; - } else { - return 0.0; - } - } - - /** Extracts request details. */ - public interface RequestExtractor { - RequestDetails extract(REQ request); - } - - /** Extracts response details. */ - public interface ResponseExtractor { - ResponseDetails extract(RESP response); - } - - /** Wraps a stream to capture metrics. */ - public interface StreamWrapper { - STREAM wrap( - STREAM stream, - java.util.function.Consumer onComplete, - java.util.function.Consumer onError); - } - - public record RequestDetails(String model, int maxTokens, double temperature) {} - - public record ResponseDetails(Usage usage) {} - - public record Usage(long promptTokens, long completionTokens, long totalTokens) {} -} diff --git a/src/main/java/dev/braintrust/claude/log/BraintrustLogger.java b/src/main/java/dev/braintrust/claude/log/BraintrustLogger.java deleted file mode 100644 index 1aa4d22..0000000 --- a/src/main/java/dev/braintrust/claude/log/BraintrustLogger.java +++ /dev/null @@ -1,257 +0,0 @@ -package dev.braintrust.claude.log; - -import java.util.concurrent.locks.ReadWriteLock; -import java.util.concurrent.locks.ReentrantReadWriteLock; -import java.util.function.Supplier; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; -import org.slf4j.event.Level; - -/** - * Centralized logging for Braintrust SDK. Provides consistent logging with SDK-specific control and - * configuration. - */ -public final class BraintrustLogger { - private static final String LOGGER_NAME = "braintrust"; - private static final String DEBUG_ENV = "BRAINTRUST_DEBUG"; - - private static final ReadWriteLock lock = new ReentrantReadWriteLock(); - private static volatile LoggerImplementation implementation; - private static volatile Level minLevel; - - static { - // Initialize with default SLF4J implementation - var defaultLogger = LoggerFactory.getLogger(LOGGER_NAME); - implementation = new Slf4jLoggerImplementation(defaultLogger); - - // Check environment for debug mode - var debugEnabled = Boolean.parseBoolean(System.getenv(DEBUG_ENV)); - minLevel = debugEnabled ? Level.DEBUG : Level.INFO; - } - - private BraintrustLogger() { - // Utility class - } - - /** Sets a custom logger implementation. */ - public static void setLogger(LoggerImplementation logger) { - lock.writeLock().lock(); - try { - implementation = logger; - } finally { - lock.writeLock().unlock(); - } - } - - /** Sets the minimum log level. */ - public static void setLevel(Level level) { - lock.writeLock().lock(); - try { - minLevel = level; - } finally { - lock.writeLock().unlock(); - } - } - - /** Enables or disables debug logging. */ - public static void setDebugEnabled(boolean enabled) { - setLevel(enabled ? Level.DEBUG : Level.INFO); - } - - /** Gets the current log level. Package-private for testing. */ - static Level getLevel() { - lock.readLock().lock(); - try { - return minLevel; - } finally { - lock.readLock().unlock(); - } - } - - /** Gets the current logger implementation. Package-private for testing. */ - static LoggerImplementation getLogger() { - lock.readLock().lock(); - try { - return implementation; - } finally { - lock.readLock().unlock(); - } - } - - /** Logs a debug message. */ - public static void debug(String message, Object... args) { - log(Level.DEBUG, message, args); - } - - /** Logs a debug message with lazy evaluation. */ - public static void debug(Supplier messageSupplier) { - if (isEnabled(Level.DEBUG)) { - log(Level.DEBUG, messageSupplier.get()); - } - } - - /** Logs an info message. */ - public static void info(String message, Object... args) { - log(Level.INFO, message, args); - } - - /** Logs a warning message. */ - public static void warn(String message, Object... args) { - log(Level.WARN, message, args); - } - - /** Logs a warning message with an exception. */ - public static void warn(String message, Throwable throwable, Object... args) { - lock.readLock().lock(); - try { - if (isEnabled(Level.WARN)) { - implementation.log(Level.WARN, message, throwable, args); - } - } finally { - lock.readLock().unlock(); - } - } - - /** Logs an error message. */ - public static void error(String message, Object... args) { - log(Level.ERROR, message, args); - } - - /** Logs an error message with an exception. */ - public static void error(String message, Throwable throwable, Object... args) { - lock.readLock().lock(); - try { - if (isEnabled(Level.ERROR)) { - implementation.log(Level.ERROR, message, throwable, args); - } - } finally { - lock.readLock().unlock(); - } - } - - private static void log(Level level, String message, Object... args) { - lock.readLock().lock(); - try { - if (isEnabled(level)) { - implementation.log(level, message, null, args); - } - } finally { - lock.readLock().unlock(); - } - } - - private static boolean isEnabled(Level level) { - // In SLF4J Level enum, severity is in reverse order: - // ERROR(0) < WARN(1) < INFO(2) < DEBUG(3) < TRACE(4) - // So we need to check if level <= minLevel - return level.compareTo(minLevel) <= 0; - } - - /** Logger implementation interface for pluggable logging. */ - public interface LoggerImplementation { - void log(Level level, String message, Throwable throwable, Object... args); - } - - /** SLF4J-based implementation. */ - public static class Slf4jLoggerImplementation implements LoggerImplementation { - private final Logger logger; - - public Slf4jLoggerImplementation(Logger logger) { - this.logger = logger; - } - - @Override - public void log(Level level, String message, Throwable throwable, Object... args) { - // Format message with prefix - var prefixedMessage = "braintrust: " + message; - - switch (level) { - case DEBUG -> { - if (throwable != null) { - if (args.length > 0) { - logger.debug(prefixedMessage, args, throwable); - } else { - logger.debug(prefixedMessage, throwable); - } - } else { - logger.debug(prefixedMessage, args); - } - } - case INFO -> { - if (throwable != null) { - if (args.length > 0) { - logger.info(prefixedMessage, args, throwable); - } else { - logger.info(prefixedMessage, throwable); - } - } else { - logger.info(prefixedMessage, args); - } - } - case WARN -> { - if (throwable != null) { - if (args.length > 0) { - logger.warn(prefixedMessage, args, throwable); - } else { - logger.warn(prefixedMessage, throwable); - } - } else { - logger.warn(prefixedMessage, args); - } - } - case ERROR -> { - if (throwable != null) { - if (args.length > 0) { - logger.error(prefixedMessage, args, throwable); - } else { - logger.error(prefixedMessage, throwable); - } - } else { - logger.error(prefixedMessage, args); - } - } - } - } - } - - /** Test implementation that captures log messages. */ - public static class TestLoggerImplementation implements LoggerImplementation { - private final java.util.List entries = - new java.util.concurrent.CopyOnWriteArrayList<>(); - - @Override - public void log(Level level, String message, Throwable throwable, Object... args) { - // Format SLF4J style messages - var formattedMessage = formatMessage(message, args); - entries.add(new LogEntry(level, formattedMessage, throwable)); - } - - private String formatMessage(String message, Object... args) { - if (args == null || args.length == 0) { - return message; - } - - var result = new StringBuilder(); - var parts = message.split("\\{\\}", -1); - - for (int i = 0; i < parts.length; i++) { - result.append(parts[i]); - if (i < args.length && i < parts.length - 1) { - result.append(args[i]); - } - } - - return result.toString(); - } - - public java.util.List getEntries() { - return new java.util.ArrayList<>(entries); - } - - public void clear() { - entries.clear(); - } - - public record LogEntry(Level level, String message, Throwable throwable) {} - } -} diff --git a/src/main/java/dev/braintrust/claude/openai/BraintrustOpenAI.java b/src/main/java/dev/braintrust/claude/openai/BraintrustOpenAI.java deleted file mode 100644 index d044abf..0000000 --- a/src/main/java/dev/braintrust/claude/openai/BraintrustOpenAI.java +++ /dev/null @@ -1,213 +0,0 @@ -package dev.braintrust.claude.openai; - -import com.fasterxml.jackson.databind.ObjectMapper; -import com.openai.client.OpenAIClient; -import com.openai.models.chat.completions.ChatCompletion; -import com.openai.models.chat.completions.ChatCompletionCreateParams; -import io.opentelemetry.api.common.AttributeKey; -import io.opentelemetry.api.trace.SpanKind; -import io.opentelemetry.api.trace.StatusCode; -import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.context.Scope; -import java.util.List; - -/** - * Wrapper for OpenAI client that adds Braintrust-specific OpenTelemetry tracing. - * - *

This follows the same semantic conventions as the Go SDK's wrap_openai functionality. - */ -public class BraintrustOpenAI { - - private final OpenAIClient client; - private final Tracer tracer; - private final ObjectMapper objectMapper; - - // Attribute keys matching Go SDK conventions - private static final AttributeKey ATTR_INPUT = - AttributeKey.stringKey("braintrust.input_json"); - private static final AttributeKey ATTR_OUTPUT = - AttributeKey.stringKey("braintrust.output_json"); - private static final AttributeKey ATTR_METADATA = - AttributeKey.stringKey("braintrust.metadata_json"); - private static final AttributeKey ATTR_METRICS = - AttributeKey.stringKey("braintrust.metrics_json"); - - // Individual metric attributes (for compatibility with Java SDK patterns) - private static final AttributeKey ATTR_PROMPT_TOKENS = - AttributeKey.longKey("braintrust.usage.prompt_tokens"); - private static final AttributeKey ATTR_COMPLETION_TOKENS = - AttributeKey.longKey("braintrust.usage.completion_tokens"); - private static final AttributeKey ATTR_TOTAL_TOKENS = - AttributeKey.longKey("braintrust.usage.total_tokens"); - - public BraintrustOpenAI(OpenAIClient client, Tracer tracer) { - this.client = client; - this.tracer = tracer; - this.objectMapper = new ObjectMapper(); - } - - public class Chat { - public Completions completions() { - return new Completions(); - } - } - - public class Completions { - /** - * Enhanced create method that adds OpenTelemetry tracing with Braintrust semantic - * conventions. - * - * @param params The chat completion parameters - * @param messages List of message contents in order (system, user, etc.) - needed because - * the SDK doesn't expose message details - * @return The chat completion response - */ - public ChatCompletion create( - ChatCompletionCreateParams params, List messages) { - // Create span with proper naming convention - var span = - tracer.spanBuilder("openai.chat.completions.create") - .setSpanKind(SpanKind.CLIENT) - .startSpan(); - - try (Scope scope = span.makeCurrent()) { - // Prepare input structure - matches Go SDK format exactly - var inputArray = objectMapper.createArrayNode(); - for (var message : messages) { - var messageNode = objectMapper.createObjectNode(); - messageNode.put("role", message.role); - messageNode.put("content", message.content); - inputArray.add(messageNode); - } - - // Set braintrust.input_json as JSON string for proper parsing - span.setAttribute(ATTR_INPUT, inputArray.toString()); - - // Prepare metadata - var metadataNode = objectMapper.createObjectNode(); - metadataNode.put("provider", "openai"); - metadataNode.put("endpoint", "/v1/chat/completions"); - metadataNode.put("model", params.model().toString()); - params.temperature().ifPresent(temp -> metadataNode.put("temperature", temp)); - params.topP().ifPresent(topP -> metadataNode.put("top_p", topP)); - params.maxTokens().ifPresent(max -> metadataNode.put("max_tokens", max)); - params.frequencyPenalty() - .ifPresent(penalty -> metadataNode.put("frequency_penalty", penalty)); - params.presencePenalty() - .ifPresent(penalty -> metadataNode.put("presence_penalty", penalty)); - params.n().ifPresent(n -> metadataNode.put("n", n)); - metadataNode.put("stream", false); // Always false for non-streaming - - // Make the actual API call - var response = client.chat().completions().create(params); - - // Add response metadata - metadataNode.put("id", response.id()); - metadataNode.put("object", "chat.completion"); - metadataNode.put("created", response.created()); - response.systemFingerprint() - .ifPresent(fp -> metadataNode.put("system_fingerprint", fp)); - - // Set braintrust.metadata_json as JSON string for proper parsing - span.setAttribute(ATTR_METADATA, metadataNode.toString()); - - // Prepare output structure - matches Go SDK format exactly - var outputArray = objectMapper.createArrayNode(); - for (var choice : response.choices()) { - var choiceNode = objectMapper.createObjectNode(); - choiceNode.put("index", choice.index()); - - var messageNode = objectMapper.createObjectNode(); - messageNode.put("role", "assistant"); - choice.message() - .content() - .ifPresent(content -> messageNode.put("content", content)); - messageNode.putNull("tool_calls"); // Explicitly null for non-tool responses - - choiceNode.set("message", messageNode); - choiceNode.putNull("logprobs"); // Explicitly null when not requested - if (choice.finishReason() != null) { - choiceNode.put( - "finish_reason", choice.finishReason().toString().toLowerCase()); - } - outputArray.add(choiceNode); - } - - // Set braintrust.output_json as JSON string for proper parsing - span.setAttribute(ATTR_OUTPUT, outputArray.toString()); - - // Set metrics - both as JSON and individual attributes - if (response.usage().isPresent()) { - var usage = response.usage().get(); - - // As JSON (matching Go SDK) - var metricsNode = objectMapper.createObjectNode(); - metricsNode.put("prompt_tokens", usage.promptTokens()); - metricsNode.put("completion_tokens", usage.completionTokens()); - metricsNode.put("tokens", usage.totalTokens()); - span.setAttribute(ATTR_METRICS, metricsNode.toString()); - - // As individual attributes (for Java SDK compatibility) - span.setAttribute(ATTR_PROMPT_TOKENS, usage.promptTokens()); - span.setAttribute(ATTR_COMPLETION_TOKENS, usage.completionTokens()); - span.setAttribute(ATTR_TOTAL_TOKENS, usage.totalTokens()); - } - - span.setStatus(StatusCode.OK); - return response; - - } catch (Exception e) { - span.recordException(e); - span.setStatus(StatusCode.ERROR, e.getMessage()); - throw new RuntimeException("Failed to call OpenAI API", e); - } finally { - span.end(); - } - } - - /** - * Simpler create method that doesn't require message info (but won't have input details). - */ - public ChatCompletion create(ChatCompletionCreateParams params) { - return create(params, List.of()); - } - } - - public Chat chat() { - return new Chat(); - } - - /** Helper class to pass message information since the SDK doesn't expose it. */ - public static class MessageInfo { - public final String role; - public final String content; - - public MessageInfo(String role, String content) { - this.role = role; - this.content = content; - } - - public static MessageInfo system(String content) { - return new MessageInfo("system", content); - } - - public static MessageInfo user(String content) { - return new MessageInfo("user", content); - } - - public static MessageInfo assistant(String content) { - return new MessageInfo("assistant", content); - } - } - - /** - * Creates a wrapped OpenAI client with Braintrust tracing. - * - * @param client The OpenAI client to wrap - * @param tracer The OpenTelemetry tracer to use - * @return A wrapped client that adds tracing - */ - public static BraintrustOpenAI wrap(OpenAIClient client, Tracer tracer) { - return new BraintrustOpenAI(client, tracer); - } -} diff --git a/src/main/java/dev/braintrust/claude/trace/BraintrustContext.java b/src/main/java/dev/braintrust/claude/trace/BraintrustContext.java deleted file mode 100644 index 267a812..0000000 --- a/src/main/java/dev/braintrust/claude/trace/BraintrustContext.java +++ /dev/null @@ -1,107 +0,0 @@ -package dev.braintrust.claude.trace; - -import io.opentelemetry.context.Context; -import io.opentelemetry.context.ContextKey; -import java.util.Optional; -import javax.annotation.Nullable; - -/** - * Context propagation for Braintrust-specific attributes. Uses OpenTelemetry's Context API for - * thread-safe propagation. - */ -public final class BraintrustContext { - private static final ContextKey KEY = ContextKey.named("braintrust-context"); - - @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) { - 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"); - } - - /** Stores this context in the given Context. */ - public Context storeInContext(Context context) { - return context.with(KEY, this); - } - - /** Retrieves a BraintrustContext from the given Context. */ - @Nullable - 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); - } - - 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/claude/trace/BraintrustSpanExporter.java b/src/main/java/dev/braintrust/claude/trace/BraintrustSpanExporter.java deleted file mode 100644 index 6be2d6d..0000000 --- a/src/main/java/dev/braintrust/claude/trace/BraintrustSpanExporter.java +++ /dev/null @@ -1,122 +0,0 @@ -package dev.braintrust.claude.trace; - -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.log.BraintrustLogger; -import io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporter; -import io.opentelemetry.sdk.common.CompletableResultCode; -import io.opentelemetry.sdk.trace.data.SpanData; -import io.opentelemetry.sdk.trace.export.SpanExporter; -import java.util.Collection; -import java.util.List; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; -import java.util.stream.Collectors; - -/** - * Custom span exporter for Braintrust that adds the x-bt-parent header dynamically based on span - * attributes. - */ -public class BraintrustSpanExporter implements SpanExporter { - - private final BraintrustConfig config; - private final String endpoint; - private final Map exporterCache = new ConcurrentHashMap<>(); - - public BraintrustSpanExporter(BraintrustConfig config) { - this.config = config; - this.endpoint = config.apiUrl() + "/otel/v1/traces"; - } - - @Override - public CompletableResultCode export(Collection spans) { - if (spans.isEmpty()) { - return CompletableResultCode.ofSuccess(); - } - - // Group spans by their parent (project or experiment) - var spansByParent = spans.stream().collect(Collectors.groupingBy(this::getParentFromSpan)); - - // Export each group with the appropriate x-bt-parent header - var results = - spansByParent.entrySet().stream() - .map(entry -> exportWithParent(entry.getKey(), entry.getValue())) - .toList(); - - // Combine all results - var combined = CompletableResultCode.ofAll(results); - - return combined; - } - - 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; - } - - // Use default project ID if configured - return config.defaultProjectId().map(id -> "project_id:" + id).orElse(""); - } - - private CompletableResultCode exportWithParent(String parent, List spans) { - try { - // Get or create exporter for this parent - var exporter = - exporterCache.computeIfAbsent( - parent, - p -> { - var exporterBuilder = - OtlpHttpSpanExporter.builder() - .setEndpoint(endpoint) - .addHeader( - "Authorization", - "Bearer " + config.apiKey()) - .setTimeout(config.requestTimeout()); - - // Add x-bt-parent header if we have a parent - if (!p.isEmpty()) { - exporterBuilder.addHeader("x-bt-parent", p); - BraintrustLogger.debug( - "Created exporter with x-bt-parent: {}", p); - } - - return exporterBuilder.build(); - }); - - BraintrustLogger.debug("Exporting {} spans with x-bt-parent: {}", spans.size(), parent); - - // Export the spans - return exporter.export(spans); - } catch (Exception e) { - BraintrustLogger.error("Failed to export spans", e); - return CompletableResultCode.ofFailure(); - } - } - - @Override - public CompletableResultCode flush() { - // Flush all cached exporters - var results = exporterCache.values().stream().map(OtlpHttpSpanExporter::flush).toList(); - return CompletableResultCode.ofAll(results); - } - - @Override - public CompletableResultCode shutdown() { - // Shutdown all cached exporters - var results = exporterCache.values().stream().map(OtlpHttpSpanExporter::shutdown).toList(); - exporterCache.clear(); - return CompletableResultCode.ofAll(results); - } -} diff --git a/src/main/java/dev/braintrust/claude/trace/BraintrustSpanProcessor.java b/src/main/java/dev/braintrust/claude/trace/BraintrustSpanProcessor.java deleted file mode 100644 index 726150b..0000000 --- a/src/main/java/dev/braintrust/claude/trace/BraintrustSpanProcessor.java +++ /dev/null @@ -1,256 +0,0 @@ -package dev.braintrust.claude.trace; - -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.log.BraintrustLogger; -import io.opentelemetry.api.common.AttributeKey; -import io.opentelemetry.context.Context; -import io.opentelemetry.sdk.common.CompletableResultCode; -import io.opentelemetry.sdk.trace.ReadWriteSpan; -import io.opentelemetry.sdk.trace.ReadableSpan; -import io.opentelemetry.sdk.trace.SpanProcessor; -import io.opentelemetry.sdk.trace.data.SpanData; -import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; -import io.opentelemetry.sdk.trace.export.SpanExporter; -import java.util.Optional; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ConcurrentMap; -import javax.annotation.Nullable; - -/** - * Custom span processor that enriches spans with Braintrust-specific attributes. Supports parent - * assignment to projects or experiments. - */ -public 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"); - - // Usage metrics - public static final AttributeKey USAGE_PROMPT_TOKENS = - AttributeKey.longKey("braintrust.usage.prompt_tokens"); - public static final AttributeKey USAGE_COMPLETION_TOKENS = - AttributeKey.longKey("braintrust.usage.completion_tokens"); - public static final AttributeKey USAGE_TOTAL_TOKENS = - AttributeKey.longKey("braintrust.usage.total_tokens"); - public static final AttributeKey USAGE_COST = - AttributeKey.doubleKey("braintrust.usage.cost"); - - // Scoring - public static final AttributeKey SCORE = AttributeKey.doubleKey("braintrust.score"); - public static final AttributeKey SCORE_NAME = - AttributeKey.stringKey("braintrust.score.name"); - - private final BraintrustConfig config; - private final SpanProcessor delegate; - private final ConcurrentMap parentContexts = new ConcurrentHashMap<>(); - - public BraintrustSpanProcessor(BraintrustConfig config, SpanProcessor delegate) { - this.config = config; - this.delegate = delegate; - } - - @Override - public void onStart(Context parentContext, ReadWriteSpan span) { - BraintrustLogger.debug("OnStart: span={}, parent={}", span.getName(), parentContext); - - // Check if span already has a parent attribute - if (span.getAttribute(PARENT) == null) { - // Check if parent context has Braintrust attributes first - var btContext = BraintrustContext.fromContext(parentContext); - if (btContext != null) { - btContext - .projectId() - .ifPresent( - id -> { - span.setAttribute(PARENT, "project_id:" + id); - BraintrustLogger.debug( - "OnStart: set parent project {} from context", id); - }); - btContext - .experimentId() - .ifPresent( - id -> { - span.setAttribute(PARENT, "experiment_id:" + id); - BraintrustLogger.debug( - "OnStart: set parent experiment {} from context", id); - }); - } else { - // Use default project ID if configured - config.defaultProjectId() - .ifPresent( - projectId -> { - span.setAttribute(PARENT, "project_id:" + projectId); - BraintrustLogger.debug( - "OnStart: set default project {} for span {}", - projectId, - span.getName()); - }); - } - } - - delegate.onStart(parentContext, span); - } - - @Override - public boolean isStartRequired() { - return true; - } - - @Override - public void onEnd(ReadableSpan span) { - if (config.debug()) { - logSpanDetails(span); - } - delegate.onEnd(span); - } - - @Override - public boolean isEndRequired() { - return delegate.isEndRequired(); - } - - @Override - public CompletableResultCode shutdown() { - return delegate.shutdown(); - } - - @Override - public CompletableResultCode forceFlush() { - return delegate.forceFlush(); - } - - /** Sets the parent context for a specific trace ID. */ - public void setParentContext(String traceId, ParentContext context) { - parentContexts.put(traceId, context); - } - - /** Gets the parent context for a specific trace ID. */ - public Optional getParentContext(String traceId) { - return Optional.ofNullable(parentContexts.get(traceId)); - } - - private void logSpanDetails(ReadableSpan span) { - var spanData = span.toSpanData(); - BraintrustLogger.debug( - "Span completed: name={}, traceId={}, spanId={}, duration={}ms, attributes={}", - spanData.getName(), - spanData.getTraceId(), - spanData.getSpanId(), - (spanData.getEndEpochNanos() - spanData.getStartEpochNanos()) / 1_000_000, - spanData.getAttributes()); - } - - /** Parent context for spans (project or experiment). */ - public record ParentContext( - @Nullable String projectId, @Nullable String experimentId, ParentType type) { - public enum ParentType { - PROJECT, - EXPERIMENT - } - - public static ParentContext project(String projectId) { - return new ParentContext(projectId, null, ParentType.PROJECT); - } - - public static ParentContext experiment(String experimentId) { - return new ParentContext(null, experimentId, ParentType.EXPERIMENT); - } - } - - /** Builder for creating span processors with various configurations. */ - public static final class Builder { - private final BraintrustConfig config; - private SpanExporter exporter; - private boolean enableConsoleLog; - - public Builder(BraintrustConfig config) { - this.config = config; - this.enableConsoleLog = config.enableTraceConsoleLog(); - } - - public Builder withExporter(SpanExporter exporter) { - this.exporter = exporter; - return this; - } - - public Builder enableConsoleLog(boolean enable) { - this.enableConsoleLog = enable; - return this; - } - - public BraintrustSpanProcessor build() { - SpanProcessor baseProcessor = - SpanProcessor.composite(createProcessors().toArray(SpanProcessor[]::new)); - - return new BraintrustSpanProcessor(config, baseProcessor); - } - - private java.util.List createProcessors() { - var processors = new java.util.ArrayList(); - - if (exporter != null) { - processors.add(BatchSpanProcessor.builder(exporter).build()); - } - - if (enableConsoleLog) { - processors.add(createConsoleProcessor()); - } - - return processors; - } - - private SpanProcessor createConsoleProcessor() { - return new SpanProcessor() { - @Override - public void onStart(Context parentContext, ReadWriteSpan span) {} - - @Override - public boolean isStartRequired() { - return false; - } - - @Override - public void onEnd(ReadableSpan span) { - System.out.println(formatSpan(span.toSpanData())); - } - - @Override - public boolean isEndRequired() { - return true; - } - - @Override - public CompletableResultCode shutdown() { - return CompletableResultCode.ofSuccess(); - } - - @Override - public CompletableResultCode forceFlush() { - return CompletableResultCode.ofSuccess(); - } - - private String formatSpan(SpanData span) { - return String.format( - "[%s] %s (duration: %dms, attributes: %s)", - span.getKind(), - span.getName(), - (span.getEndEpochNanos() - span.getStartEpochNanos()) / 1_000_000, - span.getAttributes()); - } - }; - } - } -} diff --git a/src/main/java/dev/braintrust/claude/trace/BraintrustTracing.java b/src/main/java/dev/braintrust/claude/trace/BraintrustTracing.java deleted file mode 100644 index 6a2e707..0000000 --- a/src/main/java/dev/braintrust/claude/trace/BraintrustTracing.java +++ /dev/null @@ -1,200 +0,0 @@ -package dev.braintrust.claude.trace; - -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.log.BraintrustLogger; -import io.opentelemetry.api.GlobalOpenTelemetry; -import io.opentelemetry.api.OpenTelemetry; -import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.sdk.OpenTelemetrySdk; -import io.opentelemetry.sdk.resources.Resource; -import io.opentelemetry.sdk.trace.SdkTracerProvider; -import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; -import io.opentelemetry.sdk.trace.export.SpanExporter; -import io.opentelemetry.semconv.resource.attributes.ResourceAttributes; -import java.time.Duration; -import java.util.HashMap; -import java.util.Map; -import java.util.concurrent.TimeUnit; -import java.util.function.Consumer; - -/** - * Main entry point for Braintrust tracing setup. Provides convenient methods to initialize - * OpenTelemetry with Braintrust configuration. - */ -public final class BraintrustTracing { - private static final String INSTRUMENTATION_NAME = "braintrust-java"; - private static final String INSTRUMENTATION_VERSION = "0.1.0"; - - private BraintrustTracing() { - // Utility class - } - - /** - * Quick start method that sets up OpenTelemetry with Braintrust defaults. This is the simplest - * way to get started. - */ - public static OpenTelemetry quickstart() { - return quickstart(BraintrustConfig.fromEnvironment()); - } - - /** Quick start with custom configuration. */ - public static OpenTelemetry quickstart(BraintrustConfig config) { - return quickstart(config, builder -> {}); - } - - /** Quick start with custom configuration and additional setup. */ - public static OpenTelemetry quickstart(BraintrustConfig config, Consumer customizer) { - var builder = new Builder(config); - customizer.accept(builder); - return builder.build(); - } - - /** Gets a tracer with Braintrust instrumentation scope. */ - public static Tracer getTracer() { - return GlobalOpenTelemetry.getTracer(INSTRUMENTATION_NAME, INSTRUMENTATION_VERSION); - } - - /** Gets a tracer from a specific OpenTelemetry instance. */ - public static Tracer getTracer(OpenTelemetry openTelemetry) { - return openTelemetry.getTracer(INSTRUMENTATION_NAME, INSTRUMENTATION_VERSION); - } - - /** Builder for customizing OpenTelemetry setup. */ - public static final class Builder { - private final BraintrustConfig config; - private final Map resourceAttributes = new HashMap<>(); - private String serviceName = "braintrust-app"; - private Duration exportInterval = Duration.ofSeconds(5); - private int maxQueueSize = 2048; - private int maxExportBatchSize = 512; - private boolean registerGlobal = true; - - public Builder(BraintrustConfig config) { - this.config = config; - } - - public Builder serviceName(String serviceName) { - this.serviceName = serviceName; - return this; - } - - public Builder resourceAttribute(String key, String value) { - this.resourceAttributes.put(key, value); - return this; - } - - public Builder exportInterval(Duration interval) { - this.exportInterval = interval; - return this; - } - - public Builder maxQueueSize(int size) { - this.maxQueueSize = size; - return this; - } - - public Builder maxExportBatchSize(int size) { - this.maxExportBatchSize = size; - return this; - } - - public Builder registerGlobal(boolean register) { - this.registerGlobal = register; - return this; - } - - public OpenTelemetry build() { - BraintrustLogger.info( - "Initializing Braintrust OpenTelemetry with service={}", serviceName); - - // Create OTLP HTTP exporter - // The Java OTLP HTTP exporter uses the exact endpoint we provide - var exporterEndpoint = config.apiUrl() + "/otel/v1/traces"; - - // Create the custom Braintrust exporter that handles x-bt-parent header - SpanExporter exporter = new BraintrustSpanExporter(config); - - // Create resource first so BraintrustSpanProcessor can access service.name - var resourceBuilder = - Resource.getDefault().toBuilder() - .put(ResourceAttributes.SERVICE_NAME, serviceName) - .put(ResourceAttributes.SERVICE_VERSION, INSTRUMENTATION_VERSION); - - resourceAttributes.forEach(resourceBuilder::put); - var resource = resourceBuilder.build(); - - // Create batch processor for efficient export - var batchProcessor = - BatchSpanProcessor.builder(exporter) - .setScheduleDelay(exportInterval.toMillis(), TimeUnit.MILLISECONDS) - .setMaxQueueSize(maxQueueSize) - .setMaxExportBatchSize(maxExportBatchSize) - .build(); - - // Create Braintrust span processor that wraps the batch processor - var spanProcessor = new BraintrustSpanProcessor(config, batchProcessor); - - // Create tracer provider - var tracerProvider = - SdkTracerProvider.builder() - .setResource(resource) - .addSpanProcessor(spanProcessor) - .build(); - - // Create OpenTelemetry instance - var openTelemetry = - OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).build(); - - // Register globally if requested - if (registerGlobal) { - GlobalOpenTelemetry.set(openTelemetry); - BraintrustLogger.debug("Registered OpenTelemetry globally"); - } - - // Register shutdown hook - Runtime.getRuntime() - .addShutdownHook( - new Thread( - () -> { - BraintrustLogger.debug("Shutting down tracer provider"); - tracerProvider.shutdown(); - })); - - return openTelemetry; - } - } - - /** Utility methods for working with spans. */ - public static final class SpanUtils { - private SpanUtils() {} - - /** Adds usage metrics to the current span. */ - public static void addUsageMetrics(long promptTokens, long completionTokens, double cost) { - var span = io.opentelemetry.api.trace.Span.current(); - span.setAttribute(BraintrustSpanProcessor.USAGE_PROMPT_TOKENS, promptTokens); - span.setAttribute(BraintrustSpanProcessor.USAGE_COMPLETION_TOKENS, completionTokens); - span.setAttribute( - BraintrustSpanProcessor.USAGE_TOTAL_TOKENS, promptTokens + completionTokens); - span.setAttribute(BraintrustSpanProcessor.USAGE_COST, cost); - } - - /** Adds a score to the current span. */ - public static void addScore(String name, double score) { - var span = io.opentelemetry.api.trace.Span.current(); - span.setAttribute(BraintrustSpanProcessor.SCORE_NAME, name); - span.setAttribute(BraintrustSpanProcessor.SCORE, score); - } - - /** Sets the parent project for the current span. */ - public static void setParentProject(String projectId) { - var span = io.opentelemetry.api.trace.Span.current(); - span.setAttribute(BraintrustSpanProcessor.PARENT, "project_id:" + projectId); - } - - /** Sets the parent experiment for the current span. */ - public static void setParentExperiment(String experimentId) { - var span = io.opentelemetry.api.trace.Span.current(); - span.setAttribute(BraintrustSpanProcessor.PARENT, "experiment_id:" + experimentId); - } - } -} diff --git a/src/main/java/dev/braintrust/claude/trace/DebugSpanExporter.java b/src/main/java/dev/braintrust/claude/trace/DebugSpanExporter.java deleted file mode 100644 index 467317e..0000000 --- a/src/main/java/dev/braintrust/claude/trace/DebugSpanExporter.java +++ /dev/null @@ -1,84 +0,0 @@ -package dev.braintrust.claude.trace; - -import dev.braintrust.claude.log.BraintrustLogger; -import io.opentelemetry.sdk.common.CompletableResultCode; -import io.opentelemetry.sdk.trace.data.SpanData; -import io.opentelemetry.sdk.trace.export.SpanExporter; -import java.util.Collection; - -/** Debug wrapper for SpanExporter that logs all spans being exported. */ -public class DebugSpanExporter implements SpanExporter { - private final SpanExporter delegate; - private final String configuredEndpoint; - - public DebugSpanExporter(SpanExporter delegate) { - this(delegate, null); - } - - public DebugSpanExporter(SpanExporter delegate, String configuredEndpoint) { - this.delegate = delegate; - this.configuredEndpoint = configuredEndpoint; - } - - @Override - public CompletableResultCode export(Collection spans) { - BraintrustLogger.info("braintrust: === EXPORTING {} SPANS ===", spans.size()); - BraintrustLogger.info("braintrust: Exporter type: {}", delegate.getClass().getName()); - if (configuredEndpoint != null) { - BraintrustLogger.info("braintrust: Configured endpoint: {}", configuredEndpoint); - } - - // Check if we actually have spans - if (spans.isEmpty()) { - BraintrustLogger.warn("braintrust: WARNING: No spans to export!"); - return CompletableResultCode.ofSuccess(); - } - - // Try to get more info about the exporter - BraintrustLogger.info("Exporter details: {}", delegate.toString()); - - for (SpanData span : spans) { - BraintrustLogger.info( - "Span: name={}, traceId={}, spanId={}", - span.getName(), - span.getTraceId(), - span.getSpanId()); - BraintrustLogger.info(" Attributes: {}", span.getAttributes()); - BraintrustLogger.info(" Resource: {}", span.getResource().getAttributes()); - - // Check for braintrust.parent attribute - var parentAttr = span.getAttributes().get(BraintrustSpanProcessor.PARENT); - if (parentAttr != null) { - BraintrustLogger.info(" Braintrust Parent: {}", parentAttr); - } else { - BraintrustLogger.warn(" WARNING: No braintrust.parent attribute!"); - } - } - - BraintrustLogger.info("=== END SPAN EXPORT ==="); - - // Delegate to actual exporter - var result = delegate.export(spans); - - result.whenComplete( - () -> { - if (result.isSuccess()) { - BraintrustLogger.info("Export successful"); - } else { - BraintrustLogger.error("Export failed!"); - } - }); - - return result; - } - - @Override - public CompletableResultCode flush() { - return delegate.flush(); - } - - @Override - public CompletableResultCode shutdown() { - return delegate.shutdown(); - } -} diff --git a/src/test/java/dev/braintrust/claude/api/BraintrustApiClientTest.java b/src/test/java/dev/braintrust/claude/api/BraintrustApiClientTest.java deleted file mode 100644 index 2610879..0000000 --- a/src/test/java/dev/braintrust/claude/api/BraintrustApiClientTest.java +++ /dev/null @@ -1,290 +0,0 @@ -package dev.braintrust.claude.api; - -import static org.assertj.core.api.Assertions.*; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.*; - -import dev.braintrust.claude.config.BraintrustConfig; -import java.io.IOException; -import java.net.http.HttpClient; -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.concurrent.CompletableFuture; -import java.util.concurrent.CompletionException; -import java.util.concurrent.ExecutionException; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; - -@ExtendWith(MockitoExtension.class) -class BraintrustApiClientTest { - - @Mock private HttpClient mockHttpClient; - - @Mock private HttpResponse mockResponse; - - private BraintrustConfig config; - private BraintrustApiClient client; - - @BeforeEach - void setUp() { - config = - BraintrustConfig.builder() - .apiKey("test-api-key") - .apiUrl("https://api.test.dev") - .requestTimeout(Duration.ofSeconds(10)) - .build(); - - client = new BraintrustApiClient(config, mockHttpClient); - } - - @Test - void testCreateProject() throws Exception { - // Given - var responseJson = - """ - { - "id": "project-123", - "name": "Test Project", - "created_at": "2024-01-01T00:00:00Z", - "updated_at": "2024-01-01T00:00:00Z" - } - """; - - when(mockResponse.statusCode()).thenReturn(200); - when(mockResponse.body()).thenReturn(responseJson); - when(mockHttpClient.sendAsync(any(HttpRequest.class), any())) - .thenReturn(CompletableFuture.completedFuture(mockResponse)); - - // When - var project = client.createProject("Test Project").get(); - - // Then - assertThat(project.id()).isEqualTo("project-123"); - assertThat(project.name()).isEqualTo("Test Project"); - - // Verify request - verify(mockHttpClient) - .sendAsync( - argThat( - request -> { - assertThat(request.uri().toString()) - .isEqualTo("https://api.test.dev/v1/project"); - assertThat(request.method()).isEqualTo("POST"); - assertThat(request.headers().firstValue("Authorization")) - .hasValue("Bearer test-api-key"); - assertThat(request.headers().firstValue("Content-Type")) - .hasValue("application/json"); - return true; - }), - any()); - } - - @Test - void testGetProjectNotFound() { - // Given - when(mockResponse.statusCode()).thenReturn(404); - when(mockResponse.body()).thenReturn("Not found"); - when(mockHttpClient.sendAsync(any(HttpRequest.class), any())) - .thenReturn(CompletableFuture.completedFuture(mockResponse)); - - // When/Then - assertThatThrownBy(() -> client.getProject("non-existent").get()) - .isInstanceOf(ExecutionException.class) - .hasCauseInstanceOf(CompletionException.class) - .getCause() - .getCause() - .isInstanceOf(ApiException.class) - .hasMessageContaining("404"); - } - - @Test - void testGetProjectError() { - // Given - when(mockResponse.statusCode()).thenReturn(500); - when(mockResponse.body()).thenReturn("Internal server error"); - when(mockHttpClient.sendAsync(any(HttpRequest.class), any())) - .thenReturn(CompletableFuture.completedFuture(mockResponse)); - - // When/Then - assertThatThrownBy(() -> client.getProject("project-id").get()) - .isInstanceOf(ExecutionException.class) - .hasCauseInstanceOf(CompletionException.class) - .getCause() - .getCause() - .isInstanceOf(ApiException.class) - .hasMessageContaining("500"); - } - - @Test - void testListProjects() throws Exception { - // Given - var responseJson = - """ - { - "projects": [ - { - "id": "project-1", - "name": "Project 1", - "created_at": "2024-01-01T00:00:00Z", - "updated_at": "2024-01-01T00:00:00Z" - }, - { - "id": "project-2", - "name": "Project 2", - "created_at": "2024-01-02T00:00:00Z", - "updated_at": "2024-01-02T00:00:00Z" - } - ] - } - """; - - when(mockResponse.statusCode()).thenReturn(200); - when(mockResponse.body()).thenReturn(responseJson); - when(mockHttpClient.sendAsync(any(HttpRequest.class), any())) - .thenReturn(CompletableFuture.completedFuture(mockResponse)); - - // When - var projects = client.listProjects().get(); - - // Then - assertThat(projects).hasSize(2); - assertThat(projects.get(0).id()).isEqualTo("project-1"); - assertThat(projects.get(1).id()).isEqualTo("project-2"); - } - - @Test - void testCreateExperiment() throws Exception { - // Given - var request = - new BraintrustApiClient.CreateExperimentRequest( - "project-123", - "Test Experiment", - Optional.of("Description"), - Optional.of("base-exp-123")); - - var responseJson = - """ - { - "id": "exp-123", - "project_id": "project-123", - "name": "Test Experiment", - "description": "Description", - "created_at": "2024-01-01T00:00:00Z", - "updated_at": "2024-01-01T00:00:00Z" - } - """; - - when(mockResponse.statusCode()).thenReturn(200); - when(mockResponse.body()).thenReturn(responseJson); - when(mockHttpClient.sendAsync(any(HttpRequest.class), any())) - .thenReturn(CompletableFuture.completedFuture(mockResponse)); - - // When - var experiment = client.createExperiment(request).get(); - - // Then - assertThat(experiment.id()).isEqualTo("exp-123"); - assertThat(experiment.projectId()).isEqualTo("project-123"); - assertThat(experiment.name()).isEqualTo("Test Experiment"); - assertThat(experiment.description()).hasValue("Description"); - } - - @Test - void testInsertDatasetEvents() throws Exception { - // Given - var events = - List.of( - new BraintrustApiClient.DatasetEvent("input1", "output1"), - new BraintrustApiClient.DatasetEvent( - "input2", Optional.of("output2"), Optional.of("metadata"))); - - var responseJson = - """ - { - "inserted_count": 2 - } - """; - - when(mockResponse.statusCode()).thenReturn(200); - when(mockResponse.body()).thenReturn(responseJson); - when(mockHttpClient.sendAsync(any(HttpRequest.class), any())) - .thenReturn(CompletableFuture.completedFuture(mockResponse)); - - // When - var response = client.insertDatasetEvents("dataset-123", events).get(); - - // Then - assertThat(response.insertedCount()).isEqualTo(2); - - // Verify request - verify(mockHttpClient) - .sendAsync( - argThat( - request -> { - assertThat(request.uri().toString()) - .isEqualTo( - "https://api.test.dev/v1/dataset/dataset-123/insert"); - assertThat(request.method()).isEqualTo("POST"); - return true; - }), - any()); - } - - @Test - void testRequestTimeout() { - // Given - when(mockHttpClient.sendAsync(any(HttpRequest.class), any())) - .thenReturn(CompletableFuture.failedFuture(new IOException("Request timeout"))); - - // When/Then - assertThatThrownBy(() -> client.getProject("project-id").get()) - .isInstanceOf(ExecutionException.class) - .hasCauseInstanceOf(CompletionException.class) - .getCause() - .getCause() - .isInstanceOf(IOException.class) - .hasMessageContaining("timeout"); - } - - @Test - void testInvalidJsonResponse() { - // Given - when(mockResponse.statusCode()).thenReturn(200); - when(mockResponse.body()).thenReturn("invalid json"); - when(mockHttpClient.sendAsync(any(HttpRequest.class), any())) - .thenReturn(CompletableFuture.completedFuture(mockResponse)); - - // When/Then - assertThatThrownBy(() -> client.getProject("project-id").get()) - .isInstanceOf(ExecutionException.class) - .hasCauseInstanceOf(CompletionException.class) - .getCause() - .getCause() - .isInstanceOf(ApiException.class) - .hasMessageContaining("Failed to parse response"); - } - - @Test - void testDebugLogging() throws Exception { - // Given - var debugConfig = BraintrustConfig.builder().apiKey("test-api-key").debug(true).build(); - var debugClient = new BraintrustApiClient(debugConfig, mockHttpClient); - - when(mockResponse.statusCode()).thenReturn(200); - when(mockResponse.body()).thenReturn("{}"); - when(mockHttpClient.sendAsync(any(HttpRequest.class), any())) - .thenReturn(CompletableFuture.completedFuture(mockResponse)); - - // When - debugClient.listProjects().get(); - - // Then - would see debug logs (tested via logger tests) - verify(mockHttpClient).sendAsync(any(), any()); - } -} diff --git a/src/test/java/dev/braintrust/claude/config/BraintrustConfigTest.java b/src/test/java/dev/braintrust/claude/config/BraintrustConfigTest.java deleted file mode 100644 index 611590c..0000000 --- a/src/test/java/dev/braintrust/claude/config/BraintrustConfigTest.java +++ /dev/null @@ -1,149 +0,0 @@ -package dev.braintrust.claude.config; - -import static org.assertj.core.api.Assertions.*; - -import java.net.URI; -import java.time.Duration; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.CsvSource; -import org.junit.jupiter.params.provider.ValueSource; - -class BraintrustConfigTest { - - @Test - void testBuilderWithDefaults() { - // When - var config = BraintrustConfig.builder().apiKey("test-key").build(); - - // Then - assertThat(config.apiKey()).isEqualTo("test-key"); - assertThat(config.apiUrl()).isEqualTo(URI.create("https://api.braintrust.dev")); - assertThat(config.appUrl()).isEqualTo(URI.create("https://www.braintrust.dev")); - assertThat(config.defaultProjectId()).isEmpty(); - assertThat(config.enableTraceConsoleLog()).isFalse(); - assertThat(config.debug()).isFalse(); - assertThat(config.requestTimeout()).isEqualTo(Duration.ofSeconds(30)); - } - - @Test - void testBuilderWithEnvironmentVariables() { - // Note: We can't easily mock environment variables in plain JUnit - // This test demonstrates the expected behavior if environment variables were set - // In practice, the builder pattern should be used in tests - - // When using the builder to simulate environment variables - var config = - BraintrustConfig.builder() - .apiKey("env-key") - .apiUrl("https://custom.api.dev") - .appUrl("https://custom.app.dev") - .defaultProjectId("project-123") - .enableTraceConsoleLog(true) - .debug(true) - .build(); - - // Then - assertThat(config.apiKey()).isEqualTo("env-key"); - assertThat(config.apiUrl()).isEqualTo(URI.create("https://custom.api.dev")); - assertThat(config.appUrl()).isEqualTo(URI.create("https://custom.app.dev")); - assertThat(config.defaultProjectId()).hasValue("project-123"); - assertThat(config.enableTraceConsoleLog()).isTrue(); - assertThat(config.debug()).isTrue(); - } - - @Test - void testBuilderOverridesEnvironment() { - // Builder values should override any environment values - // Since we can't mock environment variables, we test the builder directly - - // When - var config = BraintrustConfig.builder().apiKey("override-key").debug(false).build(); - - // Then - assertThat(config.apiKey()).isEqualTo("override-key"); - assertThat(config.debug()).isFalse(); - } - - @Test - void testCreateWithConsumer() { - // When - var config = - BraintrustConfig.create( - builder -> - builder.apiKey("test-key") - .apiUrl("https://test.api.dev") - .defaultProjectId("test-project") - .requestTimeout(Duration.ofSeconds(60))); - - // Then - assertThat(config.apiKey()).isEqualTo("test-key"); - assertThat(config.apiUrl()).isEqualTo(URI.create("https://test.api.dev")); - assertThat(config.defaultProjectId()).hasValue("test-project"); - assertThat(config.requestTimeout()).isEqualTo(Duration.ofSeconds(60)); - } - - @Test - void testApiKeyRequired() { - // When/Then - assertThatThrownBy(() -> BraintrustConfig.builder().build()) - .isInstanceOf(IllegalStateException.class) - .hasMessageContaining("API key is required"); - } - - @Test - void testApiKeyCannotBeBlank() { - // When/Then - assertThatThrownBy(() -> BraintrustConfig.builder().apiKey(" ").build()) - .isInstanceOf(IllegalStateException.class) - .hasMessageContaining("API key is required"); - } - - @ParameterizedTest - @CsvSource({ - "true, true", - "TRUE, true", - "True, true", - "false, false", - "FALSE, false", - "False, false", - "1, false", - "yes, false", - "no, false", - ", false" - }) - void testBooleanEnvironmentParsing(String envValue, boolean expected) { - // Test the boolean parsing logic directly - // The actual parsing in BraintrustConfig uses Boolean.parseBoolean - // which only returns true for "true" (case-insensitive) - var parsed = envValue != null ? Boolean.parseBoolean(envValue) : false; - assertThat(parsed).isEqualTo(expected); - } - - @ParameterizedTest - @ValueSource( - strings = { - "https://api.example.com", - "http://localhost:8080", - "https://api.example.com/v1" - }) - void testApiUrlParsing(String url) { - // When - var config = BraintrustConfig.builder().apiKey("test").apiUrl(url).build(); - - // Then - assertThat(config.apiUrl()).isEqualTo(URI.create(url)); - } - - @Test - void testInvalidUrlThrowsException() { - // When/Then - assertThatThrownBy( - () -> - BraintrustConfig.builder() - .apiKey("test") - .apiUrl("not a valid url") - .build()) - .isInstanceOf(IllegalArgumentException.class); - } -} diff --git a/src/test/java/dev/braintrust/claude/eval/EvaluationIntegrationTest.java b/src/test/java/dev/braintrust/claude/eval/EvaluationIntegrationTest.java deleted file mode 100644 index 8507722..0000000 --- a/src/test/java/dev/braintrust/claude/eval/EvaluationIntegrationTest.java +++ /dev/null @@ -1,420 +0,0 @@ -package dev.braintrust.claude.eval; - -import static org.assertj.core.api.Assertions.*; - -import dev.braintrust.claude.config.BraintrustConfig; -import dev.braintrust.claude.trace.BraintrustContext; -import dev.braintrust.claude.trace.BraintrustSpanProcessor; -import io.opentelemetry.api.GlobalOpenTelemetry; -import io.opentelemetry.context.Context; -import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; -import io.opentelemetry.sdk.trace.SdkTracerProvider; -import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; -import java.time.Duration; -import java.util.List; -import java.util.Map; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ForkJoinPool; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.stream.IntStream; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -class EvaluationIntegrationTest { - - private InMemorySpanExporter spanExporter; - private BraintrustConfig config; - - @BeforeEach - void setUp() { - GlobalOpenTelemetry.resetForTest(); - spanExporter = InMemorySpanExporter.create(); - - config = - BraintrustConfig.builder() - .apiKey("test-key") - .defaultProjectId("test-project") - .build(); - - // Set up tracing with in-memory exporter - var braintrustProcessor = - new BraintrustSpanProcessor(config, SimpleSpanProcessor.create(spanExporter)); - - var tracerProvider = - SdkTracerProvider.builder().addSpanProcessor(braintrustProcessor).build(); - - GlobalOpenTelemetry.resetForTest(); - GlobalOpenTelemetry.set( - io.opentelemetry.sdk.OpenTelemetrySdk.builder() - .setTracerProvider(tracerProvider) - .build()); - } - - @AfterEach - void tearDown() { - GlobalOpenTelemetry.resetForTest(); - } - - @Test - void testEvaluationWithTracing() { - // Given - var testData = List.of("test1", "test2", "test3"); - - // When - var results = - Evaluation.builder() - .name("traced-evaluation") - .data(testData) - .task(String::toUpperCase) - .scorer( - "length", - (input, output) -> output.length() == input.length() ? 1.0 : 0.0) - .run(); - - // Then - assertThat(results.results()).hasSize(3); - - // Verify spans - var spans = spanExporter.getFinishedSpanItems(); - assertThat(spans).hasSizeGreaterThan(0); - - // Find evaluation span - var evalSpan = - spans.stream().filter(s -> s.getName().equals("eval")).findFirst().orElseThrow(); - - assertThat( - evalSpan.getAttributes() - .get( - io.opentelemetry.api.common.AttributeKey.stringKey( - "braintrust.span_attributes"))) - .isEqualTo("{\"type\":\"eval\"}"); - - // Verify task spans - var taskSpans = spans.stream().filter(s -> s.getName().equals("task")).toList(); - assertThat(taskSpans).hasSize(3); - } - - @Test - void testEvaluationWithExperimentContext() { - // Given - var context = BraintrustContext.forExperiment("exp-123").storeInContext(Context.current()); - - var testData = List.of(1, 2, 3, 4, 5); - - // When - try (var scope = context.makeCurrent()) { - var results = - Evaluation.builder() - .name("experiment-evaluation") - .data(testData) - .task(n -> n * n) - .scorer( - "square_check", - (input, output) -> output == input * input ? 1.0 : 0.0) - .experimentId("exp-123") - .run(); - - // Then - assertThat(results.summary().successCount()).isEqualTo(5); - } - - // Verify experiment parent was set - var spans = spanExporter.getFinishedSpanItems(); - var evalSpan = - spans.stream().filter(s -> s.getName().equals("eval")).findFirst().orElseThrow(); - - assertThat(evalSpan.getAttributes().get(BraintrustSpanProcessor.PARENT_EXPERIMENT_ID)) - .isEqualTo("exp-123"); - assertThat(evalSpan.getAttributes().get(BraintrustSpanProcessor.PARENT_TYPE)) - .isEqualTo("experiment"); - } - - @Test - void testParallelEvaluationWithMetrics() { - // Given - var data = IntStream.range(0, 100).boxed().toList(); - var processedCount = new AtomicInteger(0); - - // When - var results = - Evaluation.builder() - .name("parallel-metrics-eval") - .data(data) - .task( - n -> { - // Simulate varying processing time - try { - Thread.sleep(n % 10); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - } - processedCount.incrementAndGet(); - return "Result-" + n; - }) - .scorer( - "format_check", - (input, output) -> output.equals("Result-" + input) ? 1.0 : 0.0) - .parallel(true) - .executor(ForkJoinPool.commonPool()) - .run(); - - // Then - assertThat(results.summary().totalCount()).isEqualTo(100); - assertThat(results.summary().successCount()).isEqualTo(100); - assertThat(processedCount.get()).isEqualTo(100); - - // Verify parallel execution via span timing - var spans = spanExporter.getFinishedSpanItems(); - var taskSpans = spans.stream().filter(s -> s.getName().equals("task")).toList(); - - assertThat(taskSpans).hasSize(100); - - // Check that some spans overlap (parallel execution) - var overlapFound = false; - for (int i = 0; i < taskSpans.size() - 1; i++) { - for (int j = i + 1; j < taskSpans.size(); j++) { - var span1 = taskSpans.get(i); - var span2 = taskSpans.get(j); - - // Check if spans overlap - if (span1.getStartEpochNanos() < span2.getEndEpochNanos() - && span2.getStartEpochNanos() < span1.getEndEpochNanos()) { - overlapFound = true; - break; - } - } - if (overlapFound) break; - } - - assertThat(overlapFound).isTrue(); - } - - @Test - void testEvaluationWithMixedResults() { - // Given - record TestCase(int number, boolean shouldFail) {} - var testCases = - List.of( - new TestCase(1, false), - new TestCase(2, true), - new TestCase(3, false), - new TestCase(4, true), - new TestCase(5, false)); - - // When - var results = - Evaluation.builder() - .name("mixed-results-eval") - .data(testCases) - .task( - tc -> { - if (tc.shouldFail) { - throw new RuntimeException( - "Simulated failure for " + tc.number); - } - return tc.number * 10; - }) - .scorer( - "validation", - (input, output) -> output == input.number * 10 ? 1.0 : 0.0) - .run(); - - // Then - assertThat(results.summary().totalCount()).isEqualTo(5); - assertThat(results.summary().successCount()).isEqualTo(3); - assertThat(results.summary().errorCount()).isEqualTo(2); - - // Verify error spans have exceptions - var spans = spanExporter.getFinishedSpanItems(); - var errorSpans = - spans.stream() - .filter(s -> s.getName().equals("task")) - .filter( - s -> - s.getStatus().getStatusCode() - == io.opentelemetry.api.trace.StatusCode.ERROR) - .toList(); - - assertThat(errorSpans).hasSize(2); - errorSpans.forEach( - span -> { - assertThat(span.getEvents()).isNotEmpty(); - assertThat(span.getEvents().get(0).getName()).isEqualTo("exception"); - }); - } - - @Test - void testEvaluationWithMultipleScorers() { - // Given - record QA(String question, String expectedAnswer) {} - var questions = - List.of( - new QA("What is 2+2?", "4"), - new QA("What is the capital of France?", "Paris"), - new QA("What color is the sky?", "Blue")); - - // Simulate an imperfect model - Map model = - Map.of( - "What is 2+2?", "4", - "What is the capital of France?", "Paris", - "What color is the sky?", "Blue or grey" // Not exact match - ); - - // When - var results = - Evaluation.builder() - .name("multi-scorer-eval") - .data(questions) - .task(qa -> model.getOrDefault(qa.question, "I don't know")) - .scorer(Scorer.StringScorers.exactMatch(qa -> qa.expectedAnswer)) - .scorer(Scorer.StringScorers.contains(qa -> qa.expectedAnswer)) - .scorer( - "length_ratio", - (qa, answer) -> { - if (answer.isEmpty()) return 0.0; - var ratio = - (double) - Math.min( - answer.length(), - qa.expectedAnswer.length()) - / Math.max( - answer.length(), - qa.expectedAnswer.length()); - return ratio; - }) - .run(); - - // Then - assertThat(results.summary().successCount()).isEqualTo(3); - - // Check score statistics - var stats = results.summary().scoreStatistics(); - assertThat(stats).containsKeys("exact_match", "contains", "length_ratio"); - - // Exact match should have lower average due to "Blue or grey" - assertThat(stats.get("exact_match").mean()).isLessThan(1.0); - - // Contains should be perfect - assertThat(stats.get("contains").mean()).isEqualTo(1.0); - - // Verify scores are recorded in spans - var spans = spanExporter.getFinishedSpanItems(); - var evalSpans = spans.stream().filter(s -> s.getName().equals("eval")).toList(); - - evalSpans.forEach( - span -> { - var attrs = span.getAttributes(); - assertThat( - attrs.asMap().keySet().stream() - .filter(key -> key.getKey().startsWith("score.")) - .count()) - .isEqualTo(3); - }); - } - - @Test - void testEvaluationTimeout() { - // Given - var data = List.of("fast", "slow", "fast"); - var latch = new CountDownLatch(1); - - // When - var results = - Evaluation.builder() - .name("timeout-eval") - .data(data) - .task( - input -> { - try { - if ("slow".equals(input)) { - latch.countDown(); - Thread.sleep(200); // Will timeout - } else { - Thread.sleep(10); - } - return input.toUpperCase(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException("Interrupted", e); - } - }) - .timeout(Duration.ofMillis(50)) - .run(); - - // Then - assertThat(results.summary().totalCount()).isEqualTo(3); - assertThat(results.summary().errorCount()).isEqualTo(1); - - // Wait for slow task to be interrupted - try { - latch.await(); - Thread.sleep(100); // Give time for cleanup - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - } - - // Verify timeout in failed result - var failedResult = - results.results().stream().filter(r -> !r.isSuccess()).findFirst().orElseThrow(); - - assertThat(failedResult.input()).isEqualTo("slow"); - assertThat(failedResult.error()).isPresent(); - assertThat(failedResult.error().get()) - .isInstanceOf(EvaluationException.class) - .hasMessageContaining("timed out"); - } - - @Test - void testEvaluationResultsAnalysis() { - // Given - var data = - IntStream.range(0, 20) - .mapToObj(i -> Map.entry(i, i % 4)) // Groups of 4 - .toList(); - - // When - var results = - Evaluation., Integer>builder() - .name("analysis-eval") - .data(data) - .task(entry -> entry.getKey() * entry.getValue()) - .scorer( - "accuracy", - (entry, result) -> { - var expected = entry.getKey() * entry.getValue(); - return result.equals(expected) ? 1.0 : 0.0; - }) - .scorer( - "group_bonus", - (entry, result) -> { - // Bonus score for group 0 - return entry.getValue() == 0 ? 1.0 : 0.5; - }) - .run(); - - // Then - test analysis methods - assertThat(results.successful().count()).isEqualTo(20); - assertThat(results.failed().count()).isEqualTo(0); - - var partitioned = results.partitionBySuccess(); - assertThat(partitioned.get(true)).hasSize(20); - assertThat(partitioned.get(false)).isEmpty(); - - var sorted = results.sortedByScore(); - assertThat(sorted.get(0).averageScore()) - .isGreaterThanOrEqualTo(sorted.get(sorted.size() - 1).averageScore()); - - var highScorers = results.withMinimumScore(0.9); - assertThat(highScorers).hasSize(5); // Group 0 entries - - // Test report generation - var report = results.generateReport(); - assertThat(report).contains("Evaluation: analysis-eval"); - assertThat(report).contains("Total: 20"); - assertThat(report).contains("Success Rate: 100.00%"); - assertThat(report).contains("accuracy:"); - assertThat(report).contains("group_bonus:"); - } -} diff --git a/src/test/java/dev/braintrust/claude/eval/EvaluationTest.java b/src/test/java/dev/braintrust/claude/eval/EvaluationTest.java deleted file mode 100644 index 6ac15c6..0000000 --- a/src/test/java/dev/braintrust/claude/eval/EvaluationTest.java +++ /dev/null @@ -1,231 +0,0 @@ -package dev.braintrust.claude.eval; - -import static org.assertj.core.api.Assertions.*; - -import io.opentelemetry.sdk.testing.junit5.OpenTelemetryExtension; -import java.time.Duration; -import java.util.List; -import java.util.stream.IntStream; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.RegisterExtension; - -class EvaluationTest { - - @RegisterExtension - static final OpenTelemetryExtension otelTesting = OpenTelemetryExtension.create(); - - @BeforeEach - void setUp() { - // OpenTelemetryExtension already sets up GlobalOpenTelemetry - } - - @Test - void testSimpleEvaluation() { - // Given - var data = List.of("hello", "world", "test"); - - // When - var results = - Evaluation.builder() - .name("test-eval") - .data(data) - .task(String::toUpperCase) - .scorer( - "exact_match", - (input, output) -> output.equals(input.toUpperCase()) ? 1.0 : 0.0) - .run(); - - // Then - assertThat(results.name()).isEqualTo("test-eval"); - assertThat(results.results()).hasSize(3); - assertThat(results.summary().totalCount()).isEqualTo(3); - assertThat(results.summary().successCount()).isEqualTo(3); - assertThat(results.summary().errorCount()).isEqualTo(0); - - // Verify all scores are perfect - results.results() - .forEach( - result -> { - assertThat(result.scores()).containsEntry("exact_match", 1.0); - assertThat(result.isSuccess()).isTrue(); - }); - } - - @Test - void testEvaluationWithErrors() { - // Given - var data = List.of(1, 0, 2); - - // When - var results = - Evaluation.builder() - .name("division-eval") - .data(data) - .task(n -> 10 / n) // Will throw for 0 - .scorer("result_check", (input, output) -> output == 10 / input ? 1.0 : 0.0) - .run(); - - // Then - assertThat(results.summary().totalCount()).isEqualTo(3); - assertThat(results.summary().successCount()).isEqualTo(2); - assertThat(results.summary().errorCount()).isEqualTo(1); - - // Verify the error case - var errorResult = - results.results().stream().filter(r -> !r.isSuccess()).findFirst().orElseThrow(); - - assertThat(errorResult.input()).isEqualTo(0); - assertThat(errorResult.error()).isPresent(); - assertThat(errorResult.error().get()).isInstanceOf(ArithmeticException.class); - } - - @Test - void testParallelEvaluation() { - // Given - var data = IntStream.range(0, 100).boxed().toList(); - - // When - var results = - Evaluation.builder() - .name("parallel-eval") - .data(data) - .task( - n -> { - try { - Thread.sleep(10); // Simulate work - return n * n; - } catch (InterruptedException e) { - throw new RuntimeException(e); - } - }) - .scorer( - "square_check", - (input, output) -> output == input * input ? 1.0 : 0.0) - .parallel(true) - .run(); - - // Then - assertThat(results.summary().totalCount()).isEqualTo(100); - assertThat(results.summary().successCount()).isEqualTo(100); - - // Verify parallel execution was faster than sequential would be - var totalDuration = - results.results().stream().mapToLong(r -> r.duration().toMillis()).sum(); - - // Total work time should be > 1000ms but wall time should be much less - assertThat(totalDuration).isGreaterThanOrEqualTo(1000); - } - - @Test - void testMultipleScorerEvaluation() { - // Given - record TestCase(String input, String expected) {} - var data = - List.of( - new TestCase("hello", "HELLO"), - new TestCase("world", "WORLD"), - new TestCase("Test", "TEST")); - - // When - var results = - Evaluation.builder() - .name("multi-scorer-eval") - .data(data) - .task(tc -> tc.input.toUpperCase()) - .scorer( - "exact_match", - (tc, output) -> output.equals(tc.expected) ? 1.0 : 0.0) - .scorer( - "length_match", - (tc, output) -> output.length() == tc.expected.length() ? 1.0 : 0.0) - .scorer( - "starts_with", - (tc, output) -> - output.startsWith(tc.expected.substring(0, 1)) ? 1.0 : 0.0) - .run(); - - // Then - assertThat(results.summary().scoreStatistics()).hasSize(3); - assertThat(results.summary().scoreStatistics()) - .containsKeys("exact_match", "length_match", "starts_with"); - - // All scorers should have perfect scores - results.summary() - .scoreStatistics() - .values() - .forEach( - stats -> { - assertThat(stats.mean()).isEqualTo(1.0); - assertThat(stats.min()).isEqualTo(1.0); - assertThat(stats.max()).isEqualTo(1.0); - }); - } - - @Test - void testEvaluationWithTimeout() { - // Given - var data = List.of("fast", "slow"); - - // When - var results = - Evaluation.builder() - .name("timeout-eval") - .data(data) - .task( - input -> { - try { - if ("slow".equals(input)) { - Thread.sleep(200); - } - return input.toUpperCase(); - } catch (InterruptedException e) { - throw new RuntimeException(e); - } - }) - .timeout(Duration.ofMillis(100)) - .run(); - - // Then - assertThat(results.summary().totalCount()).isEqualTo(2); - assertThat(results.summary().errorCount()).isEqualTo(1); - - // Fast one should succeed - var fastResult = - results.results().stream() - .filter(r -> "fast".equals(r.input())) - .findFirst() - .orElseThrow(); - assertThat(fastResult.isSuccess()).isTrue(); - - // Slow one should timeout - var slowResult = - results.results().stream() - .filter(r -> "slow".equals(r.input())) - .findFirst() - .orElseThrow(); - assertThat(slowResult.isSuccess()).isFalse(); - assertThat(slowResult.error()).isPresent(); - } - - @Test - void testStreamDataSource() { - // Given - var dataStream = IntStream.range(0, 10).mapToObj(i -> "item" + i); - - // When - var results = - Evaluation.builder() - .name("stream-eval") - .data(dataStream) - .task(String::length) - .scorer( - "length_check", - (input, output) -> output == input.length() ? 1.0 : 0.0) - .run(); - - // Then - assertThat(results.summary().totalCount()).isEqualTo(10); - assertThat(results.summary().successCount()).isEqualTo(10); - } -} diff --git a/src/test/java/dev/braintrust/claude/eval/ScorerTest.java b/src/test/java/dev/braintrust/claude/eval/ScorerTest.java deleted file mode 100644 index 0439c9f..0000000 --- a/src/test/java/dev/braintrust/claude/eval/ScorerTest.java +++ /dev/null @@ -1,225 +0,0 @@ -package dev.braintrust.claude.eval; - -import static org.assertj.core.api.Assertions.*; - -import java.util.List; -import java.util.Map; -import java.util.regex.Pattern; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.CsvSource; - -class ScorerTest { - - @Test - void testCustomScorer() { - // Given - Scorer lengthScorer = - new Scorer<>() { - @Override - public double score(String input, Integer output) { - return output == input.length() ? 1.0 : 0.0; - } - - @Override - public String name() { - return "length_scorer"; - } - }; - - // When/Then - assertThat(lengthScorer.score("hello", 5)).isEqualTo(1.0); - assertThat(lengthScorer.score("hello", 4)).isEqualTo(0.0); - assertThat(lengthScorer.name()).isEqualTo("length_scorer"); - } - - @Test - void testScorerFactoryMethod() { - // Given - var scorer = - Scorer.of( - "custom", - (String input, String output) -> - input.equalsIgnoreCase(output) ? 1.0 : 0.0); - - // When/Then - assertThat(scorer.score("Hello", "hello")).isEqualTo(1.0); - assertThat(scorer.score("Hello", "world")).isEqualTo(0.0); - assertThat(scorer.name()).isEqualTo("custom"); - } - - @ParameterizedTest - @CsvSource({ - "hello, hello, 1.0", - "hello, Hello, 0.0", - "world, world, 1.0", - "test, testing, 0.0" - }) - void testExactMatchScorer(String expected, String actual, double expectedScore) { - // Given - var scorer = Scorer.StringScorers.exactMatch((String input) -> expected); - - // When/Then - assertThat(scorer.score(null, actual)).isEqualTo(expectedScore); - } - - @ParameterizedTest - @CsvSource({ - "hello, hello, 1.0", // Exact match - "hello, hallo, 0.8", // 1 substitution - "hello, helo, 0.8", // 1 deletion - "hello, helloo, 0.833", // 1 insertion - "abc, xyz, 0.0", // Completely different - "test, , 0.0", // Empty output - ", test, 0.0" // Empty expected - }) - void testLevenshteinScorer(String expected, String actual, double expectedScore) { - // Given - var scorer = Scorer.StringScorers.levenshtein((String input) -> expected); - - // When - var score = scorer.score(null, actual); - - // Then - assertThat(score).isCloseTo(expectedScore, within(0.001)); - } - - @Test - void testContainsScorer() { - // Given - var scorer = Scorer.StringScorers.contains((String input) -> "keyword"); - - // When/Then - assertThat(scorer.score(null, "This contains keyword here")).isEqualTo(1.0); - assertThat(scorer.score(null, "This does not contain it")).isEqualTo(0.0); - assertThat(scorer.score(null, "KEYWORD in caps")).isEqualTo(0.0); // Case sensitive - assertThat(scorer.score(null, null)).isEqualTo(0.0); - } - - @Test - void testRegexScorer() { - // Given - var scorer = Scorer.StringScorers.regex((String input) -> Pattern.compile("\\d{3}-\\d{4}")); - - // When/Then - assertThat(scorer.score(null, "123-4567")).isEqualTo(1.0); - assertThat(scorer.score(null, "Call me at 555-1234")) - .isEqualTo(0.0); // Must match entire string - assertThat(scorer.score(null, "12-3456")).isEqualTo(0.0); - assertThat(scorer.score(null, null)).isEqualTo(0.0); - } - - @Test - void testAbsoluteErrorScorer() { - // Given - var scorer = - Scorer.NumericScorers.absoluteError( - (String input) -> 100.0, 10.0 // threshold - ); - - // When/Then - assertThat(scorer.score(null, 100.0)).isEqualTo(1.0); // Exact match - assertThat(scorer.score(null, 95.0)).isEqualTo(1.0); // Within threshold - assertThat(scorer.score(null, 110.0)).isEqualTo(1.0); // Within threshold - assertThat(scorer.score(null, 120.0)) - .isCloseTo(0.8, within(0.01)); // Normalized by expected - assertThat(scorer.score(null, 200.0)).isEqualTo(0.0); // Too far - assertThat(scorer.score(null, null)).isEqualTo(0.0); // Null handling - } - - @Test - void testRelativeErrorScorer() { - // Given - var scorer = - Scorer.NumericScorers.relativeError( - (String input) -> 100.0, 0.1 // max 10% error - ); - - // When/Then - assertThat(scorer.score(null, 100.0)).isEqualTo(1.0); // Exact - assertThat(scorer.score(null, 105.0)).isEqualTo(0.5); // 5% error = 50% score - assertThat(scorer.score(null, 110.0)).isEqualTo(0.0); // 10% error = 0% score - assertThat(scorer.score(null, 90.0)).isEqualTo(0.0); // -10% error = 0% score - - // Test zero expected value - var zeroScorer = Scorer.NumericScorers.relativeError((String input) -> 0.0, 0.1); - assertThat(zeroScorer.score(null, 0.0)).isEqualTo(1.0); - assertThat(zeroScorer.score(null, 1.0)).isEqualTo(0.0); - } - - @Test - void testCombinedScorer() { - // Given - var scorer1 = Scorer.of("scorer1", (String in, String out) -> 0.8); - var scorer2 = Scorer.of("scorer2", (String in, String out) -> 0.6); - var scorer3 = Scorer.of("scorer3", (String in, String out) -> 1.0); - - var avgScorer = - Scorer.combine( - "average", - List.of(scorer1, scorer2, scorer3), - scores -> - scores.stream() - .mapToDouble(Double::doubleValue) - .average() - .orElse(0.0)); - - // When/Then - assertThat(avgScorer.score("input", "output")).isCloseTo(0.8, within(0.001)); - assertThat(avgScorer.name()).isEqualTo("average"); - } - - @Test - void testWeightedAverageScorer() { - // Given - var accuracyScorer = Scorer.of("accuracy", (String in, String out) -> 0.9); - var speedScorer = Scorer.of("speed", (String in, String out) -> 0.6); - var costScorer = Scorer.of("cost", (String in, String out) -> 0.8); - - var weightedScorer = - Scorer.weightedAverage( - Map.of( - accuracyScorer, 0.5, // 50% weight - speedScorer, 0.3, // 30% weight - costScorer, 0.2 // 20% weight - )); - - // When - var score = weightedScorer.score("input", "output"); - - // Then - // (0.9 * 0.5) + (0.6 * 0.3) + (0.8 * 0.2) = 0.45 + 0.18 + 0.16 = 0.79 - assertThat(score).isCloseTo(0.79, within(0.001)); - assertThat(weightedScorer.name()).isEqualTo("weighted_average"); - } - - @Test - void testLevenshteinDistanceCalculation() { - // Test various edit distance scenarios - record TestCase(String s1, String s2, int expectedDistance) {} - - var testCases = - List.of( - new TestCase("", "", 0), - new TestCase("a", "", 1), - new TestCase("", "a", 1), - new TestCase("abc", "abc", 0), - new TestCase("abc", "abd", 1), // substitution - new TestCase("abc", "ab", 1), // deletion - new TestCase("abc", "abcd", 1), // insertion - new TestCase("kitten", "sitting", 3), // classic example - new TestCase("saturday", "sunday", 3)); - - for (var testCase : testCases) { - var scorer = Scorer.StringScorers.levenshtein(input -> testCase.s1); - var score = scorer.score(null, testCase.s2); - var maxLen = Math.max(testCase.s1.length(), testCase.s2.length()); - var expectedScore = - maxLen == 0 ? 1.0 : 1.0 - (double) testCase.expectedDistance / maxLen; - - assertThat(score) - .as("Levenshtein score for '%s' vs '%s'", testCase.s1, testCase.s2) - .isCloseTo(expectedScore, within(0.001)); - } - } -} diff --git a/src/test/java/dev/braintrust/claude/instrumentation/anthropic/AnthropicInterceptorTest.java b/src/test/java/dev/braintrust/claude/instrumentation/anthropic/AnthropicInterceptorTest.java deleted file mode 100644 index fe31ba7..0000000 --- a/src/test/java/dev/braintrust/claude/instrumentation/anthropic/AnthropicInterceptorTest.java +++ /dev/null @@ -1,430 +0,0 @@ -package dev.braintrust.claude.instrumentation.anthropic; - -import static org.assertj.core.api.Assertions.*; - -import dev.braintrust.claude.trace.BraintrustSpanProcessor; -import io.opentelemetry.api.GlobalOpenTelemetry; -import io.opentelemetry.api.trace.StatusCode; -import io.opentelemetry.sdk.testing.junit5.OpenTelemetryExtension; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Consumer; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.RegisterExtension; - -class AnthropicInterceptorTest { - - @RegisterExtension - static final OpenTelemetryExtension otelTesting = OpenTelemetryExtension.create(); - - @BeforeEach - void setUp() { - GlobalOpenTelemetry.resetForTest(); - GlobalOpenTelemetry.set(otelTesting.getOpenTelemetry()); - } - - @Test - void testTraceMessageSuccess() { - // Given - var request = new TestRequest("claude-3-opus", 2000, 0.8); - var response = new TestResponse(new TestUsage(150, 100)); - - // When - var result = - AnthropicInterceptor.traceMessage( - request, - req -> - new AnthropicInterceptor.RequestDetails( - req.model, req.maxTokens, req.temperature), - resp -> - new AnthropicInterceptor.ResponseDetails( - new AnthropicInterceptor.Usage( - resp.usage.inputTokens, resp.usage.outputTokens)), - () -> response); - - // Then - assertThat(result).isEqualTo(response); - - // Verify span - var spans = otelTesting.getSpans(); - assertThat(spans).hasSize(1); - - var span = spans.get(0); - assertThat(span.getName()).isEqualTo("anthropic.messages"); - assertThat(span.getKind()).isEqualTo(io.opentelemetry.api.trace.SpanKind.CLIENT); - assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.OK); - - // Verify attributes - var attributes = span.getAttributes(); - assertThat( - attributes.get( - io.opentelemetry.api.common.AttributeKey.stringKey( - "anthropic.model"))) - .isEqualTo("claude-3-opus"); - assertThat( - attributes.get( - io.opentelemetry.api.common.AttributeKey.longKey( - "anthropic.max_tokens"))) - .isEqualTo(2000L); - assertThat( - attributes.get( - io.opentelemetry.api.common.AttributeKey.doubleKey( - "anthropic.temperature"))) - .isEqualTo(0.8); - - // Verify usage metrics - assertThat(attributes.get(BraintrustSpanProcessor.USAGE_PROMPT_TOKENS)).isEqualTo(150L); - assertThat(attributes.get(BraintrustSpanProcessor.USAGE_COMPLETION_TOKENS)).isEqualTo(100L); - assertThat(attributes.get(BraintrustSpanProcessor.USAGE_TOTAL_TOKENS)).isEqualTo(250L); - assertThat(attributes.get(BraintrustSpanProcessor.USAGE_COST)).isNotNull(); - } - - @Test - void testTraceMessageError() { - // Given - var request = new TestRequest("claude-2", 1000, 0.5); - var expectedException = new RuntimeException("API Error"); - - // When/Then - assertThatThrownBy( - () -> - AnthropicInterceptor.traceMessage( - request, - req -> - new AnthropicInterceptor.RequestDetails( - req.model, req.maxTokens, req.temperature), - resp -> null, - () -> { - throw expectedException; - })) - .isEqualTo(expectedException); - - // Verify span - var spans = otelTesting.getSpans(); - assertThat(spans).hasSize(1); - - var span = spans.get(0); - assertThat(span.getName()).isEqualTo("anthropic.messages"); - assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.ERROR); - assertThat(span.getStatus().getDescription()).isEqualTo("API Error"); - assertThat(span.getEvents()).hasSize(1); - assertThat(span.getEvents().get(0).getName()).isEqualTo("exception"); - } - - @Test - void testTraceMessageAsync() throws Exception { - // Given - var request = new TestRequest("claude-3-sonnet", 1500, 0.7); - var response = new TestResponse(new TestUsage(200, 150)); - - // When - var future = - AnthropicInterceptor.traceMessageAsync( - request, - req -> - new AnthropicInterceptor.RequestDetails( - req.model, req.maxTokens, req.temperature), - resp -> - new AnthropicInterceptor.ResponseDetails( - new AnthropicInterceptor.Usage( - resp.usage.inputTokens, resp.usage.outputTokens)), - () -> CompletableFuture.completedFuture(response)); - - // Then - var result = future.get(); - assertThat(result).isEqualTo(response); - - // Verify span - var spans = otelTesting.getSpans(); - assertThat(spans).hasSize(1); - - var span = spans.get(0); - assertThat(span.getName()).isEqualTo("anthropic.messages"); - assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.OK); - assertThat(span.getAttributes().get(BraintrustSpanProcessor.USAGE_PROMPT_TOKENS)) - .isEqualTo(200L); - assertThat(span.getAttributes().get(BraintrustSpanProcessor.USAGE_COMPLETION_TOKENS)) - .isEqualTo(150L); - } - - @Test - void testTraceMessageAsyncError() { - // Given - var request = new TestRequest("claude-3-haiku", 500, 0.3); - var exception = new RuntimeException("Async error"); - - // When - var future = - AnthropicInterceptor.traceMessageAsync( - request, - req -> - new AnthropicInterceptor.RequestDetails( - req.model, req.maxTokens, req.temperature), - resp -> null, - () -> CompletableFuture.failedFuture(exception)); - - // Then - assertThatThrownBy(future::get) - .isInstanceOf(ExecutionException.class) - .hasCauseInstanceOf(RuntimeException.class) - .hasMessageContaining("Async error"); - - // Verify span - var spans = otelTesting.getSpans(); - assertThat(spans).hasSize(1); - - var span = spans.get(0); - assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.ERROR); - } - - @Test - void testTraceStreamingMessage() { - // Given - var request = new TestRequest("claude-3-opus", 3000, 0.9); - var stream = new TestStream(); - var capturedOnChunk = new AtomicReference>(); - var capturedOnComplete = new AtomicReference(); - var capturedOnError = new AtomicReference>(); - - // When - var result = - AnthropicInterceptor.traceStreamingMessage( - request, - req -> - new AnthropicInterceptor.RequestDetails( - req.model, req.maxTokens, req.temperature), - () -> stream, - (s, onChunk, onComplete, onError) -> { - capturedOnChunk.set(onChunk); - capturedOnComplete.set(onComplete); - capturedOnError.set(onError); - return s; - }); - - // Then - assertThat(result).isEqualTo(stream); - - // Simulate streaming chunks - capturedOnChunk.get().accept(new TestChunk("chunk1")); - capturedOnChunk.get().accept(new TestChunk("chunk2")); - capturedOnChunk.get().accept(new TestChunk("chunk3")); - - // Complete stream - capturedOnComplete.get().run(); - - // Verify span after completion - var spans = otelTesting.getSpans(); - assertThat(spans).hasSize(1); - - var span = spans.get(0); - assertThat(span.getName()).isEqualTo("anthropic.messages.stream"); - assertThat(span.hasEnded()).isTrue(); - assertThat( - span.getAttributes() - .get( - io.opentelemetry.api.common.AttributeKey.booleanKey( - "anthropic.stream"))) - .isTrue(); - assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.OK); - } - - @Test - void testTraceStreamingMessageError() { - // Given - var request = new TestRequest("claude-2", 1000, 0.6); - var stream = new TestStream(); - var capturedOnError = new AtomicReference>(); - - // When - AnthropicInterceptor.traceStreamingMessage( - request, - req -> - new AnthropicInterceptor.RequestDetails( - req.model, req.maxTokens, req.temperature), - () -> stream, - (s, onChunk, onComplete, onError) -> { - capturedOnError.set(onError); - return s; - }); - - // Simulate stream error - var error = new RuntimeException("Stream error"); - capturedOnError.get().accept(error); - - // Verify span - var spans = otelTesting.getSpans(); - assertThat(spans).hasSize(1); - - var span = spans.get(0); - assertThat(span.hasEnded()).isTrue(); - assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.ERROR); - assertThat(span.getStatus().getDescription()).isEqualTo("Stream error"); - } - - @Test - void testCostCalculation() { - // Test different Claude model pricing - var testCases = - new Object[][] { - // model, inputTokens, outputTokens, expectedCost - {"claude-3-opus-20240229", 1000L, 500L, 0.0525}, // $15/$75 per million - {"claude-3-sonnet", 1000L, 500L, 0.0105}, // $3/$15 per million - {"claude-3-haiku-20240307", 1000L, 500L, 0.00088}, // $0.25/$1.25 per million - {"claude-2.1", 1000L, 500L, 0.020}, // $8/$24 per million - }; - - for (var testCase : testCases) { - var model = (String) testCase[0]; - var inputTokens = (Long) testCase[1]; - var outputTokens = (Long) testCase[2]; - var expectedCost = (Double) testCase[3]; - - var request = new TestRequest(model, 1000, 0.7); - var response = new TestResponse(new TestUsage(inputTokens, outputTokens)); - - AnthropicInterceptor.traceMessage( - request, - req -> - new AnthropicInterceptor.RequestDetails( - req.model, req.maxTokens, req.temperature), - resp -> - new AnthropicInterceptor.ResponseDetails( - new AnthropicInterceptor.Usage( - resp.usage.inputTokens, resp.usage.outputTokens)), - () -> response); - - var span = otelTesting.getSpans().get(otelTesting.getSpans().size() - 1); - var cost = span.getAttributes().get(BraintrustSpanProcessor.USAGE_COST); - - assertThat(cost) - .as( - "Cost for model %s with %d input and %d output tokens", - model, inputTokens, outputTokens) - .isCloseTo(expectedCost, within(0.0001)); - } - } - - @Test - void testNullUsageHandling() { - // Given - var request = new TestRequest("claude-3-opus", 1000, 0.7); - var response = new TestResponse(null); // No usage data - - // When - AnthropicInterceptor.traceMessage( - request, - req -> - new AnthropicInterceptor.RequestDetails( - req.model, req.maxTokens, req.temperature), - resp -> new AnthropicInterceptor.ResponseDetails(null), - () -> response); - - // Then - var span = otelTesting.getSpans().get(0); - assertThat(span.getAttributes().get(BraintrustSpanProcessor.USAGE_PROMPT_TOKENS)).isNull(); - assertThat(span.getAttributes().get(BraintrustSpanProcessor.USAGE_COMPLETION_TOKENS)) - .isNull(); - assertThat(span.getAttributes().get(BraintrustSpanProcessor.USAGE_COST)).isNull(); - } - - @Test - void testConcurrentRequests() throws Exception { - // Given - var requestCount = 10; - var futures = new CompletableFuture[requestCount]; - - // When - make concurrent requests - for (int i = 0; i < requestCount; i++) { - final int index = i; - futures[i] = - CompletableFuture.supplyAsync( - () -> { - var request = new TestRequest("claude-3-sonnet", 1000 + index, 0.5); - var response = - new TestResponse(new TestUsage(100 + index, 50 + index)); - - return AnthropicInterceptor.traceMessage( - request, - req -> - new AnthropicInterceptor.RequestDetails( - req.model, req.maxTokens, req.temperature), - resp -> - new AnthropicInterceptor.ResponseDetails( - new AnthropicInterceptor.Usage( - resp.usage.inputTokens, - resp.usage.outputTokens)), - () -> response); - }); - } - - // Wait for all to complete - CompletableFuture.allOf(futures).get(); - - // Then - var spans = otelTesting.getSpans(); - assertThat(spans).hasSize(requestCount); - - // Verify each span has unique token counts - var tokenCounts = - spans.stream() - .map( - span -> - span.getAttributes() - .get(BraintrustSpanProcessor.USAGE_PROMPT_TOKENS)) - .toList(); - - assertThat(tokenCounts).doesNotHaveDuplicates(); - } - - @Test - void testStreamingWithTokenAccumulation() { - // Given - var request = new TestRequest("claude-3-opus", 2000, 0.7); - var stream = new TestStream(); - var onChunkRef = new AtomicReference>(); - var onCompleteRef = new AtomicReference(); - - // When - AnthropicInterceptor.traceStreamingMessage( - request, - req -> - new AnthropicInterceptor.RequestDetails( - req.model, req.maxTokens, req.temperature), - () -> stream, - (s, onChunk, onComplete, onError) -> { - onChunkRef.set(onChunk); - onCompleteRef.set(onComplete); - return s; - }); - - // Simulate chunks with usage data - onChunkRef.get().accept(new ChunkWithUsage("Hello", 10, 5)); - onChunkRef.get().accept(new ChunkWithUsage(" world", 0, 3)); - onChunkRef.get().accept(new ChunkWithUsage("!", 0, 1)); - - // Complete - onCompleteRef.get().run(); - - // Then - verify accumulated tokens - var span = otelTesting.getSpans().get(0); - // Note: The actual implementation would need to handle token accumulation - // This test shows the expected behavior - assertThat(span.hasEnded()).isTrue(); - assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.OK); - } - - // Test data classes - record TestRequest(String model, int maxTokens, double temperature) {} - - record TestResponse(TestUsage usage) {} - - record TestUsage(long inputTokens, long outputTokens) {} - - static class TestStream {} - - record TestChunk(String content) {} - - record ChunkWithUsage(String content, long inputTokens, long outputTokens) {} -} diff --git a/src/test/java/dev/braintrust/claude/instrumentation/openai/OpenAIInterceptorTest.java b/src/test/java/dev/braintrust/claude/instrumentation/openai/OpenAIInterceptorTest.java deleted file mode 100644 index 92a5142..0000000 --- a/src/test/java/dev/braintrust/claude/instrumentation/openai/OpenAIInterceptorTest.java +++ /dev/null @@ -1,271 +0,0 @@ -package dev.braintrust.claude.instrumentation.openai; - -import static org.assertj.core.api.Assertions.*; - -import dev.braintrust.claude.trace.BraintrustSpanProcessor; -import io.opentelemetry.api.trace.StatusCode; -import io.opentelemetry.sdk.testing.junit5.OpenTelemetryExtension; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Consumer; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.RegisterExtension; - -class OpenAIInterceptorTest { - - @RegisterExtension - static final OpenTelemetryExtension otelTesting = OpenTelemetryExtension.create(); - - @BeforeEach - void setUp() { - // Reset and set the test OpenTelemetry instance globally - io.opentelemetry.api.GlobalOpenTelemetry.resetForTest(); - io.opentelemetry.api.GlobalOpenTelemetry.set(otelTesting.getOpenTelemetry()); - } - - @AfterEach - void tearDown() { - io.opentelemetry.api.GlobalOpenTelemetry.resetForTest(); - } - - @Test - void testTraceCompletionSuccess() { - // Given - var request = new TestRequest("gpt-4", 1000, 0.7); - var response = new TestResponse(new TestUsage(100, 50, 150)); - - // When - var result = - OpenAIInterceptor.traceCompletion( - "chat.completion", - request, - req -> - new OpenAIInterceptor.RequestDetails( - req.model, req.maxTokens, req.temperature), - resp -> - new OpenAIInterceptor.ResponseDetails( - new OpenAIInterceptor.Usage( - resp.usage.promptTokens, - resp.usage.completionTokens, - resp.usage.totalTokens)), - () -> response); - - // Then - assertThat(result).isEqualTo(response); - - // Verify span - var spans = otelTesting.getSpans(); - assertThat(spans).hasSize(1); - - var span = spans.get(0); - assertThat(span.getName()).isEqualTo("openai.chat.completion"); - assertThat(span.getKind()).isEqualTo(io.opentelemetry.api.trace.SpanKind.CLIENT); - assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.OK); - - // Verify attributes - var attributes = span.getAttributes(); - assertThat( - attributes.get( - io.opentelemetry.api.common.AttributeKey.stringKey("openai.model"))) - .isEqualTo("gpt-4"); - assertThat( - attributes.get( - io.opentelemetry.api.common.AttributeKey.longKey( - "openai.max_tokens"))) - .isEqualTo(1000L); - assertThat( - attributes.get( - io.opentelemetry.api.common.AttributeKey.doubleKey( - "openai.temperature"))) - .isEqualTo(0.7); - - // Verify usage metrics - assertThat(attributes.get(BraintrustSpanProcessor.USAGE_PROMPT_TOKENS)).isEqualTo(100L); - assertThat(attributes.get(BraintrustSpanProcessor.USAGE_COMPLETION_TOKENS)).isEqualTo(50L); - assertThat(attributes.get(BraintrustSpanProcessor.USAGE_TOTAL_TOKENS)).isEqualTo(150L); - assertThat(attributes.get(BraintrustSpanProcessor.USAGE_COST)).isNotNull(); - } - - @Test - void testTraceCompletionError() { - // Given - var request = new TestRequest("gpt-3.5-turbo", 500, 0.5); - var expectedException = new RuntimeException("API Error"); - - // When/Then - assertThatThrownBy( - () -> - OpenAIInterceptor.traceCompletion( - "completion", - request, - req -> - new OpenAIInterceptor.RequestDetails( - req.model, req.maxTokens, req.temperature), - resp -> null, - () -> { - throw expectedException; - })) - .isEqualTo(expectedException); - - // Verify span - var spans = otelTesting.getSpans(); - assertThat(spans).hasSize(1); - - var span = spans.get(0); - assertThat(span.getName()).isEqualTo("openai.completion"); - assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.ERROR); - assertThat(span.getStatus().getDescription()).isEqualTo("API Error"); - assertThat(span.getEvents()).hasSize(1); - assertThat(span.getEvents().get(0).getName()).isEqualTo("exception"); - } - - @Test - void testTraceStreamingCompletion() { - // Given - var request = new TestRequest("gpt-4", 2000, 0.9); - var stream = new TestStream(); - var capturedOnComplete = new AtomicReference>(); - var capturedOnError = new AtomicReference>(); - - // When - var result = - OpenAIInterceptor.traceStreamingCompletion( - "chat.completion", - request, - req -> - new OpenAIInterceptor.RequestDetails( - req.model, req.maxTokens, req.temperature), - () -> stream, - (s, onComplete, onError) -> { - capturedOnComplete.set(onComplete); - capturedOnError.set(onError); - return s; - }); - - // Then - assertThat(result).isEqualTo(stream); - - // Simulate stream completion to close the span - capturedOnComplete.get().accept(new OpenAIInterceptor.Usage(200, 100, 300)); - - // Verify span after completion - var spans = otelTesting.getSpans(); - assertThat(spans).hasSize(1); - - var span = spans.get(0); - assertThat(span.getName()).isEqualTo("openai.chat.completion.stream"); - assertThat(span.hasEnded()).isTrue(); - assertThat( - span.getAttributes() - .get( - io.opentelemetry.api.common.AttributeKey.booleanKey( - "openai.stream"))) - .isTrue(); - assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.OK); - assertThat(span.getAttributes().get(BraintrustSpanProcessor.USAGE_PROMPT_TOKENS)) - .isEqualTo(200L); - assertThat(span.getAttributes().get(BraintrustSpanProcessor.USAGE_COMPLETION_TOKENS)) - .isEqualTo(100L); - assertThat(span.getAttributes().get(BraintrustSpanProcessor.USAGE_TOTAL_TOKENS)) - .isEqualTo(300L); - } - - @Test - void testTraceStreamingCompletionError() { - // Given - var request = new TestRequest("gpt-3.5-turbo", 1000, 0.8); - var stream = new TestStream(); - var capturedOnError = new AtomicReference>(); - - // When - OpenAIInterceptor.traceStreamingCompletion( - "completion", - request, - req -> - new OpenAIInterceptor.RequestDetails( - req.model, req.maxTokens, req.temperature), - () -> stream, - (s, onComplete, onError) -> { - capturedOnError.set(onError); - return s; - }); - - // Simulate stream error - var error = new IOException("Stream error"); - capturedOnError.get().accept(error); - - // Verify span - var spans = otelTesting.getSpans(); - assertThat(spans).hasSize(1); - - var span = spans.get(0); - assertThat(span.hasEnded()).isTrue(); - assertThat(span.getStatus().getStatusCode()).isEqualTo(StatusCode.ERROR); - assertThat(span.getStatus().getDescription()).isEqualTo("Stream error"); - } - - @Test - void testCostCalculation() { - // Test GPT-4 pricing - var gpt4Request = new TestRequest("gpt-4", 1000, 0.7); - var gpt4Response = new TestResponse(new TestUsage(1000, 500, 1500)); - - OpenAIInterceptor.traceCompletion( - "chat.completion", - gpt4Request, - req -> - new OpenAIInterceptor.RequestDetails( - req.model, req.maxTokens, req.temperature), - resp -> - new OpenAIInterceptor.ResponseDetails( - new OpenAIInterceptor.Usage( - resp.usage.promptTokens, - resp.usage.completionTokens, - resp.usage.totalTokens)), - () -> gpt4Response); - - var span = otelTesting.getSpans().get(0); - var cost = span.getAttributes().get(BraintrustSpanProcessor.USAGE_COST); - // 1000 prompt tokens * $0.03/1k + 500 completion tokens * $0.06/1k = $0.03 + $0.03 = $0.06 - assertThat(cost).isEqualTo(0.06); - } - - @Test - void testNullUsageHandling() { - // Given - var request = new TestRequest("gpt-4", 1000, 0.7); - var response = new TestResponse(null); // No usage data - - // When - OpenAIInterceptor.traceCompletion( - "chat.completion", - request, - req -> - new OpenAIInterceptor.RequestDetails( - req.model, req.maxTokens, req.temperature), - resp -> new OpenAIInterceptor.ResponseDetails(null), - () -> response); - - // Then - var span = otelTesting.getSpans().get(0); - assertThat(span.getAttributes().get(BraintrustSpanProcessor.USAGE_PROMPT_TOKENS)).isNull(); - assertThat(span.getAttributes().get(BraintrustSpanProcessor.USAGE_COMPLETION_TOKENS)) - .isNull(); - } - - // Test data classes - record TestRequest(String model, int maxTokens, double temperature) {} - - record TestResponse(TestUsage usage) {} - - record TestUsage(long promptTokens, long completionTokens, long totalTokens) {} - - static class TestStream {} - - static class IOException extends Exception { - IOException(String message) { - super(message); - } - } -} diff --git a/src/test/java/dev/braintrust/claude/log/BraintrustLoggerTest.java b/src/test/java/dev/braintrust/claude/log/BraintrustLoggerTest.java deleted file mode 100644 index aed0a33..0000000 --- a/src/test/java/dev/braintrust/claude/log/BraintrustLoggerTest.java +++ /dev/null @@ -1,246 +0,0 @@ -package dev.braintrust.claude.log; - -import static org.assertj.core.api.Assertions.*; - -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.TestInstance; -import org.junit.jupiter.api.parallel.Execution; -import org.junit.jupiter.api.parallel.ExecutionMode; -import org.slf4j.event.Level; - -@TestInstance(TestInstance.Lifecycle.PER_METHOD) -@Execution(ExecutionMode.SAME_THREAD) -class BraintrustLoggerTest { - - private BraintrustLogger.TestLoggerImplementation testLogger; - private BraintrustLogger.LoggerImplementation savedLogger; - private Level savedLevel; - - @BeforeEach - void setUp() { - // Save current state - savedLogger = BraintrustLogger.getLogger(); - savedLevel = BraintrustLogger.getLevel(); - - // Set test logger - testLogger = new BraintrustLogger.TestLoggerImplementation(); - BraintrustLogger.setLogger(testLogger); - // Reset to default INFO level - BraintrustLogger.setLevel(Level.INFO); - // Clear any existing entries - testLogger.clear(); - } - - @AfterEach - void tearDown() { - // Restore original state - BraintrustLogger.setLogger(savedLogger); - BraintrustLogger.setLevel(savedLevel); - } - - @Test - void testDebugLogging() { - // Given - BraintrustLogger.setDebugEnabled(true); - - // When - BraintrustLogger.debug("Debug message with {}", "parameter"); - - // Then - assertThat(testLogger.getEntries()).hasSize(1); - var entry = testLogger.getEntries().get(0); - assertThat(entry.level()).isEqualTo(Level.DEBUG); - assertThat(entry.message()).isEqualTo("Debug message with parameter"); - assertThat(entry.throwable()).isNull(); - } - - @Test - void testDebugLoggingDisabled() { - // Given - BraintrustLogger.setDebugEnabled(false); - - // When - BraintrustLogger.debug("Debug message"); - - // Then - assertThat(testLogger.getEntries()).isEmpty(); - } - - @Test - void testDebugLoggingWithSupplier() { - // Given - BraintrustLogger.setDebugEnabled(true); - var supplierCalled = new boolean[] {false}; - - // When - BraintrustLogger.debug( - () -> { - supplierCalled[0] = true; - return "Expensive debug message"; - }); - - // Then - assertThat(supplierCalled[0]).isTrue(); - assertThat(testLogger.getEntries()).hasSize(1); - assertThat(testLogger.getEntries().get(0).message()).isEqualTo("Expensive debug message"); - } - - @Test - void testDebugSupplierNotCalledWhenDisabled() { - // Given - BraintrustLogger.setDebugEnabled(false); - var supplierCalled = new boolean[] {false}; - - // When - BraintrustLogger.debug( - () -> { - supplierCalled[0] = true; - return "Expensive debug message"; - }); - - // Then - assertThat(supplierCalled[0]).isFalse(); - assertThat(testLogger.getEntries()).isEmpty(); - } - - @Test - void testInfoLogging() { - // When - BraintrustLogger.info("Info message: {}", "test"); - - // Then - assertThat(testLogger.getEntries()).hasSize(1); - var entry = testLogger.getEntries().get(0); - assertThat(entry.level()).isEqualTo(Level.INFO); - assertThat(entry.message()).isEqualTo("Info message: test"); - } - - @Test - void testWarnLogging() { - // When - BraintrustLogger.warn("Warning: {}", "something"); - - // Then - assertThat(testLogger.getEntries()).hasSize(1); - var entry = testLogger.getEntries().get(0); - assertThat(entry.level()).isEqualTo(Level.WARN); - assertThat(entry.message()).isEqualTo("Warning: something"); - } - - @Test - void testWarnLoggingWithThrowable() { - // Given - var exception = new RuntimeException("Test exception"); - - // When - BraintrustLogger.warn("Warning with exception", exception); - - // Then - assertThat(testLogger.getEntries()).hasSize(1); - var entry = testLogger.getEntries().get(0); - assertThat(entry.level()).isEqualTo(Level.WARN); - assertThat(entry.message()).isEqualTo("Warning with exception"); - assertThat(entry.throwable()).isEqualTo(exception); - } - - @Test - void testErrorLogging() { - // When - BraintrustLogger.error("Error: {}", "critical"); - - // Then - assertThat(testLogger.getEntries()).hasSize(1); - var entry = testLogger.getEntries().get(0); - assertThat(entry.level()).isEqualTo(Level.ERROR); - assertThat(entry.message()).isEqualTo("Error: critical"); - } - - @Test - void testErrorLoggingWithThrowable() { - // Given - var exception = new IllegalStateException("Test error"); - - // When - BraintrustLogger.error("Error occurred", exception); - - // Then - assertThat(testLogger.getEntries()).hasSize(1); - var entry = testLogger.getEntries().get(0); - assertThat(entry.level()).isEqualTo(Level.ERROR); - assertThat(entry.message()).isEqualTo("Error occurred"); - assertThat(entry.throwable()).isEqualTo(exception); - } - - @Test - void testLevelFiltering() { - // Given - BraintrustLogger.setLevel(Level.WARN); - - // When - BraintrustLogger.debug("Debug"); - BraintrustLogger.info("Info"); - BraintrustLogger.warn("Warn"); - BraintrustLogger.error("Error"); - - // Then - assertThat(testLogger.getEntries()).hasSize(2); - assertThat(testLogger.getEntries()) - .extracting(BraintrustLogger.TestLoggerImplementation.LogEntry::level) - .containsExactly(Level.WARN, Level.ERROR); - } - - @Test - void testMessageFormatting() { - // When - BraintrustLogger.info("Message with {} {} {}", "multiple", "parameters", 123); - - // Then - assertThat(testLogger.getEntries()).hasSize(1); - assertThat(testLogger.getEntries().get(0).message()) - .isEqualTo("Message with multiple parameters 123"); - } - - @Test - void testClearLogEntries() { - // Given - BraintrustLogger.info("Message 1"); - BraintrustLogger.info("Message 2"); - assertThat(testLogger.getEntries()).hasSize(2); - - // When - testLogger.clear(); - - // Then - assertThat(testLogger.getEntries()).isEmpty(); - } - - @Test - void testThreadSafety() throws InterruptedException { - // Given - var threads = new Thread[10]; - var messagesPerThread = 100; - - // When - for (int i = 0; i < threads.length; i++) { - final int threadId = i; - threads[i] = - new Thread( - () -> { - for (int j = 0; j < messagesPerThread; j++) { - BraintrustLogger.info("Thread {} message {}", threadId, j); - } - }); - threads[i].start(); - } - - // Wait for all threads - for (Thread thread : threads) { - thread.join(); - } - - // Then - assertThat(testLogger.getEntries()).hasSize(threads.length * messagesPerThread); - } -} diff --git a/src/test/java/dev/braintrust/claude/test/otel/SpanAssertions.java b/src/test/java/dev/braintrust/claude/test/otel/SpanAssertions.java deleted file mode 100644 index 7eb8822..0000000 --- a/src/test/java/dev/braintrust/claude/test/otel/SpanAssertions.java +++ /dev/null @@ -1,224 +0,0 @@ -package dev.braintrust.claude.test.otel; - -import io.opentelemetry.api.common.AttributeKey; -import io.opentelemetry.api.trace.SpanKind; -import io.opentelemetry.api.trace.StatusCode; -import io.opentelemetry.sdk.trace.data.SpanData; -import java.time.Duration; -import java.time.Instant; -import java.util.Map; -import java.util.stream.Collectors; -import org.assertj.core.api.AbstractAssert; - -/** - * AssertJ assertions for OpenTelemetry spans. Provides fluent API for testing span attributes, - * events, and timing. - */ -public class SpanAssertions extends AbstractAssert { - - protected SpanAssertions(SpanData span) { - super(span, SpanAssertions.class); - } - - public static SpanAssertions assertThat(SpanData span) { - return new SpanAssertions(span); - } - - public SpanAssertions hasName(String name) { - isNotNull(); - if (!actual.getName().equals(name)) { - failWithMessage("Expected span name to be <%s> but was <%s>", name, actual.getName()); - } - return this; - } - - public SpanAssertions hasKind(SpanKind kind) { - isNotNull(); - if (!actual.getKind().equals(kind)) { - failWithMessage("Expected span kind to be <%s> but was <%s>", kind, actual.getKind()); - } - return this; - } - - public SpanAssertions hasStatus(StatusCode code) { - isNotNull(); - if (!actual.getStatus().getStatusCode().equals(code)) { - failWithMessage( - "Expected status code to be <%s> but was <%s>", - code, actual.getStatus().getStatusCode()); - } - return this; - } - - public SpanAssertions hasStatusWithDescription(StatusCode code, String description) { - hasStatus(code); - if (!actual.getStatus().getDescription().equals(description)) { - failWithMessage( - "Expected status description to be <%s> but was <%s>", - description, actual.getStatus().getDescription()); - } - return this; - } - - public SpanAssertions hasAttribute(AttributeKey key, T value) { - isNotNull(); - T actualValue = actual.getAttributes().get(key); - if (!value.equals(actualValue)) { - failWithMessage( - "Expected attribute <%s> to be <%s> but was <%s>", key, value, actualValue); - } - return this; - } - - public SpanAssertions hasAttributeKey(AttributeKey key) { - isNotNull(); - if (!actual.getAttributes().asMap().containsKey(key)) { - failWithMessage("Expected span to have attribute key <%s>", key); - } - return this; - } - - public SpanAssertions doesNotHaveAttributeKey(AttributeKey key) { - isNotNull(); - if (actual.getAttributes().asMap().containsKey(key)) { - failWithMessage("Expected span not to have attribute key <%s>", key); - } - return this; - } - - public SpanAssertions hasAttributeCount(int count) { - isNotNull(); - int actualCount = actual.getAttributes().size(); - if (actualCount != count) { - failWithMessage( - "Expected span to have <%d> attributes but had <%d>", count, actualCount); - } - return this; - } - - public SpanAssertions hasEventCount(int count) { - isNotNull(); - int actualCount = actual.getEvents().size(); - if (actualCount != count) { - failWithMessage("Expected span to have <%d> events but had <%d>", count, actualCount); - } - return this; - } - - public SpanAssertions hasEvent(String name) { - isNotNull(); - boolean found = actual.getEvents().stream().anyMatch(event -> event.getName().equals(name)); - if (!found) { - failWithMessage("Expected span to have event named <%s>", name); - } - return this; - } - - public SpanAssertions hasExceptionEvent() { - return hasEvent("exception"); - } - - public SpanAssertions hasExceptionEventWithType(Class exceptionType) { - hasExceptionEvent(); - - var exceptionEvents = - actual.getEvents().stream() - .filter(e -> e.getName().equals("exception")) - .collect(Collectors.toList()); - - boolean found = - exceptionEvents.stream() - .anyMatch( - event -> { - var type = - event.getAttributes() - .get(AttributeKey.stringKey("exception.type")); - return type != null && type.equals(exceptionType.getName()); - }); - - if (!found) { - failWithMessage( - "Expected span to have exception event of type <%s>", exceptionType.getName()); - } - return this; - } - - public SpanAssertions hasDurationBetween(Duration min, Duration max) { - isNotNull(); - var duration = Duration.ofNanos(actual.getEndEpochNanos() - actual.getStartEpochNanos()); - if (duration.compareTo(min) < 0 || duration.compareTo(max) > 0) { - failWithMessage( - "Expected span duration to be between <%s> and <%s> but was <%s>", - min, max, duration); - } - return this; - } - - public SpanAssertions startedAfter(Instant time) { - isNotNull(); - var startTime = Instant.ofEpochSecond(0, actual.getStartEpochNanos()); - if (!startTime.isAfter(time)) { - failWithMessage( - "Expected span to start after <%s> but started at <%s>", time, startTime); - } - return this; - } - - public SpanAssertions endedBefore(Instant time) { - isNotNull(); - var endTime = Instant.ofEpochSecond(0, actual.getEndEpochNanos()); - if (!endTime.isBefore(time)) { - failWithMessage("Expected span to end before <%s> but ended at <%s>", time, endTime); - } - return this; - } - - public SpanAssertions hasParentSpanId(String parentSpanId) { - isNotNull(); - if (!actual.getParentSpanId().equals(parentSpanId)) { - failWithMessage( - "Expected parent span ID to be <%s> but was <%s>", - parentSpanId, actual.getParentSpanId()); - } - return this; - } - - public SpanAssertions hasNoParent() { - isNotNull(); - if (!actual.getParentSpanId().isEmpty()) { - failWithMessage( - "Expected span to have no parent but had parent ID <%s>", - actual.getParentSpanId()); - } - return this; - } - - /** Asserts that the span has all the specified attributes with their values. */ - public SpanAssertions hasAttributes(Map, Object> expectedAttributes) { - isNotNull(); - for (Map.Entry, Object> entry : expectedAttributes.entrySet()) { - @SuppressWarnings("unchecked") - AttributeKey key = (AttributeKey) entry.getKey(); - hasAttribute(key, entry.getValue()); - } - return this; - } - - /** Asserts that the span has ended. */ - public SpanAssertions hasEnded() { - isNotNull(); - if (!actual.hasEnded()) { - failWithMessage("Expected span to have ended but it hasn't"); - } - return this; - } - - /** Asserts that the span has not ended. */ - public SpanAssertions hasNotEnded() { - isNotNull(); - if (actual.hasEnded()) { - failWithMessage("Expected span not to have ended but it has"); - } - return this; - } -} diff --git a/src/test/java/dev/braintrust/claude/trace/BraintrustContextTest.java b/src/test/java/dev/braintrust/claude/trace/BraintrustContextTest.java deleted file mode 100644 index 664a35f..0000000 --- a/src/test/java/dev/braintrust/claude/trace/BraintrustContextTest.java +++ /dev/null @@ -1,179 +0,0 @@ -package dev.braintrust.claude.trace; - -import static org.assertj.core.api.Assertions.*; - -import io.opentelemetry.context.Context; -import io.opentelemetry.context.Scope; -import org.junit.jupiter.api.Test; - -class BraintrustContextTest { - - @Test - void testForProject() { - // When - var context = BraintrustContext.forProject("project-123"); - - // Then - assertThat(context.projectId()).hasValue("project-123"); - assertThat(context.experimentId()).isEmpty(); - assertThat(context.parentType()).hasValue("project"); - } - - @Test - void testForExperiment() { - // When - var context = BraintrustContext.forExperiment("exp-456"); - - // Then - assertThat(context.projectId()).isEmpty(); - assertThat(context.experimentId()).hasValue("exp-456"); - assertThat(context.parentType()).hasValue("experiment"); - } - - @Test - void testStoreInContext() { - // Given - var btContext = BraintrustContext.forProject("project-789"); - var otelContext = Context.root(); - - // When - var newContext = btContext.storeInContext(otelContext); - - // Then - var retrieved = BraintrustContext.fromContext(newContext); - assertThat(retrieved).isNotNull(); - assertThat(retrieved.projectId()).hasValue("project-789"); - } - - @Test - void testFromContextNotPresent() { - // Given - var context = Context.root(); - - // When - var btContext = BraintrustContext.fromContext(context); - - // Then - assertThat(btContext).isNull(); - } - - @Test - void testCurrent() { - // Given - no context set - assertThat(BraintrustContext.current()).isNull(); - - // When - set context - var btContext = BraintrustContext.forExperiment("exp-current"); - var context = btContext.storeInContext(Context.current()); - - try (Scope scope = context.makeCurrent()) { - // Then - var current = BraintrustContext.current(); - assertThat(current).isNotNull(); - assertThat(current.experimentId()).hasValue("exp-current"); - } - - // After scope closes - assertThat(BraintrustContext.current()).isNull(); - } - - @Test - void testBuilder() { - // When - var context = BraintrustContext.builder().projectId("project-builder").build(); - - // Then - assertThat(context.projectId()).hasValue("project-builder"); - assertThat(context.experimentId()).isEmpty(); - assertThat(context.parentType()).hasValue("project"); - } - - @Test - void testBuilderExperiment() { - // When - var context = BraintrustContext.builder().experimentId("exp-builder").build(); - - // Then - assertThat(context.projectId()).isEmpty(); - assertThat(context.experimentId()).hasValue("exp-builder"); - assertThat(context.parentType()).hasValue("experiment"); - } - - @Test - void testBuilderOverwritesProject() { - // When - experiment overwrites project - var context = - BraintrustContext.builder() - .projectId("project-first") - .experimentId("exp-second") - .build(); - - // Then - experiment takes precedence - assertThat(context.projectId()).isEmpty(); - assertThat(context.experimentId()).hasValue("exp-second"); - assertThat(context.parentType()).hasValue("experiment"); - } - - @Test - void testBuilderStoreInContext() { - // When - var newContext = - BraintrustContext.builder() - .projectId("project-store") - .storeInContext(Context.root()); - - // Then - var retrieved = BraintrustContext.fromContext(newContext); - assertThat(retrieved).isNotNull(); - assertThat(retrieved.projectId()).hasValue("project-store"); - } - - @Test - void testBuilderStoreInCurrent() { - // Given - var originalContext = Context.current(); - - // When - var newContext = BraintrustContext.builder().experimentId("exp-current").storeInCurrent(); - - try (Scope scope = newContext.makeCurrent()) { - // Then - var current = BraintrustContext.current(); - assertThat(current).isNotNull(); - assertThat(current.experimentId()).hasValue("exp-current"); - } - } - - @Test - void testContextPropagation() { - // Given - var parent = BraintrustContext.forProject("parent-project"); - var parentContext = parent.storeInContext(Context.root()); - - // When - create child context - var child = BraintrustContext.forExperiment("child-exp"); - var childContext = child.storeInContext(parentContext); - - // Then - child context overwrites parent - var retrieved = BraintrustContext.fromContext(childContext); - assertThat(retrieved).isNotNull(); - assertThat(retrieved.projectId()).isEmpty(); - assertThat(retrieved.experimentId()).hasValue("child-exp"); - - // Parent context is not affected - var parentRetrieved = BraintrustContext.fromContext(parentContext); - assertThat(parentRetrieved).isNotNull(); - assertThat(parentRetrieved.projectId()).hasValue("parent-project"); - } - - @Test - void testNullValues() { - // When - var context = new BraintrustContext.Builder().projectId(null).build(); - - // Then - assertThat(context.projectId()).isEmpty(); - assertThat(context.experimentId()).isEmpty(); - assertThat(context.parentType()).isEmpty(); - } -} diff --git a/src/test/java/dev/braintrust/claude/trace/BraintrustSpanProcessorTest.java b/src/test/java/dev/braintrust/claude/trace/BraintrustSpanProcessorTest.java deleted file mode 100644 index dcd8ee9..0000000 --- a/src/test/java/dev/braintrust/claude/trace/BraintrustSpanProcessorTest.java +++ /dev/null @@ -1,142 +0,0 @@ -package dev.braintrust.claude.trace; - -import static org.mockito.Mockito.*; - -import dev.braintrust.claude.config.BraintrustConfig; -import io.opentelemetry.context.Context; -import io.opentelemetry.sdk.testing.junit5.OpenTelemetryExtension; -import io.opentelemetry.sdk.trace.ReadWriteSpan; -import io.opentelemetry.sdk.trace.ReadableSpan; -import io.opentelemetry.sdk.trace.SpanProcessor; -import io.opentelemetry.sdk.trace.data.SpanData; -import java.util.concurrent.atomic.AtomicReference; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.RegisterExtension; - -class BraintrustSpanProcessorTest { - - @RegisterExtension - static final OpenTelemetryExtension otelTesting = OpenTelemetryExtension.create(); - - private BraintrustConfig config; - private SpanProcessor mockDelegate; - private BraintrustSpanProcessor processor; - - @BeforeEach - void setUp() { - config = - BraintrustConfig.builder() - .apiKey("test-key") - .defaultProjectId("default-project") - .build(); - - mockDelegate = mock(SpanProcessor.class); - processor = new BraintrustSpanProcessor(config, mockDelegate); - } - - @Test - void testDefaultProjectIdIsAdded() { - // Given - var span = mock(ReadWriteSpan.class); - when(span.getAttribute(BraintrustSpanProcessor.PARENT)).thenReturn(null); - - // When - processor.onStart(Context.root(), span); - - // Then - verify(span).setAttribute(BraintrustSpanProcessor.PARENT, "project_id:default-project"); - verify(mockDelegate).onStart(any(), eq(span)); - } - - @Test - void testProjectIdFromContext() { - // Given - var span = mock(ReadWriteSpan.class); - when(span.getAttribute(BraintrustSpanProcessor.PARENT)).thenReturn(null); - var context = - BraintrustContext.forProject("context-project").storeInContext(Context.root()); - - // When - processor.onStart(context, span); - - // Then - verify(span).setAttribute(BraintrustSpanProcessor.PARENT, "project_id:context-project"); - } - - @Test - void testExperimentIdFromContext() { - // Given - var span = mock(ReadWriteSpan.class); - when(span.getAttribute(BraintrustSpanProcessor.PARENT)).thenReturn(null); - var context = - BraintrustContext.forExperiment("experiment-123").storeInContext(Context.root()); - - // When - processor.onStart(context, span); - - // Then - verify(span).setAttribute(BraintrustSpanProcessor.PARENT, "experiment_id:experiment-123"); - } - - @Test - void testUsageMetricsAttributes() { - // Given - var tracer = otelTesting.getOpenTelemetry().getTracer("test"); - var capturedSpan = new AtomicReference(); - - var testProcessor = - new SpanProcessor() { - @Override - public void onStart(Context parentContext, ReadWriteSpan span) {} - - @Override - public boolean isStartRequired() { - return false; - } - - @Override - public void onEnd(ReadableSpan span) { - capturedSpan.set(span.toSpanData()); - } - - @Override - public boolean isEndRequired() { - return true; - } - }; - - var braintrustProcessor = new BraintrustSpanProcessor(config, testProcessor); - - // The test processor is already registered through the testing framework - - // When - var span = tracer.spanBuilder("test-span").startSpan(); - span.setAttribute(BraintrustSpanProcessor.USAGE_PROMPT_TOKENS, 100L); - span.setAttribute(BraintrustSpanProcessor.USAGE_COMPLETION_TOKENS, 50L); - span.setAttribute(BraintrustSpanProcessor.USAGE_TOTAL_TOKENS, 150L); - span.setAttribute(BraintrustSpanProcessor.USAGE_COST, 0.0025); - span.end(); - - // Then - the span would be processed through the registered processor - // We can verify this through the OTEL testing framework's span data - } - - @Test - void testScoringAttributes() { - // Given - var span = - otelTesting - .getOpenTelemetry() - .getTracer("test") - .spanBuilder("test-span") - .startSpan(); - - // When - try (var scope = span.makeCurrent()) { - BraintrustTracing.SpanUtils.addScore("accuracy", 0.95); - } - - span.end(); - } -} diff --git a/src/test/java/dev/braintrust/claude/trace/BraintrustTracingTest.java b/src/test/java/dev/braintrust/claude/trace/BraintrustTracingTest.java deleted file mode 100644 index 81c0302..0000000 --- a/src/test/java/dev/braintrust/claude/trace/BraintrustTracingTest.java +++ /dev/null @@ -1,274 +0,0 @@ -package dev.braintrust.claude.trace; - -import static org.assertj.core.api.Assertions.*; - -import dev.braintrust.claude.config.BraintrustConfig; -import io.opentelemetry.api.GlobalOpenTelemetry; -import io.opentelemetry.sdk.OpenTelemetrySdk; -import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; -import io.opentelemetry.sdk.trace.SdkTracerProvider; -import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; -import java.time.Duration; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -class BraintrustTracingTest { - - private InMemorySpanExporter spanExporter; - private BraintrustConfig config; - - @BeforeEach - void setUp() { - // Reset global OpenTelemetry - GlobalOpenTelemetry.resetForTest(); - - spanExporter = InMemorySpanExporter.create(); - config = - BraintrustConfig.builder() - .apiKey("test-key") - .defaultProjectId("test-project") - .build(); - } - - @AfterEach - void tearDown() { - GlobalOpenTelemetry.resetForTest(); - } - - @Test - void testQuickstartWithDefaults() { - // When - var otel = BraintrustTracing.quickstart(config); - - // Then - assertThat(otel).isNotNull(); - assertThat(GlobalOpenTelemetry.get()).isNotNull(); - - // Create a test span - var tracer = BraintrustTracing.getTracer(otel); - var span = tracer.spanBuilder("test").startSpan(); - span.end(); - - // Verify tracer name - assertThat(tracer).isNotNull(); - } - - @Test - void testQuickstartWithCustomization() { - // When - var otel = - BraintrustTracing.quickstart( - config, - builder -> - builder.serviceName("custom-service") - .resourceAttribute("custom.attribute", "value") - .exportInterval(Duration.ofSeconds(1)) - .maxQueueSize(1000) - .maxExportBatchSize(100) - .registerGlobal(false)); - - // Then - assertThat(otel).isNotNull(); - assertThat(GlobalOpenTelemetry.get()).isNotEqualTo(otel); // Not registered globally - } - - @Test - void testGetTracerWithGlobal() { - // Given - BraintrustTracing.quickstart(config); - - // When - var tracer = BraintrustTracing.getTracer(); - - // Then - assertThat(tracer).isNotNull(); - // Test by creating a span - can't directly access instrumentation scope info from Tracer - var span = tracer.spanBuilder("test").startSpan(); - span.end(); - } - - @Test - void testGetTracerWithSpecificInstance() { - // Given - var otel = BraintrustTracing.quickstart(config, builder -> builder.registerGlobal(false)); - - // When - var tracer = BraintrustTracing.getTracer(otel); - - // Then - assertThat(tracer).isNotNull(); - // Test by creating a span - var span = tracer.spanBuilder("test").startSpan(); - span.end(); - } - - @Test - void testSpanUtilsAddUsageMetrics() { - // Given - var otel = createTestOpenTelemetry(); - var tracer = otel.getTracer("test"); - - // When - var span = tracer.spanBuilder("test-span").startSpan(); - try (var scope = span.makeCurrent()) { - BraintrustTracing.SpanUtils.addUsageMetrics(100, 50, 0.025); - } finally { - span.end(); - } - - // Then - var spans = spanExporter.getFinishedSpanItems(); - assertThat(spans).hasSize(1); - - var spanData = spans.get(0); - assertThat(spanData.getAttributes().get(BraintrustSpanProcessor.USAGE_PROMPT_TOKENS)) - .isEqualTo(100L); - assertThat(spanData.getAttributes().get(BraintrustSpanProcessor.USAGE_COMPLETION_TOKENS)) - .isEqualTo(50L); - assertThat(spanData.getAttributes().get(BraintrustSpanProcessor.USAGE_TOTAL_TOKENS)) - .isEqualTo(150L); - assertThat(spanData.getAttributes().get(BraintrustSpanProcessor.USAGE_COST)) - .isEqualTo(0.025); - } - - @Test - void testSpanUtilsAddScore() { - // Given - var otel = createTestOpenTelemetry(); - var tracer = otel.getTracer("test"); - - // When - var span = tracer.spanBuilder("test-span").startSpan(); - try (var scope = span.makeCurrent()) { - BraintrustTracing.SpanUtils.addScore("accuracy", 0.95); - } finally { - span.end(); - } - - // Then - var spans = spanExporter.getFinishedSpanItems(); - assertThat(spans).hasSize(1); - - var spanData = spans.get(0); - assertThat(spanData.getAttributes().get(BraintrustSpanProcessor.SCORE_NAME)) - .isEqualTo("accuracy"); - assertThat(spanData.getAttributes().get(BraintrustSpanProcessor.SCORE)).isEqualTo(0.95); - } - - @Test - void testSpanUtilsSetParentProject() { - // Given - var otel = createTestOpenTelemetry(); - var tracer = otel.getTracer("test"); - - // When - var span = tracer.spanBuilder("test-span").startSpan(); - try (var scope = span.makeCurrent()) { - BraintrustTracing.SpanUtils.setParentProject("project-123"); - } finally { - span.end(); - } - - // Then - var spans = spanExporter.getFinishedSpanItems(); - assertThat(spans).hasSize(1); - - var spanData = spans.get(0); - assertThat(spanData.getAttributes().get(BraintrustSpanProcessor.PARENT)) - .isEqualTo("project_id:project-123"); - } - - @Test - void testSpanUtilsSetParentExperiment() { - // Given - var otel = createTestOpenTelemetry(); - var tracer = otel.getTracer("test"); - - // When - var span = tracer.spanBuilder("test-span").startSpan(); - try (var scope = span.makeCurrent()) { - BraintrustTracing.SpanUtils.setParentExperiment("exp-456"); - } finally { - span.end(); - } - - // Then - var spans = spanExporter.getFinishedSpanItems(); - assertThat(spans).hasSize(1); - - var spanData = spans.get(0); - assertThat(spanData.getAttributes().get(BraintrustSpanProcessor.PARENT)) - .isEqualTo("experiment_id:exp-456"); - } - - @Test - void testBuilderResourceAttributes() { - // When - var otel = - BraintrustTracing.quickstart( - config, - builder -> - builder.serviceName("test-service") - .resourceAttribute("environment", "test") - .resourceAttribute("version", "1.2.3") - .registerGlobal(false)); - - // Create a span to verify resource attributes - var tracer = otel.getTracer("test"); - var span = tracer.spanBuilder("test").startSpan(); - span.end(); - - // Then - assertThat(otel).isNotNull(); - } - - @Test - void testConcurrentSpanCreation() throws InterruptedException { - // Given - var otel = BraintrustTracing.quickstart(config); - var tracer = BraintrustTracing.getTracer(); - var threadCount = 10; - var spansPerThread = 100; - - // When - var threads = new Thread[threadCount]; - for (int i = 0; i < threadCount; i++) { - final int threadId = i; - threads[i] = - new Thread( - () -> { - for (int j = 0; j < spansPerThread; j++) { - var span = - tracer.spanBuilder("thread-" + threadId + "-span-" + j) - .startSpan(); - // Simulate some work - try { - Thread.sleep(1); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - } - span.end(); - } - }); - threads[i].start(); - } - - // Wait for all threads - for (Thread thread : threads) { - thread.join(); - } - - // Then - no exceptions thrown, all spans created successfully - assertThat(tracer).isNotNull(); - } - - private OpenTelemetrySdk createTestOpenTelemetry() { - var tracerProvider = - SdkTracerProvider.builder() - .addSpanProcessor(SimpleSpanProcessor.create(spanExporter)) - .build(); - - return OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).build(); - } -} diff --git a/src/test/java/dev/braintrust/config/BraintrustConfigTest.java b/src/test/java/dev/braintrust/config/BraintrustConfigTest.java new file mode 100644 index 0000000..bfece25 --- /dev/null +++ b/src/test/java/dev/braintrust/config/BraintrustConfigTest.java @@ -0,0 +1,8 @@ +package dev.braintrust.config; + +import org.junit.jupiter.api.Test; + +class BraintrustConfigTest { + @Test + void todo() {} +} From 7de7048bd97b41a637637c6abc86a5d30aaf9493 Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Mon, 22 Sep 2025 12:18:33 -0600 Subject: [PATCH 2/5] run spotless formatter to get CI working again --- .../examples/CustomOpenTelemetryExample.java | 94 +- .../examples/ExperimentExample.java | 66 +- .../OpenAIInstrumentationExample.java | 17 +- .../examples/SimpleOpenTelemetryExample.java | 7 +- .../braintrust/api/BraintrustApiClient.java | 15 +- .../braintrust/config/BraintrustConfig.java | 41 +- src/main/java/dev/braintrust/eval/Eval.java | 120 ++- .../java/dev/braintrust/eval/EvalCase.java | 16 +- src/main/java/dev/braintrust/eval/Scorer.java | 9 +- .../openai/BraintrustOpenAI.java | 3 +- .../openai/otel/ChatAttributesGetter.java | 267 +++--- .../otel/ChatCompletionEventsHelper.java | 820 +++++++++--------- .../openai/otel/CompletableFutureWrapper.java | 30 +- .../otel/DelegatingInvocationHandler.java | 42 +- .../otel/EmbeddingAttributesGetter.java | 222 ++--- .../openai/otel/GenAiAttributes.java | 22 +- .../InstrumentedChatCompletionService.java | 229 ++--- ...nstrumentedChatCompletionServiceAsync.java | 240 ++--- .../openai/otel/InstrumentedChatService.java | 59 +- .../otel/InstrumentedChatServiceAsync.java | 59 +- .../otel/InstrumentedEmbeddingService.java | 86 +- .../InstrumentedEmbeddingServiceAsync.java | 87 +- .../openai/otel/InstrumentedOpenAiClient.java | 89 +- .../otel/InstrumentedOpenAiClientAsync.java | 90 +- .../openai/otel/OpenAITelemetry.java | 81 +- .../openai/otel/OpenAITelemetryBuilder.java | 86 +- .../openai/otel/StreamListener.java | 156 ++-- .../openai/otel/StreamedMessageBuffer.java | 216 ++--- .../otel/TracingAsyncStreamedResponse.java | 86 +- .../openai/otel/TracingStreamedResponse.java | 116 +-- .../dev/braintrust/log/BraintrustLogger.java | 7 +- .../braintrust/trace/BraintrustContext.java | 6 +- .../trace/BraintrustLogExporter.java | 7 +- .../trace/BraintrustSpanExporter.java | 1 - .../trace/BraintrustSpanProcessor.java | 36 +- .../braintrust/trace/BraintrustTracing.java | 121 +-- .../braintrust/trace/DebugSpanExporter.java | 1 - 37 files changed, 1935 insertions(+), 1715 deletions(-) diff --git a/examples/src/main/java/dev/braintrust/examples/CustomOpenTelemetryExample.java b/examples/src/main/java/dev/braintrust/examples/CustomOpenTelemetryExample.java index 9efbd85..da4e3ef 100644 --- a/examples/src/main/java/dev/braintrust/examples/CustomOpenTelemetryExample.java +++ b/examples/src/main/java/dev/braintrust/examples/CustomOpenTelemetryExample.java @@ -15,41 +15,52 @@ import io.opentelemetry.sdk.metrics.export.PeriodicMetricReader; import io.opentelemetry.sdk.trace.SdkTracerProvider; import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; - import java.util.concurrent.TimeUnit; import java.util.stream.Stream; -/** - * Example showing how to add Braintrust to an existing open telemetry setup. - */ +/** Example showing how to add Braintrust to an existing open telemetry setup. */ public class CustomOpenTelemetryExample { public static void main(String[] args) throws Exception { // Not required to run this example, but this assumes you have a local collector running. // See repo root README for instructions to set up a local collector that prints to stdout var localCollector = "http://localhost:4318"; - var tracerBuilder = SdkTracerProvider.builder() - .addSpanProcessor(BatchSpanProcessor.builder(OtlpHttpSpanExporter.builder() - .setEndpoint(localCollector + "/v1/traces") - .build()).build()); - var loggerBuilder = SdkLoggerProvider.builder() - .addLogRecordProcessor(BatchLogRecordProcessor.builder(OtlpHttpLogRecordExporter.builder() - .setEndpoint(localCollector + "/v1/logs") - .build()).build()); - var meterBuilder = SdkMeterProvider.builder() - .registerMetricReader(PeriodicMetricReader.builder(OtlpHttpMetricExporter.builder() - .setEndpoint(localCollector + "/v1/metrics") - .build()).build()); + var tracerBuilder = + SdkTracerProvider.builder() + .addSpanProcessor( + BatchSpanProcessor.builder( + OtlpHttpSpanExporter.builder() + .setEndpoint(localCollector + "/v1/traces") + .build()) + .build()); + var loggerBuilder = + SdkLoggerProvider.builder() + .addLogRecordProcessor( + BatchLogRecordProcessor.builder( + OtlpHttpLogRecordExporter.builder() + .setEndpoint(localCollector + "/v1/logs") + .build()) + .build()); + var meterBuilder = + SdkMeterProvider.builder() + .registerMetricReader( + PeriodicMetricReader.builder( + OtlpHttpMetricExporter.builder() + .setEndpoint(localCollector + "/v1/metrics") + .build()) + .build()); // NOTE: there are many ways to set up otel builders, etc. - // The important line is here: call enable with your otel builders and braintrust will export open telemetry data in addition to your existing setup + // The important line is here: call enable with your otel builders and braintrust will + // export open telemetry data in addition to your existing setup var braintrustConfig = BraintrustConfig.fromEnvironment(); BraintrustTracing.enable(braintrustConfig, tracerBuilder, loggerBuilder, meterBuilder); - var openTelemetry = OpenTelemetrySdk.builder() - .setTracerProvider(tracerBuilder.build()) - .setLoggerProvider(loggerBuilder.build()) - .setMeterProvider(meterBuilder.build()) - .build(); + var openTelemetry = + OpenTelemetrySdk.builder() + .setTracerProvider(tracerBuilder.build()) + .setLoggerProvider(loggerBuilder.build()) + .setMeterProvider(meterBuilder.build()) + .build(); GlobalOpenTelemetry.set(openTelemetry); registerShutdownHook(openTelemetry); var braintrustTracer = BraintrustTracing.getTracer(openTelemetry); @@ -63,7 +74,12 @@ public static void main(String[] args) throws Exception { } finally { span.end(); } - var url = braintrustConfig.fetchProjectURI() + "/logs?r=%s&s=%s".formatted(span.getSpanContext().getSpanId(), span.getSpanContext().getSpanId()); + var url = + braintrustConfig.fetchProjectURI() + + "/logs?r=%s&s=%s" + .formatted( + span.getSpanContext().getSpanId(), + span.getSpanContext().getSpanId()); System.out.println("\n\n Example complete! View your data in Braintrust: " + url); } @@ -72,14 +88,30 @@ private static void registerShutdownHook(OpenTelemetrySdk otel) { .addShutdownHook( new Thread( () -> { - BraintrustLogger.debug("Shutting down. Force-Flushing all otel data."); - var result = CompletableResultCode.ofAll( - // run all flushes in parallel. Should (rarely) block for approx 10 seconds max - Stream.of(otel.getSdkTracerProvider().shutdown(), otel.getSdkLoggerProvider().shutdown(), otel.getSdkMeterProvider().shutdown()) - .map(operation -> operation.join(10, TimeUnit.SECONDS)) - .toList() - ); - BraintrustLogger.debug("otel shutdown complete. Flush done: %s, Flush successful: %s".formatted(result.isDone(), result.isSuccess())); + BraintrustLogger.debug( + "Shutting down. Force-Flushing all otel data."); + var result = + CompletableResultCode.ofAll( + // run all flushes in parallel. Should (rarely) + // block for approx 10 seconds max + Stream.of( + otel.getSdkTracerProvider() + .shutdown(), + otel.getSdkLoggerProvider() + .shutdown(), + otel.getSdkMeterProvider() + .shutdown()) + .map( + operation -> + operation.join( + 10, + TimeUnit + .SECONDS)) + .toList()); + BraintrustLogger.debug( + "otel shutdown complete. Flush done: %s, Flush successful: %s" + .formatted( + result.isDone(), result.isSuccess())); })); } } diff --git a/examples/src/main/java/dev/braintrust/examples/ExperimentExample.java b/examples/src/main/java/dev/braintrust/examples/ExperimentExample.java index 0a0eb68..d3690db 100644 --- a/examples/src/main/java/dev/braintrust/examples/ExperimentExample.java +++ b/examples/src/main/java/dev/braintrust/examples/ExperimentExample.java @@ -9,7 +9,6 @@ import dev.braintrust.eval.Scorer; import dev.braintrust.instrumentation.openai.BraintrustOpenAI; import dev.braintrust.trace.BraintrustTracing; - import java.util.function.Function; public class ExperimentExample { @@ -18,32 +17,45 @@ public static void main(String[] args) throws Exception { var openTelemetry = BraintrustTracing.quickstart(config, true); var openAIClient = BraintrustOpenAI.wrapOpenAI(openTelemetry, OpenAIOkHttpClient.fromEnv()); - Function getFoodType = (String food) -> { - var request = - ChatCompletionCreateParams.builder() - .model(ChatModel.GPT_4O_MINI) - .addSystemMessage("Return a one word answer") - .addUserMessage("What kind of food is " + food + "?") - .maxTokens(50L) - .temperature(0.0) - .build(); - var response = openAIClient.chat().completions().create(request); - return response.choices().get(0).message().content().orElse("").toLowerCase(); - }; + Function getFoodType = + (String food) -> { + var request = + ChatCompletionCreateParams.builder() + .model(ChatModel.GPT_4O_MINI) + .addSystemMessage("Return a one word answer") + .addUserMessage("What kind of food is " + food + "?") + .maxTokens(50L) + .temperature(0.0) + .build(); + var response = openAIClient.chat().completions().create(request); + return response.choices().get(0).message().content().orElse("").toLowerCase(); + }; - var eval = Eval.builder() - .name("java-eval-x-" + System.currentTimeMillis()) // NOTE: if you use a constant, additional runs will append new cases to the same experiment - .tracer(BraintrustTracing.getTracer(openTelemetry)) - .config(config) - .cases(EvalCase.of("strawberry", "fruit"), - EvalCase.of("asparagus", "vegetable"), - EvalCase.of("apple", "fruit"), - EvalCase.of("banana", "fruit")) - .task(getFoodType) - .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 eval = + Eval.builder() + .name("java-eval-x-" + System.currentTimeMillis()) // NOTE: if you use a + // constant, additional runs + // will append new cases to + // the same experiment + .tracer(BraintrustTracing.getTracer(openTelemetry)) + .config(config) + .cases( + EvalCase.of("strawberry", "fruit"), + EvalCase.of("asparagus", "vegetable"), + EvalCase.of("apple", "fruit"), + EvalCase.of("banana", "fruit")) + .task(getFoodType) + .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(); - System.out.println("\n\n Example complete! View your data in Braintrust: " + result.createReportString()); + System.out.println( + "\n\n Example complete! View your data in Braintrust: " + + result.createReportString()); } -} \ No newline at end of file +} diff --git a/examples/src/main/java/dev/braintrust/examples/OpenAIInstrumentationExample.java b/examples/src/main/java/dev/braintrust/examples/OpenAIInstrumentationExample.java index 8c21a53..46c7ec2 100644 --- a/examples/src/main/java/dev/braintrust/examples/OpenAIInstrumentationExample.java +++ b/examples/src/main/java/dev/braintrust/examples/OpenAIInstrumentationExample.java @@ -8,18 +8,18 @@ import dev.braintrust.instrumentation.openai.BraintrustOpenAI; import dev.braintrust.trace.BraintrustTracing; -/** - * Basic OTel + OpenAI instrumentation example - */ +/** Basic OTel + OpenAI instrumentation example */ public class OpenAIInstrumentationExample { public static void main(String[] args) throws Exception { if (null == System.getenv("OPENAI_API_KEY")) { - System.err.println("\nWARNING envar OPEN_AI_API_KEY not found. This example will likely fail.\n"); + System.err.println( + "\nWARNING envar OPEN_AI_API_KEY not found. This example will likely fail.\n"); } var braintrustConfig = BraintrustConfig.fromEnvironment(); var openTelemetry = BraintrustTracing.quickstart(braintrustConfig, true); var tracer = BraintrustTracing.getTracer(openTelemetry); - OpenAIClient openAIClient = BraintrustOpenAI.wrapOpenAI(openTelemetry, OpenAIOkHttpClient.fromEnv()); + OpenAIClient openAIClient = + BraintrustOpenAI.wrapOpenAI(openTelemetry, OpenAIOkHttpClient.fromEnv()); var rootSpan = tracer.spanBuilder("java-braintrust-example").startSpan(); try (var ignored = rootSpan.makeCurrent()) { Thread.sleep(70); // Not required. This is just to make the span look interesting @@ -38,7 +38,12 @@ public static void main(String[] args) throws Exception { } finally { rootSpan.end(); } - var url = braintrustConfig.fetchProjectURI() + "/logs?r=%s&s=%s".formatted(rootSpan.getSpanContext().getSpanId(), rootSpan.getSpanContext().getSpanId()); + var url = + braintrustConfig.fetchProjectURI() + + "/logs?r=%s&s=%s" + .formatted( + rootSpan.getSpanContext().getSpanId(), + rootSpan.getSpanContext().getSpanId()); System.out.println("\n\n Example complete! View your data in Braintrust: " + url); } } diff --git a/examples/src/main/java/dev/braintrust/examples/SimpleOpenTelemetryExample.java b/examples/src/main/java/dev/braintrust/examples/SimpleOpenTelemetryExample.java index c166338..ec2398b 100644 --- a/examples/src/main/java/dev/braintrust/examples/SimpleOpenTelemetryExample.java +++ b/examples/src/main/java/dev/braintrust/examples/SimpleOpenTelemetryExample.java @@ -17,7 +17,12 @@ public static void main(String[] args) throws Exception { } finally { span.end(); } - var url = braintrustConfig.fetchProjectURI() + "/logs?r=%s&s=%s".formatted(span.getSpanContext().getSpanId(), span.getSpanContext().getSpanId()); + var url = + braintrustConfig.fetchProjectURI() + + "/logs?r=%s&s=%s" + .formatted( + span.getSpanContext().getSpanId(), + span.getSpanContext().getSpanId()); System.out.println("\n\n Example complete! View your data in Braintrust: " + url); } } diff --git a/src/main/java/dev/braintrust/api/BraintrustApiClient.java b/src/main/java/dev/braintrust/api/BraintrustApiClient.java index d2ef611..0effd5b 100644 --- a/src/main/java/dev/braintrust/api/BraintrustApiClient.java +++ b/src/main/java/dev/braintrust/api/BraintrustApiClient.java @@ -7,7 +7,6 @@ import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; import dev.braintrust.config.BraintrustConfig; import dev.braintrust.log.BraintrustLogger; - import java.io.IOException; import java.net.URI; import java.net.http.HttpClient; @@ -58,7 +57,8 @@ public Optional getProject(String projectId) { throw new CompletionException(error); } return Optional.of(project); - }).get(); + }) + .get(); } catch (InterruptedException | ExecutionException e) { throw new RuntimeException(e); } @@ -76,7 +76,11 @@ public Experiment getOrCreateExperiment(CreateExperimentRequest request) { /** Login to get user information including organization details. */ public LoginResponse login() { try { - return postAsync("/api/apikey/login", new LoginRequest(config.apiKey()), LoginResponse.class).get(); + return postAsync( + "/api/apikey/login", + new LoginRequest(config.apiKey()), + LoginResponse.class) + .get(); } catch (InterruptedException | ExecutionException e) { throw new RuntimeException(e); } @@ -85,7 +89,7 @@ public LoginResponse login() { public Optional getProjectAndOrgInfo() { var projectId = config.defaultProjectId().orElse(null); if (null == projectId) { - projectId = getOrCreateProject(config.defaultProjectName().orElseThrow()).id(); + projectId = getOrCreateProject(config.defaultProjectName().orElseThrow()).id(); } return getProjectAndOrgInfo(projectId); } @@ -107,7 +111,8 @@ public Optional getProjectAndOrgInfo(String projectI } } if (null == orgInfo) { - throw new ApiException("Should not happen. Unable to find project's org: " + project.orgId()); + throw new ApiException( + "Should not happen. Unable to find project's org: " + project.orgId()); } return Optional.of(new OrganizationAndProjectInfo(orgInfo, project)); } diff --git a/src/main/java/dev/braintrust/config/BraintrustConfig.java b/src/main/java/dev/braintrust/config/BraintrustConfig.java index 858ca18..924b633 100644 --- a/src/main/java/dev/braintrust/config/BraintrustConfig.java +++ b/src/main/java/dev/braintrust/config/BraintrustConfig.java @@ -1,12 +1,10 @@ package dev.braintrust.config; import dev.braintrust.api.BraintrustApiClient; - import java.net.URI; import java.net.URISyntaxException; import java.time.Duration; import java.util.Optional; -import java.util.function.Consumer; import javax.annotation.Nullable; /** @@ -85,12 +83,15 @@ public Duration requestTimeout() { } /** - * The parent attribute tells braintrust where to send otel data - *

- * The otel ingestion endpoint looks for (a) braintrust.parent = project_id|project_name|experiment_id:value otel attribute and routes accordingly - *

- * (b) if a span has no parent marked explicitly, it will look to see if there's an x-bt-parent http header (with the same format marked above e.g. project_name:andrew) that parent will apply to all spans in a request that don't have one - *

+ * The parent attribute tells braintrust where to send otel data
+ *
+ * The otel ingestion endpoint looks for (a) braintrust.parent = + * project_id|project_name|experiment_id:value otel attribute and routes accordingly
+ *
+ * (b) if a span has no parent marked explicitly, it will look to see if there's an x-bt-parent + * http header (with the same format marked above e.g. project_name:andrew) that parent will + * apply to all spans in a request that don't have one
+ *
* If neither (a) nor (b) exists, the data is dropped */ public Optional getBraintrustParentValue() { @@ -117,7 +118,12 @@ public URI fetchProjectURI() { try { var client = new BraintrustApiClient(this); var orgAndProject = client.getProjectAndOrgInfo().orElseThrow(); - return new URI(appUrl() + "/app/" + orgAndProject.orgInfo().name() + "/p/" + orgAndProject.project().name()); + return new URI( + appUrl() + + "/app/" + + orgAndProject.orgInfo().name() + + "/p/" + + orgAndProject.project().name()); } catch (URISyntaxException e) { throw new RuntimeException(e); } @@ -143,17 +149,22 @@ private Builder() { this.logsPath = getEnv("BRAINTRUST_LOGS_PATH", "/otel/v1/logs"); this.appUrl = URI.create(getEnv("BRAINTRUST_APP_URL", DEFAULT_APP_URL)); this.defaultProjectId = getEnv("BRAINTRUST_DEFAULT_PROJECT_ID", null); - this.defaultProjectName = getEnv("BRAINTRUST_DEFAULT_PROJECT", DEFAULT_PROJECT_NAME).trim(); + this.defaultProjectName = + getEnv("BRAINTRUST_DEFAULT_PROJECT", DEFAULT_PROJECT_NAME).trim(); if ((null == defaultProjectId || "".equalsIgnoreCase(defaultProjectId.trim())) && "".equalsIgnoreCase(defaultProjectName)) { // NOTE: this should not happen, - // but if someone happens to export their default project to the empty string and does not set a default project ID we don't have a valid parent for otel data. - throw new RuntimeException("Missing required envars. Please export BRAINTRUST_DEFAULT_PROJECT_ID or BRAINTRUST_DEFAULT_PROJECT"); + // but if someone happens to export their default project to the empty string and + // does not set a default project ID we don't have a valid parent for otel data. + throw new RuntimeException( + "Missing required envars. Please export BRAINTRUST_DEFAULT_PROJECT_ID or" + + " BRAINTRUST_DEFAULT_PROJECT"); } this.enableTraceConsoleLog = Boolean.parseBoolean(getEnv("BRAINTRUST_ENABLE_TRACE_CONSOLE_LOG", "false")); this.debug = Boolean.parseBoolean(getEnv("BRAINTRUST_DEBUG", "false")); - this.experimentalOtelLogs = Boolean.parseBoolean(getEnv("BRAINTRUST_X_OTEL_LOGS", "false")); + this.experimentalOtelLogs = + Boolean.parseBoolean(getEnv("BRAINTRUST_X_OTEL_LOGS", "false")); } public Builder apiKey(String apiKey) { @@ -201,7 +212,9 @@ public Builder requestTimeout(Duration timeout) { public BraintrustConfig build() { if (apiKey == null || apiKey.isBlank()) { - throw new IllegalStateException("API key is required. Set BRAINTRUST_API_KEY environment variable or use apiKey() method."); + throw new IllegalStateException( + "API key is required. Set BRAINTRUST_API_KEY environment variable or use" + + " apiKey() method."); } return new BraintrustConfig(this); diff --git a/src/main/java/dev/braintrust/eval/Eval.java b/src/main/java/dev/braintrust/eval/Eval.java index b937574..7a2f31b 100644 --- a/src/main/java/dev/braintrust/eval/Eval.java +++ b/src/main/java/dev/braintrust/eval/Eval.java @@ -9,11 +9,10 @@ import dev.braintrust.trace.BraintrustTracing; import io.opentelemetry.api.trace.SpanKind; import io.opentelemetry.api.trace.Tracer; - -import javax.annotation.Nonnull; -import javax.annotation.Nullable; import java.util.*; import java.util.function.Function; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; /** * An evaluation framework for testing AI models. @@ -22,7 +21,8 @@ * @param The type of output produced by the task */ public final class Eval { - private static final ObjectMapper JSON_MAPPER = new com.fasterxml.jackson.databind.ObjectMapper(); + private static final ObjectMapper JSON_MAPPER = + new com.fasterxml.jackson.databind.ObjectMapper(); private final @Nonnull String experimentName; private final @Nonnull BraintrustConfig config; private final @Nonnull BraintrustApiClient client; @@ -39,8 +39,12 @@ private Eval(Builder builder) { if (null == builder.projectId) { this.orgAndProject = client.getProjectAndOrgInfo().orElseThrow(); } else { - this.orgAndProject = client.getProjectAndOrgInfo(builder.projectId) - .orElseThrow(() -> new RuntimeException("invalid project id: " + builder.projectId)); + this.orgAndProject = + client.getProjectAndOrgInfo(builder.projectId) + .orElseThrow( + () -> + new RuntimeException( + "invalid project id: " + builder.projectId)); } this.tracer = Objects.requireNonNull(builder.tracer); this.evalCases = List.copyOf(builder.evalCases); @@ -50,66 +54,95 @@ private Eval(Builder builder) { /** Runs the evaluation and returns results. */ public Result run() { - var experiment = client.getOrCreateExperiment(new BraintrustApiClient.CreateExperimentRequest(orgAndProject.project().id(), experimentName, Optional.empty(), Optional.empty())); + var experiment = + client.getOrCreateExperiment( + new BraintrustApiClient.CreateExperimentRequest( + orgAndProject.project().id(), + experimentName, + Optional.empty(), + Optional.empty())); var experimentID = experiment.id(); - var evalCaseResults = evalCases.stream() - .map(evalCase -> evalOne(experimentID, evalCase)) - .toList(); + var evalCaseResults = + evalCases.stream().map(evalCase -> evalOne(experimentID, evalCase)).toList(); return new Result(); } - private EvalCase.Result evalOne(String experimentId, EvalCase evalCase) { + private EvalCase.Result evalOne( + String experimentId, EvalCase evalCase) { var rootSpan = 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( + BraintrustSpanProcessor.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.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.forExperiment(experimentId, rootSpan).makeCurrent()) { final OUTPUT result; { // run task - var taskSpan = tracer.spanBuilder("task") - .setAttribute(BraintrustSpanProcessor.PARENT, "experiment_id:" + experimentId) - .setAttribute("braintrust.span_attributes", "{\"type\":\"task\"}") - .startSpan(); - try (var unused = BraintrustContext.forExperiment(experimentId, taskSpan).makeCurrent()){ + var taskSpan = + tracer.spanBuilder("task") + .setAttribute( + BraintrustSpanProcessor.PARENT, + "experiment_id:" + experimentId) + .setAttribute("braintrust.span_attributes", "{\"type\":\"task\"}") + .startSpan(); + try (var unused = + BraintrustContext.forExperiment(experimentId, taskSpan).makeCurrent()) { result = task.apply(evalCase); } finally { taskSpan.end(); } try { - rootSpan.setAttribute("braintrust.output_json", JSON_MAPPER.writeValueAsString(Map.of("output", String.valueOf(result)))); + rootSpan.setAttribute( + "braintrust.output_json", + JSON_MAPPER.writeValueAsString( + Map.of("output", String.valueOf(result)))); } catch (JsonProcessingException e) { throw new RuntimeException(e); } } { // run scorers - var scoreSpan = tracer.spanBuilder("score") - .setAttribute(BraintrustSpanProcessor.PARENT, "experiment_id:" + experimentId) - .setAttribute("braintrust.span_attributes", "{\"type\":\"score\"}") - .startSpan(); - try (var unused = BraintrustContext.forExperiment(experimentId, scoreSpan).makeCurrent()){ - // NOTE: linked hash map to preserve ordering. Not in the spec but nice user experience + var scoreSpan = + tracer.spanBuilder("score") + .setAttribute( + BraintrustSpanProcessor.PARENT, + "experiment_id:" + experimentId) + .setAttribute("braintrust.span_attributes", "{\"type\":\"score\"}") + .startSpan(); + try (var unused = + BraintrustContext.forExperiment(experimentId, scoreSpan).makeCurrent()) { + // NOTE: linked hash map to preserve ordering. Not in the spec but nice user + // experience final HashMap nameToScore = new LinkedHashMap<>(); - var scores = scorers.stream() - .map(scorer -> { - var score = scorer.score(evalCase, result); - if (score < 0.0 || score > 1.0) { - throw new RuntimeException("score must be between 0 and 1: " + scorer.getName() + " : " + score); - } - nameToScore.put(scorer.getName(), score); - return score; - }) - .toList(); + var scores = + scorers.stream() + .map( + scorer -> { + var score = scorer.score(evalCase, result); + if (score < 0.0 || score > 1.0) { + throw new RuntimeException( + "score must be between 0 and 1: " + + scorer.getName() + + " : " + + score); + } + nameToScore.put(scorer.getName(), score); + return score; + }) + .toList(); try { - scoreSpan.setAttribute("braintrust.scores", JSON_MAPPER.writeValueAsString(nameToScore)); + scoreSpan.setAttribute( + "braintrust.scores", JSON_MAPPER.writeValueAsString(nameToScore)); } catch (JsonProcessingException e) { throw new RuntimeException(e); } @@ -123,14 +156,19 @@ private EvalCase.Result evalOne(String experimentId, EvalCase(INPUT input, OUTPUT expected, @Nonnull List tags, - @Nonnull Map metadata) { +public record EvalCase( + INPUT input, + OUTPUT expected, + @Nonnull List tags, + @Nonnull Map metadata) { public EvalCase { if (!metadata.isEmpty()) { throw new RuntimeException("TODO: metadata support not yet implemented"); @@ -28,7 +30,11 @@ public static EvalCase of(INPUT input, OUTPUT exp return of(input, expected, List.of(), Map.of()); } - public static EvalCase of(INPUT input, OUTPUT expected, @Nonnull List tags, @Nonnull Map metadata) { + public static EvalCase of( + INPUT input, + OUTPUT expected, + @Nonnull List tags, + @Nonnull Map metadata) { return new EvalCase<>(input, expected, tags, metadata); } diff --git a/src/main/java/dev/braintrust/eval/Scorer.java b/src/main/java/dev/braintrust/eval/Scorer.java index 2292fda..189d5ed 100644 --- a/src/main/java/dev/braintrust/eval/Scorer.java +++ b/src/main/java/dev/braintrust/eval/Scorer.java @@ -4,7 +4,8 @@ import java.util.function.Function; /** - * A scorer evaluates the result of a test case with a score between 0 (inclusive) and 1 (inclusive). + * A scorer evaluates the result of a test case with a score between 0 (inclusive) and 1 + * (inclusive). * * @param type of the input data * @param type of the output data @@ -14,7 +15,8 @@ public interface Scorer { double score(EvalCase evalCase, OUTPUT result); - static Scorer of(String scorerName, BiFunction, OUTPUT, Double> scorerFn) { + static Scorer of( + String scorerName, BiFunction, OUTPUT, Double> scorerFn) { return new Scorer<>() { @Override public String getName() { @@ -28,7 +30,8 @@ public double score(EvalCase evalCase, OUTPUT result) { }; } - static Scorer of(String scorerName, Function scorerFn) { + static Scorer of( + String scorerName, Function scorerFn) { return of(scorerName, (evalCase, result) -> scorerFn.apply(result)); } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/BraintrustOpenAI.java b/src/main/java/dev/braintrust/instrumentation/openai/BraintrustOpenAI.java index ec088a3..1a62468 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/BraintrustOpenAI.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/BraintrustOpenAI.java @@ -7,7 +7,8 @@ public class BraintrustOpenAI { public static OpenAIClient wrapOpenAI(OpenTelemetry openTelemetry, OpenAIClient openAIClient) { if ("true".equalsIgnoreCase(System.getenv("BRAINTRUST_X_OTEL_LOGS"))) { - return io.opentelemetry.instrumentation.openai.v1_1.OpenAITelemetry.builder(openTelemetry) + return io.opentelemetry.instrumentation.openai.v1_1.OpenAITelemetry.builder( + openTelemetry) .setCaptureMessageContent(true) .build() .wrap(openAIClient); diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/ChatAttributesGetter.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/ChatAttributesGetter.java index df3031d..110a279 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/ChatAttributesGetter.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/ChatAttributesGetter.java @@ -17,138 +17,137 @@ import org.jetbrains.annotations.Nullable; enum ChatAttributesGetter - implements GenAiAttributesGetter { - INSTANCE; - - @Override - public String getOperationName(ChatCompletionCreateParams request) { - return GenAiAttributes.GenAiOperationNameIncubatingValues.CHAT; - } - - @Override - public String getSystem(ChatCompletionCreateParams request) { - return GenAiAttributes.GenAiProviderNameIncubatingValues.OPENAI; - } - - @Override - public String getRequestModel(ChatCompletionCreateParams request) { - return request.model().asString(); - } - - @Nullable - @Override - public Long getRequestSeed(ChatCompletionCreateParams request) { - return request.seed().orElse(null); - } - - @Nullable - @Override - public List getRequestEncodingFormats(ChatCompletionCreateParams request) { - return null; - } - - @Nullable - @Override - public Double getRequestFrequencyPenalty(ChatCompletionCreateParams request) { - return request.frequencyPenalty().orElse(null); - } - - @Nullable - @Override - public Long getRequestMaxTokens(ChatCompletionCreateParams request) { - return request.maxCompletionTokens().orElse(null); - } - - @Nullable - @Override - public Double getRequestPresencePenalty(ChatCompletionCreateParams request) { - return request.presencePenalty().orElse(null); - } - - @Nullable - @Override - public List getRequestStopSequences(ChatCompletionCreateParams request) { - return request - .stop() - .map( - s -> { - if (s.isString()) { - return singletonList(s.asString()); - } - if (s.isStrings()) { - return s.asStrings(); - } - return null; - }) - .orElse(null); - } - - @Nullable - @Override - public Double getRequestTemperature(ChatCompletionCreateParams request) { - return request.temperature().orElse(null); - } - - @Nullable - @Override - public Double getRequestTopK(ChatCompletionCreateParams request) { - return null; - } - - @Nullable - @Override - public Double getRequestTopP(ChatCompletionCreateParams request) { - return request.topP().orElse(null); - } - - @Override - public List getResponseFinishReasons( - ChatCompletionCreateParams request, @Nullable ChatCompletion response) { - if (response == null) { - return emptyList(); - } - return response.choices().stream() - .map(choice -> choice.finishReason().asString()) - .collect(Collectors.toList()); - } - - @Override - @Nullable - public String getResponseId( - ChatCompletionCreateParams request, @Nullable ChatCompletion response) { - if (response == null) { - return null; - } - return response.id(); - } - - @Override - @Nullable - public String getResponseModel( - ChatCompletionCreateParams request, @Nullable ChatCompletion response) { - if (response == null) { - return null; - } - return response.model(); - } - - @Override - @Nullable - public Long getUsageInputTokens( - ChatCompletionCreateParams request, @Nullable ChatCompletion response) { - if (response == null) { - return null; - } - return response.usage().map(CompletionUsage::promptTokens).orElse(null); - } - - @Override - @Nullable - public Long getUsageOutputTokens( - ChatCompletionCreateParams request, @Nullable ChatCompletion response) { - if (response == null) { - return null; - } - return response.usage().map(CompletionUsage::completionTokens).orElse(null); - } + implements GenAiAttributesGetter { + INSTANCE; + + @Override + public String getOperationName(ChatCompletionCreateParams request) { + return GenAiAttributes.GenAiOperationNameIncubatingValues.CHAT; + } + + @Override + public String getSystem(ChatCompletionCreateParams request) { + return GenAiAttributes.GenAiProviderNameIncubatingValues.OPENAI; + } + + @Override + public String getRequestModel(ChatCompletionCreateParams request) { + return request.model().asString(); + } + + @Nullable + @Override + public Long getRequestSeed(ChatCompletionCreateParams request) { + return request.seed().orElse(null); + } + + @Nullable + @Override + public List getRequestEncodingFormats(ChatCompletionCreateParams request) { + return null; + } + + @Nullable + @Override + public Double getRequestFrequencyPenalty(ChatCompletionCreateParams request) { + return request.frequencyPenalty().orElse(null); + } + + @Nullable + @Override + public Long getRequestMaxTokens(ChatCompletionCreateParams request) { + return request.maxCompletionTokens().orElse(null); + } + + @Nullable + @Override + public Double getRequestPresencePenalty(ChatCompletionCreateParams request) { + return request.presencePenalty().orElse(null); + } + + @Nullable + @Override + public List getRequestStopSequences(ChatCompletionCreateParams request) { + return request.stop() + .map( + s -> { + if (s.isString()) { + return singletonList(s.asString()); + } + if (s.isStrings()) { + return s.asStrings(); + } + return null; + }) + .orElse(null); + } + + @Nullable + @Override + public Double getRequestTemperature(ChatCompletionCreateParams request) { + return request.temperature().orElse(null); + } + + @Nullable + @Override + public Double getRequestTopK(ChatCompletionCreateParams request) { + return null; + } + + @Nullable + @Override + public Double getRequestTopP(ChatCompletionCreateParams request) { + return request.topP().orElse(null); + } + + @Override + public List getResponseFinishReasons( + ChatCompletionCreateParams request, @Nullable ChatCompletion response) { + if (response == null) { + return emptyList(); + } + return response.choices().stream() + .map(choice -> choice.finishReason().asString()) + .collect(Collectors.toList()); + } + + @Override + @Nullable + public String getResponseId( + ChatCompletionCreateParams request, @Nullable ChatCompletion response) { + if (response == null) { + return null; + } + return response.id(); + } + + @Override + @Nullable + public String getResponseModel( + ChatCompletionCreateParams request, @Nullable ChatCompletion response) { + if (response == null) { + return null; + } + return response.model(); + } + + @Override + @Nullable + public Long getUsageInputTokens( + ChatCompletionCreateParams request, @Nullable ChatCompletion response) { + if (response == null) { + return null; + } + return response.usage().map(CompletionUsage::promptTokens).orElse(null); + } + + @Override + @Nullable + public Long getUsageOutputTokens( + ChatCompletionCreateParams request, @Nullable ChatCompletion response) { + if (response == null) { + return null; + } + return response.usage().map(CompletionUsage::completionTokens).orElse(null); + } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/ChatCompletionEventsHelper.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/ChatCompletionEventsHelper.java index 6f7ce94..4926681 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/ChatCompletionEventsHelper.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/ChatCompletionEventsHelper.java @@ -40,439 +40,475 @@ final class ChatCompletionEventsHelper { - private static final AttributeKey EVENT_NAME = stringKey("event.name"); - private static final ObjectMapper JSON_MAPPER = new com.fasterxml.jackson.databind.ObjectMapper(); - - public static void emitPromptLogEvents( - Context context, - Logger eventLogger, - ChatCompletionCreateParams request, - boolean captureMessageContent) { - for (ChatCompletionMessageParam msg : request.messages()) { - String eventType; - Map> body = new HashMap<>(); - if (msg.isSystem()) { - eventType = "gen_ai.system.message"; - if (captureMessageContent) { - var content = contentToString(msg.asSystem().content()); - Span.current().setAttribute("instructions", content); - body.put("content", Value.of(content)); - } - } else if (msg.isDeveloper()) { - eventType = "gen_ai.system.message"; - body.put("role", Value.of("developer")); - if (captureMessageContent) { - var content = contentToString(msg.asDeveloper().content()); - body.put("content", Value.of(content)); - } - } else if (msg.isUser()) { - eventType = "gen_ai.user.message"; - if (captureMessageContent) { - var content = contentToString(msg.asUser().content()); - try { - Span.current().setAttribute("braintrust.input", JSON_MAPPER.writeValueAsString(content)); - } catch (JsonProcessingException e) { - BraintrustLogger.error("Error mapping json", e); + private static final AttributeKey EVENT_NAME = stringKey("event.name"); + private static final ObjectMapper JSON_MAPPER = + new com.fasterxml.jackson.databind.ObjectMapper(); + + public static void emitPromptLogEvents( + Context context, + Logger eventLogger, + ChatCompletionCreateParams request, + boolean captureMessageContent) { + for (ChatCompletionMessageParam msg : request.messages()) { + String eventType; + Map> body = new HashMap<>(); + if (msg.isSystem()) { + eventType = "gen_ai.system.message"; + if (captureMessageContent) { + var content = contentToString(msg.asSystem().content()); + Span.current().setAttribute("instructions", content); + body.put("content", Value.of(content)); + } + } else if (msg.isDeveloper()) { + eventType = "gen_ai.system.message"; + body.put("role", Value.of("developer")); + if (captureMessageContent) { + var content = contentToString(msg.asDeveloper().content()); + body.put("content", Value.of(content)); + } + } else if (msg.isUser()) { + eventType = "gen_ai.user.message"; + if (captureMessageContent) { + var content = contentToString(msg.asUser().content()); + try { + Span.current() + .setAttribute( + "braintrust.input", + JSON_MAPPER.writeValueAsString(content)); + } catch (JsonProcessingException e) { + BraintrustLogger.error("Error mapping json", e); + } + body.put("content", Value.of(content)); + } + } else if (msg.isAssistant()) { + ChatCompletionAssistantMessageParam assistantMsg = msg.asAssistant(); + eventType = "gen_ai.assistant.message"; + if (captureMessageContent) { + assistantMsg + .content() + .ifPresent( + content -> + body.put( + "content", Value.of(contentToString(content)))); + } + assistantMsg + .toolCalls() + .ifPresent( + toolCalls -> { + List> toolCallsJson = + toolCalls.stream() + .map( + call -> + buildToolCallEventObject( + call, + captureMessageContent)) + .collect(Collectors.toList()); + body.put("tool_calls", Value.of(toolCallsJson)); + }); + } else if (msg.isTool()) { + ChatCompletionToolMessageParam toolMsg = msg.asTool(); + eventType = "gen_ai.tool.message"; + if (captureMessageContent) { + body.put("content", Value.of(contentToString(toolMsg.content()))); + } + body.put("id", Value.of(toolMsg.toolCallId())); + } else { + continue; } - body.put("content", Value.of(content)); - } - } else if (msg.isAssistant()) { - ChatCompletionAssistantMessageParam assistantMsg = msg.asAssistant(); - eventType = "gen_ai.assistant.message"; - if (captureMessageContent) { - assistantMsg - .content() - .ifPresent(content -> body.put("content", Value.of(contentToString(content)))); + // newEvent(eventLogger, eventType).setContext(context).setBody(Value.of(body)).emit(); } - assistantMsg - .toolCalls() - .ifPresent( - toolCalls -> { - List> toolCallsJson = - toolCalls.stream() - .map(call -> buildToolCallEventObject(call, captureMessageContent)) - .collect(Collectors.toList()); - body.put("tool_calls", Value.of(toolCallsJson)); - }); - } else if (msg.isTool()) { - ChatCompletionToolMessageParam toolMsg = msg.asTool(); - eventType = "gen_ai.tool.message"; - if (captureMessageContent) { - body.put("content", Value.of(contentToString(toolMsg.content()))); + // "gen_ai.input.messages" + } + + private static String contentToString(ChatCompletionToolMessageParam.Content content) { + if (content.isText()) { + return content.asText(); + } else if (content.isArrayOfContentParts()) { + return joinContentParts(content.asArrayOfContentParts()); + } else { + return ""; } - body.put("id", Value.of(toolMsg.toolCallId())); - } else { - continue; - } - // newEvent(eventLogger, eventType).setContext(context).setBody(Value.of(body)).emit(); } - // "gen_ai.input.messages" - } - - private static String contentToString(ChatCompletionToolMessageParam.Content content) { - if (content.isText()) { - return content.asText(); - } else if (content.isArrayOfContentParts()) { - return joinContentParts(content.asArrayOfContentParts()); - } else { - return ""; + + private static String contentToString(ChatCompletionAssistantMessageParam.Content content) { + if (content.isText()) { + return content.asText(); + } else if (content.isArrayOfContentParts()) { + return content.asArrayOfContentParts().stream() + .map( + part -> { + if (part.isText()) { + return part.asText().text(); + } + if (part.isRefusal()) { + return part.asRefusal().refusal(); + } + return null; + }) + .filter(Objects::nonNull) + .collect(Collectors.joining()); + } else { + return ""; + } } - } - - private static String contentToString(ChatCompletionAssistantMessageParam.Content content) { - if (content.isText()) { - return content.asText(); - } else if (content.isArrayOfContentParts()) { - return content.asArrayOfContentParts().stream() - .map( - part -> { - if (part.isText()) { - return part.asText().text(); - } - if (part.isRefusal()) { - return part.asRefusal().refusal(); - } - return null; - }) - .filter(Objects::nonNull) - .collect(Collectors.joining()); - } else { - return ""; + + private static String contentToString(ChatCompletionSystemMessageParam.Content content) { + if (content.isText()) { + return content.asText(); + } else if (content.isArrayOfContentParts()) { + return joinContentParts(content.asArrayOfContentParts()); + } else { + return ""; + } } - } - - private static String contentToString(ChatCompletionSystemMessageParam.Content content) { - if (content.isText()) { - return content.asText(); - } else if (content.isArrayOfContentParts()) { - return joinContentParts(content.asArrayOfContentParts()); - } else { - return ""; + + private static String contentToString(ChatCompletionDeveloperMessageParam.Content content) { + if (content.isText()) { + return content.asText(); + } else if (content.isArrayOfContentParts()) { + return joinContentParts(content.asArrayOfContentParts()); + } else { + return ""; + } } - } - - private static String contentToString(ChatCompletionDeveloperMessageParam.Content content) { - if (content.isText()) { - return content.asText(); - } else if (content.isArrayOfContentParts()) { - return joinContentParts(content.asArrayOfContentParts()); - } else { - return ""; + + private static String contentToString(ChatCompletionUserMessageParam.Content content) { + if (content.isText()) { + return content.asText(); + } else if (content.isArrayOfContentParts()) { + return content.asArrayOfContentParts().stream() + .map(part -> part.isText() ? part.asText().text() : null) + .filter(Objects::nonNull) + .collect(Collectors.joining()); + } else { + return ""; + } } - } - - private static String contentToString(ChatCompletionUserMessageParam.Content content) { - if (content.isText()) { - return content.asText(); - } else if (content.isArrayOfContentParts()) { - return content.asArrayOfContentParts().stream() - .map(part -> part.isText() ? part.asText().text() : null) - .filter(Objects::nonNull) - .collect(Collectors.joining()); - } else { - return ""; + + private static String joinContentParts(List contentParts) { + return contentParts.stream() + .map(ChatCompletionContentPartText::text) + .collect(Collectors.joining()); } - } - - private static String joinContentParts(List contentParts) { - return contentParts.stream() - .map(ChatCompletionContentPartText::text) - .collect(Collectors.joining()); - } - - public static void emitCompletionLogEvents( - Context context, - Logger eventLogger, - ChatCompletion completion, - boolean captureMessageContent) { - try { - var completionJson = JSON_MAPPER.writeValueAsString(completion); - Span.current().setAttribute("braintrust.output", completionJson); - } catch (JsonProcessingException e) { - BraintrustLogger.error("Error mapping completion json", e); - } - for (ChatCompletion.Choice choice : completion.choices()) { - ChatCompletionMessage choiceMsg = choice.message(); - Map> message = new HashMap<>(); - if (captureMessageContent) { - choiceMsg.content().ifPresent(content -> { - message.put("content", Value.of(content)); - }); - } - choiceMsg - .toolCalls() - .ifPresent( - toolCalls -> { - message.put( - "tool_calls", - Value.of( - toolCalls.stream() - .map(call -> buildToolCallEventObject(call, captureMessageContent)) - .collect(Collectors.toList()))); - }); - emitCompletionLogEvent( - context, - eventLogger, - choice.index(), - choice.finishReason().toString(), - Value.of(message)); + + public static void emitCompletionLogEvents( + Context context, + Logger eventLogger, + ChatCompletion completion, + boolean captureMessageContent) { + try { + var completionJson = JSON_MAPPER.writeValueAsString(completion); + Span.current().setAttribute("braintrust.output", completionJson); + } catch (JsonProcessingException e) { + BraintrustLogger.error("Error mapping completion json", e); + } + for (ChatCompletion.Choice choice : completion.choices()) { + ChatCompletionMessage choiceMsg = choice.message(); + Map> message = new HashMap<>(); + if (captureMessageContent) { + choiceMsg + .content() + .ifPresent( + content -> { + message.put("content", Value.of(content)); + }); + } + choiceMsg + .toolCalls() + .ifPresent( + toolCalls -> { + message.put( + "tool_calls", + Value.of( + toolCalls.stream() + .map( + call -> + buildToolCallEventObject( + call, + captureMessageContent)) + .collect(Collectors.toList()))); + }); + emitCompletionLogEvent( + context, + eventLogger, + choice.index(), + choice.finishReason().toString(), + Value.of(message)); + } } - } - - public static void emitCompletionLogEvent( - Context context, - Logger eventLogger, - long index, - String finishReason, - Value eventMessageObject) { - Map> body = new HashMap<>(); - body.put("finish_reason", Value.of(finishReason)); - body.put("index", Value.of(index)); - body.put("message", eventMessageObject); - // newEvent(eventLogger, "gen_ai.choice").setContext(context).setBody(Value.of(body)).emit(); - } - - private static LogRecordBuilder newEvent(Logger eventLogger, String name) { - // NOTE: disabling logger events in braintrust instrumentation. We don't use these events. - // Will have to properly hanlde this if we want to merge braintrust attributes upstream into otel instrumentation - /* - return eventLogger - .logRecordBuilder() - .setAttribute(EVENT_NAME, name) - .setAttribute(GEN_AI_PROVIDER_NAME, "openai"); - */ - throw new RuntimeException("Should not invoke"); - } - - private static Value buildToolCallEventObject( - ChatCompletionMessageToolCall call, boolean captureMessageContent) { - Map> result = new HashMap<>(); - FunctionAccess functionAccess = getFunctionAccess(call); - if (functionAccess != null) { - result.put("id", Value.of(functionAccess.id())); - result.put("type", Value.of("function")); // "function" is the only currently supported type - result.put("function", buildFunctionEventObject(functionAccess, captureMessageContent)); + + public static void emitCompletionLogEvent( + Context context, + Logger eventLogger, + long index, + String finishReason, + Value eventMessageObject) { + Map> body = new HashMap<>(); + body.put("finish_reason", Value.of(finishReason)); + body.put("index", Value.of(index)); + body.put("message", eventMessageObject); + // newEvent(eventLogger, + // "gen_ai.choice").setContext(context).setBody(Value.of(body)).emit(); } - return Value.of(result); - } - - private static Value buildFunctionEventObject( - FunctionAccess functionAccess, boolean captureMessageContent) { - Map> result = new HashMap<>(); - result.put("name", Value.of(functionAccess.name())); - if (captureMessageContent) { - result.put("arguments", Value.of(functionAccess.arguments())); + + private static LogRecordBuilder newEvent(Logger eventLogger, String name) { + // NOTE: disabling logger events in braintrust instrumentation. We don't use these events. + // Will have to properly hanlde this if we want to merge braintrust attributes upstream into + // otel instrumentation + /* + return eventLogger + .logRecordBuilder() + .setAttribute(EVENT_NAME, name) + .setAttribute(GEN_AI_PROVIDER_NAME, "openai"); + */ + throw new RuntimeException("Should not invoke"); } - return Value.of(result); - } - @Nullable - private static FunctionAccess getFunctionAccess(ChatCompletionMessageToolCall call) { - if (V1FunctionAccess.isAvailable()) { - return V1FunctionAccess.create(call); + private static Value buildToolCallEventObject( + ChatCompletionMessageToolCall call, boolean captureMessageContent) { + Map> result = new HashMap<>(); + FunctionAccess functionAccess = getFunctionAccess(call); + if (functionAccess != null) { + result.put("id", Value.of(functionAccess.id())); + result.put( + "type", + Value.of("function")); // "function" is the only currently supported type + result.put("function", buildFunctionEventObject(functionAccess, captureMessageContent)); + } + return Value.of(result); } - if (V3FunctionAccess.isAvailable()) { - return V3FunctionAccess.create(call); + + private static Value buildFunctionEventObject( + FunctionAccess functionAccess, boolean captureMessageContent) { + Map> result = new HashMap<>(); + result.put("name", Value.of(functionAccess.name())); + if (captureMessageContent) { + result.put("arguments", Value.of(functionAccess.arguments())); + } + return Value.of(result); } - return null; - } + @Nullable + private static FunctionAccess getFunctionAccess(ChatCompletionMessageToolCall call) { + if (V1FunctionAccess.isAvailable()) { + return V1FunctionAccess.create(call); + } + if (V3FunctionAccess.isAvailable()) { + return V3FunctionAccess.create(call); + } - private interface FunctionAccess { - String id(); + return null; + } - String name(); + private interface FunctionAccess { + String id(); - String arguments(); - } + String name(); - private static String invokeStringHandle(@Nullable MethodHandle methodHandle, Object object) { - if (methodHandle == null) { - return ""; + String arguments(); } - try { - return (String) methodHandle.invoke(object); - } catch (Throwable ignore) { - return ""; - } - } - - private static class V1FunctionAccess implements FunctionAccess { - @Nullable private static final MethodHandle idHandle; - @Nullable private static final MethodHandle functionHandle; - @Nullable private static final MethodHandle nameHandle; - @Nullable private static final MethodHandle argumentsHandle; - - static { - MethodHandle id; - MethodHandle function; - MethodHandle name; - MethodHandle arguments; - - try { - MethodHandles.Lookup lookup = MethodHandles.lookup(); - id = - lookup.findVirtual( - ChatCompletionMessageToolCall.class, "id", MethodType.methodType(String.class)); - Class functionClass = - Class.forName( - "com.openai.models.chat.completions.ChatCompletionMessageToolCall$Function"); - function = - lookup.findVirtual( - ChatCompletionMessageToolCall.class, - "function", - MethodType.methodType(functionClass)); - name = lookup.findVirtual(functionClass, "name", MethodType.methodType(String.class)); - arguments = - lookup.findVirtual(functionClass, "arguments", MethodType.methodType(String.class)); - } catch (Exception exception) { - id = null; - function = null; - name = null; - arguments = null; - } - idHandle = id; - functionHandle = function; - nameHandle = name; - argumentsHandle = arguments; + private static String invokeStringHandle(@Nullable MethodHandle methodHandle, Object object) { + if (methodHandle == null) { + return ""; + } + + try { + return (String) methodHandle.invoke(object); + } catch (Throwable ignore) { + return ""; + } } - private final ChatCompletionMessageToolCall toolCall; - private final Object function; + private static class V1FunctionAccess implements FunctionAccess { + @Nullable private static final MethodHandle idHandle; + @Nullable private static final MethodHandle functionHandle; + @Nullable private static final MethodHandle nameHandle; + @Nullable private static final MethodHandle argumentsHandle; - V1FunctionAccess(ChatCompletionMessageToolCall toolCall, Object function) { - this.toolCall = toolCall; - this.function = function; - } + static { + MethodHandle id; + MethodHandle function; + MethodHandle name; + MethodHandle arguments; - @Nullable - static FunctionAccess create(ChatCompletionMessageToolCall toolCall) { - if (functionHandle == null) { - return null; - } + try { + MethodHandles.Lookup lookup = MethodHandles.lookup(); + id = + lookup.findVirtual( + ChatCompletionMessageToolCall.class, + "id", + MethodType.methodType(String.class)); + Class functionClass = + Class.forName( + "com.openai.models.chat.completions.ChatCompletionMessageToolCall$Function"); + function = + lookup.findVirtual( + ChatCompletionMessageToolCall.class, + "function", + MethodType.methodType(functionClass)); + name = + lookup.findVirtual( + functionClass, "name", MethodType.methodType(String.class)); + arguments = + lookup.findVirtual( + functionClass, "arguments", MethodType.methodType(String.class)); + } catch (Exception exception) { + id = null; + function = null; + name = null; + arguments = null; + } + idHandle = id; + functionHandle = function; + nameHandle = name; + argumentsHandle = arguments; + } - try { - return new V1FunctionAccess(toolCall, functionHandle.invoke(toolCall)); - } catch (Throwable ignore) { - return null; - } - } + private final ChatCompletionMessageToolCall toolCall; + private final Object function; - static boolean isAvailable() { - return idHandle != null; - } + V1FunctionAccess(ChatCompletionMessageToolCall toolCall, Object function) { + this.toolCall = toolCall; + this.function = function; + } - @Override - public String id() { - return invokeStringHandle(idHandle, toolCall); - } + @Nullable + static FunctionAccess create(ChatCompletionMessageToolCall toolCall) { + if (functionHandle == null) { + return null; + } - @Override - public String name() { - return invokeStringHandle(nameHandle, function); - } + try { + return new V1FunctionAccess(toolCall, functionHandle.invoke(toolCall)); + } catch (Throwable ignore) { + return null; + } + } - @Override - public String arguments() { - return invokeStringHandle(argumentsHandle, function); - } - } - - static class V3FunctionAccess implements FunctionAccess { - @Nullable private static final MethodHandle functionToolCallHandle; - @Nullable private static final MethodHandle idHandle; - @Nullable private static final MethodHandle functionHandle; - @Nullable private static final MethodHandle nameHandle; - @Nullable private static final MethodHandle argumentsHandle; - - static { - MethodHandle functionToolCall; - MethodHandle id; - MethodHandle function; - MethodHandle name; - MethodHandle arguments; - - try { - MethodHandles.Lookup lookup = MethodHandles.lookup(); - functionToolCall = - lookup.findVirtual( - ChatCompletionMessageToolCall.class, - "function", - MethodType.methodType(Optional.class)); - Class functionToolCallClass = - Class.forName( - "com.openai.models.chat.completions.ChatCompletionMessageFunctionToolCall"); - id = lookup.findVirtual(functionToolCallClass, "id", MethodType.methodType(String.class)); - Class functionClass = - Class.forName( - "com.openai.models.chat.completions.ChatCompletionMessageFunctionToolCall$Function"); - function = - lookup.findVirtual( - functionToolCallClass, "function", MethodType.methodType(functionClass)); - name = lookup.findVirtual(functionClass, "name", MethodType.methodType(String.class)); - arguments = - lookup.findVirtual(functionClass, "arguments", MethodType.methodType(String.class)); - } catch (Exception exception) { - functionToolCall = null; - id = null; - function = null; - name = null; - arguments = null; - } - functionToolCallHandle = functionToolCall; - idHandle = id; - functionHandle = function; - nameHandle = name; - argumentsHandle = arguments; - } + static boolean isAvailable() { + return idHandle != null; + } - private final Object functionToolCall; - private final Object function; + @Override + public String id() { + return invokeStringHandle(idHandle, toolCall); + } + + @Override + public String name() { + return invokeStringHandle(nameHandle, function); + } - V3FunctionAccess(Object functionToolCall, Object function) { - this.functionToolCall = functionToolCall; - this.function = function; + @Override + public String arguments() { + return invokeStringHandle(argumentsHandle, function); + } } - @Nullable - @SuppressWarnings("unchecked") - static FunctionAccess create(ChatCompletionMessageToolCall toolCall) { - if (functionToolCallHandle == null || functionHandle == null) { - return null; - } + static class V3FunctionAccess implements FunctionAccess { + @Nullable private static final MethodHandle functionToolCallHandle; + @Nullable private static final MethodHandle idHandle; + @Nullable private static final MethodHandle functionHandle; + @Nullable private static final MethodHandle nameHandle; + @Nullable private static final MethodHandle argumentsHandle; - try { - Optional optional = (Optional) functionToolCallHandle.invoke(toolCall); - if (!optional.isPresent()) { - return null; + static { + MethodHandle functionToolCall; + MethodHandle id; + MethodHandle function; + MethodHandle name; + MethodHandle arguments; + + try { + MethodHandles.Lookup lookup = MethodHandles.lookup(); + functionToolCall = + lookup.findVirtual( + ChatCompletionMessageToolCall.class, + "function", + MethodType.methodType(Optional.class)); + Class functionToolCallClass = + Class.forName( + "com.openai.models.chat.completions.ChatCompletionMessageFunctionToolCall"); + id = + lookup.findVirtual( + functionToolCallClass, "id", MethodType.methodType(String.class)); + Class functionClass = + Class.forName( + "com.openai.models.chat.completions.ChatCompletionMessageFunctionToolCall$Function"); + function = + lookup.findVirtual( + functionToolCallClass, + "function", + MethodType.methodType(functionClass)); + name = + lookup.findVirtual( + functionClass, "name", MethodType.methodType(String.class)); + arguments = + lookup.findVirtual( + functionClass, "arguments", MethodType.methodType(String.class)); + } catch (Exception exception) { + functionToolCall = null; + id = null; + function = null; + name = null; + arguments = null; + } + functionToolCallHandle = functionToolCall; + idHandle = id; + functionHandle = function; + nameHandle = name; + argumentsHandle = arguments; } - Object functionToolCall = optional.get(); - return new V3FunctionAccess(functionToolCall, functionHandle.invoke(functionToolCall)); - } catch (Throwable ignore) { - return null; - } - } - static boolean isAvailable() { - return idHandle != null; - } + private final Object functionToolCall; + private final Object function; - @Override - public String id() { - return invokeStringHandle(idHandle, functionToolCall); - } + V3FunctionAccess(Object functionToolCall, Object function) { + this.functionToolCall = functionToolCall; + this.function = function; + } - @Override - public String name() { - return invokeStringHandle(nameHandle, function); - } + @Nullable + @SuppressWarnings("unchecked") + static FunctionAccess create(ChatCompletionMessageToolCall toolCall) { + if (functionToolCallHandle == null || functionHandle == null) { + return null; + } + + try { + Optional optional = + (Optional) functionToolCallHandle.invoke(toolCall); + if (!optional.isPresent()) { + return null; + } + Object functionToolCall = optional.get(); + return new V3FunctionAccess( + functionToolCall, functionHandle.invoke(functionToolCall)); + } catch (Throwable ignore) { + return null; + } + } + + static boolean isAvailable() { + return idHandle != null; + } + + @Override + public String id() { + return invokeStringHandle(idHandle, functionToolCall); + } - @Override - public String arguments() { - return invokeStringHandle(argumentsHandle, function); + @Override + public String name() { + return invokeStringHandle(nameHandle, function); + } + + @Override + public String arguments() { + return invokeStringHandle(argumentsHandle, function); + } } - } - private ChatCompletionEventsHelper() {} + private ChatCompletionEventsHelper() {} } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/CompletableFutureWrapper.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/CompletableFutureWrapper.java index 229f283..9dcea97 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/CompletableFutureWrapper.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/CompletableFutureWrapper.java @@ -10,21 +10,21 @@ import java.util.concurrent.CompletableFuture; final class CompletableFutureWrapper { - private CompletableFutureWrapper() {} + private CompletableFutureWrapper() {} - static CompletableFuture wrap(CompletableFuture future, Context context) { - CompletableFuture result = new CompletableFuture<>(); - future.whenComplete( - (T value, Throwable throwable) -> { - try (Scope ignored = context.makeCurrent()) { - if (throwable != null) { - result.completeExceptionally(throwable); - } else { - result.complete(value); - } - } - }); + static CompletableFuture wrap(CompletableFuture future, Context context) { + CompletableFuture result = new CompletableFuture<>(); + future.whenComplete( + (T value, Throwable throwable) -> { + try (Scope ignored = context.makeCurrent()) { + if (throwable != null) { + result.completeExceptionally(throwable); + } else { + result.complete(value); + } + } + }); - return result; - } + return result; + } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/DelegatingInvocationHandler.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/DelegatingInvocationHandler.java index c5e3bd8..aba47c8 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/DelegatingInvocationHandler.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/DelegatingInvocationHandler.java @@ -11,32 +11,32 @@ import java.lang.reflect.Proxy; abstract class DelegatingInvocationHandler> - implements InvocationHandler { + implements InvocationHandler { - private static final ClassLoader CLASS_LOADER = - DelegatingInvocationHandler.class.getClassLoader(); + private static final ClassLoader CLASS_LOADER = + DelegatingInvocationHandler.class.getClassLoader(); - protected final T delegate; + protected final T delegate; - public DelegatingInvocationHandler(T delegate) { - this.delegate = delegate; - } + public DelegatingInvocationHandler(T delegate) { + this.delegate = delegate; + } - @Override - public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { - try { - return method.invoke(delegate, args); - } catch (InvocationTargetException e) { - throw e.getCause(); + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { + try { + return method.invoke(delegate, args); + } catch (InvocationTargetException e) { + throw e.getCause(); + } } - } - protected abstract Class getProxyType(); + protected abstract Class getProxyType(); - @SuppressWarnings("rawtypes") - public T createProxy() { - Class proxyType = getProxyType(); - Object proxy = Proxy.newProxyInstance(CLASS_LOADER, new Class[] {proxyType}, this); - return proxyType.cast(proxy); - } + @SuppressWarnings("rawtypes") + public T createProxy() { + Class proxyType = getProxyType(); + Object proxy = Proxy.newProxyInstance(CLASS_LOADER, new Class[] {proxyType}, this); + return proxyType.cast(proxy); + } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/EmbeddingAttributesGetter.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/EmbeddingAttributesGetter.java index 34a3b59..220e23f 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/EmbeddingAttributesGetter.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/EmbeddingAttributesGetter.java @@ -15,115 +15,115 @@ import javax.annotation.Nullable; enum EmbeddingAttributesGetter - implements GenAiAttributesGetter { - INSTANCE; - - @Override - public String getOperationName(EmbeddingCreateParams request) { - return GenAiAttributes.GenAiOperationNameIncubatingValues.EMBEDDINGS; - } - - @Override - public String getSystem(EmbeddingCreateParams request) { - return GenAiAttributes.GenAiProviderNameIncubatingValues.OPENAI; - } - - @Override - public String getRequestModel(EmbeddingCreateParams request) { - return request.model().asString(); - } - - @Nullable - @Override - public Long getRequestSeed(EmbeddingCreateParams request) { - return null; - } - - @Nullable - @Override - public List getRequestEncodingFormats(EmbeddingCreateParams request) { - return request.encodingFormat().map(f -> singletonList(f.asString())).orElse(null); - } - - @Nullable - @Override - public Double getRequestFrequencyPenalty(EmbeddingCreateParams request) { - return null; - } - - @Nullable - @Override - public Long getRequestMaxTokens(EmbeddingCreateParams request) { - return null; - } - - @Nullable - @Override - public Double getRequestPresencePenalty(EmbeddingCreateParams request) { - return null; - } - - @Nullable - @Override - public List getRequestStopSequences(EmbeddingCreateParams request) { - return null; - } - - @Nullable - @Override - public Double getRequestTemperature(EmbeddingCreateParams request) { - return null; - } - - @Nullable - @Override - public Double getRequestTopK(EmbeddingCreateParams request) { - return null; - } - - @Nullable - @Override - public Double getRequestTopP(EmbeddingCreateParams request) { - return null; - } - - @Override - public List getResponseFinishReasons( - EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) { - return Collections.emptyList(); - } - - @Nullable - @Override - public String getResponseId( - EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) { - return null; - } - - @Nullable - @Override - public String getResponseModel( - EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) { - if (response == null) { - return null; - } - return response.model(); - } - - @Nullable - @Override - public Long getUsageInputTokens( - EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) { - if (response == null) { - return null; - } - return response.usage().promptTokens(); - } - - @Nullable - @Override - public Long getUsageOutputTokens( - EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) { - return null; - } + implements GenAiAttributesGetter { + INSTANCE; + + @Override + public String getOperationName(EmbeddingCreateParams request) { + return GenAiAttributes.GenAiOperationNameIncubatingValues.EMBEDDINGS; + } + + @Override + public String getSystem(EmbeddingCreateParams request) { + return GenAiAttributes.GenAiProviderNameIncubatingValues.OPENAI; + } + + @Override + public String getRequestModel(EmbeddingCreateParams request) { + return request.model().asString(); + } + + @Nullable + @Override + public Long getRequestSeed(EmbeddingCreateParams request) { + return null; + } + + @Nullable + @Override + public List getRequestEncodingFormats(EmbeddingCreateParams request) { + return request.encodingFormat().map(f -> singletonList(f.asString())).orElse(null); + } + + @Nullable + @Override + public Double getRequestFrequencyPenalty(EmbeddingCreateParams request) { + return null; + } + + @Nullable + @Override + public Long getRequestMaxTokens(EmbeddingCreateParams request) { + return null; + } + + @Nullable + @Override + public Double getRequestPresencePenalty(EmbeddingCreateParams request) { + return null; + } + + @Nullable + @Override + public List getRequestStopSequences(EmbeddingCreateParams request) { + return null; + } + + @Nullable + @Override + public Double getRequestTemperature(EmbeddingCreateParams request) { + return null; + } + + @Nullable + @Override + public Double getRequestTopK(EmbeddingCreateParams request) { + return null; + } + + @Nullable + @Override + public Double getRequestTopP(EmbeddingCreateParams request) { + return null; + } + + @Override + public List getResponseFinishReasons( + EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) { + return Collections.emptyList(); + } + + @Nullable + @Override + public String getResponseId( + EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) { + return null; + } + + @Nullable + @Override + public String getResponseModel( + EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) { + if (response == null) { + return null; + } + return response.model(); + } + + @Nullable + @Override + public Long getUsageInputTokens( + EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) { + if (response == null) { + return null; + } + return response.usage().promptTokens(); + } + + @Nullable + @Override + public Long getUsageOutputTokens( + EmbeddingCreateParams request, @Nullable CreateEmbeddingResponse response) { + return null; + } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/GenAiAttributes.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/GenAiAttributes.java index 76acae2..690aca5 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/GenAiAttributes.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/GenAiAttributes.java @@ -11,20 +11,20 @@ // copied from GenAiIncubatingAttributes final class GenAiAttributes { - static final AttributeKey GEN_AI_PROVIDER_NAME = stringKey("gen_ai.provider.name"); + static final AttributeKey GEN_AI_PROVIDER_NAME = stringKey("gen_ai.provider.name"); - static final class GenAiOperationNameIncubatingValues { - static final String CHAT = "chat"; - static final String EMBEDDINGS = "embeddings"; + static final class GenAiOperationNameIncubatingValues { + static final String CHAT = "chat"; + static final String EMBEDDINGS = "embeddings"; - private GenAiOperationNameIncubatingValues() {} - } + private GenAiOperationNameIncubatingValues() {} + } - static final class GenAiProviderNameIncubatingValues { - static final String OPENAI = "openai"; + static final class GenAiProviderNameIncubatingValues { + static final String OPENAI = "openai"; - private GenAiProviderNameIncubatingValues() {} - } + private GenAiProviderNameIncubatingValues() {} + } - private GenAiAttributes() {} + private GenAiAttributes() {} } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatCompletionService.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatCompletionService.java index ce6575a..b80d1df 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatCompletionService.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatCompletionService.java @@ -18,125 +18,134 @@ import java.lang.reflect.Method; final class InstrumentedChatCompletionService - extends DelegatingInvocationHandler { - - private final Instrumenter instrumenter; - private final Logger eventLogger; - private final boolean captureMessageContent; - - InstrumentedChatCompletionService( - ChatCompletionService delegate, - Instrumenter instrumenter, - Logger eventLogger, - boolean captureMessageContent) { - super(delegate); - this.instrumenter = instrumenter; - this.eventLogger = eventLogger; - this.captureMessageContent = captureMessageContent; - } - - @Override - protected Class getProxyType() { - return ChatCompletionService.class; - } - - @Override - public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { - String methodName = method.getName(); - Class[] parameterTypes = method.getParameterTypes(); - - switch (methodName) { - case "create": - if (parameterTypes.length >= 1 && parameterTypes[0] == ChatCompletionCreateParams.class) { - if (parameterTypes.length == 1) { - return create((ChatCompletionCreateParams) args[0], RequestOptions.none()); - } else if (parameterTypes.length == 2 && parameterTypes[1] == RequestOptions.class) { - return create((ChatCompletionCreateParams) args[0], (RequestOptions) args[1]); - } - } - break; - case "createStreaming": - if (parameterTypes.length >= 1 && parameterTypes[0] == ChatCompletionCreateParams.class) { - if (parameterTypes.length == 1) { - return createStreaming((ChatCompletionCreateParams) args[0], RequestOptions.none()); - } else if (parameterTypes.length == 2 && parameterTypes[1] == RequestOptions.class) { - return createStreaming((ChatCompletionCreateParams) args[0], (RequestOptions) args[1]); - } + extends DelegatingInvocationHandler< + ChatCompletionService, InstrumentedChatCompletionService> { + + private final Instrumenter instrumenter; + private final Logger eventLogger; + private final boolean captureMessageContent; + + InstrumentedChatCompletionService( + ChatCompletionService delegate, + Instrumenter instrumenter, + Logger eventLogger, + boolean captureMessageContent) { + super(delegate); + this.instrumenter = instrumenter; + this.eventLogger = eventLogger; + this.captureMessageContent = captureMessageContent; + } + + @Override + protected Class getProxyType() { + return ChatCompletionService.class; + } + + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { + String methodName = method.getName(); + Class[] parameterTypes = method.getParameterTypes(); + + switch (methodName) { + case "create": + if (parameterTypes.length >= 1 + && parameterTypes[0] == ChatCompletionCreateParams.class) { + if (parameterTypes.length == 1) { + return create((ChatCompletionCreateParams) args[0], RequestOptions.none()); + } else if (parameterTypes.length == 2 + && parameterTypes[1] == RequestOptions.class) { + return create( + (ChatCompletionCreateParams) args[0], (RequestOptions) args[1]); + } + } + break; + case "createStreaming": + if (parameterTypes.length >= 1 + && parameterTypes[0] == ChatCompletionCreateParams.class) { + if (parameterTypes.length == 1) { + return createStreaming( + (ChatCompletionCreateParams) args[0], RequestOptions.none()); + } else if (parameterTypes.length == 2 + && parameterTypes[1] == RequestOptions.class) { + return createStreaming( + (ChatCompletionCreateParams) args[0], (RequestOptions) args[1]); + } + } + break; + default: + // fallthrough } - break; - default: - // fallthrough + + return super.invoke(proxy, method, args); } - return super.invoke(proxy, method, args); - } + private ChatCompletion create( + ChatCompletionCreateParams chatCompletionCreateParams, RequestOptions requestOptions) { + Context parentContext = Context.current(); + if (!instrumenter.shouldStart(parentContext, chatCompletionCreateParams)) { + return createWithLogs(parentContext, chatCompletionCreateParams, requestOptions); + } + + Context context = instrumenter.start(parentContext, chatCompletionCreateParams); + ChatCompletion completion; + try (Scope ignored = context.makeCurrent()) { + completion = createWithLogs(context, chatCompletionCreateParams, requestOptions); + } catch (Throwable t) { + instrumenter.end(context, chatCompletionCreateParams, null, t); + throw t; + } - private ChatCompletion create( - ChatCompletionCreateParams chatCompletionCreateParams, RequestOptions requestOptions) { - Context parentContext = Context.current(); - if (!instrumenter.shouldStart(parentContext, chatCompletionCreateParams)) { - return createWithLogs(parentContext, chatCompletionCreateParams, requestOptions); + instrumenter.end(context, chatCompletionCreateParams, completion, null); + return completion; } - Context context = instrumenter.start(parentContext, chatCompletionCreateParams); - ChatCompletion completion; - try (Scope ignored = context.makeCurrent()) { - completion = createWithLogs(context, chatCompletionCreateParams, requestOptions); - } catch (Throwable t) { - instrumenter.end(context, chatCompletionCreateParams, null, t); - throw t; + private ChatCompletion createWithLogs( + Context context, + ChatCompletionCreateParams chatCompletionCreateParams, + RequestOptions requestOptions) { + ChatCompletionEventsHelper.emitPromptLogEvents( + context, eventLogger, chatCompletionCreateParams, captureMessageContent); + ChatCompletion result = delegate.create(chatCompletionCreateParams, requestOptions); + ChatCompletionEventsHelper.emitCompletionLogEvents( + context, eventLogger, result, captureMessageContent); + return result; } - instrumenter.end(context, chatCompletionCreateParams, completion, null); - return completion; - } - - private ChatCompletion createWithLogs( - Context context, - ChatCompletionCreateParams chatCompletionCreateParams, - RequestOptions requestOptions) { - ChatCompletionEventsHelper.emitPromptLogEvents( - context, eventLogger, chatCompletionCreateParams, captureMessageContent); - ChatCompletion result = delegate.create(chatCompletionCreateParams, requestOptions); - ChatCompletionEventsHelper.emitCompletionLogEvents( - context, eventLogger, result, captureMessageContent); - return result; - } - - private StreamResponse createStreaming( - ChatCompletionCreateParams chatCompletionCreateParams, RequestOptions requestOptions) { - Context parentContext = Context.current(); - if (!instrumenter.shouldStart(parentContext, chatCompletionCreateParams)) { - return createStreamingWithLogs( - parentContext, chatCompletionCreateParams, requestOptions, false); + private StreamResponse createStreaming( + ChatCompletionCreateParams chatCompletionCreateParams, RequestOptions requestOptions) { + Context parentContext = Context.current(); + if (!instrumenter.shouldStart(parentContext, chatCompletionCreateParams)) { + return createStreamingWithLogs( + parentContext, chatCompletionCreateParams, requestOptions, false); + } + + Context context = instrumenter.start(parentContext, chatCompletionCreateParams); + try (Scope ignored = context.makeCurrent()) { + return createStreamingWithLogs( + context, chatCompletionCreateParams, requestOptions, true); + } catch (Throwable t) { + instrumenter.end(context, chatCompletionCreateParams, null, t); + throw t; + } } - Context context = instrumenter.start(parentContext, chatCompletionCreateParams); - try (Scope ignored = context.makeCurrent()) { - return createStreamingWithLogs(context, chatCompletionCreateParams, requestOptions, true); - } catch (Throwable t) { - instrumenter.end(context, chatCompletionCreateParams, null, t); - throw t; + private StreamResponse createStreamingWithLogs( + Context context, + ChatCompletionCreateParams chatCompletionCreateParams, + RequestOptions requestOptions, + boolean newSpan) { + ChatCompletionEventsHelper.emitPromptLogEvents( + context, eventLogger, chatCompletionCreateParams, captureMessageContent); + StreamResponse result = + delegate.createStreaming(chatCompletionCreateParams, requestOptions); + return new TracingStreamedResponse( + result, + new StreamListener( + context, + chatCompletionCreateParams, + instrumenter, + eventLogger, + captureMessageContent, + newSpan)); } - } - - private StreamResponse createStreamingWithLogs( - Context context, - ChatCompletionCreateParams chatCompletionCreateParams, - RequestOptions requestOptions, - boolean newSpan) { - ChatCompletionEventsHelper.emitPromptLogEvents( - context, eventLogger, chatCompletionCreateParams, captureMessageContent); - StreamResponse result = - delegate.createStreaming(chatCompletionCreateParams, requestOptions); - return new TracingStreamedResponse( - result, - new StreamListener( - context, - chatCompletionCreateParams, - instrumenter, - eventLogger, - captureMessageContent, - newSpan)); - } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatCompletionServiceAsync.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatCompletionServiceAsync.java index 68aec0e..420dec7 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatCompletionServiceAsync.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatCompletionServiceAsync.java @@ -19,131 +19,139 @@ import java.util.concurrent.CompletableFuture; final class InstrumentedChatCompletionServiceAsync - extends DelegatingInvocationHandler< - ChatCompletionServiceAsync, InstrumentedChatCompletionServiceAsync> { - - private final Instrumenter instrumenter; - private final Logger eventLogger; - private final boolean captureMessageContent; - - InstrumentedChatCompletionServiceAsync( - ChatCompletionServiceAsync delegate, - Instrumenter instrumenter, - Logger eventLogger, - boolean captureMessageContent) { - super(delegate); - this.instrumenter = instrumenter; - this.eventLogger = eventLogger; - this.captureMessageContent = captureMessageContent; - } - - @Override - protected Class getProxyType() { - return ChatCompletionServiceAsync.class; - } - - @Override - public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { - String methodName = method.getName(); - Class[] parameterTypes = method.getParameterTypes(); - - switch (methodName) { - case "create": - if (parameterTypes.length >= 1 && parameterTypes[0] == ChatCompletionCreateParams.class) { - if (parameterTypes.length == 1) { - return create((ChatCompletionCreateParams) args[0], RequestOptions.none()); - } else if (parameterTypes.length == 2 && parameterTypes[1] == RequestOptions.class) { - return create((ChatCompletionCreateParams) args[0], (RequestOptions) args[1]); - } - } - break; - case "createStreaming": - if (parameterTypes.length >= 1 && parameterTypes[0] == ChatCompletionCreateParams.class) { - if (parameterTypes.length == 1) { - return createStreaming((ChatCompletionCreateParams) args[0], RequestOptions.none()); - } else if (parameterTypes.length == 2 && parameterTypes[1] == RequestOptions.class) { - return createStreaming((ChatCompletionCreateParams) args[0], (RequestOptions) args[1]); - } + extends DelegatingInvocationHandler< + ChatCompletionServiceAsync, InstrumentedChatCompletionServiceAsync> { + + private final Instrumenter instrumenter; + private final Logger eventLogger; + private final boolean captureMessageContent; + + InstrumentedChatCompletionServiceAsync( + ChatCompletionServiceAsync delegate, + Instrumenter instrumenter, + Logger eventLogger, + boolean captureMessageContent) { + super(delegate); + this.instrumenter = instrumenter; + this.eventLogger = eventLogger; + this.captureMessageContent = captureMessageContent; + } + + @Override + protected Class getProxyType() { + return ChatCompletionServiceAsync.class; + } + + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { + String methodName = method.getName(); + Class[] parameterTypes = method.getParameterTypes(); + + switch (methodName) { + case "create": + if (parameterTypes.length >= 1 + && parameterTypes[0] == ChatCompletionCreateParams.class) { + if (parameterTypes.length == 1) { + return create((ChatCompletionCreateParams) args[0], RequestOptions.none()); + } else if (parameterTypes.length == 2 + && parameterTypes[1] == RequestOptions.class) { + return create( + (ChatCompletionCreateParams) args[0], (RequestOptions) args[1]); + } + } + break; + case "createStreaming": + if (parameterTypes.length >= 1 + && parameterTypes[0] == ChatCompletionCreateParams.class) { + if (parameterTypes.length == 1) { + return createStreaming( + (ChatCompletionCreateParams) args[0], RequestOptions.none()); + } else if (parameterTypes.length == 2 + && parameterTypes[1] == RequestOptions.class) { + return createStreaming( + (ChatCompletionCreateParams) args[0], (RequestOptions) args[1]); + } + } + break; + default: + // fallthrough } - break; - default: - // fallthrough + + return super.invoke(proxy, method, args); } - return super.invoke(proxy, method, args); - } + private CompletableFuture create( + ChatCompletionCreateParams chatCompletionCreateParams, RequestOptions requestOptions) { + Context parentContext = Context.current(); + if (!instrumenter.shouldStart(parentContext, chatCompletionCreateParams)) { + return createWithLogs(parentContext, chatCompletionCreateParams, requestOptions); + } + + Context context = instrumenter.start(parentContext, chatCompletionCreateParams); + CompletableFuture future; + try (Scope ignored = context.makeCurrent()) { + future = createWithLogs(context, chatCompletionCreateParams, requestOptions); + } catch (Throwable t) { + instrumenter.end(context, chatCompletionCreateParams, null, t); + throw t; + } - private CompletableFuture create( - ChatCompletionCreateParams chatCompletionCreateParams, RequestOptions requestOptions) { - Context parentContext = Context.current(); - if (!instrumenter.shouldStart(parentContext, chatCompletionCreateParams)) { - return createWithLogs(parentContext, chatCompletionCreateParams, requestOptions); + future = + future.whenComplete( + (res, t) -> instrumenter.end(context, chatCompletionCreateParams, res, t)); + return CompletableFutureWrapper.wrap(future, parentContext); } - Context context = instrumenter.start(parentContext, chatCompletionCreateParams); - CompletableFuture future; - try (Scope ignored = context.makeCurrent()) { - future = createWithLogs(context, chatCompletionCreateParams, requestOptions); - } catch (Throwable t) { - instrumenter.end(context, chatCompletionCreateParams, null, t); - throw t; + private CompletableFuture createWithLogs( + Context context, + ChatCompletionCreateParams chatCompletionCreateParams, + RequestOptions requestOptions) { + ChatCompletionEventsHelper.emitPromptLogEvents( + context, eventLogger, chatCompletionCreateParams, captureMessageContent); + CompletableFuture future = + delegate.create(chatCompletionCreateParams, requestOptions); + future.thenAccept( + r -> + ChatCompletionEventsHelper.emitCompletionLogEvents( + context, eventLogger, r, captureMessageContent)); + return future; } - future = - future.whenComplete( - (res, t) -> instrumenter.end(context, chatCompletionCreateParams, res, t)); - return CompletableFutureWrapper.wrap(future, parentContext); - } - - private CompletableFuture createWithLogs( - Context context, - ChatCompletionCreateParams chatCompletionCreateParams, - RequestOptions requestOptions) { - ChatCompletionEventsHelper.emitPromptLogEvents( - context, eventLogger, chatCompletionCreateParams, captureMessageContent); - CompletableFuture future = - delegate.create(chatCompletionCreateParams, requestOptions); - future.thenAccept( - r -> - ChatCompletionEventsHelper.emitCompletionLogEvents( - context, eventLogger, r, captureMessageContent)); - return future; - } - - private AsyncStreamResponse createStreaming( - ChatCompletionCreateParams chatCompletionCreateParams, RequestOptions requestOptions) { - Context parentContext = Context.current(); - if (!instrumenter.shouldStart(parentContext, chatCompletionCreateParams)) { - return createStreamingWithLogs( - parentContext, chatCompletionCreateParams, requestOptions, false); + private AsyncStreamResponse createStreaming( + ChatCompletionCreateParams chatCompletionCreateParams, RequestOptions requestOptions) { + Context parentContext = Context.current(); + if (!instrumenter.shouldStart(parentContext, chatCompletionCreateParams)) { + return createStreamingWithLogs( + parentContext, chatCompletionCreateParams, requestOptions, false); + } + + Context context = instrumenter.start(parentContext, chatCompletionCreateParams); + try (Scope ignored = context.makeCurrent()) { + return createStreamingWithLogs( + context, chatCompletionCreateParams, requestOptions, true); + } catch (Throwable t) { + instrumenter.end(context, chatCompletionCreateParams, null, t); + throw t; + } } - Context context = instrumenter.start(parentContext, chatCompletionCreateParams); - try (Scope ignored = context.makeCurrent()) { - return createStreamingWithLogs(context, chatCompletionCreateParams, requestOptions, true); - } catch (Throwable t) { - instrumenter.end(context, chatCompletionCreateParams, null, t); - throw t; + private AsyncStreamResponse createStreamingWithLogs( + Context context, + ChatCompletionCreateParams chatCompletionCreateParams, + RequestOptions requestOptions, + boolean newSpan) { + ChatCompletionEventsHelper.emitPromptLogEvents( + context, eventLogger, chatCompletionCreateParams, captureMessageContent); + AsyncStreamResponse result = + delegate.createStreaming(chatCompletionCreateParams, requestOptions); + return new TracingAsyncStreamedResponse( + result, + new StreamListener( + context, + chatCompletionCreateParams, + instrumenter, + eventLogger, + captureMessageContent, + newSpan)); } - } - - private AsyncStreamResponse createStreamingWithLogs( - Context context, - ChatCompletionCreateParams chatCompletionCreateParams, - RequestOptions requestOptions, - boolean newSpan) { - ChatCompletionEventsHelper.emitPromptLogEvents( - context, eventLogger, chatCompletionCreateParams, captureMessageContent); - AsyncStreamResponse result = - delegate.createStreaming(chatCompletionCreateParams, requestOptions); - return new TracingAsyncStreamedResponse( - result, - new StreamListener( - context, - chatCompletionCreateParams, - instrumenter, - eventLogger, - captureMessageContent, - newSpan)); - } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatService.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatService.java index 45cfc7e..7165fb1 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatService.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatService.java @@ -13,37 +13,40 @@ import java.lang.reflect.Method; final class InstrumentedChatService - extends DelegatingInvocationHandler { + extends DelegatingInvocationHandler { - private final Instrumenter instrumenter; - private final Logger eventLogger; - private final boolean captureMessageContent; + private final Instrumenter instrumenter; + private final Logger eventLogger; + private final boolean captureMessageContent; - InstrumentedChatService( - ChatService delegate, - Instrumenter instrumenter, - Logger eventLogger, - boolean captureMessageContent) { - super(delegate); - this.instrumenter = instrumenter; - this.eventLogger = eventLogger; - this.captureMessageContent = captureMessageContent; - } + InstrumentedChatService( + ChatService delegate, + Instrumenter instrumenter, + Logger eventLogger, + boolean captureMessageContent) { + super(delegate); + this.instrumenter = instrumenter; + this.eventLogger = eventLogger; + this.captureMessageContent = captureMessageContent; + } - @Override - protected Class getProxyType() { - return ChatService.class; - } + @Override + protected Class getProxyType() { + return ChatService.class; + } - @Override - public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { - String methodName = method.getName(); - Class[] parameterTypes = method.getParameterTypes(); - if (methodName.equals("completions") && parameterTypes.length == 0) { - return new InstrumentedChatCompletionService( - delegate.completions(), instrumenter, eventLogger, captureMessageContent) - .createProxy(); + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { + String methodName = method.getName(); + Class[] parameterTypes = method.getParameterTypes(); + if (methodName.equals("completions") && parameterTypes.length == 0) { + return new InstrumentedChatCompletionService( + delegate.completions(), + instrumenter, + eventLogger, + captureMessageContent) + .createProxy(); + } + return super.invoke(proxy, method, args); } - return super.invoke(proxy, method, args); - } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatServiceAsync.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatServiceAsync.java index 692ab1c..c7e0799 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatServiceAsync.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatServiceAsync.java @@ -13,37 +13,40 @@ import java.lang.reflect.Method; final class InstrumentedChatServiceAsync - extends DelegatingInvocationHandler { + extends DelegatingInvocationHandler { - private final Instrumenter instrumenter; - private final Logger eventLogger; - private final boolean captureMessageContent; + private final Instrumenter instrumenter; + private final Logger eventLogger; + private final boolean captureMessageContent; - InstrumentedChatServiceAsync( - ChatServiceAsync delegate, - Instrumenter instrumenter, - Logger eventLogger, - boolean captureMessageContent) { - super(delegate); - this.instrumenter = instrumenter; - this.eventLogger = eventLogger; - this.captureMessageContent = captureMessageContent; - } + InstrumentedChatServiceAsync( + ChatServiceAsync delegate, + Instrumenter instrumenter, + Logger eventLogger, + boolean captureMessageContent) { + super(delegate); + this.instrumenter = instrumenter; + this.eventLogger = eventLogger; + this.captureMessageContent = captureMessageContent; + } - @Override - protected Class getProxyType() { - return ChatServiceAsync.class; - } + @Override + protected Class getProxyType() { + return ChatServiceAsync.class; + } - @Override - public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { - String methodName = method.getName(); - Class[] parameterTypes = method.getParameterTypes(); - if (methodName.equals("completions") && parameterTypes.length == 0) { - return new InstrumentedChatCompletionServiceAsync( - delegate.completions(), instrumenter, eventLogger, captureMessageContent) - .createProxy(); + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { + String methodName = method.getName(); + Class[] parameterTypes = method.getParameterTypes(); + if (methodName.equals("completions") && parameterTypes.length == 0) { + return new InstrumentedChatCompletionServiceAsync( + delegate.completions(), + instrumenter, + eventLogger, + captureMessageContent) + .createProxy(); + } + return super.invoke(proxy, method, args); } - return super.invoke(proxy, method, args); - } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedEmbeddingService.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedEmbeddingService.java index 8f6faf5..7d4405c 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedEmbeddingService.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedEmbeddingService.java @@ -15,57 +15,57 @@ import java.lang.reflect.Method; final class InstrumentedEmbeddingService - extends DelegatingInvocationHandler { + extends DelegatingInvocationHandler { - private final Instrumenter instrumenter; + private final Instrumenter instrumenter; - public InstrumentedEmbeddingService( - EmbeddingService delegate, - Instrumenter instrumenter) { - super(delegate); - this.instrumenter = instrumenter; - } + public InstrumentedEmbeddingService( + EmbeddingService delegate, + Instrumenter instrumenter) { + super(delegate); + this.instrumenter = instrumenter; + } + + @Override + protected Class getProxyType() { + return EmbeddingService.class; + } - @Override - protected Class getProxyType() { - return EmbeddingService.class; - } + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { + String methodName = method.getName(); + Class[] parameterTypes = method.getParameterTypes(); - @Override - public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { - String methodName = method.getName(); - Class[] parameterTypes = method.getParameterTypes(); + if (methodName.equals("create") + && parameterTypes.length >= 1 + && parameterTypes[0] == EmbeddingCreateParams.class) { + if (parameterTypes.length == 1) { + return create((EmbeddingCreateParams) args[0], RequestOptions.none()); + } else if (parameterTypes.length == 2 && parameterTypes[1] == RequestOptions.class) { + return create((EmbeddingCreateParams) args[0], (RequestOptions) args[1]); + } + } - if (methodName.equals("create") - && parameterTypes.length >= 1 - && parameterTypes[0] == EmbeddingCreateParams.class) { - if (parameterTypes.length == 1) { - return create((EmbeddingCreateParams) args[0], RequestOptions.none()); - } else if (parameterTypes.length == 2 && parameterTypes[1] == RequestOptions.class) { - return create((EmbeddingCreateParams) args[0], (RequestOptions) args[1]); - } + return super.invoke(proxy, method, args); } - return super.invoke(proxy, method, args); - } + private CreateEmbeddingResponse create( + EmbeddingCreateParams request, RequestOptions requestOptions) { + Context parentContext = Context.current(); + if (!instrumenter.shouldStart(parentContext, request)) { + return delegate.create(request, requestOptions); + } - private CreateEmbeddingResponse create( - EmbeddingCreateParams request, RequestOptions requestOptions) { - Context parentContext = Context.current(); - if (!instrumenter.shouldStart(parentContext, request)) { - return delegate.create(request, requestOptions); - } + Context context = instrumenter.start(parentContext, request); + CreateEmbeddingResponse response; + try (Scope ignored = context.makeCurrent()) { + response = delegate.create(request, requestOptions); + } catch (Throwable t) { + instrumenter.end(context, request, null, t); + throw t; + } - Context context = instrumenter.start(parentContext, request); - CreateEmbeddingResponse response; - try (Scope ignored = context.makeCurrent()) { - response = delegate.create(request, requestOptions); - } catch (Throwable t) { - instrumenter.end(context, request, null, t); - throw t; + instrumenter.end(context, request, response, null); + return response; } - - instrumenter.end(context, request, response, null); - return response; - } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedEmbeddingServiceAsync.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedEmbeddingServiceAsync.java index f3ea2a0..432d60a 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedEmbeddingServiceAsync.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedEmbeddingServiceAsync.java @@ -16,57 +16,58 @@ import java.util.concurrent.CompletableFuture; final class InstrumentedEmbeddingServiceAsync - extends DelegatingInvocationHandler { + extends DelegatingInvocationHandler< + EmbeddingServiceAsync, InstrumentedEmbeddingServiceAsync> { - private final Instrumenter instrumenter; + private final Instrumenter instrumenter; - public InstrumentedEmbeddingServiceAsync( - EmbeddingServiceAsync delegate, - Instrumenter instrumenter) { - super(delegate); - this.instrumenter = instrumenter; - } + public InstrumentedEmbeddingServiceAsync( + EmbeddingServiceAsync delegate, + Instrumenter instrumenter) { + super(delegate); + this.instrumenter = instrumenter; + } + + @Override + protected Class getProxyType() { + return EmbeddingServiceAsync.class; + } - @Override - protected Class getProxyType() { - return EmbeddingServiceAsync.class; - } + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { + String methodName = method.getName(); + Class[] parameterTypes = method.getParameterTypes(); - @Override - public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { - String methodName = method.getName(); - Class[] parameterTypes = method.getParameterTypes(); + if (methodName.equals("create") + && parameterTypes.length >= 1 + && parameterTypes[0] == EmbeddingCreateParams.class) { + if (parameterTypes.length == 1) { + return create((EmbeddingCreateParams) args[0], RequestOptions.none()); + } else if (parameterTypes.length == 2 && parameterTypes[1] == RequestOptions.class) { + return create((EmbeddingCreateParams) args[0], (RequestOptions) args[1]); + } + } - if (methodName.equals("create") - && parameterTypes.length >= 1 - && parameterTypes[0] == EmbeddingCreateParams.class) { - if (parameterTypes.length == 1) { - return create((EmbeddingCreateParams) args[0], RequestOptions.none()); - } else if (parameterTypes.length == 2 && parameterTypes[1] == RequestOptions.class) { - return create((EmbeddingCreateParams) args[0], (RequestOptions) args[1]); - } + return super.invoke(proxy, method, args); } - return super.invoke(proxy, method, args); - } + private CompletableFuture create( + EmbeddingCreateParams request, RequestOptions requestOptions) { + Context parentContext = Context.current(); + if (!instrumenter.shouldStart(parentContext, request)) { + return delegate.create(request, requestOptions); + } - private CompletableFuture create( - EmbeddingCreateParams request, RequestOptions requestOptions) { - Context parentContext = Context.current(); - if (!instrumenter.shouldStart(parentContext, request)) { - return delegate.create(request, requestOptions); - } + Context context = instrumenter.start(parentContext, request); + CompletableFuture future; + try (Scope ignored = context.makeCurrent()) { + future = delegate.create(request, requestOptions); + } catch (Throwable t) { + instrumenter.end(context, request, null, t); + throw t; + } - Context context = instrumenter.start(parentContext, request); - CompletableFuture future; - try (Scope ignored = context.makeCurrent()) { - future = delegate.create(request, requestOptions); - } catch (Throwable t) { - instrumenter.end(context, request, null, t); - throw t; + future = future.whenComplete((res, t) -> instrumenter.end(context, request, res, t)); + return CompletableFutureWrapper.wrap(future, parentContext); } - - future = future.whenComplete((res, t) -> instrumenter.end(context, request, res, t)); - return CompletableFutureWrapper.wrap(future, parentContext); - } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedOpenAiClient.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedOpenAiClient.java index c4f2249..7396b0d 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedOpenAiClient.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedOpenAiClient.java @@ -15,53 +15,54 @@ import java.lang.reflect.Method; final class InstrumentedOpenAiClient - extends DelegatingInvocationHandler { + extends DelegatingInvocationHandler { - private final Instrumenter chatInstrumenter; - private final Instrumenter embeddingInstrumenter; - private final Logger eventLogger; - private final boolean captureMessageContent; + private final Instrumenter chatInstrumenter; + private final Instrumenter + embeddingInstrumenter; + private final Logger eventLogger; + private final boolean captureMessageContent; - InstrumentedOpenAiClient( - OpenAIClient delegate, - Instrumenter chatInstrumenter, - Instrumenter embeddingInstrumenter, - Logger eventLogger, - boolean captureMessageContent) { - super(delegate); - this.chatInstrumenter = chatInstrumenter; - this.embeddingInstrumenter = embeddingInstrumenter; - this.eventLogger = eventLogger; - this.captureMessageContent = captureMessageContent; - } - - @Override - protected Class getProxyType() { - return OpenAIClient.class; - } - - @Override - public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { - String methodName = method.getName(); - Class[] parameterTypes = method.getParameterTypes(); - if (methodName.equals("chat") && parameterTypes.length == 0) { - return new InstrumentedChatService( - delegate.chat(), chatInstrumenter, eventLogger, captureMessageContent) - .createProxy(); + InstrumentedOpenAiClient( + OpenAIClient delegate, + Instrumenter chatInstrumenter, + Instrumenter embeddingInstrumenter, + Logger eventLogger, + boolean captureMessageContent) { + super(delegate); + this.chatInstrumenter = chatInstrumenter; + this.embeddingInstrumenter = embeddingInstrumenter; + this.eventLogger = eventLogger; + this.captureMessageContent = captureMessageContent; } - if (methodName.equals("embeddings") && parameterTypes.length == 0) { - return new InstrumentedEmbeddingService(delegate.embeddings(), embeddingInstrumenter) - .createProxy(); + + @Override + protected Class getProxyType() { + return OpenAIClient.class; } - if (methodName.equals("async") && parameterTypes.length == 0) { - return new InstrumentedOpenAiClientAsync( - delegate.async(), - chatInstrumenter, - embeddingInstrumenter, - eventLogger, - captureMessageContent) - .createProxy(); + + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { + String methodName = method.getName(); + Class[] parameterTypes = method.getParameterTypes(); + if (methodName.equals("chat") && parameterTypes.length == 0) { + return new InstrumentedChatService( + delegate.chat(), chatInstrumenter, eventLogger, captureMessageContent) + .createProxy(); + } + if (methodName.equals("embeddings") && parameterTypes.length == 0) { + return new InstrumentedEmbeddingService(delegate.embeddings(), embeddingInstrumenter) + .createProxy(); + } + if (methodName.equals("async") && parameterTypes.length == 0) { + return new InstrumentedOpenAiClientAsync( + delegate.async(), + chatInstrumenter, + embeddingInstrumenter, + eventLogger, + captureMessageContent) + .createProxy(); + } + return super.invoke(proxy, method, args); } - return super.invoke(proxy, method, args); - } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedOpenAiClientAsync.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedOpenAiClientAsync.java index 218b26a..9988f8d 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedOpenAiClientAsync.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedOpenAiClientAsync.java @@ -15,53 +15,55 @@ import java.lang.reflect.Method; final class InstrumentedOpenAiClientAsync - extends DelegatingInvocationHandler { + extends DelegatingInvocationHandler { - private final Instrumenter chatInstrumenter; - private final Instrumenter embeddingInstrumenter; - private final Logger eventLogger; - private final boolean captureMessageContent; + private final Instrumenter chatInstrumenter; + private final Instrumenter + embeddingInstrumenter; + private final Logger eventLogger; + private final boolean captureMessageContent; - InstrumentedOpenAiClientAsync( - OpenAIClientAsync delegate, - Instrumenter chatInstrumenter, - Instrumenter embeddingInstrumenter, - Logger eventLogger, - boolean captureMessageContent) { - super(delegate); - this.chatInstrumenter = chatInstrumenter; - this.embeddingInstrumenter = embeddingInstrumenter; - this.eventLogger = eventLogger; - this.captureMessageContent = captureMessageContent; - } - - @Override - protected Class getProxyType() { - return OpenAIClientAsync.class; - } - - @Override - public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { - String methodName = method.getName(); - Class[] parameterTypes = method.getParameterTypes(); - if (methodName.equals("chat") && parameterTypes.length == 0) { - return new InstrumentedChatServiceAsync( - delegate.chat(), chatInstrumenter, eventLogger, captureMessageContent) - .createProxy(); + InstrumentedOpenAiClientAsync( + OpenAIClientAsync delegate, + Instrumenter chatInstrumenter, + Instrumenter embeddingInstrumenter, + Logger eventLogger, + boolean captureMessageContent) { + super(delegate); + this.chatInstrumenter = chatInstrumenter; + this.embeddingInstrumenter = embeddingInstrumenter; + this.eventLogger = eventLogger; + this.captureMessageContent = captureMessageContent; } - if (methodName.equals("embeddings") && parameterTypes.length == 0) { - return new InstrumentedEmbeddingServiceAsync(delegate.embeddings(), embeddingInstrumenter) - .createProxy(); + + @Override + protected Class getProxyType() { + return OpenAIClientAsync.class; } - if (methodName.equals("sync") && parameterTypes.length == 0) { - return new InstrumentedOpenAiClient( - delegate.sync(), - chatInstrumenter, - embeddingInstrumenter, - eventLogger, - captureMessageContent) - .createProxy(); + + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { + String methodName = method.getName(); + Class[] parameterTypes = method.getParameterTypes(); + if (methodName.equals("chat") && parameterTypes.length == 0) { + return new InstrumentedChatServiceAsync( + delegate.chat(), chatInstrumenter, eventLogger, captureMessageContent) + .createProxy(); + } + if (methodName.equals("embeddings") && parameterTypes.length == 0) { + return new InstrumentedEmbeddingServiceAsync( + delegate.embeddings(), embeddingInstrumenter) + .createProxy(); + } + if (methodName.equals("sync") && parameterTypes.length == 0) { + return new InstrumentedOpenAiClient( + delegate.sync(), + chatInstrumenter, + embeddingInstrumenter, + eventLogger, + captureMessageContent) + .createProxy(); + } + return super.invoke(proxy, method, args); } - return super.invoke(proxy, method, args); - } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/OpenAITelemetry.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/OpenAITelemetry.java index ea7aa77..db65615 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/OpenAITelemetry.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/OpenAITelemetry.java @@ -18,47 +18,56 @@ /** Entrypoint for instrumenting OpenAI clients. */ @SuppressWarnings("IdentifierName") // Want to match library's convention public final class OpenAITelemetry { - /** Returns a new {@link OpenAITelemetry} configured with the given {@link OpenTelemetry}. */ - public static OpenAITelemetry create(OpenTelemetry openTelemetry) { - return builder(openTelemetry).build(); - } + /** Returns a new {@link OpenAITelemetry} configured with the given {@link OpenTelemetry}. */ + public static OpenAITelemetry create(OpenTelemetry openTelemetry) { + return builder(openTelemetry).build(); + } - /** - * Returns a new {@link OpenAITelemetryBuilder} configured with the given {@link OpenTelemetry}. - */ - public static OpenAITelemetryBuilder builder(OpenTelemetry openTelemetry) { - return new OpenAITelemetryBuilder(openTelemetry); - } + /** + * Returns a new {@link OpenAITelemetryBuilder} configured with the given {@link OpenTelemetry}. + */ + public static OpenAITelemetryBuilder builder(OpenTelemetry openTelemetry) { + return new OpenAITelemetryBuilder(openTelemetry); + } - private final Instrumenter chatInstrumenter; - private final Instrumenter embeddingsInstrumenter; + private final Instrumenter chatInstrumenter; + private final Instrumenter + embeddingsInstrumenter; - private final Logger eventLogger; + private final Logger eventLogger; - private final boolean captureMessageContent; + private final boolean captureMessageContent; - OpenAITelemetry( - Instrumenter chatInstrumenter, - Instrumenter embeddingsInstrumenter, - Logger eventLogger, - boolean captureMessageContent) { - this.chatInstrumenter = chatInstrumenter; - this.embeddingsInstrumenter = embeddingsInstrumenter; - this.eventLogger = eventLogger; - this.captureMessageContent = captureMessageContent; - } + OpenAITelemetry( + Instrumenter chatInstrumenter, + Instrumenter embeddingsInstrumenter, + Logger eventLogger, + boolean captureMessageContent) { + this.chatInstrumenter = chatInstrumenter; + this.embeddingsInstrumenter = embeddingsInstrumenter; + this.eventLogger = eventLogger; + this.captureMessageContent = captureMessageContent; + } - /** Wraps the provided OpenAIClient, enabling telemetry for it. */ - public OpenAIClient wrap(OpenAIClient client) { - return new InstrumentedOpenAiClient( - client, chatInstrumenter, embeddingsInstrumenter, eventLogger, captureMessageContent) - .createProxy(); - } + /** Wraps the provided OpenAIClient, enabling telemetry for it. */ + public OpenAIClient wrap(OpenAIClient client) { + return new InstrumentedOpenAiClient( + client, + chatInstrumenter, + embeddingsInstrumenter, + eventLogger, + captureMessageContent) + .createProxy(); + } - /** Wraps the provided OpenAIClientAsync, enabling telemetry for it. */ - public OpenAIClientAsync wrap(OpenAIClientAsync client) { - return new InstrumentedOpenAiClientAsync( - client, chatInstrumenter, embeddingsInstrumenter, eventLogger, captureMessageContent) - .createProxy(); - } + /** Wraps the provided OpenAIClientAsync, enabling telemetry for it. */ + public OpenAIClientAsync wrap(OpenAIClientAsync client) { + return new InstrumentedOpenAiClientAsync( + client, + chatInstrumenter, + embeddingsInstrumenter, + eventLogger, + captureMessageContent) + .createProxy(); + } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/OpenAITelemetryBuilder.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/OpenAITelemetryBuilder.java index 8b8b2fa..9a25866 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/OpenAITelemetryBuilder.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/OpenAITelemetryBuilder.java @@ -21,53 +21,55 @@ /** A builder of {@link OpenAITelemetry}. */ @SuppressWarnings("IdentifierName") // Want to match library's convention public final class OpenAITelemetryBuilder { - static final String INSTRUMENTATION_NAME = "io.opentelemetry.openai-java-1.1"; + static final String INSTRUMENTATION_NAME = "io.opentelemetry.openai-java-1.1"; - private final OpenTelemetry openTelemetry; + private final OpenTelemetry openTelemetry; - private boolean captureMessageContent; + private boolean captureMessageContent; - OpenAITelemetryBuilder(OpenTelemetry openTelemetry) { - this.openTelemetry = openTelemetry; - } + OpenAITelemetryBuilder(OpenTelemetry openTelemetry) { + this.openTelemetry = openTelemetry; + } - /** - * Sets whether emitted log events include full content of user and assistant messages. - * - *

Note that full content can have data privacy and size concerns and care should be taken when - * enabling this. - */ - @CanIgnoreReturnValue - public OpenAITelemetryBuilder setCaptureMessageContent(boolean captureMessageContent) { - this.captureMessageContent = captureMessageContent; - return this; - } + /** + * Sets whether emitted log events include full content of user and assistant messages. + * + *

Note that full content can have data privacy and size concerns and care should be taken + * when enabling this. + */ + @CanIgnoreReturnValue + public OpenAITelemetryBuilder setCaptureMessageContent(boolean captureMessageContent) { + this.captureMessageContent = captureMessageContent; + return this; + } - /** - * Returns a new {@link OpenAITelemetry} with the settings of this {@link OpenAITelemetryBuilder}. - */ - public OpenAITelemetry build() { - Instrumenter chatInstrumenter = - Instrumenter.builder( - openTelemetry, - INSTRUMENTATION_NAME, - GenAiSpanNameExtractor.create(ChatAttributesGetter.INSTANCE)) - .addAttributesExtractor(GenAiAttributesExtractor.create(ChatAttributesGetter.INSTANCE)) - .addOperationMetrics(GenAiClientMetrics.get()) - .buildInstrumenter(); + /** + * Returns a new {@link OpenAITelemetry} with the settings of this {@link + * OpenAITelemetryBuilder}. + */ + public OpenAITelemetry build() { + Instrumenter chatInstrumenter = + Instrumenter.builder( + openTelemetry, + INSTRUMENTATION_NAME, + GenAiSpanNameExtractor.create(ChatAttributesGetter.INSTANCE)) + .addAttributesExtractor( + GenAiAttributesExtractor.create(ChatAttributesGetter.INSTANCE)) + .addOperationMetrics(GenAiClientMetrics.get()) + .buildInstrumenter(); - Instrumenter embeddingsInstrumenter = - Instrumenter.builder( - openTelemetry, - INSTRUMENTATION_NAME, - GenAiSpanNameExtractor.create(EmbeddingAttributesGetter.INSTANCE)) - .addAttributesExtractor( - GenAiAttributesExtractor.create(EmbeddingAttributesGetter.INSTANCE)) - .addOperationMetrics(GenAiClientMetrics.get()) - .buildInstrumenter(SpanKindExtractor.alwaysClient()); + Instrumenter embeddingsInstrumenter = + Instrumenter.builder( + openTelemetry, + INSTRUMENTATION_NAME, + GenAiSpanNameExtractor.create(EmbeddingAttributesGetter.INSTANCE)) + .addAttributesExtractor( + GenAiAttributesExtractor.create(EmbeddingAttributesGetter.INSTANCE)) + .addOperationMetrics(GenAiClientMetrics.get()) + .buildInstrumenter(SpanKindExtractor.alwaysClient()); - Logger eventLogger = openTelemetry.getLogsBridge().get(INSTRUMENTATION_NAME); - return new OpenAITelemetry( - chatInstrumenter, embeddingsInstrumenter, eventLogger, captureMessageContent); - } + Logger eventLogger = openTelemetry.getLogsBridge().get(INSTRUMENTATION_NAME); + return new OpenAITelemetry( + chatInstrumenter, embeddingsInstrumenter, eventLogger, captureMessageContent); + } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/StreamListener.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/StreamListener.java index 7187bc7..3a0c164 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/StreamListener.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/StreamListener.java @@ -20,93 +20,97 @@ final class StreamListener { - private final Context context; - private final ChatCompletionCreateParams request; - private final List choiceBuffers; + private final Context context; + private final ChatCompletionCreateParams request; + private final List choiceBuffers; - private final Instrumenter instrumenter; - private final Logger eventLogger; - private final boolean captureMessageContent; - private final boolean newSpan; - private final AtomicBoolean hasEnded; + private final Instrumenter instrumenter; + private final Logger eventLogger; + private final boolean captureMessageContent; + private final boolean newSpan; + private final AtomicBoolean hasEnded; - @Nullable private CompletionUsage usage; - @Nullable private String model; - @Nullable private String responseId; + @Nullable private CompletionUsage usage; + @Nullable private String model; + @Nullable private String responseId; - StreamListener( - Context context, - ChatCompletionCreateParams request, - Instrumenter instrumenter, - Logger eventLogger, - boolean captureMessageContent, - boolean newSpan) { - this.context = context; - this.request = request; - this.instrumenter = instrumenter; - this.eventLogger = eventLogger; - this.captureMessageContent = captureMessageContent; - this.newSpan = newSpan; - choiceBuffers = new ArrayList<>(); - hasEnded = new AtomicBoolean(); - } + StreamListener( + Context context, + ChatCompletionCreateParams request, + Instrumenter instrumenter, + Logger eventLogger, + boolean captureMessageContent, + boolean newSpan) { + this.context = context; + this.request = request; + this.instrumenter = instrumenter; + this.eventLogger = eventLogger; + this.captureMessageContent = captureMessageContent; + this.newSpan = newSpan; + choiceBuffers = new ArrayList<>(); + hasEnded = new AtomicBoolean(); + } - void onChunk(ChatCompletionChunk chunk) { - model = chunk.model(); - responseId = chunk.id(); - chunk.usage().ifPresent(u -> usage = u); + void onChunk(ChatCompletionChunk chunk) { + model = chunk.model(); + responseId = chunk.id(); + chunk.usage().ifPresent(u -> usage = u); - for (ChatCompletionChunk.Choice choice : chunk.choices()) { - while (choiceBuffers.size() <= choice.index()) { - choiceBuffers.add(null); - } - StreamedMessageBuffer buffer = choiceBuffers.get((int) choice.index()); - if (buffer == null) { - buffer = new StreamedMessageBuffer(choice.index(), captureMessageContent); - choiceBuffers.set((int) choice.index(), buffer); - } - buffer.append(choice.delta()); - if (choice.finishReason().isPresent()) { - buffer.finishReason = choice.finishReason().get().toString(); + for (ChatCompletionChunk.Choice choice : chunk.choices()) { + while (choiceBuffers.size() <= choice.index()) { + choiceBuffers.add(null); + } + StreamedMessageBuffer buffer = choiceBuffers.get((int) choice.index()); + if (buffer == null) { + buffer = new StreamedMessageBuffer(choice.index(), captureMessageContent); + choiceBuffers.set((int) choice.index(), buffer); + } + buffer.append(choice.delta()); + if (choice.finishReason().isPresent()) { + buffer.finishReason = choice.finishReason().get().toString(); - // message has ended, let's emit - ChatCompletionEventsHelper.emitCompletionLogEvent( - context, eventLogger, choice.index(), buffer.finishReason, buffer.toEventBody()); - } + // message has ended, let's emit + ChatCompletionEventsHelper.emitCompletionLogEvent( + context, + eventLogger, + choice.index(), + buffer.finishReason, + buffer.toEventBody()); + } + } } - } - void endSpan(@Nullable Throwable error) { - // Use an atomic operation since close() type of methods are exposed to the user - // and can come from any thread. - if (!hasEnded.compareAndSet(false, true)) { - return; - } + void endSpan(@Nullable Throwable error) { + // Use an atomic operation since close() type of methods are exposed to the user + // and can come from any thread. + if (!hasEnded.compareAndSet(false, true)) { + return; + } - if (model == null || responseId == null) { - // Only happens if we got no chunks, so we have no response. - if (newSpan) { - instrumenter.end(context, request, null, error); - } - return; - } + if (model == null || responseId == null) { + // Only happens if we got no chunks, so we have no response. + if (newSpan) { + instrumenter.end(context, request, null, error); + } + return; + } - ChatCompletion.Builder result = - ChatCompletion.builder() - .created(0) - .model(model) - .id(responseId) - .choices( - choiceBuffers.stream() - .map(StreamedMessageBuffer::toChoice) - .collect(Collectors.toList())); + ChatCompletion.Builder result = + ChatCompletion.builder() + .created(0) + .model(model) + .id(responseId) + .choices( + choiceBuffers.stream() + .map(StreamedMessageBuffer::toChoice) + .collect(Collectors.toList())); - if (usage != null) { - result.usage(usage); - } + if (usage != null) { + result.usage(usage); + } - if (newSpan) { - instrumenter.end(context, request, result.build(), error); + if (newSpan) { + instrumenter.end(context, request, result.build(), error); + } } - } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/StreamedMessageBuffer.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/StreamedMessageBuffer.java index a43d8ae..1c1be60 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/StreamedMessageBuffer.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/StreamedMessageBuffer.java @@ -18,127 +18,127 @@ import javax.annotation.Nullable; final class StreamedMessageBuffer { - private final long index; - private final boolean captureMessageContent; - - @Nullable String finishReason; - - @Nullable private StringBuilder message; - @Nullable private Map toolCalls; - - StreamedMessageBuffer(long index, boolean captureMessageContent) { - this.index = index; - this.captureMessageContent = captureMessageContent; - } - - ChatCompletion.Choice toChoice() { - ChatCompletion.Choice.Builder choice = - ChatCompletion.Choice.builder().index(index).logprobs(Optional.empty()); - if (finishReason != null) { - choice.finishReason(ChatCompletion.Choice.FinishReason.of(finishReason)); - } else { - // Can't happen in practice, mostly to satisfy null check - choice.finishReason(JsonField.ofNullable(null)); - } - if (message != null) { - choice.message( - ChatCompletionMessage.builder() - .content(message.toString()) - .refusal(Optional.empty()) - .build()); - } else { - choice.message(JsonField.ofNullable(null)); - } - return choice.build(); - } + private final long index; + private final boolean captureMessageContent; + + @Nullable String finishReason; - Value toEventBody() { - Map> body = new HashMap<>(); - if (message != null) { - body.put("content", Value.of(message.toString())); + @Nullable private StringBuilder message; + @Nullable private Map toolCalls; + + StreamedMessageBuffer(long index, boolean captureMessageContent) { + this.index = index; + this.captureMessageContent = captureMessageContent; } - if (toolCalls != null) { - List> toolCallsJson = - toolCalls.values().stream() - .map(StreamedMessageBuffer::buildToolCallEventObject) - .collect(Collectors.toList()); - body.put("tool_calls", Value.of(toolCallsJson)); + + ChatCompletion.Choice toChoice() { + ChatCompletion.Choice.Builder choice = + ChatCompletion.Choice.builder().index(index).logprobs(Optional.empty()); + if (finishReason != null) { + choice.finishReason(ChatCompletion.Choice.FinishReason.of(finishReason)); + } else { + // Can't happen in practice, mostly to satisfy null check + choice.finishReason(JsonField.ofNullable(null)); + } + if (message != null) { + choice.message( + ChatCompletionMessage.builder() + .content(message.toString()) + .refusal(Optional.empty()) + .build()); + } else { + choice.message(JsonField.ofNullable(null)); + } + return choice.build(); } - return Value.of(body); - } - - void append(ChatCompletionChunk.Choice.Delta delta) { - if (captureMessageContent) { - if (delta.content().isPresent()) { - if (message == null) { - message = new StringBuilder(); + + Value toEventBody() { + Map> body = new HashMap<>(); + if (message != null) { + body.put("content", Value.of(message.toString())); + } + if (toolCalls != null) { + List> toolCallsJson = + toolCalls.values().stream() + .map(StreamedMessageBuffer::buildToolCallEventObject) + .collect(Collectors.toList()); + body.put("tool_calls", Value.of(toolCallsJson)); } - message.append(delta.content().get()); - } + return Value.of(body); } - if (delta.toolCalls().isPresent()) { - if (toolCalls == null) { - toolCalls = new HashMap<>(); - } - - for (ChatCompletionChunk.Choice.Delta.ToolCall toolCall : delta.toolCalls().get()) { - ToolCallBuffer buffer = - toolCalls.computeIfAbsent( - toolCall.index(), unused -> new ToolCallBuffer(toolCall.id().orElse(""))); - toolCall.type().ifPresent(type -> buffer.type = type.toString()); - toolCall - .function() - .ifPresent( - function -> { - function.name().ifPresent(name -> buffer.function.name = name); - if (captureMessageContent) { - function - .arguments() + void append(ChatCompletionChunk.Choice.Delta delta) { + if (captureMessageContent) { + if (delta.content().isPresent()) { + if (message == null) { + message = new StringBuilder(); + } + message.append(delta.content().get()); + } + } + + if (delta.toolCalls().isPresent()) { + if (toolCalls == null) { + toolCalls = new HashMap<>(); + } + + for (ChatCompletionChunk.Choice.Delta.ToolCall toolCall : delta.toolCalls().get()) { + ToolCallBuffer buffer = + toolCalls.computeIfAbsent( + toolCall.index(), + unused -> new ToolCallBuffer(toolCall.id().orElse(""))); + toolCall.type().ifPresent(type -> buffer.type = type.toString()); + toolCall.function() .ifPresent( - args -> { - if (buffer.function.arguments == null) { - buffer.function.arguments = new StringBuilder(); - } - buffer.function.arguments.append(args); - }); - } - }); - } + function -> { + function.name().ifPresent(name -> buffer.function.name = name); + if (captureMessageContent) { + function.arguments() + .ifPresent( + args -> { + if (buffer.function.arguments == null) { + buffer.function.arguments = + new StringBuilder(); + } + buffer.function.arguments.append(args); + }); + } + }); + } + } } - } - private static Value buildToolCallEventObject(ToolCallBuffer call) { - Map> result = new HashMap<>(); - result.put("id", Value.of(call.id)); - if (call.type != null) { - result.put("type", Value.of(call.type)); - } + private static Value buildToolCallEventObject(ToolCallBuffer call) { + Map> result = new HashMap<>(); + result.put("id", Value.of(call.id)); + if (call.type != null) { + result.put("type", Value.of(call.type)); + } - Map> function = new HashMap<>(); - if (call.function.name != null) { - function.put("name", Value.of(call.function.name)); - } - if (call.function.arguments != null) { - function.put("arguments", Value.of(call.function.arguments.toString())); - } - result.put("function", Value.of(function)); + Map> function = new HashMap<>(); + if (call.function.name != null) { + function.put("name", Value.of(call.function.name)); + } + if (call.function.arguments != null) { + function.put("arguments", Value.of(call.function.arguments.toString())); + } + result.put("function", Value.of(function)); - return Value.of(result); - } + return Value.of(result); + } - private static class FunctionBuffer { - @Nullable String name; - @Nullable StringBuilder arguments; - } + private static class FunctionBuffer { + @Nullable String name; + @Nullable StringBuilder arguments; + } - private static class ToolCallBuffer { - final String id; - final FunctionBuffer function = new FunctionBuffer(); - @Nullable String type; + private static class ToolCallBuffer { + final String id; + final FunctionBuffer function = new FunctionBuffer(); + @Nullable String type; - ToolCallBuffer(String id) { - this.id = id; + ToolCallBuffer(String id) { + this.id = id; + } } - } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/TracingAsyncStreamedResponse.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/TracingAsyncStreamedResponse.java index 375c076..533dcbb 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/TracingAsyncStreamedResponse.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/TracingAsyncStreamedResponse.java @@ -13,58 +13,58 @@ final class TracingAsyncStreamedResponse implements AsyncStreamResponse { - private final AsyncStreamResponse delegate; - private final StreamListener listener; + private final AsyncStreamResponse delegate; + private final StreamListener listener; - TracingAsyncStreamedResponse( - AsyncStreamResponse delegate, StreamListener listener) { - this.delegate = delegate; - this.listener = listener; - } + TracingAsyncStreamedResponse( + AsyncStreamResponse delegate, StreamListener listener) { + this.delegate = delegate; + this.listener = listener; + } - @Override - public void close() { - listener.endSpan(null); - delegate.close(); - } + @Override + public void close() { + listener.endSpan(null); + delegate.close(); + } - @Override - public AsyncStreamResponse subscribe( - Handler handler) { - delegate.subscribe(new TracingHandler(handler)); - return this; - } + @Override + public AsyncStreamResponse subscribe( + Handler handler) { + delegate.subscribe(new TracingHandler(handler)); + return this; + } - @Override - public AsyncStreamResponse subscribe( - Handler handler, Executor executor) { - delegate.subscribe(new TracingHandler(handler), executor); - return this; - } + @Override + public AsyncStreamResponse subscribe( + Handler handler, Executor executor) { + delegate.subscribe(new TracingHandler(handler), executor); + return this; + } - @Override - public CompletableFuture onCompleteFuture() { - return delegate.onCompleteFuture(); - } + @Override + public CompletableFuture onCompleteFuture() { + return delegate.onCompleteFuture(); + } - private class TracingHandler implements Handler { + private class TracingHandler implements Handler { - private final Handler delegate; + private final Handler delegate; - private TracingHandler(Handler delegate) { - this.delegate = delegate; - } + private TracingHandler(Handler delegate) { + this.delegate = delegate; + } - @Override - public void onNext(ChatCompletionChunk chunk) { - listener.onChunk(chunk); - delegate.onNext(chunk); - } + @Override + public void onNext(ChatCompletionChunk chunk) { + listener.onChunk(chunk); + delegate.onNext(chunk); + } - @Override - public void onComplete(Optional error) { - listener.endSpan(error.orElse(null)); - delegate.onComplete(error); + @Override + public void onComplete(Optional error) { + listener.endSpan(error.orElse(null)); + delegate.onComplete(error); + } } - } } diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/TracingStreamedResponse.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/TracingStreamedResponse.java index db1810a..9a14b3b 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/TracingStreamedResponse.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/TracingStreamedResponse.java @@ -16,72 +16,72 @@ final class TracingStreamedResponse implements StreamResponse { - private final StreamResponse delegate; - private final StreamListener listener; + private final StreamResponse delegate; + private final StreamListener listener; - TracingStreamedResponse(StreamResponse delegate, StreamListener listener) { - this.delegate = delegate; - this.listener = listener; - } - - @Override - public Stream stream() { - return StreamSupport.stream(new TracingSpliterator(delegate.stream().spliterator()), false); - } - - @Override - public void close() { - listener.endSpan(null); - delegate.close(); - } - - private class TracingSpliterator implements Spliterator { - - private final Spliterator delegateSpliterator; - - private TracingSpliterator(Spliterator delegateSpliterator) { - this.delegateSpliterator = delegateSpliterator; - } - - @Override - public boolean tryAdvance(Consumer action) { - boolean chunkReceived = - delegateSpliterator.tryAdvance( - chunk -> { - listener.onChunk(chunk); - action.accept(chunk); - }); - if (!chunkReceived) { - listener.endSpan(null); - } - return chunkReceived; + TracingStreamedResponse(StreamResponse delegate, StreamListener listener) { + this.delegate = delegate; + this.listener = listener; } @Override - @Nullable - public Spliterator trySplit() { - // do not support parallelism to reliably catch the last chunk - return null; + public Stream stream() { + return StreamSupport.stream(new TracingSpliterator(delegate.stream().spliterator()), false); } @Override - public long estimateSize() { - return delegateSpliterator.estimateSize(); - } - - @Override - public long getExactSizeIfKnown() { - return delegateSpliterator.getExactSizeIfKnown(); - } - - @Override - public int characteristics() { - return delegateSpliterator.characteristics(); + public void close() { + listener.endSpan(null); + delegate.close(); } - @Override - public Comparator getComparator() { - return delegateSpliterator.getComparator(); + private class TracingSpliterator implements Spliterator { + + private final Spliterator delegateSpliterator; + + private TracingSpliterator(Spliterator delegateSpliterator) { + this.delegateSpliterator = delegateSpliterator; + } + + @Override + public boolean tryAdvance(Consumer action) { + boolean chunkReceived = + delegateSpliterator.tryAdvance( + chunk -> { + listener.onChunk(chunk); + action.accept(chunk); + }); + if (!chunkReceived) { + listener.endSpan(null); + } + return chunkReceived; + } + + @Override + @Nullable + public Spliterator trySplit() { + // do not support parallelism to reliably catch the last chunk + return null; + } + + @Override + public long estimateSize() { + return delegateSpliterator.estimateSize(); + } + + @Override + public long getExactSizeIfKnown() { + return delegateSpliterator.getExactSizeIfKnown(); + } + + @Override + public int characteristics() { + return delegateSpliterator.characteristics(); + } + + @Override + public Comparator getComparator() { + return delegateSpliterator.getComparator(); + } } - } } diff --git a/src/main/java/dev/braintrust/log/BraintrustLogger.java b/src/main/java/dev/braintrust/log/BraintrustLogger.java index 942a561..0f1a401 100644 --- a/src/main/java/dev/braintrust/log/BraintrustLogger.java +++ b/src/main/java/dev/braintrust/log/BraintrustLogger.java @@ -1,12 +1,11 @@ package dev.braintrust.log; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; -import org.slf4j.event.Level; - import java.util.concurrent.locks.ReadWriteLock; import java.util.concurrent.locks.ReentrantReadWriteLock; import java.util.function.Supplier; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.slf4j.event.Level; /** * Centralized logging for Braintrust SDK. Provides consistent logging with SDK-specific control and diff --git a/src/main/java/dev/braintrust/trace/BraintrustContext.java b/src/main/java/dev/braintrust/trace/BraintrustContext.java index 14d86f3..6a8b3b1 100644 --- a/src/main/java/dev/braintrust/trace/BraintrustContext.java +++ b/src/main/java/dev/braintrust/trace/BraintrustContext.java @@ -3,9 +3,8 @@ import io.opentelemetry.api.trace.Span; import io.opentelemetry.context.Context; import io.opentelemetry.context.ContextKey; - -import javax.annotation.Nullable; import java.util.Optional; +import javax.annotation.Nullable; /** * Context propagation for Braintrust-specific attributes. Uses OpenTelemetry's Context API for @@ -39,7 +38,8 @@ public static BraintrustContext forExperiment(String experimentId) { /** Creates a context for an experiment parent. */ public static Context forExperiment(String experimentId, Span span) { - return Context.current().with(span) + return Context.current() + .with(span) .with(KEY, new BraintrustContext(null, experimentId, "experiment")); } diff --git a/src/main/java/dev/braintrust/trace/BraintrustLogExporter.java b/src/main/java/dev/braintrust/trace/BraintrustLogExporter.java index 66d49ef..81dbc83 100644 --- a/src/main/java/dev/braintrust/trace/BraintrustLogExporter.java +++ b/src/main/java/dev/braintrust/trace/BraintrustLogExporter.java @@ -6,7 +6,6 @@ import io.opentelemetry.sdk.common.CompletableResultCode; import io.opentelemetry.sdk.logs.data.LogRecordData; import io.opentelemetry.sdk.logs.export.LogRecordExporter; - import java.util.Collection; import java.util.List; import java.util.Map; @@ -108,14 +107,16 @@ private CompletableResultCode exportWithParent(String parent, List PARENT_PROJECT_ID = AttributeKey.stringKey("braintrust.parent.project_id"); + 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"); + 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_TYPE = + AttributeKey.stringKey("braintrust.parent.type"); + // Usage metrics - public static final AttributeKey USAGE_PROMPT_TOKENS = AttributeKey.longKey("braintrust.usage.prompt_tokens"); - public static final AttributeKey USAGE_COMPLETION_TOKENS = AttributeKey.longKey("braintrust.usage.completion_tokens"); - public static final AttributeKey USAGE_TOTAL_TOKENS = AttributeKey.longKey("braintrust.usage.total_tokens"); - public static final AttributeKey USAGE_COST = AttributeKey.doubleKey("braintrust.usage.cost"); + public static final AttributeKey USAGE_PROMPT_TOKENS = + AttributeKey.longKey("braintrust.usage.prompt_tokens"); + public static final AttributeKey USAGE_COMPLETION_TOKENS = + AttributeKey.longKey("braintrust.usage.completion_tokens"); + public static final AttributeKey USAGE_TOTAL_TOKENS = + AttributeKey.longKey("braintrust.usage.total_tokens"); + public static final AttributeKey USAGE_COST = + AttributeKey.doubleKey("braintrust.usage.cost"); // Scoring public static final AttributeKey SCORE = AttributeKey.doubleKey("braintrust.score"); - public static final AttributeKey SCORE_NAME = AttributeKey.stringKey("braintrust.score.name"); + public static final AttributeKey SCORE_NAME = + AttributeKey.stringKey("braintrust.score.name"); private final BraintrustConfig config; private final SpanProcessor delegate; @@ -139,14 +147,14 @@ public Optional getParentContext(String traceId) { private void logSpanDetails(ReadableSpan span) { var spanData = span.toSpanData(); BraintrustLogger.debug( - "Span completed: name={}, traceId={}, spanId={}, duration={}ms, attributes={}, events={}", + "Span completed: name={}, traceId={}, spanId={}, duration={}ms, attributes={}," + + " events={}", spanData.getName(), spanData.getTraceId(), spanData.getSpanId(), (spanData.getEndEpochNanos() - spanData.getStartEpochNanos()) / 1_000_000, spanData.getAttributes(), - spanData.getEvents() - ); + spanData.getEvents()); } /** Parent context for spans (project or experiment). */ diff --git a/src/main/java/dev/braintrust/trace/BraintrustTracing.java b/src/main/java/dev/braintrust/trace/BraintrustTracing.java index 4d2481e..8b9a5c9 100644 --- a/src/main/java/dev/braintrust/trace/BraintrustTracing.java +++ b/src/main/java/dev/braintrust/trace/BraintrustTracing.java @@ -5,35 +5,24 @@ import io.opentelemetry.api.GlobalOpenTelemetry; import io.opentelemetry.api.OpenTelemetry; import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.exporter.otlp.http.logs.OtlpHttpLogRecordExporter; -import io.opentelemetry.exporter.otlp.http.metrics.OtlpHttpMetricExporter; -import io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporter; -import io.opentelemetry.exporter.otlp.logs.OtlpGrpcLogRecordExporter; import io.opentelemetry.sdk.OpenTelemetrySdk; -import io.opentelemetry.sdk.OpenTelemetrySdkBuilder; import io.opentelemetry.sdk.common.CompletableResultCode; import io.opentelemetry.sdk.logs.SdkLoggerProvider; import io.opentelemetry.sdk.logs.SdkLoggerProviderBuilder; import io.opentelemetry.sdk.logs.export.BatchLogRecordProcessor; import io.opentelemetry.sdk.metrics.SdkMeterProvider; import io.opentelemetry.sdk.metrics.SdkMeterProviderBuilder; -import io.opentelemetry.sdk.metrics.export.PeriodicMetricReader; import io.opentelemetry.sdk.resources.Resource; import io.opentelemetry.sdk.trace.SdkTracerProvider; import io.opentelemetry.sdk.trace.SdkTracerProviderBuilder; import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; -import io.opentelemetry.sdk.trace.export.SpanExporter; import io.opentelemetry.semconv.resource.attributes.ResourceAttributes; - -import javax.annotation.Nonnull; -import javax.annotation.Nullable; import java.time.Duration; import java.util.HashMap; -import java.util.List; import java.util.Map; import java.util.concurrent.TimeUnit; -import java.util.function.Consumer; import java.util.stream.Stream; +import javax.annotation.Nonnull; /** * Main entry point for Braintrust tracing setup. Provides convenient methods to initialize @@ -58,29 +47,33 @@ public static Tracer getTracer(OpenTelemetry openTelemetry) { } /** - * Quick start method that sets up global OpenTelemetry with Braintrust defaults. - *

- * If you're looking for more options for configuring Braintrust/OpenTelemetry, consult the `enable` method. + * Quick start method that sets up global OpenTelemetry with Braintrust defaults.
+ *
+ * If you're looking for more options for configuring Braintrust/OpenTelemetry, consult the + * `enable` method. */ public static OpenTelemetry quickstart() { return quickstart(BraintrustConfig.fromEnvironment(), true); } /** - * Quick start method that sets up OpenTelemetry with custom Braintrust and otel settings. - *

- * If you're looking for more options for configuring Braintrust and OpenTelemetry, consult the `enable` method. + * Quick start method that sets up OpenTelemetry with custom Braintrust and otel settings.
+ *
+ * If you're looking for more options for configuring Braintrust and OpenTelemetry, consult the + * `enable` method. */ - public static OpenTelemetry quickstart(@Nonnull BraintrustConfig config, boolean registerGlobal) { + public static OpenTelemetry quickstart( + @Nonnull BraintrustConfig config, boolean registerGlobal) { var tracerBuilder = SdkTracerProvider.builder(); var loggerBuilder = SdkLoggerProvider.builder(); var meterBuilder = SdkMeterProvider.builder(); enable(config, tracerBuilder, loggerBuilder, meterBuilder); - var openTelemetry = OpenTelemetrySdk.builder() - .setTracerProvider(tracerBuilder.build()) - .setLoggerProvider(loggerBuilder.build()) - .setMeterProvider(meterBuilder.build()) - .build(); + var openTelemetry = + OpenTelemetrySdk.builder() + .setTracerProvider(tracerBuilder.build()) + .setLoggerProvider(loggerBuilder.build()) + .setMeterProvider(meterBuilder.build()) + .build(); if (registerGlobal) { GlobalOpenTelemetry.set(openTelemetry); BraintrustLogger.debug("Registered OpenTelemetry globally"); @@ -89,16 +82,24 @@ public static OpenTelemetry quickstart(@Nonnull BraintrustConfig config, boolean } /** - * Add braintrust to existing open telemetry builders - *

- * This method provides the most options for configuring Braintrust and OpenTelemetry. If you're looking for a more user-friendly setup, consult the `quickstart` methods. - *

- * NOTE: if your otel setup does not have any particular builder, pass an instance of the default provider builder. E.g. `SdkMeterProvider.builder()` - *

- * NOTE: This method should only be invoked once. Enabling Braintrust multiple times is unsupported and may lead to undesired behavior + * Add braintrust to existing open telemetry builders
+ *
+ * This method provides the most options for configuring Braintrust and OpenTelemetry. If you're + * looking for a more user-friendly setup, consult the `quickstart` methods.
+ *
+ * NOTE: if your otel setup does not have any particular builder, pass an instance of the + * default provider builder. E.g. `SdkMeterProvider.builder()`
+ *
+ * NOTE: This method should only be invoked once. Enabling Braintrust multiple times is + * unsupported and may lead to undesired behavior */ - public static void enable(@Nonnull BraintrustConfig config, @Nonnull SdkTracerProviderBuilder tracerProviderBuilder, @Nonnull SdkLoggerProviderBuilder loggerProviderBuilder, @Nonnull SdkMeterProviderBuilder meterProviderBuilder) { - // NOTE: these options are hardcoded for now. We'll move them into braintrust config as-needed + public static void enable( + @Nonnull BraintrustConfig config, + @Nonnull SdkTracerProviderBuilder tracerProviderBuilder, + @Nonnull SdkLoggerProviderBuilder loggerProviderBuilder, + @Nonnull SdkMeterProviderBuilder meterProviderBuilder) { + // NOTE: these options are hardcoded for now. We'll move them into braintrust config + // as-needed final Map resourceAttributes = new HashMap<>(); final String serviceName = "braintrust-app"; final Duration exportInterval = Duration.ofSeconds(5); @@ -115,36 +116,52 @@ public static void enable(@Nonnull BraintrustConfig config, @Nonnull SdkTracerPr var resource = resourceBuilder.build(); // spans - var spanProcessor = new BraintrustSpanProcessor(config, - BatchSpanProcessor.builder(new BraintrustSpanExporter(config)) + var spanProcessor = + new BraintrustSpanProcessor( + config, + BatchSpanProcessor.builder(new BraintrustSpanExporter(config)) + .setScheduleDelay(exportInterval.toMillis(), TimeUnit.MILLISECONDS) + .setMaxQueueSize(maxQueueSize) + .setMaxExportBatchSize(maxExportBatchSize) + .build()); + tracerProviderBuilder.addResource(resource).addSpanProcessor(spanProcessor); + // logs + var logProcessor = + BatchLogRecordProcessor.builder(new BraintrustLogExporter(config)) .setScheduleDelay(exportInterval.toMillis(), TimeUnit.MILLISECONDS) .setMaxQueueSize(maxQueueSize) .setMaxExportBatchSize(maxExportBatchSize) - .build()); - tracerProviderBuilder.addResource(resource).addSpanProcessor(spanProcessor); - // logs - var logProcessor = BatchLogRecordProcessor.builder(new BraintrustLogExporter(config)) - .setScheduleDelay(exportInterval.toMillis(), TimeUnit.MILLISECONDS) - .setMaxQueueSize(maxQueueSize) - .setMaxExportBatchSize(maxExportBatchSize) - .build(); + .build(); loggerProviderBuilder.addLogRecordProcessor(logProcessor); // NOTE: we don't actually do anything with otel metrics right now, - // but it's included in the method signature so we can do so in the future without introducing a breaking change + // but it's included in the method signature so we can do so in the future without + // introducing a breaking change Runtime.getRuntime() .addShutdownHook( new Thread( () -> { - BraintrustLogger.debug("Shutting down. Force-Flushing all otel data."); - var result = CompletableResultCode.ofAll( - // run all flushes in parallel. Should (rarely) block for approx 10 seconds max - Stream.of(spanProcessor.shutdown(), logProcessor.shutdown() ) - .map(operation -> operation.join(10, TimeUnit.SECONDS)) - .toList() - ); - BraintrustLogger.debug("otel shutdown complete. Flush done: %s, Flush successful: %s".formatted(result.isDone(), result.isSuccess())); + BraintrustLogger.debug( + "Shutting down. Force-Flushing all otel data."); + var result = + CompletableResultCode.ofAll( + // run all flushes in parallel. Should (rarely) + // block for approx 10 seconds max + Stream.of( + spanProcessor.shutdown(), + logProcessor.shutdown()) + .map( + operation -> + operation.join( + 10, + TimeUnit + .SECONDS)) + .toList()); + BraintrustLogger.debug( + "otel shutdown complete. Flush done: %s, Flush successful: %s" + .formatted( + result.isDone(), result.isSuccess())); })); } } diff --git a/src/main/java/dev/braintrust/trace/DebugSpanExporter.java b/src/main/java/dev/braintrust/trace/DebugSpanExporter.java index 2e4b5a1..e93bd52 100644 --- a/src/main/java/dev/braintrust/trace/DebugSpanExporter.java +++ b/src/main/java/dev/braintrust/trace/DebugSpanExporter.java @@ -4,7 +4,6 @@ import io.opentelemetry.sdk.common.CompletableResultCode; import io.opentelemetry.sdk.trace.data.SpanData; import io.opentelemetry.sdk.trace.export.SpanExporter; - import java.util.Collection; // TODO -- is this just for unit testing? Make package-private? From 482f6df3ef449350199f80924f33c77bf80159f1 Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Mon, 22 Sep 2025 13:19:50 -0600 Subject: [PATCH 3/5] simplify CI --- .github/workflows/ci.yml | 46 +++++++++------------------------------- build.gradle | 2 +- 2 files changed, 11 insertions(+), 37 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 040197e..0f0e12f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -8,56 +8,30 @@ on: jobs: build: - runs-on: ubuntu-latest + # we want to run ubuntu-latest but we'll pin to a specific version so CI is reproducable + runs-on: ubuntu-24.04 strategy: matrix: - java: [ '17', '21' ] - + java: [ '17' ] # Java is + steps: - uses: actions/checkout@v4 - + - name: Set up JDK ${{ matrix.java }} uses: actions/setup-java@v4 with: java-version: ${{ matrix.java }} distribution: 'temurin' - + - name: Grant execute permission for gradlew run: chmod +x gradlew - - - name: Check code formatting - run: ./gradlew spotlessCheck - - - name: Build with Gradle - run: ./gradlew build - - - name: Run tests - run: ./gradlew test - + + - name: Run all checks and tests + run: ./gradlew compileTestJava check + - name: Upload test results uses: actions/upload-artifact@v4 if: always() with: name: test-results-java-${{ matrix.java }} path: build/test-results/test/ - - lint: - runs-on: ubuntu-latest - - steps: - - uses: actions/checkout@v4 - - - name: Set up JDK 17 - uses: actions/setup-java@v4 - with: - java-version: '17' - distribution: 'temurin' - - - name: Grant execute permission for gradlew - run: chmod +x gradlew - - - name: Check code formatting - run: ./gradlew spotlessCheck - - - name: Run static analysis - run: ./gradlew check \ No newline at end of file diff --git a/build.gradle b/build.gradle index 6bf7b9d..3d3be50 100644 --- a/build.gradle +++ b/build.gradle @@ -6,7 +6,7 @@ plugins { } group = 'dev.braintrust' -version = '0.1.0-SNAPSHOT' +version = '0.0.1-SNAPSHOT' java { sourceCompatibility = JavaVersion.VERSION_17 From 704ac4f76fa1b9b5ead0b46d5e0e267484954b10 Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Mon, 22 Sep 2025 13:53:08 -0600 Subject: [PATCH 4/5] simple config test --- .../braintrust/config/BraintrustConfig.java | 28 ++++++++++--------- .../config/BraintrustConfigTest.java | 22 ++++++++++++++- 2 files changed, 36 insertions(+), 14 deletions(-) diff --git a/src/main/java/dev/braintrust/config/BraintrustConfig.java b/src/main/java/dev/braintrust/config/BraintrustConfig.java index 924b633..34aa299 100644 --- a/src/main/java/dev/braintrust/config/BraintrustConfig.java +++ b/src/main/java/dev/braintrust/config/BraintrustConfig.java @@ -7,10 +7,7 @@ import java.util.Optional; import javax.annotation.Nullable; -/** - * Configuration for Braintrust SDK. Uses builder pattern with sensible defaults and environment - * variable support. - */ +/** Configuration for Braintrust SDK */ public final class BraintrustConfig { private static final String DEFAULT_PROJECT_NAME = "default-java-project"; private static final String DEFAULT_API_URL = "https://api.braintrust.dev"; @@ -40,6 +37,15 @@ private BraintrustConfig(Builder builder) { this.enableTraceConsoleLog = builder.enableTraceConsoleLog; this.debug = builder.debug; this.requestTimeout = builder.requestTimeout; + if ((null == defaultProjectId || "".equalsIgnoreCase(defaultProjectId.trim())) + && (null == defaultProjectName || "".equalsIgnoreCase(defaultProjectName.trim()))) { + // NOTE: this should not happen, + // but if someone happens to export their default project to the empty string and + // does not set a default project ID we don't have a valid parent for otel data. + throw new RuntimeException( + "Missing required envars. Please export BRAINTRUST_DEFAULT_PROJECT_ID or" + + " BRAINTRUST_DEFAULT_PROJECT"); + } } public String apiKey() { @@ -151,15 +157,6 @@ private Builder() { this.defaultProjectId = getEnv("BRAINTRUST_DEFAULT_PROJECT_ID", null); this.defaultProjectName = getEnv("BRAINTRUST_DEFAULT_PROJECT", DEFAULT_PROJECT_NAME).trim(); - if ((null == defaultProjectId || "".equalsIgnoreCase(defaultProjectId.trim())) - && "".equalsIgnoreCase(defaultProjectName)) { - // NOTE: this should not happen, - // but if someone happens to export their default project to the empty string and - // does not set a default project ID we don't have a valid parent for otel data. - throw new RuntimeException( - "Missing required envars. Please export BRAINTRUST_DEFAULT_PROJECT_ID or" - + " BRAINTRUST_DEFAULT_PROJECT"); - } this.enableTraceConsoleLog = Boolean.parseBoolean(getEnv("BRAINTRUST_ENABLE_TRACE_CONSOLE_LOG", "false")); this.debug = Boolean.parseBoolean(getEnv("BRAINTRUST_DEBUG", "false")); @@ -195,6 +192,11 @@ public Builder defaultProjectId(String projectId) { return this; } + public Builder defaultProjectName(String projectName) { + this.defaultProjectName = projectName; + return this; + } + public Builder enableTraceConsoleLog(boolean enable) { this.enableTraceConsoleLog = enable; return this; diff --git a/src/test/java/dev/braintrust/config/BraintrustConfigTest.java b/src/test/java/dev/braintrust/config/BraintrustConfigTest.java index bfece25..91487a3 100644 --- a/src/test/java/dev/braintrust/config/BraintrustConfigTest.java +++ b/src/test/java/dev/braintrust/config/BraintrustConfigTest.java @@ -1,8 +1,28 @@ package dev.braintrust.config; +import static org.junit.jupiter.api.Assertions.*; + import org.junit.jupiter.api.Test; class BraintrustConfigTest { @Test - void todo() {} + void throwsWithoutRequiredFields() { + assertThrows( + Exception.class, // don't care + () -> BraintrustConfig.builder().apiKey(null).build(), + "API key required"); + + assertThrows( + Exception.class, // don't care + () -> + BraintrustConfig.builder() + .apiKey("foo") + .defaultProjectId(null) + .defaultProjectName("") + .build(), + "Project info required"); + assertDoesNotThrow( + () -> BraintrustConfig.builder().apiKey("foobar").build(), + "Only an API key is required"); + } } From 514fef047f318237cee764dcb527e67d4045fb32 Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Mon, 22 Sep 2025 18:21:03 -0600 Subject: [PATCH 5/5] clean up comment --- .../java/dev/braintrust/trace/BraintrustSpanProcessor.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java b/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java index 23d6553..ad708ea 100644 --- a/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java +++ b/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java @@ -72,7 +72,7 @@ public void onStart(@NotNull Context parentContext, ReadWriteSpan span) { // Check if parent context has Braintrust attributes first var btContext = BraintrustContext.fromContext(parentContext); if (btContext == null) { - // Use default project ID if configured + // Get parent from the config if otel doesn't have it config.getBraintrustParentValue() .ifPresent( parentValue -> { @@ -83,7 +83,6 @@ public void onStart(@NotNull Context parentContext, ReadWriteSpan span) { span.getName()); }); } else { - // Can braintrust context have project_name instead? btContext .projectId() .ifPresent(