diff --git a/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift b/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift index b074ddca..9f7d25ff 100644 --- a/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift @@ -282,6 +282,10 @@ public struct GeminiLanguageModel: LanguageModel { var transcript = session.transcript + // The entries this call adds, which is what the response reports. `transcript` keeps the + // full conversation because each iteration rebuilds the request from it. + var entries: [Transcript.Entry] = [] + // Multi-turn conversation loop for tool calling while true { let params = try createGenerateContentParams( @@ -318,20 +322,26 @@ public struct GeminiLanguageModel: LanguageModel { switch resolution { case .stop(let calls): if !calls.isEmpty { - transcript.append(.toolCalls(Transcript.ToolCalls(calls))) + entries.append(.toolCalls(Transcript.ToolCalls(calls))) } let empty = try emptyResponseContent(for: type) return LanguageModelSession.Response( content: empty.content, rawContent: empty.rawContent, - transcriptEntries: ArraySlice(transcript) + transcriptEntries: ArraySlice(entries) ) case .invocations(let invocations): if !invocations.isEmpty { - transcript.append(.toolCalls(Transcript.ToolCalls(invocations.map(\.call)))) + let calls = Transcript.Entry.toolCalls( + Transcript.ToolCalls(invocations.map(\.call)) + ) + transcript.append(calls) + entries.append(calls) for invocation in invocations { - transcript.append(.toolOutput(invocation.output)) + let output = Transcript.Entry.toolOutput(invocation.output) + transcript.append(output) + entries.append(output) } } @@ -352,7 +362,7 @@ public struct GeminiLanguageModel: LanguageModel { return LanguageModelSession.Response( content: text as! Content, rawContent: GeneratedContent(text), - transcriptEntries: ArraySlice(transcript) + transcriptEntries: ArraySlice(entries) ) } @@ -361,7 +371,7 @@ public struct GeminiLanguageModel: LanguageModel { return LanguageModelSession.Response( content: content, rawContent: generatedContent, - transcriptEntries: ArraySlice(transcript) + transcriptEntries: ArraySlice(entries) ) } } @@ -613,7 +623,8 @@ private func resolveFunctionCalls( Transcript.ToolCall( id: callID, toolName: call.name, - arguments: args + arguments: args, + providerMetadata: call.thoughtSignature.map { [thoughtSignatureMetadataKey: $0] } ) ) } @@ -786,7 +797,13 @@ extension Transcript { // Add model's response with function calls let functionCallParts: [GeminiPart] = toolCalls.map { call in let args = try? fromGeneratedContent(call.arguments) - return .functionCall(GeminiFunctionCall(name: call.toolName, args: args)) + return .functionCall( + GeminiFunctionCall( + name: call.toolName, + args: args, + thoughtSignature: call.providerMetadata?[thoughtSignatureMetadataKey] + ) + ) } messages.append( .init( @@ -895,8 +912,11 @@ private enum GeminiPart: Codable, Sendable { let text = try container.decode(String.self, forKey: .text) self = .text(GeminiTextPart(text: text)) } else if container.contains(.functionCall) { - // Note: thoughtSignature may be present but is ignored - self = .functionCall(try container.decode(GeminiFunctionCall.self, forKey: .functionCall)) + // `thoughtSignature` is a sibling of `functionCall` within the part, not a member of it. + // Thinking models require it to be echoed back verbatim, so it travels with the call. + var call = try container.decode(GeminiFunctionCall.self, forKey: .functionCall) + call.thoughtSignature = try container.decodeIfPresent(String.self, forKey: .thoughtSignature) + self = .functionCall(call) } else if container.contains(.functionResponse) { self = .functionResponse(try container.decode(GeminiFunctionResponse.self, forKey: .functionResponse)) } else if container.contains(.inlineData) { @@ -920,6 +940,7 @@ private enum GeminiPart: Codable, Sendable { try container.encode(part.text, forKey: .text) case .functionCall(let call): try container.encode(call, forKey: .functionCall) + try container.encodeIfPresent(call.thoughtSignature, forKey: .thoughtSignature) case .functionResponse(let response): try container.encode(response, forKey: .functionResponse) case .inlineData(let data): @@ -973,10 +994,25 @@ private func convertSegmentsToGeminiParts(_ segments: [Transcript.Segment]) -> [ return parts } +/// Key under which a thought signature is kept in ``Transcript/ToolCall/providerMetadata``. +private let thoughtSignatureMetadataKey = "thoughtSignature" + private struct GeminiFunctionCall: Codable, Sendable { let name: String let args: [String: JSONValue]? + /// The opaque thought signature Gemini attached to this call, if any. + /// + /// Deliberately absent from ``CodingKeys``: on the wire the signature sits on the enclosing + /// part, so ``GeminiPart`` is what reads and writes it. + var thoughtSignature: String? + + init(name: String, args: [String: JSONValue]?, thoughtSignature: String? = nil) { + self.name = name + self.args = args + self.thoughtSignature = thoughtSignature + } + enum CodingKeys: String, CodingKey { case name case args diff --git a/Sources/AnyLanguageModel/Transcript.swift b/Sources/AnyLanguageModel/Transcript.swift index dcd1fadc..54d32da5 100644 --- a/Sources/AnyLanguageModel/Transcript.swift +++ b/Sources/AnyLanguageModel/Transcript.swift @@ -332,10 +332,26 @@ public struct Transcript: Sendable, Equatable, Codable { /// Arguments to pass to the invoked tool. public var arguments: GeneratedContent - public init(id: String, toolName: String, arguments: GeneratedContent) { + /// Opaque, provider-specific state attached to this call. + /// + /// Some providers hand back values alongside a tool call that must be echoed + /// verbatim on every later request — Gemini's thought signatures, for instance. + /// The values are meaningful only to the model that produced them; nothing else + /// should interpret or modify them. + /// + /// - Note: This property is exclusive to AnyLanguageModel + public var providerMetadata: [String: String]? + + public init( + id: String, + toolName: String, + arguments: GeneratedContent, + providerMetadata: [String: String]? = nil + ) { self.id = id self.toolName = toolName self.arguments = arguments + self.providerMetadata = providerMetadata } } diff --git a/Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift b/Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift new file mode 100644 index 00000000..c0f4c3ef --- /dev/null +++ b/Tests/AnyLanguageModelTests/GeminiThoughtSignatureTests.swift @@ -0,0 +1,136 @@ +import Foundation +import Testing + +@testable import AnyLanguageModel + +#if canImport(Darwin) && !canImport(AsyncHTTPClient) + + @Suite("GeminiLanguageModel thought signatures", .serialized) + struct GeminiThoughtSignatureTests { + private static let signature = "CvsBAdHtim8n5xQK1pVX2H0lPQeXAMPLEsignature==" + + private func makeModel() -> GeminiLanguageModel { + GeminiLanguageModel( + apiKey: "test-key", + model: "gemini-3.6-flash", + session: StubURLProtocol.makeSession() + ) + } + + private func functionCallResponse(signature: String?) -> String { + let signatureField = signature.map { ", \"thoughtSignature\": \"\($0)\"" } ?? "" + return """ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { "name": "getWeather", "args": { "city": "Paris" } }\(signatureField) + } + ] + }, + "finishReason": "STOP" + } + ] + } + """ + } + + private func textResponse(_ text: String) -> String { + """ + { + "candidates": [ + { + "content": { "role": "model", "parts": [{ "text": "\(text)" }] }, + "finishReason": "STOP" + } + ] + } + """ + } + + /// Signatures of every `functionCall` part in a request body, in order. + /// A part without a signature contributes `nil`. + private func functionCallSignatures(in body: Data) throws -> [String?] { + let json = try JSONSerialization.jsonObject(with: body) as? [String: Any] + let contents = json?["contents"] as? [[String: Any]] ?? [] + return contents.flatMap { content -> [String?] in + let parts = content["parts"] as? [[String: Any]] ?? [] + return parts.compactMap { part -> String?? in + guard part["functionCall"] != nil else { return nil } + return .some(part["thoughtSignature"] as? String) + } + } + } + + /// The number of `contents` entries in a request body. + private func contentCount(in body: Data) throws -> Int { + let json = try JSONSerialization.jsonObject(with: body) as? [String: Any] + return (json?["contents"] as? [[String: Any]] ?? []).count + } + + @Test("echoes the thought signature back with the function results") + func echoesSignatureOnFollowUpRequest() async throws { + StubURLProtocol.reset() + StubURLProtocol.enqueue(json: functionCallResponse(signature: Self.signature)) + StubURLProtocol.enqueue(json: textResponse("It is sunny in Paris.")) + + let session = LanguageModelSession(model: makeModel(), tools: [WeatherTool()]) + let response = try await session.respond(to: "What is the weather in Paris?") + + #expect(response.content == "It is sunny in Paris.") + + let bodies = StubURLProtocol.recordedBodies + try #require(bodies.count == 2) + #expect(try functionCallSignatures(in: bodies[1]) == [Self.signature]) + } + + @Test("keeps the signature on the tool call when the conversation continues") + func keepsSignatureOnLaterTurn() async throws { + StubURLProtocol.reset() + StubURLProtocol.enqueue(json: functionCallResponse(signature: Self.signature)) + StubURLProtocol.enqueue(json: textResponse("It is sunny in Paris.")) + StubURLProtocol.enqueue(json: textResponse("You asked about the weather in Paris.")) + + let session = LanguageModelSession(model: makeModel(), tools: [WeatherTool()]) + _ = try await session.respond(to: "What is the weather in Paris?") + _ = try await session.respond(to: "What did I just ask about?") + + let bodies = StubURLProtocol.recordedBodies + try #require(bodies.count == 3) + + // The tool call is replayed as history on the next turn, and Gemini rejects a request + // whose functionCall parts have lost their signatures. + let signatures = try functionCallSignatures(in: bodies[2]) + #expect(!signatures.isEmpty) + #expect(signatures.allSatisfy { $0 == Self.signature }) + } + + @Test("does not replay the conversation history on later turns") + func doesNotReplayHistoryOnLaterTurns() async throws { + StubURLProtocol.reset() + StubURLProtocol.enqueue(json: functionCallResponse(signature: Self.signature)) + StubURLProtocol.enqueue(json: textResponse("It is sunny in Paris.")) + StubURLProtocol.enqueue(json: textResponse("You asked about the weather.")) + StubURLProtocol.enqueue(json: textResponse("Paris, specifically.")) + + let session = LanguageModelSession(model: makeModel(), tools: [WeatherTool()]) + _ = try await session.respond(to: "What is the weather in Paris?") + _ = try await session.respond(to: "What did I just ask about?") + _ = try await session.respond(to: "Which city was that?") + + let bodies = StubURLProtocol.recordedBodies + try #require(bodies.count == 4) + + // One entry per turn, plus the tool call and its output. A response that reported the + // whole transcript instead of its own entries would compound: 1, 3, 6, 14. + #expect(try bodies.map { try contentCount(in: $0) } == [1, 3, 5, 7]) + + // The tool call is replayed once, still signed, however many turns later. + #expect(try functionCallSignatures(in: bodies[3]) == [Self.signature]) + } + } + +#endif diff --git a/Tests/AnyLanguageModelTests/Shared/StubURLProtocol.swift b/Tests/AnyLanguageModelTests/Shared/StubURLProtocol.swift new file mode 100644 index 00000000..49da63fc --- /dev/null +++ b/Tests/AnyLanguageModelTests/Shared/StubURLProtocol.swift @@ -0,0 +1,92 @@ +import Foundation + +@testable import AnyLanguageModel + +#if canImport(Darwin) && !canImport(AsyncHTTPClient) + + /// A `URLProtocol` that answers requests from a queue of canned responses and records + /// every request body it sees, so request/response round trips can be asserted offline. + final class StubURLProtocol: URLProtocol { + struct Exchange: Sendable { + var statusCode: Int = 200 + var body: Data + } + + private struct State: Sendable { + var pending: [Exchange] = [] + var recordedBodies: [Data] = [] + } + + private static let state = Locked(State()) + + /// Discards queued responses and recorded bodies. + static func reset() { + state.withLock { $0 = State() } + } + + /// Queues one JSON response, returned to the next request that arrives. + static func enqueue(json: String, statusCode: Int = 200) { + state.withLock { $0.pending.append(Exchange(statusCode: statusCode, body: Data(json.utf8))) } + } + + /// The bodies of the requests seen so far, in order. + static var recordedBodies: [Data] { + state.withLock { $0.recordedBodies } + } + + /// A session that routes every request to this protocol. + static func makeSession() -> URLSession { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [StubURLProtocol.self] + return URLSession(configuration: configuration) + } + + override class func canInit(with request: URLRequest) -> Bool { true } + + override class func canonicalRequest(for request: URLRequest) -> URLRequest { request } + + override func startLoading() { + // URLSession moves `httpBody` to `httpBodyStream` before the protocol sees the request. + let body = request.httpBody ?? request.httpBodyStream.map(Self.readAll) ?? Data() + + let exchange = Self.state.withLock { state -> Exchange? in + state.recordedBodies.append(body) + return state.pending.isEmpty ? nil : state.pending.removeFirst() + } + + guard let exchange, let url = request.url else { + client?.urlProtocol(self, didFailWithError: URLError(.resourceUnavailable)) + return + } + + let response = HTTPURLResponse( + url: url, + statusCode: exchange.statusCode, + httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"] + )! + + client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed) + client?.urlProtocol(self, didLoad: exchange.body) + client?.urlProtocolDidFinishLoading(self) + } + + override func stopLoading() {} + + private static func readAll(_ stream: InputStream) -> Data { + stream.open() + defer { stream.close() } + + var data = Data() + let bufferSize = 4096 + var buffer = [UInt8](repeating: 0, count: bufferSize) + while true { + let read = stream.read(&buffer, maxLength: bufferSize) + if read <= 0 { break } + data.append(buffer, count: read) + } + return data + } + } + +#endif diff --git a/Tests/AnyLanguageModelTests/TranscriptTests.swift b/Tests/AnyLanguageModelTests/TranscriptTests.swift index 037c70a0..4f3eb65e 100644 --- a/Tests/AnyLanguageModelTests/TranscriptTests.swift +++ b/Tests/AnyLanguageModelTests/TranscriptTests.swift @@ -137,4 +137,36 @@ struct TranscriptTests { ) #expect(firstToolDefinition == secondToolDefinition) } + @Test func toolCallOmitsProviderMetadataWhenAbsent() throws { + // Guards decoding of transcripts encoded before `providerMetadata` existed. Making the + // property non-optional breaks this: synthesized `Codable` demands a key for every + // non-optional property and never consults its default value. + let arguments = try GeneratedContent(json: #"{"city":"Cupertino"}"#) + let call = Transcript.ToolCall(id: "call-id", toolName: "getWeather", arguments: arguments) + + let data = try JSONEncoder().encode(call) + var object = try #require(try JSONSerialization.jsonObject(with: data) as? [String: Any]) + #expect(object["providerMetadata"] == nil) + + object.removeValue(forKey: "providerMetadata") + let encodedBeforeTheFieldExisted = try JSONSerialization.data(withJSONObject: object) + let decoded = try JSONDecoder().decode(Transcript.ToolCall.self, from: encodedBeforeTheFieldExisted) + #expect(decoded.providerMetadata == nil) + #expect(decoded.toolName == "getWeather") + } + + @Test func toolCallRoundTripsProviderMetadata() throws { + let arguments = try GeneratedContent(json: #"{"city":"Cupertino"}"#) + let call = Transcript.ToolCall( + id: "call-id", + toolName: "getWeather", + arguments: arguments, + providerMetadata: ["thoughtSignature": "opaque-signature"] + ) + + let data = try JSONEncoder().encode(call) + let decoded = try JSONDecoder().decode(Transcript.ToolCall.self, from: data) + + #expect(decoded.providerMetadata == ["thoughtSignature": "opaque-signature"]) + } }