Skip to content

Commit a369e8a

Browse files
ddobrinddobrin
committed
Added APIKey based tests
Signed-off-by: ddobrin <[email protected]>
1 parent f4a9b12 commit a369e8a

File tree

1 file changed

+184
-0
lines changed

1 file changed

+184
-0
lines changed

‎models/spring-ai-google-genai/src/test/java/org/springframework/ai/google/genai/gemini/GoogleGenAiChatModelObservationApiKeyIT.java‎

Lines changed: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,184 @@
1+
/*
2+
* Copyright 2023-2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.google.genai.gemini;
18+
19+
import static org.assertj.core.api.Assertions.assertThat;
20+
21+
import com.google.genai.Client;
22+
import io.micrometer.observation.tck.TestObservationRegistry;
23+
import io.micrometer.observation.tck.TestObservationRegistryAssert;
24+
import java.util.List;
25+
import java.util.stream.Collectors;
26+
import org.junit.jupiter.api.BeforeEach;
27+
import org.junit.jupiter.api.Test;
28+
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
29+
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
30+
import org.springframework.ai.chat.model.ChatResponse;
31+
import org.springframework.ai.chat.observation.ChatModelObservationDocumentation;
32+
import org.springframework.ai.chat.observation.DefaultChatModelObservationConvention;
33+
import org.springframework.ai.chat.prompt.Prompt;
34+
import org.springframework.ai.google.genai.GoogleGenAiChatModel;
35+
import org.springframework.ai.google.genai.GoogleGenAiChatOptions;
36+
import org.springframework.ai.observation.conventions.AiOperationType;
37+
import org.springframework.ai.observation.conventions.AiProvider;
38+
import org.springframework.beans.factory.annotation.Autowired;
39+
import org.springframework.boot.SpringBootConfiguration;
40+
import org.springframework.boot.test.context.SpringBootTest;
41+
import org.springframework.context.annotation.Bean;
42+
import reactor.core.publisher.Flux;
43+
44+
/**
45+
* @author Soby Chacko
46+
* @author Dan Dobrin
47+
*/
48+
@SpringBootTest
49+
@EnabledIfEnvironmentVariable(named = "GOOGLE_API_KEY", matches = ".*")
50+
public class GoogleGenAiChatModelObservationApiKeyIT {
51+
52+
@Autowired
53+
TestObservationRegistry observationRegistry;
54+
55+
@Autowired
56+
GoogleGenAiChatModel chatModel;
57+
58+
@BeforeEach
59+
void beforeEach() {
60+
this.observationRegistry.clear();
61+
}
62+
63+
@Test
64+
void observationForChatOperation() {
65+
66+
var options = GoogleGenAiChatOptions.builder()
67+
.model(GoogleGenAiChatModel.ChatModel.GEMINI_2_0_FLASH.getValue())
68+
.temperature(0.7)
69+
.stopSequences(List.of("this-is-the-end"))
70+
.maxOutputTokens(2048)
71+
.topP(1.0)
72+
.build();
73+
74+
Prompt prompt = new Prompt("Why does a raven look like a desk?", options);
75+
76+
ChatResponse chatResponse = this.chatModel.call(prompt);
77+
assertThat(chatResponse.getResult().getOutput().getText()).isNotEmpty();
78+
79+
ChatResponseMetadata responseMetadata = chatResponse.getMetadata();
80+
assertThat(responseMetadata).isNotNull();
81+
82+
validate(responseMetadata);
83+
}
84+
85+
@Test
86+
void observationForStreamingOperation() {
87+
88+
var options = GoogleGenAiChatOptions.builder()
89+
.model(GoogleGenAiChatModel.ChatModel.GEMINI_2_0_FLASH.getValue())
90+
.temperature(0.7)
91+
.stopSequences(List.of("this-is-the-end"))
92+
.maxOutputTokens(2048)
93+
.topP(1.0)
94+
.build();
95+
96+
Prompt prompt = new Prompt("Why does a raven look like a desk?", options);
97+
98+
Flux<ChatResponse> chatResponse = this.chatModel.stream(prompt);
99+
List<ChatResponse> responses = chatResponse.collectList().block();
100+
assertThat(responses).isNotEmpty();
101+
assertThat(responses).hasSizeGreaterThan(1);
102+
103+
String aggregatedResponse = responses.subList(0, responses.size() - 1)
104+
.stream()
105+
.map(r -> r.getResult().getOutput().getText())
106+
.collect(Collectors.joining());
107+
assertThat(aggregatedResponse).isNotEmpty();
108+
109+
ChatResponse lastChatResponse = responses.get(responses.size() - 1);
110+
111+
ChatResponseMetadata responseMetadata = lastChatResponse.getMetadata();
112+
assertThat(responseMetadata).isNotNull();
113+
114+
validate(responseMetadata);
115+
}
116+
117+
private void validate(ChatResponseMetadata responseMetadata) {
118+
TestObservationRegistryAssert.assertThat(this.observationRegistry)
119+
.doesNotHaveAnyRemainingCurrentObservation()
120+
.hasObservationWithNameEqualTo(DefaultChatModelObservationConvention.DEFAULT_NAME)
121+
.that()
122+
.hasLowCardinalityKeyValue(
123+
ChatModelObservationDocumentation.LowCardinalityKeyNames.AI_OPERATION_TYPE.asString(),
124+
AiOperationType.CHAT.value())
125+
.hasLowCardinalityKeyValue(ChatModelObservationDocumentation.LowCardinalityKeyNames.AI_PROVIDER.asString(),
126+
AiProvider.GOOGLE_GENAI_AI.value())
127+
.hasLowCardinalityKeyValue(
128+
ChatModelObservationDocumentation.LowCardinalityKeyNames.REQUEST_MODEL.asString(),
129+
GoogleGenAiChatModel.ChatModel.GEMINI_2_0_FLASH.getValue())
130+
.hasHighCardinalityKeyValue(
131+
ChatModelObservationDocumentation.HighCardinalityKeyNames.REQUEST_MAX_TOKENS.asString(), "2048")
132+
.hasHighCardinalityKeyValue(
133+
ChatModelObservationDocumentation.HighCardinalityKeyNames.REQUEST_STOP_SEQUENCES.asString(),
134+
"[\"this-is-the-end\"]")
135+
.hasHighCardinalityKeyValue(
136+
ChatModelObservationDocumentation.HighCardinalityKeyNames.REQUEST_TEMPERATURE.asString(), "0.7")
137+
.doesNotHaveHighCardinalityKeyValueWithKey(
138+
ChatModelObservationDocumentation.HighCardinalityKeyNames.REQUEST_TOP_K.asString())
139+
.hasHighCardinalityKeyValue(
140+
ChatModelObservationDocumentation.HighCardinalityKeyNames.REQUEST_TOP_P.asString(), "1.0")
141+
.hasHighCardinalityKeyValue(
142+
ChatModelObservationDocumentation.HighCardinalityKeyNames.RESPONSE_FINISH_REASONS.asString(),
143+
"[\"STOP\"]")
144+
.hasHighCardinalityKeyValue(
145+
ChatModelObservationDocumentation.HighCardinalityKeyNames.USAGE_INPUT_TOKENS.asString(),
146+
String.valueOf(responseMetadata.getUsage().getPromptTokens()))
147+
.hasHighCardinalityKeyValue(
148+
ChatModelObservationDocumentation.HighCardinalityKeyNames.USAGE_OUTPUT_TOKENS.asString(),
149+
String.valueOf(responseMetadata.getUsage().getCompletionTokens()))
150+
.hasHighCardinalityKeyValue(
151+
ChatModelObservationDocumentation.HighCardinalityKeyNames.USAGE_TOTAL_TOKENS.asString(),
152+
String.valueOf(responseMetadata.getUsage().getTotalTokens()))
153+
.hasBeenStarted()
154+
.hasBeenStopped();
155+
}
156+
157+
@SpringBootConfiguration
158+
static class Config {
159+
160+
@Bean
161+
public TestObservationRegistry observationRegistry() {
162+
return TestObservationRegistry.create();
163+
}
164+
165+
@Bean
166+
public Client genAiClient() {
167+
String apiKey = System.getenv("GOOGLE_API_KEY");
168+
return Client.builder().apiKey(apiKey).build();
169+
}
170+
171+
@Bean
172+
public GoogleGenAiChatModel vertexAiEmbedding(Client genAiClient, TestObservationRegistry observationRegistry) {
173+
174+
return GoogleGenAiChatModel.builder()
175+
.genAiClient(genAiClient)
176+
.observationRegistry(observationRegistry)
177+
.defaultOptions(
178+
GoogleGenAiChatOptions.builder().model(GoogleGenAiChatModel.ChatModel.GEMINI_2_0_FLASH).build())
179+
.build();
180+
}
181+
182+
}
183+
184+
}

0 commit comments

Comments
 (0)