diff --git a/.changeset/tall-dryers-run.md b/.changeset/tall-dryers-run.md new file mode 100644 index 000000000..bb46e6513 --- /dev/null +++ b/.changeset/tall-dryers-run.md @@ -0,0 +1,32 @@ +--- +"@voltagent/core": patch +--- + +fix: allow `andAgent` schema to accept `Output.*` specs (arrays, choices, json, text) + +`andAgent` now supports the same output flexibility as `agent.generateText`, so you can return non-object +structures from workflow steps. + +Usage: + +```ts +import { Output } from "ai"; +import { z } from "zod"; +import { Agent, createWorkflowChain } from "@voltagent/core"; + +const agent = new Agent({ + name: "Tagger", + model: "openai/gpt-4o-mini", + instructions: "Return tags only.", +}); + +const workflow = createWorkflowChain({ + id: "tag-workflow", + input: z.object({ topic: z.string() }), +}).andAgent(async ({ data }) => `List tags for ${data.topic}`, agent, { + schema: Output.array({ element: z.string() }), +}); + +const result = await workflow.run({ topic: "workflows" }); +// result: string[] +``` diff --git a/packages/core/src/workflow/chain.spec-d.ts b/packages/core/src/workflow/chain.spec-d.ts index 8be74a34a..7ff8a9c74 100644 --- a/packages/core/src/workflow/chain.spec-d.ts +++ b/packages/core/src/workflow/chain.spec-d.ts @@ -1,3 +1,4 @@ +import { Output } from "ai"; import { describe, expectTypeOf, it } from "vitest"; import { z } from "zod"; import type { Agent } from "../agent/agent"; @@ -443,6 +444,27 @@ describe("workflow chain - type inference", () => { expectTypeOf(workflow).not.toBeNever(); }); + it("should infer andAgent output from output specs", () => { + const workflow = createWorkflowChain({ + id: "test-agent-output", + name: "Test Agent Output", + input: z.object({ prompt: z.string() }), + result: z.array(z.string()), + }) + .andAgent(async ({ data }) => `List: ${data.prompt}`, mockAgent, { + schema: Output.array({ element: z.string() }), + }) + .andThen({ + id: "ensure-array", + execute: async ({ data }) => { + expectTypeOf(data).toEqualTypeOf(); + return data; + }, + }); + + expectTypeOf(workflow).not.toBeNever(); + }); + it("should allow mapping andAgent output into existing data", () => { const workflow = createWorkflowChain({ id: "test-agent-map", diff --git a/packages/core/src/workflow/chain.ts b/packages/core/src/workflow/chain.ts index 733515179..1ab100f20 100644 --- a/packages/core/src/workflow/chain.ts +++ b/packages/core/src/workflow/chain.ts @@ -3,7 +3,7 @@ import type { Logger } from "@voltagent/internal"; import type { DangerouslyAllowAny } from "@voltagent/internal/types"; import type { UIMessage } from "ai"; import type { z } from "zod"; -import type { Agent, BaseGenerationOptions } from "../agent/agent"; +import type { Agent } from "../agent/agent"; import { createWorkflow } from "./core"; import type { InternalAnyWorkflowStep, @@ -42,6 +42,7 @@ import { andWhen, andWorkflow, } from "./steps"; +import type { AgentConfig, AgentOutputSchema, InferAgentOutput } from "./steps/and-agent"; import type { InternalWorkflow } from "./steps/types"; import type { Workflow, @@ -55,23 +56,7 @@ import type { WorkflowStreamWriter, } from "./types"; -/** - * Agent configuration for the chain - */ -export type AgentConfig< - SCHEMA extends z.ZodTypeAny, - INPUT_SCHEMA extends InternalBaseWorkflowInputSchema, - CURRENT_DATA, -> = BaseGenerationOptions & { - schema: - | SCHEMA - | (( - context: Omit< - WorkflowExecuteContext, CURRENT_DATA, any, any>, - "suspend" | "writer" - >, - ) => SCHEMA | Promise); -}; +export type { AgentConfig } from "./steps/and-agent"; /** * A workflow chain that provides a fluent API for building workflows @@ -164,7 +149,7 @@ export class WorkflowChain< * @param map - Optional mapper to shape or merge the agent output with existing data * @returns A workflow step that executes the agent with the task */ - andAgent( + andAgent( task: | string | UIMessage[] @@ -177,10 +162,16 @@ export class WorkflowChain< any >, agent: Agent, - config: AgentConfig, - ): WorkflowChain, SUSPEND_SCHEMA, RESUME_SCHEMA>; + config: AgentConfig, CURRENT_DATA>, + ): WorkflowChain< + INPUT_SCHEMA, + RESULT_SCHEMA, + InferAgentOutput, + SUSPEND_SCHEMA, + RESUME_SCHEMA + >; - andAgent( + andAgent( task: | string | UIMessage[] @@ -193,9 +184,9 @@ export class WorkflowChain< any >, agent: Agent, - config: AgentConfig, + config: AgentConfig, CURRENT_DATA>, map: ( - output: z.infer, + output: InferAgentOutput, context: WorkflowExecuteContext, CURRENT_DATA, any, any>, ) => Promise | NEW_DATA, ): WorkflowChain; @@ -213,7 +204,7 @@ export class WorkflowChain< any >, agent: Agent, - config: AgentConfig, + config: AgentConfig, CURRENT_DATA>, map?: ( output: unknown, context: WorkflowExecuteContext, CURRENT_DATA, any, any>, diff --git a/packages/core/src/workflow/steps/and-agent.spec.ts b/packages/core/src/workflow/steps/and-agent.spec.ts index c80f6b49d..2cb29e265 100644 --- a/packages/core/src/workflow/steps/and-agent.spec.ts +++ b/packages/core/src/workflow/steps/and-agent.spec.ts @@ -1,4 +1,5 @@ import { safeStringify } from "@voltagent/internal"; +import { Output } from "ai"; import { describe, expect, it } from "vitest"; import { z } from "zod"; import { @@ -49,4 +50,32 @@ describe("andAgent", () => { emailType: { type: "support", priority: "high" }, }); }); + + it("supports output specs for non-object schemas", async () => { + const agent = createTestAgent({ + model: createMockLanguageModel({ + doGenerate: { + ...defaultMockResponse, + content: [ + { + type: "text" as const, + text: safeStringify({ elements: ["alpha", "beta"] }), + }, + ], + }, + }), + }); + + const step = andAgent(({ data }) => `List tags for ${data.topic}`, agent, { + schema: Output.array({ element: z.string() }), + }); + + const result = await step.execute( + createMockWorkflowExecuteContext({ + data: { topic: "workflow" }, + }), + ); + + expect(result).toEqual(["alpha", "beta"]); + }); }); diff --git a/packages/core/src/workflow/steps/and-agent.ts b/packages/core/src/workflow/steps/and-agent.ts index 0fb3312a5..a48ead09b 100644 --- a/packages/core/src/workflow/steps/and-agent.ts +++ b/packages/core/src/workflow/steps/and-agent.ts @@ -1,12 +1,25 @@ import type { ModelMessage } from "@ai-sdk/provider-utils"; -import { Output, type UIMessage } from "ai"; +import { type InferGenerateOutput, Output, type UIMessage } from "ai"; import type { z } from "zod"; import type { Agent, BaseGenerationOptions } from "../../agent/agent"; import { convertUsage } from "../../utils/usage-converter"; import type { InternalWorkflowFunc, WorkflowExecuteContext } from "../internal/types"; import type { WorkflowStepAgent } from "./types"; -export type AgentConfig = BaseGenerationOptions & { +type OutputSpec = Output.Output; + +export type AgentOutputSchema = OutputSpec | z.ZodTypeAny; + +export type InferAgentOutput = SCHEMA extends OutputSpec + ? InferGenerateOutput + : SCHEMA extends z.ZodTypeAny + ? z.infer + : never; + +export type AgentConfig = Omit< + BaseGenerationOptions, + "output" +> & { schema: | SCHEMA | (( @@ -14,11 +27,21 @@ export type AgentConfig = BaseGenerati ) => SCHEMA | Promise); }; -type AgentResultMapper = ( - output: z.infer, +type AgentResultMapper = ( + output: InferAgentOutput, context: WorkflowExecuteContext, ) => Promise | RESULT; +const isOutputSpec = (value: unknown): value is OutputSpec => { + if (!value || typeof value !== "object") return false; + const candidate = value as OutputSpec; + return ( + typeof candidate.parseCompleteOutput === "function" && + typeof candidate.parsePartialOutput === "function" && + "responseFormat" in candidate + ); +}; + /** * Creates an agent step for a workflow * @@ -39,11 +62,16 @@ type AgentResultMapper = ( * * @param task - The task (prompt) to execute for the agent, can be a string or a function that returns a string * @param agent - The agent to execute the task using `generateText` - * @param config - The config for the agent (schema) `generateText` call + * @param config - The config for the agent (schema/output) `generateText` call * @param map - Optional mapper to shape or merge the agent output with existing data * @returns A workflow step that executes the agent with the task */ -export function andAgent>( +export function andAgent< + INPUT, + DATA, + SCHEMA extends AgentOutputSchema, + RESULT = InferAgentOutput, +>( task: | UIMessage[] | ModelMessage[] @@ -64,10 +92,11 @@ export function andAgent) => { + const mapOutput = async (outputValue: InferAgentOutput) => { if (map) { return (await map(outputValue, context)) as RESULT; } @@ -99,7 +128,7 @@ export function andAgent); + return mapOutput(result.output as InferAgentOutput); } // Step start event removed - now handled by OpenTelemetry spans @@ -131,7 +160,7 @@ export function andAgent); + return mapOutput(result.output as InferAgentOutput); } catch (error) { // Check if this is a suspension, not an error if ( diff --git a/website/docs/workflows/steps/and-agent.md b/website/docs/workflows/steps/and-agent.md index 61f068d6d..5adaa522f 100644 --- a/website/docs/workflows/steps/and-agent.md +++ b/website/docs/workflows/steps/and-agent.md @@ -33,20 +33,20 @@ const result = await workflow.run({ text: "I love this!" }); ## How It Works -`andAgent` = AI prompt + structured output schema: +`andAgent` = AI prompt + structured output schema (Zod or `Output.*`): ```typescript .andAgent( prompt, // What to ask the AI agent, // Which AI to use - { schema }, // What shape the answer should be + { schema }, // Zod schema or ai-sdk Output.* spec map? // Optional: merge/shape output with existing data ) ``` If you pass a function for `prompt`, it must return a Promise (use `async`) and resolve to a string, `UIMessage[]`, or `ModelMessage[]`. -**Important:** `andAgent` uses `generateText` with `Output.object` under the hood, which means: +**Important:** `andAgent` uses `generateText`. If you pass a Zod schema, it is wrapped with `Output.object`. If you pass an `Output.*` spec, it is used directly. This means: - ✅ You get **structured, typed responses** based on your schema - ✅ The agent **can use tools** during this step @@ -90,6 +90,14 @@ By default, the step result replaces the workflow data with the agent output. If { schema: z.object({ type: z.enum(["support", "sales", "spam"]) }) }, (output, { data }) => ({ ...data, emailType: output }) ) + +// Use Output.* for non-object outputs (arrays, choices, json, text) +// (Requires: import { Output } from "ai";) +.andAgent( + async ({ data }) => `List tags for: ${data.topic}`, + agent, + { schema: Output.array({ element: z.string() }) } +) ``` ## Middleware and Retries @@ -201,6 +209,24 @@ Retry behavior: ) ``` +### Arrays and Choices + +Requires `import { Output } from "ai";` + +```typescript +.andAgent( + async ({ data }) => `Give 3 tags for: ${data.topic}`, + agent, + { schema: Output.array({ element: z.string() }) } +) + +.andAgent( + async ({ data }) => `Pick a category for: ${data.title}`, + agent, + { schema: Output.choice({ options: ["news", "blog", "doc"] }) } +) +``` + ### Merge Output With Existing Data ```typescript