mcp-run-python

from __future__ import annotations as _annotations import functools from asyncio import Lock from collections.abc import AsyncGenerator, Mapping from pathlib import Path from typing import Literal, overload import anyio.to_thread import httpx from typing_extensions import deprecated from pydantic_ai import ModelProfile from pydantic_ai.exceptions import UserError from pydantic_ai.models import cached_async_http_client from pydantic_ai.profiles.google import google_model_profile from pydantic_ai.providers import Provider try: import google.auth from google.auth.credentials import Credentials as BaseCredentials from google.auth.transport.requests import Request from google.oauth2.service_account import Credentials as ServiceAccountCredentials except ImportError as _import_error: raise ImportError( 'Please install the `google-auth` package to use the Google Vertex AI provider, ' 'you can use the `vertexai` optional group — `pip install "pydantic-ai-slim[vertexai]"`' ) from _import_error __all__ = ('GoogleVertexProvider',) @deprecated('`GoogleVertexProvider` is deprecated, use `GoogleProvider` with `GoogleModel` instead.') class GoogleVertexProvider(Provider[httpx.AsyncClient]): """Provider for Vertex AI API.""" @property def name(self) -> str: return 'google-vertex' @property def base_url(self) -> str: return ( f'https://{self.region}-aiplatform.googleapis.com/v1' f'/projects/{self.project_id}' f'/locations/{self.region}' f'/publishers/{self.model_publisher}/models/' ) @property def client(self) -> httpx.AsyncClient: return self._client def model_profile(self, model_name: str) -> ModelProfile | None: return google_model_profile(model_name) @overload def __init__( self, *, service_account_file: Path | str | None = None, project_id: str | None = None, region: VertexAiRegion = 'us-central1', model_publisher: str = 'google', http_client: httpx.AsyncClient | None = None, ) -> None: ... @overload def __init__( self, *, service_account_info: Mapping[str, str] | None = None, project_id: str | None = None, region: VertexAiRegion = 'us-central1', model_publisher: str = 'google', http_client: httpx.AsyncClient | None = None, ) -> None: ... def __init__( self, *, service_account_file: Path | str | None = None, service_account_info: Mapping[str, str] | None = None, project_id: str | None = None, region: VertexAiRegion = 'us-central1', model_publisher: str = 'google', http_client: httpx.AsyncClient | None = None, ) -> None: """Create a new Vertex AI provider. Args: service_account_file: Path to a service account file. If not provided, the service_account_info or default environment credentials will be used. service_account_info: The loaded service_account_file contents. If not provided, the service_account_file or default environment credentials will be used. project_id: The project ID to use, if not provided it will be taken from the credentials. region: The region to make requests to. model_publisher: The model publisher to use, I couldn't find a good list of available publishers, and from trial and error it seems non-google models don't work with the `generateContent` and `streamGenerateContent` functions, hence only `google` is currently supported. Please create an issue or PR if you know how to use other publishers. http_client: An existing `httpx.AsyncClient` to use for making HTTP requests. """ if service_account_file and service_account_info: raise ValueError('Only one of `service_account_file` or `service_account_info` can be provided.') self._client = http_client or cached_async_http_client(provider='google-vertex') self.service_account_file = service_account_file self.service_account_info = service_account_info self.project_id = project_id self.region = region self.model_publisher = model_publisher self._client.auth = _VertexAIAuth(service_account_file, service_account_info, project_id, region) self._client.base_url = self.base_url class _VertexAIAuth(httpx.Auth): """Auth class for Vertex AI API.""" _refresh_lock: Lock = Lock() credentials: BaseCredentials | ServiceAccountCredentials | None def __init__( self, service_account_file: Path | str | None = None, service_account_info: Mapping[str, str] | None = None, project_id: str | None = None, region: VertexAiRegion = 'us-central1', ) -> None: self.service_account_file = service_account_file self.service_account_info = service_account_info self.project_id = project_id self.region = region self.credentials = None async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]: if self.credentials is None: # pragma: no branch self.credentials = await self._get_credentials() if self.credentials.token is None: # type: ignore[reportUnknownMemberType] await self._refresh_token() request.headers['Authorization'] = f'Bearer {self.credentials.token}' # type: ignore[reportUnknownMemberType] # NOTE: This workaround is in place because we might get the project_id from the credentials. request.url = httpx.URL(str(request.url).replace('projects/None', f'projects/{self.project_id}')) response = yield request if response.status_code == 401: await self._refresh_token() request.headers['Authorization'] = f'Bearer {self.credentials.token}' # type: ignore[reportUnknownMemberType] yield request async def _get_credentials(self) -> BaseCredentials | ServiceAccountCredentials: if self.service_account_file is not None: creds = await _creds_from_file(self.service_account_file) assert creds.project_id is None or isinstance(creds.project_id, str) # type: ignore[reportUnknownMemberType] creds_project_id: str | None = creds.project_id creds_source = 'service account file' elif self.service_account_info is not None: creds = await _creds_from_info(self.service_account_info) assert creds.project_id is None or isinstance(creds.project_id, str) # type: ignore[reportUnknownMemberType] creds_project_id: str | None = creds.project_id creds_source = 'service account info' else: creds, creds_project_id = await _async_google_auth() creds_source = '`google.auth.default()`' if self.project_id is None: # pragma: no branch if creds_project_id is None: raise UserError(f'No project_id provided and none found in {creds_source}') # pragma: no cover self.project_id = creds_project_id return creds async def _refresh_token(self) -> str: # pragma: no cover async with self._refresh_lock: assert self.credentials is not None await anyio.to_thread.run_sync(self.credentials.refresh, Request()) # type: ignore[reportUnknownMemberType] assert isinstance(self.credentials.token, str), ( # type: ignore[reportUnknownMemberType] f'Expected token to be a string, got {self.credentials.token}' # type: ignore[reportUnknownMemberType] ) return self.credentials.token async def _async_google_auth() -> tuple[BaseCredentials, str | None]: return await anyio.to_thread.run_sync(google.auth.default, ['https://www.googleapis.com/auth/cloud-platform']) # type: ignore async def _creds_from_file(service_account_file: str | Path) -> ServiceAccountCredentials: service_account_credentials_from_file = functools.partial( ServiceAccountCredentials.from_service_account_file, # type: ignore[reportUnknownMemberType] scopes=['https://www.googleapis.com/auth/cloud-platform'], ) return await anyio.to_thread.run_sync(service_account_credentials_from_file, str(service_account_file)) async def _creds_from_info(service_account_info: Mapping[str, str]) -> ServiceAccountCredentials: service_account_credentials_from_string = functools.partial( ServiceAccountCredentials.from_service_account_info, # type: ignore[reportUnknownMemberType] scopes=['https://www.googleapis.com/auth/cloud-platform'], ) return await anyio.to_thread.run_sync(service_account_credentials_from_string, service_account_info) VertexAiRegion = Literal[ 'asia-east1', 'asia-east2', 'asia-northeast1', 'asia-northeast3', 'asia-south1', 'asia-southeast1', 'australia-southeast1', 'europe-central2', 'europe-north1', 'europe-southwest1', 'europe-west1', 'europe-west2', 'europe-west3', 'europe-west4', 'europe-west6', 'europe-west8', 'europe-west9', 'me-central1', 'me-central2', 'me-west1', 'northamerica-northeast1', 'southamerica-east1', 'us-central1', 'us-east1', 'us-east4', 'us-east5', 'us-south1', 'us-west1', 'us-west4', ] """Regions available for Vertex AI. More details [here](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations#genai-locations). """