From b67b7b375de3347e9833141238e5a5e5d3cc559f Mon Sep 17 00:00:00 2001
From: Nate Bosch <nbosch1@gmail.com>
Date: Thu, 23 May 2024 12:02:46 -0700
Subject: [PATCH] Add a text getter for Candidate (#170)

Closes #99

Call out that candidates may be ignored in the doc for
`GenerateContentResponse.text`.

Add a `text` getter on `Candidate`. This is the same logic that was in
the response text getter, applied to an individual candidate.

Refactor the response text implementation to use the candidate field.
---
 pkgs/google_generative_ai/CHANGELOG.md        |  2 +
 pkgs/google_generative_ai/lib/src/api.dart    | 85 +++++++++++--------
 .../test/response_parsing_test.dart           |  2 +
 3 files changed, 52 insertions(+), 37 deletions(-)

diff --git a/pkgs/google_generative_ai/CHANGELOG.md b/pkgs/google_generative_ai/CHANGELOG.md
index f5c3e61..521d4f7 100644
--- a/pkgs/google_generative_ai/CHANGELOG.md
+++ b/pkgs/google_generative_ai/CHANGELOG.md
@@ -9,6 +9,8 @@
   `systemInstruction`) in `countTokens` requests. This aligns the token count
   with the token count the backend will see in practice for a
   `generateContent` request.
+- Add a `text` getter on `Candidate` to make it easer to retrieve the text from
+  candidates other than the first in a response.
 
 ## 0.4.0
 
diff --git a/pkgs/google_generative_ai/lib/src/api.dart b/pkgs/google_generative_ai/lib/src/api.dart
index 4368baf..7b8fbf3 100644
--- a/pkgs/google_generative_ai/lib/src/api.dart
+++ b/pkgs/google_generative_ai/lib/src/api.dart
@@ -68,43 +68,24 @@ final class GenerateContentResponse {
   ///
   /// If there are no candidates, or if the first candidate does not contain any
   /// text parts, this value is `null`.
-  String? get text {
-    return switch (candidates) {
-      [] => switch (promptFeedback) {
-          PromptFeedback(
-            :final blockReason,
-            :final blockReasonMessage,
-          ) =>
-            // TODO: Add a specific subtype for this exception?
-            throw GenerativeAIException('Response was blocked'
-                '${blockReason != null ? ' due to $blockReason' : ''}'
-                '${blockReasonMessage != null ? ': $blockReasonMessage' : ''}'),
-          _ => null,
-        },
-      [
-        Candidate(
-          finishReason: (FinishReason.recitation || FinishReason.safety) &&
-              final finishReason,
-          :final finishMessage,
-        ),
-        ...
-      ] =>
-        throw GenerativeAIException(
-          // ignore: prefer_interpolation_to_compose_strings
-          'Candidate was blocked due to $finishReason' +
-              (finishMessage != null && finishMessage.isNotEmpty
-                  ? ': $finishMessage'
-                  : ''),
-        ),
-      // Special case for a single TextPart to avoid iterable chain.
-      [Candidate(content: Content(parts: [TextPart(:final text)])), ...] =>
-        text,
-      [Candidate(content: Content(:final parts)), ...]
-          when parts.any((p) => p is TextPart) =>
-        parts.whereType<TextPart>().map((p) => p.text).join(''),
-      [Candidate(), ...] => null,
-    };
-  }
+  ///
+  /// If there is more than one candidate, all but the first are ignored. See
+  /// [Candidate.text] to get the text content of candidates other than the
+  /// first.
+  String? get text => switch (candidates) {
+        [] => switch (promptFeedback) {
+            PromptFeedback(
+              :final blockReason,
+              :final blockReasonMessage,
+            ) =>
+              // TODO: Add a specific subtype for this exception?
+              throw GenerativeAIException('Response was blocked'
+                  '${blockReason != null ? ' due to $blockReason' : ''}'
+                  '${blockReasonMessage != null ? ': $blockReasonMessage' : ''}'),
+            _ => null,
+          },
+        [final candidate, ...] => candidate.text,
+      };
 
   /// The function call parts of the first candidate in [candidates], if any.
   ///
@@ -221,6 +202,36 @@ final class Candidate {
   // TODO: token count?
   Candidate(this.content, this.safetyRatings, this.citationMetadata,
       this.finishReason, this.finishMessage);
+
+  /// The concatenation of the text parts of [content], if any.
+  ///
+  /// If this candidate was finished for a reason of [FinishReason.recitation]
+  /// or [FinishReason.safety], accessing this text will throw a
+  /// [GenerativeAIException].
+  ///
+  /// If [content] contains any text parts, this value is the concatenation of
+  /// the text.
+  ///
+  /// If [content] does not contain any text parts, this value is `null`.
+  String? get text {
+    if (finishReason case FinishReason.recitation || FinishReason.safety) {
+      final String suffix;
+      if (finishMessage case final message? when message.isNotEmpty) {
+        suffix = ': $message';
+      } else {
+        suffix = '';
+      }
+      throw GenerativeAIException(
+          'Candidate was blocked due to $finishReason$suffix');
+    }
+    return switch (content.parts) {
+      // Special case for a single TextPart to avoid iterable chain.
+      [TextPart(:final text)] => text,
+      final parts when parts.any((p) => p is TextPart) =>
+        parts.whereType<TextPart>().map((p) => p.text).join(''),
+      _ => null,
+    };
+  }
 }
 
 /// Safety rating for a piece of content.
diff --git a/pkgs/google_generative_ai/test/response_parsing_test.dart b/pkgs/google_generative_ai/test/response_parsing_test.dart
index 4918b38..17edf54 100644
--- a/pkgs/google_generative_ai/test/response_parsing_test.dart
+++ b/pkgs/google_generative_ai/test/response_parsing_test.dart
@@ -457,6 +457,8 @@ void main() {
       final decoded = jsonDecode(response) as Object;
       final generateContentResponse = parseGenerateContentResponse(decoded);
       expect(generateContentResponse.text, 'Initial text And more text');
+      expect(generateContentResponse.candidates.single.text,
+          'Initial text And more text');
     });
   });