Skip to content

Commit

Permalink
added modification to text method and test for it
Browse files Browse the repository at this point in the history
  • Loading branch information
David Motsonashvili committed Jun 27, 2024
1 parent d407d68 commit 7c4c7a7
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -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 */
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 =
Expand Down

0 comments on commit 7c4c7a7

Please sign in to comment.