Skip to content

Commit 91263fe

Browse files
realarkrealark
authored
Merge pull request #52 from braintrustdata/ark/update-spring-ai-example
add oai and anthropic to spring ai example
2 parents e6381d6 + be2a102 commit 91263fe

1 file changed

Lines changed: 80 additions & 79 deletions

File tree

‎examples/src/main/java/dev/braintrust/examples/SpringAIExample.java‎

Lines changed: 80 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import dev.braintrust.Braintrust;
55
import dev.braintrust.config.BraintrustConfig;
66
import dev.braintrust.instrumentation.genai.BraintrustGenAI;
7+
import dev.braintrust.instrumentation.springai.v1_0_0.BraintrustSpringAI;
78
import io.opentelemetry.api.OpenTelemetry;
89
import io.opentelemetry.api.trace.Span;
910
import io.opentelemetry.api.trace.Tracer;
@@ -18,6 +19,7 @@
1819
import org.springframework.ai.openai.OpenAiChatModel;
1920
import org.springframework.ai.openai.OpenAiChatOptions;
2021
import org.springframework.ai.openai.api.OpenAiApi;
22+
import org.springframework.beans.factory.annotation.Qualifier;
2123
import org.springframework.boot.CommandLineRunner;
2224
import org.springframework.boot.SpringApplication;
2325
import org.springframework.boot.autoconfigure.SpringBootApplication;
@@ -39,19 +41,29 @@ public static void main(String[] args) {
3941
}
4042

