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
142 changes: 139 additions & 3 deletions Sources/AnyLanguageModel/Models/SystemLanguageModel.swift
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,9 @@
return LanguageModelSession.Response(
content: fmResponse.content as! Content,
rawContent: generatedContent,
transcriptEntries: []
transcriptEntries: ArraySlice(
fmResponse.transcriptEntries.compactMap { toolActivityEntry(from: $0) }
)
)
} else {
// For non-String types, use schema-based generation
Expand All @@ -104,6 +106,10 @@
options: fmOptions
)

let toolEntries = ArraySlice(
fmResponse.transcriptEntries.compactMap { toolActivityEntry(from: $0) }
)

func finalize(content: Content) -> LanguageModelSession.Response<Content> {
let normalizedRaw = content.generatedContent
if let jsonValue = try? JSONValue(normalizedRaw),
Expand All @@ -114,13 +120,13 @@
return LanguageModelSession.Response(
content: placeholder.content,
rawContent: placeholder.rawContent,
transcriptEntries: []
transcriptEntries: toolEntries
)
}
return LanguageModelSession.Response(
content: content,
rawContent: normalizedRaw,
transcriptEntries: []
transcriptEntries: toolEntries
)
}

Expand Down Expand Up @@ -167,9 +173,38 @@
)
)

// Entries the session was seeded with. Anything FoundationModels records beyond this
// point belongs to the turn we are about to stream.
let seededEntryCount = fmSession.transcript.count

