mirror of
https://github.com/firebase/firebase-ios-sdk
synced 2026-11-09 22:36:01 +05:45
Andrew Heard
bc9f9ea1da
Add FoundationModelsFake, a nested Swift package that fakes Apple's Foundation Models framework, so that GeminiLanguageModel (GLM) can build and run its unit tests on Linux and on OS versions without the real framework. It's used only for development and testing. - GLM depends on the fake through a `FoundationModelsFake` package trait, which is on by default where Darwin isn't available (e.g., Linux). With the trait off, the fake isn't built. The root `Package.swift` and podspecs can't reach it. - The fake mirrors Apple's public API and availability in the 27.2 SDKs, including `@Generable` macro expansions. It omits PrivateCloudComputeLanguageModel, OCRTool, BarcodeReaderTool, the adapter and feedback APIs, and deprecated APIs. SystemLanguageModel is always unavailable. Below Swift 6.4 it compiles to an empty module. - Decoding rejects transcript images with `file:` URLs and schemas whose references never resolve, instead of reading local files or overflowing the stack. - Without ImageIO, GLM passes PNG, JPEG, WebP, HEIC, and HEIF images through to Gemini after checking their signatures, and rejects images that would need re-encoding. - Add CI jobs that run both packages' tests against the fake: one on Linux in a digest-pinned `swift:6.4.2` container, and one on macOS with Xcode 27.2. Skip `sdk.ai` for changes that only touch the fake.
622 lines
25 KiB
Swift
622 lines
25 KiB
Swift
// Copyright 2026 Google LLC
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#if compiler(>=6.4) && (FoundationModelsFake || canImport(FoundationModels))
|
|
#if canImport(CoreGraphics)
|
|
import CoreGraphics
|
|
#endif
|
|
import Foundation
|
|
#if FoundationModelsFake
|
|
import FoundationModelsFake
|
|
#else
|
|
import FoundationModels
|
|
#endif
|
|
import GeminiAPIClient
|
|
import GeminiAPIDataModels
|
|
import GeminiTestUtilities
|
|
import Synchronization
|
|
import Testing
|
|
|
|
#if canImport(FoundationNetworking)
|
|
import FoundationNetworking
|
|
#endif
|
|
|
|
@testable import GeminiLanguageModel
|
|
|
|
#if FoundationModelsFake
|
|
/// Foundation Models' `Tool`, spelled out to avoid clashing with
|
|
/// `GeminiAPIDataModels.Tool`.
|
|
@available(iOS 26.0, macOS 26.0, visionOS 26.0, watchOS 27.0, *)
|
|
@available(tvOS, unavailable)
|
|
typealias FoundationModelsTool = FoundationModelsFake.Tool
|
|
#else
|
|
/// Foundation Models' `Tool`, spelled out to avoid clashing with
|
|
/// `GeminiAPIDataModels.Tool`.
|
|
@available(iOS 26.0, macOS 26.0, visionOS 26.0, watchOS 27.0, *)
|
|
@available(tvOS, unavailable)
|
|
typealias FoundationModelsTool = FoundationModels.Tool
|
|
#endif
|
|
|
|
@Suite("GeminiLanguageModel Tests", .serialized, .requireFoundationModels)
|
|
struct GeminiLanguageModelTests {
|
|
@Generable(description: "A city summary")
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
@available(tvOS, unavailable)
|
|
struct CitySummary {
|
|
var name: String
|
|
var population: Int
|
|
}
|
|
|
|
@Test
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
func modelInitializationAndCapabilities() {
|
|
let model = GeminiLanguageModel(
|
|
modelResource: .gemini35FlashLite,
|
|
endpointConfiguration: .geminiDeveloperAPI
|
|
)
|
|
|
|
#expect(model.executorConfiguration.modelResource == .gemini35FlashLite)
|
|
#expect(model.executorConfiguration.endpointConfiguration == .geminiDeveloperAPI)
|
|
#expect(model.capabilities.contains(.reasoning))
|
|
#expect(model.capabilities.contains(.guidedGeneration))
|
|
#expect(model.capabilities.contains(.toolCalling))
|
|
#expect(model.capabilities.contains(.vision))
|
|
}
|
|
|
|
@Test
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
func sessionRespondGuidedGeneration() async throws {
|
|
defer { MockHTTPURLProtocol.reset() }
|
|
let model = Self.makeMockModel()
|
|
let expectedURL = try Self.makeExpectedStreamURL()
|
|
let httpResponse = try HTTPURLResponse.mock(
|
|
url: expectedURL,
|
|
headerFields: ["Content-Type": "text/event-stream"]
|
|
)
|
|
let ssePayload = """
|
|
data: {"candidates": [{"content": {"parts": [{"text": "{\\"name\\": \\"Tokyo\\", \\"population\\": 14000000}"}], "role": "model"}, "finishReason": "STOP", "index": 0}]}
|
|
|
|
"""
|
|
let receivedRequest = Mutex<GenerateContentRequest?>(nil)
|
|
MockHTTPURLProtocol.setHandler(for: expectedURL) { request, proto in
|
|
if let body = request.httpBodyData,
|
|
let decoded = try? JSONDecoder().decode(GenerateContentRequest.self, from: body)
|
|
{
|
|
receivedRequest.withLock { $0 = decoded }
|
|
}
|
|
proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed)
|
|
proto.client?.urlProtocol(proto, didLoad: Data(ssePayload.utf8))
|
|
proto.client?.urlProtocolDidFinishLoading(proto)
|
|
}
|
|
|
|
let session = LanguageModelSession(model: model)
|
|
|
|
let response = try await session.respond(
|
|
to: "Tell me about Tokyo",
|
|
generating: CitySummary.self
|
|
)
|
|
|
|
#expect(response.content.name == "Tokyo")
|
|
#expect(response.content.population == 14_000_000)
|
|
let capturedRequest = try #require(receivedRequest.withLock { $0 })
|
|
let generationConfig = try #require(capturedRequest.generationConfig)
|
|
#expect(generationConfig.responseMIMEType == "application/json")
|
|
let schemaObject = try #require(generationConfig.responseJSONSchema)
|
|
#expect(schemaObject["x-order"] == nil)
|
|
#expect(schemaObject["propertyOrdering"] != nil)
|
|
}
|
|
|
|
@Test
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
func sessionRespondSingleTurn() async throws {
|
|
defer { MockHTTPURLProtocol.reset() }
|
|
let model = Self.makeMockModel()
|
|
let expectedURL = try Self.makeExpectedStreamURL()
|
|
let httpResponse = try HTTPURLResponse.mock(
|
|
url: expectedURL,
|
|
headerFields: ["Content-Type": "text/event-stream"]
|
|
)
|
|
let ssePayload = """
|
|
data: {"candidates": [{"content": {"parts": [{"text": "Hello world from Gemini!"}], "role": "model"}, "finishReason": "STOP", "index": 0}]}
|
|
|
|
"""
|
|
MockHTTPURLProtocol.setHandler(for: expectedURL) { request, proto in
|
|
proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed)
|
|
proto.client?.urlProtocol(proto, didLoad: Data(ssePayload.utf8))
|
|
proto.client?.urlProtocolDidFinishLoading(proto)
|
|
}
|
|
|
|
let session = LanguageModelSession(model: model)
|
|
|
|
let response = try await session.respond(to: "Hello")
|
|
|
|
#expect(response.content == "Hello world from Gemini!")
|
|
}
|
|
|
|
#if canImport(ImageIO)
|
|
@Test
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
func sessionRespondWithImageAttachmentSendsInlineData() async throws {
|
|
defer { MockHTTPURLProtocol.reset() }
|
|
let model = Self.makeMockModel()
|
|
let expectedURL = try Self.makeExpectedStreamURL()
|
|
let httpResponse = try HTTPURLResponse.mock(
|
|
url: expectedURL,
|
|
headerFields: ["Content-Type": "text/event-stream"]
|
|
)
|
|
let ssePayload = """
|
|
data: {"candidates": [{"content": {"parts": [{"text": "red"}], "role": "model"}, "finishReason": "STOP", "index": 0}]}
|
|
|
|
"""
|
|
let receivedRequest = Mutex<GenerateContentRequest?>(nil)
|
|
MockHTTPURLProtocol.setHandler(for: expectedURL) { request, proto in
|
|
if let body = request.httpBodyData,
|
|
let decoded = try? JSONDecoder().decode(GenerateContentRequest.self, from: body)
|
|
{
|
|
receivedRequest.withLock { $0 = decoded }
|
|
}
|
|
proto.client?.urlProtocol(
|
|
proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed)
|
|
proto.client?.urlProtocol(proto, didLoad: Data(ssePayload.utf8))
|
|
proto.client?.urlProtocolDidFinishLoading(proto)
|
|
}
|
|
let image = try TestImages.solid(.red, width: 32, height: 32)
|
|
let session = LanguageModelSession(model: model)
|
|
|
|
let response = try await session.respond {
|
|
"What color is this image?"
|
|
Attachment(image)
|
|
}
|
|
|
|
#expect(response.content == "red")
|
|
let capturedRequest = try #require(receivedRequest.withLock { $0 })
|
|
let parts = try #require(capturedRequest.contents?.first?.parts)
|
|
#expect(parts.count == 2)
|
|
#expect(parts[0].data == .text("What color is this image?"))
|
|
guard case .inlineData(let blob) = parts[1].data else {
|
|
Issue.record("Expected inline data, got \(String(describing: parts[1].data)).")
|
|
return
|
|
}
|
|
#expect(blob.mimeType == "image/jpeg")
|
|
let decodedImage = try TestImages.decode(try #require(blob.data))
|
|
#expect(decodedImage.width == 32)
|
|
#expect(try TestImages.quadrantColors(of: decodedImage).topLeft == .red)
|
|
}
|
|
#endif // canImport(ImageIO)
|
|
|
|
@Test
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
func sessionRespondMultiTurn() async throws {
|
|
defer { MockHTTPURLProtocol.reset() }
|
|
let model = Self.makeMockModel()
|
|
let expectedURL = try Self.makeExpectedStreamURL()
|
|
let httpResponse = try HTTPURLResponse.mock(
|
|
url: expectedURL,
|
|
headerFields: ["Content-Type": "text/event-stream"]
|
|
)
|
|
let firstPayload = """
|
|
data: {"candidates": [{"content": {"parts": [{"text": "Nice to meet you, Alice."}], "role": "model"}, "finishReason": "STOP", "index": 0}]}
|
|
|
|
"""
|
|
let secondPayload = """
|
|
data: {"candidates": [{"content": {"parts": [{"text": "Your name is Alice."}], "role": "model"}, "finishReason": "STOP", "index": 0}]}
|
|
|
|
"""
|
|
let requestCount = Mutex<Int>(0)
|
|
let lastReceivedContents = Mutex<[Content]>([])
|
|
MockHTTPURLProtocol.setHandler(for: expectedURL) { request, proto in
|
|
let currentCount = requestCount.withLock { count in
|
|
count += 1
|
|
return count
|
|
}
|
|
if let body = request.httpBodyData,
|
|
let decoded = try? JSONDecoder().decode(GenerateContentRequest.self, from: body)
|
|
{
|
|
lastReceivedContents.withLock { $0 = decoded.contents ?? [] }
|
|
}
|
|
proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed)
|
|
let payload = currentCount == 1 ? firstPayload : secondPayload
|
|
proto.client?.urlProtocol(proto, didLoad: Data(payload.utf8))
|
|
proto.client?.urlProtocolDidFinishLoading(proto)
|
|
}
|
|
|
|
let session = LanguageModelSession(model: model)
|
|
|
|
let firstResponse = try await session.respond(to: "My name is Alice.")
|
|
let secondResponse = try await session.respond(to: "What is my name?")
|
|
|
|
#expect(firstResponse.content == "Nice to meet you, Alice.")
|
|
#expect(secondResponse.content == "Your name is Alice.")
|
|
#expect(requestCount.withLock { $0 } == 2)
|
|
let contents = lastReceivedContents.withLock { $0 }
|
|
#expect(contents.count == 3)
|
|
#expect(contents.first?.parts?.first?.data == .text("My name is Alice."))
|
|
}
|
|
|
|
@Test
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
func sessionWithInstructions() async throws {
|
|
defer { MockHTTPURLProtocol.reset() }
|
|
let model = Self.makeMockModel()
|
|
let expectedURL = try Self.makeExpectedStreamURL()
|
|
let httpResponse = try HTTPURLResponse.mock(
|
|
url: expectedURL,
|
|
headerFields: ["Content-Type": "text/event-stream"]
|
|
)
|
|
let ssePayload = """
|
|
data: {"candidates": [{"content": {"parts": [{"text": "Ahoy matey!"}], "role": "model"}, "finishReason": "STOP", "index": 0}]}
|
|
|
|
"""
|
|
let receivedSystemInstruction = Mutex<Content?>(nil)
|
|
MockHTTPURLProtocol.setHandler(for: expectedURL) { request, proto in
|
|
if let body = request.httpBodyData,
|
|
let decoded = try? JSONDecoder().decode(GenerateContentRequest.self, from: body)
|
|
{
|
|
receivedSystemInstruction.withLock { $0 = decoded.systemInstruction }
|
|
}
|
|
proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed)
|
|
proto.client?.urlProtocol(proto, didLoad: Data(ssePayload.utf8))
|
|
proto.client?.urlProtocolDidFinishLoading(proto)
|
|
}
|
|
|
|
let session = LanguageModelSession(model: model, instructions: "Respond like a pirate.")
|
|
|
|
let response = try await session.respond(to: "Hello")
|
|
|
|
#expect(response.content == "Ahoy matey!")
|
|
let instruction = receivedSystemInstruction.withLock { $0 }
|
|
#expect(instruction?.parts?.first?.data == .text("Respond like a pirate."))
|
|
}
|
|
|
|
@Test
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
func sessionStreamResponse() async throws {
|
|
defer { MockHTTPURLProtocol.reset() }
|
|
let model = Self.makeMockModel()
|
|
let expectedURL = try Self.makeExpectedStreamURL()
|
|
let httpResponse = try HTTPURLResponse.mock(
|
|
url: expectedURL,
|
|
headerFields: ["Content-Type": "text/event-stream"]
|
|
)
|
|
let ssePayload = """
|
|
data: {"candidates": [{"content": {"parts": [{"text": "42", "thoughtSignature": "test-signature-123"}], "role": "model"}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"candidatesTokenCount": 2, "promptTokenCount": 5, "totalTokenCount": 7}}
|
|
|
|
"""
|
|
MockHTTPURLProtocol.setHandler(for: expectedURL) { request, proto in
|
|
proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed)
|
|
proto.client?.urlProtocol(proto, didLoad: Data(ssePayload.utf8))
|
|
proto.client?.urlProtocolDidFinishLoading(proto)
|
|
}
|
|
|
|
let session = LanguageModelSession(model: model)
|
|
|
|
let stream = session.streamResponse(to: "What is 6x7?")
|
|
var accumulated = ""
|
|
var lastUsage: LanguageModelSession.Usage?
|
|
for try await snapshot in stream {
|
|
accumulated = snapshot.content
|
|
lastUsage = snapshot.usage
|
|
}
|
|
|
|
#expect(accumulated == "42")
|
|
let usage = try #require(lastUsage)
|
|
#expect(usage.input.totalTokenCount == 5)
|
|
#expect(usage.output.totalTokenCount == 2)
|
|
#expect(usage.totalTokenCount == 7)
|
|
}
|
|
|
|
@Test
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
func rateLimitErrorMapping() async throws {
|
|
defer { MockHTTPURLProtocol.reset() }
|
|
let model = Self.makeMockModel()
|
|
let expectedURL = try Self.makeExpectedStreamURL()
|
|
let httpResponse = try HTTPURLResponse.mock(
|
|
url: expectedURL,
|
|
statusCode: 429,
|
|
headerFields: ["Content-Type": "application/json", "Retry-After": "30"]
|
|
)
|
|
let errorPayload = """
|
|
{"error": {"code": 429, "message": "Resource has been exhausted", "status": "RESOURCE_EXHAUSTED"}}
|
|
"""
|
|
MockHTTPURLProtocol.setHandler(for: expectedURL) { request, proto in
|
|
proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed)
|
|
proto.client?.urlProtocol(proto, didLoad: Data(errorPayload.utf8))
|
|
proto.client?.urlProtocolDidFinishLoading(proto)
|
|
}
|
|
|
|
let session = LanguageModelSession(model: model)
|
|
|
|
do {
|
|
_ = try await session.respond(to: "Hello")
|
|
Issue.record("Expected rateLimited error")
|
|
} catch LanguageModelError.rateLimited(let rateLimited) {
|
|
#expect(rateLimited.debugDescription.contains("Resource has been exhausted"))
|
|
} catch {
|
|
Issue.record("Unexpected error thrown: \(error)")
|
|
}
|
|
}
|
|
|
|
@Test
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
func guardrailViolationErrorMapping() async throws {
|
|
defer { MockHTTPURLProtocol.reset() }
|
|
let model = Self.makeMockModel()
|
|
let expectedURL = try Self.makeExpectedStreamURL()
|
|
let httpResponse = try HTTPURLResponse.mock(
|
|
url: expectedURL,
|
|
headerFields: ["Content-Type": "text/event-stream"]
|
|
)
|
|
let ssePayload = """
|
|
data: {"candidates": [{"finishReason": "SAFETY", "finishMessage": "Filtered for safety reasons"}]}
|
|
|
|
"""
|
|
MockHTTPURLProtocol.setHandler(for: expectedURL) { request, proto in
|
|
proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed)
|
|
proto.client?.urlProtocol(proto, didLoad: Data(ssePayload.utf8))
|
|
proto.client?.urlProtocolDidFinishLoading(proto)
|
|
}
|
|
|
|
let session = LanguageModelSession(model: model)
|
|
|
|
do {
|
|
_ = try await session.respond(to: "Harmful prompt")
|
|
Issue.record("Expected guardrailViolation error")
|
|
} catch LanguageModelError.guardrailViolation(let violation) {
|
|
#expect(violation.debugDescription.contains("Filtered for safety reasons"))
|
|
} catch {
|
|
Issue.record("Unexpected error thrown: \(error)")
|
|
}
|
|
}
|
|
|
|
@Test
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
func functionCallWithoutNameThrowsInvalidResponse() async throws {
|
|
defer { MockHTTPURLProtocol.reset() }
|
|
let model = Self.makeMockModel()
|
|
let expectedURL = try Self.makeExpectedStreamURL()
|
|
let httpResponse = try HTTPURLResponse.mock(
|
|
url: expectedURL,
|
|
headerFields: ["Content-Type": "text/event-stream"]
|
|
)
|
|
let ssePayload = """
|
|
data: {"candidates": [{"content": {"parts": [{"functionCall": {"id": "call-1", "args": {}}}], "role": "model"}, "finishReason": "STOP", "index": 0}]}
|
|
|
|
"""
|
|
MockHTTPURLProtocol.setHandler(for: expectedURL) { request, proto in
|
|
proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed)
|
|
proto.client?.urlProtocol(proto, didLoad: Data(ssePayload.utf8))
|
|
proto.client?.urlProtocolDidFinishLoading(proto)
|
|
}
|
|
let session = LanguageModelSession(model: model)
|
|
|
|
do {
|
|
_ = try await session.respond(to: "What's the weather?")
|
|
Issue.record("Expected GeminiLanguageModel.Error.invalidResponse")
|
|
} catch GeminiLanguageModel.Error.invalidResponse(let invalidResponse) {
|
|
#expect(invalidResponse.debugDescription.contains("without a name"))
|
|
} catch {
|
|
Issue.record("Unexpected error thrown: \(error)")
|
|
}
|
|
}
|
|
|
|
@Test
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
func timeoutErrorMapping() async throws {
|
|
defer { MockHTTPURLProtocol.reset() }
|
|
let model = Self.makeMockModel()
|
|
let expectedURL = try Self.makeExpectedStreamURL()
|
|
MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in
|
|
proto.client?.urlProtocol(proto, didFailWithError: URLError(.timedOut))
|
|
}
|
|
|
|
let session = LanguageModelSession(model: model)
|
|
|
|
do {
|
|
_ = try await session.respond(to: "Hello")
|
|
Issue.record("Expected timeout error")
|
|
} catch LanguageModelError.timeout {
|
|
} catch let error as URLError where error.code == .timedOut {
|
|
} catch {
|
|
Issue.record("Unexpected error thrown: \(error)")
|
|
}
|
|
}
|
|
|
|
@Test
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
func unsupportedTranscriptContentThrows() throws {
|
|
let structuredSegment = Transcript.Segment.structure(
|
|
Transcript.StructuredSegment(
|
|
id: "struct-1",
|
|
schemaName: "test",
|
|
content: GeneratedContent("test")
|
|
)
|
|
)
|
|
let transcript = Transcript(
|
|
entries: [.prompt(Transcript.Prompt(segments: [structuredSegment]))]
|
|
)
|
|
|
|
do {
|
|
_ = try GeminiTranscriptTranslator.translate(transcript)
|
|
Issue.record("Expected unsupportedTranscriptContent error")
|
|
} catch LanguageModelError.unsupportedTranscriptContent {
|
|
// Success
|
|
} catch {
|
|
Issue.record("Unexpected error thrown: \(error)")
|
|
}
|
|
}
|
|
|
|
/// A mock weather tool for testing tool call execution.
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
@available(tvOS, unavailable)
|
|
struct MockWeatherTool: FoundationModelsTool {
|
|
|
|
let name = "get_weather"
|
|
let description = "Get current weather"
|
|
|
|
@Generable
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
@available(tvOS, unavailable)
|
|
struct Arguments {
|
|
var location: String
|
|
}
|
|
|
|
func call(arguments: Arguments) async throws -> String {
|
|
"Sunny, 72°F in \(arguments.location)"
|
|
}
|
|
}
|
|
|
|
@Test
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
func sessionRespondWithToolCalling() async throws {
|
|
defer { MockHTTPURLProtocol.reset() }
|
|
let model = Self.makeMockModel()
|
|
let expectedURL = try Self.makeExpectedStreamURL()
|
|
let httpResponse = try HTTPURLResponse.mock(
|
|
url: expectedURL,
|
|
headerFields: ["Content-Type": "text/event-stream"]
|
|
)
|
|
let turn1SSE = """
|
|
data: {"candidates": [{"content": {"parts": [{"functionCall": {"id": "call-1", "name": "get_weather", "args": {"location": "Paris"}}}], "role": "model"}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"candidatesTokenCount": 5, "promptTokenCount": 10, "totalTokenCount": 15}}
|
|
|
|
"""
|
|
let turn2SSE = """
|
|
data: {"candidates": [{"content": {"parts": [{"text": "The weather in Paris is sunny, 72°F."}], "role": "model"}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"candidatesTokenCount": 8, "promptTokenCount": 20, "totalTokenCount": 28}}
|
|
|
|
"""
|
|
let requestCount = Mutex(0)
|
|
MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in
|
|
let currentCount = requestCount.withLock { count -> Int in
|
|
count += 1
|
|
return count
|
|
}
|
|
proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed)
|
|
let payload = currentCount == 1 ? turn1SSE : turn2SSE
|
|
proto.client?.urlProtocol(proto, didLoad: Data(payload.utf8))
|
|
proto.client?.urlProtocolDidFinishLoading(proto)
|
|
}
|
|
let session = LanguageModelSession(model: model, tools: [MockWeatherTool()])
|
|
|
|
let response = try await session.respond(to: "What is the weather in Paris?")
|
|
|
|
#expect(response.content == "The weather in Paris is sunny, 72°F.")
|
|
#expect(requestCount.withLock { $0 } == 2)
|
|
}
|
|
|
|
/// A mock time tool taking no arguments for testing empty argument tool calling.
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
@available(tvOS, unavailable)
|
|
struct MockTimeTool: FoundationModelsTool {
|
|
let name = "get_current_time"
|
|
let description = "Get the current time"
|
|
|
|
@Generable
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
@available(tvOS, unavailable)
|
|
struct Arguments {}
|
|
|
|
func call(arguments: Arguments) async throws -> String {
|
|
"12:00 PM"
|
|
}
|
|
}
|
|
|
|
@Test
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
func sessionRespondWithEmptyArgumentsToolCalling() async throws {
|
|
defer { MockHTTPURLProtocol.reset() }
|
|
let model = Self.makeMockModel()
|
|
let expectedURL = try Self.makeExpectedStreamURL()
|
|
let httpResponse = try HTTPURLResponse.mock(
|
|
url: expectedURL,
|
|
headerFields: ["Content-Type": "text/event-stream"]
|
|
)
|
|
let turn1SSE = """
|
|
data: {"candidates": [{"content": {"parts": [{"functionCall": {"id": "call-1", "name": "get_current_time"}}], "role": "model"}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"candidatesTokenCount": 5, "promptTokenCount": 10, "totalTokenCount": 15}}
|
|
|
|
"""
|
|
let turn2SSE = """
|
|
data: {"candidates": [{"content": {"parts": [{"text": "The time is 12:00 PM."}], "role": "model"}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"candidatesTokenCount": 8, "promptTokenCount": 20, "totalTokenCount": 28}}
|
|
|
|
"""
|
|
let requestCount = Mutex(0)
|
|
let capturedRequests = Mutex<[GenerateContentRequest]>([])
|
|
MockHTTPURLProtocol.setHandler(for: expectedURL) { request, proto in
|
|
let currentCount = requestCount.withLock { count -> Int in
|
|
count += 1
|
|
return count
|
|
}
|
|
if let body = request.httpBodyData,
|
|
let decoded = try? JSONDecoder().decode(GenerateContentRequest.self, from: body)
|
|
{
|
|
capturedRequests.withLock { $0.append(decoded) }
|
|
}
|
|
proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed)
|
|
let payload = currentCount == 1 ? turn1SSE : turn2SSE
|
|
proto.client?.urlProtocol(proto, didLoad: Data(payload.utf8))
|
|
proto.client?.urlProtocolDidFinishLoading(proto)
|
|
}
|
|
let session = LanguageModelSession(model: model, tools: [MockTimeTool()])
|
|
|
|
let response = try await session.respond(to: "What time is it?")
|
|
|
|
#expect(response.content == "The time is 12:00 PM.")
|
|
#expect(requestCount.withLock { $0 } == 2)
|
|
let requests = capturedRequests.withLock { $0 }
|
|
#expect(requests.count == 2)
|
|
let turn1Tools = try #require(requests.first?.tools)
|
|
#expect(turn1Tools.first?.functionDeclarations?.first?.name == "get_current_time")
|
|
let turn2Contents = try #require(requests.last?.contents)
|
|
let turn2UserTurn = try #require(turn2Contents.last)
|
|
guard case .functionResponse(let fr) = turn2UserTurn.parts?.first?.data else {
|
|
Issue.record("Expected functionResponse in turn 2")
|
|
return
|
|
}
|
|
#expect(fr.name == "get_current_time")
|
|
#expect(fr.response?["result"] == .string("12:00 PM"))
|
|
}
|
|
|
|
// MARK: - Helper Methods
|
|
|
|
@available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
|
|
|
|
private static func makeMockModel(
|
|
modelResource: ModelResource = .gemini38Flash,
|
|
endpointConfiguration: EndpointConfiguration = .geminiDeveloperAPI,
|
|
headerProvider: (@Sendable () async throws -> [String: String])? = nil
|
|
) -> GeminiLanguageModel {
|
|
let configuration = URLSessionConfiguration.ephemeral
|
|
configuration.protocolClasses = [MockHTTPURLProtocol.self]
|
|
return GeminiLanguageModel(
|
|
modelResource: modelResource,
|
|
endpointConfiguration: endpointConfiguration,
|
|
headerProvider: headerProvider,
|
|
configuration: configuration
|
|
)
|
|
}
|
|
|
|
private static func makeExpectedStreamURL(
|
|
host: String = EndpointConfiguration.geminiDeveloperAPIHost,
|
|
apiVersion: String = EndpointConfiguration.geminiDeveloperAPIVersion,
|
|
urlResourceName: String = ModelResource.gemini38FlashURLResourceName
|
|
) throws -> URL {
|
|
try #require(
|
|
URL(
|
|
string:
|
|
"https://\(host)/\(apiVersion)/\(urlResourceName):streamGenerateContent?alt=sse"
|
|
)
|
|
)
|
|
}
|
|
}
|
|
#endif // compiler(>=6.4) && (FoundationModelsFake || canImport(FoundationModels))
|