Skip to content

Commit b0ac8de

Browse files
mattappersonmattapperson
committed
test: Add comprehensive tests for context-aware memory features
- Add 10 new tests covering all context-aware enhancements - Test message editing and version history tracking - Test token-aware message selection within budget - Test cache management (enable, retrieve, invalidate) - Test message filtering by status and importance - Test graceful degradation when storage doesn't support features - Test contextWindow config integration - Update usage examples to demonstrate new features - All 20 tests passing
1 parent 6905221 commit b0ac8de

File tree

2 files changed

+361
-0
lines changed

2 files changed

+361
-0
lines changed

‎examples/memory-usage.ts‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,138 @@ async function main() {
107107
autoSave: true,
108108
});
109109
console.log("Custom memory config:", customMemory.getConfig());
110+
111+
// Example 9: Token-aware memory management
112+
console.log("\n=== Example 9: Token-Aware Features ===");
113+
114+
const tokenStorage = new InMemoryStorage();
115+
const tokenMemory = new Memory(tokenStorage, {
116+
maxHistoryMessages: 10,
117+
autoInject: true,
118+
autoSave: true,
119+
contextWindow: {
120+
maxTokens: 1000,
121+
strategy: "token-aware",
122+
},
123+
trackTokenUsage: true,
124+
});
125+
126+
const tokenThreadId = "token-thread-1";
127+
await tokenMemory.createThread(tokenThreadId, userId);
128+
129+
// Save messages with token counts (in real usage, these would come from API responses)
130+
const savedMsgs = await tokenMemory.saveMessages(tokenThreadId, userId, [
131+
{ role: "user" as const, content: "Tell me about AI" },
132+
{ role: "assistant" as const, content: "AI is artificial intelligence..." },
133+
]);
134+
135+
// Manually set token counts (in production, these would be from API)
136+
await tokenStorage.updateMessage(savedMsgs[0].id, {
137+
tokenCount: 50,
138+
importance: 0.8
139+
});
140+
await tokenStorage.updateMessage(savedMsgs[1].id, {
141+
tokenCount: 150,
142+
importance: 0.9
143+
});
144+
145+
// Get total token count for thread
146+
const totalTokens = await tokenMemory.getThreadTokenCount(tokenThreadId);
147+
console.log(`Total tokens in thread: ${totalTokens}`);
148+
149+
// Get messages within token budget
150+
const budgetedMessages = await tokenMemory.getMessagesWithinBudget(tokenThreadId, 500);
151+
console.log(`Messages within 500 token budget: ${budgetedMessages.length}`);
152+
153+
// Example 10: Message editing with version history
154+
console.log("\n=== Example 10: Message Editing ===");
155+
156+
const editThreadId = "edit-thread-1";
157+
await memory.createThread(editThreadId, userId);
158+
159+
const [originalMsg] = await memory.saveMessages(editThreadId, userId, [
160+
{ role: "user" as const, content: "Original message content" },
161+
]);
162+
163+
console.log("Original message:", originalMsg.message.content);
164+
165+
// Edit the message
166+
const updatedMsg = await memory.updateMessage(originalMsg.id, {
167+
content: "Updated message content",
168+
});
169+
170+
if (updatedMsg) {
171+
console.log("Updated message:", updatedMsg.message.content);
172+
console.log("Version:", updatedMsg.version);
173+
174+
// Get version history
175+
const versions = await memory.getMessageVersions(originalMsg.id);
176+
console.log(`Message has ${versions.length} versions`);
177+
versions.forEach((v) => {
178+
console.log(` v${v.version || 1}: ${v.message.content}`);
179+
});
180+
}
181+
182+
// Example 11: Cache management
183+
console.log("\n=== Example 11: Cache Management ===");
184+
185+
const cacheThreadId = "cache-thread-1";
186+
await memory.createThread(cacheThreadId, userId);
187+
188+
const [cacheMsg] = await memory.saveMessages(cacheThreadId, userId, [
189+
{ role: "system" as const, content: "System prompt that should be cached" },
190+
]);
191+
192+
// Enable caching for this message (e.g., for provider-level prompt caching)
193+
const futureDate = new Date(Date.now() + 3600000); // 1 hour from now
194+
await storage.updateMessage(cacheMsg.id, {
195+
cacheControl: { enabled: true, expiresAt: futureDate },
196+
});
197+
198+
const cachedMsgs = await memory.getCachedMessages(cacheThreadId);
199+
console.log(`Cached messages: ${cachedMsgs.length}`);
200+
201+
// Invalidate cache when needed
202+
await memory.invalidateCache(cacheThreadId);
203+
const cachedAfterInvalidate = await memory.getCachedMessages(cacheThreadId);
204+
console.log(`Cached messages after invalidation: ${cachedAfterInvalidate.length}`);
205+
206+
// Example 12: Message filtering by status and importance
207+
console.log("\n=== Example 12: Message Filtering ===");
208+
209+
const filterThreadId = "filter-thread-1";
210+
await memory.createThread(filterThreadId, userId);
211+
212+
const filterMsgs = await memory.saveMessages(filterThreadId, userId, [
213+
{ role: "user" as const, content: "Important message" },
214+
{ role: "assistant" as const, content: "Normal response" },
215+
{ role: "user" as const, content: "Low priority message" },
216+
]);
217+
218+
// Set different statuses and importance
219+
await storage.updateMessage(filterMsgs[0].id, {
220+
status: "active",
221+
importance: 0.9
222+
});
223+
await storage.updateMessage(filterMsgs[1].id, {
224+
status: "active",
225+
importance: 0.5
226+
});
227+
await storage.updateMessage(filterMsgs[2].id, {
228+
status: "archived",
229+
importance: 0.2
230+
});
231+
232+
// Filter by status
233+
const activeMsgs = await memory.getMessagesByStatus(filterThreadId, "active");
234+
console.log(`Active messages: ${activeMsgs.length}`);
235+
236+
const archivedMsgs = await memory.getMessagesByStatus(filterThreadId, "archived");
237+
console.log(`Archived messages: ${archivedMsgs.length}`);
238+
239+
// Filter by importance
240+
const importantMsgs = await memory.getMessagesByImportance(filterThreadId, 0.7);
241+
console.log(`Messages with importance >= 0.7: ${importantMsgs.length}`);
110242
}
111243

