Skip to content

Commit 23dd81b

Browse files
authored
Merge a9c16ec into 525f293
2 parents 525f293 + a9c16ec commit 23dd81b

File tree

1 file changed

+10
-7
lines changed

1 file changed

+10
-7
lines changed

firebase-ai/src/main/kotlin/com/google/firebase/ai/common/APIController.kt

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,7 @@ internal constructor(
164164
"wss://firebasevertexai.googleapis.com/ws/google.firebase.vertexai.v1beta.LlmBidiService/BidiGenerateContent/locations/$location?key=$key"
165165

166166
suspend fun getWebSocketSession(location: String): ClientWebSocketSession =
167-
client.webSocketSession(getBidiEndpoint(location))
167+
client.webSocketSession(getBidiEndpoint(location)) { applyCommonHeaders() }
168168

169169
fun generateContentStream(
170170
request: GenerateContentRequest
@@ -191,12 +191,7 @@ internal constructor(
191191
throw FirebaseCommonAIException.from(e)
192192
}
193193

194-
private fun HttpRequestBuilder.applyCommonConfiguration(request: Request) {
195-
when (request) {
196-
is GenerateContentRequest -> setBody<GenerateContentRequest>(request)
197-
is CountTokensRequest -> setBody<CountTokensRequest>(request)
198-
is GenerateImageRequest -> setBody<GenerateImageRequest>(request)
199-
}
194+
private fun HttpRequestBuilder.applyCommonHeaders() {
200195
contentType(ContentType.Application.Json)
201196
header("x-goog-api-key", key)
202197
header("x-goog-api-client", apiClient)
@@ -205,6 +200,14 @@ internal constructor(
205200
header("X-Firebase-AppVersion", appVersion)
206201
}
207202
}
203+
private fun HttpRequestBuilder.applyCommonConfiguration(request: Request) {
204+
when (request) {
205+
is GenerateContentRequest -> setBody<GenerateContentRequest>(request)
206+
is CountTokensRequest -> setBody<CountTokensRequest>(request)
207+
is GenerateImageRequest -> setBody<GenerateImageRequest>(request)
208+
}
209+
applyCommonHeaders()
210+
}
208211

209212
private suspend fun HttpRequestBuilder.applyHeaderProvider() {
210213
if (headerProvider != null) {

0 commit comments

Comments
 (0)