Skip to content

Commit f293094

Browse files
FottasFottas
committed
perf(google-genai): support HTTP client timeout configuration
Signed-off-by: yinh <fottas@163.com>
1 parent 2b0dc84 commit f293094

File tree

6 files changed

+200
-7
lines changed

6 files changed

+200
-7
lines changed

‎auto-configurations/models/spring-ai-autoconfigure-model-google-genai/src/main/java/org/springframework/ai/model/google/genai/autoconfigure/chat/GoogleGenAiChatAutoConfiguration.java‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,14 @@
1717
package org.springframework.ai.model.google.genai.autoconfigure.chat;
1818

1919
import java.io.IOException;
20+
import java.time.Duration;
21+
import java.util.Optional;
2022

2123
import com.google.auth.oauth2.GoogleCredentials;
24+
import com.google.common.collect.ImmutableMap;
2225
import com.google.genai.Client;
26+
import com.google.genai.types.ClientOptions;
27+
import com.google.genai.types.HttpOptions;
2328
import io.micrometer.observation.ObservationRegistry;
2429

2530
import org.springframework.ai.chat.observation.ChatModelObservationConvention;
@@ -88,6 +93,28 @@ public Client googleGenAiClient(GoogleGenAiConnectionProperties connectionProper
8893
// credentials are handled automatically when vertexAI is true
8994
}
9095
}
96+
HttpOptions.Builder httpOptionsBuilder = HttpOptions.builder();
97+
if (connectionProperties.getTimeout() != null) {
98+
Integer timeout = Optional.ofNullable(connectionProperties.getTimeout())
99+
.map(Duration::toMillisPart)
100+
.orElse(null);
101+
httpOptionsBuilder.timeout(timeout);
102+
}
103+
if (!connectionProperties.getCustomHeaders().isEmpty()) {
104+
httpOptionsBuilder.headers(ImmutableMap.copyOf(connectionProperties.getCustomHeaders()));
105+
}
106+
107+
ClientOptions.Builder clientOptionsBuilder = ClientOptions.builder();
108+
if (connectionProperties.getMaxConnections() != null) {
109+
clientOptionsBuilder.maxConnections(connectionProperties.getMaxConnections());
110+
}
111+
112+
if (connectionProperties.getMaxConnectionsPerHost() != null) {
113+
clientOptionsBuilder.maxConnectionsPerHost(connectionProperties.getMaxConnectionsPerHost());
114+
}
115+
116+
clientBuilder.httpOptions(httpOptionsBuilder.build());
117+
clientBuilder.clientOptions(clientOptionsBuilder.build());
91118

92119
return clientBuilder.build();
93120
}

‎auto-configurations/models/spring-ai-autoconfigure-model-google-genai/src/main/java/org/springframework/ai/model/google/genai/autoconfigure/chat/GoogleGenAiConnectionProperties.java‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,10 @@
1616

1717
package org.springframework.ai.model.google.genai.autoconfigure.chat;
1818

19+
import java.time.Duration;
20+
import java.util.HashMap;
21+
import java.util.Map;
22+
1923
import org.springframework.boot.context.properties.ConfigurationProperties;
2024
import org.springframework.core.io.Resource;
2125

@@ -56,6 +60,23 @@ public class GoogleGenAiConnectionProperties {
5660
*/
5761
private boolean vertexAi;
5862

63+
/**
64+
* Timeout for the request in milliseconds.
65+
*/
66+
private Duration timeout;
67+
68+
/**
69+
* The maximum number of connections allowed in the pool.
70+
*/
71+
private Integer maxConnections;
72+
73+
/**
74+
* The maximum number of connections allowed per host.
75+
*/
76+
private Integer maxConnectionsPerHost;
77+
78+
private Map<String, String> customHeaders = new HashMap<>();
79+
5980
public String getApiKey() {
6081
return this.apiKey;
6182
}
@@ -96,4 +117,36 @@ public void setVertexAi(boolean vertexAi) {
96117
this.vertexAi = vertexAi;
97118
}
98119

120+
public Duration getTimeout() {
121+
return this.timeout;
122+
}
123+
124+
public void setTimeout(Duration timeout) {
125+
this.timeout = timeout;
126+
}
127+
128+
public Integer getMaxConnections() {
129+
return this.maxConnections;
130+
}
131+
132+
public void setMaxConnections(Integer maxConnections) {
133+
this.maxConnections = maxConnections;
134+
}
135+
136+
public Integer getMaxConnectionsPerHost() {
137+
return this.maxConnectionsPerHost;
138+
}
139+
140+
public void setMaxConnectionsPerHost(Integer maxConnectionsPerHost) {
141+
this.maxConnectionsPerHost = maxConnectionsPerHost;
142+
}
143+
144+
public Map<String, String> getCustomHeaders() {
145+
return this.customHeaders;
146+
}
147+
148+
public void setCustomHeaders(Map<String, String> customHeaders) {
149+
this.customHeaders = customHeaders;
150+
}
151+
99152
}

