Coverage for pydantic_ai_slim/pydantic_ai/models/anthropic.py: 94.59%

1from __future__ import annotations as _annotations (empty)

2 

3import base64 (empty)

4import io (empty)

5from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator (empty)

6from contextlib import asynccontextmanager (empty)

7from dataclasses import dataclass, field (empty)

8from datetime import datetime, timezone (empty)

9from json import JSONDecodeError, loads as json_loads (empty)

10from typing import Any, Literal, Union, cast, overload (empty)

11 

12from anthropic.types import DocumentBlockParam (empty)

13from typing_extensions import assert_never (empty)

14 

15from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage (empty)

16from .._utils import guard_tool_call_id as _guard_tool_call_id (empty)

17from ..messages import ( (empty)

18 BinaryContent, 

19 DocumentUrl, 

20 ImageUrl, 

21 ModelMessage, 

22 ModelRequest, 

23 ModelResponse, 

24 ModelResponsePart, 

25 ModelResponseStreamEvent, 

26 RetryPromptPart, 

27 SystemPromptPart, 

28 TextPart, 

29 ToolCallPart, 

30 ToolReturnPart, 

31 UserPromptPart, 

32) 

33from ..providers import Provider, infer_provider (empty)

34from ..settings import ModelSettings (empty)

35from ..tools import ToolDefinition (empty)

36from . import Model, ModelRequestParameters, StreamedResponse, cached_async_http_client, check_allow_model_requests (empty)

37 

38try: (empty)

39 from anthropic import NOT_GIVEN, APIStatusError, AsyncAnthropic, AsyncStream (empty)

40 from anthropic.types import ( (empty)

41 Base64PDFSourceParam, 

42 ContentBlock, 

43 ImageBlockParam, 

44 Message as AnthropicMessage, 

45 MessageParam, 

46 MetadataParam, 

47 PlainTextSourceParam, 

48 RawContentBlockDeltaEvent, 

49 RawContentBlockStartEvent, 

50 RawContentBlockStopEvent, 

51 RawMessageDeltaEvent, 

52 RawMessageStartEvent, 

53 RawMessageStopEvent, 

54 RawMessageStreamEvent, 

55 TextBlock, 

56 TextBlockParam, 

57 TextDelta, 

58 ToolChoiceParam, 

59 ToolParam, 

60 ToolResultBlockParam, 

61 ToolUseBlock, 

62 ToolUseBlockParam, 

63 ) 

64except ImportError as _import_error: 

65 raise ImportError( 

66 'Please install `anthropic` to use the Anthropic model, ' 

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

68 ) from _import_error 

69 

70LatestAnthropicModelNames = Literal[ (empty)

71 'claude-3-7-sonnet-latest', 

72 'claude-3-5-haiku-latest', 

73 'claude-3-5-sonnet-latest', 

74 'claude-3-opus-latest', 

75] 

76"""Latest Anthropic models.""" (empty)

77 

78AnthropicModelName = Union[str, LatestAnthropicModelNames] (empty)

79"""Possible Anthropic model names. (empty)

80 

81Since Anthropic supports a variety of date-stamped models, we explicitly list the latest models but 

82allow any name in the type hints. 

83See [the Anthropic docs](https://docs.anthropic.com/en/docs/about-claude/models) for a full list. 

84""" 

85 

86 

87class AnthropicModelSettings(ModelSettings): (empty)

88 """Settings used for an Anthropic model request. 

89 

90 ALL FIELDS MUST BE `anthropic_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS. 

91 """ 

92 

93 anthropic_metadata: MetadataParam (empty)

94 """An object describing metadata about the request. (empty)

95 

96 Contains `user_id`, an external identifier for the user who is associated with the request.""" 

97 

98 

99@dataclass(init=False) (empty)

100class AnthropicModel(Model): (empty)

101 """A model that uses the Anthropic API. 

102 

103 Internally, this uses the [Anthropic Python client](https://github.com/anthropics/anthropic-sdk-python) to interact with the API. 

104 

105 Apart from `__init__`, all methods are private or match those of the base class. 

106 

107 !!! note 

108 The `AnthropicModel` class does not yet support streaming responses. 

109 We anticipate adding support for streaming responses in a near-term future release. 

110 """ 

111 

112 client: AsyncAnthropic = field(repr=False) (empty)

113 

114 _model_name: AnthropicModelName = field(repr=False) (empty)

115 _system: str = field(default='anthropic', repr=False) (empty)

116 

