feat: support human in the loop for TS

packages/server/src/handlers/chat.ts

@@ -1,8 +1,13 @@
-import type { AgentInputData } from "@llamaindex/workflow";
+import type { AgentInputData, WorkflowContext } from "@llamaindex/workflow";
 import { type Message } from "ai";
 import { IncomingMessage, ServerResponse } from "http";
 import type { MessageType } from "llamaindex";
 import { type WorkflowFactory } from "../types";
+import {
+  createWorkflowContextFromHumanResponse,
+  getHumanResponseFromMessage,
+  pauseForHumanInput,
+} from "../utils/hitl";
 import {
   parseRequestBody,
   pipeStreamToResponse,
@@ -19,6 +24,7 @@ export const handleChat = async (
   suggestNextQuestions: boolean,
 ) => {
   try {
+    const requestId = req.headers["x-request-id"] as string; // TODO: update for chat route also
     const body = await parseRequestBody(req);
     const { messages } = body as { messages: Message[] };
     const chatHistory = messages.map((message) => ({
@@ -41,14 +47,30 @@ export const handleChat = async (
     res.on("close", () => abortController.abort("Connection closed"));
 
     const workflow = await workflowFactory(body);
+    let context: WorkflowContext;
+
+    // if there is human response, we need to create a workflow context from the human response
+    // otherwise, we can start with empty context
+    const humanResponse = getHumanResponseFromMessage(lastMessage);
+    if (humanResponse) {
+      context = await createWorkflowContextFromHumanResponse(
+        workflow,
+        requestId,
+        humanResponse,
+      );
+    } else {
+      context = workflow.createContext();
+    }
+
     const workflowEventStream = await runWorkflow(
-      workflow,
+      context,
       workflowInput,
       abortController.signal,
     );
 
     const dataStream = toDataStream(workflowEventStream, {
       callbacks: {
+        onPauseForHumanInput: () => pauseForHumanInput(context, requestId),
         onFinal: async (completion, dataStreamWriter) => {
           chatHistory.push({
             role: "assistant" as MessageType,

packages/server/src/utils/hitl.ts

@@ -0,0 +1,113 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+
+import {
+  withSnapshot,
+  workflowEvent,
+  type Workflow,
+  type WorkflowContext,
+} from "@llamaindex/workflow";
+import type { Message } from "ai";
+import { z } from "zod";
+
+// @llama-flow doesn't export snapshot types, we need to infer them from the functions
+export type SnapshotWorkflow = ReturnType<typeof withSnapshot<Workflow>>;
+export type SnapshotWorkflowContext = ReturnType<
+  SnapshotWorkflow["createContext"]
+>;
+
+export const serializableMemoryMap = new Map<string, any>();
+
+export const saveSnapshot = async (requestId: string, snapshot: any) => {
+  // TODO: save to file
+  serializableMemoryMap.set(requestId, snapshot);
+};
+
+export const loadSnapshot = async (
+  requestId: string,
+): Promise<any | undefined> => {
+  return serializableMemoryMap.get(requestId);
+};
+
+export type HumanInputEventData = {
+  event_type: string; // An identifier for the input component in UI
+  data: unknown; // The data to be sent to the input component in UI
+};
+
+// humanInputEvent should be triggered when workflow need to request input from user
+// when it is emitted, workflow snapshot will be saved and stream will be paused
+// then send HumanInputEventData as annotation to UI to render the input form
+export const humanInputEvent = workflowEvent<HumanInputEventData>();
+
+// TODO: we can consider sending JSON Schema for the requested information from user
+// Then render it as a form in UI with https://github.com/rjsf-team/react-jsonschema-form
+// we can call it formInputEvent (same logic as humanInputEvent but useful when requesting multiple inputs)
+
+export type HumanResponseEventData = {
+  data: unknown; // The response data from user
+};
+
+// When user make a response to the input request, workflow will be re-created from the last snapshot
+// and then trigger humanResponseEvent to resume the workflow
+export const humanResponseEvent = workflowEvent<HumanResponseEventData>();
+
+export const humanResponseAnnotationSchema = z.object({
+  type: z.literal("human_response"),
+  data: z.any(),
+});
+
+export const getHumanResponseFromMessage = (message: Message) => {
+  if (message.annotations) {
+    for (const annotation of message.annotations) {
+      if (humanResponseAnnotationSchema.safeParse(annotation).success) {
+        return (annotation as z.infer<typeof humanResponseAnnotationSchema>)
+          .data;
+      }
+    }
+  }
+  return null;
+};
+
+export const createWorkflowContextFromHumanResponse = async (
+  workflow: Workflow,
+  requestId: string,
+  humanResponse: any,
+): Promise<SnapshotWorkflowContext> => {
+  // check workflow is snapshotable
+  if (!("resume" in workflow)) {
+    // TODO: ensure AgentWorkflow is snapshotable
+    throw new Error(
+      "Workflow is not a snapshot workflow. Please use withSnapshot() to make it snapshotable.",
+    );
+  }
+
+  // if there is no snapshot, we can't resume the workflow
+  const snapshot = await loadSnapshot(requestId);
+  if (!snapshot) {
+    throw new Error("No snapshot found for request id: " + requestId);
+  }
+
+  // resume the workflow from the snapshot with human response
+  const context = (workflow as SnapshotWorkflow).resume(
+    humanResponse,
+    snapshot,
+  );
+
+  return context;
+};
+
+export const pauseForHumanInput = async (
+  context: WorkflowContext,
+  requestId: string,
+) => {
+  if (!("snapshot" in context)) {
+    // check workflow is snapshotable
+    throw new Error(
+      "Cannot get snapshot of the workflow. Please use withSnapshot() to make it snapshotable.",
+    );
+  }
+
+  // save snapshot with the key is requestId
+  const snapshotContext = context as SnapshotWorkflowContext;
+  const snapshot = await snapshotContext.snapshot();
+  await saveSnapshot(requestId, snapshot);
+};

packages/server/src/utils/index.ts

@@ -1,6 +1,7 @@
 export * from "./events";
 export * from "./file";
 export * from "./gen-ui";
+export * from "./hitl";
 export * from "./inline";
 export * from "./prompts";
 export * from "./request";

packages/server/src/utils/stream.ts

@@ -5,6 +5,7 @@ import {
   type DataStreamWriter,
   type JSONValue,
 } from "ai";
+import { humanInputEvent, type HumanInputEventData } from "./hitl";
 
 /**
  * Configuration options and helper callback methods for stream lifecycle events.
@@ -24,6 +25,9 @@ export interface StreamCallbacks {
     text: string,
     dataStreamWriter: DataStreamWriter,
   ) => Promise<void> | void;
+
+  /** `onPauseForHumanInput`: Called when human input   event is emitted. */
+  onPauseForHumanInput?: (event: HumanInputEventData) => Promise<void> | void;
 }
 
 /**
@@ -61,6 +65,12 @@ export function toDataStream(
               await callbacks.onText(content, dataStreamWriter);
             }
           }
+        } else if (humanInputEvent.include(event)) {
+          if (callbacks?.onPauseForHumanInput) {
+            await callbacks.onPauseForHumanInput(event.data);
+          }
+          dataStreamWriter.writeMessageAnnotation(event.data as JSONValue);
+          break; // break to stop the stream
         } else {
           dataStreamWriter.writeMessageAnnotation(event.data as JSONValue);
         }

packages/server/src/utils/workflow.ts

@@ -1,12 +1,11 @@
 import {
   agentToolCallEvent,
   agentToolCallResultEvent,
-  run,
   startAgentEvent,
   stopAgentEvent,
   WorkflowStream,
   type AgentInputData,
-  type Workflow,
+  type WorkflowContext,
   type WorkflowEventData,
 } from "@llamaindex/workflow";
 import {
@@ -25,19 +24,20 @@ import { downloadFile } from "./file";
 import { toInlineAnnotationEvent } from "./inline";
 
 export async function runWorkflow(
-  workflow: Workflow,
+  context: WorkflowContext,
   input: AgentInputData,
   abortSignal?: AbortSignal,
 ): Promise<WorkflowStream<WorkflowEventData<unknown>>> {
   if (!input.userInput) {
     throw new Error("Missing user input to start the workflow");
   }
-  const workflowStream = run(workflow, [
+  context.sendEvent(
     startAgentEvent.with({
       userInput: input.userInput,
       chatHistory: input.chatHistory,
     }),
-  ]);
+  );
+  const workflowStream = context.stream;
 
   // Transform the stream to handle annotations
   return processWorkflowStream(workflowStream).until(