Skip to content

Commit e8bb417

Browse files
committed
Refactored packages to google.genai
Signed-off-by: ddobrin <[email protected]>
1 parent 53f1b90 commit e8bb417

30 files changed

+291
-300
lines changed
Lines changed: 44 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
* limitations under the License.
1515
*/
1616

17-
package org.springframework.ai.vertexai.gemini;
17+
package org.springframework.ai.google.genai;
1818

1919
import java.net.URI;
2020
import java.util.ArrayList;
@@ -25,7 +25,6 @@
2525

2626
import com.fasterxml.jackson.annotation.JsonInclude;
2727
import com.fasterxml.jackson.annotation.JsonInclude.Include;
28-
import com.fasterxml.jackson.databind.JsonNode;
2928
import com.google.genai.Client;
3029
import com.google.genai.ResponseStream;
3130
import com.google.genai.types.Candidate;
@@ -46,6 +45,7 @@
4645
import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
4746
import org.slf4j.Logger;
4847
import org.slf4j.LoggerFactory;
48+
import org.springframework.ai.google.genai.schema.GoogleGenAiToolCallingManager;
4949
import reactor.core.publisher.Flux;
5050
import reactor.core.scheduler.Schedulers;
5151

@@ -82,9 +82,8 @@
8282
import org.springframework.ai.retry.RetryUtils;
8383
import org.springframework.ai.support.UsageCalculator;
8484
import org.springframework.ai.tool.definition.ToolDefinition;
85-
import org.springframework.ai.vertexai.gemini.common.VertexAiGeminiConstants;
86-
import org.springframework.ai.vertexai.gemini.common.VertexAiGeminiSafetySetting;
87-
import org.springframework.ai.vertexai.gemini.schema.VertexToolCallingManager;
85+
import org.springframework.ai.google.genai.common.GoogleGenAiConstants;
86+
import org.springframework.ai.google.genai.common.GoogleGenAiSafetySetting;
8887
import org.springframework.beans.factory.DisposableBean;
8988
import org.springframework.lang.NonNull;
9089
import org.springframework.retry.support.RetryTemplate;
@@ -93,8 +92,8 @@
9392
import org.springframework.util.StringUtils;
9493

