wehub/deepset-ai--haystack
1
0
Fork 0
Code Issues Pull Requests Actions 8 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
deepset-ai--haystack/test/components/generators/chat/test_openai.py
T
wehub-resource-sync c56bef871b
Sync docs with Docusaurus / sync (push) Waiting to run
Details
Tests / Check if changed (push) Waiting to run
Details
Tests / format (push) Blocked by required conditions
Details
Tests / check-imports (push) Blocked by required conditions
Details
Tests / Unit / macos-latest (push) Blocked by required conditions
Details
Tests / Unit / ubuntu-latest (push) Blocked by required conditions
Details
Tests / Unit / windows-latest (push) Blocked by required conditions
Details
Tests / mypy (push) Blocked by required conditions
Details
Tests / Integration / ubuntu-latest (push) Blocked by required conditions
Details
Tests / Integration / macos-latest (push) Blocked by required conditions
Details
Tests / Integration / windows-latest (push) Blocked by required conditions
Details
Tests / notify-slack-on-failure (push) Blocked by required conditions
Details
Tests / Mark tests as completed (push) Blocked by required conditions
Details
Docker image release / Build base image (push) Waiting to run
Details
CodeQL / Analyze (python) (push) Has been cancelled
Details
Update Platform Components Table / update (push) Has been cancelled
Details
chore: import upstream snapshot with attribution
2026-07-13 13:22:28 +08:00

