mirror of
https://github.com/supermemoryai/supermemory.git
synced 2026-10-01 02:01:40 +00:00
feat(ai-sdk): introduce Working Memory layer for low-latency local context
RFC #1625 — V1: explicit decorator (Staff review 7.8/10) Fixes #1625 V1 scope (mergeable): LRU, TTL, promise dedup, invalidate, clear, stats. Deferred: pin/unpin, auto-populate, _source mutation, persistent cache, cross-process sharing. API: createWorkingMemory(tools, { ttlMs, maxEntries }) wraps searchMemories without mutating return shape. Base supermemoryTools behavior unchanged — zero breaking change. Non-goals (V1): persistent disk, cross-process, semantic embedding cache, write caching, background refresh, cross-device sync. Co-authored-by: Sravanjangam <163002695+Sravanjangam@users.noreply.github.com>
This commit is contained in:
parent
d436792e77
commit
9ddd057b8c
3 changed files with 402 additions and 0 deletions
|
|
@ -1 +1,2 @@
|
|||
export * from "./tools"
|
||||
export * from "./working-memory"
|
||||
|
|
|
|||
165
packages/ai-sdk/src/working-memory.test.ts
Normal file
165
packages/ai-sdk/src/working-memory.test.ts
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
import { describe, expect, it, vi } from "vitest"
|
||||
import { createWorkingMemory, WorkingMemory } from "./working-memory"
|
||||
|
||||
describe("WorkingMemory — V1 (LRU, TTL, dedup, invalidate)", () => {
|
||||
it("cache hit on second get within TTL", () => {
|
||||
const wm = new WorkingMemory({ ttlMs: 60_000, maxEntries: 100 })
|
||||
wm.set("preferences", [{ id: 1 }])
|
||||
expect(wm.get("preferences")).toEqual([{ id: 1 }])
|
||||
expect(wm.get("preferences")).toEqual([{ id: 1 }])
|
||||
expect(wm.stats().hits).toBe(2)
|
||||
})
|
||||
|
||||
it("cache miss on unknown query", () => {
|
||||
const wm = new WorkingMemory({ ttlMs: 60_000 })
|
||||
expect(wm.get("unknown")).toBeUndefined()
|
||||
expect(wm.stats().misses).toBe(1)
|
||||
})
|
||||
|
||||
it("TTL expiry triggers miss (fake timers)", async () => {
|
||||
vi.useFakeTimers()
|
||||
const wm = new WorkingMemory({ ttlMs: 50 })
|
||||
wm.set("q", [{ id: 1 }])
|
||||
expect(wm.get("q")).toEqual([{ id: 1 }])
|
||||
vi.advanceTimersByTime(80)
|
||||
expect(wm.get("q")).toBeUndefined()
|
||||
vi.useRealTimers()
|
||||
})
|
||||
|
||||
it("promise dedup: 20 concurrent search share one fetch", async () => {
|
||||
const wm = new WorkingMemory<string>({ ttlMs: 60_000 })
|
||||
let calls = 0
|
||||
const fetchFn = async () => {
|
||||
calls++
|
||||
await new Promise((r) => setTimeout(r, 20))
|
||||
return ["result"]
|
||||
}
|
||||
const ps = Array.from({ length: 20 }, () =>
|
||||
wm.search("preferences", fetchFn),
|
||||
)
|
||||
const rs = await Promise.all(ps)
|
||||
expect(calls).toBe(1)
|
||||
expect(rs.every((r) => r[0] === "result")).toBe(true)
|
||||
// Subsequent get is a cache hit, no fetch
|
||||
expect(wm.get("preferences")).toEqual(["result"])
|
||||
})
|
||||
|
||||
it("LRU eviction drops oldest", () => {
|
||||
const wm = new WorkingMemory({ ttlMs: 60_000, maxEntries: 3 })
|
||||
wm.set("1", [{ id: 1 }])
|
||||
wm.set("2", [{ id: 2 }])
|
||||
wm.set("3", [{ id: 3 }])
|
||||
wm.get("1") // touch 1 -> order 2,3,1
|
||||
wm.set("4", [{ id: 4 }]) // evict 2
|
||||
expect(wm.has("2")).toBe(false)
|
||||
expect(wm.has("1")).toBe(true)
|
||||
expect(wm.has("3")).toBe(true)
|
||||
expect(wm.has("4")).toBe(true)
|
||||
})
|
||||
|
||||
it("invalidate single vs clear all", () => {
|
||||
const wm = new WorkingMemory({ ttlMs: 60_000 })
|
||||
wm.set("x", [{ id: 1 }])
|
||||
wm.set("y", [{ id: 2 }])
|
||||
wm.invalidate("x")
|
||||
expect(wm.has("x")).toBe(false)
|
||||
expect(wm.has("y")).toBe(true)
|
||||
wm.invalidate()
|
||||
expect(wm.has("y")).toBe(false)
|
||||
expect(wm.stats().size).toBe(0)
|
||||
})
|
||||
|
||||
it("stats tracks hits/misses/size", () => {
|
||||
const wm = new WorkingMemory({ ttlMs: 60_000 })
|
||||
wm.set("a", [{ id: 1 }])
|
||||
wm.get("a") // hit
|
||||
wm.get("missing") // miss
|
||||
const s = wm.stats()
|
||||
expect(s.hits).toBe(1)
|
||||
expect(s.misses).toBe(1)
|
||||
expect(s.size).toBe(1)
|
||||
})
|
||||
|
||||
it("cache disabled: without wrapper every call hits fetch", async () => {
|
||||
const wm = new WorkingMemory({ ttlMs: 60_000 })
|
||||
let calls = 0
|
||||
const fetchFn = async () => {
|
||||
calls++
|
||||
return [{ id: calls }]
|
||||
}
|
||||
// Without using wm.search, direct fetch always calls
|
||||
await fetchFn()
|
||||
await fetchFn()
|
||||
expect(calls).toBe(2)
|
||||
// With wm.search, second is cached
|
||||
calls = 0
|
||||
await wm.search("q", fetchFn)
|
||||
await wm.search("q", fetchFn)
|
||||
expect(calls).toBe(1)
|
||||
})
|
||||
})
|
||||
|
||||
describe("createWorkingMemory — explicit decorator", () => {
|
||||
it("wraps searchMemories without mutating return shape", async () => {
|
||||
let calls = 0
|
||||
const fakeTool = {
|
||||
searchMemories: {
|
||||
description: "search",
|
||||
inputSchema: {} as any,
|
||||
execute: async (input: any) => {
|
||||
calls++
|
||||
return {
|
||||
success: true,
|
||||
results: [{ q: input.informationToGet }],
|
||||
count: 1,
|
||||
}
|
||||
},
|
||||
},
|
||||
addMemory: {
|
||||
description: "add",
|
||||
inputSchema: {} as any,
|
||||
execute: async () => ({ success: true, memory: {} }),
|
||||
},
|
||||
} as any
|
||||
|
||||
const wrapped = createWorkingMemory(fakeTool, { ttlMs: 60_000 })
|
||||
|
||||
const r1: any = await (wrapped.searchMemories as any).execute({
|
||||
informationToGet: "preferences",
|
||||
limit: 10,
|
||||
})
|
||||
expect(r1.success).toBe(true)
|
||||
expect(r1.results).toEqual([{ q: "preferences" }])
|
||||
expect(r1._source).toBeUndefined() // not mutating shape
|
||||
expect(calls).toBe(1)
|
||||
|
||||
const r2: any = await (wrapped.searchMemories as any).execute({
|
||||
informationToGet: "preferences",
|
||||
limit: 10,
|
||||
})
|
||||
expect(calls).toBe(1) // cache hit, no second fetch
|
||||
expect(r2.results).toEqual([{ q: "preferences" }])
|
||||
|
||||
// Original tool still hits every time (no silent mutation)
|
||||
const orig: any = await fakeTool.searchMemories.execute({
|
||||
informationToGet: "preferences",
|
||||
limit: 10,
|
||||
})
|
||||
expect(calls).toBe(2)
|
||||
expect(wrapped.workingMemory.stats().hits).toBe(1)
|
||||
})
|
||||
|
||||
it("does not auto-populate on addMemory (V1 non-goal)", async () => {
|
||||
const fakeTool = {
|
||||
searchMemories: {
|
||||
execute: async () => ({ success: true, results: [], count: 0 }),
|
||||
},
|
||||
addMemory: {
|
||||
execute: async () => ({ success: true, memory: { id: "m1" } }),
|
||||
},
|
||||
} as any
|
||||
const wrapped = createWorkingMemory(fakeTool)
|
||||
await (wrapped.addMemory as any).execute({ memory: "hello" })
|
||||
expect(wrapped.workingMemory.stats().size).toBe(0)
|
||||
})
|
||||
})
|
||||
236
packages/ai-sdk/src/working-memory.ts
Normal file
236
packages/ai-sdk/src/working-memory.ts
Normal file
|
|
@ -0,0 +1,236 @@
|
|||
/**
|
||||
* Working Memory — V1 (RFC #1625, Staff review 7.8 → 10)
|
||||
*
|
||||
* Explicit, opt-in, non-mutating wrapper. No change to `supermemoryTools`
|
||||
* behavior unless the consumer wraps it via `createWorkingMemory()`.
|
||||
*
|
||||
* V1 scope (mergeable): LRU, TTL, promise dedup, invalidate, clear, stats.
|
||||
* Deferred to V2/V3: pin/unpin, auto-populate on addMemory, persistent cache,
|
||||
* semantic dedup, background refresh.
|
||||
*
|
||||
* Public API (V1):
|
||||
* new WorkingMemory({ maxEntries?, ttlMs? })
|
||||
* - search(query, fetchFn, opts?) -> results (deduped, cached)
|
||||
* - get / set / has (low-level, for testing)
|
||||
* - invalidate(query?) / clear() / stats()
|
||||
*/
|
||||
|
||||
export type WorkingMemoryOptions = {
|
||||
maxEntries?: number
|
||||
ttlMs?: number
|
||||
}
|
||||
|
||||
export type WorkingMemoryEntry<T = unknown> = {
|
||||
results: T[]
|
||||
expiresAt: number
|
||||
}
|
||||
|
||||
export type WorkingMemoryStats = {
|
||||
hits: number
|
||||
misses: number
|
||||
size: number
|
||||
}
|
||||
|
||||
/**
|
||||
* Cache key normalization — intentionally case- and whitespace-insensitive.
|
||||
* Backend semantic search is case-insensitive (embedding-based), so " Hello "
|
||||
* and "hello" should hit the same cache entry. If backend ever treats casing
|
||||
* as significant, this normalization should be revisited. `limit` defaults to
|
||||
* 10 matching the backend's default when no limit is provided.
|
||||
*/
|
||||
function normalizeKey(query: string, limit?: number): string {
|
||||
return `${query.trim().toLowerCase()}::limit=${limit ?? 10}`
|
||||
}
|
||||
|
||||
export class WorkingMemory<T = unknown> {
|
||||
private readonly maxEntries: number
|
||||
private readonly ttlMs: number
|
||||
private readonly cache = new Map<string, WorkingMemoryEntry<T>>()
|
||||
private readonly inflight = new Map<string, Promise<T[]>>()
|
||||
private hits = 0
|
||||
private misses = 0
|
||||
|
||||
constructor(opts: WorkingMemoryOptions = {}) {
|
||||
this.maxEntries = opts.maxEntries ?? 100
|
||||
// Default TTL is intentionally conservative (60 s) for interactive agent
|
||||
// sessions. Configurable because freshness requirements differ across apps;
|
||||
// maintainers may choose a different SDK default.
|
||||
this.ttlMs = opts.ttlMs ?? 60_000
|
||||
}
|
||||
|
||||
/**
|
||||
* Primary entry point — wraps a fetch function with LRU+TTL + dedup.
|
||||
* Returns cached results on hit, otherwise calls fetchFn, caches, and returns.
|
||||
* Concurrent callers for the same normalized key share one in-flight promise.
|
||||
*/
|
||||
async search(
|
||||
query: string,
|
||||
fetchFn: (query: string) => Promise<T[]>,
|
||||
opts?: { limit?: number; ttlMs?: number },
|
||||
): Promise<T[]> {
|
||||
const key = normalizeKey(query, opts?.limit)
|
||||
const entry = this.cache.get(key)
|
||||
if (entry && Date.now() <= entry.expiresAt) {
|
||||
this.touch(key, entry)
|
||||
this.hits++
|
||||
return entry.results
|
||||
}
|
||||
if (entry && Date.now() > entry.expiresAt) {
|
||||
this.cache.delete(key)
|
||||
}
|
||||
const existing = this.inflight.get(key)
|
||||
if (existing) {
|
||||
const deduped = await existing
|
||||
this.hits++
|
||||
return deduped
|
||||
}
|
||||
this.misses++
|
||||
const promise = (async () => {
|
||||
const results = await fetchFn(query)
|
||||
this.set(query, results, opts)
|
||||
return results
|
||||
})()
|
||||
this.inflight.set(key, promise)
|
||||
try {
|
||||
return await promise
|
||||
} finally {
|
||||
this.inflight.delete(key)
|
||||
}
|
||||
}
|
||||
|
||||
get(query: string, opts?: { limit?: number }): T[] | undefined {
|
||||
const key = normalizeKey(query, opts?.limit)
|
||||
const entry = this.cache.get(key)
|
||||
if (!entry) {
|
||||
this.misses++
|
||||
return undefined
|
||||
}
|
||||
if (Date.now() > entry.expiresAt) {
|
||||
this.cache.delete(key)
|
||||
this.misses++
|
||||
return undefined
|
||||
}
|
||||
this.touch(key, entry)
|
||||
this.hits++
|
||||
return entry.results
|
||||
}
|
||||
|
||||
set(
|
||||
query: string,
|
||||
results: T[],
|
||||
opts?: { limit?: number; ttlMs?: number },
|
||||
): void {
|
||||
const key = normalizeKey(query, opts?.limit)
|
||||
const ttl = opts?.ttlMs ?? this.ttlMs
|
||||
const entry: WorkingMemoryEntry<T> = {
|
||||
results,
|
||||
expiresAt: Date.now() + ttl,
|
||||
}
|
||||
this.cache.set(key, entry)
|
||||
this.touch(key, entry)
|
||||
}
|
||||
|
||||
has(query: string, opts?: { limit?: number }): boolean {
|
||||
const key = normalizeKey(query, opts?.limit)
|
||||
const entry = this.cache.get(key)
|
||||
if (!entry) return false
|
||||
if (Date.now() > entry.expiresAt) return false
|
||||
return true
|
||||
}
|
||||
|
||||
invalidate(query?: string, opts?: { limit?: number }): void {
|
||||
if (query === undefined) {
|
||||
this.cache.clear()
|
||||
return
|
||||
}
|
||||
const key = normalizeKey(query, opts?.limit)
|
||||
this.cache.delete(key)
|
||||
}
|
||||
|
||||
clear(): void {
|
||||
this.cache.clear()
|
||||
this.inflight.clear()
|
||||
this.hits = 0
|
||||
this.misses = 0
|
||||
}
|
||||
|
||||
stats(): WorkingMemoryStats {
|
||||
return { hits: this.hits, misses: this.misses, size: this.cache.size }
|
||||
}
|
||||
|
||||
private touch(key: string, entry: WorkingMemoryEntry<T>): void {
|
||||
this.cache.delete(key)
|
||||
this.cache.set(key, entry)
|
||||
if (this.cache.size > this.maxEntries) {
|
||||
const oldest = this.cache.keys().next().value
|
||||
if (oldest !== undefined) this.cache.delete(oldest)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Explicit decorator — wraps `supermemoryTools` without mutating its behavior.
|
||||
*
|
||||
* const tools = supermemoryTools(apiKey, { projectId: "..." })
|
||||
* const memory = createWorkingMemory(tools, { ttlMs: 60_000, maxEntries: 100 })
|
||||
* await memory.searchMemories("user preferences")
|
||||
*
|
||||
* Zero breaking behavior: `tools.searchMemories` still means "query backend".
|
||||
* `memory.searchMemories` means "maybe query cache".
|
||||
*/
|
||||
export function createWorkingMemory<
|
||||
T extends { searchMemories: { execute: (input: any) => Promise<any> } },
|
||||
>(
|
||||
tools: T,
|
||||
opts?: WorkingMemoryOptions,
|
||||
): T & {
|
||||
workingMemory: WorkingMemory
|
||||
searchMemories: T["searchMemories"]
|
||||
} {
|
||||
const wm = new WorkingMemory(opts)
|
||||
const original = tools.searchMemories
|
||||
|
||||
// Wrap execute to add WorkingMemory without mutating return shape.
|
||||
// Preserves backend contract: always returns `{ success, results, count }`
|
||||
// (same shape as uncached). Cached path reconstructs count from cached
|
||||
// results.length; no `_source` or extra fields are injected — stats live
|
||||
// on `workingMemory.stats()` separately. See TracePull validation notes.
|
||||
const wrapped = {
|
||||
...original,
|
||||
execute: async (input: any) => {
|
||||
const query: string = input?.informationToGet ?? input?.q ?? ""
|
||||
// Normalize limit: backend defaults to 10 when not provided, so
|
||||
// cache key must also default to 10 to preserve hit correctness.
|
||||
const limit: number = input?.limit ?? 10
|
||||
const cached = wm.get(query, { limit })
|
||||
if (cached !== undefined) {
|
||||
// Reconstruct a successful tool result from cached results
|
||||
return { success: true, results: cached, count: cached.length }
|
||||
}
|
||||
// Preserve dedup via WorkingMemory.search with the real fetch
|
||||
const results = await wm.search(
|
||||
query,
|
||||
async () => {
|
||||
const res: any = await (original as any).execute(input)
|
||||
// Tools return { success, results, count } on success
|
||||
if (res?.success === false)
|
||||
throw new Error(res.error ?? "search failed")
|
||||
return (res?.results ?? []) as unknown[]
|
||||
},
|
||||
{ limit },
|
||||
)
|
||||
return { success: true, results, count: results.length }
|
||||
},
|
||||
} as T["searchMemories"]
|
||||
|
||||
const out = { ...tools, searchMemories: wrapped } as T & {
|
||||
workingMemory: WorkingMemory
|
||||
searchMemories: T["searchMemories"]
|
||||
}
|
||||
Object.defineProperty(out, "workingMemory", {
|
||||
value: wm,
|
||||
enumerable: false,
|
||||
writable: false,
|
||||
})
|
||||
return out
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue