mastraで型補完を効かせる方法

MastraはTypeScript向けのAIエージェント構築フレームワークです。OpenAIのAgents SDKやLangGraphに相当するものです。
Vercel AI SDKがベースになっており、VercelやNext.jsともいい感じに統合していけそうなのが嬉しいですね。

詳細についてはこちらの記事がわかりやすいです。
https://zenn.dev/yosh1/articles/mastra-ai-agent-framework-guide

ワークフロー機能

AgentsやMemoryなど色んな機能がありそうなのですが、今回はワークフローの機能を取り上げます。
Mastraでは次のように書くことでワークフローを定義することができます。 (公式ドキュメントの例)

import { Step, Workflow, Mastra } from "@mastra/core";
import { z } from "zod";
 
export const myWorkflow = new Workflow({
  name: "my-workflow",
  triggerSchema: z.object({
    inputValue: z.number(),
  }),
});
 
myWorkflow
  .step(
    new Step({
      id: "stepOne",
      outputSchema: z.object({
        doubledValue: z.number(),
      }),
      execute: async ({ context }) => ({
        doubledValue: context.triggerData.inputValue * 2,
      }),
    }),
  )
  .then(
    new Step({
      id: "stepTwo",
      outputSchema: z.object({
        incrementedValue: z.number(),
      }),
      execute: async ({ context }) => {
        if (context.steps.stepOne.status !== "success") {
          return { incrementedValue: 0 };
        }
 
        return { incrementedValue: context.steps.stepOne.output.doubledValue + 1 };
      },
    }),
  ).commit();
 
// Register the workflow with Mastra
export const mastra = new Mastra({
  workflows: { myWorkflow },
});

これを mastra dev で実行するとローカルの開発環境が立ち上がります。処理の流れを視覚的に把握したり、実行もここからできたりしてとても便利です。

型補完が効かない問題

zodも標準で組み込まれていて最高！と思ったのですが、どうやら上の方法では型補完が効いてくれません。

これじゃPython書くのと変わらんやんけ！

調べてみると、 context.steps.stepOne.output.doubledValue ではなく、stepオブジェクトを変数で定義して context.getStepResult(stepOne)?.doubledValue というように取得する必要があるようです。
ただし、triggerDataに関しては同様の方法は見つけられませんでした。StepはWorkflowに依存しないよう設計されているので、仕方ないといえばそうかもしれませんが... 😢
一旦、次のようなヘルパーを用意することで解決しました。

import { Step, Workflow, WorkflowContext } from "@mastra/core";
import { z } from "zod";

export const getWorkflowTriggerData = <TTriggerSchema extends z.ZodObject<any>>(
  context: WorkflowContext,
  workflow: Workflow<any, TTriggerSchema>
): z.infer<TTriggerSchema> => {
  return context.triggerData;
};

これを使うと、最初のワークフローもこんな感じで書き直せて、型補完が効くようになります。これで夜も安心して眠れますね 😌

import { Step, Workflow, Mastra, WorkflowContext } from "@mastra/core";
import { z } from "zod";

const getWorkflowTriggerData = <TTriggerSchema extends z.ZodObject<any>>(
  context: WorkflowContext,
  workflow: Workflow<any, TTriggerSchema>
): z.infer<TTriggerSchema> => {
  return context.triggerData;
};

export const myWorkflow = new Workflow({
  name: "my-workflow",
  triggerSchema: z.object({
    inputValue: z.number(),
  }),
});

const stepOne = new Step({
  id: "stepOne",
  outputSchema: z.object({
    doubledValue: z.number(),
  }),
  execute: async ({ context }) => ({
    doubledValue: getWorkflowTriggerData(context, myWorkflow).inputValue * 2,
  }),
});

const stepTwo = new Step({
  id: "stepTwo",
  outputSchema: z.object({
    incrementedValue: z.number(),
  }),
  execute: async ({ context }) => {
    if (context.steps.stepOne.status !== "success") {
      return { incrementedValue: 0 };
    }

    return {
      incrementedValue: context.getStepResult(stepOne)?.doubledValue + 1,
    };
  },
});

myWorkflow.step(stepOne).then(stepTwo).commit();

// Register the workflow with Mastra
export const mastra = new Mastra({
  workflows: { myWorkflow },
});

やったぜ。

感想

以前LangGraphを触ってみたことがあったのですが、Mastraの方が設計もシンプルでわかりやすく、グラフのプレビュー機能もクラウドを使わずとも手元でパッと立ち上がってくれるため、かなり扱いやすい印象を持ちました。

ただ今回見たようにうまく補完が効かない場面があったり、呼び出し時にzod.parseも入っていない (requiredなパラメータが入っているか保証されない) など、型周りはこれからな印象を受けました。
クラウド版 も公開予定らしくアップデートの勢いにも期待が持てそうなので、今後のエコシステムの発展が楽しみです。