Coverage for pydantic_ai_slim/pydantic_ai/providers/google_vertex.py: 96.04%

1from __future__ import annotations as _annotations (empty)

2 

3import functools (empty)

4from collections.abc import AsyncGenerator, Mapping (empty)

5from pathlib import Path (empty)

6from typing import Literal, overload (empty)

7 

8import anyio.to_thread (empty)

9import httpx (empty)

10 

11from pydantic_ai.exceptions import UserError (empty)

12from pydantic_ai.models import cached_async_http_client (empty)

13from pydantic_ai.providers import Provider (empty)

14 

15try: (empty)

16 import google.auth (empty)

17 from google.auth.credentials import Credentials as BaseCredentials (empty)

18 from google.auth.transport.requests import Request (empty)

19 from google.oauth2.service_account import Credentials as ServiceAccountCredentials (empty)

20except ImportError as _import_error: (empty)

21 raise ImportError( (empty)

22 'Please install the `google-auth` package to use the Google Vertex AI provider, ' 

23 'you can use the `vertexai` optional group — `pip install "pydantic-ai-slim[vertexai]"`' 

24 ) from _import_error 

25 

26 

27__all__ = ('GoogleVertexProvider',) (empty)

28 

29 

30class GoogleVertexProvider(Provider[httpx.AsyncClient]): (empty)

31 """Provider for Vertex AI API.""" 

32 

33 @property (empty)

34 def name(self) -> str: (empty)

35 return 'google-vertex' (empty)

36 

37 @property (empty)

38 def base_url(self) -> str: (empty)

39 return ( (empty)

40 f'https://{self.region}-aiplatform.googleapis.com/v1' 

41 f'/projects/{self.project_id}' 

42 f'/locations/{self.region}' 

43 f'/publishers/{self.model_publisher}/models/' 

44 ) 

45 

46 @property (empty)

47 def client(self) -> httpx.AsyncClient: (empty)

48 return self._client (empty)

49 

50 @overload (empty)

51 def __init__( (empty)

52 self, 

53 *, 

54 service_account_file: Path | str | None = None, (empty)

55 project_id: str | None = None, (empty)

56 region: VertexAiRegion = 'us-central1', (empty)

57 model_publisher: str = 'google', (empty)

58 http_client: httpx.AsyncClient | None = None, (empty)

59 ) -> None: ... (empty)

60 

61 @overload (empty)

62 def __init__( (empty)

63 self, 

64 *, 

65 service_account_info: Mapping[str, str] | None = None, (empty)

66 project_id: str | None = None, (empty)

67 region: VertexAiRegion = 'us-central1', (empty)

68 model_publisher: str = 'google', (empty)

69 http_client: httpx.AsyncClient | None = None, (empty)

70 ) -> None: ... (empty)

71 

72 def __init__( (empty)

73 self, 

74 *, 

75 service_account_file: Path | str | None = None, 

76 service_account_info: Mapping[str, str] | None = None, 

77 project_id: str | None = None, 

78 region: VertexAiRegion = 'us-central1', 

79 model_publisher: str = 'google', 

80 http_client: httpx.AsyncClient | None = None, 

81 ) -> None: 

82 """Create a new Vertex AI provider. 

83 

84 Args: 

85 service_account_file: Path to a service account file. 

86 If not provided, the service_account_info or default environment credentials will be used. 

87 service_account_info: The loaded service_account_file contents. 

88 If not provided, the service_account_file or default environment credentials will be used. 

89 project_id: The project ID to use, if not provided it will be taken from the credentials. 

90 region: The region to make requests to. 

91 model_publisher: The model publisher to use, I couldn't find a good list of available publishers, 

92 and from trial and error it seems non-google models don't work with the `generateContent` and 

93 `streamGenerateContent` functions, hence only `google` is currently supported. 

94 Please create an issue or PR if you know how to use other publishers. 

95 http_client: An existing `httpx.AsyncClient` to use for making HTTP requests. 

96 """ 

97 if service_account_file and service_account_info: (empty)

98 raise ValueError('Only one of `service_account_file` or `service_account_info` can be provided.') (empty)

99 

100 self._client = http_client or cached_async_http_client(provider='google-vertex') (empty)

101 self.service_account_file = service_account_file (empty)

102 self.service_account_info = service_account_info (empty)

103 self.project_id = project_id (empty)

104 self.region = region (empty)

105 self.model_publisher = model_publisher (empty)

106 

107 self._client.auth = _VertexAIAuth(service_account_file, service_account_info, project_id, region) (empty)

108 self._client.base_url = self.base_url (empty)

109 

110 

111class _VertexAIAuth(httpx.Auth): (empty)

112 """Auth class for Vertex AI API.""" 

113 

114 credentials: BaseCredentials | ServiceAccountCredentials | None (empty)

115 

116 def __init__( (empty)

117 self, 

118 service_account_file: Path | str | None = None, 

119 service_account_info: Mapping[str, str] | None = None, 

120 project_id: str | None = None, 

121 region: VertexAiRegion = 'us-central1', 

122 ) -> None: 

123 self.service_account_file = service_account_file (empty)

