Coverage for haystack/components/generators/chat/openai.py: 98%

227 statements  

« prev     ^ index     » next       coverage.py v7.16.1, created at 2026-10-09 14:34 +0000

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import asyncio 

6import json 

7import os 

8from datetime import datetime 

9from typing import Any, ClassVar 

10 

11from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream 

12from openai.lib._pydantic import to_strict_json_schema 

13from openai.types.chat import ( 

14 ChatCompletion, 

15 ChatCompletionChunk, 

16 ChatCompletionMessage, 

17 ChatCompletionMessageCustomToolCall, 

18 ParsedChatCompletion, 

19 ParsedChatCompletionMessage, 

20) 

21from openai.types.chat.chat_completion import Choice 

22from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice 

23from pydantic import BaseModel 

24 

25from haystack import component, default_from_dict, default_to_dict, logging 

26from haystack.components.generators.utils import ( 

27 _convert_streaming_chunks_to_chat_message, 

28 _normalize_messages, 

29 _serialize_object, 

30) 

31from haystack.dataclasses import ( 

32 ChatMessage, 

33 ComponentInfo, 

34 FinishReason, 

35 StreamingCallbackT, 

36 StreamingChunk, 

37 SyncStreamingCallbackT, 

38 ToolCall, 

39 ToolCallDelta, 

40 select_streaming_callback, 

41) 

42from haystack.dataclasses.chat_message import _parse_openai_tool_call_arguments 

43from haystack.dataclasses.streaming_chunk import _invoke_streaming_callback 

44from haystack.tools import ( 

45 ToolsType, 

46 _check_duplicate_tool_names, 

47 deserialize_tools_or_toolset_inplace, 

48 flatten_tools_or_toolsets, 

49 serialize_tools_or_toolset, 

50) 

51from haystack.utils import Secret, deserialize_callable, serialize_callable 

52from haystack.utils.http_client import init_http_client 

53 

54logger = logging.getLogger(__name__) 

55 

56 

57@component 

58class OpenAIChatGenerator: 

59 """ 

60 Completes chats using OpenAI's large language models (LLMs). 

61 

62 It works with OpenAI's GPT models through the Chat Completions API and supports streaming responses 

63 from OpenAI API. It uses [ChatMessage](https://docs.haystack.deepset.ai/docs/chatmessage) 

64 format in input and output. 

65 

66 You can customize how the text is generated by passing parameters to the 

67 OpenAI API. Use the `**generation_kwargs` argument when you initialize 

68 the component or when you run it. Any parameter that works with 

69 `openai.ChatCompletion.create` will work here too. 

70 

71 For details on OpenAI API parameters, see 

72 [OpenAI documentation](https://platform.openai.com/docs/api-reference/chat). 

73 

74 ### Usage example 

75 ```python 

76 from haystack.components.generators.chat import OpenAIChatGenerator 

77 from haystack.dataclasses import ChatMessage 

78 

79 messages = [ChatMessage.from_user("What's Natural Language Processing?")] 

80 

81 client = OpenAIChatGenerator() 

82 response = client.run(messages) 

83 print(response) 

84 ``` 

85 Output: 

86 ``` 

87 {'replies': 

88 [ChatMessage(_role=<ChatRole.ASSISTANT: 'assistant'>, _content= 

89 [TextContent(text="Natural Language Processing (NLP) is a branch of artificial intelligence 

90 that focuses on enabling computers to understand, interpret, and generate human language in 

91 a way that is meaningful and useful.")], 

92 _name=None, 

93 _meta={'model': 'gpt-5-mini', 'index': 0, 'finish_reason': 'stop', 

94 'usage': {'prompt_tokens': 15, 'completion_tokens': 36, 'total_tokens': 51}}) 

95 ] 

96 } 

97 ``` 

98 """ 

99 

100 SUPPORTED_MODELS: ClassVar[list[str]] = [ 

101 "gpt-5-mini", 

102 "gpt-5-nano", 

103 "gpt-5", 

104 "gpt-5.1", 

105 "gpt-5.2", 

106 "gpt-5.2-pro", 

107 "gpt-5.4", 

108 "gpt-5-pro", 

109 "gpt-4.1", 

110 "gpt-4.1-mini", 

111 "gpt-4.1-nano", 

112 "gpt-4o", 

113 "gpt-4o-mini", 

114 "gpt-4-turbo", 

115 "gpt-4", 

116 "gpt-3.5-turbo", 

117 ] 

