- The Batch API for Gemini lacks comprehensive documentation, making implementation challenging.
- Creating examples for the Batch API is time-consuming due to insufficient explanations.
- Advanced features of the API lack examples, hindering user understanding.
- The post includes code for using the Batch API with schemas, aiding user onboarding.
- A model class is defined to structure output from the LLM, ensuring organized data retrieval.
There isn’t any documentation available about the Batch API for Gemini in Java. There are a few code samples, but not much explanation. Additionally, some of the API’s advanced features lack examples as well. Due to the lack of documentation, it took significantly longer to create this example. In this post, I’ll share the code for using the Batch API with schemas.
Model Code
To start, here is some simple model code. I’ve created an annotation that’s used when the model is sent to the LLM. It describes what sort of thing we want to be returned with the schema.
This Joke model class will define how we receive data from the LLM with structured output.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | @Retention(RetentionPolicy.RUNTIME) @Target({ ElementType.TYPE, ElementType.FIELD }) @interface SchemaDescription { String value(); } @SchemaDescription("A joke with setup and punchline") private static class Joke { @SchemaDescription("The setup or beginning of the joke") private String setup; @SchemaDescription("The punchline or funny ending of the joke") private String punchline; public String getSetup() { return setup; } public String getPunchline() { return punchline; } } |
We also need a method that will convert our model class to the Schema.Builder that Gemini is expecting. This allows the LLM to know precisely what JSON it is expected to return.
The generateSchema uses Java’s reflection to for members (fields) that have the SchemaDescription annotation so we pull the description and member type.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 | private Schema.Builder generateSchema(Class<?> clazz) { SchemaDescription classDesc = clazz.getAnnotation(SchemaDescription.class); String description = classDesc != null ? classDesc.value() : ""; Map<String, Schema> properties = new HashMap<>(); List<String> required = new ArrayList<>(); for (Field field : clazz.getDeclaredFields()) { SchemaDescription fieldDesc = field.getAnnotation(SchemaDescription.class); if (fieldDesc != null) { String type = getType(field.getType()); Schema schema = Schema.builder() .type(type) .description(fieldDesc.value()) .build(); properties.put(field.getName(), schema); required.add(field.getName()); } } return Schema.builder() .type("object") .description(description) .properties(properties) .required(required.toArray(new String[0])); } private String getType(Class<?> type) { if (type == String.class) return "string"; // add other types if needed return "string"; // default } |
Batch Job Create
We need to create and submit the batch job. This will require filling out several objects that the API expects. In this example, we send two separate requests. Each one can have its own prompt and system prompt.
To use structured output, we request JSON as the output and pass in the Schema.Builder object we created for the Joke object.
Once all of the objects are ready, we can call the batches.create to submit everything together as a single batch.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | private BatchJob createBatchJob(Client client) { // Create the first prompt with pirate style String user1 = "Tell me a joke"; String system1 = "Always talk like a pirate."; InlinedRequest request = InlinedRequest.builder() .contents(Content.builder().parts(Part.fromText(user1))) .config( GenerateContentConfig.builder() .systemInstruction( Content.builder() .parts( Part.fromText(system1))) .temperature(0.0f) .responseMimeType("application/json") .responseSchema(generateJokeSchema())) .build(); // Create the second prompt with vampire style String user2 = "Tell me a joke"; String system2 = "Talk like a vampire."; InlinedRequest request2 = InlinedRequest.builder() .contents(Content.builder().parts(Part.fromText(user2))) .config( GenerateContentConfig.builder() .systemInstruction( Content.builder() .parts( Part.fromText(system2))) .temperature(0.0f) .responseMimeType("application/json") .responseSchema(generateJokeSchema())) .build(); // Create the batch job with the two requests BatchJobSource batchJobSource = BatchJobSource.builder() .inlinedRequests(request, request2) .build(); CreateBatchJobConfig config = CreateBatchJobConfig.builder().displayName("joke-batch-job-java").build(); // Create and submit the batch job to Gemini BatchJob batchJob1 = client.batches.create(model, batchJobSource, config); System.out.println("Created batch job: " + batchJob1); return batchJob1; } |
Listing Batch Jobs
Immediately after submitting the job, we can list it. We’ll only be able to retrieve the contents of the LLM output once the job has succeeded.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | private void listAndGetSuceeded(Client client) { ListBatchJobsConfig listConfig = ListBatchJobsConfig.builder() .pageSize(10) .build(); client.batches.list(listConfig).forEach(batchJob -> { System.out.println("Batch job: " + batchJob); if (batchJob.state().isPresent()) { // Actual enum comparison is buried in "Known" wrapper if (batchJob.state().get().knownEnum().equals(JobState.Known.JOB_STATE_SUCCEEDED) && batchJob.displayName().isPresent() && batchJob.displayName().get().equals("joke-batch-job-java")) { System.out.println("Batch job succeeded!"); createJavaObjectFromBatch(client, batchJob); } // Print the state of each batch job System.out.println("State: " + batchJob.state().get()); } }); } |
Converting LLM JSON to Java Model Object
When we find that a job has succeeded, we can output the model result and transform the JSON into a Java model object.
I couldn’t find a way to verify that a response has a schema associated with it. You’ll have to put things in a try/catch to check the LLM’s response.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 | private void createJavaObjectFromBatch(Client client, BatchJob batchJob) { // Listed jobs don't have the full response info, so we get it by name BatchJob retrievedBatchJob = client.batches.get(batchJob.name().get(), null); // Check if the destination has inlined responses if (retrievedBatchJob.dest().isPresent() && retrievedBatchJob.dest().get().inlinedResponses().isPresent()) { for (InlinedResponse response : retrievedBatchJob.dest().get().inlinedResponses().get()) { if (response.error().isPresent()) { // There was an error with this response System.out.println("Error: " + response.error().get().message()); continue; } else { // Successful response String text = response.response().get().text(); try { // There isn't a way to check if the response matches the schema, // so we just try to parse it into the Joke class ObjectMapper mapper = new ObjectMapper(); Joke joke = mapper.readValue(text, Joke.class); System.out.println("Parsed joke from Java object: setup='" + joke.getSetup() + "', punchline='" + joke.getPunchline() + "'"); } catch (Exception e) { System.out.println("Failed to parse JSON: " + e.getMessage()); System.out.println("Raw text: " + text); } } } } } |
Extras
For completeness, here is how to create the Client object for the API calls. I’m using the Builder to set the values explicitly. You can use environment variables too.
1 2 | Client client = Client.builder().apiKey(googleApiKey) .build(); |
And the import packages.
1 2 3 4 5 6 7 8 9 10 11 12 13 | import com.fasterxml.jackson.databind.ObjectMapper; import com.google.genai.Client; import com.google.genai.types.BatchJob; import com.google.genai.types.BatchJobSource; import com.google.genai.types.Content; import com.google.genai.types.CreateBatchJobConfig; import com.google.genai.types.GenerateContentConfig; import com.google.genai.types.InlinedRequest; import com.google.genai.types.InlinedResponse; import com.google.genai.types.JobState; import com.google.genai.types.ListBatchJobsConfig; import com.google.genai.types.Part; import com.google.genai.types.Schema; |
Frequently Asked Questions (AI FAQ by Summarizes)Why is the Batch API for Gemini challenging to implement?
The Batch API for Gemini lacks comprehensive documentation, making it challenging to implement.
What is a significant issue when creating examples for the Batch API?
Creating examples for the Batch API took significantly longer due to the absence of detailed explanations.
How does the post help users get started with the Batch API?
The post includes code for using the Batch API with schemas, which can help users get started.
What is the purpose of the model class defined in the post?
A model class is defined to structure the output received from the LLM, ensuring organized data retrieval.
What is crucial for proper data handling when using the Batch API?
Structured output is requested in JSON format, which is crucial for proper data handling.
What should be done if an error occurs in the API response?
If an error occurs in the response, it should be logged for debugging purposes.
What is the role of the ObjectMapper class in the Java implementation?
The ObjectMapper class is used for mapping JSON to Java objects, facilitating the conversion of LLM JSON responses into Java model objects.