diff --git a/Sources/SwiftAgentKit/Memory/Conversation.swift b/Sources/SwiftAgentKit/Memory/Conversation.swift index a831981..8d23e35 100644 --- a/Sources/SwiftAgentKit/Memory/Conversation.swift +++ b/Sources/SwiftAgentKit/Memory/Conversation.swift @@ -136,6 +136,16 @@ public class Conversation: @unchecked Sendable { messages.removeAll() } + /// Apply a pure transform to every stored NON-system message, in place. + /// For history-externalization flows (e.g. swapping image bytes for a + /// textual description). System messages pass through untouched; count + /// and order are preserved. + public func rewriteMessages(_ transform: (AgentMessage) -> AgentMessage) { + lock.lock() + defer { lock.unlock() } + messages = messages.map { $0.role == .system ? $0 : transform($0) } + } + /// Replace the system message(s) with a new one. public func setSystemMessage(_ message: AgentMessage) { lock.lock() diff --git a/Tests/SwiftAgentKitTests/ConversationRewriteTests.swift b/Tests/SwiftAgentKitTests/ConversationRewriteTests.swift new file mode 100644 index 0000000..f918c19 --- /dev/null +++ b/Tests/SwiftAgentKitTests/ConversationRewriteTests.swift @@ -0,0 +1,34 @@ +// +// ConversationRewriteTests.swift +// SwiftAgentKit +// +// Tests for Conversation.rewriteMessages — verifies that the transform +// is applied to every non-system message while system messages pass through +// untouched, and that count + order are preserved. +// + +import Testing +import Foundation +import LLMProviderKit +@testable import SwiftAgentKit + +@Test func rewriteMessagesTransformsNonSystemOnly() { + let convo = Conversation(contextWindow: 8192) + convo.setSystemMessage(AgentMessage(role: .system, content: "SYS")) + convo.append(AgentMessage(role: .user, content: "hello", + images: [LLMImage(data: Data([1, 2, 3]))])) + convo.append(AgentMessage(role: .assistant, content: "hi")) + + convo.rewriteMessages { message in + var m = message + m.images = [] + m.content += " [x]" + return m + } + + let all = convo.allMessages() + #expect(all.count == 3) + #expect(all[0].role == .system && all[0].content == "SYS") // untouched + #expect(all[1].content == "hello [x]" && all[1].images.isEmpty) // transformed + #expect(all[2].content == "hi [x]") +}