feat(google-genai): introduce batch processing, grounding metadata, and image labels (#5255)

- Implemented comprehensive Batch API support covering all major
modalities: GoogleGenAiBatchChatModel, GoogleGenAiBatchEmbeddingModel,
and GoogleGenAiBatchImageModel. These new classes support both inline
requests and file-based batch job creation, along with retrieval,
cancellation, and deletion operations.
- Exposed GroundingMetadata in GoogleGenAiChatResponseMetadata by
mapping it directly from the SDK's Candidate response in
GoogleGenAiContentMapper.
- Added support for custom labels mapping in GoogleGenAiImageModel and
GoogleGenAiBatchImageModel using GenerateContentConfig.
- Updated test suites to use newer Gemini model versions, replacing
gemini-2.0-flash with gemini-3.1-flash-lite.

---------

Co-authored-by: Akshay Dipta <akshay.dipta@gmail.com>
Co-authored-by: Dmytro Liubarskyi <ljubarskij@gmail.com>
Co-authored-by: Dmytro Skarzhynets <d.skarzh@protonmail.com>
This commit is contained in:
Guillaume Laforge
2026-06-05 16:20:37 +02:00
committed by GitHub
parent d55bef6fbd
commit 26a81e17f0
36 changed files with 4490 additions and 128 deletions
+2
View File
@@ -46,3 +46,5 @@ CLAUDE.md
### Quarkus stuff ###
.quarkus
jdk/
@@ -19,8 +19,20 @@ https://github.com/googleapis/java-genai
- [Configuring](#configuring)
- [GoogleGenAiStreamingChatModel](#googlegenaistreamingchatmodel)
- [Executor](#executor)
- [GoogleGenAiEmbeddingModel](#googlegenaiembeddingmodel)
- [GoogleGenAiImageModel](#googlegenaiimagemodel)
- [Request & Response Logging](#request--response-logging)
- [Batch API](#batch-api)
- [Tools](#tools)
- [JSON Schema / Structured Outputs](#json-schema--structured-outputs)
- [Grounding Metadata](#grounding-metadata)
- [Custom Labels](#custom-labels)
- [File API](#file-api)
- [Cached Content Support](#cached-content-support)
- [Thinking Models (Gemini 3.0+)](#thinking-models-gemini-30)
- [Multimodality (Audio, Video, PDF)](#multimodality-audio-video-pdf)
- [Token Count Estimator](#token-count-estimator)
- [Model Catalog](#model-catalog)
## Maven Dependency
@@ -32,9 +44,27 @@ https://github.com/googleapis/java-genai
</dependency>
```
## API Key
## Authentication
You can authenticate with the Gemini models using either an API key or Google Cloud Vertex AI credentials.
### Gemini Developer API (API Key)
Get an API key for free here: https://ai.google.dev/gemini-api/docs/api-key.
You can provide it to the builder using `.apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))`.
### Google Cloud Vertex AI
If you are using Vertex AI, you can authenticate using Google Credentials along with your project ID and location. The integration will automatically use Application Default Credentials (ADC) if available, or you can explicitly provide them:
```java
ChatModel gemini = GoogleGenAiChatModel.builder()
// .googleCredentials(...) // Optional: explicitly provide credentials
.projectId("your-google-cloud-project-id")
.location("us-central1")
.modelName("gemini-2.5-flash")
.build();
```
## Models available
@@ -102,11 +132,106 @@ ChatModel gemini = GoogleGenAiChatModel.builder()
.enableGoogleMaps(true)
.enableUrlContext(true)
.allowedFunctionNames(List.of("getWeather"))
.thinkingBudget(250)
.thinkingLevel("LOW")
.listeners(...)
.build();
```
## Request & Response Logging
You can enable request and response logging for debugging, troubleshooting, and audit purposes on `GoogleGenAiChatModel`, `GoogleGenAiStreamingChatModel`, `GoogleGenAiEmbeddingModel`, and `GoogleGenAiImageModel`.
To capture these logs, configure `.logRequests(true)`, `.logResponses(true)` (or both using `.logRequestsAndResponses(true)`) in your model builders.
```java
ChatModel gemini = GoogleGenAiChatModel.builder()
.apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))
.modelName("gemini-2.5-flash")
.logRequests(true)
.logResponses(true)
// Or: .logRequestsAndResponses(true)
.build();
```
### Logging Configuration Setup
All logging in the Google Gen AI integration module is routed through the standard **SLF4J** facade. To actually view the output, you must ensure that:
1. An SLF4J binding (implementation) is present in your dependencies.
2. The logging framework is configured to output logs under the `INFO` level for the package `dev.langchain4j.model.google.genai`.
Below are common setup patterns for popular logging environments:
#### 1. Setup using Logback
Add the Logback classic implementation to your project:
##### Maven
```xml
<dependency>
<groupId>ch.qos.logback</groupId>
<artifactId>logback-classic</artifactId>
<version>1.5.8</version> <!-- or your preferred version -->
</dependency>
```
##### Gradle
```groovy
implementation 'ch.qos.logback:logback-classic:1.5.8'
```
Next, configure the logging level in your `src/main/resources/logback.xml` file. For example:
```xml
<configuration>
<appender name="STDOUT" class="ch.qos.logback.core.ConsoleAppender">
<encoder>
<pattern>%d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n</pattern>
</encoder>
</appender>
<!-- Configure the package specifically for Google Gen AI logging -->
<logger name="dev.langchain4j.model.google.genai" level="INFO" />
<root level="WARN">
<appender-ref ref="STDOUT" />
</root>
</configuration>
```
#### 2. Setup in Spring Boot Applications
Spring Boot automatically provides an SLF4J provider. Simply configure the logging level in your `application.properties` (or `application.yml` equivalent):
```properties
# Enable logging for Google Gen AI models
logging.level.dev.langchain4j.model.google.genai=INFO
```
#### 3. Setup with SLF4J Simple
If you are writing a script or a simple command-line application, you can use the lightweight `slf4j-simple` backend:
##### Maven
```xml
<dependency>
<groupId>org.slf4j</groupId>
<artifactId>slf4j-simple</artifactId>
<version>2.0.13</version>
</dependency>
```
Configure SLF4J Simple via a system property when starting your application:
```bash
java -Dorg.slf4j.simpleLogger.log.dev.langchain4j.model.google.genai=INFO -jar app.jar
```
Alternatively, create a `simplelogger.properties` file in `src/main/resources/` containing:
```properties
org.slf4j.simpleLogger.log.dev.langchain4j.model.google.genai=info
```
## GoogleGenAiStreamingChatModel
The `GoogleGenAiStreamingChatModel` allows streaming the text of a response token by token.
@@ -224,3 +349,248 @@ WeatherForecast forecast = forecastAssistant.extract("""
> [!NOTE]
> The Google Gen AI API has some restrictions on advanced JSON schema features (such as `anyOf` / polymorphic typing). Simple POJOs, lists, and nested objects are fully supported.
## Cached Content Support
When working with very large context windows (like massive system prompts, large documents, or extensive codebases) that are reused across multiple requests, you can significantly reduce costs and latency by caching the content.
Once you have created the cached content using the official Google Gen AI SDK or API, you can easily pass the unique cache identifier to the LangChain4j chat model builders:
```java
// Pass your cached content URI here
String cachedContentUri = "projects/123456/locations/us-central1/cachedContents/my-cached-content-789";
ChatModel gemini = GoogleGenAiChatModel.builder()
.apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))
.modelName("gemini-2.5-pro")
.cachedContent(cachedContentUri)
.build();
// The model will automatically use the cached context!
String response = gemini.chat("Summarize the cached document in 3 bullet points.");
```
This feature is available on `GoogleGenAiChatModel`, `GoogleGenAiStreamingChatModel`, and `GoogleGenAiBatchChatModel`.
## Thinking Models (Gemini 3.0+)
Gemini 3.0 models (like `gemini-3.0-pro` and `gemini-3.0-flash`) support advanced reasoning (thinking) capabilities.
You can enable this by specifying a `thinkingLevel` during model configuration. The supported values are `"MINIMAL"`, `"LOW"`, `"MEDIUM"`, and `"HIGH"`:
```java
ChatModel gemini = GoogleGenAiChatModel.builder()
.apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))
.modelName("gemini-3.0-pro")
.thinkingLevel("MEDIUM")
.build();
```
> [!NOTE]
> Previously, thinking was configured using a token-based `thinkingBudget`. The `thinkingBudget` parameter is now considered legacy, though still supported. You cannot specify both `thinkingLevel` and `thinkingBudget` at the same time.
> [!TIP]
> The LangChain4j `google-genai` integration seamlessly manages the complex state required for multi-turn tool execution with thinking models. It automatically persists and injects the necessary hidden `thought_signature` tokens across conversation turns, ensuring robust and uninterrupted agentic workflows!
## GoogleGenAiEmbeddingModel
The `GoogleGenAiEmbeddingModel` allows you to generate embeddings for text segments using models like `gemini-embedding-2`.
```java
EmbeddingModel embeddingModel = GoogleGenAiEmbeddingModel.builder()
.apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))
.modelName("gemini-embedding-2")
.outputDimensionality(768)
.taskType(GoogleGenAiEmbeddingModel.TaskTypeEnum.RETRIEVAL_DOCUMENT)
.build();
Response<Embedding> response = embeddingModel.embed("Hello world!");
```
### Batching & Retries
When embedding multiple text segments (via `embedAll`), `GoogleGenAiEmbeddingModel` automatically manages batching and API request retries.
```java
EmbeddingModel embeddingModel = GoogleGenAiEmbeddingModel.builder()
.apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))
.modelName("gemini-embedding-2")
.maxSegmentsPerBatch(100) // Default: 100. Sets maximum segments per batch request.
.maxRetries(3) // Default: 3. Automatically retries failed requests.
.build();
```
#### Title-based Grouping Strategy
The official Google Gen AI Java SDK's `embedContent` API only supports a single common `title` per batch request. To handle this restriction cleanly and preserve document-level associations, `GoogleGenAiEmbeddingModel` implements a **group-by-title** batching strategy:
1. When `taskType` is set to `RETRIEVAL_DOCUMENT`, the model groups text segments by their document title (extracted from the segment's metadata using the key defined by `.titleMetadataKey(...)`, which defaults to `"title"`).
2. Segments sharing the same title are batched and sent together in a single API call.
3. Segments with different titles (or no title) are processed in separate, optimized batches.
4. The resulting embeddings are seamlessly reassembled and returned in their original order.
This maximizes API throughput without losing document metadata context or individual segment titles.
## GoogleGenAiImageModel
The `GoogleGenAiImageModel` allows you to generate images from text prompts. It supports custom configuration like aspect ratios, image sizes, and person generation policies.
```java
ImageModel imageModel = GoogleGenAiImageModel.builder()
.apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))
.modelName("gemini-3.1-flash-image-preview")
.aspectRatio("16:9")
.build();
Response<Image> response = imageModel.generate("A futuristic city at sunset");
```
## Batch API
The Google Gen AI integration provides support for the Batch API, allowing you to run operations asynchronously in the background. The following batch models are supported:
- `GoogleGenAiBatchChatModel`
- `GoogleGenAiBatchEmbeddingModel`
- `GoogleGenAiBatchImageModel`
You can create batch jobs inline or from an uploaded file on Google Cloud.
```java
GoogleGenAiBatchChatModel batchChatModel = GoogleGenAiBatchChatModel.builder()
.apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))
.modelName("gemini-2.5-flash")
.build();
BatchResponse<ChatResponse> batchResponse = batchChatModel.submit(
"My Batch Job",
List.of(
ChatRequest.builder().messages(UserMessage.from("What is 2+2?")).build(),
ChatRequest.builder().messages(UserMessage.from("What is the capital of France?")).build()
)
);
System.out.println("Batch Job ID: " + batchResponse.batchId());
```
You can then retrieve the status and results of the job using `batchChatModel.retrieve(batchResponse.batchId())`.
## Grounding Metadata
If you enable Google Search grounding or use a Vertex AI Search datastore, the Google Gen AI chat model exposes the native `GroundingMetadata` directly in the `ChatResponse`. You can retrieve it through the response metadata via the underlying raw `GenerateContentResponse`.
```java
ChatModel gemini = GoogleGenAiChatModel.builder()
.apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))
.modelName("gemini-2.5-flash")
.enableGoogleSearch(true)
.build();
ChatResponse response = gemini.chat(ChatRequest.builder()
.messages(UserMessage.from("Who won the super bowl in 2024?"))
.build());
GoogleGenAiChatResponseMetadata metadata =
(GoogleGenAiChatResponseMetadata) response.metadata();
if (metadata.rawResponse() != null
&& metadata.rawResponse().candidates() != null
&& !metadata.rawResponse().candidates().isEmpty()) {
var groundingMetadata = metadata.rawResponse().candidates().get(0).groundingMetadata();
if (groundingMetadata != null && groundingMetadata.webSearchQueries() != null) {
System.out.println("Search Queries: " + groundingMetadata.webSearchQueries());
}
}
```
## Custom Labels
You can apply custom key-value labels to your Google Gen AI requests, which can be useful for billing, metrics, and tracking purposes. Custom labels are supported by:
- `GoogleGenAiChatModel`
- `GoogleGenAiStreamingChatModel`
- `GoogleGenAiBatchChatModel`
- `GoogleGenAiImageModel`
- `GoogleGenAiBatchImageModel`
```java
ChatModel gemini = GoogleGenAiChatModel.builder()
.apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))
.modelName("gemini-2.5-flash")
.labels(Map.of("environment", "production", "team", "backend"))
.build();
```
## File API
The Google Gen AI integration provides the `GoogleGenAiFiles` utility to upload and manage files on Google's servers. This is particularly useful for passing large multimodal inputs (like long videos, audio files, or extensive PDFs) that might exceed standard request limits.
```java
GoogleGenAiFiles fileApi = GoogleGenAiFiles.builder()
.apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))
.build();
String uploadedFileUri = fileApi.uploadFile(
Paths.get("path/to/my-video.mp4"),
"video/mp4",
"My Video Demo"
);
// You can now use this URI in your chat requests
ChatModel gemini = GoogleGenAiChatModel.builder()
.apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))
.modelName("gemini-2.5-flash")
.build();
ChatResponse response = gemini.chat(ChatRequest.builder()
.messages(UserMessage.from(
VideoContent.from(uploadedFileUri, "video/mp4"),
TextContent.from("What happens in this video?")
))
.build());
```
## Multimodality (Audio, Video, PDF)
The integration fully supports LangChain4j's multimodal content types. The underlying `GoogleGenAiContentMapper` automatically converts them into the appropriate Gemini `Part` objects.
```java
ChatModel gemini = GoogleGenAiChatModel.builder()
.apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))
.modelName("gemini-2.5-flash")
.build();
ChatResponse response = gemini.chat(ChatRequest.builder()
.messages(UserMessage.from(
AudioContent.from("https://example.com/audio.mp3"),
PdfFileContent.from("https://example.com/document.pdf"),
TextContent.from("Summarize the document and the audio recording.")
))
.build());
```
## Token Count Estimator
You can accurately estimate the number of tokens in your prompts and messages using the `GoogleGenAiTokenCountEstimator`, which uses the official SDK's counting endpoints.
```java
TokenCountEstimator estimator = GoogleGenAiTokenCountEstimator.builder()
.apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))
.modelName("gemini-2.5-flash")
.build();
int tokenCount = estimator.estimateTokenCount("How many tokens is this sentence?");
System.out.println("Tokens: " + tokenCount);
```
## Model Catalog
You can query the list of available Gemini models programmatically using the `GoogleGenAiModelCatalog`. This is helpful for discovering model capabilities, context windows, and supported methods dynamically.
```java
GoogleGenAiModelCatalog catalog = GoogleGenAiModelCatalog.builder()
.apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))
.build();
List<Model> availableModels = catalog.listModels();
availableModels.forEach(model -> {
System.out.println("Model Name: " + model.name());
System.out.println("Supported Generation Methods: " + model.supportedGenerationMethods());
});
```
+1 -2
View File
@@ -16,11 +16,10 @@
</properties>
<dependencies>
<dependency>
<groupId>com.google.genai</groupId>
<artifactId>google-genai</artifactId>
<version>1.53.0</version>
<version>1.57.0</version>
</dependency>
<dependency>
@@ -0,0 +1,404 @@
package dev.langchain4j.model.google.genai;
import static dev.langchain4j.internal.RetryUtils.withRetryMappingExceptions;
import static dev.langchain4j.internal.Utils.copy;
import static dev.langchain4j.internal.Utils.getOrDefault;
import static dev.langchain4j.internal.ValidationUtils.ensureNotBlank;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.genai.Client;
import com.google.genai.types.BatchJob;
import com.google.genai.types.BatchJobSource;
import com.google.genai.types.CancelBatchJobConfig;
import com.google.genai.types.Content;
import com.google.genai.types.CreateBatchJobConfig;
import com.google.genai.types.DeleteBatchJobConfig;
import com.google.genai.types.File;
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.GetBatchJobConfig;
import com.google.genai.types.InlinedRequest;
import com.google.genai.types.JobState;
import com.google.genai.types.JobState.Known;
import com.google.genai.types.SafetySetting;
import dev.langchain4j.Experimental;
import dev.langchain4j.model.batch.BatchItemResult;
import dev.langchain4j.model.batch.BatchPage;
import dev.langchain4j.model.batch.BatchPagination;
import dev.langchain4j.model.batch.BatchRequest;
import dev.langchain4j.model.batch.BatchResponse;
import dev.langchain4j.model.batch.BatchState;
import dev.langchain4j.model.chat.BatchChatModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.response.ChatResponse;
import java.time.Duration;
import java.time.Instant;
import java.time.ZoneId;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;
import java.util.stream.Stream;
/**
* Provides an interface for interacting with the Google GenAI Batch API for
* Chat models.
*/
@Experimental
public final class GoogleGenAiBatchChatModel implements BatchChatModel {
private final Client client;
private final String modelName;
private final Integer maxRetries;
// Configuration parameters reused from the chat model builder
private final List<SafetySetting> safetySettings;
private final Integer thinkingBudget;
private final String thinkingLevel;
private final Integer seed;
private final boolean googleSearchEnabled;
private final boolean googleMapsEnabled;
private final boolean urlContextEnabled;
private final List<String> allowedFunctionNames;
private final String vertexSearchDatastore;
private final Map<String, String> labels;
private final String cachedContent;
private final ChatRequestParameters defaultRequestParameters;
private GoogleGenAiBatchChatModel(Builder builder) {
this.modelName = ensureNotBlank(builder.modelName, "modelName");
this.maxRetries = getOrDefault(builder.maxRetries, 3);
this.safetySettings = copy(builder.safetySettings);
this.thinkingBudget = builder.thinkingBudget;
this.thinkingLevel = builder.thinkingLevel;
this.seed = builder.seed;
this.googleSearchEnabled = getOrDefault(builder.googleSearch, false);
this.googleMapsEnabled = getOrDefault(builder.googleMaps, false);
this.urlContextEnabled = getOrDefault(builder.urlContext, false);
this.allowedFunctionNames = copy(builder.allowedFunctionNames);
this.vertexSearchDatastore = builder.vertexSearchDatastore;
this.labels = builder.labels != null ? new HashMap<>(builder.labels) : null;
this.cachedContent = builder.cachedContent;
this.defaultRequestParameters = builder.defaultRequestParameters;
this.client = builder.client != null
? builder.client
: GoogleGenAiClientFactory.createClient(
builder.apiKey,
builder.googleCredentials,
builder.projectId,
builder.location,
builder.timeout,
builder.customHeaders,
builder.apiEndpoint);
}
public static Builder builder() {
return new Builder();
}
@Override
public BatchResponse<ChatResponse> submit(BatchRequest<ChatRequest> request) {
String timestamp = DateTimeFormatter.ofPattern("yyyy-MM-dd-HH-mm-ss")
.withZone(ZoneId.systemDefault())
.format(Instant.now());
return submit("batch-chat-job-" + timestamp, request.requests());
}
@Override
public BatchResponse<ChatResponse> retrieve(String batchId) {
BatchJob batchJob =
client.batches.get(batchId, GetBatchJobConfig.builder().build());
return processResponse(batchJob);
}
@Override
public void cancel(String batchId) {
client.batches.cancel(batchId, CancelBatchJobConfig.builder().build());
}
@Override
public BatchPage<ChatResponse> list(BatchPagination pagination) {
Integer pageSize = pagination != null ? pagination.pageSize() : null;
String pageToken = pagination != null ? pagination.pageToken() : null;
return GoogleGenAiBatchUtils.listBatchJobs(client, pageSize, pageToken, this::processResponse);
}
/**
* Creates and enqueues a batch of content generation requests for asynchronous
* processing.
* All requests must use the same model.
*
* @param displayName a user-defined name for the batch
* @param requests a list of chat requests to be processed in the batch
* @return a {@link BatchResponse} representing the initial state of the batch
* operation
*/
public BatchResponse<ChatResponse> submit(String displayName, List<ChatRequest> requests) {
validateModelInChatRequests(modelName, requests);
List<InlinedRequest> inlinedRequests =
requests.stream().map(this::createInlinedRequest).collect(Collectors.toList());
BatchJobSource src =
BatchJobSource.builder().inlinedRequests(inlinedRequests).build();
CreateBatchJobConfig config =
CreateBatchJobConfig.builder().displayName(displayName).build();
BatchJob batchJob = withRetryMappingExceptions(() -> client.batches.create(modelName, src, config), maxRetries);
return processResponse(batchJob);
}
/**
* Creates a batch of chat requests from an uploaded file.
*
* @param displayName a user-defined name for the batch
* @param file the Google GenAI File object representing the uploaded
* file containing batch requests
* @return a {@link BatchResponse} representing the initial state of the batch
* operation
*/
public BatchResponse<ChatResponse> submit(String displayName, File file) {
BatchJobSource src = BatchJobSource.builder()
.fileName(file.name().isPresent() ? file.name().get() : null)
.build();
CreateBatchJobConfig config =
CreateBatchJobConfig.builder().displayName(displayName).build();
BatchJob batchJob = withRetryMappingExceptions(() -> client.batches.create(modelName, src, config), maxRetries);
return processResponse(batchJob);
}
/**
* Deletes a batch job from the system.
*/
public void deleteBatchJob(String batchId) {
client.batches.delete(batchId, DeleteBatchJobConfig.builder().build());
}
private InlinedRequest createInlinedRequest(ChatRequest request) {
Content systemInstruction = GoogleGenAiContentMapper.toSystemInstruction(request.messages());
List<Content> contents = GoogleGenAiContentMapper.toContents(request.messages());
ChatRequestParameters params = defaultRequestParameters != null
? defaultRequestParameters.overrideWith(request.parameters())
: request.parameters();
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
params,
systemInstruction,
safetySettings,
thinkingBudget,
thinkingLevel,
seed,
googleSearchEnabled,
googleMapsEnabled,
urlContextEnabled,
allowedFunctionNames,
vertexSearchDatastore,
labels,
cachedContent);
return InlinedRequest.builder().contents(contents).config(config).build();
}
private BatchResponse<ChatResponse> processResponse(BatchJob batchJob) {
String jobName = batchJob.name().orElse("unknown");
Known state = batchJob.state().map(JobState::knownEnum).orElse(Known.JOB_STATE_UNSPECIFIED);
BatchState translatedState = GoogleGenAiBatchUtils.toBatchState(state);
BatchResponse.Builder<ChatResponse> builder =
BatchResponse.<ChatResponse>builder().batchId(jobName).state(translatedState);
if (state == Known.JOB_STATE_SUCCEEDED) {
List<BatchItemResult<ChatResponse>> results = new ArrayList<>();
if (batchJob.dest().isPresent()
&& batchJob.dest().get().inlinedResponses().isPresent()) {
var inlinedResponses = batchJob.dest().get().inlinedResponses().get();
for (var inlined : inlinedResponses) {
if (inlined.response().isPresent()) {
results.add(BatchItemResult.success(GoogleGenAiContentMapper.toChatResponse(
inlined.response().get(), batchJob.model().orElse(modelName))));
} else if (inlined.error().isPresent()) {
results.add(BatchItemResult.failure(GoogleGenAiBatchUtils.toBatchError(
inlined.error().get())));
}
}
}
builder.results(results);
} else if (state == Known.JOB_STATE_FAILED) {
builder.results(List.of(BatchItemResult.failure(
GoogleGenAiBatchUtils.toBatchError(batchJob.error().orElse(null)))));
}
return builder.build();
}
private static void validateModelInChatRequests(String modelName, List<ChatRequest> requests) {
var modelNames = Stream.concat(requests.stream().map(ChatRequest::modelName), Stream.of(modelName))
.filter(Objects::nonNull)
.collect(Collectors.toSet());
if (modelNames.size() != 1) {
throw new IllegalArgumentException(
"Batch requests cannot contain ChatRequest objects with different models; "
+ "all requests must use the same model: " + modelNames);
}
}
public static class Builder {
private Client client;
private GoogleCredentials googleCredentials;
private String apiKey;
private String projectId;
private String location;
private String modelName;
private Integer maxRetries;
private Duration timeout;
private Integer thinkingBudget;
private String thinkingLevel;
private Integer seed;
private Boolean googleSearch;
private Boolean googleMaps;
private Boolean urlContext;
private List<SafetySetting> safetySettings;
private List<String> allowedFunctionNames;
private ChatRequestParameters defaultRequestParameters;
private String vertexSearchDatastore;
private Map<String, String> labels;
private String apiEndpoint;
private Map<String, String> customHeaders;
private String cachedContent;
public Builder client(Client client) {
this.client = client;
return this;
}
public Builder googleCredentials(GoogleCredentials credentials) {
this.googleCredentials = credentials;
return this;
}
public Builder apiKey(String apiKey) {
this.apiKey = apiKey;
return this;
}
public Builder projectId(String projectId) {
this.projectId = projectId;
return this;
}
public Builder location(String location) {
this.location = location;
return this;
}
public Builder modelName(String modelName) {
this.modelName = modelName;
return this;
}
public Builder maxRetries(Integer maxRetries) {
this.maxRetries = maxRetries;
return this;
}
public Builder timeout(Duration timeout) {
this.timeout = timeout;
return this;
}
/**
* The thinking budget to use. This is a legacy parameter. For Gemini 3.x
* models, use {@link #thinkingLevel(String)} instead.
*/
public Builder thinkingBudget(Integer thinkingBudget) {
this.thinkingBudget = thinkingBudget;
return this;
}
/**
* The thinking level to use. This is the recommended parameter for Gemini 3.x
* models.
* Allowed values are {@code "MINIMAL"}, {@code "LOW"}, {@code "MEDIUM"},
* {@code "HIGH"}.
* Note that this cannot be used together with {@link #thinkingBudget(Integer)}.
*/
public Builder thinkingLevel(String thinkingLevel) {
this.thinkingLevel = thinkingLevel;
return this;
}
public Builder seed(Integer seed) {
this.seed = seed;
return this;
}
public Builder googleSearch(Boolean googleSearch) {
this.googleSearch = googleSearch;
return this;
}
public Builder googleMaps(Boolean googleMaps) {
this.googleMaps = googleMaps;
return this;
}
public Builder urlContext(Boolean urlContext) {
this.urlContext = urlContext;
return this;
}
public Builder safetySettings(List<SafetySetting> safetySettings) {
this.safetySettings = safetySettings;
return this;
}
public Builder allowedFunctionNames(List<String> allowedFunctionNames) {
this.allowedFunctionNames = allowedFunctionNames;
return this;
}
public Builder defaultRequestParameters(ChatRequestParameters defaultRequestParameters) {
this.defaultRequestParameters = defaultRequestParameters;
return this;
}
public Builder vertexSearchDatastore(String vertexSearchDatastore) {
this.vertexSearchDatastore = vertexSearchDatastore;
return this;
}
public Builder labels(Map<String, String> labels) {
this.labels = labels;
return this;
}
public Builder apiEndpoint(String apiEndpoint) {
this.apiEndpoint = apiEndpoint;
return this;
}
public Builder customHeaders(Map<String, String> customHeaders) {
this.customHeaders = customHeaders;
return this;
}
public Builder cachedContent(String cachedContent) {
this.cachedContent = cachedContent;
return this;
}
public GoogleGenAiBatchChatModel build() {
return new GoogleGenAiBatchChatModel(this);
}
}
}
@@ -0,0 +1,301 @@
package dev.langchain4j.model.google.genai;
import static dev.langchain4j.internal.RetryUtils.withRetryMappingExceptions;
import static dev.langchain4j.internal.Utils.getOrDefault;
import static dev.langchain4j.internal.ValidationUtils.ensureNotBlank;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.genai.Client;
import com.google.genai.types.BatchJob;
import com.google.genai.types.CancelBatchJobConfig;
import com.google.genai.types.Content;
import com.google.genai.types.CreateEmbeddingsBatchJobConfig;
import com.google.genai.types.DeleteBatchJobConfig;
import com.google.genai.types.EmbedContentBatch;
import com.google.genai.types.EmbedContentConfig;
import com.google.genai.types.EmbeddingsBatchJobSource;
import com.google.genai.types.File;
import com.google.genai.types.GetBatchJobConfig;
import com.google.genai.types.JobState;
import com.google.genai.types.JobState.Known;
import com.google.genai.types.Part;
import dev.langchain4j.Experimental;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.batch.BatchItemResult;
import dev.langchain4j.model.batch.BatchPage;
import dev.langchain4j.model.batch.BatchPagination;
import dev.langchain4j.model.batch.BatchRequest;
import dev.langchain4j.model.batch.BatchResponse;
import dev.langchain4j.model.batch.BatchState;
import dev.langchain4j.model.embedding.BatchEmbeddingModel;
import dev.langchain4j.model.google.genai.GoogleGenAiEmbeddingModel.TaskTypeEnum;
import dev.langchain4j.model.output.Response;
import java.time.Duration;
import java.time.Instant;
import java.time.ZoneId;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
/**
* Provides an interface for interacting with the Google GenAI Batch API for Embedding models.
*/
@Experimental
public final class GoogleGenAiBatchEmbeddingModel implements BatchEmbeddingModel {
private final Client client;
private final String modelName;
private final Integer maxRetries;
private final Integer outputDimensionality;
private final TaskTypeEnum taskType;
private final String titleMetadataKey;
private GoogleGenAiBatchEmbeddingModel(Builder builder) {
this.modelName = ensureNotBlank(builder.modelName, "modelName");
this.maxRetries = getOrDefault(builder.maxRetries, 3);
this.outputDimensionality = builder.outputDimensionality;
this.taskType = builder.taskType;
this.titleMetadataKey = builder.titleMetadataKey;
this.client = builder.client != null
? builder.client
: GoogleGenAiClientFactory.createClient(
builder.apiKey,
builder.googleCredentials,
builder.projectId,
builder.location,
builder.timeout,
builder.customHeaders,
builder.apiEndpoint);
}
public static Builder builder() {
return new Builder();
}
@Override
public BatchResponse<Response<Embedding>> submit(BatchRequest<TextSegment> request) {
String timestamp = DateTimeFormatter.ofPattern("yyyy-MM-dd-HH-mm-ss")
.withZone(ZoneId.systemDefault())
.format(Instant.now());
return submit("batch-embedding-job-" + timestamp, request.requests());
}
@Override
public BatchResponse<Response<Embedding>> retrieve(String batchId) {
BatchJob batchJob =
client.batches.get(batchId, GetBatchJobConfig.builder().build());
return processResponse(batchJob);
}
@Override
public void cancel(String batchId) {
client.batches.cancel(batchId, CancelBatchJobConfig.builder().build());
}
@Override
public BatchPage<Response<Embedding>> list(BatchPagination pagination) {
Integer pageSize = pagination != null ? pagination.pageSize() : null;
String pageToken = pagination != null ? pagination.pageToken() : null;
return GoogleGenAiBatchUtils.listBatchJobs(client, pageSize, pageToken, this::processResponse);
}
/**
* Creates and enqueues a batch of embedding requests for asynchronous processing.
*
* @param displayName a user-defined name for the batch
* @param requests a list of text segments to be embedded in the batch
* @return a {@link BatchResponse} representing the initial state of the batch operation
*/
public BatchResponse<Response<Embedding>> submit(String displayName, List<TextSegment> requests) {
List<Content> contents = requests.stream()
.map(segment -> Content.builder()
.parts(List.of(Part.builder().text(segment.text()).build()))
.build())
.collect(Collectors.toList());
EmbedContentConfig.Builder configBuilder = EmbedContentConfig.builder();
if (outputDimensionality != null) {
configBuilder.outputDimensionality(outputDimensionality);
}
if (taskType != null) {
configBuilder.taskType(taskType.getSdkTaskType());
}
EmbedContentBatch inlinedRequests = EmbedContentBatch.builder()
.contents(contents)
.config(configBuilder.build())
.build();
EmbeddingsBatchJobSource src = EmbeddingsBatchJobSource.builder()
.inlinedRequests(inlinedRequests)
.build();
CreateEmbeddingsBatchJobConfig config = CreateEmbeddingsBatchJobConfig.builder()
.displayName(displayName)
.build();
BatchJob batchJob =
withRetryMappingExceptions(() -> client.batches.createEmbeddings(modelName, src, config), maxRetries);
return processResponse(batchJob);
}
/**
* Creates a batch of embedding requests from an uploaded file.
*
* @param displayName a user-defined name for the batch
* @param file the Google GenAI File object representing the uploaded file containing batch requests
* @return a {@link BatchResponse} representing the initial state of the batch operation
*/
public BatchResponse<Response<Embedding>> submit(String displayName, File file) {
EmbeddingsBatchJobSource src = EmbeddingsBatchJobSource.builder()
.fileName(file.name().isPresent() ? file.name().get() : null)
.build();
CreateEmbeddingsBatchJobConfig config = CreateEmbeddingsBatchJobConfig.builder()
.displayName(displayName)
.build();
BatchJob batchJob =
withRetryMappingExceptions(() -> client.batches.createEmbeddings(modelName, src, config), maxRetries);
return processResponse(batchJob);
}
/**
* Deletes a batch job from the system.
*/
public void deleteBatchJob(String batchId) {
client.batches.delete(batchId, DeleteBatchJobConfig.builder().build());
}
private BatchResponse<Response<Embedding>> processResponse(BatchJob batchJob) {
String jobName = batchJob.name().orElse("unknown");
Known state = batchJob.state().map(JobState::knownEnum).orElse(Known.JOB_STATE_UNSPECIFIED);
BatchState translatedState = GoogleGenAiBatchUtils.toBatchState(state);
BatchResponse.Builder<Response<Embedding>> builder =
BatchResponse.<Response<Embedding>>builder().batchId(jobName).state(translatedState);
if (state == Known.JOB_STATE_SUCCEEDED) {
List<BatchItemResult<Response<Embedding>>> results = new ArrayList<>();
if (batchJob.dest().isPresent()
&& batchJob.dest().get().inlinedEmbedContentResponses().isPresent()) {
var inlinedResponses =
batchJob.dest().get().inlinedEmbedContentResponses().get();
for (var inlined : inlinedResponses) {
if (inlined.response().isPresent()) {
var embeddingOpt = inlined.response().get().embedding();
if (embeddingOpt.isPresent()
&& embeddingOpt.get().values().isPresent()) {
var values = embeddingOpt.get().values().get();
float[] floatArray = new float[values.size()];
for (int i = 0; i < values.size(); i++) {
floatArray[i] = values.get(i).floatValue();
}
results.add(BatchItemResult.success(Response.from(Embedding.from(floatArray))));
}
} else if (inlined.error().isPresent()) {
results.add(BatchItemResult.failure(GoogleGenAiBatchUtils.toBatchError(
inlined.error().get())));
}
}
}
builder.results(results);
} else if (state == Known.JOB_STATE_FAILED) {
builder.results(List.of(BatchItemResult.failure(
GoogleGenAiBatchUtils.toBatchError(batchJob.error().orElse(null)))));
}
return builder.build();
}
public static class Builder {
private Client client;
private GoogleCredentials googleCredentials;
private String apiKey;
private String projectId;
private String location;
private String modelName;
private Integer maxRetries;
private Duration timeout;
private Integer outputDimensionality;
private TaskTypeEnum taskType;
private String titleMetadataKey;
private String apiEndpoint;
private Map<String, String> customHeaders;
public Builder client(Client client) {
this.client = client;
return this;
}
public Builder googleCredentials(GoogleCredentials credentials) {
this.googleCredentials = credentials;
return this;
}
public Builder apiKey(String apiKey) {
this.apiKey = apiKey;
return this;
}
public Builder projectId(String projectId) {
this.projectId = projectId;
return this;
}
public Builder location(String location) {
this.location = location;
return this;
}
public Builder modelName(String modelName) {
this.modelName = ensureNotBlank(modelName, "modelName");
return this;
}
public Builder maxRetries(Integer maxRetries) {
this.maxRetries = maxRetries;
return this;
}
public Builder timeout(Duration timeout) {
this.timeout = timeout;
return this;
}
public Builder outputDimensionality(Integer outputDimensionality) {
this.outputDimensionality = outputDimensionality;
return this;
}
public Builder taskType(TaskTypeEnum taskType) {
this.taskType = taskType;
return this;
}
public Builder titleMetadataKey(String titleMetadataKey) {
this.titleMetadataKey = titleMetadataKey;
return this;
}
public Builder apiEndpoint(String apiEndpoint) {
this.apiEndpoint = apiEndpoint;
return this;
}
public Builder customHeaders(Map<String, String> customHeaders) {
this.customHeaders = customHeaders;
return this;
}
public GoogleGenAiBatchEmbeddingModel build() {
return new GoogleGenAiBatchEmbeddingModel(this);
}
}
}
@@ -0,0 +1,357 @@
package dev.langchain4j.model.google.genai;
import static dev.langchain4j.internal.RetryUtils.withRetryMappingExceptions;
import static dev.langchain4j.internal.Utils.getOrDefault;
import static dev.langchain4j.internal.ValidationUtils.ensureNotBlank;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.genai.Client;
import com.google.genai.types.BatchJob;
import com.google.genai.types.BatchJobSource;
import com.google.genai.types.CancelBatchJobConfig;
import com.google.genai.types.Content;
import com.google.genai.types.CreateBatchJobConfig;
import com.google.genai.types.DeleteBatchJobConfig;
import com.google.genai.types.File;
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.GetBatchJobConfig;
import com.google.genai.types.ImageConfig;
import com.google.genai.types.InlinedRequest;
import com.google.genai.types.JobState;
import com.google.genai.types.JobState.Known;
import com.google.genai.types.Part;
import com.google.genai.types.SafetySetting;
import dev.langchain4j.Experimental;
import dev.langchain4j.data.image.Image;
import dev.langchain4j.model.batch.BatchError;
import dev.langchain4j.model.batch.BatchItemResult;
import dev.langchain4j.model.batch.BatchPage;
import dev.langchain4j.model.batch.BatchPagination;
import dev.langchain4j.model.batch.BatchRequest;
import dev.langchain4j.model.batch.BatchResponse;
import dev.langchain4j.model.batch.BatchState;
import dev.langchain4j.model.image.BatchImageModel;
import dev.langchain4j.model.output.Response;
import java.time.Duration;
import java.time.Instant;
import java.time.ZoneId;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.Base64;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
/**
* Provides an interface for interacting with the Google GenAI Batch API for Image generation models.
*/
@Experimental
public final class GoogleGenAiBatchImageModel implements BatchImageModel {
private final Client client;
private final String modelName;
private final Integer maxRetries;
private final List<SafetySetting> safetySettings;
private final String aspectRatio;
private final String imageSize;
private final String personGeneration;
private final Map<String, String> labels;
private GoogleGenAiBatchImageModel(Builder builder) {
this.modelName = ensureNotBlank(builder.modelName, "modelName");
this.maxRetries = getOrDefault(builder.maxRetries, 3);
this.safetySettings = builder.safetySettings != null ? new ArrayList<>(builder.safetySettings) : null;
this.aspectRatio = builder.aspectRatio;
this.imageSize = builder.imageSize;
this.personGeneration = builder.personGeneration;
this.labels = builder.labels != null ? new HashMap<>(builder.labels) : null;
this.client = builder.client != null
? builder.client
: GoogleGenAiClientFactory.createClient(
builder.apiKey,
builder.googleCredentials,
builder.projectId,
builder.location,
builder.timeout,
builder.customHeaders,
builder.apiEndpoint);
}
public static Builder builder() {
return new Builder();
}
@Override
public BatchResponse<Response<Image>> submit(BatchRequest<String> request) {
String timestamp = DateTimeFormatter.ofPattern("yyyy-MM-dd-HH-mm-ss")
.withZone(ZoneId.systemDefault())
.format(Instant.now());
return submit("batch-image-job-" + timestamp, request.requests());
}
@Override
public BatchResponse<Response<Image>> retrieve(String batchId) {
BatchJob batchJob =
client.batches.get(batchId, GetBatchJobConfig.builder().build());
return processResponse(batchJob);
}
@Override
public void cancel(String batchId) {
client.batches.cancel(batchId, CancelBatchJobConfig.builder().build());
}
@Override
public BatchPage<Response<Image>> list(BatchPagination pagination) {
Integer pageSize = pagination != null ? pagination.pageSize() : null;
String pageToken = pagination != null ? pagination.pageToken() : null;
return GoogleGenAiBatchUtils.listBatchJobs(client, pageSize, pageToken, this::processResponse);
}
/**
* Creates and enqueues a batch of image generation requests for asynchronous processing.
*
* @param displayName a user-defined name for the batch
* @param prompts a list of image generation prompt strings to be processed in the batch
* @return a {@link BatchResponse} representing the initial state of the batch operation
*/
public BatchResponse<Response<Image>> submit(String displayName, List<String> prompts) {
List<InlinedRequest> inlinedRequests = prompts.stream()
.map(prompt -> createInlinedRequest(new ImageGenerationRequest(prompt)))
.collect(Collectors.toList());
BatchJobSource src =
BatchJobSource.builder().inlinedRequests(inlinedRequests).build();
CreateBatchJobConfig config =
CreateBatchJobConfig.builder().displayName(displayName).build();
BatchJob batchJob = withRetryMappingExceptions(() -> client.batches.create(modelName, src, config), maxRetries);
return processResponse(batchJob);
}
/**
* Creates a batch of image generation requests from an uploaded file.
*
* @param displayName a user-defined name for the batch
* @param file the Google GenAI File object representing the uploaded file containing batch requests
* @return a {@link BatchResponse} representing the initial state of the batch operation
*/
public BatchResponse<Response<Image>> submit(String displayName, File file) {
BatchJobSource src = BatchJobSource.builder()
.fileName(file.name().isPresent() ? file.name().get() : null)
.build();
CreateBatchJobConfig config =
CreateBatchJobConfig.builder().displayName(displayName).build();
BatchJob batchJob = withRetryMappingExceptions(() -> client.batches.create(modelName, src, config), maxRetries);
return processResponse(batchJob);
}
/**
* Deletes a batch job from the system.
*/
public void deleteBatchJob(String batchId) {
client.batches.delete(batchId, DeleteBatchJobConfig.builder().build());
}
private InlinedRequest createInlinedRequest(ImageGenerationRequest request) {
Content content = Content.builder()
.parts(List.of(Part.fromText(request.prompt())))
.build();
GenerateContentConfig.Builder configBuilder =
GenerateContentConfig.builder().responseModalities(List.of("IMAGE"));
if (safetySettings != null && !safetySettings.isEmpty()) {
configBuilder.safetySettings(safetySettings);
}
if (aspectRatio != null || imageSize != null || personGeneration != null) {
ImageConfig.Builder imageConfigBuilder = ImageConfig.builder();
if (aspectRatio != null) {
imageConfigBuilder.aspectRatio(aspectRatio);
}
if (imageSize != null) {
imageConfigBuilder.imageSize(imageSize);
}
if (personGeneration != null) {
imageConfigBuilder.personGeneration(personGeneration);
}
configBuilder.imageConfig(imageConfigBuilder.build());
}
if (labels != null && !labels.isEmpty()) {
configBuilder.labels(labels);
}
return InlinedRequest.builder()
.contents(List.of(content))
.config(configBuilder.build())
.build();
}
private BatchResponse<Response<Image>> processResponse(BatchJob batchJob) {
String jobName = batchJob.name().orElse("unknown");
Known state = batchJob.state().map(JobState::knownEnum).orElse(Known.JOB_STATE_UNSPECIFIED);
BatchState translatedState = GoogleGenAiBatchUtils.toBatchState(state);
BatchResponse.Builder<Response<Image>> builder =
BatchResponse.<Response<Image>>builder().batchId(jobName).state(translatedState);
if (state == Known.JOB_STATE_SUCCEEDED) {
List<BatchItemResult<Response<Image>>> results = new ArrayList<>();
if (batchJob.dest().isPresent()
&& batchJob.dest().get().inlinedResponses().isPresent()) {
var inlinedResponses = batchJob.dest().get().inlinedResponses().get();
for (var inlined : inlinedResponses) {
if (inlined.response().isPresent()) {
var response = inlined.response().get();
boolean imageAdded = false;
if (response.parts() != null && !response.parts().isEmpty()) {
for (Part part : response.parts()) {
if (part.inlineData().isPresent()) {
var blob = part.inlineData().get();
if (blob.data().isPresent()) {
byte[] bytes = blob.data().get();
String base64Data = Base64.getEncoder().encodeToString(bytes);
String mimeType = blob.mimeType().orElse("image/png");
Image image = Image.builder()
.base64Data(base64Data)
.mimeType(mimeType)
.build();
results.add(BatchItemResult.success(Response.from(image)));
imageAdded = true;
break; // Process one image per response for now
}
}
}
}
if (!imageAdded) {
results.add(BatchItemResult.failure(
new BatchError(0, "No image data found in response", new ArrayList<>())));
}
} else if (inlined.error().isPresent()) {
results.add(BatchItemResult.failure(GoogleGenAiBatchUtils.toBatchError(
inlined.error().get())));
}
}
}
builder.results(results);
} else if (state == Known.JOB_STATE_FAILED) {
builder.results(List.of(BatchItemResult.failure(
GoogleGenAiBatchUtils.toBatchError(batchJob.error().orElse(null)))));
}
return builder.build();
}
public record ImageGenerationRequest(String prompt) {
public ImageGenerationRequest {
ensureNotBlank(prompt, "prompt");
}
}
public static class Builder {
private Client client;
private GoogleCredentials googleCredentials;
private String apiKey;
private String projectId;
private String location;
private String modelName;
private Integer maxRetries;
private Duration timeout;
private String aspectRatio;
private String imageSize;
private String personGeneration;
private List<SafetySetting> safetySettings;
private String apiEndpoint;
private Map<String, String> customHeaders;
private Map<String, String> labels;
public Builder client(Client client) {
this.client = client;
return this;
}
public Builder googleCredentials(GoogleCredentials credentials) {
this.googleCredentials = credentials;
return this;
}
public Builder apiKey(String apiKey) {
this.apiKey = apiKey;
return this;
}
public Builder projectId(String projectId) {
this.projectId = projectId;
return this;
}
public Builder location(String location) {
this.location = location;
return this;
}
public Builder modelName(String modelName) {
this.modelName = modelName;
return this;
}
public Builder maxRetries(Integer maxRetries) {
this.maxRetries = maxRetries;
return this;
}
public Builder timeout(Duration timeout) {
this.timeout = timeout;
return this;
}
public Builder aspectRatio(String aspectRatio) {
this.aspectRatio = aspectRatio;
return this;
}
public Builder imageSize(String imageSize) {
this.imageSize = imageSize;
return this;
}
public Builder personGeneration(String personGeneration) {
this.personGeneration = personGeneration;
return this;
}
public Builder safetySettings(List<SafetySetting> safetySettings) {
this.safetySettings = safetySettings;
return this;
}
public Builder apiEndpoint(String apiEndpoint) {
this.apiEndpoint = apiEndpoint;
return this;
}
public Builder customHeaders(Map<String, String> customHeaders) {
this.customHeaders = customHeaders;
return this;
}
public Builder labels(Map<String, String> labels) {
this.labels = labels;
return this;
}
public GoogleGenAiBatchImageModel build() {
return new GoogleGenAiBatchImageModel(this);
}
}
}
@@ -0,0 +1,89 @@
package dev.langchain4j.model.google.genai;
import com.google.genai.Client;
import com.google.genai.Pager;
import com.google.genai.types.BatchJob;
import com.google.genai.types.JobError;
import com.google.genai.types.JobState;
import com.google.genai.types.ListBatchJobsConfig;
import dev.langchain4j.model.batch.BatchError;
import dev.langchain4j.model.batch.BatchPage;
import dev.langchain4j.model.batch.BatchResponse;
import dev.langchain4j.model.batch.BatchState;
import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Function;
final class GoogleGenAiBatchUtils {
private GoogleGenAiBatchUtils() {}
static <T> BatchPage<T> listBatchJobs(
Client client, Integer pageSize, String pageToken, Function<BatchJob, BatchResponse<T>> mapper) {
ListBatchJobsConfig.Builder builder = ListBatchJobsConfig.builder();
if (pageSize != null) {
builder.pageSize(pageSize);
}
if (pageToken != null) {
builder.pageToken(pageToken);
}
Pager pager = client.batches.list(builder.build());
List<BatchResponse<T>> batches = new ArrayList<>();
if (pager.page() != null) {
for (Object obj : pager.page()) {
if (obj instanceof BatchJob batchJob) {
batches.add(mapper.apply(batchJob));
}
}
}
String nextPageToken = null;
try {
// The nextPageToken field on BasePager should be accessible but is currently protected
// without a public getter in the SDK, so we retrieve it using reflection.
Field field = pager.getClass().getSuperclass().getDeclaredField("nextPageToken");
field.setAccessible(true);
nextPageToken = (String) field.get(pager);
} catch (Exception e) {
// ignore/fallback
}
return new BatchPage<>(batches, nextPageToken);
}
static BatchState toBatchState(JobState.Known state) {
if (state == null) {
return BatchState.UNSPECIFIED;
}
switch (state) {
case JOB_STATE_PENDING:
return BatchState.PENDING;
case JOB_STATE_RUNNING:
case JOB_STATE_CANCELLING:
return BatchState.RUNNING;
case JOB_STATE_SUCCEEDED:
return BatchState.SUCCEEDED;
case JOB_STATE_FAILED:
return BatchState.FAILED;
case JOB_STATE_CANCELLED:
return BatchState.CANCELLED;
case JOB_STATE_EXPIRED:
return BatchState.EXPIRED;
default:
return BatchState.UNSPECIFIED;
}
}
static BatchError toBatchError(JobError error) {
Integer code = 0;
String message = "Batch job failed";
if (error != null) {
code = error.code().orElse(0);
message = error.message().orElse("Batch job failed");
}
return new BatchError(code, message, new ArrayList<>());
}
}
@@ -21,35 +21,53 @@ import dev.langchain4j.model.chat.request.DefaultChatRequestParameters;
import dev.langchain4j.model.chat.request.ResponseFormat;
import dev.langchain4j.model.chat.response.ChatResponse;
import java.time.Duration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@Experimental
public class GoogleGenAiChatModel implements ChatModel {
private static final Logger log = LoggerFactory.getLogger(GoogleGenAiChatModel.class);
private final Client client;
private final Integer maxRetries;
private final List<ChatModelListener> listeners;
private final ChatRequestParameters defaultRequestParameters;
private final boolean logRequests;
private final boolean logResponses;
private final List<SafetySetting> safetySettings;
private final Integer thinkingBudget;
private final String thinkingLevel;
private final Integer seed;
private final boolean googleSearchEnabled;
private final boolean googleMapsEnabled;
private final boolean urlContextEnabled;
private final List<String> allowedFunctionNames;
private final String vertexSearchDatastore;
private final Map<String, String> labels;
private final String cachedContent;
private GoogleGenAiChatModel(Builder builder) {
this.maxRetries = getOrDefault(builder.maxRetries, 2);
this.logRequests = getOrDefault(builder.logRequests, false);
this.logResponses = getOrDefault(builder.logResponses, false);
this.listeners = copy(builder.listeners);
this.googleSearchEnabled = getOrDefault(builder.googleSearch, false);
this.googleMapsEnabled = getOrDefault(builder.googleMaps, false);
this.urlContextEnabled = getOrDefault(builder.urlContext, false);
this.allowedFunctionNames = copy(builder.allowedFunctionNames);
this.thinkingBudget = builder.thinkingBudget;
this.thinkingLevel = builder.thinkingLevel;
this.seed = builder.seed;
this.safetySettings = copy(builder.safetySettings);
this.vertexSearchDatastore = builder.vertexSearchDatastore;
this.labels = builder.labels != null ? new HashMap<>(builder.labels) : null;
this.cachedContent = builder.cachedContent;
this.client = builder.client != null
? builder.client
@@ -58,7 +76,9 @@ public class GoogleGenAiChatModel implements ChatModel {
builder.googleCredentials,
builder.projectId,
builder.location,
builder.timeout);
builder.timeout,
builder.customHeaders,
builder.apiEndpoint);
ChatRequestParameters commonParameters =
getOrDefault(builder.defaultRequestParameters, DefaultChatRequestParameters.EMPTY);
@@ -68,6 +88,8 @@ public class GoogleGenAiChatModel implements ChatModel {
.temperature(getOrDefault(builder.temperature, commonParameters.temperature()))
.topP(getOrDefault(builder.topP, commonParameters.topP()))
.topK(getOrDefault(builder.topK, commonParameters.topK()))
.frequencyPenalty(getOrDefault(builder.frequencyPenalty, commonParameters.frequencyPenalty()))
.presencePenalty(getOrDefault(builder.presencePenalty, commonParameters.presencePenalty()))
.maxOutputTokens(getOrDefault(builder.maxOutputTokens, commonParameters.maxOutputTokens()))
.stopSequences(getOrDefault(builder.stopSequences, commonParameters.stopSequences()))
.toolSpecifications(commonParameters.toolSpecifications())
@@ -86,16 +108,34 @@ public class GoogleGenAiChatModel implements ChatModel {
systemInstruction,
safetySettings,
thinkingBudget,
thinkingLevel,
seed,
googleSearchEnabled,
googleMapsEnabled,
urlContextEnabled,
allowedFunctionNames);
allowedFunctionNames,
vertexSearchDatastore,
labels,
cachedContent);
if (logRequests) {
log.info(
"Request:\n- model: {}\n- messages: {}\n- config: {}",
chatRequest.modelName(),
chatRequest.messages(),
config);
}
var result = withRetryMappingExceptions(
() -> client.models.generateContent(chatRequest.modelName(), contents, config), maxRetries);
return GoogleGenAiContentMapper.toChatResponse(result, chatRequest.modelName());
ChatResponse response = GoogleGenAiContentMapper.toChatResponse(result, chatRequest.modelName());
if (logResponses) {
log.info("Response:\n- model: {}\n- response: {}", chatRequest.modelName(), response);
}
return response;
}
@Override
@@ -133,8 +173,11 @@ public class GoogleGenAiChatModel implements ChatModel {
private Double temperature;
private Double topP;
private Integer topK;
private Double frequencyPenalty;
private Double presencePenalty;
private Integer maxOutputTokens;
private Integer thinkingBudget;
private String thinkingLevel;
private Integer seed;
private Integer maxRetries;
private List<String> stopSequences;
@@ -147,6 +190,13 @@ public class GoogleGenAiChatModel implements ChatModel {
private List<String> allowedFunctionNames;
private List<ChatModelListener> listeners;
private ChatRequestParameters defaultRequestParameters;
private String vertexSearchDatastore;
private Map<String, String> labels;
private String apiEndpoint;
private Map<String, String> customHeaders;
private String cachedContent;
private Boolean logRequests;
private Boolean logResponses;
public Builder client(Client client) {
this.client = client;
@@ -198,16 +248,39 @@ public class GoogleGenAiChatModel implements ChatModel {
return this;
}
public Builder frequencyPenalty(Double frequencyPenalty) {
this.frequencyPenalty = frequencyPenalty;
return this;
}
public Builder presencePenalty(Double presencePenalty) {
this.presencePenalty = presencePenalty;
return this;
}
public Builder maxOutputTokens(Integer maxOutputTokens) {
this.maxOutputTokens = maxOutputTokens;
return this;
}
/**
* The thinking budget to use. This is a legacy parameter. For Gemini 3.x models, use {@link #thinkingLevel(String)} instead.
*/
public Builder thinkingBudget(Integer thinkingBudget) {
this.thinkingBudget = thinkingBudget;
return this;
}
/**
* The thinking level to use. This is the recommended parameter for Gemini 3.x models.
* Allowed values are {@code "MINIMAL"}, {@code "LOW"}, {@code "MEDIUM"}, {@code "HIGH"}.
* Note that this cannot be used together with {@link #thinkingBudget(Integer)}.
*/
public Builder thinkingLevel(String thinkingLevel) {
this.thinkingLevel = thinkingLevel;
return this;
}
public Builder seed(Integer seed) {
this.seed = seed;
return this;
@@ -263,6 +336,47 @@ public class GoogleGenAiChatModel implements ChatModel {
return this;
}
public Builder vertexSearchDatastore(String vertexSearchDatastore) {
this.vertexSearchDatastore = vertexSearchDatastore;
return this;
}
public Builder labels(Map<String, String> labels) {
this.labels = labels;
return this;
}
public Builder apiEndpoint(String apiEndpoint) {
this.apiEndpoint = apiEndpoint;
return this;
}
public Builder customHeaders(Map<String, String> customHeaders) {
this.customHeaders = customHeaders;
return this;
}
public Builder cachedContent(String cachedContent) {
this.cachedContent = cachedContent;
return this;
}
public Builder logRequests(Boolean logRequests) {
this.logRequests = logRequests;
return this;
}
public Builder logResponses(Boolean logResponses) {
this.logResponses = logResponses;
return this;
}
public Builder logRequestsAndResponses(Boolean logRequestsAndResponses) {
this.logRequests = logRequestsAndResponses;
this.logResponses = logRequestsAndResponses;
return this;
}
public GoogleGenAiChatModel build() {
return new GoogleGenAiChatModel(this);
}
@@ -4,25 +4,46 @@ import com.google.auth.oauth2.GoogleCredentials;
import com.google.genai.Client;
import com.google.genai.types.HttpOptions;
import java.time.Duration;
import java.util.HashMap;
import java.util.Map;
class GoogleGenAiClientFactory {
static Client createClient(
String apiKey, GoogleCredentials googleCredentials, String projectId, String location, Duration timeout) {
String apiKey,
GoogleCredentials googleCredentials,
String projectId,
String location,
Duration timeout,
Map<String, String> customHeaders,
String apiEndpoint) {
HttpOptions.Builder httpOptions = HttpOptions.builder();
httpOptions.headers(Map.of("User-Agent", "LangChain4j"));
Map<String, String> headers = new HashMap<>();
if (customHeaders != null) {
headers.putAll(customHeaders);
}
headers.put("User-Agent", "LangChain4j");
httpOptions.headers(headers);
if (timeout != null) {
httpOptions.timeout((int) timeout.toMillis());
}
if (apiEndpoint != null && !apiEndpoint.isEmpty()) {
httpOptions.baseUrl(apiEndpoint);
}
Client.Builder clientBuilder = Client.builder().httpOptions(httpOptions.build());
if (googleCredentials != null) {
clientBuilder.credentials(googleCredentials);
boolean isVertex = googleCredentials != null || (projectId != null && location != null);
if (isVertex) {
clientBuilder.vertexAI(true);
if (googleCredentials != null) {
clientBuilder.credentials(googleCredentials);
}
if (projectId != null) clientBuilder.project(projectId);
if (location != null) clientBuilder.location(location);
} else if (apiKey != null) {
@@ -8,12 +8,15 @@ import com.google.genai.types.FunctionDeclaration;
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.GoogleMaps;
import com.google.genai.types.GoogleSearch;
import com.google.genai.types.Retrieval;
import com.google.genai.types.SafetySetting;
import com.google.genai.types.Schema;
import com.google.genai.types.ThinkingConfig;
import com.google.genai.types.ThinkingLevel;
import com.google.genai.types.Tool;
import com.google.genai.types.ToolConfig;
import com.google.genai.types.UrlContext;
import com.google.genai.types.VertexAISearch;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.request.ResponseFormat;
@@ -21,6 +24,7 @@ import dev.langchain4j.model.chat.request.ResponseFormatType;
import dev.langchain4j.model.chat.request.ToolChoice;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
class GoogleGenAiConfigBuilder {
@@ -29,11 +33,15 @@ class GoogleGenAiConfigBuilder {
Content systemInstruction,
List<SafetySetting> safetySettings,
Integer thinkingBudget,
String thinkingLevel,
Integer seed,
boolean googleSearchEnabled,
boolean googleMapsEnabled,
boolean urlContextEnabled,
List<String> allowedFunctionNames) {
List<String> allowedFunctionNames,
String vertexSearchDatastore,
Map<String, String> labels,
String cachedContent) {
GenerateContentConfig.Builder configBuilder = GenerateContentConfig.builder();
@@ -46,6 +54,12 @@ class GoogleGenAiConfigBuilder {
if (parameters.topK() != null) {
configBuilder.topK(parameters.topK().floatValue());
}
if (parameters.frequencyPenalty() != null) {
configBuilder.frequencyPenalty(parameters.frequencyPenalty().floatValue());
}
if (parameters.presencePenalty() != null) {
configBuilder.presencePenalty(parameters.presencePenalty().floatValue());
}
if (parameters.maxOutputTokens() != null) {
configBuilder.maxOutputTokens(parameters.maxOutputTokens());
}
@@ -67,9 +81,19 @@ class GoogleGenAiConfigBuilder {
}
}
if (thinkingBudget != null) {
configBuilder.thinkingConfig(
ThinkingConfig.builder().thinkingBudget(thinkingBudget).build());
if (thinkingBudget != null && thinkingLevel != null) {
throw new IllegalArgumentException("Cannot use both thinkingBudget and thinkingLevel at the same time");
}
if (thinkingBudget != null || thinkingLevel != null) {
ThinkingConfig.Builder thinkingBuilder = ThinkingConfig.builder();
if (thinkingBudget != null) {
thinkingBuilder.thinkingBudget(thinkingBudget);
}
if (thinkingLevel != null) {
thinkingBuilder.thinkingLevel(new ThinkingLevel(thinkingLevel));
}
configBuilder.thinkingConfig(thinkingBuilder.build());
}
if (seed != null) {
@@ -80,13 +104,22 @@ class GoogleGenAiConfigBuilder {
configBuilder.systemInstruction(systemInstruction);
}
if (labels != null) {
configBuilder.labels(labels);
}
if (cachedContent != null && !cachedContent.trim().isEmpty()) {
configBuilder.cachedContent(cachedContent);
}
buildTools(
configBuilder,
parameters,
googleSearchEnabled,
googleMapsEnabled,
urlContextEnabled,
allowedFunctionNames);
allowedFunctionNames,
vertexSearchDatastore);
return configBuilder.build();
}
@@ -97,7 +130,8 @@ class GoogleGenAiConfigBuilder {
boolean googleSearchEnabled,
boolean googleMapsEnabled,
boolean urlContextEnabled,
List<String> allowedFunctionNames) {
List<String> allowedFunctionNames,
String vertexSearchDatastore) {
List<ToolSpecification> toolSpecs = parameters.toolSpecifications();
@@ -128,6 +162,15 @@ class GoogleGenAiConfigBuilder {
requestTools.add(
Tool.builder().urlContext(UrlContext.builder().build()).build());
}
if (vertexSearchDatastore != null && !vertexSearchDatastore.isEmpty()) {
requestTools.add(Tool.builder()
.retrieval(Retrieval.builder()
.vertexAiSearch(VertexAISearch.builder()
.datastore(vertexSearchDatastore)
.build())
.build())
.build());
}
if (!requestTools.isEmpty()) {
configBuilder.tools(requestTools);
@@ -28,6 +28,7 @@ import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.message.VideoContent;
import dev.langchain4j.data.pdf.PdfFile;
import dev.langchain4j.data.video.Video;
import dev.langchain4j.exception.UnsupportedFeatureException;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.output.FinishReason;
import dev.langchain4j.model.output.TokenUsage;
@@ -37,6 +38,7 @@ import java.util.Base64;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.stream.Collectors;
class GoogleGenAiContentMapper {
@@ -118,10 +120,54 @@ class GoogleGenAiContentMapper {
}
static List<Content> toContents(List<ChatMessage> messages) {
return messages.stream()
.filter(m -> !(m instanceof SystemMessage))
.map(GoogleGenAiContentMapper::toContent)
.collect(Collectors.toList());
List<Content> contents = new ArrayList<>();
List<Part> currentFunctionParts = new ArrayList<>();
for (ChatMessage message : messages) {
if (message instanceof SystemMessage) {
continue;
}
if (message instanceof ToolExecutionResultMessage toolMsg) {
String toolResult;
try {
toolResult = toolMsg.text();
} catch (IllegalStateException e) {
throw new UnsupportedFeatureException(
"Google Gen AI currently does not support non-text content in tool execution results");
}
Map<String, Object> responseMap = new HashMap<>();
responseMap.put("result", toolResult);
FunctionResponse.Builder funcRespBuilder =
FunctionResponse.builder().name(toolMsg.toolName()).response(responseMap);
if (toolMsg.id() != null) {
funcRespBuilder.id(toolMsg.id());
}
currentFunctionParts.add(
Part.builder().functionResponse(funcRespBuilder.build()).build());
} else {
if (!currentFunctionParts.isEmpty()) {
contents.add(Content.builder()
.role(USER_ROLE)
.parts(currentFunctionParts)
.build());
currentFunctionParts = new ArrayList<>();
}
contents.add(toContent(message));
}
}
if (!currentFunctionParts.isEmpty()) {
contents.add(Content.builder()
.role(USER_ROLE)
.parts(currentFunctionParts)
.build());
}
return contents;
}
static Content toContent(ChatMessage message) {
@@ -143,36 +189,26 @@ class GoogleGenAiContentMapper {
throw new RuntimeException(e);
}
}
parts.add(Part.builder()
.functionCall(FunctionCall.builder()
.name(req.name())
.args(args)
.build())
.build());
FunctionCall.Builder fcBuilder =
FunctionCall.builder().name(req.name()).args(args);
if (req.id() != null) {
fcBuilder.id(req.id());
}
Part.Builder partBuilder = Part.builder().functionCall(fcBuilder.build());
if (req.id() != null) {
String sigBase64 = aiMsg.attribute("thought_signature_" + req.id(), String.class);
if (sigBase64 != null) {
partBuilder.thoughtSignature(Base64.getDecoder().decode(sigBase64));
}
}
parts.add(partBuilder.build());
}
}
return Content.builder().role(MODEL_ROLE).parts(parts).build();
} else if (message instanceof ToolExecutionResultMessage toolMsg) {
String toolResult;
try {
toolResult = toolMsg.text();
} catch (IllegalStateException e) {
throw new dev.langchain4j.exception.UnsupportedFeatureException(
"Google Gen AI currently does not support non-text content in tool execution results");
}
Map<String, Object> responseMap = new HashMap<>();
responseMap.put("result", toolResult);
return Content.builder()
.role(FUNCTION_ROLE)
.parts(Part.builder()
.functionResponse(FunctionResponse.builder()
.name(toolMsg.toolName())
.response(responseMap)
.build())
.build())
.build();
}
throw new IllegalArgumentException("Unknown message type: " + message.type());
}
@@ -196,6 +232,7 @@ class GoogleGenAiContentMapper {
StringBuilder textBuilder = new StringBuilder();
List<ToolExecutionRequest> toolRequests = new ArrayList<>();
Map<String, Object> attributes = new HashMap<>();
if (content != null) {
List<Part> parts = content.parts().orElse(List.of());
@@ -212,8 +249,16 @@ class GoogleGenAiContentMapper {
} catch (Exception e) {
throw new RuntimeException(e);
}
String id = fc.id().orElseGet(() -> UUID.randomUUID().toString());
if (part.thoughtSignature().isPresent()) {
byte[] sig = part.thoughtSignature().get();
attributes.put(
"thought_signature_" + id, Base64.getEncoder().encodeToString(sig));
}
toolRequests.add(ToolExecutionRequest.builder()
.id(id)
.name(fnName)
.arguments(jsonArgs)
.build());
@@ -222,15 +267,22 @@ class GoogleGenAiContentMapper {
}
String text = textBuilder.toString();
AiMessage aiMessage;
AiMessage.Builder aiMessageBuilder = AiMessage.builder();
if (!toolRequests.isEmpty() && !isNullOrEmpty(text)) {
aiMessage = new AiMessage(text, toolRequests);
aiMessageBuilder.text(text);
aiMessageBuilder.toolExecutionRequests(toolRequests);
} else if (!toolRequests.isEmpty()) {
aiMessage = AiMessage.from(toolRequests);
aiMessageBuilder.toolExecutionRequests(toolRequests);
} else {
aiMessage = AiMessage.from(text);
aiMessageBuilder.text(text);
}
if (!attributes.isEmpty()) {
aiMessageBuilder.attributes(attributes);
}
AiMessage aiMessage = aiMessageBuilder.build();
TokenUsage usage = response.usageMetadata()
.map(meta -> new TokenUsage(
meta.promptTokenCount().isPresent()
@@ -241,7 +293,10 @@ class GoogleGenAiContentMapper {
: 0))
.orElse(new TokenUsage(0, 0));
FinishReason finishReason = !toolRequests.isEmpty() ? FinishReason.TOOL_EXECUTION : FinishReason.STOP;
FinishReason finishReason = candidate
.finishReason()
.map(GoogleGenAiContentMapper::mapFinishReason)
.orElseGet(() -> !toolRequests.isEmpty() ? FinishReason.TOOL_EXECUTION : FinishReason.STOP);
GoogleGenAiChatResponseMetadata metadata = GoogleGenAiChatResponseMetadata.builder()
.modelName(modelName)
@@ -314,5 +369,34 @@ class GoogleGenAiContentMapper {
return Part.fromUri(uri.toString(), mimeType);
}
static FinishReason mapFinishReason(com.google.genai.types.FinishReason finishReason) {
if (finishReason == null) {
return FinishReason.OTHER;
}
com.google.genai.types.FinishReason.Known known = finishReason.knownEnum();
if (known == null) {
return FinishReason.OTHER;
}
switch (known) {
case STOP:
return FinishReason.STOP;
case MAX_TOKENS:
return FinishReason.LENGTH;
case SAFETY:
case RECITATION:
case BLOCKLIST:
case PROHIBITED_CONTENT:
case SPII:
case IMAGE_SAFETY:
case IMAGE_PROHIBITED_CONTENT:
case IMAGE_RECITATION:
return FinishReason.CONTENT_FILTER;
default:
return FinishReason.OTHER;
}
}
private GoogleGenAiContentMapper() {}
}
@@ -0,0 +1,300 @@
package dev.langchain4j.model.google.genai;
import static dev.langchain4j.internal.RetryUtils.withRetryMappingExceptions;
import static dev.langchain4j.internal.Utils.getOrDefault;
import static dev.langchain4j.internal.ValidationUtils.ensureNotBlank;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.genai.Client;
import com.google.genai.types.EmbedContentConfig;
import com.google.genai.types.EmbedContentResponse;
import dev.langchain4j.Experimental;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.DimensionAwareEmbeddingModel;
import dev.langchain4j.model.output.Response;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@Experimental
public class GoogleGenAiEmbeddingModel extends DimensionAwareEmbeddingModel {
private static final Logger log = LoggerFactory.getLogger(GoogleGenAiEmbeddingModel.class);
public enum TaskTypeEnum {
TASK_TYPE_UNSPECIFIED("TASK_TYPE_UNSPECIFIED"),
RETRIEVAL_QUERY("RETRIEVAL_QUERY"),
RETRIEVAL_DOCUMENT("RETRIEVAL_DOCUMENT"),
SEMANTIC_SIMILARITY("SEMANTIC_SIMILARITY"),
CLASSIFICATION("CLASSIFICATION"),
CLUSTERING("CLUSTERING"),
QUESTION_ANSWERING("QUESTION_ANSWERING"),
FACT_VERIFICATION("FACT_VERIFICATION"),
CODE_RETRIEVAL_QUERY("CODE_RETRIEVAL_QUERY");
private final String sdkTaskType;
TaskTypeEnum(String sdkTaskType) {
this.sdkTaskType = sdkTaskType;
}
public String getSdkTaskType() {
return sdkTaskType;
}
}
private final Client client;
private final String modelName;
private final Integer outputDimensionality;
private final TaskTypeEnum taskType;
private final String titleMetadataKey;
private final Integer maxSegmentsPerBatch;
private final Integer maxRetries;
private final boolean logRequests;
private final boolean logResponses;
public GoogleGenAiEmbeddingModel(Builder builder) {
this.client = builder.client != null
? builder.client
: GoogleGenAiClientFactory.createClient(
builder.apiKey,
builder.googleCredentials,
builder.projectId,
builder.location,
builder.timeout,
builder.customHeaders,
builder.apiEndpoint);
this.modelName = ensureNotBlank(builder.modelName, "modelName");
this.outputDimensionality = builder.outputDimensionality;
this.taskType = builder.taskType;
this.titleMetadataKey = getOrDefault(builder.titleMetadataKey, "title");
this.maxRetries = getOrDefault(builder.maxRetries, 3);
this.maxSegmentsPerBatch = getOrDefault(builder.maxSegmentsPerBatch, 100);
this.logRequests = getOrDefault(builder.logRequests, false);
this.logResponses = getOrDefault(builder.logResponses, false);
}
@Override
public Response<Embedding> embed(TextSegment textSegment) {
return Response.from(embedAll(List.of(textSegment)).content().get(0));
}
@Override
public Response<Embedding> embed(String text) {
return embed(TextSegment.from(text));
}
@Override
public Response<List<Embedding>> embedAll(List<TextSegment> textSegments) {
if (textSegments == null || textSegments.isEmpty()) {
return Response.from(new ArrayList<>());
}
if (logRequests) {
log.info(
"Request:\n- model: {}\n- texts: {}",
modelName,
textSegments.stream().map(TextSegment::text).collect(Collectors.toList()));
}
// Group IndexedSegment objects by their segment title (or null key if none).
// Rationale: In document ingestion pipelines, multiple text segments often share the same document title.
// Since the google-genai SDK's embedContent method with list parameters shares a single EmbedContentConfig
// (which only supports a single title), grouping segments by their title allows us to batch texts sharing
// the same title together in one API request. This maximizes API throughput and fully preserves distinct
// titles.
Map<String, List<IndexedSegment>> grouped = new LinkedHashMap<>();
for (int i = 0; i < textSegments.size(); i++) {
TextSegment segment = textSegments.get(i);
String title = null;
if (TaskTypeEnum.RETRIEVAL_DOCUMENT.equals(taskType) && segment.metadata() != null) {
title = segment.metadata().getString(titleMetadataKey);
}
grouped.computeIfAbsent(title, k -> new ArrayList<>()).add(new IndexedSegment(i, segment));
}
Embedding[] embeddingsArray = new Embedding[textSegments.size()];
for (Map.Entry<String, List<IndexedSegment>> entry : grouped.entrySet()) {
String title = entry.getKey();
List<IndexedSegment> indexedSegments = entry.getValue();
int size = indexedSegments.size();
for (int i = 0; i < size; i += maxSegmentsPerBatch) {
List<IndexedSegment> batch = indexedSegments.subList(i, Math.min(i + maxSegmentsPerBatch, size));
List<String> texts = batch.stream().map(is -> is.segment.text()).collect(Collectors.toList());
EmbedContentConfig.Builder configBuilder = EmbedContentConfig.builder();
if (taskType != null) {
configBuilder.taskType(taskType.getSdkTaskType());
}
if (outputDimensionality != null) {
configBuilder.outputDimensionality(outputDimensionality);
}
if (title != null) {
configBuilder.title(title);
}
EmbedContentResponse response = withRetryMappingExceptions(
() -> client.models.embedContent(modelName, texts, configBuilder.build()), maxRetries);
if (response.embeddings().isPresent()) {
var embeddings = response.embeddings().get();
for (int j = 0; j < batch.size(); j++) {
if (j < embeddings.size() && embeddings.get(j).values().isPresent()) {
embeddingsArray[batch.get(j).index] =
Embedding.from(embeddings.get(j).values().get());
}
}
}
}
}
Response<List<Embedding>> response = Response.from(Arrays.asList(embeddingsArray));
if (logResponses) {
log.info("Response:\n- model: {}\n- response: {}", modelName, response);
}
return response;
}
private static class IndexedSegment {
final int index;
final TextSegment segment;
IndexedSegment(int index, TextSegment segment) {
this.index = index;
this.segment = segment;
}
}
@Override
public String modelName() {
return this.modelName;
}
@Override
public Integer knownDimension() {
return outputDimensionality;
}
public static Builder builder() {
return new Builder();
}
public static class Builder {
private Client client;
private String modelName;
private String apiKey;
private GoogleCredentials googleCredentials;
private String projectId;
private String location;
private Boolean logRequests;
private Boolean logResponses;
private Duration timeout;
private Integer outputDimensionality;
private TaskTypeEnum taskType;
private String titleMetadataKey;
private String apiEndpoint;
private Map<String, String> customHeaders;
private Integer maxSegmentsPerBatch = 100;
private Integer maxRetries = 3;
public Builder client(Client client) {
this.client = client;
return this;
}
public Builder modelName(String modelName) {
this.modelName = ensureNotBlank(modelName, "modelName");
return this;
}
public Builder apiKey(String apiKey) {
this.apiKey = apiKey;
return this;
}
public Builder googleCredentials(GoogleCredentials googleCredentials) {
this.googleCredentials = googleCredentials;
return this;
}
public Builder projectId(String projectId) {
this.projectId = projectId;
return this;
}
public Builder location(String location) {
this.location = location;
return this;
}
public Builder logRequests(Boolean logRequests) {
this.logRequests = logRequests;
return this;
}
public Builder logResponses(Boolean logResponses) {
this.logResponses = logResponses;
return this;
}
public Builder logRequestsAndResponses(Boolean logRequestsAndResponses) {
this.logRequests = logRequestsAndResponses;
this.logResponses = logRequestsAndResponses;
return this;
}
public Builder timeout(Duration timeout) {
this.timeout = timeout;
return this;
}
public Builder outputDimensionality(Integer outputDimensionality) {
this.outputDimensionality = outputDimensionality;
return this;
}
public Builder taskType(TaskTypeEnum taskType) {
this.taskType = taskType;
return this;
}
public Builder titleMetadataKey(String titleMetadataKey) {
this.titleMetadataKey = titleMetadataKey;
return this;
}
public Builder apiEndpoint(String apiEndpoint) {
this.apiEndpoint = apiEndpoint;
return this;
}
public Builder customHeaders(Map<String, String> customHeaders) {
this.customHeaders = customHeaders;
return this;
}
public Builder maxSegmentsPerBatch(Integer maxSegmentsPerBatch) {
this.maxSegmentsPerBatch = maxSegmentsPerBatch;
return this;
}
public Builder maxRetries(Integer maxRetries) {
this.maxRetries = maxRetries;
return this;
}
public GoogleGenAiEmbeddingModel build() {
return new GoogleGenAiEmbeddingModel(this);
}
}
}
@@ -0,0 +1,167 @@
package dev.langchain4j.model.google.genai;
import static dev.langchain4j.internal.ValidationUtils.ensureNotBlank;
import static dev.langchain4j.internal.ValidationUtils.ensureNotNull;
import com.google.genai.Client;
import com.google.genai.types.DeleteFileConfig;
import com.google.genai.types.File;
import com.google.genai.types.GetFileConfig;
import com.google.genai.types.ListFilesConfig;
import com.google.genai.types.UploadFileConfig;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
/**
* Service for uploading and managing media files with Google AI Gemini using the official com.google.genai SDK.
*
* <p>
* The Gemini models support multimodal inputs including text, images, audio, videos, and documents.
* Use this API to upload media files when the total request size exceeds 20 MB.
*
* <p>
* Files are stored for 48 hours and can be referenced in content generation requests using their URI.
* The API supports up to 20 GB of files per project, with a maximum of 2 GB per individual file.
*/
public final class GoogleGenAiFiles {
private final Client client;
private GoogleGenAiFiles(Builder builder) {
this.client = builder.client != null
? builder.client
: GoogleGenAiClientFactory.createClient(
builder.apiKey, null, null, null, null, builder.customHeaders, builder.apiEndpoint);
}
public static Builder builder() {
return new Builder();
}
/**
* Uploads a file to Gemini.
*
* <p><strong>Note:</strong> The Files API lets you store up to 20 GB of files per project, with a per-file
* maximum size of 2 GB. Files are stored for 48 hours.
*
* @param filePath path to the file to upload
* @param displayName optional display name for the file
*/
public File uploadFile(Path filePath, String displayName) {
ensureNotNull(filePath, "filePath");
try {
UploadFileConfig config = UploadFileConfig.builder()
.displayName(
displayName != null
? displayName
: filePath.getFileName().toString())
.mimeType(detectMimeType(filePath))
.build();
return client.files.upload(filePath.toFile(), config);
} catch (IOException e) {
throw new RuntimeException("Failed to upload file", e);
}
}
/**
* Uploads a file to Gemini.
*
* <p><strong>Note:</strong> The Files API lets you store up to 20 GB of files per project, with a per-file
* maximum size of 2 GB. Files are stored for 48 hours.
*
* @param fileBytes byte array that is the file to be uploaded
* @param mimeType mimetype of the file that is being uploaded
* @param name optional display name for the file
*/
public File uploadFile(byte[] fileBytes, String mimeType, String name) {
ensureNotNull(fileBytes, "fileBytes");
ensureNotNull(mimeType, "mimeType");
ensureNotNull(name, "name");
UploadFileConfig config =
UploadFileConfig.builder().displayName(name).mimeType(mimeType).build();
return client.files.upload(fileBytes, config);
}
/**
* Retrieves metadata for a specific uploaded file.
*
* @param name the name of the file to retrieve metadata for (e.g., "files/abc123")
*/
public File getMetadata(String name) {
ensureNotBlank(name, "name");
return client.files.get(name, GetFileConfig.builder().build());
}
/**
* Lists all uploaded files. Returns a list of uploaded files.
*
* <p><strong>Note:</strong> The Files API lets you store up to 20 GB of files per project, with a per-file
* maximum size of 2 GB. Files are stored for 48 hours.
*/
public List<File> listFiles() {
List<File> allFiles = new ArrayList<>();
client.files.list(ListFilesConfig.builder().build()).forEach(allFiles::add);
return allFiles;
}
/**
* Deletes an uploaded file by name.
*
* @param name the name of the file to delete (e.g., "files/abc123")
*/
public void deleteFile(String name) {
ensureNotBlank(name, "name");
client.files.delete(name, DeleteFileConfig.builder().build());
}
/**
* Detects the MIME type of a file.
*/
private String detectMimeType(Path filePath) throws IOException {
String mimeType = Files.probeContentType(filePath);
if (mimeType == null) {
try {
mimeType = GoogleGenAiContentMapper.detectMimeType(filePath.toUri());
} catch (IllegalArgumentException e) {
// Fallback to application/octet-stream if MIME type cannot be detected
mimeType = "application/octet-stream";
}
}
return mimeType;
}
public static class Builder {
private String apiKey;
private String apiEndpoint;
private java.util.Map<String, String> customHeaders;
private Client client;
public Builder apiKey(String apiKey) {
this.apiKey = apiKey;
return this;
}
public Builder apiEndpoint(String apiEndpoint) {
this.apiEndpoint = apiEndpoint;
return this;
}
public Builder customHeaders(java.util.Map<String, String> customHeaders) {
this.customHeaders = customHeaders;
return this;
}
public Builder client(Client client) {
this.client = client;
return this;
}
public GoogleGenAiFiles build() {
return new GoogleGenAiFiles(this);
}
}
}
@@ -0,0 +1,355 @@
package dev.langchain4j.model.google.genai;
import static dev.langchain4j.internal.RetryUtils.withRetryMappingExceptions;
import static dev.langchain4j.internal.Utils.copy;
import static dev.langchain4j.internal.Utils.getOrDefault;
import static dev.langchain4j.internal.ValidationUtils.ensureNotBlank;
import static dev.langchain4j.internal.ValidationUtils.ensureNotNull;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.genai.Client;
import com.google.genai.types.Candidate;
import com.google.genai.types.Content;
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.GenerateContentResponse;
import com.google.genai.types.GoogleSearch;
import com.google.genai.types.GroundingMetadata;
import com.google.genai.types.ImageConfig;
import com.google.genai.types.Part;
import com.google.genai.types.SafetySetting;
import com.google.genai.types.Tool;
import dev.langchain4j.Experimental;
import dev.langchain4j.data.image.Image;
import dev.langchain4j.model.image.ImageModel;
import dev.langchain4j.model.output.Response;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Base64;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* Represents a Google GenAI model for image generation and editing using the official com.google.genai SDK.
*/
@Experimental
public class GoogleGenAiImageModel implements ImageModel {
private static final Logger log = LoggerFactory.getLogger(GoogleGenAiImageModel.class);
private final Client client;
private final String modelName;
private final Integer maxRetries;
private final List<SafetySetting> safetySettings;
private final boolean useGoogleSearchGrounding;
private final String aspectRatio;
private final String imageSize;
private final String personGeneration;
private final Map<String, String> labels;
private final boolean logRequests;
private final boolean logResponses;
private GoogleGenAiImageModel(Builder builder) {
this.client = builder.client != null
? builder.client
: GoogleGenAiClientFactory.createClient(
builder.apiKey,
builder.googleCredentials,
builder.projectId,
builder.location,
builder.timeout,
builder.customHeaders,
builder.apiEndpoint);
this.modelName = ensureNotBlank(builder.modelName, "modelName");
this.maxRetries = getOrDefault(builder.maxRetries, 3);
this.safetySettings = copy(builder.safetySettings);
this.useGoogleSearchGrounding = getOrDefault(builder.useGoogleSearchGrounding, false);
this.aspectRatio = builder.aspectRatio;
this.imageSize = builder.imageSize;
this.personGeneration = builder.personGeneration;
this.labels = builder.labels != null ? new HashMap<>(builder.labels) : null;
this.logRequests = getOrDefault(builder.logRequests, false);
this.logResponses = getOrDefault(builder.logResponses, false);
}
public static Builder builder() {
return new Builder();
}
@Override
public Response<Image> generate(String prompt) {
ensureNotBlank(prompt, "prompt");
Content content =
Content.builder().parts(List.of(Part.fromText(prompt))).build();
return generateImageResponse(List.of(content));
}
@Override
public Response<Image> edit(Image image, String prompt) {
ensureNotNull(image, "image");
ensureNotBlank(prompt, "prompt");
List<Part> parts = new ArrayList<>();
parts.add(Part.fromText(prompt));
parts.add(createImagePart(image));
Content content = Content.builder().parts(parts).build();
return generateImageResponse(List.of(content));
}
@Override
public Response<Image> edit(Image image, Image mask, String prompt) {
ensureNotNull(image, "image");
ensureNotNull(mask, "mask");
ensureNotBlank(prompt, "prompt");
List<Part> parts = new ArrayList<>();
parts.add(Part.fromText(prompt));
parts.add(createImagePart(image));
parts.add(createImagePart(mask));
Content content = Content.builder().parts(parts).build();
return generateImageResponse(List.of(content));
}
private Response<Image> generateImageResponse(List<Content> contents) {
GenerateContentConfig config = createGenerateContentConfig();
if (logRequests) {
log.info("Request:\n- model: {}\n- contents: {}\n- config: {}", modelName, contents, config);
}
GenerateContentResponse response = withRetryMappingExceptions(
() -> client.models.generateContent(modelName, contents, config), maxRetries);
Response<Image> imageResponse = toResponse(response);
if (logResponses) {
log.info("Response:\n- model: {}\n- response: {}", modelName, imageResponse);
}
return imageResponse;
}
private GenerateContentConfig createGenerateContentConfig() {
GenerateContentConfig.Builder configBuilder =
GenerateContentConfig.builder().responseModalities(List.of("IMAGE"));
if (!safetySettings.isEmpty()) {
configBuilder.safetySettings(safetySettings);
}
if (useGoogleSearchGrounding) {
configBuilder.tools(List.of(
Tool.builder().googleSearch(GoogleSearch.builder().build()).build()));
}
if (aspectRatio != null || imageSize != null || personGeneration != null) {
ImageConfig.Builder imageConfigBuilder = ImageConfig.builder();
if (aspectRatio != null) {
imageConfigBuilder.aspectRatio(aspectRatio);
}
if (imageSize != null) {
imageConfigBuilder.imageSize(imageSize);
}
if (personGeneration != null) {
imageConfigBuilder.personGeneration(personGeneration);
}
configBuilder.imageConfig(imageConfigBuilder.build());
}
if (labels != null && !labels.isEmpty()) {
configBuilder.labels(labels);
}
return configBuilder.build();
}
private Part createImagePart(Image image) {
String base64Data = image.base64Data();
String mimeType = image.mimeType();
if (mimeType == null || mimeType.isBlank()) {
mimeType = "image/png";
}
if (base64Data == null && image.url() != null) {
return Part.fromUri(image.url().toString(), mimeType);
}
byte[] imageBytes = Base64.getDecoder().decode(ensureNotBlank(base64Data, "image.base64Data"));
return Part.fromBytes(imageBytes, mimeType);
}
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
private Response<Image> toResponse(GenerateContentResponse response) {
if (response.parts() == null || response.parts().isEmpty()) {
throw new RuntimeException("No image generated in response");
}
Map<String, Object> metadata = new HashMap<>();
if (response.candidates().isPresent() && !response.candidates().get().isEmpty()) {
Candidate candidate = response.candidates().get().get(0);
if (candidate.groundingMetadata().isPresent()) {
GroundingMetadata gm = candidate.groundingMetadata().get();
try {
Map<String, Object> groundingMap =
OBJECT_MAPPER.readValue(gm.toJson(), new TypeReference<Map<String, Object>>() {});
metadata.put("groundingMetadata", groundingMap);
} catch (Exception e) {
throw new RuntimeException("Failed to parse grounding metadata", e);
}
}
}
for (Part part : response.parts()) {
if (part.inlineData().isPresent()) {
var blob = part.inlineData().get();
if (blob.data().isPresent()) {
byte[] bytes = blob.data().get();
String base64Data = Base64.getEncoder().encodeToString(bytes);
String mimeType = blob.mimeType().orElse("image/png");
Image image = Image.builder()
.base64Data(base64Data)
.mimeType(mimeType)
.build();
return Response.from(image, null, null, metadata);
}
}
}
throw new RuntimeException("No image data found in response");
}
public static class Builder {
private Client client;
private String apiKey;
private GoogleCredentials googleCredentials;
private String projectId;
private String location;
private Duration timeout;
private String modelName;
private Integer maxRetries;
private List<SafetySetting> safetySettings;
private Boolean useGoogleSearchGrounding;
private String aspectRatio;
private String imageSize;
private String personGeneration;
private Boolean logRequests;
private Boolean logResponses;
private String apiEndpoint;
private Map<String, String> customHeaders;
private Map<String, String> labels;
public Builder client(Client client) {
this.client = client;
return this;
}
public Builder apiKey(String apiKey) {
this.apiKey = apiKey;
return this;
}
public Builder googleCredentials(GoogleCredentials googleCredentials) {
this.googleCredentials = googleCredentials;
return this;
}
public Builder projectId(String projectId) {
this.projectId = projectId;
return this;
}
public Builder location(String location) {
this.location = location;
return this;
}
public Builder timeout(Duration timeout) {
this.timeout = timeout;
return this;
}
public Builder modelName(String modelName) {
this.modelName = modelName;
return this;
}
public Builder maxRetries(Integer maxRetries) {
this.maxRetries = maxRetries;
return this;
}
public Builder safetySettings(List<SafetySetting> safetySettings) {
this.safetySettings = safetySettings;
return this;
}
public Builder useGoogleSearchGrounding(Boolean useGoogleSearchGrounding) {
this.useGoogleSearchGrounding = useGoogleSearchGrounding;
return this;
}
public Builder aspectRatio(String aspectRatio) {
this.aspectRatio = aspectRatio;
return this;
}
public Builder imageSize(String imageSize) {
this.imageSize = imageSize;
return this;
}
public Builder personGeneration(String personGeneration) {
this.personGeneration = personGeneration;
return this;
}
public Builder logRequests(Boolean logRequests) {
this.logRequests = logRequests;
return this;
}
public Builder logResponses(Boolean logResponses) {
this.logResponses = logResponses;
return this;
}
public Builder logRequestsAndResponses(Boolean logRequestsAndResponses) {
this.logRequests = logRequestsAndResponses;
this.logResponses = logRequestsAndResponses;
return this;
}
public Builder apiEndpoint(String apiEndpoint) {
this.apiEndpoint = apiEndpoint;
return this;
}
public Builder customHeaders(Map<String, String> customHeaders) {
this.customHeaders = customHeaders;
return this;
}
public Builder labels(Map<String, String> labels) {
this.labels = labels;
return this;
}
public GoogleGenAiImageModel build() {
return new GoogleGenAiImageModel(this);
}
}
}
@@ -0,0 +1,163 @@
package dev.langchain4j.model.google.genai;
import static dev.langchain4j.model.ModelProvider.GOOGLE_GENAI;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.genai.Client;
import com.google.genai.types.ListModelsConfig;
import com.google.genai.types.Model;
import dev.langchain4j.model.ModelProvider;
import dev.langchain4j.model.catalog.ModelCatalog;
import dev.langchain4j.model.catalog.ModelDescription;
import dev.langchain4j.model.catalog.ModelType;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
/**
* Google GenAI implementation of {@link ModelCatalog}.
*
* <p>Uses the Gemini Models API to dynamically discover available models.
*
* <p>Example:
* <pre>{@code
* GoogleGenAiModelCatalog catalog = GoogleGenAiModelCatalog.builder()
* .apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))
* .build();
*
* List<ModelDescription> models = catalog.listModels();
* }</pre>
*/
public class GoogleGenAiModelCatalog implements ModelCatalog {
private final Client client;
private GoogleGenAiModelCatalog(Builder builder) {
this.client = builder.client != null
? builder.client
: GoogleGenAiClientFactory.createClient(
builder.apiKey,
builder.credentials,
builder.projectId,
builder.location,
builder.timeout,
builder.customHeaders,
builder.apiEndpoint);
}
public static Builder builder() {
return new Builder();
}
@Override
public List<ModelDescription> listModels() {
List<ModelDescription> allModels = new ArrayList<>();
client.models.list(ListModelsConfig.builder().build()).forEach(modelInfo -> {
allModels.add(mapToModelDescription(modelInfo));
});
return allModels;
}
@Override
public ModelProvider provider() {
return GOOGLE_GENAI;
}
private ModelDescription mapToModelDescription(Model modelInfo) {
ModelDescription.Builder builder = ModelDescription.builder().provider(GOOGLE_GENAI);
if (modelInfo.name().isPresent()) {
String name = modelInfo.name().get();
String id = name.startsWith("models/") ? name.substring(7) : name;
builder.name(id);
}
if (modelInfo.displayName().isPresent()
&& !modelInfo.displayName().get().isEmpty()) {
builder.displayName(modelInfo.displayName().get());
}
if (modelInfo.description().isPresent()) {
builder.description(modelInfo.description().get());
}
if (modelInfo.inputTokenLimit().isPresent()) {
builder.maxInputTokens(modelInfo.inputTokenLimit().get());
}
if (modelInfo.outputTokenLimit().isPresent()) {
builder.maxOutputTokens(modelInfo.outputTokenLimit().get());
}
// Determine model type based on supported generation methods
if (modelInfo.supportedActions().isPresent()) {
List<String> actions = modelInfo.supportedActions().get();
if (actions.contains("generateContent")) {
builder.type(ModelType.CHAT);
} else if (actions.contains("embedContent")) {
builder.type(ModelType.EMBEDDING);
}
}
return builder.build();
}
public static class Builder {
private String apiKey;
private GoogleCredentials credentials;
private String projectId;
private String location;
private Duration timeout;
private String apiEndpoint;
private Map<String, String> customHeaders;
private Client client;
public Builder apiKey(String apiKey) {
this.apiKey = apiKey;
return this;
}
public Builder credentials(GoogleCredentials credentials) {
this.credentials = credentials;
return this;
}
public Builder projectId(String projectId) {
this.projectId = projectId;
return this;
}
public Builder location(String location) {
this.location = location;
return this;
}
public Builder timeout(Duration timeout) {
this.timeout = timeout;
return this;
}
public Builder apiEndpoint(String apiEndpoint) {
this.apiEndpoint = apiEndpoint;
return this;
}
public Builder customHeaders(Map<String, String> customHeaders) {
this.customHeaders = customHeaders;
return this;
}
public Builder client(Client client) {
this.client = client;
return this;
}
public GoogleGenAiModelCatalog build() {
return new GoogleGenAiModelCatalog(this);
}
}
}
@@ -30,36 +30,54 @@ import dev.langchain4j.model.output.FinishReason;
import dev.langchain4j.model.output.TokenUsage;
import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@Experimental
public class GoogleGenAiStreamingChatModel implements StreamingChatModel {
private static final Logger log = LoggerFactory.getLogger(GoogleGenAiStreamingChatModel.class);
private final Client client;
private final List<ChatModelListener> listeners;
private final ChatRequestParameters defaultRequestParameters;
private final boolean logRequests;
private final boolean logResponses;
private final List<SafetySetting> safetySettings;
private final Integer thinkingBudget;
private final String thinkingLevel;
private final Integer seed;
private final boolean googleSearchEnabled;
private final boolean googleMapsEnabled;
private final boolean urlContextEnabled;
private final List<String> allowedFunctionNames;
private final String vertexSearchDatastore;
private final Map<String, String> labels;
private final String cachedContent;
private final ExecutorService executor;
private GoogleGenAiStreamingChatModel(Builder builder) {
this.listeners = copy(builder.listeners);
this.logRequests = getOrDefault(builder.logRequests, false);
this.logResponses = getOrDefault(builder.logResponses, false);
this.googleSearchEnabled = getOrDefault(builder.googleSearch, false);
this.googleMapsEnabled = getOrDefault(builder.googleMaps, false);
this.urlContextEnabled = getOrDefault(builder.urlContext, false);
this.allowedFunctionNames = copy(builder.allowedFunctionNames);
this.thinkingBudget = builder.thinkingBudget;
this.thinkingLevel = builder.thinkingLevel;
this.seed = builder.seed;
this.safetySettings = copy(builder.safetySettings);
this.vertexSearchDatastore = builder.vertexSearchDatastore;
this.labels = builder.labels != null ? new HashMap<>(builder.labels) : null;
this.cachedContent = builder.cachedContent;
this.client = builder.client != null
? builder.client
@@ -68,7 +86,9 @@ public class GoogleGenAiStreamingChatModel implements StreamingChatModel {
builder.googleCredentials,
builder.projectId,
builder.location,
builder.timeout);
builder.timeout,
builder.customHeaders,
builder.apiEndpoint);
ChatRequestParameters commonParameters =
getOrDefault(builder.defaultRequestParameters, DefaultChatRequestParameters.EMPTY);
@@ -78,6 +98,8 @@ public class GoogleGenAiStreamingChatModel implements StreamingChatModel {
.temperature(getOrDefault(builder.temperature, commonParameters.temperature()))
.topP(getOrDefault(builder.topP, commonParameters.topP()))
.topK(getOrDefault(builder.topK, commonParameters.topK()))
.frequencyPenalty(getOrDefault(builder.frequencyPenalty, commonParameters.frequencyPenalty()))
.presencePenalty(getOrDefault(builder.presencePenalty, commonParameters.presencePenalty()))
.maxOutputTokens(getOrDefault(builder.maxOutputTokens, commonParameters.maxOutputTokens()))
.stopSequences(getOrDefault(builder.stopSequences, commonParameters.stopSequences()))
.toolSpecifications(commonParameters.toolSpecifications())
@@ -100,18 +122,32 @@ public class GoogleGenAiStreamingChatModel implements StreamingChatModel {
systemInstruction,
safetySettings,
thinkingBudget,
thinkingLevel,
seed,
googleSearchEnabled,
googleMapsEnabled,
urlContextEnabled,
allowedFunctionNames);
allowedFunctionNames,
vertexSearchDatastore,
labels,
cachedContent);
if (logRequests) {
log.info(
"Request:\n- model: {}\n- messages: {}\n- config: {}",
chatRequest.modelName(),
chatRequest.messages(),
config);
}
executor.execute(() -> {
try (ResponseStream<GenerateContentResponse> stream =
client.models.generateContentStream(modelName, contents, config)) {
try {
ResponseStream<GenerateContentResponse> stream =
client.models.generateContentStream(modelName, contents, config);
StringBuilder textBuilder = new StringBuilder();
List<ToolExecutionRequest> toolRequests = new ArrayList<>();
Map<String, Object> attributes = new java.util.HashMap<>();
TokenUsage tokenUsage = new TokenUsage();
FinishReason finishReason = null;
GenerateContentResponse lastChunk = null;
@@ -123,6 +159,11 @@ public class GoogleGenAiStreamingChatModel implements StreamingChatModel {
ChatResponse partialResponse = GoogleGenAiContentMapper.toChatResponse(chunk, modelName);
AiMessage aiMessage = partialResponse.aiMessage();
if (aiMessage.attributes() != null
&& !aiMessage.attributes().isEmpty()) {
attributes.putAll(aiMessage.attributes());
}
if (aiMessage.text() != null && !aiMessage.text().isEmpty()) {
textBuilder.append(aiMessage.text());
try {
@@ -146,9 +187,11 @@ public class GoogleGenAiStreamingChatModel implements StreamingChatModel {
if (partialResponse.tokenUsage() != null) {
tokenUsage = partialResponse.tokenUsage();
}
if (partialResponse.finishReason() != null
&& partialResponse.finishReason() != FinishReason.OTHER) {
finishReason = partialResponse.finishReason();
FinishReason partialReason = partialResponse.finishReason();
if (partialReason != null && partialReason != FinishReason.OTHER) {
if (finishReason != FinishReason.LENGTH && finishReason != FinishReason.CONTENT_FILTER) {
finishReason = partialReason;
}
}
}
@@ -161,6 +204,11 @@ public class GoogleGenAiStreamingChatModel implements StreamingChatModel {
finalAiMessage = AiMessage.from(textBuilder.toString());
}
if (!attributes.isEmpty()) {
finalAiMessage =
finalAiMessage.toBuilder().attributes(attributes).build();
}
GoogleGenAiChatResponseMetadata metadata = GoogleGenAiChatResponseMetadata.builder()
.modelName(modelName)
.tokenUsage(tokenUsage)
@@ -176,6 +224,10 @@ public class GoogleGenAiStreamingChatModel implements StreamingChatModel {
.metadata(metadata)
.build();
if (logResponses) {
log.info("Response:\n- model: {}\n- response: {}", modelName, finalChatResponse);
}
handler.onCompleteResponse(finalChatResponse);
} catch (Exception e) {
handler.onError(e);
@@ -212,8 +264,9 @@ public class GoogleGenAiStreamingChatModel implements StreamingChatModel {
private Client client;
private GoogleCredentials googleCredentials;
private String apiKey, projectId, location, modelName;
private Double temperature, topP;
private Double temperature, topP, frequencyPenalty, presencePenalty;
private Integer topK, maxOutputTokens, thinkingBudget, seed;
private String thinkingLevel;
private List<String> stopSequences;
private Duration timeout;
private Boolean googleSearch;
@@ -225,6 +278,13 @@ public class GoogleGenAiStreamingChatModel implements StreamingChatModel {
private List<ChatModelListener> listeners;
private ExecutorService executor;
private ChatRequestParameters defaultRequestParameters;
private String vertexSearchDatastore;
private Map<String, String> labels;
private String apiEndpoint;
private Map<String, String> customHeaders;
private String cachedContent;
private Boolean logRequests;
private Boolean logResponses;
public Builder client(Client client) {
this.client = client;
@@ -276,16 +336,39 @@ public class GoogleGenAiStreamingChatModel implements StreamingChatModel {
return this;
}
public Builder frequencyPenalty(Double frequencyPenalty) {
this.frequencyPenalty = frequencyPenalty;
return this;
}
public Builder presencePenalty(Double presencePenalty) {
this.presencePenalty = presencePenalty;
return this;
}
public Builder maxOutputTokens(Integer maxOutputTokens) {
this.maxOutputTokens = maxOutputTokens;
return this;
}
/**
* The thinking budget to use. This is a legacy parameter. For Gemini 3.x models, use {@link #thinkingLevel(String)} instead.
*/
public Builder thinkingBudget(Integer thinkingBudget) {
this.thinkingBudget = thinkingBudget;
return this;
}
/**
* The thinking level to use. This is the recommended parameter for Gemini 3.x models.
* Allowed values are {@code "MINIMAL"}, {@code "LOW"}, {@code "MEDIUM"}, {@code "HIGH"}.
* Note that this cannot be used together with {@link #thinkingBudget(Integer)}.
*/
public Builder thinkingLevel(String thinkingLevel) {
this.thinkingLevel = thinkingLevel;
return this;
}
public Builder seed(Integer seed) {
this.seed = seed;
return this;
@@ -350,6 +433,47 @@ public class GoogleGenAiStreamingChatModel implements StreamingChatModel {
return this;
}
public Builder cachedContent(String cachedContent) {
this.cachedContent = cachedContent;
return this;
}
public Builder vertexSearchDatastore(String vertexSearchDatastore) {
this.vertexSearchDatastore = vertexSearchDatastore;
return this;
}
public Builder labels(Map<String, String> labels) {
this.labels = labels;
return this;
}
public Builder apiEndpoint(String apiEndpoint) {
this.apiEndpoint = apiEndpoint;
return this;
}
public Builder customHeaders(Map<String, String> customHeaders) {
this.customHeaders = customHeaders;
return this;
}
public Builder logRequests(Boolean logRequests) {
this.logRequests = logRequests;
return this;
}
public Builder logResponses(Boolean logResponses) {
this.logResponses = logResponses;
return this;
}
public Builder logRequestsAndResponses(Boolean logRequestsAndResponses) {
this.logRequests = logRequestsAndResponses;
this.logResponses = logRequestsAndResponses;
return this;
}
public GoogleGenAiStreamingChatModel build() {
return new GoogleGenAiStreamingChatModel(this);
}
@@ -0,0 +1,192 @@
package dev.langchain4j.model.google.genai;
import static dev.langchain4j.internal.RetryUtils.withRetryMappingExceptions;
import static dev.langchain4j.internal.Utils.getOrDefault;
import static dev.langchain4j.internal.ValidationUtils.ensureNotBlank;
import static java.util.Collections.singletonList;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.genai.Client;
import com.google.genai.types.Content;
import com.google.genai.types.CountTokensConfig;
import com.google.genai.types.CountTokensResponse;
import com.google.genai.types.FunctionDeclaration;
import com.google.genai.types.Tool;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.TokenCountEstimator;
import java.time.Duration;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
public class GoogleGenAiTokenCountEstimator implements TokenCountEstimator {
private final Client client;
private final String modelName;
private final Integer maxRetries;
private GoogleGenAiTokenCountEstimator(Builder builder) {
this.client = builder.client != null
? builder.client
: GoogleGenAiClientFactory.createClient(
builder.apiKey,
builder.googleCredentials,
builder.projectId,
builder.location,
builder.timeout,
builder.customHeaders,
builder.apiEndpoint);
this.modelName = ensureNotBlank(builder.modelName, "modelName");
this.maxRetries = getOrDefault(builder.maxRetries, 3);
}
public static Builder builder() {
return new Builder();
}
@Override
public int estimateTokenCountInText(String text) {
return estimateTokenCountInMessages(singletonList(UserMessage.from(text)));
}
@Override
public int estimateTokenCountInMessage(ChatMessage message) {
return estimateTokenCountInMessages(singletonList(message));
}
@Override
public int estimateTokenCountInMessages(Iterable<ChatMessage> messages) {
List<ChatMessage> allMessages = new LinkedList<>();
messages.forEach(allMessages::add);
List<Content> contents = GoogleGenAiContentMapper.toContents(allMessages);
Content systemInstruction = GoogleGenAiContentMapper.toSystemInstruction(allMessages);
if (systemInstruction != null) {
// The Java SDK currently throws an exception if `systemInstruction` is passed to CountTokensConfig.
// As a workaround, we simply append the system instruction as a standard Content block to approximate
// tokens.
List<Content> merged = new ArrayList<>();
merged.add(systemInstruction);
merged.addAll(contents);
contents = merged;
}
if (contents.isEmpty()) {
return 0;
}
return estimateTokenCount(contents, null);
}
public int estimateTokenCountInToolExecutionRequests(Iterable<ToolExecutionRequest> toolExecutionRequests) {
List<ToolExecutionRequest> allToolRequests = new LinkedList<>();
toolExecutionRequests.forEach(allToolRequests::add);
return estimateTokenCountInMessage(AiMessage.from(allToolRequests));
}
public int estimateTokenCountInToolSpecifications(Iterable<ToolSpecification> toolSpecifications) {
List<FunctionDeclaration> functionDeclarations = new ArrayList<>();
for (ToolSpecification toolSpec : toolSpecifications) {
functionDeclarations.add(GoogleGenAiToolMapper.convertToGoogleFunction(toolSpec));
}
Tool tool = Tool.builder().functionDeclarations(functionDeclarations).build();
// The Java SDK currently throws an exception if `tools` are passed to CountTokensConfig.
// As a workaround, we serialize the tool declarations to a string and count the text tokens.
String toolJson = tool.toJson();
return estimateTokenCountInText(toolJson);
}
private int estimateTokenCount(List<Content> contents, CountTokensConfig config) {
CountTokensResponse response =
withRetryMappingExceptions(() -> client.models.countTokens(modelName, contents, config), maxRetries);
return response.totalTokens().orElse(0);
}
public static class Builder {
private Client client;
private String apiKey;
private GoogleCredentials googleCredentials;
private String projectId;
private String location;
private Duration timeout;
private String modelName;
private Integer maxRetries;
private Boolean logRequests;
private Boolean logResponses;
private String apiEndpoint;
private Map<String, String> customHeaders;
public Builder client(Client client) {
this.client = client;
return this;
}
public Builder apiKey(String apiKey) {
this.apiKey = apiKey;
return this;
}
public Builder googleCredentials(GoogleCredentials googleCredentials) {
this.googleCredentials = googleCredentials;
return this;
}
public Builder projectId(String projectId) {
this.projectId = projectId;
return this;
}
public Builder location(String location) {
this.location = location;
return this;
}
public Builder timeout(Duration timeout) {
this.timeout = timeout;
return this;
}
public Builder modelName(String modelName) {
this.modelName = modelName;
return this;
}
public Builder maxRetries(Integer maxRetries) {
this.maxRetries = maxRetries;
return this;
}
public Builder logRequests(Boolean logRequests) {
this.logRequests = logRequests;
return this;
}
public Builder logResponses(Boolean logResponses) {
this.logResponses = logResponses;
return this;
}
public Builder apiEndpoint(String apiEndpoint) {
this.apiEndpoint = apiEndpoint;
return this;
}
public Builder customHeaders(Map<String, String> customHeaders) {
this.customHeaders = customHeaders;
return this;
}
public GoogleGenAiTokenCountEstimator build() {
return new GoogleGenAiTokenCountEstimator(this);
}
}
}
@@ -0,0 +1,72 @@
package dev.langchain4j.model.google.genai;
import static dev.langchain4j.model.batch.BatchState.PENDING;
import static dev.langchain4j.model.batch.BatchState.RUNNING;
import static org.assertj.core.api.Assertions.assertThat;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.batch.BatchPagination;
import dev.langchain4j.model.batch.BatchRequest;
import dev.langchain4j.model.chat.request.ChatRequest;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@EnabledIfEnvironmentVariable(named = "GOOGLE_AI_GEMINI_API_KEY", matches = ".+")
class GoogleGenAiBatchChatModelIT {
private static final String GOOGLE_AI_GEMINI_API_KEY = System.getenv("GOOGLE_AI_GEMINI_API_KEY");
@Test
void test_create_and_cancel_batch() throws InterruptedException {
GoogleGenAiBatchChatModel batchModel = GoogleGenAiBatchChatModel.builder()
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.modelName("gemini-2.5-flash")
.build();
var requests = List.of(
ChatRequest.builder()
.messages(UserMessage.from("What is the capital of France?"))
.build(),
ChatRequest.builder()
.messages(UserMessage.from("What is the capital of Germany?"))
.build());
var response = batchModel.submit(new BatchRequest<>(requests));
assertThat(response).isNotNull();
assertThat(response.batchId()).startsWith("batches/");
assertThat(response.state()).isIn(PENDING, RUNNING);
String batchId = response.batchId();
// Retrieve
var retrieved = batchModel.retrieve(batchId);
assertThat(retrieved).isNotNull();
assertThat(retrieved.batchId()).isEqualTo(batchId);
// Cancel
batchModel.cancel(batchId);
// Delete
batchModel.deleteBatchJob(batchId);
}
@Test
void test_list_batch_jobs_with_pagination() {
GoogleGenAiBatchChatModel batchModel = GoogleGenAiBatchChatModel.builder()
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.modelName("gemini-2.5-flash")
.build();
var firstPage = batchModel.list(new BatchPagination(1, null));
assertThat(firstPage).isNotNull();
assertThat(firstPage.batches()).isNotNull();
if (firstPage.nextPageToken() != null) {
var secondPage = batchModel.list(new BatchPagination(1, firstPage.nextPageToken()));
assertThat(secondPage).isNotNull();
assertThat(secondPage.batches()).isNotNull();
}
}
}
@@ -0,0 +1,64 @@
package dev.langchain4j.model.google.genai;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import com.google.common.collect.ImmutableList;
import com.google.genai.Batches;
import com.google.genai.Client;
import com.google.genai.Pager;
import com.google.genai.types.BatchJob;
import com.google.genai.types.JobState;
import com.google.genai.types.JobState.Known;
import com.google.genai.types.ListBatchJobsConfig;
import dev.langchain4j.model.batch.BatchPage;
import dev.langchain4j.model.batch.BatchPagination;
import dev.langchain4j.model.chat.response.ChatResponse;
import java.lang.reflect.Field;
import java.util.Optional;
import org.junit.jupiter.api.Test;
class GoogleGenAiBatchChatModelTest {
@Test
void should_list_batch_jobs_with_pagination() throws Exception {
Client client = mock(Client.class);
Batches batchesService = mock(Batches.class);
// Use reflection to set the public final 'batches' field on the mocked Client
Field batchesField = Client.class.getDeclaredField("batches");
batchesField.setAccessible(true);
batchesField.set(client, batchesService);
Pager pager = mock(Pager.class);
when(batchesService.list(any(ListBatchJobsConfig.class))).thenReturn(pager);
BatchJob batchJob1 = mock(BatchJob.class);
when(batchJob1.name()).thenReturn(Optional.of("batches/1"));
JobState jobState = mock(JobState.class);
when(jobState.knownEnum()).thenReturn(Known.JOB_STATE_RUNNING);
when(batchJob1.state()).thenReturn(Optional.of(jobState));
when(pager.page()).thenReturn(ImmutableList.of(batchJob1));
// Use reflection to set the protected nextPageToken on the pager superclass (BasePager)
Field field = pager.getClass().getSuperclass().getDeclaredField("nextPageToken");
field.setAccessible(true);
field.set(pager, "token-123");
GoogleGenAiBatchChatModel batchModel = GoogleGenAiBatchChatModel.builder()
.client(client)
.modelName("gemini-2.5-flash")
.build();
BatchPage<ChatResponse> response = batchModel.list(new BatchPagination(10, null));
assertThat(response).isNotNull();
assertThat(response.batches()).hasSize(1);
assertThat(response.batches().get(0).batchId()).isEqualTo("batches/1");
assertThat(response.nextPageToken()).isEqualTo("token-123");
}
}
@@ -0,0 +1,48 @@
package dev.langchain4j.model.google.genai;
import static dev.langchain4j.model.batch.BatchState.PENDING;
import static dev.langchain4j.model.batch.BatchState.RUNNING;
import static org.assertj.core.api.Assertions.assertThat;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.batch.BatchRequest;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@EnabledIfEnvironmentVariable(named = "GOOGLE_AI_GEMINI_API_KEY", matches = ".+")
class GoogleGenAiBatchEmbeddingModelIT {
private static final String GOOGLE_AI_GEMINI_API_KEY = System.getenv("GOOGLE_AI_GEMINI_API_KEY");
@Test
void test_create_and_cancel_batch() throws InterruptedException {
GoogleGenAiBatchEmbeddingModel batchModel = GoogleGenAiBatchEmbeddingModel.builder()
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.modelName("gemini-embedding-2")
.build();
var requests = List.of(
TextSegment.from("What is the capital of France?"),
TextSegment.from("What is the capital of Germany?"));
var response = batchModel.submit(new BatchRequest<>(requests));
assertThat(response).isNotNull();
assertThat(response.batchId()).startsWith("batches/");
assertThat(response.state()).isIn(PENDING, RUNNING);
String batchId = response.batchId();
// Retrieve
var retrieved = batchModel.retrieve(batchId);
assertThat(retrieved).isNotNull();
assertThat(retrieved.batchId()).isEqualTo(batchId);
// Cancel
batchModel.cancel(batchId);
// Delete
batchModel.deleteBatchJob(batchId);
}
}
@@ -0,0 +1,45 @@
package dev.langchain4j.model.google.genai;
import static dev.langchain4j.model.batch.BatchState.PENDING;
import static dev.langchain4j.model.batch.BatchState.RUNNING;
import static org.assertj.core.api.Assertions.assertThat;
import dev.langchain4j.model.batch.BatchRequest;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@EnabledIfEnvironmentVariable(named = "GOOGLE_AI_GEMINI_API_KEY", matches = ".+")
class GoogleGenAiBatchImageModelIT {
private static final String GOOGLE_AI_GEMINI_API_KEY = System.getenv("GOOGLE_AI_GEMINI_API_KEY");
@Test
void test_create_and_cancel_batch() throws InterruptedException {
GoogleGenAiBatchImageModel batchModel = GoogleGenAiBatchImageModel.builder()
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.modelName("gemini-3.1-flash-image-preview")
.build();
var requests = List.of("A picture of a cat", "A picture of a dog");
var response = batchModel.submit(new BatchRequest<>(requests));
assertThat(response).isNotNull();
assertThat(response.batchId()).startsWith("batches/");
assertThat(response.state()).isIn(PENDING, RUNNING);
String batchId = response.batchId();
// Retrieve
var retrieved = batchModel.retrieve(batchId);
assertThat(retrieved).isNotNull();
assertThat(retrieved.batchId()).isEqualTo(batchId);
// Cancel
batchModel.cancel(batchId);
// Delete
batchModel.deleteBatchJob(batchId);
}
}
@@ -3,12 +3,22 @@ package dev.langchain4j.model.google.genai;
import static dev.langchain4j.internal.Utils.getOrDefault;
import static org.assertj.core.api.Assertions.assertThat;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.common.AbstractChatModelIT;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.request.ToolChoice;
import dev.langchain4j.model.chat.request.json.JsonObjectSchema;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.ChatResponseMetadata;
import dev.langchain4j.model.output.TokenUsage;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@EnabledIfEnvironmentVariable(named = "GOOGLE_AI_GEMINI_API_KEY", matches = ".+")
@@ -45,12 +55,12 @@ class GoogleGenAiChatModelIT extends AbstractChatModelIT {
@Override
protected boolean supportsToolsAndJsonResponseFormatWithSchema() {
return false; // TODO
return true;
}
@Override
protected boolean supportsJsonResponseFormatWithRawSchema() {
return false; // TODO
return false; // TCK uses raw schema string which Gemini SDK doesn't natively map in this builder logic
}
@Override
@@ -77,4 +87,50 @@ class GoogleGenAiChatModelIT extends AbstractChatModelIT {
protected void assertOutputTokenCount(TokenUsage tokenUsage, Integer maxOutputTokens) {
assertThat(tokenUsage.outputTokenCount()).isLessThanOrEqualTo(maxOutputTokens); // TODO
}
@Test
void should_persist_thought_signature_in_multi_turn_tool_execution() {
GoogleGenAiChatModel model = GoogleGenAiChatModel.builder()
.apiKey(System.getenv("GOOGLE_AI_GEMINI_API_KEY"))
.modelName("gemini-3.1-pro-preview")
.temperature(0.0)
.build();
ToolSpecification tool = ToolSpecification.builder()
.name("get_weather")
.description("Get the weather in a city")
.parameters(JsonObjectSchema.builder()
.addStringProperty("city", "the city")
.build())
.build();
UserMessage userMsg = UserMessage.from("What is the weather in London?");
ChatRequest request = ChatRequest.builder()
.messages(List.of(userMsg))
.toolSpecifications(List.of(tool))
.toolChoice(ToolChoice.REQUIRED)
.build();
ChatResponse response1 = model.chat(request);
AiMessage aiMsg = response1.aiMessage();
assertThat(aiMsg.hasToolExecutionRequests()).isTrue();
ToolExecutionRequest toolRequest = aiMsg.toolExecutionRequests().get(0);
assertThat(toolRequest.id()).isNotNull();
ToolExecutionResultMessage toolMsg =
ToolExecutionResultMessage.from(toolRequest, "The weather is 20 degrees and sunny.");
ChatRequest request2 = ChatRequest.builder()
.messages(List.of(userMsg, aiMsg, toolMsg))
.toolSpecifications(List.of(tool))
.build();
ChatResponse response2 = model.chat(request2);
// This confirms the second request succeeded without an INVALID_ARGUMENT error
assertThat(response2.aiMessage().text()).isNotBlank();
assertThat(response2.aiMessage().hasToolExecutionRequests()).isFalse();
}
}
@@ -2,6 +2,7 @@ package dev.langchain4j.model.google.genai;
import static java.util.Collections.singletonList;
import com.google.genai.errors.ClientException;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.common.AbstractChatModelListenerIT;
import dev.langchain4j.model.chat.listener.ChatModelListener;
@@ -39,7 +40,7 @@ class GoogleGenAiChatModelListenerIT extends AbstractChatModelListenerIT {
@Override
protected Class<? extends Exception> expectedExceptionClass() {
return com.google.genai.errors.ClientException.class;
return ClientException.class;
}
@Override
@@ -19,7 +19,7 @@ class GoogleGenAiChatModelTest {
void should_build_with_api_key_and_model_name() {
GoogleGenAiChatModel model = GoogleGenAiChatModel.builder()
.apiKey("test-key")
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.build();
assertThat(model).isNotNull();
@@ -32,7 +32,7 @@ class GoogleGenAiChatModelTest {
GoogleGenAiChatModel model = GoogleGenAiChatModel.builder()
.client(client)
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.build();
assertThat(model).isNotNull();
@@ -42,10 +42,12 @@ class GoogleGenAiChatModelTest {
void should_set_default_request_parameters() {
GoogleGenAiChatModel model = GoogleGenAiChatModel.builder()
.apiKey("test-key")
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.temperature(0.7)
.topP(0.9)
.topK(40)
.frequencyPenalty(0.5)
.presencePenalty(0.3)
.maxOutputTokens(512)
.stopSequences(List.of("END"))
.build();
@@ -53,6 +55,8 @@ class GoogleGenAiChatModelTest {
assertThat(model.defaultRequestParameters().temperature()).isEqualTo(0.7);
assertThat(model.defaultRequestParameters().topP()).isEqualTo(0.9);
assertThat(model.defaultRequestParameters().topK()).isEqualTo(40);
assertThat(model.defaultRequestParameters().frequencyPenalty()).isEqualTo(0.5);
assertThat(model.defaultRequestParameters().presencePenalty()).isEqualTo(0.3);
assertThat(model.defaultRequestParameters().maxOutputTokens()).isEqualTo(512);
assertThat(model.defaultRequestParameters().stopSequences()).containsExactly("END");
}
@@ -61,7 +65,7 @@ class GoogleGenAiChatModelTest {
void should_return_empty_listeners_by_default() {
GoogleGenAiChatModel model = GoogleGenAiChatModel.builder()
.apiKey("test-key")
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.build();
assertThat(model.listeners()).isEmpty();
@@ -73,7 +77,7 @@ class GoogleGenAiChatModelTest {
GoogleGenAiChatModel model = GoogleGenAiChatModel.builder()
.apiKey("test-key")
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.listeners(List.of(listener))
.build();
@@ -84,7 +88,7 @@ class GoogleGenAiChatModelTest {
void should_always_advertise_json_schema_capability() {
GoogleGenAiChatModel model = GoogleGenAiChatModel.builder()
.apiKey("test-key")
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.build();
assertThat(model.supportedCapabilities()).containsExactly(Capability.RESPONSE_FORMAT_JSON_SCHEMA);
@@ -96,10 +100,12 @@ class GoogleGenAiChatModelTest {
GoogleGenAiChatModel model = GoogleGenAiChatModel.builder()
.client(client)
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.temperature(0.5)
.topP(0.8)
.topK(30)
.frequencyPenalty(0.5)
.presencePenalty(0.3)
.maxOutputTokens(1024)
.thinkingBudget(500)
.seed(42)
@@ -122,7 +128,7 @@ class GoogleGenAiChatModelTest {
void should_build_with_null_optional_fields() {
GoogleGenAiChatModel model = GoogleGenAiChatModel.builder()
.apiKey("test-key")
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.maxRetries(null)
.listeners(null)
.safetySettings(null)
@@ -141,6 +147,8 @@ class GoogleGenAiChatModelTest {
assertThat(builder.temperature(0.5)).isSameAs(builder);
assertThat(builder.topP(0.8)).isSameAs(builder);
assertThat(builder.topK(40)).isSameAs(builder);
assertThat(builder.frequencyPenalty(0.5)).isSameAs(builder);
assertThat(builder.presencePenalty(0.3)).isSameAs(builder);
assertThat(builder.maxOutputTokens(100)).isSameAs(builder);
assertThat(builder.thinkingBudget(500)).isSameAs(builder);
assertThat(builder.seed(42)).isSameAs(builder);
@@ -36,14 +36,14 @@ class GoogleGenAiChatResponseMetadataTest {
GoogleGenAiChatResponseMetadata metadata = GoogleGenAiChatResponseMetadata.builder()
.id("test-id")
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.tokenUsage(new TokenUsage(10, 5))
.finishReason(FinishReason.STOP)
.rawResponse(rawResponse)
.build();
assertThat(metadata.id()).isEqualTo("test-id");
assertThat(metadata.modelName()).isEqualTo("gemini-2.0-flash");
assertThat(metadata.modelName()).isEqualTo("gemini-3.1-flash-lite");
assertThat(metadata.rawResponse()).isSameAs(rawResponse);
}
@@ -109,13 +109,13 @@ class GoogleGenAiChatResponseMetadataTest {
void should_have_toString() {
GoogleGenAiChatResponseMetadata metadata = GoogleGenAiChatResponseMetadata.builder()
.id("test-id")
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.finishReason(FinishReason.STOP)
.build();
String str = metadata.toString();
assertThat(str).contains("test-id");
assertThat(str).contains("gemini-2.0-flash");
assertThat(str).contains("gemini-3.1-flash-lite");
assertThat(str).contains("STOP");
}
}
@@ -10,21 +10,22 @@ class GoogleGenAiClientFactoryTest {
@Test
void should_create_client_with_api_key() {
Client client = GoogleGenAiClientFactory.createClient("test-api-key", null, null, null, null);
Client client = GoogleGenAiClientFactory.createClient("test-api-key", null, null, null, null, null, null);
assertThat(client).isNotNull();
}
@Test
void should_create_client_with_timeout() {
Client client = GoogleGenAiClientFactory.createClient("test-api-key", null, null, null, Duration.ofSeconds(30));
Client client = GoogleGenAiClientFactory.createClient(
"test-api-key", null, null, null, Duration.ofSeconds(30), null, null);
assertThat(client).isNotNull();
}
@Test
void should_create_client_without_timeout() {
Client client = GoogleGenAiClientFactory.createClient("test-api-key", null, null, null, null);
Client client = GoogleGenAiClientFactory.createClient("test-api-key", null, null, null, null, null, null);
assertThat(client).isNotNull();
}
@@ -1,6 +1,7 @@
package dev.langchain4j.model.google.genai;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import com.google.genai.types.Content;
import com.google.genai.types.GenerateContentConfig;
@@ -14,7 +15,10 @@ import dev.langchain4j.model.chat.request.ResponseFormat;
import dev.langchain4j.model.chat.request.ResponseFormatType;
import dev.langchain4j.model.chat.request.ToolChoice;
import dev.langchain4j.model.chat.request.json.JsonObjectSchema;
import dev.langchain4j.model.chat.request.json.JsonSchema;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
class GoogleGenAiConfigBuilderTest {
@@ -25,17 +29,21 @@ class GoogleGenAiConfigBuilderTest {
.temperature(0.7)
.topP(0.9)
.topK(40)
.frequencyPenalty(0.5)
.presencePenalty(0.3)
.maxOutputTokens(1024)
.stopSequences(List.of("STOP"))
.build();
GenerateContentConfig config =
GoogleGenAiConfigBuilder.buildConfig(parameters, null, null, null, null, false, false, false, null);
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, null, false, false, false, null, null, null, null);
assertThat(config).isNotNull();
assertThat(config.temperature().get()).isEqualTo(0.7f);
assertThat(config.topP().get()).isEqualTo(0.9f);
assertThat(config.topK().get()).isEqualTo(40f);
assertThat(config.frequencyPenalty().get()).isEqualTo(0.5f);
assertThat(config.presencePenalty().get()).isEqualTo(0.3f);
assertThat(config.maxOutputTokens().get()).isEqualTo(1024);
assertThat(config.stopSequences().get()).containsExactly("STOP");
}
@@ -45,8 +53,8 @@ class GoogleGenAiConfigBuilderTest {
ChatRequestParameters parameters =
DefaultChatRequestParameters.builder().build();
GenerateContentConfig config =
GoogleGenAiConfigBuilder.buildConfig(parameters, null, null, null, null, false, false, false, null);
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, null, false, false, false, null, null, null, null);
assertThat(config).isNotNull();
}
@@ -60,7 +68,7 @@ class GoogleGenAiConfigBuilderTest {
.build());
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, safetySettings, null, null, false, false, false, null);
parameters, null, safetySettings, null, null, null, false, false, false, null, null, null, null);
assertThat(config.safetySettings().get()).hasSize(1);
}
@@ -71,7 +79,7 @@ class GoogleGenAiConfigBuilderTest {
DefaultChatRequestParameters.builder().build();
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, List.of(), null, null, false, false, false, null);
parameters, null, List.of(), null, null, null, false, false, false, null, null, null, null);
assertThat(config).isNotNull();
}
@@ -83,10 +91,32 @@ class GoogleGenAiConfigBuilderTest {
ResponseFormat.builder().type(ResponseFormatType.JSON).build())
.build();
GenerateContentConfig config =
GoogleGenAiConfigBuilder.buildConfig(parameters, null, null, null, null, false, false, false, null);
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, null, false, false, false, null, null, null, null);
assertThat(config.responseMimeType().get()).isEqualTo("application/json");
assertThat(config.responseSchema().isPresent()).isFalse();
}
@Test
void should_set_json_schema_from_response_format() {
ChatRequestParameters parameters = DefaultChatRequestParameters.builder()
.responseFormat(ResponseFormat.builder()
.type(ResponseFormatType.JSON)
.jsonSchema(JsonSchema.builder()
.rootElement(JsonObjectSchema.builder()
.addStringProperty("name")
.build())
.build())
.build())
.build();
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, null, false, false, false, null, null, null, null);
assertThat(config.responseMimeType().get()).isEqualTo("application/json");
assertThat(config.responseSchema().isPresent()).isTrue();
assertThat(config.responseSchema().get().type().get().toString()).contains("OBJECT");
}
@Test
@@ -94,19 +124,66 @@ class GoogleGenAiConfigBuilderTest {
ChatRequestParameters parameters =
DefaultChatRequestParameters.builder().build();
GenerateContentConfig config =
GoogleGenAiConfigBuilder.buildConfig(parameters, null, null, 1024, null, false, false, false, null);
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, 1024, null, null, false, false, false, null, null, null, null);
assertThat(config.thinkingConfig().get().thinkingBudget().get()).isEqualTo(1024);
}
@Test
void should_set_thinking_level() {
ChatRequestParameters parameters =
DefaultChatRequestParameters.builder().build();
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, "MEDIUM", null, false, false, false, null, null, null, null);
assertThat(config.thinkingConfig().get().thinkingLevel().get().toString())
.contains("MEDIUM");
}
@Test
void should_throw_if_both_thinking_config_are_set() {
ChatRequestParameters parameters =
DefaultChatRequestParameters.builder().build();
assertThatThrownBy(() -> GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, 1024, "MEDIUM", null, false, false, false, null, null, null, null))
.isInstanceOf(IllegalArgumentException.class)
.hasMessage("Cannot use both thinkingBudget and thinkingLevel at the same time");
}
@Test
void should_set_cached_content() {
ChatRequestParameters parameters =
DefaultChatRequestParameters.builder().build();
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters,
null,
null,
null,
null,
null,
false,
false,
false,
null,
null,
null,
"projects/123/locations/us-central1/cachedContents/456");
assertThat(config.cachedContent().isPresent()).isTrue();
assertThat(config.cachedContent().get()).isEqualTo("projects/123/locations/us-central1/cachedContents/456");
}
@Test
void should_set_seed() {
ChatRequestParameters parameters =
DefaultChatRequestParameters.builder().build();
GenerateContentConfig config =
GoogleGenAiConfigBuilder.buildConfig(parameters, null, null, null, 42, false, false, false, null);
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, 42, false, false, false, null, null, null, null);
assertThat(config.seed().get()).isEqualTo(42);
}
@@ -121,7 +198,7 @@ class GoogleGenAiConfigBuilderTest {
.build();
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, systemInstruction, null, null, null, false, false, false, null);
parameters, systemInstruction, null, null, null, null, false, false, false, null, null, null, null);
assertThat(config.systemInstruction().get().parts().get().get(0).text().get())
.isEqualTo("Be helpful");
@@ -141,8 +218,8 @@ class GoogleGenAiConfigBuilderTest {
.toolSpecifications(List.of(toolSpec))
.build();
GenerateContentConfig config =
GoogleGenAiConfigBuilder.buildConfig(parameters, null, null, null, null, false, false, false, null);
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, null, false, false, false, null, null, null, null);
assertThat(config.tools().get()).isNotEmpty();
assertThat(config.toolConfig()
@@ -167,8 +244,8 @@ class GoogleGenAiConfigBuilderTest {
.toolChoice(ToolChoice.REQUIRED)
.build();
GenerateContentConfig config =
GoogleGenAiConfigBuilder.buildConfig(parameters, null, null, null, null, false, false, false, null);
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, null, false, false, false, null, null, null, null);
assertThat(config.toolConfig()
.get()
@@ -192,8 +269,8 @@ class GoogleGenAiConfigBuilderTest {
.toolChoice(ToolChoice.NONE)
.build();
GenerateContentConfig config =
GoogleGenAiConfigBuilder.buildConfig(parameters, null, null, null, null, false, false, false, null);
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, null, false, false, false, null, null, null, null);
assertThat(config.toolConfig()
.get()
@@ -217,7 +294,7 @@ class GoogleGenAiConfigBuilderTest {
.build();
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, false, false, false, List.of("getWeather"));
parameters, null, null, null, null, null, false, false, false, List.of("getWeather"), null, null, null);
assertThat(config.toolConfig()
.get()
@@ -233,8 +310,8 @@ class GoogleGenAiConfigBuilderTest {
ChatRequestParameters parameters =
DefaultChatRequestParameters.builder().build();
GenerateContentConfig config =
GoogleGenAiConfigBuilder.buildConfig(parameters, null, null, null, null, true, false, false, null);
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, null, true, false, false, null, null, null, null);
assertThat(config.tools().get()).hasSize(1);
assertThat(config.tools().get().get(0).googleSearch().isPresent()).isTrue();
@@ -245,8 +322,8 @@ class GoogleGenAiConfigBuilderTest {
ChatRequestParameters parameters =
DefaultChatRequestParameters.builder().build();
GenerateContentConfig config =
GoogleGenAiConfigBuilder.buildConfig(parameters, null, null, null, null, false, true, false, null);
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, null, false, true, false, null, null, null, null);
assertThat(config.tools().get()).hasSize(1);
assertThat(config.tools().get().get(0).googleMaps().isPresent()).isTrue();
@@ -257,8 +334,8 @@ class GoogleGenAiConfigBuilderTest {
ChatRequestParameters parameters =
DefaultChatRequestParameters.builder().build();
GenerateContentConfig config =
GoogleGenAiConfigBuilder.buildConfig(parameters, null, null, null, null, false, false, true, null);
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, null, false, false, true, null, null, null, null);
assertThat(config.tools().get()).hasSize(1);
assertThat(config.tools().get().get(0).urlContext().isPresent()).isTrue();
@@ -269,8 +346,8 @@ class GoogleGenAiConfigBuilderTest {
ChatRequestParameters parameters =
DefaultChatRequestParameters.builder().build();
GenerateContentConfig config =
GoogleGenAiConfigBuilder.buildConfig(parameters, null, null, null, null, true, true, true, null);
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, null, true, true, true, null, null, null, null);
assertThat(config.tools().get()).hasSize(3);
assertThat(config.tools().get().get(0).googleSearch().isPresent()).isTrue();
@@ -289,8 +366,8 @@ class GoogleGenAiConfigBuilderTest {
.toolSpecifications(List.of(toolSpec))
.build();
GenerateContentConfig config =
GoogleGenAiConfigBuilder.buildConfig(parameters, null, null, null, null, true, true, true, null);
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, null, true, true, true, null, null, null, null);
assertThat(config.tools().get()).hasSize(4);
assertThat(config.tools().get().get(0).functionDeclarations().get()).isNotEmpty();
@@ -310,8 +387,8 @@ class GoogleGenAiConfigBuilderTest {
.toolSpecifications(List.of(toolSpec))
.build();
GenerateContentConfig config =
GoogleGenAiConfigBuilder.buildConfig(parameters, null, null, null, null, true, false, false, null);
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, null, true, false, false, null, null, null, null);
assertThat(config.tools().get()).hasSize(2);
assertThat(config.tools().get().get(0).functionDeclarations().get()).isNotEmpty();
@@ -325,9 +402,67 @@ class GoogleGenAiConfigBuilderTest {
ResponseFormat.builder().type(ResponseFormatType.TEXT).build())
.build();
GenerateContentConfig config =
GoogleGenAiConfigBuilder.buildConfig(parameters, null, null, null, null, false, false, false, null);
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, null, false, false, false, null, null, null, null);
assertThat(config).isNotNull();
}
@Test
void should_add_vertex_search_datastore_tool() {
ChatRequestParameters parameters =
DefaultChatRequestParameters.builder().build();
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters,
null,
null,
null,
null,
null,
false,
false,
false,
null,
"projects/123/locations/global/collections/default_collection/dataStores/my-datastore",
null,
null);
assertThat(config.tools().get()).hasSize(1);
assertThat(config.tools().get().get(0).retrieval().isPresent()).isTrue();
assertThat(config.tools()
.get()
.get(0)
.retrieval()
.get()
.vertexAiSearch()
.isPresent())
.isTrue();
assertThat(config.tools()
.get()
.get(0)
.retrieval()
.get()
.vertexAiSearch()
.get()
.datastore()
.get())
.isEqualTo("projects/123/locations/global/collections/default_collection/dataStores/my-datastore");
}
@Test
void should_set_labels() {
ChatRequestParameters parameters =
DefaultChatRequestParameters.builder().build();
Map<String, String> labels = new HashMap<>();
labels.put("env", "prod");
labels.put("team", "billing");
GenerateContentConfig config = GoogleGenAiConfigBuilder.buildConfig(
parameters, null, null, null, null, null, false, false, false, null, null, labels, null);
assertThat(config.labels().isPresent()).isTrue();
assertThat(config.labels().get()).containsEntry("env", "prod").containsEntry("team", "billing");
}
}
@@ -190,14 +190,29 @@ class GoogleGenAiContentMapperTest {
}
@Test
void should_convert_tool_execution_result_message() {
ToolExecutionResultMessage message = ToolExecutionResultMessage.from("call-1", "getWeather", "Sunny, 25C");
void should_convert_parallel_tool_execution_result_messages() {
List<ChatMessage> messages = List.of(
ToolExecutionResultMessage.from("call-1", "getWeather", "Sunny, 25C"),
ToolExecutionResultMessage.from("call-2", "getWeather", "Rainy, 15C"));
Content result = GoogleGenAiContentMapper.toContent(message);
List<Content> results = GoogleGenAiContentMapper.toContents(messages);
assertThat(result.role().get()).isEqualTo("function");
assertThat(result.parts().get().get(0).functionResponse().get().name().get())
.isEqualTo("getWeather");
assertThat(results).hasSize(1);
Content result = results.get(0);
assertThat(result.role().get()).isEqualTo("user");
List<Part> parts = result.parts().get();
assertThat(parts).hasSize(2);
assertThat(parts.get(0).functionResponse().get().name().get()).isEqualTo("getWeather");
assertThat(parts.get(0).functionResponse().get().response().get().get("result"))
.isEqualTo("Sunny, 25C");
assertThat(parts.get(0).functionResponse().get().id().get()).isEqualTo("call-1");
assertThat(parts.get(1).functionResponse().get().name().get()).isEqualTo("getWeather");
assertThat(parts.get(1).functionResponse().get().response().get().get("result"))
.isEqualTo("Rainy, 15C");
assertThat(parts.get(1).functionResponse().get().id().get()).isEqualTo("call-2");
}
@Test
@@ -459,4 +474,32 @@ class GoogleGenAiContentMapperTest {
assertThat(result.parts().get()).hasSize(2);
assertThat(result.parts().get().get(0).text().get()).isEqualTo("Describe this image");
}
@Test
void should_map_finish_reason() {
assertThat(GoogleGenAiContentMapper.mapFinishReason(
new com.google.genai.types.FinishReason(com.google.genai.types.FinishReason.Known.STOP)))
.isEqualTo(FinishReason.STOP);
assertThat(GoogleGenAiContentMapper.mapFinishReason(
new com.google.genai.types.FinishReason(com.google.genai.types.FinishReason.Known.MAX_TOKENS)))
.isEqualTo(FinishReason.LENGTH);
assertThat(GoogleGenAiContentMapper.mapFinishReason(new com.google.genai.types.FinishReason(
com.google.genai.types.FinishReason.Known.IMAGE_RECITATION)))
.isEqualTo(FinishReason.CONTENT_FILTER);
assertThat(GoogleGenAiContentMapper.mapFinishReason(
new com.google.genai.types.FinishReason(com.google.genai.types.FinishReason.Known.SAFETY)))
.isEqualTo(FinishReason.CONTENT_FILTER);
assertThat(GoogleGenAiContentMapper.mapFinishReason(
new com.google.genai.types.FinishReason(com.google.genai.types.FinishReason.Known.OTHER)))
.isEqualTo(FinishReason.OTHER);
}
@Test
void should_map_null_finish_reason_to_other() {
assertThat(GoogleGenAiContentMapper.mapFinishReason(null)).isEqualTo(FinishReason.OTHER);
}
}
@@ -0,0 +1,187 @@
package dev.langchain4j.model.google.genai;
import static org.assertj.core.api.Assertions.assertThat;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.output.Response;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@EnabledIfEnvironmentVariable(named = "GOOGLE_AI_GEMINI_API_KEY", matches = ".+")
class GoogleGenAiEmbeddingModelIT {
private static final String GOOGLE_AI_GEMINI_API_KEY = System.getenv("GOOGLE_AI_GEMINI_API_KEY");
@Test
void should_embed_one_text() {
// given
GoogleGenAiEmbeddingModel embeddingModel = GoogleGenAiEmbeddingModel.builder()
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.modelName("gemini-embedding-2")
.logRequests(true)
.logResponses(false) // embeddings are huge in logs
.build();
// when
Response<Embedding> embed = embeddingModel.embed("Hello world!");
// then
Embedding content = embed.content();
assertThat(content).isNotNull();
assertThat(content.vector()).isNotNull();
// gemini-embedding-2 generates 3072-dimensional embeddings by default
assertThat(content.vector()).hasSize(3072);
assertThat(embeddingModel.dimension()).isEqualTo(3072);
}
@Test
void should_use_metadata() {
// given
GoogleGenAiEmbeddingModel embeddingModel = GoogleGenAiEmbeddingModel.builder()
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.modelName("gemini-embedding-2")
.logRequests(false) // embeddings are huge in logs
.logResponses(false)
.titleMetadataKey("title")
.taskType(GoogleGenAiEmbeddingModel.TaskTypeEnum.RETRIEVAL_DOCUMENT)
.build();
// when
TextSegment textSegment =
TextSegment.from("What is the capital of France?", Metadata.from("title", "document title"));
Response<Embedding> embed = embeddingModel.embed(textSegment);
// then
Embedding content = embed.content();
assertThat(content).isNotNull();
assertThat(content.vector()).isNotNull();
}
@Test
void should_embed_in_batch() {
// given
GoogleGenAiEmbeddingModel embeddingModel = GoogleGenAiEmbeddingModel.builder()
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.modelName("gemini-embedding-2")
.logRequests(false) // embeddings are huge in logs
.logResponses(false)
.build();
// when
List<TextSegment> textSegments = Arrays.asList(
TextSegment.from("What is the capital of France?"),
TextSegment.from("What is the capital of Germany?"));
Response<List<Embedding>> embed = embeddingModel.embedAll(textSegments);
// then
List<Embedding> embeddings = embed.content();
assertThat(embeddings).isNotNull().hasSize(2);
assertThat(embeddings.get(0).vector()).isNotNull();
assertThat(embeddings.get(1).vector()).isNotNull();
}
@Test
void should_embed_with_dimensionality_of_512() {
// given
int outputDimensionality = 512;
GoogleGenAiEmbeddingModel embeddingModel = GoogleGenAiEmbeddingModel.builder()
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.modelName("gemini-embedding-2")
.logRequests(false)
.logResponses(false)
.outputDimensionality(outputDimensionality)
.build();
// when
Response<Embedding> embed = embeddingModel.embed("What is the capital of France?");
// then
Embedding content = embed.content();
assertThat(content).isNotNull();
assertThat(content.vector()).isNotNull();
assertThat(content.vector()).hasSize(outputDimensionality);
assertThat(embeddingModel.dimension()).isEqualTo(outputDimensionality);
}
@Test
void should_embed_more_than_100() {
// given
GoogleGenAiEmbeddingModel embeddingModel = GoogleGenAiEmbeddingModel.builder()
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.modelName("gemini-embedding-2")
.build();
// when
List<TextSegment> textSegments = new ArrayList<>();
for (int i = 0; i < 150; i++) {
// Using 150 to test a moderate batch. The API has its own limits.
textSegments.add(TextSegment.from("What is the capital of France? "));
}
Response<List<Embedding>> allEmbeddings = embeddingModel.embedAll(textSegments);
// then
assertThat(allEmbeddings.content()).hasSize(150);
}
@Test
void should_embed_in_batches_of_custom_size() {
// given
GoogleGenAiEmbeddingModel embeddingModel = GoogleGenAiEmbeddingModel.builder()
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.modelName("gemini-embedding-2")
.maxSegmentsPerBatch(10)
.build();
// when
List<TextSegment> textSegments = new ArrayList<>();
for (int i = 0; i < 25; i++) {
textSegments.add(TextSegment.from("Segment " + i));
}
Response<List<Embedding>> allEmbeddings = embeddingModel.embedAll(textSegments);
// then
assertThat(allEmbeddings.content()).hasSize(25);
}
@Test
void should_embed_with_title_grouping() {
// given
GoogleGenAiEmbeddingModel embeddingModel = GoogleGenAiEmbeddingModel.builder()
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.modelName("gemini-embedding-2")
.taskType(GoogleGenAiEmbeddingModel.TaskTypeEnum.RETRIEVAL_DOCUMENT)
.titleMetadataKey("title")
.maxSegmentsPerBatch(5)
.build();
// when
List<TextSegment> textSegments = new ArrayList<>();
textSegments.add(TextSegment.from("Document 1 chunk 1", Metadata.from("title", "Doc1")));
textSegments.add(TextSegment.from("Document 1 chunk 2", Metadata.from("title", "Doc1")));
textSegments.add(TextSegment.from("Document 2 chunk 1", Metadata.from("title", "Doc2")));
textSegments.add(TextSegment.from("Document 2 chunk 2", Metadata.from("title", "Doc2")));
textSegments.add(TextSegment.from("No title chunk"));
Response<List<Embedding>> allEmbeddings = embeddingModel.embedAll(textSegments);
// then
assertThat(allEmbeddings.content()).hasSize(5);
for (Embedding embedding : allEmbeddings.content()) {
assertThat(embedding.vector()).isNotNull();
}
}
}
@@ -0,0 +1,46 @@
package dev.langchain4j.model.google.genai;
import static org.assertj.core.api.Assertions.assertThat;
import com.google.genai.types.File;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@EnabledIfEnvironmentVariable(named = "GOOGLE_AI_GEMINI_API_KEY", matches = ".+")
class GoogleGenAiFilesIT {
private static final String GOOGLE_AI_GEMINI_API_KEY = System.getenv("GOOGLE_AI_GEMINI_API_KEY");
@Test
void test_file_upload_get_list_delete() throws Exception {
GoogleGenAiFiles files =
GoogleGenAiFiles.builder().apiKey(GOOGLE_AI_GEMINI_API_KEY).build();
Path tempFile = Files.createTempFile("test-gemini", ".txt");
Files.writeString(tempFile, "Hello World Gemini File API");
try {
File file = files.uploadFile(tempFile, "My Test File");
assertThat(file).isNotNull();
assertThat(file.name().isPresent()).isTrue();
String fileName = file.name().get();
File metadata = files.getMetadata(fileName);
assertThat(metadata.displayName().orElse("")).isEqualTo("My Test File");
List<File> allFiles = files.listFiles();
assertThat(allFiles).isNotEmpty();
boolean found = allFiles.stream()
.anyMatch(f -> f.name().isPresent() && f.name().get().equals(fileName));
assertThat(found).isTrue();
files.deleteFile(fileName);
} finally {
Files.deleteIfExists(tempFile);
}
}
}
@@ -0,0 +1,183 @@
package dev.langchain4j.model.google.genai;
import static org.assertj.core.api.Assertions.assertThat;
import dev.langchain4j.data.image.Image;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Base64;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@EnabledIfEnvironmentVariable(named = "GOOGLE_AI_GEMINI_API_KEY", matches = ".+")
class GoogleGenAiImageModelIT {
private static final String GOOGLE_AI_GEMINI_API_KEY = System.getenv("GOOGLE_AI_GEMINI_API_KEY");
private static final String MODEL_NAME = "gemini-3.1-flash-image-preview";
private static final String NANO_BANANA_PRO = "gemini-3-pro-image-preview";
private static final Path OUTPUT_DIR = Paths.get("target", "test-images");
@BeforeAll
static void setUp() throws IOException {
Files.createDirectories(OUTPUT_DIR);
}
@Test
void should_generate_single_image() throws IOException {
// given
var subject = GoogleGenAiImageModel.builder()
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.modelName(MODEL_NAME)
.logRequests(true)
.aspectRatio("1:1")
.build();
// when
var response = subject.generate("Image of A simple red circle on a white background");
// then
var image = response.content();
assertThat(image).isNotNull();
assertThat(image.base64Data()).isNotBlank();
assertThat(image.mimeType()).startsWith("image/");
saveImage(image, "should_generate_single_image");
}
@Test
void should_edit_image() throws IOException {
// given
var subject = GoogleGenAiImageModel.builder()
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.modelName(MODEL_NAME)
.build();
// First generate an image to edit
var originalResponse = subject.generate("A simple blue square on a white background");
Image originalImage = originalResponse.content();
saveImage(originalImage, "should_edit_image_original");
// when - edit the generated image
var editedResponse = subject.edit(originalImage, "Change the blue square to a green triangle");
// then
assertThat(editedResponse).isNotNull();
assertThat(editedResponse.content()).isNotNull();
assertThat(editedResponse.content().base64Data()).isNotBlank();
assertThat(editedResponse.content().mimeType()).startsWith("image/");
assertThat(editedResponse.content().base64Data()).isNotEqualTo(originalImage.base64Data());
saveImage(editedResponse.content(), "should_edit_image_result");
}
@Test
@SuppressWarnings("unchecked")
void should_ground_image_in_search() throws IOException {
// given
var subject = GoogleGenAiImageModel.builder()
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.modelName(NANO_BANANA_PRO)
.useGoogleSearchGrounding(true)
.aspectRatio("1:1")
.build();
// when
var imageResponse = subject.generate("""
A kawaii illustration of the current weather forecast for Paris (France)
showing the current temperature (in Celsius)
""");
saveImage(imageResponse.content(), "paris_weather_illustration");
// then
assertThat(imageResponse).isNotNull();
assertThat(imageResponse.content()).isNotNull();
assertThat(imageResponse.content().base64Data()).isNotBlank();
assertThat(imageResponse.content().mimeType()).startsWith("image/");
assertThat(imageResponse.metadata().get("groundingMetadata")).isNotNull();
Map<String, Object> groundingMetadata =
(Map<String, Object>) imageResponse.metadata().get("groundingMetadata");
assertThat(groundingMetadata).isNotNull();
assertThat(groundingMetadata).containsKey("webSearchQueries");
List<String> webSearchQueries = (List<String>) groundingMetadata.get("webSearchQueries");
assertThat(webSearchQueries).isNotEmpty();
assertThat(groundingMetadata).containsKey("groundingChunks");
List<Map<String, Object>> groundingChunks =
(List<Map<String, Object>>) groundingMetadata.get("groundingChunks");
assertThat(groundingChunks).isNotEmpty();
groundingChunks.forEach(chunk -> {
assertThat(chunk).containsKey("web");
Map<String, Object> web = (Map<String, Object>) chunk.get("web");
assertThat(web).containsKeys("uri", "title");
});
assertThat(groundingMetadata).containsKey("searchEntryPoint");
Map<String, Object> searchEntryPoint = (Map<String, Object>) groundingMetadata.get("searchEntryPoint");
assertThat(searchEntryPoint).containsKey("renderedContent");
assertThat((String) searchEntryPoint.get("renderedContent")).isNotBlank();
}
@Test
void should_generate_high_res_image() throws IOException {
// given
var subject = GoogleGenAiImageModel.builder()
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.modelName(MODEL_NAME)
.aspectRatio("16:9")
.imageSize("4K")
.build();
// when
var response = subject.generate("""
Da Vinci style anatomical sketch of a dissected Monarch butterfly.
Detailed drawings of the head, wings, and legs on textured parchment with notes in English.
""");
// then
var image = response.content();
assertThat(image).isNotNull();
assertThat(image.base64Data()).isNotBlank();
assertThat(image.mimeType()).startsWith("image/");
saveImage(image, "should_generate_high_res_image");
}
private static void saveImage(Image image, String fileName) throws IOException {
String extension = getExtension(image.mimeType());
Path filePath = OUTPUT_DIR.resolve(fileName + "." + extension);
byte[] imageBytes = Base64.getDecoder().decode(image.base64Data());
Files.write(filePath, imageBytes);
System.out.println("Saved image to: " + filePath.toAbsolutePath());
}
private static String getExtension(String mimeType) {
if (mimeType == null) {
return "png";
}
return switch (mimeType) {
case "image/jpeg" -> "jpg";
case "image/gif" -> "gif";
case "image/webp" -> "webp";
default -> "png";
};
}
@AfterEach
void afterEach() throws InterruptedException {
String ciDelaySeconds = System.getenv("CI_DELAY_SECONDS_GOOGLE_GENAI");
if (ciDelaySeconds != null) {
Thread.sleep(Integer.parseInt(ciDelaySeconds) * 1000L);
}
}
}
@@ -0,0 +1,27 @@
package dev.langchain4j.model.google.genai;
import static org.assertj.core.api.Assertions.assertThat;
import dev.langchain4j.model.catalog.ModelDescription;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@EnabledIfEnvironmentVariable(named = "GOOGLE_AI_GEMINI_API_KEY", matches = ".+")
class GoogleGenAiModelCatalogIT {
private static final String GOOGLE_AI_GEMINI_API_KEY = System.getenv("GOOGLE_AI_GEMINI_API_KEY");
@Test
void test_list_models() {
GoogleGenAiModelCatalog catalog = GoogleGenAiModelCatalog.builder()
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.build();
List<ModelDescription> models = catalog.listModels();
assertThat(models).isNotEmpty();
boolean foundGeminiFlash = models.stream().anyMatch(m -> m.name().contains("gemini-2.5-flash"));
assertThat(foundGeminiFlash).isTrue();
}
}
@@ -0,0 +1,18 @@
package dev.langchain4j.model.google.genai;
import static org.assertj.core.api.Assertions.assertThat;
import dev.langchain4j.model.ModelProvider;
import org.junit.jupiter.api.Test;
class GoogleGenAiModelCatalogTest {
@Test
void should_return_correct_provider() {
GoogleGenAiModelCatalog catalog =
GoogleGenAiModelCatalog.builder().apiKey("test-key").build();
assertThat(catalog).isNotNull();
assertThat(catalog.provider()).isEqualTo(ModelProvider.GOOGLE_GENAI);
}
}
@@ -2,6 +2,7 @@ package dev.langchain4j.model.google.genai;
import static java.util.Collections.singletonList;
import com.google.genai.errors.ClientException;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.chat.common.AbstractStreamingChatModelListenerIT;
import dev.langchain4j.model.chat.listener.ChatModelListener;
@@ -38,7 +39,7 @@ class GoogleGenAiStreamingChatModelListenerIT extends AbstractStreamingChatModel
@Override
protected Class<? extends Exception> expectedExceptionClass() {
return com.google.genai.errors.ClientException.class;
return ClientException.class;
}
@Override
@@ -1,17 +1,37 @@
package dev.langchain4j.model.google.genai;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import com.google.genai.Client;
import com.google.genai.Models;
import com.google.genai.ResponseStream;
import com.google.genai.types.Candidate;
import com.google.genai.types.Content;
import com.google.genai.types.FunctionCall;
import com.google.genai.types.GenerateContentResponse;
import com.google.genai.types.Part;
import com.google.genai.types.SafetySetting;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.ModelProvider;
import dev.langchain4j.model.chat.Capability;
import dev.langchain4j.model.chat.listener.ChatModelListener;
import dev.langchain4j.model.chat.request.ResponseFormat;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.model.output.FinishReason;
import java.lang.reflect.Field;
import java.time.Duration;
import java.util.Base64;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import org.junit.jupiter.api.Test;
class GoogleGenAiStreamingChatModelTest {
@@ -20,7 +40,7 @@ class GoogleGenAiStreamingChatModelTest {
void should_build_with_api_key_and_model_name() {
GoogleGenAiStreamingChatModel model = GoogleGenAiStreamingChatModel.builder()
.apiKey("test-key")
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.build();
assertThat(model).isNotNull();
@@ -33,7 +53,7 @@ class GoogleGenAiStreamingChatModelTest {
GoogleGenAiStreamingChatModel model = GoogleGenAiStreamingChatModel.builder()
.client(client)
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.build();
assertThat(model).isNotNull();
@@ -43,10 +63,12 @@ class GoogleGenAiStreamingChatModelTest {
void should_set_default_request_parameters() {
GoogleGenAiStreamingChatModel model = GoogleGenAiStreamingChatModel.builder()
.apiKey("test-key")
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.temperature(0.7)
.topP(0.9)
.topK(40)
.frequencyPenalty(0.5)
.presencePenalty(0.3)
.maxOutputTokens(512)
.stopSequences(List.of("END"))
.build();
@@ -54,6 +76,8 @@ class GoogleGenAiStreamingChatModelTest {
assertThat(model.defaultRequestParameters().temperature()).isEqualTo(0.7);
assertThat(model.defaultRequestParameters().topP()).isEqualTo(0.9);
assertThat(model.defaultRequestParameters().topK()).isEqualTo(40);
assertThat(model.defaultRequestParameters().frequencyPenalty()).isEqualTo(0.5);
assertThat(model.defaultRequestParameters().presencePenalty()).isEqualTo(0.3);
assertThat(model.defaultRequestParameters().maxOutputTokens()).isEqualTo(512);
assertThat(model.defaultRequestParameters().stopSequences()).containsExactly("END");
}
@@ -62,7 +86,7 @@ class GoogleGenAiStreamingChatModelTest {
void should_return_empty_listeners_by_default() {
GoogleGenAiStreamingChatModel model = GoogleGenAiStreamingChatModel.builder()
.apiKey("test-key")
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.build();
assertThat(model.listeners()).isEmpty();
@@ -74,7 +98,7 @@ class GoogleGenAiStreamingChatModelTest {
GoogleGenAiStreamingChatModel model = GoogleGenAiStreamingChatModel.builder()
.apiKey("test-key")
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.listeners(List.of(listener))
.build();
@@ -85,7 +109,7 @@ class GoogleGenAiStreamingChatModelTest {
void should_always_advertise_json_schema_capability() {
GoogleGenAiStreamingChatModel model = GoogleGenAiStreamingChatModel.builder()
.apiKey("test-key")
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.build();
assertThat(model.supportedCapabilities()).containsExactly(Capability.RESPONSE_FORMAT_JSON_SCHEMA);
@@ -97,10 +121,12 @@ class GoogleGenAiStreamingChatModelTest {
GoogleGenAiStreamingChatModel model = GoogleGenAiStreamingChatModel.builder()
.client(client)
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.temperature(0.5)
.topP(0.8)
.topK(30)
.frequencyPenalty(0.5)
.presencePenalty(0.3)
.maxOutputTokens(1024)
.thinkingBudget(500)
.seed(42)
@@ -123,7 +149,7 @@ class GoogleGenAiStreamingChatModelTest {
void should_build_with_null_optional_fields() {
GoogleGenAiStreamingChatModel model = GoogleGenAiStreamingChatModel.builder()
.apiKey("test-key")
.modelName("gemini-2.0-flash")
.modelName("gemini-3.1-flash-lite")
.listeners(null)
.safetySettings(null)
.executor(null)
@@ -142,6 +168,8 @@ class GoogleGenAiStreamingChatModelTest {
assertThat(builder.temperature(0.5)).isSameAs(builder);
assertThat(builder.topP(0.8)).isSameAs(builder);
assertThat(builder.topK(40)).isSameAs(builder);
assertThat(builder.frequencyPenalty(0.5)).isSameAs(builder);
assertThat(builder.presencePenalty(0.3)).isSameAs(builder);
assertThat(builder.maxOutputTokens(100)).isSameAs(builder);
assertThat(builder.thinkingBudget(500)).isSameAs(builder);
assertThat(builder.seed(42)).isSameAs(builder);
@@ -165,4 +193,142 @@ class GoogleGenAiStreamingChatModelTest {
assertThat(builder.projectId("project")).isSameAs(builder);
assertThat(builder.location("us-central1")).isSameAs(builder);
}
@Test
void should_accumulate_attributes_from_streaming_chunks() throws Exception {
Client client = mock(Client.class);
Models models = mock(Models.class);
Field modelsField = Client.class.getDeclaredField("models");
modelsField.setAccessible(true);
modelsField.set(client, models);
@SuppressWarnings("unchecked")
ResponseStream<GenerateContentResponse> stream = mock(ResponseStream.class);
when(models.generateContentStream(any(String.class), any(List.class), any()))
.thenReturn(stream);
// Create a mock chunk with a function call and thought signature
Map<String, Object> args = new HashMap<>();
args.put("location", "Paris");
FunctionCall functionCall = FunctionCall.builder()
.name("get_weather")
.id("call_123")
.args(args)
.build();
Part part = Part.builder()
.functionCall(functionCall)
.thoughtSignature("signature-data".getBytes())
.build();
Content content = Content.builder().role("model").parts(List.of(part)).build();
Candidate candidate = Candidate.builder().content(content).build();
GenerateContentResponse chunk =
GenerateContentResponse.builder().candidates(List.of(candidate)).build();
when(stream.iterator()).thenReturn(List.of(chunk).iterator());
GoogleGenAiStreamingChatModel model = GoogleGenAiStreamingChatModel.builder()
.client(client)
.modelName("gemini-3.5-flash")
.build();
CompletableFuture<ChatResponse> future = new CompletableFuture<>();
model.chat(List.of(UserMessage.from("What's the weather in Paris?")), new StreamingChatResponseHandler() {
@Override
public void onPartialResponse(String partialResponse) {}
@Override
public void onCompleteResponse(ChatResponse completeResponse) {
future.complete(completeResponse);
}
@Override
public void onError(Throwable error) {
future.completeExceptionally(error);
}
});
ChatResponse response = future.get(5, TimeUnit.SECONDS);
AiMessage aiMessage = response.aiMessage();
assertThat(aiMessage.hasToolExecutionRequests()).isTrue();
assertThat(aiMessage.attributes()).isNotEmpty();
String encodedSig = Base64.getEncoder().encodeToString("signature-data".getBytes());
assertThat(aiMessage.attribute("thought_signature_call_123", String.class))
.isEqualTo(encodedSig);
}
@Test
void should_not_overwrite_truncation_finish_reason_with_stop() throws Exception {
Client client = mock(Client.class);
Models models = mock(Models.class);
Field modelsField = Client.class.getDeclaredField("models");
modelsField.setAccessible(true);
modelsField.set(client, models);
@SuppressWarnings("unchecked")
ResponseStream<GenerateContentResponse> stream = mock(ResponseStream.class);
when(models.generateContentStream(any(String.class), any(List.class), any()))
.thenReturn(stream);
// Chunk 1: FinishReason.MAX_TOKENS -> maps to LENGTH
Candidate candidate1 = Candidate.builder()
.content(Content.builder()
.role("model")
.parts(List.of(Part.builder().text("First part").build()))
.build())
.finishReason(
new com.google.genai.types.FinishReason(com.google.genai.types.FinishReason.Known.MAX_TOKENS))
.build();
GenerateContentResponse chunk1 = GenerateContentResponse.builder()
.candidates(List.of(candidate1))
.build();
// Chunk 2: FinishReason.STOP -> maps to STOP (trailing chunk)
Candidate candidate2 = Candidate.builder()
.content(Content.builder()
.role("model")
.parts(List.of(Part.builder().text("").build()))
.build())
.finishReason(new com.google.genai.types.FinishReason(com.google.genai.types.FinishReason.Known.STOP))
.build();
GenerateContentResponse chunk2 = GenerateContentResponse.builder()
.candidates(List.of(candidate2))
.build();
when(stream.iterator()).thenReturn(List.of(chunk1, chunk2).iterator());
GoogleGenAiStreamingChatModel model = GoogleGenAiStreamingChatModel.builder()
.client(client)
.modelName("gemini-3.5-flash")
.build();
CompletableFuture<ChatResponse> future = new CompletableFuture<>();
model.chat(List.of(UserMessage.from("Hello")), new StreamingChatResponseHandler() {
@Override
public void onPartialResponse(String partialResponse) {}
@Override
public void onCompleteResponse(ChatResponse completeResponse) {
future.complete(completeResponse);
}
@Override
public void onError(Throwable error) {
future.completeExceptionally(error);
}
});
ChatResponse response = future.get(5, TimeUnit.SECONDS);
assertThat(response.metadata().finishReason()).isEqualTo(FinishReason.LENGTH);
}
}
@@ -0,0 +1,176 @@
package dev.langchain4j.model.google.genai;
import static org.assertj.core.api.Assertions.assertThat;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.TokenCountEstimator;
import dev.langchain4j.model.chat.request.ResponseFormat;
import dev.langchain4j.model.chat.request.json.JsonObjectSchema;
import dev.langchain4j.service.AiServices;
import java.util.Arrays;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@EnabledIfEnvironmentVariable(named = "GOOGLE_AI_GEMINI_API_KEY", matches = ".+")
class GoogleGenAiTokenCountEstimatorIT {
private static final String GOOGLE_AI_GEMINI_API_KEY = System.getenv("GOOGLE_AI_GEMINI_API_KEY");
@Test
void should_estimate_token_count_for_text() {
// given
TokenCountEstimator tokenCountEstimator = GoogleGenAiTokenCountEstimator.builder()
.logRequests(true)
.logResponses(true)
.modelName("gemini-2.5-flash-lite")
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.build();
// when
int count = tokenCountEstimator.estimateTokenCountInText("Hello world!");
// then
assertThat(count).isEqualTo(4);
}
@Test
void should_estimate_token_count_for_a_message() {
// given
TokenCountEstimator tokenCountEstimator = GoogleGenAiTokenCountEstimator.builder()
.logRequests(true)
.logResponses(true)
.modelName("gemini-2.5-flash-lite")
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.build();
// when
int count = tokenCountEstimator.estimateTokenCountInMessage(UserMessage.from("Hello World!"));
// then
assertThat(count).isEqualTo(4);
}
@Test
void should_estimate_token_count_for_list_of_messages() {
// given
TokenCountEstimator tokenCountEstimator = GoogleGenAiTokenCountEstimator.builder()
.logRequests(true)
.logResponses(true)
.modelName("gemini-2.5-flash-lite")
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.build();
// when
int count = tokenCountEstimator.estimateTokenCountInMessages(
Arrays.asList(UserMessage.from("Hello World!"), AiMessage.from("Hi! How can I help you today?")));
// then
assertThat(count).isEqualTo(14);
}
@Test
void should_estimate_token_count_for_tool_exec_reqs() {
// given
GoogleGenAiTokenCountEstimator tokenCountEstimator = GoogleGenAiTokenCountEstimator.builder()
.logRequests(true)
.logResponses(true)
.modelName("gemini-2.5-flash-lite")
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.build();
// when
int count = tokenCountEstimator.estimateTokenCountInToolExecutionRequests(Arrays.asList(
ToolExecutionRequest.builder()
.name("weatherForecast")
.arguments("{ \"location\": \"Paris\" }")
.build(),
ToolExecutionRequest.builder()
.name("weatherForecast")
.arguments("{ \"location\": \"London\" }")
.build()));
// then
assertThat(count).isEqualTo(29);
}
@Test
void should_estimate_token_count_for_tool_specs() {
// given
GoogleGenAiTokenCountEstimator tokenCountEstimator = GoogleGenAiTokenCountEstimator.builder()
.logRequests(true)
.logResponses(true)
.modelName("gemini-2.5-flash-lite")
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.build();
// when
int count = tokenCountEstimator.estimateTokenCountInToolSpecifications(Arrays.asList(
ToolSpecification.builder()
.name("weatherForecast")
.description("Get the weather forecast for a given location on a give date")
.parameters(JsonObjectSchema.builder()
.addStringProperty("location", "the location")
.addStringProperty("date", "the date")
.required("location", "date")
.build())
.build(),
ToolSpecification.builder()
.name("convertFahrenheitToCelsius")
.description("Convert a temperature in Fahrenheit to Celsius")
.parameters(JsonObjectSchema.builder()
.addNumberProperty("fahrenheit", "the temperature in Fahrenheit")
.required("fahrenheit")
.build())
.build()));
// then
// Since we are using the new SDK, let's just make sure it returns a positive number
// Token counts may differ slightly between SDKs or implementations
assertThat(count).isGreaterThan(0);
// assertThat(count).isEqualTo(102);
}
@Test
void shouldReturnResponseWithSystemMessageRequest() {
// given
GoogleGenAiChatModel model = GoogleGenAiChatModel.builder()
.responseFormat(ResponseFormat.TEXT)
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.modelName("gemini-2.5-flash-lite")
.build();
TokenCountEstimator estimator = GoogleGenAiTokenCountEstimator.builder()
.modelName("gemini-2.5-flash-lite")
.apiKey(GOOGLE_AI_GEMINI_API_KEY)
.build();
interface Assistant {
String chat(String userMessage);
}
Assistant assistant = AiServices.builder(Assistant.class)
.chatModel(model)
.systemMessageProvider(o -> "You are a useful assistant")
.chatMemoryProvider(memoryId -> TokenWindowChatMemory.withMaxTokens(3000, estimator))
.build();
// when
String response = assistant.chat("Hello!");
// then
assertThat(response).isNotEmpty();
}
@AfterEach
void afterEach() throws InterruptedException {
String ciDelaySeconds = System.getenv("CI_DELAY_SECONDS_GOOGLE_GENAI");
if (ciDelaySeconds != null) {
Thread.sleep(Integer.parseInt(ciDelaySeconds) * 1000L);
}
}
}