본문으로 건너뛰기

[pydantic-ai] GoogleProvider에 http_client 옵션 추가 및 Vertex AI API 키 지원

PR 링크: pydantic/pydantic-ai#3217 상태: Merged | 변경: +165 / -109

들어가며

기존 GoogleProvidergoogle.genai.Client를 직접 전달하거나, API 키 또는 서비스 계정 인증 방식만 지원했습니다. 커스텀 HTTP 클라이언트를 사용하려면 Client를 직접 생성해야 했고, Vertex AI에서 API 키 인증은 지원되지 않았습니다. 이 PR은 http_client, base_url 파라미터를 추가하고, Vertex AI의 API 키 인증을 지원하며, Pydantic AI의 캐시된 httpx 클라이언트를 기본으로 사용합니다.

핵심 코드 분석

1. http_client와 base_url 파라미터 추가

Before (google.py):

def __init__(self, *, api_key: str) -> None: ...
def __init__(self, *, vertexai: bool = False) -> None: ...

After:

def __init__(self, *, api_key: str, http_client: httpx.AsyncClient | None = None, 
             base_url: str | None = None) -> None: ...
def __init__(self, *, vertexai: bool = False, api_key: str | None = None, 
             http_client: httpx.AsyncClient | None = None, base_url: str | None = None) -> None: ...

2. 캐시된 HTTP 클라이언트 기본 사용

Before:

http_options: HttpOptionsDict = {
    'headers': {'User-Agent': get_user_agent()},
    'async_client_args': {'transport': httpx.AsyncHTTPTransport()},
}

After:

http_client = http_client or cached_async_http_client(
    provider='google-vertex' if vertexai else 'google-gla'
)
http_options = HttpOptions(
    base_url=base_url,
    headers={'User-Agent': get_user_agent()},
    httpx_async_client=http_client,
    async_client_args={'transport': httpx.AsyncHTTPTransport()},
)

cached_async_http_client()로 프로바이더별 커넥션 풀을 재사용합니다. HttpOptionsDict 대신 HttpOptions 객체를 사용하여 타입 안전성도 개선됩니다.

3. Vertex AI API 키 인증 지원

if vertexai:
    if vertex_ai_args_used:
        api_key = None  # credentials 사용 시 api_key 무시
    self._client = _SafelyClosingClient(
        vertexai=True,
        api_key=api_key,  # API 키로도 Vertex AI 인증 가능
        project=project,
        location=location,
        credentials=credentials,
        http_options=http_options,
    )

4. _SafelyClosingClient — AttributeError 방지

class _SafelyClosingClient(Client):
    def close(self) -> None:
        try:
            super().close()
        except AttributeError:
            pass

Client.__del____init__ 실패 후에도 close()를 호출할 수 있어 AttributeError가 발생하는 문제를 방어합니다.

왜 이게 좋은가

  • 커넥션 풀 재사용으로 HTTP 연결 오버헤드가 감소합니다.
  • Vertex AI에서 서비스 계정 없이 API 키만으로 인증이 가능해져, 로컬 개발 환경 설정이 간소화됩니다.
  • Gateway 프로바이더에서 Client를 직접 생성하던 복잡한 코드가 http_clientbase_url 전달로 단순해졌습니다.

정리

  • HTTP 클라이언트는 재사용하라: 매 요청마다 새 클라이언트를 생성하면 TCP 핸드셰이크 비용이 누적됩니다.
  • 소멸자에서의 에러를 방어하라: __del__은 객체 초기화 실패 후에도 호출될 수 있으므로, 방어적 코딩이 필요합니다.
  • 인증 방식을 다양하게 지원하라: 서비스 계정, ADC, API 키 등 사용자 환경에 맞는 인증을 모두 지원해야 합니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글