Files
Termind/Frameworks/TermindApplication/FeatureModels/MemoryModel.swift
T

632 lines
29 KiB
Swift

import AgentKit
import Foundation
@Observable @MainActor
public final class MemoryModel {
public private(set) var isEnabled: Bool
public private(set) var isBusy = false
public private(set) var isIndexing = false
public private(set) var state = AgentMemoryState()
public private(set) var referenceDates: [UUID: Date] = [:]
public private(set) var error: String?
public var reviewEnabled: Bool {
didSet {
defaults?.set(reviewEnabled, forKey: Self.reviewKey)
if !reviewEnabled {
for task in reviews.values { task.cancel() }
for task in corrections.values { task.cancel() }
for task in correctionRetries.values { task.cancel() }
pendingReviews.removeAll()
pendingCorrections.removeAll()
queuedCorrectionIDs.removeAll()
}
}
}
public var reviewModel: AIModelRef? {
didSet { defaults?.set(try? JSONEncoder().encode(reviewModel), forKey: Self.modelKey) }
}
private static let enabledKey = "hermesMemoryEnabled"
private static let destroyingKey = "hermesMemoryDestroyPending"
private static let reviewKey = "hermesMemoryReviewEnabled"
private static let modelKey = "hermesMemoryReviewModel"
private let defaults: UserDefaults?
private let storage: (any MemoryPersisting)?
private var generation = UUID()
private var contexts: [UUID: AgentMemoryContext] = [:]
private var history: (any AgentRunPersisting)?
private var indexTask: Task<Void, Never>?
private var indexDirty = false
private var reviews: [UUID: Task<Void, Never>] = [:]
private var pendingReviews: [UUID: (@Sendable () async -> Void)] = [:]
private var corrections: [UUID: Task<Void, Never>] = [:]
private var pendingCorrections:
[UUID: [(AgentTranscriptMessage, [AgentTranscriptMessage], any AgentModelCompleting)]] = [:]
private var queuedCorrectionIDs: Set<UUID> = []
private var correctionRetries: [UUID: Task<Void, Never>] = [:]
private var correctionRetryConversations: [UUID: UUID] = [:]
private var correctionFailures: [UUID: Int] = [:]
private var correctionOrders: [UUID: [UUID: Int]] = [:]
private var nextCorrectionOrder: [UUID: Int] = [:]
private var latestSuccessfulCorrectionOrder: [UUID: Int] = [:]
private var retryAfter: [UUID: Date] = [:]
private var learningEpochs: [UUID: UUID] = [:]
private var isSuspended = false
public func learningEpoch(for conversationID: UUID) -> UUID {
if let existing = learningEpochs[conversationID] { return existing }
let fresh = UUID()
learningEpochs[conversationID] = fresh
return fresh
}
public func invalidateLearning(for conversationID: UUID) {
learningEpochs[conversationID] = UUID()
reviews[conversationID]?.cancel()
corrections[conversationID]?.cancel()
pendingReviews.removeValue(forKey: conversationID)
pendingCorrections.removeValue(forKey: conversationID)
correctionOrders.removeValue(forKey: conversationID)
nextCorrectionOrder.removeValue(forKey: conversationID)
latestSuccessfulCorrectionOrder.removeValue(forKey: conversationID)
retryAfter.removeValue(forKey: conversationID)
for id in correctionRetryConversations.filter({ $0.value == conversationID }).keys {
correctionRetries.removeValue(forKey: id)?.cancel()
correctionRetryConversations.removeValue(forKey: id)
}
}
public func learningToken(for context: AgentMemoryContext) -> UUID? {
guard isEnabled, !isSuspended, !isBusy,
contexts.values.contains(where: { $0 === context })
else { return nil }
return generation
}
public init(storage: (any MemoryPersisting)? = nil, defaults: UserDefaults? = nil) {
self.storage = storage
self.defaults = defaults
let pending = defaults?.bool(forKey: Self.destroyingKey) == true
isEnabled = !pending && storage != nil && defaults?.bool(forKey: Self.enabledKey) == true
reviewEnabled = defaults?.object(forKey: Self.reviewKey) as? Bool ?? true
reviewModel = defaults?.data(forKey: Self.modelKey).flatMap {
try? JSONDecoder().decode(AIModelRef.self, from: $0)
}
if pending { Task { await destroy() } }
}
public func attachHistory(_ history: any AgentRunPersisting) {
self.history = history
historyChanged()
}
private func invalidate() {
generation = UUID()
for context in contexts.values { context.invalidate() }
contexts.removeAll()
indexTask?.cancel()
for task in reviews.values { task.cancel() }
for task in corrections.values { task.cancel() }
for task in correctionRetries.values { task.cancel() }
pendingReviews.removeAll()
pendingCorrections.removeAll()
queuedCorrectionIDs.removeAll()
correctionRetries.removeAll()
correctionRetryConversations.removeAll()
correctionFailures.removeAll()
correctionOrders.removeAll()
nextCorrectionOrder.removeAll()
latestSuccessfulCorrectionOrder.removeAll()
retryAfter.removeAll()
learningEpochs.removeAll()
}
public func setEnabled(_ enabled: Bool) async {
guard !isBusy, let storage, defaults?.bool(forKey: Self.destroyingKey) != true else { return }
isBusy = true
defer { isBusy = false }
invalidate()
isEnabled = enabled
defaults?.set(enabled, forKey: Self.enabledKey)
error = nil
do {
if enabled {
let value = try await storage.memoryState()
let dates = try await storage.memoryReferenceDates()
state = value
referenceDates = dates
historyChanged()
} else {
try await storage.close()
}
} catch { self.error = error.localizedDescription }
}
public func refresh() async {
guard isEnabled, !isSuspended, let storage else { return }
let token = generation
do {
let value = try await storage.memoryState()
let dates = try await storage.memoryReferenceDates()
if token == generation {
state = value
referenceDates = dates
error = nil
}
} catch { if token == generation { self.error = error.localizedDescription } }
}
public func edit(_ target: AgentMemoryTarget, text: String) async {
guard isEnabled, !isSuspended, let storage else { return }
let token = generation
do {
let updated = try await storage.replaceTarget(target, text: text)
try publishCommittedState(updated, token: token)
referenceDates = try await require(token).memoryReferenceDates()
error = nil
} catch { if token == generation { self.error = error.localizedDescription } }
}
public func session(_ id: UUID) async -> (context: AgentMemoryContext, access: any AgentMemoryAccessing)? {
guard isEnabled, !isSuspended, let storage else { return nil }
let token = generation
do {
// Open and migrate the store before exposing tools to the run.
try await storage.prepareSession(id)
guard token == generation, isEnabled else { return nil }
let context = contexts[id] ?? AgentMemoryContext()
contexts[id] = context
return (context, MemoryAccess(model: self, token: token))
} catch {
if token == generation { self.error = error.localizedDescription }
return nil
}
}
private func require(_ token: UUID) throws -> any MemoryPersisting {
try Task.checkCancellation()
guard isEnabled, !isSuspended, token == generation, !isBusy, let storage else { throw CancellationError() }
return storage
}
private func publishCommittedState(_ updated: AgentMemoryState, token: UUID) throws {
do {
_ = try require(token)
} catch {
reloadStateForCurrentGeneration()
throw error
}
state = updated
}
private func reloadStateForCurrentGeneration() {
guard isEnabled, !isSuspended, !isBusy, let storage else { return }
let token = generation
Task { [weak self] in
guard let self else { return }
do {
let value = try await storage.memoryState()
let dates = try await storage.memoryReferenceDates()
guard token == generation, isEnabled, !isSuspended, !isBusy else { return }
state = value
referenceDates = dates
error = nil
} catch {
guard token == generation, isEnabled, !isSuspended, !isBusy else { return }
self.error = error.localizedDescription
}
}
}
fileprivate func read(token: UUID) async throws -> AgentMemoryState {
let result = try await require(token).memoryState()
_ = try require(token)
return result
}
@discardableResult
fileprivate func apply(_ operations: [AgentMemoryOperation], token: UUID) async throws -> AgentMemoryState {
let result = try await require(token).applyMemory(operations)
try publishCommittedState(result, token: token)
return result
}
fileprivate func search(_ request: AgentSessionSearchRequest, token: UUID) async throws -> AgentJSONValue {
_ = try require(token)
historyChanged()
await indexTask?.value
let result = try await require(token).searchSessions(request)
_ = try require(token)
return result
}
fileprivate func searchMemory(_ request: AgentMemorySearchRequest, token: UUID) async throws -> AgentJSONValue {
let result = try await require(token).searchMemories(request)
let ids =
result.objectValue?["entries"]?.arrayValue?.compactMap { entry in
entry.objectValue?["id"]?.stringValue.flatMap(UUID.init(uuidString:))
} ?? []
if !ids.isEmpty {
try await require(token).recordMemoryReferences(ids, at: .now)
let dates = try await require(token).memoryReferenceDates()
_ = try require(token)
referenceDates = dates
} else {
_ = try require(token)
}
return result
}
public func rebuildHistoryIndex() async {
guard isEnabled, !isSuspended else { return }
let token = generation
historyChanged()
await indexTask?.value
do {
try await require(token).rebuildHistoryIndex()
if token == generation { error = nil }
} catch { if token == generation { self.error = error.localizedDescription } }
}
public func historyChanged() {
guard isEnabled, !isSuspended, let storage, let history else { return }
indexDirty = true
guard indexTask == nil else { return }
let token = generation
isIndexing = true
indexTask = Task { [weak self] in
guard let self else { return }
defer {
indexTask = nil
isIndexing = false
if indexDirty, isEnabled, !isSuspended { historyChanged() }
}
do {
while indexDirty, token == generation, isEnabled, !Task.isCancelled {
indexDirty = false
let conversations = try await history.loadConversations()
_ = try require(token)
try await storage.replaceHistory(conversations)
}
} catch is CancellationError {
} catch {
if token == generation {
self.error = error.localizedDescription
indexDirty = false
}
}
}
}
public func recordActivity(
conversationID: UUID, userMessageIDs: [UUID], toolInvocations: [MemoryLearningEvent],
messages: [AgentTranscriptMessage], model: (any AgentModelCompleting)?,
sessionToken: UUID? = nil, historyEpoch: UUID? = nil
) async {
guard isEnabled, !isSuspended, reviewEnabled, let storage,
sessionToken == nil || sessionToken == generation,
historyEpoch == nil || historyEpoch == learningEpoch(for: conversationID)
else { return }
let token = generation
let epoch = learningEpoch(for: conversationID)
do {
var progress = try await storage.learningProgress(conversationID)
guard epoch == learningEpoch(for: conversationID) else { return }
for id in userMessageIDs {
_ = try require(token)
guard epoch == learningEpoch(for: conversationID) else { return }
progress = try await storage.recordActivity(
id, sourceMessageID: id, kind: .turn, conversationID: conversationID)
guard epoch == learningEpoch(for: conversationID) else { return }
}
for event in toolInvocations {
_ = try require(token)
guard epoch == learningEpoch(for: conversationID) else { return }
progress = try await storage.recordActivity(
event.id, sourceMessageID: event.sourceMessageID, kind: .tool,
conversationID: conversationID)
guard epoch == learningEpoch(for: conversationID) else { return }
}
guard let model, progress.turns >= 10 || progress.tools >= 15 else { return }
guard retryAfter[conversationID].map({ $0 <= .now }) ?? true else { return }
await review(
conversationID: conversationID, messages: messages, model: model,
force: false, expectedEpoch: epoch)
} catch is CancellationError {
} catch { if token == generation { self.error = error.localizedDescription } }
}
public nonisolated static func looksLikeCorrection(_ text: String) -> Bool {
let value = text.trimmingCharacters(in: .whitespacesAndNewlines).lowercased()
guard !value.isEmpty, value.count <= 4_000 else { return false }
if ["没关系", "不是问题", "no problem", "no worries", "actually looks great"].contains(where: value.contains) {
return false
}
let strong = [
"纠正", "更正", "不对", "别再", "以后请", "以后不要", "记住", "我说错了",
"correction", "i meant", "not what i", "from now on", "remember that",
]
if strong.contains(where: value.contains) { return true }
if value.hasPrefix("no,") || value.hasPrefix("no ") || value.hasPrefix("不,") { return true }
if value.contains("不要") && (value.contains("用") || value.contains("改")) { return true }
if value.contains("不是") && (value.contains("而是") || value.contains("应该")) { return true }
if value.contains("actually") && (value.contains("instead") || value.contains("use ")) { return true }
return false
}
public func considerCorrection(
conversationID: UUID, message: AgentTranscriptMessage, context: [AgentTranscriptMessage],
model: any AgentModelCompleting, sessionToken: UUID? = nil, historyEpoch: UUID? = nil
) {
let authored = message.authoredText ?? ContextChip.visiblePrompt(in: message.text)
guard isEnabled, !isSuspended, reviewEnabled,
sessionToken == nil || sessionToken == generation,
historyEpoch == nil || historyEpoch == learningEpoch(for: conversationID),
Self.looksLikeCorrection(authored), !queuedCorrectionIDs.contains(message.id),
(correctionOrders[conversationID]?[message.id] ?? .max)
> (latestSuccessfulCorrectionOrder[conversationID] ?? 0)
else { return }
if correctionOrders[conversationID]?[message.id] == nil {
let order = nextCorrectionOrder[conversationID, default: 0] + 1
nextCorrectionOrder[conversationID] = order
correctionOrders[conversationID, default: [:]][message.id] = order
}
queuedCorrectionIDs.insert(message.id)
var authoredMessage = message
authoredMessage.authoredText = authored
pendingCorrections[conversationID, default: []].append((authoredMessage, context, model))
reviews[conversationID]?.cancel()
startNextCorrection(conversationID)
}
private func startNextCorrection(_ conversationID: UUID) {
guard corrections[conversationID] == nil,
let next = pendingCorrections[conversationID]?.first
else { return }
pendingCorrections[conversationID]?.removeFirst()
let token = generation
let epoch = learningEpoch(for: conversationID)
let order = correctionOrders[conversationID]?[next.0.id] ?? 0
let task = Task { [weak self] in
guard let self else { return }
defer {
queuedCorrectionIDs.remove(next.0.id)
corrections.removeValue(forKey: conversationID)
startNextCorrection(conversationID)
if pendingCorrections[conversationID]?.isEmpty != false,
let pending = pendingReviews.removeValue(forKey: conversationID)
{
Task { await pending() }
}
}
do {
guard order > (latestSuccessfulCorrectionOrder[conversationID] ?? 0) else { return }
let storage = try require(token)
guard try await !storage.correctionWasReviewed(next.0.id, conversationID: conversationID) else {
return
}
if let running = reviews[conversationID] { await running.value }
_ = try require(token)
guard epoch == learningEpoch(for: conversationID) else { return }
for _ in 0..<2 {
let current = try await read(token: token)
let operations = try await timedOperations {
try await AgentMemoryLearningService().correction(
message: next.0, context: next.1, state: current, model: next.2)
}
_ = try require(token)
guard epoch == learningEpoch(for: conversationID) else { return }
if let updated = try await storage.applyCorrection(
operations, expected: current, messageID: next.0.id, conversationID: conversationID)
{
try publishCommittedState(updated, token: token)
latestSuccessfulCorrectionOrder[conversationID] = order
for id in correctionRetryConversations.filter({ $0.value == conversationID }).keys {
if let retryOrder = correctionOrders[conversationID]?[id], retryOrder < order {
correctionRetries.removeValue(forKey: id)?.cancel()
correctionRetryConversations.removeValue(forKey: id)
}
}
correctionFailures.removeValue(forKey: next.0.id)
return
}
}
scheduleCorrectionRetry(next, conversationID: conversationID, token: token, epoch: epoch)
} catch is CancellationError {
} catch {
if token == generation {
self.error = error.localizedDescription
scheduleCorrectionRetry(next, conversationID: conversationID, token: token, epoch: epoch)
}
}
}
corrections[conversationID] = task
}
private func scheduleCorrectionRetry(
_ correction: (AgentTranscriptMessage, [AgentTranscriptMessage], any AgentModelCompleting),
conversationID: UUID, token: UUID, epoch: UUID
) {
guard epoch == learningEpoch(for: conversationID),
let order = correctionOrders[conversationID]?[correction.0.id],
order > (latestSuccessfulCorrectionOrder[conversationID] ?? 0)
else { return }
let failures = correctionFailures[correction.0.id, default: 0] + 1
correctionFailures[correction.0.id] = failures
let delay = min(300, 15 * (1 << min(failures, 5)))
correctionRetries[correction.0.id]?.cancel()
correctionRetryConversations[correction.0.id] = conversationID
correctionRetries[correction.0.id] = Task { [weak self] in
try? await Task.sleep(for: .seconds(delay))
guard let self, !Task.isCancelled, token == generation,
epoch == learningEpoch(for: conversationID),
order > (latestSuccessfulCorrectionOrder[conversationID] ?? 0)
else { return }
correctionRetries.removeValue(forKey: correction.0.id)
correctionRetryConversations.removeValue(forKey: correction.0.id)
considerCorrection(
conversationID: conversationID, message: correction.0,
context: correction.1, model: correction.2, historyEpoch: epoch)
}
}
private func timedOperations(
_ work: @escaping @Sendable () async throws -> [AgentMemoryOperation]
) async throws -> [AgentMemoryOperation] {
try await withThrowingTaskGroup(of: [AgentMemoryOperation].self) { group in
group.addTask { try await work() }
group.addTask {
try await Task.sleep(for: .seconds(60))
throw MemoryReviewTimeout()
}
defer { group.cancelAll() }
return try await group.next() ?? []
}
}
public func review(
conversationID: UUID, messages: [AgentTranscriptMessage], model: any AgentModelCompleting,
wait: Bool = false, force: Bool = true, expectedEpoch: UUID? = nil
) async {
guard isEnabled, !isSuspended, reviewEnabled, let storage, !messages.isEmpty else { return }
let token = generation
let epoch = learningEpoch(for: conversationID)
guard expectedEpoch == nil || expectedEpoch == epoch else { return }
if let correction = corrections[conversationID] {
if wait {
await correction.value
} else {
pendingReviews[conversationID] = { [weak self] in
await self?.review(
conversationID: conversationID, messages: messages, model: model,
force: force, expectedEpoch: epoch)
}
return
}
}
if let running = reviews[conversationID] {
if wait {
await running.value
} else {
pendingReviews[conversationID] = { [weak self] in
await self?.review(
conversationID: conversationID, messages: messages, model: model,
force: force, expectedEpoch: epoch)
}
return
}
}
guard token == generation, epoch == learningEpoch(for: conversationID), isEnabled else { return }
let task = Task { [weak self] in
guard let self else { return }
defer {
reviews.removeValue(forKey: conversationID)
if corrections[conversationID] == nil,
let next = pendingReviews.removeValue(forKey: conversationID)
{
Task { await next() }
}
}
do {
var progress = try await storage.learningProgress(conversationID)
guard epoch == learningEpoch(for: conversationID) else { return }
guard force || progress.turns >= 10 || progress.tools >= 15 else { return }
guard force || (retryAfter[conversationID].map({ $0 <= .now }) ?? true) else { return }
while true {
let cursorIndex = progress.lastReviewedMessage.flatMap { cursor in
messages.firstIndex(where: { $0.id == cursor })
}
let pendingIndex = progress.events.compactMap { event in
messages.firstIndex(where: { $0.id == event.sourceMessageID })
}.min()
let start = min(cursorIndex.map { $0 + 1 } ?? 0, pendingIndex ?? messages.count)
guard start < messages.count else { return }
var end = start
while end < messages.count, end - start < 100 {
let candidate = end + 1
let serialized = AgentCompaction.serialized(
Array(messages[start..<candidate]), toolResultLimit: 300)
if serialized.count > 28_000, end > start { break }
end = candidate
}
let messageIDs = Set(messages[start..<end].map(\.id))
let reviewedIDs = progress.events.filter { messageIDs.contains($0.sourceMessageID) }.map(\.id)
var contextStart = max(0, start - 2)
while contextStart < start {
let withContext = AgentCompaction.serialized(
Array(messages[contextStart..<end]), toolResultLimit: 300)
if withContext.count <= 28_000 { break }
contextStart += 1
}
let delta = Array(messages[contextStart..<end])
let current = try await read(token: token)
let operations = try await timedOperations {
try await AgentMemoryLearningService().review(messages: delta, state: current, model: model)
}
_ = try require(token)
guard epoch == learningEpoch(for: conversationID) else { return }
guard
let updated = try await storage.applyReview(
operations, expected: current, messageID: messages[end - 1].id,
conversationID: conversationID, reviewedActivityIDs: reviewedIDs)
else {
retryAfter[conversationID] = .now.addingTimeInterval(30)
return
}
try publishCommittedState(updated, token: token)
retryAfter.removeValue(forKey: conversationID)
guard end < messages.count else { return }
progress = try await storage.learningProgress(conversationID)
}
} catch is CancellationError {
} catch {
if token == generation {
self.error = error.localizedDescription
retryAfter[conversationID] = .now.addingTimeInterval(30)
}
}
}
reviews[conversationID] = task
if wait { await task.value }
}
public func destroy() async {
guard !isBusy, let storage else { return }
isBusy = true
isEnabled = false
defaults?.set(false, forKey: Self.enabledKey)
defaults?.set(true, forKey: Self.destroyingKey)
invalidate()
indexDirty = false
let pending = Array(reviews.values)
let correctionPending = Array(corrections.values)
let delayed = Array(correctionRetries.values)
await indexTask?.value
for task in pending { await task.value }
for task in correctionPending { await task.value }
for task in delayed { await task.value }
do {
try await storage.destroy()
state = .init()
referenceDates = [:]
error = nil
defaults?.removeObject(forKey: Self.destroyingKey)
} catch { self.error = error.localizedDescription }
isBusy = false
}
public func suspend() {
guard !isSuspended else { return }
isSuspended = true
invalidate()
}
public func resume() {
guard isSuspended else { return }
isSuspended = false
reloadStateForCurrentGeneration()
historyChanged()
}
}
private struct MemoryReviewTimeout: Error {}
private struct MemoryAccess: AgentMemoryAccessing {
weak var model: MemoryModel?
let token: UUID
func memoryState() async throws -> AgentMemoryState {
guard let model else { throw CancellationError() }
return try await model.read(token: token)
}
func applyMemory(_ operations: [AgentMemoryOperation]) async throws -> AgentMemoryState {
guard let model else { throw CancellationError() }
return try await model.apply(operations, token: token)
}
func searchSessions(_ request: AgentSessionSearchRequest) async throws -> AgentJSONValue {
guard let model else { throw CancellationError() }
return try await model.search(request, token: token)
}
func searchMemories(_ request: AgentMemorySearchRequest) async throws -> AgentJSONValue {
guard let model else { throw CancellationError() }
return try await model.searchMemory(request, token: token)
}
}