Judgments with the AI SDK

pnpm add ai @ai-sdk/typesafe-ai

import { z } from "zod";
import { createAsyncLogic } from "xstate";
import { experimental_evaluate as evaluate, type Experimental_EvaluationModel } from "ai";
import { typeSafeAi } from "@ai-sdk/typesafe-ai";
import { setupAgent } from "@statelyai/agent";

const judgeModel: Experimental_EvaluationModel = typeSafeAi.evaluationModel("jev-latest");

/** A ticket is flagged urgent when the probability that it is blocking clears this. */
export const URGENT_THRESHOLD = 0.5;

export function createClassifyTicket(model: Experimental_EvaluationModel = judgeModel) {
  return createAsyncLogic<
    {
      answers: {
        team: { choice: "billing" | "technical" | "unclear" };
        urgent: { probability: number };
      };
    },
    { ticket: string; plan: string }
  >({
    run: async ({ input, signal }) => {
      const { answers } = await evaluate({
        model,
        state: { ticket: input.ticket, plan: input.plan },
        questions: {
          team: {
            type: "choice",
            instructions: "Which team should handle `ticket`?",
            criteria: {
              billing: "Charges, refunds, invoices, or plan changes.",
              technical: "Errors, outages, or a feature not working as documented.",
              unclear: "`ticket` does not say enough to pick a team.",
            },
          },
          urgent: {
            type: "boolean",
            instructions: "Does `ticket` describe a problem blocking the customer right now?",
          },
        },
        abortSignal: signal,
      });
      return { answers };
    },
  });
}

const agentSetup = setupAgent({
  context: z.object({
    ticket: z.string(),
    plan: z.string(),
    team: z.string().nullable(),
    urgent: z.boolean(),
  }),
  input: z.object({ ticket: z.string(), plan: z.string() }),
  output: z.object({ team: z.string().nullable(), urgent: z.boolean() }),
  actors: { classifyTicket: createClassifyTicket() },
});

export const triageMachine = agentSetup.createMachine({
  context: ({ input }) => ({ ticket: input.ticket, plan: input.plan, team: null, urgent: false }),
  initial: "classifying",
  states: {
    classifying: {
      invoke: {
        src: "classifyTicket",
        input: ({ context }) => ({ ticket: context.ticket, plan: context.plan }),
        onDone: ({ output }) => {
          const { team, urgent } = output.answers;
          return team.choice === "unclear"
            ? { target: "needsHuman" }
            : {
                target: "routed",
                context: { team: team.choice, urgent: urgent.probability >= URGENT_THRESHOLD },
              };
        },
        onError: { target: "needsHuman" },
      },
    },
    routed: {
      type: "final",
      output: ({ context }) => ({ team: context.team, urgent: context.urgent }),
    },
    needsHuman: {
      type: "final",
      output: () => ({ team: null, urgent: false }),
    },
  },
});

import { createAsyncLogic } from "xstate";
import { experimental_evaluate as evaluate, type Experimental_EvaluationModel } from "ai";

export const QUALITY_LEVELS = [
  "Off-topic or missing the requested content.",
  "On-topic but unclear or incomplete.",
  "Clear and complete with minor issues.",
  "Clear, complete, and ready to send.",
] as const;

/** Accept a draft at or above "Clear and complete with minor issues." */
export const ACCEPT_SCORE = 2;

export function createGradeDraft(model: Experimental_EvaluationModel) {
  return createAsyncLogic<
    { answers: { quality: { score: number } } },
    { brief: string; draft: string }
  >({
    run: async ({ input, signal }) => {
      const { answers } = await evaluate({
        model,
        state: { brief: input.brief, draft: input.draft },
        questions: {
          quality: {
            type: "score",
            instructions: "How well does `draft` meet `brief`?",
            criteria: QUALITY_LEVELS,
          },
        },
        abortSignal: signal,
      });
      return { answers };
    },
  });
}

// In the invoking state:
// onDone: ({ output }) =>
//   output.answers.quality.score >= ACCEPT_SCORE ? { target: "done" } : { target: "revising" },

import { experimental_evaluate as evaluate } from "ai";
import { typeSafeAi } from "@ai-sdk/typesafe-ai";

/** Route only when Jev is at least this sure of the team. */
export const ROUTING_CONFIDENCE = 0.6;

const { answers, providerMetadata } = await evaluate({
  model: typeSafeAi.evaluationModel("jev-latest"),
  state: { ticket: "I was charged twice this month." },
  questions: {
    team: {
      type: "choice",
      instructions: "Which team should handle `ticket`?",
      criteria: { billing: "Charges and refunds.", technical: "Errors and outages." },
    },
  },
});

const reported = (providerMetadata?.typesafe?.confidence as Record<string, unknown> | undefined)
  ?.team;
const confidence = typeof reported === "number" ? reported : 0;
const route = confidence >= ROUTING_CONFIDENCE ? answers.team.choice : "needsHuman";

import type { Experimental_EvaluationModel } from "ai";
import { runAgent } from "@statelyai/agent";
import { expect, test } from "vitest";

const scripted: Experimental_EvaluationModel = {
  specificationVersion: "v4",
  provider: "test",
  modelId: "scripted-judge",
  supportedQuestionTypes: ["choice", "boolean", "score"],
  doEvaluate: async () => ({
    answers: {
      team: { type: "choice", choice: "billing" },
      urgent: { type: "boolean", probability: 0.2 },
    },
    warnings: [],
  }),
};

test("routes a double charge to billing", async () => {
  const result = await runAgent(triageMachine, {
    input: { ticket: "I was charged twice this month.", plan: "pro" },
    actors: { classifyTicket: createClassifyTicket(scripted) },
  });

  expect(result.status === "done" && result.output).toEqual({ team: "billing", urgent: false });
});