118 """A non-exhaustive list of chat models supported by this component. 

119 See https://developers.openai.com/api/docs/models for the full list and snapshot IDs.""" 

120 

121 def __init__( 

122 self, 

123 api_key: Secret = Secret.from_env_var("OPENAI_API_KEY"), 

124 model: str = "gpt-5-mini", 

125 streaming_callback: StreamingCallbackT | None = None, 

126 api_base_url: str | None = None, 

127 organization: str | None = None, 

128 generation_kwargs: dict[str, Any] | None = None, 

129 timeout: float | None = None, 

130 max_retries: int | None = None, 

131 tools: ToolsType | None = None, 

132 tools_strict: bool = False, 

133 http_client_kwargs: dict[str, Any] | None = None, 

134 ) -> None: 

135 """ 

136 Creates an instance of OpenAIChatGenerator. Unless specified otherwise in `model`, uses OpenAI's gpt-5-mini 

137 

138 Before initializing the component, you can set the 'OPENAI_TIMEOUT' and 'OPENAI_MAX_RETRIES' 

139 environment variables to override the `timeout` and `max_retries` parameters respectively 

140 in the OpenAI client. 

141 

142 :param api_key: The OpenAI API key. 

143 You can set it with an environment variable `OPENAI_API_KEY`, or pass with this parameter 

144 during initialization. 

145 :param model: The name of the model to use. 

146 :param streaming_callback: A callback function that is called when a new token is received from the stream. 

147 The callback function accepts [StreamingChunk](https://docs.haystack.deepset.ai/docs/data-classes#streamingchunk) 

148 as an argument. 

149 :param api_base_url: An optional base URL. 

150 :param organization: Your organization ID, defaults to `None`. See 

151 [production best practices](https://platform.openai.com/docs/guides/production-best-practices/setting-up-your-organization). 

152 :param generation_kwargs: Other parameters to use for the model. These parameters are sent directly to 

153 the OpenAI endpoint. See OpenAI [documentation](https://platform.openai.com/docs/api-reference/chat) for 

154 more details. 

155 Some of the supported parameters: 

156 - `max_completion_tokens`: An upper bound for the number of tokens that can be generated for a completion, 

157 including visible output tokens and reasoning tokens. 

158 - `temperature`: What sampling temperature to use. Higher values mean the model will take more risks. 

159 Try 0.9 for more creative applications and 0 (argmax sampling) for ones with a well-defined answer. 

160 - `top_p`: An alternative to sampling with temperature, called nucleus sampling, where the model 

161 considers the results of the tokens with top_p probability mass. For example, 0.1 means only the tokens 

162 comprising the top 10% probability mass are considered. 

163 - `n`: How many completions to generate for each prompt. For example, if the LLM gets 3 prompts and n is 2, 

164 it will generate two completions for each of the three prompts, ending up with 6 completions in total. 

165 - `stop`: One or more sequences after which the LLM should stop generating tokens. 

166 - `presence_penalty`: What penalty to apply if a token is already present at all. Bigger values mean 

167 the model will be less likely to repeat the same token in the text. 

168 - `frequency_penalty`: What penalty to apply if a token has already been generated in the text. 

169 Bigger values mean the model will be less likely to repeat the same token in the text. 

170 - `logit_bias`: Add a logit bias to specific tokens. The keys of the dictionary are tokens, and the 

171 values are the bias to add to that token. 

172 - `response_format`: A JSON schema or a Pydantic model that enforces the structure of the model's response. 

173 If provided, the output will always be validated against this 

174 format (unless the model returns a tool call). 

175 For details, see the [OpenAI Structured Outputs documentation](https://platform.openai.com/docs/guides/structured-outputs). 

176 Notes: 

177 - This parameter accepts Pydantic models and JSON schemas for latest models starting from GPT-4o. 

178 Older models only support basic version of structured outputs through `{"type": "json_object"}`. 

179 For detailed information on JSON mode, see the [OpenAI Structured Outputs documentation](https://platform.openai.com/docs/guides/structured-outputs#json-mode). 

180 - For structured outputs with streaming, 

181 the `response_format` must be a JSON schema and not a Pydantic model. 

182 :param timeout: 

183 Timeout for OpenAI client calls. If not set, it defaults to either the 

184 `OPENAI_TIMEOUT` environment variable, or 30 seconds. 

185 :param max_retries: 

186 Maximum number of retries to contact OpenAI after an internal error. 

187 If not set, it defaults to either the `OPENAI_MAX_RETRIES` environment variable, or set to 5. 

188 :param tools: 

189 A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls. 

190 :param tools_strict: 

191 Whether to enable strict schema adherence for tool calls. If set to `True`, the model will follow exactly 

192 the schema provided in the `parameters` field of the tool definition, but this may increase latency. 

193 :param http_client_kwargs: 

194 A dictionary of keyword arguments to configure a custom `httpx.Client`or `httpx.AsyncClient`. 

195 For more information, see the [HTTPX documentation](https://www.python-httpx.org/api/#client). 

196 

197 """ 