‎auto-configurations/models/spring-ai-autoconfigure-model-google-genai/src/main/java/org/springframework/ai/model/google/genai/autoconfigure/embedding/GoogleGenAiEmbeddingConnectionAutoConfiguration.java‎

Lines changed: 36 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,10 @@
1919
import java.io.IOException;
2020

2121
import com.google.auth.oauth2.GoogleCredentials;
22+
import com.google.common.collect.ImmutableMap;
2223
import com.google.genai.Client;
24+
import com.google.genai.types.ClientOptions;
25+
import com.google.genai.types.HttpOptions;
2326

2427
import org.springframework.ai.google.genai.GoogleGenAiEmbeddingConnectionDetails;
2528
import org.springframework.boot.autoconfigure.AutoConfiguration;
@@ -49,28 +52,55 @@ public GoogleGenAiEmbeddingConnectionDetails googleGenAiEmbeddingConnectionDetai
4952
GoogleGenAiEmbeddingConnectionProperties connectionProperties) throws IOException {
5053

5154
var connectionBuilder = GoogleGenAiEmbeddingConnectionDetails.builder();
55+
Client client = buildGenAiClient(connectionProperties);
56+
connectionBuilder.genAiClient(client);
57+
return connectionBuilder.build();
58+
}
59+
60+
private Client buildGenAiClient(GoogleGenAiEmbeddingConnectionProperties connectionProperties) throws IOException {
61+
Client.Builder clientBuilder = Client.builder();
5262

5363
if (StringUtils.hasText(connectionProperties.getApiKey())) {
5464
// Gemini Developer API mode
55-
connectionBuilder.apiKey(connectionProperties.getApiKey());
65+
clientBuilder.apiKey(connectionProperties.getApiKey());
5666
}
5767
else {
5868
// Vertex AI mode
5969
Assert.hasText(connectionProperties.getProjectId(), "Google GenAI project-id must be set!");
6070
Assert.hasText(connectionProperties.getLocation(), "Google GenAI location must be set!");
6171

62-
connectionBuilder.projectId(connectionProperties.getProjectId())
63-
.location(connectionProperties.getLocation());
72+
clientBuilder.project(connectionProperties.getProjectId())
73+
.location(connectionProperties.getLocation())
74+
.vertexAI(true);
6475

6576
if (connectionProperties.getCredentialsUri() != null) {
6677
GoogleCredentials credentials = GoogleCredentials
6778
.fromStream(connectionProperties.getCredentialsUri().getInputStream());
68-
// Note: Credentials are handled automatically by the SDK when using
69-
// Vertex AI mode
79+
// Note: The new SDK doesn't have a direct setCredentials method,
80+
// credentials are handled automatically when vertexAI is true
7081
}
7182
}
83+
HttpOptions.Builder httpOptionsBuilder = HttpOptions.builder();
84+
if (connectionProperties.getTimeout() != null) {
85+
httpOptionsBuilder.timeout((int) connectionProperties.getTimeout().getSeconds());
86+
}
87+
if (!connectionProperties.getCustomHeaders().isEmpty()) {
88+
httpOptionsBuilder.headers(ImmutableMap.copyOf(connectionProperties.getCustomHeaders()));
89+
}
7290

73-
return connectionBuilder.build();
91+
ClientOptions.Builder clientOptionsBuilder = ClientOptions.builder();
92+
if (connectionProperties.getMaxConnections() != null) {
93+
clientOptionsBuilder.maxConnections(connectionProperties.getMaxConnections());
94+
}
95+
96+
if (connectionProperties.getMaxConnectionsPerHost() != null) {
97+
clientOptionsBuilder.maxConnectionsPerHost(connectionProperties.getMaxConnectionsPerHost());
98+
}
99+
100+
clientBuilder.httpOptions(httpOptionsBuilder.build());
101+
clientBuilder.clientOptions(clientOptionsBuilder.build());
102+
103+
return clientBuilder.build();
74104
}
75105

76106
}

‎auto-configurations/models/spring-ai-autoconfigure-model-google-genai/src/main/java/org/springframework/ai/model/google/genai/autoconfigure/embedding/GoogleGenAiEmbeddingConnectionProperties.java‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,10 @@
1616

1717
package org.springframework.ai.model.google.genai.autoconfigure.embedding;
1818

19+
import java.time.Duration;
20+
import java.util.HashMap;
21+
import java.util.Map;
22+
1923
import org.springframework.boot.context.properties.ConfigurationProperties;
2024
import org.springframework.core.io.Resource;
2125

