feat: add withSnapshot API (#97)

Co-authored-by: Marcus Schiesser <[email protected]>

Author: himself65

.changeset/pretty-moose-share.md

@@ -0,0 +1,12 @@
+---
+"demo": patch
+"@llama-flow/core": patch
+---
+
+feat: add `withSnapshot` middleware API
+
+Add snapshot API, for human in the loop feature. The API is designed for cross JavaScript platform, including node.js, browser, and serverless platform such as cloudflare worker and edge runtime
+
+- `workflow.createContext(): Context`
+- `context.snapshot(): Promise<[requestEvent, snapshot]>`
+- `workflow.resume(data, snapshot)`

demo/hono/app.ts

@@ -18,6 +18,54 @@ app.post(
   ),
 );
 
+const serializableMemoryMap = new Map<string, any>();
+
+app.post("/human-in-the-loop", async (ctx) => {
+  const { workflow, stopEvent, startEvent, humanInteractionEvent } =
+    await import("../workflows/human-in-the-loop");
+  const json = await ctx.req.json();
+  let context: ReturnType<typeof workflow.createContext>;
+  if (json.requestId) {
+    const data = json.data;
+    const serializable = serializableMemoryMap.get(json.requestId);
+    context = workflow.resume(data, serializable);
+  } else {
+    context = workflow.createContext();
+    context.sendEvent(startEvent.with(json.data));
+  }
+
+  const { onRequest, stream } = context;
+  return new Promise<Response>((resolve) => {
+    // listen to human interaction
+    onRequest(humanInteractionEvent, async (reason) => {
+      context.snapshot().then(([re, sd]) => {
+        const requestId = crypto.randomUUID();
+        serializableMemoryMap.set(requestId, sd);
+        resolve(
+          Response.json({
+            requestId: requestId,
+            reason: reason,
+            data: re.map((r) =>
+              r === humanInteractionEvent
+                ? "request human in the loop"
+                : "UNKNOWN",
+            ),
+          }),
+        );
+      });
+    });
+
+    // consume stream
+    stream
+      .until(stopEvent)
+      .toArray()
+      .then((events) => {
+        const stopEvent = events.at(-1)!;
+        resolve(Response.json(stopEvent.data));
+      });
+  });
+});
+
 serve(app, ({ port }) => {
   console.log(`Server started at http://localhost:${port}`);
 });

demo/node/name-ask-readline.ts

@@ -0,0 +1,29 @@
+import { input } from "@inquirer/prompts";
+import {
+  workflow,
+  stopEvent,
+  startEvent,
+  humanInteractionEvent,
+} from "../workflows/human-in-the-loop";
+
+const name = await input({
+  message: "What is your name?",
+});
+const { onRequest, stream, sendEvent } = workflow.createContext();
+
+sendEvent(startEvent.with(name));
+
+onRequest(humanInteractionEvent, async (reason) => {
+  console.log("Requesting human interaction...");
+  const name = await input({
+    message: JSON.parse(reason).message,
+  });
+  console.log("Human interaction completed.");
+  sendEvent(humanInteractionEvent.with(name));
+});
+
+stream.on(stopEvent, ({ data }) => {
+  console.log("AI analysis: ", data);
+});
+
+await stream.until(stopEvent).toArray();

demo/package.json

