Support parsing the altname of citationSources (#74)

It can also be present in the response as `citations`
google-gemini · Feb 29, 2024 · 9bd6231 · 9bd6231
1 parent 3264d2b
commit 9bd6231
Show file tree

Hide file tree

Showing 5 changed files with 110 additions and 1 deletion.
diff --git a/generativeai/src/main/java/com/google/ai/client/generativeai/internal/api/server/Types.kt b/generativeai/src/main/java/com/google/ai/client/generativeai/internal/api/server/Types.kt
@@ -19,9 +19,11 @@ package com.google.ai.client.generativeai.internal.api.server
 import com.google.ai.client.generativeai.internal.api.shared.Content
 import com.google.ai.client.generativeai.internal.api.shared.HarmCategory
 import com.google.ai.client.generativeai.internal.util.FirstOrdinalSerializer
+import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.KSerializer
 import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
+import kotlinx.serialization.json.JsonNames
 
 internal object BlockReasonSerializer :
   KSerializer<BlockReason> by FirstOrdinalSerializer(BlockReason::class)
@@ -54,7 +56,10 @@ internal data class Candidate(
   val citationMetadata: CitationMetadata? = null
 )
 
-@Serializable internal data class CitationMetadata(val citationSources: List<CitationSources>)
+@Serializable
+internal data class CitationMetadata
+@OptIn(ExperimentalSerializationApi::class)
+constructor(@JsonNames("citations") val citationSources: List<CitationSources>)
 
 @Serializable
 internal data class CitationSources(

diff --git a/generativeai/src/test/java/com/google/ai/client/generativeai/StreamingSnapshotTests.kt b/generativeai/src/test/java/com/google/ai/client/generativeai/StreamingSnapshotTests.kt
@@ -142,6 +142,17 @@ internal class StreamingSnapshotTests {
       }
     }
 
+  @Test
+  fun `citation returns correctly when using alternative name`() =
+    goldenStreamingFile("success-citations-altname.txt") {
+      val responses = model.generateContentStream()
+
+      withTimeout(testTimeout) {
+        val responseList = responses.toList()
+        responseList.any { it.candidates.any { it.citationMetadata.isNotEmpty() } } shouldBe true
+      }
+    }
+
   @Test
   fun `stopped for recitation`() =
     goldenStreamingFile("failure-recitation-no-content.txt") {

diff --git a/generativeai/src/test/java/com/google/ai/client/generativeai/UnarySnapshotTests.kt b/generativeai/src/test/java/com/google/ai/client/generativeai/UnarySnapshotTests.kt
@@ -124,6 +124,17 @@ internal class UnarySnapshotTests {
       }
     }
 
+  @Test
+  fun `citation returns correctly when using alternative name`() =
+    goldenUnaryFile("success-citations-altname.json") {
+      withTimeout(testTimeout) {
+        val response = model.generateContent()
+
+        response.candidates.isEmpty() shouldBe false
+        response.candidates.first().citationMetadata.isNotEmpty() shouldBe true
+      }
+    }
+
   @Test
   fun `invalid response`() =
     goldenUnaryFile("failure-invalid-response.json") {

diff --git a/generativeai/src/test/resources/golden-files/streaming/success-citations-altname.txt b/generativeai/src/test/resources/golden-files/streaming/success-citations-altname.txt
@@ -0,0 +1,12 @@
+data: {"candidates": [{"content": {"parts": [{"text": "placeholder"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}],"promptFeedback": {"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}}
+
+data: {"candidates": [{"content": {"parts": [{"text": "placeholder"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]}
+
+data: {"candidates": [{"content": {"parts": [{"text": "placeholder"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]}
+
+data: {"candidates": [{"content": {"parts": [{"text": "placeholder"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}],"citationMetadata": {"citations": [{"startIndex": 574,"endIndex": 705,"uri": "https://example.com","license": ""},{"startIndex": 899,"endIndex": 1026,"uri": "https://example.com","license": ""}]}}]}
+
+data: {"candidates": [{"content": {"parts": [{"text": "placeholder"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}],"citationMetadata": {"citations": [{"startIndex": 574,"endIndex": 705,"uri": "https://example.com","license": ""},{"startIndex": 899,"endIndex": 1026,"uri": "https://example.com","license": ""}]}}]}
+
+data: {"candidates": [{"content": {"parts": [{"text": "placeholder"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}],"citationMetadata": {"citations": [{"startIndex": 574,"endIndex": 705,"uri": "https://example.com","license": ""},{"startIndex": 899,"endIndex": 1026,"uri": "https://example.com","license": ""}]}}]}
+
diff --git a/generativeai/src/test/resources/golden-files/unary/success-citations-altname.json b/generativeai/src/test/resources/golden-files/unary/success-citations-altname.json
@@ -0,0 +1,70 @@
+{
+  "candidates": [
+    {
+      "content": {
+        "parts": [
+          {
+            "text": "placeholder"
+          }
+        ],
+        "role": "model"
+      },
+      "finishReason": "STOP",
+      "index": 0,
+      "safetyRatings": [
+        {
+          "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+          "probability": "NEGLIGIBLE"
+        },
+        {
+          "category": "HARM_CATEGORY_HATE_SPEECH",
+          "probability": "NEGLIGIBLE"
+        },
+        {
+          "category": "HARM_CATEGORY_HARASSMENT",
+          "probability": "NEGLIGIBLE"
+        },
+        {
+          "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+          "probability": "NEGLIGIBLE"
+        }
+      ],
+      "citationMetadata": {
+        "citations": [
+          {
+            "startIndex": 574,
+            "endIndex": 705,
+            "uri": "https://example.com/",
+            "license": ""
+          },
+          {
+            "startIndex": 899,
+            "endIndex": 1026,
+            "uri": "https://example.com/",
+            "license": ""
+          }
+        ]
+      }
+    }
+  ],
+  "promptFeedback": {
+    "safetyRatings": [
+      {
+        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+        "probability": "NEGLIGIBLE"
+      },
+      {
+        "category": "HARM_CATEGORY_HATE_SPEECH",
+        "probability": "NEGLIGIBLE"
+      },
+      {
+        "category": "HARM_CATEGORY_HARASSMENT",
+        "probability": "NEGLIGIBLE"
+      },
+      {
+        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+        "probability": "NEGLIGIBLE"
+      }
+    ]
+  }
+}