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

1from __future__ import annotations as _annotations (empty)

2 

3from collections.abc import AsyncIterable, AsyncIterator (empty)

4from contextlib import asynccontextmanager (empty)

5from dataclasses import dataclass, field (empty)

6from datetime import datetime, timezone (empty)

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

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

9 

10from httpx import AsyncClient as AsyncHTTPClient (empty)

11from typing_extensions import assert_never (empty)

12 

13from .. import UnexpectedModelBehavior, _utils, usage (empty)

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

15from ..messages import ( (empty)

16 ModelMessage, 

17 ModelRequest, 

18 ModelResponse, 

19 ModelResponsePart, 

20 ModelResponseStreamEvent, 

21 RetryPromptPart, 

22 SystemPromptPart, 

23 TextPart, 

24 ToolCallPart, 

25 ToolReturnPart, 

26 UserPromptPart, 

27) 

28from ..settings import ModelSettings (empty)

29from ..tools import ToolDefinition (empty)

30from . import ( (empty)

31 AgentModel, 

32 Model, 

33 StreamedResponse, 

34 cached_async_http_client, 

35 check_allow_model_requests, 

36) 

37 

38try: (empty)

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

40 from anthropic.types import ( (empty)

41 Message as AnthropicMessage, 

42 MessageParam, 

43 MetadataParam, 

44 RawContentBlockDeltaEvent, 

45 RawContentBlockStartEvent, 

46 RawContentBlockStopEvent, 

47 RawMessageDeltaEvent, 

48 RawMessageStartEvent, 

49 RawMessageStopEvent, 

50 RawMessageStreamEvent, 

51 TextBlock, 

52 TextBlockParam, 

53 TextDelta, 

54 ToolChoiceParam, 

55 ToolParam, 

56 ToolResultBlockParam, 

57 ToolUseBlock, 

58 ToolUseBlockParam, 

59 ) 

60except ImportError as _import_error: (empty)

61 raise ImportError( (empty)

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

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

64 ) from _import_error 

65 

66LatestAnthropicModelNames = Literal[ (empty)

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

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

69 'claude-3-opus-latest', 

70] 

71"""Latest named Anthropic models.""" (empty)

72 

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

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

75 

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

77allow any name in the type hints. 

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

79""" 

80 

81 

82class AnthropicModelSettings(ModelSettings): (empty)

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

84 

85 anthropic_metadata: MetadataParam (empty)

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

87 

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

89 

90 

91@dataclass(init=False) (empty)

92class AnthropicModel(Model): (empty)

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

94 

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

96 

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

98 

99 !!! note 

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

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

102 """ 

103 

104 model_name: AnthropicModelName (empty)

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

106 

107 def __init__( (empty)

108 self, 

109 model_name: AnthropicModelName, 

110 *, 

111 api_key: str | None = None, 

112 anthropic_client: AsyncAnthropic | None = None, 

113 http_client: AsyncHTTPClient | None = None, 

114 ): 

115 """Initialize an Anthropic model. 

116 

117 Args: 

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

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

120 api_key: The API key to use for authentication, if not provided, the `ANTHROPIC_API_KEY` environment variable 

121 will be used if available. 

122 anthropic_client: An existing 

123 [`AsyncAnthropic`](https://github.com/anthropics/anthropic-sdk-python?tab=readme-ov-file#async-usage) 

124 client to use, if provided, `api_key` and `http_client` must be `None`. 

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

126 """ 

127 self.model_name = model_name (empty)

128 if anthropic_client is not None: (empty)

129 assert http_client is None, 'Cannot provide both `anthropic_client` and `http_client`' (empty)

130 assert api_key is None, 'Cannot provide both `anthropic_client` and `api_key`' (empty)

131 self.client = anthropic_client (empty)

132 elif http_client is not None: 132 ↛ 133line 132 didn't jump to line 133 because the condition on line 132 was never true(empty)

133 self.client = AsyncAnthropic(api_key=api_key, http_client=http_client) 

134 else: 