117 def __init__( (empty)

118 self, 

119 model_name: AnthropicModelName, 

120 *, 

121 provider: Literal['anthropic'] | Provider[AsyncAnthropic] = 'anthropic', 

122 ): 

123 """Initialize an Anthropic model. 

124 

125 Args: 

126 model_name: The name of the Anthropic model to use. List of model names available 

127 [here](https://docs.anthropic.com/en/docs/about-claude/models). 

128 provider: The provider to use for the Anthropic API. Can be either the string 'anthropic' or an 

129 instance of `Provider[AsyncAnthropic]`. If not provided, the other parameters will be used. 

130 """ 

131 self._model_name = model_name (empty)

132 

133 if isinstance(provider, str): (empty)

134 provider = infer_provider(provider) (empty)

135 self.client = provider.client (empty)

136 

137 @property (empty)

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

139 return str(self.client.base_url) (empty)

140 

141 async def request( (empty)

142 self, 

143 messages: list[ModelMessage], 

144 model_settings: ModelSettings | None, 

145 model_request_parameters: ModelRequestParameters, 

146 ) -> tuple[ModelResponse, usage.Usage]: 

147 check_allow_model_requests() (empty)

148 response = await self._messages_create( (empty)

149 messages, False, cast(AnthropicModelSettings, model_settings or {}), model_request_parameters 

150 ) 

151 return self._process_response(response), _map_usage(response) (empty)

152 

153 @asynccontextmanager (empty)

154 async def request_stream( (empty)

155 self, 

156 messages: list[ModelMessage], 

157 model_settings: ModelSettings | None, 

158 model_request_parameters: ModelRequestParameters, 

159 ) -> AsyncIterator[StreamedResponse]: 

160 check_allow_model_requests() (empty)

161 response = await self._messages_create( (empty)

162 messages, True, cast(AnthropicModelSettings, model_settings or {}), model_request_parameters 

163 ) 

164 async with response: (empty)

165 yield await self._process_streamed_response(response) (empty)

166 

167 @property (empty)

168 def model_name(self) -> AnthropicModelName: (empty)

169 """The model name.""" 

170 return self._model_name (empty)

171 

172 @property (empty)

173 def system(self) -> str: (empty)

174 """The system / model provider.""" 

175 return self._system (empty)

176 

177 @overload (empty)

178 async def _messages_create( (empty)

179 self, 

180 messages: list[ModelMessage], (empty)

181 stream: Literal[True], (empty)

182 model_settings: AnthropicModelSettings, (empty)

183 model_request_parameters: ModelRequestParameters, (empty)

184 ) -> AsyncStream[RawMessageStreamEvent]: (empty)

185 pass 

186 

187 @overload (empty)

188 async def _messages_create( (empty)

189 self, 

190 messages: list[ModelMessage], (empty)

191 stream: Literal[False], (empty)

192 model_settings: AnthropicModelSettings, (empty)

193 model_request_parameters: ModelRequestParameters, (empty)

194 ) -> AnthropicMessage: (empty)

195 pass 

196 

197 async def _messages_create( (empty)

198 self, 

199 messages: list[ModelMessage], 

200 stream: bool, 

201 model_settings: AnthropicModelSettings, 

202 model_request_parameters: ModelRequestParameters, 

203 ) -> AnthropicMessage | AsyncStream[RawMessageStreamEvent]: 

204 # standalone function to make it easier to override 

205 tools = self._get_tools(model_request_parameters) (empty)

206 tool_choice: ToolChoiceParam | None 

207 

208 if not tools: (empty)

209 tool_choice = None (empty)

210 else: 

211 if not model_request_parameters.allow_text_result: (empty)

212 tool_choice = {'type': 'any'} (empty)

213 else: 

214 tool_choice = {'type': 'auto'} (empty)

215 

216 if (allow_parallel_tool_calls := model_settings.get('parallel_tool_calls')) is not None: (empty)

217 tool_choice['disable_parallel_tool_use'] = not allow_parallel_tool_calls (empty)

218 

219 system_prompt, anthropic_messages = await self._map_message(messages) (empty)

220 

221 try: (empty)

222 return await self.client.messages.create( (empty)

223 max_tokens=model_settings.get('max_tokens', 1024), 

224 system=system_prompt or NOT_GIVEN, 

225 messages=anthropic_messages, 

226 model=self._model_name, 

227 tools=tools or NOT_GIVEN, 

228 tool_choice=tool_choice or NOT_GIVEN, 

229 stream=stream, 

230 temperature=model_settings.get('temperature', NOT_GIVEN), 

231 top_p=model_settings.get('top_p', NOT_GIVEN), 

232 timeout=model_settings.get('timeout', NOT_GIVEN), 

233 metadata=model_settings.get('anthropic_metadata', NOT_GIVEN), 

234 ) 

