From 03e4768732da91128a9da9581fc9e701b33aaa89 Mon Sep 17 00:00:00 2001 From: Tarun Sukhani Date: Mon, 9 Feb 2026 22:19:00 +0800 Subject: [PATCH] =?UTF-8?q?memory-neo4j:=20code=20review=20fixes=20?= =?UTF-8?q?=E2=80=94=20search,=20decay,=20dedup,=20retry,=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Search: fix entity classification order (proper nouns before word count), BM25 min-max normalization with floor, empty query guard. Decay: retrieval-reinforced half-life with effective age anchored to lastRetrievedAt, parameterized category curves (no string interpolation). Dedup: transfer TAGGED relationships to survivor during merge. Orphans: use EXISTS pattern instead of stale mentionCount. Embeddings: Ollama retry with exponential backoff (2 retries, 1s base). Config: resolve env vars in neo4j.uri, re-export MemoryCategory from schema. Extractor: abort-aware batch delay, anonymize prompt examples. Tests: add 80 tests for index.ts (attention gates, message extraction, wrapper stripping). Full suite: 480 tests across 8 files, all passing. Co-Authored-By: Claude Opus 4.6 --- extensions/memory-neo4j/config.ts | 21 +- extensions/memory-neo4j/embeddings.ts | 22 + extensions/memory-neo4j/extractor.ts | 19 +- extensions/memory-neo4j/index.test.ts | 751 +++++++++++++++++++ extensions/memory-neo4j/index.ts | 5 +- extensions/memory-neo4j/neo4j-client.test.ts | 6 +- extensions/memory-neo4j/neo4j-client.ts | 74 +- extensions/memory-neo4j/search.test.ts | 7 +- extensions/memory-neo4j/search.ts | 47 +- 9 files changed, 884 insertions(+), 68 deletions(-) create mode 100644 extensions/memory-neo4j/index.test.ts diff --git a/extensions/memory-neo4j/config.ts b/extensions/memory-neo4j/config.ts index c9a729034b0..ea068287d65 100644 --- a/extensions/memory-neo4j/config.ts +++ b/extensions/memory-neo4j/config.ts @@ -5,6 +5,12 @@ * Provides runtime parsing with env var resolution and defaults. */ +import type { MemoryCategory } from "./schema.js"; +import { MEMORY_CATEGORIES } from "./schema.js"; + +export type { MemoryCategory }; +export { MEMORY_CATEGORIES }; + export type EmbeddingProvider = "openai" | "ollama"; export type MemoryNeo4jConfig = { @@ -65,17 +71,6 @@ export type ExtractionConfig = { maxRetries: number; }; -export const MEMORY_CATEGORIES = [ - "core", - "preference", - "fact", - "decision", - "entity", - "other", -] as const; - -export type MemoryCategory = (typeof MEMORY_CATEGORIES)[number]; - export const EMBEDDING_DIMENSIONS: Record = { // OpenAI models "text-embedding-3-small": 1536, @@ -231,7 +226,7 @@ export const memoryNeo4jConfigSchema = { if (typeof neo4jRaw.uri !== "string" || !neo4jRaw.uri) { throw new Error("neo4j.uri is required"); } - const neo4jUri = neo4jRaw.uri; + const neo4jUri = resolveEnvVars(neo4jRaw.uri); // Validate URI scheme — must be a valid Neo4j connection protocol const VALID_NEO4J_SCHEMES = [ "bolt://", @@ -362,7 +357,7 @@ export const memoryNeo4jConfigSchema = { return { neo4j: { - uri: neo4jRaw.uri, + uri: neo4jUri, username: neo4jUsername, password: neo4jPassword, }, diff --git a/extensions/memory-neo4j/embeddings.ts b/extensions/memory-neo4j/embeddings.ts index 7328ee9e1ff..53803e3b29a 100644 --- a/extensions/memory-neo4j/embeddings.ts +++ b/extensions/memory-neo4j/embeddings.ts @@ -255,8 +255,30 @@ export class Embeddings { // Timeout for Ollama embedding fetch calls to prevent hanging indefinitely private static readonly EMBED_TIMEOUT_MS = 30_000; + // Retry configuration for transient Ollama errors (model loading, GPU pressure) + private static readonly OLLAMA_MAX_RETRIES = 2; + private static readonly OLLAMA_RETRY_BASE_DELAY_MS = 1000; private async embedOllama(text: string): Promise { + let lastError: unknown; + for (let attempt = 0; attempt <= Embeddings.OLLAMA_MAX_RETRIES; attempt++) { + try { + return await this.fetchOllamaEmbedding(text); + } catch (err) { + lastError = err; + if (attempt < Embeddings.OLLAMA_MAX_RETRIES) { + const delay = Embeddings.OLLAMA_RETRY_BASE_DELAY_MS * Math.pow(2, attempt); + this.logger?.warn?.( + `memory-neo4j: Ollama embedding failed (attempt ${attempt + 1}/${Embeddings.OLLAMA_MAX_RETRIES + 1}), retrying in ${delay}ms: ${String(err)}`, + ); + await new Promise((resolve) => setTimeout(resolve, delay)); + } + } + } + throw lastError; + } + + private async fetchOllamaEmbedding(text: string): Promise { const url = `${this.baseUrl}/api/embed`; const response = await fetch(url, { method: "POST", diff --git a/extensions/memory-neo4j/extractor.ts b/extensions/memory-neo4j/extractor.ts index f1fc26ebe01..871a134f2fe 100644 --- a/extensions/memory-neo4j/extractor.ts +++ b/extensions/memory-neo4j/extractor.ts @@ -38,10 +38,10 @@ Return JSON: { "category": "preference|fact|decision|entity|other", "entities": [ - {"name": "tarun", "type": "person", "aliases": ["boss"], "description": "brief description"} + {"name": "alice", "type": "person", "aliases": ["manager"], "description": "brief description"} ], "relationships": [ - {"source": "tarun", "target": "abundent", "type": "WORKS_AT", "confidence": 0.95} + {"source": "alice", "target": "acme corp", "type": "WORKS_AT", "confidence": 0.95} ], "tags": [ {"name": "neo4j", "category": "technology"} @@ -1073,9 +1073,20 @@ export async function runSleepCycle( } } - // Delay between batches + // Delay between batches (abort-aware) if (hasMore && !abortSignal?.aborted) { - await new Promise((resolve) => setTimeout(resolve, extractionDelayMs)); + await new Promise((resolve) => { + const timer = setTimeout(resolve, extractionDelayMs); + // If abort fires during delay, resolve immediately + abortSignal?.addEventListener( + "abort", + () => { + clearTimeout(timer); + resolve(); + }, + { once: true }, + ); + }); } } } diff --git a/extensions/memory-neo4j/index.test.ts b/extensions/memory-neo4j/index.test.ts new file mode 100644 index 00000000000..e6dd66149d4 --- /dev/null +++ b/extensions/memory-neo4j/index.test.ts @@ -0,0 +1,751 @@ +/** + * Tests for the memory-neo4j plugin entry point. + * + * Covers: + * 1. Attention gates (user and assistant) — re-exported from attention-gate.ts + * 2. Message extraction — extractUserMessages, extractAssistantMessages from message-utils.ts + * 3. Strip wrappers — stripMessageWrappers, stripAssistantWrappers from message-utils.ts + * + * Does NOT test the plugin registration or CLI commands (those require the + * full OpenClaw SDK runtime). Focuses on pure functions and the behavioral + * contracts of the auto-capture pipeline helpers. + */ + +import { describe, it, expect } from "vitest"; +import { passesAttentionGate, passesAssistantAttentionGate } from "./attention-gate.js"; +import { + extractUserMessages, + extractAssistantMessages, + stripMessageWrappers, + stripAssistantWrappers, +} from "./message-utils.js"; + +// ============================================================================ +// Test Helpers +// ============================================================================ + +/** Generate a string of a specific length using a repeating word pattern. */ +function makeText(wordCount: number, word = "lorem"): string { + return Array.from({ length: wordCount }, () => word).join(" "); +} + +/** Generate a string of a specific character length. */ +function makeChars(charCount: number, char = "x"): string { + return char.repeat(charCount); +} + +// ============================================================================ +// passesAttentionGate() — User Attention Gate +// ============================================================================ + +describe("passesAttentionGate", () => { + // ----------------------------------------------------------------------- + // Length bounds + // ----------------------------------------------------------------------- + + describe("length bounds", () => { + it("should reject messages shorter than 30 characters", () => { + expect(passesAttentionGate("too short")).toBe(false); + expect(passesAttentionGate("a".repeat(29))).toBe(false); + }); + + it("should reject messages longer than 2000 characters", () => { + // 2001 chars — exceeds MAX_CAPTURE_CHARS + const longText = makeText(300, "longword"); + expect(longText.length).toBeGreaterThan(2000); + expect(passesAttentionGate(longText)).toBe(false); + }); + + it("should accept messages at exactly 30 characters with sufficient words", () => { + // 30 chars, 5 words: "abcde abcde abcde abcde abcde" = 29 chars (5*5 + 4 spaces) + // Need 30+ chars and 5+ words + const text = "abcdef abcdef abcdef abcdef ab"; + expect(text.length).toBe(30); + expect(text.split(/\s+/).length).toBeGreaterThanOrEqual(5); + expect(passesAttentionGate(text)).toBe(true); + }); + + it("should accept messages at exactly 2000 characters with sufficient words", () => { + // Build exactly 2000 chars: repeated "testing " (8 chars each) = 250 words + // 250 * 8 = 2000, but join adds spaces between (not after last), so 250 * 7 + 249 = 1999 + // Use a padded approach: fill with "testing " then pad to exactly 2000 + const base = "testing ".repeat(249) + "testing"; // 249*8 + 7 = 1999 + const text = base + "s"; // 2000 chars + expect(text.length).toBe(2000); + expect(passesAttentionGate(text)).toBe(true); + }); + }); + + // ----------------------------------------------------------------------- + // Word count + // ----------------------------------------------------------------------- + + describe("word count", () => { + it("should reject messages with fewer than 5 words", () => { + // 4 words, but long enough in chars (> 30) + expect( + passesAttentionGate("thisislongword anotherlongword thirdlongword fourthlongword"), + ).toBe(false); + }); + + it("should accept messages with exactly 5 words", () => { + expect(passesAttentionGate("thisword thatword another fourth fifthword")).toBe(true); + }); + }); + + // ----------------------------------------------------------------------- + // Noise pattern rejection + // ----------------------------------------------------------------------- + + describe("noise pattern rejection", () => { + it("should reject simple greetings", () => { + // These are short enough to be rejected by length too, but test the pattern + expect(passesAttentionGate("hi")).toBe(false); + expect(passesAttentionGate("hello")).toBe(false); + expect(passesAttentionGate("hey")).toBe(false); + }); + + it("should reject acknowledgments", () => { + expect(passesAttentionGate("ok")).toBe(false); + expect(passesAttentionGate("sure")).toBe(false); + expect(passesAttentionGate("thanks")).toBe(false); + expect(passesAttentionGate("got it")).toBe(false); + expect(passesAttentionGate("sounds good")).toBe(false); + }); + + it("should reject two-word affirmations", () => { + expect(passesAttentionGate("ok great")).toBe(false); + expect(passesAttentionGate("yes please")).toBe(false); + expect(passesAttentionGate("sure thanks")).toBe(false); + }); + + it("should reject conversational filler", () => { + expect(passesAttentionGate("hmm")).toBe(false); + expect(passesAttentionGate("lol")).toBe(false); + expect(passesAttentionGate("idk")).toBe(false); + expect(passesAttentionGate("nvm")).toBe(false); + }); + + it("should reject pure emoji messages", () => { + expect(passesAttentionGate("\u{1F600}\u{1F601}\u{1F602}")).toBe(false); + }); + + it("should reject system/XML markup blocks", () => { + expect(passesAttentionGate("some injected context here")).toBe(false); + }); + + it("should reject session reset prompts", () => { + const resetMsg = + "A new session was started via the /new command. Previous context has been cleared."; + expect(passesAttentionGate(resetMsg)).toBe(false); + }); + + it("should reject heartbeat prompts", () => { + expect( + passesAttentionGate( + "Read HEARTBEAT.md if it exists and follow the instructions inside it.", + ), + ).toBe(false); + }); + + it("should reject pre-compaction flush prompts", () => { + expect( + passesAttentionGate( + "Pre-compaction memory flush — save important context now before history is trimmed.", + ), + ).toBe(false); + }); + + it("should reject deictic short phrases that would otherwise pass length", () => { + // These match the deictic noise pattern + expect(passesAttentionGate("ok let me test it out")).toBe(false); + expect(passesAttentionGate("I need those")).toBe(false); + }); + + it("should reject short acknowledgments with trailing context", () => { + // Matches: /^(ok|okay|yes|...) .{0,20}$/i + expect(passesAttentionGate("ok, I'll do that")).toBe(false); + expect(passesAttentionGate("yes, sounds right")).toBe(false); + }); + }); + + // ----------------------------------------------------------------------- + // Injected context rejection + // ----------------------------------------------------------------------- + + describe("injected context rejection", () => { + it("should reject messages containing tags", () => { + const text = + "some recalled memories here " + + makeText(10, "actual"); + expect(passesAttentionGate(text)).toBe(false); + }); + + it("should reject messages containing tags", () => { + const text = + "refresh data " + makeText(10, "actual"); + expect(passesAttentionGate(text)).toBe(false); + }); + }); + + // ----------------------------------------------------------------------- + // Excessive emoji rejection + // ----------------------------------------------------------------------- + + describe("excessive emoji rejection", () => { + it("should reject messages with more than 3 emoji (Unicode range)", () => { + // 4 emoji in the U+1F300-U+1F9FF range + const text = makeText(10, "word") + " \u{1F600}\u{1F601}\u{1F602}\u{1F603}"; + expect(passesAttentionGate(text)).toBe(false); + }); + + it("should accept messages with 3 or fewer emoji", () => { + const text = makeText(10, "testing") + " \u{1F600}\u{1F601}\u{1F602}"; + expect(passesAttentionGate(text)).toBe(true); + }); + }); + + // ----------------------------------------------------------------------- + // Substantive messages that should pass + // ----------------------------------------------------------------------- + + describe("substantive messages", () => { + it("should accept a clear factual statement", () => { + expect(passesAttentionGate("I prefer dark mode for all my code editors and terminals")).toBe( + true, + ); + }); + + it("should accept a preference statement", () => { + expect( + passesAttentionGate( + "My favorite programming language is TypeScript because of its type system", + ), + ).toBe(true); + }); + + it("should accept a decision statement", () => { + expect( + passesAttentionGate( + "We decided to use Neo4j for the knowledge graph instead of PostgreSQL", + ), + ).toBe(true); + }); + + it("should accept a multi-sentence message", () => { + expect( + passesAttentionGate( + "The deployment pipeline uses GitHub Actions. It builds and tests on every push to main.", + ), + ).toBe(true); + }); + + it("should handle leading/trailing whitespace via trimming", () => { + expect( + passesAttentionGate(" I prefer using vitest for testing my TypeScript projects "), + ).toBe(true); + }); + }); +}); + +// ============================================================================ +// passesAssistantAttentionGate() — Assistant Attention Gate +// ============================================================================ + +describe("passesAssistantAttentionGate", () => { + // ----------------------------------------------------------------------- + // Length bounds (stricter than user) + // ----------------------------------------------------------------------- + + describe("length bounds", () => { + it("should reject messages shorter than 30 characters", () => { + expect(passesAssistantAttentionGate("short msg")).toBe(false); + }); + + it("should reject messages longer than 1000 characters", () => { + const longText = makeText(200, "wordword"); + expect(longText.length).toBeGreaterThan(1000); + expect(passesAssistantAttentionGate(longText)).toBe(false); + }); + }); + + // ----------------------------------------------------------------------- + // Word count (higher threshold — 10 words minimum) + // ----------------------------------------------------------------------- + + describe("word count", () => { + it("should reject messages with fewer than 10 words", () => { + // 9 words, each 5 chars + space = more than 30 chars total + const nineWords = "alpha bravo charm delta eerie found ghost horse india"; + expect(nineWords.split(/\s+/).length).toBe(9); + expect(nineWords.length).toBeGreaterThan(30); + expect(passesAssistantAttentionGate(nineWords)).toBe(false); + }); + + it("should accept messages with exactly 10 words", () => { + const tenWords = "alpha bravo charm delta eerie found ghost horse india julep"; + expect(tenWords.split(/\s+/).length).toBe(10); + expect(tenWords.length).toBeGreaterThan(30); + expect(passesAssistantAttentionGate(tenWords)).toBe(true); + }); + }); + + // ----------------------------------------------------------------------- + // Code-heavy message rejection (> 50% fenced code) + // ----------------------------------------------------------------------- + + describe("code-heavy rejection", () => { + it("should reject messages that are more than 50% fenced code blocks", () => { + // ~60 chars of prose + ~200 chars of code block => code > 50% + const text = + "Here is some explanation for the code below that follows.\n" + + "```typescript\n" + + "function example() {\n" + + " const x = 1;\n" + + " const y = 2;\n" + + " return x + y;\n" + + "}\n" + + "function another() {\n" + + " const a = 3;\n" + + " return a * 2;\n" + + "}\n" + + "```"; + expect(passesAssistantAttentionGate(text)).toBe(false); + }); + + it("should accept messages with less than 50% code", () => { + const text = + "The configuration requires setting up the environment variables correctly. " + + "You need to set NEO4J_URI, NEO4J_USER, and NEO4J_PASSWORD. " + + "Make sure the password is at least 8 characters long for security. " + + "```\nNEO4J_URI=bolt://localhost:7687\n```"; + expect(passesAssistantAttentionGate(text)).toBe(true); + }); + }); + + // ----------------------------------------------------------------------- + // Tool output rejection + // ----------------------------------------------------------------------- + + describe("tool output rejection", () => { + it("should reject messages containing tags", () => { + const text = + "Here is the result of the search query across all the relevant documents " + + "some result data here"; + expect(passesAssistantAttentionGate(text)).toBe(false); + }); + + it("should reject messages containing tags", () => { + const text = + "I will use this tool to help answer your question about the system setup " + + "tool invocation here"; + expect(passesAssistantAttentionGate(text)).toBe(false); + }); + + it("should reject messages containing tags", () => { + const text = + "Calling the function to retrieve the relevant data from the database now " + + "fn call here"; + expect(passesAssistantAttentionGate(text)).toBe(false); + }); + }); + + // ----------------------------------------------------------------------- + // Injected context rejection + // ----------------------------------------------------------------------- + + describe("injected context rejection", () => { + it("should reject messages with tags", () => { + const text = + "cached recall data " + makeText(15, "answer"); + expect(passesAssistantAttentionGate(text)).toBe(false); + }); + + it("should reject messages with tags", () => { + const text = + "identity refresh " + makeText(15, "answer"); + expect(passesAssistantAttentionGate(text)).toBe(false); + }); + }); + + // ----------------------------------------------------------------------- + // Noise patterns and emoji (shared with user gate) + // ----------------------------------------------------------------------- + + describe("noise patterns", () => { + it("should reject greeting noise", () => { + expect(passesAssistantAttentionGate("hello")).toBe(false); + }); + + it("should reject excessive emoji", () => { + const text = makeText(15, "answer") + " \u{1F600}\u{1F601}\u{1F602}\u{1F603}"; + expect(passesAssistantAttentionGate(text)).toBe(false); + }); + }); + + // ----------------------------------------------------------------------- + // Substantive assistant messages that should pass + // ----------------------------------------------------------------------- + + describe("substantive assistant messages", () => { + it("should accept a clear explanatory response", () => { + expect( + passesAssistantAttentionGate( + "The Neo4j database uses a property graph model where nodes represent entities and edges represent relationships between them.", + ), + ).toBe(true); + }); + + it("should accept a recommendation response", () => { + expect( + passesAssistantAttentionGate( + "Based on your requirements, I recommend using vitest for unit testing because it has native TypeScript support and fast execution times.", + ), + ).toBe(true); + }); + }); +}); + +// ============================================================================ +// extractUserMessages() +// ============================================================================ + +describe("extractUserMessages", () => { + it("should extract text from string content format", () => { + const messages = [{ role: "user", content: "This is a substantive user message for testing" }]; + const result = extractUserMessages(messages); + expect(result).toEqual(["This is a substantive user message for testing"]); + }); + + it("should extract text from content block array format", () => { + const messages = [ + { + role: "user", + content: [{ type: "text", text: "This is a substantive user message from a block array" }], + }, + ]; + const result = extractUserMessages(messages); + expect(result).toEqual(["This is a substantive user message from a block array"]); + }); + + it("should extract multiple text blocks from a single message", () => { + const messages = [ + { + role: "user", + content: [ + { type: "text", text: "First text block with enough characters" }, + { type: "image", url: "http://example.com/img.png" }, + { type: "text", text: "Second text block with enough characters" }, + ], + }, + ]; + const result = extractUserMessages(messages); + expect(result).toHaveLength(2); + expect(result[0]).toBe("First text block with enough characters"); + expect(result[1]).toBe("Second text block with enough characters"); + }); + + it("should ignore non-user messages", () => { + const messages = [ + { role: "assistant", content: "I am the assistant response message here" }, + { role: "system", content: "This is the system prompt configuration text" }, + { role: "user", content: "This is the actual user message text here" }, + ]; + const result = extractUserMessages(messages); + expect(result).toEqual(["This is the actual user message text here"]); + }); + + it("should filter out messages shorter than 10 characters after stripping", () => { + const messages = [ + { role: "user", content: "short" }, + { role: "user", content: "This is a long enough message to pass the filter" }, + ]; + const result = extractUserMessages(messages); + expect(result).toHaveLength(1); + expect(result[0]).toBe("This is a long enough message to pass the filter"); + }); + + it("should strip Telegram wrappers before returning", () => { + const messages = [ + { + role: "user", + content: + "[Telegram @user123 in group] The actual user message is right here\n[message_id: 456]", + }, + ]; + const result = extractUserMessages(messages); + expect(result).toEqual(["The actual user message is right here"]); + }); + + it("should strip Slack wrappers before returning", () => { + const messages = [ + { + role: "user", + content: + "[Slack workspace #channel @user] The actual user message text goes here\n[slack message id: abc123]", + }, + ]; + const result = extractUserMessages(messages); + expect(result).toEqual(["The actual user message text goes here"]); + }); + + it("should strip injected context", () => { + const messages = [ + { + role: "user", + content: + "recalled: user likes dark mode What editor do you recommend for me?", + }, + ]; + const result = extractUserMessages(messages); + expect(result).toEqual(["What editor do you recommend for me?"]); + }); + + it("should handle null and non-object entries gracefully", () => { + const messages = [ + null, + undefined, + 42, + "string", + { role: "user", content: "This is a valid message with enough text" }, + ]; + const result = extractUserMessages(messages as unknown[]); + expect(result).toEqual(["This is a valid message with enough text"]); + }); + + it("should handle empty messages array", () => { + expect(extractUserMessages([])).toEqual([]); + }); + + it("should ignore content blocks that are not type 'text'", () => { + const messages = [ + { + role: "user", + content: [ + { type: "image", url: "http://example.com/photo.jpg" }, + { type: "audio", data: "base64data..." }, + ], + }, + ]; + const result = extractUserMessages(messages); + expect(result).toEqual([]); + }); +}); + +// ============================================================================ +// extractAssistantMessages() +// ============================================================================ + +describe("extractAssistantMessages", () => { + it("should extract text from string content format", () => { + const messages = [ + { role: "assistant", content: "Here is a substantive assistant response text" }, + ]; + const result = extractAssistantMessages(messages); + expect(result).toEqual(["Here is a substantive assistant response text"]); + }); + + it("should extract text from content block array format", () => { + const messages = [ + { + role: "assistant", + content: [{ type: "text", text: "The assistant provides an answer to your question here" }], + }, + ]; + const result = extractAssistantMessages(messages); + expect(result).toEqual(["The assistant provides an answer to your question here"]); + }); + + it("should ignore non-assistant messages", () => { + const messages = [ + { role: "user", content: "This is a user message that should be ignored" }, + { role: "assistant", content: "This is the assistant response message here" }, + ]; + const result = extractAssistantMessages(messages); + expect(result).toEqual(["This is the assistant response message here"]); + }); + + it("should filter out messages shorter than 10 characters after stripping", () => { + const messages = [ + { role: "assistant", content: "short" }, + { role: "assistant", content: "This is a long enough assistant response message" }, + ]; + const result = extractAssistantMessages(messages); + expect(result).toHaveLength(1); + expect(result[0]).toBe("This is a long enough assistant response message"); + }); + + it("should strip tool-use blocks from assistant messages", () => { + const messages = [ + { + role: "assistant", + content: + "search function call parametersHere is the answer to your question about configuration", + }, + ]; + const result = extractAssistantMessages(messages); + expect(result).toEqual(["Here is the answer to your question about configuration"]); + }); + + it("should strip tool_result blocks from assistant messages", () => { + const messages = [ + { + role: "assistant", + content: + "The query returned: raw database output here which means the config is correct and working.", + }, + ]; + const result = extractAssistantMessages(messages); + expect(result).toEqual(["The query returned: which means the config is correct and working."]); + }); + + it("should strip thinking blocks from assistant messages", () => { + const messages = [ + { + role: "assistant", + content: + "I need to figure out the best approach hereThe best approach is to use a hybrid search combining vector and BM25 signals.", + }, + ]; + const result = extractAssistantMessages(messages); + expect(result).toEqual([ + "The best approach is to use a hybrid search combining vector and BM25 signals.", + ]); + }); + + it("should strip code_output blocks from assistant messages", () => { + const messages = [ + { + role: "assistant", + content: + "I ran the code: stdout: success and it completed without any errors at all.", + }, + ]; + const result = extractAssistantMessages(messages); + expect(result).toEqual(["I ran the code: and it completed without any errors at all."]); + }); + + it("should handle null and non-object entries gracefully", () => { + const messages = [ + null, + undefined, + { role: "assistant", content: "This is a valid assistant response text" }, + ]; + const result = extractAssistantMessages(messages as unknown[]); + expect(result).toEqual(["This is a valid assistant response text"]); + }); + + it("should handle empty messages array", () => { + expect(extractAssistantMessages([])).toEqual([]); + }); +}); + +// ============================================================================ +// stripMessageWrappers() +// ============================================================================ + +describe("stripMessageWrappers", () => { + it("should strip tags and content", () => { + const input = + "user likes dark mode What editor should I use?"; + expect(stripMessageWrappers(input)).toBe("What editor should I use?"); + }); + + it("should strip tags and content", () => { + const input = + "identity: Tarun How do I configure this?"; + expect(stripMessageWrappers(input)).toBe("How do I configure this?"); + }); + + it("should strip tags and content", () => { + const input = "You are a helpful assistant. What is the weather?"; + expect(stripMessageWrappers(input)).toBe("What is the weather?"); + }); + + it("should strip attachment tags", () => { + const input = 'base64content Summarize this document for me.'; + expect(stripMessageWrappers(input)).toBe("Summarize this document for me."); + }); + + it("should strip Telegram wrapper and message_id", () => { + const input = "[Telegram @john in private] Please remember my preference\n[message_id: 12345]"; + expect(stripMessageWrappers(input)).toBe("Please remember my preference"); + }); + + it("should strip Slack wrapper and slack message id", () => { + const input = + "[Slack acme-corp #general @alice] Please deploy the latest build\n[slack message id: ts-123]"; + expect(stripMessageWrappers(input)).toBe("Please deploy the latest build"); + }); + + it("should strip media attachment preamble", () => { + const input = + "[media attached: image/jpeg]\nTo send an image reply with...\n[Telegram @user in private] What is this picture?"; + expect(stripMessageWrappers(input)).toBe("What is this picture?"); + }); + + it("should strip System exec output blocks before Telegram wrapper", () => { + const input = + "System: [2024-01-01] exec completed\n[Telegram @user in private] What happened with the deploy?"; + expect(stripMessageWrappers(input)).toBe("What happened with the deploy?"); + }); + + it("should handle multiple wrappers in one message", () => { + const input = + "recalled facts You are helpful. [Telegram @user in group] What is up?"; + const result = stripMessageWrappers(input); + expect(result).toBe("What is up?"); + }); + + it("should return trimmed text when no wrappers are present", () => { + expect(stripMessageWrappers(" Just a plain message ")).toBe("Just a plain message"); + }); +}); + +// ============================================================================ +// stripAssistantWrappers() +// ============================================================================ + +describe("stripAssistantWrappers", () => { + it("should strip blocks", () => { + const input = "call searchThe answer is 42."; + expect(stripAssistantWrappers(input)).toBe("The answer is 42."); + }); + + it("should strip blocks", () => { + const input = "Result: raw output processed successfully."; + // The regex consumes trailing whitespace after the closing tag + expect(stripAssistantWrappers(input)).toBe("Result: processed successfully."); + }); + + it("should strip blocks", () => { + const input = "fn(args)Done with the operation."; + expect(stripAssistantWrappers(input)).toBe("Done with the operation."); + }); + + it("should strip blocks", () => { + const input = "Let me consider...I recommend using vitest."; + expect(stripAssistantWrappers(input)).toBe("I recommend using vitest."); + }); + + it("should strip blocks", () => { + const input = "analyzing the requestHere is the analysis."; + expect(stripAssistantWrappers(input)).toBe("Here is the analysis."); + }); + + it("should strip blocks", () => { + const input = "Output: success everything worked."; + // The regex consumes trailing whitespace after the closing tag + expect(stripAssistantWrappers(input)).toBe("Output: everything worked."); + }); + + it("should strip multiple wrapper types in one message", () => { + const input = + "hmmsearchThe final answer is here.data"; + expect(stripAssistantWrappers(input)).toBe("The final answer is here."); + }); + + it("should return trimmed text when no wrappers are present", () => { + expect(stripAssistantWrappers(" Plain assistant text ")).toBe("Plain assistant text"); + }); +}); diff --git a/extensions/memory-neo4j/index.ts b/extensions/memory-neo4j/index.ts index ac30a0a7012..46411a06ea5 100644 --- a/extensions/memory-neo4j/index.ts +++ b/extensions/memory-neo4j/index.ts @@ -426,14 +426,15 @@ const memoryNeo4jPlugin = { .description("Search memories") .argument("", "Search query") .option("--limit ", "Max results", "5") - .action(async (query: string, opts: { limit: string }) => { + .option("--agent ", "Agent id (default: default)") + .action(async (query: string, opts: { limit: string; agent?: string }) => { try { const results = await hybridSearch( db, embeddings, query, parseInt(opts.limit, 10), - "default", + opts.agent ?? "default", extractionConfig.enabled, { graphSearchDepth: cfg.graphSearchDepth }, ); diff --git a/extensions/memory-neo4j/neo4j-client.test.ts b/extensions/memory-neo4j/neo4j-client.test.ts index 8f5e59115cf..81fc4b06e79 100644 --- a/extensions/memory-neo4j/neo4j-client.test.ts +++ b/extensions/memory-neo4j/neo4j-client.test.ts @@ -1255,9 +1255,9 @@ describe("Neo4jMemoryClient", () => { it("should count memories by extraction status", async () => { mockSession.run.mockResolvedValue({ records: [ - { get: vi.fn((key) => (key === "status" ? "pending" : { toNumber: () => 5 })) }, - { get: vi.fn((key) => (key === "status" ? "complete" : { toNumber: () => 10 })) }, - { get: vi.fn((key) => (key === "status" ? "failed" : { toNumber: () => 2 })) }, + { get: vi.fn((key) => (key === "status" ? "pending" : 5)) }, + { get: vi.fn((key) => (key === "status" ? "complete" : 10)) }, + { get: vi.fn((key) => (key === "status" ? "failed" : 2)) }, ], }); diff --git a/extensions/memory-neo4j/neo4j-client.ts b/extensions/memory-neo4j/neo4j-client.ts index f3a3814d51f..484d762e3d5 100644 --- a/extensions/memory-neo4j/neo4j-client.ts +++ b/extensions/memory-neo4j/neo4j-client.ts @@ -4,7 +4,7 @@ * Handles connection management, index creation, CRUD operations, * and the three search signals (vector, BM25, graph). * - * Patterns adapted from ~/Downloads/ontology/app/services/neo4j_client.py + * Patterns adapted from ontology project Neo4j client * with retry-on-transient and MERGE idempotency. */ @@ -160,6 +160,11 @@ export class Neo4jMemoryClient { session, "CREATE INDEX memory_agent_category_index IF NOT EXISTS FOR (m:Memory) ON (m.agentId, m.category)", ); + // Extraction status index for listPendingExtractions (sleep cycle) + await this.runSafe( + session, + "CREATE INDEX memory_extraction_status_index IF NOT EXISTS FOR (m:Memory) ON (m.extractionStatus)", + ); this.logger.info("memory-neo4j: indexes ensured"); } finally { @@ -488,10 +493,15 @@ export class Neo4jMemoryClient { if (records.length === 0) { return []; } - const maxScore = records[0].rawScore || 1; + // Min-max normalization with a floor: prevents a single weak BM25 + // match from getting score 1.0 and inflating its RRF contribution. + const maxScore = records[0].rawScore; + const minScore = records[records.length - 1].rawScore; + const range = maxScore - minScore; + const FLOOR = 0.3; // Minimum normalized score for the lowest-ranked result return records.map((r) => ({ ...r, - score: r.rawScore / maxScore, + score: range > 0 ? FLOOR + ((1 - FLOOR) * (r.rawScore - minScore)) / range : 1.0, // All scores identical → all get 1.0 })); } finally { await session.close(); @@ -1082,7 +1092,7 @@ export class Neo4jMemoryClient { }; for (const record of result.records) { const status = record.get("status") as string; - const count = (record.get("count") as { toNumber?: () => number })?.toNumber?.() ?? 0; + const count = (record.get("count") as number) ?? 0; if (status in counts) { counts[status] = count; } @@ -1335,6 +1345,16 @@ export class Neo4jMemoryClient { { toDelete, survivorId }, ); + // Transfer TAGGED relationships from deleted memories to survivor + await tx.run( + `UNWIND $toDelete AS deadId + MATCH (dead:Memory {id: deadId})-[r:TAGGED]->(t:Tag) + MATCH (survivor:Memory {id: $survivorId}) + MERGE (survivor)-[:TAGGED]->(t) + DELETE r`, + { toDelete, survivorId }, + ); + // Delete the duplicate memories await tx.run( `UNWIND $toDelete AS deadId @@ -1398,18 +1418,25 @@ export class Neo4jMemoryClient { try { const agentFilter = agentId ? "AND m.agentId = $agentId" : ""; - // Build per-category half-life CASE expression if curves are configured - let halfLifeExpr = "$baseHalfLife"; - if (decayCurves && Object.keys(decayCurves).length > 0) { - const cases = Object.entries(decayCurves) - .map( - ([cat, { halfLifeDays }]) => - `WHEN m.category = '${cat.replace(/'/g, "\\'")}' THEN ${halfLifeDays}`, - ) - .join(" "); - halfLifeExpr = `CASE ${cases} ELSE $baseHalfLife END`; + // Build per-category half-life using parameterized map lookup instead of + // string interpolation, avoiding any injection risk from category names. + const curveEntries = decayCurves ? Object.entries(decayCurves) : []; + const hasCurves = curveEntries.length > 0; + + // Pass category→halfLife mapping as a Cypher map parameter + const curveMap: Record = {}; + for (const [cat, { halfLifeDays }] of curveEntries) { + curveMap[cat] = halfLifeDays; } + const halfLifeExpr = hasCurves + ? "CASE WHEN $curveMap[m.category] IS NOT NULL THEN $curveMap[m.category] ELSE $baseHalfLife END" + : "$baseHalfLife"; + + // Decay formula uses retrieval reinforcement: memories that are frequently + // accessed decay slower. The effective age is anchored to the most recent + // of createdAt or lastRetrievedAt, so recently recalled memories get a + // recency boost even if they were created long ago. const result = await session.run( `MATCH (m:Memory) WHERE m.createdAt IS NOT NULL @@ -1417,11 +1444,17 @@ export class Neo4jMemoryClient { ${agentFilter} WITH m, duration.between(datetime(m.createdAt), datetime()).days AS ageDays, - m.importance AS importance - WITH m, ageDays, importance, - ${halfLifeExpr} * (1.0 + importance * $importanceMult) AS halfLife + CASE + WHEN m.lastRetrievedAt IS NOT NULL + THEN duration.between(datetime(m.lastRetrievedAt), datetime()).days + ELSE duration.between(datetime(m.createdAt), datetime()).days + END AS effectiveAgeDays, + m.importance AS importance, + coalesce(m.retrievalCount, 0) AS retrievalCount + WITH m, ageDays, effectiveAgeDays, importance, retrievalCount, + ${halfLifeExpr} * (1.0 + importance * $importanceMult) * (1.0 + log(1.0 + retrievalCount) * 0.2) AS halfLife WITH m, ageDays, importance, halfLife, - importance * exp(-1.0 * ageDays / halfLife) AS decayScore + importance * exp(-1.0 * effectiveAgeDays / halfLife) AS decayScore WHERE decayScore < $threshold RETURN m.id AS id, m.text AS text, importance, ageDays, decayScore ORDER BY decayScore ASC @@ -1430,6 +1463,7 @@ export class Neo4jMemoryClient { threshold: retentionThreshold, baseHalfLife: baseHalfLifeDays, importanceMult: importanceMultiplier, + curveMap, agentId, limit: neo4j.int(limit), }, @@ -1490,9 +1524,11 @@ export class Neo4jMemoryClient { await this.ensureInitialized(); const session = this.driver!.session(); try { + // Use EXISTS check as the authoritative source — mentionCount can go + // stale if crashes occur between decrement and delete operations. const result = await session.run( `MATCH (e:Entity) - WHERE coalesce(e.mentionCount, 0) <= 0 + WHERE NOT EXISTS { MATCH (:Memory)-[:MENTIONS]->(e) } RETURN e.id AS id, e.name AS name, e.type AS type LIMIT $limit`, { limit: neo4j.int(limit) }, diff --git a/extensions/memory-neo4j/search.test.ts b/extensions/memory-neo4j/search.test.ts index f3ff09a53f0..46b6822dad5 100644 --- a/extensions/memory-neo4j/search.test.ts +++ b/extensions/memory-neo4j/search.test.ts @@ -25,16 +25,15 @@ describe("classifyQuery", () => { expect(classifyQuery("best coffee")).toBe("short"); }); - it("should classify a single capitalized word as 'short' (word count takes priority)", () => { - expect(classifyQuery("TypeScript")).toBe("short"); - }); - it("should handle whitespace-padded short queries", () => { expect(classifyQuery(" hello ")).toBe("short"); }); }); describe("entity queries (proper nouns)", () => { + it("should classify a single capitalized word as 'entity' (proper noun detection)", () => { + expect(classifyQuery("TypeScript")).toBe("entity"); + }); it("should classify query with proper noun as 'entity'", () => { expect(classifyQuery("tell me about Tarun")).toBe("entity"); }); diff --git a/extensions/memory-neo4j/search.ts b/extensions/memory-neo4j/search.ts index 528c8015847..113b2c197ec 100644 --- a/extensions/memory-neo4j/search.ts +++ b/extensions/memory-neo4j/search.ts @@ -9,7 +9,7 @@ * Fused using confidence-weighted Reciprocal Rank Fusion (RRF) * with query-adaptive signal weights. * - * Adapted from ~/Downloads/ontology/app/services/rrf.py + * Adapted from ontology project RRF implementation. */ import type { Embeddings } from "./embeddings.js"; @@ -34,35 +34,32 @@ export function classifyQuery(query: string): QueryType { const words = query.trim().split(/\s+/); const wordCount = words.length; - // Short queries: 1-2 words → boost BM25 - if (wordCount <= 2) { - return "short"; - } - - // Long queries: 5+ words → boost vector - if (wordCount >= 5) { - return "long"; - } - // Entity detection: check for capitalized words (proper nouns) - // Heuristic: if more than half of non-first words are capitalized - const capitalizedWords = words - .slice(1) // skip first word (often capitalized anyway) - .filter( - (w) => - /^[A-Z]/.test(w) && - !/^(I|A|An|The|Is|Are|Was|Were|What|Who|Where|When|How|Why|Do|Does|Did)$/.test(w), - ); + // Runs before word count so "John" or "TypeScript" are classified as entity + const commonWords = + /^(I|A|An|The|Is|Are|Was|Were|What|Who|Where|When|How|Why|Do|Does|Did|Find|Show|Get|Tell|Me|My|About|For)$/; + const capitalizedWords = words.filter((w) => /^[A-Z]/.test(w) && !commonWords.test(w)); if (capitalizedWords.length > 0) { return "entity"; } - // Check for question patterns targeting entities - if (/^(who|where|what)\s+(is|does|did|was|were)\s/i.test(query)) { + // Short queries: 1-2 words → boost BM25 + if (wordCount <= 2) { + return "short"; + } + + // Question patterns targeting entities (3-4 word queries only, + // so generic long questions like "what is the best framework" fall through to "long") + if (wordCount <= 4 && /^(who|where|what)\s+(is|does|did|was|were)\s/i.test(query)) { return "entity"; } + // Long queries: 5+ words → boost vector + if (wordCount >= 5) { + return "long"; + } + return "default"; } @@ -121,8 +118,7 @@ type FusedCandidate = { * score magnitude: rank-1 with score 0.99 contributes more than * rank-1 with score 0.55. * - * Reference: Cormack et al. (2009), extended with confidence weighting - * from ~/Downloads/ontology/app/services/rrf.py + * Reference: Cormack et al. (2009), extended with confidence weighting. */ function fuseWithConfidenceRRF( signals: SearchSignalResult[][], @@ -220,6 +216,11 @@ export async function hybridSearch( graphSearchDepth?: number; } = {}, ): Promise { + // Guard against empty queries + if (!query.trim()) { + return []; + } + const { rrfK = 60, candidateMultiplier = 4,