@workglow/ai 0.0.89 → 0.0.91

package/dist/bun.js

@@ -12,6 +12,17 @@ import { globalServiceRegistry, WORKER_MANAGER } from "@workglow/util";
 
 class AiProviderRegistry {
   runFnRegistry = new Map;
+  streamFnRegistry = new Map;
+  providers = new Map;
+  registerProvider(provider) {
+    this.providers.set(provider.name, provider);
+  }
+  getProvider(name) {
+    return this.providers.get(name);
+  }
+  getProviders() {
+    return new Map(this.providers);
+  }
   registerRunFn(modelProvider, taskType, runFn) {
     if (!this.runFnRegistry.has(taskType)) {
       this.runFnRegistry.set(taskType, new Map);
@@ -29,6 +40,23 @@ class AiProviderRegistry {
     };
     this.registerRunFn(modelProvider, taskType, workerFn);
   }
+  registerStreamFn(modelProvider, taskType, streamFn) {
+    if (!this.streamFnRegistry.has(taskType)) {
+      this.streamFnRegistry.set(taskType, new Map);
+    }
+    this.streamFnRegistry.get(taskType).set(modelProvider, streamFn);
+  }
+  registerAsWorkerStreamFn(modelProvider, taskType) {
+    const streamFn = async function* (input, model, signal) {
+      const workerManager = globalServiceRegistry.get(WORKER_MANAGER);
+      yield* workerManager.callWorkerStreamFunction(modelProvider, taskType, [input, model], { signal });
+    };
+    this.registerStreamFn(modelProvider, taskType, streamFn);
+  }
+  getStreamFn(modelProvider, taskType) {
+    const taskTypeMap = this.streamFnRegistry.get(taskType);
+    return taskTypeMap?.get(modelProvider);
+  }
   getDirectRunFn(modelProvider, taskType) {
     const taskTypeMap = this.runFnRegistry.get(taskType);
     const runFn = taskTypeMap?.get(modelProvider);
@@ -82,6 +110,19 @@ class AiJob extends Job {
       }
     }
   }
+  async* executeStream(input, context) {
+    if (context.signal.aborted || this.status === JobStatus.ABORTING) {
+      throw new AbortSignalJobError("Abort signal aborted before streaming execution of job");
+    }
+    const streamFn = getAiProviderRegistry().getStreamFn(input.aiProvider, input.taskType);
+    if (!streamFn) {
+      const result = await this.execute(input, context);
+      yield { type: "finish", data: result };
+      return;
+    }
+    const model = input.taskInput.model;
+    yield* streamFn(input.taskInput, model, context.signal);
+  }
 }
 // src/model/InMemoryModelRepository.ts
 import { InMemoryTabularStorage } from "@workglow/storage";
@@ -239,6 +280,100 @@ async function resolveModelFromRegistry(id, format, registry) {
   return model;
 }
 registerInputResolver("model", resolveModelFromRegistry);
