Add isArray and isMap to tree-agent context (#25725)

This gives the LLM an alternative to `Array.isArray` and `instanceof
Map` that will work on our node types, since `Array.isArray` does not
work with our non-pojo array classes.

The provided functions also work for JS arrays and Maps, so the LLM
doesn't have to distinguish between when to use the built-in
Array.isArray vs. when to use our provided `isArray` helper - it can
always use the latter. This also gives us the option in the future to
force it to use ours by making `Array.isArray` error within the
boundaries of the eval.

This also adds some explicit tests in SharedTree for `Array.isArray`
support.

Author: noencke

packages/dds/tree/src/test/simple-tree/node-kinds/arrayNode.spec.ts

@@ -72,6 +72,18 @@ describe("ArrayNode", () => {
 			assert.equal(n.y, 3);
 			assert.deepEqual(thisList, [n, n]);
 		});
+
+		it("does not pass Array.isArray", () => {
+			const array = init(CustomizableNumberArray, [1, 2, 3]);
+			assert.equal(Array.isArray(array), false);
+		});
+	});
+
+	describeHydration("pojo-emulation", (init) => {
+		it("passes Array.isArray", () => {
+			const array = init(PojoEmulationNumberArray, [1, 2, 3]);
+			assert.equal(Array.isArray(array), true);
+		});
 	});
 
 	describe("insertable types", () => {

packages/framework/tree-agent/src/agent.ts

@@ -8,7 +8,7 @@ import type {
 	TreeFieldFromImplicitField,
 	TreeNodeSchema,
 } from "@fluidframework/tree";
-import { TreeNode } from "@fluidframework/tree";
+import { NodeKind, TreeNode } from "@fluidframework/tree";
 import type {
 	ReadableField,
 	FactoryContentObject,
@@ -313,6 +313,26 @@ function bindEditorToSubtree<TSchema extends ImplicitFieldSchema>(
 		},
 		create,
 		is,
+		isArray(node) {
+			if (Array.isArray(node)) {
+				return true;
+			}
+			if (node instanceof TreeNode) {
+				const schema = Tree.schema(node);
+				return schema.kind === NodeKind.Array;
+			}
+			return false;
+		},
+		isMap(node) {
+			if (node instanceof Map) {
+				return true;
+			}
+			if (node instanceof TreeNode) {
+				const schema = Tree.schema(node);
+				return schema.kind === NodeKind.Map;
+			}
+			return false;
+		},
 		parent: (child: TreeNode): TreeNode | undefined => Tree.parent(child),
 		key: (child: TreeNode): string | number => Tree.key(child),
 	} satisfies Context<TSchema>;

packages/framework/tree-agent/src/api.ts

@@ -49,8 +49,6 @@ export interface Context<TSchema extends ImplicitFieldSchema> {
 	 * Always use these builder functions when creating new nodes rather than plain JavaScript objects.
 	 *
 	 * Example: Create a new Person node with `context.create.Person({ name: "Alice", age: 30 })`
-	 *
-	 * Example: Create a new Task node with `context.create.Task({ title: "Buy groceries", completed: false })`
 	 */
 	create: Record<string, (input: FactoryContentObject) => TreeNode>;
 
@@ -65,6 +63,16 @@ export interface Context<TSchema extends ImplicitFieldSchema> {
 	 */
 	is: Record<string, <T extends TreeNode>(input: T) => input is T>;
 
+	/**
+	 * Checks if the given node is an array.
+	 */
+	isArray(value: unknown): boolean;
+
+	/**
+	 * Checks if the given node is a map.
+	 */
+	isMap(value: unknown): boolean;
+
 	/**
 	 * Returns the parent node of the given child node.
 	 * @param child - The node whose parent you want to find.

packages/framework/tree-agent/src/prompt.ts

@@ -79,8 +79,15 @@ export function getPrompt<TRoot extends ImplicitFieldSchema>(args: {
 	const stringified = stringifyTree(field);
 	const details: SchemaDetails = { hasHelperMethods: false };
 	const typescriptSchemaTypes = getZodSchemaAsTypeScript(domainTypes, details);
-
-	const create =
+	const exampleTypeName =
+		nodeTypeUnion === undefined
+			? undefined
+			: nodeTypeUnion
+					.split("|")
+					.map((part) => part.trim())
+					.find((part) => part.length > 0);
+
+	const createDocs =
 		exampleObjectName === undefined
 			? ""
 			: `\n	/**
@@ -102,7 +109,7 @@ export function getPrompt<TRoot extends ImplicitFieldSchema>(args: {
 	create: Record<string, <T extends TreeData>(input: T) => T>;\n`;
 
 	const isDocs =
-		exampleObjectName === undefined
+		exampleTypeName === undefined
 			? ""
 			: `\n	/**
 	 * A collection of type-guard functions for data in the tree.
@@ -111,9 +118,31 @@ export function getPrompt<TRoot extends ImplicitFieldSchema>(args: {
 	 * Call the corresponding function to check if a node is of that specific type.
 	 * This is useful when working with nodes that could be one of multiple types.
 	 *
-	 * ${`Example: Check if a node is a ${exampleObjectName} with \`if (context.is.${exampleObjectName}(node)) {}\``}
+	 * ${`Example: Check if a node is a ${exampleTypeName} with \`if (context.is.${exampleTypeName}(node)) {}\``}
+	 */
+	is: Record<string, <T extends TreeData>(data: unknown) => data is T>;
+	
+	/**
+	 * Checks if the provided data is an array.
+	 * @remarks
+	 * DO NOT use \`Array.isArray\` to check if tree data is an array - use this function instead.
+	 * 
+	 * This function will also work for native JavaScript arrays.
+	 *
+	 * ${`Example: \`if (context.isArray(node)) {}\``}
+	 */
+	isArray(data: any): boolean;
+	
+	/**
+	 * Checks if the provided data is a map.
+	 * @remarks
+	 * DO NOT use \`instanceof Map\` to check if tree data is a map - use this function instead.
+	 * 
+	 * This function will also work for native JavaScript Map instances.
+	 *
+	 * ${`Example: \`if (context.isMap(node)) {}\``}
 	 */
-	is: Record<string, <T extends TreeData>(input: unknown) => input is T>;\n`;
+	isMap(data: any): boolean;\n`;
 
 	const context = `\`\`\`typescript
 	${nodeTypeUnion === undefined ? "" : `type TreeData = ${nodeTypeUnion};\n\n`}	/**
@@ -130,7 +159,7 @@ export function getPrompt<TRoot extends ImplicitFieldSchema>(args: {
 	 * Example: Read the current root with \`const currentRoot = context.root;\`
 	 *${rootTypes.length > 0 ? ` Example: Replace the entire root with \`context.root = context.create.${getFriendlyName(rootTypes[0] ?? oob())}({ });\`\n	 *` : ""}/
 	root: ReadableField<TSchema>;
-	${create}
+	${createDocs}
 	${isDocs}
 	/**
 	 * Returns the parent object/array/map of the given object/array/map, if there is one.

packages/framework/tree-agent/src/test/__snapshots__/prompt.md

@@ -88,9 +88,31 @@ Here is the definition of the `Context` interface:
 	 * Call the corresponding function to check if a node is of that specific type.
 	 * This is useful when working with nodes that could be one of multiple types.
 	 *
-	 * Example: Check if a node is a TestArrayItem with `if (context.is.TestArrayItem(node)) {}`
+	 * Example: Check if a node is a TestMap with `if (context.is.TestMap(node)) {}`
 	 */
-	is: Record<string, <T extends TreeData>(input: unknown) => input is T>;
+	is: Record<string, <T extends TreeData>(data: unknown) => data is T>;
+	
+	/**
+	 * Checks if the provided data is an array.
+	 * @remarks
+	 * DO NOT use `Array.isArray` to check if tree data is an array - use this function instead.
+	 * 
+	 * This function will also work for native JavaScript arrays.
+	 *
+	 * Example: `if (context.isArray(node)) {}`
+	 */
+	isArray(data: any): boolean;
+	
+	/**
+	 * Checks if the provided data is a map.
+	 * @remarks
+	 * DO NOT use `instanceof Map` to check if tree data is a map - use this function instead.
+	 * 
+	 * This function will also work for native JavaScript Map instances.
+	 *
+	 * Example: `if (context.isMap(node)) {}`
+	 */
+	isMap(data: any): boolean;
 
 	/**
 	 * Returns the parent object/array/map of the given object/array/map, if there is one.

packages/framework/tree-agent/src/test/agent.spec.ts

@@ -248,10 +248,6 @@ describe("Semantic Agent", () => {
 		assert.equal(view.root, "Initial", "Tree should not have changed");
 	});
 
-	it("passes constructors to edit code", async () => {
-		// Moved to the dedicated context helpers describe block below.
-	});
-
 	it("supplies the system prompt as context", async () => {
 		const view = independentView(new TreeViewConfiguration({ schema: sf.string }), {});
 		view.initialize("X");
@@ -365,6 +361,46 @@ describe("Semantic Agent", () => {
 			assert.equal(await agent.query("Key Edit"), "Done");
 			assert.equal(view.root.child.value, "KeyWorked");
 		});
+
+		it("provides working isArray helper", async () => {
+			const sfLocal = new SchemaFactory("TestIsArray");
+			const NumberArray = sfLocal.array(sfLocal.number);
+			const view = independentView(new TreeViewConfiguration({ schema: NumberArray }), {});
+			view.initialize([1, 2, 3]);
+			const model: SharedTreeChatModel = {
+				editToolName,
+				async query({ edit }) {
+					const result = await edit(
+						`if (!context.isArray([]) ||!context.isArray(context.root)) { throw new Error('Expected array root'); } context.root.insertAt(0, 99);`,
+					);
+					assert.equal(result.type, "success", result.message);
+					return "ArrayOK";
+				},
+			};
+			const agent = new SharedTreeSemanticAgent(model, view);
+			assert.equal(await agent.query("Array Edit"), "ArrayOK");
+			assert.equal(view.root[0], 99);
+		});
+
+		it("provides working isMap helper", async () => {
+			const sfLocal = new SchemaFactory("TestIsMap");
+			class NumberMap extends sfLocal.map("NumberMap", sfLocal.number) {}
+			const view = independentView(new TreeViewConfiguration({ schema: NumberMap }), {});
+			view.initialize(new NumberMap(new Map([["x", 1]])));
+			const model: SharedTreeChatModel = {
+				editToolName,
+				async query({ edit }) {
+					const result = await edit(
+						`if (!context.isMap(new Map()) || !context.isMap(context.root)) { throw new Error('Expected map root'); } context.root.set('y', 2);`,
+					);
+					assert.equal(result.type, "success", result.message);
+					return "MapOK";
+				},
+			};
+			const agent = new SharedTreeSemanticAgent(model, view);
+			assert.equal(await agent.query("Map Edit"), "MapOK");
+			assert.equal(view.root.get("y"), 2);
+		});
 	});
 
 	it("supplies additional context if the tree changes between queries", async () => {
@@ -433,7 +469,6 @@ describe("Semantic Agent", () => {
 	});
 
 	describe("bindEditor", () => {
-		/* eslint-disable @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-call */
 		interface TestContext {
 			root: string;
 		}
@@ -485,6 +520,5 @@ describe("Semantic Agent", () => {
 			await assert.rejects(promise, /asyncBoom/);
 			assert.equal(view.root, "Initial", "Tree should not have changed after async error");
 		});
-		/* eslint-enable @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-call */
 	});
 });