Conditional branches

import {
  createConditionStep,
  createStep,
  createWorkflow,
} from "@ai_kit/core";

const calculateScore = createStep<{ score: number }, { score: number }>({
  id: "calculate-score",
  handler: ({ input }) => ({ score: Math.max(0, Math.min(1, input.score)) }),
});

const evaluateRisk = createConditionStep<{ score: number }, { score: number }>({
  id: "evaluate-risk",
  resolveBranch: ({ output }) => {
    if (output.score >= 0.8) return "high";
    if (output.score >= 0.5) return "medium";
    if (output.score >= 0.2) return "low";
    return undefined;
  },
});

const handleLow = createStep<{ score: number }, string>({
  id: "handle-low",
  handler: ({ input }) => `LOW:${input.score}`,
});

const handleMedium = createStep<{ score: number }, string>({
  id: "handle-medium",
  handler: ({ input }) => `MEDIUM:${input.score}`,
});

const handleHigh = createStep<{ score: number }, string>({
  id: "handle-high",
  handler: ({ input }) => `HIGH:${input.score}`,
});

export const riskWorkflow = createWorkflow({ id: "risk" })
  .then(calculateScore)
  .conditions(evaluateRisk)
  .then({
    low: handleLow,
    medium: handleMedium,
    high: handleHigh,
  })
  .commit();

createWorkflow({ id: "risk" })
  .then(calculateScore)
  .conditions(
    createConditionStep<{ score: number }, { score: number }>({
      id: "evaluate-risk-index",
      resolveBranch: ({ output }) => (output.score >= 0.5 ? 0 : 1),
    }),
  )
  .then(handleHigh, handleLow)
  .commit();