135 self.client = AsyncAnthropic(api_key=api_key, http_client=cached_async_http_client()) (empty)

136 

137 async def agent_model( (empty)

138 self, 

139 *, 

140 function_tools: list[ToolDefinition], 

141 allow_text_result: bool, 

142 result_tools: list[ToolDefinition], 

143 ) -> AgentModel: 

144 check_allow_model_requests() (empty)

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

146 if result_tools: (empty)

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

148 return AnthropicAgentModel( (empty)

149 self.client, 

150 self.model_name, 

151 allow_text_result, 

152 tools, 

153 ) 

154 

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

156 return f'anthropic:{self.model_name}' (empty)

157 

158 @staticmethod (empty)

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

160 return { (empty)

161 'name': f.name, 

162 'description': f.description, 

163 'input_schema': f.parameters_json_schema, 

164 } 

165 

166 

167@dataclass (empty)

168class AnthropicAgentModel(AgentModel): (empty)

169 """Implementation of `AgentModel` for Anthropic models.""" 

170 

171 client: AsyncAnthropic (empty)

172 model_name: AnthropicModelName (empty)

173 allow_text_result: bool (empty)

174 tools: list[ToolParam] (empty)

175 

176 async def request( (empty)

177 self, messages: list[ModelMessage], model_settings: ModelSettings | None 

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

179 response = await self._messages_create(messages, False, cast(AnthropicModelSettings, model_settings or {})) (empty)

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

181 

182 @asynccontextmanager (empty)

183 async def request_stream( (empty)

184 self, messages: list[ModelMessage], model_settings: ModelSettings | None 

185 ) -> AsyncIterator[StreamedResponse]: 

186 response = await self._messages_create(messages, True, cast(AnthropicModelSettings, model_settings or {})) (empty)

187 async with response: (empty)

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

189 

190 @overload (empty)

191 async def _messages_create( (empty)

192 self, messages: list[ModelMessage], stream: Literal[True], model_settings: AnthropicModelSettings (empty)

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

194 pass 

195 

196 @overload (empty)

197 async def _messages_create( (empty)

198 self, messages: list[ModelMessage], stream: Literal[False], model_settings: AnthropicModelSettings (empty)

199 ) -> AnthropicMessage: (empty)

200 pass 

201 

202 async def _messages_create( (empty)

203 self, messages: list[ModelMessage], stream: bool, model_settings: AnthropicModelSettings 

204 ) -> AnthropicMessage | AsyncStream[RawMessageStreamEvent]: 

205 # standalone function to make it easier to override 

206 tool_choice: ToolChoiceParam | None 

207 

208 if not self.tools: (empty)

209 tool_choice = None (empty)

210 else: 

211 if not self.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 = self._map_message(messages) (empty)

220 

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

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

223 system=system_prompt or NOT_GIVEN, 

224 messages=anthropic_messages, 

225 model=self.model_name, 

226 tools=self.tools or NOT_GIVEN, 

227 tool_choice=tool_choice or NOT_GIVEN, 

228 stream=stream, 

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

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

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

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

233 ) 

234 

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

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

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

238 for item in response.content: (empty)

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

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

241 else: 

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

243 items.append( (empty)

244 ToolCallPart( 

245 tool_name=item.name, 

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

247 tool_call_id=item.id, 

248 ) 

249 ) 

250 

251 return ModelResponse(items, model_name=self.model_name) (empty)

252 

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

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

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

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

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

258 

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

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

261 return AnthropicStreamedResponse(_model_name=self.model_name, _response=peekable_response, _timestamp=timestamp) (empty)

262 

263 @staticmethod (empty)

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

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

266 system_prompt: str = '' (empty)

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

268 for m in messages: (empty)

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

270 for part in m.parts: (empty)

271 if isinstance(part, SystemPromptPart): (empty)

272 system_prompt += part.content (empty)

273 elif isinstance(part, UserPromptPart): (empty)

274 anthropic_messages.append(MessageParam(role='user', content=part.content)) (empty)

275 elif isinstance(part, ToolReturnPart): (empty)

276 anthropic_messages.append( (empty)

277 MessageParam( 

278 role='user', 

279 content=[ 

280 ToolResultBlockParam( 

281 tool_use_id=_guard_tool_call_id(t=part, model_source='Anthropic'), 

282 type='tool_result', 

283 content=part.model_response_str(), 

284 is_error=False, 

285 ) 

286 ], 

287 ) 

288 ) 

289 elif isinstance(part, RetryPromptPart): 289 ↛ 270line 289 didn't jump to line 270 because the condition on line 289 was always true(empty)

290 if part.tool_name is None: 290 ↛ 291line 290 didn't jump to line 291 because the condition on line 290 was never true(empty)

291 anthropic_messages.append(MessageParam(role='user', content=part.model_response())) 

292 else: 

293 anthropic_messages.append( (empty)

294 MessageParam( 

295 role='user', 

296 content=[ 

297 ToolResultBlockParam( 

298 tool_use_id=_guard_tool_call_id(t=part, model_source='Anthropic'), 

299 type='tool_result', 

300 content=part.model_response(), 

301 is_error=True, 

302 ), 

303 ], 

304 ) 

305 ) 

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

307 content: list[TextBlockParam | ToolUseBlockParam] = [] (empty)

308 for item in m.parts: (empty)

309 if isinstance(item, TextPart): (empty)

310 content.append(TextBlockParam(text=item.content, type='text')) (empty)

311 else: 

312 assert isinstance(item, ToolCallPart) (empty)

313 content.append(_map_tool_call(item)) (empty)

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

315 else: 

316 assert_never(m) 

317 return system_prompt, anthropic_messages (empty)

318 

319 

320def _map_tool_call(t: ToolCallPart) -> ToolUseBlockParam: (empty)

321 return ToolUseBlockParam( (empty)

322 id=_guard_tool_call_id(t=t, model_source='Anthropic'), 

323 type='tool_use', 

324 name=t.tool_name, 

325 input=t.args_as_dict(), 

326 ) 

327 

328 

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

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

331 response_usage = message.usage (empty)

332 else: 

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

334 response_usage = message.message.usage (empty)

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

336 response_usage = message.usage (empty)

337 else: 

338 # No usage information provided in: 

339 # - RawMessageStopEvent 

340 # - RawContentBlockStartEvent 

341 # - RawContentBlockDeltaEvent 

342 # - RawContentBlockStopEvent 

343 response_usage = None (empty)

344 

345 if response_usage is None: (empty)

346 return usage.Usage() (empty)

347 

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

349 

350 return usage.Usage( (empty)

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

352 request_tokens=request_tokens, 

353 response_tokens=response_usage.output_tokens, 

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

355 ) 

356 

357 

358@dataclass (empty)

359class AnthropicStreamedResponse(StreamedResponse): (empty)

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

361 

362 _response: AsyncIterable[RawMessageStreamEvent] (empty)

363 _timestamp: datetime (empty)

364 

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

366 current_block: TextBlock | ToolUseBlock | None = None (empty)

367 current_json: str = '' (empty)

368 

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

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

371 

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

373 current_block = event.content_block (empty)

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

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

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

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

378 vendor_part_id=current_block.id, 

379 tool_name=current_block.name, 

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

381 tool_call_id=current_block.id, 

382 ) 

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

384 yield maybe_event (empty)

385 

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

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

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

389 elif ( 389 ↛ 369line 389 didn't jump to line 369(empty)

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

391 ): 

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

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

394 try: (empty)

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

396 current_json = '' (empty)

397 except JSONDecodeError: (empty)

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

399 continue (empty)

400 

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

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

403 vendor_part_id=current_block.id, 

404 tool_name='', 

405 args=parsed_args, 

406 tool_call_id=current_block.id, 

407 ) 

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

409 yield maybe_event (empty)

410 

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

412 current_block = None (empty)

413 

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

415 return self._timestamp (empty)