Files
wehub-resource-sync c56bef871b
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
chore: import upstream snapshot with attribution
2026-07-13 13:22:28 +08:00

737 lines
34 KiB
Python

# 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