198 self.api_key = api_key 

199 self.model = model 

200 self.generation_kwargs = generation_kwargs or {} 

201 self.streaming_callback = streaming_callback 

202 self.api_base_url = api_base_url 

203 self.organization = organization 

204 self.timeout = timeout 

205 self.max_retries = max_retries 

206 self.tools = tools # Store tools as-is, whether it's a list or a Toolset 

207 self.tools_strict = tools_strict 

208 self.http_client_kwargs = http_client_kwargs 

209 # Check for duplicate tool names 

210 _check_duplicate_tool_names(flatten_tools_or_toolsets(self.tools)) 

211 

212 self.client: OpenAI | None = None 

213 self.async_client: AsyncOpenAI | None = None 

214 

215 def _client_kwargs(self) -> dict[str, Any]: 

216 timeout = self.timeout if self.timeout is not None else float(os.environ.get("OPENAI_TIMEOUT", "30.0")) 

217 max_retries = ( 

218 self.max_retries if self.max_retries is not None else int(os.environ.get("OPENAI_MAX_RETRIES", "5")) 

219 ) 

220 return { 

221 "api_key": self.api_key.resolve_value(), 

222 "organization": self.organization, 

223 "base_url": self.api_base_url, 

224 "timeout": timeout, 

225 "max_retries": max_retries, 

226 } 

227 

228 def warm_up(self) -> None: 

229 """ 

230 Initialize the synchronous OpenAI client. 

231 """ 

232 if self.client is None: 

233 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. 

234 # https://github.com/openai/openai-python/blob/main/httpx2.md 

235 http_client = init_http_client(self.http_client_kwargs, async_client=False) 

236 self.client = OpenAI( 

237 http_client=http_client, # type: ignore[arg-type] 

238 **self._client_kwargs(), 

239 ) 

240 

241 async def warm_up_async(self) -> None: # noqa: RUF029 

242 """ 

243 Initialize the asynchronous OpenAI client on the serving event loop. 

244 """ 

245 if self.async_client is None: 

246 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. 

247 # https://github.com/openai/openai-python/blob/main/httpx2.md 

248 http_client = init_http_client(self.http_client_kwargs, async_client=True) 

249 self.async_client = AsyncOpenAI( 

250 http_client=http_client, # type: ignore[arg-type] 

251 **self._client_kwargs(), 

252 ) 

253 

254 def close(self) -> None: 

255 """ 

256 Releases the synchronous OpenAI client. 

257 """ 

258 if self.client is not None: 

259 self.client.close() 

260 self.client = None 

261 

262 async def close_async(self) -> None: 

263 """ 

264 Releases the asynchronous OpenAI client. 

265 """ 

266 if self.async_client is not None: 

267 await self.async_client.close() 

268 self.async_client = None 

269 

270 def _get_telemetry_data(self) -> dict[str, Any]: 

271 """ 

272 Data that is sent to Posthog for usage analytics. 

273 """ 

274 return {"model": self.model} 

275 

276 def to_dict(self) -> dict[str, Any]: 

277 """ 

278 Serialize this component to a dictionary. 

279 

280 :returns: 

281 The serialized component as a dictionary. 

282 """ 

283 callback_name = serialize_callable(self.streaming_callback) if self.streaming_callback else None 

284 generation_kwargs = self.generation_kwargs.copy() 

285 response_format = generation_kwargs.get("response_format") 

286 

287 # If the response format is a Pydantic model, it's converted to openai's json schema format 

288 # If it's already a json schema, it's left as is 

289 if response_format and isinstance(response_format, type) and issubclass(response_format, BaseModel): 

290 json_schema = { 

291 "type": "json_schema", 

292 "json_schema": { 

293 "name": response_format.__name__, 

294 "strict": True, 

295 "schema": to_strict_json_schema(response_format), 

296 }, 

297 } 