4143
@Bean
42-
public CommandLineRunner run(ChatModel chatModel, Tracer tracer, Braintrust braintrust) {
44+
public CommandLineRunner run(
45+
@Qualifier("openAIChatModel") ChatModel openAIChatModel,
46+
@Qualifier("anthropicChatModel") ChatModel anthropicChatModel,
47+
@Qualifier("googleChatModel") ChatModel googleChatModel,
48+
Tracer tracer,
49+
Braintrust braintrust) {
4350
return args -> {
4451
Span rootSpan = tracer.spanBuilder("spring-ai-example").startSpan();
4552
try (Scope scope = rootSpan.makeCurrent()) {
4653
System.out.println("\n=== Running Spring Boot Example ===\n");
4754

4855
// Make a simple chat call
4956
var prompt = new Prompt("what's the name of the most popular java DI framework?");
50-
var response = chatModel.call(prompt);
57+
var oaiResponse = openAIChatModel.call(prompt);
58+
var anthropicResponse = anthropicChatModel.call(prompt);
59+
var googleResponse = googleChatModel.call(prompt);
5160

5261
System.out.println(
53-
"~~~ SPRING AI CHAT RESPONSE: %s\n"
54-
.formatted(response.getResult().getOutput().getText()));
62+
"~~~ SPRING AI CHAT RESPONSES: \noat: %s\nanthropic: %s\ngoogle: %s\n"
63+
.formatted(
64+
oaiResponse.getResult().getOutput().getText(),
65+
anthropicResponse.getResult().getOutput().getText(),
66+
googleResponse.getResult().getOutput().getText()));
5567
} finally {
5668
rootSpan.end();
5769
}
@@ -84,84 +96,73 @@ public Tracer tracer(OpenTelemetry openTelemetry) {
8496
}
8597

8698
@Bean
87-
public String aiProvider() {
88-
var provider = System.getenv("SPRING_AI_EXAMPLE_PROVIDER");
89-
if (provider == null || provider.isBlank()) {
90-
return "openai";
99+
public ChatModel openAIChatModel(OpenTelemetry openTelemetry) {
100+
if (null == System.getenv("OPENAI_API_KEY")) {
101+
System.err.println(
102+
"\n"
103+
+ "WARNING: OPENAI_API_KEY not found. This example will likely"
104+
+ " fail.\n"
105+
+ "Set it with: export OPENAI_API_KEY='your-key'\n");
91106
}
92-
return switch (provider) {
93-
case "openai", "anthropic", "google" -> provider;
94-
default ->
95-
throw new RuntimeException(
96-
"unsupported SPRING_AI_EXAMPLE_PROVIDER: '%s'. Allowed values: openai, anthropic, google"
97-
.formatted(provider));
98-
};
107+
return BraintrustSpringAI.wrap(
108+
openTelemetry,
109+
OpenAiChatModel.builder()
110+
.openAiApi(
111+
OpenAiApi.builder()
112+
.apiKey(System.getenv("OPENAI_API_KEY"))
113+
.build())
114+
.defaultOptions(
115+
OpenAiChatOptions.builder()
116+
.model("gpt-4o-mini")
117+
.temperature(0.0)
118+
.maxTokens(50)
119+
.build()))
120+
.build();
99121
}
100122

101123
@Bean
102-
public ChatModel chatModel(String aiProvider, OpenTelemetry openTelemetry) {
103-
return switch (aiProvider) {
104-
case "openai" -> {
105-
if (null == System.getenv("OPENAI_API_KEY")) {
106-
System.err.println(
107-
"\n"
108-
+ "WARNING: OPENAI_API_KEY not found. This example will likely"
109-
+ " fail.\n"
110-
+ "Set it with: export OPENAI_API_KEY='your-key'\n");
111-
}
112-
var openAiApi = OpenAiApi.builder().apiKey(System.getenv("OPENAI_API_KEY")).build();
113-
yield OpenAiChatModel.builder()
114-
.openAiApi(openAiApi)
115-
.defaultOptions(
116-
OpenAiChatOptions.builder()
117-
.model("gpt-4o-mini")
118-
.temperature(0.0)
119-
.maxTokens(50)
120-
.build())
121-
.build();
122-
}
123-
case "anthropic" -> {
124-
if (null == System.getenv("ANTHROPIC_API_KEY")) {
125-
System.err.println(
126-
"\n"
127-
+ "WARNING: ANTHROPIC_API_KEY not found. This example will"
128-
+ " likely fail.\n"
129-
+ "Set it with: export ANTHROPIC_API_KEY='your-key'\n");
130-
}
131-
var anthropicApi =
132-
AnthropicApi.builder().apiKey(System.getenv("ANTHROPIC_API_KEY")).build();
133-
yield AnthropicChatModel.builder()
134-
.anthropicApi(anthropicApi)
135-
.defaultOptions(
136-
AnthropicChatOptions.builder()
137-
.model("claude-3-haiku-20240307")
138-
.temperature(0.0)
139-
.maxTokens(50)
140-
.build())
141-
.build();
142-
}
143-
case "google" -> {
144-
if (null == System.getenv("GOOGLE_API_KEY")
145-
&& null == System.getenv("GEMINI_API_KEY")) {
146-
System.err.println(
147-
"\n"
148-
+ "WARNING: Neither GOOGLE_API_KEY nor GEMINI_API_KEY found. This"
149-
+ " example will likely fail.\n"
150-
+ "Set either: export GOOGLE_API_KEY='your-key' (recommended) or"
151-
+ " export GEMINI_API_KEY='your-key'\n");
152-
}
153-
Client genAIClient = BraintrustGenAI.wrap(openTelemetry, new Client.Builder());
154-
yield GoogleGenAiChatModel.builder()
155-
.genAiClient(genAIClient)
156-
.defaultOptions(
157-
GoogleGenAiChatOptions.builder()
158-
.model("gemini-2.0-flash-lite")
159-
.temperature(0.0)
160-
.maxOutputTokens(50)
161-
.build())
162-
.build();
163-
}
164-
default -> throw new RuntimeException("unsupported provider: " + aiProvider);
165-
};
124+
public ChatModel anthropicChatModel(OpenTelemetry openTelemetry) {
125+
if (null == System.getenv("ANTHROPIC_API_KEY")) {
126+
System.err.println(
127+
"\n"
128+
+ "WARNING: ANTHROPIC_API_KEY not found. This example will"
129+
+ " likely fail.\n"
130+
+ "Set it with: export ANTHROPIC_API_KEY='your-key'\n");
131+
}
132+
return BraintrustSpringAI.wrap(
133+
openTelemetry,
134+
AnthropicChatModel.builder()
135+
.anthropicApi(
136+
AnthropicApi.builder()
137+
.apiKey(System.getenv("ANTHROPIC_API_KEY"))
138+
.build())
139+
.defaultOptions(
140+
AnthropicChatOptions.builder()
141+
.model("claude-3-haiku-20240307")
142+
.temperature(0.0)
143+
.maxTokens(50)
144+
.build()))
145+
.build();
146+
}
147+
148+
@Bean
149+
public ChatModel googleChatModel(OpenTelemetry openTelemetry) {
150+
if (null == System.getenv("GOOGLE_API_KEY") && null == System.getenv("GEMINI_API_KEY")) {
151+
System.err.println(
152+
"\n"
153+
+ "WARNING: Neither GOOGLE_API_KEY nor GEMINI_API_KEY found. This"
154+
+ " example will likely fail.\n"
155+
+ "Set either: export GOOGLE_API_KEY='your-key' (recommended) or"
156+
+ " export GEMINI_API_KEY='your-key'\n");
157+
}
158+
return GoogleGenAiChatModel.builder()
159+
.genAiClient(BraintrustGenAI.wrap(openTelemetry, new Client.Builder()))
160+
.defaultOptions(
161+
GoogleGenAiChatOptions.builder()
162+
.model("gemini-2.0-flash-lite")
163+
.temperature(0.0)
164+
.maxOutputTokens(50)
165+
.build())
166+
.build();
166167
}
167168
}

0 commit comments

Comments
 (0)