@workglow/task-graph 0.0.90 → 0.0.91

package/dist/bun.js

@@ -7,6 +7,7 @@ var TaskStatus = {
   PENDING: "PENDING",
   DISABLED: "DISABLED",
   PROCESSING: "PROCESSING",
+  STREAMING: "STREAMING",
   COMPLETED: "COMPLETED",
   ABORTING: "ABORTING",
   FAILED: "FAILED"
@@ -36,10 +37,62 @@ class Dataflow {
   value = undefined;
   status = TaskStatus.PENDING;
   error;
+  stream = undefined;
+  setStream(stream) {
+    this.stream = stream;
+  }
+  getStream() {
+    return this.stream;
+  }
+  async awaitStreamValue() {
+    if (!this.stream)
+      return;
+    const reader = this.stream.getReader();
+    let accumulatedText = "";
+    let lastSnapshotData = undefined;
+    let finishData = undefined;
+    let hasTextDelta = false;
+    try {
+      while (true) {
+        const { done, value: event } = await reader.read();
+        if (done)
+          break;
+        switch (event.type) {
+          case "text-delta":
+            hasTextDelta = true;
+            accumulatedText += event.textDelta;
+            break;
+          case "snapshot":
+            lastSnapshotData = event.data;
+            break;
+          case "finish":
+            finishData = event.data;
+            break;
+          case "error":
+            break;
+        }
+      }
+    } finally {
+      reader.releaseLock();
+      this.stream = undefined;
+    }
+    if (lastSnapshotData !== undefined) {
+      this.setPortData(lastSnapshotData);
+    } else if (finishData && Object.keys(finishData).length > 0) {
+      this.setPortData(finishData);
+    } else if (hasTextDelta) {
+      if (this.sourceTaskPortId === DATAFLOW_ALL_PORTS) {
+        this.value = { text: accumulatedText };
+      } else {
+        this.value = accumulatedText;
+      }
+    }
+  }
   reset() {
     this.status = TaskStatus.PENDING;
     this.error = undefined;
     this.value = undefined;
+    this.stream = undefined;
     this.emit("reset");
     this.emit("status", this.status);
   }
@@ -51,6 +104,9 @@ class Dataflow {
       case TaskStatus.PROCESSING:
         this.emit("start");
         break;
+      case TaskStatus.STREAMING:
+        this.emit("streaming");
+        break;
       case TaskStatus.COMPLETED:
         this.emit("complete");
         break;
@@ -398,6 +454,49 @@ async function resolveSchemaInputs(input, schema, config) {
   return resolved;
 }
 
+// src/task/StreamTypes.ts
+function getPortStreamMode(schema, portId) {
+  if (typeof schema === "boolean")
+    return "none";
+  const prop = schema.properties?.[portId];
+  if (!prop || typeof prop === "boolean")
+    return "none";
+  const xStream = prop["x-stream"];
+  if (xStream === "append" || xStream === "replace")
+    return xStream;
+  return "none";
+}
+function getOutputStreamMode(outputSchema) {
+  if (typeof outputSchema === "boolean")
+    return "none";
+  const props = outputSchema.properties;
+  if (!props)
+    return "none";
+  let found = "none";
+  for (const prop of Object.values(props)) {
+    if (!prop || typeof prop === "boolean")
+      continue;
+    const xStream = prop["x-stream"];
+    if (xStream === "append")
+      return "append";
+    if (xStream === "replace")
+      found = "replace";
+  }
+  return found;
+}
+function isTaskStreamable(task) {
+  if (typeof task.executeStream !== "function")
+    return false;
+  return getOutputStreamMode(task.outputSchema()) !== "none";
+}
+function edgeNeedsAccumulation(sourceSchema, sourcePort, targetSchema, targetPort) {
+  const sourceMode = getPortStreamMode(sourceSchema, sourcePort);
+  if (sourceMode === "none")
+    return false;
+  const targetMode = getPortStreamMode(targetSchema, targetPort);
+  return sourceMode !== targetMode;
+}
+
 // src/task/TaskRunner.ts
 class TaskRunner {
   running = false;
@@ -427,14 +526,27 @@ class TaskRunner {
       }
       const inputs = this.task.runInputData;
       let outputs;
+      const isStreamable = isTaskStreamable(this.task);
       if (this.task.cacheable) {
         outputs = await this.outputCache?.getOutput(this.task.type, inputs);
         if (outputs) {
-          this.task.runOutputData = await this.executeTaskReactive(inputs, outputs);
+          if (isStreamable) {
+            this.task.runOutputData = outputs;
+            this.task.emit("stream_start");
+            this.task.emit("stream_chunk", { type: "finish", data: outputs });
+            this.task.emit("stream_end", outputs);
+            this.task.runOutputData = await this.executeTaskReactive(inputs, outputs);
+          } else {
+            this.task.runOutputData = await this.executeTaskReactive(inputs, outputs);
+          }
         }
       }
       if (!outputs) {
-        outputs = await this.executeTask(inputs);
+        if (isStreamable) {
+          outputs = await this.executeStreamingTask(inputs);
+        } else {
+          outputs = await this.executeTask(inputs);
+        }
         if (this.task.cacheable && outputs !== undefined) {
           await this.outputCache?.saveOutput(this.task.type, inputs, outputs);
         }
@@ -493,6 +605,64 @@ class TaskRunner {
     const reactiveResult = await this.task.executeReactive(input, output, { own: this.own });
     return Object.assign({}, output, reactiveResult ?? {});
   }
+  async executeStreamingTask(input) {
+    const streamMode = getOutputStreamMode(this.task.outputSchema());
+    let accumulated = "";
+    let chunkCount = 0;
+    let finalOutput;
+    this.task.emit("stream_start");
+    const stream = this.task.executeStream(input, {
+      signal: this.abortController.signal,
+      updateProgress: this.handleProgress.bind(this),
+      own: this.own,
+      registry: this.registry
+    });
+    for await (const event of stream) {
+      chunkCount++;
+      if (chunkCount === 1) {
+        this.task.status = TaskStatus.STREAMING;
+        this.task.emit("status", this.task.status);
+      }
+      if (event.type === "snapshot") {
+        this.task.runOutputData = event.data;
+      }
+      this.task.emit("stream_chunk", event);
+      switch (event.type) {
+        case "text-delta": {
+          accumulated += event.textDelta;
+          const progress = Math.min(99, Math.round(100 * (1 - Math.exp(-0.02 * chunkCount))));
+          await this.handleProgress(progress);
+          break;
+        }
+        case "snapshot": {
+          const progress = Math.min(99, Math.round(100 * (1 - Math.exp(-0.02 * chunkCount))));
+          await this.handleProgress(progress);
+          break;
+        }
+        case "finish": {
+          if (streamMode === "append") {
+            const text = accumulated.length > 0 ? accumulated : event.data?.text ?? "";
+            finalOutput = { ...event.data || {}, text };
+          } else if (streamMode === "replace") {
+            finalOutput = event.data;
+          }
+          break;
+        }
+        case "error": {
+          throw event.error;
+        }
+      }
+    }
+    if (this.abortController?.signal.aborted) {
+      throw new TaskAbortedError("Task aborted during streaming");
+    }
+    if (finalOutput !== undefined) {
+      this.task.runOutputData = finalOutput;
+    }
+    this.task.emit("stream_end", this.task.runOutputData);
+    const reactiveResult = await this.executeTaskReactive(input, this.task.runOutputData || {});
+    return reactiveResult;
+  }
   async handleStart(config = {}) {
     if (this.task.status === TaskStatus.PROCESSING)
       return;
@@ -1186,6 +1356,7 @@ class TopologicalScheduler {
     }
   }
   onTaskCompleted(taskId) {}
+  onTaskStreaming(taskId) {}
   reset() {
     this.sortedNodes = this.dag.topologicallySortedNodes();
     this.currentIndex = 0;
@@ -1195,11 +1366,13 @@ class TopologicalScheduler {
 class DependencyBasedScheduler {
   dag;
   completedTasks;
+  streamingTasks;
   pendingTasks;
   nextResolver = null;
   constructor(dag) {
     this.dag = dag;
     this.completedTasks = new Set;
+    this.streamingTasks = new Set;
     this.pendingTasks = new Set;
     this.reset();
   }
@@ -1214,8 +1387,23 @@ class DependencyBasedScheduler {
         return false;
       }
     }
-    const dependencies = sourceDataflows.filter((df) => df.status !== TaskStatus.DISABLED).map((dataflow) => dataflow.sourceTaskId);
-    return dependencies.every((dep) => this.completedTasks.has(dep));
+    const activeDataflows = sourceDataflows.filter((df) => df.status !== TaskStatus.DISABLED);
+    return activeDataflows.every((df) => {
+      const depId = df.sourceTaskId;
+      if (this.completedTasks.has(depId))
+        return true;
+      if (this.streamingTasks.has(depId)) {
+        const sourceTask = this.dag.getTask(depId);
+        if (sourceTask) {
+          const sourceMode = getPortStreamMode(sourceTask.outputSchema(), df.sourceTaskPortId);
+          const targetMode = getPortStreamMode(task.inputSchema(), df.targetTaskPortId);
+          if (sourceMode !== "none" && sourceMode === targetMode) {
+            return true;
+          }
+        }
+      }
+      return false;
+    });
   }
   async waitForNextTask() {
     if (this.pendingTasks.size === 0)
@@ -1270,8 +1458,26 @@ class DependencyBasedScheduler {
       }
     }
   }
+  onTaskStreaming(taskId) {
+    this.streamingTasks.add(taskId);
+    for (const task of Array.from(this.pendingTasks)) {
+      if (task.status === TaskStatus.DISABLED) {
+        this.pendingTasks.delete(task);
+      }
+    }
+    if (this.nextResolver) {
+      const readyTask = Array.from(this.pendingTasks).find((task) => this.isTaskReady(task));
+      if (readyTask) {
+        this.pendingTasks.delete(readyTask);
+        const resolver = this.nextResolver;
+        this.nextResolver = null;
+        resolver(readyTask);
+      }
+    }
+  }
   reset() {
     this.completedTasks.clear();
+    this.streamingTasks.clear();
     this.pendingTasks = new Set(this.dag.topologicallySortedNodes());
     this.nextResolver = null;
   }
@@ -1508,7 +1714,12 @@ class TaskGraphRunner {
     }
   }
   async runTask(task, input) {
+    const isStreamable = isTaskStreamable(task);
+    await this.awaitStreamInputs(task);
     this.copyInputFromEdgesToNode(task);
+    if (isStreamable) {
+      return this.runStreamingTask(task, input);
+    }
     const results = await task.runner.run(input, {
       outputCache: this.outputCache,
       updateProgress: async (task2, progress, message, ...args) => await this.handleProgress(task2, progress, message, ...args),
@@ -1521,6 +1732,93 @@ class TaskGraphRunner {
       data: results
     };
   }
+  async awaitStreamInputs(task) {
+    const dataflows = this.graph.getSourceDataflows(task.config.id);
+    const streamPromises = dataflows.filter((df) => df.stream !== undefined).map((df) => df.awaitStreamValue());
+    if (streamPromises.length > 0) {
+      await Promise.all(streamPromises);
+    }
+  }
+  async runStreamingTask(task, input) {
+    const streamMode = getOutputStreamMode(task.outputSchema());
+    let streamingNotified = false;
+    const onStatus = (status) => {
+      if (status === TaskStatus.STREAMING && !streamingNotified) {
+        streamingNotified = true;
+        this.pushStatusFromNodeToEdges(this.graph, task, TaskStatus.STREAMING);
+        this.pushStreamToEdges(task, streamMode);
+        this.processScheduler.onTaskStreaming(task.config.id);
+      }
+    };
+    const onStreamStart = () => {
+      this.graph.emit("task_stream_start", task.config.id);
+    };
+    const onStreamChunk = (event) => {
+      this.graph.emit("task_stream_chunk", task.config.id, event);
+    };
+    const onStreamEnd = (output) => {
+      this.graph.emit("task_stream_end", task.config.id, output);
+    };
+    task.on("status", onStatus);
+    task.on("stream_start", onStreamStart);
+    task.on("stream_chunk", onStreamChunk);
+    task.on("stream_end", onStreamEnd);
+    try {
+      const results = await task.runner.run(input, {
+        outputCache: this.outputCache,
+        updateProgress: async (task2, progress, message, ...args) => await this.handleProgress(task2, progress, message, ...args),
+        registry: this.registry
+      });
+      await this.pushOutputFromNodeToEdges(task, results);
+      return {
+        id: task.config.id,
+        type: task.constructor.runtype || task.constructor.type,
+        data: results
+      };
+    } finally {
+      task.off("status", onStatus);
+      task.off("stream_start", onStreamStart);
+      task.off("stream_chunk", onStreamChunk);
+      task.off("stream_end", onStreamEnd);
+    }
+  }
+  pushStreamToEdges(task, streamMode) {
+    const targetDataflows = this.graph.getTargetDataflows(task.config.id);
+    if (targetDataflows.length === 0)
+      return;
+    const stream = new ReadableStream({
+      start: (controller) => {
+        const onChunk = (event) => {
+          try {
+            controller.enqueue(event);
+          } catch {}
+        };
+        const onEnd = () => {
+          try {
+            controller.close();
+          } catch {}
+          task.off("stream_chunk", onChunk);
+          task.off("stream_end", onEnd);
+        };
+        task.on("stream_chunk", onChunk);
+        task.on("stream_end", onEnd);
+      }
+    });
+    if (targetDataflows.length === 1) {
+      targetDataflows[0].setStream(stream);
+    } else {
+      let currentStream = stream;
+      for (let i = 0;i < targetDataflows.length; i++) {
+        if (i === targetDataflows.length - 1) {
+          targetDataflows[i].setStream(currentStream);
+        } else {
+          const [s1, s2] = currentStream.tee();
+          targetDataflows[i].setStream(s1);
+          currentStream = s2;
+        }
+      }
+    }
+  }
   resetTask(graph, task, runId) {
     task.status = TaskStatus.PENDING;
     task.resetInputData();
@@ -1604,7 +1902,7 @@ class TaskGraphRunner {
   }
   async handleError(error) {
     await Promise.allSettled(this.graph.getTasks().map(async (task) => {
-      if (task.status === TaskStatus.PROCESSING) {
+      if (task.status === TaskStatus.PROCESSING || task.status === TaskStatus.STREAMING) {
         task.abort();
       }
     }));
@@ -1616,7 +1914,7 @@ class TaskGraphRunner {
   }
   async handleAbort() {
     this.graph.getTasks().map(async (task) => {
-      if (task.status === TaskStatus.PROCESSING) {
+      if (task.status === TaskStatus.PROCESSING || task.status === TaskStatus.STREAMING) {
         task.abort();
       }
     });
@@ -1893,6 +2191,54 @@ function CreateEndLoopWorkflow(methodName) {
     return this.finalizeAndReturn();
   };
 }
+var TYPED_ARRAY_FORMAT_PREFIX = "TypedArray";
+function schemaHasTypedArrayFormat(schema) {
+  if (typeof schema === "boolean")
+    return false;
+  if (!schema || typeof schema !== "object" || Array.isArray(schema))
+    return false;
+  const s = schema;
+  if (typeof s.format === "string" && s.format.startsWith(TYPED_ARRAY_FORMAT_PREFIX)) {
+    return true;
+  }
+  const checkUnion = (schemas) => {
+    if (!Array.isArray(schemas))
+      return false;
+    return schemas.some((sub) => schemaHasTypedArrayFormat(sub));
+  };
+  if (checkUnion(s.oneOf) || checkUnion(s.anyOf))
+    return true;
+  const items = s.items;
+  if (items && typeof items === "object" && !Array.isArray(items)) {
+    if (schemaHasTypedArrayFormat(items))
+      return true;
+  }
+  return false;
+}
+function hasVectorOutput(task) {
+  const outputSchema = task.outputSchema();
+  if (typeof outputSchema === "boolean" || !outputSchema?.properties)
+    return false;
+  return Object.values(outputSchema.properties).some((prop) => schemaHasTypedArrayFormat(prop));
+}
+function hasVectorLikeInput(input) {
+  if (!input || typeof input !== "object")
+    return false;
+  const v = input.vectors;
+  return Array.isArray(v) && v.length > 0 && typeof v[0] === "object" && v[0] !== null && ArrayBuffer.isView(v[0]);
+}
+function CreateAdaptiveWorkflow(scalarClass, vectorClass) {
+  const scalarHelper = Workflow.createWorkflow(scalarClass);
+  const vectorHelper = Workflow.createWorkflow(vectorClass);
+  return function(input = {}, config = {}) {
+    const parent = getLastTask(this);
+    const useVector = parent !== undefined && hasVectorOutput(parent) || hasVectorLikeInput(input);
+    if (useVector) {
+      return vectorHelper.call(this, input, config);
+    }
+    return scalarHelper.call(this, input, config);
+  };
+}
 
 class WorkflowTask extends GraphAsTask {
   static type = "Workflow";
@@ -2013,6 +2359,11 @@ class Workflow {
     }
     this.events.emit("start");
     this._abortController = new AbortController;
+    const unsubStreaming = this.graph.subscribeToTaskStreaming({
+      onStreamStart: (taskId) => this.events.emit("stream_start", taskId),
+      onStreamChunk: (taskId, event) => this.events.emit("stream_chunk", taskId, event),
+      onStreamEnd: (taskId, output) => this.events.emit("stream_end", taskId, output)
+    });
     try {
       const output = await this.graph.run(input, {
         parentSignal: this._abortController.signal,
@@ -2025,6 +2376,7 @@ class Workflow {
       this.events.emit("error", String(error));
       throw error;
     } finally {
+      unsubStreaming();
       this._abortController = undefined;
     }
   }
@@ -2304,13 +2656,26 @@ class Workflow {
       if (typeof fromSchema === "boolean" || typeof toSchema === "boolean") {
         return;
       }
-      for (const [fromOutputPortId, fromPortOutputSchema] of Object.entries(fromSchema.properties || {})) {
-        for (const [toInputPortId, toPortInputSchema] of Object.entries(toSchema.properties || {})) {
-          if (!matches.has(toInputPortId) && comparator([fromOutputPortId, fromPortOutputSchema], [toInputPortId, toPortInputSchema])) {
-            matches.set(toInputPortId, fromOutputPortId);
-            graph.addDataflow(new Dataflow(fromTaskId, fromOutputPortId, toTaskId, toInputPortId));
+      for (const [toInputPortId, toPortInputSchema] of Object.entries(toSchema.properties || {})) {
+        if (matches.has(toInputPortId))
+          continue;
+        const candidates = [];
+        for (const [fromOutputPortId, fromPortOutputSchema] of Object.entries(fromSchema.properties || {})) {
+          if (comparator([fromOutputPortId, fromPortOutputSchema], [toInputPortId, toPortInputSchema])) {
+            candidates.push(fromOutputPortId);
           }
         }
+        if (candidates.length === 0)
+          continue;
+        let winner = candidates[0];
+        if (candidates.length > 1) {
+          const targetStreamMode = getPortStreamMode(toSchema, toInputPortId);
+          const streamMatch = candidates.find((portId) => getPortStreamMode(fromSchema, portId) === targetStreamMode);
+          if (streamMatch)
+            winner = streamMatch;
+        }
+        matches.set(toInputPortId, winner);
+        graph.addDataflow(new Dataflow(fromTaskId, winner, toTaskId, toInputPortId));
       }
     };
     makeMatch(sourceSchema, targetSchema, sourceTask.config.id, targetTask.config.id, ([fromOutputPortId, fromPortOutputSchema], [toInputPortId, toPortInputSchema]) => {
@@ -2758,6 +3123,24 @@ class TaskGraph {
       unsubscribes.forEach((unsub) => unsub());
     };
   }
+  subscribeToTaskStreaming(callbacks) {
+    const unsubscribes = [];
+    if (callbacks.onStreamStart) {
+      const unsub = this.subscribe("task_stream_start", callbacks.onStreamStart);
+      unsubscribes.push(unsub);
+    }
+    if (callbacks.onStreamChunk) {
+      const unsub = this.subscribe("task_stream_chunk", callbacks.onStreamChunk);
+      unsubscribes.push(unsub);
+    }
+    if (callbacks.onStreamEnd) {
+      const unsub = this.subscribe("task_stream_end", callbacks.onStreamEnd);
+      unsubscribes.push(unsub);
+    }
+    return () => {
+      unsubscribes.forEach((unsub) => unsub());
+    };
+  }
   on(name, fn) {
     const dagEvent = EventTaskGraphToDagMapping[name];
     if (dagEvent) {
@@ -4515,10 +4898,15 @@ export {
   pipe,
   parallel,
   mergeChainedOutputToInput,
+  isTaskStreamable,
   isStrictArraySchema,
   isIterationProperty,
   isFlexibleSchema,
+  hasVectorOutput,
+  hasVectorLikeInput,
   getTaskQueueRegistry,
+  getPortStreamMode,
+  getOutputStreamMode,
   getNestedValue,
   getLastTask,
   getJobQueueFactory,
@@ -4530,6 +4918,7 @@ export {
   extractBaseSchema,
   evaluateCondition,
   ensureTask,
+  edgeNeedsAccumulation,
   createTaskFromGraphJSON,
   createTaskFromDependencyJSON,
   createJobQueueFactoryWithOptions,
@@ -4588,7 +4977,8 @@ export {
   CreateWorkflow,
   CreateLoopWorkflow,
   CreateEndLoopWorkflow,
+  CreateAdaptiveWorkflow,
   ConditionalTask
 };
 
-//# debugId=4D22DE67A27A853E64756E2164756E21
+//# debugId=65C6DA6762951D4864756E2164756E21