From dd22f4e708fc93e93ecc054b3368fa89748c3dff Mon Sep 17 00:00:00 2001 From: Zac White Date: Sun, 12 Jul 2026 15:41:31 -0700 Subject: [PATCH 1/5] Add NaturalLanguage entity tagging to the memory write path - New EntityTagger provider protocol on MemoryConfiguration with an NLTagger-backed NLEntityTagger default under the MemoryNaturalLanguage trait (nil otherwise), mirroring the query-side NLQueryAnalyzer. - Tagger results merge additively into heuristic extraction: they upgrade generic entity labels to person/location/organization and append entities the heuristics missed; heuristic behavior remains the floor. Facet inference now sees merged entities, so location/person facets fire for tagger-recognized names beyond the built-in lexicon. - Ambiguous self-reported locations ("i'm in lisbon") are accepted when the tagger confirms the phrase names a place, validated by title-casing inside a carrier sentence where NLTagger place recognition is precise. Travel qualifiers ("i'm in boston this week") mark the mention transient and skip the residence write. Residence statements ("i live in ...") never require tagger confirmation, keeping eval gates OS-stable. - ingest/save entity normalization consults the tagger when candidates supply no entities. --- README.md | 6 +- .../MemoryExtractionHeuristics.swift | 128 +++++++++++++++--- Sources/AgentMemory/MemoryIndex.swift | 13 +- .../NLContextualEmbeddingProvider.swift | 4 +- Sources/AgentMemory/NLEntityTagger.swift | 97 +++++++++++++ Sources/AgentMemory/ProtocolsAndErrors.swift | 25 ++++ .../MemoryTests/MemoryExternalAPITests.swift | 121 +++++++++++++++++ 7 files changed, 373 insertions(+), 21 deletions(-) create mode 100644 Sources/AgentMemory/NLEntityTagger.swift diff --git a/README.md b/README.md index 578401e..95ce340 100644 --- a/README.md +++ b/README.md @@ -108,7 +108,7 @@ Most integrations only need: - `save`, `extract`, `ingest`, and `recall` for agent memory workflows - `capture`, `prepareContext`, `recordSignal`, and `runMaintenance` for higher-level agent memory workflows - `memorySearch` and `memoryGet` for tool-style retrieval -- customization protocols (`EmbeddingProvider`, `Reranker`, `StructuredQueryExpander`, `MemoryExtractor`, `RecallPlanner`) only when you are swapping in your own providers +- customization protocols (`EmbeddingProvider`, `Reranker`, `StructuredQueryExpander`, `MemoryExtractor`, `RecallPlanner`, `EntityTagger`) only when you are swapping in your own providers ## Agent Memory Workflows @@ -131,6 +131,10 @@ let capture = try await index.capture( Captured `MemoryCandidate` and stored `MemoryRecord` values can include a `subject` and original-message `evidence`. The default agent policy focuses on user-authored durable facts, rejects assistant capability/refusal text, and keeps embedded declarations separate from raw questions. +### Linguistic Entity Tagging + +When the `MemoryNaturalLanguage` trait is enabled (the default), extraction runs `NLEntityTagger` — an `NLTagger`-backed named entity recognizer — alongside the heuristics. Tagger results are strictly additive: they upgrade generic entity labels to `person`/`location`/`organization`, add entities the heuristics missed in well-capitalized text, and confirm ambiguous self-reported locations ("i'm in portland") before they become durable profile memories. Travel-qualified phrases ("i'm in boston this week") are treated as transient and not recorded as residence. Heuristic behavior is the floor: with `entityTagger: nil` (or the trait disabled) extraction is unchanged, and residence statements ("i live in ...") never require tagger confirmation. Provide a custom `EntityTagger` through `MemoryConfiguration` to swap in your own recognizer. + Context preparation retrieves bounded memory context for the next model turn: ```swift diff --git a/Sources/AgentMemory/MemoryExtractionHeuristics.swift b/Sources/AgentMemory/MemoryExtractionHeuristics.swift index 11b113e..78a4110 100644 --- a/Sources/AgentMemory/MemoryExtractionHeuristics.swift +++ b/Sources/AgentMemory/MemoryExtractionHeuristics.swift @@ -27,7 +27,8 @@ internal enum MemoryExtractionHeuristics { messages: [ConversationMessage], limit: Int, canonicalKey: CanonicalKeyResolver, - proposedAction: ProposedActionResolver + proposedAction: ProposedActionResolver, + entityTagger: (any EntityTagger)? = nil ) -> MemoryExtractionResult { var extracted: [MemoryCandidate] = [] extracted.reserveCapacity(limit) @@ -41,7 +42,11 @@ internal enum MemoryExtractionHeuristics { .trimmingCharacters(in: .whitespacesAndNewlines) guard !normalized.isEmpty else { continue } - let focusedSegments = focusedUserProfileSegments(from: normalized, role: message.role) + let focusedSegments = focusedUserProfileSegments( + from: normalized, + role: message.role, + entityTagger: entityTagger + ) let rawSegments = focusedSegments.map(\.text) + splitExtractionSegments(normalized) for rawSegment in rawSegments { @@ -75,12 +80,15 @@ internal enum MemoryExtractionHeuristics { let importance = inferredImportance(for: kind) let confidence = inferredConfidence(for: kind) let tags = inferredTags(forExtractedText: segment) - let facetTags = inferFacetTags(forExtractedText: segment, kind: kind) - .union(focusedSegment?.facetTags ?? []) let entities = mergeEntities( pinned: focusedSegment?.entities ?? [], - inferred: inferEntities(forExtractedText: segment) + inferred: mergeTaggerEntities( + entityTagger?.recognizeEntities(in: segment) ?? [], + into: inferEntities(forExtractedText: segment) + ) ) + let facetTags = inferFacetTags(forExtractedText: segment, kind: kind, knownEntities: entities) + .union(focusedSegment?.facetTags ?? []) let topics = inferTopics(forExtractedText: segment, seedTags: tags) let subject = inferSubject(forExtractedText: segment, role: message.role, kind: kind) let evidence = MemoryEvidence( @@ -155,9 +163,13 @@ internal enum MemoryExtractionHeuristics { return tags } - internal static func inferFacetTags(forExtractedText text: String, kind: MemoryKind) -> Set { + internal static func inferFacetTags( + forExtractedText text: String, + kind: MemoryKind, + knownEntities: [MemoryEntity]? = nil + ) -> Set { let normalizedText = phraseEnvelope(for: text) - let knownEntities = inferKnownEntities(forExtractedText: text) + let knownEntities = knownEntities ?? inferKnownEntities(forExtractedText: text) var facets: Set = [] if containsAnyNormalizedPhrase(normalizedText, phrases: ["prefer", "prefers", "preferred", "preference", "favorite", "likes", "dislikes"]) @@ -588,11 +600,15 @@ internal enum MemoryExtractionHeuristics { var facetTags: Set } - private static func focusedUserProfileSegments(from text: String, role: ConversationRole) -> [FocusedProfileSegment] { + private static func focusedUserProfileSegments( + from text: String, + role: ConversationRole, + entityTagger: (any EntityTagger)? = nil + ) -> [FocusedProfileSegment] { guard role == .user else { return [] } var segments: [FocusedProfileSegment] = [] - if let location = selfReportedLocation(in: text) { + if let location = selfReportedLocation(in: text, entityTagger: entityTagger) { let city = location.split(separator: ",").first.map(String.init) ?? location segments.append( FocusedProfileSegment( @@ -624,9 +640,55 @@ internal enum MemoryExtractionHeuristics { return merged } + /// Merges tagger-recognized entities into heuristic results additively: + /// heuristic entities are the floor, tagger results upgrade `.other` + /// labels and append entities the heuristics missed. + internal static func mergeTaggerEntities( + _ tagged: [MemoryEntity], + into inferred: [MemoryEntity] + ) -> [MemoryEntity] { + guard !tagged.isEmpty else { return inferred } + + var merged = inferred + var indexByValue: [String: Int] = [:] + for (index, entity) in merged.enumerated() { + indexByValue[normalizeEntityValue(entity.normalizedValue)] = index + } + + for entity in tagged { + let normalizedValue = normalizeEntityValue(entity.normalizedValue) + guard !normalizedValue.isEmpty, !isGenericExtractedEntity(normalizedValue) else { continue } + if let index = indexByValue[normalizedValue] { + let existing = merged[index] + if existing.label == .other { + merged[index] = MemoryEntity( + label: entity.label, + value: existing.value, + normalizedValue: normalizedValue, + confidence: max(existing.confidence ?? 0, entity.confidence ?? 0) + ) + } + } else if merged.count < 8 { + merged.append( + MemoryEntity( + label: entity.label, + value: entity.value, + normalizedValue: normalizedValue, + confidence: entity.confidence + ) + ) + indexByValue[normalizedValue] = merged.count - 1 + } + } + + return merged + } + // Residence statements ("i live in ...") accept any captured place phrase. - // The ambiguous forms ("i'm in ...") only fire for locations in the alias table, - // since they commonly describe transient states ("i'm in a meeting"). + // The ambiguous forms ("i'm in ...") commonly describe transient states + // ("i'm in a meeting"), so they require confirmation: an alias-table hit, + // or an entity tagger recognizing the phrase as a place — and never with a + // travel qualifier ("i'm in boston this week"). private static let strongLocationTriggerPattern = #"\b(?:i\s+live\s+in|i\s+am\s+living\s+in|i['’]m\s+living\s+in|i\s+am\s+based\s+in|i['’]m\s+based\s+in|i\s+(?:just\s+)?moved\s+to|my\s+city\s+is|my\s+location\s+is)\s+([^.!?;\n]+)"# private static let weakLocationTriggerPattern = @@ -649,26 +711,38 @@ internal enum MemoryExtractionHeuristics { private static let locationStopTokens: Set = [ "and", "but", "or", "so", "which", "where", "when", "while", "because", "since", "near", "with", "for", "though", "although", "what", "what's", - "what’s", "who", "how", "why", "that", "then", "if", "as", "at" + "what’s", "who", "how", "why", "that", "then", "if", "as", "at", "about", + "this", "next", "until" ] private static let locationTrailingNoiseTokens: Set = [ "now", "currently", "atm", "too", "btw", "there", "here", "tonight", "today" ] + // Tokens that mark an ambiguous "i'm in " as a visit rather than + // residence, either trailing the place ("i'm in boston tonight") or + // immediately following it ("i'm in boston this week"). + private static let transientLocationQualifiers: Set = [ + "this", "next", "until", "for", "tonight", "today", "tomorrow", + "visiting", "traveling", "travelling", "temporarily" + ] + private static let locationLeadingRejectTokens: Set = [ "a", "an", "my", "our", "your", "his", "her", "their", "this", "that", "these", "those", "some", "any", "no", "one", "it", "front", "between", "town", "meetings", "meeting" ] - private static func selfReportedLocation(in text: String) -> String? { + private static func selfReportedLocation( + in text: String, + entityTagger: (any EntityTagger)? = nil + ) -> String? { if let raw = firstRegexCapture(of: strongLocationTriggerPattern, in: text), - let location = parseLocationPhrase(raw, requireKnownLocation: false) { + let location = parseLocationPhrase(raw, isAmbiguousTrigger: false, entityTagger: entityTagger) { return location } if let raw = firstRegexCapture(of: weakLocationTriggerPattern, in: text), - let location = parseLocationPhrase(raw, requireKnownLocation: true) { + let location = parseLocationPhrase(raw, isAmbiguousTrigger: true, entityTagger: entityTagger) { return location } return nil @@ -686,10 +760,15 @@ internal enum MemoryExtractionHeuristics { return nsText.substring(with: match.range(at: 1)) } - private static func parseLocationPhrase(_ raw: String, requireKnownLocation: Bool) -> String? { + private static func parseLocationPhrase( + _ raw: String, + isAmbiguousTrigger: Bool, + entityTagger: (any EntityTagger)? = nil + ) -> String? { let tokens = raw.split(whereSeparator: \.isWhitespace).map(String.init) var cityTokens: [String] = [] var sawComma = false + var sawTransientQualifier = false var index = 0 while index < tokens.count, cityTokens.count < 4 { @@ -711,8 +790,20 @@ internal enum MemoryExtractionHeuristics { } while let last = cityTokens.last, locationTrailingNoiseTokens.contains(last.lowercased()) { + if transientLocationQualifiers.contains(last.lowercased()) { + sawTransientQualifier = true + } cityTokens.removeLast() } + if index < tokens.count { + let following = tokens[index].trimmingCharacters(in: CharacterSet.punctuationCharacters).lowercased() + if transientLocationQualifiers.contains(following) { + sawTransientQualifier = true + } + } + if isAmbiguousTrigger, sawTransientQualifier { + return nil + } var regionToken: String? if sawComma, index < tokens.count { @@ -748,7 +839,10 @@ internal enum MemoryExtractionHeuristics { } } - guard !requireKnownLocation, city.count >= 2 else { return nil } + guard city.count >= 2 else { return nil } + if isAmbiguousTrigger { + guard entityTagger?.isLikelyPlaceName(city) == true else { return nil } + } let displayCity = titleCasedLocation(cityTokens) guard let regionToken else { return displayCity } diff --git a/Sources/AgentMemory/MemoryIndex.swift b/Sources/AgentMemory/MemoryIndex.swift index 01a52fd..d9359a9 100644 --- a/Sources/AgentMemory/MemoryIndex.swift +++ b/Sources/AgentMemory/MemoryIndex.swift @@ -1961,7 +1961,8 @@ public actor MemoryIndex { }, proposedAction: { candidate in self.proposedWriteAction(for: candidate) - } + }, + entityTagger: configuration.entityTagger ) } @@ -2077,7 +2078,15 @@ public actor MemoryIndex { } private func normalizeEntities(_ supplied: [MemoryEntity], text: String) -> [MemoryEntity] { - let preferred = supplied.isEmpty ? inferEntities(forExtractedText: text) : supplied + let preferred: [MemoryEntity] + if supplied.isEmpty { + preferred = MemoryExtractionHeuristics.mergeTaggerEntities( + configuration.entityTagger?.recognizeEntities(in: text) ?? [], + into: inferEntities(forExtractedText: text) + ) + } else { + preferred = supplied + } var normalized: [MemoryEntity] = [] var seen: Set = [] normalized.reserveCapacity(min(preferred.count, 8)) diff --git a/Sources/AgentMemory/NLContextualEmbeddingProvider.swift b/Sources/AgentMemory/NLContextualEmbeddingProvider.swift index cf95122..3b1f53f 100644 --- a/Sources/AgentMemory/NLContextualEmbeddingProvider.swift +++ b/Sources/AgentMemory/NLContextualEmbeddingProvider.swift @@ -279,6 +279,7 @@ public extension MemoryConfiguration { structuredQueryExpander: (any StructuredQueryExpander)? = GenericStructuredQueryExpander(), reranker: (any Reranker)? = nil, contentTagger: (any ContentTagger)? = nil, + entityTagger: (any EntityTagger)? = NLEntityTagger(), memoryExtractor: (any MemoryExtractor)? = nil, recallPlanner: (any RecallPlanner)? = nil, queryAnalyzer: (any QueryAnalyzer)? = NLQueryAnalyzer(), @@ -307,7 +308,8 @@ public extension MemoryConfiguration { lexicalCandidateLimit: lexicalCandidateLimit, fusionK: fusionK, positionAwareBlending: positionAwareBlending, - ftsTokenizer: ftsTokenizer + ftsTokenizer: ftsTokenizer, + entityTagger: entityTagger ) } } diff --git a/Sources/AgentMemory/NLEntityTagger.swift b/Sources/AgentMemory/NLEntityTagger.swift new file mode 100644 index 0000000..3f61403 --- /dev/null +++ b/Sources/AgentMemory/NLEntityTagger.swift @@ -0,0 +1,97 @@ +#if MEMORY_NATURAL_LANGUAGE +import Foundation +import NaturalLanguage + +/// NaturalLanguage-backed named entity recognition for the memory write path. +/// +/// `NLTagger` name recognition depends on capitalization: it stays silent on +/// all-lowercase chat text, so `recognizeEntities` degrades to a no-op there +/// and heuristic extraction remains the behavior floor. `isLikelyPlaceName` +/// validates short phrases by title-casing them inside a natural carrier +/// sentence, where place recognition is markedly more precise than tagging +/// the bare phrase. +public struct NLEntityTagger: EntityTagger { + public let identifier = "nl-entity-tagger" + + public init() {} + + public func recognizeEntities(in text: String) -> [MemoryEntity] { + let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return [] } + + var entities: [MemoryEntity] = [] + var seen: Set = [] + enumerateNameTags(in: trimmed) { value, label in + let normalizedValue = MemoryExtractionHeuristics.normalizeEntityValue(value) + guard !normalizedValue.isEmpty, seen.insert(normalizedValue).inserted else { return } + entities.append( + MemoryEntity( + label: label, + value: value, + normalizedValue: normalizedValue, + confidence: 0.8 + ) + ) + } + return entities + } + + public func isLikelyPlaceName(_ phrase: String) -> Bool { + let trimmed = phrase.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty, trimmed.count <= 60 else { return false } + + let candidate = titleCased(trimmed) + let carrier = "They moved to \(candidate) last month." + var foundPlace = false + enumerateNameTags(in: carrier) { value, label in + guard label == .location else { return } + if candidate.range(of: value, options: [.caseInsensitive]) != nil { + foundPlace = true + } + } + return foundPlace + } + + private func enumerateNameTags(in text: String, handler: (String, EntityLabel) -> Void) { + let tagger = NLTagger(tagSchemes: [.nameType]) + tagger.string = text + tagger.enumerateTags( + in: text.startIndex.. EntityLabel? { + switch tag { + case .personalName: + .person + case .placeName: + .location + case .organizationName: + .organization + default: + nil + } + } + + private func titleCased(_ phrase: String) -> String { + phrase + .split(separator: " ", omittingEmptySubsequences: false) + .map { word -> String in + guard let first = word.first, first.isLowercase else { return String(word) } + return first.uppercased() + word.dropFirst() + } + .joined(separator: " ") + } +} +#endif diff --git a/Sources/AgentMemory/ProtocolsAndErrors.swift b/Sources/AgentMemory/ProtocolsAndErrors.swift index ad22617..2b4bb7a 100644 --- a/Sources/AgentMemory/ProtocolsAndErrors.swift +++ b/Sources/AgentMemory/ProtocolsAndErrors.swift @@ -231,6 +231,20 @@ public protocol Chunker: Sendable { func chunk(text: String, kind: DocumentKind, sourceURL: URL?) -> [Chunk] } +public protocol EntityTagger: Sendable { + var identifier: String { get } + + /// Recognizes named entities (people, places, organizations) in prose. + /// Implementations should be conservative: returned entities are merged + /// additively into heuristic extraction results. + func recognizeEntities(in text: String) -> [MemoryEntity] + + /// Whether a short phrase plausibly names a real-world place. + /// Used to confirm ambiguous self-reported location phrases before they + /// become durable profile memories. + func isLikelyPlaceName(_ phrase: String) -> Bool +} + public struct MemoryConfiguration: Sendable { public var databaseURL: URL public var embeddingProvider: any EmbeddingProvider @@ -249,6 +263,7 @@ public struct MemoryConfiguration: Sendable { public var positionAwareBlending: PositionAwareBlending public var groundedQueryExpansion: GroundedQueryExpansionConfiguration public var ftsTokenizer: (any Tokenizer)? + public var entityTagger: (any EntityTagger)? /// How long recorded memory signals are kept before being pruned. /// Values <= 0 disable pruning. public var memorySignalRetention: TimeInterval @@ -271,6 +286,7 @@ public struct MemoryConfiguration: Sendable { positionAwareBlending: PositionAwareBlending = .default, groundedQueryExpansion: GroundedQueryExpansionConfiguration = .conservativeDefault, ftsTokenizer: (any Tokenizer)? = nil, + entityTagger: (any EntityTagger)? = Self.defaultEntityTagger, memorySignalRetention: TimeInterval = Self.defaultMemorySignalRetention ) { self.databaseURL = databaseURL @@ -290,9 +306,18 @@ public struct MemoryConfiguration: Sendable { self.positionAwareBlending = positionAwareBlending self.groundedQueryExpansion = groundedQueryExpansion self.ftsTokenizer = ftsTokenizer + self.entityTagger = entityTagger self.memorySignalRetention = memorySignalRetention } + public static var defaultEntityTagger: (any EntityTagger)? { + #if MEMORY_NATURAL_LANGUAGE + NLEntityTagger() + #else + nil + #endif + } + public static let defaultMemorySignalRetention: TimeInterval = 90 * 24 * 60 * 60 public static var defaultSupportedExtensions: Set { diff --git a/Tests/MemoryTests/MemoryExternalAPITests.swift b/Tests/MemoryTests/MemoryExternalAPITests.swift index 5111945..c098332 100644 --- a/Tests/MemoryTests/MemoryExternalAPITests.swift +++ b/Tests/MemoryTests/MemoryExternalAPITests.swift @@ -889,6 +889,127 @@ struct MemoryExternalAPITests { #expect(maintenance.consideredSignalCount == 1) } + @Test + func ambiguousLocationRequiresConfirmationWithoutTagger() async throws { + let root = try makeTemporaryDirectory() + let dbURL = root.appendingPathComponent("index.sqlite") + + let index = try MemoryIndex( + configuration: MemoryConfiguration( + databaseURL: dbURL, + embeddingProvider: MockEmbeddingProvider(), + entityTagger: nil + ) + ) + + let extracted = try await index.extract( + from: [ + ConversationMessage(role: .user, content: "btw i'm in lisbon, any dinner recs?"), + ], + limit: 10 + ) + + #expect(!extracted.contains { $0.text.contains("lives in") }) + } + +#if MEMORY_NATURAL_LANGUAGE + @Test + func nlEntityTaggerConfirmsPlacesAndRejectsCommonNouns() { + let tagger = NLEntityTagger() + + #expect(tagger.isLikelyPlaceName("lisbon")) + #expect(tagger.isLikelyPlaceName("tokyo")) + #expect(tagger.isLikelyPlaceName("new york")) + #expect(!tagger.isLikelyPlaceName("a meeting")) + #expect(!tagger.isLikelyPlaceName("fear")) + #expect(!tagger.isLikelyPlaceName("my apartment")) + + let entities = tagger.recognizeEntities(in: "Annie retired from Google last year.") + #expect(entities.contains { $0.label == .person && $0.normalizedValue == "annie" }) + #expect(entities.contains { $0.label == .organization && $0.normalizedValue == "google" }) + + // Lowercase chat text yields nothing; heuristics remain the floor. + #expect(tagger.recognizeEntities(in: "annie retired from google last year").isEmpty) + } + + @Test + func heuristicExtractAcceptsTaggerConfirmedAmbiguousLocation() async throws { + let root = try makeTemporaryDirectory() + let dbURL = root.appendingPathComponent("index.sqlite") + + let index = try MemoryIndex( + configuration: MemoryConfiguration( + databaseURL: dbURL, + embeddingProvider: MockEmbeddingProvider() + ) + ) + + let extracted = try await index.extract( + from: [ + ConversationMessage(role: .user, content: "btw i'm in lisbon, any dinner recs?"), + ], + limit: 10 + ) + + let profile = try #require(extracted.first { $0.text == "The user lives in Lisbon." }) + #expect(profile.kind == .profile) + #expect(profile.subject == .user) + #expect(profile.canonicalKey == "profile:user:location") + #expect(profile.facetTags.contains(.location)) + #expect(profile.entities.contains { entity in + entity.label == .location && entity.normalizedValue == "lisbon" + }) + } + + @Test + func heuristicExtractRejectsTransientTravelLocation() async throws { + let root = try makeTemporaryDirectory() + let dbURL = root.appendingPathComponent("index.sqlite") + + let index = try MemoryIndex( + configuration: MemoryConfiguration( + databaseURL: dbURL, + embeddingProvider: MockEmbeddingProvider() + ) + ) + + let extracted = try await index.extract( + from: [ + ConversationMessage(role: .user, content: "i'm in boston this week, remind me to pack"), + ConversationMessage(role: .user, content: "i'm in tokyo tonight, any dinner recs?"), + ], + limit: 10 + ) + + #expect(!extracted.contains { $0.text.contains("lives in") }) + } + + @Test + func heuristicExtractUpgradesEntityLabelsFromTagger() async throws { + let root = try makeTemporaryDirectory() + let dbURL = root.appendingPathComponent("index.sqlite") + + let index = try MemoryIndex( + configuration: MemoryConfiguration( + databaseURL: dbURL, + embeddingProvider: MockEmbeddingProvider() + ) + ) + + let extracted = try await index.extract( + from: [ + ConversationMessage(role: .user, content: "Priya is the maintainer of the ingest pipeline."), + ], + limit: 10 + ) + + let profile = try #require(extracted.first) + #expect(profile.entities.contains { entity in + entity.label == .person && entity.normalizedValue == "priya" + }) + } +#endif + @Test func capturePreviewAndIngestUseSubjectAwareEvidence() async throws { let root = try makeTemporaryDirectory() From 07402aeea98dddcace57185d6a3d8cd7f99f7b58 Mon Sep 17 00:00:00 2001 From: Zac White Date: Sun, 12 Jul 2026 16:31:38 -0700 Subject: [PATCH 2/5] Add CI workflow for PR commits and main pushes - ci.yml runs on every pull_request event and pushes to main: build and test with default traits, benchmark-leakage check, strict dataset validation, plus a separate CoreML+NaturalLanguage trait build job. SwiftPM build directories are cached per trait flavor; in-progress runs for a ref are cancelled on new commits. - evals.yml is workflow_dispatch-only: runs the extraction-affected eval suites under coreml_default, gates the pressure and smoke baselines, and uploads run reports. Manual because eval results are OS-sensitive; the release gate stays local per AGENTS.md. --- .github/workflows/ci.yml | 63 +++++++++++++++++++++++++++++++++++++ .github/workflows/evals.yml | 53 +++++++++++++++++++++++++++++++ 2 files changed, 116 insertions(+) create mode 100644 .github/workflows/ci.yml create mode 100644 .github/workflows/evals.yml diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..83d3193 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,63 @@ +name: CI + +on: + pull_request: + push: + branches: [main] + +concurrency: + group: ci-${{ github.ref }} + cancel-in-progress: true + +jobs: + test: + name: Build and test (default traits) + runs-on: macos-15 + timeout-minutes: 45 + steps: + - uses: actions/checkout@v4 + + - name: Select latest Xcode + run: | + sudo xcode-select -s "$(ls -d /Applications/Xcode*.app | sort -V | tail -n 1)/Contents/Developer" + swift --version + + - name: Cache SwiftPM build + uses: actions/cache@v4 + with: + path: .build + key: ${{ runner.os }}-spm-default-${{ hashFiles('Package.swift', 'Package.resolved') }} + restore-keys: | + ${{ runner.os }}-spm-default- + + - name: Run tests + run: swift test + + - name: Benchmark leakage check + run: python3 Scripts/check_benchmark_leakage.py + + - name: Validate eval datasets + run: swift run memory_eval validate-datasets --strict + + trait-build: + name: Build (CoreML + NaturalLanguage traits) + runs-on: macos-15 + timeout-minutes: 45 + steps: + - uses: actions/checkout@v4 + + - name: Select latest Xcode + run: | + sudo xcode-select -s "$(ls -d /Applications/Xcode*.app | sort -V | tail -n 1)/Contents/Developer" + swift --version + + - name: Cache SwiftPM build + uses: actions/cache@v4 + with: + path: .build + key: ${{ runner.os }}-spm-traits-${{ hashFiles('Package.swift', 'Package.resolved') }} + restore-keys: | + ${{ runner.os }}-spm-traits- + + - name: Build with traits + run: swift build --traits CoreMLEmbedding,MemoryNaturalLanguage diff --git a/.github/workflows/evals.yml b/.github/workflows/evals.yml new file mode 100644 index 0000000..ee4b6d7 --- /dev/null +++ b/.github/workflows/evals.yml @@ -0,0 +1,53 @@ +name: Evals + +# Manual-only: eval results are OS-sensitive (NaturalLanguage models vary by +# macOS version), so these run on demand for diagnostics rather than as a +# PR-blocking gate. The authoritative release gate remains local runs against +# Evals/baselines/current.json per AGENTS.md. +on: + workflow_dispatch: + +jobs: + extraction-suites: + name: Extraction-affected eval suites (coreml_default) + runs-on: macos-15 + timeout-minutes: 90 + steps: + - uses: actions/checkout@v4 + + - name: Select latest Xcode + run: | + sudo xcode-select -s "$(ls -d /Applications/Xcode*.app | sort -V | tail -n 1)/Contents/Developer" + swift --version + + - name: Cache SwiftPM build + uses: actions/cache@v4 + with: + path: .build + key: ${{ runner.os }}-spm-traits-${{ hashFiles('Package.swift', 'Package.resolved') }} + restore-keys: | + ${{ runner.os }}-spm-traits- + + - name: Run extraction-affected suites + run: | + for ds in memory_schema_gold_v2 agent_memory_gold_v1 agent_memory_pressure_v1 eval_quality_smoke_v1 storage_heldout_v1; do + echo "=== $ds ===" + swift run --traits CoreMLEmbedding,MemoryNaturalLanguage memory_eval run \ + --profile coreml_default --dataset-root "./Evals/$ds" --no-cache --no-index-cache + done + + - name: Gate focused baselines + run: | + pressure_run=$(ls -t Evals/agent_memory_pressure_v1/runs/*-coreml_default.json | head -1) + smoke_run=$(ls -t Evals/eval_quality_smoke_v1/runs/*-coreml_default.json | head -1) + swift run --traits CoreMLEmbedding,MemoryNaturalLanguage memory_eval gate \ + --baseline Evals/baselines/pressure.json "$pressure_run" + swift run --traits CoreMLEmbedding,MemoryNaturalLanguage memory_eval gate \ + --baseline Evals/baselines/eval_quality_smoke.json "$smoke_run" + + - name: Upload run reports + uses: actions/upload-artifact@v4 + with: + name: eval-runs + path: Evals/*/runs/*coreml_default.* + if-no-files-found: warn From 308d16e96ff8c8577174db7e79520849fe70eb08 Mon Sep 17 00:00:00 2001 From: Zac White Date: Sun, 12 Jul 2026 16:36:35 -0700 Subject: [PATCH 3/5] Run scoped-FTS caller closure outside the seeding transaction Swift 6.2 (the CI toolchain) rejects capturing the caller's closure inside the actor-isolated transaction closure with a sending-risk diagnostic; Swift 6.3 accepts it. The temp FTS table is session-scoped, so seeding commits first and the caller's queries run after, with identical observable behavior. --- Sources/MemoryStorage/MemoryStorage.swift | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/Sources/MemoryStorage/MemoryStorage.swift b/Sources/MemoryStorage/MemoryStorage.swift index 374ffa5..9f9a171 100644 --- a/Sources/MemoryStorage/MemoryStorage.swift +++ b/Sources/MemoryStorage/MemoryStorage.swift @@ -1957,7 +1957,10 @@ public actor MemoryStorage { sql: "CREATE VIRTUAL TABLE IF NOT EXISTS temp.\(Self.scopedLexicalFTSTableName) USING fts5(content)" ) - return try database.transaction { + // body must stay outside the transaction closure: Swift 6.2 rejects + // capturing the caller's closure there (sending-risk diagnostic), and + // the temp table is session-scoped, so its rows survive the commit. + try database.transaction { try database.execute(sql: "DELETE FROM temp.\(Self.scopedLexicalFTSTableName)") for metadata in metadataRows { try database.execute( @@ -1968,8 +1971,8 @@ public actor MemoryStorage { arguments: [metadata.chunkID, metadata.content] ) } - return try body() } + return try body() } private static func runPreparedScopedLexicalSearchQuery( From b197ce05096e911a183449755d450498e74da9a9 Mon Sep 17 00:00:00 2001 From: Zac White Date: Sun, 12 Jul 2026 16:40:08 -0700 Subject: [PATCH 4/5] Honor DEVELOPER_DIR for developer framework search paths The Swift Testing workaround hardcoded /Applications/Xcode.app framework paths. On CI runners the selected Xcode differs from that symlink, so the @Test macro from the selected toolchain expanded against another Xcode's Testing framework and failed to build. When DEVELOPER_DIR is set, use only the selected toolchain's paths; CI now exports it from the Xcode selection step. --- .github/workflows/ci.yml | 8 ++++++-- .github/workflows/evals.yml | 4 +++- Package.swift | 41 +++++++++++++++++++++++++------------ 3 files changed, 37 insertions(+), 16 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 83d3193..ca1376f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -19,7 +19,9 @@ jobs: - name: Select latest Xcode run: | - sudo xcode-select -s "$(ls -d /Applications/Xcode*.app | sort -V | tail -n 1)/Contents/Developer" + XCODE_DEV="$(ls -d /Applications/Xcode*.app | sort -V | tail -n 1)/Contents/Developer" + sudo xcode-select -s "$XCODE_DEV" + echo "DEVELOPER_DIR=$XCODE_DEV" >> "$GITHUB_ENV" swift --version - name: Cache SwiftPM build @@ -48,7 +50,9 @@ jobs: - name: Select latest Xcode run: | - sudo xcode-select -s "$(ls -d /Applications/Xcode*.app | sort -V | tail -n 1)/Contents/Developer" + XCODE_DEV="$(ls -d /Applications/Xcode*.app | sort -V | tail -n 1)/Contents/Developer" + sudo xcode-select -s "$XCODE_DEV" + echo "DEVELOPER_DIR=$XCODE_DEV" >> "$GITHUB_ENV" swift --version - name: Cache SwiftPM build diff --git a/.github/workflows/evals.yml b/.github/workflows/evals.yml index ee4b6d7..691224a 100644 --- a/.github/workflows/evals.yml +++ b/.github/workflows/evals.yml @@ -17,7 +17,9 @@ jobs: - name: Select latest Xcode run: | - sudo xcode-select -s "$(ls -d /Applications/Xcode*.app | sort -V | tail -n 1)/Contents/Developer" + XCODE_DEV="$(ls -d /Applications/Xcode*.app | sort -V | tail -n 1)/Contents/Developer" + sudo xcode-select -s "$XCODE_DEV" + echo "DEVELOPER_DIR=$XCODE_DEV" >> "$GITHUB_ENV" swift --version - name: Cache SwiftPM build diff --git a/Package.swift b/Package.swift index 3dec64e..dcd6b59 100644 --- a/Package.swift +++ b/Package.swift @@ -3,27 +3,42 @@ import Foundation import PackageDescription // Some Apple Swift CLT/snapshot installs place Swift Testing in Developer -// frameworks instead of the default SwiftPM search paths. -let developerFrameworkSwiftSettings: [SwiftSetting] = [ - "/Library/Developer/CommandLineTools/Library/Developer/Frameworks", - "/Applications/Xcode.app/Contents/Developer/Library/Developer/Frameworks", -].compactMap { path in +// frameworks instead of the default SwiftPM search paths. When DEVELOPER_DIR +// is set, use only that toolchain's paths: mixing another Xcode's Testing +// framework with the selected compiler's @Test macro fails to build. +let developerDir = ProcessInfo.processInfo.environment["DEVELOPER_DIR"] + +let developerFrameworkPaths: [String] = { + if let developerDir, !developerDir.isEmpty { + return [developerDir + "/Library/Developer/Frameworks"] + } + return [ + "/Library/Developer/CommandLineTools/Library/Developer/Frameworks", + "/Applications/Xcode.app/Contents/Developer/Library/Developer/Frameworks", + ] +}() + +let developerLibraryPaths: [String] = { + if let developerDir, !developerDir.isEmpty { + return [developerDir + "/Library/Developer/usr/lib"] + } + return [ + "/Library/Developer/CommandLineTools/Library/Developer/usr/lib", + "/Applications/Xcode.app/Contents/Developer/Library/Developer/usr/lib", + ] +}() + +let developerFrameworkSwiftSettings: [SwiftSetting] = developerFrameworkPaths.compactMap { path in guard FileManager.default.fileExists(atPath: path) else { return nil } return .unsafeFlags(["-F", path], .when(platforms: [.macOS])) } -let developerFrameworkLinkerSettings: [LinkerSetting] = [ - "/Library/Developer/CommandLineTools/Library/Developer/Frameworks", - "/Applications/Xcode.app/Contents/Developer/Library/Developer/Frameworks", -].compactMap { path in +let developerFrameworkLinkerSettings: [LinkerSetting] = developerFrameworkPaths.compactMap { path in guard FileManager.default.fileExists(atPath: path) else { return nil } return .unsafeFlags(["-F", path, "-Xlinker", "-rpath", "-Xlinker", path], .when(platforms: [.macOS])) } -let developerLibraryLinkerSettings: [LinkerSetting] = [ - "/Library/Developer/CommandLineTools/Library/Developer/usr/lib", - "/Applications/Xcode.app/Contents/Developer/Library/Developer/usr/lib", -].compactMap { path in +let developerLibraryLinkerSettings: [LinkerSetting] = developerLibraryPaths.compactMap { path in guard FileManager.default.fileExists(atPath: path) else { return nil } return .unsafeFlags(["-L", path, "-Xlinker", "-rpath", "-Xlinker", path], .when(platforms: [.macOS])) } From da3dab404bf7ec5c19d9609d8e1513157c2bd2ff Mon Sep 17 00:00:00 2001 From: Zac White Date: Sun, 12 Jul 2026 17:26:11 -0700 Subject: [PATCH 5/5] Tighten presence location handling in memory extraction --- Evals/agent_memory_gold_v1/scenarios.jsonl | 2 + README.md | 4 +- .../MemoryExtractionHeuristics.swift | 165 ++++++++++-------- Sources/AgentMemory/MemoryIndex.swift | 43 +++-- Sources/AgentMemory/NLEntityTagger.swift | 78 +++------ Sources/AgentMemory/NLQueryAnalyzer.swift | 49 +----- Sources/AgentMemory/ProtocolsAndErrors.swift | 10 +- .../MemoryTests/MemoryExternalAPITests.swift | 101 +++++++++-- Tests/MemoryTests/TestSupport.swift | 20 +++ 9 files changed, 262 insertions(+), 210 deletions(-) diff --git a/Evals/agent_memory_gold_v1/scenarios.jsonl b/Evals/agent_memory_gold_v1/scenarios.jsonl index 7d75de2..51128f3 100644 --- a/Evals/agent_memory_gold_v1/scenarios.jsonl +++ b/Evals/agent_memory_gold_v1/scenarios.jsonl @@ -34,3 +34,5 @@ {"id": "context-prep-location-protects-profile", "source_family": "context_preparation", "difficulty": "medium", "generation_method": "seed", "setup_memories": [{"text": "The user lives in San Francisco, CA.", "kind": "profile", "status": "active", "canonical_key": "profile:user:location", "facet_tags": ["fact_about_user", "location"], "entity_values": ["san francisco"], "topics": []}, {"text": "The user previously researched Austin restaurants.", "kind": "episode", "status": "active", "canonical_key": "episode:austin-restaurants", "facet_tags": [], "entity_values": ["austin"], "topics": ["austin restaurants"]}], "setup_context_hints": [{"path_prefix": "memory://", "context": "Memory records are durable user-facing facts."}], "messages": [{"role": "user", "content": "What should I do tonight in San Francisco?"}], "expected_write_count": 0, "expected_memories": [], "context_expectations": [{"query": "What should I do tonight in San Francisco?", "max_references": 4, "max_tokens": 256, "expected_text_contains": ["San Francisco"], "forbidden_text_contains": ["Austin"], "expected_hint_contains": ["durable user-facing facts"], "require_untrusted_framing": true}]} {"id": "maintenance-repeated-recall-promotes-profile", "source_family": "maintenance", "difficulty": "medium", "generation_method": "seed", "setup_memories": [{"text": "The user prefers ramen for casual dinners.", "kind": "profile", "status": "active", "canonical_key": "profile:user:preference:ramen", "facet_tags": ["fact_about_user", "preference"], "entity_values": [], "topics": ["ramen"]}], "messages": [{"role": "user", "content": "Please do not save anything from this turn."}], "expected_write_count": 0, "expected_memories": [], "maintenance_expectation": {"signal_memory_canonical_key": "profile:user:preference:ramen", "signal_queries": ["dinner ideas", "casual food", "dinner ideas"], "signal_confidence": 0.9, "min_signal_count": 3, "min_distinct_queries": 2, "min_confidence": 0.75, "expected_proposal_text_contains": ["ramen"], "forbidden_proposal_text_contains": ["do not save"]}} {"id": "maintenance-threshold-blocks-single-query", "source_family": "maintenance", "difficulty": "medium", "generation_method": "seed", "setup_memories": [{"text": "The user prefers ramen for casual dinners.", "kind": "profile", "status": "active", "canonical_key": "profile:user:preference:ramen", "facet_tags": ["fact_about_user", "preference"], "entity_values": [], "topics": ["ramen"]}], "messages": [{"role": "user", "content": "Thanks, no memory update."}], "expected_write_count": 0, "expected_memories": [], "maintenance_expectation": {"signal_memory_canonical_key": "profile:user:preference:ramen", "signal_queries": ["dinner ideas", "dinner ideas"], "signal_confidence": 0.9, "min_signal_count": 3, "min_distinct_queries": 2, "min_confidence": 0.75, "expected_proposal_count": 0, "forbidden_proposal_text_contains": ["ramen"]}} +{"id":"presence-location-does-not-become-residence","case_category":"adversarial","source_family":"no_write","difficulty":"hard","generation_method":"seed","review_status":"curated","synthetic_status":"synthetic","messages":[{"role":"user","content":"I'm in Lisbon, any dinner recommendations?"}],"expected_write_count":0,"expected_memories":[],"recall_queries":[]} +{"id":"qualified-presence-location-does-not-become-residence","case_category":"known_regression","source_family":"no_write","difficulty":"hard","generation_method":"seed","review_status":"curated","synthetic_status":"synthetic","messages":[{"role":"user","content":"I'm in Tokyo currently for work."}],"expected_write_count":0,"expected_memories":[],"recall_queries":[]} diff --git a/README.md b/README.md index 95ce340..7a43d5a 100644 --- a/README.md +++ b/README.md @@ -133,7 +133,9 @@ Captured `MemoryCandidate` and stored `MemoryRecord` values can include a `subje ### Linguistic Entity Tagging -When the `MemoryNaturalLanguage` trait is enabled (the default), extraction runs `NLEntityTagger` — an `NLTagger`-backed named entity recognizer — alongside the heuristics. Tagger results are strictly additive: they upgrade generic entity labels to `person`/`location`/`organization`, add entities the heuristics missed in well-capitalized text, and confirm ambiguous self-reported locations ("i'm in portland") before they become durable profile memories. Travel-qualified phrases ("i'm in boston this week") are treated as transient and not recorded as residence. Heuristic behavior is the floor: with `entityTagger: nil` (or the trait disabled) extraction is unchanged, and residence statements ("i live in ...") never require tagger confirmation. Provide a custom `EntityTagger` through `MemoryConfiguration` to swap in your own recognizer. +When the `MemoryNaturalLanguage` trait is enabled (the default), every extraction provider is followed by a shared entity-enrichment stage using `NLEntityTagger`, an `NLTagger`-backed named entity recognizer. Extractor-supplied and heuristic entities remain the floor. Linguistic annotations can conservatively specialize generic labels or add missed entities when the surrounding prose supports the proposed person/location/organization label, but they do not independently create candidates, facets, or durable write decisions. NaturalLanguage does not expose calibrated confidence for name tags, so its annotations are intentionally treated as advisory. + +Residence extraction requires an explicit durable relation such as "I live in", "I'm based in", or "I moved to". Presence statements such as "I'm in Portland" may describe travel and are not rewritten as residence even when Portland is recognized as a place. Set `entityTagger: nil` to disable linguistic enrichment, or provide a custom `EntityTagger`; high-confidence custom annotations can be accepted without the default contextual corroboration. Context preparation retrieves bounded memory context for the next model turn: diff --git a/Sources/AgentMemory/MemoryExtractionHeuristics.swift b/Sources/AgentMemory/MemoryExtractionHeuristics.swift index 78a4110..5179299 100644 --- a/Sources/AgentMemory/MemoryExtractionHeuristics.swift +++ b/Sources/AgentMemory/MemoryExtractionHeuristics.swift @@ -27,8 +27,7 @@ internal enum MemoryExtractionHeuristics { messages: [ConversationMessage], limit: Int, canonicalKey: CanonicalKeyResolver, - proposedAction: ProposedActionResolver, - entityTagger: (any EntityTagger)? = nil + proposedAction: ProposedActionResolver ) -> MemoryExtractionResult { var extracted: [MemoryCandidate] = [] extracted.reserveCapacity(limit) @@ -42,11 +41,7 @@ internal enum MemoryExtractionHeuristics { .trimmingCharacters(in: .whitespacesAndNewlines) guard !normalized.isEmpty else { continue } - let focusedSegments = focusedUserProfileSegments( - from: normalized, - role: message.role, - entityTagger: entityTagger - ) + let focusedSegments = focusedUserProfileSegments(from: normalized, role: message.role) let rawSegments = focusedSegments.map(\.text) + splitExtractionSegments(normalized) for rawSegment in rawSegments { @@ -82,10 +77,7 @@ internal enum MemoryExtractionHeuristics { let tags = inferredTags(forExtractedText: segment) let entities = mergeEntities( pinned: focusedSegment?.entities ?? [], - inferred: mergeTaggerEntities( - entityTagger?.recognizeEntities(in: segment) ?? [], - into: inferEntities(forExtractedText: segment) - ) + inferred: inferEntities(forExtractedText: segment) ) let facetTags = inferFacetTags(forExtractedText: segment, kind: kind, knownEntities: entities) .union(focusedSegment?.facetTags ?? []) @@ -533,6 +525,12 @@ internal enum MemoryExtractionHeuristics { lower, needles: ["remember that", "please remember", "for future reference", "note that", "keep in mind"] ) + if kind == .profile, + isAmbiguousPresenceStatement(lower), + !durableMemoryRequest, + !containsDurableProfileSignal(lower) { + return false + } if isQuestionLikeNonMemorySegment(lower), !durableMemoryRequest { return false } @@ -569,6 +567,25 @@ internal enum MemoryExtractionHeuristics { } } + private static func isAmbiguousPresenceStatement(_ lower: String) -> Bool { + lower.range( + of: #"\b(?:i\s+am|i['’]m)\s+in\s+"#, + options: .regularExpression + ) != nil + } + + private static func containsDurableProfileSignal(_ lower: String) -> Bool { + containsAny( + lower, + needles: [ + "prefer", "favorite", "usually", "works closely", "timezone", + "my name", "my role", "role is", "maintainer", "owner", + "i live in", "living in", "based in", "moved to", + "my city is", "my location is", + ] + ) + } + private static func isQuestionLikeNonMemorySegment(_ lower: String) -> Bool { let questionPrefixes = [ "what ", "when ", "where ", "which ", "who ", "why ", "how ", @@ -602,13 +619,12 @@ internal enum MemoryExtractionHeuristics { private static func focusedUserProfileSegments( from text: String, - role: ConversationRole, - entityTagger: (any EntityTagger)? = nil + role: ConversationRole ) -> [FocusedProfileSegment] { guard role == .user else { return [] } var segments: [FocusedProfileSegment] = [] - if let location = selfReportedLocation(in: text, entityTagger: entityTagger) { + if let location = selfReportedLocation(in: text) { let city = location.split(separator: ",").first.map(String.init) ?? location segments.append( FocusedProfileSegment( @@ -640,16 +656,18 @@ internal enum MemoryExtractionHeuristics { return merged } - /// Merges tagger-recognized entities into heuristic results additively: - /// heuristic entities are the floor, tagger results upgrade `.other` - /// labels and append entities the heuristics missed. - internal static func mergeTaggerEntities( - _ tagged: [MemoryEntity], - into inferred: [MemoryEntity] + /// Reconciles extractor-, heuristic-, and tagger-supplied entities. Explicit + /// extractor values remain the floor. Low-confidence linguistic labels only + /// specialize an entity when the surrounding prose supports that label. + internal static func enrichedEntities( + supplied: [MemoryEntity], + tagged: [MemoryEntity], + text: String ) -> [MemoryEntity] { - guard !tagged.isEmpty else { return inferred } + let inferred = inferEntities(forExtractedText: text) + var merged = mergeEntities(pinned: supplied, inferred: inferred) + guard !tagged.isEmpty else { return merged } - var merged = inferred var indexByValue: [String: Int] = [:] for (index, entity) in merged.enumerated() { indexByValue[normalizeEntityValue(entity.normalizedValue)] = index @@ -658,6 +676,7 @@ internal enum MemoryExtractionHeuristics { for entity in tagged { let normalizedValue = normalizeEntityValue(entity.normalizedValue) guard !normalizedValue.isEmpty, !isGenericExtractedEntity(normalizedValue) else { continue } + guard isSupportedTaggedEntity(entity, in: text) else { continue } if let index = indexByValue[normalizedValue] { let existing = merged[index] if existing.label == .other { @@ -668,7 +687,8 @@ internal enum MemoryExtractionHeuristics { confidence: max(existing.confidence ?? 0, entity.confidence ?? 0) ) } - } else if merged.count < 8 { + } else if merged.count < 8, + !merged.contains(where: { valuesOverlap($0.normalizedValue, normalizedValue) }) { merged.append( MemoryEntity( label: entity.label, @@ -684,15 +704,56 @@ internal enum MemoryExtractionHeuristics { return merged } - // Residence statements ("i live in ...") accept any captured place phrase. - // The ambiguous forms ("i'm in ...") commonly describe transient states - // ("i'm in a meeting"), so they require confirmation: an alias-table hit, - // or an entity tagger recognizing the phrase as a place — and never with a - // travel qualifier ("i'm in boston this week"). + private static func isSupportedTaggedEntity(_ entity: MemoryEntity, in text: String) -> Bool { + if entity.confidence.map({ $0 >= 0.9 }) == true { + return true + } + + let escapedValue = NSRegularExpression.escapedPattern(for: entity.value) + guard !escapedValue.isEmpty else { return false } + + let patterns: [String] + switch entity.label { + case .person: + patterns = [ + #"\b\#(escapedValue)\b\s+(?:is|was|said|asked|joined|reviewed|prefers|works|retired|captured|noted|owns|leads)\b"#, + #"\b(?:met|with|ask|asked|contact|manager|coworker|teammate)\s+\#(escapedValue)\b"#, + ] + case .location: + patterns = [ + #"\b(?:in|at|near|from|to)\s+\#(escapedValue)\b"#, + #"\b(?:live|lives|living|based|moved)\s+(?:in|to)\s+\#(escapedValue)\b"#, + ] + case .organization: + patterns = [ + #"\b(?:at|for|from|joined|maintainer\s+of|works\s+at)\s+\#(escapedValue)\b"#, + #"\b\#(escapedValue)\b\s+(?:team|company|organization)\b"#, + ] + case .product, .project, .tool, .date, .other: + return true + } + + return patterns.contains { pattern in + text.range(of: pattern, options: [.regularExpression, .caseInsensitive]) != nil + } + } + + private static func valuesOverlap(_ lhs: String, _ rhs: String) -> Bool { + let left = normalizeEntityValue(lhs) + let right = normalizeEntityValue(rhs) + guard !left.isEmpty, !right.isEmpty else { return false } + return left == right + || left.hasPrefix(right + " ") + || left.hasSuffix(" " + right) + || right.hasPrefix(left + " ") + || right.hasSuffix(" " + left) + } + + // Only explicit residence statements create durable profile locations. + // Ambiguous presence statements such as "I'm in Lisbon" may describe + // travel and must not be rewritten as residence. private static let strongLocationTriggerPattern = #"\b(?:i\s+live\s+in|i\s+am\s+living\s+in|i['’]m\s+living\s+in|i\s+am\s+based\s+in|i['’]m\s+based\s+in|i\s+(?:just\s+)?moved\s+to|my\s+city\s+is|my\s+location\s+is)\s+([^.!?;\n]+)"# - private static let weakLocationTriggerPattern = - #"\b(?:i\s+am\s+in|i['’]m\s+in)\s+([^.!?;\n]+)"# private static let knownLocationAliases: [String: String] = [ "sf": "San Francisco, CA", @@ -711,38 +772,22 @@ internal enum MemoryExtractionHeuristics { private static let locationStopTokens: Set = [ "and", "but", "or", "so", "which", "where", "when", "while", "because", "since", "near", "with", "for", "though", "although", "what", "what's", - "what’s", "who", "how", "why", "that", "then", "if", "as", "at", "about", - "this", "next", "until" + "what’s", "who", "how", "why", "that", "then", "if", "as", "at", "about" ] private static let locationTrailingNoiseTokens: Set = [ "now", "currently", "atm", "too", "btw", "there", "here", "tonight", "today" ] - // Tokens that mark an ambiguous "i'm in " as a visit rather than - // residence, either trailing the place ("i'm in boston tonight") or - // immediately following it ("i'm in boston this week"). - private static let transientLocationQualifiers: Set = [ - "this", "next", "until", "for", "tonight", "today", "tomorrow", - "visiting", "traveling", "travelling", "temporarily" - ] - private static let locationLeadingRejectTokens: Set = [ "a", "an", "my", "our", "your", "his", "her", "their", "this", "that", "these", "those", "some", "any", "no", "one", "it", "front", "between", "town", "meetings", "meeting" ] - private static func selfReportedLocation( - in text: String, - entityTagger: (any EntityTagger)? = nil - ) -> String? { + private static func selfReportedLocation(in text: String) -> String? { if let raw = firstRegexCapture(of: strongLocationTriggerPattern, in: text), - let location = parseLocationPhrase(raw, isAmbiguousTrigger: false, entityTagger: entityTagger) { - return location - } - if let raw = firstRegexCapture(of: weakLocationTriggerPattern, in: text), - let location = parseLocationPhrase(raw, isAmbiguousTrigger: true, entityTagger: entityTagger) { + let location = parseLocationPhrase(raw) { return location } return nil @@ -760,15 +805,10 @@ internal enum MemoryExtractionHeuristics { return nsText.substring(with: match.range(at: 1)) } - private static func parseLocationPhrase( - _ raw: String, - isAmbiguousTrigger: Bool, - entityTagger: (any EntityTagger)? = nil - ) -> String? { + private static func parseLocationPhrase(_ raw: String) -> String? { let tokens = raw.split(whereSeparator: \.isWhitespace).map(String.init) var cityTokens: [String] = [] var sawComma = false - var sawTransientQualifier = false var index = 0 while index < tokens.count, cityTokens.count < 4 { @@ -790,20 +830,8 @@ internal enum MemoryExtractionHeuristics { } while let last = cityTokens.last, locationTrailingNoiseTokens.contains(last.lowercased()) { - if transientLocationQualifiers.contains(last.lowercased()) { - sawTransientQualifier = true - } cityTokens.removeLast() } - if index < tokens.count { - let following = tokens[index].trimmingCharacters(in: CharacterSet.punctuationCharacters).lowercased() - if transientLocationQualifiers.contains(following) { - sawTransientQualifier = true - } - } - if isAmbiguousTrigger, sawTransientQualifier { - return nil - } var regionToken: String? if sawComma, index < tokens.count { @@ -840,9 +868,6 @@ internal enum MemoryExtractionHeuristics { } guard city.count >= 2 else { return nil } - if isAmbiguousTrigger { - guard entityTagger?.isLikelyPlaceName(city) == true else { return nil } - } let displayCity = titleCasedLocation(cityTokens) guard let regionToken else { return displayCity } diff --git a/Sources/AgentMemory/MemoryIndex.swift b/Sources/AgentMemory/MemoryIndex.swift index d9359a9..0fc5444 100644 --- a/Sources/AgentMemory/MemoryIndex.swift +++ b/Sources/AgentMemory/MemoryIndex.swift @@ -1152,11 +1152,14 @@ public actor MemoryIndex { guard limit > 0 else { return MemoryExtractionResult() } guard !messages.isEmpty else { return MemoryExtractionResult() } + let extraction: MemoryExtractionResult if let extractor = configuration.memoryExtractor { - return try await extractor.extract(messages: messages, limit: limit) + extraction = try await extractor.extract(messages: messages, limit: limit) + } else { + extraction = heuristicExtract(messages: messages, limit: limit) } - return heuristicExtract(messages: messages, limit: limit) + return enrichExtractionResult(extraction) } public func capture(_ request: MemoryCaptureRequest) async throws -> MemoryCaptureResult { @@ -1961,11 +1964,29 @@ public actor MemoryIndex { }, proposedAction: { candidate in self.proposedWriteAction(for: candidate) - }, - entityTagger: configuration.entityTagger + } ) } + /// Applies the same entity reconciliation to every extraction provider. + /// Linguistic annotations enrich candidate metadata but do not create + /// candidates, facets, or durable write decisions on their own. + private func enrichExtractionResult(_ result: MemoryExtractionResult) -> MemoryExtractionResult { + guard let entityTagger = configuration.entityTagger else { return result } + + var enriched = result + enriched.candidates = result.candidates.map { candidate in + var candidate = candidate + candidate.entities = MemoryExtractionHeuristics.enrichedEntities( + supplied: candidate.entities, + tagged: entityTagger.recognizeEntities(in: candidate.text), + text: candidate.text + ) + return candidate + } + return enriched + } + private func inferredTags(forExtractedText text: String) -> [String] { MemoryExtractionHeuristics.inferredTags(forExtractedText: text) } @@ -2078,15 +2099,11 @@ public actor MemoryIndex { } private func normalizeEntities(_ supplied: [MemoryEntity], text: String) -> [MemoryEntity] { - let preferred: [MemoryEntity] - if supplied.isEmpty { - preferred = MemoryExtractionHeuristics.mergeTaggerEntities( - configuration.entityTagger?.recognizeEntities(in: text) ?? [], - into: inferEntities(forExtractedText: text) - ) - } else { - preferred = supplied - } + let preferred = MemoryExtractionHeuristics.enrichedEntities( + supplied: supplied, + tagged: configuration.entityTagger?.recognizeEntities(in: text) ?? [], + text: text + ) var normalized: [MemoryEntity] = [] var seen: Set = [] normalized.reserveCapacity(min(preferred.count, 8)) diff --git a/Sources/AgentMemory/NLEntityTagger.swift b/Sources/AgentMemory/NLEntityTagger.swift index 3f61403..7cabffa 100644 --- a/Sources/AgentMemory/NLEntityTagger.swift +++ b/Sources/AgentMemory/NLEntityTagger.swift @@ -6,72 +6,54 @@ import NaturalLanguage /// /// `NLTagger` name recognition depends on capitalization: it stays silent on /// all-lowercase chat text, so `recognizeEntities` degrades to a no-op there -/// and heuristic extraction remains the behavior floor. `isLikelyPlaceName` -/// validates short phrases by title-casing them inside a natural carrier -/// sentence, where place recognition is markedly more precise than tagging -/// the bare phrase. +/// and heuristic extraction remains the behavior floor. NaturalLanguage does +/// not expose calibrated confidence for name tags, so annotations use a +/// deliberately modest confidence and downstream reconciliation requires +/// contextual support before specializing heuristic labels. public struct NLEntityTagger: EntityTagger { public let identifier = "nl-entity-tagger" public init() {} public func recognizeEntities(in text: String) -> [MemoryEntity] { + NLNamedEntityRecognition.recognize(in: text, confidence: 0.55) + } +} + +internal enum NLNamedEntityRecognition { + internal static func recognize(in text: String, confidence: Double) -> [MemoryEntity] { let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) guard !trimmed.isEmpty else { return [] } var entities: [MemoryEntity] = [] var seen: Set = [] - enumerateNameTags(in: trimmed) { value, label in - let normalizedValue = MemoryExtractionHeuristics.normalizeEntityValue(value) - guard !normalizedValue.isEmpty, seen.insert(normalizedValue).inserted else { return } - entities.append( - MemoryEntity( - label: label, - value: value, - normalizedValue: normalizedValue, - confidence: 0.8 - ) - ) - } - return entities - } - - public func isLikelyPlaceName(_ phrase: String) -> Bool { - let trimmed = phrase.trimmingCharacters(in: .whitespacesAndNewlines) - guard !trimmed.isEmpty, trimmed.count <= 60 else { return false } - - let candidate = titleCased(trimmed) - let carrier = "They moved to \(candidate) last month." - var foundPlace = false - enumerateNameTags(in: carrier) { value, label in - guard label == .location else { return } - if candidate.range(of: value, options: [.caseInsensitive]) != nil { - foundPlace = true - } - } - return foundPlace - } - - private func enumerateNameTags(in text: String, handler: (String, EntityLabel) -> Void) { let tagger = NLTagger(tagSchemes: [.nameType]) - tagger.string = text + tagger.string = trimmed tagger.enumerateTags( - in: text.startIndex.. EntityLabel? { + private static func entityLabel(for tag: NLTag) -> EntityLabel? { switch tag { case .personalName: .person @@ -83,15 +65,5 @@ public struct NLEntityTagger: EntityTagger { nil } } - - private func titleCased(_ phrase: String) -> String { - phrase - .split(separator: " ", omittingEmptySubsequences: false) - .map { word -> String in - guard let first = word.first, first.isLowercase else { return String(word) } - return first.uppercased() + word.dropFirst() - } - .joined(separator: " ") - } } #endif diff --git a/Sources/AgentMemory/NLQueryAnalyzer.swift b/Sources/AgentMemory/NLQueryAnalyzer.swift index 6632eb0..06dabcf 100644 --- a/Sources/AgentMemory/NLQueryAnalyzer.swift +++ b/Sources/AgentMemory/NLQueryAnalyzer.swift @@ -13,33 +13,7 @@ public struct NLQueryAnalyzer: QueryAnalyzer, Sendable { let lowered = trimmed.lowercased() - var entities: [MemoryEntity] = [] - let nerTagger = NLTagger(tagSchemes: [.nameType]) - nerTagger.string = trimmed - nerTagger.enumerateTags( - in: trimmed.startIndex.. = [] @@ -151,27 +125,6 @@ public struct NLQueryAnalyzer: QueryAnalyzer, Sendable { return base } - private func normalizeEntityValue(_ raw: String) -> String { - raw - .trimmingCharacters(in: .whitespacesAndNewlines) - .split(whereSeparator: \.isWhitespace) - .joined(separator: " ") - .lowercased() - } - - private func entityLabel(for tag: NLTag) -> EntityLabel { - switch tag { - case .personalName: - return .person - case .organizationName: - return .organization - case .placeName: - return .location - default: - return .other - } - } - private let facetKeywords: [FacetTag: [String]] = [ .preference: ["prefer", "favorite", "like", "dislike"], .person: ["who", "person", "people", "contact"], diff --git a/Sources/AgentMemory/ProtocolsAndErrors.swift b/Sources/AgentMemory/ProtocolsAndErrors.swift index 2b4bb7a..13c731b 100644 --- a/Sources/AgentMemory/ProtocolsAndErrors.swift +++ b/Sources/AgentMemory/ProtocolsAndErrors.swift @@ -235,14 +235,10 @@ public protocol EntityTagger: Sendable { var identifier: String { get } /// Recognizes named entities (people, places, organizations) in prose. - /// Implementations should be conservative: returned entities are merged - /// additively into heuristic extraction results. + /// Results are treated as linguistic annotations and reconciled with + /// extractor- and heuristic-supplied entities. They do not independently + /// authorize durable memories or semantic facets. func recognizeEntities(in text: String) -> [MemoryEntity] - - /// Whether a short phrase plausibly names a real-world place. - /// Used to confirm ambiguous self-reported location phrases before they - /// become durable profile memories. - func isLikelyPlaceName(_ phrase: String) -> Bool } public struct MemoryConfiguration: Sendable { diff --git a/Tests/MemoryTests/MemoryExternalAPITests.swift b/Tests/MemoryTests/MemoryExternalAPITests.swift index c098332..479451a 100644 --- a/Tests/MemoryTests/MemoryExternalAPITests.swift +++ b/Tests/MemoryTests/MemoryExternalAPITests.swift @@ -890,7 +890,7 @@ struct MemoryExternalAPITests { } @Test - func ambiguousLocationRequiresConfirmationWithoutTagger() async throws { + func ambiguousLocationDoesNotBecomeResidenceWithoutTagger() async throws { let root = try makeTemporaryDirectory() let dbURL = root.appendingPathComponent("index.sqlite") @@ -914,26 +914,20 @@ struct MemoryExternalAPITests { #if MEMORY_NATURAL_LANGUAGE @Test - func nlEntityTaggerConfirmsPlacesAndRejectsCommonNouns() { + func nlEntityTaggerRecognizesNamedEntitiesConservatively() { let tagger = NLEntityTagger() - #expect(tagger.isLikelyPlaceName("lisbon")) - #expect(tagger.isLikelyPlaceName("tokyo")) - #expect(tagger.isLikelyPlaceName("new york")) - #expect(!tagger.isLikelyPlaceName("a meeting")) - #expect(!tagger.isLikelyPlaceName("fear")) - #expect(!tagger.isLikelyPlaceName("my apartment")) - let entities = tagger.recognizeEntities(in: "Annie retired from Google last year.") #expect(entities.contains { $0.label == .person && $0.normalizedValue == "annie" }) #expect(entities.contains { $0.label == .organization && $0.normalizedValue == "google" }) + #expect(entities.allSatisfy { ($0.confidence ?? 1) < 0.9 }) // Lowercase chat text yields nothing; heuristics remain the floor. #expect(tagger.recognizeEntities(in: "annie retired from google last year").isEmpty) } @Test - func heuristicExtractAcceptsTaggerConfirmedAmbiguousLocation() async throws { + func heuristicExtractDoesNotRewritePresenceAsResidence() async throws { let root = try makeTemporaryDirectory() let dbURL = root.appendingPathComponent("index.sqlite") @@ -947,18 +941,14 @@ struct MemoryExternalAPITests { let extracted = try await index.extract( from: [ ConversationMessage(role: .user, content: "btw i'm in lisbon, any dinner recs?"), + ConversationMessage(role: .user, content: "i'm in sf now, what should we do?"), + ConversationMessage(role: .user, content: "i am in tokyo currently for work."), + ConversationMessage(role: .user, content: "i'm in new york atm."), ], limit: 10 ) - let profile = try #require(extracted.first { $0.text == "The user lives in Lisbon." }) - #expect(profile.kind == .profile) - #expect(profile.subject == .user) - #expect(profile.canonicalKey == "profile:user:location") - #expect(profile.facetTags.contains(.location)) - #expect(profile.entities.contains { entity in - entity.label == .location && entity.normalizedValue == "lisbon" - }) + #expect(extracted.isEmpty) } @Test @@ -1008,8 +998,83 @@ struct MemoryExternalAPITests { entity.label == .person && entity.normalizedValue == "priya" }) } + + @Test + func linguisticLabelsRequireContextBeforeChangingFacets() async throws { + let root = try makeTemporaryDirectory() + let dbURL = root.appendingPathComponent("index.sqlite") + + let index = try MemoryIndex( + configuration: MemoryConfiguration( + databaseURL: dbURL, + embeddingProvider: MockEmbeddingProvider() + ) + ) + + let extracted = try await index.extract( + from: [ + ConversationMessage(role: .user, content: "On Monday, Rowan reviewed the TrailMap export failure."), + ConversationMessage(role: .user, content: "Mina prefers Linear for AtlasKit project triage."), + ConversationMessage(role: .user, content: "After the incident, Mina captured a lesson about Redis backoff limits."), + ], + limit: 10 + ) + + #expect(extracted.count == 3) + #expect(extracted.allSatisfy { !$0.facetTags.contains(.location) }) + #expect(!extracted.flatMap(\.entities).contains { entity in + entity.label == .location && ["mina", "rowan"].contains(entity.normalizedValue) + }) + #expect(!extracted.flatMap(\.entities).contains { entity in + entity.label == .person && entity.normalizedValue == "redis" + }) + } #endif + @Test + func customExtractorCandidatesReceiveCentralEntityEnrichment() async throws { + let root = try makeTemporaryDirectory() + let dbURL = root.appendingPathComponent("index.sqlite") + let candidate = MemoryCandidate( + text: "Priya is the maintainer for AtlasKit.", + kind: .profile, + entities: [ + MemoryEntity( + label: .project, + value: "AtlasKit", + normalizedValue: "atlaskit", + confidence: 0.98 + ), + ] + ) + + let index = try MemoryIndex( + configuration: MemoryConfiguration( + databaseURL: dbURL, + embeddingProvider: MockEmbeddingProvider(), + memoryExtractor: StaticMemoryExtractor( + result: MemoryExtractionResult(candidates: [candidate]) + ), + entityTagger: StaticEntityTagger( + entities: [ + MemoryEntity( + label: .person, + value: "Priya", + normalizedValue: "priya", + confidence: 0.95 + ), + ] + ) + ) + ) + + let extracted = try await index.extract(from: "ignored") + let enriched = try #require(extracted.first) + #expect(enriched.entities.contains { $0.label == .project && $0.normalizedValue == "atlaskit" }) + #expect(enriched.entities.contains { $0.label == .person && $0.normalizedValue == "priya" }) + #expect(enriched.facetTags.isEmpty) + } + @Test func capturePreviewAndIngestUseSubjectAwareEvidence() async throws { let root = try makeTemporaryDirectory() diff --git a/Tests/MemoryTests/TestSupport.swift b/Tests/MemoryTests/TestSupport.swift index ad380af..a7bdb99 100644 --- a/Tests/MemoryTests/TestSupport.swift +++ b/Tests/MemoryTests/TestSupport.swift @@ -181,6 +181,26 @@ actor StaticContentTagger: ContentTagger { } } +struct StaticMemoryExtractor: MemoryExtractor { + let identifier = "static-memory-extractor" + let result: MemoryExtractionResult + + func extract(messages: [ConversationMessage], limit: Int) async throws -> MemoryExtractionResult { + var limited = result + limited.candidates = Array(result.candidates.prefix(limit)) + return limited + } +} + +struct StaticEntityTagger: EntityTagger { + let identifier = "static-entity-tagger" + let entities: [MemoryEntity] + + func recognizeEntities(in text: String) -> [MemoryEntity] { + entities + } +} + func makeTemporaryDirectory(function: String = #function) throws -> URL { let root = FileManager.default.temporaryDirectory .appendingPathComponent("memory-tests")