[feat] Add streaming support for AI tasks

- Introduced a new base class, StreamingAiTask, to facilitate streaming output for AI tasks, enhancing the execution model for tasks like text generation, rewriting, and summarization.
- Implemented streaming functions for existing tasks such as TextGenerationTask, TextRewriterTask, and TextSummaryTask, allowing for real-time output delivery.
- Updated AI providers (Anthropic, Google Gemini, Ollama, and Hugging Face) to support streaming capabilities, improving responsiveness and user experience.
- Enhanced the AiJob class with a new executeStream method to handle streaming execution, providing a fallback to non-streaming execution when necessary.
- Refactored task registration to include streaming functions, ensuring seamless integration with the existing architecture.

Author: sroussey

packages/ai-provider/src/anthropic/AnthropicProvider.ts

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { AiProvider, type AiProviderRunFn } from "@workglow/ai";
+import { AiProvider, type AiProviderRunFn, type AiProviderStreamFn } from "@workglow/ai";
 import { ANTHROPIC } from "./common/Anthropic_Constants";
 import type { AnthropicModelConfig } from "./common/Anthropic_ModelSchema";
 
@@ -39,7 +39,10 @@ export class AnthropicProvider extends AiProvider<AnthropicModelConfig> {
 
   readonly taskTypes = ["TextGenerationTask", "TextRewriterTask", "TextSummaryTask"] as const;
 
-  constructor(tasks?: Record<string, AiProviderRunFn<any, any, AnthropicModelConfig>>) {
-    super(tasks);
+  constructor(
+    tasks?: Record<string, AiProviderRunFn<any, any, AnthropicModelConfig>>,
+    streamTasks?: Record<string, AiProviderStreamFn<any, any, AnthropicModelConfig>>
+  ) {
+    super(tasks, streamTasks);
   }
 }

packages/ai-provider/src/anthropic/Anthropic_Worker.ts

@@ -5,12 +5,14 @@
  */
 
 import { globalServiceRegistry, parentPort, WORKER_SERVER } from "@workglow/util";
-import { ANTHROPIC_TASKS } from "./common/Anthropic_JobRunFns";
 import { AnthropicProvider } from "./AnthropicProvider";
+import { ANTHROPIC_STREAM_TASKS, ANTHROPIC_TASKS } from "./common/Anthropic_JobRunFns";
 
 export function ANTHROPIC_WORKER_JOBRUN_REGISTER() {
   const workerServer = globalServiceRegistry.get(WORKER_SERVER);
-  new AnthropicProvider(ANTHROPIC_TASKS).registerOnWorkerServer(workerServer);
+  new AnthropicProvider(ANTHROPIC_TASKS, ANTHROPIC_STREAM_TASKS).registerOnWorkerServer(
+    workerServer
+  );
   parentPort.postMessage({ type: "ready" });
   console.log("ANTHROPIC_WORKER_JOBRUN registered");
 }

packages/ai-provider/src/anthropic/common/Anthropic_JobRunFns.ts

@@ -7,13 +7,15 @@
 import Anthropic from "@anthropic-ai/sdk";
 import type {
   AiProviderRunFn,
+  AiProviderStreamFn,
   TextGenerationTaskInput,
   TextGenerationTaskOutput,
   TextRewriterTaskInput,
   TextRewriterTaskOutput,
   TextSummaryTaskInput,
   TextSummaryTaskOutput,
 } from "@workglow/ai";
+import type { StreamEvent } from "@workglow/task-graph";
 import type { AnthropicModelConfig } from "./Anthropic_ModelSchema";
 
 function getClient(model: AnthropicModelConfig | undefined): Anthropic {
@@ -123,8 +125,104 @@ export const Anthropic_TextSummary: AiProviderRunFn<
   return { text };
 };
 
+// ========================================================================
+// Streaming implementations (append mode)
+// ========================================================================
+
+export const Anthropic_TextGeneration_Stream: AiProviderStreamFn<
+  TextGenerationTaskInput,
+  TextGenerationTaskOutput,
+  AnthropicModelConfig
+> = async function* (input, model, signal): AsyncIterable<StreamEvent<TextGenerationTaskOutput>> {
+  const client = getClient(model);
+  const modelName = getModelName(model);
+
+  const stream = client.messages.stream(
+    {
+      model: modelName,
+      messages: [{ role: "user", content: input.prompt }],
+      max_tokens: getMaxTokens(input, model),
+      temperature: input.temperature,
+      top_p: input.topP,
+    },
+    { signal }
+  );
+
+  for await (const event of stream) {
+    if (event.type === "content_block_delta" && event.delta.type === "text_delta") {
+      yield { type: "text-delta", textDelta: event.delta.text };
+    }
+  }
+  yield { type: "finish", data: {} as TextGenerationTaskOutput };
+};
+
+export const Anthropic_TextRewriter_Stream: AiProviderStreamFn<
+  TextRewriterTaskInput,
+  TextRewriterTaskOutput,
+  AnthropicModelConfig
+> = async function* (input, model, signal): AsyncIterable<StreamEvent<TextRewriterTaskOutput>> {
+  const client = getClient(model);
+  const modelName = getModelName(model);
+
+  const stream = client.messages.stream(
+    {
+      model: modelName,
+      system: input.prompt,
+      messages: [{ role: "user", content: input.text }],
+      max_tokens: getMaxTokens({}, model),
+    },
+    { signal }
+  );
+
+  for await (const event of stream) {
+    if (event.type === "content_block_delta" && event.delta.type === "text_delta") {
+      yield { type: "text-delta", textDelta: event.delta.text };
+    }
+  }
+  yield { type: "finish", data: {} as TextRewriterTaskOutput };
+};
+
+export const Anthropic_TextSummary_Stream: AiProviderStreamFn<
+  TextSummaryTaskInput,
+  TextSummaryTaskOutput,
+  AnthropicModelConfig
+> = async function* (input, model, signal): AsyncIterable<StreamEvent<TextSummaryTaskOutput>> {
+  const client = getClient(model);
+  const modelName = getModelName(model);
+
+  const stream = client.messages.stream(
+    {
+      model: modelName,
+      system: "Summarize the following text concisely.",
+      messages: [{ role: "user", content: input.text }],
+      max_tokens: getMaxTokens({}, model),
+    },
+    { signal }
+  );
+
+  for await (const event of stream) {
+    if (event.type === "content_block_delta" && event.delta.type === "text_delta") {
+      yield { type: "text-delta", textDelta: event.delta.text };
+    }
+  }
+  yield { type: "finish", data: {} as TextSummaryTaskOutput };
+};
+
+// ========================================================================
+// Task registries
+// ========================================================================
+
 export const ANTHROPIC_TASKS: Record<string, AiProviderRunFn<any, any, AnthropicModelConfig>> = {
   TextGenerationTask: Anthropic_TextGeneration,
   TextRewriterTask: Anthropic_TextRewriter,
   TextSummaryTask: Anthropic_TextSummary,
 };
+
+export const ANTHROPIC_STREAM_TASKS: Record<
+  string,
+  AiProviderStreamFn<any, any, AnthropicModelConfig>
+> = {
+  TextGenerationTask: Anthropic_TextGeneration_Stream,
+  TextRewriterTask: Anthropic_TextRewriter_Stream,
+  TextSummaryTask: Anthropic_TextSummary_Stream,
+};

packages/ai-provider/src/google-gemini/Gemini_Worker.ts

@@ -5,12 +5,12 @@
  */
 
 import { globalServiceRegistry, parentPort, WORKER_SERVER } from "@workglow/util";
-import { GEMINI_TASKS } from "./common/Gemini_JobRunFns";
+import { GEMINI_STREAM_TASKS, GEMINI_TASKS } from "./common/Gemini_JobRunFns";
 import { GoogleGeminiProvider } from "./GoogleGeminiProvider";
 
 export function GEMINI_WORKER_JOBRUN_REGISTER() {
   const workerServer = globalServiceRegistry.get(WORKER_SERVER);
-  new GoogleGeminiProvider(GEMINI_TASKS).registerOnWorkerServer(workerServer);
+  new GoogleGeminiProvider(GEMINI_TASKS, GEMINI_STREAM_TASKS).registerOnWorkerServer(workerServer);
   parentPort.postMessage({ type: "ready" });
   console.log("GEMINI_WORKER_JOBRUN registered");
 }

packages/ai-provider/src/google-gemini/GoogleGeminiProvider.ts

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { AiProvider, type AiProviderRunFn } from "@workglow/ai";
+import { AiProvider, type AiProviderRunFn, type AiProviderStreamFn } from "@workglow/ai";
 import { GOOGLE_GEMINI } from "./common/Gemini_Constants";
 import type { GeminiModelConfig } from "./common/Gemini_ModelSchema";
 
@@ -41,7 +41,10 @@ export class GoogleGeminiProvider extends AiProvider<GeminiModelConfig> {
     "TextSummaryTask",
   ] as const;
 
-  constructor(tasks?: Record<string, AiProviderRunFn<any, any, GeminiModelConfig>>) {
-    super(tasks);
+  constructor(
+    tasks?: Record<string, AiProviderRunFn<any, any, GeminiModelConfig>>,
+    streamTasks?: Record<string, AiProviderStreamFn<any, any, GeminiModelConfig>>
+  ) {
+    super(tasks, streamTasks);
   }
 }

packages/ai-provider/src/google-gemini/common/Gemini_JobRunFns.ts

@@ -7,6 +7,7 @@
 import { GoogleGenerativeAI, type TaskType } from "@google/generative-ai";
 import type {
   AiProviderRunFn,
+  AiProviderStreamFn,
   TextEmbeddingTaskInput,
   TextEmbeddingTaskOutput,
   TextGenerationTaskInput,
@@ -16,6 +17,7 @@ import type {
   TextSummaryTaskInput,
   TextSummaryTaskOutput,
 } from "@workglow/ai";
+import type { StreamEvent } from "@workglow/task-graph";
 import type { GeminiModelConfig } from "./Gemini_ModelSchema";
 
 function getApiKey(model: GeminiModelConfig | undefined): string {
@@ -143,9 +145,102 @@ export const Gemini_TextSummary: AiProviderRunFn<
   return { text };
 };
 
+// ========================================================================
+// Streaming implementations (append mode)
+// ========================================================================
+
+export const Gemini_TextGeneration_Stream: AiProviderStreamFn<
+  TextGenerationTaskInput,
+  TextGenerationTaskOutput,
+  GeminiModelConfig
+> = async function* (input, model, signal): AsyncIterable<StreamEvent<TextGenerationTaskOutput>> {
+  const genAI = new GoogleGenerativeAI(getApiKey(model));
+  const genModel = genAI.getGenerativeModel({
+    model: getModelName(model),
+    generationConfig: {
+      maxOutputTokens: input.maxTokens,
+      temperature: input.temperature,
+      topP: input.topP,
+    },
+  });
+
+  const result = await genModel.generateContentStream({
+    contents: [{ role: "user", parts: [{ text: input.prompt }] }],
+  });
+
+  for await (const chunk of result.stream) {
+    const text = chunk.text();
+    if (text) {
+      yield { type: "text-delta", textDelta: text };
+    }
+  }
+  yield { type: "finish", data: {} as TextGenerationTaskOutput };
+};
+
+export const Gemini_TextRewriter_Stream: AiProviderStreamFn<
+  TextRewriterTaskInput,
+  TextRewriterTaskOutput,
+  GeminiModelConfig
+> = async function* (input, model, signal): AsyncIterable<StreamEvent<TextRewriterTaskOutput>> {
+  const genAI = new GoogleGenerativeAI(getApiKey(model));
+  const genModel = genAI.getGenerativeModel({
+    model: getModelName(model),
+    systemInstruction: input.prompt,
+  });
+
+  const result = await genModel.generateContentStream({
+    contents: [{ role: "user", parts: [{ text: input.text }] }],
+  });
+
+  for await (const chunk of result.stream) {
+    const text = chunk.text();
+    if (text) {
+      yield { type: "text-delta", textDelta: text };
+    }
+  }
+  yield { type: "finish", data: {} as TextRewriterTaskOutput };
+};
+
+export const Gemini_TextSummary_Stream: AiProviderStreamFn<
+  TextSummaryTaskInput,
+  TextSummaryTaskOutput,
+  GeminiModelConfig
+> = async function* (input, model, signal): AsyncIterable<StreamEvent<TextSummaryTaskOutput>> {
+  const genAI = new GoogleGenerativeAI(getApiKey(model));
+  const genModel = genAI.getGenerativeModel({
+    model: getModelName(model),
+    systemInstruction: "Summarize the following text concisely.",
+  });
+
+  const result = await genModel.generateContentStream({
+    contents: [{ role: "user", parts: [{ text: input.text }] }],
+  });
+
+  for await (const chunk of result.stream) {
+    const text = chunk.text();
+    if (text) {
+      yield { type: "text-delta", textDelta: text };
+    }
+  }
+  yield { type: "finish", data: {} as TextSummaryTaskOutput };
+};
+
+// ========================================================================
+// Task registries
+// ========================================================================
+
 export const GEMINI_TASKS: Record<string, AiProviderRunFn<any, any, GeminiModelConfig>> = {
   TextGenerationTask: Gemini_TextGeneration,
   TextEmbeddingTask: Gemini_TextEmbedding,
   TextRewriterTask: Gemini_TextRewriter,
   TextSummaryTask: Gemini_TextSummary,
 };
+
+export const GEMINI_STREAM_TASKS: Record<
+  string,
+  AiProviderStreamFn<any, any, GeminiModelConfig>
+> = {
+  TextGenerationTask: Gemini_TextGeneration_Stream,
+  TextRewriterTask: Gemini_TextRewriter_Stream,
+  TextSummaryTask: Gemini_TextSummary_Stream,
+};

packages/ai-provider/src/hf-transformers/HFT_Worker.ts

@@ -5,12 +5,14 @@
  */
 
 import { globalServiceRegistry, parentPort, WORKER_SERVER } from "@workglow/util";
-import { HFT_TASKS } from "./common/HFT_JobRunFns";
+import { HFT_STREAM_TASKS, HFT_TASKS } from "./common/HFT_JobRunFns";
 import { HuggingFaceTransformersProvider } from "./HuggingFaceTransformersProvider";
 
 export function HFT_WORKER_JOBRUN_REGISTER() {
   const workerServer = globalServiceRegistry.get(WORKER_SERVER);
-  new HuggingFaceTransformersProvider(HFT_TASKS).registerOnWorkerServer(workerServer);
+  new HuggingFaceTransformersProvider(HFT_TASKS, HFT_STREAM_TASKS).registerOnWorkerServer(
+    workerServer
+  );
   parentPort.postMessage({ type: "ready" });
   console.log("HFT_WORKER_JOBRUN registered");
 }

packages/ai-provider/src/hf-transformers/HuggingFaceTransformersProvider.ts

@@ -4,7 +4,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { AiProvider, type AiProviderRegisterOptions, type AiProviderRunFn } from "@workglow/ai";
+import {
+  AiProvider,
+  type AiProviderRegisterOptions,
+  type AiProviderRunFn,
+  type AiProviderStreamFn,
+} from "@workglow/ai";
 import { HF_TRANSFORMERS_ONNX } from "./common/HFT_Constants";
 import type { HfTransformersOnnxModelConfig } from "./common/HFT_ModelSchema";
 
@@ -58,8 +63,11 @@ export class HuggingFaceTransformersProvider extends AiProvider<HfTransformersOn
     "ObjectDetectionTask",
   ] as const;
 
-  constructor(tasks?: Record<string, AiProviderRunFn<any, any, HfTransformersOnnxModelConfig>>) {
-    super(tasks);
+  constructor(
+    tasks?: Record<string, AiProviderRunFn<any, any, HfTransformersOnnxModelConfig>>,
+    streamTasks?: Record<string, AiProviderStreamFn<any, any, HfTransformersOnnxModelConfig>>
+  ) {
+    super(tasks, streamTasks);
   }
 
   protected override async onInitialize(options: AiProviderRegisterOptions): Promise<void> {