+// src/provider/AiProvider.ts
+import { globalServiceRegistry as globalServiceRegistry3, WORKER_MANAGER as WORKER_MANAGER2 } from "@workglow/util";
+
+// src/queue/createDefaultQueue.ts
+import { ConcurrencyLimiter, JobQueueClient, JobQueueServer } from "@workglow/job-queue";
+import { InMemoryQueueStorage } from "@workglow/storage";
+import { getTaskQueueRegistry } from "@workglow/task-graph";
+async function createDefaultQueue(providerName, concurrency) {
+  const storage = new InMemoryQueueStorage(providerName);
+  await storage.setupDatabase();
+  const server = new JobQueueServer(AiJob, {
+    storage,
+    queueName: providerName,
+    limiter: new ConcurrencyLimiter(concurrency, 100)
+  });
+  const client = new JobQueueClient({
+    storage,
+    queueName: providerName
+  });
+  client.attach(server);
+  getTaskQueueRegistry().registerQueue({ server, client, storage });
+  await server.start();
+}
+
+// src/provider/AiProvider.ts
+class AiProvider {
+  tasks;
+  streamTasks;
+  constructor(tasks, streamTasks) {
+    this.tasks = tasks;
+    this.streamTasks = streamTasks;
+  }
+  get supportedTaskTypes() {
+    return this.taskTypes;
+  }
+  getRunFn(taskType) {
+    return this.tasks?.[taskType];
+  }
+  getStreamFn(taskType) {
+    return this.streamTasks?.[taskType];
+  }
+  async register(options = { mode: "inline" }) {
+    await this.onInitialize(options);
+    if (options.mode === "worker") {
+      if (!options.worker) {
+        throw new Error(`AiProvider "${this.name}": worker is required when mode is "worker". ` + `Pass a Web Worker instance, e.g. register({ mode: "worker", worker: new Worker(...) }).`);
+      }
+    } else {
+      if (!this.tasks) {
+        throw new Error(`AiProvider "${this.name}": tasks must be provided via the constructor for inline mode. ` + `Pass the tasks record when constructing the provider, e.g. new MyProvider(MY_TASKS).`);
+      }
+    }
+    const registry = getAiProviderRegistry();
+    if (options.mode === "worker" && options.worker) {
+      const workerManager = globalServiceRegistry3.get(WORKER_MANAGER2);
+      workerManager.registerWorker(this.name, options.worker);
+      for (const taskType of this.taskTypes) {
+        registry.registerAsWorkerRunFn(this.name, taskType);
+        registry.registerAsWorkerStreamFn(this.name, taskType);
+      }
+    } else {
+      for (const [taskType, fn] of Object.entries(this.tasks)) {
+        registry.registerRunFn(this.name, taskType, fn);
+      }
+      if (this.streamTasks) {
+        for (const [taskType, fn] of Object.entries(this.streamTasks)) {
+          registry.registerStreamFn(this.name, taskType, fn);
+        }
+      }
+    }
+    registry.registerProvider(this);
+    if (options.queue?.autoCreate !== false) {
+      await this.createQueue(options.queue?.concurrency ?? 1);
+    }
+  }
+  registerOnWorkerServer(workerServer) {
+    if (!this.tasks) {
+      throw new Error(`AiProvider "${this.name}": tasks must be provided via the constructor for worker server registration. ` + `Pass the tasks record when constructing the provider, e.g. new MyProvider(MY_TASKS).`);
+    }
+    for (const [taskType, fn] of Object.entries(this.tasks)) {
+      workerServer.registerFunction(taskType, fn);
+    }
+    if (this.streamTasks) {
+      for (const [taskType, fn] of Object.entries(this.streamTasks)) {
+        workerServer.registerStreamFunction(taskType, fn);
+      }
+    }
+  }
+  async onInitialize(_options) {}
+  async dispose() {}
+  async createQueue(concurrency) {
+    await createDefaultQueue(this.name, concurrency);
+  }
+}
 // src/task/index.ts
 import { TaskRegistry } from "@workglow/task-graph";
 
