@fluidframework/ai-collab 2.22.1 → 2.23.0

package/src/explicit-strategy/index.ts

@@ -3,6 +3,7 @@
  * Licensed under the MIT License.
  */
 
+import { UsageError } from "@fluidframework/telemetry-utils/internal";
 import {
 	getSimpleSchema,
 	Tree,
@@ -15,12 +16,34 @@ import type {
 	ChatCompletionCreateParams,
 	// eslint-disable-next-line import/no-internal-modules
 } from "openai/resources/index.mjs";
+import { v4 as uuidv4 } from "uuid";
 import { z } from "zod";
 
-import type { OpenAiClientOptions, TokenLimits, TokenUsage } from "../aiCollabApi.js";
+import type {
+	DebugEventLogHandler,
+	OpenAiClientOptions,
+	TokenLimits,
+	TokenUsage,
+} from "../aiCollabApi.js";
 
 import { applyAgentEdit } from "./agentEditReducer.js";
 import type { EditWrapper, TreeEdit } from "./agentEditTypes.js";
+import {
+	type ApplyEditFailure,
+	type ApplyEditSuccess,
+	type GenerateTreeEditCompleted,
+	type GenerateTreeEditStarted,
+	type FinalReviewCompleted,
+	type FinalReviewStarted,
+	type LlmApiCallDebugEvent,
+	type PlanningPromptCompleted,
+	type CoreEventLoopStarted,
+	type CoreEventLoopCompleted,
+	generateDebugEvent,
+	type PlanningPromptStarted,
+	type EventFlowDebugName,
+	EventFlowDebugNames,
+} from "./debugEvents.js";
 import { IdGenerator } from "./idGenerator.js";
 import {
 	getEditingSystemPrompt,
@@ -32,7 +55,23 @@ import {
 import { generateGenericEditTypes } from "./typeGeneration.js";
 import { fail } from "./utils.js";
 
-const DEBUG_LOG: string[] = [];
+// TODO: Create a proper index file and move the logic of this file to a new location
+export type {
+	ApplyEditFailure,
+	ApplyEditSuccess,
+	CoreEventLoopCompleted,
+	CoreEventLoopStarted,
+	FinalReviewCompleted,
+	FinalReviewStarted,
+	GenerateTreeEditCompleted,
+	GenerateTreeEditStarted,
+	LlmApiCallDebugEvent,
+	PlanningPromptCompleted,
+	PlanningPromptStarted,
+	LlmTreeEdit,
+	EventFlowDebugName,
+	EventFlowDebugNames,
+} from "./debugEvents.js";
 
 /**
  * {@link generateTreeEdits} options.
@@ -54,7 +93,7 @@ export interface GenerateTreeEditsOptions {
 	};
 	finalReviewStep?: boolean;
 	validator?: (newContent: TreeNode) => void;
-	dumpDebugLog?: boolean;
+	debugEventLogHandler?: DebugEventLogHandler;
 	planningStep?: boolean;
 }
 
@@ -65,7 +104,12 @@ interface GenerateTreeEditsSuccessResponse {
 
 interface GenerateTreeEditsErrorResponse {
 	status: "failure" | "partial-failure";
-	errorMessage: "tokenLimitExceeded" | "tooManyErrors" | "tooManyModelCalls" | "aborted";
+	errorMessage:
+		| "tokenLimitExceeded"
+		| "tooManyErrors"
+		| "tooManyModelCalls"
+		| "aborted"
+		| "unexpectedError";
 	tokensUsed: TokenUsage;
 }
 
@@ -91,18 +135,32 @@ export async function generateTreeEdits(
 
 	const tokensUsed = { inputTokens: 0, outputTokens: 0 };
 
+	const debugLogTraceId = uuidv4();
+
+	const coreEventFlowTraceId = uuidv4();
+	options.debugEventLogHandler?.({
+		...generateDebugEvent("CORE_EVENT_LOOP_STARTED", debugLogTraceId),
+		eventFlowName: EventFlowDebugNames.CORE_EVENT_LOOP,
+		eventFlowStatus: "STARTED",
+		eventFlowTraceId: coreEventFlowTraceId,
+	} satisfies CoreEventLoopStarted);
+
 	try {
-		for await (const edit of generateEdits(
+		for await (const generateEditResult of generateEdits(
 			options,
 			simpleSchema,
 			idGenerator,
 			editLog,
 			options.limiters?.tokenLimits,
 			tokensUsed,
+			options.debugEventLogHandler && {
+				eventLogHandler: options.debugEventLogHandler,
+				traceId: debugLogTraceId,
+			},
 		)) {
 			try {
 				const result = applyAgentEdit(
-					edit,
+					generateEditResult.edit,
 					idGenerator,
 					simpleSchema.definitions,
 					options.validator,
@@ -110,55 +168,81 @@ export async function generateTreeEdits(
 				const explanation = result.explanation;
 				editLog.push({ edit: { ...result, explanation } });
 				sequentialErrorCount = 0;
+
+				options.debugEventLogHandler?.({
+					...generateDebugEvent("APPLIED_EDIT_SUCCESS", debugLogTraceId),
+					eventFlowName: EventFlowDebugNames.GENERATE_AND_APPLY_TREE_EDIT,
+					eventFlowStatus: "IN_PROGRESS",
+					eventFlowTraceId: generateEditResult.eventFlowTraceId,
+					edit: generateEditResult.edit as unknown as Record<string, unknown>,
+				} satisfies ApplyEditSuccess);
 			} catch (error: unknown) {
-				if (error instanceof Error) {
+				options.debugEventLogHandler?.({
+					...generateDebugEvent("APPLIED_EDIT_FAILURE", debugLogTraceId),
+					eventFlowName: EventFlowDebugNames.GENERATE_AND_APPLY_TREE_EDIT,
+					eventFlowStatus: "IN_PROGRESS",
+					eventFlowTraceId: generateEditResult.eventFlowTraceId,
+					edit: generateEditResult.edit as unknown as Record<string, unknown>,
+					errorMessage: (error as Error)?.message,
+					sequentialErrorCount,
+				} satisfies ApplyEditFailure);
+
+				if (error instanceof UsageError) {
 					sequentialErrorCount += 1;
-					editLog.push({ edit, error: error.message });
-					DEBUG_LOG?.push(`Error: ${error.message}`);
+					editLog.push({ edit: generateEditResult.edit, error: error.message });
 				} else {
 					throw error;
 				}
 			}
 
-			const responseStatus =
-				editCount > 0 && sequentialErrorCount < editCount ? "partial-failure" : "failure";
+			let shouldExitEarly = false;
+			const completionResponse: GenerateTreeEditsErrorResponse = {
+				status:
+					editCount > 0 && sequentialErrorCount < editCount ? "partial-failure" : "failure",
+				errorMessage: "unexpectedError",
+				tokensUsed,
+			};
 
 			if (options.limiters?.abortController?.signal.aborted === true) {
-				return {
-					status: responseStatus,
-					errorMessage: "aborted",
-					tokensUsed,
-				};
-			}
-
-			if (
+				completionResponse.errorMessage = "aborted";
+				shouldExitEarly = true;
+			} else if (
 				sequentialErrorCount >
 				(options.limiters?.maxSequentialErrors ?? Number.POSITIVE_INFINITY)
 			) {
-				return {
-					status: responseStatus,
-					errorMessage: "tooManyErrors",
-					tokensUsed,
-				};
+				completionResponse.errorMessage = "tooManyErrors";
+				shouldExitEarly = true;
+			} else if (
+				++editCount >= (options.limiters?.maxModelCalls ?? Number.POSITIVE_INFINITY)
+			) {
+				completionResponse.errorMessage = "tooManyModelCalls";
+				shouldExitEarly = true;
 			}
 
-			if (++editCount >= (options.limiters?.maxModelCalls ?? Number.POSITIVE_INFINITY)) {
-				return {
-					status: responseStatus,
-					errorMessage: "tooManyModelCalls",
-					tokensUsed,
-				};
+			if (shouldExitEarly) {
+				options.debugEventLogHandler?.({
+					...generateDebugEvent("CORE_EVENT_LOOP_COMPLETED", debugLogTraceId),
+					eventFlowName: EventFlowDebugNames.CORE_EVENT_LOOP,
+					eventFlowStatus: "COMPLETED",
+					status: "failure",
+					failureReason: completionResponse.errorMessage,
+					eventFlowTraceId: coreEventFlowTraceId,
+				} satisfies CoreEventLoopCompleted);
+
+				return completionResponse;
 			}
 		}
 	} catch (error: unknown) {
-		if (error instanceof Error) {
-			DEBUG_LOG?.push(`Error: ${error.message}`);
-		}
-
-		if (options.dumpDebugLog ?? false) {
-			console.log(DEBUG_LOG.join("\n\n"));
-			DEBUG_LOG.length = 0;
-		}
+		options.debugEventLogHandler?.({
+			...generateDebugEvent("CORE_EVENT_LOOP_COMPLETED", debugLogTraceId),
+			eventFlowName: EventFlowDebugNames.CORE_EVENT_LOOP,
+			eventFlowStatus: "COMPLETED",
+			status: "failure",
+			eventFlowTraceId: coreEventFlowTraceId,
+			failureReason:
+				error instanceof TokenLimitExceededError ? "tokenLimitExceeded" : "unexpectedError",
+			errorMessage: (error as Error)?.message,
+		} satisfies CoreEventLoopCompleted);
 
 		if (error instanceof TokenLimitExceededError) {
 			return {
@@ -171,10 +255,13 @@ export async function generateTreeEdits(
 		throw error;
 	}
 
-	if (options.dumpDebugLog ?? false) {
-		console.log(DEBUG_LOG.join("\n\n"));
-		DEBUG_LOG.length = 0;
-	}
+	options.debugEventLogHandler?.({
+		...generateDebugEvent("CORE_EVENT_LOOP_COMPLETED", debugLogTraceId),
+		eventFlowName: EventFlowDebugNames.CORE_EVENT_LOOP,
+		eventFlowStatus: "COMPLETED",
+		eventFlowTraceId: coreEventFlowTraceId,
+		status: "success",
+	} satisfies CoreEventLoopCompleted);
 
 	return {
 		status: "success",
@@ -202,19 +289,48 @@ async function* generateEdits(
 	editLog: EditLog,
 	tokenLimits: TokenLimits | undefined,
 	tokensUsed: TokenUsage,
-): AsyncGenerator<TreeEdit> {
+	debugOptions?: {
+		eventLogHandler: DebugEventLogHandler;
+		traceId: string;
+	},
+): AsyncGenerator<{ edit: TreeEdit; eventFlowTraceId: string }> {
 	const [types, rootTypeName] = generateGenericEditTypes(simpleSchema, true);
 
 	let plan: string | undefined;
 	if (options.planningStep !== undefined) {
-		const planningPromt = getPlanningSystemPrompt(
+		const planningPrompt = getPlanningSystemPrompt(
 			options.treeNode,
 			options.prompt.userAsk,
 			options.prompt.systemRoleContext,
 		);
-		DEBUG_LOG?.push(planningPromt);
-		plan = await getStringFromLlm(planningPromt, options.openAI, tokensUsed);
-		DEBUG_LOG?.push(`AI Generated the following plan: ${planningPromt}`);
+
+		const generatePlanningPromptEventFlowId = uuidv4();
+		debugOptions?.eventLogHandler?.({
+			...generateDebugEvent("GENERATE_PLANNING_PROMPT_STARTED", debugOptions.traceId),
+			eventFlowName: EventFlowDebugNames.GENERATE_PLANNING_PROMPT,
+			eventFlowTraceId: generatePlanningPromptEventFlowId,
+			eventFlowStatus: "STARTED",
+		} satisfies PlanningPromptStarted);
+
+		plan = await getStringFromLlm(
+			planningPrompt,
+			options.openAI,
+			tokensUsed,
+			debugOptions && {
+				...debugOptions,
+				triggeringEventFlowName: EventFlowDebugNames.GENERATE_PLANNING_PROMPT,
+				eventFlowTraceId: generatePlanningPromptEventFlowId,
+			},
+		);
+
+		debugOptions?.eventLogHandler?.({
+			...generateDebugEvent("GENERATE_PLANNING_PROMPT_COMPLETED", debugOptions.traceId),
+			eventFlowName: EventFlowDebugNames.GENERATE_PLANNING_PROMPT,
+			eventFlowStatus: "COMPLETED",
+			eventFlowTraceId: generatePlanningPromptEventFlowId,
+			isLlmResponseValid: plan !== undefined,
+			llmGeneratedPlan: plan,
+		} satisfies PlanningPromptCompleted);
 	}
 
 	const originalDecoratedJson =
@@ -223,7 +339,9 @@ async function* generateEdits(
 			: undefined;
 	// reviewed is implicitly true if finalReviewStep is false
 	let hasReviewed = (options.finalReviewStep ?? false) ? false : true;
-	async function getNextEdit(): Promise<TreeEdit | undefined> {
+	async function getNextEdit(): Promise<
+		{ edit: TreeEdit; eventFlowTraceId: string } | undefined
+	> {
 		const systemPrompt = getEditingSystemPrompt(
 			options.prompt.userAsk,
 			idGenerator,
@@ -233,30 +351,47 @@ async function* generateEdits(
 			plan,
 		);
 
-		DEBUG_LOG?.push(systemPrompt);
-
 		const schema = types[rootTypeName] ?? fail("Root type not found.");
+
+		const generateTreeEditEventFlowId = uuidv4();
+		debugOptions?.eventLogHandler?.({
+			...generateDebugEvent("GENERATE_TREE_EDIT_STARTED", debugOptions.traceId),
+			eventFlowName: EventFlowDebugNames.GENERATE_AND_APPLY_TREE_EDIT,
+			eventFlowStatus: "STARTED",
+			eventFlowTraceId: generateTreeEditEventFlowId,
+			llmPrompt: systemPrompt,
+		} satisfies GenerateTreeEditStarted);
+
 		const wrapper = await getStructuredOutputFromLlm<EditWrapper>(
 			systemPrompt,
 			options.openAI,
 			schema,
 			"A JSON object that represents an edit to a JSON tree.",
 			tokensUsed,
+			debugOptions && {
+				...debugOptions,
+				triggeringEventFlowName: EventFlowDebugNames.GENERATE_AND_APPLY_TREE_EDIT,
+				eventFlowTraceId: generateTreeEditEventFlowId,
+			},
 		);
 
-		// eslint-disable-next-line unicorn/no-null
-		DEBUG_LOG?.push(JSON.stringify(wrapper, null, 2));
+		debugOptions?.eventLogHandler?.({
+			...generateDebugEvent("GENERATE_TREE_EDIT_COMPLETED", debugOptions.traceId),
+			eventFlowName: EventFlowDebugNames.GENERATE_AND_APPLY_TREE_EDIT,
+			eventFlowStatus: "COMPLETED",
+			eventFlowTraceId: generateTreeEditEventFlowId,
+			isLlmResponseValid: wrapper?.edit !== undefined,
+			llmGeneratedEdit: wrapper?.edit as Record<string, unknown> | null,
+		} satisfies GenerateTreeEditCompleted);
+
 		if (wrapper === undefined) {
-			DEBUG_LOG?.push("Failed to get response");
 			return undefined;
 		}
 
 		if (wrapper.edit === null) {
-			DEBUG_LOG?.push("No more edits.");
 			if ((options.finalReviewStep ?? false) && !hasReviewed) {
 				const reviewResult = await reviewGoal();
 				if (reviewResult === undefined) {
-					DEBUG_LOG?.push("Failed to get review response");
 					return undefined;
 				}
 				// eslint-disable-next-line require-atomic-updates
@@ -270,7 +405,7 @@ async function* generateEdits(
 				}
 			}
 		} else {
-			return wrapper.edit;
+			return { edit: wrapper.edit, eventFlowTraceId: generateTreeEditEventFlowId };
 		}
 	}
 
@@ -283,14 +418,46 @@ async function* generateEdits(
 			options.prompt.systemRoleContext,
 		);
 
-		DEBUG_LOG?.push(systemPrompt);
-
 		const schema = z.object({
 			goalAccomplished: z
 				.enum(["yes", "no"])
 				.describe('Whether the user\'s goal was met in the "after" tree.'),
 		});
-		return getStructuredOutputFromLlm<ReviewResult>(systemPrompt, options.openAI, schema);
+
+		const finalReviewEventFlowTraceId = uuidv4();
+		debugOptions?.eventLogHandler?.({
+			...generateDebugEvent("FINAL_REVIEW_STARTED", debugOptions.traceId),
+			eventFlowName: EventFlowDebugNames.FINAL_REVIEW,
+			eventFlowStatus: "STARTED",
+			eventFlowTraceId: finalReviewEventFlowTraceId,
+			llmPrompt: systemPrompt,
+		} satisfies FinalReviewStarted);
+
+		// TODO: In the future, when using structured output isn't guarenteed, we will
+		// need to add a custom type guard to ensure that output is in the right shape.
+		const output = await getStructuredOutputFromLlm<ReviewResult>(
+			systemPrompt,
+			options.openAI,
+			schema,
+			undefined,
+			tokensUsed,
+			debugOptions && {
+				...debugOptions,
+				triggeringEventFlowName: EventFlowDebugNames.FINAL_REVIEW,
+				eventFlowTraceId: finalReviewEventFlowTraceId,
+			},
+		);
+
+		debugOptions?.eventLogHandler?.({
+			...generateDebugEvent("FINAL_REVIEW_COMPLETED", debugOptions.traceId),
+			eventFlowName: EventFlowDebugNames.FINAL_REVIEW,
+			eventFlowStatus: "COMPLETED",
+			eventFlowTraceId: finalReviewEventFlowTraceId,
+			isLlmResponseValid: output !== undefined,
+			didLlmAccomplishGoal: output?.goalAccomplished,
+		} satisfies FinalReviewCompleted);
+
+		return output;
 	}
 
 	let edit = await getNextEdit();
@@ -315,6 +482,12 @@ async function getStructuredOutputFromLlm<T>(
 	structuredOutputSchema: Zod.ZodTypeAny,
 	description?: string,
 	tokensUsed?: TokenUsage,
+	debugOptions?: {
+		eventLogHandler: DebugEventLogHandler;
+		traceId: string;
+		triggeringEventFlowName: EventFlowDebugName;
+		eventFlowTraceId: string;
+	},
 ): Promise<T | undefined> {
 	const response_format = zodResponseFormat(structuredOutputSchema, "SharedTreeAI", {
 		description,
@@ -328,6 +501,21 @@ async function getStructuredOutputFromLlm<T>(
 
 	const result = await openAi.client.beta.chat.completions.parse(body);
 
+	debugOptions?.eventLogHandler?.({
+		...generateDebugEvent("LLM_API_CALL", debugOptions.traceId),
+		triggeringEventFlowName: debugOptions.triggeringEventFlowName,
+		eventFlowTraceId: debugOptions.eventFlowTraceId,
+		modelName: openAi.modelName ?? "gpt-4o",
+		requestParams: body,
+		response: { ...result },
+		...(result.usage && {
+			tokenUsage: {
+				promptTokens: result.usage.prompt_tokens,
+				completionTokens: result.usage.completion_tokens,
+			},
+		}),
+	} satisfies LlmApiCallDebugEvent);
+
 	if (result.usage !== undefined && tokensUsed !== undefined) {
 		tokensUsed.inputTokens += result.usage?.prompt_tokens;
 		tokensUsed.outputTokens += result.usage?.completion_tokens;
@@ -335,6 +523,7 @@ async function getStructuredOutputFromLlm<T>(
 
 	// TODO: fix types so this isn't null and doesn't need a cast
 	// The type should be derived from the zod schema
+	// TODO: Determine why this value would be undefined.
 	return result.choices[0]?.message.parsed as T | undefined;
 }
 
@@ -345,6 +534,12 @@ async function getStringFromLlm(
 	prompt: string,
 	openAi: OpenAiClientOptions,
 	tokensUsed?: TokenUsage,
+	debugOptions?: {
+		eventLogHandler: DebugEventLogHandler;
+		traceId: string;
+		triggeringEventFlowName: EventFlowDebugName;
+		eventFlowTraceId: string;
+	},
 ): Promise<string | undefined> {
 	const body: ChatCompletionCreateParams = {
 		messages: [{ role: "system", content: prompt }],
@@ -353,6 +548,21 @@ async function getStringFromLlm(
 
 	const result = await openAi.client.chat.completions.create(body);
 
+	debugOptions?.eventLogHandler?.({
+		...generateDebugEvent("LLM_API_CALL", debugOptions.traceId),
+		triggeringEventFlowName: debugOptions.triggeringEventFlowName,
+		eventFlowTraceId: debugOptions.eventFlowTraceId,
+		modelName: openAi.modelName ?? "gpt-4o",
+		requestParams: body,
+		response: { ...result },
+		...(result.usage && {
+			tokenUsage: {
+				promptTokens: result.usage.prompt_tokens,
+				completionTokens: result.usage.completion_tokens,
+			},
+		}),
+	} satisfies LlmApiCallDebugEvent);
+
 	if (result.usage !== undefined && tokensUsed !== undefined) {
 		tokensUsed.inputTokens += result.usage?.prompt_tokens;
 		tokensUsed.outputTokens += result.usage?.completion_tokens;

package/src/index.ts

@@ -27,6 +27,23 @@ export {
 	sharedTreeTraverse,
 } from "./implicit-strategy/index.js";
 
+export type {
+	ApplyEditFailure,
+	ApplyEditSuccess,
+	CoreEventLoopCompleted,
+	CoreEventLoopStarted,
+	FinalReviewCompleted,
+	FinalReviewStarted,
+	GenerateTreeEditCompleted,
+	GenerateTreeEditStarted,
+	LlmApiCallDebugEvent,
+	PlanningPromptCompleted,
+	PlanningPromptStarted,
+	LlmTreeEdit,
+	EventFlowDebugName,
+	EventFlowDebugNames,
+} from "./explicit-strategy/index.js";
+
 export {
 	type AiCollabOptions,
 	type AiCollabSuccessResponse,
@@ -34,6 +51,9 @@ export {
 	type TokenUsage,
 	type TokenLimits,
 	type OpenAiClientOptions,
+	type DebugEvent,
+	type DebugEventLogHandler,
+	type EventFlowDebugEvent,
 } from "./aiCollabApi.js";
 
 export { aiCollab } from "./aiCollab.js";