14
14
* limitations under the License.
15
15
*/
16
16
17
- package org .springframework .ai .vertexai . gemini ;
17
+ package org .springframework .ai .google . genai ;
18
18
19
19
import java .net .URI ;
20
20
import java .util .ArrayList ;
25
25
26
26
import com .fasterxml .jackson .annotation .JsonInclude ;
27
27
import com .fasterxml .jackson .annotation .JsonInclude .Include ;
28
- import com .fasterxml .jackson .databind .JsonNode ;
29
28
import com .google .genai .Client ;
30
29
import com .google .genai .ResponseStream ;
31
30
import com .google .genai .types .Candidate ;
46
45
import io .micrometer .observation .contextpropagation .ObservationThreadLocalAccessor ;
47
46
import org .slf4j .Logger ;
48
47
import org .slf4j .LoggerFactory ;
48
+ import org .springframework .ai .google .genai .schema .GoogleGenAiToolCallingManager ;
49
49
import reactor .core .publisher .Flux ;
50
50
import reactor .core .scheduler .Schedulers ;
51
51
82
82
import org .springframework .ai .retry .RetryUtils ;
83
83
import org .springframework .ai .support .UsageCalculator ;
84
84
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 ;
88
87
import org .springframework .beans .factory .DisposableBean ;
89
88
import org .springframework .lang .NonNull ;
90
89
import org .springframework .retry .support .RetryTemplate ;
93
92
import org .springframework .util .StringUtils ;
94
93
95
94
/**
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.
98
97
*
99
98
* <p>
100
99
* Key features include:
103
102
* Gemini 1.5/2.0 Flash variants</li>
104
103
* <li>Tool/Function calling capabilities through {@link ToolCallingManager}</li>
105
104
* <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>
107
106
* <li>Support for system messages and multi-modal content (text and images)</li>
108
107
* <li>Built-in retry mechanism and observability through Micrometer</li>
109
108
* <li>Google Search Retrieval integration</li>
112
111
* <p>
113
112
* The model can be configured with various options including temperature, top-k, top-p
114
113
* sampling, maximum output tokens, and candidate count through
115
- * {@link VertexAiGeminiChatOptions }.
114
+ * {@link GoogleGenAiChatOptions }.
116
115
*
117
116
* <p>
118
117
* Use the {@link Builder} to create instances with custom configurations:
119
118
*
120
119
* <pre>{@code
121
- * VertexAiGeminiChatModel model = VertexAiGeminiChatModel .builder()
120
+ * GoogleGenAiChatModel model = GoogleGenAiChatModel .builder()
122
121
* .genAiClient(genAiClient)
123
122
* .defaultOptions(options)
124
123
* .toolCallingManager(toolManager)
135
134
* @author Alexandros Pappas
136
135
* @author Ilayaperumal Gopinathan
137
136
* @since 0.8.1
138
- * @see VertexAiGeminiChatOptions
137
+ * @see GoogleGenAiChatOptions
139
138
* @see ToolCallingManager
140
139
* @see ChatModel
141
140
*/
142
- public class VertexAiGeminiChatModel implements ChatModel , DisposableBean {
141
+ public class GoogleGenAiChatModel implements ChatModel , DisposableBean {
143
142
144
143
private static final ChatModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultChatModelObservationConvention ();
145
144
@@ -149,7 +148,7 @@ public class VertexAiGeminiChatModel implements ChatModel, DisposableBean {
149
148
150
149
private final Client genAiClient ;
151
150
152
- private final VertexAiGeminiChatOptions defaultOptions ;
151
+ private final GoogleGenAiChatOptions defaultOptions ;
153
152
154
153
/**
155
154
* The retry template used to retry the API calls.
@@ -180,40 +179,40 @@ public class VertexAiGeminiChatModel implements ChatModel, DisposableBean {
180
179
private ChatModelObservationConvention observationConvention = DEFAULT_OBSERVATION_CONVENTION ;
181
180
182
181
/**
183
- * Creates a new instance of VertexAiGeminiChatModel .
182
+ * Creates a new instance of GoogleGenAiChatModel .
184
183
* @param genAiClient the GenAI Client instance to use
185
184
* @param defaultOptions the default options to use
186
185
* @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.
189
188
* @param retryTemplate the retry template to use
190
189
* @param observationRegistry the observation registry to use
191
190
*/
192
- public VertexAiGeminiChatModel (Client genAiClient , VertexAiGeminiChatOptions defaultOptions ,
191
+ public GoogleGenAiChatModel (Client genAiClient , GoogleGenAiChatOptions defaultOptions ,
193
192
ToolCallingManager toolCallingManager , RetryTemplate retryTemplate ,
194
193
ObservationRegistry observationRegistry ) {
195
194
this (genAiClient , defaultOptions , toolCallingManager , retryTemplate , observationRegistry ,
196
195
new DefaultToolExecutionEligibilityPredicate ());
197
196
}
198
197
199
198
/**
200
- * Creates a new instance of VertexAiGeminiChatModel .
199
+ * Creates a new instance of GoogleGenAiChatModel .
201
200
* @param genAiClient the GenAI Client instance to use
202
201
* @param defaultOptions the default options to use
203
202
* @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.
206
205
* @param retryTemplate the retry template to use
207
206
* @param observationRegistry the observation registry to use
208
207
* @param toolExecutionEligibilityPredicate the tool execution eligibility predicate
209
208
*/
210
- public VertexAiGeminiChatModel (Client genAiClient , VertexAiGeminiChatOptions defaultOptions ,
209
+ public GoogleGenAiChatModel (Client genAiClient , GoogleGenAiChatOptions defaultOptions ,
211
210
ToolCallingManager toolCallingManager , RetryTemplate retryTemplate , ObservationRegistry observationRegistry ,
212
211
ToolExecutionEligibilityPredicate toolExecutionEligibilityPredicate ) {
213
212
214
213
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" );
217
216
Assert .notNull (retryTemplate , "RetryTemplate must not be null" );
218
217
Assert .notNull (toolCallingManager , "ToolCallingManager must not be null" );
219
218
Assert .notNull (toolExecutionEligibilityPredicate , "ToolExecutionEligibilityPredicate must not be null" );
@@ -225,14 +224,14 @@ public VertexAiGeminiChatModel(Client genAiClient, VertexAiGeminiChatOptions def
225
224
this .observationRegistry = observationRegistry ;
226
225
this .toolExecutionEligibilityPredicate = toolExecutionEligibilityPredicate ;
227
226
228
- // Wrap the provided tool calling manager in a VertexToolCallingManager to
227
+ // Wrap the provided tool calling manager in a GoogleGenAiToolCallingManager to
229
228
// ensure
230
229
// compatibility with Vertex AI's OpenAPI schema format.
231
- if (toolCallingManager instanceof VertexToolCallingManager ) {
230
+ if (toolCallingManager instanceof GoogleGenAiToolCallingManager ) {
232
231
this .toolCallingManager = toolCallingManager ;
233
232
}
234
233
else {
235
- this .toolCallingManager = new VertexToolCallingManager (toolCallingManager );
234
+ this .toolCallingManager = new GoogleGenAiToolCallingManager (toolCallingManager );
236
235
}
237
236
}
238
237
@@ -393,7 +392,7 @@ private ChatResponse internalCall(Prompt prompt, ChatResponse previousChatRespon
393
392
394
393
ChatModelObservationContext observationContext = ChatModelObservationContext .builder ()
395
394
.prompt (prompt )
396
- .provider (VertexAiGeminiConstants .PROVIDER_NAME )
395
+ .provider (GoogleGenAiConstants .PROVIDER_NAME )
397
396
.build ();
398
397
399
398
ChatResponse response = ChatModelObservationDocumentation .CHAT_MODEL_OPERATION
@@ -444,21 +443,21 @@ private ChatResponse internalCall(Prompt prompt, ChatResponse previousChatRespon
444
443
445
444
Prompt buildRequestPrompt (Prompt prompt ) {
446
445
// Process runtime options
447
- VertexAiGeminiChatOptions runtimeOptions = null ;
446
+ GoogleGenAiChatOptions runtimeOptions = null ;
448
447
if (prompt .getOptions () != null ) {
449
448
if (prompt .getOptions () instanceof ToolCallingChatOptions toolCallingChatOptions ) {
450
449
runtimeOptions = ModelOptionsUtils .copyToTarget (toolCallingChatOptions , ToolCallingChatOptions .class ,
451
- VertexAiGeminiChatOptions .class );
450
+ GoogleGenAiChatOptions .class );
452
451
}
453
452
else {
454
453
runtimeOptions = ModelOptionsUtils .copyToTarget (prompt .getOptions (), ChatOptions .class ,
455
- VertexAiGeminiChatOptions .class );
454
+ GoogleGenAiChatOptions .class );
456
455
}
457
456
}
458
457
459
458
// 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 );
462
461
463
462
// Merge @JsonIgnore-annotated options explicitly since they are ignored by
464
463
// Jackson, used by ModelOptionsUtils.
@@ -504,7 +503,7 @@ public Flux<ChatResponse> internalStream(Prompt prompt, ChatResponse previousCha
504
503
505
504
ChatModelObservationContext observationContext = ChatModelObservationContext .builder ()
506
505
.prompt (prompt )
507
- .provider (VertexAiGeminiConstants .PROVIDER_NAME )
506
+ .provider (GoogleGenAiConstants .PROVIDER_NAME )
508
507
.build ();
509
508
510
509
Observation observation = ChatModelObservationDocumentation .CHAT_MODEL_OPERATION .observation (
@@ -636,7 +635,7 @@ private DefaultUsage getDefaultUsage(com.google.genai.types.GenerateContentRespo
636
635
637
636
GeminiRequest createGeminiRequest (Prompt prompt ) {
638
637
639
- VertexAiGeminiChatOptions requestOptions = (VertexAiGeminiChatOptions ) prompt .getOptions ();
638
+ GoogleGenAiChatOptions requestOptions = (GoogleGenAiChatOptions ) prompt .getOptions ();
640
639
641
640
// Build GenerateContentConfig
642
641
GenerateContentConfig .Builder configBuilder = GenerateContentConfig .builder ();
@@ -692,7 +691,7 @@ GeminiRequest createGeminiRequest(Prompt prompt) {
692
691
tools .add (Tool .builder ().functionDeclarations (functionDeclarations ).build ());
693
692
}
694
693
695
- if (prompt .getOptions () instanceof VertexAiGeminiChatOptions options && options .getGoogleSearchRetrieval ()) {
694
+ if (prompt .getOptions () instanceof GoogleGenAiChatOptions options && options .getGoogleSearchRetrieval ()) {
696
695
var googleSearch = GoogleSearch .builder ().build ();
697
696
final var googleSearchRetrievalTool = Tool .builder ().googleSearch (googleSearch ).build ();
698
697
tools .add (googleSearchRetrievalTool );
@@ -721,7 +720,7 @@ GeminiRequest createGeminiRequest(Prompt prompt) {
721
720
722
721
// Helper methods for mapping safety settings enums
723
722
private static com .google .genai .types .HarmCategory mapToGenAiHarmCategory (
724
- VertexAiGeminiSafetySetting .HarmCategory category ) {
723
+ GoogleGenAiSafetySetting .HarmCategory category ) {
725
724
switch (category ) {
726
725
case HARM_CATEGORY_UNSPECIFIED :
727
726
return new com .google .genai .types .HarmCategory (
@@ -744,7 +743,7 @@ private static com.google.genai.types.HarmCategory mapToGenAiHarmCategory(
744
743
}
745
744
746
745
private static com .google .genai .types .HarmBlockThreshold mapToGenAiHarmBlockThreshold (
747
- VertexAiGeminiSafetySetting .HarmBlockThreshold threshold ) {
746
+ GoogleGenAiSafetySetting .HarmBlockThreshold threshold ) {
748
747
switch (threshold ) {
749
748
case HARM_BLOCK_THRESHOLD_UNSPECIFIED :
750
749
return new com .google .genai .types .HarmBlockThreshold (
@@ -781,7 +780,7 @@ private List<Content> toGeminiContent(List<Message> instructions) {
781
780
return contents ;
782
781
}
783
782
784
- private List <SafetySetting > toGeminiSafetySettings (List <VertexAiGeminiSafetySetting > safetySettings ) {
783
+ private List <SafetySetting > toGeminiSafetySettings (List <GoogleGenAiSafetySetting > safetySettings ) {
785
784
return safetySettings .stream ()
786
785
.map (safetySetting -> SafetySetting .builder ()
787
786
.category (mapToGenAiHarmCategory (safetySetting .getCategory ()))
@@ -808,7 +807,7 @@ GenerateContentResponse getContentResponse(GeminiRequest request) {
808
807
809
808
@ Override
810
809
public ChatOptions getDefaultOptions () {
811
- return VertexAiGeminiChatOptions .fromOptions (this .defaultOptions );
810
+ return GoogleGenAiChatOptions .fromOptions (this .defaultOptions );
812
811
}
813
812
814
813
@ Override
@@ -833,10 +832,10 @@ public static final class Builder {
833
832
834
833
private Client genAiClient ;
835
834
836
- private VertexAiGeminiChatOptions defaultOptions = VertexAiGeminiChatOptions .builder ()
835
+ private GoogleGenAiChatOptions defaultOptions = GoogleGenAiChatOptions .builder ()
837
836
.temperature (0.7 )
838
837
.topP (1.0 )
839
- .model (VertexAiGeminiChatModel .ChatModel .GEMINI_2_0_FLASH )
838
+ .model (GoogleGenAiChatModel .ChatModel .GEMINI_2_0_FLASH )
840
839
.build ();
841
840
842
841
private ToolCallingManager toolCallingManager ;
@@ -855,7 +854,7 @@ public Builder genAiClient(Client genAiClient) {
855
854
return this ;
856
855
}
857
856
858
- public Builder defaultOptions (VertexAiGeminiChatOptions defaultOptions ) {
857
+ public Builder defaultOptions (GoogleGenAiChatOptions defaultOptions ) {
859
858
this .defaultOptions = defaultOptions ;
860
859
return this ;
861
860
}
@@ -881,12 +880,12 @@ public Builder observationRegistry(ObservationRegistry observationRegistry) {
881
880
return this ;
882
881
}
883
882
884
- public VertexAiGeminiChatModel build () {
883
+ public GoogleGenAiChatModel build () {
885
884
if (this .toolCallingManager != null ) {
886
- return new VertexAiGeminiChatModel (this .genAiClient , this .defaultOptions , this .toolCallingManager ,
885
+ return new GoogleGenAiChatModel (this .genAiClient , this .defaultOptions , this .toolCallingManager ,
887
886
this .retryTemplate , this .observationRegistry , this .toolExecutionEligibilityPredicate );
888
887
}
889
- return new VertexAiGeminiChatModel (this .genAiClient , this .defaultOptions , DEFAULT_TOOL_CALLING_MANAGER ,
888
+ return new GoogleGenAiChatModel (this .genAiClient , this .defaultOptions , DEFAULT_TOOL_CALLING_MANAGER ,
890
889
this .retryTemplate , this .observationRegistry , this .toolExecutionEligibilityPredicate );
891
890
}
892
891
0 commit comments