Human in the Loop

/**
 * Human-in-the-loop using AG-Kit Agents (two-turn flow)
 */

import {
  Agent,
  OpenAIProvider,
  ControlFlowDecision
} from "@ag-kit/agents";
import { z } from "zod";
import { tool } from "@ag-kit/tools";

// Business state for Human-in-the-loop example
class HitLState {
  steps: Array<{ description: string; status: string }> = [];
  userResponse?: string;
  [key: string]: unknown; // Add index signature to satisfy StateConstraint
}

const planExecutionSteps = tool(
  async (input: unknown, context: any) => {
    const { steps } = input as {
      steps: Array<{ description: string; status: string }>;
    };

    // Direct state modification - update the business state directly
    if (context.state) {
      context.state.steps = steps;
      // Clear any existing userResponse to allow new interrupts
      delete context.state.userResponse;
    }

    // Return simple business result
    return {
      success: true,
      steps,
    };
  },
  {
    name: "plan_execution_steps",
    description:
      "Make up 10 steps (only a couple of words per step) that are required for a task. The step should be in imperative form (i.e. Dig hole, Open door, ...)",
    schema: z.object({
      steps: z
        .array(
          z.object({
            description: z
              .string()
              .describe("The text of the step in imperative form"),
            status: z
              .literal("enabled")
              .describe("The status of the step, always 'enabled'"),
          })
        )
        .describe(
          "An array of 10 step objects, each containing text and status"
        ),
    }),
  }
);

export function createHumanInTheLoopAgent(): Agent<HitLState> {
  if (!process.env.OPENAI_API_KEY)
    throw new Error("OPENAI_API_KEY is required");

  const provider = new OpenAIProvider({
    apiKey: process.env.OPENAI_API_KEY!,
    defaultModel: process.env.OPENAI_MODEL || "gpt-4o-mini",
    baseURL: process.env.OPENAI_BASE_URL,
  });

  return new Agent<HitLState>({
    name: "human-in-the-loop-agent",
    description: "Human-in-the-loop planning with two-turn flow",
    model: provider,
    stateType: HitLState,
    instructions:
      "You are a helpful assistant. When the user asks to perform a task, you MUST call the `plan_execution_steps` function to generate 10 imperative steps and then wait for the user's selection in a follow-up turn. In the second turn, produce a short (<= 3 sentences) creative description of how you will perform the task. Do not include any disabled steps; if an essential step is disabled, you may use light humor.",
    modelSettings: { temperature: 0.6, maxTokens: 4096 },
    tools: [planExecutionSteps],
    controlFlow: {
      maxSteps: 10,
      errorRetryLimit: 0,
      customHandler: {
        async handleNextStep(_context: any, state?: HitLState) {
          // Phase inference logic:
          // - No steps = planning phase
          // - Has steps but no user response = waiting phase
          // - Has steps and user response = executing phase

          if (!state) {
            return { action: "continue", nextStep: null } as const;
          }

          const steps = state.steps || [];
          const userResponse = state.userResponse;

          if (steps.length === 0) {
            // Planning phase: continue execution, wait for step generation
            return { action: "continue", nextStep: null } as const;
          }

          if (steps.length > 0 && !userResponse) {
            // Waiting phase: wait for user selection
            return {
              action: "interrupt",
              reason: "waiting_for_user_selection",
              payload: { steps },
            } as ControlFlowDecision;
          }

          if (steps.length > 0 && userResponse) {
            // Executing phase: continue execution
            return { action: "continue", nextStep: null } as const;
          }

          return { action: "continue", nextStep: null } as const;
        },

        async handleToolCall(toolCall: any, _context: any, state?: HitLState) {
          // Tool execution is handled by the tool handler itself
          return { action: "execute", immediate: true } as const;
        },

        async handleError(error: any) {
          return { action: "abort", reason: error.message } as const;
        },
      },
    },
  });
}