Skip to content

Commit 59d8956

Browse files
andrewheardgemini-code-assist[bot]
andrewheard
and
gemini-code-assist[bot]
authored
[AI] Add hybrid tag in x-goog-api-client header (#16118)
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
1 parent b1c9e27 commit 59d8956

4 files changed

Lines changed: 160 additions & 65 deletions

File tree

‎FirebaseAI/Sources/GenerativeAIService.swift‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -181,8 +181,12 @@ struct GenerativeAIService {
181181
if let bundleID = Bundle.main.bundleIdentifier {
182182
urlRequest.setValue(bundleID, forHTTPHeaderField: "x-ios-bundle-identifier")
183183
}
184+
var apiClientHeaders = [GenerativeAIService.languageTag, GenerativeAIService.firebaseVersionTag]
185+
if TaskLocals.isHybridRequest {
186+
apiClientHeaders.append("hybrid")
187+
}
184188
urlRequest.setValue(
185-
"\(GenerativeAIService.languageTag) \(GenerativeAIService.firebaseVersionTag)",
189+
apiClientHeaders.joined(separator: " "),
186190
forHTTPHeaderField: "x-goog-api-client"
187191
)
188192
urlRequest.setValue("application/json", forHTTPHeaderField: "Content-Type")

‎FirebaseAI/Sources/TaskLocals.swift‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
import Foundation
16+
17+
/// A container for task-local values used within the Firebase AI Logic SDK.
18+
///
19+
/// See https://developer.apple.com/documentation/swift/tasklocal for more details about
20+
/// `TaskLocal` values in Swift.
21+
enum TaskLocals {
22+
/// A task-local value indicating whether the current request is a hybrid request.
23+
///
24+
/// This is used to pass context down the call stack without modifying function signatures.
25+
@TaskLocal static var isHybridRequest = false
26+
}

‎FirebaseAI/Sources/Types/Internal/HybridModelSession.swift‎

Lines changed: 68 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -42,38 +42,40 @@
4242
func _respond(to prompt: [any Part], schema: FirebaseAI.GenerationSchema?,
4343
includeSchemaInPrompt: Bool, options: any GenerationOptionsRepresentable)
4444
async throws -> _ModelSessionResponse {
45-
// If the secondary session contains history then a previous fallback occurred.
46-
// Stick with the secondary session to maintain conversation consistency.
47-
if secondary._hasHistory {
48-
return try await secondary._respond(
49-
to: prompt,
50-
schema: schema,
51-
includeSchemaInPrompt: includeSchemaInPrompt,
52-
options: options
53-
)
54-
}
55-
56-
do {
57-
// First try the primary session.
58-
return try await primary._respond(
59-
to: prompt,
60-
schema: schema,
61-
includeSchemaInPrompt: includeSchemaInPrompt,
62-
options: options
63-
)
64-
} catch {
65-
// Do not fallback to second session if the primary session contains history.
66-
if primary._hasHistory {
67-
throw error
45+
return try await TaskLocals.$isHybridRequest.withValue(true) {
46+
// If the secondary session contains history then a previous fallback occurred.
47+
// Stick with the secondary session to maintain conversation consistency.
48+
if secondary._hasHistory {
49+
return try await secondary._respond(
50+
to: prompt,
51+
schema: schema,
52+
includeSchemaInPrompt: includeSchemaInPrompt,
53+
options: options
54+
)
6855
}
6956

70-
// Fallback to the second session if the first fails or is unavailable.
71-
return try await secondary._respond(
72-
to: prompt,
73-
schema: schema,
74-
includeSchemaInPrompt: includeSchemaInPrompt,
75-
options: options
76-
)
57+
do {
58+
// First try the primary session.
59+
return try await primary._respond(
60+
to: prompt,
61+
schema: schema,
62+
includeSchemaInPrompt: includeSchemaInPrompt,
63+
options: options
64+
)
65+
} catch {
66+
// Do not fallback to second session if the primary session contains history.
67+
if primary._hasHistory {
68+
throw error
69+
}
70+
71+
// Fallback to the second session if the first fails or is unavailable.
72+
return try await secondary._respond(
73+
to: prompt,
74+
schema: schema,
75+
includeSchemaInPrompt: includeSchemaInPrompt,
76+
options: options
77+
)
78+
}
7779
}
7880
}
7981

@@ -90,61 +92,63 @@
9092
includeSchemaInPrompt: Bool,
9193
options: any GenerationOptionsRepresentable)
9294
-> sending AsyncThrowingStream<_ModelSessionResponse, any Error> {
93-
// If the secondary session contains history then a previous fallback occurred.
94-
// Stick with the secondary session to maintain conversation consistency.
95-
if secondary._hasHistory {
96-
return secondary._streamResponse(
97-
to: prompt,
98-
schema: schema,
99-
includeSchemaInPrompt: includeSchemaInPrompt,
100-
options: options
101-
)
102-
}
103-
104-
return AsyncThrowingStream { continuation in
105-
let task = Task {
106-
// First try the primary session.
107-
let stream = primary._streamResponse(
95+
return TaskLocals.$isHybridRequest.withValue(true) {
96+
// If the secondary session contains history then a previous fallback occurred.
97+
// Stick with the secondary session to maintain conversation consistency.
98+
if secondary._hasHistory {
99+
return secondary._streamResponse(
108100
to: prompt,
109101
schema: schema,
110102
includeSchemaInPrompt: includeSchemaInPrompt,
111103
options: options
112104
)
105+
}
113106

114-
var didYield = false
115-
do {
116-
for try await snapshot in stream {
117-
didYield = true
118-
continuation.yield(snapshot)
119-
}
120-
continuation.finish()
121-
} catch {
122-
// Do not fallback to second session if the primary session contains history or has
123-
// already yielded data.
124-
if didYield || primary._hasHistory {
125-
continuation.finish(throwing: error)
126-
return
127-
}
128-
129-
// Fallback to the second session if the first fails or is unavailable.
130-
let stream = secondary._streamResponse(
107+
return AsyncThrowingStream { continuation in
108+
let task = Task {
109+
// First try the primary session.
110+
let stream = primary._streamResponse(
131111
to: prompt,
132112
schema: schema,
133113
includeSchemaInPrompt: includeSchemaInPrompt,
134114
options: options
135115
)
136116

117+
var didYield = false
137118
do {
138119
for try await snapshot in stream {
120+
didYield = true
139121
continuation.yield(snapshot)
140122
}
141123
continuation.finish()
142124
} catch {
143-
continuation.finish(throwing: error)
125+
// Do not fallback to second session if the primary session contains history or has
126+
// already yielded data.
127+
if didYield || primary._hasHistory {
128+
continuation.finish(throwing: error)
129+
return
130+
}
131+
132+
// Fallback to the second session if the first fails or is unavailable.
133+
let stream = secondary._streamResponse(
134+
to: prompt,
135+
schema: schema,
136+
includeSchemaInPrompt: includeSchemaInPrompt,
137+
options: options
138+
)
139+
140+
do {
141+
for try await snapshot in stream {
142+
continuation.yield(snapshot)
143+
}
144+
continuation.finish()
145+
} catch {
146+
continuation.finish(throwing: error)
147+
}
144148
}
145149
}
150+
continuation.onTermination = { _ in task.cancel() }
146151
}
147-
continuation.onTermination = { _ in task.cancel() }
148152
}
149153
}
150154
}

‎FirebaseAI/Tests/Unit/GenerativeModelSessionTests.swift‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -556,6 +556,67 @@
556556
XCTAssertEqual(functionResponse.response, ["result": .string(CurrentTimeTool.currentTime)])
557557
}
558558

559+
func testRespondTo_hybridHeader() async throws {
560+
let model1 = try mockGeminiModel(modelName: "test-model-1")
561+
let model2 = try mockGeminiModel(modelName: "test-model-2")
562+
let session = GenerativeModelSession(model: HybridModel(primary: model1, secondary: model2))
563+
564+
let bundle = BundleTestUtil.bundle()
565+
let fileURL = try XCTUnwrap(bundle.url(
566+
forResource: "unary-success-thinking-reply-thought-summary",
567+
withExtension: "json",
568+
subdirectory: googleAISubdirectory
569+
))
570+
571+
MockURLProtocol.requestHandler = { request in
572+
let apiClientHeader = try XCTUnwrap(request.value(forHTTPHeaderField: "x-goog-api-client"))
573+
let apiClientTags = apiClientHeader.components(separatedBy: " ")
574+
575+
XCTAssertTrue(apiClientTags.contains("hybrid"), "Header was: \(apiClientTags)")
576+
577+
let requestURL = try XCTUnwrap(request.url)
578+
let response = try XCTUnwrap(HTTPURLResponse(
579+
url: requestURL,
580+
statusCode: 200,
581+
httpVersion: nil,
582+
headerFields: nil
583+
))
584+
return (response, fileURL.lines)
585+
}
586+
587+
_ = try await session.respond(to: testPrompt)
588+
}
589+
590+
func testRespondTo_noHybridHeader() async throws {
591+
let model = try mockGeminiModel()
592+
let session = GenerativeModelSession(model: model)
593+
594+
let bundle = BundleTestUtil.bundle()
595+
let fileURL = try XCTUnwrap(bundle.url(
596+
forResource: "unary-success-thinking-reply-thought-summary",
597+
withExtension: "json",
598+
subdirectory: googleAISubdirectory
599+
))
600+
601+
MockURLProtocol.requestHandler = { request in
602+
let apiClientHeader = try XCTUnwrap(request.value(forHTTPHeaderField: "x-goog-api-client"))
603+
let apiClientTags = apiClientHeader.components(separatedBy: " ")
604+
605+
XCTAssertFalse(apiClientTags.contains("hybrid"), "Header was: \(apiClientTags)")
606+
607+
let requestURL = try XCTUnwrap(request.url)
608+
let response = try XCTUnwrap(HTTPURLResponse(
609+
url: requestURL,
610+
statusCode: 200,
611+
httpVersion: nil,
612+
headerFields: nil
613+
))
614+
return (response, fileURL.lines)
615+
}
616+
617+
_ = try await session.respond(to: testPrompt)
618+
}
619+
559620
// MARK: - Helper Utilities
560621

561622
func mockGeminiModel(modelName: String? = nil, modelResourceName: String? = nil,

0 commit comments

Comments
 (0)