Coverage for haystack/components/generators/chat/openai.py: 98%
232 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import asyncio
6import json
7import os
8from datetime import datetime
9from typing import Any, ClassVar
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
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.streaming_chunk import _invoke_streaming_callback
43from haystack.tools import (
44 ToolsType,
45 _check_duplicate_tool_names,
46 deserialize_tools_or_toolset_inplace,
47 flatten_tools_or_toolsets,
48 serialize_tools_or_toolset,
49 warm_up_tools,
50)
51from haystack.utils import Secret, deserialize_callable, serialize_callable
52from haystack.utils.http_client import init_http_client
54logger = logging.getLogger(__name__)
57@component
58class OpenAIChatGenerator:
59 """
60 Completes chats using OpenAI's large language models (LLMs).
62 It works with the gpt-4 and gpt-5 series models and supports streaming responses
63 from OpenAI API. It uses [ChatMessage](https://docs.haystack.deepset.ai/docs/chatmessage)
64 format in input and output.
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.
71 For details on OpenAI API parameters, see
72 [OpenAI documentation](https://platform.openai.com/docs/api-reference/chat).
74 ### Usage example
75 ```python
76 from haystack.components.generators.chat import OpenAIChatGenerator
77 from haystack.dataclasses import ChatMessage
79 messages = [ChatMessage.from_user("What's Natural Language Processing?")]
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 """
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."""
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
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.
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).
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))
212 self.client: OpenAI | None = None
213 self.async_client: AsyncOpenAI | None = None
214 self._tools_warmed_up = False
216 def _client_kwargs(self) -> dict[str, Any]:
217 timeout = self.timeout if self.timeout is not None else float(os.environ.get("OPENAI_TIMEOUT", "30.0"))
218 max_retries = (
219 self.max_retries if self.max_retries is not None else int(os.environ.get("OPENAI_MAX_RETRIES", "5"))
220 )
221 return {
222 "api_key": self.api_key.resolve_value(),
223 "organization": self.organization,
224 "base_url": self.api_base_url,
225 "timeout": timeout,
226 "max_retries": max_retries,
227 }
229 def _warm_up_tools(self) -> None:
230 if not self._tools_warmed_up:
231 warm_up_tools(self.tools)
232 self._tools_warmed_up = True
234 def warm_up(self) -> None:
235 """
236 Warm up the tools and initialize the synchronous OpenAI client.
237 """
238 self._warm_up_tools()
239 if self.client is None:
240 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
241 # https://github.com/openai/openai-python/blob/main/httpx2.md
242 http_client = init_http_client(self.http_client_kwargs, async_client=False)
243 self.client = OpenAI(
244 http_client=http_client, # type: ignore[arg-type]
245 **self._client_kwargs(),
246 )
248 async def warm_up_async(self) -> None: # noqa: RUF029
249 """
250 Warm up the tools and initialize the asynchronous OpenAI client on the serving event loop.
251 """
252 self._warm_up_tools()
253 if self.async_client is None:
254 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
255 # https://github.com/openai/openai-python/blob/main/httpx2.md
256 http_client = init_http_client(self.http_client_kwargs, async_client=True)
257 self.async_client = AsyncOpenAI(
258 http_client=http_client, # type: ignore[arg-type]
259 **self._client_kwargs(),
260 )
262 def close(self) -> None:
263 """
264 Releases the synchronous OpenAI client.
265 """
266 if self.client is not None:
267 self.client.close()
268 self.client = None
270 async def close_async(self) -> None:
271 """
272 Releases the asynchronous OpenAI client.
273 """
274 if self.async_client is not None:
275 await self.async_client.close()
276 self.async_client = None
278 def _get_telemetry_data(self) -> dict[str, Any]:
279 """
280 Data that is sent to Posthog for usage analytics.
281 """
282 return {"model": self.model}
284 def to_dict(self) -> dict[str, Any]:
285 """
286 Serialize this component to a dictionary.
288 :returns:
289 The serialized component as a dictionary.
290 """
291 callback_name = serialize_callable(self.streaming_callback) if self.streaming_callback else None
292 generation_kwargs = self.generation_kwargs.copy()
293 response_format = generation_kwargs.get("response_format")
295 # If the response format is a Pydantic model, it's converted to openai's json schema format
296 # If it's already a json schema, it's left as is
297 if response_format and isinstance(response_format, type) and issubclass(response_format, BaseModel):
298 json_schema = {
299 "type": "json_schema",
300 "json_schema": {
301 "name": response_format.__name__,
302 "strict": True,
303 "schema": to_strict_json_schema(response_format),
304 },
305 }
306 generation_kwargs["response_format"] = json_schema
308 return default_to_dict(
309 self,
310 model=self.model,
311 streaming_callback=callback_name,
312 api_base_url=self.api_base_url,
313 organization=self.organization,
314 generation_kwargs=generation_kwargs,
315 api_key=self.api_key,
316 timeout=self.timeout,
317 max_retries=self.max_retries,
318 tools=serialize_tools_or_toolset(self.tools),
319 tools_strict=self.tools_strict,
320 http_client_kwargs=self.http_client_kwargs,
321 )
323 @classmethod
324 def from_dict(cls, data: dict[str, Any]) -> "OpenAIChatGenerator":
325 """
326 Deserialize this component from a dictionary.
328 :param data: The dictionary representation of this component.
329 :returns:
330 The deserialized component instance.
331 """
332 deserialize_tools_or_toolset_inplace(data["init_parameters"], key="tools")
333 init_params = data.get("init_parameters", {})
334 serialized_callback_handler = init_params.get("streaming_callback")
336 if serialized_callback_handler:
337 data["init_parameters"]["streaming_callback"] = deserialize_callable(serialized_callback_handler)
338 return default_from_dict(cls, data)
340 @component.output_types(replies=list[ChatMessage])
341 def run(
342 self,
343 messages: list[ChatMessage] | str,
344 streaming_callback: StreamingCallbackT | None = None,
345 generation_kwargs: dict[str, Any] | None = None,
346 *,
347 tools: ToolsType | None = None,
348 tools_strict: bool | None = None,
349 ) -> dict[str, list[ChatMessage]]:
350 """
351 Invokes chat completion based on the provided messages and generation parameters.
353 :param messages:
354 A list of ChatMessage instances representing the input messages. If a string is provided, it is converted
355 to a list containing a ChatMessage with user role.
356 :param streaming_callback:
357 A callback function that is called when a new token is received from the stream.
358 :param generation_kwargs:
359 Additional keyword arguments for text generation. These are merged per key with the
360 `generation_kwargs` passed at initialization: keys provided here take precedence, keys set
361 only at initialization are kept.
362 For details on OpenAI API parameters, see [OpenAI documentation](https://platform.openai.com/docs/api-reference/chat/create).
363 :param tools:
364 A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls.
365 If set, it will override the `tools` parameter provided during initialization.
366 :param tools_strict:
367 Whether to enable strict schema adherence for tool calls. If set to `True`, the model will follow exactly
368 the schema provided in the `parameters` field of the tool definition, but this may increase latency.
369 If set, it will override the `tools_strict` parameter set during component initialization.
371 :returns:
372 A dictionary with the following key:
373 - `replies`: A list containing the generated responses as ChatMessage instances.
374 """
375 self.warm_up()
377 messages = _normalize_messages(messages)
379 if len(messages) == 0:
380 return {"replies": []}
382 streaming_callback = select_streaming_callback(
383 init_callback=self.streaming_callback, runtime_callback=streaming_callback, requires_async=False
384 )
385 chat_completion: Stream[ChatCompletionChunk] | ChatCompletion | ParsedChatCompletion
387 api_args = self._prepare_api_call(
388 messages=messages,
389 streaming_callback=streaming_callback,
390 generation_kwargs=generation_kwargs,
391 tools=tools,
392 tools_strict=tools_strict,
393 )
394 openai_endpoint = api_args.pop("openai_endpoint")
395 assert self.client is not None # mypy: client is built by warm_up above
396 openai_endpoint_method = getattr(self.client.chat.completions, openai_endpoint)
397 chat_completion = openai_endpoint_method(**api_args)
399 if streaming_callback is not None:
400 completions = self._handle_stream_response(
401 # we cannot check isinstance(chat_completion, Stream) because some observability tools wrap Stream
402 # and return a different type. See https://github.com/deepset-ai/haystack/issues/9014.
403 chat_completion, # type: ignore
404 streaming_callback,
405 )
407 else:
408 assert isinstance(chat_completion, ChatCompletion), "Unexpected response type for non-streaming request."
409 completions = [
410 _convert_chat_completion_to_chat_message(chat_completion, choice) for choice in chat_completion.choices
411 ]
413 # before returning, do post-processing of the completions
414 for message in completions:
415 _check_finish_reason(message.meta)
417 return {"replies": completions}
419 @component.output_types(replies=list[ChatMessage])
420 async def run_async(
421 self,
422 messages: list[ChatMessage] | str,
423 streaming_callback: StreamingCallbackT | None = None,
424 generation_kwargs: dict[str, Any] | None = None,
425 *,
426 tools: ToolsType | None = None,
427 tools_strict: bool | None = None,
428 ) -> dict[str, list[ChatMessage]]:
429 """
430 Asynchronously invokes chat completion based on the provided messages and generation parameters.
432 This is the asynchronous version of the `run` method. It has the same parameters and return values
433 but can be used with `await` in async code.
435 :param messages:
436 A list of ChatMessage instances representing the input messages. If a string is provided, it is converted
437 to a list containing a ChatMessage with user role.
438 :param streaming_callback:
439 A callback function that is called when a new token is received from the stream. Async callbacks are
440 preferred; a sync callback is accepted but will run synchronously on the event loop and may block it.
441 :param generation_kwargs:
442 Additional keyword arguments for text generation. These are merged per key with the
443 `generation_kwargs` passed at initialization: keys provided here take precedence, keys set
444 only at initialization are kept.
445 For details on OpenAI API parameters, see [OpenAI documentation](https://platform.openai.com/docs/api-reference/chat/create).
446 :param tools: A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls.
447 If set, it will override the `tools` parameter provided during initialization.
448 :param tools_strict:
449 Whether to enable strict schema adherence for tool calls. If set to `True`, the model will follow exactly
450 the schema provided in the `parameters` field of the tool definition, but this may increase latency.
451 If set, it will override the `tools_strict` parameter set during component initialization.
453 :returns:
454 A dictionary with the following key:
455 - `replies`: A list containing the generated responses as ChatMessage instances.
456 """
457 await self.warm_up_async()
459 messages = _normalize_messages(messages)
461 # validate and select the streaming callback
462 streaming_callback = select_streaming_callback(
463 init_callback=self.streaming_callback, runtime_callback=streaming_callback, requires_async=True
464 )
465 chat_completion: AsyncStream[ChatCompletionChunk] | ChatCompletion | ParsedChatCompletion
467 if len(messages) == 0:
468 return {"replies": []}
470 api_args = self._prepare_api_call(
471 messages=messages,
472 streaming_callback=streaming_callback,
473 generation_kwargs=generation_kwargs,
474 tools=tools,
475 tools_strict=tools_strict,
476 )
478 openai_endpoint = api_args.pop("openai_endpoint")
479 assert self.async_client is not None # mypy: async_client is built by warm_up_async above
480 openai_endpoint_method = getattr(self.async_client.chat.completions, openai_endpoint)
481 chat_completion = await openai_endpoint_method(**api_args)
483 if streaming_callback is not None:
484 completions = await self._handle_async_stream_response(
485 # we cannot check isinstance(chat_completion, AsyncStream) because some observability tools wrap
486 # AsyncStream and return a different type. See https://github.com/deepset-ai/haystack/issues/9014.
487 chat_completion, # type: ignore
488 streaming_callback,
489 )
491 else:
492 assert isinstance(chat_completion, ChatCompletion), "Unexpected response type for non-streaming request."
493 completions = [
494 _convert_chat_completion_to_chat_message(chat_completion, choice) for choice in chat_completion.choices
495 ]
497 # before returning, do post-processing of the completions
498 for message in completions:
499 _check_finish_reason(message.meta)
501 return {"replies": completions}
503 def _prepare_api_call( # noqa: PLR0913
504 self,
505 *,
506 messages: list[ChatMessage],
507 streaming_callback: StreamingCallbackT | None = None,
508 generation_kwargs: dict[str, Any] | None = None,
509 tools: ToolsType | None = None,
510 tools_strict: bool | None = None,
511 ) -> dict[str, Any]:
512 # update generation kwargs by merging with the generation kwargs passed to the run method
513 generation_kwargs = {**self.generation_kwargs, **(generation_kwargs or {})}
515 is_streaming = streaming_callback is not None
516 num_responses = generation_kwargs.pop("n", 1)
518 if is_streaming and num_responses > 1:
519 raise ValueError("Cannot stream multiple responses, please set n=1.")
520 response_format = generation_kwargs.pop("response_format", None)
522 # adapt ChatMessage(s) to the format expected by the OpenAI API
523 openai_formatted_messages = [message.to_openai_dict_format() for message in messages]
525 flattened_tools = flatten_tools_or_toolsets(tools or self.tools)
526 tools_strict = tools_strict if tools_strict is not None else self.tools_strict
527 _check_duplicate_tool_names(flattened_tools)
529 openai_tools = {}
530 if flattened_tools:
531 tool_definitions = []
532 for t in flattened_tools:
533 function_spec = {**t.tool_spec}
534 if tools_strict:
535 function_spec["strict"] = True
536 function_spec["parameters"] = _make_schema_strict(function_spec["parameters"])
537 tool_definitions.append({"type": "function", "function": function_spec})
538 openai_tools = {"tools": tool_definitions}
540 base_args = {
541 "model": self.model,
542 "messages": openai_formatted_messages,
543 "n": num_responses,
544 **openai_tools,
545 **generation_kwargs,
546 }
548 if response_format and not is_streaming:
549 # for structured outputs without streaming, we use openai's parse endpoint
550 # Note: `stream` cannot be passed to chat.completions.parse
551 # we pass a key `openai_endpoint` as a hint to the run method to use the parse endpoint
552 # this key will be removed before the API call is made
553 return {**base_args, "response_format": response_format, "openai_endpoint": "parse"}
555 # for structured outputs with streaming, we use openai's create endpoint
556 # we pass a key `openai_endpoint` as a hint to the run method to use the create endpoint
557 # this key will be removed before the API call is made
558 final_args = {**base_args, "stream": is_streaming, "openai_endpoint": "create"}
560 # We only set the response_format parameter if it's not None since None is not a valid value in the API.
561 if response_format:
562 final_args["response_format"] = response_format
563 return final_args
565 def _handle_stream_response(self, chat_completion: Stream, callback: SyncStreamingCallbackT) -> list[ChatMessage]:
566 component_info = ComponentInfo.from_component(self)
567 chunks: list[StreamingChunk] = []
568 for chunk in chat_completion:
569 assert len(chunk.choices) <= 1, "Streaming responses should have at most one choice."
570 chunk_delta = _convert_chat_completion_chunk_to_streaming_chunk(
571 chunk=chunk, previous_chunks=chunks, component_info=component_info
572 )
573 chunks.append(chunk_delta)
574 callback(chunk_delta)
575 return [_convert_streaming_chunks_to_chat_message(chunks=chunks)]
577 async def _handle_async_stream_response(
578 self, chat_completion: AsyncStream, callback: StreamingCallbackT
579 ) -> list[ChatMessage]:
580 component_info = ComponentInfo.from_component(self)
581 chunks: list[StreamingChunk] = []
582 try:
583 async for chunk in chat_completion:
584 assert len(chunk.choices) <= 1, "Streaming responses should have at most one choice."
585 chunk_delta = _convert_chat_completion_chunk_to_streaming_chunk(
586 chunk=chunk, previous_chunks=chunks, component_info=component_info
587 )
588 chunks.append(chunk_delta)
589 await _invoke_streaming_callback(callback, chunk_delta)
591 except asyncio.CancelledError:
592 await asyncio.shield(chat_completion.close())
593 # close the stream when task is cancelled
594 # asyncio.shield ensures the close operation completes
595 # https://docs.python.org/3/library/asyncio-task.html#shielding-from-cancellation
596 raise # Re-raise to propagate cancellation
598 return [_convert_streaming_chunks_to_chat_message(chunks=chunks)]
601def _make_schema_strict(schema: dict[str, Any]) -> dict[str, Any]:
602 """
603 Recursively transform a JSON schema to be OpenAI strict-mode compliant.
605 Sets `additionalProperties: false` on all objects and ensures every defined
606 property is listed in `required`. Walks into nested properties, `$defs`,
607 array `items`, and `anyOf`/`oneOf`/`allOf` combinators.
609 See https://platform.openai.com/docs/guides/structured-outputs#supported-schemas
610 """
611 schema = {**schema}
613 schema_type = schema.get("type")
615 if schema_type == "object" or "properties" in schema:
616 schema["additionalProperties"] = False
617 if "properties" in schema:
618 schema["required"] = list(schema["properties"].keys())
619 schema["properties"] = {k: _make_schema_strict(v) for k, v in schema["properties"].items()}
621 if "items" in schema:
622 schema["items"] = _make_schema_strict(schema["items"])
624 if "$defs" in schema:
625 schema["$defs"] = {k: _make_schema_strict(v) for k, v in schema["$defs"].items()}
627 for combinator in ("anyOf", "oneOf", "allOf"):
628 if combinator in schema:
629 schema[combinator] = [_make_schema_strict(s) for s in schema[combinator]]
631 return schema
634def _check_finish_reason(meta: dict[str, Any]) -> None:
635 if meta["finish_reason"] == "length":
636 logger.warning(
637 "The completion for index {index} has been truncated before reaching a natural stopping point. "
638 "Increase the max_completion_tokens parameter to allow for longer completions.",
639 index=meta["index"],
640 finish_reason=meta["finish_reason"],
641 )
642 if meta["finish_reason"] == "content_filter":
643 logger.warning(
644 "The completion for index {index} has been truncated due to the content filter.",
645 index=meta["index"],
646 finish_reason=meta["finish_reason"],
647 )
650def _convert_chat_completion_to_chat_message(
651 completion: ChatCompletion | ParsedChatCompletion, choice: Choice
652) -> ChatMessage:
653 """
654 Converts the non-streaming response from the OpenAI API to a ChatMessage.
656 :param completion: The completion returned by the OpenAI API.
657 :param choice: The choice returned by the OpenAI API.
658 :return: The ChatMessage.
659 """
660 message: ChatCompletionMessage | ParsedChatCompletionMessage = choice.message
661 text = message.content
662 tool_calls = []
663 if message.tool_calls:
664 # we currently only support function tools (not custom tools)
665 # https://platform.openai.com/docs/guides/function-calling#custom-tools
666 openai_tool_calls = [tc for tc in message.tool_calls if not isinstance(tc, ChatCompletionMessageCustomToolCall)]
667 for openai_tc in openai_tool_calls:
668 arguments_str = openai_tc.function.arguments
669 try:
670 arguments = json.loads(arguments_str)
671 tool_calls.append(ToolCall(id=openai_tc.id, tool_name=openai_tc.function.name, arguments=arguments))
672 except json.JSONDecodeError:
673 logger.warning(
674 "OpenAI returned a malformed JSON string for tool call arguments. This tool call "
675 "will be skipped. To always generate a valid JSON, set `tools_strict` to `True`. "
676 "Tool call ID: {_id}, Tool name: {_name}, Arguments: {_arguments}",
677 _id=openai_tc.id,
678 _name=openai_tc.function.name,
679 _arguments=arguments_str,
680 )
682 logprobs = _serialize_object(choice.logprobs) if choice.logprobs else None
683 meta = {
684 "model": completion.model,
685 "index": choice.index,
686 "finish_reason": choice.finish_reason,
687 "usage": _serialize_object(completion.usage),
688 }
689 if logprobs:
690 meta["logprobs"] = logprobs
692 return ChatMessage.from_assistant(text=text, tool_calls=tool_calls, meta=meta)
695def _convert_chat_completion_chunk_to_streaming_chunk(
696 chunk: ChatCompletionChunk, previous_chunks: list[StreamingChunk], component_info: ComponentInfo | None = None
697) -> StreamingChunk:
698 """
699 Converts the streaming response chunk from the OpenAI API to a StreamingChunk.
701 :param chunk: The chunk returned by the OpenAI API.
702 :param previous_chunks: A list of previously received StreamingChunks.
703 :param component_info: An optional `ComponentInfo` object containing information about the component that
704 generated the chunk, such as the component name and type.
706 :returns:
707 A StreamingChunk object representing the content of the chunk from the OpenAI API.
708 """
709 finish_reason_mapping: dict[str, FinishReason] = {
710 "stop": "stop",
711 "length": "length",
712 "content_filter": "content_filter",
713 "tool_calls": "tool_calls",
714 "function_call": "tool_calls",
715 }
716 # On very first chunk so len(previous_chunks) == 0, the Choices field only provides role info (e.g. "assistant")
717 # Choices is empty if include_usage is set to True where the usage information is returned.
718 if len(chunk.choices) == 0:
719 return StreamingChunk(
720 content="",
721 component_info=component_info,
722 # Index is None since it's only set to an int when a content block is present
723 index=None,
724 finish_reason=None,
725 meta={
726 "model": chunk.model,
727 "received_at": datetime.now().isoformat(),
728 "usage": _serialize_object(chunk.usage),
729 },
730 )
732 choice: ChunkChoice = chunk.choices[0]
734 # create a list of ToolCallDelta objects from the tool calls
735 if choice.delta and choice.delta.tool_calls:
736 tool_calls_deltas = []
737 for tool_call in choice.delta.tool_calls:
738 function = tool_call.function
739 tool_calls_deltas.append(
740 ToolCallDelta(
741 index=tool_call.index,
742 id=tool_call.id,
743 tool_name=function.name if function else None,
744 arguments=function.arguments if function and function.arguments else None,
745 )
746 )
747 return StreamingChunk(
748 content=choice.delta.content or "",
749 component_info=component_info,
750 # We adopt the first tool_calls_deltas.index as the overall index of the chunk.
751 index=tool_calls_deltas[0].index,
752 tool_calls=tool_calls_deltas,
753 start=tool_calls_deltas[0].tool_name is not None,
754 finish_reason=finish_reason_mapping.get(choice.finish_reason) if choice.finish_reason else None,
755 meta={
756 "model": chunk.model,
757 "index": choice.index,
758 "tool_calls": choice.delta.tool_calls,
759 "finish_reason": choice.finish_reason,
760 "received_at": datetime.now().isoformat(),
761 "usage": _serialize_object(chunk.usage),
762 },
763 )
765 # On very first chunk the choice field only provides role info (e.g. "assistant") so we set index to None
766 # We set all chunks missing the content field to index of None. E.g. can happen if chunk only contains finish
767 # reason.
768 if choice.delta and (choice.delta.content is None or choice.delta.role is not None):
769 resolved_index = None
770 else:
771 # We set the index to be 0 since if text content is being streamed then no tool calls are being streamed
772 # NOTE: We may need to revisit this if OpenAI allows planning/thinking content before tool calls like
773 # Anthropic Claude
774 resolved_index = 0
776 # Initialize meta dictionary
777 meta = {
778 "model": chunk.model,
779 "index": choice.index,
780 "tool_calls": choice.delta.tool_calls if choice.delta and choice.delta.tool_calls else None,
781 "finish_reason": choice.finish_reason,
782 "received_at": datetime.now().isoformat(),
783 "usage": _serialize_object(chunk.usage),
784 }
786 # check if logprobs are present
787 # logprobs are returned only for text content
788 logprobs = _serialize_object(choice.logprobs) if choice.logprobs else None
789 if logprobs:
790 meta["logprobs"] = logprobs
792 content = ""
793 if choice.delta and choice.delta.content:
794 content = choice.delta.content
796 return StreamingChunk(
797 content=content,
798 component_info=component_info,
799 index=resolved_index,
800 # The first chunk is always a start message chunk that only contains role information, so if we reach here
801 # and previous_chunks is length 1 then this is the start of text content.
802 start=len(previous_chunks) == 1,
803 finish_reason=finish_reason_mapping.get(choice.finish_reason) if choice.finish_reason else None,
804 meta=meta,
805 )