diff --git a/.changeset/memory-3953-lightweight-reads.md b/.changeset/memory-3953-lightweight-reads.md new file mode 100644 index 0000000000..33ffc03834 --- /dev/null +++ b/.changeset/memory-3953-lightweight-reads.md @@ -0,0 +1,5 @@ +--- +"@modelcontextprotocol/server-memory": minor +--- + +Add optional lightweight read modes, so a large graph need not be returned whole: `observationLimit` on `read_graph`, `search_nodes` and `open_nodes` caps the observations returned per entity and adds each entity's total as `observationCount` (`0` gives an overview of names, types and counts), `entityTypes` on `read_graph` keeps only entities of those types, and `limit` on `search_nodes` caps the number of matches. Calls without them return what they always have. The `read_graph` description now says when to use it rather than `search_nodes` or `open_nodes`. diff --git a/src/memory/README.md b/src/memory/README.md index f56252a06c..a235e2fe57 100644 --- a/src/memory/README.md +++ b/src/memory/README.md @@ -109,8 +109,11 @@ Example: - **read_graph** - Read the entire knowledge graph - - No input required + - Optional input: + - `entityTypes` (string[], at least one): return only entities of these types, and the relations that touch one of them + - `observationLimit` (integer, 0 or more): see below - Returns complete graph structure with all entities and relations + - With `observationLimit`, each entity carries at most that many observations (the first ones, in the order they were added) and an `observationCount` with its total. `observationLimit: 0` returns names, types and counts only: a cheap overview of a large graph, from which `open_nodes` reads the entities you need in full - **search_nodes** - Search for nodes based on query @@ -120,6 +123,9 @@ Example: - Entity types - Observation content - Returns matching entities and their relations + - Optional input: + - `limit` (integer, 1 or more): return at most this many matches (the first, in stored order), with the relations that touch them + - `observationLimit` (integer, 0 or more): as for `read_graph` - **open_nodes** - Retrieve specific nodes by name @@ -128,6 +134,7 @@ Example: - Requested entities - Relations between requested entities - Silently skips non-existent nodes + - Optional input: `observationLimit` (integer, 0 or more), as for `read_graph` ### Resources diff --git a/src/memory/__tests__/read-options.test.ts b/src/memory/__tests__/read-options.test.ts new file mode 100644 index 0000000000..ec0435bf38 --- /dev/null +++ b/src/memory/__tests__/read-options.test.ts @@ -0,0 +1,249 @@ +// Tests for the read tools' optional inputs (#3953), driven through an SDK +// Client over an in-memory transport: read_graph's entityTypes filter, the +// observationLimit every read tool takes (with the observationCount it adds), +// and search_nodes' limit. Each option is optional, so the default results +// are pinned too: a call without them returns what it always has, and the +// knowledge-graph resource never carries observationCount. +import { describe, it, expect, beforeEach, afterEach } from "vitest"; +import type { Client } from "@modelcontextprotocol/client"; +import { call, connect, makeTempGraph, textOf } from "./helpers.js"; +import type { Connection } from "./helpers.js"; + +const alice = { + name: "Alice", + entityType: "person", + observations: ["works at Acme Corp", "likes tea", "plays chess"], +}; +const bob = { + name: "Bob", + entityType: "person", + observations: ["likes programming"], +}; +const acme = { + name: "Acme Corp", + entityType: "organization", + observations: ["makes anvils", "founded in 1900"], +}; +const empty = { name: "Nobody", entityType: "placeholder", observations: [] }; +const aliceKnowsBob = { from: "Alice", to: "Bob", relationType: "knows" }; +const aliceWorksAtAcme = { + from: "Alice", + to: "Acme Corp", + relationType: "works_at", +}; +const bobKnowsAlice = { from: "Bob", to: "Alice", relationType: "knows" }; + +describe("read tool options (#3953)", () => { + let conn: Connection; + let client: Client; + let cleanup: () => Promise; + + beforeEach(async () => { + const graph = await makeTempGraph(); + cleanup = graph.cleanup; + conn = await connect(graph.filePath); + client = conn.client; + await call(client, "create_entities", { + entities: [alice, bob, acme, empty], + }); + await call(client, "create_relations", { + relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice], + }); + }); + + afterEach(async () => { + await conn.close(); + await cleanup(); + }); + + // Call a tool and return its structured result, checking that the text + // block carries the same graph. + async function read(name: string, args: Record) { + const result = await call(client, name, args); + expect(result.isError).toBeFalsy(); + expect(textOf(result)).toBe( + JSON.stringify(result.structuredContent, null, 2), + ); + return result.structuredContent; + } + + // The text of a call that must fail input validation. + async function rejected(name: string, args: Record) { + const result = await call(client, name, args); + expect(result.isError).toBe(true); + return textOf(result); + } + + describe("defaults", () => { + it("read_graph without options returns every observation and no counts", async () => { + expect(await read("read_graph", {})).toEqual({ + entities: [alice, bob, acme, empty], + relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice], + }); + }); + + it("search_nodes and open_nodes without options return full entities", async () => { + expect(await read("search_nodes", { query: "acme" })).toEqual({ + entities: [alice, acme], + relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice], + }); + expect(await read("open_nodes", { names: ["Bob"] })).toEqual({ + entities: [bob], + relations: [aliceKnowsBob, bobKnowsAlice], + }); + }); + + it("the knowledge-graph resource is unaffected", async () => { + const { contents } = await client.readResource({ + uri: "memory://knowledge-graph", + }); + const [content] = contents; + expect("text" in content && JSON.parse(content.text)).toEqual({ + entities: [alice, bob, acme, empty], + relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice], + }); + }); + }); + + describe("observationLimit", () => { + it("0 gives read_graph an overview of names, types and counts", async () => { + expect(await read("read_graph", { observationLimit: 0 })).toEqual({ + entities: [ + { ...alice, observations: [], observationCount: 3 }, + { ...bob, observations: [], observationCount: 1 }, + { ...acme, observations: [], observationCount: 2 }, + { ...empty, observationCount: 0 }, + ], + relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice], + }); + }); + + it("keeps the first observations, in the order they were added", async () => { + expect(await read("read_graph", { observationLimit: 2 })).toEqual({ + entities: [ + { + ...alice, + observations: ["works at Acme Corp", "likes tea"], + observationCount: 3, + }, + { ...bob, observationCount: 1 }, + { ...acme, observationCount: 2 }, + { ...empty, observationCount: 0 }, + ], + relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice], + }); + }); + + it("applies to search_nodes", async () => { + expect( + await read("search_nodes", { query: "acme", observationLimit: 1 }), + ).toEqual({ + entities: [ + { + ...alice, + observations: ["works at Acme Corp"], + observationCount: 3, + }, + { ...acme, observations: ["makes anvils"], observationCount: 2 }, + ], + relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice], + }); + }); + + it("applies to open_nodes", async () => { + expect( + await read("open_nodes", { + names: ["Alice", "Missing"], + observationLimit: 1, + }), + ).toEqual({ + entities: [ + { + ...alice, + observations: ["works at Acme Corp"], + observationCount: 3, + }, + ], + relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice], + }); + }); + + it.each([ + ["read_graph", {}], + ["search_nodes", { query: "a" }], + ["open_nodes", { names: ["Alice"] }], + ])("%s rejects a negative or fractional limit", async (name, args) => { + expect(await rejected(name, { ...args, observationLimit: -1 })).toContain( + "observationLimit", + ); + expect( + await rejected(name, { ...args, observationLimit: 1.5 }), + ).toContain("observationLimit"); + }); + }); + + describe("read_graph entityTypes", () => { + it("keeps entities of the given types and the relations that touch them", async () => { + expect( + await read("read_graph", { entityTypes: ["organization"] }), + ).toEqual({ + entities: [acme], + relations: [aliceWorksAtAcme], + }); + }); + + it("accepts several types, and combines with observationLimit", async () => { + expect( + await read("read_graph", { + entityTypes: ["organization", "placeholder"], + observationLimit: 0, + }), + ).toEqual({ + entities: [ + { ...acme, observations: [], observationCount: 2 }, + { ...empty, observationCount: 0 }, + ], + relations: [aliceWorksAtAcme], + }); + }); + + it("matches types exactly", async () => { + expect(await read("read_graph", { entityTypes: ["Person"] })).toEqual({ + entities: [], + relations: [], + }); + }); + + it("rejects an empty list", async () => { + expect(await rejected("read_graph", { entityTypes: [] })).toContain( + "entityTypes", + ); + }); + }); + + describe("search_nodes limit", () => { + it("keeps the first matches and the relations that touch them", async () => { + expect(await read("search_nodes", { query: "acme", limit: 1 })).toEqual({ + entities: [alice], + relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice], + }); + expect(await read("search_nodes", { query: "o", limit: 2 })).toEqual({ + entities: [alice, bob], + relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice], + }); + }); + + it("returns every match when the limit exceeds them", async () => { + expect(await read("search_nodes", { query: "acme", limit: 10 })).toEqual({ + entities: [alice, acme], + relations: [aliceKnowsBob, aliceWorksAtAcme, bobKnowsAlice], + }); + }); + + it("rejects a limit below 1", async () => { + expect( + await rejected("search_nodes", { query: "a", limit: 0 }), + ).toContain("limit"); + }); + }); +}); diff --git a/src/memory/__tests__/tools.test.ts b/src/memory/__tests__/tools.test.ts index 170d5933e5..cf13f626a3 100644 --- a/src/memory/__tests__/tools.test.ts +++ b/src/memory/__tests__/tools.test.ts @@ -60,6 +60,39 @@ const graphOutputSchema = { additionalProperties: false, }; +// The read tools' output (#3953): an entity may also carry observationCount. +const readGraphOutputSchema = { + ...graphOutputSchema, + properties: { + ...graphOutputSchema.properties, + entities: { + type: "array", + items: { + ...entityItem(true), + properties: { + ...entityItem(true).properties, + observationCount: { + type: "integer", + minimum: 0, + maximum: Number.MAX_SAFE_INTEGER, + description: + "How many observations the entity holds in total; present only when observationLimit was given", + }, + }, + }, + }, + }, +}; + +// The observationLimit input every read tool takes (#3953). +const observationLimitInput = { + type: "integer", + minimum: 0, + maximum: Number.MAX_SAFE_INTEGER, + description: + "Return at most this many observations per entity (the first ones, in the order they were added) and add each entity's total as observationCount. Use 0 for names, types and counts only: a cheap overview of a large graph, after which open_nodes reads the entities you need in full. Omit to return every observation.", +}; + const statusOutputSchema = { type: "object", properties: { success: { type: "boolean" }, message: { type: "string" } }, @@ -270,13 +303,23 @@ const expectedTools = [ { name: "read_graph", title: "Read Graph", - description: "Read the entire knowledge graph", + description: + "Read the entire knowledge graph: every entity with its observations, and every relation. The result grows with the graph, so use this only when you need everything; prefer search_nodes to find entities by keyword and open_nodes to read entities whose names you know.", inputSchema: { type: "object", - properties: {}, + properties: { + entityTypes: { + type: "array", + items: { type: "string" }, + minItems: 1, + description: + "Return only entities whose entityType is one of these (exact match), and only the relations that touch one of them. Omit to return every entity.", + }, + observationLimit: observationLimitInput, + }, $schema: JSON_SCHEMA_2020_12, }, - outputSchema: graphOutputSchema, + outputSchema: readGraphOutputSchema, annotations: readAnnotations, execution, }, @@ -293,11 +336,19 @@ const expectedTools = [ description: "The search query to match against entity names, types, and observation content", }, + limit: { + type: "integer", + minimum: 1, + maximum: Number.MAX_SAFE_INTEGER, + description: + "Return at most this many matching entities (the first matches, in stored order), with the relations that touch them. Omit to return every match.", + }, + observationLimit: observationLimitInput, }, required: ["query"], $schema: JSON_SCHEMA_2020_12, }, - outputSchema: graphOutputSchema, + outputSchema: readGraphOutputSchema, annotations: readAnnotations, execution, }, @@ -313,11 +364,12 @@ const expectedTools = [ items: { type: "string" }, description: "An array of entity names to retrieve", }, + observationLimit: observationLimitInput, }, required: ["names"], $schema: JSON_SCHEMA_2020_12, }, - outputSchema: graphOutputSchema, + outputSchema: readGraphOutputSchema, annotations: readAnnotations, execution, }, diff --git a/src/memory/index.ts b/src/memory/index.ts index 4892ee85a2..d6cf296c63 100644 --- a/src/memory/index.ts +++ b/src/memory/index.ts @@ -707,23 +707,50 @@ export class KnowledgeGraphManager { }); } - async readGraph(): Promise { - return this.loadGraph(); + // With no options, the whole graph, as the knowledge-graph resource serves + // it. entityTypes keeps only entities of those types, and the relations + // that touch one of them; observationLimit caps each entity's observations + // (see limitObservations). + async readGraph(options: ReadGraphOptions = {}): Promise { + const graph = await this.loadGraph(); + const { entityTypes, observationLimit } = options; + if (entityTypes === undefined) { + return { + entities: limitObservations(graph.entities, observationLimit), + relations: graph.relations, + }; + } + const entities = graph.entities.filter((e) => + entityTypes.includes(e.entityType), + ); + const names = new Set(entities.map((e) => e.name)); + return { + entities: limitObservations(entities, observationLimit), + relations: graph.relations.filter( + (r) => names.has(r.from) || names.has(r.to), + ), + }; } - // Very basic search function - async searchNodes(query: string): Promise { + // Very basic search function. limit keeps the first matches in stored + // order; observationLimit caps each entity's observations. + async searchNodes( + query: string, + options: SearchNodesOptions = {}, + ): Promise { const graph = await this.loadGraph(); // Filter entities - const filteredEntities = graph.entities.filter( - (e) => - e.name.toLowerCase().includes(query.toLowerCase()) || - e.entityType.toLowerCase().includes(query.toLowerCase()) || - e.observations.some((o) => - o.toLowerCase().includes(query.toLowerCase()), - ), - ); + const filteredEntities = graph.entities + .filter( + (e) => + e.name.toLowerCase().includes(query.toLowerCase()) || + e.entityType.toLowerCase().includes(query.toLowerCase()) || + e.observations.some((o) => + o.toLowerCase().includes(query.toLowerCase()), + ), + ) + .slice(0, options.limit); // Create a Set of filtered entity names for quick lookup const filteredEntityNames = new Set(filteredEntities.map((e) => e.name)); @@ -734,15 +761,16 @@ export class KnowledgeGraphManager { (r) => filteredEntityNames.has(r.from) || filteredEntityNames.has(r.to), ); - const filteredGraph: KnowledgeGraph = { - entities: filteredEntities, + return { + entities: limitObservations(filteredEntities, options.observationLimit), relations: filteredRelations, }; - - return filteredGraph; } - async openNodes(names: string[]): Promise { + async openNodes( + names: string[], + options: OpenNodesOptions = {}, + ): Promise { const graph = await this.loadGraph(); // Filter entities @@ -761,13 +789,57 @@ export class KnowledgeGraphManager { (r) => filteredEntityNames.has(r.from) || filteredEntityNames.has(r.to), ); - const filteredGraph: KnowledgeGraph = { - entities: filteredEntities, + return { + entities: limitObservations(filteredEntities, options.observationLimit), relations: filteredRelations, }; + } +} + +// Options of the read tools (#3953), all optional, so a call without them +// returns what it always has. They let a client read a large graph in +// pieces: an overview of names, types and observation counts first +// (observationLimit 0), then the entities it needs in full. +export interface OpenNodesOptions { + observationLimit?: number; +} + +export interface ReadGraphOptions extends OpenNodesOptions { + entityTypes?: string[]; +} + +export interface SearchNodesOptions extends OpenNodesOptions { + limit?: number; +} + +// An entity as a read tool returns it. observationCount is set only when +// observationLimit was given, so a count above observations.length tells the +// caller that observations were left out. +export interface ReadEntity extends Entity { + observationCount?: number; +} + +export interface ReadGraphResult { + entities: ReadEntity[]; + relations: Relation[]; +} - return filteredGraph; +// Keep each entity's first observationLimit observations, in the order they +// were added, and record how many it holds. Without a limit the entities are +// returned as they are. +function limitObservations( + entities: Entity[], + observationLimit: number | undefined, +): ReadEntity[] { + if (observationLimit === undefined) { + return entities; } + return entities.map((e) => ({ + name: e.name, + entityType: e.entityType, + observations: e.observations.slice(0, observationLimit), + observationCount: e.observations.length, + })); } // Zod schemas for entities and relations @@ -796,6 +868,47 @@ export const SearchNodesQuerySchema = z "The search query to match against entity names, types, and observation content", ); +// The read tools' optional inputs (#3953). Their descriptions carry the +// guidance on when to use them, so the tool descriptions need not. +const ObservationLimitSchema = z + .number() + .int() + .min(0) + .optional() + .describe( + "Return at most this many observations per entity (the first ones, in the order they were added) and add each entity's total as observationCount. Use 0 for names, types and counts only: a cheap overview of a large graph, after which open_nodes reads the entities you need in full. Omit to return every observation.", + ); + +const ReadGraphEntityTypesSchema = z + .array(z.string()) + .min(1) + .optional() + .describe( + "Return only entities whose entityType is one of these (exact match), and only the relations that touch one of them. Omit to return every entity.", + ); + +const SearchNodesLimitSchema = z + .number() + .int() + .min(1) + .optional() + .describe( + "Return at most this many matching entities (the first matches, in stored order), with the relations that touch them. Omit to return every match.", + ); + +// An entity in a read tool's result: observationCount appears only when the +// call gave observationLimit. +const ReadEntitySchema = EntitySchema.extend({ + observationCount: z + .number() + .int() + .min(0) + .optional() + .describe( + "How many observations the entity holds in total; present only when observationLimit was given", + ), +}); + // Build a memory server over the graph file at memoryFilePath: the nine // tools, the knowledge-graph resource and its subscriptions. main() serves the // result on stdio; tests connect it to an in-memory transport, so each call @@ -1092,10 +1205,14 @@ export function createServer( "read_graph", { title: "Read Graph", - description: "Read the entire knowledge graph", - inputSchema: z.object({}), + description: + "Read the entire knowledge graph: every entity with its observations, and every relation. The result grows with the graph, so use this only when you need everything; prefer search_nodes to find entities by keyword and open_nodes to read entities whose names you know.", + inputSchema: z.object({ + entityTypes: ReadGraphEntityTypesSchema, + observationLimit: ObservationLimitSchema, + }), outputSchema: z.object({ - entities: z.array(EntitySchema), + entities: z.array(ReadEntitySchema), relations: z.array(RelationSchema), }), annotations: { @@ -1105,8 +1222,11 @@ export function createServer( openWorldHint: false, }, }, - async () => { - const graph = await knowledgeGraphManager.readGraph(); + async ({ entityTypes, observationLimit }) => { + const graph = await knowledgeGraphManager.readGraph({ + entityTypes, + observationLimit, + }); return { content: [ { type: "text" as const, text: JSON.stringify(graph, null, 2) }, @@ -1124,9 +1244,11 @@ export function createServer( description: "Search for nodes in the knowledge graph based on a query", inputSchema: z.object({ query: SearchNodesQuerySchema, + limit: SearchNodesLimitSchema, + observationLimit: ObservationLimitSchema, }), outputSchema: z.object({ - entities: z.array(EntitySchema), + entities: z.array(ReadEntitySchema), relations: z.array(RelationSchema), }), annotations: { @@ -1136,8 +1258,11 @@ export function createServer( openWorldHint: false, }, }, - async ({ query }) => { - const graph = await knowledgeGraphManager.searchNodes(query); + async ({ query, limit, observationLimit }) => { + const graph = await knowledgeGraphManager.searchNodes(query, { + limit, + observationLimit, + }); return { content: [ { type: "text" as const, text: JSON.stringify(graph, null, 2) }, @@ -1157,9 +1282,10 @@ export function createServer( names: z .array(z.string()) .describe("An array of entity names to retrieve"), + observationLimit: ObservationLimitSchema, }), outputSchema: z.object({ - entities: z.array(EntitySchema), + entities: z.array(ReadEntitySchema), relations: z.array(RelationSchema), }), annotations: { @@ -1169,8 +1295,10 @@ export function createServer( openWorldHint: false, }, }, - async ({ names }) => { - const graph = await knowledgeGraphManager.openNodes(names); + async ({ names, observationLimit }) => { + const graph = await knowledgeGraphManager.openNodes(names, { + observationLimit, + }); return { content: [ { type: "text" as const, text: JSON.stringify(graph, null, 2) },