Coverage for pydantic_ai_slim/pydantic_ai/models/openai.py: 98.12%

1from __future__ import annotations as _annotations (empty)

2 

3import base64 (empty)

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

5from contextlib import asynccontextmanager (empty)

6from dataclasses import dataclass, field (empty)

7from datetime import datetime, timezone (empty)

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

9 

10from typing_extensions import assert_never (empty)

11 

12from pydantic_ai.providers import Provider, infer_provider (empty)

13 

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

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

16from ..messages import ( (empty)

17 AudioUrl, 

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 ..settings import ModelSettings (empty)

34from ..tools import ToolDefinition (empty)

35from . import ( (empty)

36 Model, 

37 ModelRequestParameters, 

38 StreamedResponse, 

39 cached_async_http_client, 

40 check_allow_model_requests, 

41) 

42 

43try: (empty)

44 from openai import NOT_GIVEN, APIStatusError, AsyncOpenAI, AsyncStream (empty)

45 from openai.types import ChatModel, chat (empty)

46 from openai.types.chat import ( (empty)

47 ChatCompletionChunk, 

48 ChatCompletionContentPartImageParam, 

49 ChatCompletionContentPartInputAudioParam, 

50 ChatCompletionContentPartParam, 

51 ChatCompletionContentPartTextParam, 

52 ) 

53 from openai.types.chat.chat_completion_content_part_image_param import ImageURL (empty)

54 from openai.types.chat.chat_completion_content_part_input_audio_param import InputAudio (empty)

55except ImportError as _import_error: (empty)

56 raise ImportError( (empty)

57 'Please install `openai` to use the OpenAI model, ' 

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

59 ) from _import_error 

60 

61OpenAIModelName = Union[str, ChatModel] (empty)

62""" (empty)

63Possible OpenAI model names. 

64 

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

66allow any name in the type hints. 

67See [the OpenAI docs](https://platform.openai.com/docs/models) for a full list. 

68 

69Using this more broad type for the model name instead of the ChatModel definition 

70allows this model to be used more easily with other model types (ie, Ollama, Deepseek). 

71""" 

72 

73OpenAISystemPromptRole = Literal['system', 'developer', 'user'] (empty)

74 

75 

76class OpenAIModelSettings(ModelSettings, total=False): (empty)

77 """Settings used for an OpenAI model request. 

78 

79 ALL FIELDS MUST BE `openai_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS. 

80 """ 

81 

82 openai_reasoning_effort: chat.ChatCompletionReasoningEffort (empty)

83 """ (empty)

84 Constrains effort on reasoning for [reasoning models](https://platform.openai.com/docs/guides/reasoning). 

85 Currently supported values are `low`, `medium`, and `high`. Reducing reasoning effort can 

86 result in faster responses and fewer tokens used on reasoning in a response. 

87 """ 

88 

89 openai_user: str (empty)

90 """A unique identifier representing the end-user, which can help OpenAI monitor and detect abuse. (empty)

91 

92 See [OpenAI's safety best practices](https://platform.openai.com/docs/guides/safety-best-practices#end-user-ids) for more details. 

93 """ 

94 

95 

96@dataclass(init=False) (empty)

97class OpenAIModel(Model): (empty)

98 """A model that uses the OpenAI API. 

99 

100 Internally, this uses the [OpenAI Python client](https://github.com/openai/openai-python) to interact with the API. 

101 

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

103 """ 

104 

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

106 system_prompt_role: OpenAISystemPromptRole | None = field(default=None) (empty)

107 

108 _model_name: OpenAIModelName = field(repr=False) (empty)

109 _system: str = field(default='openai', repr=False) (empty)

110 

111 def __init__( (empty)

112 self, 

113 model_name: OpenAIModelName, 

114 *, 

115 provider: Literal['openai', 'deepseek', 'azure'] | Provider[AsyncOpenAI] = 'openai', 

116 system_prompt_role: OpenAISystemPromptRole | None = None, 

117 ): 

118 """Initialize an OpenAI model. 

119 

120 Args: 

121 model_name: The name of the OpenAI model to use. List of model names available 

122 [here](https://github.com/openai/openai-python/blob/v1.54.3/src/openai/types/chat_model.py#L7) 

123 (Unfortunately, despite being ask to do so, OpenAI do not provide `.inv` files for their API). 

124 provider: The provider to use. Defaults to `'openai'`. 

125 system_prompt_role: The role to use for the system prompt message. If not provided, defaults to `'system'`. 

126 In the future, this may be inferred from the model name. 

127 """ 

128 self._model_name = model_name (empty)

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

130 provider = infer_provider(provider) (empty)

131 self.client = provider.client (empty)

132 self.system_prompt_role = system_prompt_role (empty)

133 

134 @property (empty)

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

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

137 

138 async def request( (empty)

139 self, 

140 messages: list[ModelMessage], 

141 model_settings: ModelSettings | None, 

142 model_request_parameters: ModelRequestParameters, 

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

144 check_allow_model_requests() (empty)

145 response = await self._completions_create( (empty)

146 messages, False, cast(OpenAIModelSettings, model_settings or {}), model_request_parameters 

147 ) 

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

149 

150 @asynccontextmanager (empty)

151 async def request_stream( (empty)

152 self, 

153 messages: list[ModelMessage], 

154 model_settings: ModelSettings | None, 

155 model_request_parameters: ModelRequestParameters, 

156 ) -> AsyncIterator[StreamedResponse]: 

157 check_allow_model_requests() (empty)

158 response = await self._completions_create( (empty)

159 messages, True, cast(OpenAIModelSettings, model_settings or {}), model_request_parameters 

160 ) 

161 async with response: (empty)

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

163 

164 @property (empty)

165 def model_name(self) -> OpenAIModelName: (empty)

166 """The model name.""" 

167 return self._model_name (empty)

168 

169 @property (empty)

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

171 """The system / model provider.""" 

172 return self._system (empty)

173 

174 @overload (empty)

175 async def _completions_create( (empty)

176 self, 

177 messages: list[ModelMessage], (empty)

178 stream: Literal[True], (empty)

179 model_settings: OpenAIModelSettings, (empty)

180 model_request_parameters: ModelRequestParameters, (empty)

181 ) -> AsyncStream[ChatCompletionChunk]: (empty)

182 pass 

183 

184 @overload (empty)

185 async def _completions_create( (empty)

186 self, 

187 messages: list[ModelMessage], (empty)

188 stream: Literal[False], (empty)

189 model_settings: OpenAIModelSettings, (empty)

190 model_request_parameters: ModelRequestParameters, (empty)

191 ) -> chat.ChatCompletion: (empty)

192 pass 

193 

194 async def _completions_create( (empty)

195 self, 

196 messages: list[ModelMessage], 

197 stream: bool, 

198 model_settings: OpenAIModelSettings, 

199 model_request_parameters: ModelRequestParameters, 

200 ) -> chat.ChatCompletion | AsyncStream[ChatCompletionChunk]: 

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

202 

203 # standalone function to make it easier to override 

204 if not tools: (empty)

205 tool_choice: Literal['none', 'required', 'auto'] | None = None (empty)

206 elif not model_request_parameters.allow_text_result: (empty)

207 tool_choice = 'required' (empty)

208 else: 

209 tool_choice = 'auto' (empty)

210 

211 openai_messages: list[chat.ChatCompletionMessageParam] = [] (empty)

212 for m in messages: (empty)

213 async for msg in self._map_message(m): (empty)

214 openai_messages.append(msg) (empty)

215 

216 try: (empty)

217 return await self.client.chat.completions.create( (empty)

218 model=self._model_name, 

219 messages=openai_messages, 

220 n=1, 

221 parallel_tool_calls=model_settings.get('parallel_tool_calls', NOT_GIVEN), 

222 tools=tools or NOT_GIVEN, 

223 tool_choice=tool_choice or NOT_GIVEN, 

224 stream=stream, 

225 stream_options={'include_usage': True} if stream else NOT_GIVEN, 

226 max_completion_tokens=model_settings.get('max_tokens', NOT_GIVEN), 

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

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

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

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

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

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

233 logit_bias=model_settings.get('logit_bias', NOT_GIVEN), 

234 reasoning_effort=model_settings.get('openai_reasoning_effort', NOT_GIVEN), 

235 user=model_settings.get('openai_user', NOT_GIVEN), 

236 ) 

237 except APIStatusError as e: (empty)

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

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

240 raise (empty)

241 

242 def _process_response(self, response: chat.ChatCompletion) -> ModelResponse: (empty)

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

244 timestamp = datetime.fromtimestamp(response.created, tz=timezone.utc) (empty)

245 choice = response.choices[0] (empty)

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

247 if choice.message.content is not None: (empty)

248 items.append(TextPart(choice.message.content)) (empty)

249 if choice.message.tool_calls is not None: (empty)

250 for c in choice.message.tool_calls: (empty)

251 items.append(ToolCallPart(c.function.name, c.function.arguments, c.id)) (empty)

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

253 

254 async def _process_streamed_response(self, response: AsyncStream[ChatCompletionChunk]) -> OpenAIStreamedResponse: (empty)

255 """Process a streamed response, and prepare a streaming response to return.""" 

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

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

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

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

260 

261 return OpenAIStreamedResponse( (empty)

262 _model_name=self._model_name, 

263 _response=peekable_response, 

264 _timestamp=datetime.fromtimestamp(first_chunk.created, tz=timezone.utc), 

265 ) 

266 

267 def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[chat.ChatCompletionToolParam]: (empty)

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

269 if model_request_parameters.result_tools: (empty)

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

271 return tools (empty)

272 

273 async def _map_message(self, message: ModelMessage) -> AsyncIterable[chat.ChatCompletionMessageParam]: (empty)

274 """Just maps a `pydantic_ai.Message` to a `openai.types.ChatCompletionMessageParam`.""" 

275 if isinstance(message, ModelRequest): (empty)

276 async for item in self._map_user_message(message): (empty)

277 yield item (empty)

278 elif isinstance(message, ModelResponse): (empty)

279 texts: list[str] = [] (empty)

280 tool_calls: list[chat.ChatCompletionMessageToolCallParam] = [] (empty)

281 for item in message.parts: (empty)

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

283 texts.append(item.content) (empty)

284 elif isinstance(item, ToolCallPart): (empty)

285 tool_calls.append(self._map_tool_call(item)) (empty)

286 else: 

287 assert_never(item) 

288 message_param = chat.ChatCompletionAssistantMessageParam(role='assistant') (empty)

289 if texts: (empty)

290 # Note: model responses from this model should only have one text item, so the following 

291 # shouldn't merge multiple texts into one unless you switch models between runs: 

292 message_param['content'] = '\n\n'.join(texts) (empty)

293 if tool_calls: (empty)

294 message_param['tool_calls'] = tool_calls (empty)

295 yield message_param (empty)

296 else: 

297 assert_never(message) 

298 

299 @staticmethod (empty)

300 def _map_tool_call(t: ToolCallPart) -> chat.ChatCompletionMessageToolCallParam: (empty)

301 return chat.ChatCompletionMessageToolCallParam( (empty)

302 id=_guard_tool_call_id(t=t), 

303 type='function', 

304 function={'name': t.tool_name, 'arguments': t.args_as_json_str()}, 

305 ) 

306 

307 @staticmethod (empty)

308 def _map_tool_definition(f: ToolDefinition) -> chat.ChatCompletionToolParam: (empty)

309 return { (empty)

310 'type': 'function', 

311 'function': { 

312 'name': f.name, 

313 'description': f.description, 

314 'parameters': f.parameters_json_schema, 

315 }, 

316 } 

317 

318 async def _map_user_message(self, message: ModelRequest) -> AsyncIterable[chat.ChatCompletionMessageParam]: (empty)

319 for part in message.parts: (empty)

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

321 if self.system_prompt_role == 'developer': (empty)

322 yield chat.ChatCompletionDeveloperMessageParam(role='developer', content=part.content) (empty)

323 elif self.system_prompt_role == 'user': (empty)

324 yield chat.ChatCompletionUserMessageParam(role='user', content=part.content) (empty)

325 else: 

326 yield chat.ChatCompletionSystemMessageParam(role='system', content=part.content) (empty)

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

328 yield await self._map_user_prompt(part) (empty)

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

330 yield chat.ChatCompletionToolMessageParam( (empty)

331 role='tool', 

332 tool_call_id=_guard_tool_call_id(t=part), 

333 content=part.model_response_str(), 

334 ) 

335 elif isinstance(part, RetryPromptPart): (empty)

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

337 yield chat.ChatCompletionUserMessageParam(role='user', content=part.model_response()) 

338 else: 

339 yield chat.ChatCompletionToolMessageParam( (empty)

340 role='tool', 

341 tool_call_id=_guard_tool_call_id(t=part), 

342 content=part.model_response(), 

343 ) 

344 else: 

345 assert_never(part) 

346 

347 @staticmethod (empty)

348 async def _map_user_prompt(part: UserPromptPart) -> chat.ChatCompletionUserMessageParam: (empty)

349 content: str | list[ChatCompletionContentPartParam] 

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

351 content = part.content (empty)

352 else: 

353 content = [] (empty)

354 for item in part.content: (empty)

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

356 content.append(ChatCompletionContentPartTextParam(text=item, type='text')) (empty)

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

358 image_url = ImageURL(url=item.url) (empty)

359 content.append(ChatCompletionContentPartImageParam(image_url=image_url, type='image_url')) (empty)

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

361 base64_encoded = base64.b64encode(item.data).decode('utf-8') (empty)

362 if item.is_image: (empty)

363 image_url = ImageURL(url=f'data:{item.media_type};base64,{base64_encoded}') (empty)

364 content.append(ChatCompletionContentPartImageParam(image_url=image_url, type='image_url')) (empty)

365 elif item.is_audio: (empty)

366 assert item.format in ('wav', 'mp3') (empty)

367 audio = InputAudio(data=base64_encoded, format=item.format) (empty)

368 content.append(ChatCompletionContentPartInputAudioParam(input_audio=audio, type='input_audio')) (empty)

369 else: # pragma: no cover 

370 raise RuntimeError(f'Unsupported binary content type: {item.media_type}') 

371 elif isinstance(item, AudioUrl): # pragma: no cover 

372 client = cached_async_http_client() 

373 response = await client.get(item.url) 

374 response.raise_for_status() 

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

376 audio = InputAudio(data=base64_encoded, format=response.headers.get('content-type')) 

377 content.append(ChatCompletionContentPartInputAudioParam(input_audio=audio, type='input_audio')) 

378 elif isinstance(item, DocumentUrl): # pragma: no cover 

379 raise NotImplementedError('DocumentUrl is not supported for OpenAI') 

380 # The following implementation should have worked, but it seems we have the following error: 

381 # pydantic_ai.exceptions.ModelHTTPError: status_code: 400, model_name: gpt-4o, body: 

382 # { 

383 # 'message': "Unknown parameter: 'messages[1].content[1].file.data'.", 

384 # 'type': 'invalid_request_error', 

385 # 'param': 'messages[1].content[1].file.data', 

386 # 'code': 'unknown_parameter' 

387 # } 

388 # 

389 # client = cached_async_http_client() 

390 # response = await client.get(item.url) 

391 # response.raise_for_status() 

392 # base64_encoded = base64.b64encode(response.content).decode('utf-8') 

393 # media_type = response.headers.get('content-type').split(';')[0] 

394 # file_data = f'data:{media_type};base64,{base64_encoded}' 

395 # file = File(file={'file_data': file_data, 'file_name': item.url, 'file_id': item.url}, type='file') 

396 # content.append(file) 

397 else: 

398 assert_never(item) 

399 return chat.ChatCompletionUserMessageParam(role='user', content=content) (empty)

400 

401 

402@dataclass (empty)

403class OpenAIStreamedResponse(StreamedResponse): (empty)

404 """Implementation of `StreamedResponse` for OpenAI models.""" 

405 

406 _model_name: OpenAIModelName (empty)

407 _response: AsyncIterable[ChatCompletionChunk] (empty)

408 _timestamp: datetime (empty)

409 

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

411 async for chunk in self._response: (empty)

412 self._usage += _map_usage(chunk) (empty)

413 

414 try: (empty)

415 choice = chunk.choices[0] (empty)

416 except IndexError: (empty)

417 continue (empty)

418 

419 # Handle the text part of the response 

420 content = choice.delta.content (empty)

421 if content is not None: (empty)

422 yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=content) (empty)

423 

424 for dtc in choice.delta.tool_calls or []: (empty)

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

426 vendor_part_id=dtc.index, 

427 tool_name=dtc.function and dtc.function.name, 

428 args=dtc.function and dtc.function.arguments, 

429 tool_call_id=dtc.id, 

430 ) 

431 if maybe_event is not None: (empty)

432 yield maybe_event (empty)

433 

434 @property (empty)

435 def model_name(self) -> OpenAIModelName: (empty)

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

437 return self._model_name (empty)

438 

439 @property (empty)

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

441 """Get the timestamp of the response.""" 

442 return self._timestamp (empty)

443 

444 

445def _map_usage(response: chat.ChatCompletion | ChatCompletionChunk) -> usage.Usage: (empty)

446 response_usage = response.usage (empty)

447 if response_usage is None: (empty)

448 return usage.Usage() (empty)

449 else: 

450 details: dict[str, int] = {} (empty)

451 if response_usage.completion_tokens_details is not None: (empty)

452 details.update(response_usage.completion_tokens_details.model_dump(exclude_none=True)) (empty)

453 if response_usage.prompt_tokens_details is not None: (empty)

454 details.update(response_usage.prompt_tokens_details.model_dump(exclude_none=True)) (empty)

455 return usage.Usage( (empty)

456 request_tokens=response_usage.prompt_tokens, 

457 response_tokens=response_usage.completion_tokens, 

458 total_tokens=response_usage.total_tokens, 

459 details=details, 

460 )