블로그 목록
AI

Vercel AI SDK prepareStep으로 에이전트 단계마다 모델 갈아끼우기 — 초안은 싸게, 마무리는 세게

멀티스텝 에이전트를 굴리다 보면 이런 생각이 든다. 모든 단계에 굳이 제일 비싼 모델을 써야 하나?

내가 만든 건 리서치 에이전트였다. 검색 툴을 여러 번 돌려서 자료를 모으고, 마지막에 그걸 정리해서 답을 쓴다. 근데 검색 쿼리를 만들고 툴 호출을 반복하는 중간 단계는 사실 판단이 그렇게 어렵지 않다. 진짜 실력이 필요한 건 모아온 자료를 종합해서 최종 답을 쓰는 마지막 한 스텝이다.

streamText에 모델을 하나 박아두면 모든 스텝이 그 모델로 돈다. 중간까지 강한 모델을 쓰는 건 돈 낭비고, 그렇다고 싼 모델로 통일하면 마무리 품질이 떨어진다.

prepareStep은 스텝 시작 전에 불린다

Vercel AI SDK의 prepareStep은 멀티스텝 루프에서 매 스텝이 시작되기 직전에 호출되는 콜백이다. 여기서 이번 스텝에 쓸 모델, 툴, 시스템 프롬프트를 그때그때 바꿔 끼울 수 있다. 리턴한 값이 그 스텝에만 적용된다.

import { streamText, stepCountIs } from 'ai';
import { anthropic } from '@ai-sdk/anthropic';

const result = streamText({
  model: anthropic('claude-haiku-4-5'), // 기본값
  tools: { search },
  stopWhen: stepCountIs(6),
  prepareStep: async ({ stepNumber, steps }) => {
    // 마지막 스텝이라고 판단되면 강한 모델로 교체
    const isFinal = stepNumber >= 5 || steps.at(-1)?.toolCalls?.length === 0;
    if (isFinal) {
      return {
        model: anthropic('claude-opus-4-8'),
        // 마무리 단계엔 툴을 아예 빼서 딴짓 못 하게
        activeTools: [],
        system: '모아온 자료만으로 근거를 붙여 최종 답변을 작성해.',
      };
    }
    // 중간 단계는 기본값 그대로 (아무것도 리턴 안 하면 유지)
  },
  messages,
});

stepNumber로 몇 번째 스텝인지 알 수 있고, steps에 지금까지의 스텝 기록이 다 들어있다. 직전 스텝에서 툴 호출이 없었다면(toolCalls가 비었다면) 모델이 더 이상 검색할 게 없다고 판단한 거니까, 그때를 마무리 신호로 잡아도 된다.

툴을 스텝마다 열고 닫기

activeTools도 유용하다. 전체 툴은 tools에 다 정의해두되, 특정 스텝에서 실제로 쓸 수 있는 툴만 이름으로 골라 여는 방식이다.

prepareStep: async ({ stepNumber }) => {
  if (stepNumber === 0) {
    // 첫 스텝은 검색만 허용
    return { activeTools: ['search'] };
  }
  if (stepNumber >= 3) {
    // 후반엔 저장 툴만 열기
    return { activeTools: ['saveResult'] };
  }
},

툴이 열 개씩 붙은 에이전트는 매 스텝 프롬프트에 툴 스키마가 다 실려서 토큰도 먹고, 모델이 엉뚱한 툴을 부르기도 한다. 스텝마다 필요한 것만 열어두면 프롬프트도 가벼워지고 오작동도 준다.

messages를 리턴하면 그 스텝에 넘길 대화 내용 자체를 갈아끼울 수도 있어서, 오래된 툴 결과를 압축해 넣는 컨텍스트 관리에도 쓴다.

실제로 얼마나 아꼈나

내 리서치 에이전트는 보통 5~6스텝을 돈다. 그중 마지막 한 스텝만 Opus, 나머지는 Haiku로 돌렸더니 응답 품질은 거의 그대로인데 호출당 비용이 눈에 띄게 내려갔다. 강한 모델은 정말 판단이 필요한 순간에만 부르면 된다.

핵심은 하나다. 에이전트라고 처음부터 끝까지 같은 모델을 쓸 이유는 없다. prepareStep으로 스텝마다 "지금은 뭐가 필요한가"를 정해주면, 비용과 품질을 굳이 하나만 고를 필요가 없어진다.