235 except APIStatusError as e: (empty)

236 if (status_code := e.status_code) >= 400: 236 ↛ 238line 236 didn't jump to line 238 because the condition on line 236 was always true(empty)

237 raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e (empty)

238 raise (empty)

239 

240 def _process_response(self, response: AnthropicMessage) -> ModelResponse: (empty)

241 """Process a non-streamed response, and prepare a message to return.""" 

242 items: list[ModelResponsePart] = [] (empty)

243 for item in response.content: (empty)

244 if isinstance(item, TextBlock): (empty)

245 items.append(TextPart(content=item.text)) (empty)

246 else: 

247 assert isinstance(item, ToolUseBlock), 'unexpected item type' (empty)

248 items.append( (empty)

249 ToolCallPart( 

250 tool_name=item.name, 

251 args=cast(dict[str, Any], item.input), 

252 tool_call_id=item.id, 

253 ) 

254 ) 

255 

256 return ModelResponse(items, model_name=response.model) (empty)

257 

258 async def _process_streamed_response(self, response: AsyncStream[RawMessageStreamEvent]) -> StreamedResponse: (empty)

259 peekable_response = _utils.PeekableAsyncStream(response) (empty)

260 first_chunk = await peekable_response.peek() (empty)

261 if isinstance(first_chunk, _utils.Unset): 261 ↛ 262line 261 didn't jump to line 262 because the condition on line 261 was never true(empty)

262 raise UnexpectedModelBehavior('Streamed response ended without content or tool calls') 

263 

264 # Since Anthropic doesn't provide a timestamp in the message, we'll use the current time 

265 timestamp = datetime.now(tz=timezone.utc) (empty)

266 return AnthropicStreamedResponse( (empty)

267 _model_name=self._model_name, _response=peekable_response, _timestamp=timestamp 

268 ) 

269 

270 def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[ToolParam]: (empty)

271 tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools] (empty)

272 if model_request_parameters.result_tools: (empty)

273 tools += [self._map_tool_definition(r) for r in model_request_parameters.result_tools] (empty)

274 return tools (empty)

275 

276 async def _map_message(self, messages: list[ModelMessage]) -> tuple[str, list[MessageParam]]: (empty)

277 """Just maps a `pydantic_ai.Message` to a `anthropic.types.MessageParam`.""" 

278 system_prompt: str = '' (empty)

279 anthropic_messages: list[MessageParam] = [] (empty)

280 for m in messages: (empty)

281 if isinstance(m, ModelRequest): (empty)

282 user_content_params: list[ (empty)

283 ToolResultBlockParam | TextBlockParam | ImageBlockParam | DocumentBlockParam 

284 ] = [] 

285 for request_part in m.parts: (empty)

286 if isinstance(request_part, SystemPromptPart): (empty)

287 system_prompt += request_part.content (empty)

288 elif isinstance(request_part, UserPromptPart): (empty)

289 async for content in self._map_user_prompt(request_part): (empty)

290 user_content_params.append(content) (empty)

291 elif isinstance(request_part, ToolReturnPart): (empty)

292 tool_result_block_param = ToolResultBlockParam( (empty)

293 tool_use_id=_guard_tool_call_id(t=request_part), 

294 type='tool_result', 

295 content=request_part.model_response_str(), 

296 is_error=False, 

297 ) 

298 user_content_params.append(tool_result_block_param) (empty)

299 elif isinstance(request_part, RetryPromptPart): 299 ↛ 285line 299 didn't jump to line 285 because the condition on line 299 was always true(empty)

300 if request_part.tool_name is None: 300 ↛ 301line 300 didn't jump to line 301 because the condition on line 300 was never true(empty)

301 retry_param = TextBlockParam(type='text', text=request_part.model_response()) 

302 else: 

303 retry_param = ToolResultBlockParam( (empty)

304 tool_use_id=_guard_tool_call_id(t=request_part), 

305 type='tool_result', 

306 content=request_part.model_response(), 

307 is_error=True, 

308 ) 

309 user_content_params.append(retry_param) (empty)

310 anthropic_messages.append(MessageParam(role='user', content=user_content_params)) (empty)

311 elif isinstance(m, ModelResponse): (empty)

312 assistant_content_params: list[TextBlockParam | ToolUseBlockParam] = [] (empty)