@@ -4,6 +4,7 @@
   "private": true,
   "dependencies": {
     "@hono/node-server": "^1.14.1",
+    "@inquirer/prompts": "^7.5.0",
     "@llama-flow/core": "latest",
     "@modelcontextprotocol/sdk": "^1.10.1",
     "hono": "^4.7.7",

demo/workflows/human-in-the-loop.ts

@@ -0,0 +1,73 @@
+import { withSnapshot, request } from "@llama-flow/core/middleware/snapshot";
+import { createWorkflow, workflowEvent, getContext } from "@llama-flow/core";
+import { OpenAI } from "openai";
+
+const openai = new OpenAI();
+
+const workflow = withSnapshot(createWorkflow());
+
+const startEvent = workflowEvent<string>({
+  debugLabel: "start",
+});
+const humanInteractionEvent = workflowEvent<string>({
+  debugLabel: "humanInteraction",
+});
+const stopEvent = workflowEvent<string>({
+  debugLabel: "stop",
+});
+
+workflow.handle([startEvent], async ({ data }) => {
+  const response = await openai.chat.completions.create({
+    stream: false,
+    model: "gpt-4.1",
+    messages: [
+      {
+        role: "system",
+        content: `You are a helpful assistant.
+If user doesn't provide his/her name, call ask_name tool to ask for user's name.
+Otherwise, analyze user's name with a good meaning and return the analysis.
+
+For example, alex is from "Alexander the Great", who was a king of the ancient Greek kingdom of Macedon and one of history's greatest military minds.`,
+      },
+      {
+        role: "user",
+        content: data,
+      },
+    ],
+    tools: [
+      {
+        type: "function",
+        function: {
+          name: "ask_name",
+          description: "Ask for user's name",
+          parameters: {
+            type: "object",
+            properties: {
+              message: {
+                type: "string",
+                description: "The message to ask for user's name",
+              },
+            },
+            required: ["message"],
+          },
+        },
+      },
+    ],
+  });
+  const tools = response.choices[0].message.tool_calls;
+  if (tools && tools.length > 0) {
+    const askName = tools.find((tool) => tool.function.name === "ask_name");
+    if (askName) {
+      return request(humanInteractionEvent, askName.function.arguments);
+    }
+  }
+  return stopEvent.with(response.choices[0].message.content!);
+});
+
+workflow.handle([humanInteractionEvent], async ({ data }) => {
+  const { sendEvent } = getContext();
+  // going back to the start event
+  sendEvent(startEvent.with(data));
+});
+
+export { workflow, startEvent, humanInteractionEvent, stopEvent };

packages/core/package.json

@@ -68,6 +68,10 @@
       "types": "./middleware/validation.d.ts",
       "default": "./middleware/validation.js"
     },
