@polka-codes/core 0.9.37 → 0.9.38

package/README.md

@@ -53,6 +53,51 @@ class MyTool extends ToolBase {
 }
 ```
 
+### Nested workflows
+
+You can reuse existing workflows inside other workflows with the `workflow` step. This is helpful when you want to repeat a
+common CLI pattern such as `custom → agent → custom`.
+
+```typescript
+import {
+  builder,
+  combineHandlers,
+  customStepSpecHandler,
+  makeAgentStepSpecHandler,
+  sequentialStepSpecHandler,
+  workflowStepSpecHandler,
+} from '@polka-codes/core/workflow';
+
+const reviewWorkflow = builder<{ code: string }>()
+  .custom('prepare', async ({ code }) => ({ prompt: `Review this snippet:\n${code}` }))
+  .agent('agent', {
+    messages: [{ type: 'function', fn: (input) => input.prompt }],
+    systemPrompt: 'You are a careful reviewer.',
+    parseOutput: (raw) => ({ success: true, data: { review: raw } }),
+  })
+  .custom('summarize', async (input) => ({ summary: input.review }))
+  .build();
+
+const mainWorkflow = builder<{ code: string; notify: boolean }>()
+  .workflow('review', {
+    workflow: reviewWorkflow,
+    mapInput: (input) => ({ code: input.code }),
+    mapOutput: ({ workflowOutput, input }) => ({ summary: workflowOutput.summary, notify: input.notify }),
+  })
+  .custom('finalize', async ({ summary, notify }) => ({ summary, notified: notify }))
+  .build();
+
+const handlers = combineHandlers(
+  sequentialStepSpecHandler,
+  customStepSpecHandler,
+  workflowStepSpecHandler,
+  makeAgentStepSpecHandler(async () => /* return a LanguageModelV2 */ throw new Error('model not configured')),
+);
+```
+
+The optional `mapInput` and `mapOutput` helpers let you shape data for the child workflow and massage the result before the
+next parent step runs.
+
 ## Development
 
 ### Building

package/dist/_tsup-dts-rollup.d.ts

@@ -235,6 +235,27 @@ export { BaseStepSpec }
 export { BaseStepSpec as BaseStepSpec_alias_1 }
 export { BaseStepSpec as BaseStepSpec_alias_2 }
 
+declare interface BranchStepSpec<TInput extends Record<string, Json> = Record<string, Json>, TOutput extends Record<string, Json> = Record<string, Json>> extends BaseStepSpec<TInput, TOutput> {
+    type: 'branch';
+    branches: {
+        id: string;
+        when: (input: TInput & {
+            $: Record<string, Record<string, Json>>;
+        }, context: WorkflowContext) => boolean | Promise<boolean>;
+        step: BaseStepSpec<TInput, TOutput>;
+    }[];
+    otherwise?: BaseStepSpec<TInput, TOutput>;
+}
+export { BranchStepSpec }
+export { BranchStepSpec as BranchStepSpec_alias_1 }
+export { BranchStepSpec as BranchStepSpec_alias_2 }
+
+declare const branchStepSpecHandler: StepSpecHandler;
+export { branchStepSpecHandler }
+export { branchStepSpecHandler as branchStepSpecHandler_alias_1 }
+export { branchStepSpecHandler as branchStepSpecHandler_alias_2 }
+export { branchStepSpecHandler as branchStepSpecHandler_alias_3 }
+
 declare const builder: <TInput extends Record<string, Json>>() => StepsBuilder<TInput, TInput>;
 export { builder }
 export { builder as builder_alias_1 }
@@ -1010,6 +1031,55 @@ export { Json }
 export { Json as Json_alias_1 }
 export { Json as Json_alias_2 }
 
+declare interface LoopAccumulatorSpec<TInput extends Record<string, Json>, TIterationOutput extends Record<string, Json>, TAccumulator extends Record<string, Json>> {
+    initial: TAccumulator | (() => TAccumulator);
+    update: (state: LoopStepState<TInput, TIterationOutput, TAccumulator> & {
+        result: TIterationOutput;
+    }, context: WorkflowContext) => TAccumulator | Promise<TAccumulator>;
+}
+export { LoopAccumulatorSpec }
+export { LoopAccumulatorSpec as LoopAccumulatorSpec_alias_1 }
+export { LoopAccumulatorSpec as LoopAccumulatorSpec_alias_2 }
+
+declare type LoopStepOutput<TIterationOutput extends Record<string, Json> = Record<string, Json>, TAccumulator extends Record<string, Json> = Record<string, Json>> = {
+    results: TIterationOutput[];
+    last?: TIterationOutput;
+    accumulator?: TAccumulator;
+};
+export { LoopStepOutput }
+export { LoopStepOutput as LoopStepOutput_alias_1 }
+export { LoopStepOutput as LoopStepOutput_alias_2 }
+
+declare interface LoopStepSpec<TInput extends Record<string, Json> = Record<string, Json>, TIterationInput extends Record<string, Json> = TInput, TIterationOutput extends Record<string, Json> = Record<string, Json>, TAccumulator extends Record<string, Json> = Record<string, Json>> extends BaseStepSpec<TInput, LoopStepOutput<TIterationOutput, TAccumulator>> {
+    type: 'loop';
+    step: BaseStepSpec<TIterationInput, TIterationOutput>;
+    while?: (state: LoopStepState<TInput, TIterationOutput, TAccumulator>, context: WorkflowContext) => boolean | Promise<boolean>;
+    until?: (state: LoopStepState<TInput, TIterationOutput, TAccumulator>, context: WorkflowContext) => boolean | Promise<boolean>;
+    maxIterations?: number;
+    mapInput?: (state: LoopStepState<TInput, TIterationOutput, TAccumulator>, context: WorkflowContext) => TIterationInput | Promise<TIterationInput>;
+    accumulator?: LoopAccumulatorSpec<TInput, TIterationOutput, TAccumulator>;
+}
+export { LoopStepSpec }
+export { LoopStepSpec as LoopStepSpec_alias_1 }
+export { LoopStepSpec as LoopStepSpec_alias_2 }
+
+declare const loopStepSpecHandler: StepSpecHandler;
+export { loopStepSpecHandler }
+export { loopStepSpecHandler as loopStepSpecHandler_alias_1 }
+export { loopStepSpecHandler as loopStepSpecHandler_alias_2 }
+export { loopStepSpecHandler as loopStepSpecHandler_alias_3 }
+
+declare type LoopStepState<TInput extends Record<string, Json> = Record<string, Json>, TIterationOutput extends Record<string, Json> = Record<string, Json>, TAccumulator extends Record<string, Json> = Record<string, Json>> = {
+    iteration: number;
+    input: TInput;
+    results: TIterationOutput[];
+    last?: TIterationOutput;
+    accumulator?: TAccumulator;
+};
+export { LoopStepState }
+export { LoopStepState as LoopStepState_alias_1 }
+export { LoopStepState as LoopStepState_alias_2 }
+
 declare const makeAgentStepSpecHandler: (getModelFn: (step: AgentStepSpec, context: WorkflowContext) => Promise<LanguageModelV2>) => StepSpecHandler;
 export { makeAgentStepSpecHandler }
 export { makeAgentStepSpecHandler as makeAgentStepSpecHandler_alias_1 }
@@ -1311,9 +1381,12 @@ declare class StepsBuilder<TInput extends Record<string, Json> = Record<string,
     build(): BaseStepSpec<TInput, TOutput>;
     step<TStepOutput extends Record<string, Json>, TStepSpec extends BaseStepSpec<TOutput, TStepOutput>>(step: TStepSpec): StepsBuilder<TInput, TStepOutput>;
     parallel<TStepOutput extends Record<string, Json>>(id: string, step: BaseStepSpec<TOutput, TStepOutput>): StepsBuilder<TInput, TStepOutput>;
+    loop<TIterationInput extends Record<string, Json> = TOutput, TIterationOutput extends Record<string, Json> = Record<string, Json>, TAccumulator extends Record<string, Json> = Record<string, Json>>(id: string, config: Omit<LoopStepSpec<TOutput, TIterationInput, TIterationOutput, TAccumulator>, 'id' | 'type'>): StepsBuilder<TInput, LoopStepOutput<TIterationOutput, TAccumulator>>;
     custom<TStepOutput extends Record<string, Json>>(spec: CustomStepSpec<TOutput, TStepOutput>): StepsBuilder<TInput, TStepOutput>;
     custom<TStepOutput extends Record<string, Json>>(id: string, run: CustomStepSpec<TOutput, TStepOutput>['run']): StepsBuilder<TInput, TStepOutput>;
+    workflow<TWorkflowInput extends Record<string, Json> = TOutput, TWorkflowOutput extends Record<string, Json> = Record<string, Json>, TStepOutput extends Record<string, Json> = TWorkflowOutput>(id: string, config: Omit<WorkflowStepSpec<TOutput, TWorkflowInput, TWorkflowOutput, TStepOutput>, 'id' | 'type'>): StepsBuilder<TInput, TStepOutput>;
     agent<TStepOutput extends Record<string, Json>>(id: string, spec: Omit<AgentStepSpec<TOutput, TStepOutput>, 'id' | 'type'>): StepsBuilder<TInput, TStepOutput>;
+    branch<TStepOutput extends Record<string, Json>>(id: string, config: Omit<BranchStepSpec<TOutput, TStepOutput>, 'id' | 'type'>): StepsBuilder<TInput, TStepOutput>;
 }
 
 declare type StepSpecHandler = {
@@ -2013,4 +2086,28 @@ export { WorkflowSpecRaw }
 export { WorkflowSpecRaw as WorkflowSpecRaw_alias_1 }
 export { WorkflowSpecRaw as WorkflowSpecRaw_alias_2 }
 
+declare interface WorkflowStepSpec<TInput extends Record<string, Json> = Record<string, Json>, TWorkflowInput extends Record<string, Json> = TInput, TWorkflowOutput extends Record<string, Json> = Record<string, Json>, TOutput extends Record<string, Json> = TWorkflowOutput> extends BaseStepSpec<TInput, TOutput> {
+    type: 'workflow';
+    workflow: WorkflowSpec<TWorkflowInput, TWorkflowOutput>;
+    mapInput?: (input: TInput & {
+        $: Record<string, Record<string, Json>>;
+    }, context: WorkflowContext) => TWorkflowInput | Promise<TWorkflowInput>;
+    mapOutput?: (params: {
+        input: TInput & {
+            $: Record<string, Record<string, Json>>;
+        };
+        workflowInput: TWorkflowInput;
+        workflowOutput: TWorkflowOutput;
+    }, context: WorkflowContext) => TOutput | Promise<TOutput>;
+}
+export { WorkflowStepSpec }
+export { WorkflowStepSpec as WorkflowStepSpec_alias_1 }
+export { WorkflowStepSpec as WorkflowStepSpec_alias_2 }
+
+declare const workflowStepSpecHandler: StepSpecHandler;
+export { workflowStepSpecHandler }
+export { workflowStepSpecHandler as workflowStepSpecHandler_alias_1 }
+export { workflowStepSpecHandler as workflowStepSpecHandler_alias_2 }
+export { workflowStepSpecHandler as workflowStepSpecHandler_alias_3 }
+
 export { }

package/dist/index.d.ts

@@ -118,9 +118,12 @@ export { builder } from './_tsup-dts-rollup.js';
 export { CommandStepSpec } from './_tsup-dts-rollup.js';
 export { commandStepSpecHandler } from './_tsup-dts-rollup.js';
 export { command } from './_tsup-dts-rollup.js';
+export { branchStepSpecHandler } from './_tsup-dts-rollup.js';
 export { customStepSpecHandler } from './_tsup-dts-rollup.js';
+export { loopStepSpecHandler } from './_tsup-dts-rollup.js';
 export { parallelStepSpecHandler } from './_tsup-dts-rollup.js';
 export { sequentialStepSpecHandler } from './_tsup-dts-rollup.js';
+export { workflowStepSpecHandler } from './_tsup-dts-rollup.js';
 export { step } from './_tsup-dts-rollup.js';
 export { combineHandlers } from './_tsup-dts-rollup.js';
 export { Json } from './_tsup-dts-rollup.js';
@@ -130,6 +133,12 @@ export { WorkflowContext } from './_tsup-dts-rollup.js';
 export { StepSpecRaw } from './_tsup-dts-rollup.js';
 export { SequentialStepSpec } from './_tsup-dts-rollup.js';
 export { ParallelStepSpec } from './_tsup-dts-rollup.js';
+export { LoopStepState } from './_tsup-dts-rollup.js';
+export { LoopAccumulatorSpec } from './_tsup-dts-rollup.js';
+export { LoopStepOutput } from './_tsup-dts-rollup.js';
+export { LoopStepSpec } from './_tsup-dts-rollup.js';
+export { WorkflowStepSpec } from './_tsup-dts-rollup.js';
+export { BranchStepSpec } from './_tsup-dts-rollup.js';
 export { CustomStepSpec } from './_tsup-dts-rollup.js';
 export { StepRunResultSuccess } from './_tsup-dts-rollup.js';
 export { StepRunResultPaused } from './_tsup-dts-rollup.js';

package/dist/index.js

@@ -3727,7 +3727,13 @@ var makeAgentStepSpecHandler = (getModelFn) => ({
           const handleExitReason = async (reason) => {
             switch (reason.type) {
               case "Pause":
-                return { type: "paused", state: { messages: agent.messages } };
+                return {
+                  type: "paused",
+                  state: {
+                    messages: agent.messages,
+                    usage: usageMeter.usage
+                  }
+                };
               case "UsageExceeded":
                 return { type: "error", error: new Error("Usage limit exceeded") };
               case "Exit" /* Exit */: {
@@ -3837,6 +3843,13 @@ var StepsBuilder = class {
       step: step2
     });
   }
+  loop(id, config) {
+    return this.step({
+      ...config,
+      id,
+      type: "loop"
+    });
+  }
   custom(idOrSpec, run2) {
     if (typeof idOrSpec === "string") {
       if (!run2) {
@@ -3850,6 +3863,13 @@ var StepsBuilder = class {
     }
     return this.step(idOrSpec);
   }
+  workflow(id, config) {
+    return this.step({
+      ...config,
+      id,
+      type: "workflow"
+    });
+  }
   agent(id, spec) {
     return this.step({
       ...spec,
@@ -3857,6 +3877,13 @@ var StepsBuilder = class {
       type: "agent"
     });
   }
+  branch(id, config) {
+    return this.step({
+      ...config,
+      id,
+      type: "branch"
+    });
+  }
 };
 var builder = () => {
   return new StepsBuilder();
@@ -3934,6 +3961,50 @@ var runStep = async (step2, input, context, resumedState, allOutputs) => {
   }
 };
 
+// src/workflow/workflow.ts
+var run = async (workflow, context, handler15, input) => {
+  const rootStep = handler15(workflow.step, handler15);
+  const result = await runStep(rootStep, input, context, void 0, {});
+  switch (result.type) {
+    case "paused":
+      return {
+        type: "paused",
+        state: result.state
+      };
+    case "error":
+      return {
+        type: "error",
+        error: result.error
+      };
+    case "success":
+      return {
+        type: "success",
+        output: result.output
+      };
+  }
+};
+var resume = async (workflow, context, handler15, state, input) => {
+  const rootStep = handler15(workflow.step, handler15);
+  const result = await runStep(rootStep, input, context, state, {});
+  switch (result.type) {
+    case "paused":
+      return {
+        type: "paused",
+        state: result.state
+      };
+    case "error":
+      return {
+        type: "error",
+        error: result.error
+      };
+    case "success":
+      return {
+        type: "success",
+        output: result.output
+      };
+  }
+};
+
 // src/workflow/steps.ts
 var combineHandlers = (...handlers) => {
   const allHandlers = {};
@@ -4023,6 +4094,196 @@ var parallelStepSpecHandler = {
     };
   }
 };
+var loopStepSpecHandler = {
+  type: "loop",
+  handler(step2, rootHandler) {
+    const innerStep = rootHandler(step2.step, rootHandler);
+    return {
+      ...step2,
+      async run(input, context, resumedState) {
+        const allOutputs = input.$ ?? {};
+        const initialInput = { ...input };
+        delete initialInput.$;
+        const results = resumedState ? [...resumedState.results] : [];
+        let iteration = resumedState?.iteration ?? results.length;
+        let accumulator = resumedState?.accumulator;
+        if (step2.accumulator && accumulator === void 0) {
+          const initialAccumulator = step2.accumulator.initial;
+          accumulator = typeof initialAccumulator === "function" ? initialAccumulator() : initialAccumulator;
+        }
+        const hasPredicates = Boolean(step2.while || step2.until);
+        const effectiveMaxIterations = step2.maxIterations ?? (hasPredicates ? void 0 : 1);
+        const createState = () => ({
+          iteration,
+          input: initialInput,
+          results,
+          last: results[results.length - 1],
+          accumulator
+        });
+        const completeSuccess = () => {
+          const output = { results };
+          if (results.length > 0) {
+            output.last = results[results.length - 1];
+          }
+          if (accumulator !== void 0 && step2.accumulator) {
+            output.accumulator = accumulator;
+          }
+          return { type: "success", output };
+        };
+        const shouldContinue = async () => {
+          const state = createState();
+          if (step2.until) {
+            const shouldStop = await step2.until(state, context);
+            if (shouldStop) {
+              return false;
+            }
+          }
+          if (step2.while) {
+            return Boolean(await step2.while(state, context));
+          }
+          return true;
+        };
+        let pendingResumeState = resumedState;
+        while (true) {
+          if (!hasPredicates && effectiveMaxIterations !== void 0 && iteration >= effectiveMaxIterations) {
+            return completeSuccess();
+          }
+          if (hasPredicates) {
+            const canContinue = await shouldContinue();
+            if (!canContinue) {
+              return completeSuccess();
+            }
+            if (step2.maxIterations !== void 0 && iteration >= step2.maxIterations) {
+              return {
+                type: "error",
+                error: new Error(`Loop exceeded max iterations: ${step2.maxIterations}`)
+              };
+            }
+          }
+          const stateBeforeRun = createState();
+          const iterationInput = step2.mapInput ? await step2.mapInput(stateBeforeRun, context) : stateBeforeRun.last ?? initialInput;
+          const resumeStateForInner = pendingResumeState && pendingResumeState.iteration === iteration ? pendingResumeState.stepState : void 0;
+          const result = await runStep(innerStep, iterationInput, context, resumeStateForInner, allOutputs);
+          if (result.type === "paused") {
+            return {
+              type: "paused",
+              state: {
+                iteration,
+                stepState: result.state,
+                results,
+                accumulator
+              }
+            };
+          }
+          if (result.type === "error") {
+            return result;
+          }
+          results[iteration] = result.output;
+          if (step2.accumulator) {
+            const updated = await step2.accumulator.update(
+              {
+                ...stateBeforeRun,
+                results,
+                last: result.output,
+                accumulator,
+                result: result.output
+              },
+              context
+            );
+            accumulator = updated;
+          }
+          iteration += 1;
+          pendingResumeState = void 0;
+        }
+      }
+    };
+  }
+};
+var workflowStepSpecHandler = {
+  type: "workflow",
+  handler(step2, rootHandler) {
+    return {
+      ...step2,
+      async run(input, context, resumedState) {
+        const { $, ...rest } = input;
+        void $;
+        const defaultWorkflowInput = rest;
+        const workflowInput = step2.mapInput ? await step2.mapInput(input, context) : defaultWorkflowInput;
+        const workflowState = resumedState && "workflowState" in resumedState ? resumedState.workflowState : resumedState;
+        const result = workflowState ? await resume(step2.workflow, context, rootHandler, workflowState, workflowInput) : await run(step2.workflow, context, rootHandler, workflowInput);
+        if (result.type === "paused") {
+          return { type: "paused", state: { workflowState: result.state, workflowInput } };
+        }
+        if (result.type === "error") {
+          return { type: "error", error: result.error };
+        }
+        const finalOutput = step2.mapOutput ? await step2.mapOutput(
+          {
+            input,
+            workflowInput,
+            workflowOutput: result.output
+          },
+          context
+        ) : result.output;
+        return { type: "success", output: finalOutput };
+      }
+    };
+  }
+};
+var branchStepSpecHandler = {
+  type: "branch",
+  handler(step2, rootHandler) {
+    const branches = step2.branches.map((branch) => ({
+      ...branch,
+      step: rootHandler(branch.step, rootHandler)
+    }));
+    const otherwise = step2.otherwise ? rootHandler(step2.otherwise, rootHandler) : void 0;
+    return {
+      ...step2,
+      async run(input, context, resumedState) {
+        const allOutputs = input.$ ?? {};
+        let branchIndex = resumedState?.branchIndex;
+        let stepToRun;
+        if (branchIndex === void 0) {
+          for (let i = 0; i < branches.length; i++) {
+            const shouldRun = await branches[i].when(input, context);
+            if (shouldRun) {
+              branchIndex = i;
+              break;
+            }
+          }
+          if (branchIndex === void 0 && otherwise) {
+            branchIndex = branches.length;
+          }
+        }
+        if (branchIndex !== void 0) {
+          if (branchIndex < branches.length) {
+            stepToRun = branches[branchIndex].step;
+          } else if (branchIndex === branches.length && otherwise) {
+            stepToRun = otherwise;
+          }
+        }
+        if (!stepToRun) {
+          return {
+            type: "error",
+            error: new Error("No branch predicate matched and no otherwise branch provided")
+          };
+        }
+        const result = await runStep(
+          stepToRun,
+          input,
+          context,
+          branchIndex === resumedState?.branchIndex ? resumedState?.stepState : void 0,
+          allOutputs
+        );
+        if (result.type === "paused") {
+          return { type: "paused", state: { branchIndex, stepState: result.state } };
+        }
+        return result;
+      }
+    };
+  }
+};
 var customStepSpecHandler = {
   type: "custom",
   handler(step2) {
@@ -4037,50 +4298,6 @@ var customStepSpecHandler = {
 function step(id, run2) {
   return { id, type: "custom", run: run2 };
 }
-
-// src/workflow/workflow.ts
-var run = async (workflow, context, handler15, input) => {
-  const rootStep = handler15(workflow.step, handler15);
-  const result = await runStep(rootStep, input, context, void 0, {});
-  switch (result.type) {
-    case "paused":
-      return {
-        type: "paused",
-        state: result.state
-      };
-    case "error":
-      return {
-        type: "error",
-        error: result.error
-      };
-    case "success":
-      return {
-        type: "success",
-        output: result.output
-      };
-  }
-};
-var resume = async (workflow, context, handler15, state, input) => {
-  const rootStep = handler15(workflow.step, handler15);
-  const result = await runStep(rootStep, input, context, state, {});
-  switch (result.type) {
-    case "paused":
-      return {
-        type: "paused",
-        state: result.state
-      };
-    case "error":
-      return {
-        type: "error",
-        error: result.error
-      };
-    case "success":
-      return {
-        type: "success",
-        output: result.output
-      };
-  }
-};
 export {
   AgentBase,
   AnalyzerAgent,
@@ -4101,6 +4318,7 @@ export {
   architectAgentInfo,
   askFollowupQuestion_default as askFollowupQuestion,
   attemptCompletion_default as attemptCompletion,
+  branchStepSpecHandler,
   builder,
   capabilities,
   codeFixerAgentInfo,
@@ -4123,6 +4341,7 @@ export {
   getAvailableTools,
   handOver_default as handOver,
   listFiles_default as listFiles,
+  loopStepSpecHandler,
   makeAgentStepSpecHandler,
   makeAgentTool,
   makeMultiAgentTool,
@@ -4144,5 +4363,6 @@ export {
   step,
   systemInformation,
   toolUsePrompt,
+  workflowStepSpecHandler,
   writeToFile_default as writeToFile
 };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@polka-codes/core",
-  "version": "0.9.37",
+  "version": "0.9.38",
   "license": "AGPL-3.0",
   "author": "github@polka.codes",
   "type": "module",