298 generation_kwargs["response_format"] = json_schema 

299 

300 return default_to_dict( 

301 self, 

302 model=self.model, 

303 streaming_callback=callback_name, 

304 api_base_url=self.api_base_url, 

305 organization=self.organization, 

306 generation_kwargs=generation_kwargs, 

307 api_key=self.api_key, 

308 timeout=self.timeout, 

309 max_retries=self.max_retries, 

310 tools=serialize_tools_or_toolset(self.tools), 

311 tools_strict=self.tools_strict, 

312 http_client_kwargs=self.http_client_kwargs, 

313 ) 

314 

315 @classmethod 

316 def from_dict(cls, data: dict[str, Any]) -> "OpenAIChatGenerator": 

317 """ 

318 Deserialize this component from a dictionary. 

319 

320 :param data: The dictionary representation of this component. 

321 :returns: 

322 The deserialized component instance. 

323 """ 

324 deserialize_tools_or_toolset_inplace(data["init_parameters"], key="tools") 

325 init_params = data.get("init_parameters", {}) 

326 serialized_callback_handler = init_params.get("streaming_callback") 

327 

328 if serialized_callback_handler: 

329 data["init_parameters"]["streaming_callback"] = deserialize_callable(serialized_callback_handler) 

330 return default_from_dict(cls, data) 

331 

332 @component.output_types(replies=list[ChatMessage]) 

333 def run( 

334 self, 

335 messages: list[ChatMessage] | str, 

336 streaming_callback: StreamingCallbackT | None = None, 

337 generation_kwargs: dict[str, Any] | None = None, 

338 *, 

339 tools: ToolsType | None = None, 

340 tools_strict: bool | None = None, 

341 ) -> dict[str, list[ChatMessage]]: 

342 """ 

343 Invokes chat completion based on the provided messages and generation parameters. 

344 

345 :param messages: 

346 A list of ChatMessage instances representing the input messages. If a string is provided, it is converted 

347 to a list containing a ChatMessage with user role. 

348 :param streaming_callback: 

349 A callback function that is called when a new token is received from the stream. 

350 :param generation_kwargs: 

351 Additional keyword arguments for text generation. These are merged per key with the 

352 `generation_kwargs` passed at initialization: keys provided here take precedence, keys set 

353 only at initialization are kept. 

354 For details on OpenAI API parameters, see [OpenAI documentation](https://platform.openai.com/docs/api-reference/chat/create). 

355 :param tools: 

356 A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls. 

357 If set, it will override the `tools` parameter provided during initialization. 

358 :param tools_strict: 

359 Whether to enable strict schema adherence for tool calls. If set to `True`, the model will follow exactly 

360 the schema provided in the `parameters` field of the tool definition, but this may increase latency. 

361 If set, it will override the `tools_strict` parameter set during component initialization. 

362 

363 :returns: 

364 A dictionary with the following key: 

365 - `replies`: A list containing the generated responses as ChatMessage instances. 

366 """ 

367 self.warm_up() 

368 

369 messages = _normalize_messages(messages) 

370 

371 if len(messages) == 0: 

372 return {"replies": []} 

373 

374 streaming_callback = select_streaming_callback( 

375 init_callback=self.streaming_callback, runtime_callback=streaming_callback, requires_async=False 

376 ) 

377 chat_completion: Stream[ChatCompletionChunk] | ChatCompletion | ParsedChatCompletion 

378 

379 api_args = self._prepare_api_call( 

380 messages=messages, 

381 streaming_callback=streaming_callback, 

382 generation_kwargs=generation_kwargs, 

383 tools=tools, 

384 tools_strict=tools_strict, 

385 ) 

386 openai_endpoint = api_args.pop("openai_endpoint") 

387 assert self.client is not None # mypy: client is built by warm_up above 

388 openai_endpoint_method = getattr(self.client.chat.completions, openai_endpoint) 

389 chat_completion = openai_endpoint_method(**api_args) 

390 

391 if streaming_callback is not None: 

392 completions = self._handle_stream_response( 

393 # we cannot check isinstance(chat_completion, Stream) because some observability tools wrap Stream 

394 # and return a different type. See https://github.com/deepset-ai/haystack/issues/9014. 

395 chat_completion, # type: ignore 

396 streaming_callback, 

397 ) 

398 

399 else: 

400 assert isinstance(chat_completion, ChatCompletion), "Unexpected response type for non-streaming request." 