313 for response_part in m.parts: (empty)

314 if isinstance(response_part, TextPart): (empty)

315 assistant_content_params.append(TextBlockParam(text=response_part.content, type='text')) (empty)

316 else: 

317 tool_use_block_param = ToolUseBlockParam( (empty)

318 id=_guard_tool_call_id(t=response_part), 

319 type='tool_use', 

320 name=response_part.tool_name, 

321 input=response_part.args_as_dict(), 

322 ) 

323 assistant_content_params.append(tool_use_block_param) (empty)

324 anthropic_messages.append(MessageParam(role='assistant', content=assistant_content_params)) (empty)

325 else: 

326 assert_never(m) 

327 return system_prompt, anthropic_messages (empty)

328 

329 @staticmethod (empty)

330 async def _map_user_prompt( (empty)

331 part: UserPromptPart, 

332 ) -> AsyncGenerator[ImageBlockParam | TextBlockParam | DocumentBlockParam]: 

333 if isinstance(part.content, str): (empty)

334 yield TextBlockParam(text=part.content, type='text') (empty)

335 else: 

336 for item in part.content: (empty)

337 if isinstance(item, str): (empty)

338 yield TextBlockParam(text=item, type='text') (empty)

339 elif isinstance(item, BinaryContent): (empty)

340 if item.is_image: 340 ↛ 341line 340 didn't jump to line 341 because the condition on line 340 was never true(empty)

341 yield ImageBlockParam( 

342 source={'data': io.BytesIO(item.data), 'media_type': item.media_type, 'type': 'base64'}, # type: ignore 

343 type='image', 

344 ) 

345 elif item.media_type == 'application/pdf': (empty)

346 yield DocumentBlockParam( (empty)

347 source=Base64PDFSourceParam( 

348 data=io.BytesIO(item.data), 

349 media_type='application/pdf', 

350 type='base64', 

351 ), 

352 type='document', 

353 ) 

354 else: 

355 raise RuntimeError('Only images and PDFs are supported for binary content') (empty)

356 elif isinstance(item, ImageUrl): (empty)

357 try: (empty)

358 response = await cached_async_http_client().get(item.url) (empty)

359 response.raise_for_status() (empty)

360 yield ImageBlockParam( (empty)

361 source={ 

362 'data': io.BytesIO(response.content), 

363 'media_type': item.media_type, 

364 'type': 'base64', 

365 }, 

366 type='image', 

367 ) 

368 except ValueError: (empty)

369 # Download the file if can't find the mime type. 

370 client = cached_async_http_client() (empty)

371 response = await client.get(item.url, follow_redirects=True) (empty)

372 response.raise_for_status() (empty)

373 base64_encoded = base64.b64encode(response.content).decode('utf-8') (empty)

374 if (mime_type := response.headers['Content-Type']) in ( (empty)

375 'image/jpeg', 

376 'image/png', 

377 'image/gif', 

378 'image/webp', 

379 ): 

380 yield ImageBlockParam( (empty)

381 source={'data': base64_encoded, 'media_type': mime_type, 'type': 'base64'}, 

382 type='image', 

383 ) 

384 else: # pragma: no cover 

385 raise RuntimeError(f'Unsupported image type: {mime_type}') 

386 elif isinstance(item, DocumentUrl): 386 ↛ 406line 386 didn't jump to line 406 because the condition on line 386 was always true(empty)

387 response = await cached_async_http_client().get(item.url) (empty)

388 response.raise_for_status() (empty)

389 if item.media_type == 'application/pdf': (empty)

390 yield DocumentBlockParam( (empty)

391 source=Base64PDFSourceParam( 

392 data=io.BytesIO(response.content), 

393 media_type=item.media_type, 

394 type='base64', 

395 ), 

396 type='document', 

397 ) 

398 elif item.media_type == 'text/plain': (empty)

399 yield DocumentBlockParam( (empty)

400 source=PlainTextSourceParam(data=response.text, media_type=item.media_type, type='text'), 

401 type='document', 

402 ) 

403 else: # pragma: no cover 

404 raise RuntimeError(f'Unsupported media type: {item.media_type}') 

405 else: 

406 raise RuntimeError(f'Unsupported content type: {type(item)}') 

407 

408 @staticmethod (empty)

409 def _map_tool_definition(f: ToolDefinition) -> ToolParam: (empty)

410 return { (empty)

411 'name': f.name, 

412 'description': f.description, 

413 'input_schema': f.parameters_json_schema, 

414 } 

