diff --git a/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift b/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift index 5724a449..14d486c2 100644 --- a/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift @@ -395,60 +395,122 @@ public struct GeminiLanguageModel: LanguageModel { let geminiTools = try buildTools(from: session.tools, serverTools: effectiveServerTools) - let params = try createGenerateContentParams( - contents: session.transcript.toGeminiContent(), - tools: geminiTools, - generating: type, - options: options, - thinking: effectiveThinking, - jsonMode: effectiveJsonMode - ) - - let body = try JSONEncoder().encode(params) - - let stream: AsyncThrowingStream = - httpSession - .fetchEventStream( - .post, - url: url, - headers: headers, - body: body + var contents: [GeminiContent] = session.transcript.toGeminiContent() + + // Multi-turn conversation loop for tool calling. + turnLoop: while true { + let params = try createGenerateContentParams( + contents: contents, + tools: geminiTools, + generating: type, + options: options, + thinking: effectiveThinking, + jsonMode: effectiveJsonMode ) - var accumulatedText = "" - - for try await chunk in stream { - guard let candidate = chunk.candidates.first else { continue } - - if let parts = candidate.content.parts { - for part in parts { - if case .text(let textPart) = part { - accumulatedText += textPart.text - - var raw: GeneratedContent - let content: Content.PartiallyGenerated? - - if type == String.self { - raw = GeneratedContent(accumulatedText) - content = (accumulatedText as! Content).asPartiallyGenerated() - } else { - raw = - (try? GeneratedContent(json: accumulatedText)) - ?? GeneratedContent(accumulatedText) - if let parsed = try? type.init(raw) { - content = parsed.asPartiallyGenerated() + let body = try JSONEncoder().encode(params) + + let events: AsyncThrowingStream = + httpSession + .fetchEventStream( + .post, + url: url, + headers: headers, + body: body + ) + + var accumulatedText = "" + var functionCalls: [GeminiFunctionCall] = [] + + for try await chunk in events { + guard let candidate = chunk.candidates.first else { continue } + + if let parts = candidate.content.parts { + for part in parts { + switch part { + case .functionCall(let call): + // Gemini delivers function calls as whole parts, so they only + // need to be collected until the turn ends. + functionCalls.append(call) + case .text(let textPart): + accumulatedText += textPart.text + + // Grow the observable transcript so a Transcript-driven UI updates live. + session.growStreamingTranscript(text: accumulatedText) + + var raw: GeneratedContent + let content: Content.PartiallyGenerated? + + if type == String.self { + raw = GeneratedContent(accumulatedText) + content = (accumulatedText as! Content).asPartiallyGenerated() } else { - // Skip invalid partial JSON until it parses cleanly. - content = nil + raw = + (try? GeneratedContent(json: accumulatedText)) + ?? GeneratedContent(accumulatedText) + if let parsed = try? type.init(raw) { + content = parsed.asPartiallyGenerated() + } else { + // Skip invalid partial JSON until it parses cleanly. + content = nil + } } - } - if let content { - continuation.yield(.init(content: content, rawContent: raw)) + if let content { + continuation.yield(.init(content: content, rawContent: raw)) + } + case .functionResponse, .inlineData, .fileData: + continue } } } } + + // The turn finished without the model asking for a tool, so the response is complete. + guard !functionCalls.isEmpty else { break turnLoop } + + let resolution = try await resolveFunctionCalls(functionCalls, session: session) + switch resolution { + case .stop(let calls): + if !calls.isEmpty { + session.appendTranscriptEntry(.toolCalls(Transcript.ToolCalls(calls))) + } + continuation.finish() + return + case .invocations(let invocations): + // Nothing was executed, so there is no new information to send back. + guard !invocations.isEmpty else { break turnLoop } + + // Tool calls must be recorded before their outputs. + session.appendTranscriptEntry( + .toolCalls(Transcript.ToolCalls(invocations.map(\.call))) + ) + + var modelParts: [GeminiPart] = [] + if !accumulatedText.isEmpty { + modelParts.append(.text(GeminiTextPart(text: accumulatedText))) + } + modelParts.append(contentsOf: functionCalls.map { GeminiPart.functionCall($0) }) + contents.append(GeminiContent(role: .model, parts: modelParts)) + + var responseParts: [GeminiPart] = [] + responseParts.reserveCapacity(invocations.count) + for invocation in invocations { + session.appendTranscriptEntry(.toolOutput(invocation.output)) + + responseParts.append( + .functionResponse( + GeminiFunctionResponse( + name: invocation.output.toolName, + response: try toJSONValue(invocation.output) + ) + ) + ) + } + + // Gemini expects function responses to come back from the user role. + contents.append(GeminiContent(role: .user, parts: responseParts)) + } } continuation.finish() @@ -796,8 +858,10 @@ 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)) + var call = try container.decode(GeminiFunctionCall.self, forKey: .functionCall) + // Carried alongside the call so it can be echoed back on the next turn. + 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) { @@ -821,6 +885,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): @@ -878,6 +943,13 @@ private struct GeminiFunctionCall: Codable, Sendable { let name: String let args: [String: JSONValue]? + /// An opaque signature newer models attach to a function call. + /// + /// It must be echoed back verbatim when the call is sent again on the follow-up turn, or the + /// API rejects the request with `INVALID_ARGUMENT`. It sits beside `functionCall` rather than + /// inside it, so ``GeminiPart`` reads and writes it and it stays out of `CodingKeys` here. + var thoughtSignature: String? + enum CodingKeys: String, CodingKey { case name case args diff --git a/Tests/AnyLanguageModelTests/GeminiLanguageModelTests.swift b/Tests/AnyLanguageModelTests/GeminiLanguageModelTests.swift index 9375851e..258800b8 100644 --- a/Tests/AnyLanguageModelTests/GeminiLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/GeminiLanguageModelTests.swift @@ -10,7 +10,7 @@ private let geminiAPIKey: String? = ProcessInfo.processInfo.environment["GEMINI_ struct GeminiLanguageModelTests { let model = GeminiLanguageModel( apiKey: geminiAPIKey!, - model: "gemini-2.5-flash" + model: "gemini-flash-latest" ) @Test func customHost() throws { @@ -104,6 +104,46 @@ struct GeminiLanguageModelTests { #expect(foundToolOutput) } + @Test func streamWithTools() async throws { + let weatherTool = WeatherTool() + let session = LanguageModelSession(model: model, tools: [weatherTool]) + + let stream = session.streamResponse(to: "How's the weather in San Francisco?") + + var snapshots: [LanguageModelSession.ResponseStream.Snapshot] = [] + + var toolAppearedInTranscript: Bool = false + var toolResponseAppearedInTranscript: Bool = false + + for try await snapshot in stream { + snapshots.append(snapshot) + + for entry in session.transcript { + switch entry { + case .toolCalls: + toolAppearedInTranscript = true + case .toolOutput: + toolResponseAppearedInTranscript = true + default: break + } + } + } + + #expect(toolAppearedInTranscript, "Expected a tool call to appear in the transcript during streaming.") + #expect( + toolResponseAppearedInTranscript, + "Expected a tool output to appear in the transcript during streaming." + ) + + var foundToolOutput = false + for case let .toolOutput(toolOutput) in session.transcript { + #expect(!toolOutput.id.isEmpty) + #expect(toolOutput.toolName == "getWeather") + foundToolOutput = true + } + #expect(foundToolOutput, "Expected the 'getWeather' tool to exist in the final transcript.") + } + @Test func withServerTools() async throws { let session = LanguageModelSession(model: model) @@ -206,7 +246,7 @@ struct GeminiLanguageModelTests { } private var model: GeminiLanguageModel { - GeminiLanguageModel(apiKey: geminiAPIKey!, model: "gemini-2.5-flash") + GeminiLanguageModel(apiKey: geminiAPIKey!, model: "gemini-flash-latest") } @Test func basicStructuredOutput() async throws {