401 completions = [ 

402 _convert_chat_completion_to_chat_message(chat_completion, choice) for choice in chat_completion.choices 

403 ] 

404 

405 # before returning, do post-processing of the completions 

406 for message in completions: 

407 _check_finish_reason(message.meta) 

408 

409 return {"replies": completions} 

410 

411 @component.output_types(replies=list[ChatMessage]) 

412 async def run_async( 

413 self, 

414 messages: list[ChatMessage] | str, 

415 streaming_callback: StreamingCallbackT | None = None, 

416 generation_kwargs: dict[str, Any] | None = None, 

417 *, 

418 tools: ToolsType | None = None, 

419 tools_strict: bool | None = None, 

420 ) -> dict[str, list[ChatMessage]]: 

421 """ 

422 Asynchronously invokes chat completion based on the provided messages and generation parameters. 

423 

424 This is the asynchronous version of the `run` method. It has the same parameters and return values 

425 but can be used with `await` in async code. 

426 

427 :param messages: 

428 A list of ChatMessage instances representing the input messages. If a string is provided, it is converted 

429 to a list containing a ChatMessage with user role. 

430 :param streaming_callback: 

431 A callback function that is called when a new token is received from the stream. Async callbacks are 

432 preferred; a sync callback is accepted but will run synchronously on the event loop and may block it. 

433 :param generation_kwargs: 

434 Additional keyword arguments for text generation. These are merged per key with the 

435 `generation_kwargs` passed at initialization: keys provided here take precedence, keys set 

436 only at initialization are kept. 

437 For details on OpenAI API parameters, see [OpenAI documentation](https://platform.openai.com/docs/api-reference/chat/create). 

438 :param tools: A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls. 

439 If set, it will override the `tools` parameter provided during initialization. 

440 :param tools_strict: 

441 Whether to enable strict schema adherence for tool calls. If set to `True`, the model will follow exactly 

442 the schema provided in the `parameters` field of the tool definition, but this may increase latency. 

443 If set, it will override the `tools_strict` parameter set during component initialization. 

444 

445 :returns: 

446 A dictionary with the following key: 

447 - `replies`: A list containing the generated responses as ChatMessage instances. 

448 """ 

449 await self.warm_up_async() 

450 

451 messages = _normalize_messages(messages) 

452 

453 # validate and select the streaming callback 

454 streaming_callback = select_streaming_callback( 

455 init_callback=self.streaming_callback, runtime_callback=streaming_callback, requires_async=True 

456 ) 

457 chat_completion: AsyncStream[ChatCompletionChunk] | ChatCompletion | ParsedChatCompletion 

458 

459 if len(messages) == 0: 

460 return {"replies": []} 

461 

462 api_args = self._prepare_api_call( 

463 messages=messages, 

464 streaming_callback=streaming_callback, 

465 generation_kwargs=generation_kwargs, 

466 tools=tools, 

467 tools_strict=tools_strict, 

468 ) 

469 

470 openai_endpoint = api_args.pop("openai_endpoint") 

471 assert self.async_client is not None # mypy: async_client is built by warm_up_async above 

472 openai_endpoint_method = getattr(self.async_client.chat.completions, openai_endpoint) 

473 chat_completion = await openai_endpoint_method(**api_args) 

474 

475 if streaming_callback is not None: 

476 completions = await self._handle_async_stream_response( 

477 # we cannot check isinstance(chat_completion, AsyncStream) because some observability tools wrap 

478 # AsyncStream and return a different type. See https://github.com/deepset-ai/haystack/issues/9014. 

479 chat_completion, # type: ignore 

480 streaming_callback, 

481 ) 

482 

483 else: 

484 assert isinstance(chat_completion, ChatCompletion), "Unexpected response type for non-streaming request." 

485 completions = [ 

486 _convert_chat_completion_to_chat_message(chat_completion, choice) for choice in chat_completion.choices 

487 ] 

488 

489 # before returning, do post-processing of the completions 

490 for message in completions: 

491 _check_finish_reason(message.meta) 

492 

493 return {"replies": completions} 

494 

495 def _prepare_api_call( # noqa: PLR0913 

496 self, 

497 *, 

498 messages: list[ChatMessage], 

499 streaming_callback: StreamingCallbackT | None = None, 

500 generation_kwargs: dict[str, Any] | None = None, 

501 tools: ToolsType | None = None, 

502 tools_strict: bool | None = None, 

503 ) -> dict[str, Any]: 