let stream: AsyncThrowingStream<LanguageModelSession.ResponseStream<Content>.Snapshot, Error> =
AsyncThrowingStream { continuation in

/// Copies newly recorded FoundationModels tool entries into `session.transcript`.
///
/// FoundationModels runs the tool-calling loop itself, and its
/// `ResponseStream.Snapshot` carries only `content` and `rawContent` — it has no
/// tool channel at all. The session's own transcript is therefore the only place
/// tool activity is observable, so we replay new entries from it as the turn
/// progresses. Replaying in FoundationModels' recorded order is what keeps a
/// `.toolCalls` entry ahead of its matching `.toolOutput`.
///
/// - Note: This runs on snapshot boundaries rather than the instant
/// FoundationModels records a call, because polling the session concurrently
/// would mean sharing a non-`Sendable` `FoundationModels.LanguageModelSession`
/// across tasks. Tool entries therefore land with the first snapshot generated
/// after the tool ran, which is still while the stream is in progress.
func mirrorToolEntries(mirroredEntryCount: inout Int) {
let fmEntries = Array(fmSession.transcript)
guard fmEntries.count > mirroredEntryCount else { return }
for entry in fmEntries[mirroredEntryCount...] {
if let toolEntry = toolActivityEntry(from: entry) {
session.appendTranscriptEntry(toolEntry)
}
}
mirroredEntryCount = fmEntries.count
}

func accumulateText(
_ chunkText: String,
accumulatedText: inout String,
Expand Down Expand Up @@ -197,6 +232,7 @@
fmSession.streamResponse(to: fmPrompt, options: fmOptions)

var accumulatedText = ""
var mirroredEntryCount = seededEntryCount
do {
var lastLength = 0
for try await snapshot in fmStream {
Expand All @@ -213,10 +249,18 @@
lastLength: &lastLength
)

// Surface tool activity before any response text so the transcript
// keeps `.toolCalls` -> `.toolOutput` -> `.response` ordering.
mirrorToolEntries(mirroredEntryCount: &mirroredEntryCount)
if !accumulatedText.isEmpty {
session.growStreamingTranscript(text: accumulatedText)
}

let raw = GeneratedContent(accumulatedText)
let snapshotContent = (accumulatedText as! Content).asPartiallyGenerated()
continuation.yield(.init(content: snapshotContent, rawContent: raw))
}
mirrorToolEntries(mirroredEntryCount: &mirroredEntryCount)
continuation.finish()
} catch {
continuation.finish(throwing: error)
Expand All @@ -233,6 +277,8 @@
options: fmOptions
)

var mirroredEntryCount = seededEntryCount

func processTextFallback() async {
let fmTextStream: FoundationModels.LanguageModelSession.ResponseStream<String> =
fmSession.streamResponse(to: fmPrompt, options: fmOptions)
Expand All @@ -253,6 +299,11 @@
lastLength: &lastLength
)

mirrorToolEntries(mirroredEntryCount: &mirroredEntryCount)
if !accumulatedText.isEmpty {
session.growStreamingTranscript(text: accumulatedText)
}

let jsonString = accumulatedText
if let partialContent = try? partialDecoder.decode(
GeneratedContent.self,
Expand All @@ -271,6 +322,7 @@
.init(content: placeholder.content, rawContent: placeholder.rawContent)
)
}
mirrorToolEntries(mirroredEntryCount: &mirroredEntryCount)
continuation.finish()
} catch {
if !didYield, let placeholder = placeholderPartialContent(for: type) {
Expand All @@ -286,6 +338,12 @@
do {
for try await snapshot in fmStream {
let jsonString = snapshot.content.jsonString

mirrorToolEntries(mirroredEntryCount: &mirroredEntryCount)
if !jsonString.isEmpty, jsonString != "null" {
session.growStreamingTranscript(text: jsonString)
}

let raw =
(try? GeneratedContent(snapshot.content))
?? (try? GeneratedContent(json: jsonString))
Expand Down Expand Up @@ -317,6 +375,7 @@
.init(content: placeholder.content, rawContent: placeholder.rawContent)
)
}
mirrorToolEntries(mirroredEntryCount: &mirroredEntryCount)
continuation.finish()
} catch {
if didYield {
Expand Down Expand Up @@ -442,6 +501,27 @@
}

/// A type-erased wrapper that bridges any `Tool` to `FoundationModels.Tool`.
///
/// - Note: Unlike every other provider, this bridge does not consult
/// `session.toolExecutionDelegate`, and that is a deliberate gap rather than an oversight.
/// FoundationModels owns the tool-calling loop and invokes a tool through
/// `FoundationModels.Tool.call(arguments:)`, which receives the decoded arguments and nothing
/// else — no call identifier, no session, no view of the sibling calls in the same batch.
/// That makes the delegate contract unrepresentable here:
///
/// - `didGenerateToolCalls` is defined over the whole batch of calls the model produced,
/// before any of them run. FoundationModels only ever surfaces a single invocation, and
/// only once it has already decided to make it.
/// - `toolCallDecision`, `didExecuteToolCall`, and `didFailToolCall` all take a
/// `Transcript.ToolCall`, whose `id` is assigned by FoundationModels and is not visible
/// from inside the tool. Synthesizing an id would hand the delegate an identifier that
/// never matches the `.toolCalls` entry that later appears in the transcript.
/// - `.stop` has no expression at all. FoundationModels offers no way to halt a turn from
/// inside a tool; throwing would surface an error to the caller instead of the contract's
/// clean stop with the tool calls recorded.
///
/// Honoring the delegate partially would be worse than not honoring it: a caller returning
/// `.stop` would have its tool executed anyway, silently. So the delegate is not consulted.
@available(macOS 26.0, iOS 26.0, watchOS 26.0, tvOS 26.0, visionOS 26.0, *)
private struct AnyToolWrapper: FoundationModels.Tool {
typealias Arguments = FoundationModels.GeneratedContent
Expand Down Expand Up @@ -762,6 +842,62 @@
}
}

// MARK: - FoundationModels to AnyLanguageModel Conversions

/// Converts a FoundationModels transcript entry into its AnyLanguageModel equivalent if it
/// records tool activity.
///
/// Returns `nil` for every other entry kind. Instructions, prompts, and responses are owned by
/// ``LanguageModelSession`` itself — it appends the prompt before calling the model and the
/// response after — so mirroring them here would duplicate transcript entries.
@available(macOS 26.0, iOS 26.0, watchOS 26.0, tvOS 26.0, visionOS 26.0, *)
private func toolActivityEntry(from entry: FoundationModels.Transcript.Entry) -> Transcript.Entry? {
if case .toolCalls(let fmToolCalls) = entry {
let calls = fmToolCalls.compactMap { call -> Transcript.ToolCall? in
guard let arguments = try? AnyLanguageModel.GeneratedContent(call.arguments) else { return nil }
return Transcript.ToolCall(id: call.id, toolName: call.toolName, arguments: arguments)
}
guard !calls.isEmpty else { return nil }
return .toolCalls(Transcript.ToolCalls(id: fmToolCalls.id, calls))
}

if case .toolOutput(let fmToolOutput) = entry {
return .toolOutput(
Transcript.ToolOutput(
id: fmToolOutput.id,
toolName: fmToolOutput.toolName,
segments: transcriptSegments(from: fmToolOutput.segments)
)
)
}

return nil
}

/// Converts FoundationModels transcript segments into their AnyLanguageModel equivalents.
@available(macOS 26.0, iOS 26.0, watchOS 26.0, tvOS 26.0, visionOS 26.0, *)
private func transcriptSegments(
from segments: [FoundationModels.Transcript.Segment]
) -> [Transcript.Segment] {
segments.compactMap { segment -> Transcript.Segment? in
if case .text(let textSegment) = segment {
return .text(.init(id: textSegment.id, content: textSegment.content))
}
if case .structure(let structuredSegment) = segment,
let content = try? AnyLanguageModel.GeneratedContent(structuredSegment.content)
{
return .structure(
.init(
id: structuredSegment.id,
source: structuredSegment.source,
content: content
)
)
}
return nil
}
}

@available(macOS 26.0, iOS 26.0, watchOS 26.0, tvOS 26.0, visionOS 26.0, *)
extension Array where Element == Transcript.ToolDefinition {
fileprivate func toFoundationModels() -> [FoundationModels.Transcript.ToolDefinition] {
Expand Down
77 changes: 67 additions & 10 deletions Tests/AnyLanguageModelTests/SystemLanguageModelTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -174,18 +174,75 @@ import Testing

let response = try await session.respond(to: "How's the weather in San Francisco?")

#if false // Disabled for now because transcript entries are not converted from FoundationModels for now
var foundToolOutput = false
for case let .toolOutput(toolOutput) in response.transcriptEntries {
#expect(toolOutput.id == "getWeather")
foundToolOutput = true
}
#expect(foundToolOutput)
#endif
var foundToolOutput = false
for case let .toolOutput(toolOutput) in response.transcriptEntries {
#expect(!toolOutput.id.isEmpty, "Expected the tool output to carry the id FoundationModels assigned.")
#expect(toolOutput.toolName == "getWeather", "Expected the output to name the tool that produced it.")
foundToolOutput = true
}
#expect(foundToolOutput, "Expected a tool output among the response's transcript entries.")

let content = response.content
#expect(content.contains("San Francisco"))
#expect(content.contains("72°F"))
#expect(
content.contains("San Francisco"),
"Expected the answer to mention the city the tool was asked about."
)
#expect(content.contains("72°F"), "Expected the tool's result to reach the model's answer.")
}

@available(macOS 26.0, iOS 26.0, watchOS 26.0, tvOS 26.0, visionOS 26.0, *)
@Test func streamWithTools() async throws {
let weatherTool = WeatherTool()
let session = LanguageModelSession(model: SystemLanguageModel.default, 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(!snapshots.isEmpty, "Expected the stream to yield at least one snapshot.")
#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."
)

// The tool call must be recorded before the output it produced.
let toolCallIndex = session.transcript.firstIndex { entry in
if case .toolCalls = entry { return true } else { return false }
}
let toolOutputIndex = session.transcript.firstIndex { entry in
if case .toolOutput = entry { return true } else { return false }
}
#expect(toolCallIndex != nil, "Expected a .toolCalls entry in the final transcript.")
#expect(toolOutputIndex != nil, "Expected a .toolOutput entry in the final transcript.")
if let toolCallIndex, let toolOutputIndex {
#expect(toolCallIndex < toolOutputIndex, "Expected .toolCalls to precede .toolOutput.")
}

var foundToolOutput = false
for case let .toolOutput(toolOutput) in session.transcript {
#expect(!toolOutput.id.isEmpty, "Expected the tool output to carry the id FoundationModels assigned.")
#expect(toolOutput.toolName == "getWeather", "Expected the output to name the tool that produced it.")
foundToolOutput = true
}
#expect(foundToolOutput, "Expected the 'getWeather' tool to exist in the final transcript.")
}

@available(macOS 26.0, iOS 26.0, watchOS 26.0, tvOS 26.0, visionOS 26.0, *)
Expand Down