Skip to content

Commit 4c29e0b

Browse files
authored
Add headers to be able to track bidi usage (#6939)
Originally we were not tracking bidi usage in android. This PR enables it by adding certain headers when the initial websocket connection is made. There was some refractoring done wrt to extract out the part which just adds the headers from the applyCommonConfiguration function
1 parent 525f293 commit 4c29e0b

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)