[pydantic-ai] 병렬 tool call 제한 적용 방식 개선 — 사전 검증으로 전환

PR 링크: pydantic/pydantic-ai#2978 상태: Merged | 변경: +137 / -34

들어가며

Pydantic AI는 LLM이 여러 tool call을 한 번에 반환하면 병렬로 실행합니다. 기존에는 각 tool 실행 직전에 usage_limits.check_before_tool_call(usage)를 호출하여 제한을 검사했는데, 병렬 실행 시 여러 task가 동시에 검사를 통과하여 제한을 초과할 수 있었습니다. 이 PR은 tool batch 실행 전에 전체 호출 수를 미리 계산하여 제한을 검증하는 방식으로 전환합니다.

핵심 코드 분석

1. 사전 검증으로 전환

Before (_tool_manager.py):

async def _call_tool(self, call, allow_partial, wrap_validation_errors, 
                     usage_limits=None, count_tool_usage=True):
    # 개별 tool 실행 직전 검사
    if usage_limits is not None and count_tool_usage:
        usage_limits.check_before_tool_call(self.ctx.usage)
    result = await self.toolset.call_tool(name, args_dict, ctx, tool)
    if count_tool_usage:
        self.ctx.usage.tool_calls += 1

After (_agent_graph.py):

async def _call_tools(...):
    # batch 실행 전 전체 호출 수를 미리 계산하여 검증
    if usage_limits.tool_calls_limit is not None:
        projected_usage = deepcopy(usage)
        projected_usage.tool_calls += len(tool_calls)
        usage_limits.check_before_tool_call(projected_usage)
    
    # 검증 통과 후 실행
    for call in tool_calls:
        yield _messages.FunctionToolCallEvent(call)

2. tool_calls 카운터 갱신 위치 이동

Before: _call_tool() 내부에서 usage.tool_calls += 1

After: _call_function_tool() 트레이싱 래퍼에서 실행 성공 후 갱신

async def _call_function_tool(self, ...):
    with tracer.start_as_current_span(...) as span:
        try:
            tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors)
            usage.tool_calls += 1  # 트레이싱 컨텍스트 내에서 갱신

3. 에러 메시지 개선

Before:

The next tool call would exceed the tool_calls_limit of 1 (tool_calls=0)

After:

The next tool call(s) would exceed the tool_calls_limit of 1 (tool_calls=2).

복수형 표현과 실제 projected 값을 포함하여 더 명확한 에러 메시지를 제공합니다.

왜 이게 좋은가

• race condition 근본 해결: Lock 없이, tool 실행 전에 전체 batch를 한번에 검증하므로 병렬 실행에서도 제한이 정확히 적용됩니다.
• output tool은 제한에서 제외: check_before_tool_call()이 _call_tools() 레벨에서 호출되므로, output tool(결과 반환용)은 카운터에 포함되지 않습니다.
• 관심사 분리: _call_tool()에서 usage 관련 로직이 제거되어, 순수하게 tool 실행만 담당합니다.

정리

• 병렬 실행의 제한 검사는 실행 전에 일괄 수행하라: 개별 task 내부에서 검사하면 동시성 문제가 불가피합니다. 실행 전에 projected 값으로 검증하면 깔끔합니다.
• deepcopy로 projected usage를 만들어라: 원본 usage를 수정하지 않고 시뮬레이션하여 side effect를 방지합니다.

참고 자료

• pydantic/pydantic-ai#2978 — PR 전체 diff
• pydantic/pydantic-ai#3133 — 이전 Lock 기반 수정 (이후 revert됨)

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.