feat: prepare architecture for parallelism and add helpers

example/src/surprise_trip.ts

@@ -68,7 +68,7 @@ const researchTripWorkflow = workflow({
   `,
   output: `
     Comprehensive day-by-day itinerary for the trip to Wrocław, Poland.
-    Ensure the itinerary integrates flights, hotel information, and all planned activities and dining experiences.
+    Ensure the itinerary includes flights, hotel information, and all planned activities and dining experiences.
   `,
   snapshot: logger,
 })

packages/framework/src/agent.ts

@@ -3,7 +3,7 @@ import { zodFunction, zodResponseFormat } from 'openai/helpers/zod.js'
 import { z } from 'zod'
 
 import { openai, Provider } from './models.js'
-import { WorkflowState } from './state.js'
+import { finish, request, response, WorkflowState } from './state.js'
 import { Tool } from './tool.js'
 import { Message } from './types.js'
 import { Workflow } from './workflow.js'
@@ -39,7 +39,7 @@ export const agent = (options: AgentOptions = {}): Agent => {
             )
           : []
 
-        const response = await provider.completions({
+        const res = await provider.completions({
           messages: [
             {
               role: 'system',
@@ -56,18 +56,9 @@ export const agent = (options: AgentOptions = {}): Agent => {
                 Try to complete the task on your own.
               `,
             },
-            {
-              role: 'assistant',
-              content: 'What have been done so far?',
-            },
-            {
-              role: 'user',
-              content: `Here is all the work done so far by other agents: ${JSON.stringify(context)}`,
-            },
-            {
-              role: 'assistant',
-              content: 'What do you want me to do now?',
-            },
+            response('What have been done so far?'),
+            request(`Here is all the work done so far by other agents: ${JSON.stringify(context)}`),
+            response('What do you want me to do now?'),
             ...state.messages,
           ],
           tools: mappedTools.length > 0 ? mappedTools : undefined,
@@ -94,48 +85,34 @@ export const agent = (options: AgentOptions = {}): Agent => {
           ),
         })
 
-        if (response.choices[0].message.tool_calls.length > 0) {
+        if (res.choices[0].message.tool_calls.length > 0) {
           return {
             ...state,
             status: 'paused',
-            messages: state.messages.concat(response.choices[0].message),
+            messages: state.messages.concat(res.choices[0].message),
           }
         }
 
-        const result = response.choices[0].message.parsed
-        if (!result) {
+        const message = res.choices[0].message.parsed
+        if (!message) {
           throw new Error('No parsed response received')
         }
 
-        if (result.response.kind === 'error') {
-          throw new Error(result.response.reasoning)
+        if (message.response.kind === 'error') {
+          throw new Error(message.response.reasoning)
         }
 
-        const agentResponse = {
-          role: 'assistant' as const,
-          content: result.response.result,
-        }
+        const agentResponse = response(message.response.result)
 
-        if (result.response.nextStep) {
+        if (message.response.nextStep) {
           return {
             ...state,
             status: 'running',
-            messages: [
-              ...state.messages,
-              agentResponse,
-              {
-                role: 'user',
-                content: result.response.nextStep,
-              },
-            ],
+            messages: [...state.messages, agentResponse, request(message.response.nextStep)],
           }
         }
 
-        return {
-          ...state,
-          status: 'finished',
-          messages: [agentResponse],
-        }
+        return finish(state, agentResponse)
       }),
   }
 }

packages/framework/src/agents/final_boss.ts

@@ -3,10 +3,11 @@ import { zodResponseFormat } from 'openai/helpers/zod'
 import { z } from 'zod'
 
 import { agent, AgentOptions } from '../agent.js'
