chore: import upstream snapshot with attribution
CodeQL / Analyze (python) (push) Has been cancelled
Update Platform Components Table / update (push) Has been cancelled
Docker image release / Build base image (push) Has been cancelled
Sync docs with Docusaurus / sync (push) Has been cancelled
Tests / Check if changed (push) Has been cancelled
Tests / format (push) Has been cancelled
Tests / check-imports (push) Has been cancelled
Tests / Unit / macos-latest (push) Has been cancelled
Tests / Unit / ubuntu-latest (push) Has been cancelled
Tests / Unit / windows-latest (push) Has been cancelled
Tests / mypy (push) Has been cancelled
Tests / Integration / ubuntu-latest (push) Has been cancelled
Tests / Integration / macos-latest (push) Has been cancelled
Tests / Integration / windows-latest (push) Has been cancelled
Tests / notify-slack-on-failure (push) Has been cancelled
Tests / Mark tests as completed (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
Update Platform Components Table / update (push) Has been cancelled
Docker image release / Build base image (push) Has been cancelled
Sync docs with Docusaurus / sync (push) Has been cancelled
Tests / Check if changed (push) Has been cancelled
Tests / format (push) Has been cancelled
Tests / check-imports (push) Has been cancelled
Tests / Unit / macos-latest (push) Has been cancelled
Tests / Unit / ubuntu-latest (push) Has been cancelled
Tests / Unit / windows-latest (push) Has been cancelled
Tests / mypy (push) Has been cancelled
Tests / Integration / ubuntu-latest (push) Has been cancelled
Tests / Integration / macos-latest (push) Has been cancelled
Tests / Integration / windows-latest (push) Has been cancelled
Tests / notify-slack-on-failure (push) Has been cancelled
Tests / Mark tests as completed (push) Has been cancelled
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,3 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
@@ -0,0 +1,736 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from openai import OpenAIError
|
||||
from pydantic import BaseModel
|
||||
|
||||
import haystack.components.generators.chat.azure as azure_chat_module
|
||||
from haystack import Pipeline, component
|
||||
from haystack.components.generators.chat import AzureOpenAIChatGenerator
|
||||
from haystack.components.generators.utils import print_streaming_chunk
|
||||
from haystack.dataclasses import ChatMessage, ToolCall
|
||||
from haystack.tools import ComponentTool, Tool
|
||||
from haystack.tools.toolset import Toolset
|
||||
from haystack.utils.auth import Secret
|
||||
from haystack.utils.azure import default_azure_ad_token_provider
|
||||
|
||||
|
||||
class CalendarEvent(BaseModel):
|
||||
event_name: str
|
||||
event_date: str
|
||||
event_location: str
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def calendar_event_model():
|
||||
return CalendarEvent
|
||||
|
||||
|
||||
def get_weather(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"})
|
||||
|
||||
|
||||
@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=get_weather,
|
||||
)
|
||||
# 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 TestAzureOpenAIChatGenerator:
|
||||
def test_supported_models(self) -> None:
|
||||
"""SUPPORTED_MODELS is a non-empty list of strings."""
|
||||
models = AzureOpenAIChatGenerator.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("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
component = AzureOpenAIChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
assert component.api_key == Secret.from_env_var("AZURE_OPENAI_API_KEY", strict=False)
|
||||
assert component.azure_deployment == "gpt-4.1-mini"
|
||||
assert component.streaming_callback is None
|
||||
assert not component.generation_kwargs
|
||||
assert component.client is None
|
||||
assert component.async_client is None
|
||||
|
||||
def test_init_does_not_fail_wo_api_key(self, monkeypatch):
|
||||
monkeypatch.delenv("AZURE_OPENAI_API_KEY", raising=False)
|
||||
monkeypatch.delenv("AZURE_OPENAI_AD_TOKEN", raising=False)
|
||||
component = AzureOpenAIChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
assert component.client is None
|
||||
assert component.async_client is None
|
||||
|
||||
def test_init_with_parameters(self, tools):
|
||||
component = AzureOpenAIChatGenerator(
|
||||
api_key=Secret.from_token("test-api-key"),
|
||||
azure_endpoint="some-non-existing-endpoint",
|
||||
streaming_callback=print_streaming_chunk,
|
||||
generation_kwargs={"max_completion_tokens": 10, "some_test_param": "test-params"},
|
||||
tools=tools,
|
||||
tools_strict=True,
|
||||
azure_ad_token_provider=default_azure_ad_token_provider,
|
||||
)
|
||||
assert component.api_key == Secret.from_token("test-api-key")
|
||||
assert component.azure_deployment == "gpt-4.1-mini"
|
||||
assert component.streaming_callback is print_streaming_chunk
|
||||
assert component.generation_kwargs == {"max_completion_tokens": 10, "some_test_param": "test-params"}
|
||||
assert component.tools == tools
|
||||
assert component.tools_strict
|
||||
assert component.azure_ad_token_provider is not None
|
||||
assert component.max_retries is None
|
||||
assert component.client is None
|
||||
assert component.async_client is None
|
||||
|
||||
def test_init_with_0_max_retries(self, tools):
|
||||
"""Tests that the max_retries init param is set correctly if equal 0"""
|
||||
component = AzureOpenAIChatGenerator(
|
||||
api_key=Secret.from_token("test-api-key"),
|
||||
azure_endpoint="some-non-existing-endpoint",
|
||||
streaming_callback=print_streaming_chunk,
|
||||
generation_kwargs={"max_completion_tokens": 10, "some_test_param": "test-params"},
|
||||
tools=tools,
|
||||
tools_strict=True,
|
||||
azure_ad_token_provider=default_azure_ad_token_provider,
|
||||
max_retries=0,
|
||||
)
|
||||
assert component.api_key == Secret.from_token("test-api-key")
|
||||
assert component.azure_deployment == "gpt-4.1-mini"
|
||||
assert component.streaming_callback is print_streaming_chunk
|
||||
assert component.generation_kwargs == {"max_completion_tokens": 10, "some_test_param": "test-params"}
|
||||
assert component.tools == tools
|
||||
assert component.tools_strict
|
||||
assert component.azure_ad_token_provider is not None
|
||||
assert component.max_retries == 0
|
||||
assert component.client is None
|
||||
assert component.async_client is None
|
||||
|
||||
def test_init_with_secret_azure_endpoint_and_api_version(self, monkeypatch):
|
||||
"""`azure_endpoint` and `api_version` accept a Secret that is resolved from an environment variable."""
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://test-resource.azure.openai.com/")
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_VERSION", "2024-08-01-preview")
|
||||
component = AzureOpenAIChatGenerator(
|
||||
azure_endpoint=Secret.from_env_var("AZURE_OPENAI_ENDPOINT"),
|
||||
api_version=Secret.from_env_var("AZURE_OPENAI_API_VERSION"),
|
||||
)
|
||||
# The Secret objects are kept on the instance so they can be serialized
|
||||
assert component.azure_endpoint == Secret.from_env_var("AZURE_OPENAI_ENDPOINT")
|
||||
assert component.api_version == Secret.from_env_var("AZURE_OPENAI_API_VERSION")
|
||||
|
||||
def test_init_fail_with_unset_secret_azure_endpoint(self, monkeypatch):
|
||||
"""A Secret azure_endpoint that resolves to nothing raises the same error as a missing endpoint."""
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
monkeypatch.delenv("AZURE_OPENAI_ENDPOINT", raising=False)
|
||||
with pytest.raises(ValueError, match="Azure endpoint"):
|
||||
AzureOpenAIChatGenerator(azure_endpoint=Secret.from_env_var("AZURE_OPENAI_ENDPOINT", strict=False))
|
||||
|
||||
def test_to_dict_with_secret_azure_endpoint_and_api_version(self, monkeypatch):
|
||||
"""Secret `azure_endpoint` and `api_version` are serialized as Secret dictionaries."""
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://test-resource.azure.openai.com/")
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_VERSION", "2024-08-01-preview")
|
||||
component = AzureOpenAIChatGenerator(
|
||||
azure_endpoint=Secret.from_env_var("AZURE_OPENAI_ENDPOINT"),
|
||||
api_version=Secret.from_env_var("AZURE_OPENAI_API_VERSION"),
|
||||
)
|
||||
init_params = component.to_dict()["init_parameters"]
|
||||
assert init_params["azure_endpoint"] == {
|
||||
"type": "env_var",
|
||||
"env_vars": ["AZURE_OPENAI_ENDPOINT"],
|
||||
"strict": True,
|
||||
}
|
||||
assert init_params["api_version"] == {
|
||||
"type": "env_var",
|
||||
"env_vars": ["AZURE_OPENAI_API_VERSION"],
|
||||
"strict": True,
|
||||
}
|
||||
|
||||
def test_secret_azure_endpoint_and_api_version_roundtrip(self, monkeypatch):
|
||||
"""Serializing and deserializing a component with Secret endpoint/version restores the Secrets."""
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://test-resource.azure.openai.com/")
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_VERSION", "2024-08-01-preview")
|
||||
component = AzureOpenAIChatGenerator(
|
||||
azure_endpoint=Secret.from_env_var("AZURE_OPENAI_ENDPOINT"),
|
||||
api_version=Secret.from_env_var("AZURE_OPENAI_API_VERSION"),
|
||||
)
|
||||
deserialized = AzureOpenAIChatGenerator.from_dict(component.to_dict())
|
||||
assert deserialized.azure_endpoint == Secret.from_env_var("AZURE_OPENAI_ENDPOINT")
|
||||
assert deserialized.api_version == Secret.from_env_var("AZURE_OPENAI_API_VERSION")
|
||||
deserialized.warm_up()
|
||||
assert str(deserialized.client._azure_endpoint) == "https://test-resource.azure.openai.com/"
|
||||
assert deserialized.client._api_version == "2024-08-01-preview"
|
||||
|
||||
def test_from_dict_with_secret_azure_endpoint_and_api_version(self, monkeypatch):
|
||||
"""from_dict deserializes Secret azure_endpoint/api_version dicts and resolves them for the client."""
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://test-resource.azure.openai.com/")
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_VERSION", "2024-08-01-preview")
|
||||
data = {
|
||||
"type": "haystack.components.generators.chat.azure.AzureOpenAIChatGenerator",
|
||||
"init_parameters": {
|
||||
"api_key": {"env_vars": ["AZURE_OPENAI_API_KEY"], "strict": False, "type": "env_var"},
|
||||
"azure_ad_token": {"env_vars": ["AZURE_OPENAI_AD_TOKEN"], "strict": False, "type": "env_var"},
|
||||
"azure_endpoint": {"env_vars": ["AZURE_OPENAI_ENDPOINT"], "strict": True, "type": "env_var"},
|
||||
"api_version": {"env_vars": ["AZURE_OPENAI_API_VERSION"], "strict": True, "type": "env_var"},
|
||||
"azure_deployment": "gpt-4.1-mini",
|
||||
"organization": None,
|
||||
"streaming_callback": None,
|
||||
"generation_kwargs": {},
|
||||
"timeout": None,
|
||||
"max_retries": None,
|
||||
"default_headers": {},
|
||||
"tools": None,
|
||||
"tools_strict": False,
|
||||
"azure_ad_token_provider": None,
|
||||
"http_client_kwargs": None,
|
||||
},
|
||||
}
|
||||
generator = AzureOpenAIChatGenerator.from_dict(data)
|
||||
# The Secret dicts are deserialized back into Secret objects
|
||||
assert generator.azure_endpoint == Secret.from_env_var("AZURE_OPENAI_ENDPOINT")
|
||||
assert generator.api_version == Secret.from_env_var("AZURE_OPENAI_API_VERSION")
|
||||
# And they are resolved to the string values the client expects
|
||||
generator.warm_up()
|
||||
assert str(generator.client._azure_endpoint) == "https://test-resource.azure.openai.com/"
|
||||
assert generator.client._api_version == "2024-08-01-preview"
|
||||
|
||||
def test_to_dict_default(self, monkeypatch):
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
component = AzureOpenAIChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
data = component.to_dict()
|
||||
assert data == {
|
||||
"type": "haystack.components.generators.chat.azure.AzureOpenAIChatGenerator",
|
||||
"init_parameters": {
|
||||
"api_key": {"env_vars": ["AZURE_OPENAI_API_KEY"], "strict": False, "type": "env_var"},
|
||||
"azure_ad_token": {"env_vars": ["AZURE_OPENAI_AD_TOKEN"], "strict": False, "type": "env_var"},
|
||||
"api_version": "2024-12-01-preview",
|
||||
"azure_endpoint": "some-non-existing-endpoint",
|
||||
"azure_deployment": "gpt-4.1-mini",
|
||||
"organization": None,
|
||||
"streaming_callback": None,
|
||||
"generation_kwargs": {},
|
||||
"timeout": None,
|
||||
"max_retries": None,
|
||||
"default_headers": {},
|
||||
"tools": None,
|
||||
"tools_strict": False,
|
||||
"azure_ad_token_provider": None,
|
||||
"http_client_kwargs": None,
|
||||
},
|
||||
}
|
||||
|
||||
def test_to_dict_with_parameters(self, monkeypatch, calendar_event_model):
|
||||
monkeypatch.setenv("ENV_VAR", "test-api-key")
|
||||
component = AzureOpenAIChatGenerator(
|
||||
api_key=Secret.from_env_var("ENV_VAR", strict=False),
|
||||
azure_ad_token=Secret.from_env_var("ENV_VAR1", strict=False),
|
||||
azure_endpoint="some-non-existing-endpoint",
|
||||
streaming_callback=print_streaming_chunk,
|
||||
timeout=2.5,
|
||||
max_retries=10,
|
||||
generation_kwargs={
|
||||
"max_completion_tokens": 10,
|
||||
"some_test_param": "test-params",
|
||||
"response_format": calendar_event_model,
|
||||
},
|
||||
azure_ad_token_provider=default_azure_ad_token_provider,
|
||||
http_client_kwargs={"proxy": "http://localhost:8080"},
|
||||
)
|
||||
data = component.to_dict()
|
||||
assert data == {
|
||||
"type": "haystack.components.generators.chat.azure.AzureOpenAIChatGenerator",
|
||||
"init_parameters": {
|
||||
"api_key": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
|
||||
"azure_ad_token": {"env_vars": ["ENV_VAR1"], "strict": False, "type": "env_var"},
|
||||
"api_version": "2024-12-01-preview",
|
||||
"azure_endpoint": "some-non-existing-endpoint",
|
||||
"azure_deployment": "gpt-4.1-mini",
|
||||
"organization": None,
|
||||
"streaming_callback": "haystack.components.generators.utils.print_streaming_chunk",
|
||||
"timeout": 2.5,
|
||||
"max_retries": 10,
|
||||
"generation_kwargs": {
|
||||
"max_completion_tokens": 10,
|
||||
"some_test_param": "test-params",
|
||||
"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": None,
|
||||
"tools_strict": False,
|
||||
"default_headers": {},
|
||||
"azure_ad_token_provider": "haystack.utils.azure.default_azure_ad_token_provider",
|
||||
"http_client_kwargs": {"proxy": "http://localhost:8080"},
|
||||
},
|
||||
}
|
||||
|
||||
def test_from_dict(self, monkeypatch):
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
monkeypatch.setenv("AZURE_OPENAI_AD_TOKEN", "test-ad-token")
|
||||
data = {
|
||||
"type": "haystack.components.generators.chat.azure.AzureOpenAIChatGenerator",
|
||||
"init_parameters": {
|
||||
"api_key": {"env_vars": ["AZURE_OPENAI_API_KEY"], "strict": False, "type": "env_var"},
|
||||
"azure_ad_token": {"env_vars": ["AZURE_OPENAI_AD_TOKEN"], "strict": False, "type": "env_var"},
|
||||
"api_version": "2024-12-01-preview",
|
||||
"azure_endpoint": "some-non-existing-endpoint",
|
||||
"azure_deployment": "gpt-4.1-mini",
|
||||
"organization": None,
|
||||
"streaming_callback": None,
|
||||
"generation_kwargs": {},
|
||||
"timeout": 30.0,
|
||||
"max_retries": 5,
|
||||
"default_headers": {},
|
||||
"tools": [
|
||||
{
|
||||
"type": "haystack.tools.tool.Tool",
|
||||
"data": {
|
||||
"description": "description",
|
||||
"function": "builtins.print",
|
||||
"name": "name",
|
||||
"parameters": {"x": {"type": "string"}},
|
||||
},
|
||||
}
|
||||
],
|
||||
"tools_strict": False,
|
||||
"http_client_kwargs": None,
|
||||
},
|
||||
}
|
||||
|
||||
generator = AzureOpenAIChatGenerator.from_dict(data)
|
||||
assert isinstance(generator, AzureOpenAIChatGenerator)
|
||||
|
||||
assert generator.api_key == Secret.from_env_var("AZURE_OPENAI_API_KEY", strict=False)
|
||||
assert generator.azure_ad_token == Secret.from_env_var("AZURE_OPENAI_AD_TOKEN", strict=False)
|
||||
assert generator.api_version == "2024-12-01-preview"
|
||||
assert generator.azure_endpoint == "some-non-existing-endpoint"
|
||||
assert generator.azure_deployment == "gpt-4.1-mini"
|
||||
assert generator.organization is None
|
||||
assert generator.streaming_callback is None
|
||||
assert generator.generation_kwargs == {}
|
||||
assert generator.timeout == 30.0
|
||||
assert generator.max_retries == 5
|
||||
assert generator.default_headers == {}
|
||||
assert generator.tools == [
|
||||
Tool(name="name", description="description", parameters={"x": {"type": "string"}}, function=print)
|
||||
]
|
||||
assert generator.tools_strict is False
|
||||
assert generator.http_client_kwargs is None
|
||||
|
||||
def test_pipeline_serialization_deserialization(self, tmp_path, monkeypatch):
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
generator = AzureOpenAIChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
p = Pipeline()
|
||||
p.add_component(instance=generator, name="generator")
|
||||
|
||||
assert p.to_dict() == {
|
||||
"metadata": {},
|
||||
"max_runs_per_component": 100,
|
||||
"connection_type_validation": True,
|
||||
"components": {
|
||||
"generator": {
|
||||
"type": "haystack.components.generators.chat.azure.AzureOpenAIChatGenerator",
|
||||
"init_parameters": {
|
||||
"azure_endpoint": "some-non-existing-endpoint",
|
||||
"azure_deployment": "gpt-4.1-mini",
|
||||
"organization": None,
|
||||
"api_version": "2024-12-01-preview",
|
||||
"streaming_callback": None,
|
||||
"generation_kwargs": {},
|
||||
"timeout": None,
|
||||
"max_retries": None,
|
||||
"api_key": {"type": "env_var", "env_vars": ["AZURE_OPENAI_API_KEY"], "strict": False},
|
||||
"azure_ad_token": {"type": "env_var", "env_vars": ["AZURE_OPENAI_AD_TOKEN"], "strict": False},
|
||||
"default_headers": {},
|
||||
"tools": None,
|
||||
"tools_strict": False,
|
||||
"azure_ad_token_provider": None,
|
||||
"http_client_kwargs": None,
|
||||
},
|
||||
}
|
||||
},
|
||||
"connections": [],
|
||||
}
|
||||
p_str = p.dumps()
|
||||
q = Pipeline.loads(p_str)
|
||||
assert p.to_dict() == q.to_dict(), "Pipeline serialization/deserialization w/ AzureOpenAIChatGenerator failed."
|
||||
|
||||
def test_azure_chat_generator_with_toolset_initialization(self, tools, monkeypatch):
|
||||
"""Test that the AzureOpenAIChatGenerator can be initialized with a Toolset."""
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
toolset = Toolset(tools)
|
||||
generator = AzureOpenAIChatGenerator(azure_endpoint="some-non-existing-endpoint", tools=toolset)
|
||||
assert generator.tools == toolset
|
||||
|
||||
def test_from_dict_with_toolset(self, tools, monkeypatch):
|
||||
"""Test that the AzureOpenAIChatGenerator can be deserialized from a dictionary with a Toolset."""
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
toolset = Toolset(tools)
|
||||
component = AzureOpenAIChatGenerator(azure_endpoint="some-non-existing-endpoint", tools=toolset)
|
||||
data = component.to_dict()
|
||||
|
||||
deserialized_component = AzureOpenAIChatGenerator.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.integration
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("AZURE_OPENAI_API_KEY", None) or not os.environ.get("AZURE_OPENAI_ENDPOINT", None),
|
||||
reason=(
|
||||
"Please export env variables called AZURE_OPENAI_API_KEY containing "
|
||||
"the Azure OpenAI key, AZURE_OPENAI_ENDPOINT containing "
|
||||
"the Azure OpenAI endpoint URL to run this test."
|
||||
),
|
||||
)
|
||||
def test_live_run(self):
|
||||
chat_messages = [ChatMessage.from_user("What's the capital of France")]
|
||||
component = AzureOpenAIChatGenerator(organization="HaystackCI")
|
||||
results = component.run(chat_messages)
|
||||
assert len(results["replies"]) == 1
|
||||
message: ChatMessage = results["replies"][0]
|
||||
assert "Paris" in message.text
|
||||
assert "gpt-4.1-mini" in message.meta["model"]
|
||||
assert message.meta["finish_reason"] == "stop"
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("AZURE_OPENAI_API_KEY", None) or not os.environ.get("AZURE_OPENAI_ENDPOINT", None),
|
||||
reason=(
|
||||
"Please export env variables called AZURE_OPENAI_API_KEY containing "
|
||||
"the Azure OpenAI key, AZURE_OPENAI_ENDPOINT containing "
|
||||
"the Azure OpenAI endpoint URL to run this test."
|
||||
),
|
||||
)
|
||||
def test_live_run_with_tools(self, tools):
|
||||
chat_messages = [ChatMessage.from_user("What's the weather like in Paris?")]
|
||||
component = AzureOpenAIChatGenerator(organization="HaystackCI", tools=tools)
|
||||
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 == {"city": "Paris"}
|
||||
assert message.meta["finish_reason"] == "tool_calls"
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("AZURE_OPENAI_API_KEY", None),
|
||||
reason="Export an env var called AZURE_OPENAI_API_KEY containing the Azure OpenAI API key to run this test.",
|
||||
)
|
||||
@pytest.mark.integration
|
||||
def test_live_run_with_response_format(self):
|
||||
class CalendarEvent(BaseModel):
|
||||
event_name: str
|
||||
event_date: str
|
||||
event_location: str
|
||||
|
||||
chat_messages = [
|
||||
ChatMessage.from_user("The marketing summit takes place on October12th at the Hilton Hotel downtown.")
|
||||
]
|
||||
component = AzureOpenAIChatGenerator(
|
||||
api_version="2024-08-01-preview", generation_kwargs={"response_format": CalendarEvent}
|
||||
)
|
||||
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)
|
||||
|
||||
assert message.meta["finish_reason"] == "stop"
|
||||
|
||||
def test_to_dict_with_toolset(self, tools, monkeypatch):
|
||||
"""Test that the AzureOpenAIChatGenerator can be serialized to a dictionary with a Toolset."""
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
toolset = Toolset(tools[:1])
|
||||
component = AzureOpenAIChatGenerator(azure_endpoint="some-non-existing-endpoint", tools=toolset)
|
||||
data = component.to_dict()
|
||||
|
||||
expected_tools_data = {
|
||||
"type": "haystack.tools.toolset.Toolset",
|
||||
"data": {
|
||||
"tools": [
|
||||
{
|
||||
"type": "haystack.tools.tool.Tool",
|
||||
"data": {
|
||||
"name": "weather",
|
||||
"description": "useful to determine the weather in a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
"function": "generators.chat.test_azure.get_weather",
|
||||
"async_function": None,
|
||||
"outputs_to_string": None,
|
||||
"inputs_from_state": None,
|
||||
"outputs_to_state": None,
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
assert data["init_parameters"]["tools"] == expected_tools_data
|
||||
|
||||
|
||||
class TestAzureOpenAIChatGeneratorAsync:
|
||||
async def test_warm_up_async_builds_async_client(self, tools):
|
||||
component = AzureOpenAIChatGenerator(
|
||||
api_key=Secret.from_token("test-api-key"),
|
||||
azure_endpoint="some-non-existing-endpoint",
|
||||
streaming_callback=print_streaming_chunk,
|
||||
generation_kwargs={"max_completion_tokens": 10, "some_test_param": "test-params"},
|
||||
tools=tools,
|
||||
tools_strict=True,
|
||||
)
|
||||
assert component.async_client is None
|
||||
await component.warm_up_async()
|
||||
assert component.async_client.api_key == "test-api-key"
|
||||
assert component.client is None
|
||||
assert component.azure_deployment == "gpt-4.1-mini"
|
||||
assert component.streaming_callback is print_streaming_chunk
|
||||
assert component.generation_kwargs == {"max_completion_tokens": 10, "some_test_param": "test-params"}
|
||||
assert component.tools == tools
|
||||
assert component.tools_strict
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("AZURE_OPENAI_API_KEY", None) or not os.environ.get("AZURE_OPENAI_ENDPOINT", None),
|
||||
reason=(
|
||||
"Please export env variables called AZURE_OPENAI_API_KEY containing "
|
||||
"the Azure OpenAI key, AZURE_OPENAI_ENDPOINT containing "
|
||||
"the Azure OpenAI endpoint URL to run this test."
|
||||
),
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_live_run_async(self):
|
||||
component = AzureOpenAIChatGenerator(generation_kwargs={"n": 1})
|
||||
chat_messages = [ChatMessage.from_user("What's the capital of France")]
|
||||
results = await component.run_async(chat_messages)
|
||||
assert len(results["replies"]) == 1
|
||||
message: ChatMessage = results["replies"][0]
|
||||
assert "Paris" in message.text
|
||||
assert "gpt-4.1-mini" in message.meta["model"]
|
||||
assert message.meta["finish_reason"] == "stop"
|
||||
await component.close_async()
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("AZURE_OPENAI_API_KEY", None) or not os.environ.get("AZURE_OPENAI_ENDPOINT", None),
|
||||
reason=(
|
||||
"Please export env variables called AZURE_OPENAI_API_KEY containing "
|
||||
"the Azure OpenAI key, AZURE_OPENAI_ENDPOINT containing "
|
||||
"the Azure OpenAI endpoint URL to run this test."
|
||||
),
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_live_run_with_tools_async(self, tools):
|
||||
component = AzureOpenAIChatGenerator(tools=tools)
|
||||
chat_messages = [ChatMessage.from_user("What's the weather like in Paris?")]
|
||||
results = await component.run_async(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 == {"city": "Paris"}
|
||||
assert message.meta["finish_reason"] == "tool_calls"
|
||||
|
||||
await component.close_async()
|
||||
|
||||
# additional tests intentionally omitted as they are covered by test_openai.py
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_azure_clients(monkeypatch):
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake")
|
||||
sync_cls = MagicMock(name="AzureOpenAI")
|
||||
async_cls = MagicMock(name="AsyncAzureOpenAI")
|
||||
async_cls.return_value.close = AsyncMock()
|
||||
monkeypatch.setattr(azure_chat_module, "AzureOpenAI", sync_cls)
|
||||
monkeypatch.setattr(azure_chat_module, "AsyncAzureOpenAI", async_cls)
|
||||
return sync_cls, async_cls
|
||||
|
||||
|
||||
class TestComponentLifecycle:
|
||||
def test_warm_up_uses_default_timeout_and_max_retries(self, monkeypatch):
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake-api-key")
|
||||
generator = AzureOpenAIChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
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 = AzureOpenAIChatGenerator(
|
||||
api_key=Secret.from_token("fake-api-key"),
|
||||
azure_endpoint="some-non-existing-endpoint",
|
||||
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 = AzureOpenAIChatGenerator(
|
||||
api_key=Secret.from_token("fake-api-key"), azure_endpoint="some-non-existing-endpoint"
|
||||
)
|
||||
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("AZURE_OPENAI_API_KEY", raising=False)
|
||||
monkeypatch.delenv("AZURE_OPENAI_AD_TOKEN", raising=False)
|
||||
generator = AzureOpenAIChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
with pytest.raises(OpenAIError):
|
||||
generator.warm_up()
|
||||
|
||||
def test_warm_up_warms_tools_once(self, monkeypatch):
|
||||
monkeypatch.setenv("AZURE_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 = AzureOpenAIChatGenerator(
|
||||
azure_endpoint="some-non-existing-endpoint", 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("AZURE_OPENAI_API_KEY", "fake-api-key")
|
||||
generator = AzureOpenAIChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
generator.warm_up()
|
||||
assert generator._tools_warmed_up
|
||||
|
||||
def test_sync_lifecycle(self, mock_azure_clients):
|
||||
sync_cls, _ = mock_azure_clients
|
||||
generator = AzureOpenAIChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
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_azure_clients):
|
||||
_, async_cls = mock_azure_clients
|
||||
generator = AzureOpenAIChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
|
||||
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_azure_clients):
|
||||
generator = AzureOpenAIChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
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_azure_clients):
|
||||
generator = AzureOpenAIChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
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
|
||||
@@ -0,0 +1,607 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from haystack import Pipeline, component
|
||||
from haystack.components.generators.chat import AzureOpenAIResponsesChatGenerator
|
||||
from haystack.components.generators.utils import print_streaming_chunk
|
||||
from haystack.dataclasses import ChatMessage, ToolCall
|
||||
from haystack.tools import ComponentTool, Tool
|
||||
from haystack.tools.toolset import Toolset
|
||||
from haystack.utils.auth import Secret
|
||||
from haystack.utils.azure import default_azure_ad_token_provider
|
||||
|
||||
|
||||
class CalendarEvent(BaseModel):
|
||||
event_name: str
|
||||
event_date: str
|
||||
event_location: str
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def calendar_event_model():
|
||||
return CalendarEvent
|
||||
|
||||
|
||||
def get_weather(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"})
|
||||
|
||||
|
||||
@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=get_weather,
|
||||
)
|
||||
# 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 TestInitialization:
|
||||
def test_supported_models(self) -> None:
|
||||
"""SUPPORTED_MODELS is a non-empty list of strings."""
|
||||
models = AzureOpenAIResponsesChatGenerator.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("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
component = AzureOpenAIResponsesChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
assert component.client is None
|
||||
assert component.async_client is None
|
||||
assert component._azure_deployment == "gpt-5-mini"
|
||||
assert component.streaming_callback is None
|
||||
assert not component.generation_kwargs
|
||||
|
||||
def test_init_fail_wo_azure_endpoint(self, monkeypatch):
|
||||
monkeypatch.delenv("AZURE_OPENAI_ENDPOINT", raising=False)
|
||||
with pytest.raises(ValueError):
|
||||
AzureOpenAIResponsesChatGenerator()
|
||||
|
||||
def test_init_with_parameters(self, tools):
|
||||
component = AzureOpenAIResponsesChatGenerator(
|
||||
api_key=Secret.from_token("test-api-key"),
|
||||
azure_endpoint="some-non-existing-endpoint",
|
||||
streaming_callback=print_streaming_chunk,
|
||||
generation_kwargs={"max_completion_tokens": 10, "some_test_param": "test-params"},
|
||||
tools=tools,
|
||||
tools_strict=True,
|
||||
)
|
||||
assert component.client is None
|
||||
assert component.async_client is None
|
||||
assert component._azure_deployment == "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.tools == tools
|
||||
assert component.tools_strict
|
||||
assert component.max_retries is None
|
||||
|
||||
def test_init_with_toolset(self, tools, monkeypatch):
|
||||
"""Test that the AzureOpenAIChatGenerator can be initialized with a Toolset."""
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
toolset = Toolset(tools)
|
||||
generator = AzureOpenAIResponsesChatGenerator(azure_endpoint="some-non-existing-endpoint", tools=toolset)
|
||||
assert generator.tools == toolset
|
||||
|
||||
|
||||
class TestSerDe:
|
||||
def test_to_dict_default(self, monkeypatch):
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
component = AzureOpenAIResponsesChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
data = component.to_dict()
|
||||
assert data == {
|
||||
"type": "haystack.components.generators.chat.azure_responses.AzureOpenAIResponsesChatGenerator",
|
||||
"init_parameters": {
|
||||
"api_key": {"env_vars": ["AZURE_OPENAI_API_KEY"], "strict": False, "type": "env_var"},
|
||||
"azure_endpoint": "some-non-existing-endpoint",
|
||||
"azure_deployment": "gpt-5-mini",
|
||||
"organization": None,
|
||||
"streaming_callback": None,
|
||||
"generation_kwargs": {},
|
||||
"timeout": None,
|
||||
"max_retries": None,
|
||||
"tools": None,
|
||||
"tools_strict": False,
|
||||
"http_client_kwargs": None,
|
||||
},
|
||||
}
|
||||
|
||||
def test_to_dict_with_parameters(self, monkeypatch, calendar_event_model):
|
||||
monkeypatch.setenv("ENV_VAR", "test-api-key")
|
||||
component = AzureOpenAIResponsesChatGenerator(
|
||||
api_key=Secret.from_env_var("ENV_VAR", strict=False),
|
||||
azure_endpoint="some-non-existing-endpoint",
|
||||
streaming_callback=print_streaming_chunk,
|
||||
timeout=2.5,
|
||||
max_retries=10,
|
||||
generation_kwargs={
|
||||
"max_completion_tokens": 10,
|
||||
"some_test_param": "test-params",
|
||||
"text_format": calendar_event_model,
|
||||
},
|
||||
http_client_kwargs={"proxy": "http://localhost:8080"},
|
||||
)
|
||||
data = component.to_dict()
|
||||
assert data == {
|
||||
"type": "haystack.components.generators.chat.azure_responses.AzureOpenAIResponsesChatGenerator",
|
||||
"init_parameters": {
|
||||
"api_key": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"},
|
||||
"azure_endpoint": "some-non-existing-endpoint",
|
||||
"azure_deployment": "gpt-5-mini",
|
||||
"organization": None,
|
||||
"streaming_callback": "haystack.components.generators.utils.print_streaming_chunk",
|
||||
"timeout": 2.5,
|
||||
"max_retries": 10,
|
||||
"generation_kwargs": {
|
||||
"max_completion_tokens": 10,
|
||||
"some_test_param": "test-params",
|
||||
"text": {
|
||||
"format": {
|
||||
"type": "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": None,
|
||||
"tools_strict": False,
|
||||
"http_client_kwargs": {"proxy": "http://localhost:8080"},
|
||||
},
|
||||
}
|
||||
|
||||
def test_to_dict_with_ad_token_provider(self):
|
||||
component = AzureOpenAIResponsesChatGenerator(
|
||||
api_key=default_azure_ad_token_provider, azure_endpoint="some-non-existing-endpoint"
|
||||
)
|
||||
data = component.to_dict()
|
||||
assert data == {
|
||||
"type": "haystack.components.generators.chat.azure_responses.AzureOpenAIResponsesChatGenerator",
|
||||
"init_parameters": {
|
||||
"api_key": "haystack.utils.azure.default_azure_ad_token_provider",
|
||||
"azure_endpoint": "some-non-existing-endpoint",
|
||||
"azure_deployment": "gpt-5-mini",
|
||||
"organization": None,
|
||||
"streaming_callback": None,
|
||||
"generation_kwargs": {},
|
||||
"timeout": None,
|
||||
"max_retries": None,
|
||||
"tools": None,
|
||||
"tools_strict": False,
|
||||
"http_client_kwargs": None,
|
||||
},
|
||||
}
|
||||
|
||||
def test_to_dict_with_toolset(self, tools, monkeypatch):
|
||||
"""Test that the AzureOpenAIChatGenerator can be serialized to a dictionary with a Toolset."""
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
toolset = Toolset(tools[:1])
|
||||
component = AzureOpenAIResponsesChatGenerator(azure_endpoint="some-non-existing-endpoint", tools=toolset)
|
||||
data = component.to_dict()
|
||||
|
||||
expected_tools_data = {
|
||||
"type": "haystack.tools.toolset.Toolset",
|
||||
"data": {
|
||||
"tools": [
|
||||
{
|
||||
"type": "haystack.tools.tool.Tool",
|
||||
"data": {
|
||||
"name": "weather",
|
||||
"description": "useful to determine the weather in a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
"function": "generators.chat.test_azure_responses.get_weather",
|
||||
"async_function": None,
|
||||
"outputs_to_string": None,
|
||||
"inputs_from_state": None,
|
||||
"outputs_to_state": None,
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
assert data["init_parameters"]["tools"] == expected_tools_data
|
||||
|
||||
def test_from_dict(self, monkeypatch):
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
monkeypatch.setenv("AZURE_OPENAI_AD_TOKEN", "test-ad-token")
|
||||
data = {
|
||||
"type": "haystack.components.generators.chat.azure_responses.AzureOpenAIResponsesChatGenerator",
|
||||
"init_parameters": {
|
||||
"api_key": {"env_vars": ["AZURE_OPENAI_API_KEY"], "strict": False, "type": "env_var"},
|
||||
"azure_endpoint": "some-non-existing-endpoint",
|
||||
"azure_deployment": "gpt-5-mini",
|
||||
"organization": None,
|
||||
"streaming_callback": None,
|
||||
"generation_kwargs": {},
|
||||
"timeout": 30.0,
|
||||
"max_retries": 5,
|
||||
"tools": [
|
||||
{
|
||||
"type": "haystack.tools.tool.Tool",
|
||||
"data": {
|
||||
"description": "description",
|
||||
"function": "builtins.print",
|
||||
"name": "name",
|
||||
"parameters": {"x": {"type": "string"}},
|
||||
},
|
||||
}
|
||||
],
|
||||
"tools_strict": False,
|
||||
"http_client_kwargs": None,
|
||||
},
|
||||
}
|
||||
|
||||
generator = AzureOpenAIResponsesChatGenerator.from_dict(data)
|
||||
assert isinstance(generator, AzureOpenAIResponsesChatGenerator)
|
||||
|
||||
assert generator.api_key == Secret.from_env_var("AZURE_OPENAI_API_KEY", strict=False)
|
||||
assert generator._azure_endpoint == "some-non-existing-endpoint"
|
||||
assert generator._azure_deployment == "gpt-5-mini"
|
||||
assert generator.organization is None
|
||||
assert generator.streaming_callback is None
|
||||
assert generator.generation_kwargs == {}
|
||||
assert generator.timeout == 30.0
|
||||
assert generator.max_retries == 5
|
||||
assert generator.tools == [
|
||||
Tool(name="name", description="description", parameters={"x": {"type": "string"}}, function=print)
|
||||
]
|
||||
assert generator.tools_strict is False
|
||||
assert generator.http_client_kwargs is None
|
||||
|
||||
def test_from_dict_with_ad_token_provider(self):
|
||||
data = {
|
||||
"type": "haystack.components.generators.chat.azure_responses.AzureOpenAIResponsesChatGenerator",
|
||||
"init_parameters": {
|
||||
"api_key": "haystack.utils.azure.default_azure_ad_token_provider",
|
||||
"azure_endpoint": "some-non-existing-endpoint",
|
||||
"azure_deployment": "gpt-5-mini",
|
||||
"organization": None,
|
||||
"streaming_callback": None,
|
||||
"generation_kwargs": {},
|
||||
"timeout": None,
|
||||
"max_retries": None,
|
||||
"tools": None,
|
||||
"tools_strict": False,
|
||||
"http_client_kwargs": None,
|
||||
},
|
||||
}
|
||||
|
||||
generator = AzureOpenAIResponsesChatGenerator.from_dict(data)
|
||||
assert isinstance(generator, AzureOpenAIResponsesChatGenerator)
|
||||
|
||||
assert generator.api_key == default_azure_ad_token_provider
|
||||
assert generator._azure_endpoint == "some-non-existing-endpoint"
|
||||
assert generator._azure_deployment == "gpt-5-mini"
|
||||
assert generator.organization is None
|
||||
assert generator.streaming_callback is None
|
||||
assert generator.generation_kwargs == {}
|
||||
assert generator.timeout is None
|
||||
assert generator.max_retries is None
|
||||
assert generator.tools is None
|
||||
assert generator.tools_strict is False
|
||||
assert generator.http_client_kwargs is None
|
||||
|
||||
def test_from_dict_with_toolset(self, tools, monkeypatch):
|
||||
"""Test that the AzureOpenAIChatGenerator can be deserialized from a dictionary with a Toolset."""
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
toolset = Toolset(tools)
|
||||
component = AzureOpenAIResponsesChatGenerator(azure_endpoint="some-non-existing-endpoint", tools=toolset)
|
||||
data = component.to_dict()
|
||||
|
||||
deserialized_component = AzureOpenAIResponsesChatGenerator.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)
|
||||
|
||||
def test_pipeline_serialization_deserialization(self, tmp_path, monkeypatch):
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
generator = AzureOpenAIResponsesChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
p = Pipeline()
|
||||
p.add_component(instance=generator, name="generator")
|
||||
|
||||
assert p.to_dict() == {
|
||||
"metadata": {},
|
||||
"max_runs_per_component": 100,
|
||||
"connection_type_validation": True,
|
||||
"components": {
|
||||
"generator": {
|
||||
"type": "haystack.components.generators.chat.azure_responses.AzureOpenAIResponsesChatGenerator",
|
||||
"init_parameters": {
|
||||
"azure_endpoint": "some-non-existing-endpoint",
|
||||
"azure_deployment": "gpt-5-mini",
|
||||
"organization": None,
|
||||
"streaming_callback": None,
|
||||
"generation_kwargs": {},
|
||||
"timeout": None,
|
||||
"max_retries": None,
|
||||
"api_key": {"type": "env_var", "env_vars": ["AZURE_OPENAI_API_KEY"], "strict": False},
|
||||
"tools": None,
|
||||
"tools_strict": False,
|
||||
"http_client_kwargs": None,
|
||||
},
|
||||
}
|
||||
},
|
||||
"connections": [],
|
||||
}
|
||||
p_str = p.dumps()
|
||||
q = Pipeline.loads(p_str)
|
||||
assert p.to_dict() == q.to_dict()
|
||||
|
||||
|
||||
class TestComponentLifecycle:
|
||||
def test_warm_up_warms_tools_once(self, monkeypatch):
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-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)
|
||||
|
||||
component = AzureOpenAIResponsesChatGenerator(
|
||||
azure_endpoint="some-non-existing-endpoint", tools=[MockTool("tool1"), MockTool("tool2")]
|
||||
)
|
||||
assert not component._tools_warmed_up
|
||||
|
||||
component.warm_up()
|
||||
assert sorted(warm_up_calls) == ["tool1", "tool2"]
|
||||
assert component._tools_warmed_up
|
||||
|
||||
component.warm_up()
|
||||
assert sorted(warm_up_calls) == ["tool1", "tool2"]
|
||||
|
||||
def test_warm_up_with_no_tools_does_not_raise(self, monkeypatch):
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
component = AzureOpenAIResponsesChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
component.warm_up()
|
||||
assert component._tools_warmed_up
|
||||
|
||||
def test_sync_lifecycle(self, monkeypatch):
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
component = AzureOpenAIResponsesChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
assert component.client is None
|
||||
assert component.async_client is None
|
||||
|
||||
component.warm_up()
|
||||
assert component.client is not None
|
||||
assert component.async_client is None
|
||||
|
||||
component.close()
|
||||
assert component.client is None
|
||||
|
||||
async def test_async_lifecycle(self, monkeypatch):
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
component = AzureOpenAIResponsesChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
|
||||
await component.warm_up_async()
|
||||
assert component.async_client is not None
|
||||
assert component.client is None
|
||||
|
||||
await component.close_async()
|
||||
assert component.async_client is None
|
||||
|
||||
async def test_close_is_safe_without_warm_up(self, monkeypatch):
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-api-key")
|
||||
component = AzureOpenAIResponsesChatGenerator(azure_endpoint="some-non-existing-endpoint")
|
||||
component.close()
|
||||
await component.close_async()
|
||||
assert component.client is None
|
||||
assert component.async_client is None
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("AZURE_OPENAI_API_KEY", None) or not os.environ.get("AZURE_OPENAI_ENDPOINT", None),
|
||||
reason=(
|
||||
"Please export env variables called AZURE_OPENAI_API_KEY containing "
|
||||
"the Azure OpenAI key, AZURE_OPENAI_ENDPOINT containing "
|
||||
"the Azure OpenAI endpoint URL to run this test."
|
||||
),
|
||||
)
|
||||
class TestIntegration:
|
||||
def test_live_run(self):
|
||||
chat_messages = [ChatMessage.from_user("What's the capital of France")]
|
||||
component = AzureOpenAIResponsesChatGenerator(azure_deployment="gpt-4o-mini")
|
||||
results = component.run(chat_messages)
|
||||
assert len(results["replies"]) == 1
|
||||
message: ChatMessage = results["replies"][0]
|
||||
assert "paris" in message.text.lower()
|
||||
assert "gpt-4o-mini" in message.meta["model"]
|
||||
assert message.meta["status"] == "completed"
|
||||
|
||||
def test_live_run_with_tools(self, tools):
|
||||
chat_messages = [ChatMessage.from_user("What's the weather like in Paris?")]
|
||||
component = AzureOpenAIResponsesChatGenerator(
|
||||
organization="HaystackCI", tools=tools, azure_deployment="gpt-4o-mini"
|
||||
)
|
||||
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 "city" in tool_call.arguments
|
||||
assert "paris" in tool_call.arguments["city"].lower()
|
||||
assert message.meta["status"] == "completed"
|
||||
|
||||
def test_live_run_with_text_format(self, calendar_event_model):
|
||||
chat_messages = [
|
||||
ChatMessage.from_user("The marketing summit takes place on October12th at the Hilton Hotel downtown.")
|
||||
]
|
||||
component = AzureOpenAIResponsesChatGenerator(
|
||||
azure_deployment="gpt-4o-mini", generation_kwargs={"text_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"].lower()
|
||||
assert isinstance(msg["event_date"], str)
|
||||
assert isinstance(msg["event_location"], str)
|
||||
assert message.meta["status"] == "completed"
|
||||
|
||||
# So far from documentation, responses.parse only supports BaseModel
|
||||
def test_live_run_with_text_format_json_schema(self):
|
||||
json_schema = {
|
||||
"format": {
|
||||
"type": "json_schema",
|
||||
"name": "person",
|
||||
"strict": True,
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string", "minLength": 1},
|
||||
"age": {"type": "number", "minimum": 0, "maximum": 130},
|
||||
},
|
||||
"required": ["name", "age"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
}
|
||||
}
|
||||
chat_messages = [ChatMessage.from_user("Jane 54 years old")]
|
||||
component = AzureOpenAIResponsesChatGenerator(
|
||||
azure_deployment="gpt-4o-mini", generation_kwargs={"text": json_schema}
|
||||
)
|
||||
results = component.run(chat_messages)
|
||||
assert len(results["replies"]) == 1
|
||||
message: ChatMessage = results["replies"][0]
|
||||
msg = json.loads(message.text)
|
||||
assert "jane" in msg["name"].lower()
|
||||
assert msg["age"] == 54
|
||||
assert message.meta["status"] == "completed"
|
||||
assert message.meta["usage"]["output_tokens"] > 0
|
||||
|
||||
|
||||
class TestAzureOpenAIResponsesChatGeneratorAsync:
|
||||
async def test_warm_up_async_creates_async_client_with_expected_args(self, tools):
|
||||
component = AzureOpenAIResponsesChatGenerator(
|
||||
api_key=Secret.from_token("test-api-key"),
|
||||
azure_endpoint="some-non-existing-endpoint",
|
||||
streaming_callback=print_streaming_chunk,
|
||||
generation_kwargs={"max_completion_tokens": 10, "some_test_param": "test-params"},
|
||||
tools=tools,
|
||||
tools_strict=True,
|
||||
)
|
||||
assert component.async_client is None
|
||||
|
||||
await component.warm_up_async()
|
||||
|
||||
assert component.async_client.api_key == "test-api-key"
|
||||
assert component._azure_deployment == "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.tools == tools
|
||||
assert component.tools_strict
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("AZURE_OPENAI_API_KEY", None) or not os.environ.get("AZURE_OPENAI_ENDPOINT", None),
|
||||
reason=(
|
||||
"Please export env variables called AZURE_OPENAI_API_KEY containing "
|
||||
"the Azure OpenAI key, AZURE_OPENAI_ENDPOINT containing "
|
||||
"the Azure OpenAI endpoint URL to run this test."
|
||||
),
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_live_run_async(self):
|
||||
chat_messages = [ChatMessage.from_user("What's the capital of France")]
|
||||
component = AzureOpenAIResponsesChatGenerator(azure_deployment="gpt-4o-mini")
|
||||
results = await component.run_async(chat_messages)
|
||||
assert len(results["replies"]) == 1
|
||||
message: ChatMessage = results["replies"][0]
|
||||
assert "paris" in message.text.lower()
|
||||
assert "gpt-4o-mini" in message.meta["model"]
|
||||
assert message.meta["status"] == "completed"
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("AZURE_OPENAI_API_KEY", None) or not os.environ.get("AZURE_OPENAI_ENDPOINT", None),
|
||||
reason=(
|
||||
"Please export env variables called AZURE_OPENAI_API_KEY containing "
|
||||
"the Azure OpenAI key, AZURE_OPENAI_ENDPOINT containing "
|
||||
"the Azure OpenAI endpoint URL to run this test."
|
||||
),
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_live_run_with_tools_async(self, tools):
|
||||
chat_messages = [ChatMessage.from_user("What's the weather like in Paris?")]
|
||||
component = AzureOpenAIResponsesChatGenerator(tools=tools, azure_deployment="gpt-4o-mini")
|
||||
results = await component.run_async(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 "city" in tool_call.arguments
|
||||
assert "paris" in tool_call.arguments["city"].lower()
|
||||
assert message.meta["status"] == "completed"
|
||||
|
||||
# additional tests intentionally omitted as they are covered by test_openai_responses.py
|
||||
# and test_openai_responses_conversion.py
|
||||
@@ -0,0 +1,489 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
from urllib.error import HTTPError as URLLibHTTPError
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack import component, default_from_dict, default_to_dict
|
||||
from haystack.components.generators.chat.fallback import FallbackChatGenerator
|
||||
from haystack.core.errors import SerializationError
|
||||
from haystack.dataclasses import ChatMessage, StreamingCallbackT
|
||||
from haystack.tools import ToolsType
|
||||
|
||||
|
||||
@component
|
||||
class _DummySuccessGen:
|
||||
def __init__(self, text: str = "ok", delay: float = 0.0, streaming_callback: StreamingCallbackT | None = None):
|
||||
self.text = text
|
||||
self.delay = delay
|
||||
self.streaming_callback = streaming_callback
|
||||
self.received_messages: list[list[ChatMessage]] = []
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return default_to_dict(self, text=self.text, delay=self.delay, streaming_callback=None)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "_DummySuccessGen":
|
||||
return default_from_dict(cls, data)
|
||||
|
||||
def run(
|
||||
self,
|
||||
messages: list[ChatMessage],
|
||||
generation_kwargs: dict[str, Any] | None = None,
|
||||
tools: ToolsType | None = None,
|
||||
streaming_callback: StreamingCallbackT | None = None,
|
||||
) -> dict[str, Any]:
|
||||
self.received_messages.append(messages)
|
||||
if self.delay:
|
||||
time.sleep(self.delay)
|
||||
if streaming_callback:
|
||||
streaming_callback({"dummy": True}) # type: ignore[arg-type]
|
||||
return {"replies": [ChatMessage.from_assistant(self.text)], "meta": {"dummy_meta": True}}
|
||||
|
||||
async def run_async(
|
||||
self,
|
||||
messages: list[ChatMessage],
|
||||
generation_kwargs: dict[str, Any] | None = None,
|
||||
tools: ToolsType | None = None,
|
||||
streaming_callback: StreamingCallbackT | None = None,
|
||||
) -> dict[str, Any]:
|
||||
self.received_messages.append(messages)
|
||||
if self.delay:
|
||||
await asyncio.sleep(self.delay)
|
||||
if streaming_callback:
|
||||
await asyncio.sleep(0)
|
||||
streaming_callback({"dummy": True}) # type: ignore[arg-type]
|
||||
return {"replies": [ChatMessage.from_assistant(self.text)], "meta": {"dummy_meta": True}}
|
||||
|
||||
|
||||
@component
|
||||
class _DummyFailGen:
|
||||
def __init__(self, exc: Exception | None = None, delay: float = 0.0):
|
||||
self.exc = exc or RuntimeError("boom")
|
||||
self.delay = delay
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return default_to_dict(self, exc={"message": str(self.exc)}, delay=self.delay)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "_DummyFailGen":
|
||||
init = data.get("init_parameters", {})
|
||||
msg = None
|
||||
if isinstance(init.get("exc"), dict):
|
||||
msg = init.get("exc", {}).get("message")
|
||||
return cls(exc=RuntimeError(msg or "boom"), delay=init.get("delay", 0.0))
|
||||
|
||||
def run(
|
||||
self,
|
||||
messages: list[ChatMessage],
|
||||
generation_kwargs: dict[str, Any] | None = None,
|
||||
tools: ToolsType | None = None,
|
||||
streaming_callback: StreamingCallbackT | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if self.delay:
|
||||
time.sleep(self.delay)
|
||||
raise self.exc
|
||||
|
||||
async def run_async(
|
||||
self,
|
||||
messages: list[ChatMessage],
|
||||
generation_kwargs: dict[str, Any] | None = None,
|
||||
tools: ToolsType | None = None,
|
||||
streaming_callback: StreamingCallbackT | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if self.delay:
|
||||
await asyncio.sleep(self.delay)
|
||||
raise self.exc
|
||||
|
||||
|
||||
def test_init_validation():
|
||||
with pytest.raises(ValueError):
|
||||
FallbackChatGenerator(chat_generators=[])
|
||||
|
||||
gen = FallbackChatGenerator(chat_generators=[_DummySuccessGen(text="A")])
|
||||
assert len(gen.chat_generators) == 1
|
||||
|
||||
|
||||
def test_sequential_first_success():
|
||||
gen = FallbackChatGenerator(chat_generators=[_DummySuccessGen(text="A")])
|
||||
res = gen.run([ChatMessage.from_user("hi")])
|
||||
assert res["replies"][0].text == "A"
|
||||
assert res["meta"]["successful_chat_generator_index"] == 0
|
||||
assert res["meta"]["total_attempts"] == 1
|
||||
|
||||
|
||||
def test_run_with_string_input():
|
||||
inner = _DummySuccessGen()
|
||||
gen = FallbackChatGenerator(chat_generators=[inner])
|
||||
res = gen.run("hi")
|
||||
assert inner.received_messages[0] == [ChatMessage.from_user("hi")]
|
||||
assert isinstance(res["replies"][0], ChatMessage)
|
||||
|
||||
|
||||
async def test_run_async_with_string_input():
|
||||
inner = _DummySuccessGen()
|
||||
gen = FallbackChatGenerator(chat_generators=[inner])
|
||||
res = await gen.run_async("hi")
|
||||
assert inner.received_messages[0] == [ChatMessage.from_user("hi")]
|
||||
assert isinstance(res["replies"][0], ChatMessage)
|
||||
|
||||
|
||||
def test_sequential_second_success_after_failure():
|
||||
gen = FallbackChatGenerator(chat_generators=[_DummyFailGen(), _DummySuccessGen(text="B")])
|
||||
res = gen.run([ChatMessage.from_user("hi")])
|
||||
assert res["replies"][0].text == "B"
|
||||
assert res["meta"]["successful_chat_generator_index"] == 1
|
||||
assert res["meta"]["failed_chat_generators"]
|
||||
|
||||
|
||||
def test_all_fail_raises():
|
||||
gen = FallbackChatGenerator(chat_generators=[_DummyFailGen(), _DummyFailGen()])
|
||||
with pytest.raises(RuntimeError):
|
||||
gen.run([ChatMessage.from_user("hi")])
|
||||
|
||||
|
||||
def test_timeout_handling_sync():
|
||||
slow = _DummySuccessGen(text="slow", delay=0.01)
|
||||
fast = _DummySuccessGen(text="fast", delay=0.0)
|
||||
gen = FallbackChatGenerator(chat_generators=[slow, fast])
|
||||
res = gen.run([ChatMessage.from_user("hi")])
|
||||
assert res["replies"][0].text == "slow"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timeout_handling_async():
|
||||
slow = _DummySuccessGen(text="slow", delay=0.01)
|
||||
fast = _DummySuccessGen(text="fast", delay=0.0)
|
||||
gen = FallbackChatGenerator(chat_generators=[slow, fast])
|
||||
res = await gen.run_async([ChatMessage.from_user("hi")])
|
||||
assert res["replies"][0].text == "slow"
|
||||
|
||||
|
||||
def test_streaming_callback_forwarding_sync():
|
||||
calls: list[Any] = []
|
||||
|
||||
def cb(x: Any) -> None:
|
||||
calls.append(x)
|
||||
|
||||
gen = FallbackChatGenerator(chat_generators=[_DummySuccessGen(text="A")])
|
||||
_ = gen.run([ChatMessage.from_user("hi")], streaming_callback=cb)
|
||||
assert calls
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_callback_forwarding_async():
|
||||
calls: list[Any] = []
|
||||
|
||||
def cb(x: Any) -> None:
|
||||
calls.append(x)
|
||||
|
||||
gen = FallbackChatGenerator(chat_generators=[_DummySuccessGen(text="A")])
|
||||
_ = await gen.run_async([ChatMessage.from_user("hi")], streaming_callback=cb)
|
||||
assert calls
|
||||
|
||||
|
||||
def test_serialization_roundtrip():
|
||||
original = FallbackChatGenerator(chat_generators=[_DummySuccessGen(text="hello")])
|
||||
data = original.to_dict()
|
||||
restored = FallbackChatGenerator.from_dict(data)
|
||||
assert isinstance(restored, FallbackChatGenerator)
|
||||
assert len(restored.chat_generators) == 1
|
||||
res = restored.run([ChatMessage.from_user("hi")])
|
||||
assert res["replies"][0].text == "hello"
|
||||
|
||||
original = FallbackChatGenerator(chat_generators=[_DummySuccessGen(text="hello"), _DummySuccessGen(text="world")])
|
||||
data = original.to_dict()
|
||||
restored = FallbackChatGenerator.from_dict(data)
|
||||
assert isinstance(restored, FallbackChatGenerator)
|
||||
assert len(restored.chat_generators) == 2
|
||||
res = restored.run([ChatMessage.from_user("hi")])
|
||||
assert res["replies"][0].text == "hello"
|
||||
|
||||
|
||||
def test_automatic_completion_mode_without_streaming():
|
||||
gen = FallbackChatGenerator(chat_generators=[_DummySuccessGen(text="completion")])
|
||||
res = gen.run([ChatMessage.from_user("hi")])
|
||||
assert res["replies"][0].text == "completion"
|
||||
assert res["meta"]["successful_chat_generator_index"] == 0
|
||||
|
||||
|
||||
def test_automatic_ttft_mode_with_streaming():
|
||||
calls: list[Any] = []
|
||||
|
||||
def cb(x: Any) -> None:
|
||||
calls.append(x)
|
||||
|
||||
gen = FallbackChatGenerator(chat_generators=[_DummySuccessGen(text="streaming")])
|
||||
res = gen.run([ChatMessage.from_user("hi")], streaming_callback=cb)
|
||||
assert res["replies"][0].text == "streaming"
|
||||
assert calls
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_automatic_ttft_mode_with_streaming_async():
|
||||
calls: list[Any] = []
|
||||
|
||||
def cb(x: Any) -> None:
|
||||
calls.append(x)
|
||||
|
||||
gen = FallbackChatGenerator(chat_generators=[_DummySuccessGen(text="streaming_async")])
|
||||
res = await gen.run_async([ChatMessage.from_user("hi")], streaming_callback=cb)
|
||||
assert res["replies"][0].text == "streaming_async"
|
||||
assert calls
|
||||
|
||||
|
||||
def create_http_error(status_code: int, message: str) -> URLLibHTTPError:
|
||||
return URLLibHTTPError("", status_code, message, {}, None)
|
||||
|
||||
|
||||
@component
|
||||
class _DummyHTTPErrorGen:
|
||||
def __init__(self, text: str = "success", error: Exception | None = None):
|
||||
self.text = text
|
||||
self.error = error
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return default_to_dict(self, text=self.text, error=str(self.error) if self.error else None)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "_DummyHTTPErrorGen":
|
||||
init = data.get("init_parameters", {})
|
||||
error = None
|
||||
if init.get("error"):
|
||||
error = RuntimeError(init["error"])
|
||||
return cls(text=init.get("text", "success"), error=error)
|
||||
|
||||
def run(
|
||||
self,
|
||||
messages: list[ChatMessage],
|
||||
generation_kwargs: dict[str, Any] | None = None,
|
||||
tools: ToolsType | None = None,
|
||||
streaming_callback: StreamingCallbackT | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if self.error:
|
||||
raise self.error
|
||||
return {
|
||||
"replies": [ChatMessage.from_assistant(self.text)],
|
||||
"meta": {"error_type": type(self.error).__name__ if self.error else None},
|
||||
}
|
||||
|
||||
|
||||
def test_failover_trigger_429_rate_limit():
|
||||
rate_limit_gen = _DummyHTTPErrorGen(text="rate_limited", error=create_http_error(429, "Rate limit exceeded"))
|
||||
success_gen = _DummySuccessGen(text="success_after_rate_limit")
|
||||
|
||||
fallback = FallbackChatGenerator(chat_generators=[rate_limit_gen, success_gen])
|
||||
result = fallback.run([ChatMessage.from_user("test")])
|
||||
|
||||
assert result["replies"][0].text == "success_after_rate_limit"
|
||||
assert result["meta"]["successful_chat_generator_index"] == 1
|
||||
assert result["meta"]["failed_chat_generators"] == ["_DummyHTTPErrorGen"]
|
||||
|
||||
|
||||
def test_failover_trigger_401_authentication():
|
||||
auth_error_gen = _DummyHTTPErrorGen(text="auth_failed", error=create_http_error(401, "Authentication failed"))
|
||||
success_gen = _DummySuccessGen(text="success_after_auth")
|
||||
|
||||
fallback = FallbackChatGenerator(chat_generators=[auth_error_gen, success_gen])
|
||||
result = fallback.run([ChatMessage.from_user("test")])
|
||||
|
||||
assert result["replies"][0].text == "success_after_auth"
|
||||
assert result["meta"]["successful_chat_generator_index"] == 1
|
||||
assert result["meta"]["failed_chat_generators"] == ["_DummyHTTPErrorGen"]
|
||||
|
||||
|
||||
def test_failover_trigger_400_bad_request():
|
||||
bad_request_gen = _DummyHTTPErrorGen(text="bad_request", error=create_http_error(400, "Context length exceeded"))
|
||||
success_gen = _DummySuccessGen(text="success_after_bad_request")
|
||||
|
||||
fallback = FallbackChatGenerator(chat_generators=[bad_request_gen, success_gen])
|
||||
result = fallback.run([ChatMessage.from_user("test")])
|
||||
|
||||
assert result["replies"][0].text == "success_after_bad_request"
|
||||
assert result["meta"]["successful_chat_generator_index"] == 1
|
||||
assert result["meta"]["failed_chat_generators"] == ["_DummyHTTPErrorGen"]
|
||||
|
||||
|
||||
def test_failover_trigger_500_server_error():
|
||||
server_error_gen = _DummyHTTPErrorGen(text="server_error", error=create_http_error(500, "Internal server error"))
|
||||
success_gen = _DummySuccessGen(text="success_after_server_error")
|
||||
|
||||
fallback = FallbackChatGenerator(chat_generators=[server_error_gen, success_gen])
|
||||
result = fallback.run([ChatMessage.from_user("test")])
|
||||
|
||||
assert result["replies"][0].text == "success_after_server_error"
|
||||
assert result["meta"]["successful_chat_generator_index"] == 1
|
||||
assert result["meta"]["failed_chat_generators"] == ["_DummyHTTPErrorGen"]
|
||||
|
||||
|
||||
def test_failover_trigger_multiple_errors():
|
||||
rate_limit_gen = _DummyHTTPErrorGen(text="rate_limited", error=create_http_error(429, "Rate limit exceeded"))
|
||||
auth_error_gen = _DummyHTTPErrorGen(text="auth_failed", error=create_http_error(401, "Authentication failed"))
|
||||
server_error_gen = _DummyHTTPErrorGen(text="server_error", error=create_http_error(500, "Internal server error"))
|
||||
success_gen = _DummySuccessGen(text="success_after_all_errors")
|
||||
|
||||
fallback = FallbackChatGenerator(chat_generators=[rate_limit_gen, auth_error_gen, server_error_gen, success_gen])
|
||||
result = fallback.run([ChatMessage.from_user("test")])
|
||||
|
||||
assert result["replies"][0].text == "success_after_all_errors"
|
||||
assert result["meta"]["successful_chat_generator_index"] == 3
|
||||
assert len(result["meta"]["failed_chat_generators"]) == 3
|
||||
|
||||
|
||||
def test_failover_trigger_all_generators_fail():
|
||||
rate_limit_gen = _DummyHTTPErrorGen(text="rate_limited", error=create_http_error(429, "Rate limit exceeded"))
|
||||
auth_error_gen = _DummyHTTPErrorGen(text="auth_failed", error=create_http_error(401, "Authentication failed"))
|
||||
server_error_gen = _DummyHTTPErrorGen(text="server_error", error=create_http_error(500, "Internal server error"))
|
||||
|
||||
fallback = FallbackChatGenerator(chat_generators=[rate_limit_gen, auth_error_gen, server_error_gen])
|
||||
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
fallback.run([ChatMessage.from_user("test")])
|
||||
|
||||
error_msg = str(exc_info.value)
|
||||
assert "All 3 chat generators failed" in error_msg
|
||||
assert "Failed chat generators: [_DummyHTTPErrorGen, _DummyHTTPErrorGen, _DummyHTTPErrorGen]" in error_msg
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failover_trigger_429_rate_limit_async():
|
||||
rate_limit_gen = _DummyHTTPErrorGen(text="rate_limited", error=create_http_error(429, "Rate limit exceeded"))
|
||||
success_gen = _DummySuccessGen(text="success_after_rate_limit")
|
||||
|
||||
fallback = FallbackChatGenerator(chat_generators=[rate_limit_gen, success_gen])
|
||||
result = await fallback.run_async([ChatMessage.from_user("test")])
|
||||
|
||||
assert result["replies"][0].text == "success_after_rate_limit"
|
||||
assert result["meta"]["successful_chat_generator_index"] == 1
|
||||
assert result["meta"]["failed_chat_generators"] == ["_DummyHTTPErrorGen"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failover_trigger_401_authentication_async():
|
||||
auth_error_gen = _DummyHTTPErrorGen(text="auth_failed", error=create_http_error(401, "Authentication failed"))
|
||||
success_gen = _DummySuccessGen(text="success_after_auth")
|
||||
|
||||
fallback = FallbackChatGenerator(chat_generators=[auth_error_gen, success_gen])
|
||||
result = await fallback.run_async([ChatMessage.from_user("test")])
|
||||
|
||||
assert result["replies"][0].text == "success_after_auth"
|
||||
assert result["meta"]["successful_chat_generator_index"] == 1
|
||||
assert result["meta"]["failed_chat_generators"] == ["_DummyHTTPErrorGen"]
|
||||
|
||||
|
||||
class TestComponentLifecycle:
|
||||
def test_warm_up_delegates_to_every_generator(self):
|
||||
gens = [Mock(spec=["run", "warm_up"]) for _ in range(3)]
|
||||
fallback = FallbackChatGenerator(chat_generators=gens)
|
||||
fallback.warm_up()
|
||||
for gen in gens:
|
||||
gen.warm_up.assert_called_once()
|
||||
|
||||
async def test_warm_up_async_delegates_to_every_generator(self):
|
||||
gens = [Mock(spec=["run", "warm_up_async"]) for _ in range(3)]
|
||||
for gen in gens:
|
||||
gen.warm_up_async = AsyncMock()
|
||||
fallback = FallbackChatGenerator(chat_generators=gens)
|
||||
await fallback.warm_up_async()
|
||||
for gen in gens:
|
||||
gen.warm_up_async.assert_awaited_once()
|
||||
|
||||
async def test_warm_up_async_falls_back_to_sync_warm_up(self):
|
||||
gens = [Mock(spec=["run", "warm_up"]) for _ in range(3)]
|
||||
fallback = FallbackChatGenerator(chat_generators=gens)
|
||||
await fallback.warm_up_async()
|
||||
for gen in gens:
|
||||
gen.warm_up.assert_called_once()
|
||||
|
||||
def test_close_delegates_to_every_generator(self):
|
||||
gens = [Mock(spec=["run", "close"]) for _ in range(3)]
|
||||
fallback = FallbackChatGenerator(chat_generators=gens)
|
||||
fallback.close()
|
||||
for gen in gens:
|
||||
gen.close.assert_called_once()
|
||||
|
||||
async def test_close_async_delegates_to_every_generator(self):
|
||||
gens = [Mock(spec=["run", "close_async"]) for _ in range(3)]
|
||||
for gen in gens:
|
||||
gen.close_async = AsyncMock()
|
||||
fallback = FallbackChatGenerator(chat_generators=gens)
|
||||
await fallback.close_async()
|
||||
for gen in gens:
|
||||
gen.close_async.assert_awaited_once()
|
||||
|
||||
async def test_close_async_falls_back_to_sync_close(self):
|
||||
gens = [Mock(spec=["run", "close"]) for _ in range(3)]
|
||||
fallback = FallbackChatGenerator(chat_generators=gens)
|
||||
await fallback.close_async()
|
||||
for gen in gens:
|
||||
gen.close.assert_called_once()
|
||||
|
||||
def test_lifecycle_is_safe_when_generators_lack_methods(self):
|
||||
gens = [Mock(spec=["run"]) for _ in range(3)]
|
||||
fallback = FallbackChatGenerator(chat_generators=gens)
|
||||
fallback.warm_up()
|
||||
fallback.close()
|
||||
|
||||
|
||||
@component
|
||||
class CustomGeneratorWithoutSerDe:
|
||||
def __init__(self, text: str = "custom_ok"):
|
||||
self.text = text
|
||||
|
||||
def run(
|
||||
self,
|
||||
messages: list[ChatMessage],
|
||||
generation_kwargs: dict[str, Any] | None = None,
|
||||
tools: ToolsType | None = None,
|
||||
streaming_callback: StreamingCallbackT | None = None,
|
||||
) -> dict[str, Any]:
|
||||
return {"replies": [ChatMessage.from_assistant(self.text)], "meta": {}}
|
||||
|
||||
|
||||
@component
|
||||
class NonSerializableGenerator:
|
||||
def __init__(self, non_serializable_arg: Any):
|
||||
self.non_serializable_arg = non_serializable_arg
|
||||
|
||||
def run(self, messages: list[ChatMessage]) -> dict[str, Any]:
|
||||
return {"replies": []}
|
||||
|
||||
|
||||
def test_serialization_with_custom_generators_without_to_dict():
|
||||
# 1. Test mixed chain serialization, order preservation, execution, and round-trip
|
||||
gen0 = _DummySuccessGen(text="dummy_has_dict")
|
||||
gen1 = CustomGeneratorWithoutSerDe(text="custom_no_dict_1")
|
||||
gen2 = CustomGeneratorWithoutSerDe(text="custom_no_dict_2")
|
||||
|
||||
original = FallbackChatGenerator(chat_generators=[gen0, gen1, gen2])
|
||||
data = original.to_dict()
|
||||
|
||||
# Ensure all three components are serialized and not silently omitted
|
||||
assert len(data["init_parameters"]["chat_generators"]) == 3
|
||||
|
||||
# Reconstruct/Deserialize
|
||||
restored = FallbackChatGenerator.from_dict(data)
|
||||
assert isinstance(restored, FallbackChatGenerator)
|
||||
assert len(restored.chat_generators) == 3
|
||||
|
||||
# Assert fallback order is exactly preserved
|
||||
assert restored.chat_generators[0].text == "dummy_has_dict"
|
||||
assert restored.chat_generators[1].text == "custom_no_dict_1"
|
||||
assert restored.chat_generators[2].text == "custom_no_dict_2"
|
||||
assert isinstance(restored.chat_generators[1], CustomGeneratorWithoutSerDe)
|
||||
assert isinstance(restored.chat_generators[2], CustomGeneratorWithoutSerDe)
|
||||
|
||||
# Verify pipeline execution on the restored instance
|
||||
res = restored.run([ChatMessage.from_user("hi")])
|
||||
assert res["replies"][0].text == "dummy_has_dict"
|
||||
|
||||
# 2. Test failure path (fail loud) when a component is not serializable
|
||||
non_serializable_fallback = FallbackChatGenerator(chat_generators=[NonSerializableGenerator(object())])
|
||||
with pytest.raises(SerializationError, match="unsupported value of type"):
|
||||
non_serializable_fallback.to_dict()
|
||||
@@ -0,0 +1,369 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack import Document, Pipeline, component
|
||||
from haystack.components.agents.agent import Agent
|
||||
from haystack.components.generators.chat import LLM
|
||||
from haystack.components.generators.chat.openai import OpenAIChatGenerator
|
||||
from haystack.components.joiners.branch import BranchJoiner
|
||||
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
|
||||
from haystack.components.routers.conditional_router import ConditionalRouter
|
||||
from haystack.core.component.types import InputSocket, OutputSocket
|
||||
from haystack.dataclasses import ChatMessage
|
||||
from haystack.dataclasses.chat_message import ChatRole
|
||||
from haystack.dataclasses.streaming_chunk import StreamingChunk
|
||||
from haystack.document_stores.in_memory import InMemoryDocumentStore
|
||||
from haystack.tools import Tool
|
||||
from haystack.tools.toolset import Toolset
|
||||
|
||||
|
||||
def sync_streaming_callback(chunk: StreamingChunk) -> None:
|
||||
pass
|
||||
|
||||
|
||||
@component
|
||||
class MockChatGeneratorWithTools:
|
||||
"""A mock chat generator that accepts a tools parameter."""
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {"type": "test_llm.MockChatGeneratorWithTools", "data": {}}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "MockChatGeneratorWithTools":
|
||||
return cls()
|
||||
|
||||
@component.output_types(replies=list[ChatMessage])
|
||||
def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs) -> dict[str, Any]:
|
||||
return {"replies": [ChatMessage.from_assistant("Reply with tools support")]}
|
||||
|
||||
@component.output_types(replies=list[ChatMessage])
|
||||
async def run_async(
|
||||
self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs
|
||||
) -> dict[str, Any]:
|
||||
return {"replies": [ChatMessage.from_assistant("Async reply with tools support")]}
|
||||
|
||||
|
||||
@component
|
||||
class MockChatGenerator:
|
||||
"""A mock chat generator that does NOT accept a tools parameter."""
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {"type": "test_llm.MockChatGenerator", "data": {}}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "MockChatGenerator":
|
||||
return cls()
|
||||
|
||||
@component.output_types(replies=list[ChatMessage])
|
||||
def run(self, messages: list[ChatMessage], **kwargs) -> dict[str, Any]:
|
||||
return {"replies": [ChatMessage.from_assistant("Sync reply")]}
|
||||
|
||||
@component.output_types(replies=list[ChatMessage])
|
||||
async def run_async(self, messages: list[ChatMessage], **kwargs) -> dict[str, Any]:
|
||||
return {"replies": [ChatMessage.from_assistant("Async reply")]}
|
||||
|
||||
|
||||
class TestLLM:
|
||||
class TestInit:
|
||||
USER_PROMPT = '{% message role="user" %}{{ query }}{% endmessage %}'
|
||||
|
||||
def test_is_subclass_of_agent(self):
|
||||
assert issubclass(LLM, Agent)
|
||||
|
||||
def test_defaults(self):
|
||||
llm = LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT)
|
||||
assert llm.chat_generator is not None
|
||||
assert llm.tools == []
|
||||
assert llm.system_prompt is None
|
||||
assert llm.user_prompt == self.USER_PROMPT
|
||||
assert llm.required_variables == "*"
|
||||
assert llm.streaming_callback is None
|
||||
|
||||
def test_output_sockets(self):
|
||||
llm = LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT)
|
||||
assert llm.__haystack_output__._sockets_dict == {
|
||||
"messages": OutputSocket(name="messages", type=list[ChatMessage], receivers=[]),
|
||||
"last_message": OutputSocket(name="last_message", type=ChatMessage, receivers=[]),
|
||||
"token_usage": OutputSocket(name="token_usage", type=dict[str, Any], receivers=[]),
|
||||
}
|
||||
|
||||
def test_detects_no_tools_support(self):
|
||||
llm = LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT)
|
||||
assert llm._chat_generator_supports_tools is False
|
||||
|
||||
def test_detects_tools_support(self):
|
||||
llm = LLM(chat_generator=MockChatGeneratorWithTools(), user_prompt=self.USER_PROMPT)
|
||||
assert llm._chat_generator_supports_tools is True
|
||||
|
||||
def test_messages_required_when_no_prompt_variables(self):
|
||||
llm = LLM(
|
||||
chat_generator=MockChatGenerator(), user_prompt='{% message role="user" %}Hello world{% endmessage %}'
|
||||
)
|
||||
messages_socket = llm.__haystack_input__._sockets_dict["messages"]
|
||||
assert isinstance(messages_socket, InputSocket)
|
||||
assert messages_socket.is_mandatory
|
||||
|
||||
def test_messages_optional_when_prompt_has_variables(self):
|
||||
llm = LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT)
|
||||
messages_socket = llm.__haystack_input__._sockets_dict["messages"]
|
||||
assert isinstance(messages_socket, InputSocket)
|
||||
assert not messages_socket.is_mandatory
|
||||
|
||||
def test_messages_optional_when_plain_prompt_has_variables(self):
|
||||
llm = LLM(chat_generator=MockChatGenerator(), user_prompt="Question: {{ query }}")
|
||||
messages_socket = llm.__haystack_input__._sockets_dict["messages"]
|
||||
assert isinstance(messages_socket, InputSocket)
|
||||
assert not messages_socket.is_mandatory
|
||||
assert "query" in llm.__haystack_input__._sockets_dict
|
||||
|
||||
def test_runtime_prompt_overrides_not_component_inputs(self):
|
||||
llm = LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT)
|
||||
assert "system_prompt" not in llm.__haystack_input__._sockets_dict
|
||||
assert "user_prompt" not in llm.__haystack_input__._sockets_dict
|
||||
|
||||
def test_raises_if_required_variables_empty(self):
|
||||
with pytest.raises(ValueError, match="required_variables must not be empty"):
|
||||
LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT, required_variables=[])
|
||||
|
||||
class TestSerialization:
|
||||
def test_to_dict_excludes_agent_only_params(self, monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
|
||||
user_prompt = '{% message role="user" %}{{ query }}{% endmessage %}'
|
||||
llm = LLM(chat_generator=OpenAIChatGenerator(), system_prompt="You are helpful.", user_prompt=user_prompt)
|
||||
|
||||
serialized = llm.to_dict()
|
||||
|
||||
assert serialized["type"] == "haystack.components.generators.chat.llm.LLM"
|
||||
assert "chat_generator" in serialized["init_parameters"]
|
||||
assert serialized["init_parameters"]["system_prompt"] == "You are helpful."
|
||||
|
||||
agent_only_params = [
|
||||
"tools",
|
||||
"exit_conditions",
|
||||
"max_agent_steps",
|
||||
"raise_on_tool_invocation_failure",
|
||||
"tool_concurrency_limit",
|
||||
"tool_streaming_callback_passthrough",
|
||||
"confirmation_strategies",
|
||||
"state_schema",
|
||||
]
|
||||
for param in agent_only_params:
|
||||
assert param not in serialized["init_parameters"], (
|
||||
f"Agent-only param '{param}' should not be serialized"
|
||||
)
|
||||
|
||||
def test_to_dict_includes_llm_params(self, monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
|
||||
llm = LLM(
|
||||
chat_generator=OpenAIChatGenerator(),
|
||||
system_prompt="Be concise.",
|
||||
user_prompt='{% message role="user" %}{{ query }}{% endmessage %}',
|
||||
required_variables=["query"],
|
||||
)
|
||||
|
||||
serialized = llm.to_dict()
|
||||
|
||||
assert serialized["init_parameters"]["system_prompt"] == "Be concise."
|
||||
assert "{{ query }}" in serialized["init_parameters"]["user_prompt"]
|
||||
assert serialized["init_parameters"]["required_variables"] == ["query"]
|
||||
assert serialized["init_parameters"]["streaming_callback"] is None
|
||||
|
||||
def test_from_dict(self, monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
|
||||
data = {
|
||||
"type": "haystack.components.generators.chat.llm.LLM",
|
||||
"init_parameters": {
|
||||
"chat_generator": {
|
||||
"type": "haystack.components.generators.chat.openai.OpenAIChatGenerator",
|
||||
"init_parameters": {
|
||||
"model": "gpt-4o-mini",
|
||||
"streaming_callback": None,
|
||||
"api_base_url": None,
|
||||
"organization": None,
|
||||
"generation_kwargs": {},
|
||||
"api_key": {"type": "env_var", "env_vars": ["OPENAI_API_KEY"], "strict": True},
|
||||
"timeout": None,
|
||||
"max_retries": None,
|
||||
"tools": None,
|
||||
"tools_strict": False,
|
||||
"http_client_kwargs": None,
|
||||
},
|
||||
},
|
||||
"system_prompt": "You are helpful.",
|
||||
"user_prompt": '{% message role="user" %}{{ query }}{% endmessage %}',
|
||||
"required_variables": "*",
|
||||
"streaming_callback": None,
|
||||
},
|
||||
}
|
||||
|
||||
llm = LLM.from_dict(data)
|
||||
|
||||
assert isinstance(llm, LLM)
|
||||
assert isinstance(llm.chat_generator, OpenAIChatGenerator)
|
||||
assert llm.system_prompt == "You are helpful."
|
||||
assert llm.tools == []
|
||||
|
||||
def test_roundtrip(self, monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
|
||||
user_prompt = '{% message role="user" %}{{ query }}{% endmessage %}'
|
||||
original = LLM(
|
||||
chat_generator=OpenAIChatGenerator(), system_prompt="You are a poet.", user_prompt=user_prompt
|
||||
)
|
||||
|
||||
restored = LLM.from_dict(original.to_dict())
|
||||
|
||||
assert isinstance(restored, LLM)
|
||||
assert isinstance(restored.chat_generator, OpenAIChatGenerator)
|
||||
assert restored.system_prompt == original.system_prompt
|
||||
assert restored.tools == []
|
||||
|
||||
class TestRun:
|
||||
USER_PROMPT = '{% message role="user" %}{{ query }}{% endmessage %}'
|
||||
|
||||
def test_run_accepts_messages_via_kwargs(self):
|
||||
llm = LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT)
|
||||
prior_message = ChatMessage.from_user("Some prior context")
|
||||
result = llm.run(query="What is 2+2?", messages=[prior_message])
|
||||
assert result["last_message"].text == "Sync reply"
|
||||
assert prior_message in result["messages"]
|
||||
|
||||
def test_run_without_messages(self):
|
||||
llm = LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT)
|
||||
result = llm.run(query="What is 2+2?")
|
||||
assert result["last_message"].text == "Sync reply"
|
||||
user_messages = [m for m in result["messages"] if m.is_from(ChatRole.USER)]
|
||||
assert any("What is 2+2?" in m.text for m in user_messages)
|
||||
|
||||
def test_run_with_plain_user_prompt(self):
|
||||
llm = LLM(chat_generator=MockChatGenerator(), user_prompt="Question: {{ query }}")
|
||||
result = llm.run(query="What is 2+2?")
|
||||
assert result["last_message"].text == "Sync reply"
|
||||
user_messages = [m for m in result["messages"] if m.is_from(ChatRole.USER)]
|
||||
assert any("Question: What is 2+2?" in m.text for m in user_messages)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_accepts_messages_via_kwargs(self):
|
||||
llm = LLM(chat_generator=MockChatGenerator(), user_prompt=self.USER_PROMPT)
|
||||
prior_message = ChatMessage.from_user("Some prior context")
|
||||
result = await llm.run_async(query="What is 2+2?", messages=[prior_message])
|
||||
assert result["last_message"].text == "Async reply"
|
||||
assert prior_message in result["messages"]
|
||||
|
||||
class TestPipelineIntegration:
|
||||
@pytest.fixture()
|
||||
def document_store_with_docs(self):
|
||||
store = InMemoryDocumentStore()
|
||||
store.write_documents(
|
||||
[
|
||||
Document(content="The Eiffel Tower is located in Paris."),
|
||||
Document(content="The Brandenburg Gate is in Berlin."),
|
||||
Document(content="The Colosseum is in Rome."),
|
||||
]
|
||||
)
|
||||
return store
|
||||
|
||||
def test_rag_pipeline(self, document_store_with_docs):
|
||||
user_prompt = (
|
||||
'{% message role="user" %}'
|
||||
"Use the following documents to answer the question.\n"
|
||||
"Documents:\n{% for doc in documents %}{{ doc.content }}\n{% endfor %}"
|
||||
"Question: {{ query }}"
|
||||
"{% endmessage %}"
|
||||
)
|
||||
llm = LLM(
|
||||
chat_generator=MockChatGenerator(),
|
||||
system_prompt="You are a knowledgeable assistant.",
|
||||
user_prompt=user_prompt,
|
||||
required_variables=["query", "documents"],
|
||||
)
|
||||
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("retriever", InMemoryBM25Retriever(document_store=document_store_with_docs))
|
||||
pipe.add_component("llm", llm)
|
||||
pipe.connect("retriever.documents", "llm.documents")
|
||||
|
||||
query = "Where is the Colosseum?"
|
||||
result = pipe.run(data={"retriever": {"query": query}, "llm": {"query": query}})
|
||||
|
||||
assert "llm" in result
|
||||
llm_output = result["llm"]
|
||||
assert "messages" in llm_output
|
||||
assert "last_message" in llm_output
|
||||
|
||||
messages = llm_output["messages"]
|
||||
|
||||
assert messages[0].is_from(ChatRole.SYSTEM)
|
||||
assert messages[0].text == "You are a knowledgeable assistant."
|
||||
|
||||
user_messages = [m for m in messages if m.is_from(ChatRole.USER)]
|
||||
assert len(user_messages) == 1
|
||||
rendered = user_messages[0].text
|
||||
assert "Question: Where is the Colosseum?" in rendered
|
||||
assert "Documents:" in rendered
|
||||
assert "Colosseum" in rendered
|
||||
|
||||
assert llm_output["last_message"].is_from(ChatRole.ASSISTANT)
|
||||
assert llm_output["last_message"].text == "Sync reply"
|
||||
|
||||
|
||||
class TestLLMNotTriggeredByInjectedInput:
|
||||
"""
|
||||
Regression guard for the optional-messages scheduling hazard described in
|
||||
https://github.com/deepset-ai/haystack/issues/11109.
|
||||
|
||||
When `user_prompt` contains template variables, `messages` is optional on the LLM.
|
||||
An optional input with `sender=None` (i.e., injected directly via `pipeline.run`)
|
||||
would flip `has_user_input()` to True and incorrectly trigger the component even
|
||||
when its required inputs (e.g. `query`) never arrive.
|
||||
"""
|
||||
|
||||
def test_llm_not_triggered_by_injected_streaming_callback(self):
|
||||
|
||||
@component
|
||||
class Planner:
|
||||
@component.output_types(messages=list[ChatMessage], last_role=str)
|
||||
def run(self) -> dict:
|
||||
return {"messages": [ChatMessage.from_user("hello")], "last_role": "assistant"}
|
||||
|
||||
chat_generator = MockChatGenerator()
|
||||
llm = LLM(chat_generator=chat_generator)
|
||||
chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("x")]})
|
||||
|
||||
router = ConditionalRouter(
|
||||
routes=[
|
||||
{
|
||||
"condition": "{{ last_role == 'tool' }}",
|
||||
"output": "{{ messages }}",
|
||||
"output_name": "processing",
|
||||
"output_type": list[ChatMessage],
|
||||
},
|
||||
{
|
||||
"condition": "{{ True }}",
|
||||
"output": "{{ messages }}",
|
||||
"output_name": "planning",
|
||||
"output_type": list[ChatMessage],
|
||||
},
|
||||
],
|
||||
unsafe=True,
|
||||
)
|
||||
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("planner", Planner())
|
||||
pipeline.add_component("router", router)
|
||||
pipeline.add_component("branch_joiner", BranchJoiner(type_=list[ChatMessage]))
|
||||
pipeline.add_component("llm", llm)
|
||||
pipeline.connect("planner.messages", "router.messages")
|
||||
pipeline.connect("planner.last_role", "router.last_role")
|
||||
pipeline.connect("router.processing", "branch_joiner.value")
|
||||
pipeline.connect("branch_joiner.value", "llm.messages")
|
||||
|
||||
result = pipeline.run(data={"llm": {"streaming_callback": sync_streaming_callback}})
|
||||
|
||||
assert "llm" not in result
|
||||
chat_generator.run.assert_not_called()
|
||||
@@ -0,0 +1,204 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import inspect
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack import Pipeline
|
||||
from haystack.components.generators.chat import MockChatGenerator
|
||||
from haystack.dataclasses import ChatMessage, StreamingChunk, ToolCall
|
||||
|
||||
|
||||
def _exclaim(messages: list[ChatMessage]) -> str:
|
||||
"""Module-level response function (returns a string) used to test `response_fn` and its serialization."""
|
||||
return f"{messages[-1].text}!"
|
||||
|
||||
|
||||
def _assistant_reply(messages: list[ChatMessage]) -> ChatMessage:
|
||||
"""Module-level response function that returns a full ChatMessage."""
|
||||
return ChatMessage.from_assistant("canned message")
|
||||
|
||||
|
||||
def _noop_callback(chunk: StreamingChunk) -> None:
|
||||
"""Module-level streaming callback used to test init-level callback serialization."""
|
||||
|
||||
|
||||
class TestMockChatGenerator:
|
||||
@pytest.mark.parametrize(
|
||||
("args", "kwargs", "exception", "match"),
|
||||
[
|
||||
(("a",), {"response_fn": _exclaim}, ValueError, "either 'responses' or 'response_fn'"),
|
||||
(([],), {}, ValueError, "must not be an empty list"),
|
||||
((123,), {}, TypeError, "must be a string, ChatMessage, or a sequence"),
|
||||
(([123],), {}, TypeError, "Each response must be a string or ChatMessage"),
|
||||
((ChatMessage.from_user("hi"),), {}, ValueError, "must have the 'assistant' role"),
|
||||
],
|
||||
)
|
||||
def test_init_rejects_invalid_config(self, args, kwargs, exception, match):
|
||||
with pytest.raises(exception, match=match):
|
||||
MockChatGenerator(*args, **kwargs)
|
||||
|
||||
def test_fixed_response(self):
|
||||
gen = MockChatGenerator("the same answer")
|
||||
for _ in range(3):
|
||||
result = gen.run([ChatMessage.from_user("anything")])
|
||||
assert result["replies"][0].text == "the same answer"
|
||||
|
||||
def test_cycling_responses(self):
|
||||
# a mix of strings and ChatMessage objects, returned in order and wrapping around
|
||||
gen = MockChatGenerator(["one", ChatMessage.from_assistant("two"), "three"])
|
||||
texts = [gen.run([ChatMessage.from_user("hi")])["replies"][0].text for _ in range(4)]
|
||||
assert texts == ["one", "two", "three", "one"]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("messages", "expected"),
|
||||
[
|
||||
(
|
||||
[ChatMessage.from_system("sys"), ChatMessage.from_user("first"), ChatMessage.from_user("second")],
|
||||
"second",
|
||||
),
|
||||
([ChatMessage.from_system("only system")], "only system"), # falls back to the last message with text
|
||||
([], None), # nothing to echo
|
||||
],
|
||||
)
|
||||
def test_echo_default(self, messages, expected):
|
||||
replies = MockChatGenerator().run(messages)["replies"]
|
||||
if expected is None:
|
||||
assert replies == []
|
||||
else:
|
||||
assert replies[0].text == expected
|
||||
|
||||
@pytest.mark.parametrize(("fn", "expected"), [(_exclaim, "hello!"), (_assistant_reply, "canned message")])
|
||||
def test_response_fn(self, fn, expected):
|
||||
result = MockChatGenerator(response_fn=fn).run([ChatMessage.from_user("hello")])
|
||||
assert result["replies"][0].text == expected
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("fn", "exception", "match"),
|
||||
[
|
||||
(lambda messages: 123, TypeError, "must return a string or ChatMessage"),
|
||||
(lambda messages: ChatMessage.from_user("nope"), ValueError, "must return an assistant ChatMessage"),
|
||||
],
|
||||
)
|
||||
def test_response_fn_invalid_return_raises(self, fn, exception, match):
|
||||
with pytest.raises(exception, match=match):
|
||||
MockChatGenerator(response_fn=fn).run([ChatMessage.from_user("hi")])
|
||||
|
||||
def test_string_input_is_normalized(self):
|
||||
gen = MockChatGenerator(response_fn=_exclaim)
|
||||
assert gen.run("plain string")["replies"][0].text == "plain string!"
|
||||
|
||||
def test_tool_call_response(self):
|
||||
tool_call = ToolCall(tool_name="search", arguments={"query": "Haystack"})
|
||||
gen = MockChatGenerator(ChatMessage.from_assistant(tool_calls=[tool_call]))
|
||||
reply = gen.run([ChatMessage.from_user("search for Haystack")])["replies"][0]
|
||||
assert reply.tool_calls == [tool_call]
|
||||
assert reply.meta["finish_reason"] == "tool_calls"
|
||||
|
||||
def test_meta_defaults(self):
|
||||
meta = MockChatGenerator("hello world").run([ChatMessage.from_user("a b c")])["replies"][0].meta
|
||||
assert meta["model"] == "mock-model"
|
||||
assert meta["finish_reason"] == "stop"
|
||||
assert meta["usage"] == {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}
|
||||
|
||||
def test_meta_merging_precedence(self):
|
||||
# init meta overrides defaults; per-response meta overrides init meta
|
||||
response = ChatMessage.from_assistant("hi", meta={"custom": "from-response", "finish_reason": "length"})
|
||||
gen = MockChatGenerator(response, model="custom-model", meta={"custom": "from-init", "extra": "init"})
|
||||
meta = gen.run([ChatMessage.from_user("x")])["replies"][0].meta
|
||||
assert meta["model"] == "custom-model"
|
||||
assert meta["custom"] == "from-response"
|
||||
assert meta["finish_reason"] == "length"
|
||||
assert meta["extra"] == "init"
|
||||
|
||||
def test_does_not_mutate_stored_responses(self):
|
||||
gen = MockChatGenerator("hello")
|
||||
gen.run([ChatMessage.from_user("a b")])
|
||||
# the stored response keeps its original (empty) meta, untouched by the per-run meta
|
||||
assert gen._responses[0].meta == {}
|
||||
|
||||
async def test_run_async(self):
|
||||
gen = MockChatGenerator(["one", "two"])
|
||||
assert (await gen.run_async([ChatMessage.from_user("hi")]))["replies"][0].text == "one"
|
||||
assert (await gen.run_async([ChatMessage.from_user("hi")]))["replies"][0].text == "two"
|
||||
# echo mode with empty input returns no replies (async path)
|
||||
assert (await MockChatGenerator().run_async([]))["replies"] == []
|
||||
|
||||
def test_streaming_callback_sync(self):
|
||||
chunks: list[StreamingChunk] = []
|
||||
result = MockChatGenerator("hello there friend").run(
|
||||
[ChatMessage.from_user("hi")], streaming_callback=chunks.append
|
||||
)
|
||||
assert "".join(chunk.content for chunk in chunks) == "hello there friend"
|
||||
assert chunks[0].start is True
|
||||
assert chunks[-1].finish_reason == "stop"
|
||||
# the returned reply matches the predefined response
|
||||
assert result["replies"][0].text == "hello there friend"
|
||||
|
||||
def test_run_signature_matches_openai_order(self):
|
||||
# run()/run_async() must mirror OpenAIChatGenerator's parameter order so the mock is a positional drop-in.
|
||||
expected = [
|
||||
("self", inspect.Parameter.POSITIONAL_OR_KEYWORD),
|
||||
("messages", inspect.Parameter.POSITIONAL_OR_KEYWORD),
|
||||
("streaming_callback", inspect.Parameter.POSITIONAL_OR_KEYWORD),
|
||||
("generation_kwargs", inspect.Parameter.POSITIONAL_OR_KEYWORD),
|
||||
("tools", inspect.Parameter.KEYWORD_ONLY),
|
||||
("tools_strict", inspect.Parameter.KEYWORD_ONLY),
|
||||
]
|
||||
for method in ("run", "run_async"):
|
||||
params = list(inspect.signature(getattr(MockChatGenerator, method)).parameters.values())
|
||||
assert [(p.name, p.kind) for p in params] == expected
|
||||
|
||||
# passing the callback as the 2nd positional arg must be treated as streaming_callback, not generation_kwargs
|
||||
chunks: list[StreamingChunk] = []
|
||||
MockChatGenerator("hi").run([ChatMessage.from_user("x")], chunks.append)
|
||||
assert chunks
|
||||
|
||||
async def test_streaming_callback_async(self):
|
||||
chunks: list[StreamingChunk] = []
|
||||
|
||||
async def callback(chunk: StreamingChunk) -> None:
|
||||
chunks.append(chunk)
|
||||
|
||||
await MockChatGenerator("hello world").run_async([ChatMessage.from_user("hi")], streaming_callback=callback)
|
||||
assert "".join(chunk.content for chunk in chunks) == "hello world"
|
||||
assert chunks[-1].finish_reason == "stop"
|
||||
|
||||
def test_streaming_empty_reply(self):
|
||||
chunks: list[StreamingChunk] = []
|
||||
MockChatGenerator("").run([ChatMessage.from_user("hi")], streaming_callback=chunks.append)
|
||||
assert chunks[-1].finish_reason == "stop"
|
||||
|
||||
def test_streaming_callback_with_tool_call(self):
|
||||
chunks: list[StreamingChunk] = []
|
||||
tool_call = ToolCall(tool_name="search", arguments={"query": "x"})
|
||||
gen = MockChatGenerator(ChatMessage.from_assistant(tool_calls=[tool_call]))
|
||||
gen.run([ChatMessage.from_user("hi")], streaming_callback=chunks.append)
|
||||
assert any(chunk.tool_calls for chunk in chunks)
|
||||
assert chunks[-1].finish_reason == "tool_calls"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"generator",
|
||||
[
|
||||
MockChatGenerator(["a", ChatMessage.from_assistant("b")], model="m", meta={"k": "v"}),
|
||||
MockChatGenerator(response_fn=_exclaim),
|
||||
MockChatGenerator(), # echo mode
|
||||
MockChatGenerator("hi", streaming_callback=_noop_callback), # serialized init-level callback
|
||||
],
|
||||
ids=["responses", "response_fn", "echo", "streaming_callback"],
|
||||
)
|
||||
def test_serialization_roundtrip(self, generator):
|
||||
restored = MockChatGenerator.from_dict(generator.to_dict())
|
||||
assert isinstance(restored, MockChatGenerator)
|
||||
# behavior is preserved across the roundtrip
|
||||
messages = [ChatMessage.from_user("hi")]
|
||||
assert restored.run(messages)["replies"][0].text == generator.run(messages)["replies"][0].text
|
||||
|
||||
def test_in_pipeline(self):
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("generator", MockChatGenerator("from the pipeline"))
|
||||
restored = Pipeline.from_dict(pipeline.to_dict())
|
||||
result = restored.run({"generator": {"messages": [ChatMessage.from_user("hi")]}})
|
||||
assert result["generator"]["replies"][0].text == "from the pipeline"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,495 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import os
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from openai import AsyncOpenAI, AsyncStream, OpenAIError
|
||||
from openai.types.chat import (
|
||||
ChatCompletion,
|
||||
ChatCompletionChunk,
|
||||
ChatCompletionMessage,
|
||||
ChatCompletionMessageFunctionToolCall,
|
||||
chat_completion_chunk,
|
||||
)
|
||||
from openai.types.chat.chat_completion import Choice
|
||||
from openai.types.chat.chat_completion_message_function_tool_call import Function
|
||||
from openai.types.completion_usage import CompletionTokensDetails, CompletionUsage, PromptTokensDetails
|
||||
|
||||
from haystack.components.generators.chat.openai import OpenAIChatGenerator
|
||||
from haystack.dataclasses import ChatMessage, StreamingChunk, ToolCall
|
||||
from haystack.tools import Tool
|
||||
from haystack.utils.auth import Secret
|
||||
|
||||
|
||||
@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_async):
|
||||
"""
|
||||
Mock the OpenAI API completion chunk response and reuse it for tests
|
||||
"""
|
||||
|
||||
with patch(
|
||||
"openai.resources.chat.completions.AsyncCompletions.create", new_callable=AsyncMock
|
||||
) 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()),
|
||||
usage=None,
|
||||
)
|
||||
mock_chat_completion_create.return_value = openai_mock_stream_async(completion)
|
||||
yield mock_chat_completion_create
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tools():
|
||||
tool_parameters = {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}
|
||||
tool = Tool(
|
||||
name="weather",
|
||||
description="useful to determine the weather in a given location",
|
||||
parameters=tool_parameters,
|
||||
function=lambda x: x,
|
||||
)
|
||||
return [tool]
|
||||
|
||||
|
||||
class TestOpenAIChatGeneratorAsync:
|
||||
async def test_warm_up_async_should_create_async_client_with_same_args(self, monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
|
||||
component = OpenAIChatGenerator(
|
||||
api_key=Secret.from_token("test-api-key"),
|
||||
api_base_url="test-base-url",
|
||||
organization="test-organization",
|
||||
timeout=30,
|
||||
max_retries=5,
|
||||
)
|
||||
await component.warm_up_async()
|
||||
|
||||
assert isinstance(component.async_client, AsyncOpenAI)
|
||||
assert component.async_client.api_key == "test-api-key"
|
||||
assert component.async_client.organization == "test-organization"
|
||||
assert component.async_client.base_url == "test-base-url/"
|
||||
assert component.async_client.timeout == 30
|
||||
assert component.async_client.max_retries == 5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async(self, chat_messages, openai_mock_async_chat_completion):
|
||||
component = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"))
|
||||
response = await component.run_async(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"]]
|
||||
|
||||
async def test_run_async_with_string_input(self, openai_mock_async_chat_completion):
|
||||
component = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"))
|
||||
response = await component.run_async("What's the capital of France?")
|
||||
|
||||
_, kwargs = openai_mock_async_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)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_with_params_async(self, chat_messages, openai_mock_async_chat_completion):
|
||||
component = OpenAIChatGenerator(
|
||||
api_key=Secret.from_token("test-api-key"),
|
||||
generation_kwargs={"max_completion_tokens": 10, "temperature": 0.5},
|
||||
)
|
||||
response = await component.run_async(chat_messages)
|
||||
|
||||
# check that the component calls the OpenAI API with the correct parameters
|
||||
_, kwargs = openai_mock_async_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"]]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_with_params_streaming_async(self, chat_messages, openai_mock_async_chat_completion_chunk):
|
||||
streaming_callback_called = False
|
||||
|
||||
async 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 = await component.run_async(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
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_with_streaming_callback_in_run_method_async(
|
||||
self, chat_messages, openai_mock_async_chat_completion_chunk
|
||||
):
|
||||
streaming_callback_called = False
|
||||
|
||||
async 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 = await component.run_async(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
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_with_tools_async(self, tools):
|
||||
with patch(
|
||||
"openai.resources.chat.completions.AsyncCompletions.create", new_callable=AsyncMock
|
||||
) 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, tools_strict=True)
|
||||
response = await component.run_async([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
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_with_tools_streaming_async(self, mock_chat_completion_chunk_with_tools, tools):
|
||||
streaming_callback_called = False
|
||||
|
||||
async 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 = await component.run_async(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"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_stream_closes_on_cancellation(self, monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
|
||||
generator = OpenAIChatGenerator(
|
||||
api_key=Secret.from_token("test-api-key"),
|
||||
api_base_url="test-base-url",
|
||||
organization="test-organization",
|
||||
timeout=30,
|
||||
max_retries=5,
|
||||
)
|
||||
|
||||
# mocked the async stream that will be passed to the _handle_async_stream_response() method
|
||||
mock_stream = AsyncMock(spec=AsyncStream)
|
||||
mock_stream.close = AsyncMock()
|
||||
|
||||
async def mock_chunk_generator():
|
||||
for i in range(10):
|
||||
yield MagicMock(
|
||||
choices=[
|
||||
MagicMock(
|
||||
index=0,
|
||||
delta=MagicMock(content=f"chunk{i}", role=None, tool_calls=None),
|
||||
finish_reason=None,
|
||||
logprobs=None,
|
||||
)
|
||||
],
|
||||
model="gpt-4",
|
||||
usage=None,
|
||||
)
|
||||
await asyncio.sleep(0.005) # delay between chunks
|
||||
|
||||
mock_stream.__aiter__ = lambda _: mock_chunk_generator()
|
||||
|
||||
received_chunks = []
|
||||
|
||||
async def test_callback(chunk: StreamingChunk):
|
||||
received_chunks.append(chunk)
|
||||
|
||||
# the task that will be cancelled
|
||||
task = asyncio.create_task(generator._handle_async_stream_response(mock_stream, test_callback))
|
||||
|
||||
# trigger the task, process a few chunks, then cancel
|
||||
await asyncio.sleep(0.01)
|
||||
task.cancel()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
mock_stream.close.assert_awaited_once()
|
||||
|
||||
# we received some chunks before cancellation but not all of them
|
||||
assert len(received_chunks) > 0
|
||||
assert len(received_chunks) < 10
|
||||
|
||||
@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
|
||||
@pytest.mark.asyncio
|
||||
async def test_live_run_async(self):
|
||||
component = OpenAIChatGenerator(model="gpt-4.1-nano", generation_kwargs={"n": 1})
|
||||
chat_messages = [ChatMessage.from_user("What's the capital of France")]
|
||||
results = await component.run_async(chat_messages)
|
||||
assert len(results["replies"]) == 1
|
||||
message: ChatMessage = results["replies"][0]
|
||||
assert "Paris" in message.text
|
||||
assert message.meta["model"]
|
||||
assert message.meta["finish_reason"] == "stop"
|
||||
# Close async client; suppress RuntimeError if the event loop is already closed
|
||||
with contextlib.suppress(RuntimeError):
|
||||
await component.async_client.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_with_wrong_model_async(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.async_client = mock_client
|
||||
|
||||
with pytest.raises(OpenAIError):
|
||||
await generator.run_async([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
|
||||
@pytest.mark.asyncio
|
||||
async def test_live_run_streaming_async(self):
|
||||
counter = 0
|
||||
responses = ""
|
||||
|
||||
async def callback(chunk: StreamingChunk):
|
||||
nonlocal counter
|
||||
nonlocal responses
|
||||
counter += 1
|
||||
responses += chunk.content if chunk.content else ""
|
||||
|
||||
component = OpenAIChatGenerator(
|
||||
model="gpt-4.1-nano",
|
||||
generation_kwargs={"stream_options": {"include_usage": True}},
|
||||
streaming_callback=callback,
|
||||
)
|
||||
results = await component.run_async([ChatMessage.from_user("What's the capital of France?")])
|
||||
|
||||
assert len(results["replies"]) == 1
|
||||
message: ChatMessage = results["replies"][0]
|
||||
assert "Paris" in message.text
|
||||
|
||||
assert message.meta["model"]
|
||||
assert message.meta["finish_reason"] == "stop"
|
||||
|
||||
assert counter > 1
|
||||
assert "Paris" in responses
|
||||
|
||||
# check that the completion_start_time is set and valid ISO format
|
||||
assert "completion_start_time" in message.meta
|
||||
assert datetime.fromisoformat(message.meta["completion_start_time"]) <= datetime.now()
|
||||
|
||||
assert isinstance(message.meta["usage"], dict)
|
||||
assert message.meta["usage"]["prompt_tokens"] > 0
|
||||
assert message.meta["usage"]["completion_tokens"] > 0
|
||||
assert message.meta["usage"]["total_tokens"] > 0
|
||||
|
||||
# Close async client; suppress RuntimeError if the event loop is already closed
|
||||
with contextlib.suppress(RuntimeError):
|
||||
await component.async_client.close()
|
||||
|
||||
@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
|
||||
@pytest.mark.asyncio
|
||||
async def test_live_run_with_tools_async(self, tools):
|
||||
component = OpenAIChatGenerator(model="gpt-4.1-nano", tools=tools)
|
||||
chat_messages = [ChatMessage.from_user("What's the weather like in Paris?")]
|
||||
results = await component.run_async(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"
|
||||
# Check that Paris is in the city argument (case-insensitive, allowing for variations like "Paris, France")
|
||||
assert "paris" in tool_call.arguments["city"].lower()
|
||||
assert message.meta["finish_reason"] == "tool_calls"
|
||||
|
||||
# Close async client; suppress RuntimeError if the event loop is already closed
|
||||
with contextlib.suppress(RuntimeError):
|
||||
await component.async_client.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_with_wrapped_stream_simulation_async(self, chat_messages, openai_mock_stream_async):
|
||||
streaming_callback_called = False
|
||||
|
||||
async 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 async stream in an AsyncMock
|
||||
# This is to simulate the behavior of some tools like Weave (https://github.com/wandb/weave)
|
||||
# which wrap the OpenAI async stream in their own stream
|
||||
wrapped_openai_async_stream = AsyncMock()
|
||||
wrapped_openai_async_stream.__aiter__.return_value = iter([chunk])
|
||||
|
||||
component = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"))
|
||||
await component.warm_up_async()
|
||||
|
||||
# Patch the async client's create method
|
||||
with patch.object(
|
||||
component.async_client.chat.completions,
|
||||
"create",
|
||||
return_value=wrapped_openai_async_stream,
|
||||
new_callable=AsyncMock,
|
||||
) as mock_create:
|
||||
response = await component.run_async(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
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user