Ai collab explicit #22836

seanimam · 2024-10-17T18:11:01Z

Adds the Explicit Strategy to the ai-collab package, exported under a new, shared API surface.

Description

The following PR implements the explicit strategy to the ai-collab package and exports it under a new, simple API surface. (See aiCollabApi.ts).

Folder Structure

/explicit-strategy: The new explicit strategy, utilizing the prototype built during the fall FHL, with a few adjustments.
- agentEditReducer: This file houses the logic for taking in a TreeEdit, which the LLM produces, and applying said edit to the
- actual SharedTree.
- agentEditTypes.ts: The types of edits an LLM is prompted to produce in order to modify a SharedTree.
- idGenerator.ts: `A manager for producing and mapping simple id's in place of UUID hashes when generating prompts for an LLM
- jsonTypes.ts: utility JSON related types used in parsing LLM response and generating LLM prompts.
- promptGeneration.ts: Logic for producing the different types of prompts sent to an LLM in order to edit a SharedTree.
- typeGeneration.ts: Generates serialized(/able) representations of a SharedTree Schema which is used within prompts and the generated of the structured output JSON schema
- utils.ts: Utilities for interacting with a SharedTree
/implicit-strategy: (... original implicit strategy code, unmodified and move to a new subfolder)

You'll see that the ai-collab package has been restructured to contain explicit-strategy and implicit-strategy subfolders while the root has the aiCollab.ts file which is the actual exported API surface of this package.

There are many new @internal exports that needed to be exposed from SharedTree to get this package working properly.

In a follow up PR, all exports from implicit-strategy/index.js will be removed. They are left in place to minimize the blast radius of this PR (which involves editing the ai sample application to use the new code.

Usage

Your SharedTree types file

This file is where we define the types of our task management application's SharedTree data

//  --------- File name: "types.ts" ---------
import { SchemaFactory } from "@fluidframework/tree";

const sf = new SchemaFactory("ai-collab-sample-application");

export class Task extends sf.object("Task", {
	title: sf.required(sf.string, {
		metadata: { description: `The title of the task` },
	}),
	id: sf.identifier,
	description: sf.required(sf.string, {
		metadata: { description: `The description of the task` },
	}),
	priority: sf.required(sf.string, {
		metadata: { description: `The priority of the task in three levels, "low", "medium", "high"` },
	}),
	complexity: sf.required(sf.number, {
		metadata: { description: `The complexity of the task as a fibonacci number` },
	}),
	status: sf.required(sf.string, {
		metadata: { description: `The status of the task as either "todo", "in-progress", or "done"` },
	}),
	assignee: sf.required(sf.string, {
		metadata: { description: `The name of the tasks assignee e.g. "Bob" or "Alice"` },
	}),
}) {}

export class TaskList extends sf.array("TaskList", SharedTreeTask) {}

export class Engineer extends sf.object("Engineer", {
	name: sf.required(sf.string, {
		metadata: { description: `The name of an engineer whom can be assigned to a task` },
	}),
	id: sf.identifier,
	skills: sf.required(sf.string, {
		metadata: { description: `A description of the engineers skills which influence what types of tasks they should be assigned to.` },
	}),
	maxCapacity: sf.required(sf.number, {
		metadata: { description: `The maximum capacity of tasks this engineer can handle measured in in task complexity points.` },
	}),
}) {}

export class EngineerList extends sf.array("EngineerList", SharedTreeEngineer) {}

export class TaskGroup extends sf.object("TaskGroup", {
	description: sf.required(sf.string, {
		metadata: { description: `The description of the task group, which is a collection of tasks and engineers that can be assigned to said tasks.` },
	}),
	id: sf.identifier,
	title: sf.required(sf.string, {
		metadata: { description: `The title of the task group.` },
	}),
	tasks: sf.required(SharedTreeTaskList, {
		metadata: { description: `The lists of tasks within this task group.` },
	}),
	engineers: sf.required(SharedTreeEngineerList, {
		metadata: { description: `The lists of engineers within this task group which can be assigned to tasks.` },
	}),
}) {}

export class TaskGroupList extends sf.array("TaskGroupList", SharedTreeTaskGroup) {}

export class PlannerAppState extends sf.object("PlannerAppState", {
	taskGroups: sf.required(SharedTreeTaskGroupList, {
		metadata: { description: `The list of task groups that are being managed by this task management application.` },
	}),
}) {}

Example 1: Collaborate with AI

import { aiCollab } from "@fluid-experimental/ai-collab";
import { PlannerAppState } from "./types.ts"
// This is not a real file, this is meant to represent how you initialize your app data.
import { initializeAppState } from "./yourAppInitializationFile.ts"

//  --------- File name: "app.ts" ---------

// Initialize your app state somehow
const appState: PlannerAppState = initializeAppState({
		taskGroups: [
		{
			title: "My First Task Group",
			description: "Placeholder for first task group",
			tasks: [
				{
					assignee: "Alice",
					title: "Task #1",
					description:
						"This is the first Sample task.",
					priority: "low",
					complexity: 1,
					status: "todo",
				},
			],
			engineers: [
				{
					name: "Alice",
					maxCapacity: 15,
					skills:
						"Senior engineer capable of handling complex tasks. Versed in most languages",
				},
				{
					name: "Charlie",
					maxCapacity: 7,
					skills: "Junior engineer capable of handling simple tasks. Versed in Node.JS",
				},
			],
		},
	],
})

// Typically, the user would input this through a UI form/input of some sort.
const userAsk = "Update the task group description to be a about creating a new Todo list application. Create a set of tasks to accomplish this and assign them to the available engineers. Keep in mind the max capacity of each engineer as you assign tasks."

// Collaborate with AI one function call.
const response = await aiCollab<typeof PlannerAppState>({
		openAI: {
			client: new OpenAI({
				apiKey: OPENAI_API_KEY,
			}),
			modelName: "gpt-4o",
		},
		treeView: view,
		treeNode: view.root.taskGroups[0],
		prompt: {
			systemRoleContext:
				"You are a manager that is helping out with a project management tool. You have been asked to edit a group of tasks.",
			userAsk: userAsk,
		},
		planningStep: true,
		finalReviewStep: true,
		dumpDebugLog: true,
	});

if (response.status === 'sucess') {
	// Render the UI view of your task groups.
	window.alert(`The AI has successfully completed your request.`);
} else {
	window.alert(`Something went wrong! response status: ${response.status}, error message: ${response.errorMessage}`);
}

Once the aiCollab function call is initiated, an LLM will immediately begin attempting to make changes to your Shared Tree using the provided user prompt, the types of your SharedTree and the provided app guidance. The LLM produces multiple changes, in a loop asynchronously. Meaning, you will immediatley see changes if your UI's render loop is connected to your SharedTree App State.

How the explicit strategy has changed since FHL

The logic remains intact with small adjustments:

The generateTreeEdits and related API's have been slightly adjusted, with the only logic change being they now accepts a TreeNode instead of the entire tree, allowing the LLM to work on subtree's and omitting data from parts of the tree that are not intended to be worked on. This also means that in places where the tree's schema must be accessed, there is branching logic based on whether the root or a child node is passed to the function, which determines whether root tree's schema needs to be generated or just the schema of the tree node.
The API's now allow token limitations. You now have the ability to pass an object that will be updated with token usage numbers as various functions calls are made to an LLM
Recent changes to move away from the json handler streaming approach have been merged in.

Known Issues & limitations

Union types for a TreeNode are not present when generating App Schema. This will require extracting a field schema instead of TreeNodeSchema when passed a non root node.
The Editing System prompt & structured out schema currently provide array related edits even when there are no arrays. This forces you to have an array in your schema to produce a valid json schema
Optional roots are not allowed, This is because if you pass undefined as your treeNode to the API, we cannot disambiguate whether you passed the root or not.
Primitive root nodes are not allowed to be passed to the API. You must use an object or array as your root.
Optional nodes are not supported -- when we use optional nodes, the OpenAI API returns an error complaining that the structured output JSON schema is invalid. I have introduced a fix that should work upon manual validation of the json schema, but there looks to be an issue with their API. I have filed a ticket with OpenAI to address this
The current scheme does not allow manipulation of arrays of primitive values because you cannot refer to them. We could accomplish this via a path (probably JSON Pointer or JSONPath) from a possibly-null objectId, or wrap arrays in an identified object.
Only 100 object fields total are allowed by OpenAI right now, so larger schemas will fail faster if we have a bunch of schema types generated for type-specific edits.
We don't support nested arrays yet.
Handle 429 rate limit error in streamFromLlm.
Top level arrays are not supported with current DSL.
Structured Output fails when multiple schema types have the same first field name (e.g. id: sf.identifier on multiple types).
Pass descriptions from schema metadata to the generated TS types that we put in the prompt.

Reviewer Guidance

SharedTree team please look over new internal exports from SharedTree and the explicit-strategy folder code.
Please give feedback on the API surface for aiCollab. See aiCollabApi.ts
Tests for the token limits have yet to be added

…r one shared api surface

…to aiCollabExplicit

…tId parameter in diffs

msfluid-bot

Code Coverage Summary

↑ packages.dds.tree.src.simple-tree:
Line Coverage Change: 0.07% Branch Coverage Change: 0.20%

Metric Name	Baseline coverage	PR coverage	Coverage Diff
Branch Coverage	93.29%	93.49%	↑ 0.20%
Line Coverage	97.13%	97.20%	↑ 0.07%

↑ packages.dds.tree.src.simple-tree.api:
Line Coverage Change: 0.01% Branch Coverage Change: 1.14%

Metric Name	Baseline coverage	PR coverage	Coverage Diff
Branch Coverage	86.53%	87.67%	↑ 1.14%
Line Coverage	82.12%	82.13%	↑ 0.01%

↑ packages.framework.ai-collab.src:
Line Coverage Change: 30.72% Branch Coverage Change: No change

Metric Name	Baseline coverage	PR coverage	Coverage Diff
Branch Coverage	0.00%	0.00%	→ No change
Line Coverage	0.00%	30.72%	↑ 30.72%

↑ packages.framework.ai-collab.src.explicit-strategy:
Line Coverage Change: 57.98% Branch Coverage Change: 75.00%

Metric Name	Baseline coverage	PR coverage	Coverage Diff
Branch Coverage	0.00%	75.00%	↑ 75.00%
Line Coverage	0.00%	57.98%	↑ 57.98%

↑ packages.framework.ai-collab.src.explicit-strategy:
Line Coverage Change: 57.98% Branch Coverage Change: 75.00%

Metric Name	Baseline coverage	PR coverage	Coverage Diff
Branch Coverage	0.00%	75.00%	↑ 75.00%
Line Coverage	0.00%	57.98%	↑ 57.98%

Baseline commit: 199b9d0
Baseline build: 302875
Happy Coding!!

Code coverage comparison check passed!!

…n exported from tree

…logic because its not in use.

packages/framework/ai-collab/api-report/ai-collab.alpha.api.md

packages/framework/ai-collab/src/README.md

Josmithr · 2024-10-29T22:50:36Z

packages/framework/ai-collab/src/explicit-strategy/agentEditTypes.ts

+ */
+
+/**
+ * TBD


@taylorsw04 @noencke Suggestions for this? (and others?)

I want to make sure these get documented before we merge into main, otherwise they'll probably be lost 🫤

I've given a go at adding jsdoc for these types. I'd like the SharedTree team to still review and confirm the jsdoc. There are also some remarks asking what different thing are.

Josmithr · 2024-10-29T22:51:45Z

packages/framework/ai-collab/src/explicit-strategy/agentEditTypes.ts

+ */
+export const typeField = "__fluid_type";
+/**
+ * TBD


This is auto-generated and injected into nodes before passing data to the LLM to ensure the LLM can refer to nodes in a stable way.

Josmithr · 2024-10-29T22:53:43Z

packages/framework/ai-collab/src/explicit-strategy/agentEditTypes.ts

+import type { JsonPrimitive } from "./jsonTypes.js";
+
+/**
+ * TODO: The current scheme does not allow manipulation of arrays of primitive values because you cannot refer to them.


We should probably file follow-up tasks for many of these. I suspect most/all aren't blockers for the alpha release, but they should still be on our radar. Some will certainly be blocked on missing SharedTree functionality - identifying those blockers would be useful.

@taylorsw04 maybe you can help point us to the relevant backlog items that would be blockers for these.

packages/framework/ai-collab/src/explicit-strategy/idGenerator.ts

Josmithr · 2024-10-29T22:56:29Z

packages/framework/ai-collab/src/explicit-strategy/idGenerator.ts

+					this.assignIds(element);
+				});
+			} else {
+				// TODO: SharedTree Team needs to either publish TreeNode as a class to use .instanceof() or a typeguard.


@taylorsw04 are you guys tracking this?

Josmithr · 2024-10-29T22:57:09Z

packages/framework/ai-collab/src/explicit-strategy/index.ts

+import { generateGenericEditTypes } from "./typeGeneration.js";
+import { fail } from "./utils.js";
+
+const DEBUG_LOG: string[] = [];


Probably best to remove this?

Josmithr · 2024-10-29T22:59:28Z

packages/framework/ai-collab/src/explicit-strategy/index.ts

+
+	const isRootNode = Tree.parent(options.treeNode) === undefined;
+	const simpleSchema = isRootNode
+		? getSimpleSchema(normalizeFieldSchema(options.treeView.schema).allowedTypes)


I believe the following should work and is a bit simpler:

Suggested change

? getSimpleSchema(normalizeFieldSchema(options.treeView.schema).allowedTypes)

? getSimpleSchema(options.treeView.schema)

Josmithr · 2024-10-29T23:00:28Z

packages/framework/ai-collab/src/explicit-strategy/index.ts

+export interface GenerateTreeEditsOptions<TSchema extends ImplicitFieldSchema> {
+	openAI: OpenAiClientOptions;
+	treeView: TreeView<TSchema>;
+	treeNode: TreeNode;


It would be good to document how treeNode and treeView relate here. Is treeView needed?

packages/framework/ai-collab/src/explicit-strategy/index.ts

Josmithr · 2024-10-29T23:05:19Z

packages/framework/ai-collab/src/explicit-strategy/promptGeneration.ts

+	const schema = isRootNode
+		? normalizeFieldSchema(view.schema)
+		: normalizeFieldSchema(Tree.schema(treeNode));
+	const promptFriendlySchema = getPromptFriendlyTreeSchema(getJsonSchema(schema.allowedTypes));


I believe you should be able to simplify this to

Suggested change

const promptFriendlySchema = getPromptFriendlyTreeSchema(getJsonSchema(schema.allowedTypes));

const promptFriendlySchema = getPromptFriendlyTreeSchema(getJsonSchema(schema));

Josmithr · 2024-10-29T23:05:46Z

packages/framework/ai-collab/src/explicit-strategy/promptGeneration.ts

+}
+
+/**
+ * TBD


@noencke @taylorsw04 recommendations?

(and for the others)

packages/framework/ai-collab/src/explicit-strategy/typeGeneration.ts

Josmithr · 2024-10-29T23:12:17Z

packages/framework/ai-collab/src/explicit-strategy/utils.ts

+ * Licensed under the MIT License.
+ */
+
+/**


Do we have anything tracking moving these utilities to a shared location? fail should probably live next to our assert. The map utilities could probably go in core-utils.

@noencke thoughts?

…rom integ test, adds baseline comments to agentEditTypes.ts

…unused tests

…hemaMetadata

…ailure error message as intented to generateTreeEdits

seanimam · 2024-10-31T16:41:31Z

Blocking on the changes to tree's public API surface. I don't think we should be adding a omitFromJson property. I left a comment with my suggested alternative pattern (that is more user-customizable and avoids changes to the public API.

I've removed omitFromJson. It's unnecessary for our demo applications anyways. Please check out the updates.

Josmithr · 2024-10-31T17:12:05Z

packages/dds/tree/src/index.ts

+	numberSchema,
+	stringSchema,
+	booleanSchema,
+	handleSchema,
+	nullSchema,


Are these needed? Can we not use the schemaFactory.number, etc. pattern? Not super concerned about exposing these as @internal, just wondering if it's strictly necessary.

seanimam added 4 commits October 15, 2024 22:02

Adds explicit strategy with a few unresovled type errors

5f912d8

New SharedTree exports for aiCollab library

ab3b543

Updates ai-collab to include the explicit strategy and export it unde…

822092e

…r one shared api surface

Merge branch 'main' of https://github.com/microsoft/FluidFramework in…

478c5b2

…to aiCollabExplicit

seanimam added 3 commits October 18, 2024 19:30

Adds back accidentlaly removed dependencies from tree

4b5bc91

fixes broken tests for implicit strategy not accounting for new objec…

c8b6684

…tId parameter in diffs

WIP adding test for explicit-strategy

3466402

msfluid-bot reviewed Oct 18, 2024

View reviewed changes

seanimam added 10 commits October 18, 2024 21:36

comments out TreeNode typeguarding as we don't have a working solutio…

62b3783

…n exported from tree

Adds agentEditingReducer tests, 2 failing currently

ecd9746

Adds remaining tests with two failing

ccfe5b6

small import warning fix

2be8309

In progress adding planning step and other new updates from taylors fork

a01a52a

Updates package with new planning step and removal of json streaming …

a29d48a

…logic because its not in use.

Adds jsdoc, resolves build errors, work in progress integration tests

0cf8cfe

skips integ tests

232d335

Adds more documented bug cases to integration testing

e5e3b2d

API Updated to accept a treeNode instead of the whole tree

b60ecf3

seanimam marked this pull request as ready for review October 28, 2024 21:10

seanimam requested review from a team as code owners October 28, 2024 21:10

seanimam requested review from a team and noencke October 28, 2024 21:13

Small integ test update

c3c50d8