diff --git a/haystack/components/embedders/azure_document_embedder.py b/haystack/components/embedders/azure_document_embedder.py index 38ebc03838..04ab8c5d78 100644 --- a/haystack/components/embedders/azure_document_embedder.py +++ b/haystack/components/embedders/azure_document_embedder.py @@ -173,10 +173,7 @@ def warm_up(self) -> None: # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. # https://github.com/openai/openai-python/blob/main/httpx2.md http_client = init_http_client(self.http_client_kwargs, async_client=False) - self.client = AzureOpenAI( - http_client=http_client, # type: ignore[arg-type] - **self._client_kwargs(), - ) + self.client = AzureOpenAI(http_client=http_client, **self._client_kwargs()) async def warm_up_async(self) -> None: # noqa: RUF029 """ @@ -186,10 +183,7 @@ async def warm_up_async(self) -> None: # noqa: RUF029 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. # https://github.com/openai/openai-python/blob/main/httpx2.md http_client = init_http_client(self.http_client_kwargs, async_client=True) - self.async_client = AsyncAzureOpenAI( - http_client=http_client, # type: ignore[arg-type] - **self._client_kwargs(), - ) + self.async_client = AsyncAzureOpenAI(http_client=http_client, **self._client_kwargs()) def close(self) -> None: """ diff --git a/haystack/components/embedders/azure_text_embedder.py b/haystack/components/embedders/azure_text_embedder.py index 76d4c2a1b7..f96bed364b 100644 --- a/haystack/components/embedders/azure_text_embedder.py +++ b/haystack/components/embedders/azure_text_embedder.py @@ -154,10 +154,7 @@ def warm_up(self) -> None: # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. # https://github.com/openai/openai-python/blob/main/httpx2.md http_client = init_http_client(self.http_client_kwargs, async_client=False) - self.client = AzureOpenAI( - http_client=http_client, # type: ignore[arg-type] - **self._client_kwargs(), - ) + self.client = AzureOpenAI(http_client=http_client, **self._client_kwargs()) async def warm_up_async(self) -> None: # noqa: RUF029 """ @@ -167,10 +164,7 @@ async def warm_up_async(self) -> None: # noqa: RUF029 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. # https://github.com/openai/openai-python/blob/main/httpx2.md http_client = init_http_client(self.http_client_kwargs, async_client=True) - self.async_client = AsyncAzureOpenAI( - http_client=http_client, # type: ignore[arg-type] - **self._client_kwargs(), - ) + self.async_client = AsyncAzureOpenAI(http_client=http_client, **self._client_kwargs()) def close(self) -> None: """ diff --git a/haystack/components/embedders/openai_document_embedder.py b/haystack/components/embedders/openai_document_embedder.py index e9911c36c5..f54f4d8495 100644 --- a/haystack/components/embedders/openai_document_embedder.py +++ b/haystack/components/embedders/openai_document_embedder.py @@ -146,10 +146,7 @@ def warm_up(self) -> None: # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. # https://github.com/openai/openai-python/blob/main/httpx2.md http_client = init_http_client(self.http_client_kwargs, async_client=False) - self.client = OpenAI( - http_client=http_client, # type: ignore[arg-type] - **self._client_kwargs(), - ) + self.client = OpenAI(http_client=http_client, **self._client_kwargs()) async def warm_up_async(self) -> None: # noqa: RUF029 """ @@ -159,10 +156,7 @@ async def warm_up_async(self) -> None: # noqa: RUF029 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. # https://github.com/openai/openai-python/blob/main/httpx2.md http_client = init_http_client(self.http_client_kwargs, async_client=True) - self.async_client = AsyncOpenAI( - http_client=http_client, # type: ignore[arg-type] - **self._client_kwargs(), - ) + self.async_client = AsyncOpenAI(http_client=http_client, **self._client_kwargs()) def close(self) -> None: """ diff --git a/haystack/components/embedders/openai_text_embedder.py b/haystack/components/embedders/openai_text_embedder.py index 22d0b75a96..25113843bc 100644 --- a/haystack/components/embedders/openai_text_embedder.py +++ b/haystack/components/embedders/openai_text_embedder.py @@ -121,10 +121,7 @@ def warm_up(self) -> None: # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. # https://github.com/openai/openai-python/blob/main/httpx2.md http_client = init_http_client(self.http_client_kwargs, async_client=False) - self.client = OpenAI( - http_client=http_client, # type: ignore[arg-type] - **self._client_kwargs(), - ) + self.client = OpenAI(http_client=http_client, **self._client_kwargs()) async def warm_up_async(self) -> None: # noqa: RUF029 """ @@ -134,10 +131,7 @@ async def warm_up_async(self) -> None: # noqa: RUF029 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. # https://github.com/openai/openai-python/blob/main/httpx2.md http_client = init_http_client(self.http_client_kwargs, async_client=True) - self.async_client = AsyncOpenAI( - http_client=http_client, # type: ignore[arg-type] - **self._client_kwargs(), - ) + self.async_client = AsyncOpenAI(http_client=http_client, **self._client_kwargs()) def close(self) -> None: """ diff --git a/haystack/components/generators/chat/__init__.py b/haystack/components/generators/chat/__init__.py index 7542db6d48..0c05ad94db 100644 --- a/haystack/components/generators/chat/__init__.py +++ b/haystack/components/generators/chat/__init__.py @@ -9,6 +9,7 @@ _import_structure = { "openai": ["OpenAIChatGenerator"], + "openai_batch": ["OpenAIBatchChatGenerator"], "openai_responses": ["OpenAIResponsesChatGenerator"], "azure": ["AzureOpenAIChatGenerator"], "azure_responses": ["AzureOpenAIResponsesChatGenerator"], @@ -24,6 +25,7 @@ from .llm import LLM as LLM from .mock import MockChatGenerator as MockChatGenerator from .openai import OpenAIChatGenerator as OpenAIChatGenerator + from .openai_batch import OpenAIBatchChatGenerator as OpenAIBatchChatGenerator from .openai_responses import OpenAIResponsesChatGenerator as OpenAIResponsesChatGenerator else: diff --git a/haystack/components/generators/chat/azure.py b/haystack/components/generators/chat/azure.py index f5907fcdbe..ea62b132db 100644 --- a/haystack/components/generators/chat/azure.py +++ b/haystack/components/generators/chat/azure.py @@ -274,10 +274,7 @@ def warm_up(self) -> None: # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. # https://github.com/openai/openai-python/blob/main/httpx2.md http_client = init_http_client(self.http_client_kwargs, async_client=False) - self.client = AzureOpenAI( - http_client=http_client, # type: ignore[arg-type] - **self._client_kwargs(), - ) + self.client = AzureOpenAI(http_client=http_client, **self._client_kwargs()) async def warm_up_async(self) -> None: # noqa: RUF029 """ @@ -288,10 +285,7 @@ async def warm_up_async(self) -> None: # noqa: RUF029 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. # https://github.com/openai/openai-python/blob/main/httpx2.md http_client = init_http_client(self.http_client_kwargs, async_client=True) - self.async_client = AsyncAzureOpenAI( - http_client=http_client, # type: ignore[arg-type] - **self._client_kwargs(), - ) + self.async_client = AsyncAzureOpenAI(http_client=http_client, **self._client_kwargs()) def close(self) -> None: """ diff --git a/haystack/components/generators/chat/openai.py b/haystack/components/generators/chat/openai.py index 4cf1823df2..a65192cc18 100644 --- a/haystack/components/generators/chat/openai.py +++ b/haystack/components/generators/chat/openai.py @@ -240,10 +240,7 @@ def warm_up(self) -> None: # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. # https://github.com/openai/openai-python/blob/main/httpx2.md http_client = init_http_client(self.http_client_kwargs, async_client=False) - self.client = OpenAI( - http_client=http_client, # type: ignore[arg-type] - **self._client_kwargs(), - ) + self.client = OpenAI(http_client=http_client, **self._client_kwargs()) async def warm_up_async(self) -> None: # noqa: RUF029 """ @@ -254,10 +251,7 @@ async def warm_up_async(self) -> None: # noqa: RUF029 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. # https://github.com/openai/openai-python/blob/main/httpx2.md http_client = init_http_client(self.http_client_kwargs, async_client=True) - self.async_client = AsyncOpenAI( - http_client=http_client, # type: ignore[arg-type] - **self._client_kwargs(), - ) + self.async_client = AsyncOpenAI(http_client=http_client, **self._client_kwargs()) def close(self) -> None: """ diff --git a/haystack/components/generators/chat/openai_batch.py b/haystack/components/generators/chat/openai_batch.py new file mode 100644 index 0000000000..c3f333574f --- /dev/null +++ b/haystack/components/generators/chat/openai_batch.py @@ -0,0 +1,484 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import io +import json +import os +import time +from typing import Any, Final + +from openai import AsyncOpenAI, OpenAI +from openai.types import Batch + +from haystack import component, default_from_dict, default_to_dict, logging +from haystack.dataclasses import ChatMessage, ToolCall +from haystack.utils import Secret +from haystack.utils.http_client import init_http_client + +logger = logging.getLogger(__name__) + +# The only batch-supported endpoint for chat completions +_BATCH_ENDPOINT: Final = "/v1/chat/completions" + +# Once a batch reaches one of these, there's nothing more to wait for +_TERMINAL_STATUSES = frozenset({"completed", "failed", "expired", "cancelled"}) + + +@component +class OpenAIBatchChatGenerator: + """ + Submits multiple conversations to OpenAI's Batch API for asynchronous processing. + + Unlike `OpenAIChatGenerator` which handles one request at a time, this component + accepts *multiple* conversations, bundles them into a single batch job, and polls + until OpenAI finishes processing. The trade-off is latency for cost: batch requests + are 50% cheaper and enjoy higher rate limits, but can take up to 24 hours. + + Best suited for large-scale, non-latency-critical workloads like classification, + summarization, or translation over thousands of inputs. + + ### Usage example + ```python + from haystack.components.generators.chat import OpenAIBatchChatGenerator + from haystack.dataclasses import ChatMessage + + conversations = [ + [ChatMessage.from_user("Summarize: The quick brown fox...")], + [ChatMessage.from_user("Summarize: To be or not to be...")], + [ChatMessage.from_user("Summarize: It was the best of times...")], + ] + + generator = OpenAIBatchChatGenerator() + result = generator.run(message_sets=conversations) + + for reply_set in result["replies"]: + print(reply_set[0].text) + ``` + """ + + def __init__( + self, + api_key: Secret = Secret.from_env_var("OPENAI_API_KEY"), + model: str = "gpt-5-mini", + api_base_url: str | None = None, + organization: str | None = None, + generation_kwargs: dict[str, Any] | None = None, + timeout: float | None = None, + max_retries: int | None = None, + poll_interval: float = 30.0, + max_wait_seconds: float = 86400.0, + completion_window: str = "24h", + http_client_kwargs: dict[str, Any] | None = None, + ) -> None: + """ + Creates an instance of OpenAIBatchChatGenerator. + + Before initializing the component, you can set the 'OPENAI_TIMEOUT' and 'OPENAI_MAX_RETRIES' + environment variables to override the `timeout` and `max_retries` parameters respectively + in the OpenAI client. + + :param api_key: The OpenAI API key. + You can set it with an environment variable `OPENAI_API_KEY`, or pass with this parameter + during initialization. + :param model: The name of the model to use for all requests in the batch. + :param api_base_url: An optional base URL for the OpenAI API. + :param organization: Your OpenAI organization ID, defaults to `None`. See + [production best practices](https://platform.openai.com/docs/guides/production-best-practices/setting-up-your-organization). + :param generation_kwargs: Default generation parameters (temperature, max_completion_tokens, etc.) + applied to every request in the batch. These parameters are sent directly to the OpenAI endpoint. + Can be overridden at runtime via the `generation_kwargs` parameter in `run()`. + :param timeout: Timeout for the OpenAI SDK HTTP client (not the batch job itself). + Defaults to the `OPENAI_TIMEOUT` env var, or 30 seconds. + :param max_retries: Max retries for transient SDK errors. + Defaults to the `OPENAI_MAX_RETRIES` env var, or 5. + :param poll_interval: Seconds between batch status checks. Default: 30. + :param max_wait_seconds: Maximum seconds to wait for batch completion before raising + a TimeoutError. Default: 86400 (24 hours), matching OpenAI's completion window. + :param completion_window: OpenAI's batch completion window. Currently only ``"24h"`` is supported + by the API. + :param http_client_kwargs: Keyword arguments for a custom ``httpx.Client`` or ``httpx.AsyncClient``. + For more information, see the `HTTPX documentation `_. + """ + self.api_key = api_key + self.model = model + self.api_base_url = api_base_url + self.organization = organization + self.generation_kwargs = generation_kwargs or {} + self.timeout = timeout + self.max_retries = max_retries + self.poll_interval = poll_interval + self.max_wait_seconds = max_wait_seconds + self.completion_window = completion_window + self.http_client_kwargs = http_client_kwargs + + self.client: OpenAI | None = None + self.async_client: AsyncOpenAI | None = None + + # ------------------------------------------------------------------ # + # Client lifecycle — same pattern as OpenAIChatGenerator # + # ------------------------------------------------------------------ # + + def _client_kwargs(self) -> dict[str, Any]: + """Shared config for both sync and async OpenAI clients.""" + timeout = self.timeout if self.timeout is not None else float(os.environ.get("OPENAI_TIMEOUT", "30.0")) + max_retries = ( + self.max_retries if self.max_retries is not None else int(os.environ.get("OPENAI_MAX_RETRIES", "5")) + ) + return { + "api_key": self.api_key.resolve_value(), + "organization": self.organization, + "base_url": self.api_base_url, + "timeout": timeout, + "max_retries": max_retries, + } + + def warm_up(self) -> None: + """Initialize the synchronous OpenAI client.""" + if self.client is None: + self.client = OpenAI( + http_client=init_http_client(self.http_client_kwargs, async_client=False), **self._client_kwargs() + ) + + async def warm_up_async(self) -> None: # noqa: RUF029 + """Initialize the asynchronous OpenAI client on the serving event loop.""" + if self.async_client is None: + self.async_client = AsyncOpenAI( + http_client=init_http_client(self.http_client_kwargs, async_client=True), **self._client_kwargs() + ) + + def close(self) -> None: + """Release the synchronous OpenAI client.""" + if self.client is not None: + self.client.close() + self.client = None + + async def close_async(self) -> None: + """Release the asynchronous OpenAI client.""" + if self.async_client is not None: + await self.async_client.close() + self.async_client = None + + def _get_telemetry_data(self) -> dict[str, Any]: + """Data that is sent to Posthog for usage analytics.""" + return {"model": self.model} + + def to_dict(self) -> dict[str, Any]: + """ + Serialize this component to a dictionary. + + :returns: + The serialized component as a dictionary. + """ + return default_to_dict( + self, + model=self.model, + api_base_url=self.api_base_url, + organization=self.organization, + generation_kwargs=self.generation_kwargs, + api_key=self.api_key, + timeout=self.timeout, + max_retries=self.max_retries, + poll_interval=self.poll_interval, + max_wait_seconds=self.max_wait_seconds, + completion_window=self.completion_window, + http_client_kwargs=self.http_client_kwargs, + ) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "OpenAIBatchChatGenerator": + """ + Deserialize this component from a dictionary. + + :param data: The dictionary representation of this component. + :returns: + The deserialized component instance. + """ + return default_from_dict(cls, data) + + @component.output_types(replies=list[list[ChatMessage]], meta=dict[str, Any]) + def run( + self, message_sets: list[list[ChatMessage]], generation_kwargs: dict[str, Any] | None = None + ) -> dict[str, Any]: + """ + Submit multiple conversations to OpenAI's Batch API and wait for results. + + Each item in ``message_sets`` is a separate conversation (a list of ChatMessage objects) + that becomes one request in the batch. The method blocks until the batch reaches a + terminal status or the ``max_wait_seconds`` timeout is reached. + + :param message_sets: + A list of conversations. Each conversation is a list of ChatMessage instances + representing the full message history (system prompt, user messages, etc.). + :param generation_kwargs: + Runtime overrides for generation parameters. Merged with (and takes precedence over) + the init-time ``generation_kwargs``. + :returns: + A dictionary with: + - ``replies``: A list of lists, where ``replies[i]`` contains the ChatMessage + response(s) for ``message_sets[i]``. + - ``meta``: Batch-level metadata including ``batch_id``, ``request_counts``, + and timestamps. + """ + self.warm_up() + assert self.client is not None # mypy: guaranteed by warm_up + + if not message_sets: + return {"replies": [], "meta": {}} + + merged_kwargs = {**self.generation_kwargs, **(generation_kwargs or {})} + + # Build the JSONL payload, upload it, and kick off the batch + jsonl_content = _build_jsonl(message_sets, self.model, merged_kwargs) + input_file = self.client.files.create(file=("batch_input.jsonl", jsonl_content), purpose="batch") + + batch = self.client.batches.create( + input_file_id=input_file.id, + endpoint=_BATCH_ENDPOINT, + completion_window=self.completion_window, # type: ignore + ) + logger.info("Batch {batch_id} created with {n} requests.", batch_id=batch.id, n=len(message_sets)) + + # Wait for the batch to finish + batch = self._poll_sync(batch.id) + _raise_on_failure(batch) + + # Download and parse the output + assert batch.output_file_id is not None # guaranteed when status == "completed" + output_content = self.client.files.content(batch.output_file_id) + replies = _parse_results(output_content.text, len(message_sets)) + + return {"replies": replies, "meta": _batch_meta(batch)} + + # ------------------------------------------------------------------ # + # run_async() — asynchronous entry point # + # ------------------------------------------------------------------ # + + @component.output_types(replies=list[list[ChatMessage]], meta=dict[str, Any]) + async def run_async( + self, message_sets: list[list[ChatMessage]], generation_kwargs: dict[str, Any] | None = None + ) -> dict[str, Any]: + """ + Async version of ``run()``. Uses ``asyncio.sleep`` for polling instead of blocking. + + See ``run()`` for full parameter documentation. + """ + await self.warm_up_async() + assert self.async_client is not None # mypy: guaranteed by warm_up_async + + if not message_sets: + return {"replies": [], "meta": {}} + + merged_kwargs = {**self.generation_kwargs, **(generation_kwargs or {})} + + jsonl_content = _build_jsonl(message_sets, self.model, merged_kwargs) + input_file = await self.async_client.files.create(file=("batch_input.jsonl", jsonl_content), purpose="batch") + + batch = await self.async_client.batches.create( + input_file_id=input_file.id, + endpoint=_BATCH_ENDPOINT, + completion_window=self.completion_window, # type: ignore + ) + logger.info("Batch {batch_id} created with {n} requests.", batch_id=batch.id, n=len(message_sets)) + + batch = await self._poll_async(batch.id) + _raise_on_failure(batch) + + assert batch.output_file_id is not None + output_response = await self.async_client.files.content(batch.output_file_id) + replies = _parse_results(output_response.text, len(message_sets)) + + return {"replies": replies, "meta": _batch_meta(batch)} + + def _poll_sync(self, batch_id: str) -> Batch: + """Poll the batch status synchronously until it reaches a terminal state.""" + start = time.monotonic() + + while True: + batch = self.client.batches.retrieve(batch_id) # type: ignore + + if batch.status in _TERMINAL_STATUSES: + logger.info("Batch {batch_id} reached status '{status}'.", batch_id=batch_id, status=batch.status) + return batch + + elapsed = time.monotonic() - start + if elapsed >= self.max_wait_seconds: + raise TimeoutError( + f"Batch {batch_id} did not complete within {self.max_wait_seconds}s. Last status: '{batch.status}'." + ) + + logger.debug( + "Batch {batch_id} status: '{status}'. Checking again in {interval}s.", + batch_id=batch_id, + status=batch.status, + interval=self.poll_interval, + ) + time.sleep(self.poll_interval) + + async def _poll_async(self, batch_id: str) -> Batch: + """Poll the batch status asynchronously, yielding to the event loop between checks.""" + start = time.monotonic() + + while True: + batch = await self.async_client.batches.retrieve(batch_id) # type: ignore + + if batch.status in _TERMINAL_STATUSES: + logger.info("Batch {batch_id} reached status '{status}'.", batch_id=batch_id, status=batch.status) + return batch + + elapsed = time.monotonic() - start + if elapsed >= self.max_wait_seconds: + raise TimeoutError( + f"Batch {batch_id} did not complete within {self.max_wait_seconds}s. Last status: '{batch.status}'." + ) + + logger.debug( + "Batch {batch_id} status: '{status}'. Checking again in {interval}s.", + batch_id=batch_id, + status=batch.status, + interval=self.poll_interval, + ) + await asyncio.sleep(self.poll_interval) + + +def _build_jsonl(message_sets: list[list[ChatMessage]], model: str, generation_kwargs: dict[str, Any]) -> io.BytesIO: + """ + Convert a list of conversations into a JSONL payload for the Batch API. + + Each conversation gets a ``custom_id`` like ``"request-0"``, ``"request-1"``, etc. + so we can match results back to the original order — OpenAI doesn't guarantee + output ordering. + """ + lines: list[str] = [] + for idx, messages in enumerate(message_sets): + openai_messages = [msg.to_openai_dict_format() for msg in messages] + request_body = {"model": model, "messages": openai_messages, **generation_kwargs} + line = {"custom_id": f"request-{idx}", "method": "POST", "url": _BATCH_ENDPOINT, "body": request_body} + lines.append(json.dumps(line)) + + return io.BytesIO("\n".join(lines).encode("utf-8")) + + +def _parse_results(output_text: str, expected_count: int) -> list[list[ChatMessage]]: + """ + Parse the batch output JSONL into ChatMessage objects, ordered by request index. + + The Batch API returns results as raw JSON dicts — not pydantic ``ChatCompletion`` objects — + so we extract the fields manually. The logic mirrors ``_convert_chat_completion_to_chat_message`` + from ``openai.py``, but works on plain dicts instead of SDK models. + """ + # Index results by custom_id for O(1) lookup + results_by_id: dict[str, dict[str, Any]] = {} + for line in output_text.strip().split("\n"): + if not line: + continue + entry = json.loads(line) + results_by_id[entry["custom_id"]] = entry + + # Reassemble in the original request order + replies: list[list[ChatMessage]] = [] + for idx in range(expected_count): + custom_id = f"request-{idx}" + entry = results_by_id.get(custom_id) + + if entry is None: + # Shouldn't happen with a completed batch, but don't crash + logger.warning("No result found for {custom_id} in batch output.", custom_id=custom_id) + replies.append([]) + continue + + # Per-request errors (e.g., content filter, invalid model) + if entry.get("error") is not None: + error_info = entry["error"] + logger.warning("Request {custom_id} failed: {error}", custom_id=custom_id, error=error_info) + replies.append([]) + continue + + response = entry.get("response", {}) + if response.get("status_code") != 200: + logger.warning( + "Request {custom_id} returned HTTP {status_code}.", + custom_id=custom_id, + status_code=response.get("status_code"), + ) + replies.append([]) + continue + + response_body = response.get("body", {}) + conversation_replies = _parse_choices(response_body) + replies.append(conversation_replies) + + return replies + + +def _parse_choices(response_body: dict[str, Any]) -> list[ChatMessage]: + """ + Turn the ``choices`` array from a raw completion dict into ChatMessage objects. + + This is the dict-based equivalent of ``_convert_chat_completion_to_chat_message`` + in openai.py. We can't reuse that function directly because it expects pydantic + ``ChatCompletion``/``Choice`` objects, but batch output is plain JSON. + """ + choices = response_body.get("choices", []) + model = response_body.get("model", "") + usage = response_body.get("usage") + + messages: list[ChatMessage] = [] + for choice in choices: + message_data = choice.get("message", {}) + text = message_data.get("content") + + # Parse tool calls if present — not in MVP scope, but costs nothing + # to handle correctly and avoids data loss if someone passes tool-enabled + # generation_kwargs anyway + tool_calls: list[ToolCall] = [] + for tc in message_data.get("tool_calls") or []: + fn = tc.get("function", {}) + try: + arguments = json.loads(fn.get("arguments", "{}")) + except json.JSONDecodeError: + logger.warning( + "Malformed JSON in tool call arguments. Tool call ID: {tc_id}, " + "Tool name: {tc_name}, Arguments: {tc_args}", + tc_id=tc.get("id"), + tc_name=fn.get("name"), + tc_args=fn.get("arguments"), + ) + arguments = {} + tool_calls.append(ToolCall(id=tc.get("id", ""), tool_name=fn.get("name", ""), arguments=arguments)) + + meta: dict[str, Any] = { + "model": model, + "index": choice.get("index", 0), + "finish_reason": choice.get("finish_reason"), + "usage": usage, + } + + messages.append(ChatMessage.from_assistant(text=text, tool_calls=tool_calls or None, meta=meta)) + + return messages + + +def _raise_on_failure(batch: Batch) -> None: + """Raise a RuntimeError if the batch didn't complete successfully.""" + if batch.status == "completed": + return + + error_messages: list[str] = [] + if batch.errors and batch.errors.data: + error_messages = [e.message for e in batch.errors.data if e.message] + + detail = "; ".join(error_messages) if error_messages else "No error details available." + raise RuntimeError(f"Batch {batch.id} finished with status '{batch.status}'. {detail}") + + +def _batch_meta(batch: Batch) -> dict[str, Any]: + """Extract batch-level metadata into a plain dict for the output.""" + return { + "batch_id": batch.id, + "status": batch.status, + "created_at": batch.created_at, + "completed_at": batch.completed_at, + "request_counts": batch.request_counts.model_dump() if batch.request_counts else None, + } diff --git a/haystack/components/generators/chat/openai_responses.py b/haystack/components/generators/chat/openai_responses.py index 45c11b7c7a..344a640018 100644 --- a/haystack/components/generators/chat/openai_responses.py +++ b/haystack/components/generators/chat/openai_responses.py @@ -264,10 +264,7 @@ def warm_up(self) -> None: # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. # https://github.com/openai/openai-python/blob/main/httpx2.md http_client = init_http_client(self.http_client_kwargs, async_client=False) - self.client = OpenAI( - http_client=http_client, # type: ignore[arg-type] - **self._client_kwargs(), - ) + self.client = OpenAI(http_client=http_client, **self._client_kwargs()) async def warm_up_async(self) -> None: # noqa: RUF029 """ @@ -278,10 +275,7 @@ async def warm_up_async(self) -> None: # noqa: RUF029 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. # https://github.com/openai/openai-python/blob/main/httpx2.md http_client = init_http_client(self.http_client_kwargs, async_client=True) - self.async_client = AsyncOpenAI( - http_client=http_client, # type: ignore[arg-type] - **self._client_kwargs(), - ) + self.async_client = AsyncOpenAI(http_client=http_client, **self._client_kwargs()) def close(self) -> None: """ diff --git a/haystack/components/generators/openai_image_generator.py b/haystack/components/generators/openai_image_generator.py index 1dc47dc79c..48e2529b88 100644 --- a/haystack/components/generators/openai_image_generator.py +++ b/haystack/components/generators/openai_image_generator.py @@ -108,10 +108,7 @@ def warm_up(self) -> None: # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. # https://github.com/openai/openai-python/blob/main/httpx2.md http_client = init_http_client(self.http_client_kwargs, async_client=False) - self.client = OpenAI( - http_client=http_client, # type: ignore[arg-type] - **self._client_kwargs(), - ) + self.client = OpenAI(http_client=http_client, **self._client_kwargs()) async def warm_up_async(self) -> None: # noqa: RUF029 """ @@ -121,10 +118,7 @@ async def warm_up_async(self) -> None: # noqa: RUF029 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime. # https://github.com/openai/openai-python/blob/main/httpx2.md http_client = init_http_client(self.http_client_kwargs, async_client=True) - self.async_client = AsyncOpenAI( - http_client=http_client, # type: ignore[arg-type] - **self._client_kwargs(), - ) + self.async_client = AsyncOpenAI(http_client=http_client, **self._client_kwargs()) def close(self) -> None: """ diff --git a/haystack/token_counters/openai_counter.py b/haystack/token_counters/openai_counter.py index f01ca19113..a7b99d9b2a 100644 --- a/haystack/token_counters/openai_counter.py +++ b/haystack/token_counters/openai_counter.py @@ -87,7 +87,7 @@ def warm_up(self) -> None: base_url=self.api_base_url, timeout=timeout, max_retries=max_retries, - http_client=http_client, # type: ignore[arg-type] + http_client=http_client, ) def count(self, messages: list[ChatMessage], tools: ToolsType | None = None) -> int: diff --git a/releasenotes/notes/Add-OpenAIBatchChatGenerator-for-processing-batch-api-requests-53eb6eded8107ebc.yaml b/releasenotes/notes/Add-OpenAIBatchChatGenerator-for-processing-batch-api-requests-53eb6eded8107ebc.yaml new file mode 100644 index 0000000000..460cb7e9ba --- /dev/null +++ b/releasenotes/notes/Add-OpenAIBatchChatGenerator-for-processing-batch-api-requests-53eb6eded8107ebc.yaml @@ -0,0 +1,4 @@ +--- +features: + - | + Added ``OpenAIBatchChatGenerator``, a new component that allows users to submit multiple chat completion requests to OpenAI's Batch API for asynchronous processing, enabling up to 50% cost savings for high-throughput, latency-insensitive workloads. diff --git a/test/components/generators/chat/test_openai.py b/test/components/generators/chat/test_openai.py index c33c1a1ed9..ce9584e260 100644 --- a/test/components/generators/chat/test_openai.py +++ b/test/components/generators/chat/test_openai.py @@ -1520,15 +1520,9 @@ def chat_completion_chunks(): prompt_tokens=282, total_tokens=324, completion_tokens_details=CompletionTokensDetails( - accepted_prediction_tokens=0, - audio_tokens=0, - reasoning_tokens=0, - rejected_prediction_tokens=0, - text_tokens=42, - ), - prompt_tokens_details=PromptTokensDetails( - audio_tokens=0, cached_tokens=0, cache_write_tokens=0, image_tokens=0, text_tokens=282 + accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0 ), + prompt_tokens_details=PromptTokensDetails(audio_tokens=0, cached_tokens=0), ), ), ] @@ -1734,15 +1728,8 @@ def streaming_chunks(): "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0, - "text_tokens": 42, - }, - "prompt_tokens_details": { - "audio_tokens": 0, - "cached_tokens": 0, - "cache_write_tokens": 0, - "image_tokens": 0, - "text_tokens": 282, }, + "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0, "cache_write_tokens": None}, }, }, ), @@ -2024,15 +2011,8 @@ def test_handle_stream_response( "audio_tokens": 0, "reasoning_tokens": 0, "rejected_prediction_tokens": 0, - "text_tokens": 42, - }, - "prompt_tokens_details": { - "audio_tokens": 0, - "cached_tokens": 0, - "cache_write_tokens": 0, - "image_tokens": 0, - "text_tokens": 282, }, + "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0, "cache_write_tokens": None}, } def test_convert_usage_chunk_to_streaming_chunk(self) -> None: diff --git a/test/components/generators/chat/test_openai_batch.py b/test/components/generators/chat/test_openai_batch.py new file mode 100644 index 0000000000..eba9d00d34 --- /dev/null +++ b/test/components/generators/chat/test_openai_batch.py @@ -0,0 +1,703 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import json +from typing import Literal +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from openai.types import Batch, BatchError, BatchRequestCounts +from openai.types.batch import Errors + +from haystack.components.generators.chat.openai_batch import ( + OpenAIBatchChatGenerator, + _batch_meta, + _build_jsonl, + _parse_choices, + _parse_results, + _raise_on_failure, +) +from haystack.dataclasses import ChatMessage +from haystack.utils.auth import Secret + + +@pytest.fixture +def component(): + """A default-configured component for testing.""" + return OpenAIBatchChatGenerator(api_key=Secret.from_token("test-api-key")) + + +@pytest.fixture +def custom_component(): + """A fully-customized component to verify all init params are stored.""" + return OpenAIBatchChatGenerator( + api_key=Secret.from_token("custom-key"), + model="gpt-5", + api_base_url="https://custom.openai.com", + organization="org-123", + generation_kwargs={"temperature": 0.7, "max_completion_tokens": 500}, + timeout=60.0, + max_retries=3, + poll_interval=10.0, + max_wait_seconds=3600.0, + completion_window="24h", + http_client_kwargs={"verify": False}, + ) + + +@pytest.fixture +def single_conversation(): + """One conversation with a system prompt and a user message.""" + return [ + ChatMessage.from_system("You are a helpful assistant."), + ChatMessage.from_user("What is the capital of France?"), + ] + + +@pytest.fixture +def two_conversations(single_conversation): + """Two separate conversations for batch processing.""" + return [single_conversation, [ChatMessage.from_user("What is 2+2?")]] + + +BatchStatus = Literal[ + "validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled" +] + + +def _make_batch( + batch_id: str = "batch_abc", + status: BatchStatus = "completed", + output_file_id: str | None = "file-out-123", + error_file_id: str | None = None, + errors: Errors | None = None, +) -> Batch: + """Build a Batch object for mocking, without hitting the API.""" + return Batch( + id=batch_id, + object="batch", + endpoint="/v1/chat/completions", + input_file_id="file-in-456", + completion_window="24h", + status=status, + output_file_id=output_file_id, + error_file_id=error_file_id, + errors=errors, + created_at=1700000000, + completed_at=1700001000 if status == "completed" else None, + request_counts=BatchRequestCounts(completed=2, failed=0, total=2), + ) + + +def _make_output_jsonl(*responses: dict) -> str: + """ + Build a batch output JSONL string from a list of response body dicts. + + Each dict should look like a chat completion body: + {"choices": [{"message": {"content": "Paris"}, ...}], "model": "gpt-5-mini", ...} + """ + lines = [] + for idx, body in enumerate(responses): + entry = { + "id": f"batch_req_{idx}", + "custom_id": f"request-{idx}", + "response": {"status_code": 200, "request_id": f"req_{idx}", "body": body}, + "error": None, + } + lines.append(json.dumps(entry)) + return "\n".join(lines) + + +def _simple_completion_body(text: str, model: str = "gpt-5-mini") -> dict: + """A minimal chat completion body dict for testing.""" + return { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1700000500, + "model": model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + + +class TestOpenAIBatchChatGeneratorInit: + def test_default_params(self, component): + assert component.model == "gpt-5-mini" + assert component.generation_kwargs == {} + assert component.poll_interval == 30.0 + assert component.max_wait_seconds == 86400.0 + assert component.completion_window == "24h" + assert component.api_base_url is None + assert component.organization is None + assert component.timeout is None + assert component.max_retries is None + assert component.http_client_kwargs is None + assert component.client is None + assert component.async_client is None + + def test_custom_params(self, custom_component): + assert custom_component.model == "gpt-5" + assert custom_component.api_base_url == "https://custom.openai.com" + assert custom_component.organization == "org-123" + assert custom_component.generation_kwargs == {"temperature": 0.7, "max_completion_tokens": 500} + assert custom_component.timeout == 60.0 + assert custom_component.max_retries == 3 + assert custom_component.poll_interval == 10.0 + assert custom_component.max_wait_seconds == 3600.0 + assert custom_component.http_client_kwargs == {"verify": False} + + +class TestOpenAIBatchChatGeneratorSerialization: + def test_to_dict(self): + # Token-based secrets can't be serialized (by design), so use env var + gen = OpenAIBatchChatGenerator(api_key=Secret.from_env_var("OPENAI_API_KEY")) + result = gen.to_dict() + assert result["type"] == "haystack.components.generators.chat.openai_batch.OpenAIBatchChatGenerator" + params = result["init_parameters"] + assert params["model"] == "gpt-5-mini" + assert params["poll_interval"] == 30.0 + assert params["max_wait_seconds"] == 86400.0 + assert params["completion_window"] == "24h" + assert params["generation_kwargs"] == {} + assert params["api_base_url"] is None + + def test_to_dict_custom(self): + gen = OpenAIBatchChatGenerator( + api_key=Secret.from_env_var("MY_KEY"), + model="gpt-5", + api_base_url="https://custom.openai.com", + organization="org-123", + generation_kwargs={"temperature": 0.7, "max_completion_tokens": 500}, + timeout=60.0, + max_retries=3, + poll_interval=10.0, + max_wait_seconds=3600.0, + ) + result = gen.to_dict() + params = result["init_parameters"] + assert params["model"] == "gpt-5" + assert params["api_base_url"] == "https://custom.openai.com" + assert params["organization"] == "org-123" + assert params["generation_kwargs"] == {"temperature": 0.7, "max_completion_tokens": 500} + assert params["timeout"] == 60.0 + assert params["max_retries"] == 3 + assert params["poll_interval"] == 10.0 + assert params["max_wait_seconds"] == 3600.0 + + def test_from_dict_round_trip(self): + original = OpenAIBatchChatGenerator(api_key=Secret.from_env_var("OPENAI_API_KEY")) + serialized = original.to_dict() + restored = OpenAIBatchChatGenerator.from_dict(serialized) + assert restored.model == original.model + assert restored.poll_interval == original.poll_interval + assert restored.max_wait_seconds == original.max_wait_seconds + assert restored.completion_window == original.completion_window + assert restored.generation_kwargs == original.generation_kwargs + + +class TestOpenAIBatchChatGeneratorLifecycle: + def test_warm_up_creates_client(self, component): + assert component.client is None + with patch("haystack.components.generators.chat.openai_batch.OpenAI"): + component.warm_up() + assert component.client is not None + + def test_warm_up_idempotent(self, component): + with patch("haystack.components.generators.chat.openai_batch.OpenAI") as mock_cls: + component.warm_up() + first_client = component.client + component.warm_up() + # Should not create a second client + assert component.client is first_client + mock_cls.assert_called_once() + + def test_close_releases_client(self, component): + with patch("haystack.components.generators.chat.openai_batch.OpenAI"): + component.warm_up() + assert component.client is not None + component.close() + assert component.client is None + + def test_close_when_no_client(self, component): + # Should not raise + component.close() + + @pytest.mark.asyncio + async def test_close_async_releases_client(self, component): + with patch("haystack.components.generators.chat.openai_batch.AsyncOpenAI"): + await component.warm_up_async() + assert component.async_client is not None + + # Ensure we mock the close method as AsyncMock + component.async_client.close = AsyncMock() + await component.close_async() + assert component.async_client is None + + @pytest.mark.asyncio + async def test_close_async_when_no_client(self, component): + # Should not raise + await component.close_async() + + def test_get_telemetry_data(self, component): + assert component._get_telemetry_data() == {"model": "gpt-5-mini"} + + +class TestBuildJsonl: + def test_single_conversation(self, single_conversation): + result = _build_jsonl([single_conversation], "gpt-5-mini", {}) + content = result.read().decode("utf-8") + lines = content.strip().split("\n") + assert len(lines) == 1 + + parsed = json.loads(lines[0]) + assert parsed["custom_id"] == "request-0" + assert parsed["method"] == "POST" + assert parsed["url"] == "/v1/chat/completions" + assert parsed["body"]["model"] == "gpt-5-mini" + assert len(parsed["body"]["messages"]) == 2 + + def test_multiple_conversations(self, two_conversations): + result = _build_jsonl(two_conversations, "gpt-5", {"temperature": 0.5}) + content = result.read().decode("utf-8") + lines = content.strip().split("\n") + assert len(lines) == 2 + + # First conversation + line0 = json.loads(lines[0]) + assert line0["custom_id"] == "request-0" + assert line0["body"]["model"] == "gpt-5" + assert line0["body"]["temperature"] == 0.5 + assert len(line0["body"]["messages"]) == 2 + + # Second conversation + line1 = json.loads(lines[1]) + assert line1["custom_id"] == "request-1" + assert len(line1["body"]["messages"]) == 1 + + def test_generation_kwargs_included(self, single_conversation): + kwargs = {"temperature": 0.3, "max_completion_tokens": 100} + result = _build_jsonl([single_conversation], "gpt-5-mini", kwargs) + parsed = json.loads(result.read().decode("utf-8")) + assert parsed["body"]["temperature"] == 0.3 + assert parsed["body"]["max_completion_tokens"] == 100 + + +class TestParseResults: + def test_single_result(self): + output = _make_output_jsonl(_simple_completion_body("Paris")) + replies = _parse_results(output, expected_count=1) + + assert len(replies) == 1 + assert len(replies[0]) == 1 + assert replies[0][0].text == "Paris" + assert replies[0][0].meta["model"] == "gpt-5-mini" + assert replies[0][0].meta["finish_reason"] == "stop" + assert replies[0][0].meta["usage"]["total_tokens"] == 15 + + def test_multiple_results(self): + output = _make_output_jsonl(_simple_completion_body("Paris"), _simple_completion_body("4")) + replies = _parse_results(output, expected_count=2) + + assert len(replies) == 2 + assert replies[0][0].text == "Paris" + assert replies[1][0].text == "4" + + def test_results_reordered(self): + """The batch output might not be in the same order as the input.""" + line0 = { + "id": "req_1", + "custom_id": "request-1", + "response": {"status_code": 200, "request_id": "r1", "body": _simple_completion_body("second")}, + "error": None, + } + line1 = { + "id": "req_0", + "custom_id": "request-0", + "response": {"status_code": 200, "request_id": "r0", "body": _simple_completion_body("first")}, + "error": None, + } + output = json.dumps(line0) + "\n" + json.dumps(line1) + replies = _parse_results(output, expected_count=2) + + # Should be re-sorted by custom_id index + assert replies[0][0].text == "first" + assert replies[1][0].text == "second" + + def test_missing_result_returns_empty(self): + """If a result is missing for a custom_id, we get an empty list — not a crash.""" + output = _make_output_jsonl(_simple_completion_body("Paris")) + replies = _parse_results(output, expected_count=2) + + assert len(replies) == 2 + assert replies[0][0].text == "Paris" + assert replies[1] == [] # missing request-1 + + def test_per_request_error(self): + """Per-request errors should produce empty replies, not crash.""" + error_line = { + "id": "req_0", + "custom_id": "request-0", + "response": None, + "error": {"code": "content_filter", "message": "Content was blocked."}, + } + output = json.dumps(error_line) + replies = _parse_results(output, expected_count=1) + + assert len(replies) == 1 + assert replies[0] == [] + + def test_results_with_empty_lines(self): + entry1 = {"custom_id": "request-0", "response": {"status_code": 200, "body": _simple_completion_body("A")}} + entry2 = {"custom_id": "request-1", "response": {"status_code": 200, "body": _simple_completion_body("B")}} + + # Put an empty line right in the middle + output = json.dumps(entry1) + "\n\n\n" + json.dumps(entry2) + + replies = _parse_results(output, expected_count=2) + assert len(replies) == 2 + assert replies[0][0].text == "A" + assert replies[1][0].text == "B" + + def test_per_request_http_error(self): + # A valid JSONL entry but with an HTTP non-200 status_code in the 'response' + entry = {"custom_id": "request-0", "response": {"status_code": 500, "body": {}}} + output = json.dumps(entry) + "\n" + replies = _parse_results(output, expected_count=1) + assert len(replies) == 1 + assert replies[0] == [] + + +class TestParseChoices: + def test_basic_text(self): + body = _simple_completion_body("Hello!") + messages = _parse_choices(body) + + assert len(messages) == 1 + assert messages[0].text == "Hello!" + assert messages[0].meta["finish_reason"] == "stop" + + def test_with_tool_calls(self): + """Tool calls in the response should be parsed even though tools aren't in MVP scope.""" + body = { + "model": "gpt-5-mini", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + }, + "finish_reason": "tool_calls", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}, + } + messages = _parse_choices(body) + + assert len(messages) == 1 + assert messages[0].tool_calls is not None + assert len(messages[0].tool_calls) == 1 + assert messages[0].tool_calls[0].tool_name == "get_weather" + assert messages[0].tool_calls[0].arguments == {"city": "Paris"} + + def test_multiple_choices(self): + """When n>1, the API returns multiple choices per request.""" + body = { + "model": "gpt-5-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "Answer A"}, "finish_reason": "stop"}, + {"index": 1, "message": {"role": "assistant", "content": "Answer B"}, "finish_reason": "stop"}, + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 10, "total_tokens": 20}, + } + messages = _parse_choices(body) + assert len(messages) == 2 + assert messages[0].text == "Answer A" + assert messages[1].text == "Answer B" + + def test_malformed_tool_call_json(self): + body = { + "model": "gpt-5-mini", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"'}, # Missing brace + } + ], + }, + "finish_reason": "tool_calls", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}, + } + messages = _parse_choices(body) + assert len(messages) == 1 + assert messages[0].tool_calls is not None + assert len(messages[0].tool_calls) == 1 + assert messages[0].tool_calls[0].arguments == {} + + +class TestRaiseOnFailure: + def test_completed_does_not_raise(self): + batch = _make_batch(status="completed") + _raise_on_failure(batch) # should not raise + + def test_failed_raises(self): + batch = _make_batch(status="failed") + with pytest.raises(RuntimeError, match="finished with status 'failed'"): + _raise_on_failure(batch) + + def test_expired_raises(self): + batch = _make_batch(status="expired") + with pytest.raises(RuntimeError, match="finished with status 'expired'"): + _raise_on_failure(batch) + + def test_error_details_included(self): + errors = Errors( + object="list", data=[BatchError(code="invalid_model", message="Model not found", line=1, param="model")] + ) + batch = _make_batch(status="failed", errors=errors) + with pytest.raises(RuntimeError, match="Model not found"): + _raise_on_failure(batch) + + +class TestBatchMeta: + def test_extracts_metadata(self): + batch = _make_batch() + meta = _batch_meta(batch) + assert meta["batch_id"] == "batch_abc" + assert meta["status"] == "completed" + assert meta["created_at"] == 1700000000 + assert meta["completed_at"] == 1700001000 + assert meta["request_counts"]["total"] == 2 + + +class TestOpenAIBatchChatGeneratorRun: + def test_run_single_conversation(self, component, single_conversation): + output_jsonl = _make_output_jsonl(_simple_completion_body("Paris")) + + mock_client = MagicMock() + mock_client.files.create.return_value = MagicMock(id="file-in-123") + mock_client.batches.create.return_value = _make_batch(status="validating") + mock_client.batches.retrieve.return_value = _make_batch(status="completed") + mock_client.files.content.return_value = MagicMock(text=output_jsonl) + + with patch("haystack.components.generators.chat.openai_batch.OpenAI", return_value=mock_client): + component.warm_up() + + result = component.run(message_sets=[single_conversation]) + + assert len(result["replies"]) == 1 + assert result["replies"][0][0].text == "Paris" + assert result["meta"]["batch_id"] == "batch_abc" + assert result["meta"]["status"] == "completed" + + def test_run_multiple_conversations(self, component, two_conversations): + output_jsonl = _make_output_jsonl(_simple_completion_body("Paris"), _simple_completion_body("4")) + + mock_client = MagicMock() + mock_client.files.create.return_value = MagicMock(id="file-in-123") + mock_client.batches.create.return_value = _make_batch(status="validating") + mock_client.batches.retrieve.return_value = _make_batch(status="completed") + mock_client.files.content.return_value = MagicMock(text=output_jsonl) + + with patch("haystack.components.generators.chat.openai_batch.OpenAI", return_value=mock_client): + component.warm_up() + + result = component.run(message_sets=two_conversations) + + assert len(result["replies"]) == 2 + assert result["replies"][0][0].text == "Paris" + assert result["replies"][1][0].text == "4" + + def test_run_empty_input(self, component): + result = component.run(message_sets=[]) + assert result == {"replies": [], "meta": {}} + + def test_run_merges_generation_kwargs(self, component, single_conversation): + """Runtime generation_kwargs should merge with and override init kwargs.""" + component.generation_kwargs = {"temperature": 0.5, "max_completion_tokens": 100} + output_jsonl = _make_output_jsonl(_simple_completion_body("test")) + + mock_client = MagicMock() + mock_client.files.create.return_value = MagicMock(id="file-in-123") + mock_client.batches.create.return_value = _make_batch(status="completed") + mock_client.batches.retrieve.return_value = _make_batch(status="completed") + mock_client.files.content.return_value = MagicMock(text=output_jsonl) + + with patch("haystack.components.generators.chat.openai_batch.OpenAI", return_value=mock_client): + component.warm_up() + + # Runtime kwarg overrides temperature but keeps max_completion_tokens + component.run(message_sets=[single_conversation], generation_kwargs={"temperature": 0.9}) + + # Inspect the JSONL that was uploaded + upload_call = mock_client.files.create.call_args + file_tuple = upload_call.kwargs.get("file") or upload_call[1].get("file") + jsonl_content = file_tuple[1].read().decode("utf-8") + parsed = json.loads(jsonl_content) + + assert parsed["body"]["temperature"] == 0.9 # overridden + assert parsed["body"]["max_completion_tokens"] == 100 # preserved from init + + def test_run_batch_failed_raises(self, component, single_conversation): + mock_client = MagicMock() + mock_client.files.create.return_value = MagicMock(id="file-in-123") + mock_client.batches.create.return_value = _make_batch(status="validating") + mock_client.batches.retrieve.return_value = _make_batch(status="failed") + + with patch("haystack.components.generators.chat.openai_batch.OpenAI", return_value=mock_client): + component.warm_up() + + with pytest.raises(RuntimeError, match="finished with status 'failed'"): + component.run(message_sets=[single_conversation]) + + def test_run_batch_expired_raises(self, component, single_conversation): + mock_client = MagicMock() + mock_client.files.create.return_value = MagicMock(id="file-in-123") + mock_client.batches.create.return_value = _make_batch(status="validating") + mock_client.batches.retrieve.return_value = _make_batch(status="expired") + + with patch("haystack.components.generators.chat.openai_batch.OpenAI", return_value=mock_client): + component.warm_up() + + with pytest.raises(RuntimeError, match="finished with status 'expired'"): + component.run(message_sets=[single_conversation]) + + def test_run_timeout_raises(self, component, single_conversation): + """When max_wait_seconds is exceeded, a TimeoutError should be raised.""" + component.max_wait_seconds = 0.0 # Immediately times out + component.poll_interval = 0.01 + + mock_client = MagicMock() + mock_client.files.create.return_value = MagicMock(id="file-in-123") + mock_client.batches.create.return_value = _make_batch(status="in_progress") + # Never reaches a terminal status + mock_client.batches.retrieve.return_value = _make_batch(status="in_progress") + + with patch("haystack.components.generators.chat.openai_batch.OpenAI", return_value=mock_client): + component.warm_up() + + with pytest.raises(TimeoutError, match="did not complete within"): + component.run(message_sets=[single_conversation]) + + def test_run_polls_until_completed(self, component, single_conversation): + """Verify polling continues through non-terminal statuses.""" + component.poll_interval = 0.01 # Speed up tests + + output_jsonl = _make_output_jsonl(_simple_completion_body("result")) + + mock_client = MagicMock() + mock_client.files.create.return_value = MagicMock(id="file-in-123") + mock_client.batches.create.return_value = _make_batch(status="validating") + # Simulate: validating → in_progress → finalizing → completed + mock_client.batches.retrieve.side_effect = [ + _make_batch(status="validating"), + _make_batch(status="in_progress"), + _make_batch(status="finalizing"), + _make_batch(status="completed"), + ] + mock_client.files.content.return_value = MagicMock(text=output_jsonl) + + with patch("haystack.components.generators.chat.openai_batch.OpenAI", return_value=mock_client): + component.warm_up() + + result = component.run(message_sets=[single_conversation]) + assert result["replies"][0][0].text == "result" + assert mock_client.batches.retrieve.call_count == 4 + + +class TestOpenAIBatchChatGeneratorRunAsync: + @pytest.mark.asyncio + async def test_run_async_basic(self, component, single_conversation): + output_jsonl = _make_output_jsonl(_simple_completion_body("Paris")) + + mock_client = AsyncMock() + mock_client.files.create.return_value = MagicMock(id="file-in-123") + mock_client.batches.create.return_value = _make_batch(status="validating") + mock_client.batches.retrieve.return_value = _make_batch(status="completed") + mock_client.files.content.return_value = MagicMock(text=output_jsonl) + + with patch("haystack.components.generators.chat.openai_batch.AsyncOpenAI", return_value=mock_client): + await component.warm_up_async() + + result = await component.run_async(message_sets=[single_conversation]) + + assert len(result["replies"]) == 1 + assert result["replies"][0][0].text == "Paris" + assert result["meta"]["batch_id"] == "batch_abc" + + @pytest.mark.asyncio + async def test_run_async_empty_input(self, component): + result = await component.run_async(message_sets=[]) + assert result == {"replies": [], "meta": {}} + + @pytest.mark.asyncio + async def test_run_async_batch_failed_raises(self, component, single_conversation): + mock_client = AsyncMock() + mock_client.files.create.return_value = MagicMock(id="file-in-123") + mock_client.batches.create.return_value = _make_batch(status="validating") + mock_client.batches.retrieve.return_value = _make_batch(status="failed") + + with patch("haystack.components.generators.chat.openai_batch.AsyncOpenAI", return_value=mock_client): + await component.warm_up_async() + + with pytest.raises(RuntimeError, match="finished with status 'failed'"): + await component.run_async(message_sets=[single_conversation]) + + @pytest.mark.asyncio + async def test_run_async_timeout_raises(self, component, single_conversation): + component.max_wait_seconds = 0.1 + component.poll_interval = 0.5 # larger than timeout so it fails on first sleep check + + mock_client = AsyncMock() + mock_client.files.create.return_value = MagicMock(id="file-in-123") + mock_client.batches.create.return_value = _make_batch(status="validating") + mock_client.batches.retrieve.return_value = _make_batch(status="in_progress") + + with patch("haystack.components.generators.chat.openai_batch.AsyncOpenAI", return_value=mock_client): + await component.warm_up_async() + + with pytest.raises(TimeoutError, match="did not complete within"): + await component.run_async(message_sets=[single_conversation]) + + @pytest.mark.asyncio + @patch("asyncio.sleep") + async def test_run_async_polls_until_completed(self, mock_sleep, component, single_conversation): + output_jsonl = _make_output_jsonl(_simple_completion_body("result")) + + mock_client = AsyncMock() + mock_client.files.create.return_value = MagicMock(id="file-in-123") + mock_client.batches.create.return_value = _make_batch(status="validating") + + # Simulate: validating -> in_progress -> completed + mock_client.batches.retrieve.side_effect = [ + _make_batch(status="validating"), + _make_batch(status="in_progress"), + _make_batch(status="completed"), + ] + mock_client.files.content.return_value = MagicMock(text=output_jsonl) + + with patch("haystack.components.generators.chat.openai_batch.AsyncOpenAI", return_value=mock_client): + await component.warm_up_async() + + result = await component.run_async(message_sets=[single_conversation]) + + assert result["replies"][0][0].text == "result" + assert mock_client.batches.retrieve.call_count == 3 + assert mock_sleep.call_count == 2