[pydantic-ai] Bedrock CachePoint가 여러 trailing 문서 사이에 잘못 배치되는 버그 수정

PR 링크: pydantic/pydantic-ai#4207 상태: Merged | 변경: +240 / -45

들어가며

AWS Bedrock API는 CachePoint가 문서/비디오 바로 뒤에 오는 것을 거부합니다. pydantic-ai는 이를 우회하기 위해 CachePoint를 문서 앞으로 이동시키는 로직을 가지고 있었지만, 문서가 여러 개 연속될 때 마지막 문서 바로 앞에만 배치하여 여전히 다른 문서 뒤에 위치하는 문제가 있었습니다.

핵심 코드 분석

중복 로직을 함수로 추출

Before (두 곳에서 중복):

multimodal_keys = ['document', 'video']
last_multimodal_index = next(
    (i for i, block in reversed(list(enumerate(last_user_content)))
     if any(key in block for key in multimodal_keys)),
    None,
)
if last_multimodal_index is not None and last_multimodal_index > 0:
    prev_block = last_user_content[last_multimodal_index - 1]
    if not (isinstance(prev_block, dict) and 'cachePoint' in prev_block):
        last_user_content.insert(last_multimodal_index, {'cachePoint': {'type': 'default'}})

이 코드는 reversed로 마지막 문서의 인덱스만 찾았으므로, [text, doc1, doc2, CachePoint]에서 doc2 앞에 CachePoint를 넣어 [text, doc1, CachePoint, doc2]가 되었습니다. doc1 뒤에 CachePoint가 오므로 API가 거부합니다.

After:

def _insert_cache_point_before_trailing_documents(
    content: list[Any],
    *,
    raise_if_cannot_insert: bool = False,
) -> bool:
    multimodal_keys = ['document', 'video']
    # Find where the trailing contiguous group starts
    trailing_start: int | None = None
    for i in range(len(content) - 1, -1, -1):
        if any(key in content[i] for key in multimodal_keys):
            trailing_start = i
        else:
            break

    if trailing_start is not None and trailing_start > 0:
        prev_block = content[trailing_start - 1]
        if isinstance(prev_block, dict) and 'cachePoint' in prev_block:
            return False
        content.insert(trailing_start, {'cachePoint': {'type': 'default'}})
        return True

역순 순회하면서 trailing_start를 계속 갱신하여 연속된 문서/비디오 그룹의 시작점을 찾습니다. [text, doc1, doc2]에서 trailing_start는 doc1의 인덱스가 되어 [text, CachePoint, doc1, doc2]로 올바르게 배치됩니다. 중간에 이미지가 있으면 그룹이 끊어집니다(이미지는 제한 없음).

왜 이게 좋은가

이 수정은 세 가지 측면에서 우수합니다. 첫째, 실제 Bedrock API 제약사항("문서/비디오 바로 뒤에 CachePoint 불가")을 정확히 이해하고 trailing contiguous group 전체를 고려합니다. 둘째, 두 곳에서 중복되던 로직을 하나의 함수로 추출하여 향후 유지보수성을 높였습니다. 셋째, raise_if_cannot_insert 파라미터로 자동 캐싱(조용히 스킵)과 명시적 CachePoint(에러 발생) 두 가지 사용 맥락을 하나의 함수로 지원합니다.

정리

항목	내용
문제	복수 trailing 문서 시 CachePoint가 중간에 배치됨
해결	trailing contiguous group의 시작점 탐색으로 변경
부가	중복 코드 함수 추출 + 테스트 추가

참고 자료

• AWS Bedrock Prompt Caching
• pydantic-ai Bedrock 문서

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