@@ -1706,6 +1841,26 @@ Workflow9.prototype.textNamedEntityRecognition = CreateWorkflow9(TextNamedEntity
 
 // src/task/TextSummaryTask.ts
 import { CreateWorkflow as CreateWorkflow10, Workflow as Workflow10 } from "@workglow/task-graph";
+
+// src/task/base/StreamingAiTask.ts
+class StreamingAiTask extends AiTask {
+  static type = "StreamingAiTask";
+  async* executeStream(input, context) {
+    const jobInput = await this.getJobInput(input);
+    const queueName = await this.getDefaultQueueName(input);
+    const job = new AiJob({
+      queueName: queueName ?? this.type,
+      jobRunId: this.config.runnerId,
+      input: jobInput
+    });
+    yield* job.executeStream(jobInput, {
+      signal: context.signal,
+      updateProgress: context.updateProgress.bind(this)
+    });
+  }
+}
+
+// src/task/TextSummaryTask.ts
 var modelSchema4 = TypeModel("model:TextSummaryTask");
 var TextSummaryInputSchema = {
   type: "object",
@@ -1726,14 +1881,15 @@ var TextSummaryOutputSchema = {
     text: {
       type: "string",
       title: "Text",
-      description: "The summarized text"
+      description: "The summarized text",
+      "x-stream": "append"
     }
   },
   required: ["text"],
   additionalProperties: false
 };
 
-class TextSummaryTask extends AiTask {
+class TextSummaryTask extends StreamingAiTask {
   static type = "TextSummaryTask";
   static category = "AI Text Model";
   static title = "Text Summary";
@@ -4494,7 +4650,8 @@ import { CreateWorkflow as CreateWorkflow31, Workflow as Workflow31 } from "@wor
 var generatedTextSchema2 = {
   type: "string",
   title: "Text",
-  description: "The generated text"
+  description: "The generated text",
+  "x-stream": "append"
 };
 var modelSchema18 = TypeModel("model:TextGenerationTask");
 var TextGenerationInputSchema = {
@@ -4559,7 +4716,7 @@ var TextGenerationOutputSchema = {
   additionalProperties: false
 };
 
-class TextGenerationTask extends AiTask {
+class TextGenerationTask extends StreamingAiTask {
   static type = "TextGenerationTask";
   static category = "AI Text Model";
   static title = "Text Generation";
@@ -4662,7 +4819,8 @@ var questionSchema = {
 var textSchema = {
   type: "string",
   title: "Text",
-  description: "The generated text"
+  description: "The generated text",
+  "x-stream": "append"
 };
 var modelSchema20 = TypeModel("model:TextQuestionAnswerTask");
 var TextQuestionAnswerInputSchema = {
@@ -4684,7 +4842,7 @@ var TextQuestionAnswerOutputSchema = {
   additionalProperties: false
 };
 
-class TextQuestionAnswerTask extends AiTask {
+class TextQuestionAnswerTask extends StreamingAiTask {
   static type = "TextQuestionAnswerTask";
   static category = "AI Text Model";
   static title = "Text Question Answer";
@@ -4728,14 +4886,15 @@ var TextRewriterOutputSchema = {
     text: {
       type: "string",
       title: "Text",
-      description: "The rewritten text"
+      description: "The rewritten text",
+      "x-stream": "append"
     }
   },
   required: ["text"],
   additionalProperties: false
 };
 
-class TextRewriterTask extends AiTask {
+class TextRewriterTask extends StreamingAiTask {
   static type = "TextRewriterTask";
   static category = "AI Text Model";
   static title = "Text Rewriter";
@@ -4758,7 +4917,8 @@ var modelSchema22 = TypeModel("model:TextTranslationTask");
 var translationTextSchema = {
   type: "string",
   title: "Text",
-  description: "The translated text"
+  description: "The translated text",
+  "x-stream": "replace"
 };
 var TextTranslationInputSchema = {
   type: "object",
@@ -4800,7 +4960,7 @@ var TextTranslationOutputSchema = {
   additionalProperties: false
 };
 
-class TextTranslationTask extends AiTask {
+class TextTranslationTask extends StreamingAiTask {
   static type = "TextTranslationTask";
   static category = "AI Text Model";
   static title = "Text Translation";
@@ -5225,8 +5385,8 @@ var outputSchema15 = {
 
 class VectorQuantizeTask extends Task15 {
   static type = "VectorQuantizeTask";
-  static category = "Vector Processing";
-  static title = "Quantize Vector";
+  static category = "Vector";
+  static title = "Quantize";
   static description = "Quantize vectors to reduce storage and improve performance";
   static cacheable = true;
   static inputSchema() {
@@ -5570,6 +5730,7 @@ export {
   TextClassificationInputSchema,
   TextChunkerTask,
   StructuralParserTask,
+  StreamingAiTask,
   SimilarityFn,
   SegmentationMethod,
   RerankerTask,
@@ -5628,7 +5789,8 @@ export {
   BackgroundRemovalInputSchema,
   AiTask,
   AiProviderRegistry,
+  AiProvider,
   AiJob
 };
 
-//# debugId=2BE731EEB6308C6864756E2164756E21
+//# debugId=DA970B9D5234412964756E2164756E21