@@ -52,6 +56,23 @@ public class GoogleGenAiEmbeddingConnectionProperties {
5256
*/
5357
private Resource credentialsUri;
5458

59+
/**
60+
* Timeout for the request in milliseconds.
61+
*/
62+
private Duration timeout;
63+
64+
/**
65+
* The maximum number of connections allowed in the pool.
66+
*/
67+
private Integer maxConnections;
68+
69+
/**
70+
* The maximum number of connections allowed per host.
71+
*/
72+
private Integer maxConnectionsPerHost;
73+
74+
private Map<String, String> customHeaders = new HashMap<>();
75+
5576
/**
5677
* Whether to use Vertex AI mode. If false, uses Gemini Developer API mode. This is
5778
* automatically determined based on whether apiKey or projectId is set.
@@ -98,4 +119,36 @@ public void setVertexAi(boolean vertexAi) {
98119
this.vertexAi = vertexAi;
99120
}
100121

122+
public Duration getTimeout() {
123+
return this.timeout;
124+
}
125+
126+
public void setTimeout(Duration timeout) {
127+
this.timeout = timeout;
128+
}
129+
130+
public Integer getMaxConnections() {
131+
return this.maxConnections;
132+
}
133+
134+
public void setMaxConnections(Integer maxConnections) {
135+
this.maxConnections = maxConnections;
136+
}
137+
138+
public Integer getMaxConnectionsPerHost() {
139+
return this.maxConnectionsPerHost;
140+
}
141+
142+
public void setMaxConnectionsPerHost(Integer maxConnectionsPerHost) {
143+
this.maxConnectionsPerHost = maxConnectionsPerHost;
144+
}
145+
146+
public Map<String, String> getCustomHeaders() {
147+
return this.customHeaders;
148+
}
149+
150+
public void setCustomHeaders(Map<String, String> customHeaders) {
151+
this.customHeaders = customHeaders;
152+
}
153+
101154
}

‎auto-configurations/models/spring-ai-autoconfigure-model-google-genai/src/test/java/org/springframework/ai/model/google/genai/autoconfigure/chat/GoogleGenAiChatAutoConfigurationIT.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,4 +120,20 @@ void generateStreamingWithVertexAi() {
120120
});
121121
}
122122

123+
@Test
124+
@EnabledIfEnvironmentVariable(named = "GOOGLE_API_KEY", matches = ".*")
125+
void generateWithTimeout() {
126+
ApplicationContextRunner contextRunner = new ApplicationContextRunner()
127+
.withPropertyValues("spring.ai.google.genai.api-key=" + System.getenv("GOOGLE_API_KEY"),
128+
"spring.ai.google.genai.timeout=1s")
129+
.withConfiguration(SpringAiTestAutoConfigurations.of(GoogleGenAiChatAutoConfiguration.class));
130+
131+
contextRunner.run(context -> {
132+
GoogleGenAiChatModel chatModel = context.getBean(GoogleGenAiChatModel.class);
133+
String response = chatModel.call("Hello");
134+
assertThat(response).isNotEmpty();
135+
logger.info("Response: " + response);
136+
});
137+
}
138+
123139
}

‎auto-configurations/models/spring-ai-autoconfigure-model-google-genai/src/test/java/org/springframework/ai/model/google/genai/autoconfigure/chat/GoogleGenAiPropertiesTests.java‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
2424
import org.springframework.context.annotation.Configuration;
2525

26+
import java.time.Duration;
27+
2628
import static org.assertj.core.api.Assertions.assertThat;
2729

2830
/**
@@ -37,13 +39,15 @@ public class GoogleGenAiPropertiesTests {
3739
void connectionPropertiesBinding() {
3840
this.contextRunner
3941
.withPropertyValues("spring.ai.google.genai.api-key=test-key",
40-
"spring.ai.google.genai.project-id=test-project", "spring.ai.google.genai.location=us-central1")
42+
"spring.ai.google.genai.project-id=test-project", "spring.ai.google.genai.location=us-central1",
43+
"spring.ai.google.genai.timeout=1ms")
4144
.run(context -> {
4245
GoogleGenAiConnectionProperties connectionProperties = context
4346
.getBean(GoogleGenAiConnectionProperties.class);
4447
assertThat(connectionProperties.getApiKey()).isEqualTo("test-key");
4548
assertThat(connectionProperties.getProjectId()).isEqualTo("test-project");
4649
assertThat(connectionProperties.getLocation()).isEqualTo("us-central1");
50+
assertThat(connectionProperties.getTimeout()).isEqualTo(Duration.ofMillis(1));
4751
});
4852
}
4953

@@ -131,6 +135,16 @@ void extendedUsageMetadataDefaultBinding() {
131135
});
132136
}
133137

138+
@Test
139+
void extendedUsageCustomTimeoutPropertiesBinding() {
140+
this.contextRunner
141+
.withPropertyValues("spring.ai.google.genai.chat.options.include-extended-usage-metadata=true")
142+
.run(context -> {
143+
GoogleGenAiChatProperties chatProperties = context.getBean(GoogleGenAiChatProperties.class);
144+
assertThat(chatProperties.getOptions().getIncludeExtendedUsageMetadata()).isTrue();
145+
});
146+
}
147+
134148
@Configuration
135149
@EnableConfigurationProperties({ GoogleGenAiConnectionProperties.class, GoogleGenAiChatProperties.class,
136150
GoogleGenAiEmbeddingConnectionProperties.class })

0 commit comments

Comments
 (0)