415 

416 

417def _map_usage(message: AnthropicMessage | RawMessageStreamEvent) -> usage.Usage: (empty)

418 if isinstance(message, AnthropicMessage): (empty)

419 response_usage = message.usage (empty)

420 else: 

421 if isinstance(message, RawMessageStartEvent): (empty)

422 response_usage = message.message.usage (empty)

423 elif isinstance(message, RawMessageDeltaEvent): (empty)

424 response_usage = message.usage (empty)

425 else: 

426 # No usage information provided in: 

427 # - RawMessageStopEvent 

428 # - RawContentBlockStartEvent 

429 # - RawContentBlockDeltaEvent 

430 # - RawContentBlockStopEvent 

431 response_usage = None (empty)

432 

433 if response_usage is None: (empty)

434 return usage.Usage() (empty)

435 

436 request_tokens = getattr(response_usage, 'input_tokens', None) (empty)

437 

438 return usage.Usage( (empty)

439 # Usage coming from the RawMessageDeltaEvent doesn't have input token data, hence this getattr 

440 request_tokens=request_tokens, 

441 response_tokens=response_usage.output_tokens, 

442 total_tokens=(request_tokens or 0) + response_usage.output_tokens, 

443 ) 

444 

445 

446@dataclass (empty)

447class AnthropicStreamedResponse(StreamedResponse): (empty)

448 """Implementation of `StreamedResponse` for Anthropic models.""" 

449 

450 _model_name: AnthropicModelName (empty)

451 _response: AsyncIterable[RawMessageStreamEvent] (empty)

452 _timestamp: datetime (empty)

453 

454 async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]: (empty)

455 current_block: ContentBlock | None = None (empty)

456 current_json: str = '' (empty)

457 

458 async for event in self._response: (empty)

459 self._usage += _map_usage(event) (empty)

460 

461 if isinstance(event, RawContentBlockStartEvent): (empty)

462 current_block = event.content_block (empty)

463 if isinstance(current_block, TextBlock) and current_block.text: (empty)

464 yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=current_block.text) (empty)

465 elif isinstance(current_block, ToolUseBlock): 465 ↛ 458line 465 didn't jump to line 458 because the condition on line 465 was always true(empty)

466 maybe_event = self._parts_manager.handle_tool_call_delta( (empty)

467 vendor_part_id=current_block.id, 

468 tool_name=current_block.name, 

469 args=cast(dict[str, Any], current_block.input), 

470 tool_call_id=current_block.id, 

471 ) 

472 if maybe_event is not None: 472 ↛ 458line 472 didn't jump to line 458 because the condition on line 472 was always true(empty)

473 yield maybe_event (empty)

474 

475 elif isinstance(event, RawContentBlockDeltaEvent): (empty)

476 if isinstance(event.delta, TextDelta): 476 ↛ 477line 476 didn't jump to line 477 because the condition on line 476 was never true(empty)

477 yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=event.delta.text) 

478 elif ( 478 ↛ 458line 478 didn't jump to line 458 because the condition on line 478 was always true(empty)

479 current_block and event.delta.type == 'input_json_delta' and isinstance(current_block, ToolUseBlock) 

480 ): 

481 # Try to parse the JSON immediately, otherwise cache the value for later. This handles 

482 # cases where the JSON is not currently valid but will be valid once we stream more tokens. 

483 try: (empty)

484 parsed_args = json_loads(current_json + event.delta.partial_json) (empty)

485 current_json = '' (empty)

486 except JSONDecodeError: (empty)

487 current_json += event.delta.partial_json (empty)

488 continue (empty)

489 

490 # For tool calls, we need to handle partial JSON updates 

491 maybe_event = self._parts_manager.handle_tool_call_delta( (empty)

492 vendor_part_id=current_block.id, 

493 tool_name='', 

494 args=parsed_args, 

495 tool_call_id=current_block.id, 

496 ) 

497 if maybe_event is not None: 497 ↛ 458line 497 didn't jump to line 458 because the condition on line 497 was always true(empty)

498 yield maybe_event (empty)

499 

500 elif isinstance(event, (RawContentBlockStopEvent, RawMessageStopEvent)): (empty)

501 current_block = None (empty)

502 

503 @property (empty)

504 def model_name(self) -> AnthropicModelName: (empty)

505 """Get the model name of the response.""" 

506 return self._model_name (empty)

507 

508 @property (empty)

509 def timestamp(self) -> datetime: (empty)

510 """Get the timestamp of the response.""" 

511 return self._timestamp (empty)