Conversation.rewriteMessages

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()

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]")
+}