diff --git a/src/memory/__tests__/knowledge-graph.test.ts b/src/memory/__tests__/knowledge-graph.test.ts index 7edab5e42c..eebf2319c9 100644 --- a/src/memory/__tests__/knowledge-graph.test.ts +++ b/src/memory/__tests__/knowledge-graph.test.ts @@ -336,16 +336,22 @@ describe('KnowledgeGraphManager', () => { expect(result.entities.map(e => e.name)).toContain('Bob'); }); - it('should include relations between opened nodes', async () => { + it('should include all relations connected to opened nodes', async () => { const result = await manager.openNodes(['Alice', 'Bob']); - expect(result.relations).toHaveLength(1); - expect(result.relations[0].from).toBe('Alice'); - expect(result.relations[0].to).toBe('Bob'); + expect(result.relations).toHaveLength(2); + // Alice -> Bob (Alice and Bob both opened) + expect(result.relations).toContainEqual({ from: 'Alice', to: 'Bob', relationType: 'knows' }); + // Bob -> Charlie (Bob is opened, so its outgoing relations are included) + expect(result.relations).toContainEqual({ from: 'Bob', to: 'Charlie', relationType: 'knows' }); }); - it('should exclude relations to unopened nodes', async () => { + it('should include both incoming and outgoing relations', async () => { const result = await manager.openNodes(['Bob']); - expect(result.relations).toHaveLength(0); + expect(result.relations).toHaveLength(2); + // Alice -> Bob (incoming) + expect(result.relations).toContainEqual({ from: 'Alice', to: 'Bob', relationType: 'knows' }); + // Bob -> Charlie (outgoing) + expect(result.relations).toContainEqual({ from: 'Bob', to: 'Charlie', relationType: 'knows' }); }); it('should handle opening non-existent nodes', async () => { diff --git a/src/memory/index.ts b/src/memory/index.ts index 600a7edcc8..69f4df0b6f 100644 --- a/src/memory/index.ts +++ b/src/memory/index.ts @@ -212,23 +212,24 @@ export class KnowledgeGraphManager { async openNodes(names: string[]): Promise { const graph = await this.loadGraph(); - + // Filter entities const filteredEntities = graph.entities.filter(e => names.includes(e.name)); - + // Create a Set of filtered entity names for quick lookup const filteredEntityNames = new Set(filteredEntities.map(e => e.name)); - - // Filter relations to only include those between filtered entities - const filteredRelations = graph.relations.filter(r => - filteredEntityNames.has(r.from) && filteredEntityNames.has(r.to) + + // Filter relations to include those connected to any of the filtered entities + // This includes both outgoing relations (where from matches) and incoming relations (where to matches) + const filteredRelations = graph.relations.filter(r => + filteredEntityNames.has(r.from) || filteredEntityNames.has(r.to) ); - + const filteredGraph: KnowledgeGraph = { entities: filteredEntities, relations: filteredRelations, }; - + return filteredGraph; } }