@fluidframework/ai-collab 2.10.0-306579

package/src/explicit-strategy/promptGeneration.ts

@@ -0,0 +1,294 @@
+/*!
+ * Copyright (c) Microsoft Corporation and contributors. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+import { assert } from "@fluidframework/core-utils/internal";
+import {
+	NodeKind,
+	type ImplicitFieldSchema,
+	type TreeFieldFromImplicitField,
+	getJsonSchema,
+	type JsonFieldSchema,
+	type JsonNodeSchema,
+	type JsonSchemaRef,
+	type JsonTreeSchema,
+	getSimpleSchema,
+	Tree,
+	type TreeNode,
+} from "@fluidframework/tree/internal";
+// eslint-disable-next-line import/no-internal-modules
+import { createZodJsonValidator } from "typechat/zod";
+
+import { objectIdKey, type TreeEdit } from "./agentEditTypes.js";
+import type { IdGenerator } from "./idGenerator.js";
+import { generateGenericEditTypes } from "./typeGeneration.js";
+import { fail } from "./utils.js";
+
+/**
+ * A log of edits that have been made to a tree.
+ * @remarks This is primarily used to help an LLM keep track of the active changes it has made.
+ */
+export type EditLog = {
+	edit: TreeEdit;
+	error?: string;
+}[];
+
+/**
+ * TBD
+ */
+export function toDecoratedJson(
+	idGenerator: IdGenerator,
+	root: TreeFieldFromImplicitField<ImplicitFieldSchema>,
+): string {
+	idGenerator.assignIds(root);
+	const stringified: string = JSON.stringify(root, (_, value) => {
+		if (typeof value === "object" && !Array.isArray(value) && value !== null) {
+			// TODO: SharedTree Team needs to either publish TreeNode as a class to use .instanceof() or a typeguard.
+			// Uncomment this assertion back once we have a typeguard ready.
+			// assert(isTreeNode(node), "Non-TreeNode value in tree.");
+			const objId =
+				// eslint-disable-next-line @typescript-eslint/no-unsafe-argument
+				idGenerator.getId(value) ?? fail("ID of new node should have been assigned.");
+			assert(
+				!Object.prototype.hasOwnProperty.call(value, objectIdKey),
+				`Collision of object id property.`,
+			);
+			return {
+				[objectIdKey]: objId,
+				...value,
+			} as unknown;
+		}
+		return value as unknown;
+	});
+	return stringified;
+}
+
+/**
+ * Generates a prompt designed to make an LLM produce a plan to edit the SharedTree to accomplish a user-specified goal.
+ */
+export function getPlanningSystemPrompt(
+	treeNode: TreeNode,
+	userPrompt: string,
+	systemRoleContext?: string,
+): string {
+	const schema = Tree.schema(treeNode);
+
+	const promptFriendlySchema = getPromptFriendlyTreeSchema(getJsonSchema(schema));
+	const role = `I'm an agent who makes plans for another agent to achieve a user-specified goal to update the state of an application.${
+		systemRoleContext === undefined
+			? ""
+			: `
+			The other agent follows this guidance: ${systemRoleContext}`
+	}`;
+
+	const systemPrompt = `
+	${role}
+	The application state tree is a JSON object with the following schema: ${promptFriendlySchema}
+	The current state is: ${JSON.stringify(treeNode)}.
+	The user requested that I accomplish the following goal:
+	"${userPrompt}"
+	I've made a plan to accomplish this goal by doing a sequence of edits to the tree.
+	Edits can include setting the root, inserting, modifying, removing, or moving elements in the tree.
+	Here is my plan:`;
+
+	return systemPrompt;
+}
+
+/**
+ * Generates a prompt that provides a history of the edits an LLM has made to a SharedTree as well as any errors that occured from attemping to apply each respsecitve edit to the tree.
+ */
+export function createEditListHistoryPrompt(edits: EditLog): string {
+	return edits
+		.map((edit, index) => {
+			const error =
+				edit.error === undefined
+					? ""
+					: ` This edit produced an error, and was discarded. The error message was: "${edit.error}"`;
+			return `${index + 1}. ${JSON.stringify(edit.edit)}${error}`;
+		})
+		.join("\n");
+}
+
+/**
+ * Generates the main prompt of this explicit strategy.
+ * This prompt is designed to give an LLM instructions on how it can modify a SharedTree using specific types of {@link TreeEdit}'s
+ * and provides with both a serialized version of the current state of the provided tree node as well as  the interfaces that compromise said tree nodes data.
+ */
+export function getEditingSystemPrompt(
+	userPrompt: string,
+	idGenerator: IdGenerator,
+	treeNode: TreeNode,
+	log: EditLog,
+	appGuidance?: string,
+	plan?: string,
+): string {
+	const schema = Tree.schema(treeNode);
+	const promptFriendlySchema = getPromptFriendlyTreeSchema(getJsonSchema(schema));
+	const decoratedTreeJson = toDecoratedJson(idGenerator, treeNode);
+
+	const role = `You are a collaborative agent who interacts with a JSON tree by performing edits to achieve a user-specified goal.${
+		appGuidance === undefined
+			? ""
+			: `
+			The application that owns the JSON tree has the following guidance about your role: ${appGuidance}`
+	}`;
+
+	const treeSchemaString = createZodJsonValidator(
+		...generateGenericEditTypes(getSimpleSchema(schema), false),
+	).getSchemaText();
+
+	// TODO: security: user prompt in system prompt
+	const systemPrompt = `
+	${role}
+	Edits are JSON objects that conform to the following schema.
+	The top level object you produce is an "EditWrapper" object which contains one of "Insert", "Modify", "Remove", "Move", or null.
+	${treeSchemaString}
+	The tree is a JSON object with the following schema: ${promptFriendlySchema}
+	${plan === undefined ? "" : `You have made a plan to accomplish the user's goal. The plan is: "${plan}". You will perform one or more edits that correspond to that plan to accomplish the goal.`}
+	${
+		log.length === 0
+			? ""
+			: `You have already performed the following edits:
+			${createEditListHistoryPrompt(log)}
+			This means that the current state of the tree reflects these changes.`
+	}
+	The current state of the tree is: ${decoratedTreeJson}.
+	${log.length > 0 ? "Before you made the above edits t" : "T"}he user requested you accomplish the following goal:
+	"${userPrompt}"
+	If the goal is now completed or is impossible, you should return null.
+	Otherwise, you should create an edit that makes progress towards the goal. It should have an English description ("explanation") of which edit to perform (specifying one of the allowed edit types).`;
+	return systemPrompt;
+}
+
+/**
+ * Generates a prompt designed to make an LLM review the edits it created and applied to a SharedTree based
+ * on a user-specified goal. This prompt is designed to give the LLM's ability to correct for mistakes and improve the accuracy/fidelity of its final set of tree edits
+ */
+export function getReviewSystemPrompt(
+	userPrompt: string,
+	idGenerator: IdGenerator,
+	treeNode: TreeNode,
+	originalDecoratedJson: string,
+	appGuidance?: string,
+): string {
+	const schema = Tree.schema(treeNode);
+	const promptFriendlySchema = getPromptFriendlyTreeSchema(getJsonSchema(schema));
+	const decoratedTreeJson = toDecoratedJson(idGenerator, treeNode);
+
+	const role = `You are a collaborative agent who interacts with a JSON tree by performing edits to achieve a user-specified goal.${
+		appGuidance === undefined
+			? ""
+			: `
+			The application that owns the JSON tree has the following guidance: ${appGuidance}`
+	}`;
+
+	// TODO: security: user prompt in system prompt
+	const systemPrompt = `
+	${role}
+	You have performed a number of actions already to accomplish a user request.
+	You must review the resulting state to determine if the actions you performed successfully accomplished the user's goal.
+	The tree is a JSON object with the following schema: ${promptFriendlySchema}
+	The state of the tree BEFORE changes was: ${originalDecoratedJson}.
+	The state of the tree AFTER changes is: ${decoratedTreeJson}.
+	The user requested that the following goal should be accomplished:
+	${userPrompt}
+	Was the goal accomplished?`;
+	return systemPrompt;
+}
+
+/**
+ * Converts a fully-qualified SharedTree schema name to a single-word name for use in textual TypeScript-style types.
+ *
+ * @remarks
+ * - TODO: Determine what to do with user-provided names that include periods (e.g. "Foo.Bar").
+ * - TODO: Should probably ensure name starts with an uppercase character.
+ */
+export function getPromptFriendlyTreeSchema(jsonSchema: JsonTreeSchema): string {
+	let stringifiedSchema = "";
+	for (const [name, def] of Object.entries(jsonSchema.$defs)) {
+		if (def.type !== "object" || def._treeNodeSchemaKind === NodeKind.Map) {
+			continue;
+		}
+
+		let stringifiedEntry = `interface ${getFriendlySchemaName(name)} {`;
+
+		for (const [fieldName, fieldSchema] of Object.entries(def.properties)) {
+			let typeString: string;
+			if (isJsonSchemaRef(fieldSchema)) {
+				const nextFieldName = fieldSchema.$ref;
+				const nextDef = getDef(jsonSchema.$defs, nextFieldName);
+				typeString = `${getTypeString(jsonSchema.$defs, [nextFieldName, nextDef])}`;
+			} else {
+				typeString = `${getAnyOfTypeString(jsonSchema.$defs, fieldSchema.anyOf, true)}`;
+			}
+			if (def.required && !def.required.includes(fieldName)) {
+				typeString = `${typeString} | undefined`;
+			}
+			stringifiedEntry += ` ${fieldName}: ${typeString};`;
+		}
+
+		stringifiedEntry += " }";
+
+		stringifiedSchema += (stringifiedSchema === "" ? "" : " ") + stringifiedEntry;
+	}
+	return stringifiedSchema;
+}
+
+function getTypeString(
+	defs: Record<string, JsonNodeSchema>,
+	[name, currentDef]: [string, JsonNodeSchema],
+): string {
+	const { _treeNodeSchemaKind } = currentDef;
+	if (_treeNodeSchemaKind === NodeKind.Leaf) {
+		return currentDef.type;
+	}
+	if (_treeNodeSchemaKind === NodeKind.Object) {
+		return getFriendlySchemaName(name);
+	}
+	if (_treeNodeSchemaKind === NodeKind.Array) {
+		const items = currentDef.items;
+		const innerType = isJsonSchemaRef(items)
+			? getTypeString(defs, [items.$ref, getDef(defs, items.$ref)])
+			: getAnyOfTypeString(defs, items.anyOf);
+		return `${innerType}[]`;
+	}
+	fail("Non-object, non-leaf, non-array schema type.");
+}
+
+function getAnyOfTypeString(
+	defs: Record<string, JsonNodeSchema>,
+	refList: JsonSchemaRef[],
+	topLevel = false,
+): string {
+	const typeNames: string[] = [];
+	for (const ref of refList) {
+		typeNames.push(getTypeString(defs, [ref.$ref, getDef(defs, ref.$ref)]));
+	}
+	const typeString = typeNames.join(" | ");
+	return topLevel ? typeString : `(${typeString})`;
+}
+
+function isJsonSchemaRef(field: JsonFieldSchema): field is JsonSchemaRef {
+	return (field as JsonSchemaRef).$ref !== undefined;
+}
+
+function getDef(defs: Record<string, JsonNodeSchema>, ref: string): JsonNodeSchema {
+	// strip the "#/$defs/" prefix
+	const strippedRef = ref.slice(8);
+	const nextDef = defs[strippedRef];
+	assert(nextDef !== undefined, "Ref not found.");
+	return nextDef;
+}
+
+/**
+ * TBD
+ */
+export function getFriendlySchemaName(schemaName: string): string {
+	const matches = schemaName.match(/[^.]+$/);
+	if (matches === null) {
+		// empty scope
+		return schemaName;
+	}
+	return matches[0];
+}

package/src/explicit-strategy/typeGeneration.ts

@@ -0,0 +1,374 @@
+/*!
+ * Copyright (c) Microsoft Corporation and contributors. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+import { assert } from "@fluidframework/core-utils/internal";
+import { FieldKind, NodeKind, ValueSchema } from "@fluidframework/tree/internal";
+import type {
+	SimpleFieldSchema,
+	SimpleNodeSchema,
+	SimpleTreeSchema,
+} from "@fluidframework/tree/internal";
+import { z } from "zod";
+
+import { objectIdKey, typeField } from "./agentEditTypes.js";
+import { fail, getOrCreate, mapIterable } from "./utils.js";
+
+/**
+ * Zod Object type used to represent & validate the ObjectTarget type within a {@link TreeEdit}.
+ * @remarks this is used as a component with {@link generateGenericEditTypes} to produce the final zod validation objects.
+ */
+const objectTarget = z.object({
+	target: z
+		.string()
+		.describe(
+			`The id of the object (as specified by the object's ${objectIdKey} property) that is being referenced`,
+		),
+});
+/**
+ * Zod Object type used to represent & validate the ObjectPlace type within a {@link TreeEdit}.
+ * @remarks this is used as a component with {@link generateGenericEditTypes} to produce the final zod validation objects.
+ */
+const objectPlace = z
+	.object({
+		type: z.enum(["objectPlace"]),
+		target: z
+			.string()
+			.describe(
+				`The id (${objectIdKey}) of the object that the new/moved object should be placed relative to. This must be the id of an object that already existed in the tree content that was originally supplied.`,
+			),
+		place: z
+			.enum(["before", "after"])
+			.describe(
+				"Where the new/moved object will be relative to the target object - either just before or just after",
+			),
+	})
+	.describe(
+		"A pointer to a location either just before or just after an object that is in an array",
+	);
+/**
+ * Zod Object type used to represent & validate the ArrayPlace type within a {@link TreeEdit}.
+ * @remarks this is used as a component with {@link generateGenericEditTypes} to produce the final zod validation objects.
+ */
+const arrayPlace = z
+	.object({
+		type: z.enum(["arrayPlace"]),
+		parentId: z
+			.string()
+			.describe(
+				`The id (${objectIdKey}) of the parent object of the array. This must be the id of an object that already existed in the tree content that was originally supplied.`,
+			),
+		field: z.string().describe("The key of the array to insert into"),
+		location: z
+			.enum(["start", "end"])
+			.describe("Where to insert into the array - either the start or the end"),
+	})
+	.describe(
+		`either the "start" or "end" of an array, as specified by a "parent" ObjectTarget and a "field" name under which the array is stored (useful for prepending or appending)`,
+	);
+/**
+ * Zod Object type used to represent & validate the Range type within a {@link TreeEdit}.
+ * @remarks this is used as a component with {@link generateGenericEditTypes} to produce the final zod validation objects.
+ */
+const range = z
+	.object({
+		from: objectPlace,
+		to: objectPlace,
+	})
+	.describe(
+		'A range of objects in the same array specified by a "from" and "to" Place. The "to" and "from" objects MUST be in the same array.',
+	);
+/**
+ * Cache used to prevent repeatedly generating the same Zod validation objects for the same {@link SimpleTreeSchema} as generate propts for repeated calls to an LLM
+ */
+const cache = new WeakMap<SimpleTreeSchema, ReturnType<typeof generateGenericEditTypes>>();
+
+/**
+ * Generates a set of ZOD validation objects for the various types of data that can be put into the provided {@link SimpleTreeSchema}
+ * and then uses those sets to generate an all-encompassing ZOD object for each type of {@link TreeEdit} that can validate any of the types of data that can be put into the tree.
+ *
+ * @returns a Record of schema names to Zod validation objects, and the name of the root schema used to encompass all of the other schemas.
+ *
+ * @remarks The return type of this function is designed to work with Typechat's createZodJsonValidator as well as be used as the JSON schema for OpenAi's structured output response format.
+ */
+export function generateGenericEditTypes(
+	schema: SimpleTreeSchema,
+	generateDomainTypes: boolean,
+): [Record<string, Zod.ZodTypeAny>, root: string] {
+	return getOrCreate(cache, schema, () => {
+		const insertSet = new Set<string>();
+		const modifyFieldSet = new Set<string>();
+		const modifyTypeSet = new Set<string>();
+		const typeMap = new Map<string, Zod.ZodTypeAny>();
+		for (const name of schema.definitions.keys()) {
+			getOrCreateType(
+				schema.definitions,
+				typeMap,
+				insertSet,
+				modifyFieldSet,
+				modifyTypeSet,
+				name,
+			);
+		}
+		function getType(allowedTypes: ReadonlySet<string>): Zod.ZodTypeAny {
+			switch (allowedTypes.size) {
+				case 0: {
+					return z.never();
+				}
+				case 1: {
+					return (
+						typeMap.get(tryGetSingleton(allowedTypes) ?? fail("Expected singleton")) ??
+						fail("Unknown type")
+					);
+				}
+				default: {
+					const types = Array.from(
+						allowedTypes,
+						(name) => typeMap.get(name) ?? fail("Unknown type"),
+					);
+					assert(hasAtLeastTwo(types), "Expected at least two types");
+					return z.union(types);
+				}
+			}
+		}
+		const insert = z
+			.object({
+				type: z.literal("insert"),
+				explanation: z.string().describe(editDescription),
+				content: generateDomainTypes
+					? getType(insertSet)
+					: z.any().describe("Domain-specific content here"),
+				destination: z.union([arrayPlace, objectPlace]),
+			})
+			.describe("Inserts a new object at a specific Place or ArrayPlace.");
+		const remove = z
+			.object({
+				type: z.literal("remove"),
+				explanation: z.string().describe(editDescription),
+				source: z.union([objectTarget, range]),
+			})
+			.describe("Deletes an object or Range of objects from the tree.");
+		const move = z
+			.object({
+				type: z.literal("move"),
+				explanation: z.string().describe(editDescription),
+				source: z.union([objectTarget, range]),
+				destination: z.union([arrayPlace, objectPlace]),
+			})
+			.describe("Moves an object or Range of objects to a new Place or ArrayPlace.");
+		const modify = z
+			.object({
+				type: z.enum(["modify"]),
+				explanation: z.string().describe(editDescription),
+				target: objectTarget,
+				field: z.enum([...modifyFieldSet] as [string, ...string[]]), // Modify with appropriate fields
+				modification: generateDomainTypes
+					? getType(modifyTypeSet)
+					: z.any().describe("Domain-specific content here"),
+			})
+			.describe("Sets a field on a specific ObjectTarget.");
+		const editTypes = [insert, remove, move, modify, z.null()] as const;
+		const editWrapper = z.object({
+			edit: z
+				.union(editTypes)
+				.describe("The next edit to apply to the tree, or null if the task is complete."),
+		});
+		const typeRecord: Record<string, Zod.ZodTypeAny> = {
+			ObjectTarget: objectTarget,
+			ObjectPlace: objectPlace,
+			ArrayPlace: arrayPlace,
+			Range: range,
+			Insert: insert,
+			Remove: remove,
+			Move: move,
+			Modify: modify,
+			EditWrapper: editWrapper,
+		};
+		return [typeRecord, "EditWrapper"];
+	});
+}
+const editDescription =
+	"A description of what this edit is meant to accomplish in human readable English";
+function getOrCreateType(
+	definitionMap: ReadonlyMap<string, SimpleNodeSchema>,
+	typeMap: Map<string, Zod.ZodTypeAny>,
+	insertSet: Set<string>,
+	modifyFieldSet: Set<string>,
+	modifyTypeSet: Set<string>,
+	definition: string,
+): Zod.ZodTypeAny {
+	return getOrCreate(typeMap, definition, () => {
+		const nodeSchema = definitionMap.get(definition) ?? fail("Unexpected definition");
+		switch (nodeSchema.kind) {
+			case NodeKind.Object: {
+				for (const [key, field] of Object.entries(nodeSchema.fields)) {
+					// TODO: Remove when AI better
+					if (
+						Array.from(
+							field.allowedTypes,
+							(n) => definitionMap.get(n) ?? fail("Unknown definition"),
+						).some((n) => n.kind === NodeKind.Array)
+					) {
+						continue;
+					}
+					modifyFieldSet.add(key);
+					for (const type of field.allowedTypes) {
+						modifyTypeSet.add(type);
+					}
+				}
+				// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+				const properties = Object.fromEntries(
+					Object.entries(nodeSchema.fields)
+						.map(([key, field]) => {
+							return [
+								key,
+								getOrCreateTypeForField(
+									definitionMap,
+									typeMap,
+									insertSet,
+									modifyFieldSet,
+									modifyTypeSet,
+									field,
+								),
+							];
+						})
+						.filter(([, value]) => value !== undefined),
+				);
+				// eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+				properties[typeField] = z.enum([definition]);
+				return z.object(properties);
+			}
+			case NodeKind.Array: {
+				for (const [name] of Array.from(
+					nodeSchema.allowedTypes,
+					(n): [string, SimpleNodeSchema] => [
+						n,
+						definitionMap.get(n) ?? fail("Unknown definition"),
+					],
+				).filter(
+					([_, schema]) => schema.kind === NodeKind.Object || schema.kind === NodeKind.Leaf,
+				)) {
+					insertSet.add(name);
+				}
+				return z.array(
+					getTypeForAllowedTypes(
+						definitionMap,
+						typeMap,
+						insertSet,
+						modifyFieldSet,
+						modifyTypeSet,
+						nodeSchema.allowedTypes,
+					),
+				);
+			}
+			case NodeKind.Leaf: {
+				switch (nodeSchema.leafKind) {
+					case ValueSchema.Boolean: {
+						return z.boolean();
+					}
+					case ValueSchema.Number: {
+						return z.number();
+					}
+					case ValueSchema.String: {
+						return z.string();
+					}
+					case ValueSchema.Null: {
+						return z.null();
+					}
+					default: {
+						throw new Error(`Unsupported leaf kind ${NodeKind[nodeSchema.leafKind]}.`);
+					}
+				}
+			}
+			default: {
+				throw new Error(`Unsupported node kind ${NodeKind[nodeSchema.kind]}.`);
+			}
+		}
+	});
+}
+function getOrCreateTypeForField(
+	definitionMap: ReadonlyMap<string, SimpleNodeSchema>,
+	typeMap: Map<string, Zod.ZodTypeAny>,
+	insertSet: Set<string>,
+	modifyFieldSet: Set<string>,
+	modifyTypeSet: Set<string>,
+	fieldSchema: SimpleFieldSchema,
+): Zod.ZodTypeAny | undefined {
+	switch (fieldSchema.kind) {
+		case FieldKind.Required: {
+			return getTypeForAllowedTypes(
+				definitionMap,
+				typeMap,
+				insertSet,
+				modifyFieldSet,
+				modifyTypeSet,
+				fieldSchema.allowedTypes,
+			);
+		}
+		case FieldKind.Optional: {
+			return z.union([
+				z.null(),
+				getTypeForAllowedTypes(
+					definitionMap,
+					typeMap,
+					insertSet,
+					modifyFieldSet,
+					modifyTypeSet,
+					fieldSchema.allowedTypes,
+				),
+			]);
+		}
+		case FieldKind.Identifier: {
+			return undefined;
+		}
+		default: {
+			throw new Error(`Unsupported field kind ${NodeKind[fieldSchema.kind]}.`);
+		}
+	}
+}
+function getTypeForAllowedTypes(
+	definitionMap: ReadonlyMap<string, SimpleNodeSchema>,
+	typeMap: Map<string, Zod.ZodTypeAny>,
+	insertSet: Set<string>,
+	modifyFieldSet: Set<string>,
+	modifyTypeSet: Set<string>,
+	allowedTypes: ReadonlySet<string>,
+): Zod.ZodTypeAny {
+	const single = tryGetSingleton(allowedTypes);
+	if (single === undefined) {
+		const types = [
+			...mapIterable(allowedTypes, (name) => {
+				return getOrCreateType(
+					definitionMap,
+					typeMap,
+					insertSet,
+					modifyFieldSet,
+					modifyTypeSet,
+					name,
+				);
+			}),
+		];
+		assert(hasAtLeastTwo(types), "Expected at least two types");
+		return z.union(types);
+	} else {
+		return getOrCreateType(
+			definitionMap,
+			typeMap,
+			insertSet,
+			modifyFieldSet,
+			modifyTypeSet,
+			single,
+		);
+	}
+}
+function tryGetSingleton<T>(set: ReadonlySet<T>): T | undefined {
+	if (set.size === 1) {
+		for (const item of set) {
+			return item;
+		}
+	}
+}
+function hasAtLeastTwo<T>(array: T[]): array is [T, T, ...T[]] {
+	return array.length >= 2;
+}