Skip to content

Support tool calling while streaming Gemini responses #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ActuallyTaylor merged 3 commits into from
Aug 3, 2026
Merged

Support tool calling while streaming Gemini responses #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0ba54b6
Support tool calling while streaming Gemini responses
claude Aug 3, 2026
328953a
Add streamWithTools test for Gemini
claude Aug 3, 2026
3107112
Echo Gemini thought signatures back with function calls
claude Aug 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
166 changes: 119 additions & 47 deletions Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<GeminiGenerateContentResponse, any Error> =
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<GeminiGenerateContentResponse, any Error> =
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()
Expand Down Expand Up @@ -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) {
Expand All @@ -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):
Expand Down Expand Up @@ -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
Expand Down
44 changes: 42 additions & 2 deletions Tests/AnyLanguageModelTests/GeminiLanguageModelTests.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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<String>.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)

Expand Down Expand Up @@ -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 {
Expand Down