2201 lines
88 KiB
Python
Raw Permalink Blame History

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
import json
import logging
import os
from datetime import datetime
from typing import Any
from unittest.mock import ANY, AsyncMock, MagicMock, patch
import pytest
from openai import OpenAIError
from openai.types.chat import (
ChatCompletion,
ChatCompletionChunk,
ChatCompletionMessage,
ChatCompletionMessageFunctionToolCall,
ParsedChatCompletion,
ParsedChatCompletionMessage,
ParsedChoice,
ParsedFunction,
ParsedFunctionToolCall,
chat_completion_chunk,
)
from openai.types.chat.chat_completion import Choice
from openai.types.chat.chat_completion_chunk import ChoiceDelta, ChoiceDeltaToolCall, ChoiceDeltaToolCallFunction
from openai.types.chat.chat_completion_message_function_tool_call import Function
from openai.types.completion_usage import CompletionTokensDetails, CompletionUsage, PromptTokensDetails
from pydantic import BaseModel
import haystack.components.generators.chat.openai as openai_chat_module
from haystack import component
from haystack.components.generators.chat.openai import (
OpenAIChatGenerator,
_check_finish_reason,
_convert_chat_completion_chunk_to_streaming_chunk,
_make_schema_strict,
)
from haystack.components.generators.utils import print_streaming_chunk
from haystack.dataclasses import (
ChatMessage,
ChatRole,
FileContent,
ImageContent,
StreamingChunk,
ToolCall,
ToolCallDelta,
)
from haystack.tools import ComponentTool, Tool
from haystack.tools.toolset import Toolset
from haystack.utils.auth import Secret
class CalendarEvent(BaseModel):
event_name: str
event_date: str
event_location: str
@pytest.fixture
def calendar_event_model():
return CalendarEvent
@pytest.fixture
def chat_messages():
return [
ChatMessage.from_system("You are a helpful assistant"),
ChatMessage.from_user("What's the capital of France"),
]
@pytest.fixture
def mock_chat_completion_chunk_with_tools(openai_mock_stream):
"""
Mock the OpenAI API completion chunk response and reuse it for tests
"""
with patch("openai.resources.chat.completions.Completions.create") as mock_chat_completion_create:
completion = ChatCompletionChunk(
id="foo",
model="gpt-4",
object="chat.completion.chunk",
choices=[
chat_completion_chunk.Choice(
finish_reason="tool_calls",
logprobs=None,
index=0,
delta=chat_completion_chunk.ChoiceDelta(
role="assistant",
tool_calls=[
chat_completion_chunk.ChoiceDeltaToolCall(
index=0,
id="123",
type="function",
function=chat_completion_chunk.ChoiceDeltaToolCallFunction(
name="weather", arguments='{"city": "Paris"}'
),
)
],
),
)
],
created=int(datetime.now().timestamp()),
)
mock_chat_completion_create.return_value = openai_mock_stream(
completion, cast_to=None, response=None, client=None
)
yield mock_chat_completion_create
def weather_function(city: str) -> dict[str, Any]:
weather_info = {
"Berlin": {"weather": "mostly sunny", "temperature": 7, "unit": "celsius"},
"Paris": {"weather": "mostly cloudy", "temperature": 8, "unit": "celsius"},
"Rome": {"weather": "sunny", "temperature": 14, "unit": "celsius"},
}
return weather_info.get(city, {"weather": "unknown", "temperature": 0, "unit": "celsius"})
# mock chat completions with structured outputs
@pytest.fixture
def mock_parsed_chat_completion():
with patch("openai.resources.chat.completions.Completions.parse") as mock_chat_completion_parse:
completion = ParsedChatCompletion[CalendarEvent](
id="json_foo",
model="gpt-5-mini",
object="chat.completion",
choices=[
ParsedChoice[CalendarEvent](
finish_reason="stop",
index=0,
message=ParsedChatCompletionMessage[CalendarEvent](
content='{"event_name":"Team Meeting","event_date":"2024-03-15",'
'"event_location":"Conference Room A"}',
refusal=None,
role="assistant",
annotations=[],
audio=None,
function_call=None,
tool_calls=None,
parsed=CalendarEvent(
event_name="Team Meeting", event_date="2024-03-15", event_location="Conference Room A"
),
),
)
],
created=1757328264,
usage=CompletionUsage(completion_tokens=29, prompt_tokens=86, total_tokens=115),
)
mock_chat_completion_parse.return_value = completion
yield mock_chat_completion_parse
@component
class MessageExtractor:
@component.output_types(messages=list[str], meta=dict[str, Any])
def run(self, messages: list[ChatMessage], meta: dict[str, Any] | None = None) -> dict[str, Any]:
"""
Extracts the text content of ChatMessage objects
:param messages: List of Haystack ChatMessage objects
:param meta: Optional metadata to include in the response.
:returns:
A dictionary with keys "messages" and "meta".
"""
if meta is None:
meta = {}
return {"messages": [m.text for m in messages], "meta": meta}
@pytest.fixture
def tools():
weather_tool = Tool(
name="weather",
description="useful to determine the weather in a given location",
parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
function=weather_function,
)
# We add a tool that has a more complex parameter signature
message_extractor_tool = ComponentTool(
component=MessageExtractor(),
name="message_extractor",
description="Useful for returning the text content of ChatMessage objects",
)
return [weather_tool, message_extractor_tool]
class TestOpenAIChatGenerator:
def test_supported_models(self):
"""SUPPORTED_MODELS is a non-empty list of strings."""
models = OpenAIChatGenerator.SUPPORTED_MODELS
assert isinstance(models, list)
assert len(models) > 0
assert all(isinstance(m, str) for m in models)
def test_init_default(self, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
component = OpenAIChatGenerator()
assert component.api_key.resolve_value() == "test-api-key"
assert component.model == "gpt-5-mini"
assert component.streaming_callback is None
assert not component.generation_kwargs
assert component.timeout is None
assert component.max_retries is None
assert component.tools is None
assert not component.tools_strict
assert component.http_client_kwargs is None
assert component.client is None
assert component.async_client is None
def test_init_fail_with_duplicate_tool_names(self, monkeypatch, tools):
monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
duplicate_tools = [tools[0], tools[0]]
with pytest.raises(ValueError):
OpenAIChatGenerator(tools=duplicate_tools)
def test_init_with_parameters(self, monkeypatch):
tool = Tool(name="name", description="description", parameters={"x": {"type": "string"}}, function=lambda x: x)
monkeypatch.setenv("OPENAI_TIMEOUT", "100")
monkeypatch.setenv("OPENAI_MAX_RETRIES", "10")
component = OpenAIChatGenerator(
api_key=Secret.from_token("test-api-key"),
streaming_callback=print_streaming_chunk,
api_base_url="test-base-url",
generation_kwargs={"max_completion_tokens": 10, "some_test_param": "test-params"},
timeout=40.0,
max_retries=1,
tools=[tool],
tools_strict=True,
http_client_kwargs={"proxy": "http://example.com:8080", "verify": False},
)
assert component.api_key.resolve_value() == "test-api-key"
assert component.model == "gpt-5-mini"
assert component.streaming_callback is print_streaming_chunk
assert component.generation_kwargs == {"max_completion_tokens": 10, "some_test_param": "test-params"}
assert component.timeout == 40.0
assert component.max_retries == 1
assert component.tools == [tool]
assert component.tools_strict
assert component.http_client_kwargs == {"proxy": "http://example.com:8080", "verify": False}
assert component.client is None
assert component.async_client is None
def test_init_with_parameters_and_env_vars(self, monkeypatch):
monkeypatch.setenv("OPENAI_TIMEOUT", "100")
monkeypatch.setenv("OPENAI_MAX_RETRIES", "10")
component = OpenAIChatGenerator(
api_key=Secret.from_token("test-api-key"),
streaming_callback=print_streaming_chunk,
api_base_url="test-base-url",
generation_kwargs={"max_completion_tokens": 10, "some_test_param": "test-params"},
)
assert component.api_key.resolve_value() == "test-api-key"
assert component.model == "gpt-5-mini"
assert component.streaming_callback is print_streaming_chunk
assert component.generation_kwargs == {"max_completion_tokens": 10, "some_test_param": "test-params"}
assert component.timeout is None
assert component.max_retries is None
assert component.client is None
assert component.async_client is None
def test_to_dict_default(self, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
component = OpenAIChatGenerator()
data = component.to_dict()
assert data == {
"type": "haystack.components.generators.chat.openai.OpenAIChatGenerator",
"init_parameters": {
"api_key": {"env_vars": ["OPENAI_API_KEY"], "strict": True, "type": "env_var"},
"model": "gpt-5-mini",
"organization": None,
"streaming_callback": None,
"api_base_url": None,
"generation_kwargs": {},
"tools": None,
"tools_strict": False,
"max_retries": None,
"timeout": None,
"http_client_kwargs": None,
},
}
def test_to_dict_with_parameters(self, monkeypatch, calendar_event_model):
tool = Tool(name="name", description="description", parameters={"x": {"type": "string"}}, function=print)
monkeypatch.setenv("ENV_VAR", "test-api-key")
component = OpenAIChatGenerator(
api_key=Secret.from_env_var("ENV_VAR"),
streaming_callback=print_streaming_chunk,
api_base_url="test-base-url",
generation_kwargs={
"max_completion_tokens": 10,
"some_test_param": "test-params",
"response_format": calendar_event_model,
"logprobs": True,
},
tools=[tool],
tools_strict=True,
max_retries=10,
timeout=100.0,
http_client_kwargs={"proxy": "http://example.com:8080", "verify": False},
)
data = component.to_dict()
assert data == {
"type": "haystack.components.generators.chat.openai.OpenAIChatGenerator",
"init_parameters": {
"api_key": {"env_vars": ["ENV_VAR"], "strict": True, "type": "env_var"},
"model": "gpt-5-mini",
"organization": None,
"api_base_url": "test-base-url",
"max_retries": 10,
"timeout": 100.0,
"streaming_callback": "haystack.components.generators.utils.print_streaming_chunk",
"generation_kwargs": {
"max_completion_tokens": 10,
"some_test_param": "test-params",
"logprobs": True,
"response_format": {
"type": "json_schema",
"json_schema": {
"name": "CalendarEvent",
"strict": True,
"schema": {
"properties": {
"event_name": {"title": "Event Name", "type": "string"},
"event_date": {"title": "Event Date", "type": "string"},
"event_location": {"title": "Event Location", "type": "string"},
},
"required": ["event_name", "event_date", "event_location"],
"title": "CalendarEvent",
"type": "object",
"additionalProperties": False,
},
},
},
},
"tools": [
{
"type": "haystack.tools.tool.Tool",
"data": {
"async_function": None,
"description": "description",
"function": "builtins.print",
"inputs_from_state": None,
"name": "name",
"outputs_to_state": None,
"outputs_to_string": None,
"parameters": {"x": {"type": "string"}},
},
}
],
"tools_strict": True,
"http_client_kwargs": {"proxy": "http://example.com:8080", "verify": False},
},
}
def test_to_dict_with_response_format_json_object(self, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
component = OpenAIChatGenerator(
api_key=Secret.from_env_var("OPENAI_API_KEY"),
generation_kwargs={"response_format": {"type": "json_object"}},
)
data = component.to_dict()
assert data == {
"type": "haystack.components.generators.chat.openai.OpenAIChatGenerator",
"init_parameters": {
"api_key": {"env_vars": ["OPENAI_API_KEY"], "strict": True, "type": "env_var"},
"model": "gpt-5-mini",
"api_base_url": None,
"organization": None,
"streaming_callback": None,
"generation_kwargs": {"response_format": {"type": "json_object"}},
"tools": None,
"tools_strict": False,
"max_retries": None,
"timeout": None,
"http_client_kwargs": None,
},
}
def test_from_dict(self, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
data = {
"type": "haystack.components.generators.chat.openai.OpenAIChatGenerator",
"init_parameters": {
"api_key": {"env_vars": ["OPENAI_API_KEY"], "strict": True, "type": "env_var"},
"model": "gpt-5-mini",
"api_base_url": "test-base-url",
"streaming_callback": "haystack.components.generators.utils.print_streaming_chunk",
"max_retries": 10,
"timeout": 100.0,
"generation_kwargs": {"max_completion_tokens": 10, "some_test_param": "test-params"},
"tools": [
{
"type": "haystack.tools.tool.Tool",
"data": {
"description": "description",
"function": "builtins.print",
"name": "name",
"parameters": {"x": {"type": "string"}},
},
}
],
"tools_strict": True,
"http_client_kwargs": {"proxy": "http://example.com:8080", "verify": False},
},
}
component = OpenAIChatGenerator.from_dict(data)
assert isinstance(component, OpenAIChatGenerator)
assert component.model == "gpt-5-mini"
assert component.streaming_callback is print_streaming_chunk
assert component.api_base_url == "test-base-url"
assert component.generation_kwargs == {"max_completion_tokens": 10, "some_test_param": "test-params"}
assert component.api_key == Secret.from_env_var("OPENAI_API_KEY")
assert component.tools == [
Tool(name="name", description="description", parameters={"x": {"type": "string"}}, function=print)
]
assert component.tools_strict
assert component.timeout == 100.0
assert component.max_retries == 10
assert component.http_client_kwargs == {"proxy": "http://example.com:8080", "verify": False}
def test_from_dict_wo_env_var_does_not_fail(self, monkeypatch):
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
data = {
"type": "haystack.components.generators.chat.openai.OpenAIChatGenerator",
"init_parameters": {
"api_key": {"env_vars": ["OPENAI_API_KEY"], "strict": True, "type": "env_var"},
"model": "gpt-4",
"organization": None,
"api_base_url": "test-base-url",
"streaming_callback": "haystack.components.generators.utils.print_streaming_chunk",
"generation_kwargs": {"max_completion_tokens": 10, "some_test_param": "test-params"},
"tools": None,
},
}
component = OpenAIChatGenerator.from_dict(data)
assert component.client is None
assert component.async_client is None
def test_run(self, chat_messages, openai_mock_chat_completion):
component = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"))
response = component.run(chat_messages)
# check that the component returns the correct ChatMessage response
assert isinstance(response, dict)
assert "replies" in response
assert isinstance(response["replies"], list)
assert len(response["replies"]) == 1
assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
def test_run_with_string_input(self, openai_mock_chat_completion):
component = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"))
response = component.run("What's the capital of France?")
_, kwargs = openai_mock_chat_completion.call_args
assert kwargs["messages"] == [{"role": "user", "content": "What's the capital of France?"}]
assert isinstance(response["replies"], list)
assert len(response["replies"]) == 1
assert isinstance(response["replies"][0], ChatMessage)
def test_run_with_params(self, chat_messages, openai_mock_chat_completion):
component = OpenAIChatGenerator(
api_key=Secret.from_token("test-api-key"),
generation_kwargs={"max_completion_tokens": 10, "temperature": 0.5},
)
response = component.run(chat_messages)
# check that the component calls the OpenAI API with the correct parameters
_, kwargs = openai_mock_chat_completion.call_args
assert kwargs["max_completion_tokens"] == 10
assert kwargs["temperature"] == 0.5
# check that the tools are not passed to the OpenAI API (the generator is initialized without tools)
assert "tools" not in kwargs
# check that the component returns the correct response
assert isinstance(response, dict)
assert "replies" in response
assert isinstance(response["replies"], list)
assert len(response["replies"]) == 1
assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
def test_run_with_params_streaming(self, chat_messages, openai_mock_chat_completion_chunk):
streaming_callback_called = False
def streaming_callback(chunk: StreamingChunk) -> None:
nonlocal streaming_callback_called
streaming_callback_called = True
component = OpenAIChatGenerator(
api_key=Secret.from_token("test-api-key"), streaming_callback=streaming_callback
)
response = component.run(chat_messages)
# check we called the streaming callback
assert streaming_callback_called
# check that the component still returns the correct response
assert isinstance(response, dict)
assert "replies" in response
assert isinstance(response["replies"], list)
assert len(response["replies"]) == 1
assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
assert "Hello" in response["replies"][0].text # see openai_mock_chat_completion_chunk
def test_run_with_streaming_callback_in_run_method(self, chat_messages, openai_mock_chat_completion_chunk):
streaming_callback_called = False
def streaming_callback(chunk: StreamingChunk) -> None:
nonlocal streaming_callback_called
streaming_callback_called = True
component = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"))
response = component.run(chat_messages, streaming_callback=streaming_callback)
# check we called the streaming callback
assert streaming_callback_called
# check that the component still returns the correct response
assert isinstance(response, dict)
assert "replies" in response
assert isinstance(response["replies"], list)
assert len(response["replies"]) == 1
assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
assert "Hello" in response["replies"][0].text # see openai_mock_chat_completion_chunk
def test_run_with_response_format(self, chat_messages, mock_parsed_chat_completion):
component = OpenAIChatGenerator(
api_key=Secret.from_token("test-api-key"), generation_kwargs={"response_format": CalendarEvent}
)
response = component.run(chat_messages)
assert isinstance(response, dict)
assert "replies" in response
assert isinstance(response["replies"], list)
assert len(response["replies"]) == 1
assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
assert "Team Meeting" in response["replies"][0].text # see mock_parsed_chat_completion
def test_run_with_response_format_in_run_method(self, chat_messages, mock_parsed_chat_completion):
component = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"))
response = component.run(chat_messages, generation_kwargs={"response_format": CalendarEvent})
assert isinstance(response, dict)
assert "replies" in response
assert isinstance(response["replies"], list)
assert len(response["replies"]) == 1
assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
assert "Team Meeting" in response["replies"][0].text # see mock_parsed_chat_completion
def test_run_with_wrapped_stream_simulation(self, chat_messages, openai_mock_stream):
streaming_callback_called = False
def streaming_callback(chunk: StreamingChunk) -> None:
nonlocal streaming_callback_called
streaming_callback_called = True
assert isinstance(chunk, StreamingChunk)
chunk = ChatCompletionChunk(
id="id",
model="gpt-4",
object="chat.completion.chunk",
choices=[chat_completion_chunk.Choice(index=0, delta=chat_completion_chunk.ChoiceDelta(content="Hello"))],
created=int(datetime.now().timestamp()),
)
# Here we wrap the OpenAI stream in a MagicMock
# This is to simulate the behavior of some tools like Weave (https://github.com/wandb/weave)
# which wrap the OpenAI stream in their own stream
wrapped_openai_stream = MagicMock()
wrapped_openai_stream.__iter__.return_value = iter([chunk])
component = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"))
component.warm_up()
with patch.object(
component.client.chat.completions, "create", return_value=wrapped_openai_stream
) as mock_create:
response = component.run(chat_messages, streaming_callback=streaming_callback)
mock_create.assert_called_once()
assert streaming_callback_called
assert "replies" in response
assert "Hello" in response["replies"][0].text
def test_check_abnormal_completions(self, caplog):
caplog.set_level(logging.INFO)
messages = [
ChatMessage.from_assistant(
"", meta={"finish_reason": "content_filter" if i % 2 == 0 else "length", "index": i}
)
for i, _ in enumerate(range(4))
]
for m in messages:
_check_finish_reason(m.meta)
# check truncation warning
message_template = (
"The completion for index {index} has been truncated before reaching a natural stopping point. "
"Increase the max_completion_tokens parameter to allow for longer completions."
)
for index in [1, 3]:
assert caplog.records[index].message == message_template.format(index=index)
# check content filter warning
message_template = "The completion for index {index} has been truncated due to the content filter."
for index in [0, 2]:
assert caplog.records[index].message == message_template.format(index=index)
def test_run_with_tools(self, tools):
with patch("openai.resources.chat.completions.Completions.create") as mock_chat_completion_create:
completion = ChatCompletion(
id="foo",
model="gpt-4",
object="chat.completion",
choices=[
Choice(
finish_reason="tool_calls",
logprobs=None,
index=0,
message=ChatCompletionMessage(
role="assistant",
tool_calls=[
ChatCompletionMessageFunctionToolCall(
id="123",
type="function",
function=Function(name="weather", arguments='{"city": "Paris"}'),
)
],
),
)
],
created=int(datetime.now().timestamp()),
usage=CompletionUsage(
completion_tokens=40,
prompt_tokens=57,
total_tokens=97,
completion_tokens_details=CompletionTokensDetails(
accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0
),
prompt_tokens_details=PromptTokensDetails(audio_tokens=0, cached_tokens=0),
),
)
mock_chat_completion_create.return_value = completion
component = OpenAIChatGenerator(
api_key=Secret.from_token("test-api-key"), tools=tools[:1], tools_strict=True
)
response = component.run([ChatMessage.from_user("What's the weather like in Paris?")])
# ensure that the tools are passed to the OpenAI API
function_spec = {**tools[0].tool_spec}
function_spec["strict"] = True
function_spec["parameters"]["additionalProperties"] = False
assert mock_chat_completion_create.call_args[1]["tools"] == [{"type": "function", "function": function_spec}]
assert len(response["replies"]) == 1
message = response["replies"][0]
assert not message.texts
assert not message.text
assert message.tool_calls
tool_call = message.tool_call
assert isinstance(tool_call, ToolCall)
assert tool_call.tool_name == "weather"
assert tool_call.arguments == {"city": "Paris"}
assert message.meta["finish_reason"] == "tool_calls"
assert message.meta["usage"]["completion_tokens"] == 40
def test_run_with_tools_and_response_format(self, tools, mock_parsed_chat_completion):
"""
Test the run method with tools and response format
When tools are used, the function call overrides the schema passed in response_format
"""
with patch("openai.resources.chat.completions.Completions.parse") as mock_chat_completion_parse:
completion = ParsedChatCompletion[CalendarEvent](
id="foo",
model="gpt-4",
object="chat.completion",
choices=[
ParsedChoice[CalendarEvent](
finish_reason="tool_calls",
logprobs=None,
index=0,
message=ParsedChatCompletionMessage[CalendarEvent](
role="assistant",
tool_calls=[
ParsedFunctionToolCall(
id="123",
type="function",
function=ParsedFunction(name="weather", arguments='{"city": "Paris"}'),
)
],
),
)
],
created=int(datetime.now().timestamp()),
usage=CompletionUsage(
completion_tokens=40,
prompt_tokens=57,
total_tokens=97,
completion_tokens_details=CompletionTokensDetails(
accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0
),
prompt_tokens_details=PromptTokensDetails(audio_tokens=0, cached_tokens=0),
),
)
mock_chat_completion_parse.return_value = completion
component = OpenAIChatGenerator(
api_key=Secret.from_token("test-api-key"), tools=tools[:1], tools_strict=True
)
response_with_format = component.run(
[ChatMessage.from_user("What's the weather like in Paris?")],
generation_kwargs={"response_format": CalendarEvent},
)
assert len(response_with_format["replies"]) == 1
message_with_format = response_with_format["replies"][0]
assert not message_with_format.texts
assert not message_with_format.text
assert message_with_format.tool_calls
tool_call = message_with_format.tool_call
assert isinstance(tool_call, ToolCall)
assert tool_call.tool_name == "weather"
assert tool_call.arguments == {"city": "Paris"}
assert message_with_format.meta["finish_reason"] == "tool_calls"
assert message_with_format.meta["usage"]["completion_tokens"] == 40
def test_run_with_tools_streaming(self, mock_chat_completion_chunk_with_tools, tools):
streaming_callback_called = False
def streaming_callback(chunk: StreamingChunk) -> None:
nonlocal streaming_callback_called
streaming_callback_called = True
component = OpenAIChatGenerator(
api_key=Secret.from_token("test-api-key"), streaming_callback=streaming_callback
)
chat_messages = [ChatMessage.from_user("What's the weather like in Paris?")]
response = component.run(chat_messages, tools=tools)
# check we called the streaming callback
assert streaming_callback_called
# check that the component still returns the correct response
assert isinstance(response, dict)
assert "replies" in response
assert isinstance(response["replies"], list)
assert len(response["replies"]) == 1
assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
message = response["replies"][0]
assert message.tool_calls
tool_call = message.tool_call
assert isinstance(tool_call, ToolCall)
assert tool_call.tool_name == "weather"
assert tool_call.arguments == {"city": "Paris"}
assert message.meta["finish_reason"] == "tool_calls"
def test_invalid_tool_call_json(self, tools, caplog):
caplog.set_level(logging.WARNING)
with patch("openai.resources.chat.completions.Completions.create") as mock_create:
mock_create.return_value = ChatCompletion(
id="test",
model="gpt-5-mini",
object="chat.completion",
choices=[
Choice(
finish_reason="tool_calls",
index=0,
message=ChatCompletionMessage(
role="assistant",
tool_calls=[
ChatCompletionMessageFunctionToolCall(
id="1",
type="function",
function=Function(name="weather", arguments='"invalid": "json"'),
)
],
),
)
],
created=1234567890,
usage=CompletionUsage(
completion_tokens=47,
prompt_tokens=540,
total_tokens=587,
completion_tokens_details=CompletionTokensDetails(
accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0
),
prompt_tokens_details=PromptTokensDetails(audio_tokens=0, cached_tokens=0),
),
)
component = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"), tools=tools)
response = component.run([ChatMessage.from_user("What's the weather in Paris?")])
assert len(response["replies"]) == 1
message = response["replies"][0]
assert len(message.tool_calls) == 0
assert "OpenAI returned a malformed JSON string for tool call arguments" in caplog.text
assert message.meta["finish_reason"] == "tool_calls"
assert message.meta["usage"]["completion_tokens"] == 47
def test_run_with_response_format_and_streaming_pydantic_model(self, calendar_event_model):
chat_messages = [
ChatMessage.from_user("The marketing summit takes place on October12th at the Hilton Hotel downtown.")
]
component = OpenAIChatGenerator(
api_key=Secret.from_token("test-api-key"),
generation_kwargs={"response_format": calendar_event_model},
streaming_callback=print_streaming_chunk,
)
with pytest.raises(TypeError):
component.run(chat_messages)
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY", None),
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
)
@pytest.mark.integration
def test_live_run(self):
chat_messages = [ChatMessage.from_user("What's the capital of France")]
component = OpenAIChatGenerator(model="gpt-4.1-nano", generation_kwargs={"n": 1})
results = component.run(chat_messages)
assert len(results["replies"]) == 1
message: ChatMessage = results["replies"][0]
assert "Paris" in message.text
assert "gpt-4.1-nano" in message.meta["model"]
assert message.meta["finish_reason"] == "stop"
assert message.meta["usage"]["prompt_tokens"] > 0
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY", None),
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
)
@pytest.mark.integration
def test_live_run_with_response_format_pydantic_model(self, calendar_event_model):
chat_messages = [
ChatMessage.from_user("The marketing summit takes place on October12th at the Hilton Hotel downtown.")
]
component = OpenAIChatGenerator(
model="gpt-4.1-nano", generation_kwargs={"response_format": calendar_event_model}
)
results = component.run(chat_messages)
assert len(results["replies"]) == 1
message: ChatMessage = results["replies"][0]
msg = json.loads(message.text)
assert "Marketing Summit" in msg["event_name"]
assert isinstance(msg["event_date"], str)
assert isinstance(msg["event_location"], str)
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY", None),
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
)
@pytest.mark.integration
def test_live_run_with_response_format_json_object(self):
chat_messages = [
ChatMessage.from_user(
'Answer in JSON: What\'s the capital of France? Please respond with a JSON object with the key "city". '
'For example: {"city": "Paris"}'
)
]
comp = OpenAIChatGenerator(model="gpt-4.1-nano", generation_kwargs={"response_format": {"type": "json_object"}})
results = comp.run(chat_messages)
assert len(results["replies"]) == 1
message: ChatMessage = results["replies"][0]
msg = json.loads(message.text)
assert "paris" in msg["city"].lower()
assert message.meta["finish_reason"] == "stop"
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY", None),
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
)
@pytest.mark.integration
def test_live_run_with_response_format_json_object_streaming(self):
streaming_callback_called = False
def streaming_callback(chunk: StreamingChunk) -> None:
nonlocal streaming_callback_called
streaming_callback_called = True
chat_messages = [
ChatMessage.from_user(
'Answer in JSON: What\'s the capital of France? Please respond with a JSON object with the key "city". '
'For example: {"city": "Paris"}'
)
]
comp = OpenAIChatGenerator(
model="gpt-4.1-nano",
generation_kwargs={"response_format": {"type": "json_object"}},
streaming_callback=streaming_callback,
)
results = comp.run(chat_messages)
assert len(results["replies"]) == 1
message: ChatMessage = results["replies"][0]
msg = json.loads(message.text)
assert "paris" in msg["city"].lower()
assert message.meta["finish_reason"] == "stop"
assert streaming_callback_called is True
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY", None),
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
)
@pytest.mark.integration
def test_live_run_with_response_format_json_schema(self):
response_schema = {
"type": "json_schema",
"json_schema": {
"name": "CapitalCity",
"strict": True,
"schema": {
"title": "CapitalCity",
"type": "object",
"properties": {
"city": {"title": "City", "type": "string"},
"country": {"title": "Country", "type": "string"},
},
"required": ["city", "country"],
"additionalProperties": False,
},
},
}
chat_messages = [ChatMessage.from_user("What's the capital of France?")]
comp = OpenAIChatGenerator(model="gpt-4.1-nano", generation_kwargs={"response_format": response_schema})
results = comp.run(chat_messages)
assert len(results["replies"]) == 1
message: ChatMessage = results["replies"][0]
msg = json.loads(message.text)
assert "Paris" in msg["city"]
assert isinstance(msg["country"], str)
assert "France" in msg["country"]
assert message.meta["finish_reason"] == "stop"
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY", None),
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
)
@pytest.mark.integration
def test_live_run_with_response_format_json_schema_streaming(self):
streaming_callback_called = False
def streaming_callback(chunk: StreamingChunk) -> None:
nonlocal streaming_callback_called
streaming_callback_called = True
response_schema = {
"type": "json_schema",
"json_schema": {
"name": "CapitalCity",
"strict": True,
"schema": {
"title": "CapitalCity",
"type": "object",
"properties": {
"city": {"title": "City", "type": "string"},
"country": {"title": "Country", "type": "string"},
},
"required": ["city", "country"],
"additionalProperties": False,
},
},
}
chat_messages = [ChatMessage.from_user("What's the capital of France?")]
comp = OpenAIChatGenerator(
model="gpt-4.1-nano",
generation_kwargs={"response_format": response_schema},
streaming_callback=streaming_callback,
)
results = comp.run(chat_messages)
assert len(results["replies"]) == 1
message: ChatMessage = results["replies"][0]
msg = json.loads(message.text)
assert "Paris" in msg["city"]
assert isinstance(msg["country"], str)
assert "France" in msg["country"]
assert message.meta["finish_reason"] == "stop"
assert streaming_callback_called is True
def test_run_with_wrong_model(self):
mock_client = MagicMock()
mock_client.chat.completions.create.side_effect = OpenAIError("Invalid model name")
generator = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"), model="something-obviously-wrong")
generator.client = mock_client
with pytest.raises(OpenAIError):
generator.run([ChatMessage.from_user("irrelevant")])
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY", None),
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
)
@pytest.mark.integration
def test_live_run_streaming(self):
class Callback:
def __init__(self):
self.responses = ""
self.counter = 0
def __call__(self, chunk: StreamingChunk) -> None:
self.counter += 1
self.responses += chunk.content if chunk.content else ""
callback = Callback()
component = OpenAIChatGenerator(
model="gpt-4.1-nano",
streaming_callback=callback,
generation_kwargs={"stream_options": {"include_usage": True}},
)
results = component.run([ChatMessage.from_user("What's the capital of France?")])
# Basic response checks
assert "replies" in results
assert len(results["replies"]) == 1
message: ChatMessage = results["replies"][0]
assert "Paris" in message.text
assert isinstance(message.meta, dict)
# Metadata checks
metadata = message.meta
assert "gpt-4.1-nano" in metadata["model"]
assert metadata["finish_reason"] == "stop"
# Usage information checks
assert isinstance(metadata.get("usage"), dict), "meta.usage not a dict"
usage = metadata["usage"]
assert "prompt_tokens" in usage and usage["prompt_tokens"] > 0
assert "completion_tokens" in usage and usage["completion_tokens"] > 0
# Detailed token information checks
assert isinstance(usage.get("completion_tokens_details"), dict), "usage.completion_tokens_details not a dict"
assert isinstance(usage.get("prompt_tokens_details"), dict), "usage.prompt_tokens_details not a dict"
# Streaming callback verification
assert callback.counter > 1
assert "Paris" in callback.responses
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY", None),
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
)
@pytest.mark.integration
def test_live_run_with_tools_streaming(self, tools):
chat_messages = [ChatMessage.from_user("What's the weather like in Paris and Berlin?")]
component = OpenAIChatGenerator(
model="gpt-4.1-nano",
tools=tools,
streaming_callback=print_streaming_chunk,
generation_kwargs={"stream_options": {"include_usage": True}},
)
results = component.run(chat_messages)
assert len(results["replies"]) == 1
message = results["replies"][0]
assert not message.texts
assert not message.text
assert message.tool_calls
tool_calls = message.tool_calls
assert len(tool_calls) == 2
for tool_call in tool_calls:
assert isinstance(tool_call, ToolCall)
assert tool_call.tool_name == "weather"
arguments = [tool_call.arguments for tool_call in tool_calls]
# Check that both cities are present (case-insensitive, allowing for variations like "Paris, France")
city_values = [arg["city"].lower() for arg in arguments]
assert any("berlin" in city for city in city_values)
assert any("paris" in city for city in city_values)
assert message.meta["finish_reason"] == "tool_calls"
def test_openai_chat_generator_with_toolset_initialization(self, tools, monkeypatch):
"""Test that the OpenAIChatGenerator can be initialized with a Toolset."""
monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
toolset = Toolset(tools)
generator = OpenAIChatGenerator(tools=toolset)
assert generator.tools == toolset
def test_from_dict_with_toolset(self, tools, monkeypatch):
"""Test that the OpenAIChatGenerator can be deserialized from a dictionary with a Toolset."""
monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
toolset = Toolset(tools)
component = OpenAIChatGenerator(tools=toolset)
data = component.to_dict()
deserialized_component = OpenAIChatGenerator.from_dict(data)
assert isinstance(deserialized_component.tools, Toolset)
assert len(deserialized_component.tools) == len(tools)
assert all(isinstance(tool, Tool) for tool in deserialized_component.tools)
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY", None),
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
)
@pytest.mark.integration
def test_live_run_with_toolset(self, tools):
chat_messages = [ChatMessage.from_user("What's the weather like in Paris?")]
toolset = Toolset(tools)
component = OpenAIChatGenerator(model="gpt-4.1-nano", tools=toolset)
results = component.run(chat_messages)
assert len(results["replies"]) == 1
message = results["replies"][0]
assert not message.texts
assert not message.text
assert message.tool_calls
tool_call = message.tool_call
assert isinstance(tool_call, ToolCall)
assert tool_call.tool_name == "weather"
assert tool_call.arguments.keys() == {"city"}
assert "Paris" in tool_call.arguments["city"]
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY", None),
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
)
@pytest.mark.integration
def test_live_run_multimodal(self, test_files_path):
image_path = test_files_path / "images" / "apple.jpg"
# we resize the image to keep this test fast (around 1s) - increase the size in case of errors
image_content = ImageContent.from_file_path(file_path=image_path, size=(100, 100), detail="low")
chat_messages = [ChatMessage.from_user(content_parts=["What does this image show? Max 5 words", image_content])]
generator = OpenAIChatGenerator(model="gpt-4.1-nano")
results = generator.run(chat_messages)
assert len(results["replies"]) == 1
message: ChatMessage = results["replies"][0]
assert message.text
assert "apple" in message.text.lower()
assert message.is_from(ChatRole.ASSISTANT)
assert not message.tool_calls
assert not message.tool_call_results
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY", None),
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
)
@pytest.mark.integration
def test_live_run_with_file_content(self, test_files_path):
pdf_path = test_files_path / "pdf" / "sample_pdf_3.pdf"
file_content = FileContent.from_file_path(file_path=pdf_path)
chat_messages = [
ChatMessage.from_user(
content_parts=[file_content, "Is this document a paper about LLMs? Respond with 'yes' or 'no' only."]
)
]
generator = OpenAIChatGenerator(model="gpt-4.1-nano")
results = generator.run(chat_messages)
assert len(results["replies"]) == 1
message: ChatMessage = results["replies"][0]
assert message.is_from(ChatRole.ASSISTANT)
assert message.text
assert "no" in message.text.lower()
def test_init_with_list_of_toolsets(self, monkeypatch, tools):
"""Test initialization with a list of Toolsets."""
monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
toolset1 = Toolset([tools[0]])
toolset2 = Toolset([tools[1]])
component = OpenAIChatGenerator(tools=[toolset1, toolset2])
assert component.tools == [toolset1, toolset2]
assert isinstance(component.tools, list)
assert len(component.tools) == 2
assert all(isinstance(ts, Toolset) for ts in component.tools)
def test_serde_with_list_of_toolsets(self, monkeypatch, tools):
"""Test serialization and deserialization with a list of Toolsets."""
monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
toolset1 = Toolset([tools[0]])
toolset2 = Toolset([tools[1]])
component = OpenAIChatGenerator(tools=[toolset1, toolset2])
data = component.to_dict()
# Verify serialization preserves list[Toolset] structure
tools_data = data["init_parameters"]["tools"]
assert isinstance(tools_data, list)
assert len(tools_data) == 2
assert all(isinstance(ts, dict) for ts in tools_data)
assert tools_data[0]["type"] == "haystack.tools.toolset.Toolset"
assert tools_data[1]["type"] == "haystack.tools.toolset.Toolset"
# Deserialize and verify
deserialized = OpenAIChatGenerator.from_dict(data)
assert isinstance(deserialized.tools, list)
assert len(deserialized.tools) == 2
assert all(isinstance(ts, Toolset) for ts in deserialized.tools)
@pytest.fixture
def chat_completion_chunks():
return [
ChatCompletionChunk(
id="chatcmpl-BZdwjFecdcaQfCf7bn319vRp6fY8F",
choices=[chat_completion_chunk.Choice(delta=ChoiceDelta(role="assistant"), index=0)],
created=1747834733,
model="gpt-5-mini",
object="chat.completion.chunk",
service_tier="default",
system_fingerprint="fp_54eb4bd693",
),
ChatCompletionChunk(
id="chatcmpl-BZdwjFecdcaQfCf7bn319vRp6fY8F",
choices=[
chat_completion_chunk.Choice(
delta=ChoiceDelta(
tool_calls=[
ChoiceDeltaToolCall(
index=0,
id="call_zcvlnVaTeJWRjLAFfYxX69z4",
function=ChoiceDeltaToolCallFunction(arguments="", name="weather"),
type="function",
)
]
),
index=0,
)
],
created=1747834733,
model="gpt-5-mini",
object="chat.completion.chunk",
service_tier="default",
system_fingerprint="fp_54eb4bd693",
),
ChatCompletionChunk(
id="chatcmpl-BZdwjFecdcaQfCf7bn319vRp6fY8F",
choices=[
chat_completion_chunk.Choice(
delta=ChoiceDelta(
tool_calls=[
ChoiceDeltaToolCall(index=0, function=ChoiceDeltaToolCallFunction(arguments='{"ci'))
]
),
index=0,
)
],
created=1747834733,
model="gpt-5-mini",
object="chat.completion.chunk",
service_tier="default",
system_fingerprint="fp_54eb4bd693",
),
ChatCompletionChunk(
id="chatcmpl-BZdwjFecdcaQfCf7bn319vRp6fY8F",
choices=[
chat_completion_chunk.Choice(
delta=ChoiceDelta(
tool_calls=[
ChoiceDeltaToolCall(index=0, function=ChoiceDeltaToolCallFunction(arguments='ty": '))
]
),
index=0,
)
],
created=1747834733,
model="gpt-5-mini",
object="chat.completion.chunk",
service_tier="default",
system_fingerprint="fp_54eb4bd693",
),
ChatCompletionChunk(
id="chatcmpl-BZdwjFecdcaQfCf7bn319vRp6fY8F",
choices=[
chat_completion_chunk.Choice(
delta=ChoiceDelta(
tool_calls=[
ChoiceDeltaToolCall(index=0, function=ChoiceDeltaToolCallFunction(arguments='"Paris'))
]
),
index=0,
)
],
created=1747834733,
model="gpt-5-mini",
object="chat.completion.chunk",
service_tier="default",
system_fingerprint="fp_54eb4bd693",
),
ChatCompletionChunk(
id="chatcmpl-BZdwjFecdcaQfCf7bn319vRp6fY8F",
choices=[
chat_completion_chunk.Choice(
delta=ChoiceDelta(
tool_calls=[ChoiceDeltaToolCall(index=0, function=ChoiceDeltaToolCallFunction(arguments='"}'))]
),
index=0,
)
],
created=1747834733,
model="gpt-5-mini",
object="chat.completion.chunk",
service_tier="default",
system_fingerprint="fp_54eb4bd693",
),
ChatCompletionChunk(
id="chatcmpl-BZdwjFecdcaQfCf7bn319vRp6fY8F",
choices=[
chat_completion_chunk.Choice(
delta=ChoiceDelta(
tool_calls=[
ChoiceDeltaToolCall(
index=1,
id="call_C88m67V16CrETq6jbNXjdZI9",
function=ChoiceDeltaToolCallFunction(arguments="", name="weather"),
type="function",
)
]
),
index=0,
)
],
created=1747834733,
model="gpt-5-mini",
object="chat.completion.chunk",
service_tier="default",
system_fingerprint="fp_54eb4bd693",
),
ChatCompletionChunk(
id="chatcmpl-BZdwjFecdcaQfCf7bn319vRp6fY8F",
choices=[
chat_completion_chunk.Choice(
delta=ChoiceDelta(
tool_calls=[
ChoiceDeltaToolCall(index=1, function=ChoiceDeltaToolCallFunction(arguments='{"ci'))
]
),
index=0,
)
],
created=1747834733,
model="gpt-5-mini",
object="chat.completion.chunk",
service_tier="default",
system_fingerprint="fp_54eb4bd693",
),
ChatCompletionChunk(
id="chatcmpl-BZdwjFecdcaQfCf7bn319vRp6fY8F",
choices=[
chat_completion_chunk.Choice(
delta=ChoiceDelta(
tool_calls=[
ChoiceDeltaToolCall(index=1, function=ChoiceDeltaToolCallFunction(arguments='ty": '))
]
),
index=0,
)
],
created=1747834733,
model="gpt-5-mini",
object="chat.completion.chunk",
service_tier="default",
system_fingerprint="fp_54eb4bd693",
),
ChatCompletionChunk(
id="chatcmpl-BZdwjFecdcaQfCf7bn319vRp6fY8F",
choices=[
chat_completion_chunk.Choice(
delta=ChoiceDelta(
tool_calls=[
ChoiceDeltaToolCall(index=1, function=ChoiceDeltaToolCallFunction(arguments='"Berli'))
]
),
index=0,
)
],
created=1747834733,
model="gpt-5-mini",
object="chat.completion.chunk",
service_tier="default",
system_fingerprint="fp_54eb4bd693",
),
ChatCompletionChunk(
id="chatcmpl-BZdwjFecdcaQfCf7bn319vRp6fY8F",
choices=[
chat_completion_chunk.Choice(
delta=ChoiceDelta(
tool_calls=[ChoiceDeltaToolCall(index=1, function=ChoiceDeltaToolCallFunction(arguments='n"}'))]
),
index=0,
)
],
created=1747834733,
model="gpt-5-mini",
object="chat.completion.chunk",
service_tier="default",
system_fingerprint="fp_54eb4bd693",
),
ChatCompletionChunk(
id="chatcmpl-BZdwjFecdcaQfCf7bn319vRp6fY8F",
choices=[chat_completion_chunk.Choice(delta=ChoiceDelta(), finish_reason="tool_calls", index=0)],
created=1747834733,
model="gpt-5-mini",
object="chat.completion.chunk",
service_tier="default",
system_fingerprint="fp_54eb4bd693",
),
ChatCompletionChunk(
id="chatcmpl-BZdwjFecdcaQfCf7bn319vRp6fY8F",
choices=[],
created=1747834733,
model="gpt-5-mini",
object="chat.completion.chunk",
service_tier="default",
system_fingerprint="fp_54eb4bd693",
usage=CompletionUsage(
completion_tokens=42,
prompt_tokens=282,
total_tokens=324,
completion_tokens_details=CompletionTokensDetails(
accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0
),
prompt_tokens_details=PromptTokensDetails(audio_tokens=0, cached_tokens=0),
),
),
]
@pytest.fixture
def chat_completion_chunk_delta_none():
chunk = ChatCompletionChunk(
id="chatcmpl-BC1y4wqIhe17R8sv3lgLcWlB4tXCw",
choices=[chat_completion_chunk.Choice(delta=ChoiceDelta(), index=0)],
created=1742207200,
model="gpt-5-mini",
object="chat.completion.chunk",
)
# pydantic complains if we set delta to None at initialization
chunk.choices[0].delta = None
return chunk
@pytest.fixture
def streaming_chunks():
return [
StreamingChunk(
content="",
meta={
"model": "gpt-5-mini",
"index": 0,
"tool_calls": None,
"finish_reason": None,
"received_at": ANY,
"usage": None,
},
),
StreamingChunk(
content="",
meta={
"model": "gpt-5-mini",
"index": 0,
"tool_calls": [
ChoiceDeltaToolCall(
index=0,
id="call_zcvlnVaTeJWRjLAFfYxX69z4",
function=ChoiceDeltaToolCallFunction(arguments="", name="weather"),
type="function",
)
],
"finish_reason": None,
"received_at": ANY,
"usage": None,
},
index=0,
tool_calls=[ToolCallDelta(tool_name="weather", id="call_zcvlnVaTeJWRjLAFfYxX69z4", index=0)],
start=True,
),
StreamingChunk(
content="",
meta={
"model": "gpt-5-mini",
"index": 0,
"tool_calls": [ChoiceDeltaToolCall(index=0, function=ChoiceDeltaToolCallFunction(arguments='{"ci'))],
"finish_reason": None,
"received_at": ANY,
"usage": None,
},
index=0,
tool_calls=[ToolCallDelta(arguments='{"ci', index=0)],
),
StreamingChunk(
content="",
meta={
"model": "gpt-5-mini",
"index": 0,
"tool_calls": [ChoiceDeltaToolCall(index=0, function=ChoiceDeltaToolCallFunction(arguments='ty": '))],
"finish_reason": None,
"received_at": ANY,
"usage": None,
},
index=0,
tool_calls=[ToolCallDelta(arguments='ty": ', index=0)],
),
StreamingChunk(
content="",
meta={
"model": "gpt-5-mini",
"index": 0,
"tool_calls": [ChoiceDeltaToolCall(index=0, function=ChoiceDeltaToolCallFunction(arguments='"Paris'))],
"finish_reason": None,
"received_at": ANY,
"usage": None,
},
index=0,
tool_calls=[ToolCallDelta(arguments='"Paris', index=0)],
),
StreamingChunk(
content="",
meta={
"model": "gpt-5-mini",
"index": 0,
"tool_calls": [ChoiceDeltaToolCall(index=0, function=ChoiceDeltaToolCallFunction(arguments='"}'))],
"finish_reason": None,
"received_at": ANY,
"usage": None,
},
index=0,
tool_calls=[ToolCallDelta(arguments='"}', index=0)],
),
StreamingChunk(
content="",
meta={
"model": "gpt-5-mini",
"index": 0,
"tool_calls": [
ChoiceDeltaToolCall(
index=1,
id="call_C88m67V16CrETq6jbNXjdZI9",
function=ChoiceDeltaToolCallFunction(arguments="", name="weather"),
type="function",
)
],
"finish_reason": None,
"received_at": ANY,
"usage": None,
},
index=1,
tool_calls=[ToolCallDelta(tool_name="weather", id="call_C88m67V16CrETq6jbNXjdZI9", index=1)],
start=True,
),
StreamingChunk(
content="",
meta={
"model": "gpt-5-mini",
"index": 0,
"tool_calls": [ChoiceDeltaToolCall(index=1, function=ChoiceDeltaToolCallFunction(arguments='{"ci'))],
"finish_reason": None,
"received_at": ANY,
"usage": None,
},
index=1,
tool_calls=[ToolCallDelta(arguments='{"ci', index=1)],
),
StreamingChunk(
content="",
meta={
"model": "gpt-5-mini",
"index": 0,
"tool_calls": [ChoiceDeltaToolCall(index=1, function=ChoiceDeltaToolCallFunction(arguments='ty": '))],
"finish_reason": None,
"received_at": ANY,
"usage": None,
},
index=1,
tool_calls=[ToolCallDelta(arguments='ty": ', index=1)],
),
StreamingChunk(
content="",
meta={
"model": "gpt-5-mini",
"index": 0,
"tool_calls": [ChoiceDeltaToolCall(index=1, function=ChoiceDeltaToolCallFunction(arguments='"Berli'))],
"finish_reason": None,
"received_at": ANY,
"usage": None,
},
index=1,
tool_calls=[ToolCallDelta(arguments='"Berli', index=1)],
),
StreamingChunk(
content="",
meta={
"model": "gpt-5-mini",
"index": 0,
"tool_calls": [ChoiceDeltaToolCall(index=1, function=ChoiceDeltaToolCallFunction(arguments='n"}'))],
"finish_reason": None,
"received_at": ANY,
"usage": None,
},
index=1,
tool_calls=[ToolCallDelta(arguments='n"}', index=1)],
),
StreamingChunk(
content="",
meta={
"model": "gpt-5-mini",
"index": 0,
"tool_calls": None,
"finish_reason": "tool_calls",
"received_at": ANY,
"usage": None,
},
finish_reason="tool_calls",
),
StreamingChunk(
content="",
meta={
"model": "gpt-5-mini",
"received_at": ANY,
"usage": {
"completion_tokens": 42,
"prompt_tokens": 282,
"total_tokens": 324,
"completion_tokens_details": {
"accepted_prediction_tokens": 0,
"audio_tokens": 0,
"reasoning_tokens": 0,
"rejected_prediction_tokens": 0,
},
"prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0},
},
},
),
]
@pytest.fixture
def mock_openai_clients(monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "fake")
sync_cls = MagicMock(name="OpenAI")
async_cls = MagicMock(name="AsyncOpenAI")
async_cls.return_value.close = AsyncMock()
monkeypatch.setattr(openai_chat_module, "OpenAI", sync_cls)
monkeypatch.setattr(openai_chat_module, "AsyncOpenAI", async_cls)
return sync_cls, async_cls
class TestComponentLifecycle:
def test_warm_up_uses_default_timeout_and_max_retries(self, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
generator = OpenAIChatGenerator()
generator.warm_up()
assert generator.client.max_retries == 5
assert generator.client.timeout == 30.0
def test_warm_up_uses_timeout_and_max_retries_from_parameters(self):
generator = OpenAIChatGenerator(api_key=Secret.from_token("fake-api-key"), timeout=40.0, max_retries=1)
generator.warm_up()
assert generator.client.max_retries == 1
assert generator.client.timeout == 40.0
def test_warm_up_uses_timeout_and_max_retries_from_env_vars(self, monkeypatch):
monkeypatch.setenv("OPENAI_TIMEOUT", "100")
monkeypatch.setenv("OPENAI_MAX_RETRIES", "10")
generator = OpenAIChatGenerator(api_key=Secret.from_token("fake-api-key"))
generator.warm_up()
assert generator.client.max_retries == 10
assert generator.client.timeout == 100.0
def test_key_resolved_at_warm_up_not_init(self, monkeypatch):
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
generator = OpenAIChatGenerator()
with pytest.raises(ValueError, match="None of the .* environment variables are set"):
generator.warm_up()
def test_warm_up_warms_tools_once(self, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
warm_up_calls = []
class MockTool(Tool):
def __init__(self, tool_name):
super().__init__(
name=tool_name,
description=f"Mock tool {tool_name}",
parameters={"type": "object", "properties": {"x": {"type": "string"}}, "required": ["x"]},
function=lambda x: x,
)
def warm_up(self):
warm_up_calls.append(self.name)
generator = OpenAIChatGenerator(tools=[MockTool("tool1"), MockTool("tool2")])
assert not generator._tools_warmed_up
generator.warm_up()
assert sorted(warm_up_calls) == ["tool1", "tool2"]
assert generator._tools_warmed_up
generator.warm_up()
assert sorted(warm_up_calls) == ["tool1", "tool2"]
def test_warm_up_with_no_tools_does_not_raise(self, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
generator = OpenAIChatGenerator()
generator.warm_up()
assert generator._tools_warmed_up
def test_sync_lifecycle(self, mock_openai_clients):
sync_cls, _ = mock_openai_clients
generator = OpenAIChatGenerator()
assert generator.client is None
assert generator.async_client is None
generator.warm_up()
assert generator.client is sync_cls.return_value
assert generator.async_client is None
generator.close()
sync_cls.return_value.close.assert_called_once()
assert generator.client is None
async def test_async_lifecycle(self, mock_openai_clients):
_, async_cls = mock_openai_clients
generator = OpenAIChatGenerator()
await generator.warm_up_async()
assert generator.async_client is async_cls.return_value
assert generator.client is None
await generator.close_async()
async_cls.return_value.close.assert_awaited_once()
assert generator.async_client is None
async def test_close_is_safe_without_warm_up(self, mock_openai_clients):
generator = OpenAIChatGenerator()
generator.close()
await generator.close_async()
assert generator.client is None
assert generator.async_client is None
async def test_close_and_close_async_are_independent(self, mock_openai_clients):
generator = OpenAIChatGenerator()
generator.warm_up()
await generator.warm_up_async()
generator.close()
assert generator.client is None
assert generator.async_client is not None
await generator.close_async()
assert generator.async_client is None
class TestChatCompletionChunkConversion:
def test_convert_chat_completion_chunk_to_streaming_chunk(self, chat_completion_chunks, streaming_chunks):
previous_chunks = []
for openai_chunk, haystack_chunk in zip(chat_completion_chunks, streaming_chunks, strict=True):
stream_chunk = _convert_chat_completion_chunk_to_streaming_chunk(
chunk=openai_chunk, previous_chunks=previous_chunks
)
assert stream_chunk == haystack_chunk
previous_chunks.append(stream_chunk)
def test_convert_chat_completion_chunk_with_empty_tool_calls(self):
# This can happen with some LLM providers where tool calls are not present but the pydantic models are still
# initialized.
chunk = ChatCompletionChunk(
id="chatcmpl-BC1y4wqIhe17R8sv3lgLcWlB4tXCw",
choices=[
chat_completion_chunk.Choice(
delta=chat_completion_chunk.ChoiceDelta(
tool_calls=[ChoiceDeltaToolCall(index=0, function=ChoiceDeltaToolCallFunction())]
),
index=0,
)
],
created=1742207200,
model="gpt-5-mini",
object="chat.completion.chunk",
)
result = _convert_chat_completion_chunk_to_streaming_chunk(chunk=chunk, previous_chunks=[])
assert result.content == ""
assert result.start is False
assert result.tool_calls == [ToolCallDelta(index=0)]
assert result.tool_call_result is None
assert result.index == 0
assert result.meta["model"] == "gpt-5-mini"
assert result.meta["received_at"] is not None
def test_convert_chat_completion_chunk_with_delta_none(self, chat_completion_chunk_delta_none):
"""
Test that a chat completion chunk with a delta set to None is converted to a streaming chunk properly.
This should not happen, but some OpenAI-compatible providers sometimes return a delta set to None.
"""
result = _convert_chat_completion_chunk_to_streaming_chunk(
chunk=chat_completion_chunk_delta_none, previous_chunks=[]
)
assert result.content == ""
assert result.start is False
assert result.tool_calls is None
assert result.tool_call_result is None
assert result.index == 0
assert result.component_info is None
assert result.finish_reason is None
assert result.reasoning is None
assert result.meta["model"] == "gpt-5-mini"
assert result.meta["received_at"] is not None
assert result.meta["index"] == 0
assert result.meta["finish_reason"] is None
assert result.meta["usage"] is None
assert result.meta["tool_calls"] is None
def test_handle_stream_response(self, chat_completion_chunks, chat_completion_chunk_delta_none):
openai_chunks = [chat_completion_chunk_delta_none] + chat_completion_chunks
comp = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"))
result = comp._handle_stream_response(openai_chunks, callback=lambda _: None)[0] # type: ignore
assert not result.texts
assert not result.text
# Verify both tool calls were found and processed
assert len(result.tool_calls) == 2
assert result.tool_calls[0].id == "call_zcvlnVaTeJWRjLAFfYxX69z4"
assert result.tool_calls[0].tool_name == "weather"
assert result.tool_calls[0].arguments == {"city": "Paris"}
assert result.tool_calls[1].id == "call_C88m67V16CrETq6jbNXjdZI9"
assert result.tool_calls[1].tool_name == "weather"
assert result.tool_calls[1].arguments == {"city": "Berlin"}
# Verify meta information
assert result.meta["model"] == "gpt-5-mini"
assert result.meta["finish_reason"] == "tool_calls"
assert result.meta["index"] == 0
assert result.meta["completion_start_time"] is not None
assert result.meta["usage"] == {
"completion_tokens": 42,
"prompt_tokens": 282,
"total_tokens": 324,
"completion_tokens_details": {
"accepted_prediction_tokens": 0,
"audio_tokens": 0,
"reasoning_tokens": 0,
"rejected_prediction_tokens": 0,
},
"prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0},
}
def test_convert_usage_chunk_to_streaming_chunk(self):
usage_chunk = ChatCompletionChunk(
id="chatcmpl-BC1y4wqIhe17R8sv3lgLcWlB4tXCw",
choices=[],
created=1742207200,
model="gpt-5-mini",
object="chat.completion.chunk",
service_tier="default",
system_fingerprint="fp_06737a9306",
usage=CompletionUsage(
completion_tokens=8,
prompt_tokens=13,
total_tokens=21,
completion_tokens_details=CompletionTokensDetails(
accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0
),
prompt_tokens_details=PromptTokensDetails(audio_tokens=0, cached_tokens=0),
),
)
result = _convert_chat_completion_chunk_to_streaming_chunk(chunk=usage_chunk, previous_chunks=[])
assert result.content == ""
assert result.start is False
assert result.tool_calls is None
assert result.tool_call_result is None
assert result.meta["model"] == "gpt-5-mini"
assert result.meta["received_at"] is not None
class TestMakeSchemaStrict:
def test_flat_object(self):
schema = {"type": "object", "properties": {"name": {"type": "string"}}}
result = _make_schema_strict(schema)
assert result == {
"type": "object",
"properties": {"name": {"type": "string"}},
"additionalProperties": False,
"required": ["name"],
}
def test_nested_object(self):
schema = {
"type": "object",
"properties": {
"person": {"type": "object", "properties": {"name": {"type": "string"}, "age": {"type": "integer"}}}
},
}
result = _make_schema_strict(schema)
assert result == {
"type": "object",
"properties": {
"person": {
"type": "object",
"properties": {"name": {"type": "string"}, "age": {"type": "integer"}},
"additionalProperties": False,
"required": ["name", "age"],
}
},
"additionalProperties": False,
"required": ["person"],
}
def test_defs_and_ref(self):
schema = {
"type": "object",
"properties": {"address": {"$ref": "#/$defs/Address"}},
"$defs": {
"Address": {"type": "object", "properties": {"street": {"type": "string"}, "city": {"type": "string"}}}
},
}
result = _make_schema_strict(schema)
assert result == {
"type": "object",
"properties": {"address": {"$ref": "#/$defs/Address"}},
"$defs": {
"Address": {
"type": "object",
"properties": {"street": {"type": "string"}, "city": {"type": "string"}},
"additionalProperties": False,
"required": ["street", "city"],
}
},
"additionalProperties": False,
"required": ["address"],
}
def test_array_items(self):
schema = {
"type": "object",
"properties": {
"people": {"type": "array", "items": {"type": "object", "properties": {"name": {"type": "string"}}}}
},
}
result = _make_schema_strict(schema)
assert result == {
"type": "object",
"properties": {
"people": {
"type": "array",
"items": {
"type": "object",
"properties": {"name": {"type": "string"}},
"additionalProperties": False,
"required": ["name"],
},
}
},
"additionalProperties": False,
"required": ["people"],
}
def test_anyof(self):
schema = {
"type": "object",
"properties": {
"value": {"anyOf": [{"type": "string"}, {"type": "object", "properties": {"x": {"type": "integer"}}}]}
},
}
result = _make_schema_strict(schema)
assert result == {
"type": "object",
"properties": {
"value": {
"anyOf": [
{"type": "string"},
{
"type": "object",
"properties": {"x": {"type": "integer"}},
"additionalProperties": False,
"required": ["x"],
},
]
}
},
"additionalProperties": False,
"required": ["value"],
}
def test_does_not_mutate_original(self):
schema = {"type": "object", "properties": {"a": {"type": "string"}}}
result = _make_schema_strict(schema)
assert "additionalProperties" not in schema
assert "required" not in schema
assert result == {
"type": "object",
"properties": {"a": {"type": "string"}},
"additionalProperties": False,
"required": ["a"],
}
def test_preserves_existing_required(self):
schema = {
"type": "object",
"properties": {"a": {"type": "string"}, "b": {"type": "integer"}},
"required": ["a"],
}
result = _make_schema_strict(schema)
assert result == {
"type": "object",
"properties": {"a": {"type": "string"}, "b": {"type": "integer"}},
"additionalProperties": False,
"required": ["a", "b"],
}
def test_complex_schema_with_defs_and_combinators(self):
schema = {
"type": "object",
"properties": {
"messages": {"type": "array", "items": {"$ref": "#/$defs/ChatMessage"}},
"config": {
"oneOf": [
{"type": "null"},
{
"type": "object",
"properties": {"temperature": {"type": "number"}, "max_tokens": {"type": "integer"}},
},
]
},
},
"$defs": {
"ChatMessage": {
"type": "object",
"properties": {
"role": {"type": "string"},
"content": {"anyOf": [{"type": "string"}, {"type": "null"}]},
"meta": {
"type": "object",
"properties": {
"model": {"type": "string"},
"usage": {
"type": "object",
"properties": {
"prompt_tokens": {"type": "integer"},
"completion_tokens": {"type": "integer"},
},
},
},
},
},
}
},
}
result = _make_schema_strict(schema)
assert result == {
"type": "object",
"properties": {
"messages": {"type": "array", "items": {"$ref": "#/$defs/ChatMessage"}},
"config": {
"oneOf": [
{"type": "null"},
{
"type": "object",
"properties": {"temperature": {"type": "number"}, "max_tokens": {"type": "integer"}},
"additionalProperties": False,
"required": ["temperature", "max_tokens"],
},
]
},
},
"$defs": {
"ChatMessage": {
"type": "object",
"properties": {
"role": {"type": "string"},
"content": {"anyOf": [{"type": "string"}, {"type": "null"}]},
"meta": {
"type": "object",
"properties": {
"model": {"type": "string"},
"usage": {
"type": "object",
"properties": {
"prompt_tokens": {"type": "integer"},
"completion_tokens": {"type": "integer"},
},
"additionalProperties": False,
"required": ["prompt_tokens", "completion_tokens"],
},
},
"additionalProperties": False,
"required": ["model", "usage"],
},
},
"additionalProperties": False,
"required": ["role", "content", "meta"],
}
},
"additionalProperties": False,
"required": ["messages", "config"],
}
def test_prepare_api_call_strict_nested_tool(self):
nested_tool = Tool(
name="create_person",
description="Create a person record",
parameters={
"type": "object",
"properties": {
"name": {"type": "string"},
"address": {
"type": "object",
"properties": {"street": {"type": "string"}, "city": {"type": "string"}},
},
},
"required": ["name"],
},
function=lambda name, address: f"{name} at {address}",
)
component = OpenAIChatGenerator(api_key=Secret.from_token("test-key"), tools_strict=True)
api_args = component._prepare_api_call(messages=[ChatMessage.from_user("test")], tools=[nested_tool])
tool_def = api_args["tools"][0]["function"]
assert tool_def["strict"] is True
assert tool_def["parameters"] == {
"type": "object",
"properties": {
"name": {"type": "string"},
"address": {
"type": "object",
"properties": {"street": {"type": "string"}, "city": {"type": "string"}},
"additionalProperties": False,
"required": ["street", "city"],
},
},
"additionalProperties": False,
"required": ["name", "address"],
}
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY", None),
reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
)
@pytest.mark.integration
def test_live_run_strict_nested_tool(self):
tool = Tool(
name="create_person",
description="Create a person record with an address",
parameters={
"type": "object",
"properties": {
"name": {"type": "string", "description": "Full name"},
"address": {
"type": "object",
"properties": {
"street": {"type": "string", "description": "Street address"},
"city": {"type": "string", "description": "City name"},
},
},
},
},
function=lambda name, address: f"{name} at {address}",
)
component = OpenAIChatGenerator(model="gpt-4.1-nano", tools_strict=True)
results = component.run(
messages=[ChatMessage.from_user("Create a person named John at 123 Main St, Springfield")], tools=[tool]
)
assert len(results["replies"]) == 1
message = results["replies"][0]
assert message.tool_calls
tool_call = message.tool_call
assert tool_call.tool_name == "create_person"
assert "name" in tool_call.arguments
assert "address" in tool_call.arguments
assert "street" in tool_call.arguments["address"]
assert "city" in tool_call.arguments["address"]
Reference in New Issue View Git Blame Copy Permalink