GeneratedContent: add finishReason

add more safety to parsing response body

Author: michael-ameri

gemini-api/src/main/java/swiss/ameri/gemini/api/GenAi.java

@@ -210,18 +210,7 @@ public Stream<GeneratedContent> generateContentStream(GenerativeModel model) {
             );
             return response.body()
                     .filter(l -> l.length() > STREAM_LINE_PREFIX_LENGTH)
-                    .map(line -> {
-                        try {
-                            var gcr = jsonParser.fromJson(line.substring(STREAM_LINE_PREFIX_LENGTH), GenerateContentResponse.class);
-                            // each element can just replace the previous one
-                            this.responseById.put(uuid, gcr);
-                            // todo catch safety error
-                            //if ("SAFETY".equals(gcr.candidates().get(0).finishReason())) {
-                            return new GeneratedContent(uuid, gcr.candidates().get(0).content().parts().get(0).text());
-                        } catch (Exception e) {
-                            throw new RuntimeException("Unexpected line:\n" + line, e);
-                        }
-                    });
+                    .map(line -> parse(line.substring(STREAM_LINE_PREFIX_LENGTH), uuid));
         });
     }
 
@@ -249,15 +238,7 @@ public CompletableFuture<GeneratedContent> generateContent(GenerativeModel model
             );
             return response
                     .thenApply(HttpResponse::body)
-                    .thenApply(body -> {
-                        try {
-                            var gcr = jsonParser.fromJson(body, GenerateContentResponse.class);
-                            responseById.put(uuid, gcr);
-                            return new GeneratedContent(uuid, gcr.candidates().get(0).content().parts().get(0).text());
-                        } catch (Exception e) {
-                            throw new RuntimeException("Unexpected body:\n" + body, e);
-                        }
-                    });
+                    .thenApply(body -> parse(body, uuid));
         });
     }
 
@@ -418,12 +399,14 @@ public void close() {
     /**
      * Content generated by Gemini API.
      *
-     * @param id   the id of the request, for subsequent queries regarding metadata of the query
-     * @param text of the generated content
+     * @param id           the id of the request, for subsequent queries regarding metadata of the query
+     * @param text         of the generated content
+     * @param finishReason the reason generation was finished, according to <a href="https://ai.google.dev/api/generate-content#FinishReason">FinishReason</a>
      */
     public record GeneratedContent(
             UUID id,
-            String text
+            String text,
+            String finishReason
     ) {
     }
 
@@ -542,6 +525,22 @@ public record Model(
     ) {
     }
 
+    private GeneratedContent parse(String body, UUID uuid) {
+        try {
+            var gcr = jsonParser.fromJson(body, GenerateContentResponse.class);
+            // each element can just replace the previous one
+            this.responseById.put(uuid, gcr);
+            // we assume we always get a candidate. Otherwise, there is probably something wrong with the input
+            var candidate = gcr.candidates().get(0);
+            if (candidate.content() == null) {
+                return new GeneratedContent(uuid, "", candidate.finishReason());
+            }
+            return new GeneratedContent(uuid, candidate.content().parts().get(0).text(), candidate.finishReason());
+        } catch (Exception e) {
+            throw new RuntimeException("Unexpected body:\n" + body, e);
+        }
+    }
+
     private record BatchEmbedContentRequest(
             List<EmbedContentRequest> requests
     ) {