+import { finish, request, response } from '../state.js'
 
 const defaults: AgentOptions = {
   run: async (state, context, workflow) => {
-    const response = await workflow.team[state.agent].provider.completions({
+    const res = await workflow.team[state.agent].provider.completions({
       messages: [
         {
           role: 'system',
@@ -15,13 +16,10 @@ const defaults: AgentOptions = {
           `,
         },
         ...context,
-        {
-          role: 'user',
-          content: s`
-            Please summarize all executed steps and do your best to achieve 
-            the main goal while responding with the final answer
-          `,
-        },
+        request(s`
+          Please summarize all executed steps and do your best to achieve 
+          the main goal while responding with the final answer
+        `),
       ],
       response_format: zodResponseFormat(
         z.object({
@@ -30,15 +28,11 @@ const defaults: AgentOptions = {
         'task_result'
       ),
     })
-    const result = response.choices[0].message.parsed
-    if (!result) {
+    const message = res.choices[0].message.parsed
+    if (!message) {
       throw new Error('No parsed response received')
     }
-    return {
-      ...state,
-      status: 'finished',
-      messages: [...state.messages, { role: 'assistant', content: result.finalAnswer }],
-    }
+    return finish(state, response(message.finalAnswer))
   },
 }
 

packages/framework/src/agents/resource_planner.ts

@@ -3,11 +3,11 @@ import { zodResponseFormat } from 'openai/helpers/zod'
 import { z } from 'zod'
 
 import { agent, AgentOptions } from '../agent.js'
-import { childState } from '../state.js'
+import { handoff, request, response } from '../state.js'
 
 const defaults: AgentOptions = {
   run: async (state, context, workflow) => {
-    const response = await workflow.team[state.agent].provider.completions({
+    const res = await workflow.team[state.agent].provider.completions({
       messages: [
         {
           role: 'system',
@@ -22,21 +22,14 @@ const defaults: AgentOptions = {
             4. Previous task context if available  
           `,
         },
-        {
-          role: 'user',
-          content: s`
+        request(s`
           Here are the available agents:
           <agents>
             ${Object.entries(workflow.team).map(([name, agent]) =>
               agent.description ? `<agent name="${name}">${agent.description}</agent>` : ''
             )}
-          </agents>
-        `,
-        },
-        {
-          role: 'assistant',
-          content: 'What is the task?',
-        },
+          </agents>`),
+        response('What is the task?'),
         ...state.messages,
       ],
       temperature: 0.1,
@@ -49,15 +42,12 @@ const defaults: AgentOptions = {
       ),
     })
 
-    const content = response.choices[0].message.parsed
-    if (!content) {
+    const message = res.choices[0].message.parsed
+    if (!message) {
       throw new Error('No content in response')
     }
 
-    return childState({
-      agent: content.agent,
-      messages: state.messages,
-    })
+    return handoff(state, message.agent, state.messages)
   },
 }
 

packages/framework/src/agents/supervisor.ts

@@ -3,18 +3,22 @@ import { zodResponseFormat } from 'openai/helpers/zod'
 import { z } from 'zod'
 
 import { agent, AgentOptions } from '../agent.js'
-import { childState } from '../state.js'
+import { getSteps } from '../messages.js'
+import { delegate, request, response } from '../state.js'
 
 const defaults: AgentOptions = {
   run: async (state, context, workflow) => {
-    const response = await workflow.team[state.agent].provider.completions({
+    const [workflowRequest, ...messages] = state.messages
+
+    const res = await workflow.team[state.agent].provider.completions({
       messages: [
         {
           role: 'system',
           content: s`
             You are a planner that breaks down complex workflows into smaller, actionable steps.
-            Your job is to determine the next task that needs to be done based on the original workflow and what has been completed so far.
-            If all required tasks are completed, return null.
+            Your job is to determine the next task that needs to be done based on the <workflow> and what has been completed so far.
+
+            You can run tasks in parallel, if they do not depend on each other. Otherwise, run them sequentially.
 
             Rules:
             1. Each task should be self-contained and achievable
@@ -24,55 +28,44 @@ const defaults: AgentOptions = {
             5. Use context from completed tasks to inform next steps
           `,
         },
-        {
-          role: 'assistant',
-          content: 'What is the request?',
-        },
-        ...context,
-        ...state.messages,
+        response('What is the request?'),
+        workflowRequest,
+        response('What has been completed so far?'),
+        ...getSteps(messages),
       ],
       temperature: 0.2,
       response_format: zodResponseFormat(
         z.object({
-          task: z
-            .string()
-            .describe('The next task to be completed or null if the workflow is complete')
-            .nullable(),
-          reasoning: z
-            .string()
-            .describe('The reasoning for selecting the next task or why the workflow is complete'),
+          tasks: z
+            .array(
+              z.object({
+                task: z.string().describe('The next task to be completed'),
+                reasoning: z.string().describe('The reasoning for selecting the next task'),
+              })
+            )
+            .describe('Next tasks, or empty array if the workflow is complete'),
         }),
-        'next_task'
+        'next_tasks'
       ),
     })
 
     try {
-      const content = response.choices[0].message.parsed
+      const content = res.choices[0].message.parsed
       if (!content) {
         throw new Error('No content in response')
       }
 
-      if (!content.task) {
+      if (content.tasks.length === 0) {
         return {
           ...state,
           status: 'finished',
         }
       }
 
-      const agentRequest = {
-        role: 'user' as const,
-        content: content.task,
-      }
-
-      return {
-        ...state,
-        status: 'running',
-        messages: [...state.messages, agentRequest],
-        child: childState({
-          agent: 'resourcePlanner',
-          messages: [agentRequest],
-        }),
-      }
+      return delegate(
+        state,
+        content.tasks.map((item) => ['resourcePlanner', request(item.task)])
+      )
     } catch (error) {
       throw new Error('Failed to determine next task')
     }

packages/framework/src/iterate.ts

@@ -1,4 +1,4 @@
-import { childState, WorkflowState } from './state.js'
+import { childState, finish, response, WorkflowState } from './state.js'
 import { runTools } from './tool_calls.js'
 import { Message } from './types.js'
 import { Workflow } from './workflow.js'
@@ -16,18 +16,20 @@ export async function run(
     })
   }
 
-  if (state.child) {
-    const child = await run(state.child, context.concat(state.messages), workflow)
-    if (child.status === 'finished') {
+  if (state.children.length > 0) {
+    const children = await Promise.all(
+      state.children.map((child) => run(child, context.concat(state.messages), workflow))
+    )
+    if (children.every((child) => child.status === 'finished')) {
       return {
         ...state,
-        messages: state.messages.concat(child.messages),
-        child: null,
+        messages: state.messages.concat(children.flatMap((child) => child.messages)),
+        children: [],
       }
     }
     return {
       ...state,
-      child,
+      children,
     }
   }
 
@@ -43,13 +45,10 @@ export async function run(
   }
 
   if (state.status === 'running' || state.status === 'idle') {
-    return agent.run(state, context, workflow)
-  }
-
-  if (state.status === 'failed') {
-    return {
-      ...state,
-      status: 'finished',
+    try {
+      return agent.run(state, context, workflow)
+    } catch (error) {
+      return finish(state, response(error instanceof Error ? error.message : 'Unknown error'))
     }
   }
 

packages/framework/src/messages.ts

@@ -0,0 +1,29 @@
+import s from 'dedent'
+
+import { Message } from './types.js'
+
+const groupMessagePairs = <T>(messages: T[]): T[][] => {
+  return messages.reduce((pairs: T[][], message: T, index: number) => {
+    if (index % 2 === 0) {
+      pairs.push([message])
+    } else {
+      pairs[pairs.length - 1].push(message)
+    }
+    return pairs
+  }, [])
+}
+
+/**
+ * Chat conversation is stored as an array of [task, result, task, result...] pairs.
+ * This function groups them and returns them as an array of steps.
+ */
+export const getSteps = (conversation: Message[]): Message[] => {
+  const steps = groupMessagePairs(conversation)
+  return steps.map(([task, result]) => ({
+    role: 'user' as const,
+    content: s`
+      Step name: ${task.content}
+      Step result: ${result.content}
+    `,
+  }))
+}