Add property binding to tree-agent

packages/framework/tree-agent/api-report/tree-agent.alpha.api.md

@@ -44,6 +44,11 @@ export interface EditResult {
     type: "success" | "disabledError" | "editingError" | "tooManyEditsError" | "expiredError";
 }
 
+// @alpha
+export type ExposableKeys<T> = {
+    [K in keyof T]?: T[K] extends (...args: any[]) => any ? never : K;
+}[keyof T];
+
 // @alpha
 export interface ExposedMethods {
     // (undocumented)
@@ -53,9 +58,23 @@ export interface ExposedMethods {
     instanceOf<T extends TreeNodeSchemaClass>(schema: T): z.ZodType<InstanceType<T>, z.ZodTypeDef, InstanceType<T>>;
 }
 
+// @alpha
+export interface ExposedProperties {
+    // (undocumented)
+    exposeProperty<S extends BindableSchema & Ctor, K extends string & ExposableKeys<InstanceType<S>>, TZ extends ZodTypeAny>(schema: S, name: K, def: {
+        schema: TZ;
+        description?: string;
+    } & ReadOnlyRequirement<InstanceType<S>, K> & TypeMatchOrError<InstanceType<S>[K], infer<TZ>>): void;
+    // (undocumented)
+    instanceOf<T extends TreeNodeSchemaClass>(schema: T): ZodType<InstanceType<T>, ZodTypeDef, InstanceType<T>>;
+}
+
 // @alpha
 export const exposeMethodsSymbol: unique symbol;
 
+// @alpha
+export const exposePropertiesSymbol: unique symbol;
+
 // @alpha
 export interface FunctionDef<Args extends readonly Arg[], Return extends z.ZodTypeAny, Rest extends z.ZodTypeAny | null = null> {
     // (undocumented)
@@ -74,6 +93,15 @@ export interface IExposedMethods {
     [exposeMethodsSymbol](methods: ExposedMethods): void;
 }
 
+// @alpha
+export interface IExposedProperties {
+    // (undocumented)
+    [exposePropertiesSymbol]?(properties: ExposedProperties): void;
+}
+
+// @alpha
+export type IfEquals<X, Y, A = true, B = false> = (<T>() => T extends X ? 1 : 2) extends <T>() => T extends Y ? 1 : 2 ? A : B;
+
 // @alpha
 export type Infer<T> = T extends FunctionDef<infer Args, infer Return, infer Rest> ? z.infer<z.ZodFunction<z.ZodTuple<ArgsTuple<Args>, Rest>, Return>> : never;
 
@@ -90,6 +118,37 @@ export type MethodKeys<T> = {
     [K in keyof T]: T[K] extends (...args: any[]) => any ? K : never;
 };
 
+// @alpha
+export class PropertyDef {
+    constructor(name: string, description: string | undefined, schema: ZodTypeAny, readOnly: boolean);
+    // (undocumented)
+    readonly description: string | undefined;
+    // (undocumented)
+    readonly name: string;
+    // (undocumented)
+    readonly readOnly: boolean;
+    // (undocumented)
+    readonly schema: ZodTypeAny;
+}
+
+// @alpha
+export type ReadonlyKeys<T> = {
+    [P in keyof T]-?: IfEquals<{
+        [Q in P]: T[P];
+    }, {
+        -readonly [Q in P]: T[P];
+    }, never, P>;
+}[keyof T];
+
+// @alpha
+export type ReadOnlyRequirement<TObj, K extends keyof TObj> = {
+    [P in K]-?: P extends ReadonlyKeys<TObj> ? {
+        readOnly: true;
+    } : {
+        readOnly?: false;
+    };
+}[K];
+
 // @alpha
 export interface SemanticAgentOptions {
     domainHints?: string;
@@ -124,4 +183,11 @@ export type SynchronousEditor = (context: Record<string, unknown>, code: string)
 // @alpha
 export type TreeView<TRoot extends ImplicitFieldSchema | UnsafeUnknownSchema> = Pick<TreeViewAlpha<TRoot>, "root" | "fork" | "merge" | "rebaseOnto" | "schema" | "events"> & TreeBranchAlpha;
 
+// @alpha
+export type TypeMatchOrError<Expected, Received> = [Received] extends [Expected] ? unknown : {
+    __error__: "Zod schema value type does not match the property's declared type";
+    expected: Expected;
+    received: Received;
+};
+
 ```

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

@@ -38,3 +38,14 @@ export {
 	type Infer,
 	type IExposedMethods,
 } from "./methodBinding.js";
+export type {
+	exposePropertiesSymbol,
+	PropertyDef,
+	ExposedProperties,
+	IExposedProperties,
+	ExposableKeys,
+	ReadOnlyRequirement,
+	ReadonlyKeys,
+	TypeMatchOrError,
+	IfEquals,
+} from "./propertyBinding.js";

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

@@ -77,7 +77,7 @@ export function getPrompt<TRoot extends ImplicitFieldSchema>(args: {
 	}
 
 	const stringified = stringifyTree(field);
-	const details: SchemaDetails = { hasHelperMethods: false };
+	const details: SchemaDetails = { hasHelperMethods: false, hasHelperProperties: false };
 	const typescriptSchemaTypes = getZodSchemaAsTypeScript(domainTypes, details);
 	const exampleTypeName =
 		nodeTypeUnion === undefined
@@ -121,23 +121,23 @@ export function getPrompt<TRoot extends ImplicitFieldSchema>(args: {
 	 * ${`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)) {}\``}
