Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions .changeset/feat-agent-prepare-step.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
---
"@voltagent/core": minor
---

feat(core): add `prepareStep` to AgentOptions for per-step tool control

Surfaces the AI SDK's `prepareStep` callback as a top-level `AgentOptions` property so users can set a default step preparation callback at agent creation time. Per-call `prepareStep` in method options overrides the agent-level default.

This enables controlling tool availability, tool choice, and other step settings on a per-step basis without passing `prepareStep` on every call.
21 changes: 20 additions & 1 deletion packages/core/src/agent/agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,7 @@ export type {
SemanticMemoryOptions,
} from "./types";
import { P, match } from "ts-pattern";
import type { StopWhen } from "../ai-types";
import type { PrepareStep, StopWhen } from "../ai-types";
import type { SamplingPolicy } from "../eval/runtime";
import type { ConversationStepRecord } from "../memory/types";
import { applySummarization } from "./apply-summarization";
Expand Down Expand Up @@ -923,6 +923,13 @@ export interface BaseGenerationOptions<TProviderOptions extends ProviderOptions
* Tool choice strategy for AI SDK calls.
*/
toolChoice?: ToolChoice<Record<string, unknown>>;

/**
* Step preparation callback (ai-sdk `prepareStep`).
* Called before each step to control tool availability, tool choice, etc.
* Overrides the agent-level `prepareStep` if provided.
*/
prepareStep?: PrepareStep;
}