504 # update generation kwargs by merging with the generation kwargs passed to the run method 

505 generation_kwargs = {**self.generation_kwargs, **(generation_kwargs or {})} 

506 

507 is_streaming = streaming_callback is not None 

508 num_responses = generation_kwargs.pop("n", 1) 

509 

510 if is_streaming and num_responses > 1: 

511 raise ValueError("Cannot stream multiple responses, please set n=1.") 

512 response_format = generation_kwargs.pop("response_format", None) 

513 

514 # adapt ChatMessage(s) to the format expected by the OpenAI API 

515 openai_formatted_messages = [message.to_openai_dict_format() for message in messages] 

516 

517 resolved_tools = tools if tools is not None else self.tools 

518 flattened_tools = flatten_tools_or_toolsets(resolved_tools) 

519 tools_strict = tools_strict if tools_strict is not None else self.tools_strict 

520 _check_duplicate_tool_names(flattened_tools) 

521 

522 openai_tools = {} 

523 if flattened_tools: 

524 tool_definitions = [] 

525 for t in flattened_tools: 

526 function_spec = {**t.tool_spec} 

527 if tools_strict: 

528 function_spec["strict"] = True 

529 function_spec["parameters"] = _make_schema_strict(function_spec["parameters"]) 

530 tool_definitions.append({"type": "function", "function": function_spec}) 

531 openai_tools = {"tools": tool_definitions} 

532 

533 base_args = { 

534 "model": self.model, 

535 "messages": openai_formatted_messages, 

536 "n": num_responses, 

537 **openai_tools, 

538 **generation_kwargs, 

539 } 

540 

541 if response_format and not is_streaming: 

542 # for structured outputs without streaming, we use openai's parse endpoint 

543 # Note: `stream` cannot be passed to chat.completions.parse 

544 # we pass a key `openai_endpoint` as a hint to the run method to use the parse endpoint 

545 # this key will be removed before the API call is made 

546 return {**base_args, "response_format": response_format, "openai_endpoint": "parse"} 

547 

548 # for structured outputs with streaming, we use openai's create endpoint 

549 # we pass a key `openai_endpoint` as a hint to the run method to use the create endpoint 

550 # this key will be removed before the API call is made 

551 final_args = {**base_args, "stream": is_streaming, "openai_endpoint": "create"} 

552 

553 # We only set the response_format parameter if it's not None since None is not a valid value in the API. 

554 if response_format: 

555 final_args["response_format"] = response_format 

556 return final_args 

557 

558 def _handle_stream_response(self, chat_completion: Stream, callback: SyncStreamingCallbackT) -> list[ChatMessage]: 

559 component_info = ComponentInfo.from_component(self) 

560 chunks: list[StreamingChunk] = [] 

561 for chunk in chat_completion: 

562 assert len(chunk.choices) <= 1, "Streaming responses should have at most one choice." 

563 chunk_delta = _convert_chat_completion_chunk_to_streaming_chunk( 

564 chunk=chunk, previous_chunks=chunks, component_info=component_info 

565 ) 

566 chunks.append(chunk_delta) 

567 callback(chunk_delta) 

568 return [_convert_streaming_chunks_to_chat_message(chunks=chunks)] 

569 

570 async def _handle_async_stream_response( 

571 self, chat_completion: AsyncStream, callback: StreamingCallbackT 

572 ) -> list[ChatMessage]: 

573 component_info = ComponentInfo.from_component(self) 

574 chunks: list[StreamingChunk] = [] 

575 try: 

576 async for chunk in chat_completion: 

577 assert len(chunk.choices) <= 1, "Streaming responses should have at most one choice." 

578 chunk_delta = _convert_chat_completion_chunk_to_streaming_chunk( 

579 chunk=chunk, previous_chunks=chunks, component_info=component_info 

580 ) 

581 chunks.append(chunk_delta) 

582 await _invoke_streaming_callback(callback, chunk_delta) 

583 

584 except asyncio.CancelledError: 

585 await asyncio.shield(chat_completion.close()) 

586 # close the stream when task is cancelled 

587 # asyncio.shield ensures the close operation completes 

588 # https://docs.python.org/3/library/asyncio-task.html#shielding-from-cancellation 

589 raise # Re-raise to propagate cancellation 

590 

591 return [_convert_streaming_chunks_to_chat_message(chunks=chunks)] 

592 

593 

594def _make_schema_strict(schema: dict[str, Any]) -> dict[str, Any]: 

