diff --git a/README.md b/README.md index 15363f3..8b616be 100644 --- a/README.md +++ b/README.md @@ -2,35 +2,39 @@ An OpenTelemetry-based Braintrust SDK for Java 17 and up. Contains instrumentation for major AI vendors (OpenAI, Anthropic, etc) -# Basic Usage +**WARNING (HERE BE DRAGONS):** *This is an experimental SDK for Braintrust. APIs are in active development and cannot be relied upon.* -## Using the SDK in your code +# Quickstart (Running examples and reporting data to Braintrust) -Add to your `build.gradle`: +Prerequisites: +- A Braintrust account and a valid BRAINTRUST_API_KEY +- Java 17 (or greater): + - macOS: `brew install openjdk@17` + - Ubuntu: `sudo apt install openjdk-17-jdk` -```gradle -dependencies { - implementation 'dev.braintrust:braintrust-x-java:0.1.0' -} +Then, from the repo root: +``` +# assumes you have BRAINTRUST_API_KEY exported +./gradlew :examples:runSimpleOpenTelemetry ``` -## Running examples +If you wish to see all examples run: +``` +./gradlew :examples:tasks --group="Braintrust SDK Examples" +``` -NOTE: if you wish to develop the SDK you can skip this section (developer setup includes a better way to set up the JDK). +Source code for all examples can be found here: [./examples/src/main/java/dev/braintrust/examples](./examples/src/main/java/dev/braintrust/examples) -First, install java 17 or greater +If you wish to hack around with the examples you can modify source then re-run the gradle task. -- **Java 17** (LTS) - Required for building and running - - macOS: `brew install openjdk@17` - - Ubuntu/Debian: `sudo apt install openjdk-17-jdk` - - Windows: Download from [Adoptium](https://adoptium.net/) +# Using the SDK in your code -``` -# assumes you have BRAINTRUST_API_KEY and OPENAI_API_KEY exported -./gradlew :examples:runOpenAIInstrumentation -``` +TODO: update this section when we have published our first release + +# Developer Setup + +This section documents how to set up your machine to develop the SDK. This is not required if you simply wish to use the SDK or run examples. -# Developer Docs ## Setup TODO -- document sdkman, gradle, etc diff --git a/examples/README.md b/examples/README.md deleted file mode 100644 index 5ffeb69..0000000 --- a/examples/README.md +++ /dev/null @@ -1,319 +0,0 @@ -# Braintrust Java SDK Examples - -This example demonstrates how to use Braintrust with OpenTelemetry to instrument OpenAI API calls in Java. It showcases various features including streaming/non-streaming completions, custom scoring, error handling, and complex multi-step operations. - -## Prerequisites - -1. **Java 17 or higher** - The example uses modern Java features like records and var -2. **Gradle** - For building and running the example -3. **API Keys** - You'll need both OpenAI and Braintrust API keys - -## Setup - -### 1. Get Your API Keys - -#### OpenAI API Key -- Sign up or log in at [OpenAI Platform](https://platform.openai.com/) -- Navigate to [API Keys](https://platform.openai.com/api-keys) -- Create a new secret key and save it securely - -#### Braintrust API Key -- Sign up or log in at [Braintrust](https://www.braintrust.dev/) -- Go to your project settings -- Create a new API key and save it securely - -### 2. Set Environment Variables - -Export your API keys as environment variables: - -```bash -export OPENAI_API_KEY="sk-..." -export BRAINTRUST_API_KEY="sk-..." -``` - -For Windows: -```cmd -set OPENAI_API_KEY=sk-... -set BRAINTRUST_API_KEY=sk-... -``` - -### 3. Build the Project - -From the root directory of the braintrust-x-java project: - -```bash -./gradlew build -``` - -## Available Examples - -### 1. OpenTelemetry Example (`OpenTelemetryExample.java`) -Full example using OpenAI's GPT-4o model with real API calls. - -```bash -./gradlew :examples:runOpenTelemetry -``` - -### 2. Simple Example (`SimpleExample.java`) -Demonstrates basic functionality without requiring API keys. - -```bash -./gradlew :examples:runSimple -``` - -### 3. Project Example (`ProjectExample.java`) -Shows how Braintrust organizes traces into projects. - -```bash -./gradlew :examples:runProject -``` - -## What the Example Demonstrates - -### 1. Basic Setup -The example shows how to initialize Braintrust with OpenTelemetry: - -```java -var config = BraintrustConfig.builder() - .apiKey(BRAINTRUST_API_KEY) - .build(); - -var openTelemetry = BraintrustTracing.quickstart(config, builder -> builder - .serviceName("opentelemetry-example") - .resourceAttribute("environment", "development") - .resourceAttribute("team", "engineering") - .exportInterval(Duration.ofSeconds(2)) -); -``` - -### 2. Non-Streaming Completion -Demonstrates a simple chat completion with instrumentation: - -```java -var response = OpenAIInterceptor.traceCompletion( - "chat.completion", - request, - requestExtractor, - responseExtractor, - () -> service.createChatCompletion(request) -); -``` - -Features shown: -- Creating spans for API calls -- Automatic token usage tracking -- Cost estimation -- Custom metadata and scoring - -### 3. Streaming Completion -Shows how to handle streaming responses: - -```java -var stream = OpenAIInterceptor.traceStreamingCompletion( - "chat.completion", - request, - requestExtractor, - () -> service.streamChatCompletion(request), - streamWrapper -); -``` - -Features shown: -- Wrapping streaming responses -- Collecting chunks for analysis -- Token usage estimation for streams - -### 4. Complex Multi-Step Operations -Demonstrates a complex workflow with nested spans: - -```java -// Root span for the entire operation -var rootSpan = tracer.spanBuilder("example.complex_operation") - .setSpanKind(SpanKind.INTERNAL) - .startSpan(); - -// Child spans for each step -var outlineSpan = tracer.spanBuilder("generate_outline") - .setSpanKind(SpanKind.INTERNAL) - .startSpan(); -``` - -Features shown: -- Nested span hierarchies -- Step-by-step operation tracking -- Custom scoring at each step -- Aggregated metrics - -### 5. Error Handling -Shows proper error handling and recovery: - -```java -try { - // Intentionally invalid request - var response = OpenAIInterceptor.traceCompletion(...); -} catch (Exception e) { - span.recordException(e); - span.setAttribute("error.handled", true); - // Recovery logic... -} -``` - -Features shown: -- Exception recording in spans -- Error attributes -- Recovery strategies -- Proper span status setting - -## Understanding the Output - -When you run the example, you'll see: - -1. **Console Output** - Real-time feedback showing: - - Generated haiku - - Streaming text as it arrives - - Story generation progress - - Error handling demonstration - -2. **Braintrust Dashboard** - After running, check your Braintrust dashboard to see: - - Trace visualization - - Token usage metrics - - Cost tracking - - Custom scores - - Error details - - Span hierarchy - -## Key Features Demonstrated - -### Custom Scoring -```java -BraintrustTracing.SpanUtils.addScore("haiku_quality", evaluateHaikuQuality(completion)); -``` - -### Metadata Addition -```java -currentSpan.setAttribute("example.haiku", completion); -currentSpan.setAttribute("example.quality_score", score); -``` - -### Token Usage Tracking -Automatically captured from OpenAI responses: -- Prompt tokens -- Completion tokens -- Total tokens -- Estimated cost - -### Error Recovery -```java -span.addEvent("Attempting recovery with valid parameters"); -// Recovery logic... -span.setAttribute("recovery.success", true); -``` - -## Customization - -You can customize the example by: - -1. **Changing Models** - Try different OpenAI models: - ```java - .model("gpt-4") // or "gpt-3.5-turbo-16k" - ``` - -2. **Adjusting Parameters**: - ```java - .temperature(0.9) // More creative - .maxTokens(500) // Longer responses - ``` - -3. **Adding Custom Attributes**: - ```java - span.setAttribute("custom.metric", value); - span.setAttribute("user.id", userId); - ``` - -4. **Implementing Custom Scoring**: - ```java - BraintrustTracing.SpanUtils.addScore("custom_score", calculateScore(response)); - ``` - -## Troubleshooting - -### Connection Issues -If you see connection errors: -- Check your internet connection -- Verify API keys are correctly set -- Ensure Braintrust API endpoint is accessible - -### Token Limit Errors -If you hit token limits: -- Reduce `maxTokens` in requests -- Use shorter prompts -- Consider using a model with higher limits - -### Export Delays -Spans are exported in batches. If you don't see traces immediately: -- Wait a few seconds for batch export -- Check the export interval setting -- Ensure the application runs long enough for export - -## Next Steps - -1. **Explore Other Examples** - Check out the evaluation and dataset examples -2. **Integrate Into Your App** - Use this pattern in your own applications -3. **Custom Instrumentation** - Create spans for your own operations -4. **Advanced Features** - Explore Braintrust's evaluation and experimentation features - -## Project Organization in Braintrust - -In Braintrust, all traces are organized into projects. The project determines where your traces appear in the Braintrust dashboard. - -### How Project Names are Determined - -1. **Default: Service Name** - If no project is specified, Braintrust uses the service name from OpenTelemetry: - ```java - BraintrustTracing.quickstart(config, builder -> builder - .serviceName("my-ai-app") // This becomes the project name - ) - ``` - -2. **Explicit Project ID** - Set a project for specific spans: - ```java - BraintrustTracing.SpanUtils.setParentProject("customer-support"); - ``` - -3. **Via Context** - Set a project for multiple spans: - ```java - var context = BraintrustContext.builder() - .projectId("recommendation-engine") - .build() - .storeInCurrent(); - ``` - -4. **Environment Variable** - Set `BRAINTRUST_PROJECT_ID` to override the default - -### Best Practices - -- Use meaningful project names that reflect your application's purpose -- Group related traces in the same project for easier analysis -- Use different projects for different environments (dev, staging, prod) -- Consider using experiments for A/B testing and model comparisons - -## Model Updates - -The examples use OpenAI's GPT-4o model, which offers: -- Better performance than GPT-3.5-turbo -- More accurate responses -- Support for vision and other multimodal features -- Pricing: $5/1M input tokens, $15/1M output tokens - -To use a different model, update the `ChatModel` parameter: -```java -.model(ChatModel.GPT_4O) // GPT-4o (recommended) -.model(ChatModel.GPT_4_TURBO) // GPT-4 Turbo -.model(ChatModel.GPT_3_5_TURBO) // GPT-3.5 Turbo -``` - -## Additional Resources - -- [Braintrust Documentation](https://docs.braintrust.dev/) -- [OpenTelemetry Java Guide](https://opentelemetry.io/docs/instrumentation/java/) -- [OpenAI API Reference](https://platform.openai.com/docs/api-reference) \ No newline at end of file diff --git a/examples/build.gradle b/examples/build.gradle index e684494..6159077 100644 --- a/examples/build.gradle +++ b/examples/build.gradle @@ -14,6 +14,8 @@ repositories { dependencies { implementation project(':') + // TODO: not sure why I have to do this. I would have expected this to come over in the project root transitive deps + implementation "io.opentelemetry:opentelemetry-exporter-otlp:${otelVersion}" // Official OpenAI Java SDK implementation 'com.openai:openai-java:2.8.1' @@ -23,33 +25,45 @@ dependencies { } application { - mainClass = 'dev.braintrust.examples.OpenTelemetryTracingExample' + mainClass = 'dev.braintrust.examples.SimpleOpenTelemetryExample' } -task runOpenTelemetryTracing(type: JavaExec) { +task runSimpleOpenTelemetry(type: JavaExec) { + group = 'Braintrust SDK Examples' + description = 'Run the simple OpenTelemetry example' classpath = sourceSets.main.runtimeClasspath - mainClass = 'dev.braintrust.examples.OpenTelemetryTracingExample' + mainClass = 'dev.braintrust.examples.SimpleOpenTelemetryExample' } -task runOpenAIInstrumentation(type: JavaExec) { - classpath = sourceSets.main.runtimeClasspath - mainClass = 'dev.braintrust.examples.OpenAIInstrumentationExample' - debugOptions { - enabled = true - port = 5566 - server = true - suspend = false - } +task runCustomOpenTelemetry(type: JavaExec) { + group = 'Braintrust SDK Examples' + description = 'Run the custom OpenTelemetry example' + classpath = sourceSets.main.runtimeClasspath + mainClass = 'dev.braintrust.examples.CustomOpenTelemetryExample' } +task runOpenAIInstrumentation(type: JavaExec) { + group = 'Braintrust SDK Examples' + description = 'Run the OpenAI instrumentation example. NOTE: this requires OPENAI_API_KEY to be exported and will make a small call to openai, using your tokens' + classpath = sourceSets.main.runtimeClasspath + mainClass = 'dev.braintrust.examples.OpenAIInstrumentationExample' + debugOptions { + enabled = true + port = 5566 + server = true + suspend = false + } +} task runExperiment(type: JavaExec) { - classpath = sourceSets.main.runtimeClasspath - mainClass = 'dev.braintrust.examples.ExperimentExample' - debugOptions { - enabled = true - port = 5566 - server = true - suspend = false - } + group = 'Braintrust SDK Examples' + description = 'Run the experiment example' + classpath = sourceSets.main.runtimeClasspath + mainClass = 'dev.braintrust.examples.ExperimentExample' + debugOptions { + enabled = true + port = 5566 + server = true + suspend = false + } } diff --git a/examples/src/main/java/dev/braintrust/examples/CustomOpenTelemetryExample.java b/examples/src/main/java/dev/braintrust/examples/CustomOpenTelemetryExample.java new file mode 100644 index 0000000..49aeb47 --- /dev/null +++ b/examples/src/main/java/dev/braintrust/examples/CustomOpenTelemetryExample.java @@ -0,0 +1,84 @@ +package dev.braintrust.examples; + +import dev.braintrust.config.BraintrustConfig; +import dev.braintrust.log.BraintrustLogger; +import dev.braintrust.trace.BraintrustTracing; +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.OpenTelemetry; +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.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.common.CompletableResultCode; +import io.opentelemetry.sdk.logs.SdkLoggerProvider; +import io.opentelemetry.sdk.logs.export.BatchLogRecordProcessor; +import io.opentelemetry.sdk.metrics.SdkMeterProvider; +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. + */ +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()); + + // 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 + BraintrustTracing.enable(BraintrustConfig.fromEnvironment(), tracerBuilder, loggerBuilder, meterBuilder); + + var openTelemetry = OpenTelemetrySdk.builder() + .setTracerProvider(tracerBuilder.build()) + .setLoggerProvider(loggerBuilder.build()) + .setMeterProvider(meterBuilder.build()) + .build(); + GlobalOpenTelemetry.set(openTelemetry); + registerShutdownHook(openTelemetry); + var braintrustTracer = BraintrustTracing.getTracer(openTelemetry); + + // Now otel data will be exported to the local collector AND to braintrust + var span = braintrustTracer.spanBuilder("hello-java").startSpan(); + try (var ignored = span.makeCurrent()) { + System.out.println("Performing simple operation..."); + span.setAttribute("some boolean attribute", true); + Thread.sleep(100); // Not required. This is just to make the span look interesting + } finally { + span.end(); + } + System.out.println("custom example completed!"); + } + + private static void registerShutdownHook(OpenTelemetrySdk otel) { + 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(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 b7c7438..295d5c2 100644 --- a/examples/src/main/java/dev/braintrust/examples/ExperimentExample.java +++ b/examples/src/main/java/dev/braintrust/examples/ExperimentExample.java @@ -5,6 +5,8 @@ import com.openai.models.chat.completions.ChatCompletionCreateParams; import dev.braintrust.config.BraintrustConfig; import dev.braintrust.eval.Eval; +import dev.braintrust.eval.EvalCase; +import dev.braintrust.eval.Scorer; import dev.braintrust.instrumentation.openai.BraintrustOpenAI; import dev.braintrust.trace.BraintrustTracing; @@ -13,7 +15,7 @@ public class ExperimentExample { public static void main(String[] args) throws Exception { var config = BraintrustConfig.fromEnvironment(); - var openTelemetry = BraintrustTracing.quickstart(config); + var openTelemetry = BraintrustTracing.quickstart(config, true); var openAIClient = BraintrustOpenAI.wrapOpenAI(openTelemetry, OpenAIOkHttpClient.fromEnv()); Function getFoodType = (String food) -> { @@ -29,17 +31,17 @@ public static void main(String[] args) throws Exception { return response.choices().get(0).message().content().orElse("").toLowerCase(); }; - var eval = Eval.builder() + 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(Eval.EvalCase.of("strawberry", "fruit"), - Eval.EvalCase.of("asparagus", "vegetable"), - Eval.EvalCase.of("apple", "fruit"), - Eval.EvalCase.of("banana", "fruit")) + .cases(EvalCase.of("strawberry", "fruit"), + EvalCase.of("asparagus", "vegetable"), + EvalCase.of("apple", "fruit"), + EvalCase.of("banana", "fruit")) .task(getFoodType) - .scorers(Eval.Scorer.of("fruit_scorer", result -> "fruit".equals(result) ? 1.0 : 0.0), - Eval.Scorer.of("vegetable_scorer", result -> "vegetable".equals(result) ? 1.0 : 0.0)) + .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(result.createReportString()); diff --git a/examples/src/main/java/dev/braintrust/examples/OpenAIInstrumentationExample.java b/examples/src/main/java/dev/braintrust/examples/OpenAIInstrumentationExample.java index 75cc0d2..93a2efe 100644 --- a/examples/src/main/java/dev/braintrust/examples/OpenAIInstrumentationExample.java +++ b/examples/src/main/java/dev/braintrust/examples/OpenAIInstrumentationExample.java @@ -4,7 +4,6 @@ import com.openai.client.okhttp.OpenAIOkHttpClient; import com.openai.models.ChatModel; import com.openai.models.chat.completions.ChatCompletionCreateParams; -import dev.braintrust.config.BraintrustConfig; import dev.braintrust.instrumentation.openai.BraintrustOpenAI; import dev.braintrust.trace.BraintrustTracing; @@ -13,13 +12,15 @@ */ public class OpenAIInstrumentationExample { public static void main(String[] args) throws Exception { - var config = BraintrustConfig.builder().build(); - var openTelemetry = BraintrustTracing.quickstart(config); + if (null == System.getenv("OPENAI_API_KEY")) { + System.err.println("\nWARNING envar OPEN_AI_API_KEY not found. This example will likely fail.\n"); + } + var openTelemetry = BraintrustTracing.quickstart(); var tracer = BraintrustTracing.getTracer(openTelemetry); OpenAIClient openAIClient = BraintrustOpenAI.wrapOpenAI(openTelemetry, OpenAIOkHttpClient.fromEnv()); var rootSpan = tracer.spanBuilder("java-braintrust-example").startSpan(); try (var ignored = rootSpan.makeCurrent()) { - Thread.sleep(70); + Thread.sleep(70); // Not required. This is just to make the span look interesting var request = ChatCompletionCreateParams.builder() .model(ChatModel.GPT_4O_MINI) @@ -31,7 +32,7 @@ public static void main(String[] args) throws Exception { var response = openAIClient.chat().completions().create(request); System.out.println("~~~ GOT RESPONSE: " + response); openAIClient.completions(); - Thread.sleep(30); + Thread.sleep(30); // Not required. This is just to make the span look interesting } finally { rootSpan.end(); } diff --git a/examples/src/main/java/dev/braintrust/examples/OpenTelemetryTracingExample.java b/examples/src/main/java/dev/braintrust/examples/SimpleOpenTelemetryExample.java similarity index 54% rename from examples/src/main/java/dev/braintrust/examples/OpenTelemetryTracingExample.java rename to examples/src/main/java/dev/braintrust/examples/SimpleOpenTelemetryExample.java index 28a09d6..c42c49c 100644 --- a/examples/src/main/java/dev/braintrust/examples/OpenTelemetryTracingExample.java +++ b/examples/src/main/java/dev/braintrust/examples/SimpleOpenTelemetryExample.java @@ -1,25 +1,20 @@ package dev.braintrust.examples; -import dev.braintrust.config.BraintrustConfig; import dev.braintrust.trace.BraintrustTracing; -/** - * Basic OTel tracing example - */ -public class OpenTelemetryTracingExample { +public class SimpleOpenTelemetryExample { public static void main(String[] args) throws Exception { - var config = BraintrustConfig.builder().build(); - var openTelemetry = BraintrustTracing.quickstart(config); + var openTelemetry = BraintrustTracing.quickstart(); var tracer = BraintrustTracing.getTracer(openTelemetry); var span = tracer.spanBuilder("hello-java").startSpan(); try (var ignored = span.makeCurrent()) { System.out.println("Performing simple operation..."); - span.setAttribute("is-java-span", true); - Thread.sleep(100); + span.setAttribute("some boolean attribute", true); + Thread.sleep(100); // Not required. This is just to make the span look interesting } finally { span.end(); } - System.out.println("example completed!"); + System.out.println("simple example completed!"); } } diff --git a/src/main/java/dev/braintrust/eval/Eval.java b/src/main/java/dev/braintrust/eval/Eval.java index a238445..4e5a5a8 100644 --- a/src/main/java/dev/braintrust/eval/Eval.java +++ b/src/main/java/dev/braintrust/eval/Eval.java @@ -3,40 +3,36 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import dev.braintrust.api.BraintrustApiClient; -import dev.braintrust.claude.trace.BraintrustSpanProcessor; import dev.braintrust.config.BraintrustConfig; import dev.braintrust.trace.BraintrustContext; +import dev.braintrust.trace.BraintrustSpanProcessor; import dev.braintrust.trace.BraintrustTracing; import io.opentelemetry.api.trace.SpanKind; import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.context.Context; import javax.annotation.Nonnull; import javax.annotation.Nullable; import java.util.*; -import java.util.concurrent.ExecutionException; -import java.util.function.BiFunction; import java.util.function.Function; /** - * Main evaluation framework for testing AI models. Uses Java generics for type safety and streams - * for functional processing. + * An evaluation framework for testing AI models. * * @param The type of input data for the evaluation - * @param The type of output produced by the task + * @param The type of output produced by the task */ -public final class Eval { +public final class Eval { private static final ObjectMapper JSON_MAPPER = new com.fasterxml.jackson.databind.ObjectMapper(); private final @Nonnull String experimentName; private final @Nonnull String projectName; private final @Nonnull BraintrustConfig config; private final @Nonnull BraintrustApiClient client; private final @Nonnull Tracer tracer; - private final @Nonnull List> evalCases; - private final @Nonnull Task task; - private final @Nonnull List> scorers; + private final @Nonnull List> evalCases; + private final @Nonnull Task task; + private final @Nonnull List> scorers; - private Eval(Builder builder) { + private Eval(Builder builder) { this.experimentName = builder.experimentName; this.projectName = Objects.requireNonNull(builder.projectName); this.config = Objects.requireNonNull(builder.config); @@ -48,16 +44,16 @@ private Eval(Builder builder) { } /** Runs the evaluation and returns results. */ - public EvalResult run() { + public Result run() { var experiment = client.createExperiment(new BraintrustApiClient.CreateExperimentRequest(projectName, experimentName, Optional.empty(), Optional.empty())); var experimentID = experiment.id(); var evalCaseResults = evalCases.stream() .map(evalCase -> evalOne(experimentID, evalCase)) .toList(); - return new EvalResult<>(this); + return new Result(); } - private EvalCaseResult 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 @@ -72,7 +68,7 @@ private EvalCaseResult evalOne(String experimentId, Eva .setAttribute(BraintrustSpanProcessor.PARENT_TYPE, "experiment") .startSpan(); try (var rootScope = BraintrustContext.forExperiment(experimentId, rootSpan).makeCurrent()) { - final RESULT result; + final OUTPUT result; { // run task var taskSpan = tracer.spanBuilder("task") .setAttribute(BraintrustSpanProcessor.PARENT, "experiment_id:" + experimentId) @@ -116,7 +112,7 @@ private EvalCaseResult evalOne(String experimentId, Eva scoreSpan.end(); } } - return new EvalCaseResult<>(evalCase, result); + return new EvalCase.Result<>(evalCase, result); } finally { rootSpan.end(); } @@ -125,18 +121,12 @@ private EvalCaseResult evalOne(String experimentId, Eva /** * Results of all eval cases of an experiment. */ - public static class EvalResult { - private final Eval eval; - - private EvalResult(Eval eval) { - this.eval = eval; - } - + public class Result { public String createReportString() { try { - var project = eval.client.getProject(eval.projectName).get().orElseThrow(); + var project = client.getProject(projectName).get().orElseThrow(); // FIXME: don't hardcore the first part of the url. Need to check for staging and also get the actual org - var experimentUrl = "https://www.braintrust.dev/app/braintrustdata.com/p/" + project.name() + "/experiments/" + eval.experimentName; + var experimentUrl = "https://www.braintrust.dev/app/braintrustdata.com/p/" + project.name() + "/experiments/" + experimentName; return "Experiment complete. View results in braintrust: " + experimentUrl; } catch (Exception e) { throw new RuntimeException(e); @@ -144,65 +134,22 @@ public String createReportString() { } } - public record EvalCaseResult(EvalCase evalCase, RESULT result) {} - - public record EvalCase(INPUT input, EXPECTED expected, @Nonnull List tags, @Nonnull Metadata metadata) { - public static EvalCase of(INPUT input, EXPECTED expected) { - return of(input, expected, List.of(), new Metadata()); - } - public static EvalCase of(INPUT input, EXPECTED expected, @Nonnull List tags, @Nonnull Metadata metadata) { - return new EvalCase<>(input, expected, tags, metadata); - } - } - - public record Metadata() { - // TODO implement: arbitrary map of string->json - } - - public interface Task { - RESULT apply(EvalCase evalCase); - } - - public interface Scorer { - String getName(); - - double score(EvalCase evalCase, RESULT result); - - static Scorer of(String scorerName, BiFunction, RESULT, Double> scorerFn) { - return new Scorer<>() { - @Override - public String getName() { - return scorerName; - } - - @Override - public double score(EvalCase evalCase, RESULT result) { - return scorerFn.apply(evalCase, result); - } - }; - } - - static Scorer of(String scorerName, Function scorerFn) { - return of(scorerName, (evalCase, result) -> scorerFn.apply(result)); - } - } - /** Creates a new eval builder. */ - public static Builder builder() { + public static Builder builder() { return new Builder<>(); } /** Builder for creating evaluations with fluent API. */ - public static final class Builder { + public static final class Builder { private @Nonnull String experimentName = "unnamed-java-eval"; private @Nullable BraintrustConfig config; private @Nullable String projectName; private @Nullable Tracer tracer = null; - private @Nonnull List> evalCases = List.of(); - private @Nullable Task task; - private @Nonnull List> scorers = List.of(); + private @Nonnull List> evalCases = List.of(); + private @Nullable Task task; + private @Nonnull List> scorers = List.of(); - public Eval build() { + public Eval build() { if (config == null) { config = BraintrustConfig.fromEnvironment(); } @@ -222,43 +169,43 @@ public Eval build() { return new Eval<>(this); } - public Builder name(String name) { + public Builder name(String name) { this.experimentName = name; return this; } - public Builder projectName(@Nonnull String projectName) { + public Builder projectName(@Nonnull String projectName) { this.projectName = Objects.requireNonNull(projectName); return this; } - public Builder config(BraintrustConfig config) { + public Builder config(BraintrustConfig config) { this.config = config; return this; } - public Builder tracer(Tracer tracer) { + public Builder tracer(Tracer tracer) { this.tracer = tracer; return this; } @SafeVarargs - public final Builder cases(EvalCase... cases) { + public final Builder cases(EvalCase... cases) { this.evalCases = List.of(cases); return this; } - public Builder task(Task task) { + public Builder task(Task task) { this.task = task; return this; } - public Builder task(Function taskFn) { - return task((Task) evalCase -> taskFn.apply(evalCase.input())); + public Builder task(Function taskFn) { + return task((Task) evalCase -> taskFn.apply(evalCase.input())); } @SafeVarargs - public final Builder scorers(Scorer... scorers) { + public final Builder scorers(Scorer... scorers) { this.scorers = List.of(scorers); return this; } diff --git a/src/main/java/dev/braintrust/eval/EvalCase.java b/src/main/java/dev/braintrust/eval/EvalCase.java new file mode 100644 index 0000000..1c26b19 --- /dev/null +++ b/src/main/java/dev/braintrust/eval/EvalCase.java @@ -0,0 +1,36 @@ +package dev.braintrust.eval; + +import javax.annotation.Nonnull; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +/** + * A single test case in an LLM eval. + * + * @param input test input + * @param expected expected value + * @param tags additional tags to apply in the braintrust UI + * @param metadata additional key-value data to apply in the braintrust UI + */ +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"); + } + if (!tags.isEmpty()) { + throw new RuntimeException("TODO: tags support not yet implemented"); + } + } + + public static EvalCase of(INPUT input, OUTPUT expected) { + return of(input, expected, List.of(), Map.of()); + } + + public static EvalCase of(INPUT input, OUTPUT expected, @Nonnull List tags, @Nonnull Map metadata) { + return new EvalCase<>(input, expected, tags, metadata); + } + + public record Result(EvalCase evalCase, OUTPUT result) {} +} diff --git a/src/main/java/dev/braintrust/eval/Scorer.java b/src/main/java/dev/braintrust/eval/Scorer.java new file mode 100644 index 0000000..2292fda --- /dev/null +++ b/src/main/java/dev/braintrust/eval/Scorer.java @@ -0,0 +1,34 @@ +package dev.braintrust.eval; + +import java.util.function.BiFunction; +import java.util.function.Function; + +/** + * 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 + */ +public interface Scorer { + String getName(); + + double score(EvalCase evalCase, OUTPUT result); + + static Scorer of(String scorerName, BiFunction, OUTPUT, Double> scorerFn) { + return new Scorer<>() { + @Override + public String getName() { + return scorerName; + } + + @Override + public double score(EvalCase evalCase, OUTPUT result) { + return scorerFn.apply(evalCase, result); + } + }; + } + + static Scorer of(String scorerName, Function scorerFn) { + return of(scorerName, (evalCase, result) -> scorerFn.apply(result)); + } +} diff --git a/src/main/java/dev/braintrust/eval/Task.java b/src/main/java/dev/braintrust/eval/Task.java new file mode 100644 index 0000000..0807eae --- /dev/null +++ b/src/main/java/dev/braintrust/eval/Task.java @@ -0,0 +1,11 @@ +package dev.braintrust.eval; + +/** + * A task executes an EvalCase and returns a result + * + * @param type of the input data + * @param type of the output data + */ +public interface Task { + OUTPUT apply(EvalCase evalCase); +} diff --git a/src/main/java/dev/braintrust/trace/BraintrustTracing.java b/src/main/java/dev/braintrust/trace/BraintrustTracing.java index 67d3456..4d2481e 100644 --- a/src/main/java/dev/braintrust/trace/BraintrustTracing.java +++ b/src/main/java/dev/braintrust/trace/BraintrustTracing.java @@ -6,23 +6,34 @@ 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; /** * Main entry point for Braintrust tracing setup. Provides convenient methods to initialize @@ -30,33 +41,12 @@ */ public final class BraintrustTracing { private static final String INSTRUMENTATION_NAME = "braintrust-java"; - private static final String INSTRUMENTATION_VERSION = "0.1.0"; + private static final String INSTRUMENTATION_VERSION = "0.0.1"; 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 -> {}); - } - - // TODO -- remove this constructor? - /** 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 getTracer(GlobalOpenTelemetry.get()); @@ -64,159 +54,97 @@ public static Tracer getTracer() { /** Gets a tracer from a specific OpenTelemetry instance. */ public static Tracer getTracer(OpenTelemetry openTelemetry) { - // TODO: will otel return some dummy object if we don't have a braintrust tracer? Are we call with that? 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(); - - var logExporter = new BraintrustLogExporter(config); - var loggerProvider = SdkLoggerProvider.builder() - .addLogRecordProcessor(BatchLogRecordProcessor.builder(logExporter).build()) - .build(); - - // Create OpenTelemetry instance - var openTelemetry = OpenTelemetrySdk.builder() - .setTracerProvider(tracerProvider) - .setLoggerProvider(loggerProvider) - .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. Force-Flushing all otel data."); - var result = CompletableResultCode.ofAll(List.of( - openTelemetry.getSdkLoggerProvider().forceFlush().join(10, TimeUnit.SECONDS), - openTelemetry.getSdkMeterProvider().forceFlush().join(10, TimeUnit.SECONDS), - openTelemetry.getSdkTracerProvider().forceFlush().join(10, TimeUnit.SECONDS) - )); - BraintrustLogger.debug("tracer shutdown complete. Flush successful? " + result.isSuccess()); - tracerProvider.shutdown(); - })); - - return 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. + */ + public static OpenTelemetry quickstart() { + return quickstart(BraintrustConfig.fromEnvironment(), true); } - /** 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); + /** + * 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) { + 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(); + if (registerGlobal) { + GlobalOpenTelemetry.set(openTelemetry); + BraintrustLogger.debug("Registered OpenTelemetry globally"); } + return openTelemetry; + } - /** 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); - } + /** + * 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 + final Map resourceAttributes = new HashMap<>(); + final String serviceName = "braintrust-app"; + final Duration exportInterval = Duration.ofSeconds(5); + final int maxQueueSize = 2048; + final int maxExportBatchSize = 512; + BraintrustLogger.info("Initializing Braintrust OpenTelemetry with service={}", serviceName); + + // 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(); + + // spans + 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(); + 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 + + 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())); + })); } }