@tanstack/ai 0.58.0 → 0.61.0

package/src/activities/chat/agents/route.ts

@@ -0,0 +1,115 @@
+import { boolean, choice } from '../../evaluate/index'
+import type { DefinedAgent } from './define-agent'
+import type { SubagentOrder, SubagentRouterPick } from './spawn'
+
+const ORDER_KEY = 'order'
+
+export interface SubagentRouteOptions<
+  TAgents extends ReadonlyArray<DefinedAgent>,
+> {
+  /**
+   * Question text for every agent. The key is the agent name.
+   * Omit this and each question uses that agent's description.
+   */
+  when?: { [K in TAgents[number]['name']]: string }
+  /**
+   * Agents that run after the other selected agents. The lead group starts
+   * together. The `then` agents then run one after another in list order, and
+   * each reads the text so far. Used only when the router picks at least one
+   * agent from each group. Otherwise `pick` returns `{ names, order }`.
+   */
+  then?: ReadonlyArray<TAgents[number]['name']>
+}
+
+type RouteQuestions<TAgents extends ReadonlyArray<DefinedAgent>> = {
+  [K in TAgents[number]['name']]: ReturnType<typeof boolean>
+} & {
+  order: ReturnType<
+    typeof choice<{
+      parallel: string
+      sequence: string
+    }>
+  >
+}
+
+type RouteAnswers<TAgents extends ReadonlyArray<DefinedAgent>> = {
+  [K in TAgents[number]['name']]: { value: boolean }
+} & {
+  order: { value: SubagentOrder }
+}
+
+/**
+ * Build `decide()` questions for a subagent router.
+ *
+ * One yes/no question per agent, plus an `order` choice.
+ * `pick` returns `main`, one name, `{ names, order }`, or `{ steps }`.
+ * Names follow the `agents` array order.
+ * `{ names, order }` overrides `subagents.order` for that turn.
+ * `then`: agents that run after the other selected agents. The lead group
+ * starts together. The `then` agents then run one after another in list
+ * order, and each reads the text so far. Used only when the router picks at
+ * least one agent from each group. Otherwise `pick` returns `{ names, order }`.
+ */
+export function subagentRoute<
+  const TAgents extends ReadonlyArray<DefinedAgent>,
+>(agents: TAgents, options?: SubagentRouteOptions<TAgents>) {
+  const namesInList = new Set(agents.map((agent) => agent.name))
+  for (const agent of agents) {
+    if (agent.name === ORDER_KEY) {
+      throw new Error(
+        'subagentRoute cannot use an agent named "order". Rename that agent.',
+      )
+    }
+  }
+  for (const name of options?.then ?? []) {
+    if (!namesInList.has(name)) {
+      throw new Error(`subagentRoute then includes unknown agent "${name}".`)
+    }
+  }
+
+  const questions: Record<string, unknown> = {
+    order: choice({
+      instructions: 'When more than one agent runs, how must they run?',
+      options: {
+        parallel:
+          'Start them together. Use this when no agent must read text from another agent. Working on the same topic is not a reason to wait.',
+        sequence:
+          'Run them in agent-list order. Use this only when a later agent must read the earlier agent text, such as research notes and then a draft article.',
+      },
+    }),
+  }
+
+  for (const agent of agents) {
+    const when = options?.when?.[agent.name as TAgents[number]['name']]
+    questions[agent.name] = boolean({
+      instructions: when ?? agent.description,
+    })
+  }
+
+  function pick(result: RouteAnswers<TAgents>): SubagentRouterPick {
+    const names = agents
+      .map((agent) => agent.name)
+      .filter((name) => result[name as TAgents[number]['name']].value)
+    if (names.length === 0) return 'main'
+    const only = names.length === 1 ? names[0] : undefined
+    if (only !== undefined) return only
+    const later = new Set(options?.then ?? [])
+    const lead = names.filter((name) => !later.has(name))
+    const tail = names.filter((name) => later.has(name))
+    if (lead.length === 0 || tail.length === 0) {
+      return { names, order: result.order.value }
+    }
+    return {
+      steps: [
+        lead.length > 1
+          ? { names: lead, order: 'parallel' as const }
+          : { names: lead },
+        tail.length > 1
+          ? { names: tail, order: 'sequence' as const }
+          : { names: tail },
+      ],
+    }
+  }
+
+  return { questions: questions as RouteQuestions<TAgents>, pick }
+}