+    "./middleware/snapshot": {
+      "types": "./middleware/snapshot.d.ts",
+      "default": "./middleware/snapshot.js"
+    },
     "./util/p-retry": {
       "types": "./util/p-retry.d.ts",
       "default": "./util/p-retry.js"

packages/core/src/core/context.ts

@@ -156,11 +156,11 @@ export const createContext = ({
             // return value is a special event
             if (isPromiseLike(result)) {
               (handlerContext as any).async = true;
-              (handlerContext as any).pending = result;
-              result.then((event) => {
+              (handlerContext as any).pending = result.then((event) => {
                 if (isEventData(event)) {
                   workflowContext.sendEvent(event);
                 }
+                return event;
               });
             } else if (isEventData(result)) {
               workflowContext.sendEvent(result);
@@ -196,44 +196,50 @@ export const createContext = ({
   };
   const createWorkflowContext = (
     handlerContext: HandlerContext,
-  ): WorkflowContext => ({
-    get stream() {
-      const subscribable = createSubscribable<
-        [event: WorkflowEventData<any>],
-        void
-      >();
-      rootWorkflowContext.__internal__call_send_event.subscribe(
-        (newEvent: WorkflowEventData<any>) => {
-          let currentEventContext = eventContextWeakMap.get(newEvent);
-          while (currentEventContext) {
-            if (currentEventContext === handlerContext) {
-              subscribable.publish(newEvent);
-              break;
-            }
-            currentEventContext = currentEventContext.prev;
-          }
-        },
-      );
-      return new WorkflowStream<WorkflowEventData<any>>(subscribable, null);
-    },
-    get signal() {
-      return handlerContext.abortController.signal;
-    },
-    sendEvent: (...events) => {
-      events.forEach((event) => {
-        eventContextWeakMap.set(event, handlerContext);
-        handlerContext.outputs.push(event);
-        queue.push(event);
-        rootWorkflowContext.__internal__call_send_event.publish(
-          event,
-          handlerContext,
-        );
-        queueUpdateCallback(handlerContext);
-      });
-    },
-    __internal__call_context: createSubscribable(),
-    __internal__call_send_event: createSubscribable(),
-  });
+  ): WorkflowContext => {
+    let lazyLoadStream: WorkflowStream | null = null;
+    return {
+      get stream() {
+        if (!lazyLoadStream) {
+          const subscribable = createSubscribable<
+            [event: WorkflowEventData<any>],
+            void
+          >();
+          rootWorkflowContext.__internal__call_send_event.subscribe(
+            (newEvent: WorkflowEventData<any>) => {
+              let currentEventContext = eventContextWeakMap.get(newEvent);
+              while (currentEventContext) {
+                if (currentEventContext === handlerContext) {
+                  subscribable.publish(newEvent);
+                  break;
+                }
+                currentEventContext = currentEventContext.prev;
+              }
+            },
+          );
+          lazyLoadStream = new WorkflowStream(subscribable, null);
+        }
+        return lazyLoadStream;
+      },
+      get signal() {
+        return handlerContext.abortController.signal;
+      },
+      sendEvent: (...events) => {
+        events.forEach((event) => {
+          eventContextWeakMap.set(event, handlerContext);
+          handlerContext.outputs.push(event);
+          queue.push(event);
+          rootWorkflowContext.__internal__call_send_event.publish(
+            event,
+            handlerContext,
+          );
+          queueUpdateCallback(handlerContext);
+        });
+      },
+      __internal__call_context: createSubscribable(),
+      __internal__call_send_event: createSubscribable(),
+    };
+  };
 
   let rootAbortController = new AbortController();
   const handlerRootContext: HandlerContext = {

packages/core/src/core/stream.ts

@@ -76,9 +76,9 @@ export class WorkflowStream<R = any>
   #stream: ReadableStream<R>;
   #subscribable: Subscribable<[data: R], void>;
 
-  on(
-    event: WorkflowEvent<any>,
-    handler: (event: WorkflowEventData<any>) => void,
+  on<T>(
+    event: WorkflowEvent<T>,
+    handler: (event: WorkflowEventData<T>) => void,
   ): () => void {
     return this.#subscribable.subscribe((ev) => {
       if (event.include(ev)) {

packages/core/src/core/utils.ts

@@ -29,7 +29,7 @@ export function flattenEvents(
 
 export type Subscribable<Args extends any[], R> = {
   subscribe: (callback: (...args: Args) => R) => () => void;
-  publish: (...args: Args) => void;
+  publish: (...args: Args) => unknown[];
 };
 
 const __internal__subscribesSourcemap = new WeakMap<
@@ -49,24 +49,30 @@ export function getSubscribers<Args extends any[], R>(
 /**
  * @internal
  */
-export function createSubscribable<Args extends any[], R>(): Subscribable<
-  Args,
-  R
-> {
-  const subscribers = new Set<(...args: Args) => R>();
+export function createSubscribable<
+  FnOrArgs extends ((...args: any[]) => any) | any[],
+  R = unknown,
+>(): FnOrArgs extends (...args: any[]) => any
+  ? Subscribable<Parameters<FnOrArgs>, ReturnType<FnOrArgs>>
+  : FnOrArgs extends any[]
+    ? Subscribable<FnOrArgs, R>
+    : never {
+  const subscribers = new Set<(...args: any) => any>();
   const obj = {
-    subscribe: (callback: (...args: Args) => R) => {
+    subscribe: (callback: (...args: any) => any) => {
       subscribers.add(callback);
       return () => {
         subscribers.delete(callback);
       };
     },
-    publish: (...args: Args) => {
+    publish: (...args: any) => {
+      const results: unknown[] = [];
       for (const callback of subscribers) {
-        callback(...args);
+        results.push(callback(...args));
       }
+      return results;
     },
   };
   __internal__subscribesSourcemap.set(obj, subscribers);
-  return obj;
+  return obj as any;
 }