@brainbank/memory 0.1.0 → 0.2.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/README.md CHANGED
@@ -1,13 +1,14 @@
1
1
  # @brainbank/memory
2
2
 
3
- Deterministic memory extraction and deduplication for LLM conversations. Inspired by [mem0](https://github.com/mem0ai/mem0)'s pipeline.
3
+ Deterministic memory extraction, deduplication, and entity graph for LLM conversations. Framework-agnostic works with any LLM provider.
4
4
 
5
5
  After every conversation turn, automatically:
6
6
 
7
- 1. **Extract** atomic facts via LLM call
7
+ 1. **Extract** atomic facts + entities + relationships via LLM call
8
8
  2. **Search** existing memories for duplicates
9
9
  3. **Decide** ADD / UPDATE / NONE per fact
10
- 4. **Execute** the operations
10
+ 4. **Upsert** entities and relationships into the knowledge graph
11
+ 5. **Execute** the operations
11
12
 
12
13
  No function calling. No relying on the model to "remember" to save.
13
14
 
@@ -31,11 +32,11 @@ const memory = new Memory(brain.collection('memories'), {
31
32
  });
32
33
 
33
34
  // After every conversation turn — deterministic, automatic
34
- const ops = await memory.process(
35
+ const result = await memory.process(
35
36
  'My name is Berna, I prefer TypeScript',
36
37
  'Nice to meet you Berna!'
37
38
  );
38
- // ops → [
39
+ // result.operations → [
39
40
  // { fact: "User's name is Berna", action: "ADD", reason: "no similar memories" },
40
41
  // { fact: "User prefers TypeScript", action: "ADD", reason: "no similar memories" }
41
42
  // ]
@@ -45,7 +46,7 @@ await memory.process(
45
46
  'I like TypeScript a lot',
46
47
  'TypeScript is great!'
47
48
  );
48
- // → [{ fact: "User likes TypeScript", action: "NONE", reason: "already captured" }]
49
+ // → operations: [{ fact: "User likes TypeScript", action: "NONE", reason: "already captured" }]
49
50
 
50
51
  // Build system prompt context
51
52
  const context = memory.buildContext();
@@ -55,6 +56,72 @@ const context = memory.buildContext();
55
56
  const results = await memory.search('what language does user prefer');
56
57
  ```
57
58
 
59
+ ## Entity Extraction (Knowledge Graph)
60
+
61
+ Opt-in entity and relationship extraction from the same LLM call — zero extra cost:
62
+
63
+ ```typescript
64
+ import { Memory, EntityStore, OpenAIProvider } from '@brainbank/memory';
65
+
66
+ const entityStore = new EntityStore({
67
+ entityCollection: brain.collection('entities'),
68
+ relationCollection: brain.collection('relationships'),
69
+ });
70
+
71
+ const memory = new Memory(brain.collection('memories'), {
72
+ llm: new OpenAIProvider({ model: 'gpt-4.1-nano' }),
73
+ entityStore, // opt-in — omit for facts-only mode
74
+ });
75
+
76
+ // Process extracts facts + entities + relationships in one LLM call
77
+ const result = await memory.process(
78
+ 'Tell Juan to migrate payments to Stripe before Friday',
79
+ "I'll let Juan know about the Stripe migration deadline."
80
+ );
81
+ // result.operations → [{ fact: "deadline for Stripe migration is Friday", action: "ADD" }]
82
+ // result.entities → { entitiesProcessed: 2, relationshipsProcessed: 1 }
83
+
84
+ // Query entities
85
+ const related = await entityStore.getRelated('Juan');
86
+ // → [{ source: "Juan", target: "Stripe", relation: "migrating_to" }]
87
+
88
+ // Build context includes entities
89
+ const context = memory.buildContext();
90
+ // → "## Memories\n- ...\n\n## Known Entities\n- Juan (person, 2x)\n- Stripe (service, 1x)\n\n## Relationships\n- Juan → migrating_to → Stripe"
91
+ ```
92
+
93
+ ### EntityStore API
94
+
95
+ | Method | Description |
96
+ |--------|-------------|
97
+ | `upsert(entity)` | Add or update entity (increments mention count) |
98
+ | `relate(source, target, relation, context?)` | Add a relationship |
99
+ | `findEntity(name)` | Search entities by name (semantic) |
100
+ | `getRelated(entityName)` | Get all relationships for an entity |
101
+ | `relationsOf(entityName)` | Shorthand for `getRelated()` |
102
+ | `listEntities({ type?, limit? })` | List entities, optionally filtered by type |
103
+ | `listRelationships()` | List all relationships |
104
+ | `traverse(entity, maxDepth?)` | Multi-hop BFS graph traversal (default: 2 hops) |
105
+ | `entityCount()` | Total entity count |
106
+ | `relationCount()` | Total relationship count |
107
+ | `buildContext(entityName?)` | Build markdown context (all or specific entity) |
108
+ | `processExtraction(entities, relationships)` | Batch process from LLM response |
109
+
110
+ #### Graph Traversal
111
+
112
+ ```typescript
113
+ // Explore the entity graph from a starting point
114
+ const graph = await entityStore.traverse('Juan', 2);
115
+ // graph.nodes → [
116
+ // { entity: "Stripe", relation: "migrating_to", depth: 1, path: ["Juan", "Stripe"] },
117
+ // { entity: "Payments", relation: "uses", depth: 2, path: ["Juan", "Stripe", "Payments"] }
118
+ // ]
119
+
120
+ // Filter entities by type
121
+ const people = entityStore.listEntities({ type: 'person' });
122
+ const services = entityStore.listEntities({ type: 'service' });
123
+ ```
124
+
58
125
  ## Framework Integration
59
126
 
60
127
  The `LLMProvider` interface is framework-agnostic. Bring your own LLM:
@@ -134,6 +201,7 @@ const memory = new Memory(store, { llm });
134
201
  ```typescript
135
202
  new Memory(store, {
136
203
  llm: provider, // required — LLM provider
204
+ entityStore: entityStore, // optional — enables entity extraction
137
205
  maxFacts: 5, // max facts to extract per turn (default: 5)
138
206
  maxMemories: 50, // max existing memories to load for dedup (default: 50)
139
207
  dedupTopK: 3, // similar memories to compare against (default: 3)
@@ -149,11 +217,12 @@ new Memory(store, {
149
217
 
150
218
  | Method | Description |
151
219
  |--------|-------------|
152
- | `process(userMsg, assistantMsg)` | Run the full pipeline: extract → dedup → execute. Returns `MemoryOperation[]` |
220
+ | `process(userMsg, assistantMsg)` | Full pipeline: extract → dedup → execute. Returns `ProcessResult` |
153
221
  | `search(query, k?)` | Semantic search across memories |
154
222
  | `recall(limit?)` | Get all memories (for system prompt injection) |
155
223
  | `count()` | Total stored memories |
156
- | `buildContext(limit?)` | Build a markdown section for system prompt injection |
224
+ | `buildContext(limit?)` | Build markdown context (memories + entities if enabled) |
225
+ | `getEntityStore()` | Get the entity store instance (if enabled) |
157
226
 
158
227
  ## How it works
159
228
 
@@ -161,27 +230,27 @@ new Memory(store, {
161
230
  User message + Assistant response
162
231
 
163
232
 
164
- ┌─── Extract (LLM) ───┐
165
- "User's name is X"
166
- "Prefers TypeScript"
167
- └──────────┬───────────┘
168
- for each fact:
169
-
170
- ┌─── Search (semantic) ─┐
171
- Find similar existing
172
- │ memories (top-K) │
173
- └──────────┬────────────┘
174
-
175
-
176
- ┌─── Dedup (LLM) ──────┐
177
- │ Compare new vs existing│
178
- │ → ADD / UPDATE / NONE │
179
- └──────────┬────────────┘
233
+ ┌─── Extract (LLM) ──────────┐
234
+ Facts:
235
+ "User's name is X"
236
+ │ "Prefers TypeScript" │
237
+ Entities:
238
+ │ X (person), TypeScript │
239
+ Relationships: │
240
+ X prefers → TypeScript
241
+ └──────────┬──────────────────┘
180
242
 
181
- ┌───────┼───────┐
182
-
183
- ADD UPDATE NONE
184
- (store) (replace) (skip)
243
+ ┌───────┴───────┐
244
+
245
+ Facts Entities
246
+ │ │
247
+ ▼ ▼
248
+ ┌─ Dedup ──┐ ┌─ Upsert ─┐
249
+ │ ADD │ │ name │
250
+ │ UPDATE │ │ type │
251
+ │ NONE │ │ mentions │
252
+ └──────────┘ │ relate │
253
+ └──────────┘
185
254
  ```
186
255
 
187
256
  ## License
package/dist/index.d.ts CHANGED
@@ -13,6 +13,8 @@ interface GenerateOptions {
13
13
  json?: boolean;
14
14
  /** Max tokens for response */
15
15
  maxTokens?: number;
16
+ /** Temperature for response (0 = deterministic) */
17
+ temperature?: number;
16
18
  }
17
19
  /**
18
20
  * LLM provider interface. Implement this to bring your own model.
@@ -65,10 +67,183 @@ declare class OpenAIProvider implements LLMProvider {
65
67
  generate(messages: ChatMessage[], options?: GenerateOptions): Promise<string>;
66
68
  }
67
69
 
70
+ /**
71
+ * @brainbank/memory — Entity Store
72
+ *
73
+ * Manages entities and relationships extracted from conversations.
74
+ * Uses BrainBank collections (SQLite) — no Neo4j or external graph DB needed.
75
+ * Optional LLM for intelligent entity resolution (e.g. merging "TS" with "TypeScript").
76
+ */
77
+
78
+ interface Entity {
79
+ name: string;
80
+ type: 'person' | 'service' | 'project' | 'organization' | 'concept' | string;
81
+ attributes?: Record<string, any>;
82
+ firstSeen?: number;
83
+ lastSeen?: number;
84
+ mentionCount?: number;
85
+ }
86
+ interface Relationship {
87
+ source: string;
88
+ target: string;
89
+ relation: string;
90
+ context?: string;
91
+ timestamp?: number;
92
+ }
93
+ interface TraversalNode {
94
+ entity: string;
95
+ relation: string;
96
+ depth: number;
97
+ path: string[];
98
+ }
99
+ interface TraversalResult {
100
+ start: string;
101
+ maxDepth: number;
102
+ nodes: TraversalNode[];
103
+ }
104
+ /** Object with a .collection() method — satisfied by BrainBank */
105
+ interface CollectionProvider {
106
+ collection(name: string): MemoryStore;
107
+ }
108
+ interface EntityStoreConfig {
109
+ /** Optional LLM for intelligent entity resolution */
110
+ llm?: LLMProvider;
111
+ /** Callback fired for each entity operation */
112
+ onEntity?: (op: {
113
+ action: 'NEW' | 'UPDATED' | 'RELATED';
114
+ name: string;
115
+ type?: string;
116
+ detail?: string;
117
+ }) => void;
118
+ /** Custom entity collection name. Default: 'entities' */
119
+ entityCollectionName?: string;
120
+ /** Custom relationship collection name. Default: 'relationships' */
121
+ relationCollectionName?: string;
122
+ }
123
+ /**
124
+ * @deprecated Use `new EntityStore(brain, config?)` instead.
125
+ */
126
+ interface EntityStoreOptions {
127
+ /** Collection for entities */
128
+ entityCollection: MemoryStore;
129
+ /** Collection for relationships */
130
+ relationCollection: MemoryStore;
131
+ /** Optional LLM for intelligent entity resolution */
132
+ llm?: LLMProvider;
133
+ /** Callback fired for each entity operation */
134
+ onEntity?: (op: {
135
+ action: 'NEW' | 'UPDATED' | 'RELATED';
136
+ name: string;
137
+ type?: string;
138
+ detail?: string;
139
+ }) => void;
140
+ }
141
+ declare class EntityStore {
142
+ private readonly entities;
143
+ private readonly relations;
144
+ private llm?;
145
+ private readonly onEntity?;
146
+ /**
147
+ * Create an EntityStore.
148
+ *
149
+ * @example Simple (recommended)
150
+ * ```typescript
151
+ * const entityStore = new EntityStore(brain);
152
+ * ```
153
+ *
154
+ * @example With config
155
+ * ```typescript
156
+ * const entityStore = new EntityStore(brain, {
157
+ * onEntity: (op) => console.log(op),
158
+ * });
159
+ * ```
160
+ */
161
+ constructor(provider: CollectionProvider, config?: EntityStoreConfig);
162
+ /** @deprecated Pass a CollectionProvider (brain) instead */
163
+ constructor(options: EntityStoreOptions);
164
+ /** @internal — used by Memory to share its LLM if EntityStore doesn't have one */
165
+ setLLM(llm: LLMProvider): void;
166
+ /**
167
+ * Upsert an entity — create if new, update mention count if exists.
168
+ */
169
+ upsert(entity: Entity): Promise<void>;
170
+ /**
171
+ * Add a relationship between two entities.
172
+ */
173
+ relate(source: string, target: string, relation: string, context?: string): Promise<void>;
174
+ /**
175
+ * Find an entity by name.
176
+ * 1. Exact name match (case-insensitive)
177
+ * 2. LLM resolution if available (e.g. "TS" → "TypeScript")
178
+ */
179
+ findEntity(name: string): Promise<(MemoryItem & {
180
+ metadata?: Record<string, any>;
181
+ }) | null>;
182
+ /**
183
+ * Ask LLM if a new entity matches any existing entity.
184
+ * Returns the matching entity name or null.
185
+ */
186
+ private resolveEntity;
187
+ /**
188
+ * Get all relationships for an entity (as source or target).
189
+ */
190
+ getRelated(entityName: string): Promise<Relationship[]>;
191
+ /**
192
+ * Get all relationships for an entity (shorthand for getRelated).
193
+ */
194
+ relationsOf(entityName: string): Promise<Relationship[]>;
195
+ /**
196
+ * List all entities, optionally filtered by type.
197
+ */
198
+ listEntities(options?: {
199
+ type?: string;
200
+ limit?: number;
201
+ }): MemoryItem[];
202
+ /**
203
+ * List all relationships.
204
+ */
205
+ listRelationships(): MemoryItem[];
206
+ /**
207
+ * Get entity count.
208
+ */
209
+ entityCount(): number;
210
+ /**
211
+ * Get relationship count.
212
+ */
213
+ relationCount(): number;
214
+ /**
215
+ * Traverse the entity graph — multi-hop BFS from a starting entity.
216
+ * Returns all reachable entities within the given depth.
217
+ */
218
+ traverse(startEntity: string, maxDepth?: number): Promise<TraversalResult>;
219
+ /**
220
+ * Build markdown context for system prompt injection.
221
+ */
222
+ buildContext(entityName?: string): string;
223
+ /**
224
+ * Process raw entities and relationships from LLM extraction.
225
+ */
226
+ processExtraction(entities: Array<{
227
+ name: string;
228
+ type: string;
229
+ attributes?: Record<string, any>;
230
+ }>, relationships: Array<{
231
+ source: string;
232
+ target: string;
233
+ relation: string;
234
+ }>, context?: string): Promise<{
235
+ entitiesProcessed: number;
236
+ relationshipsProcessed: number;
237
+ }>;
238
+ private serializeEntity;
239
+ private extractName;
240
+ }
241
+
68
242
  /**
69
243
  * @brainbank/memory — Deterministic Memory Pipeline
70
244
  *
71
245
  * Automatic fact extraction and deduplication for LLM conversations.
246
+ * Optionally extracts entities and relationships (knowledge graph).
72
247
  * Runs after every turn: extract → search → dedup → ADD/UPDATE/NONE.
73
248
  */
74
249
 
@@ -84,6 +259,14 @@ interface MemoryOperation {
84
259
  action: MemoryAction;
85
260
  reason: string;
86
261
  }
262
+ interface EntityOperation {
263
+ entitiesProcessed: number;
264
+ relationshipsProcessed: number;
265
+ }
266
+ interface ProcessResult {
267
+ operations: MemoryOperation[];
268
+ entities?: EntityOperation;
269
+ }
87
270
  /**
88
271
  * Collection interface — matches BrainBank's collection API.
89
272
  * Implement this to use a different storage backend.
@@ -105,6 +288,8 @@ interface MemoryStore {
105
288
  interface MemoryOptions {
106
289
  /** LLM provider for extraction and dedup */
107
290
  llm: LLMProvider;
291
+ /** Entity store for knowledge graph (opt-in) */
292
+ entityStore?: EntityStore;
108
293
  /** Max facts to extract per turn. Default: 5 */
109
294
  maxFacts?: number;
110
295
  /** Max existing memories to compare against for dedup. Default: 50 */
@@ -117,22 +302,39 @@ interface MemoryOptions {
117
302
  dedupPrompt?: string;
118
303
  /** Called for each memory operation */
119
304
  onOperation?: (op: MemoryOperation) => void;
305
+ /** Custom collection name for memories. Default: 'memories' */
306
+ collectionName?: string;
120
307
  }
121
308
  declare class Memory {
122
309
  private readonly store;
123
310
  private readonly llm;
311
+ private readonly entityStore?;
124
312
  private readonly maxFacts;
125
313
  private readonly maxMemories;
126
314
  private readonly dedupTopK;
127
315
  private readonly extractPrompt;
128
316
  private readonly dedupPrompt;
129
317
  private readonly onOperation?;
318
+ /**
319
+ * Create a Memory instance.
320
+ *
321
+ * @example Recommended (pass brain directly)
322
+ * ```typescript
323
+ * const memory = new Memory(brain, { llm });
324
+ * ```
325
+ *
326
+ * @example Legacy (pass collection directly)
327
+ * ```typescript
328
+ * const memory = new Memory(brain.collection('memories'), { llm });
329
+ * ```
330
+ */
331
+ constructor(provider: CollectionProvider, options: MemoryOptions);
130
332
  constructor(store: MemoryStore, options: MemoryOptions);
131
333
  /**
132
- * Process a conversation turn — extract facts and store/update memories.
334
+ * Process a conversation turn — extract facts (and optionally entities) and store/update memories.
133
335
  * This is the main entry point. Call after every user↔assistant exchange.
134
336
  */
135
- process(userMessage: string, assistantMessage: string): Promise<MemoryOperation[]>;
337
+ process(userMessage: string, assistantMessage: string): Promise<ProcessResult>;
136
338
  /**
137
339
  * Search memories semantically.
138
340
  */
@@ -146,12 +348,16 @@ declare class Memory {
146
348
  */
147
349
  count(): number;
148
350
  /**
149
- * Build a system prompt section with all memories.
351
+ * Build a system prompt section with all memories (and entities if enabled).
150
352
  * Drop this into your system prompt.
151
353
  */
152
354
  buildContext(limit?: number): string;
355
+ /**
356
+ * Get entity store (if enabled).
357
+ */
358
+ getEntityStore(): EntityStore | undefined;
153
359
  private extract;
154
360
  private dedup;
155
361
  }
156
362
 
157
- export { type ChatMessage, type GenerateOptions, type LLMProvider, Memory, type MemoryAction, type MemoryItem, type MemoryOperation, type MemoryOptions, type MemoryStore, OpenAIProvider, type OpenAIProviderOptions };
363
+ export { type ChatMessage, type CollectionProvider, type Entity, type EntityOperation, EntityStore, type EntityStoreConfig, type EntityStoreOptions, type GenerateOptions, type LLMProvider, Memory, type MemoryAction, type MemoryItem, type MemoryOperation, type MemoryOptions, type MemoryStore, OpenAIProvider, type OpenAIProviderOptions, type ProcessResult, type Relationship, type TraversalNode, type TraversalResult };
package/dist/index.js CHANGED
@@ -15,6 +15,29 @@ Rules:
15
15
  - Be specific ("prefers TypeScript" not "has programming preferences")
16
16
  - Skip trivial info ("said hello", "asked a question")
17
17
  - Max 5 facts per turn`;
18
+ var EXTRACT_WITH_ENTITIES_PROMPT = `You are a memory extraction engine. Given a conversation turn, extract:
19
+ 1. Atomic facts worth remembering
20
+ 2. Named entities (people, services, projects, organizations, concepts)
21
+ 3. Relationships between entities
22
+
23
+ Respond with JSON:
24
+ {
25
+ "facts": ["fact1", "fact2"],
26
+ "entities": [
27
+ { "name": "EntityName", "type": "person|service|project|organization|concept", "attributes": {} }
28
+ ],
29
+ "relationships": [
30
+ { "source": "EntityA", "target": "EntityB", "relation": "verb_phrase" }
31
+ ]
32
+ }
33
+
34
+ Rules:
35
+ - Facts: single self-contained sentences, be specific, max 5 per turn
36
+ - Entities: only extract clearly named entities, not generic nouns
37
+ - Relationships: use lowercase verb phrases ("works_on", "prefers", "depends_on", "migrating_to")
38
+ - Entity types: person, service, project, organization, concept (or custom)
39
+ - If nothing found, return empty arrays
40
+ - Skip trivial info`;
18
41
  var DEDUP_PROMPT = `You are a memory deduplication engine. Given a NEW fact and a list of EXISTING memories, decide what action to take.
19
42
 
20
43
  Respond with JSON: { "action": "ADD" | "UPDATE" | "NONE", "reason": "brief reason" }
@@ -29,32 +52,41 @@ Be conservative \u2014 if in doubt, say NONE.`;
29
52
  var Memory = class {
30
53
  store;
31
54
  llm;
55
+ entityStore;
32
56
  maxFacts;
33
57
  maxMemories;
34
58
  dedupTopK;
35
59
  extractPrompt;
36
60
  dedupPrompt;
37
61
  onOperation;
38
- constructor(store, options) {
39
- this.store = store;
62
+ constructor(providerOrStore, options) {
63
+ if ("collection" in providerOrStore && typeof providerOrStore.collection === "function") {
64
+ this.store = providerOrStore.collection(options.collectionName ?? "memories");
65
+ } else {
66
+ this.store = providerOrStore;
67
+ }
40
68
  this.llm = options.llm;
69
+ this.entityStore = options.entityStore;
41
70
  this.maxFacts = options.maxFacts ?? 5;
42
71
  this.maxMemories = options.maxMemories ?? 50;
43
72
  this.dedupTopK = options.dedupTopK ?? 3;
44
- this.extractPrompt = options.extractPrompt ?? EXTRACT_PROMPT;
73
+ this.extractPrompt = options.extractPrompt ?? (options.entityStore ? EXTRACT_WITH_ENTITIES_PROMPT : EXTRACT_PROMPT);
45
74
  this.dedupPrompt = options.dedupPrompt ?? DEDUP_PROMPT;
46
75
  this.onOperation = options.onOperation;
76
+ if (this.entityStore) this.entityStore.setLLM(this.llm);
47
77
  }
48
78
  /**
49
- * Process a conversation turn — extract facts and store/update memories.
79
+ * Process a conversation turn — extract facts (and optionally entities) and store/update memories.
50
80
  * This is the main entry point. Call after every user↔assistant exchange.
51
81
  */
52
82
  async process(userMessage, assistantMessage) {
53
- const facts = await this.extract(userMessage, assistantMessage);
54
- if (facts.length === 0) return [];
83
+ const extraction = await this.extract(userMessage, assistantMessage);
84
+ if (extraction.facts.length === 0 && extraction.entities.length === 0) {
85
+ return { operations: [] };
86
+ }
55
87
  const existing = this.store.list({ limit: this.maxMemories });
56
88
  const operations = [];
57
- for (const fact of facts) {
89
+ for (const fact of extraction.facts) {
58
90
  const op = await this.dedup(fact, existing);
59
91
  operations.push(op);
60
92
  this.onOperation?.(op);
@@ -75,7 +107,16 @@ var Memory = class {
75
107
  break;
76
108
  }
77
109
  }
78
- return operations;
110
+ let entityOps;
111
+ if (this.entityStore && (extraction.entities.length > 0 || extraction.relationships.length > 0)) {
112
+ const context = `${userMessage} \u2014 ${assistantMessage}`.slice(0, 200);
113
+ entityOps = await this.entityStore.processExtraction(
114
+ extraction.entities,
115
+ extraction.relationships,
116
+ context
117
+ );
118
+ }
119
+ return { operations, entities: entityOps };
79
120
  }
80
121
  /**
81
122
  * Search memories semantically.
@@ -96,13 +137,26 @@ var Memory = class {
96
137
  return this.store.count();
97
138
  }
98
139
  /**
99
- * Build a system prompt section with all memories.
140
+ * Build a system prompt section with all memories (and entities if enabled).
100
141
  * Drop this into your system prompt.
101
142
  */
102
143
  buildContext(limit = 20) {
144
+ const parts = [];
103
145
  const items = this.store.list({ limit });
104
- if (items.length === 0) return "";
105
- return "## Memories\n" + items.map((m) => `- ${m.content}`).join("\n");
146
+ if (items.length > 0) {
147
+ parts.push("## Memories\n" + items.map((m) => `- ${m.content}`).join("\n"));
148
+ }
149
+ if (this.entityStore) {
150
+ const entityCtx = this.entityStore.buildContext();
151
+ if (entityCtx) parts.push(entityCtx);
152
+ }
153
+ return parts.join("\n\n");
154
+ }
155
+ /**
156
+ * Get entity store (if enabled).
157
+ */
158
+ getEntityStore() {
159
+ return this.entityStore;
106
160
  }
107
161
  // ─── Internal ───────────────────────────────────
108
162
  async extract(userMsg, assistantMsg) {
@@ -111,13 +165,16 @@ var Memory = class {
111
165
  { role: "user", content: `User: ${userMsg}
112
166
 
113
167
  Assistant: ${assistantMsg}` }
114
- ], { json: true, maxTokens: 300 });
168
+ ], { json: true, maxTokens: 500 });
115
169
  try {
116
170
  const parsed = JSON.parse(response);
117
- const facts = parsed.facts ?? [];
118
- return facts.slice(0, this.maxFacts);
171
+ return {
172
+ facts: (parsed.facts ?? []).slice(0, this.maxFacts),
173
+ entities: parsed.entities ?? [],
174
+ relationships: parsed.relationships ?? []
175
+ };
119
176
  } catch {
120
- return [];
177
+ return { facts: [], entities: [], relationships: [] };
121
178
  }
122
179
  }
123
180
  async dedup(fact, _existing) {
@@ -146,6 +203,287 @@ ${context}` }
146
203
  }
147
204
  };
148
205
 
206
+ // src/entities.ts
207
+ var ENTITY_RESOLVE_PROMPT = `You are an entity resolution engine. Given a NEW entity and a list of EXISTING entities, determine if the new entity refers to the same real-world thing as any existing one.
208
+
209
+ Consider aliases, abbreviations, alternate names, and typos:
210
+ - "TS" = "TypeScript"
211
+ - "JS" = "JavaScript"
212
+ - "berna" = "Berna"
213
+ - "GCP" = "Google Cloud Platform"
214
+ - "React.js" = "React"
215
+
216
+ Respond with ONLY the matching entity name (exactly as listed) or "NONE" if no match.
217
+
218
+ Existing entities:
219
+ {entities}
220
+
221
+ New entity: {newEntity}
222
+
223
+ Match:`;
224
+ var EntityStore = class {
225
+ entities;
226
+ relations;
227
+ llm;
228
+ onEntity;
229
+ constructor(providerOrOptions, config) {
230
+ if ("collection" in providerOrOptions && typeof providerOrOptions.collection === "function") {
231
+ const c = config ?? {};
232
+ this.entities = providerOrOptions.collection(c.entityCollectionName ?? "entities");
233
+ this.relations = providerOrOptions.collection(c.relationCollectionName ?? "relationships");
234
+ this.llm = c.llm;
235
+ this.onEntity = c.onEntity;
236
+ } else {
237
+ const opts = providerOrOptions;
238
+ this.entities = opts.entityCollection;
239
+ this.relations = opts.relationCollection;
240
+ this.llm = opts.llm;
241
+ this.onEntity = opts.onEntity;
242
+ }
243
+ }
244
+ /** @internal — used by Memory to share its LLM if EntityStore doesn't have one */
245
+ setLLM(llm) {
246
+ if (!this.llm) this.llm = llm;
247
+ }
248
+ /**
249
+ * Upsert an entity — create if new, update mention count if exists.
250
+ */
251
+ async upsert(entity) {
252
+ const now = Date.now();
253
+ const existing = await this.findEntity(entity.name);
254
+ if (existing) {
255
+ if (existing.id != null) await this.entities.remove(existing.id);
256
+ await this.entities.add(this.serializeEntity(entity), {
257
+ metadata: {
258
+ type: entity.type,
259
+ attributes: { ...existing.metadata?.attributes ?? {}, ...entity.attributes ?? {} },
260
+ firstSeen: existing.metadata?.firstSeen ?? now,
261
+ lastSeen: now,
262
+ mentionCount: (existing.metadata?.mentionCount ?? 1) + 1
263
+ }
264
+ });
265
+ this.onEntity?.({ action: "UPDATED", name: entity.name, type: entity.type, detail: `${(existing.metadata?.mentionCount ?? 1) + 1}x` });
266
+ } else {
267
+ await this.entities.add(this.serializeEntity(entity), {
268
+ metadata: {
269
+ type: entity.type,
270
+ attributes: entity.attributes ?? {},
271
+ firstSeen: now,
272
+ lastSeen: now,
273
+ mentionCount: 1
274
+ }
275
+ });
276
+ this.onEntity?.({ action: "NEW", name: entity.name, type: entity.type });
277
+ }
278
+ }
279
+ /**
280
+ * Add a relationship between two entities.
281
+ */
282
+ async relate(source, target, relation, context) {
283
+ const now = Date.now();
284
+ const content = `${source} \u2192 ${relation} \u2192 ${target}`;
285
+ await this.relations.add(content, {
286
+ metadata: { source, target, relation, context, timestamp: now }
287
+ });
288
+ this.onEntity?.({ action: "RELATED", name: source, detail: `${source} \u2192 ${relation} \u2192 ${target}` });
289
+ }
290
+ /**
291
+ * Find an entity by name.
292
+ * 1. Exact name match (case-insensitive)
293
+ * 2. LLM resolution if available (e.g. "TS" → "TypeScript")
294
+ */
295
+ async findEntity(name) {
296
+ const results = await this.entities.search(name, { k: 5 });
297
+ for (const r of results) {
298
+ if (this.extractName(r.content).toLowerCase() === name.toLowerCase()) {
299
+ return r;
300
+ }
301
+ }
302
+ if (this.llm && results.length > 0) {
303
+ const candidateNames = results.map((r) => this.extractName(r.content));
304
+ const resolved = await this.resolveEntity(name, candidateNames);
305
+ if (resolved) {
306
+ return results.find((r) => this.extractName(r.content) === resolved) ?? null;
307
+ }
308
+ }
309
+ return null;
310
+ }
311
+ /**
312
+ * Ask LLM if a new entity matches any existing entity.
313
+ * Returns the matching entity name or null.
314
+ */
315
+ async resolveEntity(newEntity, existing) {
316
+ if (!this.llm || existing.length === 0) return null;
317
+ const prompt = ENTITY_RESOLVE_PROMPT.replace("{entities}", existing.map((e) => `- ${e}`).join("\n")).replace("{newEntity}", newEntity);
318
+ try {
319
+ const response = await this.llm.generate(
320
+ [{ role: "user", content: prompt }],
321
+ { maxTokens: 50, temperature: 0 }
322
+ );
323
+ const match = response.trim();
324
+ if (match === "NONE" || match === "none") return null;
325
+ const found = existing.find((e) => e.toLowerCase() === match.toLowerCase());
326
+ return found ?? null;
327
+ } catch {
328
+ return null;
329
+ }
330
+ }
331
+ /**
332
+ * Get all relationships for an entity (as source or target).
333
+ */
334
+ async getRelated(entityName) {
335
+ const results = await this.relations.search(entityName, { k: 20 });
336
+ return results.filter((r) => r.metadata?.source === entityName || r.metadata?.target === entityName).map((r) => ({
337
+ source: r.metadata?.source ?? "",
338
+ target: r.metadata?.target ?? "",
339
+ relation: r.metadata?.relation ?? "",
340
+ context: r.metadata?.context,
341
+ timestamp: r.metadata?.timestamp
342
+ }));
343
+ }
344
+ /**
345
+ * Get all relationships for an entity (shorthand for getRelated).
346
+ */
347
+ async relationsOf(entityName) {
348
+ return this.getRelated(entityName);
349
+ }
350
+ /**
351
+ * List all entities, optionally filtered by type.
352
+ */
353
+ listEntities(options) {
354
+ const all = this.entities.list({ limit: options?.limit ?? 100 });
355
+ if (options?.type) {
356
+ return all.filter((e) => e.metadata?.type === options.type);
357
+ }
358
+ return all;
359
+ }
360
+ /**
361
+ * List all relationships.
362
+ */
363
+ listRelationships() {
364
+ return this.relations.list({ limit: 200 });
365
+ }
366
+ /**
367
+ * Get entity count.
368
+ */
369
+ entityCount() {
370
+ return this.entities.count();
371
+ }
372
+ /**
373
+ * Get relationship count.
374
+ */
375
+ relationCount() {
376
+ return this.relations.count();
377
+ }
378
+ /**
379
+ * Traverse the entity graph — multi-hop BFS from a starting entity.
380
+ * Returns all reachable entities within the given depth.
381
+ */
382
+ async traverse(startEntity, maxDepth = 2) {
383
+ const visited = /* @__PURE__ */ new Set();
384
+ const queue = [
385
+ { entity: startEntity, depth: 0, path: [startEntity], relation: "" }
386
+ ];
387
+ const nodes = [];
388
+ while (queue.length > 0) {
389
+ const current = queue.shift();
390
+ if (current.depth > maxDepth || visited.has(current.entity)) continue;
391
+ visited.add(current.entity);
392
+ const rels = await this.getRelated(current.entity);
393
+ for (const rel of rels) {
394
+ const next = rel.source === current.entity ? rel.target : rel.source;
395
+ const nextPath = [...current.path, next];
396
+ if (!visited.has(next)) {
397
+ nodes.push({
398
+ entity: next,
399
+ relation: rel.relation,
400
+ depth: current.depth + 1,
401
+ path: nextPath
402
+ });
403
+ queue.push({
404
+ entity: next,
405
+ depth: current.depth + 1,
406
+ path: nextPath,
407
+ relation: rel.relation
408
+ });
409
+ }
410
+ }
411
+ }
412
+ return { start: startEntity, maxDepth, nodes };
413
+ }
414
+ /**
415
+ * Build markdown context for system prompt injection.
416
+ */
417
+ buildContext(entityName) {
418
+ const parts = [];
419
+ if (entityName) {
420
+ const entities = this.entities.list({ limit: 100 });
421
+ const match = entities.find((e) => this.extractName(e.content).toLowerCase() === entityName.toLowerCase());
422
+ if (match) {
423
+ parts.push(`## Entity: ${this.extractName(match.content)}`);
424
+ if (match.metadata?.type) parts.push(`Type: ${match.metadata.type}`);
425
+ if (match.metadata?.mentionCount) parts.push(`Mentions: ${match.metadata.mentionCount}`);
426
+ }
427
+ const rels = this.relations.list({ limit: 200 });
428
+ const related = rels.filter(
429
+ (r) => r.metadata?.source === entityName || r.metadata?.target === entityName
430
+ );
431
+ if (related.length > 0) {
432
+ parts.push("### Relationships");
433
+ for (const r of related) {
434
+ parts.push(`- ${r.metadata?.source} \u2192 ${r.metadata?.relation} \u2192 ${r.metadata?.target}`);
435
+ }
436
+ }
437
+ } else {
438
+ const entities = this.entities.list({ limit: 50 });
439
+ if (entities.length === 0) return "";
440
+ parts.push("## Known Entities");
441
+ for (const e of entities) {
442
+ const name = this.extractName(e.content);
443
+ const type = e.metadata?.type ?? "unknown";
444
+ const mentions = e.metadata?.mentionCount ?? 1;
445
+ parts.push(`- ${name} (${type}, ${mentions}x)`);
446
+ }
447
+ const rels = this.relations.list({ limit: 50 });
448
+ if (rels.length > 0) {
449
+ parts.push("\n## Relationships");
450
+ for (const r of rels) {
451
+ parts.push(`- ${r.metadata?.source} \u2192 ${r.metadata?.relation} \u2192 ${r.metadata?.target}`);
452
+ }
453
+ }
454
+ }
455
+ return parts.join("\n");
456
+ }
457
+ /**
458
+ * Process raw entities and relationships from LLM extraction.
459
+ */
460
+ async processExtraction(entities, relationships, context) {
461
+ let entitiesProcessed = 0;
462
+ let relationshipsProcessed = 0;
463
+ for (const entity of entities) {
464
+ await this.upsert(entity);
465
+ entitiesProcessed++;
466
+ }
467
+ for (const rel of relationships) {
468
+ await this.relate(rel.source, rel.target, rel.relation, context);
469
+ relationshipsProcessed++;
470
+ }
471
+ return { entitiesProcessed, relationshipsProcessed };
472
+ }
473
+ // ─── Internal ───────────────────────────────────
474
+ serializeEntity(entity) {
475
+ const parts = [entity.name];
476
+ if (entity.type) parts.push(`(${entity.type})`);
477
+ if (entity.attributes && Object.keys(entity.attributes).length > 0) {
478
+ parts.push(JSON.stringify(entity.attributes));
479
+ }
480
+ return parts.join(" ");
481
+ }
482
+ extractName(content) {
483
+ return content.split(/\s*\(/)[0].trim();
484
+ }
485
+ };
486
+
149
487
  // src/llm.ts
150
488
  var OpenAIProvider = class {
151
489
  apiKey;
@@ -170,6 +508,7 @@ var OpenAIProvider = class {
170
508
  model: this.model,
171
509
  messages,
172
510
  max_tokens: options?.maxTokens ?? 500,
511
+ ...options?.temperature != null ? { temperature: options.temperature } : {},
173
512
  ...options?.json ? { response_format: { type: "json_object" } } : {}
174
513
  })
175
514
  });
@@ -181,6 +520,7 @@ var OpenAIProvider = class {
181
520
  }
182
521
  };
183
522
  export {
523
+ EntityStore,
184
524
  Memory,
185
525
  OpenAIProvider
186
526
  };
package/package.json CHANGED
@@ -1,22 +1,29 @@
1
1
  {
2
2
  "name": "@brainbank/memory",
3
- "version": "0.1.0",
3
+ "version": "0.2.0",
4
4
  "description": "Deterministic memory extraction and deduplication for LLM conversations — extract, dedup, ADD/UPDATE/NONE",
5
5
  "type": "module",
6
6
  "main": "dist/index.js",
7
7
  "types": "dist/index.d.ts",
8
8
  "exports": {
9
- ".": { "import": "./dist/index.js", "types": "./dist/index.d.ts" }
9
+ ".": {
10
+ "import": "./dist/index.js",
11
+ "types": "./dist/index.d.ts"
12
+ }
10
13
  },
11
- "files": ["dist/"],
14
+ "files": [
15
+ "dist/"
16
+ ],
12
17
  "scripts": {
13
18
  "build": "tsup"
14
19
  },
15
20
  "peerDependencies": {
16
- "brainbank": ">=0.2.0"
21
+ "brainbank": ">=0.1.0"
17
22
  },
18
23
  "peerDependenciesMeta": {
19
- "brainbank": { "optional": true }
24
+ "brainbank": {
25
+ "optional": true
26
+ }
20
27
  },
21
28
  "repository": {
22
29
  "type": "git",
@@ -24,8 +31,16 @@
24
31
  "directory": "packages/memory"
25
32
  },
26
33
  "keywords": [
27
- "memory", "llm", "ai", "agent", "langchain", "deduplication",
28
- "fact-extraction", "conversation-memory", "rag", "brainbank"
34
+ "memory",
35
+ "llm",
36
+ "ai",
37
+ "agent",
38
+ "langchain",
39
+ "deduplication",
40
+ "fact-extraction",
41
+ "conversation-memory",
42
+ "rag",
43
+ "brainbank"
29
44
  ],
30
45
  "author": "Bernardo Castro <bernardo@pinecall.io>",
31
46
  "license": "MIT"