Skip to content

Commit fedf744

Browse files
authored
Merge branch 'main' into rl.vertexai.changelog
2 parents 0a7775e + b1da895 commit fedf744

File tree

3 files changed

+97
-1
lines changed

3 files changed

+97
-1
lines changed

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/GenerativeModel.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -231,7 +231,7 @@ internal constructor(
231231
generationConfig?.toInternal(),
232232
tools?.map { it.toInternal() },
233233
toolConfig?.toInternal(),
234-
systemInstruction?.toInternal()
234+
systemInstruction?.copy(role = "system")?.toInternal()
235235
)
236236

237237
private fun constructCountTokensRequest(vararg prompt: Content) =

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/Content.kt

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,10 @@ public class Content
3030
@JvmOverloads
3131
constructor(public val role: String? = "user", public val parts: List<Part>) {
3232

33+
public fun copy(role: String? = this.role, parts: List<Part> = this.parts): Content {
34+
return Content(role, parts)
35+
}
36+
3337
/** Builder class to facilitate constructing complex [Content] objects. */
3438
public class Builder {
3539

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.firebase.vertexai
18+
19+
import com.google.firebase.vertexai.common.APIController
20+
import com.google.firebase.vertexai.common.GenerateContentResponse
21+
import com.google.firebase.vertexai.common.JSON
22+
import com.google.firebase.vertexai.common.server.Candidate
23+
import com.google.firebase.vertexai.common.shared.Content
24+
import com.google.firebase.vertexai.common.shared.TextPart
25+
import com.google.firebase.vertexai.common.util.doBlocking
26+
import com.google.firebase.vertexai.type.RequestOptions
27+
import com.google.firebase.vertexai.type.content
28+
import io.kotest.assertions.json.shouldContainJsonKey
29+
import io.kotest.assertions.json.shouldContainJsonKeyValue
30+
import io.kotest.matchers.collections.shouldNotBeEmpty
31+
import io.kotest.matchers.types.shouldBeInstanceOf
32+
import io.ktor.client.engine.mock.MockEngine
33+
import io.ktor.client.engine.mock.respond
34+
import io.ktor.http.HttpHeaders
35+
import io.ktor.http.HttpStatusCode
36+
import io.ktor.http.content.TextContent
37+
import io.ktor.http.headersOf
38+
import kotlin.time.Duration.Companion.seconds
39+
import kotlinx.coroutines.withTimeout
40+
import kotlinx.serialization.encodeToString
41+
import org.junit.Test
42+
43+
internal class GenerativeModelTesting {
44+
private val TEST_CLIENT_ID = "test"
45+
46+
@Test
47+
fun addition() = doBlocking {
48+
val mockEngine = MockEngine {
49+
respond(
50+
generateContentResponseAsJsonString("text response"),
51+
HttpStatusCode.OK,
52+
headersOf(HttpHeaders.ContentType, "application/json")
53+
)
54+
}
55+
56+
val apiController =
57+
APIController(
58+
"super_cool_test_key",
59+
"gemini-1.5-flash",
60+
RequestOptions(timeout = 5.seconds, endpoint = "https://my.custom.endpoint"),
61+
mockEngine,
62+
TEST_CLIENT_ID,
63+
null,
64+
)
65+
66+
val generativeModel =
67+
GenerativeModel(
68+
"gemini-1.5-flash",
69+
systemInstruction = content { text("system instruction") },
70+
controller = apiController
71+
)
72+
73+
withTimeout(5.seconds) { generativeModel.generateContent("my test prompt") }
74+
75+
mockEngine.requestHistory.shouldNotBeEmpty()
76+
77+
val request = mockEngine.requestHistory.first().body
78+
request.shouldBeInstanceOf<TextContent>()
79+
80+
request.text.let {
81+
it shouldContainJsonKey "system_instruction"
82+
it.shouldContainJsonKeyValue("$.system_instruction.role", "system")
83+
it.shouldContainJsonKeyValue("$.system_instruction.parts[0].text", "system instruction")
84+
}
85+
}
86+
87+
private fun generateContentResponseAsJsonString(text: String): String {
88+
return JSON.encodeToString(
89+
GenerateContentResponse(listOf(Candidate(Content(parts = listOf(TextPart(text))))))
90+
)
91+
}
92+
}

0 commit comments

Comments
 (0)