diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..ca1376f --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,67 @@ +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: | + 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 + 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: | + 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 + 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..691224a --- /dev/null +++ b/.github/workflows/evals.yml @@ -0,0 +1,55 @@ +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: | + 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 + 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 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/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])) } diff --git a/README.md b/README.md index 578401e..7a43d5a 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,12 @@ 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), 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: ```swift diff --git a/Sources/AgentMemory/MemoryExtractionHeuristics.swift b/Sources/AgentMemory/MemoryExtractionHeuristics.swift index 11b113e..5179299 100644 --- a/Sources/AgentMemory/MemoryExtractionHeuristics.swift +++ b/Sources/AgentMemory/MemoryExtractionHeuristics.swift @@ -75,12 +75,12 @@ 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) ) + 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 +155,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"]) @@ -521,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 } @@ -557,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 ", @@ -588,7 +617,10 @@ 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 + ) -> [FocusedProfileSegment] { guard role == .user else { return [] } var segments: [FocusedProfileSegment] = [] @@ -624,13 +656,104 @@ internal enum MemoryExtractionHeuristics { 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"). + /// 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] { + let inferred = inferEntities(forExtractedText: text) + var merged = mergeEntities(pinned: supplied, inferred: inferred) + guard !tagged.isEmpty else { return merged } + + 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 } + guard isSupportedTaggedEntity(entity, in: text) 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.contains(where: { valuesOverlap($0.normalizedValue, normalizedValue) }) { + merged.append( + MemoryEntity( + label: entity.label, + value: entity.value, + normalizedValue: normalizedValue, + confidence: entity.confidence + ) + ) + indexByValue[normalizedValue] = merged.count - 1 + } + } + + return merged + } + + 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", @@ -649,7 +772,7 @@ 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" ] private static let locationTrailingNoiseTokens: Set = [ @@ -664,11 +787,7 @@ internal enum MemoryExtractionHeuristics { private static func selfReportedLocation(in text: String) -> String? { if let raw = firstRegexCapture(of: strongLocationTriggerPattern, in: text), - let location = parseLocationPhrase(raw, requireKnownLocation: false) { - return location - } - if let raw = firstRegexCapture(of: weakLocationTriggerPattern, in: text), - let location = parseLocationPhrase(raw, requireKnownLocation: true) { + let location = parseLocationPhrase(raw) { return location } return nil @@ -686,7 +805,7 @@ 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) -> String? { let tokens = raw.split(whereSeparator: \.isWhitespace).map(String.init) var cityTokens: [String] = [] var sawComma = false @@ -748,7 +867,7 @@ internal enum MemoryExtractionHeuristics { } } - guard !requireKnownLocation, city.count >= 2 else { return nil } + guard city.count >= 2 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..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 { @@ -1965,6 +1968,25 @@ public actor MemoryIndex { ) } + /// 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) } @@ -2077,7 +2099,11 @@ public actor MemoryIndex { } private func normalizeEntities(_ supplied: [MemoryEntity], text: String) -> [MemoryEntity] { - let preferred = supplied.isEmpty ? inferEntities(forExtractedText: text) : 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/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..7cabffa --- /dev/null +++ b/Sources/AgentMemory/NLEntityTagger.swift @@ -0,0 +1,69 @@ +#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. 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 = [] + let tagger = NLTagger(tagSchemes: [.nameType]) + tagger.string = trimmed + tagger.enumerateTags( + in: trimmed.startIndex.. EntityLabel? { + switch tag { + case .personalName: + .person + case .placeName: + .location + case .organizationName: + .organization + default: + nil + } + } +} +#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 ad22617..13c731b 100644 --- a/Sources/AgentMemory/ProtocolsAndErrors.swift +++ b/Sources/AgentMemory/ProtocolsAndErrors.swift @@ -231,6 +231,16 @@ 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. + /// 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] +} + public struct MemoryConfiguration: Sendable { public var databaseURL: URL public var embeddingProvider: any EmbeddingProvider @@ -249,6 +259,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 +282,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 +302,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/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( diff --git a/Tests/MemoryTests/MemoryExternalAPITests.swift b/Tests/MemoryTests/MemoryExternalAPITests.swift index 5111945..479451a 100644 --- a/Tests/MemoryTests/MemoryExternalAPITests.swift +++ b/Tests/MemoryTests/MemoryExternalAPITests.swift @@ -889,6 +889,192 @@ struct MemoryExternalAPITests { #expect(maintenance.consideredSignalCount == 1) } + @Test + func ambiguousLocationDoesNotBecomeResidenceWithoutTagger() 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 nlEntityTaggerRecognizesNamedEntitiesConservatively() { + let tagger = NLEntityTagger() + + 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 heuristicExtractDoesNotRewritePresenceAsResidence() 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?"), + 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 + ) + + #expect(extracted.isEmpty) + } + + @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" + }) + } + + @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")