@@ -179,6 +179,10 @@ export function getPrompt<TRoot extends ImplicitFieldSchema>(args: {
 }
 \`\`\``;
 
+	const helperPropertyExplanation = details.hasHelperProperties
+		? "Some schema types expose additional helper properties directly on the objects (including readonly properties). When these properties are available, you may read them and use them in your logic."
+		: "";
+
 	const helperMethodExplanation = details.hasHelperMethods
 		? `Manipulating the data using the APIs described below is allowed, but when possible ALWAYS prefer to use any application helper methods exposed on the schema TypeScript types if the goal can be accomplished that way.
 It will often not be possible to fully accomplish the goal using those helpers. When this is the case, mutate the objects as normal, taking into account the following guidance.`
@@ -276,6 +280,7 @@ There are other additional helper functions available on the \`context\` object
 Here is the definition of the \`Context\` interface:
 ${context}
 ${helperMethodExplanation}
+${helperPropertyExplanation}
 ${hasArrays ? arrayEditing : ""}${hasMaps ? mapEditing : ""}#### Additional Notes
 
 Before outputting the edit function, you should check that it is valid according to both the application tree's schema and any restrictions of the editing APIs described above.

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

@@ -0,0 +1,181 @@
+/*!
+ * Copyright (c) Microsoft Corporation and contributors. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+/* eslint-disable @typescript-eslint/no-explicit-any */
+import type { TreeNodeSchema, TreeNodeSchemaClass } from "@fluidframework/tree";
+import type { ZodType, ZodTypeAny, ZodTypeDef, infer as ZodInfer } from "zod";
+
+import type { BindableSchema, Ctor } from "./methodBinding.js";
+import { instanceOf } from "./utils.js";
+
+/**
+ * A symbol used to expose properties to the LLM.
+ * @alpha
+ */
+export const exposePropertiesSymbol: unique symbol = Symbol.for(
+	"@fluidframework/tree-agent/exposeProperties",
+);
+
+/**
+ * Set of property keys from `T` that are not methods.
+ * @alpha
+ */
+export type ExposableKeys<T> = {
+	[K in keyof T]?: T[K] extends (...args: any[]) => any ? never : K;
+}[keyof T];
+
+/**
+ * Type-level equality test used as a helper to evaluate readonly keys.
+ * - If X and Y are the same type, it evaluates to A.
+ * - If X and Y are different, it evaluates to B.
+ * @alpha
+ */
+export type IfEquals<X, Y, A = true, B = false> = (<T>() => T extends X ? 1 : 2) extends <
+	T,
+>() => T extends Y ? 1 : 2
+	? A
+	: B;
+
+/**
+ * Produces a union of keys of `T` which are readonly.
+ * @alpha
+ */
+export type ReadonlyKeys<T> = {
+	[P in keyof T]-?: IfEquals<{ [Q in P]: T[P] }, { -readonly [Q in P]: T[P] }, never, P>;
+}[keyof T];
+
+/**
+ * Type to enforce `readOnly: true` for readonly properties.
+ * @alpha
+ */
+export type ReadOnlyRequirement<TObj, K extends keyof TObj> = {
+	[P in K]-?: P extends ReadonlyKeys<TObj> ? { readOnly: true } : { readOnly?: false };
+}[K];
+
+/**
+ * Emits compile-time error when there is a type mismatch.
+ * @alpha
+ */
+export type TypeMatchOrError<Expected, Received> = [Received] extends [Expected]
+	? unknown
+	: {
+			__error__: "Zod schema value type does not match the property's declared type";
+			expected: Expected;
+			received: Received;
+		};
+
+/**
+ * A property definition class that describes the structure of the property
+ * @alpha
+ */
+export class PropertyDef {
+	public constructor(
+		public readonly name: string,
+		public readonly description: string | undefined,
+		public readonly schema: ZodTypeAny,
+		public readonly readOnly: boolean,
+	) {}
+}
+
+/**
+ * An interface for exposing properties of schema classes to an agent.
+ * @alpha
+ */
+export interface ExposedProperties {
+	exposeProperty<
+		S extends BindableSchema & Ctor,
+		K extends string & ExposableKeys<InstanceType<S>>,
+		TZ extends ZodTypeAny,
+	>(
+		schema: S,
+		name: K,
+		def: { schema: TZ; description?: string } & ReadOnlyRequirement<InstanceType<S>, K> &
+			TypeMatchOrError<InstanceType<S>[K], ZodInfer<TZ>>,
+	): void;
+
+	instanceOf<T extends TreeNodeSchemaClass>(
+		schema: T,
+	): ZodType<InstanceType<T>, ZodTypeDef, InstanceType<T>>;
+}
+
+/**
+ * An interface that SharedTree schema classes should implement to expose their properties to the LLM.
+ *
+ * @remarks
+ * The `getExposedProperties` free function will cause the method here to be called on the class passed to it.
+ *
+ * @privateremarks
+ * Implementing this interface correctly seems tricky?
+ * To actually implement it in a way that satisfies TypeScript,
+ * classes need to declare both a static version and an instance version of the method
+ * (the instance one can just delegate to the static one).
+ *
+ * @alpha
+ */
+export interface IExposedProperties {
+	[exposePropertiesSymbol]?(properties: ExposedProperties): void;
+}
+
+class ExposedPropertiesI implements ExposedProperties {
+	private readonly properties: Record<string, PropertyDef> = {};
+	private readonly referencedTypes = new Set<TreeNodeSchema>();
+
+	public constructor(private readonly schemaClass: BindableSchema) {}
+
+	public exposeProperty<
+		S extends BindableSchema & Ctor,
+		K extends string & ExposableKeys<InstanceType<S>>,
+		TZ extends ZodTypeAny,
+	>(
+		schema: S,
+		name: K,
+		def: { schema: TZ; description?: string } & ReadOnlyRequirement<InstanceType<S>, K> &
+			TypeMatchOrError<InstanceType<S>[K], ZodInfer<TZ>>,
+	): void {
+		if (schema !== this.schemaClass) {
+			throw new Error('Must expose properties on the "this" schema class');
+		}
+		this.properties[name] = new PropertyDef(
+			name,
+			def.description,
+			def.schema,
+			def.readOnly === true,
+		);
+	}
+
+	public instanceOf<T extends TreeNodeSchemaClass>(
+		schema: T,
+	): ZodType<InstanceType<T>, ZodTypeDef, InstanceType<T>> {
+		this.referencedTypes.add(schema);
+		return instanceOf(schema);
+	}
+
+	public static getExposedProperties(schemaClass: BindableSchema): {
+		properties: Record<string, PropertyDef>;
+		referencedTypes: Set<TreeNodeSchema>;
+	} {
+		const exposed = new ExposedPropertiesI(schemaClass);
+		const extractable = schemaClass as unknown as IExposedProperties;
+		if (extractable[exposePropertiesSymbol] !== undefined) {
+			extractable[exposePropertiesSymbol](exposed);
+		}
+		return {
+			properties: exposed.properties,
+			referencedTypes: exposed.referencedTypes,
+		};
+	}
+}
+
+/**
+ * Get the exposed properties of a schema class.
+ * @param schemaClass - The schema class to extract properties from.
+ * @returns A record of property names and their corresponding Zod types.
+ */
+export function getExposedProperties(schemaClass: BindableSchema): {
+	properties: Record<string, PropertyDef>;
+	referencedTypes: Set<TreeNodeSchema>;
+} {
+	return ExposedPropertiesI.getExposedProperties(schemaClass);
+}