diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index 50ea2ec3..6996b187 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -194,9 +194,25 @@ public final class LanguageModelSession: @unchecked Sendable { options: options ) - // Add response entries to transcript + // Add response entry to transcript + let textContent: String + if case .string(let str) = response.rawContent.kind { + textContent = str + } else { + textContent = response.rawContent.jsonString + } + + let responseEntry = Transcript.Entry.response( + Transcript.Response( + assetIDs: [], + segments: [.text(.init(content: textContent))] + ) + ) + + // Add tool entries and response to transcript await MainActor.run { self.transcript.append(contentsOf: response.transcriptEntries) + self.transcript.append(responseEntry) } return response @@ -525,9 +541,25 @@ extension LanguageModelSession { options: options ) - // Add response entries to transcript + // Add response entry to transcript + let textContent: String + if case .string(let str) = response.rawContent.kind { + textContent = str + } else { + textContent = response.rawContent.jsonString + } + + let responseEntry = Transcript.Entry.response( + Transcript.Response( + assetIDs: [], + segments: [.text(.init(content: textContent))] + ) + ) + + // Add tool entries and response to transcript await MainActor.run { self.transcript.append(contentsOf: response.transcriptEntries) + self.transcript.append(responseEntry) } return response diff --git a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift index 5ec4ac82..a5c162e9 100644 --- a/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/MockLanguageModelTests.swift @@ -13,9 +13,8 @@ struct MockLanguageModelTests { let response = try await session.respond(to: "Say hello") #expect(response.content == "Hello, World!") - // Verify transcript was updated + // Verify transcript was updated (prompt + response) #expect(session.transcript.count == 2) - #expect(response.transcriptEntries.count > 0) } @Test func echoResponse() async throws { @@ -99,7 +98,6 @@ struct MockLanguageModelTests { try await Task.sleep(for: .milliseconds(10)) #expect(asyncSession.isResponding == false) #expect(asyncSession.transcript.count == 2) - #expect(response.transcriptEntries.count > 0) // Test streaming response with isResponding state let streamModel = MockLanguageModel.streamingMock() diff --git a/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift b/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift index 3c96cf45..bcb2bd6f 100644 --- a/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift +++ b/Tests/AnyLanguageModelTests/Shared/MockLanguageModel.swift @@ -36,17 +36,10 @@ struct MockLanguageModel: LanguageModel { let promptWithInstructions = Prompt("Instructions: \(session.instructions?.description ?? "N/A")\n\(prompt)") let text = try await responseProvider(promptWithInstructions, options) - let responseEntry = Transcript.Entry.response( - Transcript.Response( - assetIDs: [], - segments: [.text(.init(content: text))] - ) - ) - return LanguageModelSession.Response( content: text as! Content, rawContent: GeneratedContent(text), - transcriptEntries: [responseEntry] + transcriptEntries: [] ) }