595 """ 

596 Recursively transform a JSON schema to be OpenAI strict-mode compliant. 

597 

598 Sets `additionalProperties: false` on all objects and ensures every defined 

599 property is listed in `required`. Walks into nested properties, `$defs`, 

600 array `items`, and `anyOf`/`oneOf`/`allOf` combinators. 

601 

602 See https://platform.openai.com/docs/guides/structured-outputs#supported-schemas 

603 """ 

604 schema = {**schema} 

605 

606 schema_type = schema.get("type") 

607 

608 if schema_type == "object" or "properties" in schema: 

609 schema["additionalProperties"] = False 

610 if "properties" in schema: 

611 schema["required"] = list(schema["properties"].keys()) 

612 schema["properties"] = {k: _make_schema_strict(v) for k, v in schema["properties"].items()} 

613 

614 if "items" in schema: 

615 schema["items"] = _make_schema_strict(schema["items"]) 

616 

617 if "$defs" in schema: 

618 schema["$defs"] = {k: _make_schema_strict(v) for k, v in schema["$defs"].items()} 

619 

620 for combinator in ("anyOf", "oneOf", "allOf"): 

621 if combinator in schema: 

622 schema[combinator] = [_make_schema_strict(s) for s in schema[combinator]] 

623 

624 return schema 

625 

626 

627def _check_finish_reason(meta: dict[str, Any]) -> None: 

628 if meta["finish_reason"] == "length": 

629 logger.warning( 

630 "The completion for index {index} has been truncated before reaching a natural stopping point. " 

631 "Increase the max_completion_tokens parameter to allow for longer completions.", 

632 index=meta["index"], 

633 finish_reason=meta["finish_reason"], 

634 ) 

635 if meta["finish_reason"] == "content_filter": 

636 logger.warning( 

637 "The completion for index {index} has been truncated due to the content filter.", 

638 index=meta["index"], 

639 finish_reason=meta["finish_reason"], 

640 ) 

641 

642 

643def _convert_chat_completion_to_chat_message( 

644 completion: ChatCompletion | ParsedChatCompletion, choice: Choice 

645) -> ChatMessage: 

646 """ 

647 Converts the non-streaming response from the OpenAI API to a ChatMessage. 

648 

649 :param completion: The completion returned by the OpenAI API. 

650 :param choice: The choice returned by the OpenAI API. 

651 :return: The ChatMessage. 

652 """ 

653 message: ChatCompletionMessage | ParsedChatCompletionMessage = choice.message 

654 text = message.content 

655 tool_calls = [] 

656 if message.tool_calls: 

657 # we currently only support function tools (not custom tools) 

658 # https://platform.openai.com/docs/guides/function-calling#custom-tools 

659 openai_tool_calls = [tc for tc in message.tool_calls if not isinstance(tc, ChatCompletionMessageCustomToolCall)] 

660 for openai_tc in openai_tool_calls: 

661 arguments_str = openai_tc.function.arguments 

662 try: 

663 arguments = _parse_openai_tool_call_arguments(arguments_str) 

664 tool_calls.append(ToolCall(id=openai_tc.id, tool_name=openai_tc.function.name, arguments=arguments)) 

665 except json.JSONDecodeError: 

666 logger.warning( 

667 "OpenAI returned a malformed JSON string for tool call arguments. This tool call " 

668 "will be skipped. To always generate a valid JSON, set `tools_strict` to `True`. " 

669 "Tool call ID: {_id}, Tool name: {_name}, Arguments: {_arguments}", 

670 _id=openai_tc.id, 

671 _name=openai_tc.function.name, 

672 _arguments=arguments_str, 

673 ) 

674 

675 logprobs = _serialize_object(choice.logprobs) if choice.logprobs else None 

676 meta = { 

677 "model": completion.model, 

678 "index": choice.index, 

679 "finish_reason": choice.finish_reason, 

680 "usage": _serialize_object(completion.usage), 

681 } 

682 if logprobs: 

683 meta["logprobs"] = logprobs 

684 

685 return ChatMessage.from_assistant(text=text, tool_calls=tool_calls, meta=meta) 

686 

687 

688def _convert_chat_completion_chunk_to_streaming_chunk( 

689 chunk: ChatCompletionChunk, previous_chunks: list[StreamingChunk], component_info: ComponentInfo | None = None 

690) -> StreamingChunk: 