9594
/**
96-
* Vertex AI Gemini Chat Model implementation that provides access to Google's Gemini
97-
* language models.
95+
* Google GenAI Chat Model implementation that provides access to Google's Gemini language
96+
* models.
9897
*
9998
* <p>
10099
* Key features include:
@@ -103,7 +102,7 @@
103102
* Gemini 1.5/2.0 Flash variants</li>
104103
* <li>Tool/Function calling capabilities through {@link ToolCallingManager}</li>
105104
* <li>Streaming support via {@link #stream(Prompt)} method</li>
106-
* <li>Configurable safety settings through {@link VertexAiGeminiSafetySetting}</li>
105+
* <li>Configurable safety settings through {@link GoogleGenAiSafetySetting}</li>
107106
* <li>Support for system messages and multi-modal content (text and images)</li>
108107
* <li>Built-in retry mechanism and observability through Micrometer</li>
109108
* <li>Google Search Retrieval integration</li>
@@ -112,13 +111,13 @@
112111
* <p>
113112
* The model can be configured with various options including temperature, top-k, top-p
114113
* sampling, maximum output tokens, and candidate count through
115-
* {@link VertexAiGeminiChatOptions}.
114+
* {@link GoogleGenAiChatOptions}.
116115
*
117116
* <p>
118117
* Use the {@link Builder} to create instances with custom configurations:
119118
*
120119
* <pre>{@code
121-
* VertexAiGeminiChatModel model = VertexAiGeminiChatModel.builder()
120+
* GoogleGenAiChatModel model = GoogleGenAiChatModel.builder()
122121
* .genAiClient(genAiClient)
123122
* .defaultOptions(options)
124123
* .toolCallingManager(toolManager)
@@ -135,11 +134,11 @@
135134
* @author Alexandros Pappas
136135
* @author Ilayaperumal Gopinathan
137136
* @since 0.8.1
138-
* @see VertexAiGeminiChatOptions
137+
* @see GoogleGenAiChatOptions
139138
* @see ToolCallingManager
140139
* @see ChatModel
141140
*/
142-
public class VertexAiGeminiChatModel implements ChatModel, DisposableBean {
141+
public class GoogleGenAiChatModel implements ChatModel, DisposableBean {
143142

144143
private static final ChatModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultChatModelObservationConvention();
145144

@@ -149,7 +148,7 @@ public class VertexAiGeminiChatModel implements ChatModel, DisposableBean {
149148

150149
private final Client genAiClient;
151150

152-
private final VertexAiGeminiChatOptions defaultOptions;
151+
private final GoogleGenAiChatOptions defaultOptions;
153152

154153
/**
155154
* The retry template used to retry the API calls.
@@ -180,40 +179,40 @@ public class VertexAiGeminiChatModel implements ChatModel, DisposableBean {
180179
private ChatModelObservationConvention observationConvention = DEFAULT_OBSERVATION_CONVENTION;
181180

182181
/**
183-
* Creates a new instance of VertexAiGeminiChatModel.
182+
* Creates a new instance of GoogleGenAiChatModel.
184183
* @param genAiClient the GenAI Client instance to use
185184
* @param defaultOptions the default options to use
186185
* @param toolCallingManager the tool calling manager to use. It is wrapped in a
187-
* {@link VertexToolCallingManager} to ensure compatibility with Vertex AI's OpenAPI
188-
* schema format.
186+
* {@link GoogleGenAiToolCallingManager} to ensure compatibility with Vertex AI's
187+
* OpenAPI schema format.
189188
* @param retryTemplate the retry template to use
190189
* @param observationRegistry the observation registry to use
191190
*/
192-
public VertexAiGeminiChatModel(Client genAiClient, VertexAiGeminiChatOptions defaultOptions,
191+
public GoogleGenAiChatModel(Client genAiClient, GoogleGenAiChatOptions defaultOptions,
193192
ToolCallingManager toolCallingManager, RetryTemplate retryTemplate,
194193
ObservationRegistry observationRegistry) {
195194
this(genAiClient, defaultOptions, toolCallingManager, retryTemplate, observationRegistry,
196195
new DefaultToolExecutionEligibilityPredicate());
197196
}
198197

199198
/**
200-
* Creates a new instance of VertexAiGeminiChatModel.
199+
* Creates a new instance of GoogleGenAiChatModel.
201200
* @param genAiClient the GenAI Client instance to use
202201
* @param defaultOptions the default options to use
203202
* @param toolCallingManager the tool calling manager to use. It is wrapped in a
204-
* {@link VertexToolCallingManager} to ensure compatibility with Vertex AI's OpenAPI
205-
* schema format.
203+
* {@link GoogleGenAiToolCallingManager} to ensure compatibility with Vertex AI's
204+
* OpenAPI schema format.
206205
* @param retryTemplate the retry template to use
207206
* @param observationRegistry the observation registry to use
208207
* @param toolExecutionEligibilityPredicate the tool execution eligibility predicate
209208
*/
210-
public VertexAiGeminiChatModel(Client genAiClient, VertexAiGeminiChatOptions defaultOptions,
209+
public GoogleGenAiChatModel(Client genAiClient, GoogleGenAiChatOptions defaultOptions,
211210
ToolCallingManager toolCallingManager, RetryTemplate retryTemplate, ObservationRegistry observationRegistry,
212211
ToolExecutionEligibilityPredicate toolExecutionEligibilityPredicate) {
213212

214213
Assert.notNull(genAiClient, "GenAI Client must not be null");
215-
Assert.notNull(defaultOptions, "VertexAiGeminiChatOptions must not be null");
216-
Assert.notNull(defaultOptions.getModel(), "VertexAiGeminiChatOptions.modelName must not be null");
214+
Assert.notNull(defaultOptions, "GoogleGenAiChatOptions must not be null");
215+
Assert.notNull(defaultOptions.getModel(), "GoogleGenAiChatOptions.modelName must not be null");
217216
Assert.notNull(retryTemplate, "RetryTemplate must not be null");
218217
Assert.notNull(toolCallingManager, "ToolCallingManager must not be null");
219218
Assert.notNull(toolExecutionEligibilityPredicate, "ToolExecutionEligibilityPredicate must not be null");
@@ -225,14 +224,14 @@ public VertexAiGeminiChatModel(Client genAiClient, VertexAiGeminiChatOptions def
225224
this.observationRegistry = observationRegistry;
226225
this.toolExecutionEligibilityPredicate = toolExecutionEligibilityPredicate;
227226

228-
// Wrap the provided tool calling manager in a VertexToolCallingManager to
227+
// Wrap the provided tool calling manager in a GoogleGenAiToolCallingManager to
229228
// ensure
230229
// compatibility with Vertex AI's OpenAPI schema format.
231-
if (toolCallingManager instanceof VertexToolCallingManager) {
230+
if (toolCallingManager instanceof GoogleGenAiToolCallingManager) {
232231
this.toolCallingManager = toolCallingManager;
233232
}
234233
else {
235-
this.toolCallingManager = new VertexToolCallingManager(toolCallingManager);
234+
this.toolCallingManager = new GoogleGenAiToolCallingManager(toolCallingManager);
236235
}
237236
}
238237

@@ -393,7 +392,7 @@ private ChatResponse internalCall(Prompt prompt, ChatResponse previousChatRespon
393392

394393
ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
395394
.prompt(prompt)
396-
.provider(VertexAiGeminiConstants.PROVIDER_NAME)
395+
.provider(GoogleGenAiConstants.PROVIDER_NAME)
397396
.build();
398397

399398
ChatResponse response = ChatModelObservationDocumentation.CHAT_MODEL_OPERATION
@@ -444,21 +443,21 @@ private ChatResponse internalCall(Prompt prompt, ChatResponse previousChatRespon
444443

445444
Prompt buildRequestPrompt(Prompt prompt) {
446445
// Process runtime options
447-
VertexAiGeminiChatOptions runtimeOptions = null;
446+
GoogleGenAiChatOptions runtimeOptions = null;
448447
if (prompt.getOptions() != null) {
449448
if (prompt.getOptions() instanceof ToolCallingChatOptions toolCallingChatOptions) {
450449
runtimeOptions = ModelOptionsUtils.copyToTarget(toolCallingChatOptions, ToolCallingChatOptions.class,
451-
VertexAiGeminiChatOptions.class);
450+
GoogleGenAiChatOptions.class);
452451
}
453452
else {
454453
runtimeOptions = ModelOptionsUtils.copyToTarget(prompt.getOptions(), ChatOptions.class,
455-
VertexAiGeminiChatOptions.class);
454+
GoogleGenAiChatOptions.class);
456455
}
457456
}
458457

459458
// Define request options by merging runtime options and default options
460-
VertexAiGeminiChatOptions requestOptions = ModelOptionsUtils.merge(runtimeOptions, this.defaultOptions,
461-
VertexAiGeminiChatOptions.class);
459+
GoogleGenAiChatOptions requestOptions = ModelOptionsUtils.merge(runtimeOptions, this.defaultOptions,
460+
GoogleGenAiChatOptions.class);
462461

463462
// Merge @JsonIgnore-annotated options explicitly since they are ignored by
464463
// Jackson, used by ModelOptionsUtils.
@@ -504,7 +503,7 @@ public Flux<ChatResponse> internalStream(Prompt prompt, ChatResponse previousCha
504503

505504
ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
506505
.prompt(prompt)
507-
.provider(VertexAiGeminiConstants.PROVIDER_NAME)
506+
.provider(GoogleGenAiConstants.PROVIDER_NAME)
508507
.build();
509508

510509
Observation observation = ChatModelObservationDocumentation.CHAT_MODEL_OPERATION.observation(
@@ -636,7 +635,7 @@ private DefaultUsage getDefaultUsage(com.google.genai.types.GenerateContentRespo
636635

637636
GeminiRequest createGeminiRequest(Prompt prompt) {
638637

639-
VertexAiGeminiChatOptions requestOptions = (VertexAiGeminiChatOptions) prompt.getOptions();
638+
GoogleGenAiChatOptions requestOptions = (GoogleGenAiChatOptions) prompt.getOptions();
640639

641640
// Build GenerateContentConfig
642641
GenerateContentConfig.Builder configBuilder = GenerateContentConfig.builder();
@@ -692,7 +691,7 @@ GeminiRequest createGeminiRequest(Prompt prompt) {
692691
tools.add(Tool.builder().functionDeclarations(functionDeclarations).build());
693692
}
694693

695-
if (prompt.getOptions() instanceof VertexAiGeminiChatOptions options && options.getGoogleSearchRetrieval()) {
694+
if (prompt.getOptions() instanceof GoogleGenAiChatOptions options && options.getGoogleSearchRetrieval()) {
696695
var googleSearch = GoogleSearch.builder().build();
697696
final var googleSearchRetrievalTool = Tool.builder().googleSearch(googleSearch).build();
698697
tools.add(googleSearchRetrievalTool);
@@ -721,7 +720,7 @@ GeminiRequest createGeminiRequest(Prompt prompt) {
721720

722721
// Helper methods for mapping safety settings enums
723722
private static com.google.genai.types.HarmCategory mapToGenAiHarmCategory(
724-
VertexAiGeminiSafetySetting.HarmCategory category) {
723+
GoogleGenAiSafetySetting.HarmCategory category) {
725724
switch (category) {
726725
case HARM_CATEGORY_UNSPECIFIED:
727726
return new com.google.genai.types.HarmCategory(
@@ -744,7 +743,7 @@ private static com.google.genai.types.HarmCategory mapToGenAiHarmCategory(
744743
}
745744

746745
private static com.google.genai.types.HarmBlockThreshold mapToGenAiHarmBlockThreshold(
747-
VertexAiGeminiSafetySetting.HarmBlockThreshold threshold) {
746+
GoogleGenAiSafetySetting.HarmBlockThreshold threshold) {
748747
switch (threshold) {
749748
case HARM_BLOCK_THRESHOLD_UNSPECIFIED:
750749
return new com.google.genai.types.HarmBlockThreshold(
@@ -781,7 +780,7 @@ private List<Content> toGeminiContent(List<Message> instructions) {
781780
return contents;
782781
}
783782

784-
private List<SafetySetting> toGeminiSafetySettings(List<VertexAiGeminiSafetySetting> safetySettings) {
783+
private List<SafetySetting> toGeminiSafetySettings(List<GoogleGenAiSafetySetting> safetySettings) {
785784
return safetySettings.stream()
786785
.map(safetySetting -> SafetySetting.builder()
787786
.category(mapToGenAiHarmCategory(safetySetting.getCategory()))
@@ -808,7 +807,7 @@ GenerateContentResponse getContentResponse(GeminiRequest request) {
808807

809808
@Override
810809
public ChatOptions getDefaultOptions() {
811-
return VertexAiGeminiChatOptions.fromOptions(this.defaultOptions);
810+
return GoogleGenAiChatOptions.fromOptions(this.defaultOptions);
812811
}
813812

814813
@Override
@@ -833,10 +832,10 @@ public static final class Builder {
833832

834833
private Client genAiClient;
835834

836-
private VertexAiGeminiChatOptions defaultOptions = VertexAiGeminiChatOptions.builder()
835+
private GoogleGenAiChatOptions defaultOptions = GoogleGenAiChatOptions.builder()
837836
.temperature(0.7)
838837
.topP(1.0)
839-
.model(VertexAiGeminiChatModel.ChatModel.GEMINI_2_0_FLASH)
838+
.model(GoogleGenAiChatModel.ChatModel.GEMINI_2_0_FLASH)
840839
.build();
841840

842841
private ToolCallingManager toolCallingManager;
@@ -855,7 +854,7 @@ public Builder genAiClient(Client genAiClient) {
855854
return this;
856855
}
857856

858-
public Builder defaultOptions(VertexAiGeminiChatOptions defaultOptions) {
857+
public Builder defaultOptions(GoogleGenAiChatOptions defaultOptions) {
859858
this.defaultOptions = defaultOptions;
860859
return this;
861860
}
@@ -881,12 +880,12 @@ public Builder observationRegistry(ObservationRegistry observationRegistry) {
881880
return this;
882881
}
883882

884-
public VertexAiGeminiChatModel build() {
883+
public GoogleGenAiChatModel build() {
885884
if (this.toolCallingManager != null) {
886-
return new VertexAiGeminiChatModel(this.genAiClient, this.defaultOptions, this.toolCallingManager,
885+
return new GoogleGenAiChatModel(this.genAiClient, this.defaultOptions, this.toolCallingManager,
887886
this.retryTemplate, this.observationRegistry, this.toolExecutionEligibilityPredicate);
888887
}
889-
return new VertexAiGeminiChatModel(this.genAiClient, this.defaultOptions, DEFAULT_TOOL_CALLING_MANAGER,
888+
return new GoogleGenAiChatModel(this.genAiClient, this.defaultOptions, DEFAULT_TOOL_CALLING_MANAGER,
890889
this.retryTemplate, this.observationRegistry, this.toolExecutionEligibilityPredicate);
891890
}
892891

0 commit comments

Comments
 (0)