112244
main().catch(console.error);

‎tests/e2e/memory.test.ts‎

Lines changed: 229 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,4 +189,233 @@ describe("Memory Integration E2E Tests", () => {
189189
// Should have 10 unique messages
190190
expect(new Set(contents).size).toBe(10);
191191
});
192+
193+
describe("Context-Aware Memory Features", () => {
194+
it("should update a message and track version history", async () => {
195+
await memory.createThread("thread-1", "user-1");
196+
197+
const saved = await memory.saveMessages("thread-1", "user-1", [
198+
{ role: "user" as const, content: "Original message" },
199+
]);
200+
const messageId = saved[0].id;
201+
202+
// Update the message
203+
const updated = await memory.updateMessage(messageId, {
204+
content: "Updated message",
205+
});
206+
207+
expect(updated).toBeDefined();
208+
expect(updated?.message.content).toBe("Updated message");
209+
expect(updated?.version).toBe(2);
210+
211+
// Get version history
212+
const versions = await memory.getMessageVersions(messageId);
213+
expect(versions).toHaveLength(2);
214+
expect(versions[0].message.content).toBe("Original message");
215+
expect(versions[1].message.content).toBe("Updated message");
216+
});
217+
218+
it("should track token counts and get thread token total", async () => {
219+
await memory.createThread("thread-1", "user-1");
220+
221+
// Save messages with token counts
222+
const messages = await memory.getMessages("thread-1");
223+
const saved = await memory.saveMessages("thread-1", "user-1", [
224+
{ role: "user" as const, content: "Hello" },
225+
{ role: "assistant" as const, content: "Hi there" },
226+
]);
227+
228+
// Manually add token counts (in real usage these would come from API)
229+
await storage.updateMessage(saved[0].id, { tokenCount: 10 });
230+
await storage.updateMessage(saved[1].id, { tokenCount: 15 });
231+
232+
const tokenCount = await memory.getThreadTokenCount("thread-1");
233+
expect(tokenCount).toBe(25);
234+
});
235+
236+
it("should get messages within token budget", async () => {
237+
await memory.createThread("thread-1", "user-1");
238+
239+
const saved = await memory.saveMessages("thread-1", "user-1", [
240+
{ role: "user" as const, content: "Message 1" },
241+
{ role: "assistant" as const, content: "Response 1" },
242+
{ role: "user" as const, content: "Message 2" },
243+
{ role: "assistant" as const, content: "Response 2" },
244+
]);
245+
246+
// Add token counts
247+
await storage.updateMessage(saved[0].id, { tokenCount: 20, importance: 0.5 });
248+
await storage.updateMessage(saved[1].id, { tokenCount: 30, importance: 0.5 });
249+
await storage.updateMessage(saved[2].id, { tokenCount: 25, importance: 0.8 });
250+
await storage.updateMessage(saved[3].id, { tokenCount: 35, importance: 0.8 });
251+
252+
// Get messages within 70 token budget
253+
const messages = await memory.getMessagesWithinBudget("thread-1", 70);
254+
255+
// Should get the two highest importance messages (Message 2 + Response 2 = 60 tokens)
256+
expect(messages.length).toBeGreaterThan(0);
257+
expect(messages.length).toBeLessThanOrEqual(4);
258+
});
259+
260+
it("should manage cache control for messages", async () => {
261+
await memory.createThread("thread-1", "user-1");
262+
263+
const saved = await memory.saveMessages("thread-1", "user-1", [
264+
{ role: "user" as const, content: "Cached message" },
265+
{ role: "assistant" as const, content: "Not cached" },
266+
]);
267+
268+
// Enable cache for first message
269+
const futureDate = new Date(Date.now() + 60000); // 1 minute from now
270+
await storage.updateMessage(saved[0].id, {
271+
cacheControl: { enabled: true, expiresAt: futureDate },
272+
});
273+
274+
const cachedMessages = await memory.getCachedMessages("thread-1");
275+
expect(cachedMessages).toHaveLength(1);
276+
expect(cachedMessages[0].message.content).toBe("Cached message");
277+
});
278+
279+
it("should invalidate cache for messages", async () => {
280+
await memory.createThread("thread-1", "user-1");
281+
282+
const saved = await memory.saveMessages("thread-1", "user-1", [
283+
{ role: "user" as const, content: "Cached message" },
284+
]);
285+
286+
// Enable cache
287+
const futureDate = new Date(Date.now() + 60000);
288+
await storage.updateMessage(saved[0].id, {
289+
cacheControl: { enabled: true, expiresAt: futureDate },
290+
});
291+
292+
// Verify cache is active
293+
let cachedMessages = await memory.getCachedMessages("thread-1");
294+
expect(cachedMessages).toHaveLength(1);
295+
296+
// Invalidate cache
297+
await memory.invalidateCache("thread-1");
298+
299+
// Verify cache is invalidated
300+
cachedMessages = await memory.getCachedMessages("thread-1");
301+
expect(cachedMessages).toHaveLength(0);
302+
});
303+
304+
it("should filter messages by status", async () => {
305+
await memory.createThread("thread-1", "user-1");
306+
307+
const saved = await memory.saveMessages("thread-1", "user-1", [
308+
{ role: "user" as const, content: "Active message" },
309+
{ role: "assistant" as const, content: "Archived message" },
310+
{ role: "user" as const, content: "Deleted message" },
311+
]);
312+
313+
// Set different statuses
314+
await storage.updateMessage(saved[0].id, { status: "active" });
315+
await storage.updateMessage(saved[1].id, { status: "archived" });
316+
await storage.updateMessage(saved[2].id, { status: "deleted" });
317+
318+
const activeMessages = await memory.getMessagesByStatus("thread-1", "active");
319+
expect(activeMessages).toHaveLength(1);
320+
expect(activeMessages[0].message.content).toBe("Active message");
321+
322+
const archivedMessages = await memory.getMessagesByStatus("thread-1", "archived");
323+
expect(archivedMessages).toHaveLength(1);
324+
expect(archivedMessages[0].message.content).toBe("Archived message");
325+
});
326+
327+
it("should filter messages by importance threshold", async () => {
328+
await memory.createThread("thread-1", "user-1");
329+
330+
const saved = await memory.saveMessages("thread-1", "user-1", [
331+
{ role: "user" as const, content: "Low importance" },
332+
{ role: "assistant" as const, content: "Medium importance" },
333+
{ role: "user" as const, content: "High importance" },
334+
]);
335+
336+
// Set importance scores
337+
await storage.updateMessage(saved[0].id, { importance: 0.3 });
338+
await storage.updateMessage(saved[1].id, { importance: 0.6 });
339+
await storage.updateMessage(saved[2].id, { importance: 0.9 });
340+
341+
// Get messages with importance >= 0.5
342+
const importantMessages = await memory.getMessagesByImportance("thread-1", 0.5);
343+
expect(importantMessages).toHaveLength(2);
344+
expect(importantMessages[0].importance).toBeGreaterThanOrEqual(0.5);
345+
expect(importantMessages[1].importance).toBeGreaterThanOrEqual(0.5);
346+
});
347+
348+
it("should handle graceful degradation when storage doesn't support features", async () => {
349+
// Create a storage that doesn't implement optional methods
350+
class BasicStorage extends InMemoryStorage {
351+
updateMessage = undefined;
352+
getMessageHistory = undefined;
353+
}
354+
355+
const basicStorage = new BasicStorage();
356+
const basicMemory = new Memory(basicStorage);
357+
358+
await basicMemory.createThread("thread-1", "user-1");
359+
const saved = await basicMemory.saveMessages("thread-1", "user-1", [
360+
{ role: "user" as const, content: "Test" },
361+
]);
362+
363+
// These should return null/empty without errors
364+
const updated = await basicMemory.updateMessage(saved[0].id, { content: "Updated" });
365+
expect(updated).toBeNull();
366+
367+
const versions = await basicMemory.getMessageVersions(saved[0].id);
368+
expect(versions).toEqual([]);
369+
});
370+
371+
it("should use contextWindow config for token-aware selection", async () => {
372+
const configuredStorage = new InMemoryStorage();
373+
const configuredMemory = new Memory(configuredStorage, {
374+
contextWindow: {
375+
maxTokens: 100,
376+
strategy: "token-aware",
377+
},
378+
});
379+
380+
await configuredMemory.createThread("thread-1", "user-1");
381+
const saved = await configuredMemory.saveMessages("thread-1", "user-1", [
382+
{ role: "user" as const, content: "Message 1" },
383+
{ role: "assistant" as const, content: "Response 1" },
384+
{ role: "user" as const, content: "Message 2" },
385+
]);
386+
387+
// Add token counts
388+
await configuredStorage.updateMessage(saved[0].id, { tokenCount: 40 });
389+
await configuredStorage.updateMessage(saved[1].id, { tokenCount: 50 });
390+
await configuredStorage.updateMessage(saved[2].id, { tokenCount: 30 });
391+
392+
// Should use contextWindow.maxTokens from config
393+
const messages = await configuredMemory.getMessagesWithinBudget("thread-1");
394+
expect(messages.length).toBeGreaterThan(0);
395+
});
396+
397+
it("should preserve message order when sorting by importance and recency", async () => {
398+
await memory.createThread("thread-1", "user-1");
399+
400+
const saved = await memory.saveMessages("thread-1", "user-1", [
401+
{ role: "user" as const, content: "First" },
402+
{ role: "assistant" as const, content: "Second" },
403+
{ role: "user" as const, content: "Third" },
404+
]);
405+
406+
// Set same importance, different times
407+
await storage.updateMessage(saved[0].id, { importance: 0.5, tokenCount: 10 });
408+
await new Promise(resolve => setTimeout(resolve, 10));
409+
await storage.updateMessage(saved[1].id, { importance: 0.5, tokenCount: 10 });
410+
await new Promise(resolve => setTimeout(resolve, 10));
411+
await storage.updateMessage(saved[2].id, { importance: 0.5, tokenCount: 10 });
412+
413+
const messages = await memory.getMessagesWithinBudget("thread-1", 30);
414+
415+
// Should maintain chronological order in result
416+
expect(messages[0].content).toBe("First");
417+
expect(messages[1].content).toBe("Second");
418+
expect(messages[2].content).toBe("Third");
419+
});
420+
});
192421
});

0 commit comments

Comments
 (0)