diff --git a/src/memory/__tests__/knowledge-graph.test.ts b/src/memory/__tests__/knowledge-graph.test.ts index 7edab5e42c..68007bdbcf 100644 --- a/src/memory/__tests__/knowledge-graph.test.ts +++ b/src/memory/__tests__/knowledge-graph.test.ts @@ -336,16 +336,19 @@ describe('KnowledgeGraphManager', () => { expect(result.entities.map(e => e.name)).toContain('Bob'); }); - it('should include relations between opened nodes', async () => { + it('should include relations between opened nodes and outgoing relations', 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'); + // Should include: Alice->Bob (between opened nodes) and Bob->Charlie (outgoing from opened node) + expect(result.relations).toHaveLength(2); + expect(result.relations.map(r => r.from).sort()).toEqual(['Alice', 'Bob']); }); - it('should exclude relations to unopened nodes', async () => { + it('should return outgoing relations from opened nodes', async () => { const result = await manager.openNodes(['Bob']); - expect(result.relations).toHaveLength(0); + // Opening 'Bob' should return its outgoing relation to 'Charlie' + expect(result.relations).toHaveLength(1); + expect(result.relations[0].from).toBe('Bob'); + expect(result.relations[0].to).toBe('Charlie'); }); it('should handle opening non-existent nodes', async () => { diff --git a/src/memory/index.ts b/src/memory/index.ts index 600a7edcc8..edd2bbb505 100644 --- a/src/memory/index.ts +++ b/src/memory/index.ts @@ -219,9 +219,12 @@ export class KnowledgeGraphManager { // 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 + // Filter relations to include: + // 1. Relations between filtered entities (both endpoints in the list) + // 2. Outgoing relations from filtered entities (from is in the list) const filteredRelations = graph.relations.filter(r => - filteredEntityNames.has(r.from) && filteredEntityNames.has(r.to) + (filteredEntityNames.has(r.from) && filteredEntityNames.has(r.to)) || + filteredEntityNames.has(r.from) ); const filteredGraph: KnowledgeGraph = { @@ -447,7 +450,7 @@ server.registerTool( "open_nodes", { title: "Open Nodes", - description: "Open specific nodes in the knowledge graph by their names", + description: "Open specific nodes in the knowledge graph by their names. Returns entities and their outgoing relations (relations where the entity is the 'from' endpoint).", inputSchema: { names: z.array(z.string()).describe("An array of entity names to retrieve") },