Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
113 changes: 6 additions & 107 deletions Sources/AnyLanguageModel/Models/AnthropicLanguageModel.swift
Original file line number Diff line number Diff line change
Expand Up @@ -867,16 +867,6 @@ private func createMessageParams(

// MARK: - Tool Invocation Handling

private struct ToolInvocationResult {
let call: Transcript.ToolCall
let output: Transcript.ToolOutput
}

private enum ToolResolutionOutcome {
case stop(calls: [Transcript.ToolCall])
case invocations([ToolInvocationResult])
}

private func emptyResponseContent<Content: Generable>(
for type: Content.Type
) throws -> (content: Content, rawContent: GeneratedContent) {
Expand Down Expand Up @@ -916,105 +906,14 @@ private func resolveToolUses(
_ toolUses: [AnthropicToolUse],
session: LanguageModelSession
) async throws -> ToolResolutionOutcome {
if toolUses.isEmpty { return .invocations([]) }

var toolsByName: [String: any Tool] = [:]
for tool in session.tools {
if toolsByName[tool.name] == nil {
toolsByName[tool.name] = tool
}
}

var transcriptCalls: [Transcript.ToolCall] = []
transcriptCalls.reserveCapacity(toolUses.count)
for use in toolUses {
let args = try toGeneratedContent(use.input)
let callID = use.id
transcriptCalls.append(
Transcript.ToolCall(
id: callID,
toolName: use.name,
arguments: args
)
let calls = try toolUses.map { use in
Transcript.ToolCall(
id: use.id,
toolName: use.name,
arguments: try toGeneratedContent(use.input)
)
}

if let delegate = session.toolExecutionDelegate {
await delegate.didGenerateToolCalls(transcriptCalls, in: session)
}

guard !transcriptCalls.isEmpty else { return .invocations([]) }

var decisions: [ToolExecutionDecision] = []
decisions.reserveCapacity(transcriptCalls.count)

if let delegate = session.toolExecutionDelegate {
for call in transcriptCalls {
let decision = await delegate.toolCallDecision(for: call, in: session)
if case .stop = decision {
return .stop(calls: transcriptCalls)
}
decisions.append(decision)
}
} else {
decisions = Array(repeating: .execute, count: transcriptCalls.count)
}

var results: [ToolInvocationResult] = []
results.reserveCapacity(transcriptCalls.count)

for (index, call) in transcriptCalls.enumerated() {
switch decisions[index] {
case .stop:
// This branch should be unreachable because `.stop` returns during decision collection.
// Keep it as a defensive guard in case that logic changes.
return .stop(calls: transcriptCalls)
case .provideOutput(let segments):
let output = Transcript.ToolOutput(
id: call.id,
toolName: call.toolName,
segments: segments
)
if let delegate = session.toolExecutionDelegate {
await delegate.didExecuteToolCall(call, output: output, in: session)
}
results.append(ToolInvocationResult(call: call, output: output))
case .execute:
guard let tool = toolsByName[call.toolName] else {
let message = Transcript.Segment.text(.init(content: "Tool not found: \(call.toolName)"))
let output = Transcript.ToolOutput(
id: call.id,
toolName: call.toolName,
segments: [message]
)
if let delegate = session.toolExecutionDelegate {
await delegate.didExecuteToolCall(call, output: output, in: session)
}
results.append(ToolInvocationResult(call: call, output: output))
continue
}

do {
let segments = try await tool.makeOutputSegments(from: call.arguments)
let output = Transcript.ToolOutput(
id: call.id,
toolName: tool.name,
segments: segments
)
if let delegate = session.toolExecutionDelegate {
await delegate.didExecuteToolCall(call, output: output, in: session)
}
results.append(ToolInvocationResult(call: call, output: output))
} catch {
if let delegate = session.toolExecutionDelegate {
await delegate.didFailToolCall(call, error: error, in: session)
}
throw LanguageModelSession.ToolCallError(tool: tool, underlyingError: error)
}
}
}

return .invocations(results)
return try await resolveToolCalls(calls, session: session)
}

// Convert our GenerationSchema into Anthropic's expected JSON Schema payload
Expand Down
113 changes: 6 additions & 107 deletions Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift
Original file line number Diff line number Diff line change
Expand Up @@ -581,119 +581,18 @@ private func createGenerateContentParams<Content: Generable>(
return params
}

private struct ToolInvocationResult {
let call: Transcript.ToolCall
let output: Transcript.ToolOutput
}

private enum ToolResolutionOutcome {
case stop(calls: [Transcript.ToolCall])
case invocations([ToolInvocationResult])
}

private func resolveFunctionCalls(
_ functionCalls: [GeminiFunctionCall],
session: LanguageModelSession
) async throws -> ToolResolutionOutcome {
if functionCalls.isEmpty { return .invocations([]) }

var toolsByName: [String: any Tool] = [:]
for tool in session.tools {
if toolsByName[tool.name] == nil {
toolsByName[tool.name] = tool
}
}

var transcriptCalls: [Transcript.ToolCall] = []
transcriptCalls.reserveCapacity(functionCalls.count)
for call in functionCalls {
let args = try toGeneratedContent(call.args)
let callID = UUID().uuidString
transcriptCalls.append(
Transcript.ToolCall(
id: callID,
toolName: call.name,
arguments: args
)
let calls = try functionCalls.map { call in
Transcript.ToolCall(
id: UUID().uuidString,
toolName: call.name,
arguments: try toGeneratedContent(call.args)
)
}

if let delegate = session.toolExecutionDelegate {
await delegate.didGenerateToolCalls(transcriptCalls, in: session)
}

guard !transcriptCalls.isEmpty else { return .invocations([]) }

var decisions: [ToolExecutionDecision] = []
decisions.reserveCapacity(transcriptCalls.count)

if let delegate = session.toolExecutionDelegate {
for call in transcriptCalls {
let decision = await delegate.toolCallDecision(for: call, in: session)
if case .stop = decision {
return .stop(calls: transcriptCalls)
}
decisions.append(decision)
}
} else {
decisions = Array(repeating: .execute, count: transcriptCalls.count)
}

var results: [ToolInvocationResult] = []
results.reserveCapacity(transcriptCalls.count)

for (index, call) in transcriptCalls.enumerated() {
switch decisions[index] {
case .stop:
// This branch should be unreachable because `.stop` returns during decision collection.
// Keep it as a defensive guard in case that logic changes.
return .stop(calls: transcriptCalls)
case .provideOutput(let segments):
let output = Transcript.ToolOutput(
id: call.id,
toolName: call.toolName,
segments: segments
)
if let delegate = session.toolExecutionDelegate {
await delegate.didExecuteToolCall(call, output: output, in: session)
}
results.append(ToolInvocationResult(call: call, output: output))
case .execute:
guard let tool = toolsByName[call.toolName] else {
let message = Transcript.Segment.text(.init(content: "Tool not found: \(call.toolName)"))
let output = Transcript.ToolOutput(
id: call.id,
toolName: call.toolName,
segments: [message]
)
if let delegate = session.toolExecutionDelegate {
await delegate.didExecuteToolCall(call, output: output, in: session)
}
results.append(ToolInvocationResult(call: call, output: output))
continue
}

do {
let segments = try await tool.makeOutputSegments(from: call.arguments)
let output = Transcript.ToolOutput(
id: call.id,
toolName: tool.name,
segments: segments
)
if let delegate = session.toolExecutionDelegate {
await delegate.didExecuteToolCall(call, output: output, in: session)
}
results.append(ToolInvocationResult(call: call, output: output))
} catch {
if let delegate = session.toolExecutionDelegate {
await delegate.didFailToolCall(call, error: error, in: session)
}
throw LanguageModelSession.ToolCallError(tool: tool, underlyingError: error)
}
}
}

return .invocations(results)
return try await resolveToolCalls(calls, session: session)
}

private func emptyResponseContent<Content: Generable>(
Expand Down
103 changes: 4 additions & 99 deletions Sources/AnyLanguageModel/Models/MLXLanguageModel.swift
Original file line number Diff line number Diff line change
Expand Up @@ -1002,7 +1002,7 @@ import Foundation
}
previousToolCallSignature = signature

let resolution = try await resolveToolCalls(collectedToolCalls, session: session)
let resolution = try await resolveMLXToolCalls(collectedToolCalls, session: session)
switch resolution {
case .stop(let calls):
if !calls.isEmpty {
Expand Down Expand Up @@ -1438,16 +1438,6 @@ import Foundation

// MARK: - Tool Invocation Handling

private struct ToolInvocationResult {
let call: Transcript.ToolCall
let output: Transcript.ToolOutput
}

private enum ToolResolutionOutcome {
case stop(calls: [Transcript.ToolCall])
case invocations([ToolInvocationResult])
}

private func makeTranscriptToolCalls(
from toolCalls: [MLXLMCommon.ToolCall]
) throws -> [Transcript.ToolCall] {
Expand All @@ -1467,97 +1457,12 @@ import Foundation
return transcriptCalls
}

private func resolveToolCalls(
private func resolveMLXToolCalls(
_ toolCalls: [MLXLMCommon.ToolCall],
session: LanguageModelSession
) async throws -> ToolResolutionOutcome {
if toolCalls.isEmpty { return .invocations([]) }

var toolsByName: [String: any Tool] = [:]
for tool in session.tools {
if toolsByName[tool.name] == nil {
toolsByName[tool.name] = tool
}
}

let transcriptCalls = try makeTranscriptToolCalls(from: toolCalls)

if let delegate = session.toolExecutionDelegate {
await delegate.didGenerateToolCalls(transcriptCalls, in: session)
}

guard !transcriptCalls.isEmpty else { return .invocations([]) }

var decisions: [ToolExecutionDecision] = []
decisions.reserveCapacity(transcriptCalls.count)

if let delegate = session.toolExecutionDelegate {
for call in transcriptCalls {
let decision = await delegate.toolCallDecision(for: call, in: session)
if case .stop = decision {
return .stop(calls: transcriptCalls)
}
decisions.append(decision)
}
} else {
decisions = Array(repeating: .execute, count: transcriptCalls.count)
}

var results: [ToolInvocationResult] = []
results.reserveCapacity(transcriptCalls.count)

for (index, call) in transcriptCalls.enumerated() {
switch decisions[index] {
case .stop:
// This branch should be unreachable because `.stop` returns during decision collection.
// Keep it as a defensive guard in case that logic changes.
return .stop(calls: transcriptCalls)
case .provideOutput(let segments):
let output = Transcript.ToolOutput(
id: call.id,
toolName: call.toolName,
segments: segments
)
if let delegate = session.toolExecutionDelegate {
await delegate.didExecuteToolCall(call, output: output, in: session)
}
results.append(ToolInvocationResult(call: call, output: output))
case .execute:
guard let tool = toolsByName[call.toolName] else {
let message = Transcript.Segment.text(.init(content: "Tool not found: \(call.toolName)"))
let output = Transcript.ToolOutput(
id: call.id,
toolName: call.toolName,
segments: [message]
)
if let delegate = session.toolExecutionDelegate {
await delegate.didExecuteToolCall(call, output: output, in: session)
}
results.append(ToolInvocationResult(call: call, output: output))
continue
}

do {
let segments = try await tool.makeOutputSegments(from: call.arguments)
let output = Transcript.ToolOutput(
id: call.id,
toolName: tool.name,
segments: segments
)
if let delegate = session.toolExecutionDelegate {
await delegate.didExecuteToolCall(call, output: output, in: session)
}
results.append(ToolInvocationResult(call: call, output: output))
} catch {
if let delegate = session.toolExecutionDelegate {
await delegate.didFailToolCall(call, error: error, in: session)
}
throw LanguageModelSession.ToolCallError(tool: tool, underlyingError: error)
}
}
}

return .invocations(results)
let calls = try makeTranscriptToolCalls(from: toolCalls)
return try await AnyLanguageModel.resolveToolCalls(calls, session: session)
}

private func toGeneratedContent(_ args: [String: MLXLMCommon.JSONValue]) throws -> GeneratedContent {
Expand Down
Loading