From 7c4c7a77718e3cf85dfc78133c3120b1c44948b4 Mon Sep 17 00:00:00 2001
From: David Motsonashvili <davidmotson@google.com>
Date: Thu, 27 Jun 2024 09:44:00 -0700
Subject: [PATCH] added modification to text method and test for it
---
.../type/GenerateContentResponse.kt | 14 +++++++-
.../GenerateContentResponseTest.kt | 35 +++++++++++++++++++
2 files changed, 48 insertions(+), 1 deletion(-)
diff --git a/generativeai/src/main/java/com/google/ai/client/generativeai/type/GenerateContentResponse.kt b/generativeai/src/main/java/com/google/ai/client/generativeai/type/GenerateContentResponse.kt
index fad24fd6..bbaa3f23 100644
--- a/generativeai/src/main/java/com/google/ai/client/generativeai/type/GenerateContentResponse.kt
+++ b/generativeai/src/main/java/com/google/ai/client/generativeai/type/GenerateContentResponse.kt
@@ -32,7 +32,19 @@ class GenerateContentResponse(
) {
/** Convenience field representing all the text parts in the response, if they exists. */
val text: String? by lazy {
- candidates.first().content.parts.filterIsInstance<TextPart>().joinToString(" ") { it.text }
+ candidates
+ .first()
+ .content
+ .parts
+ .filter { it is TextPart || it is ExecutableCodePart || it is CodeExecutionResultPart }
+ .joinToString(" ") {
+ when (it) {
+ is TextPart -> it.text
+ is ExecutableCodePart -> "\n```${it.language.lowercase()}\n${it.code}\n```"
+ is CodeExecutionResultPart -> "\n```\n${it.output}\n```"
+ else -> throw RuntimeException("unreachable")
+ }
+ }
}
/** Convenience field representing the first function call part in the request, if it exists */
diff --git a/generativeai/src/test/java/com/google/ai/client/generativeai/GenerateContentResponseTest.kt b/generativeai/src/test/java/com/google/ai/client/generativeai/GenerateContentResponseTest.kt
index 335e81d5..fa9f6d58 100644
--- a/generativeai/src/test/java/com/google/ai/client/generativeai/GenerateContentResponseTest.kt
+++ b/generativeai/src/test/java/com/google/ai/client/generativeai/GenerateContentResponseTest.kt
@@ -17,6 +17,9 @@
package com.google.ai.client.generativeai
import com.google.ai.client.generativeai.type.Candidate
+import com.google.ai.client.generativeai.type.CodeExecutionResultPart
+import com.google.ai.client.generativeai.type.ExecutableCodePart
+import com.google.ai.client.generativeai.type.ExecutionOutcome
import com.google.ai.client.generativeai.type.FunctionCallPart
import com.google.ai.client.generativeai.type.GenerateContentResponse
import com.google.ai.client.generativeai.type.content
@@ -74,6 +77,38 @@ internal class GenerateContentResponseTest {
response.text shouldBe "This is a textPart"
}
+ @Test
+ fun `generate response should add generated code to the response`() {
+ val response =
+ GenerateContentResponse(
+ candidates =
+ listOf(
+ Candidate(
+ content {
+ text("I can calculate that for you!")
+ part(ExecutableCodePart("python", "print(\"hello world\")"))
+ part(CodeExecutionResultPart(ExecutionOutcome.OK, "hello world"))
+ },
+ listOf(),
+ listOf(),
+ null,
+ )
+ ),
+ null,
+ null,
+ )
+
+ response.text shouldBe """
+ I can calculate that for you!
+ ```python
+ print("hello world")
+ ```
+ ```
+ hello world
+ ```
+ """.trimIndent()
+ }
+
@Test
fun `generate response should get strings and concatenate them together`() {
val response =