691 """ 

692 Converts the streaming response chunk from the OpenAI API to a StreamingChunk. 

693 

694 :param chunk: The chunk returned by the OpenAI API. 

695 :param previous_chunks: A list of previously received StreamingChunks. 

696 :param component_info: An optional `ComponentInfo` object containing information about the component that 

697 generated the chunk, such as the component name and type. 

698 

699 :returns: 

700 A StreamingChunk object representing the content of the chunk from the OpenAI API. 

701 """ 

702 finish_reason_mapping: dict[str, FinishReason] = { 

703 "stop": "stop", 

704 "length": "length", 

705 "content_filter": "content_filter", 

706 "tool_calls": "tool_calls", 

707 "function_call": "tool_calls", 

708 } 

709 # On very first chunk so len(previous_chunks) == 0, the Choices field only provides role info (e.g. "assistant") 

710 # Choices is empty if include_usage is set to True where the usage information is returned. 

711 if len(chunk.choices) == 0: 

712 return StreamingChunk( 

713 content="", 

714 component_info=component_info, 

715 # Index is None since it's only set to an int when a content block is present 

716 index=None, 

717 finish_reason=None, 

718 meta={ 

719 "model": chunk.model, 

720 "received_at": datetime.now().isoformat(), 

721 "usage": _serialize_object(chunk.usage), 

722 }, 

723 ) 

724 

725 choice: ChunkChoice = chunk.choices[0] 

726 

727 # create a list of ToolCallDelta objects from the tool calls 

728 if choice.delta and choice.delta.tool_calls: 

729 tool_calls_deltas = [] 

730 for tool_call in choice.delta.tool_calls: 

731 function = tool_call.function 

732 tool_calls_deltas.append( 

733 ToolCallDelta( 

734 index=tool_call.index, 

735 id=tool_call.id, 

736 tool_name=function.name if function else None, 

737 arguments=function.arguments if function and function.arguments else None, 

738 ) 

739 ) 

740 return StreamingChunk( 

741 content=choice.delta.content or "", 

742 component_info=component_info, 

743 # We adopt the first tool_calls_deltas.index as the overall index of the chunk. 

744 index=tool_calls_deltas[0].index, 

745 tool_calls=tool_calls_deltas, 

746 start=tool_calls_deltas[0].tool_name is not None, 

747 finish_reason=finish_reason_mapping.get(choice.finish_reason) if choice.finish_reason else None, 

748 meta={ 

749 "model": chunk.model, 

750 "index": choice.index, 

751 "tool_calls": choice.delta.tool_calls, 

752 "finish_reason": choice.finish_reason, 

753 "received_at": datetime.now().isoformat(), 

754 "usage": _serialize_object(chunk.usage), 

755 }, 

756 ) 

757 

758 # On very first chunk the choice field only provides role info (e.g. "assistant") so we set index to None 

759 # We set all chunks missing the content field to index of None. E.g. can happen if chunk only contains finish 

760 # reason. 

761 if choice.delta and (choice.delta.content is None or choice.delta.role is not None): 

762 resolved_index = None 

763 else: 

764 # We set the index to be 0 since if text content is being streamed then no tool calls are being streamed 

765 # NOTE: We may need to revisit this if OpenAI allows planning/thinking content before tool calls like 

766 # Anthropic Claude 

767 resolved_index = 0 

768 

769 # Initialize meta dictionary 

770 meta = { 

771 "model": chunk.model, 

772 "index": choice.index, 

773 "tool_calls": choice.delta.tool_calls if choice.delta and choice.delta.tool_calls else None, 

774 "finish_reason": choice.finish_reason, 

775 "received_at": datetime.now().isoformat(), 

776 "usage": _serialize_object(chunk.usage), 

777 } 

778 

779 # check if logprobs are present 

780 # logprobs are returned only for text content 

781 logprobs = _serialize_object(choice.logprobs) if choice.logprobs else None 

782 if logprobs: 

783 meta["logprobs"] = logprobs 

784 

785 content = "" 

786 if choice.delta and choice.delta.content: 

787 content = choice.delta.content 

788 

789 return StreamingChunk( 

790 content=content, 

791 component_info=component_info, 

792 index=resolved_index, 

793 # The first chunk is always a start message chunk that only contains role information, so if we reach here 

794 # and previous_chunks is length 1 then this is the start of text content. 

795 start=len(previous_chunks) == 1, 

796 finish_reason=finish_reason_mapping.get(choice.finish_reason) if choice.finish_reason else None, 

797 meta=meta, 

798 )