Skip to content

Commit a894fbb

Browse files
committed
refactor: Move serialization to storage layer
Serialization is a storage concern, not a memory concern. This refactoring: - Moves serialize/hydrate methods from Memory class to MemoryStorage interface - Implements serialize/hydrate/serializeThread/hydrateThread in InMemoryStorage - Updates tests to use storage.serialize() instead of memory.serialize() - Updates usage example to show proper storage serialization pattern This makes the architecture cleaner and allows different storage backends to implement their own serialization strategies (e.g., direct DB export, file formats, etc.)
1 parent 7e16bda commit a894fbb

File tree

5 files changed

+91
-128
lines changed

5 files changed

+91
-128
lines changed

examples/memory-usage.ts

Lines changed: 15 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,11 @@
88
import { OpenRouter, Memory, InMemoryStorage } from "@openrouter/sdk";
99

1010
async function main() {
11-
// Create a memory instance with in-memory storage
12-
const memory = new Memory(new InMemoryStorage(), {
11+
// Create storage instance
12+
const storage = new InMemoryStorage();
13+
14+
// Create a memory instance with the storage
15+
const memory = new Memory(storage, {
1316
maxHistoryMessages: 10, // Keep last 10 messages in context
1417
autoInject: true, // Automatically inject history
1518
autoSave: true, // Automatically save messages
@@ -30,58 +33,23 @@ async function main() {
3033
// First message
3134
const response1 = client.getResponse({
3235
model: "meta-llama/llama-3.2-1b-instruct",
33-
input: [{ role: "user", content: "My name is Alice." }],
36+
input: "My name is Alice.",
3437
threadId,
3538
resourceId: userId,
3639
});
3740
const text1 = await response1.text;
38-
console.log("Assistant:", text1);
41+
console.log("AI Response:", text1);
3942

4043
// Second message - history is automatically injected
4144
const response2 = client.getResponse({
4245
model: "meta-llama/llama-3.2-1b-instruct",
43-
input: [{ role: "user", content: "What's my name?" }],
46+
input: "What's my name?",
4447
threadId,
4548
resourceId: userId,
4649
});
4750
const text2 = await response2.text;
48-
console.log("Assistant:", text2); // Should remember the name is Alice
49-
50-
// Example 2: Working with thread working memory
51-
console.log("\n=== Example 2: Thread Working Memory ===");
52-
53-
await memory.updateThreadWorkingMemory(threadId, {
54-
topic: "Introduction",
55-
lastActivity: new Date().toISOString(),
56-
messageCount: 2,
57-
});
58-
59-
const threadMemory = await memory.getThreadWorkingMemory(threadId);
60-
console.log("Thread working memory:", threadMemory?.data);
61-
62-
// Example 3: Working with resource (user) working memory
63-
console.log("\n=== Example 3: Resource Working Memory ===");
64-
65-
await memory.updateResourceWorkingMemory(userId, {
66-
name: "Alice",
67-
preferences: {
68-
theme: "dark",
69-
language: "en",
70-
},
71-
createdAt: new Date().toISOString(),
72-
});
73-
74-
const userMemory = await memory.getResourceWorkingMemory(userId);
75-
console.log("User working memory:", userMemory?.data);
76-
77-
// Example 4: Managing multiple threads for a user
78-
console.log("\n=== Example 4: Multiple Threads ===");
51+
console.log("AI Response:", text2); // Should remember the name is Alice
7952

80-
const thread2Id = "conversation-789";
81-
await memory.createThread(thread2Id, userId, "Second Conversation");
82-
await memory.saveMessages(thread2Id, userId, [
83-
{ role: "user", content: "Hello from second thread!" },
84-
]);
8553

8654
const userThreads = await memory.getThreadsByResource(userId);
8755
console.log(`User has ${userThreads.length} threads:`,
@@ -90,8 +58,8 @@ async function main() {
9058
// Example 5: Serialization and persistence
9159
console.log("\n=== Example 5: Serialization ===");
9260

93-
// Serialize entire memory state
94-
const memoryState = await memory.serialize();
61+
// Serialize entire memory state using storage
62+
const memoryState = await storage.serialize();
9563
console.log("Serialized state:", {
9664
threads: memoryState.threads.length,
9765
messages: memoryState.messages.length,
@@ -102,12 +70,13 @@ async function main() {
10270
// For example: fs.writeFileSync('memory-state.json', JSON.stringify(memoryState));
10371

10472
// Later, you can restore the state
105-
const newMemory = new Memory(new InMemoryStorage());
106-
await newMemory.hydrate(memoryState);
73+
const newStorage = new InMemoryStorage();
74+
await newStorage.hydrate(memoryState);
75+
const newMemory = new Memory(newStorage);
10776
console.log("Memory restored successfully!");
10877

10978
// Example 6: Serialize a single thread
110-
const threadState = await memory.serializeThread(threadId);
79+
const threadState = await storage.serializeThread(threadId);
11180
if (threadState) {
11281
console.log("Thread state:", {
11382
threadId: threadState.thread.id,

src/lib/memory/memory.ts

Lines changed: 0 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,6 @@ import type {
1111
MemoryMessage,
1212
Resource,
1313
ResourceWorkingMemory,
14-
SerializedMemoryState,
15-
SerializedThreadState,
1614
Thread,
1715
ThreadWorkingMemory,
1816
WorkingMemoryData,
@@ -261,69 +259,6 @@ export class Memory {
261259
return await this.storage.getResourceWorkingMemory(resourceId);
262260
}
263261

264-
// ===== Serialization =====
265-
266-
/**
267-
* Serialize the entire memory state to JSON
268-
* @returns The serialized state
269-
*/
270-
async serialize(): Promise<SerializedMemoryState> {
271-
return await this.storage.exportState();
272-
}
273-
274-
/**
275-
* Serialize a single thread to JSON
276-
* @param threadId The thread ID
277-
* @returns The serialized thread state, or null if thread not found
278-
*/
279-
async serializeThread(threadId: string): Promise<SerializedThreadState | null> {
280-
const thread = await this.storage.getThread(threadId);
281-
if (!thread) {
282-
return null;
283-
}
284-
285-
const messages = await this.storage.getMessages(threadId);
286-
const threadWorkingMemory =
287-
await this.storage.getThreadWorkingMemory(threadId);
288-
289-
return {
290-
version: "1.0.0",
291-
thread,
292-
messages,
293-
...(threadWorkingMemory !== null && { threadWorkingMemory }),
294-
serializedAt: new Date(),
295-
};
296-
}
297-
298-
/**
299-
* Hydrate (restore) the entire memory state from JSON
300-
* Warning: This will replace all existing data in memory
301-
* @param state The serialized state to restore
302-
*/
303-
async hydrate(state: SerializedMemoryState): Promise<void> {
304-
await this.storage.importState(state);
305-
}
306-
307-
/**
308-
* Hydrate (restore) a single thread from JSON
309-
* @param threadState The serialized thread state to restore
310-
*/
311-
async hydrateThread(threadState: SerializedThreadState): Promise<void> {
312-
// Save the thread
313-
await this.storage.saveThread(threadState.thread);
314-
315-
// Save all messages
316-
await this.storage.saveMessages(threadState.messages);
317-
318-
// Save thread working memory if present
319-
if (threadState.threadWorkingMemory) {
320-
await this.storage.updateThreadWorkingMemory(
321-
threadState.thread.id,
322-
threadState.threadWorkingMemory.data,
323-
);
324-
}
325-
}
326-
327262
// ===== Utility Methods =====
328263

329264
/**

src/lib/memory/storage/in-memory.ts

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import type {
99
Resource,
1010
ResourceWorkingMemory,
1111
SerializedMemoryState,
12+
SerializedThreadState,
1213
Thread,
1314
ThreadWorkingMemory,
1415
WorkingMemoryData,
@@ -196,7 +197,7 @@ export class InMemoryStorage implements MemoryStorage {
196197

197198
// ===== Serialization Operations =====
198199

199-
async exportState(): Promise<SerializedMemoryState> {
200+
async serialize(): Promise<SerializedMemoryState> {
200201
return {
201202
version: "1.0.0",
202203
threads: Array.from(this.threads.values()),
@@ -210,7 +211,25 @@ export class InMemoryStorage implements MemoryStorage {
210211
};
211212
}
212213

213-
async importState(state: SerializedMemoryState): Promise<void> {
214+
async serializeThread(threadId: string): Promise<SerializedThreadState | null> {
215+
const thread = await this.getThread(threadId);
216+
if (!thread) {
217+
return null;
218+
}
219+
220+
const messages = await this.getMessages(threadId);
221+
const threadWorkingMemory = await this.getThreadWorkingMemory(threadId);
222+
223+
return {
224+
version: "1.0.0",
225+
thread,
226+
messages,
227+
...(threadWorkingMemory !== null && { threadWorkingMemory }),
228+
serializedAt: new Date(),
229+
};
230+
}
231+
232+
async hydrate(state: SerializedMemoryState): Promise<void> {
214233
// Clear existing data
215234
this.threads.clear();
216235
this.messages.clear();
@@ -242,4 +261,20 @@ export class InMemoryStorage implements MemoryStorage {
242261
this.resourceWorkingMemories.set(rwm.resourceId, rwm);
243262
}
244263
}
264+
265+
async hydrateThread(threadState: SerializedThreadState): Promise<void> {
266+
// Save the thread
267+
await this.saveThread(threadState.thread);
268+
269+
// Save all messages
270+
await this.saveMessages(threadState.messages);
271+
272+
// Save thread working memory if present
273+
if (threadState.threadWorkingMemory) {
274+
await this.updateThreadWorkingMemory(
275+
threadState.thread.id,
276+
threadState.threadWorkingMemory.data,
277+
);
278+
}
279+
}
245280
}

src/lib/memory/storage/interface.ts

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import type {
99
Resource,
1010
ResourceWorkingMemory,
1111
SerializedMemoryState,
12+
SerializedThreadState,
1213
Thread,
1314
ThreadWorkingMemory,
1415
WorkingMemoryData,
@@ -135,14 +136,28 @@ export interface MemoryStorage {
135136
// ===== Serialization Operations =====
136137

137138
/**
138-
* Export the entire storage state
139+
* Serialize the entire storage state
139140
* @returns Serialized state of all data in storage
140141
*/
141-
exportState(): Promise<SerializedMemoryState>;
142+
serialize(): Promise<SerializedMemoryState>;
142143

143144
/**
144-
* Import a complete storage state
145-
* @param state The state to import
145+
* Serialize a single thread and its data
146+
* @param threadId The thread ID to serialize
147+
* @returns The serialized thread state, or null if not found
146148
*/
147-
importState(state: SerializedMemoryState): Promise<void>;
149+
serializeThread(threadId: string): Promise<SerializedThreadState | null>;
150+
151+
/**
152+
* Hydrate (restore) the entire storage state
153+
* Warning: This will replace all existing data in storage
154+
* @param state The state to restore
155+
*/
156+
hydrate(state: SerializedMemoryState): Promise<void>;
157+
158+
/**
159+
* Hydrate (restore) a single thread
160+
* @param threadState The thread state to restore
161+
*/
162+
hydrateThread(threadState: SerializedThreadState): Promise<void>;
148163
}

tests/e2e/memory.test.ts

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,12 @@ import { OpenRouter, Memory, InMemoryStorage } from "../../src/index.js";
44
describe("Memory Integration E2E Tests", () => {
55
const apiKey = process.env.OPENROUTER_API_KEY;
66
let memory: Memory;
7+
let storage: InMemoryStorage;
78
let client: OpenRouter;
89

910
beforeEach(() => {
10-
memory = new Memory(new InMemoryStorage());
11+
storage = new InMemoryStorage();
12+
memory = new Memory(storage);
1113
client = new OpenRouter({
1214
apiKey,
1315
memory,
@@ -85,16 +87,19 @@ describe("Memory Integration E2E Tests", () => {
8587
]);
8688
await memory.updateThreadWorkingMemory("thread-1", { test: "data" });
8789

88-
// Serialize
89-
const state = await memory.serialize();
90+
// Serialize using storage
91+
const state = await storage.serialize();
9092
expect(state.threads).toHaveLength(1);
9193
expect(state.messages).toHaveLength(1);
9294
expect(state.resources).toHaveLength(1);
9395
expect(state.threadWorkingMemories).toHaveLength(1);
9496

95-
// Create new memory and hydrate
96-
const newMemory = new Memory(new InMemoryStorage());
97-
await newMemory.hydrate(state);
97+
// Create new storage and hydrate
98+
const newStorage = new InMemoryStorage();
99+
await newStorage.hydrate(state);
100+
101+
// Create new memory with hydrated storage
102+
const newMemory = new Memory(newStorage);
98103

99104
// Verify data was restored
100105
const thread = await newMemory.getThread("thread-1");
@@ -115,14 +120,18 @@ describe("Memory Integration E2E Tests", () => {
115120
{ role: "assistant" as const, content: "Hi!" },
116121
]);
117122

118-
const threadState = await memory.serializeThread("thread-1");
123+
// Serialize using storage
124+
const threadState = await storage.serializeThread("thread-1");
119125
expect(threadState).toBeDefined();
120126
expect(threadState?.thread.id).toBe("thread-1");
121127
expect(threadState?.messages).toHaveLength(2);
122128

123-
// Hydrate into new memory
124-
const newMemory = new Memory(new InMemoryStorage());
125-
await newMemory.hydrateThread(threadState!);
129+
// Hydrate into new storage
130+
const newStorage = new InMemoryStorage();
131+
await newStorage.hydrateThread(threadState!);
132+
133+
// Create new memory with hydrated storage
134+
const newMemory = new Memory(newStorage);
126135

127136
const thread = await newMemory.getThread("thread-1");
128137
expect(thread?.title).toBe("Test Thread");

0 commit comments

Comments
 (0)