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
105 changes: 103 additions & 2 deletions Sources/AnyLanguageModel/Models/GeminiLanguageModel.swift
Original file line number Diff line number Diff line change
Expand Up @@ -349,7 +349,8 @@ public struct GeminiLanguageModel: LanguageModel, ProfileRequestIdentifiedLangua
generating: type,
provider: "Gemini",
transcriptEntries: entries,
requests: usage
requests: usage,
providerMetadata: stoppedResponseContentMetadata
)
case .invocations(let invocations):
entries.append(.toolCalls(callsEntry))
Expand Down Expand Up @@ -850,7 +851,7 @@ private func geminiFunctionResponseData(from url: URL) throws -> GeminiInlineDat
extension Transcript {
fileprivate func toGeminiContent() throws -> [GeminiContent] {
var messages = [GeminiContent]()
for item in self {
for item in geminiReplayEntries {
switch item {
case .instructions(let instructions):
messages.append(
Expand Down Expand Up @@ -1090,6 +1091,106 @@ private func textPartMetadata(
return [textPartsMetadataKey: String(decoding: data, as: UTF8.self)]
}

/// The text of the stored text parts in `metadata`, if there are any.
private func storedText(in metadata: [String: String]?) -> String? {
guard let json = metadata?[textPartsMetadataKey],
let textParts = try? JSONDecoder().decode([GeminiTextHistoryPart].self, from: Data(json.utf8))
else { return nil }
return textParts.map(\.part.text).joined()
}

extension Transcript {
/// The entries to send on replay, with each turn's stored text parts resolved.
///
/// A string response shows the caller the text of its turn's tool rounds,
/// and a caller can edit, rebuild, or remove the response but can't see the metadata.
/// So stored parts apply only when the response's segments still contain their text,
/// and replay sends only text that is in the current segments.
/// Otherwise, the segments are the turn's text, without the old signatures,
/// and the rounds send only their calls, which keep their signatures.
/// The same goes for a finished turn without a response.
fileprivate var geminiReplayEntries: [Entry] {
var entries = Array(self)
var removed: Set<Int> = []
for turn in assistantTurns(entries) {
if case .prompt(let prompt) = entries[turn.range.lowerBound], prompt.responseFormat != nil {
resolveStructuredTurn(in: &entries, turn: turn.range)
} else {
resolveStringTurn(in: &entries, turn: turn.range, isLast: turn.isLast, removed: &removed)
}
}
return entries.indices.filter { !removed.contains($0) }.map { entries[$0] }
}

private func resolveStringTurn(
in entries: inout [Entry],
turn: Range<Int>,
isLast: Bool,
removed: inout Set<Int>
) {
let roundText = entries[turn].compactMap { entry -> String? in
if case .toolCalls(let calls) = entry { storedText(in: calls.providerMetadata) } else { nil }
}.joined()
var hasResponse = false
var isCurrent = true
var carried = roundText
for index in turn {
guard case .response(var response) = entries[index] else { continue }
hasResponse = true
let text = response.segments.singleText
if response.providerMetadata == stoppedResponseContentMetadata, text == "" {
removed.insert(index)
} else if let text, text.hasPrefix(carried) {
let rest = String(text.dropFirst(carried.count))
if rest.isEmpty, !carried.isEmpty {
removed.insert(index)
} else if storedText(in: response.providerMetadata) != rest {
response.providerMetadata?[textPartsMetadataKey] = nil
response.segments = [.text(.init(content: rest))]
entries[index] = .response(response)
}
} else {
isCurrent = false
response.providerMetadata?[textPartsMetadataKey] = nil
entries[index] = .response(response)
}
carried = ""
}
guard !(hasResponse ? isCurrent : isLast) else { return }
for index in turn {
guard case .toolCalls(var calls) = entries[index] else { continue }
calls.providerMetadata?[textPartsMetadataKey] = nil
entries[index] = .toolCalls(calls)
}
}

/// A structured response doesn't include its rounds' text, so they keep it.
/// Its stored parts apply only when they are still the response's content.
private func resolveStructuredTurn(in entries: inout [Entry], turn: Range<Int>) {
for index in turn {
guard case .response(var response) = entries[index] else { continue }
let current = response.segments.singleText.flatMap { try? GeneratedContent(json: $0) }
let stored = storedText(in: response.providerMetadata).flatMap { try? GeneratedContent(json: $0) }
if current == nil || current != stored {
response.providerMetadata?[textPartsMetadataKey] = nil
entries[index] = .response(response)
}
}
}
}

extension [Transcript.Segment] {
/// The segments' text, if they are all text.
fileprivate var singleText: String? {
var result = ""
for segment in self {
guard case .text(let text) = segment else { return nil }
result += text.content
}
return result
}
}

private func restoreTextParts(
from metadata: [String: String]?,
around functionCalls: [GeminiPart] = []
Expand Down
104 changes: 62 additions & 42 deletions Sources/AnyLanguageModel/Models/OpenAILanguageModel.swift
Original file line number Diff line number Diff line change
Expand Up @@ -635,7 +635,8 @@ public struct OpenAILanguageModel: LanguageModel, ProfileRequestIdentifiedLangua
generating: type,
provider: "OpenAI",
transcriptEntries: entries,
requests: usage
requests: usage,
providerMetadata: stoppedResponseContentMetadata
)
case .invocations(let invocations):
if !invocations.isEmpty {
Expand Down Expand Up @@ -700,6 +701,8 @@ public struct OpenAILanguageModel: LanguageModel, ProfileRequestIdentifiedLangua
var text = ""
// The text of earlier tool rounds, which string responses include.
var earlierText = ""
// The text of the final round, which replay sends as the response.
var roundText = ""
var lastOutput: [JSONValue]?
var inFlightMessages: [OpenAIMessage] = []

Expand Down Expand Up @@ -762,6 +765,7 @@ public struct OpenAILanguageModel: LanguageModel, ProfileRequestIdentifiedLangua

let toolCalls = extractToolCallsFromOutput(resp.output)
lastOutput = resp.output
roundText = resp.outputText ?? extractTextFromOutput(resp.output) ?? ""
if !toolCalls.isEmpty {
if let output = resp.output {
for msg in output {
Expand All @@ -773,6 +777,7 @@ public struct OpenAILanguageModel: LanguageModel, ProfileRequestIdentifiedLangua
toolCalls,
requestContext: requestContext,
currentRoundEntries: entries,
providerMetadata: assistantContentMetadata(roundText),
session: session
)
switch resolution {
Expand All @@ -784,7 +789,8 @@ public struct OpenAILanguageModel: LanguageModel, ProfileRequestIdentifiedLangua
generating: type,
provider: "OpenAI",
transcriptEntries: entries,
requests: usage
requests: usage,
providerMetadata: stoppedResponseContentMetadata
)
case .invocations(let invocations):
if !invocations.isEmpty {
Expand All @@ -800,15 +806,13 @@ public struct OpenAILanguageModel: LanguageModel, ProfileRequestIdentifiedLangua
)
)
}
if type == String.self {
earlierText += resp.outputText ?? extractTextFromOutput(resp.output) ?? ""
}
if type == String.self { earlierText += roundText }
continue
}
}
}

text = earlierText + (resp.outputText ?? extractTextFromOutput(resp.output) ?? "")
text = earlierText + roundText

break
}
Expand All @@ -818,7 +822,8 @@ public struct OpenAILanguageModel: LanguageModel, ProfileRequestIdentifiedLangua
content: text as! Content,
rawContent: GeneratedContent(text),
transcriptEntries: ArraySlice(entries),
requests: usage
requests: usage,
providerMetadata: responseContentMetadata(text, roundText: roundText)
)
}

Expand Down Expand Up @@ -968,7 +973,10 @@ public struct OpenAILanguageModel: LanguageModel, ProfileRequestIdentifiedLangua
switch event {
case .outputTextDelta(let delta):
state.text += delta
if let snapshot = state.snapshot() { continuation.yield(snapshot) }
let metadata = responseContentMetadata(state.responseText, roundText: state.text)
if let snapshot = state.snapshot(providerMetadata: metadata) {
continuation.yield(snapshot)
}
case .completed(let response):
state.usage.merge(response?.usage?.reportedUsage)
// The completed response contains full tool arguments and
Expand All @@ -981,7 +989,10 @@ public struct OpenAILanguageModel: LanguageModel, ProfileRequestIdentifiedLangua
)
}
}
if let snapshot = state.snapshot() { continuation.yield(snapshot) }
let metadata = responseContentMetadata(state.responseText, roundText: state.text)
if let snapshot = state.snapshot(providerMetadata: metadata) {
continuation.yield(snapshot)
}
break responseEvents
case .failed(let failure):
throw OpenAILanguageModelError.streamFailed(
Expand Down Expand Up @@ -1048,14 +1059,18 @@ public struct OpenAILanguageModel: LanguageModel, ProfileRequestIdentifiedLangua
toolCalls,
requestContext: requestContext,
currentRoundEntries: state.entries,
providerMetadata: apiVariant == .chatCompletions
? assistantContentMetadata(state.text) : nil,
providerMetadata: assistantContentMetadata(state.text),
session: session
)
switch resolution {
case .stop:
state.entries.append(.toolCalls(callsEntry))
continuation.yield(try state.stoppedSnapshot(provider: "OpenAI"))
continuation.yield(
try state.stoppedSnapshot(
provider: "OpenAI",
providerMetadata: responseContentMetadata(state.responseText, roundText: "")
)
)
continuation.finish()
return
case .invocations(let invocations):
Expand All @@ -1082,9 +1097,7 @@ public struct OpenAILanguageModel: LanguageModel, ProfileRequestIdentifiedLangua
}
// This round's text now goes with its tool calls, so a response that
// ends here has no text of its own to replay.
let metadata =
apiVariant == .chatCompletions
? responseContentMetadata(state.responseText, roundText: "") : nil
let metadata = responseContentMetadata(state.responseText, roundText: "")
if let snapshot = state.snapshot(providerMetadata: metadata) { continuation.yield(snapshot) }
state.beginNextRound()
}
Expand Down Expand Up @@ -1362,6 +1375,10 @@ private enum Responses {
messageRole == "assistant",
case .array(let assistantToolCalls) = assistantMessageObject["tool_calls"]
{
// The round's text goes before its calls, as a message of its own.
if case .string(let text) = assistantMessageObject["content"], !text.isEmpty {
outputs.append(assistantInputItem(text: text))
}
for assistantToolCall in assistantToolCalls {
if case .object(let toolCallObject) = assistantToolCall,
case .string(let toolCallID) = toolCallObject["id"],
Expand Down Expand Up @@ -1397,7 +1414,8 @@ private enum Responses {
}

case .assistant:
break
let text = msg.responsesAssistantText
if !text.isEmpty { outputs.append(assistantInputItem(text: text)) }
}
}
body["input"] = .array(outputs)
Expand Down Expand Up @@ -1501,6 +1519,16 @@ private enum Responses {
return .object(body)
}

/// An assistant message as an input item, in the `EasyInputMessage` form.
/// This form takes string content, without the `id` and `status` of an output message.
private static func assistantInputItem(text: String) -> JSONValue {
.object([
"type": .string("message"),
"role": .string("assistant"),
"content": .string(text),
])
}

struct Response: Decodable, Sendable {
let id: String
let output: [JSONValue]?
Expand Down Expand Up @@ -1528,26 +1556,6 @@ private enum Responses {

// MARK: - Supporting Types

/// Key under which the assistant's text for a Chat Completions tool-call round is kept in
/// ``Transcript/ToolCalls/providerMetadata``.
private let assistantContentMetadataKey = "openai.content"

private func assistantContentMetadata(_ text: String?) -> [String: String]? {
guard let text, !text.isEmpty else { return nil }
return [assistantContentMetadataKey: text]
}

/// The text that a Chat Completions string response replays, kept in
/// ``Transcript/Response/providerMetadata`` when it differs from the response's content.
///
/// A string response includes the text of every tool round,
/// but replay already sends each round's text with that round's calls.
/// The response keeps only the final round's text,
/// so replay sends each piece of text once, as the requests in flight did.
private func responseContentMetadata(_ responseText: String, roundText: String) -> [String: String]? {
responseText == roundText ? nil : [assistantContentMetadataKey: roundText]
}

/// The assistant message that carries a round of tool calls.
///
/// Built from the transcript's calls both for the round in flight and when the history is
Expand Down Expand Up @@ -1712,7 +1720,7 @@ extension Transcript {
// A Chat Completions reasoning entry goes back on the assistant message right after
// it, when the request goes to the host that wrote it.
var pendingReasoning: Transcript.Reasoning?
for item in self {
for item in replayEntries {
let reasoning = pendingReasoning
pendingReasoning = nil
switch item {
Expand All @@ -1735,14 +1743,10 @@ extension Transcript {
pendingReasoning = entry
continue
case .response(let response):
let segments =
response.providerMetadata?[assistantContentMetadataKey].map {
[Transcript.Segment.text(.init(content: $0))]
} ?? response.segments
messages.append(
.init(
role: .assistant,
content: .blocks(convertSegmentsToOpenAIBlocks(segments)),
content: .blocks(convertSegmentsToOpenAIBlocks(response.segments)),
reasoning: ChatCompletionsReasoning(reasoning, sentTo: host).fields
)
)
Expand Down Expand Up @@ -1786,6 +1790,22 @@ private struct OpenAIMessage: Hashable, Codable, Sendable {
/// Reasoning fields sent back on an assistant message (``ChatCompletionsReasoning``).
var reasoning: [String: JSONValue] = [:]

/// The message's content as the Responses API's text-only assistant content,
/// with an `<image>` placeholder for each image, as in tool outputs.
var responsesAssistantText: String {
switch content {
case .text(let text):
return text
case .blocks(let blocks):
return blocks.map { block in
switch block {
case .text(let text): return text
case .imageURL: return "<image>"
}
}.joined(separator: "\n")
}
}

func contentAsJsonValue(for apiVariant: OpenAILanguageModel.APIVariant) -> JSONValue {
switch content {
case .text(let text):
Expand Down
Loading
Loading