export type GenerateTextOptions<
Expand Down Expand Up @@ -964,6 +971,7 @@ export class Agent {
readonly maxSteps: number;
readonly maxRetries: number;
readonly stopWhen?: StopWhen;
readonly prepareStep?: PrepareStep;
readonly markdown: boolean;
readonly inheritParentSpan: boolean;
readonly voice?: Voice;
Expand Down Expand Up @@ -1022,6 +1030,7 @@ export class Agent {
this.maxSteps = options.maxSteps ?? defaultMaxSteps;
this.maxRetries = options.maxRetries ?? DEFAULT_LLM_MAX_RETRIES;
this.stopWhen = options.stopWhen;
this.prepareStep = options.prepareStep;
this.markdown = options.markdown ?? false;
this.inheritParentSpan = options.inheritParentSpan ?? true;
this.voice = options.voice;
Expand Down Expand Up @@ -1265,6 +1274,11 @@ export class Agent {
...aiSDKOptions
} = options || {};

// Apply agent-level prepareStep as default (per-call overrides)
if (this.prepareStep && !aiSDKOptions.prepareStep) {
aiSDKOptions.prepareStep = this.prepareStep as AITextCallOptions["prepareStep"];
}

const forcedToolChoice = oc.systemContext.get(FORCED_TOOL_CHOICE_CONTEXT_KEY) as
| ToolChoice<Record<string, unknown>>
| undefined;
Expand Down Expand Up @@ -1879,6 +1893,11 @@ export class Agent {
...aiSDKOptions
} = options || {};

// Apply agent-level prepareStep as default (per-call overrides)
if (this.prepareStep && !aiSDKOptions.prepareStep) {
aiSDKOptions.prepareStep = this.prepareStep as AITextCallOptions["prepareStep"];
}

const forcedToolChoice = oc.systemContext.get(FORCED_TOOL_CHOICE_CONTEXT_KEY) as
| ToolChoice<Record<string, unknown>>
| undefined;
Expand Down
122 changes: 122 additions & 0 deletions packages/core/src/agent/prepare-step.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,122 @@
import { describe, expect, it, vi } from "vitest";
import { Agent } from "./agent";
import { createMockLanguageModel, defaultMockResponse } from "./test-utils";

describe("prepareStep", () => {
it("should accept prepareStep in AgentOptions", () => {
const prepareStep = vi.fn(() => ({}));
const model = createMockLanguageModel();

const agent = new Agent({
name: "test-agent",
instructions: "test",
model,
prepareStep,
});

expect(agent.prepareStep).toBe(prepareStep);
});

it("should default to undefined when prepareStep is not provided", () => {
const model = createMockLanguageModel();

const agent = new Agent({
name: "test-agent",
instructions: "test",
model,
});

expect(agent.prepareStep).toBeUndefined();
});

it("should pass agent-level prepareStep to generateText", async () => {
const prepareStep = vi.fn(() => ({}));
const model = createMockLanguageModel({
doGenerate: {
...defaultMockResponse,
content: [{ type: "text", text: "done" }],
},
});

const agent = new Agent({
name: "test-agent",
instructions: "test",
model,
prepareStep,
});

await agent.generateText("hello");

// prepareStep is called by the AI SDK on each step
expect(prepareStep).toHaveBeenCalled();
});

it("should pass agent-level prepareStep to streamText", async () => {
const prepareStep = vi.fn(() => ({}));
const model = createMockLanguageModel();

const agent = new Agent({
name: "test-agent",
instructions: "test",
model,
prepareStep,
});

const result = await agent.streamText("hello");
// consume the stream to completion
for await (const _part of result.textStream) {
// drain
}

expect(prepareStep).toHaveBeenCalled();
});

it("should allow per-call prepareStep to override agent-level", async () => {
const agentPrepareStep = vi.fn(() => ({}));
const callPrepareStep = vi.fn(() => ({}));
const model = createMockLanguageModel({
doGenerate: {
...defaultMockResponse,
content: [{ type: "text", text: "done" }],
},
});

const agent = new Agent({
name: "test-agent",
instructions: "test",
model,
prepareStep: agentPrepareStep,
});

await agent.generateText("hello", {
prepareStep: callPrepareStep,
});

// per-call should be used, not agent-level
expect(callPrepareStep).toHaveBeenCalled();
expect(agentPrepareStep).not.toHaveBeenCalled();
});

it("should allow per-call prepareStep to override agent-level in streamText", async () => {
const agentPrepareStep = vi.fn(() => ({}));
const callPrepareStep = vi.fn(() => ({}));
const model = createMockLanguageModel();

const agent = new Agent({
name: "test-agent",
instructions: "test",
model,
prepareStep: agentPrepareStep,
});

const result = await agent.streamText("hello", {
prepareStep: callPrepareStep,
});
for await (const _part of result.textStream) {
// drain
}

expect(callPrepareStep).toHaveBeenCalled();
expect(agentPrepareStep).not.toHaveBeenCalled();
});
});
13 changes: 12 additions & 1 deletion packages/core/src/agent/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ import type {
ProviderTextResponse,
ProviderTextStreamResponse,
} from "../agent/providers/base/types";
import type { StopWhen } from "../ai-types";
import type { PrepareStep, StopWhen } from "../ai-types";

import type { LanguageModel, TextStreamPart, UIMessage } from "ai";
import type { Memory } from "../memory";
Expand Down Expand Up @@ -723,6 +723,17 @@ export type AgentOptions = {
* Per-call `stopWhen` in method options overrides this.
*/
stopWhen?: StopWhen;
/**
* Default step preparation callback (ai-sdk `prepareStep`).
* Called before each step to control tool availability, tool choice, etc.
* Per-call `prepareStep` in method options overrides this.
*
* @example
* ```ts
* prepareStep: ({ steps }) => (steps.length > 0 ? { toolChoice: 'none' } : {}),
* ```
*/
prepareStep?: PrepareStep;
markdown?: boolean;
/**
* When true, use the active VoltAgent span as the parent if parentSpan is not provided.
Expand Down
3 changes: 3 additions & 0 deletions packages/core/src/ai-types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,3 +3,6 @@ import type { generateText } from "ai";

// StopWhen predicate type used by ai-sdk generate/stream functions
export type StopWhen = Parameters<typeof generateText>[0]["stopWhen"];

// PrepareStep callback type used by ai-sdk generate/stream functions
export type PrepareStep = Parameters<typeof generateText>[0]["prepareStep"];
2 changes: 1 addition & 1 deletion packages/core/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -302,7 +302,7 @@ export { createAsyncIterableStream, type AsyncIterableStream } from "@voltagent/
// Convenience re-exports from ai-sdk so apps need only @voltagent/core
export { stepCountIs, hasToolCall } from "ai";
export type { LanguageModel } from "ai";
export type { StopWhen } from "./ai-types";
export type { PrepareStep, StopWhen } from "./ai-types";

export type {
ManagedMemoryStatus,
Expand Down
13 changes: 13 additions & 0 deletions website/docs/getting-started/migration-guide.md
Original file line number Diff line number Diff line change
Expand Up @@ -610,6 +610,19 @@ console.log(out.context); // VoltAgent Map
- This overrides VoltAgent's default `stepCountIs(maxSteps)` guard.
- Be cautious: permissive predicates can lead to long-running or looping generations; overly strict ones may stop before tools complete.

### prepareStep callback (advanced)

- You can pass an ai-sdk `prepareStep` callback in `AgentOptions` or in per-call method options to control tool availability, tool choice, and other settings before each step.
- Per-call `prepareStep` overrides the agent-level default.
- Example: force text-only output after the first step:
```ts
const agent = new Agent({
name: "my-agent",
model,
prepareStep: ({ steps }) => (steps.length > 0 ? { toolChoice: "none" } : {}),
});
```

### Built-in server removed; use `@voltagent/server-hono`

VoltAgent 1.x decouples the HTTP server from `@voltagent/core`. The built-in server is removed in favor of pluggable server providers. The recommended provider is `@voltagent/server-hono` (powered by Hono). Default port remains `3141`.
Expand Down
Loading