124 self.service_account_info = service_account_info (empty)

125 self.project_id = project_id (empty)

126 self.region = region (empty)

127 

128 self.credentials = None (empty)

129 

130 async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]: (empty)

131 if self.credentials is None: 131 ↛ 133line 131 didn't jump to line 133 because the condition on line 131 was always true(empty)

132 self.credentials = await self._get_credentials() (empty)

133 if self.credentials.token is None: # type: ignore[reportUnknownMemberType] (empty)

134 await self._refresh_token() (empty)

135 request.headers['Authorization'] = f'Bearer {self.credentials.token}' # type: ignore[reportUnknownMemberType] (empty)

136 # NOTE: This workaround is in place because we might get the project_id from the credentials. 

137 request.url = httpx.URL(str(request.url).replace('projects/None', f'projects/{self.project_id}')) (empty)

138 response = yield request (empty)

139 

140 if response.status_code == 401: (empty)

141 await self._refresh_token() (empty)

142 request.headers['Authorization'] = f'Bearer {self.credentials.token}' # type: ignore[reportUnknownMemberType] (empty)

143 yield request (empty)

144 

145 async def _get_credentials(self) -> BaseCredentials | ServiceAccountCredentials: (empty)

146 if self.service_account_file is not None: (empty)

147 creds = await _creds_from_file(self.service_account_file) (empty)

148 assert creds.project_id is None or isinstance(creds.project_id, str) # type: ignore[reportUnknownMemberType] (empty)

149 creds_project_id: str | None = creds.project_id (empty)

150 creds_source = 'service account file' (empty)

151 elif self.service_account_info is not None: (empty)

152 creds = await _creds_from_info(self.service_account_info) (empty)

153 assert creds.project_id is None or isinstance(creds.project_id, str) # type: ignore[reportUnknownMemberType] (empty)

154 creds_project_id: str | None = creds.project_id (empty)

155 creds_source = 'service account info' (empty)

156 else: 

157 creds, creds_project_id = await _async_google_auth() (empty)

158 creds_source = '`google.auth.default()`' (empty)

159 

160 if self.project_id is None: 160 ↛ 164line 160 didn't jump to line 164 because the condition on line 160 was always true(empty)

161 if creds_project_id is None: 161 ↛ 162line 161 didn't jump to line 162 because the condition on line 161 was never true(empty)

162 raise UserError(f'No project_id provided and none found in {creds_source}') 

163 self.project_id = creds_project_id (empty)

164 return creds (empty)

165 

166 async def _refresh_token(self) -> str: # pragma: no cover (empty)

167 assert self.credentials is not None 

168 await anyio.to_thread.run_sync(self.credentials.refresh, Request()) # type: ignore[reportUnknownMemberType] 

169 assert isinstance(self.credentials.token, str), f'Expected token to be a string, got {self.credentials.token}' # type: ignore[reportUnknownMemberType] 

170 return self.credentials.token 

171 

172 

173async def _async_google_auth() -> tuple[BaseCredentials, str | None]: (empty)

174 return await anyio.to_thread.run_sync(google.auth.default, ['https://www.googleapis.com/auth/cloud-platform']) # type: ignore (empty)

175 

176 

177async def _creds_from_file(service_account_file: str | Path) -> ServiceAccountCredentials: (empty)

178 service_account_credentials_from_file = functools.partial( (empty)

179 ServiceAccountCredentials.from_service_account_file, # type: ignore[reportUnknownMemberType] 

180 scopes=['https://www.googleapis.com/auth/cloud-platform'], 

181 ) 

182 return await anyio.to_thread.run_sync(service_account_credentials_from_file, str(service_account_file)) (empty)

183 

184 

185async def _creds_from_info(service_account_info: Mapping[str, str]) -> ServiceAccountCredentials: (empty)

186 service_account_credentials_from_string = functools.partial( (empty)

187 ServiceAccountCredentials.from_service_account_info, # type: ignore[reportUnknownMemberType] 

188 scopes=['https://www.googleapis.com/auth/cloud-platform'], 

189 ) 

190 return await anyio.to_thread.run_sync(service_account_credentials_from_string, service_account_info) (empty)

191 

192 

193VertexAiRegion = Literal[ (empty)

194 'asia-east1', 

195 'asia-east2', 

196 'asia-northeast1', 

197 'asia-northeast3', 

198 'asia-south1', 

199 'asia-southeast1', 

200 'australia-southeast1', 

201 'europe-central2', 

202 'europe-north1', 

203 'europe-southwest1', 

204 'europe-west1', 

205 'europe-west2', 

206 'europe-west3', 

207 'europe-west4', 

208 'europe-west6', 

209 'europe-west8', 

210 'europe-west9', 

211 'me-central1', 

212 'me-central2', 

213 'me-west1', 

214 'northamerica-northeast1', 

215 'southamerica-east1', 

216 'us-central1', 

217 'us-east1', 

218 'us-east4', 

219 'us-east5', 

220 'us-south1', 

221 'us-west1', 

222 'us-west4', 

223] 

224"""Regions available for Vertex AI. (empty)

225 

226More details [here](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations#genai-locations). 

227"""