chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,400 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
from collections.abc import AsyncGenerator
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from anthropic import AsyncAnthropic
|
||||
from anthropic.lib.streaming import TextEvent
|
||||
from anthropic.lib.streaming._types import InputJsonEvent
|
||||
from anthropic.types import (
|
||||
ContentBlockStopEvent,
|
||||
InputJSONDelta,
|
||||
Message,
|
||||
MessageDeltaUsage,
|
||||
MessageStopEvent,
|
||||
RawContentBlockDeltaEvent,
|
||||
RawContentBlockStartEvent,
|
||||
RawMessageDeltaEvent,
|
||||
RawMessageStartEvent,
|
||||
TextBlock,
|
||||
TextDelta,
|
||||
ToolUseBlock,
|
||||
Usage,
|
||||
)
|
||||
from anthropic.types.raw_message_delta_event import Delta
|
||||
|
||||
from semantic_kernel.connectors.ai.anthropic.prompt_execution_settings.anthropic_prompt_execution_settings import (
|
||||
AnthropicChatPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.contents.chat_message_content import (
|
||||
ChatMessageContent,
|
||||
FunctionCallContent,
|
||||
FunctionResultContent,
|
||||
TextContent,
|
||||
)
|
||||
from semantic_kernel.contents.const import ContentTypes
|
||||
from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent, StreamingTextContent
|
||||
from semantic_kernel.contents.utils.author_role import AuthorRole
|
||||
from semantic_kernel.contents.utils.finish_reason import FinishReason
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tool_calls_message() -> ChatMessageContent:
|
||||
return ChatMessageContent(
|
||||
ai_model_id="claude-3-opus-20240229",
|
||||
metadata={},
|
||||
content_type="message",
|
||||
role=AuthorRole.ASSISTANT,
|
||||
name=None,
|
||||
items=[
|
||||
TextContent(
|
||||
inner_content=None,
|
||||
ai_model_id=None,
|
||||
metadata={},
|
||||
content_type="text",
|
||||
text="<thinking></thinking>",
|
||||
encoding=None,
|
||||
),
|
||||
FunctionCallContent(
|
||||
inner_content=None,
|
||||
ai_model_id=None,
|
||||
metadata={},
|
||||
content_type=ContentTypes.FUNCTION_CALL_CONTENT,
|
||||
id="test_function_call_content",
|
||||
index=1,
|
||||
name="math-Add",
|
||||
function_name="Add",
|
||||
plugin_name="math",
|
||||
arguments={"input": 3, "amount": 3},
|
||||
),
|
||||
],
|
||||
encoding=None,
|
||||
finish_reason=FinishReason.TOOL_CALLS,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_parallel_tool_calls_message() -> ChatMessageContent:
|
||||
return ChatMessageContent(
|
||||
ai_model_id="claude-3-opus-20240229",
|
||||
metadata={},
|
||||
content_type="message",
|
||||
role=AuthorRole.ASSISTANT,
|
||||
name=None,
|
||||
items=[
|
||||
TextContent(
|
||||
inner_content=None,
|
||||
ai_model_id=None,
|
||||
metadata={},
|
||||
content_type="text",
|
||||
text="<thinking></thinking>",
|
||||
encoding=None,
|
||||
),
|
||||
FunctionCallContent(
|
||||
inner_content=None,
|
||||
ai_model_id=None,
|
||||
metadata={},
|
||||
content_type=ContentTypes.FUNCTION_CALL_CONTENT,
|
||||
id="test_function_call_content_1",
|
||||
index=1,
|
||||
name="math-Add",
|
||||
function_name="Add",
|
||||
plugin_name="math",
|
||||
arguments={"input": 3, "amount": 3},
|
||||
),
|
||||
FunctionCallContent(
|
||||
inner_content=None,
|
||||
ai_model_id=None,
|
||||
metadata={},
|
||||
content_type=ContentTypes.FUNCTION_CALL_CONTENT,
|
||||
id="test_function_call_content_2",
|
||||
index=1,
|
||||
name="math-Subtract",
|
||||
function_name="Subtract",
|
||||
plugin_name="math",
|
||||
arguments={"input": 6, "amount": 3},
|
||||
),
|
||||
],
|
||||
encoding=None,
|
||||
finish_reason=FinishReason.TOOL_CALLS,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_streaming_tool_calls_message() -> list:
|
||||
stream_events = [
|
||||
RawMessageStartEvent(
|
||||
message=Message(
|
||||
id="test_message_id",
|
||||
content=[],
|
||||
model="claude-3-opus-20240229",
|
||||
role="assistant",
|
||||
stop_reason=None,
|
||||
stop_sequence=None,
|
||||
type="message",
|
||||
usage=Usage(input_tokens=1720, output_tokens=2),
|
||||
),
|
||||
type="message_start",
|
||||
),
|
||||
RawContentBlockStartEvent(content_block=TextBlock(text="", type="text"), index=0, type="content_block_start"),
|
||||
RawContentBlockDeltaEvent(
|
||||
delta=TextDelta(text="<thinking>", type="text_delta"), index=0, type="content_block_delta"
|
||||
),
|
||||
TextEvent(type="text", text="<thinking>", snapshot="<thinking>"),
|
||||
RawContentBlockDeltaEvent(
|
||||
delta=TextDelta(text="</thinking>", type="text_delta"), index=0, type="content_block_delta"
|
||||
),
|
||||
TextEvent(type="text", text="</thinking>", snapshot="<thinking></thinking>"),
|
||||
ContentBlockStopEvent(
|
||||
index=0, type="content_block_stop", content_block=TextBlock(text="<thinking></thinking>", type="text")
|
||||
),
|
||||
RawContentBlockStartEvent(
|
||||
content_block=ToolUseBlock(id="test_tool_use_message_id", input={}, name="math-Add", type="tool_use"),
|
||||
index=1,
|
||||
type="content_block_start",
|
||||
),
|
||||
RawContentBlockDeltaEvent(
|
||||
delta=InputJSONDelta(partial_json='{"input": 3, "amount": 3}', type="input_json_delta"),
|
||||
index=1,
|
||||
type="content_block_delta",
|
||||
),
|
||||
InputJsonEvent(type="input_json", partial_json='{"input": 3, "amount": 3}', snapshot={"input": 3, "amount": 3}),
|
||||
ContentBlockStopEvent(
|
||||
index=1,
|
||||
type="content_block_stop",
|
||||
content_block=ToolUseBlock(
|
||||
id="test_tool_use_block_id", input={"input": 3, "amount": 3}, name="math-Add", type="tool_use"
|
||||
),
|
||||
),
|
||||
RawMessageDeltaEvent(
|
||||
delta=Delta(stop_reason="tool_use", stop_sequence=None),
|
||||
type="message_delta",
|
||||
usage=MessageDeltaUsage(output_tokens=159),
|
||||
),
|
||||
MessageStopEvent(
|
||||
type="message_stop",
|
||||
message=Message(
|
||||
id="test_message_id",
|
||||
content=[
|
||||
TextBlock(text="<thinking></thinking>", type="text"),
|
||||
ToolUseBlock(
|
||||
id="test_tool_use_block_id", input={"input": 3, "amount": 3}, name="math-Add", type="tool_use"
|
||||
),
|
||||
],
|
||||
model="claude-3-opus-20240229",
|
||||
role="assistant",
|
||||
stop_reason="tool_use",
|
||||
stop_sequence=None,
|
||||
type="message",
|
||||
usage=Usage(input_tokens=100, output_tokens=100),
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
async def async_generator():
|
||||
for event in stream_events:
|
||||
yield event
|
||||
|
||||
stream_mock = AsyncMock()
|
||||
stream_mock.__aenter__.return_value = async_generator()
|
||||
|
||||
return stream_mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tool_call_result_message() -> ChatMessageContent:
|
||||
return ChatMessageContent(
|
||||
inner_content=None,
|
||||
ai_model_id=None,
|
||||
metadata={},
|
||||
content_type="message",
|
||||
role=AuthorRole.TOOL,
|
||||
name=None,
|
||||
items=[
|
||||
FunctionResultContent(
|
||||
id="test_function_call_content",
|
||||
result=6,
|
||||
)
|
||||
],
|
||||
encoding=None,
|
||||
finish_reason=FinishReason.TOOL_CALLS,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_parallel_tool_call_result_message() -> ChatMessageContent:
|
||||
return ChatMessageContent(
|
||||
inner_content=None,
|
||||
ai_model_id=None,
|
||||
metadata={},
|
||||
content_type="message",
|
||||
role=AuthorRole.TOOL,
|
||||
name=None,
|
||||
items=[
|
||||
FunctionResultContent(
|
||||
id="test_function_call_content_1",
|
||||
result=6,
|
||||
),
|
||||
FunctionResultContent(
|
||||
id="test_function_call_content_2",
|
||||
result=3,
|
||||
),
|
||||
],
|
||||
encoding=None,
|
||||
finish_reason=FinishReason.TOOL_CALLS,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_streaming_chat_message_content() -> StreamingChatMessageContent:
|
||||
return StreamingChatMessageContent(
|
||||
choice_index=0,
|
||||
ai_model_id="claude-3-opus-20240229",
|
||||
metadata={},
|
||||
role=AuthorRole.ASSISTANT,
|
||||
name=None,
|
||||
items=[
|
||||
StreamingTextContent(
|
||||
inner_content=None,
|
||||
ai_model_id=None,
|
||||
metadata={},
|
||||
content_type="text",
|
||||
text="<thinking></thinking>",
|
||||
encoding=None,
|
||||
choice_index=0,
|
||||
),
|
||||
FunctionCallContent(
|
||||
inner_content=None,
|
||||
ai_model_id=None,
|
||||
metadata={},
|
||||
content_type=ContentTypes.FUNCTION_CALL_CONTENT,
|
||||
id="tool_id",
|
||||
index=0,
|
||||
name="math-Add",
|
||||
function_name="Add",
|
||||
plugin_name="math",
|
||||
arguments='{"input": 3, "amount": 3}',
|
||||
),
|
||||
],
|
||||
encoding=None,
|
||||
finish_reason=FinishReason.TOOL_CALLS,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_settings() -> AnthropicChatPromptExecutionSettings:
|
||||
return AnthropicChatPromptExecutionSettings()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_chat_message_response() -> Message:
|
||||
return Message(
|
||||
id="test_message_id",
|
||||
content=[TextBlock(text="Hello, how are you?", type="text")],
|
||||
model="claude-3-opus-20240229",
|
||||
role="assistant",
|
||||
stop_reason="end_turn",
|
||||
stop_sequence=None,
|
||||
type="message",
|
||||
usage=Usage(input_tokens=10, output_tokens=10),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_streaming_message_response() -> AsyncGenerator:
|
||||
raw_message_start_event = RawMessageStartEvent(
|
||||
message=Message(
|
||||
id="test_message_id",
|
||||
content=[],
|
||||
model="claude-3-opus-20240229",
|
||||
role="assistant",
|
||||
stop_reason=None,
|
||||
stop_sequence=None,
|
||||
type="message",
|
||||
usage=Usage(input_tokens=41, output_tokens=3),
|
||||
),
|
||||
type="message_start",
|
||||
)
|
||||
|
||||
raw_content_block_start_event = RawContentBlockStartEvent(
|
||||
content_block=TextBlock(text="", type="text"),
|
||||
index=0,
|
||||
type="content_block_start",
|
||||
)
|
||||
|
||||
raw_content_block_delta_event = RawContentBlockDeltaEvent(
|
||||
delta=TextDelta(text="Hello! It", type="text_delta"),
|
||||
index=0,
|
||||
type="content_block_delta",
|
||||
)
|
||||
|
||||
text_event = TextEvent(
|
||||
type="text",
|
||||
text="Hello! It",
|
||||
snapshot="Hello! It",
|
||||
)
|
||||
|
||||
content_block_stop_event = ContentBlockStopEvent(
|
||||
index=0,
|
||||
type="content_block_stop",
|
||||
content_block=TextBlock(text="Hello! It's nice to meet you.", type="text"),
|
||||
)
|
||||
|
||||
raw_message_delta_event = RawMessageDeltaEvent(
|
||||
delta=Delta(stop_reason="end_turn", stop_sequence=None),
|
||||
type="message_delta",
|
||||
usage=MessageDeltaUsage(output_tokens=84),
|
||||
)
|
||||
|
||||
message_stop_event = MessageStopEvent(
|
||||
type="message_stop",
|
||||
message=Message(
|
||||
id="test_message_stop_id",
|
||||
content=[TextBlock(text="Hello! It's nice to meet you.", type="text")],
|
||||
model="claude-3-opus-20240229",
|
||||
role="assistant",
|
||||
stop_reason="end_turn",
|
||||
stop_sequence=None,
|
||||
type="message",
|
||||
usage=Usage(input_tokens=41, output_tokens=84),
|
||||
),
|
||||
)
|
||||
|
||||
# Combine all mock events into a list
|
||||
stream_events = [
|
||||
raw_message_start_event,
|
||||
raw_content_block_start_event,
|
||||
raw_content_block_delta_event,
|
||||
text_event,
|
||||
content_block_stop_event,
|
||||
raw_message_delta_event,
|
||||
message_stop_event,
|
||||
]
|
||||
|
||||
async def async_generator():
|
||||
for event in stream_events:
|
||||
yield event
|
||||
|
||||
# Create an AsyncMock for the stream
|
||||
stream_mock = AsyncMock()
|
||||
stream_mock.__aenter__.return_value = async_generator()
|
||||
|
||||
return stream_mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_anthropic_client_completion(mock_chat_message_response: Message) -> AsyncAnthropic:
|
||||
client = MagicMock(spec=AsyncAnthropic)
|
||||
messages_mock = MagicMock()
|
||||
messages_mock.create = AsyncMock(return_value=mock_chat_message_response)
|
||||
client.messages = messages_mock
|
||||
return client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_anthropic_client_completion_stream(mock_streaming_message_response: AsyncGenerator) -> AsyncAnthropic:
|
||||
client = MagicMock(spec=AsyncAnthropic)
|
||||
messages_mock = MagicMock()
|
||||
messages_mock.stream.return_value = mock_streaming_message_response
|
||||
client.messages = messages_mock
|
||||
return client
|
||||
@@ -0,0 +1,549 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from anthropic import AsyncAnthropic
|
||||
from anthropic.types import Message
|
||||
|
||||
from semantic_kernel.connectors.ai.anthropic.prompt_execution_settings.anthropic_prompt_execution_settings import (
|
||||
AnthropicChatPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.anthropic.services.anthropic_chat_completion import AnthropicChatCompletion
|
||||
from semantic_kernel.connectors.ai.chat_completion_client_base import ChatCompletionClientBase
|
||||
from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior
|
||||
from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.open_ai_prompt_execution_settings import (
|
||||
OpenAIChatPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.contents.chat_history import ChatHistory
|
||||
from semantic_kernel.contents.chat_message_content import ChatMessageContent, FunctionCallContent, TextContent
|
||||
from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent
|
||||
from semantic_kernel.contents.utils.author_role import AuthorRole
|
||||
from semantic_kernel.exceptions import (
|
||||
ServiceInitializationError,
|
||||
ServiceInvalidExecutionSettingsError,
|
||||
ServiceResponseException,
|
||||
)
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInvalidRequestError
|
||||
from semantic_kernel.functions.kernel_arguments import KernelArguments
|
||||
from semantic_kernel.functions.kernel_function_decorator import kernel_function
|
||||
from semantic_kernel.kernel import Kernel
|
||||
|
||||
|
||||
async def test_complete_chat_contents(
|
||||
kernel: Kernel,
|
||||
mock_settings: AnthropicChatPromptExecutionSettings,
|
||||
mock_chat_message_response: Message,
|
||||
):
|
||||
client = MagicMock(spec=AsyncAnthropic)
|
||||
messages_mock = MagicMock()
|
||||
messages_mock.create = AsyncMock(return_value=mock_chat_message_response)
|
||||
client.messages = messages_mock
|
||||
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_user_message("test_user_message")
|
||||
|
||||
arguments = KernelArguments()
|
||||
chat_completion_base = AnthropicChatCompletion(
|
||||
ai_model_id="test_model_id", service_id="test", api_key="", async_client=client
|
||||
)
|
||||
|
||||
content: list[ChatMessageContent] = await chat_completion_base.get_chat_message_contents(
|
||||
chat_history=chat_history, settings=mock_settings, kernel=kernel, arguments=arguments
|
||||
)
|
||||
|
||||
assert len(content) > 0
|
||||
assert content[0].content != ""
|
||||
assert content[0].role == AuthorRole.ASSISTANT
|
||||
|
||||
|
||||
mock_message_text_content = ChatMessageContent(role=AuthorRole.ASSISTANT, items=[TextContent(text="test")])
|
||||
|
||||
mock_message_function_call = ChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT,
|
||||
items=[
|
||||
FunctionCallContent(
|
||||
name="test",
|
||||
arguments={"key": "test"},
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"function_choice_behavior,model_responses,expected_result",
|
||||
[
|
||||
pytest.param(
|
||||
FunctionChoiceBehavior.Auto(),
|
||||
[[mock_message_function_call], [mock_message_text_content]],
|
||||
TextContent,
|
||||
id="auto",
|
||||
),
|
||||
pytest.param(
|
||||
FunctionChoiceBehavior.Auto(auto_invoke=False),
|
||||
[[mock_message_function_call]],
|
||||
FunctionCallContent,
|
||||
id="auto_none_invoke",
|
||||
),
|
||||
pytest.param(
|
||||
FunctionChoiceBehavior.Required(auto_invoke=False),
|
||||
[[mock_message_function_call]],
|
||||
FunctionCallContent,
|
||||
id="required_none_invoke",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_complete_chat_contents_function_call_behavior_tool_call(
|
||||
kernel: Kernel,
|
||||
mock_settings: AnthropicChatPromptExecutionSettings,
|
||||
function_choice_behavior: FunctionChoiceBehavior,
|
||||
model_responses,
|
||||
expected_result,
|
||||
):
|
||||
kernel.add_function("test", kernel_function(lambda key: "test", name="test"))
|
||||
mock_settings.function_choice_behavior = function_choice_behavior
|
||||
|
||||
arguments = KernelArguments()
|
||||
chat_completion_base = AnthropicChatCompletion(ai_model_id="test_model_id", service_id="test", api_key="")
|
||||
|
||||
with (
|
||||
patch.object(chat_completion_base, "_inner_get_chat_message_contents", side_effect=model_responses),
|
||||
):
|
||||
response: list[ChatMessageContent] = await chat_completion_base.get_chat_message_contents(
|
||||
chat_history=ChatHistory(system_message="Test"), settings=mock_settings, kernel=kernel, arguments=arguments
|
||||
)
|
||||
|
||||
assert all(isinstance(content, expected_result) for content in response[0].items)
|
||||
|
||||
|
||||
async def test_complete_chat_contents_function_call_behavior_without_kernel(
|
||||
mock_settings: AnthropicChatPromptExecutionSettings,
|
||||
mock_anthropic_client_completion: AsyncAnthropic,
|
||||
):
|
||||
chat_history = MagicMock()
|
||||
chat_completion_base = AnthropicChatCompletion(
|
||||
ai_model_id="test_model_id", service_id="test", api_key="", async_client=mock_anthropic_client_completion
|
||||
)
|
||||
|
||||
mock_settings.function_choice_behavior = FunctionChoiceBehavior.Auto()
|
||||
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
await chat_completion_base.get_chat_message_contents(chat_history=chat_history, settings=mock_settings)
|
||||
|
||||
|
||||
async def test_complete_chat_stream_contents(
|
||||
kernel: Kernel,
|
||||
mock_settings: AnthropicChatPromptExecutionSettings,
|
||||
mock_streaming_message_response,
|
||||
):
|
||||
client = MagicMock(spec=AsyncAnthropic)
|
||||
messages_mock = MagicMock()
|
||||
messages_mock.stream.return_value = mock_streaming_message_response
|
||||
client.messages = messages_mock
|
||||
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_user_message("test_user_message")
|
||||
|
||||
arguments = KernelArguments()
|
||||
|
||||
chat_completion_base = AnthropicChatCompletion(
|
||||
ai_model_id="test_model_id",
|
||||
service_id="test",
|
||||
api_key="",
|
||||
async_client=client,
|
||||
)
|
||||
|
||||
async for content in chat_completion_base.get_streaming_chat_message_contents(
|
||||
chat_history, mock_settings, kernel=kernel, arguments=arguments
|
||||
):
|
||||
assert content is not None
|
||||
|
||||
|
||||
mock_message_function_call = StreamingChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT, items=[FunctionCallContent(name="test")], choice_index="0"
|
||||
)
|
||||
|
||||
mock_message_text_content = StreamingChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT, items=[TextContent(text="test")], choice_index="0"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"function_choice_behavior,model_responses,expected_result",
|
||||
[
|
||||
pytest.param(
|
||||
FunctionChoiceBehavior.Auto(),
|
||||
[[mock_message_function_call], [mock_message_text_content]],
|
||||
TextContent,
|
||||
id="auto",
|
||||
),
|
||||
pytest.param(
|
||||
FunctionChoiceBehavior.Auto(auto_invoke=False),
|
||||
[[mock_message_function_call]],
|
||||
FunctionCallContent,
|
||||
id="auto_none_invoke",
|
||||
),
|
||||
pytest.param(
|
||||
FunctionChoiceBehavior.Required(auto_invoke=False),
|
||||
[[mock_message_function_call]],
|
||||
FunctionCallContent,
|
||||
id="required_none_invoke",
|
||||
),
|
||||
pytest.param(FunctionChoiceBehavior.NoneInvoke(), [[mock_message_text_content]], TextContent, id="none"),
|
||||
],
|
||||
)
|
||||
async def test_complete_chat_contents_streaming_function_call_behavior_tool_call(
|
||||
kernel: Kernel,
|
||||
mock_settings: AnthropicChatPromptExecutionSettings,
|
||||
function_choice_behavior: FunctionChoiceBehavior,
|
||||
model_responses,
|
||||
expected_result,
|
||||
):
|
||||
mock_settings.function_choice_behavior = function_choice_behavior
|
||||
|
||||
# Mock sequence of model responses
|
||||
generator_mocks = []
|
||||
for mock_message in model_responses:
|
||||
generator_mock = MagicMock()
|
||||
generator_mock.__aiter__.return_value = [mock_message]
|
||||
generator_mocks.append(generator_mock)
|
||||
|
||||
arguments = KernelArguments()
|
||||
chat_completion_base = AnthropicChatCompletion(ai_model_id="test_model_id", service_id="test", api_key="")
|
||||
|
||||
with patch.object(chat_completion_base, "_inner_get_streaming_chat_message_contents", side_effect=generator_mocks):
|
||||
messages = []
|
||||
async for chunk in chat_completion_base.get_streaming_chat_message_contents(
|
||||
chat_history=ChatHistory(system_message="Test"), settings=mock_settings, kernel=kernel, arguments=arguments
|
||||
):
|
||||
messages.append(chunk)
|
||||
|
||||
response = messages[-1]
|
||||
assert all(isinstance(content, expected_result) for content in response[0].items)
|
||||
|
||||
|
||||
async def test_anthropic_sdk_exception(kernel: Kernel, mock_settings: AnthropicChatPromptExecutionSettings):
|
||||
client = MagicMock(spec=AsyncAnthropic)
|
||||
messages_mock = MagicMock()
|
||||
messages_mock.create.side_effect = Exception("Test Exception")
|
||||
client.messages = messages_mock
|
||||
|
||||
chat_history = MagicMock()
|
||||
arguments = KernelArguments()
|
||||
|
||||
chat_completion_base = AnthropicChatCompletion(
|
||||
ai_model_id="test_model_id", service_id="test", api_key="", async_client=client
|
||||
)
|
||||
|
||||
with pytest.raises(ServiceResponseException):
|
||||
await chat_completion_base.get_chat_message_contents(
|
||||
chat_history=chat_history, settings=mock_settings, kernel=kernel, arguments=arguments
|
||||
)
|
||||
|
||||
|
||||
async def test_anthropic_sdk_exception_streaming(kernel: Kernel, mock_settings: AnthropicChatPromptExecutionSettings):
|
||||
client = MagicMock(spec=AsyncAnthropic)
|
||||
messages_mock = MagicMock()
|
||||
messages_mock.stream.side_effect = Exception("Test Exception")
|
||||
client.messages = messages_mock
|
||||
|
||||
chat_history = MagicMock()
|
||||
arguments = KernelArguments()
|
||||
|
||||
chat_completion_base = AnthropicChatCompletion(
|
||||
ai_model_id="test_model_id", service_id="test", api_key="", async_client=client
|
||||
)
|
||||
|
||||
with pytest.raises(ServiceResponseException):
|
||||
async for content in chat_completion_base.get_streaming_chat_message_contents(
|
||||
chat_history, mock_settings, kernel=kernel, arguments=arguments
|
||||
):
|
||||
assert content is not None
|
||||
|
||||
|
||||
def test_anthropic_chat_completion_init(anthropic_unit_test_env) -> None:
|
||||
# Test successful initialization
|
||||
anthropic_chat_completion = AnthropicChatCompletion()
|
||||
|
||||
assert anthropic_chat_completion.ai_model_id == anthropic_unit_test_env["ANTHROPIC_CHAT_MODEL_ID"]
|
||||
assert isinstance(anthropic_chat_completion, ChatCompletionClientBase)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["ANTHROPIC_API_KEY"]], indirect=True)
|
||||
def test_anthropic_chat_completion_init_with_empty_api_key(anthropic_unit_test_env) -> None:
|
||||
ai_model_id = "test_model_id"
|
||||
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AnthropicChatCompletion(
|
||||
ai_model_id=ai_model_id,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["ANTHROPIC_CHAT_MODEL_ID"]], indirect=True)
|
||||
def test_anthropic_chat_completion_init_with_empty_model_id(anthropic_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AnthropicChatCompletion(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
def test_prompt_execution_settings_class(anthropic_unit_test_env):
|
||||
anthropic_chat_completion = AnthropicChatCompletion()
|
||||
prompt_execution_settings = anthropic_chat_completion.get_prompt_execution_settings_class()
|
||||
assert prompt_execution_settings == AnthropicChatPromptExecutionSettings
|
||||
|
||||
|
||||
async def test_with_different_execution_settings(kernel: Kernel, mock_anthropic_client_completion: MagicMock):
|
||||
chat_history = MagicMock()
|
||||
settings = OpenAIChatPromptExecutionSettings(temperature=0.2)
|
||||
arguments = KernelArguments()
|
||||
chat_completion_base = AnthropicChatCompletion(
|
||||
ai_model_id="test_model_id", service_id="test", api_key="", async_client=mock_anthropic_client_completion
|
||||
)
|
||||
|
||||
await chat_completion_base.get_chat_message_contents(
|
||||
chat_history=chat_history, settings=settings, kernel=kernel, arguments=arguments
|
||||
)
|
||||
|
||||
assert mock_anthropic_client_completion.messages.create.call_args.kwargs["temperature"] == 0.2
|
||||
|
||||
|
||||
async def test_with_different_execution_settings_stream(
|
||||
kernel: Kernel, mock_anthropic_client_completion_stream: MagicMock
|
||||
):
|
||||
chat_history = MagicMock()
|
||||
settings = OpenAIChatPromptExecutionSettings(temperature=0.2, seed=2)
|
||||
arguments = KernelArguments()
|
||||
chat_completion_base = AnthropicChatCompletion(
|
||||
ai_model_id="test_model_id",
|
||||
service_id="test",
|
||||
api_key="",
|
||||
async_client=mock_anthropic_client_completion_stream,
|
||||
)
|
||||
|
||||
async for chunk in chat_completion_base.get_streaming_chat_message_contents(
|
||||
chat_history, settings, kernel=kernel, arguments=arguments
|
||||
):
|
||||
assert chunk is not None
|
||||
assert mock_anthropic_client_completion_stream.messages.stream.call_args.kwargs["temperature"] == 0.2
|
||||
|
||||
|
||||
async def test_prepare_chat_history_for_request_with_system_message(mock_anthropic_client_completion_stream: MagicMock):
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_system_message("System message")
|
||||
chat_history.add_user_message("User message")
|
||||
chat_history.add_assistant_message("Assistant message")
|
||||
chat_history.add_system_message("Another system message")
|
||||
|
||||
chat_completion_base = AnthropicChatCompletion(
|
||||
ai_model_id="test_model_id",
|
||||
service_id="test",
|
||||
api_key="",
|
||||
async_client=mock_anthropic_client_completion_stream,
|
||||
)
|
||||
|
||||
remaining_messages, system_message_content = chat_completion_base._prepare_chat_history_for_request(
|
||||
chat_history, role_key="role", content_key="content"
|
||||
)
|
||||
|
||||
assert system_message_content == "System message"
|
||||
assert remaining_messages == [
|
||||
{"role": AuthorRole.USER, "content": "User message"},
|
||||
{"role": AuthorRole.ASSISTANT, "content": [{"type": "text", "text": "Assistant message"}]},
|
||||
]
|
||||
assert not any(msg["role"] == AuthorRole.SYSTEM for msg in remaining_messages)
|
||||
|
||||
|
||||
async def test_prepare_chat_history_for_request_with_tool_message(
|
||||
mock_anthropic_client_completion_stream: MagicMock,
|
||||
mock_tool_calls_message: ChatMessageContent,
|
||||
mock_tool_call_result_message: ChatMessageContent,
|
||||
):
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_user_message("What is 3+3?")
|
||||
chat_history.add_message(mock_tool_calls_message)
|
||||
chat_history.add_message(mock_tool_call_result_message)
|
||||
|
||||
chat_completion_client = AnthropicChatCompletion(
|
||||
ai_model_id="test_model_id",
|
||||
service_id="test",
|
||||
api_key="",
|
||||
async_client=mock_anthropic_client_completion_stream,
|
||||
)
|
||||
|
||||
remaining_messages, system_message_content = chat_completion_client._prepare_chat_history_for_request(
|
||||
chat_history, role_key="role", content_key="content"
|
||||
)
|
||||
|
||||
assert system_message_content is None
|
||||
assert remaining_messages == [
|
||||
{"role": AuthorRole.USER, "content": "What is 3+3?"},
|
||||
{
|
||||
"role": AuthorRole.ASSISTANT,
|
||||
"content": [
|
||||
{"type": "text", "text": mock_tool_calls_message.items[0].text},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": mock_tool_calls_message.items[1].id,
|
||||
"name": mock_tool_calls_message.items[1].name,
|
||||
"input": mock_tool_calls_message.items[1].arguments,
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": AuthorRole.USER,
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": mock_tool_call_result_message.items[0].id,
|
||||
"content": str(mock_tool_call_result_message.items[0].result),
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
async def test_prepare_chat_history_for_request_with_parallel_tool_message(
|
||||
mock_anthropic_client_completion_stream: MagicMock,
|
||||
mock_parallel_tool_calls_message: ChatMessageContent,
|
||||
mock_parallel_tool_call_result_message: ChatMessageContent,
|
||||
):
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_user_message("What is 3+3?")
|
||||
chat_history.add_message(mock_parallel_tool_calls_message)
|
||||
chat_history.add_message(mock_parallel_tool_call_result_message)
|
||||
|
||||
chat_completion_client = AnthropicChatCompletion(
|
||||
ai_model_id="test_model_id",
|
||||
service_id="test",
|
||||
api_key="",
|
||||
async_client=mock_anthropic_client_completion_stream,
|
||||
)
|
||||
|
||||
remaining_messages, system_message_content = chat_completion_client._prepare_chat_history_for_request(
|
||||
chat_history, role_key="role", content_key="content"
|
||||
)
|
||||
|
||||
assert system_message_content is None
|
||||
assert remaining_messages == [
|
||||
{"role": AuthorRole.USER, "content": "What is 3+3?"},
|
||||
{
|
||||
"role": AuthorRole.ASSISTANT,
|
||||
"content": [
|
||||
{"type": "text", "text": mock_parallel_tool_calls_message.items[0].text},
|
||||
*[
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": function_call_content.id,
|
||||
"name": function_call_content.name,
|
||||
"input": function_call_content.arguments,
|
||||
}
|
||||
for function_call_content in mock_parallel_tool_calls_message.items[1:]
|
||||
],
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": AuthorRole.USER,
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": function_result_content.id,
|
||||
"content": str(function_result_content.result),
|
||||
}
|
||||
for function_result_content in mock_parallel_tool_call_result_message.items
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
async def test_prepare_chat_history_for_request_with_tool_message_right_after_user_message(
|
||||
mock_anthropic_client_completion_stream: MagicMock,
|
||||
mock_tool_call_result_message: ChatMessageContent,
|
||||
):
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_user_message("What is 3+3?")
|
||||
chat_history.add_message(mock_tool_call_result_message)
|
||||
|
||||
chat_completion_client = AnthropicChatCompletion(
|
||||
ai_model_id="test_model_id",
|
||||
service_id="test",
|
||||
api_key="",
|
||||
async_client=mock_anthropic_client_completion_stream,
|
||||
)
|
||||
|
||||
with pytest.raises(ServiceInvalidRequestError, match="Tool message found after a user or system message."):
|
||||
chat_completion_client._prepare_chat_history_for_request(chat_history, role_key="role", content_key="content")
|
||||
|
||||
|
||||
async def test_prepare_chat_history_for_request_with_tool_message_as_the_first_message(
|
||||
mock_anthropic_client_completion_stream: MagicMock,
|
||||
mock_tool_call_result_message: ChatMessageContent,
|
||||
):
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_message(mock_tool_call_result_message)
|
||||
|
||||
chat_completion_client = AnthropicChatCompletion(
|
||||
ai_model_id="test_model_id",
|
||||
service_id="test",
|
||||
api_key="",
|
||||
async_client=mock_anthropic_client_completion_stream,
|
||||
)
|
||||
|
||||
with pytest.raises(ServiceInvalidRequestError, match="Tool message found without a preceding message."):
|
||||
chat_completion_client._prepare_chat_history_for_request(chat_history, role_key="role", content_key="content")
|
||||
|
||||
|
||||
async def test_send_chat_stream_request_tool_calls(
|
||||
mock_streaming_tool_calls_message: MagicMock,
|
||||
mock_streaming_chat_message_content: StreamingChatMessageContent,
|
||||
):
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_user_message("What is 3+3?")
|
||||
chat_history.add_message(mock_streaming_chat_message_content)
|
||||
|
||||
settings = AnthropicChatPromptExecutionSettings(
|
||||
temperature=0.2,
|
||||
max_tokens=100,
|
||||
top_p=1.0,
|
||||
frequency_penalty=0.0,
|
||||
presence_penalty=0.0,
|
||||
chat_history=chat_history,
|
||||
)
|
||||
|
||||
client = MagicMock(spec=AsyncAnthropic)
|
||||
messages_mock = MagicMock()
|
||||
messages_mock.stream.return_value = mock_streaming_tool_calls_message
|
||||
client.messages = messages_mock
|
||||
|
||||
chat_completion = AnthropicChatCompletion(
|
||||
ai_model_id="test_model_id",
|
||||
service_id="test",
|
||||
api_key="",
|
||||
async_client=client,
|
||||
)
|
||||
|
||||
response = chat_completion._send_chat_stream_request(settings)
|
||||
async for message in response:
|
||||
assert message is not None
|
||||
|
||||
|
||||
def test_client_base_url(mock_anthropic_client_completion: MagicMock):
|
||||
chat_completion_base = AnthropicChatCompletion(
|
||||
ai_model_id="test_model_id", service_id="test", api_key="", async_client=mock_anthropic_client_completion
|
||||
)
|
||||
|
||||
assert chat_completion_base.service_url() is not None
|
||||
|
||||
|
||||
def test_chat_completion_reset_settings(
|
||||
mock_anthropic_client_completion: MagicMock,
|
||||
):
|
||||
chat_completion = AnthropicChatCompletion(
|
||||
ai_model_id="test_model_id", service_id="test", api_key="", async_client=mock_anthropic_client_completion
|
||||
)
|
||||
|
||||
settings = AnthropicChatPromptExecutionSettings(tools=[{"name": "test"}], tool_choice={"type": "any"})
|
||||
chat_completion._reset_function_choice_settings(settings)
|
||||
|
||||
assert settings.tools is None
|
||||
assert settings.tool_choice is None
|
||||
@@ -0,0 +1,129 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.ai.anthropic.prompt_execution_settings.anthropic_prompt_execution_settings import (
|
||||
AnthropicChatPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
from semantic_kernel.exceptions import ServiceInvalidExecutionSettingsError
|
||||
|
||||
|
||||
def test_default_anthropic_chat_prompt_execution_settings():
|
||||
settings = AnthropicChatPromptExecutionSettings()
|
||||
assert settings.temperature is None
|
||||
assert settings.top_p is None
|
||||
assert settings.max_tokens == 1024
|
||||
assert settings.messages is None
|
||||
|
||||
|
||||
def test_custom_anthropic_chat_prompt_execution_settings():
|
||||
settings = AnthropicChatPromptExecutionSettings(
|
||||
temperature=0.5,
|
||||
top_p=0.5,
|
||||
max_tokens=128,
|
||||
messages=[{"role": "system", "content": "Hello"}],
|
||||
)
|
||||
assert settings.temperature == 0.5
|
||||
assert settings.top_p == 0.5
|
||||
assert settings.max_tokens == 128
|
||||
assert settings.messages == [{"role": "system", "content": "Hello"}]
|
||||
|
||||
|
||||
def test_anthropic_chat_prompt_execution_settings_from_default_completion_config():
|
||||
settings = PromptExecutionSettings(service_id="test_service")
|
||||
chat_settings = AnthropicChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
assert chat_settings.service_id == "test_service"
|
||||
assert chat_settings.temperature is None
|
||||
assert chat_settings.top_p is None
|
||||
assert chat_settings.max_tokens == 1024
|
||||
|
||||
|
||||
def test_anthropic_chat_prompt_execution_settings_from_openai_prompt_execution_settings():
|
||||
chat_settings = AnthropicChatPromptExecutionSettings(service_id="test_service", temperature=1.0)
|
||||
new_settings = AnthropicChatPromptExecutionSettings(service_id="test_2", temperature=0.0)
|
||||
chat_settings.update_from_prompt_execution_settings(new_settings)
|
||||
assert chat_settings.service_id == "test_2"
|
||||
assert chat_settings.temperature == 0.0
|
||||
|
||||
|
||||
def test_anthropic_chat_prompt_execution_settings_from_custom_completion_config():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"max_tokens": 128,
|
||||
"messages": [{"role": "system", "content": "Hello"}],
|
||||
},
|
||||
)
|
||||
chat_settings = AnthropicChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
assert chat_settings.temperature == 0.5
|
||||
assert chat_settings.top_p == 0.5
|
||||
assert chat_settings.max_tokens == 128
|
||||
|
||||
|
||||
def test_openai_chat_prompt_execution_settings_from_custom_completion_config_with_none():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"max_tokens": 128,
|
||||
"messages": [{"role": "system", "content": "Hello"}],
|
||||
},
|
||||
)
|
||||
chat_settings = AnthropicChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
assert chat_settings.temperature == 0.5
|
||||
assert chat_settings.top_p == 0.5
|
||||
assert chat_settings.max_tokens == 128
|
||||
|
||||
|
||||
def test_openai_chat_prompt_execution_settings_from_custom_completion_config_with_functions():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"max_tokens": 128,
|
||||
"tools": [{}],
|
||||
"messages": [{"role": "system", "content": "Hello"}],
|
||||
},
|
||||
)
|
||||
chat_settings = AnthropicChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
assert chat_settings.temperature == 0.5
|
||||
assert chat_settings.top_p == 0.5
|
||||
assert chat_settings.max_tokens == 128
|
||||
|
||||
|
||||
def test_create_options():
|
||||
settings = AnthropicChatPromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"max_tokens": 128,
|
||||
"tools": [{}],
|
||||
"messages": [{"role": "system", "content": "Hello"}],
|
||||
},
|
||||
)
|
||||
options = settings.prepare_settings_dict()
|
||||
assert options["temperature"] == 0.5
|
||||
assert options["top_p"] == 0.5
|
||||
assert options["max_tokens"] == 128
|
||||
|
||||
|
||||
def test_tool_choice_none():
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError, match="Tool choice 'none' is not supported by Anthropic."):
|
||||
AnthropicChatPromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"max_tokens": 128,
|
||||
"tool_choice": {"type": "none"},
|
||||
"messages": [{"role": "system", "content": "Hello"}],
|
||||
},
|
||||
function_choice_behavior=FunctionChoiceBehavior.NoneInvoke(),
|
||||
)
|
||||
@@ -0,0 +1,278 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import datetime
|
||||
from collections.abc import AsyncGenerator, AsyncIterator
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from azure.ai.inference.aio import ChatCompletionsClient, EmbeddingsClient
|
||||
from azure.ai.inference.models import (
|
||||
ChatChoice,
|
||||
ChatCompletions,
|
||||
ChatCompletionsToolCall,
|
||||
ChatResponseMessage,
|
||||
CompletionsUsage,
|
||||
FunctionCall,
|
||||
StreamingChatChoiceUpdate,
|
||||
StreamingChatCompletionsUpdate,
|
||||
StreamingChatResponseToolCallUpdate,
|
||||
)
|
||||
from azure.core.credentials import AzureKeyCredential
|
||||
|
||||
from semantic_kernel.connectors.ai.azure_ai_inference import (
|
||||
AzureAIInferenceChatCompletion,
|
||||
AzureAIInferenceTextEmbedding,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def model_id() -> str:
|
||||
return "test_model_id"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def service_id() -> str:
|
||||
return "test_service_id"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def azure_ai_inference_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for Azure AI Inference Unit Tests."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"AZURE_AI_INFERENCE_API_KEY": "test-api-key",
|
||||
"AZURE_AI_INFERENCE_ENDPOINT": "https://test-endpoint.com",
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def model_diagnostics_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for Azure AI Inference Unit Tests."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS": "true",
|
||||
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE": "true",
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def disabled_model_diagnostics_test_env(monkeypatch):
|
||||
"""Fixture to disable diagnostics for tests that use mocking.
|
||||
|
||||
This is needed because AIInferenceInstrumentor's instrument/uninstrument
|
||||
cycle interferes with class-level mocking of ChatCompletionsClient.complete.
|
||||
"""
|
||||
monkeypatch.setenv("SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS", "false")
|
||||
monkeypatch.setenv("SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE", "false")
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def azure_ai_inference_client(azure_ai_inference_unit_test_env, request) -> ChatCompletionsClient | EmbeddingsClient:
|
||||
"""Fixture to create Azure AI Inference client for unit tests."""
|
||||
endpoint = azure_ai_inference_unit_test_env["AZURE_AI_INFERENCE_ENDPOINT"]
|
||||
api_key = azure_ai_inference_unit_test_env["AZURE_AI_INFERENCE_API_KEY"]
|
||||
credential = AzureKeyCredential(api_key)
|
||||
|
||||
if request.param == AzureAIInferenceChatCompletion.__name__:
|
||||
return ChatCompletionsClient(endpoint=endpoint, credential=credential)
|
||||
if request.param == AzureAIInferenceTextEmbedding.__name__:
|
||||
return EmbeddingsClient(endpoint=endpoint, credential=credential)
|
||||
|
||||
raise ValueError(f"Service {request.param} not supported.")
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def azure_ai_inference_service(azure_ai_inference_unit_test_env, model_id, request):
|
||||
"""Fixture to create Azure AI Inference service for unit tests.
|
||||
|
||||
This is required because the Azure AI Inference services require a client to be created,
|
||||
and the client will be talking to the endpoint at creation time.
|
||||
"""
|
||||
|
||||
endpoint = azure_ai_inference_unit_test_env["AZURE_AI_INFERENCE_ENDPOINT"]
|
||||
api_key = azure_ai_inference_unit_test_env["AZURE_AI_INFERENCE_API_KEY"]
|
||||
|
||||
if request.param == AzureAIInferenceChatCompletion.__name__:
|
||||
return AzureAIInferenceChatCompletion(model_id, api_key=api_key, endpoint=endpoint)
|
||||
if request.param == AzureAIInferenceTextEmbedding.__name__:
|
||||
return AzureAIInferenceTextEmbedding(model_id, api_key=api_key, endpoint=endpoint)
|
||||
|
||||
raise ValueError(f"Service {request.param} not supported.")
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_azure_ai_inference_chat_completion_response(model_id) -> ChatCompletions:
|
||||
return ChatCompletions(
|
||||
id="test_id",
|
||||
created=datetime.datetime.now(),
|
||||
model=model_id,
|
||||
usage=CompletionsUsage(
|
||||
completion_tokens=0,
|
||||
prompt_tokens=0,
|
||||
total_tokens=0,
|
||||
),
|
||||
choices=[
|
||||
ChatChoice(
|
||||
index=0,
|
||||
finish_reason="stop",
|
||||
message=ChatResponseMessage(
|
||||
role="assistant",
|
||||
content="Hello",
|
||||
),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_azure_ai_inference_chat_completion_response_with_tool_call(model_id) -> ChatCompletions:
|
||||
return ChatCompletions(
|
||||
id="test_id",
|
||||
created=datetime.datetime.now(),
|
||||
model=model_id,
|
||||
usage=CompletionsUsage(
|
||||
completion_tokens=0,
|
||||
prompt_tokens=0,
|
||||
total_tokens=0,
|
||||
),
|
||||
choices=[
|
||||
ChatChoice(
|
||||
index=0,
|
||||
finish_reason="tool_calls",
|
||||
message=ChatResponseMessage(
|
||||
role="assistant",
|
||||
tool_calls=[
|
||||
ChatCompletionsToolCall(
|
||||
id="test_id",
|
||||
function=FunctionCall(
|
||||
name="getLightStatus",
|
||||
arguments='{"arg1": "test_value"}',
|
||||
),
|
||||
),
|
||||
],
|
||||
),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_azure_ai_inference_streaming_chat_completion_response(model_id) -> AsyncIterator:
|
||||
streaming_chat_response = MagicMock(spec=AsyncGenerator)
|
||||
streaming_chat_response.__aiter__.return_value = [
|
||||
StreamingChatCompletionsUpdate(
|
||||
id="test_id",
|
||||
created=datetime.datetime.now(),
|
||||
model=model_id,
|
||||
usage=CompletionsUsage(
|
||||
completion_tokens=0,
|
||||
prompt_tokens=0,
|
||||
total_tokens=0,
|
||||
),
|
||||
choices=[
|
||||
# Empty choice
|
||||
],
|
||||
),
|
||||
StreamingChatCompletionsUpdate(
|
||||
id="test_id",
|
||||
created=datetime.datetime.now(),
|
||||
model=model_id,
|
||||
usage=CompletionsUsage(
|
||||
completion_tokens=0,
|
||||
prompt_tokens=0,
|
||||
total_tokens=0,
|
||||
),
|
||||
choices=[
|
||||
StreamingChatChoiceUpdate(
|
||||
index=0,
|
||||
finish_reason="stop",
|
||||
delta=ChatResponseMessage(
|
||||
role="assistant",
|
||||
content="Hello",
|
||||
),
|
||||
)
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
return streaming_chat_response
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_azure_ai_inference_streaming_chat_completion_response_with_tool_call(model_id) -> AsyncIterator:
|
||||
streaming_chat_response = MagicMock(spec=AsyncGenerator)
|
||||
streaming_chat_response.__aiter__.return_value = [
|
||||
StreamingChatCompletionsUpdate(
|
||||
id="test_id",
|
||||
created=datetime.datetime.now(),
|
||||
model=model_id,
|
||||
usage=CompletionsUsage(
|
||||
completion_tokens=0,
|
||||
prompt_tokens=0,
|
||||
total_tokens=0,
|
||||
),
|
||||
choices=[
|
||||
# Empty choice
|
||||
],
|
||||
),
|
||||
StreamingChatCompletionsUpdate(
|
||||
id="test_id",
|
||||
created=datetime.datetime.now(),
|
||||
model=model_id,
|
||||
usage=CompletionsUsage(
|
||||
completion_tokens=0,
|
||||
prompt_tokens=0,
|
||||
total_tokens=0,
|
||||
),
|
||||
choices=[
|
||||
StreamingChatChoiceUpdate(
|
||||
index=0,
|
||||
finish_reason="tool_calls",
|
||||
delta=ChatResponseMessage(
|
||||
role="assistant",
|
||||
tool_calls=[
|
||||
StreamingChatResponseToolCallUpdate(
|
||||
id="test_id",
|
||||
function=FunctionCall(
|
||||
name="getLightStatus",
|
||||
arguments='{"arg1": "test_value"}',
|
||||
),
|
||||
),
|
||||
],
|
||||
),
|
||||
)
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
return streaming_chat_response
|
||||
+710
@@ -0,0 +1,710 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from azure.ai.inference.aio import ChatCompletionsClient
|
||||
from azure.ai.inference.models import JsonSchemaFormat, UserMessage
|
||||
from azure.core.credentials import AzureKeyCredential
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from semantic_kernel.connectors.ai.azure_ai_inference import (
|
||||
AzureAIInferenceChatCompletion,
|
||||
AzureAIInferenceChatPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.azure_ai_inference.azure_ai_inference_settings import AzureAIInferenceSettings
|
||||
from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior
|
||||
from semantic_kernel.contents.chat_history import ChatHistory
|
||||
from semantic_kernel.contents.utils.finish_reason import FinishReason
|
||||
from semantic_kernel.exceptions.service_exceptions import (
|
||||
ServiceInitializationError,
|
||||
ServiceInvalidExecutionSettingsError,
|
||||
)
|
||||
from semantic_kernel.functions.kernel_arguments import KernelArguments
|
||||
from semantic_kernel.kernel import Kernel
|
||||
from semantic_kernel.utils.telemetry.user_agent import SEMANTIC_KERNEL_USER_AGENT
|
||||
|
||||
|
||||
# region init
|
||||
def test_azure_ai_inference_chat_completion_init(azure_ai_inference_unit_test_env, model_id) -> None:
|
||||
"""Test initialization of AzureAIInferenceChatCompletion"""
|
||||
azure_ai_inference = AzureAIInferenceChatCompletion(model_id, instruction_role="developer")
|
||||
|
||||
assert azure_ai_inference.ai_model_id == model_id
|
||||
assert azure_ai_inference.service_id == model_id
|
||||
assert isinstance(azure_ai_inference.client, ChatCompletionsClient)
|
||||
assert azure_ai_inference.instruction_role == "developer"
|
||||
|
||||
|
||||
@patch("azure.ai.inference.aio.ChatCompletionsClient.__init__", return_value=None)
|
||||
def test_azure_ai_inference_chat_completion_client_init(
|
||||
mock_client, azure_ai_inference_unit_test_env, model_id
|
||||
) -> None:
|
||||
"""Test initialization of the Azure AI Inference client"""
|
||||
endpoint = azure_ai_inference_unit_test_env["AZURE_AI_INFERENCE_ENDPOINT"]
|
||||
api_key = azure_ai_inference_unit_test_env["AZURE_AI_INFERENCE_API_KEY"]
|
||||
settings = AzureAIInferenceSettings(endpoint=endpoint, api_key=api_key)
|
||||
|
||||
_ = AzureAIInferenceChatCompletion(model_id)
|
||||
|
||||
assert mock_client.call_count == 1
|
||||
assert isinstance(mock_client.call_args.kwargs["endpoint"], str)
|
||||
assert mock_client.call_args.kwargs["endpoint"] == str(settings.endpoint)
|
||||
assert isinstance(mock_client.call_args.kwargs["credential"], AzureKeyCredential)
|
||||
assert mock_client.call_args.kwargs["credential"].key == settings.api_key.get_secret_value()
|
||||
assert mock_client.call_args.kwargs["user_agent"] == SEMANTIC_KERNEL_USER_AGENT
|
||||
|
||||
|
||||
def test_azure_ai_inference_chat_completion_init_with_service_id(
|
||||
azure_ai_inference_unit_test_env, model_id, service_id
|
||||
) -> None:
|
||||
"""Test initialization of AzureAIInferenceChatCompletion with service_id"""
|
||||
azure_ai_inference = AzureAIInferenceChatCompletion(model_id, service_id=service_id)
|
||||
|
||||
assert azure_ai_inference.ai_model_id == model_id
|
||||
assert azure_ai_inference.service_id == service_id
|
||||
assert isinstance(azure_ai_inference.client, ChatCompletionsClient)
|
||||
|
||||
|
||||
def test_azure_ai_inference_chat_completion_init_with_api_version(azure_ai_inference_unit_test_env, model_id) -> None:
|
||||
"""Test initialization of AzureAIInferenceChatCompletion with api_version"""
|
||||
azure_ai_inference = AzureAIInferenceChatCompletion(model_id, api_version="2024-02-15-test")
|
||||
|
||||
assert azure_ai_inference.ai_model_id == model_id
|
||||
assert isinstance(azure_ai_inference.client, ChatCompletionsClient)
|
||||
assert azure_ai_inference.client._config.api_version == "2024-02-15-test"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_client",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
def test_azure_ai_inference_chat_completion_init_with_custom_client(azure_ai_inference_client, model_id) -> None:
|
||||
"""Test initialization of AzureAIInferenceChatCompletion with custom client"""
|
||||
client = azure_ai_inference_client
|
||||
azure_ai_inference = AzureAIInferenceChatCompletion(model_id, client=client)
|
||||
|
||||
assert azure_ai_inference.ai_model_id == model_id
|
||||
assert azure_ai_inference.service_id == model_id
|
||||
assert azure_ai_inference.client == client
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_AI_INFERENCE_ENDPOINT"]], indirect=True)
|
||||
def test_azure_ai_inference_chat_completion_init_with_empty_endpoint(
|
||||
azure_ai_inference_unit_test_env, model_id
|
||||
) -> None:
|
||||
"""Test initialization of AzureAIInferenceChatCompletion with empty endpoint"""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureAIInferenceChatCompletion(model_id, env_file_path="fake_path")
|
||||
|
||||
|
||||
def test_prompt_execution_settings_class(azure_ai_inference_unit_test_env, model_id) -> None:
|
||||
azure_ai_inference = AzureAIInferenceChatCompletion(model_id)
|
||||
assert azure_ai_inference.get_prompt_execution_settings_class() == AzureAIInferenceChatPromptExecutionSettings
|
||||
|
||||
|
||||
# endregion init
|
||||
|
||||
|
||||
# region chat completion
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
async def test_azure_ai_inference_chat_completion(
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
model_id,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test completion of AzureAIInferenceChatCompletion"""
|
||||
user_message_content: str = "Hello"
|
||||
chat_history.add_user_message(user_message_content)
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings()
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_chat_completion_response
|
||||
|
||||
responses = await azure_ai_inference_service.get_chat_message_contents(chat_history=chat_history, settings=settings)
|
||||
|
||||
mock_complete.assert_awaited_once_with(
|
||||
messages=[UserMessage(content=user_message_content)],
|
||||
model=model_id,
|
||||
model_extras=None,
|
||||
**settings.prepare_settings_dict(),
|
||||
)
|
||||
assert len(responses) == 1
|
||||
assert responses[0].role == "assistant"
|
||||
assert responses[0].content == "Hello"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
async def test_azure_ai_inference_chat_completion_with_standard_parameters(
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
model_id,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test completion of AzureAIInferenceChatCompletion with standard OpenAI parameters"""
|
||||
user_message_content: str = "Hello"
|
||||
chat_history.add_user_message(user_message_content)
|
||||
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings(
|
||||
frequency_penalty=0.5,
|
||||
max_tokens=100,
|
||||
presence_penalty=0.5,
|
||||
seed=123,
|
||||
stop="stop",
|
||||
temperature=0.5,
|
||||
top_p=0.5,
|
||||
)
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_chat_completion_response
|
||||
|
||||
responses = await azure_ai_inference_service.get_chat_message_contents(chat_history=chat_history, settings=settings)
|
||||
|
||||
mock_complete.assert_awaited_once_with(
|
||||
messages=[UserMessage(content=user_message_content)],
|
||||
model=model_id,
|
||||
model_extras=None,
|
||||
frequency_penalty=settings.frequency_penalty,
|
||||
max_tokens=settings.max_tokens,
|
||||
presence_penalty=settings.presence_penalty,
|
||||
seed=settings.seed,
|
||||
stop=settings.stop,
|
||||
temperature=settings.temperature,
|
||||
top_p=settings.top_p,
|
||||
)
|
||||
assert len(responses) == 1
|
||||
assert responses[0].role == "assistant"
|
||||
assert responses[0].content == "Hello"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
async def test_azure_ai_inference_chat_completion_with_extra_parameters(
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
model_id,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test completion of AzureAIInferenceChatCompletion with extra parameters"""
|
||||
user_message_content: str = "Hello"
|
||||
chat_history.add_user_message(user_message_content)
|
||||
extra_parameters = {"test_key": "test_value"}
|
||||
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings(extra_parameters=extra_parameters)
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_chat_completion_response
|
||||
|
||||
responses = await azure_ai_inference_service.get_chat_message_contents(chat_history=chat_history, settings=settings)
|
||||
|
||||
mock_complete.assert_awaited_once_with(
|
||||
messages=[UserMessage(content=user_message_content)],
|
||||
model=model_id,
|
||||
model_extras=settings.extra_parameters,
|
||||
)
|
||||
assert len(responses) == 1
|
||||
assert responses[0].role == "assistant"
|
||||
assert responses[0].content == "Hello"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_azure_ai_inference_chat_completion_with_function_choice_behavior_fail_verification(
|
||||
azure_ai_inference_service,
|
||||
kernel,
|
||||
chat_history: ChatHistory,
|
||||
) -> None:
|
||||
"""Test completion of AzureAIInferenceChatCompletion with function choice behavior expect verification failure"""
|
||||
|
||||
# Missing kernel
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
await azure_ai_inference_service.get_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=settings,
|
||||
arguments=KernelArguments(),
|
||||
)
|
||||
|
||||
# More than 1 responses
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
extra_parameters={"n": 2},
|
||||
)
|
||||
await azure_ai_inference_service.get_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=settings,
|
||||
kernel=kernel,
|
||||
arguments=KernelArguments(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
async def test_azure_ai_inference_chat_completion_with_function_choice_behavior(
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
kernel: Kernel,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_chat_completion_response_with_tool_call,
|
||||
mock_azure_ai_inference_chat_completion_response,
|
||||
decorated_native_function,
|
||||
disabled_model_diagnostics_test_env,
|
||||
) -> None:
|
||||
"""Test completion of AzureAIInferenceChatCompletion with function choice behavior"""
|
||||
user_message_content: str = "Hello"
|
||||
chat_history.add_user_message(user_message_content)
|
||||
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
settings.function_choice_behavior.maximum_auto_invoke_attempts = 1
|
||||
|
||||
# First call returns tool call, second call returns final response
|
||||
mock_complete.side_effect = [
|
||||
mock_azure_ai_inference_chat_completion_response_with_tool_call,
|
||||
mock_azure_ai_inference_chat_completion_response,
|
||||
]
|
||||
|
||||
kernel.add_function(plugin_name="TestPlugin", function=decorated_native_function)
|
||||
|
||||
responses = await azure_ai_inference_service.get_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=settings,
|
||||
kernel=kernel,
|
||||
arguments=KernelArguments(),
|
||||
)
|
||||
|
||||
# Completion should be called twice:
|
||||
# One for the tool call and one for the last completion
|
||||
# after the maximum_auto_invoke_attempts is reached
|
||||
assert mock_complete.call_count == 2
|
||||
assert len(responses) == 1
|
||||
assert responses[0].role == "assistant"
|
||||
assert responses[0].content == "Hello"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
async def test_azure_ai_inference_chat_completion_with_function_choice_behavior_no_tool_call(
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
model_id,
|
||||
kernel,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test completion of AzureAIInferenceChatCompletion with function choice behavior but no tool call"""
|
||||
user_message_content: str = "Hello"
|
||||
chat_history.add_user_message(user_message_content)
|
||||
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_chat_completion_response
|
||||
|
||||
responses = await azure_ai_inference_service.get_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=settings,
|
||||
kernel=kernel,
|
||||
arguments=KernelArguments(),
|
||||
)
|
||||
|
||||
mock_complete.assert_awaited_once_with(
|
||||
messages=[UserMessage(content=user_message_content)],
|
||||
model=model_id,
|
||||
model_extras=None,
|
||||
**settings.prepare_settings_dict(),
|
||||
)
|
||||
assert len(responses) == 1
|
||||
assert responses[0].role == "assistant"
|
||||
assert responses[0].content == "Hello"
|
||||
|
||||
|
||||
class MockResponseModel(BaseModel):
|
||||
a: int = Field(..., description="The a field")
|
||||
b: str = Field(..., description="The b field")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
async def test_azure_ai_inference_response_format_json_schema(
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
model_id,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_chat_completion_response,
|
||||
):
|
||||
chat_history.add_user_message("Return an object with fields a and b.")
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings(
|
||||
response_format=MockResponseModel,
|
||||
structured_json_response=True,
|
||||
)
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_chat_completion_response
|
||||
|
||||
_ = await azure_ai_inference_service.get_chat_message_contents(chat_history=chat_history, settings=settings)
|
||||
|
||||
assert mock_complete.call_count == 1
|
||||
kwargs = mock_complete.call_args.kwargs
|
||||
|
||||
assert "response_format" in kwargs
|
||||
response_format = kwargs["response_format"]
|
||||
assert isinstance(response_format, JsonSchemaFormat)
|
||||
assert response_format.name == "MockResponseModel"
|
||||
assert response_format.strict is True
|
||||
|
||||
schema = response_format.schema
|
||||
assert schema["title"] == "MockResponseModel"
|
||||
assert "properties" in schema
|
||||
assert "a" in schema["properties"]
|
||||
assert schema["properties"]["a"]["type"] == "integer"
|
||||
assert "b" in schema["properties"]
|
||||
assert schema["properties"]["b"]["type"] == "string"
|
||||
|
||||
assert kwargs["messages"][0].content == "Return an object with fields a and b."
|
||||
assert kwargs["model"] == model_id
|
||||
|
||||
|
||||
# endregion chat completion
|
||||
|
||||
|
||||
# region streaming chat completion
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
async def test_azure_ai_inference_streaming_chat_completion(
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
model_id,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_streaming_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test streaming completion of AzureAIInferenceChatCompletion"""
|
||||
user_message_content: str = "Hello"
|
||||
chat_history.add_user_message(user_message_content)
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings()
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_streaming_chat_completion_response
|
||||
|
||||
async for messages in azure_ai_inference_service.get_streaming_chat_message_contents(chat_history, settings):
|
||||
assert len(messages) == 1
|
||||
assert messages[0].role == "assistant"
|
||||
assert messages[0].content == "Hello"
|
||||
|
||||
mock_complete.assert_awaited_once_with(
|
||||
stream=True,
|
||||
messages=[UserMessage(content=user_message_content)],
|
||||
model=model_id,
|
||||
model_extras=None,
|
||||
**settings.prepare_settings_dict(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
async def test_azure_ai_inference_chat_streaming_completion_with_standard_parameters(
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
model_id,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_streaming_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test streaming completion of AzureAIInferenceChatCompletion with standard OpenAI parameters"""
|
||||
user_message_content: str = "Hello"
|
||||
chat_history.add_user_message(user_message_content)
|
||||
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings(
|
||||
frequency_penalty=0.5,
|
||||
max_tokens=100,
|
||||
presence_penalty=0.5,
|
||||
seed=123,
|
||||
stop="stop",
|
||||
temperature=0.5,
|
||||
top_p=0.5,
|
||||
)
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_streaming_chat_completion_response
|
||||
|
||||
async for messages in azure_ai_inference_service.get_streaming_chat_message_contents(chat_history, settings):
|
||||
assert len(messages) == 1
|
||||
assert messages[0].role == "assistant"
|
||||
assert messages[0].content == "Hello"
|
||||
|
||||
mock_complete.assert_awaited_once_with(
|
||||
stream=True,
|
||||
messages=[UserMessage(content=user_message_content)],
|
||||
model=model_id,
|
||||
model_extras=None,
|
||||
frequency_penalty=settings.frequency_penalty,
|
||||
max_tokens=settings.max_tokens,
|
||||
presence_penalty=settings.presence_penalty,
|
||||
seed=settings.seed,
|
||||
stop=settings.stop,
|
||||
temperature=settings.temperature,
|
||||
top_p=settings.top_p,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
async def test_azure_ai_inference_streaming_chat_completion_with_extra_parameters(
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
model_id,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_streaming_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test streaming completion of AzureAIInferenceChatCompletion with extra parameters"""
|
||||
user_message_content: str = "Hello"
|
||||
chat_history.add_user_message(user_message_content)
|
||||
extra_parameters = {"test_key": "test_value"}
|
||||
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings(extra_parameters=extra_parameters)
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_streaming_chat_completion_response
|
||||
|
||||
async for messages in azure_ai_inference_service.get_streaming_chat_message_contents(chat_history, settings):
|
||||
assert len(messages) == 1
|
||||
assert messages[0].role == "assistant"
|
||||
assert messages[0].content == "Hello"
|
||||
|
||||
mock_complete.assert_awaited_once_with(
|
||||
stream=True,
|
||||
messages=[UserMessage(content=user_message_content)],
|
||||
model=model_id,
|
||||
model_extras=settings.extra_parameters,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_azure_ai_inference_streaming_chat_completion_with_function_choice_behavior_fail_verification(
|
||||
azure_ai_inference_service,
|
||||
kernel,
|
||||
chat_history: ChatHistory,
|
||||
) -> None:
|
||||
"""Test completion of AzureAIInferenceChatCompletion with function choice behavior expect verification failure"""
|
||||
|
||||
# Missing kernel
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
async for _ in azure_ai_inference_service.get_streaming_chat_message_contents(
|
||||
chat_history,
|
||||
settings,
|
||||
arguments=KernelArguments(),
|
||||
):
|
||||
pass
|
||||
|
||||
# More than 1 responses
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
extra_parameters={"n": 2},
|
||||
)
|
||||
async for _ in azure_ai_inference_service.get_streaming_chat_message_contents(
|
||||
chat_history,
|
||||
settings,
|
||||
kernel=kernel,
|
||||
arguments=KernelArguments(),
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
async def test_azure_ai_inference_streaming_chat_completion_with_function_choice_behavior(
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
kernel: Kernel,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_streaming_chat_completion_response_with_tool_call,
|
||||
decorated_native_function,
|
||||
) -> None:
|
||||
"""Test streaming completion of AzureAIInferenceChatCompletion with function choice behavior."""
|
||||
user_message_content: str = "Hello"
|
||||
chat_history.add_user_message(user_message_content)
|
||||
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
settings.function_choice_behavior.maximum_auto_invoke_attempts = 1
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_streaming_chat_completion_response_with_tool_call
|
||||
|
||||
kernel.add_function(plugin_name="TestPlugin", function=decorated_native_function)
|
||||
|
||||
all_messages = []
|
||||
async for messages in azure_ai_inference_service.get_streaming_chat_message_contents(
|
||||
chat_history,
|
||||
settings,
|
||||
kernel=kernel,
|
||||
arguments=KernelArguments(),
|
||||
):
|
||||
all_messages.extend(messages)
|
||||
|
||||
# Assert the number of total messages
|
||||
assert len(all_messages) == 2, f"Expected 2 messages, got {len(all_messages)}"
|
||||
|
||||
# Validate the first message
|
||||
assert all_messages[0].role == "assistant", f"Unexpected role for first message: {all_messages[0].role}"
|
||||
assert all_messages[0].content == "", f"Unexpected content for first message: {all_messages[0].content}"
|
||||
assert all_messages[0].finish_reason == FinishReason.TOOL_CALLS, (
|
||||
f"Unexpected finish reason for first message: {all_messages[0].finish_reason}"
|
||||
)
|
||||
|
||||
# Validate the second message
|
||||
assert all_messages[1].role == "tool", f"Unexpected role for second message: {all_messages[1].role}"
|
||||
assert all_messages[1].content == "", f"Unexpected content for second message: {all_messages[1].content}"
|
||||
assert all_messages[1].finish_reason is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
async def test_azure_ai_inference_streaming_chat_completion_with_function_choice_behavior_no_tool_call(
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
model_id,
|
||||
kernel,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_streaming_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test streaming completion of AzureAIInferenceChatCompletion with function choice behavior but no tool call"""
|
||||
user_message_content: str = "Hello"
|
||||
chat_history.add_user_message(user_message_content)
|
||||
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_streaming_chat_completion_response
|
||||
|
||||
async for messages in azure_ai_inference_service.get_streaming_chat_message_contents(
|
||||
chat_history,
|
||||
settings,
|
||||
kernel=kernel,
|
||||
arguments=KernelArguments(),
|
||||
):
|
||||
assert len(messages) == 1
|
||||
assert messages[0].role == "assistant"
|
||||
assert messages[0].content == "Hello"
|
||||
|
||||
mock_complete.assert_awaited_once_with(
|
||||
stream=True,
|
||||
messages=[UserMessage(content=user_message_content)],
|
||||
model=model_id,
|
||||
model_extras=None,
|
||||
**settings.prepare_settings_dict(),
|
||||
)
|
||||
|
||||
|
||||
class MockStreamingResponseModel(BaseModel):
|
||||
foo: float = Field(..., description="Foo value")
|
||||
bar: bool = Field(..., description="Bar value")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
async def test_azure_ai_inference_streaming_response_format_json_schema(
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
model_id,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_streaming_chat_completion_response,
|
||||
):
|
||||
chat_history.add_user_message("Stream a response with foo and bar.")
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings(
|
||||
response_format=MockStreamingResponseModel,
|
||||
structured_json_response=True,
|
||||
)
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_streaming_chat_completion_response
|
||||
|
||||
messages = []
|
||||
async for chunk in azure_ai_inference_service.get_streaming_chat_message_contents(chat_history, settings):
|
||||
messages.extend(chunk)
|
||||
|
||||
assert mock_complete.call_count == 1
|
||||
kwargs = mock_complete.call_args.kwargs
|
||||
assert "response_format" in kwargs
|
||||
response_format = kwargs["response_format"]
|
||||
assert isinstance(response_format, JsonSchemaFormat)
|
||||
assert response_format.name == "MockStreamingResponseModel"
|
||||
assert response_format.strict is True
|
||||
schema = response_format.schema
|
||||
assert schema["title"] == "MockStreamingResponseModel"
|
||||
assert "foo" in schema["properties"]
|
||||
assert schema["properties"]["foo"]["type"] == "number"
|
||||
assert "bar" in schema["properties"]
|
||||
assert schema["properties"]["bar"]["type"] == "boolean"
|
||||
assert kwargs["stream"] is True
|
||||
assert kwargs["model"] == model_id
|
||||
|
||||
|
||||
# endregion streaming chat completion
|
||||
+163
@@ -0,0 +1,163 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from azure.ai.inference.aio import EmbeddingsClient
|
||||
from azure.core.credentials import AzureKeyCredential
|
||||
|
||||
from semantic_kernel.connectors.ai.azure_ai_inference import (
|
||||
AzureAIInferenceEmbeddingPromptExecutionSettings,
|
||||
AzureAIInferenceTextEmbedding,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.azure_ai_inference.azure_ai_inference_settings import AzureAIInferenceSettings
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
|
||||
from semantic_kernel.utils.telemetry.user_agent import SEMANTIC_KERNEL_USER_AGENT
|
||||
|
||||
|
||||
def test_azure_ai_inference_text_embedding_init(azure_ai_inference_unit_test_env, model_id) -> None:
|
||||
"""Test initialization of AzureAIInferenceTextEmbedding"""
|
||||
azure_ai_inference = AzureAIInferenceTextEmbedding(model_id)
|
||||
|
||||
assert azure_ai_inference.ai_model_id == model_id
|
||||
assert azure_ai_inference.service_id == model_id
|
||||
assert isinstance(azure_ai_inference.client, EmbeddingsClient)
|
||||
|
||||
|
||||
@patch("azure.ai.inference.aio.EmbeddingsClient.__init__", return_value=None)
|
||||
def test_azure_ai_inference_text_embedding_client_init(mock_client, azure_ai_inference_unit_test_env, model_id) -> None:
|
||||
"""Test initialization of the Azure AI Inference client"""
|
||||
endpoint = azure_ai_inference_unit_test_env["AZURE_AI_INFERENCE_ENDPOINT"]
|
||||
api_key = azure_ai_inference_unit_test_env["AZURE_AI_INFERENCE_API_KEY"]
|
||||
settings = AzureAIInferenceSettings(endpoint=endpoint, api_key=api_key)
|
||||
|
||||
_ = AzureAIInferenceTextEmbedding(model_id)
|
||||
|
||||
assert mock_client.call_count == 1
|
||||
assert isinstance(mock_client.call_args.kwargs["endpoint"], str)
|
||||
assert mock_client.call_args.kwargs["endpoint"] == str(settings.endpoint)
|
||||
assert isinstance(mock_client.call_args.kwargs["credential"], AzureKeyCredential)
|
||||
assert mock_client.call_args.kwargs["credential"].key == settings.api_key.get_secret_value()
|
||||
assert mock_client.call_args.kwargs["user_agent"] == SEMANTIC_KERNEL_USER_AGENT
|
||||
|
||||
|
||||
def test_azure_ai_inference_text_embedding_init_with_service_id(
|
||||
azure_ai_inference_unit_test_env, model_id, service_id
|
||||
) -> None:
|
||||
"""Test initialization of AzureAIInferenceTextEmbedding"""
|
||||
azure_ai_inference = AzureAIInferenceTextEmbedding(model_id, service_id=service_id)
|
||||
|
||||
assert azure_ai_inference.ai_model_id == model_id
|
||||
assert azure_ai_inference.service_id == service_id
|
||||
assert isinstance(azure_ai_inference.client, EmbeddingsClient)
|
||||
|
||||
|
||||
def test_azure_ai_inference_text_embedding_init_with_api_version(azure_ai_inference_unit_test_env, model_id) -> None:
|
||||
"""Test initialization of AzureAIInferenceTextEmbedding with api_version"""
|
||||
azure_ai_inference = AzureAIInferenceTextEmbedding(model_id, api_version="2024-02-15-test")
|
||||
|
||||
assert azure_ai_inference.ai_model_id == model_id
|
||||
assert isinstance(azure_ai_inference.client, EmbeddingsClient)
|
||||
assert azure_ai_inference.client._config.api_version == "2024-02-15-test"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_client",
|
||||
[AzureAIInferenceTextEmbedding.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
def test_azure_ai_inference_chat_completion_init_with_custom_client(azure_ai_inference_client, model_id) -> None:
|
||||
"""Test initialization of AzureAIInferenceTextEmbedding with custom client"""
|
||||
client = azure_ai_inference_client
|
||||
azure_ai_inference = AzureAIInferenceTextEmbedding(model_id, client=client)
|
||||
|
||||
assert azure_ai_inference.ai_model_id == model_id
|
||||
assert azure_ai_inference.service_id == model_id
|
||||
assert azure_ai_inference.client == client
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_AI_INFERENCE_ENDPOINT"]], indirect=True)
|
||||
def test_azure_ai_inference_text_embedding_init_with_empty_endpoint(azure_ai_inference_unit_test_env, model_id) -> None:
|
||||
"""Test initialization of AzureAIInferenceTextEmbedding with empty endpoint"""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureAIInferenceTextEmbedding(model_id, env_file_path="fake_path")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceTextEmbedding.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(EmbeddingsClient, "embed", new_callable=AsyncMock)
|
||||
async def test_azure_ai_inference_text_embedding(
|
||||
mock_embed,
|
||||
azure_ai_inference_service,
|
||||
model_id,
|
||||
) -> None:
|
||||
"""Test text embedding generation of AzureAIInferenceTextEmbedding without settings"""
|
||||
texts = ["hello", "world"]
|
||||
await azure_ai_inference_service.generate_embeddings(texts)
|
||||
|
||||
mock_embed.assert_awaited_once_with(
|
||||
input=texts,
|
||||
model=model_id,
|
||||
model_extras=None,
|
||||
dimensions=None,
|
||||
encoding_format=None,
|
||||
input_type=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceTextEmbedding.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(EmbeddingsClient, "embed", new_callable=AsyncMock)
|
||||
async def test_azure_ai_inference_text_embedding_with_standard_settings(
|
||||
mock_embed,
|
||||
azure_ai_inference_service,
|
||||
model_id,
|
||||
) -> None:
|
||||
"""Test text embedding generation of AzureAIInferenceTextEmbedding with standard settings"""
|
||||
texts = ["hello", "world"]
|
||||
settings = AzureAIInferenceEmbeddingPromptExecutionSettings(
|
||||
dimensions=1024, encoding_format="float", input_type="text"
|
||||
)
|
||||
await azure_ai_inference_service.generate_embeddings(texts, settings)
|
||||
|
||||
mock_embed.assert_awaited_once_with(
|
||||
input=texts,
|
||||
model=model_id,
|
||||
model_extras=None,
|
||||
dimensions=settings.dimensions,
|
||||
encoding_format=settings.encoding_format,
|
||||
input_type=settings.input_type,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceTextEmbedding.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(EmbeddingsClient, "embed", new_callable=AsyncMock)
|
||||
async def test_azure_ai_inference_text_embedding_with_extra_parameters(
|
||||
mock_embed,
|
||||
azure_ai_inference_service,
|
||||
model_id,
|
||||
) -> None:
|
||||
"""Test text embedding generation of AzureAIInferenceTextEmbedding with extra parameters"""
|
||||
texts = ["hello", "world"]
|
||||
extra_parameters = {"test_key": "test_value"}
|
||||
settings = AzureAIInferenceEmbeddingPromptExecutionSettings(extra_parameters=extra_parameters)
|
||||
await azure_ai_inference_service.generate_embeddings(texts, settings)
|
||||
|
||||
mock_embed.assert_awaited_once_with(
|
||||
input=texts,
|
||||
model=model_id,
|
||||
model_extras=extra_parameters,
|
||||
dimensions=settings.dimensions,
|
||||
encoding_format=settings.encoding_format,
|
||||
input_type=settings.input_type,
|
||||
)
|
||||
+239
@@ -0,0 +1,239 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from azure.ai.inference.aio import ChatCompletionsClient
|
||||
|
||||
from semantic_kernel.connectors.ai.azure_ai_inference.azure_ai_inference_prompt_execution_settings import (
|
||||
AzureAIInferenceChatPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.azure_ai_inference.services.azure_ai_inference_chat_completion import (
|
||||
AzureAIInferenceChatCompletion,
|
||||
)
|
||||
from semantic_kernel.contents.chat_history import ChatHistory
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.uninstrument")
|
||||
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.instrument")
|
||||
async def test_azure_ai_inference_chat_completion_instrumentation(
|
||||
mock_instrument,
|
||||
mock_uninstrument,
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_chat_completion_response,
|
||||
model_diagnostics_test_env,
|
||||
) -> None:
|
||||
"""Test completion of AzureAIInferenceChatCompletion"""
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings()
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_chat_completion_response
|
||||
|
||||
await azure_ai_inference_service.get_chat_message_contents(chat_history=chat_history, settings=settings)
|
||||
|
||||
mock_instrument.assert_called_once_with(enable_content_recording=True)
|
||||
mock_uninstrument.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[
|
||||
AzureAIInferenceChatCompletion.__name__,
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"override_env_param_dict",
|
||||
[
|
||||
{
|
||||
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS": "False",
|
||||
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE": "False",
|
||||
},
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.uninstrument")
|
||||
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.instrument")
|
||||
async def test_azure_ai_inference_chat_completion_not_instrumentation(
|
||||
mock_instrument,
|
||||
mock_uninstrument,
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_chat_completion_response,
|
||||
model_diagnostics_test_env,
|
||||
) -> None:
|
||||
"""Test completion of AzureAIInferenceChatCompletion"""
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings()
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_chat_completion_response
|
||||
|
||||
await azure_ai_inference_service.get_chat_message_contents(chat_history=chat_history, settings=settings)
|
||||
|
||||
mock_instrument.assert_not_called()
|
||||
mock_uninstrument.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[
|
||||
AzureAIInferenceChatCompletion.__name__,
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"override_env_param_dict",
|
||||
[
|
||||
{
|
||||
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS": "True",
|
||||
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE": "False",
|
||||
},
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.uninstrument")
|
||||
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.instrument")
|
||||
async def test_azure_ai_inference_chat_completion_instrumentation_without_sensitive(
|
||||
mock_instrument,
|
||||
mock_uninstrument,
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_chat_completion_response,
|
||||
model_diagnostics_test_env,
|
||||
) -> None:
|
||||
"""Test completion of AzureAIInferenceChatCompletion"""
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings()
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_chat_completion_response
|
||||
|
||||
await azure_ai_inference_service.get_chat_message_contents(chat_history=chat_history, settings=settings)
|
||||
|
||||
mock_instrument.assert_called_once_with(enable_content_recording=False)
|
||||
mock_uninstrument.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[AzureAIInferenceChatCompletion.__name__],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.uninstrument")
|
||||
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.instrument")
|
||||
async def test_azure_ai_inference_streaming_chat_completion_instrumentation(
|
||||
mock_instrument,
|
||||
mock_uninstrument,
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_streaming_chat_completion_response,
|
||||
model_diagnostics_test_env,
|
||||
) -> None:
|
||||
"""Test completion of AzureAIInferenceChatCompletion"""
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings()
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_streaming_chat_completion_response
|
||||
|
||||
async for _ in azure_ai_inference_service.get_streaming_chat_message_contents(
|
||||
chat_history=chat_history, settings=settings
|
||||
):
|
||||
pass
|
||||
|
||||
mock_instrument.assert_called_once_with(enable_content_recording=True)
|
||||
mock_uninstrument.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[
|
||||
AzureAIInferenceChatCompletion.__name__,
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"override_env_param_dict",
|
||||
[
|
||||
{
|
||||
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS": "False",
|
||||
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE": "False",
|
||||
},
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.uninstrument")
|
||||
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.instrument")
|
||||
async def test_azure_ai_inference_streaming_chat_completion_not_instrumentation(
|
||||
mock_instrument,
|
||||
mock_uninstrument,
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_streaming_chat_completion_response,
|
||||
model_diagnostics_test_env,
|
||||
) -> None:
|
||||
"""Test completion of AzureAIInferenceChatCompletion"""
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings()
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_streaming_chat_completion_response
|
||||
|
||||
async for _ in azure_ai_inference_service.get_streaming_chat_message_contents(
|
||||
chat_history=chat_history, settings=settings
|
||||
):
|
||||
pass
|
||||
|
||||
mock_instrument.assert_not_called()
|
||||
mock_uninstrument.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"azure_ai_inference_service",
|
||||
[
|
||||
AzureAIInferenceChatCompletion.__name__,
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"override_env_param_dict",
|
||||
[
|
||||
{
|
||||
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS": "True",
|
||||
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE": "False",
|
||||
},
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
|
||||
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.uninstrument")
|
||||
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.instrument")
|
||||
async def test_azure_ai_inference_streaming_chat_completion_instrumentation_without_sensitive(
|
||||
mock_instrument,
|
||||
mock_uninstrument,
|
||||
mock_complete,
|
||||
azure_ai_inference_service,
|
||||
chat_history: ChatHistory,
|
||||
mock_azure_ai_inference_streaming_chat_completion_response,
|
||||
model_diagnostics_test_env,
|
||||
) -> None:
|
||||
"""Test completion of AzureAIInferenceChatCompletion"""
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings()
|
||||
|
||||
mock_complete.return_value = mock_azure_ai_inference_streaming_chat_completion_response
|
||||
|
||||
async for _ in azure_ai_inference_service.get_streaming_chat_message_contents(
|
||||
chat_history=chat_history, settings=settings
|
||||
):
|
||||
pass
|
||||
|
||||
mock_instrument.assert_called_once_with(enable_content_recording=False)
|
||||
mock_uninstrument.assert_called_once()
|
||||
+169
@@ -0,0 +1,169 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
import pytest
|
||||
from azure.ai.inference.models import (
|
||||
AssistantMessage,
|
||||
ImageContentItem,
|
||||
SystemMessage,
|
||||
TextContentItem,
|
||||
ToolMessage,
|
||||
UserMessage,
|
||||
)
|
||||
|
||||
from semantic_kernel.connectors.ai.azure_ai_inference.services.utils import MESSAGE_CONVERTERS
|
||||
from semantic_kernel.contents.chat_message_content import ChatMessageContent
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
from semantic_kernel.contents.function_result_content import FunctionResultContent
|
||||
from semantic_kernel.contents.image_content import ImageContent
|
||||
from semantic_kernel.contents.text_content import TextContent
|
||||
from semantic_kernel.contents.utils.author_role import AuthorRole
|
||||
|
||||
|
||||
def test_message_convertors_contain_all_author_roles() -> None:
|
||||
"""Test that all AuthorRoles are present in the MESSAGE_CONVERTERS dict."""
|
||||
for role in AuthorRole:
|
||||
assert role in MESSAGE_CONVERTERS
|
||||
|
||||
|
||||
def test_format_system_message() -> None:
|
||||
"""Test that a system message is formatted correctly."""
|
||||
message = ChatMessageContent(role=AuthorRole.SYSTEM, content="test content")
|
||||
system_message = MESSAGE_CONVERTERS[message.role](message)
|
||||
|
||||
assert isinstance(system_message, SystemMessage)
|
||||
assert system_message.content == message.content
|
||||
|
||||
|
||||
def test_format_user_message_with_no_image() -> None:
|
||||
"""Test that a user message with no image items is formatted correctly."""
|
||||
message = ChatMessageContent(role=AuthorRole.USER, content="test content")
|
||||
user_message = MESSAGE_CONVERTERS[message.role](message)
|
||||
|
||||
assert isinstance(user_message, UserMessage)
|
||||
assert user_message.content == message.content
|
||||
|
||||
|
||||
def test_format_user_message_with_image() -> None:
|
||||
"""Test that a user message with image items is formatted correctly"""
|
||||
message = ChatMessageContent(
|
||||
role=AuthorRole.USER,
|
||||
items=[
|
||||
TextContent(text="test text"),
|
||||
ImageContent(uri="https://test.com/image.jpg"),
|
||||
],
|
||||
)
|
||||
user_message = MESSAGE_CONVERTERS[message.role](message)
|
||||
|
||||
assert isinstance(user_message, UserMessage)
|
||||
assert len(user_message.content) == 2
|
||||
assert isinstance(user_message.content[0], TextContentItem)
|
||||
assert isinstance(user_message.content[1], ImageContentItem)
|
||||
|
||||
|
||||
def test_format_user_message_with_unsupported_items() -> None:
|
||||
"""Test that a user message with unsupported items is formatted correctly"""
|
||||
message = ChatMessageContent(
|
||||
role=AuthorRole.USER,
|
||||
items=[
|
||||
TextContent(text="test text"),
|
||||
ImageContent(), # ImageContent without uri or data_uri is unsupported
|
||||
FunctionCallContent(id="test function"), # FunctionCallContent unsupported
|
||||
],
|
||||
)
|
||||
user_message = MESSAGE_CONVERTERS[message.role](message)
|
||||
|
||||
assert isinstance(user_message, UserMessage)
|
||||
assert len(user_message.content) == 1
|
||||
assert isinstance(user_message.content[0], TextContentItem)
|
||||
|
||||
|
||||
def test_format_assistant_message() -> None:
|
||||
"""Test that an assistant message is formatted correctly."""
|
||||
message = ChatMessageContent(role=AuthorRole.ASSISTANT, content="test content")
|
||||
assistant_message = MESSAGE_CONVERTERS[message.role](message)
|
||||
|
||||
assert isinstance(assistant_message, AssistantMessage)
|
||||
assert assistant_message.content == message.content
|
||||
|
||||
|
||||
def test_format_assistant_message_with_tool_call() -> None:
|
||||
"""Test that an assistant message with a tool call is formatted correctly."""
|
||||
function_call_content = FunctionCallContent(id="test function")
|
||||
|
||||
message = ChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT,
|
||||
items=[function_call_content],
|
||||
)
|
||||
assistant_message = MESSAGE_CONVERTERS[message.role](message)
|
||||
|
||||
assert isinstance(assistant_message, AssistantMessage)
|
||||
assert assistant_message.content == message.content
|
||||
assert len(assistant_message.tool_calls) == 1
|
||||
assert assistant_message.tool_calls[0].id == function_call_content.id
|
||||
|
||||
|
||||
def test_format_assistant_message_with_unsupported_items() -> None:
|
||||
"""Test that an assistant message with unsupported items is formatted correctly."""
|
||||
text_content = TextContent(text="test text")
|
||||
|
||||
message = ChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT,
|
||||
items=[
|
||||
text_content,
|
||||
ImageContent(), # ImageContent is unsupported
|
||||
],
|
||||
)
|
||||
assistant_message = MESSAGE_CONVERTERS[message.role](message)
|
||||
|
||||
assert isinstance(assistant_message, AssistantMessage)
|
||||
assert assistant_message.content == message.content
|
||||
|
||||
|
||||
def test_format_tool_message() -> None:
|
||||
"""Test that a tool message is formatted correctly."""
|
||||
function_result_content = FunctionResultContent(id="test function", result="test result")
|
||||
|
||||
message = ChatMessageContent(
|
||||
role=AuthorRole.TOOL,
|
||||
items=[function_result_content],
|
||||
)
|
||||
tool_message = MESSAGE_CONVERTERS[message.role](message)
|
||||
|
||||
assert isinstance(tool_message, ToolMessage)
|
||||
assert tool_message.content == function_result_content.result
|
||||
assert tool_message.tool_call_id == function_result_content.id
|
||||
|
||||
|
||||
def test_format_tool_message_item_not_found_as_the_first_item() -> None:
|
||||
"""Test that formatting a tool message where the function result item is not the first item."""
|
||||
function_result_content = FunctionResultContent(id="test function", result="test result")
|
||||
|
||||
message = ChatMessageContent(
|
||||
role=AuthorRole.TOOL,
|
||||
items=[
|
||||
TextContent(text="test text"),
|
||||
function_result_content,
|
||||
],
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
MESSAGE_CONVERTERS[message.role](message)
|
||||
|
||||
|
||||
def test_format_tool_message_with_more_than_one_items() -> None:
|
||||
"""Test that a tool message with more than one item is formatted correctly."""
|
||||
function_result_content = FunctionResultContent(id="test function", result="test result")
|
||||
|
||||
message = ChatMessageContent(
|
||||
role=AuthorRole.TOOL,
|
||||
items=[
|
||||
function_result_content,
|
||||
TextContent(text="test text"),
|
||||
],
|
||||
)
|
||||
tool_message = MESSAGE_CONVERTERS[message.role](message)
|
||||
|
||||
assert isinstance(tool_message, ToolMessage)
|
||||
assert tool_message.content == function_result_content.result
|
||||
assert tool_message.tool_call_id == function_result_content.id
|
||||
+143
@@ -0,0 +1,143 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from semantic_kernel.connectors.ai.azure_ai_inference import (
|
||||
AzureAIInferenceChatPromptExecutionSettings,
|
||||
AzureAIInferenceEmbeddingPromptExecutionSettings,
|
||||
AzureAIInferencePromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
|
||||
|
||||
def test_default_azure_ai_inference_prompt_execution_settings():
|
||||
settings = AzureAIInferencePromptExecutionSettings()
|
||||
|
||||
assert settings.frequency_penalty is None
|
||||
assert settings.max_tokens is None
|
||||
assert settings.presence_penalty is None
|
||||
assert settings.seed is None
|
||||
assert settings.stop is None
|
||||
assert settings.temperature is None
|
||||
assert settings.top_p is None
|
||||
assert settings.extra_parameters is None
|
||||
|
||||
|
||||
def test_custom_azure_ai_inference_prompt_execution_settings():
|
||||
settings = AzureAIInferencePromptExecutionSettings(
|
||||
frequency_penalty=0.5,
|
||||
max_tokens=128,
|
||||
presence_penalty=0.5,
|
||||
seed=1,
|
||||
stop="world",
|
||||
temperature=0.5,
|
||||
top_p=0.5,
|
||||
extra_parameters={"key": "value"},
|
||||
)
|
||||
|
||||
assert settings.frequency_penalty == 0.5
|
||||
assert settings.max_tokens == 128
|
||||
assert settings.presence_penalty == 0.5
|
||||
assert settings.seed == 1
|
||||
assert settings.stop == "world"
|
||||
assert settings.temperature == 0.5
|
||||
assert settings.top_p == 0.5
|
||||
assert settings.extra_parameters == {"key": "value"}
|
||||
|
||||
|
||||
def test_azure_ai_inference_prompt_execution_settings_from_default_completion_config():
|
||||
settings = PromptExecutionSettings(service_id="test_service")
|
||||
chat_settings = AzureAIInferenceChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
|
||||
assert chat_settings.service_id == "test_service"
|
||||
assert chat_settings.frequency_penalty is None
|
||||
assert chat_settings.max_tokens is None
|
||||
assert chat_settings.presence_penalty is None
|
||||
assert chat_settings.seed is None
|
||||
assert chat_settings.stop is None
|
||||
assert chat_settings.temperature is None
|
||||
assert chat_settings.top_p is None
|
||||
assert chat_settings.extra_parameters is None
|
||||
|
||||
|
||||
def test_azure_ai_inference_prompt_execution_settings_from_openai_prompt_execution_settings():
|
||||
chat_settings = AzureAIInferenceChatPromptExecutionSettings(service_id="test_service", temperature=1.0)
|
||||
new_settings = AzureAIInferencePromptExecutionSettings(service_id="test_2", temperature=0.0)
|
||||
chat_settings.update_from_prompt_execution_settings(new_settings)
|
||||
|
||||
assert chat_settings.service_id == "test_2"
|
||||
assert chat_settings.temperature == 0.0
|
||||
|
||||
|
||||
def test_azure_ai_inference_prompt_execution_settings_from_custom_completion_config():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"frequency_penalty": 0.5,
|
||||
"max_tokens": 128,
|
||||
"presence_penalty": 0.5,
|
||||
"seed": 1,
|
||||
"stop": "world",
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"extra_parameters": {"key": "value"},
|
||||
},
|
||||
)
|
||||
chat_settings = AzureAIInferenceChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
|
||||
assert chat_settings.service_id == "test_service"
|
||||
assert chat_settings.frequency_penalty == 0.5
|
||||
assert chat_settings.max_tokens == 128
|
||||
assert chat_settings.presence_penalty == 0.5
|
||||
assert chat_settings.seed == 1
|
||||
assert chat_settings.stop == "world"
|
||||
assert chat_settings.temperature == 0.5
|
||||
assert chat_settings.top_p == 0.5
|
||||
assert chat_settings.extra_parameters == {"key": "value"}
|
||||
|
||||
|
||||
def test_azure_ai_inference_chat_prompt_execution_settings_from_custom_completion_config_with_functions():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"tools": [{"function": {}}],
|
||||
},
|
||||
)
|
||||
chat_settings = AzureAIInferenceChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
|
||||
assert chat_settings.tools == [{"function": {}}]
|
||||
|
||||
|
||||
def test_create_options():
|
||||
settings = AzureAIInferenceChatPromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"frequency_penalty": 0.5,
|
||||
"max_tokens": 128,
|
||||
"presence_penalty": 0.5,
|
||||
"seed": 1,
|
||||
"stop": "world",
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"extra_parameters": {"key": "value"},
|
||||
},
|
||||
)
|
||||
options = settings.prepare_settings_dict()
|
||||
|
||||
assert options["frequency_penalty"] == 0.5
|
||||
assert options["max_tokens"] == 128
|
||||
assert options["presence_penalty"] == 0.5
|
||||
assert options["seed"] == 1
|
||||
assert options["stop"] == "world"
|
||||
assert options["temperature"] == 0.5
|
||||
assert options["top_p"] == 0.5
|
||||
assert options["extra_parameters"] == {"key": "value"}
|
||||
assert "tools" not in options
|
||||
assert "tool_config" not in options
|
||||
|
||||
|
||||
def test_default_azure_ai_inference_embedding_prompt_execution_settings():
|
||||
settings = AzureAIInferenceEmbeddingPromptExecutionSettings()
|
||||
|
||||
assert settings.dimensions is None
|
||||
assert settings.encoding_format is None
|
||||
assert settings.input_type is None
|
||||
assert settings.extra_parameters is None
|
||||
@@ -0,0 +1,320 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import json
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.ai.bedrock.services.model_provider.bedrock_model_provider import BedrockModelProvider
|
||||
from semantic_kernel.contents.chat_history import ChatHistory
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def model_id(request) -> str:
|
||||
if hasattr(request, "param"):
|
||||
return request.param
|
||||
return "test_model_id"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def service_id() -> str:
|
||||
return "test_service_id"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def chat_history() -> ChatHistory:
|
||||
chat_history = ChatHistory(system_message="You are a helpful assistant.")
|
||||
|
||||
chat_history.add_user_message("Hello!")
|
||||
chat_history.add_assistant_message("Hi! How can I help you today?")
|
||||
chat_history.add_system_message("Be polite and respectful.")
|
||||
chat_history.add_user_message("I need help with a task.")
|
||||
|
||||
return chat_history
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def bedrock_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for Amazon Bedrock AI connector unit tests."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"BEDROCK_TEXT_MODEL_ID": "env_test_text_model_id",
|
||||
"BEDROCK_CHAT_MODEL_ID": "env_test_chat_model_id",
|
||||
"BEDROCK_EMBEDDING_MODEL_ID": "env_test_embedding_model_id",
|
||||
"BEDROCK_MODEL_PROVIDER": "amazon",
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
class MockBedrockClient(Mock):
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def get_foundation_model(self, *args, **kwargs):
|
||||
return {
|
||||
"modelDetails": {
|
||||
"responseStreamingSupported": True,
|
||||
"inputModalities": ["TEXT"],
|
||||
"outputModalities": ["TEXT", "EMBEDDING"],
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class MockBedrockRuntimeClient(Mock):
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def converse(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def converse_stream(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def invoke_model(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def invoke_model_with_response_stream(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
|
||||
# region mock chat completion responses
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_bedrock_chat_completion_response():
|
||||
# https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference-call.html#conversation-inference-call-response
|
||||
return {
|
||||
"output": {
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"text": "Hi! How can I help you today?",
|
||||
}
|
||||
],
|
||||
}
|
||||
},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {
|
||||
"inputTokens": 125,
|
||||
"outputTokens": 60,
|
||||
"totalTokens": 185,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_bedrock_streaming_chat_completion_response():
|
||||
# https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference-call.html#conversation-inference-call-response
|
||||
events = [
|
||||
{"messageStart": {"role": "assistant"}},
|
||||
{"contentBlockStart": {"contentBlockIndex": 0, "start": {}}},
|
||||
{"contentBlockDelta": {"contentBlockIndex": 0, "delta": {"text": "Hi! "}}},
|
||||
{"contentBlockDelta": {"contentBlockIndex": 0, "delta": {"text": "How can "}}},
|
||||
{"contentBlockDelta": {"contentBlockIndex": 0, "delta": {"text": "I help you today?"}}},
|
||||
{"contentBlockStop": {"contentBlockIndex": 0}},
|
||||
{"messageStop": {"stopReason": "end_turn"}},
|
||||
{
|
||||
"metadata": {
|
||||
"metrics": {"latencyMs": 1000},
|
||||
"usage": {"inputTokens": 125, "outputTokens": 60, "totalTokens": 185},
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
def event_stream(events):
|
||||
yield from events
|
||||
|
||||
return {"stream": event_stream(events)}
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_bedrock_streaming_chat_completion_invalid_response():
|
||||
events = [
|
||||
{"unknown": {}},
|
||||
]
|
||||
|
||||
def event_stream(events):
|
||||
yield from events
|
||||
|
||||
return {"stream": event_stream(events)}
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
# region mock text completion responses
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def output_text():
|
||||
return "Hi! How can I help you today?"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def model_provider():
|
||||
return BedrockModelProvider.AMAZON
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_bedrock_text_completion_response(
|
||||
model_id: str,
|
||||
output_text: str,
|
||||
request,
|
||||
):
|
||||
# Check if model_provider fixture is requested by the test
|
||||
model_provider = None
|
||||
if "model_provider" in request.fixturenames:
|
||||
model_provider = request.getfixturevalue("model_provider")
|
||||
else:
|
||||
model_provider = BedrockModelProvider.to_model_provider(model_id)
|
||||
|
||||
match model_provider:
|
||||
case BedrockModelProvider.AMAZON:
|
||||
body = {
|
||||
"inputTextTokenCount": 10,
|
||||
"results": [
|
||||
{
|
||||
"tokenCount": 10,
|
||||
"outputText": output_text,
|
||||
"completionReason": "FINISHED ",
|
||||
}
|
||||
],
|
||||
}
|
||||
case BedrockModelProvider.ANTHROPIC:
|
||||
body = {
|
||||
"completion": output_text,
|
||||
"stop_reason": "stop_sequence",
|
||||
"stop": "",
|
||||
}
|
||||
case BedrockModelProvider.COHERE:
|
||||
body = {
|
||||
"generations": [
|
||||
{
|
||||
"text": output_text,
|
||||
}
|
||||
],
|
||||
}
|
||||
case BedrockModelProvider.AI21LABS:
|
||||
body = {
|
||||
"completions": [
|
||||
{
|
||||
"data": {
|
||||
"text": output_text,
|
||||
}
|
||||
}
|
||||
],
|
||||
}
|
||||
case BedrockModelProvider.META:
|
||||
body = {
|
||||
"generation": output_text,
|
||||
"prompt_token_count": 10,
|
||||
"generation_token_count": 10,
|
||||
}
|
||||
case BedrockModelProvider.MISTRALAI:
|
||||
body = {"outputs": [{"text": output_text}]}
|
||||
|
||||
mock = Mock()
|
||||
mock.read.return_value = json.dumps(body)
|
||||
|
||||
return {"body": mock}
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_bedrock_streaming_text_completion_response(
|
||||
model_id: str,
|
||||
output_text: str,
|
||||
request,
|
||||
):
|
||||
# Check if model_provider fixture is requested by the test
|
||||
model_provider = None
|
||||
if "model_provider" in request.fixturenames:
|
||||
model_provider = request.getfixturevalue("model_provider")
|
||||
else:
|
||||
model_provider = BedrockModelProvider.to_model_provider(model_id)
|
||||
|
||||
match model_provider:
|
||||
case BedrockModelProvider.AMAZON:
|
||||
chunks = [
|
||||
{
|
||||
"chunk": {
|
||||
"bytes": json.dumps({
|
||||
"inputTextTokenCount": 10,
|
||||
"totalOutputTextTokenCount": 10,
|
||||
"outputText": chunk,
|
||||
}).encode(),
|
||||
}
|
||||
}
|
||||
for chunk in [output_text[i : i + 3] for i in range(0, len(output_text), 3)]
|
||||
]
|
||||
|
||||
def event_stream(events):
|
||||
yield from events
|
||||
|
||||
return {"body": event_stream(chunks)}
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
# region mock text embedding responses
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_bedrock_text_embedding_response(
|
||||
model_id: str,
|
||||
request,
|
||||
):
|
||||
# Check if model_provider fixture is requested by the test
|
||||
model_provider = None
|
||||
if "model_provider" in request.fixturenames:
|
||||
model_provider = request.getfixturevalue("model_provider")
|
||||
else:
|
||||
model_provider = BedrockModelProvider.to_model_provider(model_id)
|
||||
|
||||
match model_provider:
|
||||
case BedrockModelProvider.AMAZON:
|
||||
body = {
|
||||
"embedding": [0.1, 0.2, 0.3],
|
||||
}
|
||||
case BedrockModelProvider.COHERE:
|
||||
body = {
|
||||
"embeddings": [[0.1, 0.2, 0.3]],
|
||||
}
|
||||
|
||||
mock = Mock()
|
||||
mock.read.return_value = json.dumps(body)
|
||||
|
||||
return {"body": mock}
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_bedrock_text_embedding_invalid_response(model_id: str):
|
||||
model_provider = BedrockModelProvider.to_model_provider(model_id)
|
||||
|
||||
match model_provider:
|
||||
case BedrockModelProvider.AMAZON:
|
||||
body = {"embedding": 0.1}
|
||||
case BedrockModelProvider.COHERE:
|
||||
body = {"embeddings": 0.1}
|
||||
|
||||
mock = Mock()
|
||||
mock.read.return_value = json.dumps(body)
|
||||
|
||||
return {"body": mock}
|
||||
|
||||
|
||||
# endregion
|
||||
@@ -0,0 +1,371 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from functools import reduce
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import boto3
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.ai.bedrock.bedrock_prompt_execution_settings import BedrockChatPromptExecutionSettings
|
||||
from semantic_kernel.connectors.ai.bedrock.services.bedrock_chat_completion import BedrockChatCompletion
|
||||
from semantic_kernel.connectors.ai.bedrock.services.model_provider.bedrock_model_provider import BedrockModelProvider
|
||||
from semantic_kernel.connectors.ai.completion_usage import CompletionUsage
|
||||
from semantic_kernel.contents.chat_history import ChatHistory
|
||||
from semantic_kernel.contents.chat_message_content import ChatMessageContent
|
||||
from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent
|
||||
from semantic_kernel.contents.text_content import TextContent
|
||||
from semantic_kernel.contents.utils.author_role import AuthorRole
|
||||
from semantic_kernel.contents.utils.finish_reason import FinishReason
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidResponseError
|
||||
from tests.unit.connectors.ai.bedrock.conftest import MockBedrockClient, MockBedrockRuntimeClient
|
||||
|
||||
# region init
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_chat_completion_init(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service"""
|
||||
bedrock_chat_completion = BedrockChatCompletion()
|
||||
|
||||
assert bedrock_chat_completion.ai_model_id == bedrock_unit_test_env["BEDROCK_CHAT_MODEL_ID"]
|
||||
assert bedrock_chat_completion.service_id == bedrock_unit_test_env["BEDROCK_CHAT_MODEL_ID"]
|
||||
|
||||
assert bedrock_chat_completion.bedrock_model_provider == BedrockModelProvider(
|
||||
bedrock_unit_test_env["BEDROCK_MODEL_PROVIDER"]
|
||||
)
|
||||
assert bedrock_chat_completion.bedrock_client is not None
|
||||
assert bedrock_chat_completion.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_chat_completion_init_model_id_override(mock_client, bedrock_unit_test_env, model_id) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service"""
|
||||
bedrock_chat_completion = BedrockChatCompletion(model_id=model_id)
|
||||
|
||||
assert bedrock_chat_completion.ai_model_id == model_id
|
||||
assert bedrock_chat_completion.service_id == model_id
|
||||
|
||||
assert bedrock_chat_completion.bedrock_client is not None
|
||||
assert bedrock_chat_completion.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_chat_completion_init_custom_service_id(mock_client, bedrock_unit_test_env, service_id) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service"""
|
||||
bedrock_chat_completion = BedrockChatCompletion(service_id=service_id)
|
||||
|
||||
assert bedrock_chat_completion.service_id == service_id
|
||||
|
||||
assert bedrock_chat_completion.bedrock_client is not None
|
||||
assert bedrock_chat_completion.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
def test_bedrock_chat_completion_init_custom_clients(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service"""
|
||||
bedrock_chat_completion = BedrockChatCompletion(
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
assert isinstance(bedrock_chat_completion.bedrock_client, MockBedrockClient)
|
||||
assert isinstance(bedrock_chat_completion.bedrock_runtime_client, MockBedrockRuntimeClient)
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_chat_completion_init_custom_client(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service"""
|
||||
bedrock_chat_completion = BedrockChatCompletion(
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
assert isinstance(bedrock_chat_completion.bedrock_client, MockBedrockClient)
|
||||
assert bedrock_chat_completion.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_chat_completion_init_custom_runtime_client(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service"""
|
||||
bedrock_chat_completion = BedrockChatCompletion(
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
)
|
||||
|
||||
assert bedrock_chat_completion.bedrock_client is not None
|
||||
assert isinstance(bedrock_chat_completion.bedrock_runtime_client, MockBedrockRuntimeClient)
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_chat_completion_init_custom_bedrock_model_provider(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service"""
|
||||
bedrock_chat_completion = BedrockChatCompletion(
|
||||
model_provider=BedrockModelProvider.AMAZON,
|
||||
)
|
||||
|
||||
assert bedrock_chat_completion.bedrock_model_provider == BedrockModelProvider.AMAZON
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["BEDROCK_CHAT_MODEL_ID"]], indirect=True)
|
||||
def test_bedrock_chat_completion_client_init_with_empty_model_id(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service with empty model id"""
|
||||
with pytest.raises(ServiceInitializationError, match="The Amazon Bedrock Chat Model ID is missing."):
|
||||
BedrockChatCompletion(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
def test_bedrock_chat_completion_client_init_invalid_settings(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service with invalid settings"""
|
||||
with pytest.raises(
|
||||
ServiceInitializationError, match="Failed to initialize the Amazon Bedrock Chat Completion Service."
|
||||
):
|
||||
BedrockChatCompletion(model_id=123) # Model ID must be a string
|
||||
|
||||
|
||||
def test_bedrock_chat_completion_client_init_invalid_model_provider(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service with invalid settings"""
|
||||
with pytest.raises(
|
||||
ServiceInitializationError, match="Failed to initialize the Amazon Bedrock Chat Completion Service."
|
||||
):
|
||||
BedrockChatCompletion(model_provider="invalid_provider")
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_prompt_execution_settings_class(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test getting prompt execution settings class"""
|
||||
bedrock_completion_client = BedrockChatCompletion()
|
||||
assert bedrock_completion_client.get_prompt_execution_settings_class() == BedrockChatPromptExecutionSettings
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
# region private methods
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_prepare_chat_history_for_request(mock_client, bedrock_unit_test_env, chat_history) -> None:
|
||||
"""Test preparing chat history for request"""
|
||||
bedrock_chat_completion = BedrockChatCompletion()
|
||||
parsed_chat_history = bedrock_chat_completion._prepare_chat_history_for_request(chat_history)
|
||||
|
||||
assert isinstance(parsed_chat_history, list)
|
||||
assert len(parsed_chat_history) == len(chat_history) - 2 # Exclude system message
|
||||
assert all([item["role"] in ["user", "assistant"] for item in parsed_chat_history])
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_prepare_system_message_for_request(mock_client, bedrock_unit_test_env, chat_history) -> None:
|
||||
"""Test preparing system message for request"""
|
||||
bedrock_chat_completion = BedrockChatCompletion()
|
||||
parsed_system_message = bedrock_chat_completion._prepare_system_messages_for_request(chat_history)
|
||||
|
||||
assert isinstance(parsed_system_message, list)
|
||||
assert len(parsed_system_message) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_id",
|
||||
[
|
||||
"amazon.titan",
|
||||
"anthropic.claude",
|
||||
"cohere.command",
|
||||
"ai21.jamba",
|
||||
"meta.llama",
|
||||
"mistral.ai",
|
||||
],
|
||||
)
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_prepare_settings_for_request(mock_client, model_id, chat_history) -> None:
|
||||
"""Test preparing settings for request"""
|
||||
bedrock_chat_completion = BedrockChatCompletion(model_id=model_id)
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
parsed_settings = bedrock_chat_completion._prepare_settings_for_request(chat_history, settings)
|
||||
|
||||
assert isinstance(parsed_settings, dict)
|
||||
assert parsed_settings["modelId"] == bedrock_chat_completion.ai_model_id
|
||||
assert parsed_settings["messages"] == bedrock_chat_completion._prepare_chat_history_for_request(chat_history)
|
||||
assert parsed_settings["system"] == bedrock_chat_completion._prepare_system_messages_for_request(chat_history)
|
||||
assert isinstance(parsed_settings["inferenceConfig"], dict)
|
||||
assert all([parsed_settings["inferenceConfig"].values()])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_id",
|
||||
[
|
||||
"arn:aws:bedrock:us-east-1:972143716085:application-inference-profile/123456",
|
||||
],
|
||||
)
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_prepare_settings_for_request_with_application_inference_profile(mock_client, model_id, chat_history) -> None:
|
||||
"""Test preparing settings for request"""
|
||||
# Without a valid model provider, it should raise an error
|
||||
bedrock_chat_completion = BedrockChatCompletion(model_id=model_id)
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=f"Model ID {model_id} does not contain a valid model provider name.",
|
||||
):
|
||||
bedrock_chat_completion._prepare_settings_for_request(chat_history, settings)
|
||||
|
||||
# With a valid model provider, it should not raise an error
|
||||
bedrock_chat_completion = BedrockChatCompletion(model_id=model_id, model_provider=BedrockModelProvider.AMAZON)
|
||||
parsed_settings = bedrock_chat_completion._prepare_settings_for_request(chat_history, settings)
|
||||
|
||||
assert isinstance(parsed_settings, dict)
|
||||
assert parsed_settings["modelId"] == bedrock_chat_completion.ai_model_id
|
||||
assert parsed_settings["messages"] == bedrock_chat_completion._prepare_chat_history_for_request(chat_history)
|
||||
assert parsed_settings["system"] == bedrock_chat_completion._prepare_system_messages_for_request(chat_history)
|
||||
assert isinstance(parsed_settings["inferenceConfig"], dict)
|
||||
assert all([parsed_settings["inferenceConfig"].values()])
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
# region chat completion
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"amazon.titan",
|
||||
"anthropic.claude",
|
||||
"cohere.command",
|
||||
"ai21.jamba",
|
||||
"meta.llama",
|
||||
"mistral.ai",
|
||||
],
|
||||
)
|
||||
async def test_bedrock_chat_completion(
|
||||
model_id,
|
||||
chat_history: ChatHistory,
|
||||
mock_bedrock_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test Amazon Bedrock Chat Completion complete method"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient, "converse", return_value=mock_bedrock_chat_completion_response
|
||||
) as mock_converse:
|
||||
# Setup
|
||||
bedrock_chat_completion = BedrockChatCompletion(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
response = await bedrock_chat_completion.get_chat_message_contents(chat_history=chat_history, settings=settings)
|
||||
|
||||
# Assert
|
||||
mock_converse.assert_called_once_with(
|
||||
**(bedrock_chat_completion._prepare_settings_for_request(chat_history, settings))
|
||||
)
|
||||
|
||||
assert isinstance(response, list)
|
||||
assert len(response) == 1
|
||||
assert isinstance(response[0], ChatMessageContent)
|
||||
assert response[0].ai_model_id == model_id
|
||||
assert response[0].role == AuthorRole.ASSISTANT
|
||||
assert len(response[0].items) == 1
|
||||
assert isinstance(response[0].items[0], TextContent)
|
||||
assert response[0].finish_reason == FinishReason.STOP
|
||||
assert response[0].metadata["usage"] == CompletionUsage(
|
||||
prompt_tokens=mock_bedrock_chat_completion_response["usage"]["inputTokens"],
|
||||
completion_tokens=mock_bedrock_chat_completion_response["usage"]["outputTokens"],
|
||||
)
|
||||
assert (
|
||||
response[0].items[0].text
|
||||
== mock_bedrock_chat_completion_response["output"]["message"]["content"][0]["text"]
|
||||
)
|
||||
assert response[0].inner_content == mock_bedrock_chat_completion_response
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"amazon.titan",
|
||||
"anthropic.claude",
|
||||
"cohere.command",
|
||||
"ai21.jamba",
|
||||
"meta.llama",
|
||||
"mistral.ai",
|
||||
],
|
||||
)
|
||||
async def test_bedrock_streaming_chat_completion(
|
||||
model_id,
|
||||
chat_history: ChatHistory,
|
||||
mock_bedrock_streaming_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test Amazon Bedrock Streaming Chat Completion complete method"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient, "converse_stream", return_value=mock_bedrock_streaming_chat_completion_response
|
||||
) as mock_converse_stream:
|
||||
# Setup
|
||||
bedrock_chat_completion = BedrockChatCompletion(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
chunks: list[StreamingChatMessageContent] = []
|
||||
async for streaming_messages in bedrock_chat_completion.get_streaming_chat_message_contents(
|
||||
chat_history=chat_history, settings=settings
|
||||
):
|
||||
chunks.extend(streaming_messages)
|
||||
response = reduce(lambda p, r: p + r, chunks)
|
||||
|
||||
# Assert
|
||||
mock_converse_stream.assert_called_once_with(
|
||||
**(bedrock_chat_completion._prepare_settings_for_request(chat_history, settings))
|
||||
)
|
||||
|
||||
assert isinstance(response, StreamingChatMessageContent)
|
||||
assert response.ai_model_id == model_id
|
||||
assert response.role == AuthorRole.ASSISTANT
|
||||
assert len(response.items) == 1
|
||||
assert isinstance(response.inner_content, list)
|
||||
assert len(response.inner_content) == 7
|
||||
assert response.finish_reason == FinishReason.STOP
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"amazon.titan",
|
||||
"anthropic.claude",
|
||||
"cohere.command",
|
||||
"ai21.jamba",
|
||||
"meta.llama",
|
||||
"mistral.ai",
|
||||
],
|
||||
)
|
||||
async def test_bedrock_streaming_chat_completion_invalid_event(
|
||||
model_id,
|
||||
chat_history: ChatHistory,
|
||||
mock_bedrock_streaming_chat_completion_invalid_response,
|
||||
) -> None:
|
||||
"""Test Amazon Bedrock Streaming Chat Completion complete method"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient,
|
||||
"converse_stream",
|
||||
return_value=mock_bedrock_streaming_chat_completion_invalid_response,
|
||||
):
|
||||
# Setup
|
||||
bedrock_chat_completion = BedrockChatCompletion(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
with pytest.raises(ServiceInvalidResponseError):
|
||||
async for chunk in bedrock_chat_completion.get_streaming_chat_message_contents(
|
||||
chat_history=chat_history, settings=settings
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
# endregion
|
||||
+443
@@ -0,0 +1,443 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.ai.bedrock.bedrock_prompt_execution_settings import BedrockChatPromptExecutionSettings
|
||||
from semantic_kernel.connectors.ai.bedrock.services.model_provider.bedrock_model_provider import (
|
||||
BedrockModelProvider,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.bedrock.services.model_provider.utils import (
|
||||
MESSAGE_CONVERTERS,
|
||||
finish_reason_from_bedrock_to_semantic_kernel,
|
||||
remove_none_recursively,
|
||||
update_settings_from_function_choice_configuration,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.function_call_choice_configuration import FunctionCallChoiceConfiguration
|
||||
from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior
|
||||
from semantic_kernel.connectors.ai.function_choice_type import FunctionChoiceType
|
||||
from semantic_kernel.contents.chat_message_content import ChatMessageContent
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
from semantic_kernel.contents.function_result_content import FunctionResultContent
|
||||
from semantic_kernel.contents.image_content import ImageContent
|
||||
from semantic_kernel.contents.text_content import TextContent
|
||||
from semantic_kernel.contents.utils.author_role import AuthorRole
|
||||
from semantic_kernel.contents.utils.finish_reason import FinishReason
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInvalidRequestError
|
||||
from semantic_kernel.kernel import Kernel
|
||||
|
||||
|
||||
def test_remove_none_recursively():
|
||||
data = {
|
||||
"a": 1,
|
||||
"b": None,
|
||||
"c": {
|
||||
"d": 2,
|
||||
"e": None,
|
||||
"f": {
|
||||
"g": 3,
|
||||
"h": None,
|
||||
},
|
||||
},
|
||||
}
|
||||
expected = {
|
||||
"a": 1,
|
||||
"c": {
|
||||
"d": 2,
|
||||
"f": {
|
||||
"g": 3,
|
||||
},
|
||||
},
|
||||
}
|
||||
assert remove_none_recursively(data) == expected
|
||||
|
||||
|
||||
def test_remove_recursively_max_depth():
|
||||
data = {
|
||||
"a": {"b": None},
|
||||
}
|
||||
|
||||
assert remove_none_recursively(data, max_depth=1) == data
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_auto(kernel: Kernel, custom_plugin_class) -> None:
|
||||
kernel.add_plugin(plugin=custom_plugin_class(), plugin_name="custom_plugin")
|
||||
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
|
||||
auto_function_choice_behavior = FunctionChoiceBehavior.Auto()
|
||||
auto_function_choice_behavior.configure(
|
||||
kernel,
|
||||
update_settings_from_function_choice_configuration,
|
||||
settings,
|
||||
)
|
||||
|
||||
assert "auto" in settings.tool_choice
|
||||
assert len(settings.tools) == 1
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_auto_without_plugin(kernel: Kernel) -> None:
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
|
||||
auto_function_choice_behavior = FunctionChoiceBehavior.Auto()
|
||||
auto_function_choice_behavior.configure(
|
||||
kernel,
|
||||
update_settings_from_function_choice_configuration,
|
||||
settings,
|
||||
)
|
||||
|
||||
assert settings.tool_choice is None
|
||||
assert settings.tools is None
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_none(kernel: Kernel) -> None:
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
|
||||
auto_function_choice_behavior = FunctionChoiceBehavior.NoneInvoke()
|
||||
auto_function_choice_behavior.configure(
|
||||
kernel,
|
||||
update_settings_from_function_choice_configuration,
|
||||
settings,
|
||||
)
|
||||
|
||||
assert settings.tool_choice is None
|
||||
assert settings.tools is None
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_required_with_one_function(
|
||||
kernel: Kernel,
|
||||
custom_plugin_class,
|
||||
) -> None:
|
||||
kernel.add_plugin(plugin=custom_plugin_class(), plugin_name="custom_plugin")
|
||||
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
|
||||
auto_function_choice_behavior = FunctionChoiceBehavior.Required()
|
||||
auto_function_choice_behavior.configure(
|
||||
kernel,
|
||||
update_settings_from_function_choice_configuration,
|
||||
settings,
|
||||
)
|
||||
|
||||
assert "tool" in settings.tool_choice
|
||||
assert len(settings.tools) == 1
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_required_with_more_than_one_functions(
|
||||
kernel: Kernel,
|
||||
custom_plugin_class,
|
||||
experimental_plugin_class,
|
||||
) -> None:
|
||||
kernel.add_plugin(plugin=custom_plugin_class(), plugin_name="custom_plugin")
|
||||
kernel.add_plugin(plugin=experimental_plugin_class(), plugin_name="experimental_plugin")
|
||||
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
|
||||
auto_function_choice_behavior = FunctionChoiceBehavior.Required()
|
||||
auto_function_choice_behavior.configure(
|
||||
kernel,
|
||||
update_settings_from_function_choice_configuration,
|
||||
settings,
|
||||
)
|
||||
|
||||
assert "any" in settings.tool_choice
|
||||
assert len(settings.tools) == 2
|
||||
|
||||
|
||||
def test_inference_profile_with_bedrock_model() -> None:
|
||||
"""Test the BedrockModelProvider class returns the correct model for a given inference profile."""
|
||||
|
||||
us_amazon_inference_profile = "us.amazon.nova-lite-v1:0"
|
||||
assert BedrockModelProvider.to_model_provider(us_amazon_inference_profile) == BedrockModelProvider.AMAZON
|
||||
|
||||
us_anthropic_inference_profile = "us.anthropic.claude-3-sonnet-20240229-v1:0"
|
||||
assert BedrockModelProvider.to_model_provider(us_anthropic_inference_profile) == BedrockModelProvider.ANTHROPIC
|
||||
|
||||
eu_meta_inference_profile = "eu.meta.llama3-2-3b-instruct-v1:0"
|
||||
assert BedrockModelProvider.to_model_provider(eu_meta_inference_profile) == BedrockModelProvider.META
|
||||
|
||||
unknown_inference_profile = "unknown"
|
||||
with pytest.raises(ValueError, match="Model ID unknown does not contain a valid model provider name."):
|
||||
BedrockModelProvider.to_model_provider(unknown_inference_profile)
|
||||
|
||||
|
||||
def test_remove_none_recursively_empty_dict() -> None:
|
||||
"""Test that an empty dict returns an empty dict."""
|
||||
assert remove_none_recursively({}) == {}
|
||||
|
||||
|
||||
def test_remove_none_recursively_no_none() -> None:
|
||||
"""Test that a dict with no None values remains the same."""
|
||||
original = {"a": 1, "b": 2}
|
||||
result = remove_none_recursively(original)
|
||||
assert result == {"a": 1, "b": 2}
|
||||
|
||||
|
||||
def test_remove_none_recursively_with_none() -> None:
|
||||
"""Test that dict values of None are removed."""
|
||||
original = {"a": 1, "b": None, "c": {"d": None, "e": 3}}
|
||||
result = remove_none_recursively(original)
|
||||
# 'b' should be removed and 'd' inside nested dict should be removed
|
||||
assert result == {"a": 1, "c": {"e": 3}}
|
||||
|
||||
|
||||
def test_remove_none_recursively_max_depth() -> None:
|
||||
"""Test that the function respects max_depth."""
|
||||
original = {"a": {"b": {"c": None}}}
|
||||
# If max_depth=1, it won't go deep enough to remove 'c'.
|
||||
result = remove_none_recursively(original, max_depth=1)
|
||||
assert result == {"a": {"b": {"c": None}}}
|
||||
|
||||
# If max_depth=3, it should remove 'c'.
|
||||
result = remove_none_recursively(original, max_depth=3)
|
||||
assert result == {"a": {"b": {}}}
|
||||
|
||||
|
||||
def test_format_system_message() -> None:
|
||||
"""Test that system message is formatted correctly."""
|
||||
content = ChatMessageContent(role=AuthorRole.SYSTEM, content="System message")
|
||||
formatted = MESSAGE_CONVERTERS[AuthorRole.SYSTEM](content)
|
||||
assert formatted == {"text": "System message"}
|
||||
|
||||
|
||||
def test_format_user_message_text_only() -> None:
|
||||
"""Test user message with only text content."""
|
||||
text_item = TextContent(text="Hello!")
|
||||
user_message = ChatMessageContent(role=AuthorRole.USER, items=[text_item])
|
||||
|
||||
formatted = MESSAGE_CONVERTERS[AuthorRole.USER](user_message)
|
||||
assert formatted["role"] == "user"
|
||||
assert len(formatted["content"]) == 1
|
||||
assert formatted["content"][0] == {"text": "Hello!"}
|
||||
|
||||
|
||||
def test_format_user_message_image_only() -> None:
|
||||
"""Test user message with only image content."""
|
||||
img_item = ImageContent(data=b"abc", mime_type="image/png")
|
||||
user_message = ChatMessageContent(role=AuthorRole.USER, items=[img_item])
|
||||
|
||||
formatted = MESSAGE_CONVERTERS[AuthorRole.USER](user_message)
|
||||
assert formatted["role"] == "user"
|
||||
assert len(formatted["content"]) == 1
|
||||
image_section = formatted["content"][0].get("image")
|
||||
assert image_section["format"] == "png"
|
||||
assert image_section["source"]["bytes"] == b"abc"
|
||||
|
||||
|
||||
def test_format_user_message_unsupported_content() -> None:
|
||||
"""Test user message raises error with unsupported content type."""
|
||||
# We can simulate an unsupported content type by using FunctionCallContent.
|
||||
func_call_item = FunctionCallContent(id="123", function_name="test_function", arguments="{}")
|
||||
user_message = ChatMessageContent(role=AuthorRole.USER, items=[func_call_item])
|
||||
|
||||
with pytest.raises(ServiceInvalidRequestError) as exc:
|
||||
MESSAGE_CONVERTERS[AuthorRole.USER](user_message)
|
||||
|
||||
assert "Only text and image content are supported in a user message." in str(exc.value)
|
||||
|
||||
|
||||
def test_format_assistant_message_text_content() -> None:
|
||||
"""Test assistant message with text content."""
|
||||
text_item = TextContent(text="Assistant response")
|
||||
assistant_message = ChatMessageContent(role=AuthorRole.ASSISTANT, items=[text_item])
|
||||
|
||||
formatted = MESSAGE_CONVERTERS[AuthorRole.ASSISTANT](assistant_message)
|
||||
assert formatted["role"] == "assistant"
|
||||
assert formatted["content"] == [{"text": "Assistant response"}]
|
||||
|
||||
|
||||
def test_format_assistant_message_function_call_content() -> None:
|
||||
"""Test assistant message with function call content."""
|
||||
func_item = FunctionCallContent(
|
||||
id="fc1", plugin_name="plugin", function_name="function", arguments='{"param": "value"}'
|
||||
)
|
||||
assistant_message = ChatMessageContent(role=AuthorRole.ASSISTANT, items=[func_item])
|
||||
|
||||
formatted = MESSAGE_CONVERTERS[AuthorRole.ASSISTANT](assistant_message)
|
||||
assert len(formatted["content"]) == 1
|
||||
tool_use = formatted["content"][0].get("toolUse")
|
||||
assert tool_use
|
||||
assert tool_use["toolUseId"] == "fc1"
|
||||
assert tool_use["name"] == "plugin-function"
|
||||
assert tool_use["input"] == {"param": "value"}
|
||||
|
||||
|
||||
def test_format_assistant_message_image_content_raises() -> None:
|
||||
"""Test assistant message with image raises error."""
|
||||
img_item = ImageContent(data=b"abc", mime_type="image/jpeg")
|
||||
assistant_message = ChatMessageContent(role=AuthorRole.ASSISTANT, items=[img_item])
|
||||
|
||||
with pytest.raises(ServiceInvalidRequestError) as exc:
|
||||
MESSAGE_CONVERTERS[AuthorRole.ASSISTANT](assistant_message)
|
||||
|
||||
assert "Image content is not supported in an assistant message." in str(exc.value)
|
||||
|
||||
|
||||
def test_format_assistant_message_unsupported_type() -> None:
|
||||
"""Test assistant message with unsupported item content type."""
|
||||
func_res_item = FunctionResultContent(id="res1", function_name="some_function", result="some_result")
|
||||
assistant_message = ChatMessageContent(role=AuthorRole.ASSISTANT, items=[func_res_item])
|
||||
|
||||
with pytest.raises(ServiceInvalidRequestError) as exc:
|
||||
MESSAGE_CONVERTERS[AuthorRole.ASSISTANT](assistant_message)
|
||||
assert "Unsupported content type in an assistant message:" in str(exc.value)
|
||||
|
||||
|
||||
def test_format_tool_message_text() -> None:
|
||||
"""Test tool message with text content."""
|
||||
text_item = TextContent(text="Some text")
|
||||
tool_message = ChatMessageContent(role=AuthorRole.TOOL, items=[text_item])
|
||||
|
||||
formatted = MESSAGE_CONVERTERS[AuthorRole.TOOL](tool_message)
|
||||
assert formatted["role"] == "user" # note that for a tool message, role set to 'user'
|
||||
assert formatted["content"] == [{"text": "Some text"}]
|
||||
|
||||
|
||||
def test_format_tool_message_function_result() -> None:
|
||||
"""Test tool message with function result content."""
|
||||
func_result_item = FunctionResultContent(id="res_id", function_name="test_function", result="some result")
|
||||
tool_message = ChatMessageContent(role=AuthorRole.TOOL, items=[func_result_item])
|
||||
|
||||
formatted = MESSAGE_CONVERTERS[AuthorRole.TOOL](tool_message)
|
||||
assert formatted["role"] == "user"
|
||||
content = formatted["content"][0]
|
||||
assert content.get("toolResult")
|
||||
assert content["toolResult"]["toolUseId"] == "res_id"
|
||||
assert content["toolResult"]["content"] == [{"text": "some result"}]
|
||||
|
||||
|
||||
def test_format_tool_message_image_raises() -> None:
|
||||
"""Test tool message with image content raises an error."""
|
||||
img_item = ImageContent(data=b"xyz", mime_type="image/jpeg")
|
||||
tool_message = ChatMessageContent(role=AuthorRole.TOOL, items=[img_item])
|
||||
|
||||
with pytest.raises(ServiceInvalidRequestError) as exc:
|
||||
MESSAGE_CONVERTERS[AuthorRole.TOOL](tool_message)
|
||||
assert "Image content is not supported in a tool message." in str(exc.value)
|
||||
|
||||
|
||||
def test_finish_reason_from_bedrock_to_semantic_kernel_stop() -> None:
|
||||
"""Test that 'stop_sequence' maps to FinishReason.STOP"""
|
||||
reason = finish_reason_from_bedrock_to_semantic_kernel("stop_sequence")
|
||||
assert reason == FinishReason.STOP
|
||||
|
||||
reason = finish_reason_from_bedrock_to_semantic_kernel("end_turn")
|
||||
assert reason == FinishReason.STOP
|
||||
|
||||
|
||||
def test_finish_reason_from_bedrock_to_semantic_kernel_length() -> None:
|
||||
"""Test that 'max_tokens' maps to FinishReason.LENGTH"""
|
||||
reason = finish_reason_from_bedrock_to_semantic_kernel("max_tokens")
|
||||
assert reason == FinishReason.LENGTH
|
||||
|
||||
|
||||
def test_finish_reason_from_bedrock_to_semantic_kernel_content_filtered() -> None:
|
||||
"""Test that 'content_filtered' maps to FinishReason.CONTENT_FILTER"""
|
||||
reason = finish_reason_from_bedrock_to_semantic_kernel("content_filtered")
|
||||
assert reason == FinishReason.CONTENT_FILTER
|
||||
|
||||
|
||||
def test_finish_reason_from_bedrock_to_semantic_kernel_tool_use() -> None:
|
||||
"""Test that 'tool_use' maps to FinishReason.TOOL_CALLS"""
|
||||
reason = finish_reason_from_bedrock_to_semantic_kernel("tool_use")
|
||||
assert reason == FinishReason.TOOL_CALLS
|
||||
|
||||
|
||||
def test_finish_reason_from_bedrock_to_semantic_kernel_unknown() -> None:
|
||||
"""Test that unknown finish reason returns None"""
|
||||
reason = finish_reason_from_bedrock_to_semantic_kernel("something_unknown")
|
||||
assert reason is None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_bedrock_settings() -> BedrockChatPromptExecutionSettings:
|
||||
"""Helper fixture for BedrockChatPromptExecutionSettings."""
|
||||
return BedrockChatPromptExecutionSettings()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_function_choice_config() -> FunctionCallChoiceConfiguration:
|
||||
"""Helper fixture for a sample FunctionCallChoiceConfiguration."""
|
||||
|
||||
# We'll create mock kernel functions with metadata
|
||||
mock_func_1 = MagicMock()
|
||||
mock_func_1.fully_qualified_name = "plugin-function1"
|
||||
mock_func_1.description = "Function 1 description"
|
||||
|
||||
param1 = MagicMock()
|
||||
param1.name = "param1"
|
||||
param1.schema_data = {"type": "string"}
|
||||
param1.is_required = True
|
||||
|
||||
param2 = MagicMock()
|
||||
param2.name = "param2"
|
||||
param2.schema_data = {"type": "integer"}
|
||||
param2.is_required = False
|
||||
|
||||
mock_func_1.parameters = [
|
||||
param1,
|
||||
param2,
|
||||
]
|
||||
mock_func_2 = MagicMock()
|
||||
mock_func_2.fully_qualified_name = "plugin-function2"
|
||||
mock_func_2.description = "Function 2 description"
|
||||
mock_func_2.parameters = []
|
||||
|
||||
config = FunctionCallChoiceConfiguration()
|
||||
config.available_functions = [mock_func_1, mock_func_2]
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_none_type(
|
||||
mock_function_choice_config, mock_bedrock_settings
|
||||
) -> None:
|
||||
"""Test that if the FunctionChoiceType is NONE it doesn't modify settings."""
|
||||
update_settings_from_function_choice_configuration(
|
||||
mock_function_choice_config, mock_bedrock_settings, FunctionChoiceType.NONE
|
||||
)
|
||||
assert mock_bedrock_settings.tool_choice is None
|
||||
assert mock_bedrock_settings.tools is None
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_auto_two_tools(
|
||||
mock_function_choice_config, mock_bedrock_settings
|
||||
) -> None:
|
||||
"""Test that AUTO sets tool_choice to {"auto": {}} and sets tools list"""
|
||||
update_settings_from_function_choice_configuration(
|
||||
mock_function_choice_config, mock_bedrock_settings, FunctionChoiceType.AUTO
|
||||
)
|
||||
assert mock_bedrock_settings.tool_choice == {"auto": {}}
|
||||
assert len(mock_bedrock_settings.tools) == 2
|
||||
# Validate structure of first tool
|
||||
tool_spec_1 = mock_bedrock_settings.tools[0].get("toolSpec")
|
||||
assert tool_spec_1["name"] == "plugin-function1"
|
||||
assert tool_spec_1["description"] == "Function 1 description"
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_required_many(
|
||||
mock_function_choice_config, mock_bedrock_settings
|
||||
) -> None:
|
||||
"""Test that REQUIRED with more than one function sets tool_choice to {"any": {}}."""
|
||||
update_settings_from_function_choice_configuration(
|
||||
mock_function_choice_config, mock_bedrock_settings, FunctionChoiceType.REQUIRED
|
||||
)
|
||||
assert mock_bedrock_settings.tool_choice == {"any": {}}
|
||||
assert len(mock_bedrock_settings.tools) == 2
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_required_one(mock_bedrock_settings) -> None:
|
||||
"""Test that REQUIRED with a single function picks "tool" with that function name."""
|
||||
single_func = MagicMock()
|
||||
single_func.fully_qualified_name = "plugin-function"
|
||||
single_func.description = "Only function"
|
||||
single_func.parameters = []
|
||||
|
||||
config = FunctionCallChoiceConfiguration()
|
||||
config.available_functions = [single_func]
|
||||
|
||||
update_settings_from_function_choice_configuration(config, mock_bedrock_settings, FunctionChoiceType.REQUIRED)
|
||||
assert mock_bedrock_settings.tool_choice == {"tool": {"name": "plugin-function"}}
|
||||
assert len(mock_bedrock_settings.tools) == 1
|
||||
assert mock_bedrock_settings.tools[0]["toolSpec"]["name"] == "plugin-function"
|
||||
@@ -0,0 +1,324 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import json
|
||||
from functools import reduce
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import boto3
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.ai.bedrock.bedrock_prompt_execution_settings import BedrockTextPromptExecutionSettings
|
||||
from semantic_kernel.connectors.ai.bedrock.services.bedrock_text_completion import BedrockTextCompletion
|
||||
from semantic_kernel.connectors.ai.bedrock.services.model_provider.bedrock_model_provider import (
|
||||
BedrockModelProvider,
|
||||
get_text_completion_request_body,
|
||||
)
|
||||
from semantic_kernel.contents.streaming_text_content import StreamingTextContent
|
||||
from semantic_kernel.contents.text_content import TextContent
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
|
||||
from tests.unit.connectors.ai.bedrock.conftest import MockBedrockClient, MockBedrockRuntimeClient
|
||||
|
||||
# region init
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_completion_init(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service"""
|
||||
bedrock_text_completion = BedrockTextCompletion()
|
||||
|
||||
assert bedrock_text_completion.ai_model_id == bedrock_unit_test_env["BEDROCK_TEXT_MODEL_ID"]
|
||||
assert bedrock_text_completion.service_id == bedrock_unit_test_env["BEDROCK_TEXT_MODEL_ID"]
|
||||
|
||||
assert bedrock_text_completion.bedrock_model_provider == BedrockModelProvider(
|
||||
bedrock_unit_test_env["BEDROCK_MODEL_PROVIDER"]
|
||||
)
|
||||
assert bedrock_text_completion.bedrock_client is not None
|
||||
assert bedrock_text_completion.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_completion_init_model_id_override(mock_client, bedrock_unit_test_env, model_id) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service"""
|
||||
bedrock_text_completion = BedrockTextCompletion(model_id=model_id)
|
||||
|
||||
assert bedrock_text_completion.ai_model_id == model_id
|
||||
assert bedrock_text_completion.service_id == model_id
|
||||
|
||||
assert bedrock_text_completion.bedrock_client is not None
|
||||
assert bedrock_text_completion.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_completion_init_custom_service_id(mock_client, bedrock_unit_test_env, service_id) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service"""
|
||||
bedrock_text_completion = BedrockTextCompletion(service_id=service_id)
|
||||
|
||||
assert bedrock_text_completion.service_id == service_id
|
||||
|
||||
assert bedrock_text_completion.bedrock_client is not None
|
||||
assert bedrock_text_completion.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
def test_bedrock_text_completion_init_custom_clients(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service"""
|
||||
bedrock_text_completion = BedrockTextCompletion(
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
assert isinstance(bedrock_text_completion.bedrock_client, MockBedrockClient)
|
||||
assert isinstance(bedrock_text_completion.bedrock_runtime_client, MockBedrockRuntimeClient)
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_completion_init_custom_client(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service"""
|
||||
bedrock_text_completion = BedrockTextCompletion(
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
assert isinstance(bedrock_text_completion.bedrock_client, MockBedrockClient)
|
||||
assert bedrock_text_completion.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_completion_init_custom_runtime_client(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service"""
|
||||
bedrock_text_completion = BedrockTextCompletion(
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
)
|
||||
|
||||
assert bedrock_text_completion.bedrock_client is not None
|
||||
assert isinstance(bedrock_text_completion.bedrock_runtime_client, MockBedrockRuntimeClient)
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_completion_init_custom_bedrock_model_provider(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service"""
|
||||
bedrock_text_completion = BedrockTextCompletion(
|
||||
model_provider=BedrockModelProvider.AMAZON,
|
||||
)
|
||||
|
||||
assert bedrock_text_completion.bedrock_model_provider == BedrockModelProvider.AMAZON
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["BEDROCK_TEXT_MODEL_ID"]], indirect=True)
|
||||
def test_bedrock_text_completion_client_init_with_empty_model_id(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service with empty model id"""
|
||||
with pytest.raises(ServiceInitializationError, match="The Amazon Bedrock Text Model ID is missing."):
|
||||
BedrockTextCompletion(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
def test_bedrock_text_completion_client_init_invalid_settings(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service with invalid settings"""
|
||||
with pytest.raises(
|
||||
ServiceInitializationError, match="Failed to initialize the Amazon Bedrock Text Completion Service."
|
||||
):
|
||||
BedrockTextCompletion(model_id=123) # Model ID must be a string
|
||||
|
||||
|
||||
def test_bedrock_text_completion_client_init_invalid_model_provider(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service with invalid settings"""
|
||||
with pytest.raises(
|
||||
ServiceInitializationError, match="Failed to initialize the Amazon Bedrock Text Completion Service."
|
||||
):
|
||||
BedrockTextCompletion(model_provider="invalid_provider")
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_prompt_execution_settings_class(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test getting prompt execution settings class"""
|
||||
bedrock_completion_client = BedrockTextCompletion()
|
||||
assert bedrock_completion_client.get_prompt_execution_settings_class() == BedrockTextPromptExecutionSettings
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
# region text completion
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"amazon.titan",
|
||||
"anthropic.claude",
|
||||
"cohere.command",
|
||||
"ai21.jamba",
|
||||
"meta.llama",
|
||||
"mistral.ai",
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_bedrock_text_completion(
|
||||
model_id,
|
||||
mock_bedrock_text_completion_response,
|
||||
output_text,
|
||||
) -> None:
|
||||
"""Test Amazon Bedrock Text Completion complete method"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient, "invoke_model", return_value=mock_bedrock_text_completion_response
|
||||
) as mock_model_invoke:
|
||||
# Setup
|
||||
bedrock_text_completion = BedrockTextCompletion(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockTextPromptExecutionSettings()
|
||||
response = await bedrock_text_completion.get_text_contents("Hello!", settings=settings)
|
||||
|
||||
# Assert
|
||||
mock_model_invoke.assert_called_once_with(
|
||||
body=json.dumps(get_text_completion_request_body(model_id, "Hello!", settings)),
|
||||
modelId=model_id,
|
||||
accept="application/json",
|
||||
contentType="application/json",
|
||||
)
|
||||
|
||||
assert isinstance(response, list)
|
||||
assert len(response) == 1
|
||||
assert isinstance(response[0], TextContent)
|
||||
assert response[0].ai_model_id == model_id
|
||||
assert response[0].text == output_text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"arn:aws:bedrock:us-east-1:972143716085:application-inference-profile/123456",
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_bedrock_text_completion_with_application_inference_profile(
|
||||
model_id,
|
||||
mock_bedrock_text_completion_response,
|
||||
output_text,
|
||||
model_provider,
|
||||
) -> None:
|
||||
"""Test Amazon Bedrock Text Completion complete method"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient,
|
||||
"invoke_model",
|
||||
return_value=mock_bedrock_text_completion_response,
|
||||
) as mock_model_invoke:
|
||||
# Setup
|
||||
bedrock_text_completion = BedrockTextCompletion(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
model_provider=model_provider,
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockTextPromptExecutionSettings()
|
||||
await bedrock_text_completion.get_text_contents("Hello!", settings=settings)
|
||||
|
||||
# Assert
|
||||
mock_model_invoke.assert_called_once_with(
|
||||
body=json.dumps(get_text_completion_request_body(model_id, "Hello!", settings, model_provider)),
|
||||
modelId=model_id,
|
||||
accept="application/json",
|
||||
contentType="application/json",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"amazon.titan",
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_bedrock_streaming_text_completion(
|
||||
model_id,
|
||||
mock_bedrock_streaming_text_completion_response,
|
||||
output_text,
|
||||
) -> None:
|
||||
"""Test Amazon Bedrock Text Completion complete method"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient,
|
||||
"invoke_model_with_response_stream",
|
||||
return_value=mock_bedrock_streaming_text_completion_response,
|
||||
) as mock_invoke_model_with_response_stream:
|
||||
# Setup
|
||||
bedrock_text_completion = BedrockTextCompletion(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockTextPromptExecutionSettings()
|
||||
chunks: list[StreamingTextContent] = []
|
||||
async for streaming_responses in bedrock_text_completion.get_streaming_text_contents(
|
||||
"Hello!", settings=settings
|
||||
):
|
||||
chunks.extend(streaming_responses)
|
||||
response = reduce(lambda p, r: p + r, chunks)
|
||||
|
||||
# Assert
|
||||
mock_invoke_model_with_response_stream.assert_called_once_with(
|
||||
body=json.dumps(get_text_completion_request_body(model_id, "Hello!", settings)),
|
||||
modelId=model_id,
|
||||
accept="application/json",
|
||||
contentType="application/json",
|
||||
)
|
||||
assert isinstance(response, StreamingTextContent)
|
||||
assert response.ai_model_id == model_id
|
||||
assert response.text == output_text
|
||||
assert response.choice_index == 0
|
||||
assert isinstance(response.inner_content, list)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"arn:aws:bedrock:us-east-1:972143716085:application-inference-profile/123456",
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_bedrock_streaming_text_completion_with_application_inference_profile(
|
||||
model_id,
|
||||
mock_bedrock_streaming_text_completion_response,
|
||||
output_text,
|
||||
model_provider,
|
||||
) -> None:
|
||||
"""Test Amazon Bedrock Chat Completion complete method"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient,
|
||||
"invoke_model_with_response_stream",
|
||||
return_value=mock_bedrock_streaming_text_completion_response,
|
||||
) as mock_invoke_model_with_response_stream:
|
||||
# Setup
|
||||
bedrock_text_completion = BedrockTextCompletion(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
model_provider=model_provider,
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockTextPromptExecutionSettings()
|
||||
chunks: list[StreamingTextContent] = []
|
||||
async for streaming_responses in bedrock_text_completion.get_streaming_text_contents(
|
||||
"Hello!", settings=settings
|
||||
):
|
||||
chunks.extend(streaming_responses)
|
||||
|
||||
# Assert
|
||||
mock_invoke_model_with_response_stream.assert_called_once_with(
|
||||
body=json.dumps(get_text_completion_request_body(model_id, "Hello!", settings, model_provider)),
|
||||
modelId=model_id,
|
||||
accept="application/json",
|
||||
contentType="application/json",
|
||||
)
|
||||
|
||||
|
||||
# endregion
|
||||
+235
@@ -0,0 +1,235 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import ANY, Mock, patch
|
||||
|
||||
import boto3
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.ai.bedrock.bedrock_prompt_execution_settings import (
|
||||
BedrockEmbeddingPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.bedrock.services.bedrock_text_embedding import BedrockTextEmbedding
|
||||
from semantic_kernel.connectors.ai.bedrock.services.model_provider.bedrock_model_provider import BedrockModelProvider
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidResponseError
|
||||
from tests.unit.connectors.ai.bedrock.conftest import MockBedrockClient, MockBedrockRuntimeClient
|
||||
|
||||
# region init
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_embedding_init(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service"""
|
||||
bedrock_text_embedding = BedrockTextEmbedding()
|
||||
|
||||
assert bedrock_text_embedding.ai_model_id == bedrock_unit_test_env["BEDROCK_EMBEDDING_MODEL_ID"]
|
||||
assert bedrock_text_embedding.service_id == bedrock_unit_test_env["BEDROCK_EMBEDDING_MODEL_ID"]
|
||||
|
||||
assert bedrock_text_embedding.bedrock_model_provider == BedrockModelProvider(
|
||||
bedrock_unit_test_env["BEDROCK_MODEL_PROVIDER"]
|
||||
)
|
||||
assert bedrock_text_embedding.bedrock_client is not None
|
||||
assert bedrock_text_embedding.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_embedding_init_model_id_override(mock_client, bedrock_unit_test_env, model_id) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service"""
|
||||
bedrock_text_embedding = BedrockTextEmbedding(model_id=model_id)
|
||||
|
||||
assert bedrock_text_embedding.ai_model_id == model_id
|
||||
assert bedrock_text_embedding.service_id == model_id
|
||||
|
||||
assert bedrock_text_embedding.bedrock_client is not None
|
||||
assert bedrock_text_embedding.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_embedding_init_custom_service_id(mock_client, bedrock_unit_test_env, service_id) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service"""
|
||||
bedrock_text_embedding = BedrockTextEmbedding(service_id=service_id)
|
||||
|
||||
assert bedrock_text_embedding.service_id == service_id
|
||||
|
||||
assert bedrock_text_embedding.bedrock_client is not None
|
||||
assert bedrock_text_embedding.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
def test_bedrock_text_embedding_init_custom_clients(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service"""
|
||||
bedrock_text_embedding = BedrockTextEmbedding(
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
assert isinstance(bedrock_text_embedding.bedrock_client, MockBedrockClient)
|
||||
assert isinstance(bedrock_text_embedding.bedrock_runtime_client, MockBedrockRuntimeClient)
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_embedding_init_custom_client(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service"""
|
||||
bedrock_text_embedding = BedrockTextEmbedding(
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
assert isinstance(bedrock_text_embedding.bedrock_client, MockBedrockClient)
|
||||
assert bedrock_text_embedding.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_embedding_init_custom_runtime_client(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service"""
|
||||
bedrock_text_embedding = BedrockTextEmbedding(
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
)
|
||||
|
||||
assert bedrock_text_embedding.bedrock_client is not None
|
||||
assert isinstance(bedrock_text_embedding.bedrock_runtime_client, MockBedrockRuntimeClient)
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_embedding_init_custom_bedrock_model_provider(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service"""
|
||||
bedrock_text_embedding = BedrockTextEmbedding(
|
||||
model_provider=BedrockModelProvider.AMAZON,
|
||||
)
|
||||
|
||||
assert bedrock_text_embedding.bedrock_model_provider == BedrockModelProvider.AMAZON
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["BEDROCK_EMBEDDING_MODEL_ID"]], indirect=True)
|
||||
def test_bedrock_text_embedding_client_init_with_empty_model_id(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service with empty model id"""
|
||||
with pytest.raises(ServiceInitializationError, match="The Amazon Bedrock Text Embedding Model ID is missing."):
|
||||
BedrockTextEmbedding(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
def test_bedrock_text_embedding_client_init_invalid_settings(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service with invalid settings"""
|
||||
with pytest.raises(
|
||||
ServiceInitializationError, match="Failed to initialize the Amazon Bedrock Text Embedding Service."
|
||||
):
|
||||
BedrockTextEmbedding(model_id=123) # Model ID must be a string
|
||||
|
||||
|
||||
def test_bedrock_text_embedding_client_init_invalid_model_provider(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service with invalid settings"""
|
||||
with pytest.raises(
|
||||
ServiceInitializationError, match="Failed to initialize the Amazon Bedrock Text Embedding Service."
|
||||
):
|
||||
BedrockTextEmbedding(model_provider="invalid_provider")
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_prompt_execution_settings_class(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test getting prompt execution settings class"""
|
||||
bedrock_completion_client = BedrockTextEmbedding()
|
||||
assert bedrock_completion_client.get_prompt_execution_settings_class() == BedrockEmbeddingPromptExecutionSettings
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"amazon.titan",
|
||||
"cohere.command",
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_bedrock_text_embedding(model_id, mock_bedrock_text_embedding_response) -> None:
|
||||
"""Test Bedrock text embedding generation"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient, "invoke_model", return_value=mock_bedrock_text_embedding_response
|
||||
) as mock_model_invoke:
|
||||
# Setup
|
||||
bedrock_text_embedding = BedrockTextEmbedding(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockEmbeddingPromptExecutionSettings()
|
||||
response = await bedrock_text_embedding.generate_embeddings(["hello", "world"], settings)
|
||||
|
||||
# Assert
|
||||
mock_model_invoke.assert_called_with(
|
||||
body=ANY,
|
||||
modelId=model_id,
|
||||
accept="application/json",
|
||||
contentType="application/json",
|
||||
)
|
||||
assert mock_model_invoke.call_count == 2
|
||||
|
||||
assert len(response) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"arn:aws:bedrock:us-east-1:972143716085:application-inference-profile/123456",
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_bedrock_text_embedding_with_application_inference_profile(
|
||||
model_id,
|
||||
mock_bedrock_text_embedding_response,
|
||||
model_provider,
|
||||
) -> None:
|
||||
"""Test Bedrock text embedding generation"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient, "invoke_model", return_value=mock_bedrock_text_embedding_response
|
||||
) as mock_model_invoke:
|
||||
# Setup
|
||||
bedrock_text_embedding = BedrockTextEmbedding(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
model_provider=BedrockModelProvider.AMAZON,
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockEmbeddingPromptExecutionSettings()
|
||||
response = await bedrock_text_embedding.generate_embeddings(["hello", "world"], settings)
|
||||
|
||||
# Assert
|
||||
mock_model_invoke.assert_called_with(
|
||||
body=ANY,
|
||||
modelId=model_id,
|
||||
accept="application/json",
|
||||
contentType="application/json",
|
||||
)
|
||||
assert mock_model_invoke.call_count == 2
|
||||
|
||||
assert len(response) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"amazon.titan",
|
||||
"cohere.command",
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_bedrock_text_embedding_with_invalid_response(
|
||||
model_id, mock_bedrock_text_embedding_invalid_response
|
||||
) -> None:
|
||||
"""Test Bedrock text embedding generation with invalid response"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient, "invoke_model", return_value=mock_bedrock_text_embedding_invalid_response
|
||||
):
|
||||
# Setup
|
||||
bedrock_text_embedding = BedrockTextEmbedding(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
with pytest.raises(ServiceInvalidResponseError):
|
||||
await bedrock_text_embedding.generate_embeddings(["hello", "world"])
|
||||
@@ -0,0 +1,104 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from semantic_kernel.connectors.ai.bedrock import BedrockChatPromptExecutionSettings, BedrockPromptExecutionSettings
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
|
||||
|
||||
def test_default_bedrock_prompt_execution_settings():
|
||||
settings = BedrockPromptExecutionSettings()
|
||||
|
||||
assert settings.temperature is None
|
||||
assert settings.top_p is None
|
||||
assert settings.top_k is None
|
||||
assert settings.max_tokens is None
|
||||
assert settings.stop == []
|
||||
|
||||
|
||||
def test_custom_bedrock_prompt_execution_settings():
|
||||
settings = BedrockPromptExecutionSettings(
|
||||
temperature=0.5,
|
||||
top_p=0.5,
|
||||
top_k=10,
|
||||
max_tokens=128,
|
||||
stop=["world"],
|
||||
)
|
||||
|
||||
assert settings.temperature == 0.5
|
||||
assert settings.top_p == 0.5
|
||||
assert settings.top_k == 10
|
||||
assert settings.max_tokens == 128
|
||||
assert settings.stop == ["world"]
|
||||
|
||||
|
||||
def test_bedrock_prompt_execution_settings_from_default_completion_config():
|
||||
settings = PromptExecutionSettings(service_id="test_service")
|
||||
chat_settings = BedrockChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
|
||||
assert chat_settings.service_id == "test_service"
|
||||
assert chat_settings.temperature is None
|
||||
assert chat_settings.top_p is None
|
||||
assert chat_settings.top_k is None
|
||||
assert chat_settings.max_tokens is None
|
||||
assert chat_settings.stop == []
|
||||
|
||||
|
||||
def test_bedrock_prompt_execution_settings_from_openai_prompt_execution_settings():
|
||||
chat_settings = BedrockChatPromptExecutionSettings(service_id="test_service", temperature=1.0)
|
||||
new_settings = BedrockPromptExecutionSettings(service_id="test_2", temperature=0.0)
|
||||
chat_settings.update_from_prompt_execution_settings(new_settings)
|
||||
|
||||
assert chat_settings.service_id == "test_2"
|
||||
assert chat_settings.temperature == 0.0
|
||||
|
||||
|
||||
def test_bedrock_prompt_execution_settings_from_custom_completion_config():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"top_k": 10,
|
||||
"max_tokens": 128,
|
||||
"stop": ["world"],
|
||||
},
|
||||
)
|
||||
chat_settings = BedrockChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
|
||||
assert chat_settings.temperature == 0.5
|
||||
assert chat_settings.top_p == 0.5
|
||||
assert chat_settings.top_k == 10
|
||||
assert chat_settings.max_tokens == 128
|
||||
assert chat_settings.stop == ["world"]
|
||||
|
||||
|
||||
def test_bedrock_chat_prompt_execution_settings_from_custom_completion_config_with_functions():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"tools": [{"function": {}}],
|
||||
},
|
||||
)
|
||||
chat_settings = BedrockChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
|
||||
assert chat_settings.tools == [{"function": {}}]
|
||||
|
||||
|
||||
def test_create_options():
|
||||
settings = BedrockPromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"top_k": 10,
|
||||
"max_tokens": 128,
|
||||
"stop": ["world"],
|
||||
},
|
||||
)
|
||||
options = settings.prepare_settings_dict()
|
||||
|
||||
assert options["temperature"] == 0.5
|
||||
assert options["top_p"] == 0.5
|
||||
assert options["top_k"] == 10
|
||||
assert options["max_tokens"] == 128
|
||||
assert options["stop"] == ["world"]
|
||||
@@ -0,0 +1,23 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.contents.chat_history import ChatHistory
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def service_id() -> str:
|
||||
return "test_service_id"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def chat_history() -> ChatHistory:
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_system_message("system_prompt")
|
||||
chat_history.add_user_message("test_prompt")
|
||||
return chat_history
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def prompt() -> str:
|
||||
return "test_prompt"
|
||||
@@ -0,0 +1,174 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from collections.abc import AsyncGenerator, AsyncIterator
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from google.genai.types import (
|
||||
Candidate,
|
||||
Content,
|
||||
FinishReason,
|
||||
GenerateContentResponse,
|
||||
GenerateContentResponseUsageMetadata,
|
||||
Part,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def google_ai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for Google AI Unit Tests."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"GOOGLE_AI_GEMINI_MODEL_ID": "test-gemini-model-id",
|
||||
"GOOGLE_AI_EMBEDDING_MODEL_ID": "test-embedding-model-id",
|
||||
"GOOGLE_AI_API_KEY": "test-api-key",
|
||||
"GOOGLE_AI_CLOUD_PROJECT_ID": "test-project-id",
|
||||
"GOOGLE_AI_CLOUD_REGION": "test-region",
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_google_ai_chat_completion_response() -> GenerateContentResponse:
|
||||
"""Mock Google AI Chat Completion response."""
|
||||
candidate = Candidate()
|
||||
candidate.index = 0
|
||||
candidate.content = Content(role="user", parts=[Part.from_text(text="Test content")])
|
||||
candidate.finish_reason = FinishReason.STOP
|
||||
|
||||
response = GenerateContentResponse()
|
||||
response.candidates = [candidate]
|
||||
response.usage_metadata = GenerateContentResponseUsageMetadata(
|
||||
prompt_token_count=0,
|
||||
cached_content_token_count=0,
|
||||
candidates_token_count=0,
|
||||
total_token_count=0,
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_google_ai_chat_completion_response_with_tool_call() -> GenerateContentResponse:
|
||||
"""Mock Google AI Chat Completion response."""
|
||||
candidate = Candidate()
|
||||
candidate.index = 0
|
||||
candidate.content = Content(
|
||||
role="user",
|
||||
parts=[
|
||||
Part.from_function_call(
|
||||
name="test_function",
|
||||
args={"test_arg": "test_value"},
|
||||
)
|
||||
],
|
||||
)
|
||||
candidate.finish_reason = FinishReason.STOP
|
||||
|
||||
response = GenerateContentResponse()
|
||||
response.candidates = [candidate]
|
||||
response.usage_metadata = GenerateContentResponseUsageMetadata(
|
||||
prompt_token_count=0,
|
||||
cached_content_token_count=0,
|
||||
candidates_token_count=0,
|
||||
total_token_count=0,
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def mock_google_ai_streaming_chat_completion_response() -> AsyncIterator[GenerateContentResponse]:
|
||||
"""Mock Google AI streaming Chat Completion response."""
|
||||
candidate = Candidate()
|
||||
candidate.index = 0
|
||||
candidate.content = Content(role="user", parts=[Part.from_text(text="Test content")])
|
||||
candidate.finish_reason = FinishReason.STOP
|
||||
|
||||
response = GenerateContentResponse()
|
||||
response.candidates = [candidate]
|
||||
response.usage_metadata = GenerateContentResponseUsageMetadata(
|
||||
prompt_token_count=0,
|
||||
cached_content_token_count=0,
|
||||
candidates_token_count=0,
|
||||
total_token_count=0,
|
||||
)
|
||||
|
||||
iterable = MagicMock(spec=AsyncGenerator)
|
||||
iterable.__aiter__.return_value = [response]
|
||||
|
||||
return iterable
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def mock_google_ai_streaming_chat_completion_response_with_tool_call() -> AsyncIterator[GenerateContentResponse]:
|
||||
"""Mock Google AI streaming Chat Completion response with tool call."""
|
||||
candidate = Candidate()
|
||||
candidate.index = 0
|
||||
candidate.content = Content(
|
||||
role="user",
|
||||
parts=[
|
||||
Part.from_function_call(
|
||||
name="getLightStatus",
|
||||
args={"arg1": "test_value"},
|
||||
)
|
||||
],
|
||||
)
|
||||
candidate.finish_reason = FinishReason.STOP
|
||||
|
||||
response = GenerateContentResponse()
|
||||
response.candidates = [candidate]
|
||||
response.usage_metadata = GenerateContentResponseUsageMetadata(
|
||||
prompt_token_count=0,
|
||||
cached_content_token_count=0,
|
||||
candidates_token_count=0,
|
||||
total_token_count=0,
|
||||
)
|
||||
|
||||
iterable = MagicMock(spec=AsyncGenerator)
|
||||
iterable.__aiter__.return_value = [response]
|
||||
|
||||
return iterable
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_google_ai_text_completion_response() -> GenerateContentResponse:
|
||||
"""Mock Google AI Text Completion response."""
|
||||
candidate = Candidate()
|
||||
candidate.index = 0
|
||||
candidate.content = Content(parts=[Part.from_text(text="Test content")])
|
||||
|
||||
response = GenerateContentResponse()
|
||||
response.candidates = [candidate]
|
||||
|
||||
return response
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def mock_google_ai_streaming_text_completion_response() -> AsyncIterator[GenerateContentResponse]:
|
||||
"""Mock Google AI streaming Text Completion response."""
|
||||
candidate = Candidate()
|
||||
candidate.index = 0
|
||||
candidate.content = Content(parts=[Part.from_text(text="Test content")])
|
||||
|
||||
response = GenerateContentResponse()
|
||||
response.candidates = [candidate]
|
||||
|
||||
iterable = MagicMock(spec=AsyncGenerator)
|
||||
iterable.__aiter__.return_value = [response]
|
||||
|
||||
return iterable
|
||||
+650
@@ -0,0 +1,650 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from google.genai import Client
|
||||
from google.genai.models import AsyncModels
|
||||
from google.genai.types import Content, GenerateContentConfigDict
|
||||
|
||||
from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior
|
||||
from semantic_kernel.connectors.ai.google.google_ai.google_ai_prompt_execution_settings import (
|
||||
GoogleAIChatPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.google.google_ai.google_ai_settings import GoogleAISettings
|
||||
from semantic_kernel.connectors.ai.google.google_ai.services.google_ai_chat_completion import GoogleAIChatCompletion
|
||||
from semantic_kernel.connectors.ai.google.shared_utils import filter_system_message
|
||||
from semantic_kernel.contents.chat_history import ChatHistory
|
||||
from semantic_kernel.contents.chat_message_content import ChatMessageContent
|
||||
from semantic_kernel.contents.utils.finish_reason import FinishReason
|
||||
from semantic_kernel.exceptions.service_exceptions import (
|
||||
ServiceInitializationError,
|
||||
ServiceInvalidExecutionSettingsError,
|
||||
)
|
||||
from semantic_kernel.kernel import Kernel
|
||||
|
||||
|
||||
# region init
|
||||
def test_google_ai_chat_completion_init(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAIChatCompletion"""
|
||||
model_id = google_ai_unit_test_env["GOOGLE_AI_GEMINI_MODEL_ID"]
|
||||
api_key = google_ai_unit_test_env["GOOGLE_AI_API_KEY"]
|
||||
google_ai_chat_completion = GoogleAIChatCompletion()
|
||||
|
||||
assert google_ai_chat_completion.ai_model_id == model_id
|
||||
assert google_ai_chat_completion.service_id == model_id
|
||||
|
||||
assert isinstance(google_ai_chat_completion.service_settings, GoogleAISettings)
|
||||
assert google_ai_chat_completion.service_settings.gemini_model_id == model_id
|
||||
assert google_ai_chat_completion.service_settings.api_key.get_secret_value() == api_key
|
||||
|
||||
|
||||
def test_google_ai_chat_completion_init_with_service_id(google_ai_unit_test_env, service_id) -> None:
|
||||
"""Test initialization of GoogleAIChatCompletion with a service_id that is not the model_id"""
|
||||
google_ai_chat_completion = GoogleAIChatCompletion(service_id=service_id)
|
||||
|
||||
assert google_ai_chat_completion.service_id == service_id
|
||||
|
||||
|
||||
def test_google_ai_chat_completion_init_with_model_id_in_argument(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAIChatCompletion with model_id in argument"""
|
||||
google_ai_chat_completion = GoogleAIChatCompletion(gemini_model_id="custom_model_id")
|
||||
|
||||
assert google_ai_chat_completion.ai_model_id == "custom_model_id"
|
||||
assert google_ai_chat_completion.service_id == "custom_model_id"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["GOOGLE_AI_GEMINI_MODEL_ID"]], indirect=True)
|
||||
def test_google_ai_chat_completion_init_with_empty_model_id(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAIChatCompletion with an empty model_id"""
|
||||
with pytest.raises(ServiceInitializationError, match="The Google AI Gemini model ID is required."):
|
||||
GoogleAIChatCompletion(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["GOOGLE_AI_API_KEY"]], indirect=True)
|
||||
def test_google_ai_chat_completion_init_with_empty_api_key(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAIChatCompletion with an empty api_key"""
|
||||
with pytest.raises(ServiceInitializationError, match="API key is required when use_vertexai is False."):
|
||||
GoogleAIChatCompletion(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["GOOGLE_AI_CLOUD_PROJECT_ID"]], indirect=True)
|
||||
def test_google_ai_chat_completion_init_with_use_vertexai_missing_project_id(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAIChatCompletion with use_vertexai true but missing project_id"""
|
||||
with pytest.raises(ServiceInitializationError, match="Project ID must be provided when use_vertexai is True."):
|
||||
GoogleAIChatCompletion(use_vertexai=True, env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["GOOGLE_AI_CLOUD_REGION"]], indirect=True)
|
||||
def test_google_ai_chat_completion_init_with_use_vertexai_missing_region(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAIChatCompletion with use_vertexai true but missing region"""
|
||||
with pytest.raises(ServiceInitializationError, match="Region must be provided when use_vertexai is True."):
|
||||
GoogleAIChatCompletion(use_vertexai=True, env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["GOOGLE_AI_API_KEY"]], indirect=True)
|
||||
def test_google_ai_chat_completion_init_with_use_vertexai_no_api_key(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAIChatCompletion succeeds with use_vertexai=True and no api_key"""
|
||||
chat_completion = GoogleAIChatCompletion(use_vertexai=True)
|
||||
assert chat_completion.service_settings.use_vertexai is True
|
||||
|
||||
|
||||
def test_prompt_execution_settings_class(google_ai_unit_test_env) -> None:
|
||||
google_ai_chat_completion = GoogleAIChatCompletion()
|
||||
assert google_ai_chat_completion.get_prompt_execution_settings_class() == GoogleAIChatPromptExecutionSettings
|
||||
|
||||
|
||||
# endregion init
|
||||
|
||||
|
||||
# region chat completion
|
||||
|
||||
|
||||
@patch.object(AsyncModels, "generate_content", new_callable=AsyncMock)
|
||||
async def test_google_ai_chat_completion(
|
||||
mock_google_ai_model_generate_content,
|
||||
google_ai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_google_ai_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test chat completion with GoogleAIChatCompletion"""
|
||||
settings = GoogleAIChatPromptExecutionSettings(top_k=5, temperature=0.7)
|
||||
|
||||
mock_google_ai_model_generate_content.return_value = mock_google_ai_chat_completion_response
|
||||
|
||||
google_ai_chat_completion = GoogleAIChatCompletion()
|
||||
responses: list[ChatMessageContent] = await google_ai_chat_completion.get_chat_message_contents(
|
||||
chat_history, settings
|
||||
)
|
||||
|
||||
mock_google_ai_model_generate_content.assert_called_once_with(
|
||||
model=google_ai_chat_completion.service_settings.gemini_model_id,
|
||||
contents=google_ai_chat_completion._prepare_chat_history_for_request(chat_history),
|
||||
config=GenerateContentConfigDict(
|
||||
**settings.prepare_settings_dict(),
|
||||
system_instruction=filter_system_message(chat_history),
|
||||
),
|
||||
)
|
||||
assert len(responses) == 1
|
||||
assert responses[0].role == "assistant"
|
||||
assert responses[0].content == mock_google_ai_chat_completion_response.candidates[0].content.parts[0].text
|
||||
assert responses[0].finish_reason == FinishReason.STOP
|
||||
assert "usage" in responses[0].metadata
|
||||
assert "prompt_feedback" in responses[0].metadata
|
||||
assert responses[0].inner_content == mock_google_ai_chat_completion_response
|
||||
|
||||
|
||||
async def test_google_ai_chat_completion_with_custom_client(
|
||||
chat_history: ChatHistory,
|
||||
google_ai_unit_test_env,
|
||||
mock_google_ai_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test chat completion with GoogleAIChatCompletion using a custom client"""
|
||||
# Create a custom client with a fake API key for testing
|
||||
custom_client = Client(api_key="fake-api-key-for-testing")
|
||||
|
||||
# Mock the custom client's generate_content method
|
||||
mock_generate_content = AsyncMock(return_value=mock_google_ai_chat_completion_response)
|
||||
custom_client.aio.models.generate_content = mock_generate_content
|
||||
|
||||
google_ai_chat_completion = GoogleAIChatCompletion(client=custom_client)
|
||||
responses: list[ChatMessageContent] = await google_ai_chat_completion.get_chat_message_contents(
|
||||
chat_history,
|
||||
GoogleAIChatPromptExecutionSettings(),
|
||||
)
|
||||
|
||||
custom_client.close()
|
||||
|
||||
# Verify that the custom client was used and returned the expected response
|
||||
assert len(responses) == 1
|
||||
assert responses[0].role == "assistant"
|
||||
assert responses[0].content == mock_google_ai_chat_completion_response.candidates[0].content.parts[0].text
|
||||
assert responses[0].finish_reason == FinishReason.STOP
|
||||
|
||||
# Verify that the custom client's method was called
|
||||
mock_generate_content.assert_called_once()
|
||||
|
||||
|
||||
async def test_google_ai_chat_completion_with_function_choice_behavior_fail_verification(
|
||||
chat_history: ChatHistory,
|
||||
google_ai_unit_test_env,
|
||||
) -> None:
|
||||
"""Test completion of GoogleAIChatCompletion with function choice behavior expect verification failure"""
|
||||
|
||||
# Missing kernel
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
settings = GoogleAIChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
|
||||
google_ai_chat_completion = GoogleAIChatCompletion()
|
||||
|
||||
await google_ai_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=settings,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncModels, "generate_content", new_callable=AsyncMock)
|
||||
async def test_google_ai_chat_completion_with_function_choice_behavior(
|
||||
mock_google_ai_model_generate_content,
|
||||
google_ai_unit_test_env,
|
||||
kernel,
|
||||
chat_history: ChatHistory,
|
||||
mock_google_ai_chat_completion_response_with_tool_call,
|
||||
) -> None:
|
||||
"""Test completion of GoogleAIChatCompletion with function choice behavior"""
|
||||
mock_google_ai_model_generate_content.return_value = mock_google_ai_chat_completion_response_with_tool_call
|
||||
|
||||
settings = GoogleAIChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
settings.function_choice_behavior.maximum_auto_invoke_attempts = 1
|
||||
|
||||
google_ai_chat_completion = GoogleAIChatCompletion()
|
||||
|
||||
responses = await google_ai_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=settings,
|
||||
kernel=kernel,
|
||||
)
|
||||
|
||||
# The function should be called twice:
|
||||
# One for the tool call and one for the last completion
|
||||
# after the maximum_auto_invoke_attempts is reached
|
||||
assert mock_google_ai_model_generate_content.call_count == 2
|
||||
assert len(responses) == 1
|
||||
assert responses[0].role == "assistant"
|
||||
# Google doesn't return STOP as the finish reason for tool calls
|
||||
assert responses[0].finish_reason == FinishReason.STOP
|
||||
|
||||
|
||||
@patch.object(AsyncModels, "generate_content", new_callable=AsyncMock)
|
||||
async def test_google_ai_chat_completion_with_function_choice_behavior_no_tool_call(
|
||||
mock_google_ai_model_generate_content,
|
||||
google_ai_unit_test_env,
|
||||
kernel,
|
||||
chat_history: ChatHistory,
|
||||
mock_google_ai_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test completion of GoogleAIChatCompletion with function choice behavior but no tool call returned"""
|
||||
mock_google_ai_model_generate_content.return_value = mock_google_ai_chat_completion_response
|
||||
|
||||
settings = GoogleAIChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
settings.function_choice_behavior.maximum_auto_invoke_attempts = 1
|
||||
|
||||
google_ai_chat_completion = GoogleAIChatCompletion()
|
||||
|
||||
responses = await google_ai_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=settings,
|
||||
kernel=kernel,
|
||||
)
|
||||
|
||||
mock_google_ai_model_generate_content.assert_awaited_once_with(
|
||||
model=google_ai_chat_completion.service_settings.gemini_model_id,
|
||||
contents=google_ai_chat_completion._prepare_chat_history_for_request(chat_history),
|
||||
config=GenerateContentConfigDict(
|
||||
**settings.prepare_settings_dict(),
|
||||
system_instruction=filter_system_message(chat_history),
|
||||
),
|
||||
)
|
||||
assert len(responses) == 1
|
||||
assert responses[0].role == "assistant"
|
||||
assert responses[0].content == mock_google_ai_chat_completion_response.candidates[0].content.parts[0].text
|
||||
|
||||
|
||||
# endregion chat completion
|
||||
|
||||
|
||||
# region streaming chat completion
|
||||
|
||||
|
||||
@patch.object(AsyncModels, "generate_content_stream", new_callable=AsyncMock)
|
||||
async def test_google_ai_streaming_chat_completion(
|
||||
mock_google_ai_model_generate_content_stream,
|
||||
google_ai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_google_ai_streaming_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test streaming chat completion with GoogleAIChatCompletion"""
|
||||
settings = GoogleAIChatPromptExecutionSettings()
|
||||
|
||||
mock_google_ai_model_generate_content_stream.return_value = mock_google_ai_streaming_chat_completion_response
|
||||
|
||||
google_ai_chat_completion = GoogleAIChatCompletion()
|
||||
async for messages in google_ai_chat_completion.get_streaming_chat_message_contents(chat_history, settings):
|
||||
assert len(messages) == 1
|
||||
assert messages[0].role == "assistant"
|
||||
assert messages[0].finish_reason == FinishReason.STOP
|
||||
assert "usage" in messages[0].metadata
|
||||
assert "prompt_feedback" in messages[0].metadata
|
||||
|
||||
mock_google_ai_model_generate_content_stream.assert_called_once_with(
|
||||
model=google_ai_chat_completion.service_settings.gemini_model_id,
|
||||
contents=google_ai_chat_completion._prepare_chat_history_for_request(chat_history),
|
||||
config=GenerateContentConfigDict(
|
||||
**settings.prepare_settings_dict(),
|
||||
system_instruction=filter_system_message(chat_history),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def test_google_ai_streaming_chat_completion_with_custom_client(
|
||||
chat_history: ChatHistory,
|
||||
google_ai_unit_test_env,
|
||||
mock_google_ai_streaming_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test streaming chat completion with GoogleAIChatCompletion using a custom client"""
|
||||
# Create a custom client with a fake API key for testing
|
||||
custom_client = Client(api_key="fake-api-key-for-testing")
|
||||
|
||||
# Mock the custom client's generate_content_stream method
|
||||
mock_generate_content_stream = AsyncMock(return_value=mock_google_ai_streaming_chat_completion_response)
|
||||
custom_client.aio.models.generate_content_stream = mock_generate_content_stream
|
||||
|
||||
google_ai_chat_completion = GoogleAIChatCompletion(client=custom_client)
|
||||
|
||||
all_messages = []
|
||||
async for messages in google_ai_chat_completion.get_streaming_chat_message_contents(
|
||||
chat_history, GoogleAIChatPromptExecutionSettings()
|
||||
):
|
||||
all_messages.extend(messages)
|
||||
assert len(messages) == 1
|
||||
assert messages[0].role == "assistant"
|
||||
assert messages[0].finish_reason == FinishReason.STOP
|
||||
assert "usage" in messages[0].metadata
|
||||
assert "prompt_feedback" in messages[0].metadata
|
||||
|
||||
custom_client.close()
|
||||
|
||||
# Verify that the custom client was used and returned the expected response
|
||||
assert len(all_messages) > 0
|
||||
|
||||
# Verify that the custom client's method was called
|
||||
mock_generate_content_stream.assert_called_once()
|
||||
|
||||
|
||||
async def test_google_ai_streaming_chat_completion_with_function_choice_behavior_fail_verification(
|
||||
chat_history: ChatHistory,
|
||||
google_ai_unit_test_env,
|
||||
) -> None:
|
||||
"""Test streaming chat completion of GoogleAIChatCompletion with function choice
|
||||
behavior expect verification failure"""
|
||||
|
||||
# Missing kernel
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
settings = GoogleAIChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
|
||||
google_ai_chat_completion = GoogleAIChatCompletion()
|
||||
|
||||
async for _ in google_ai_chat_completion.get_streaming_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=settings,
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
@patch.object(AsyncModels, "generate_content_stream", new_callable=AsyncMock)
|
||||
async def test_google_ai_streaming_chat_completion_with_function_choice_behavior(
|
||||
mock_google_ai_model_generate_content_stream,
|
||||
google_ai_unit_test_env,
|
||||
kernel: Kernel,
|
||||
chat_history: ChatHistory,
|
||||
mock_google_ai_streaming_chat_completion_response_with_tool_call,
|
||||
decorated_native_function,
|
||||
) -> None:
|
||||
"""Test streaming chat completion of GoogleAIChatCompletion with function choice behavior"""
|
||||
mock_google_ai_model_generate_content_stream.return_value = (
|
||||
mock_google_ai_streaming_chat_completion_response_with_tool_call
|
||||
)
|
||||
|
||||
settings = GoogleAIChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
settings.function_choice_behavior.maximum_auto_invoke_attempts = 1
|
||||
|
||||
google_ai_chat_completion = GoogleAIChatCompletion()
|
||||
|
||||
kernel.add_function(plugin_name="TestPlugin", function=decorated_native_function)
|
||||
|
||||
all_messages = []
|
||||
async for messages in google_ai_chat_completion.get_streaming_chat_message_contents(
|
||||
chat_history,
|
||||
settings,
|
||||
kernel=kernel,
|
||||
):
|
||||
all_messages.extend(messages)
|
||||
|
||||
assert len(all_messages) == 2, f"Expected 2 messages, got {len(all_messages)}"
|
||||
|
||||
# Validate the first message
|
||||
assert all_messages[0].role == "assistant", f"Unexpected role for first message: {all_messages[0].role}"
|
||||
assert all_messages[0].content == "", f"Unexpected content for first message: {all_messages[0].content}"
|
||||
assert all_messages[0].finish_reason == FinishReason.STOP, (
|
||||
f"Unexpected finish reason for first message: {all_messages[0].finish_reason}"
|
||||
)
|
||||
|
||||
# Validate the second message
|
||||
assert all_messages[1].role == "tool", f"Unexpected role for second message: {all_messages[1].role}"
|
||||
assert all_messages[1].content == "", f"Unexpected content for second message: {all_messages[1].content}"
|
||||
assert all_messages[1].finish_reason is None
|
||||
|
||||
|
||||
@patch.object(AsyncModels, "generate_content_stream", new_callable=AsyncMock)
|
||||
async def test_google_ai_streaming_chat_completion_with_function_choice_behavior_no_tool_call(
|
||||
mock_google_ai_model_generate_content_stream,
|
||||
google_ai_unit_test_env,
|
||||
kernel,
|
||||
chat_history: ChatHistory,
|
||||
mock_google_ai_streaming_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test completion of GoogleAIChatCompletion with function choice behavior but no tool call returned"""
|
||||
mock_google_ai_model_generate_content_stream.return_value = mock_google_ai_streaming_chat_completion_response
|
||||
|
||||
settings = GoogleAIChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
settings.function_choice_behavior.maximum_auto_invoke_attempts = 1
|
||||
|
||||
google_ai_chat_completion = GoogleAIChatCompletion()
|
||||
|
||||
async for messages in google_ai_chat_completion.get_streaming_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=settings,
|
||||
kernel=kernel,
|
||||
):
|
||||
assert len(messages) == 1
|
||||
assert messages[0].role == "assistant"
|
||||
assert messages[0].content == "Test content"
|
||||
|
||||
mock_google_ai_model_generate_content_stream.assert_awaited_once_with(
|
||||
model=google_ai_chat_completion.service_settings.gemini_model_id,
|
||||
contents=google_ai_chat_completion._prepare_chat_history_for_request(chat_history),
|
||||
config=GenerateContentConfigDict(
|
||||
**settings.prepare_settings_dict(),
|
||||
system_instruction=filter_system_message(chat_history),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# endregion streaming chat completion
|
||||
|
||||
|
||||
def test_google_ai_chat_completion_parse_chat_history_correctly(google_ai_unit_test_env) -> None:
|
||||
"""Test _prepare_chat_history_for_request method"""
|
||||
google_ai_chat_completion = GoogleAIChatCompletion()
|
||||
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_system_message("test_system_message")
|
||||
chat_history.add_user_message("test_user_message")
|
||||
chat_history.add_assistant_message("test_assistant_message")
|
||||
|
||||
parsed_chat_history = google_ai_chat_completion._prepare_chat_history_for_request(chat_history)
|
||||
|
||||
assert isinstance(parsed_chat_history, list)
|
||||
# System message should be ignored
|
||||
assert len(parsed_chat_history) == 2
|
||||
assert all(isinstance(message, Content) for message in parsed_chat_history)
|
||||
assert parsed_chat_history[0].role == "user"
|
||||
assert parsed_chat_history[0].parts[0].text == "test_user_message"
|
||||
assert parsed_chat_history[1].role == "model"
|
||||
assert parsed_chat_history[1].parts[0].text == "test_assistant_message"
|
||||
|
||||
|
||||
# region deserialization (Part → FunctionCallContent round-trip)
|
||||
|
||||
|
||||
def test_create_chat_message_content_with_thought_signature(google_ai_unit_test_env) -> None:
|
||||
"""Test that thought_signature from a Part is deserialized into FunctionCallContent.metadata."""
|
||||
from google.genai.types import (
|
||||
Candidate,
|
||||
Content,
|
||||
GenerateContentResponse,
|
||||
GenerateContentResponseUsageMetadata,
|
||||
Part,
|
||||
)
|
||||
from google.genai.types import (
|
||||
FinishReason as GFinishReason,
|
||||
)
|
||||
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
|
||||
thought_sig_value = b"test-thought-sig-bytes"
|
||||
part = Part.from_function_call(name="test_function", args={"key": "value"})
|
||||
part.thought_signature = thought_sig_value
|
||||
|
||||
candidate = Candidate()
|
||||
candidate.index = 0
|
||||
candidate.content = Content(role="user", parts=[part])
|
||||
candidate.finish_reason = GFinishReason.STOP
|
||||
|
||||
response = GenerateContentResponse()
|
||||
response.candidates = [candidate]
|
||||
response.usage_metadata = GenerateContentResponseUsageMetadata(
|
||||
prompt_token_count=0, cached_content_token_count=0, candidates_token_count=0, total_token_count=0
|
||||
)
|
||||
|
||||
completion = GoogleAIChatCompletion()
|
||||
result = completion._create_chat_message_content(response, candidate)
|
||||
|
||||
fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
|
||||
assert len(fc_items) == 1
|
||||
assert fc_items[0].metadata is not None
|
||||
assert fc_items[0].metadata["thought_signature"] == thought_sig_value
|
||||
|
||||
|
||||
def test_create_chat_message_content_without_thought_signature(google_ai_unit_test_env) -> None:
|
||||
"""Test that FunctionCallContent works when Part has no thought_signature."""
|
||||
from google.genai.types import (
|
||||
Candidate,
|
||||
Content,
|
||||
GenerateContentResponse,
|
||||
GenerateContentResponseUsageMetadata,
|
||||
Part,
|
||||
)
|
||||
from google.genai.types import (
|
||||
FinishReason as GFinishReason,
|
||||
)
|
||||
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
|
||||
part = Part.from_function_call(name="test_function", args={"key": "value"})
|
||||
|
||||
candidate = Candidate()
|
||||
candidate.index = 0
|
||||
candidate.content = Content(role="user", parts=[part])
|
||||
candidate.finish_reason = GFinishReason.STOP
|
||||
|
||||
response = GenerateContentResponse()
|
||||
response.candidates = [candidate]
|
||||
response.usage_metadata = GenerateContentResponseUsageMetadata(
|
||||
prompt_token_count=0, cached_content_token_count=0, candidates_token_count=0, total_token_count=0
|
||||
)
|
||||
|
||||
completion = GoogleAIChatCompletion()
|
||||
result = completion._create_chat_message_content(response, candidate)
|
||||
|
||||
fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
|
||||
assert len(fc_items) == 1
|
||||
assert "thought_signature" not in fc_items[0].metadata
|
||||
|
||||
|
||||
def test_create_streaming_chat_message_content_with_thought_signature(google_ai_unit_test_env) -> None:
|
||||
"""Test that thought_signature from a Part is deserialized in streaming path."""
|
||||
from google.genai.types import (
|
||||
Candidate,
|
||||
Content,
|
||||
GenerateContentResponse,
|
||||
GenerateContentResponseUsageMetadata,
|
||||
Part,
|
||||
)
|
||||
from google.genai.types import (
|
||||
FinishReason as GFinishReason,
|
||||
)
|
||||
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
|
||||
thought_sig_value = b"streaming-thought-sig"
|
||||
part = Part.from_function_call(name="stream_func", args={"a": "b"})
|
||||
part.thought_signature = thought_sig_value
|
||||
|
||||
candidate = Candidate()
|
||||
candidate.index = 0
|
||||
candidate.content = Content(role="user", parts=[part])
|
||||
candidate.finish_reason = GFinishReason.STOP
|
||||
|
||||
chunk = GenerateContentResponse()
|
||||
chunk.candidates = [candidate]
|
||||
chunk.usage_metadata = GenerateContentResponseUsageMetadata(
|
||||
prompt_token_count=0, cached_content_token_count=0, candidates_token_count=0, total_token_count=0
|
||||
)
|
||||
|
||||
completion = GoogleAIChatCompletion()
|
||||
result = completion._create_streaming_chat_message_content(chunk, candidate)
|
||||
|
||||
fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
|
||||
assert len(fc_items) == 1
|
||||
assert fc_items[0].metadata is not None
|
||||
assert fc_items[0].metadata["thought_signature"] == thought_sig_value
|
||||
|
||||
|
||||
def test_create_streaming_chat_message_content_without_thought_signature(google_ai_unit_test_env) -> None:
|
||||
"""Test that streaming FunctionCallContent works when Part lacks thought_signature."""
|
||||
from google.genai.types import (
|
||||
Candidate,
|
||||
Content,
|
||||
GenerateContentResponse,
|
||||
GenerateContentResponseUsageMetadata,
|
||||
Part,
|
||||
)
|
||||
from google.genai.types import (
|
||||
FinishReason as GFinishReason,
|
||||
)
|
||||
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
|
||||
part = Part.from_function_call(name="stream_func", args={"a": "b"})
|
||||
|
||||
candidate = Candidate()
|
||||
candidate.index = 0
|
||||
candidate.content = Content(role="user", parts=[part])
|
||||
candidate.finish_reason = GFinishReason.STOP
|
||||
|
||||
chunk = GenerateContentResponse()
|
||||
chunk.candidates = [candidate]
|
||||
chunk.usage_metadata = GenerateContentResponseUsageMetadata(
|
||||
prompt_token_count=0, cached_content_token_count=0, candidates_token_count=0, total_token_count=0
|
||||
)
|
||||
|
||||
completion = GoogleAIChatCompletion()
|
||||
result = completion._create_streaming_chat_message_content(chunk, candidate)
|
||||
|
||||
fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
|
||||
assert len(fc_items) == 1
|
||||
assert "thought_signature" not in fc_items[0].metadata
|
||||
|
||||
|
||||
def test_create_chat_message_content_getattr_guard_on_missing_attribute(google_ai_unit_test_env) -> None:
|
||||
"""Test that getattr guard handles SDK versions where thought_signature doesn't exist on Part."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from google.genai.types import (
|
||||
GenerateContentResponse,
|
||||
GenerateContentResponseUsageMetadata,
|
||||
)
|
||||
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
|
||||
# Create a mock Part that lacks 'thought_signature' attribute entirely
|
||||
mock_part = MagicMock()
|
||||
mock_part.text = None
|
||||
mock_part.function_call.name = "test_func"
|
||||
mock_part.function_call.args = {"x": "y"}
|
||||
del mock_part.thought_signature # simulate older SDK without the field
|
||||
|
||||
# Use a fully-mocked candidate to avoid Content pydantic validation
|
||||
mock_candidate = MagicMock()
|
||||
mock_candidate.index = 0
|
||||
mock_candidate.content.parts = [mock_part]
|
||||
mock_candidate.finish_reason = 1 # STOP
|
||||
|
||||
response = GenerateContentResponse()
|
||||
response.candidates = [mock_candidate]
|
||||
response.usage_metadata = GenerateContentResponseUsageMetadata(
|
||||
prompt_token_count=0, cached_content_token_count=0, candidates_token_count=0, total_token_count=0
|
||||
)
|
||||
|
||||
completion = GoogleAIChatCompletion()
|
||||
result = completion._create_chat_message_content(response, mock_candidate)
|
||||
|
||||
fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
|
||||
assert len(fc_items) == 1
|
||||
assert "thought_signature" not in fc_items[0].metadata
|
||||
|
||||
|
||||
# endregion deserialization
|
||||
+212
@@ -0,0 +1,212 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from google.genai import Client
|
||||
from google.genai.models import AsyncModels
|
||||
from google.genai.types import GenerateContentConfigDict
|
||||
|
||||
from semantic_kernel.connectors.ai.google.google_ai.google_ai_prompt_execution_settings import (
|
||||
GoogleAITextPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.google.google_ai.google_ai_settings import GoogleAISettings
|
||||
from semantic_kernel.connectors.ai.google.google_ai.services.google_ai_text_completion import GoogleAITextCompletion
|
||||
from semantic_kernel.contents.text_content import TextContent
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
|
||||
|
||||
|
||||
# region init
|
||||
def test_google_ai_text_completion_init(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAITextCompletion"""
|
||||
model_id = google_ai_unit_test_env["GOOGLE_AI_GEMINI_MODEL_ID"]
|
||||
api_key = google_ai_unit_test_env["GOOGLE_AI_API_KEY"]
|
||||
google_ai_text_completion = GoogleAITextCompletion()
|
||||
|
||||
assert google_ai_text_completion.ai_model_id == model_id
|
||||
assert google_ai_text_completion.service_id == model_id
|
||||
|
||||
assert isinstance(google_ai_text_completion.service_settings, GoogleAISettings)
|
||||
assert google_ai_text_completion.service_settings.gemini_model_id == model_id
|
||||
assert google_ai_text_completion.service_settings.api_key.get_secret_value() == api_key
|
||||
|
||||
|
||||
def test_google_ai_text_completion_init_with_service_id(google_ai_unit_test_env, service_id) -> None:
|
||||
"""Test initialization of GoogleAITextCompletion with a service_id that is not the model_id"""
|
||||
google_ai_text_completion = GoogleAITextCompletion(service_id=service_id)
|
||||
|
||||
assert google_ai_text_completion.service_id == service_id
|
||||
|
||||
|
||||
def test_google_ai_text_completion_init_with_model_id_in_argument(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAIChatCompletion with model_id in argument"""
|
||||
google_ai_text_completion = GoogleAITextCompletion(gemini_model_id="custom_model_id")
|
||||
|
||||
assert google_ai_text_completion.ai_model_id == "custom_model_id"
|
||||
assert google_ai_text_completion.service_id == "custom_model_id"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["GOOGLE_AI_GEMINI_MODEL_ID"]], indirect=True)
|
||||
def test_google_ai_text_completion_init_with_empty_model_id(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAITextCompletion with an empty model_id"""
|
||||
with pytest.raises(ServiceInitializationError, match="The Google AI Gemini model ID is required."):
|
||||
GoogleAITextCompletion(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["GOOGLE_AI_API_KEY"]], indirect=True)
|
||||
def test_google_ai_text_completion_init_with_empty_api_key(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAITextCompletion with an empty api_key"""
|
||||
with pytest.raises(ServiceInitializationError, match="The API key is required when use_vertexai is False."):
|
||||
GoogleAITextCompletion(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", ["GOOGLE_AI_CLOUD_PROJECT_ID"], indirect=True)
|
||||
def test_google_ai_text_completion_init_with_use_vertexai_missing_project_id(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAITextCompletion with use_vertexai true but missing project_id"""
|
||||
with pytest.raises(ServiceInitializationError, match="Project ID must be provided when use_vertexai is True."):
|
||||
GoogleAITextCompletion(use_vertexai=True, env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["GOOGLE_AI_CLOUD_REGION"]], indirect=True)
|
||||
def test_google_ai_text_completion_init_with_use_vertexai_missing_region(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAITextCompletion with use_vertexai true but missing region"""
|
||||
with pytest.raises(ServiceInitializationError, match="Region must be provided when use_vertexai is True."):
|
||||
GoogleAITextCompletion(use_vertexai=True, env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["GOOGLE_AI_API_KEY"]], indirect=True)
|
||||
def test_google_ai_text_completion_init_with_use_vertexai_no_api_key(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAITextCompletion succeeds with use_vertexai=True and no api_key"""
|
||||
text_completion = GoogleAITextCompletion(use_vertexai=True)
|
||||
assert text_completion.service_settings.use_vertexai is True
|
||||
|
||||
|
||||
def test_prompt_execution_settings_class(google_ai_unit_test_env) -> None:
|
||||
google_ai_text_completion = GoogleAITextCompletion()
|
||||
assert google_ai_text_completion.get_prompt_execution_settings_class() == GoogleAITextPromptExecutionSettings
|
||||
|
||||
|
||||
# endregion init
|
||||
|
||||
|
||||
# region text completion
|
||||
|
||||
|
||||
@patch.object(AsyncModels, "generate_content", new_callable=AsyncMock)
|
||||
async def test_google_ai_text_completion(
|
||||
mock_google_model_generate_content,
|
||||
google_ai_unit_test_env,
|
||||
prompt: str,
|
||||
mock_google_ai_text_completion_response,
|
||||
) -> None:
|
||||
"""Test text completion with GoogleAITextCompletion"""
|
||||
settings = GoogleAITextPromptExecutionSettings()
|
||||
|
||||
mock_google_model_generate_content.return_value = mock_google_ai_text_completion_response
|
||||
|
||||
google_ai_text_completion = GoogleAITextCompletion()
|
||||
responses: list[TextContent] = await google_ai_text_completion.get_text_contents(prompt, settings)
|
||||
|
||||
mock_google_model_generate_content.assert_called_once_with(
|
||||
model=google_ai_text_completion.service_settings.gemini_model_id,
|
||||
contents=prompt,
|
||||
config=GenerateContentConfigDict(**settings.prepare_settings_dict()),
|
||||
)
|
||||
assert len(responses) == 1
|
||||
assert responses[0].text == mock_google_ai_text_completion_response.candidates[0].content.parts[0].text
|
||||
assert "usage" in responses[0].metadata
|
||||
assert "prompt_feedback" in responses[0].metadata
|
||||
assert responses[0].inner_content == mock_google_ai_text_completion_response
|
||||
|
||||
|
||||
async def test_google_ai_text_completion_with_custom_client(
|
||||
prompt: str,
|
||||
google_ai_unit_test_env,
|
||||
mock_google_ai_text_completion_response,
|
||||
) -> None:
|
||||
"""Test text completion with GoogleAITextCompletion using a custom client"""
|
||||
# Create a custom client with a fake API key for testing
|
||||
custom_client = Client(api_key="fake-api-key-for-testing")
|
||||
|
||||
# Mock the custom client's generate_content method
|
||||
mock_generate_content = AsyncMock(return_value=mock_google_ai_text_completion_response)
|
||||
custom_client.aio.models.generate_content = mock_generate_content
|
||||
|
||||
google_ai_text_completion = GoogleAITextCompletion(client=custom_client)
|
||||
responses: list[TextContent] = await google_ai_text_completion.get_text_contents(
|
||||
prompt,
|
||||
GoogleAITextPromptExecutionSettings(),
|
||||
)
|
||||
|
||||
custom_client.close()
|
||||
|
||||
# Verify that the custom client was used and returned the expected response
|
||||
assert len(responses) == 1
|
||||
assert responses[0].text == mock_google_ai_text_completion_response.candidates[0].content.parts[0].text
|
||||
assert "usage" in responses[0].metadata
|
||||
assert "prompt_feedback" in responses[0].metadata
|
||||
assert responses[0].inner_content == mock_google_ai_text_completion_response
|
||||
|
||||
# Verify that the custom client's method was called
|
||||
mock_generate_content.assert_called_once()
|
||||
|
||||
|
||||
# endregion text completion
|
||||
|
||||
|
||||
# region streaming text completion
|
||||
|
||||
|
||||
@patch.object(AsyncModels, "generate_content_stream", new_callable=AsyncMock)
|
||||
async def test_google_ai_streaming_text_completion(
|
||||
mock_google_model_generate_content_stream,
|
||||
google_ai_unit_test_env,
|
||||
prompt: str,
|
||||
mock_google_ai_streaming_text_completion_response,
|
||||
) -> None:
|
||||
"""Test streaming text completion with GoogleAITextCompletion"""
|
||||
settings = GoogleAITextPromptExecutionSettings()
|
||||
|
||||
mock_google_model_generate_content_stream.return_value = mock_google_ai_streaming_text_completion_response
|
||||
|
||||
google_ai_text_completion = GoogleAITextCompletion()
|
||||
async for chunks in google_ai_text_completion.get_streaming_text_contents(prompt, settings):
|
||||
assert len(chunks) == 1
|
||||
assert "usage" in chunks[0].metadata
|
||||
assert "prompt_feedback" in chunks[0].metadata
|
||||
|
||||
mock_google_model_generate_content_stream.assert_called_once_with(
|
||||
model=google_ai_text_completion.service_settings.gemini_model_id,
|
||||
contents=prompt,
|
||||
config=GenerateContentConfigDict(**settings.prepare_settings_dict()),
|
||||
)
|
||||
|
||||
|
||||
async def test_google_ai_streaming_text_completion_with_custom_client(
|
||||
prompt: str,
|
||||
google_ai_unit_test_env,
|
||||
mock_google_ai_streaming_text_completion_response,
|
||||
) -> None:
|
||||
"""Test streaming text completion with GoogleAITextCompletion using a custom client"""
|
||||
# Create a custom client with a fake API key for testing
|
||||
custom_client = Client(api_key="fake-api-key-for-testing")
|
||||
|
||||
# Mock the custom client's generate_content_stream method
|
||||
mock_generate_content_stream = AsyncMock(return_value=mock_google_ai_streaming_text_completion_response)
|
||||
custom_client.aio.models.generate_content_stream = mock_generate_content_stream
|
||||
|
||||
google_ai_text_completion = GoogleAITextCompletion(client=custom_client)
|
||||
async for chunks in google_ai_text_completion.get_streaming_text_contents(
|
||||
prompt, GoogleAITextPromptExecutionSettings()
|
||||
):
|
||||
assert len(chunks) == 1
|
||||
assert "usage" in chunks[0].metadata
|
||||
assert "prompt_feedback" in chunks[0].metadata
|
||||
|
||||
custom_client.close()
|
||||
|
||||
# Verify that the custom client's method was called
|
||||
mock_generate_content_stream.assert_called_once()
|
||||
|
||||
|
||||
# endregion streaming text completion
|
||||
+237
@@ -0,0 +1,237 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from google.genai import Client
|
||||
from google.genai.models import AsyncModels
|
||||
from google.genai.types import ContentEmbedding, EmbedContentConfigDict, EmbedContentResponse
|
||||
from numpy import array, ndarray
|
||||
|
||||
from semantic_kernel.connectors.ai.google.google_ai.google_ai_prompt_execution_settings import (
|
||||
GoogleAIEmbeddingPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.google.google_ai.google_ai_settings import GoogleAISettings
|
||||
from semantic_kernel.connectors.ai.google.google_ai.services.google_ai_text_embedding import GoogleAITextEmbedding
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
|
||||
|
||||
|
||||
# region init
|
||||
def test_google_ai_text_embedding_init(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAITextEmbedding"""
|
||||
model_id = google_ai_unit_test_env["GOOGLE_AI_EMBEDDING_MODEL_ID"]
|
||||
api_key = google_ai_unit_test_env["GOOGLE_AI_API_KEY"]
|
||||
google_ai_text_embedding = GoogleAITextEmbedding()
|
||||
|
||||
assert google_ai_text_embedding.ai_model_id == model_id
|
||||
assert google_ai_text_embedding.service_id == model_id
|
||||
|
||||
assert isinstance(google_ai_text_embedding.service_settings, GoogleAISettings)
|
||||
assert google_ai_text_embedding.service_settings.embedding_model_id == model_id
|
||||
assert google_ai_text_embedding.service_settings.api_key.get_secret_value() == api_key
|
||||
|
||||
|
||||
def test_google_ai_text_embedding_init_with_service_id(google_ai_unit_test_env, service_id) -> None:
|
||||
"""Test initialization of GoogleAITextEmbedding with a service_id that is not the model_id"""
|
||||
google_ai_text_embedding = GoogleAITextEmbedding(service_id=service_id)
|
||||
|
||||
assert google_ai_text_embedding.service_id == service_id
|
||||
|
||||
|
||||
def test_google_ai_text_embedding_init_with_model_id_in_argument(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAITextEmbedding with model_id in argument"""
|
||||
google_ai_chat_completion = GoogleAITextEmbedding(embedding_model_id="custom_model_id")
|
||||
|
||||
assert google_ai_chat_completion.ai_model_id == "custom_model_id"
|
||||
assert google_ai_chat_completion.service_id == "custom_model_id"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["GOOGLE_AI_EMBEDDING_MODEL_ID"]], indirect=True)
|
||||
def test_google_ai_text_embedding_init_with_empty_model_id(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAITextEmbedding with an empty model_id"""
|
||||
with pytest.raises(ServiceInitializationError, match="The Google AI embedding model ID is required."):
|
||||
GoogleAITextEmbedding(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["GOOGLE_AI_API_KEY"]], indirect=True)
|
||||
def test_google_ai_text_embedding_init_with_empty_api_key(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAITextEmbedding with an empty api_key"""
|
||||
with pytest.raises(ServiceInitializationError, match="The API key is required when use_vertexai is False."):
|
||||
GoogleAITextEmbedding(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["GOOGLE_AI_API_KEY"]], indirect=True)
|
||||
def test_google_ai_text_embedding_init_with_use_vertexai_no_api_key(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAITextEmbedding succeeds with use_vertexai=True and no api_key"""
|
||||
embedding = GoogleAITextEmbedding(use_vertexai=True)
|
||||
assert embedding.service_settings.use_vertexai is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["GOOGLE_AI_CLOUD_PROJECT_ID"]], indirect=True)
|
||||
def test_google_ai_text_embedding_init_with_use_vertexai_missing_project_id(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAITextEmbedding with use_vertexai true but missing project ID"""
|
||||
with pytest.raises(ServiceInitializationError, match="Project ID must be provided when use_vertexai is True."):
|
||||
GoogleAITextEmbedding(use_vertexai=True, env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["GOOGLE_AI_CLOUD_REGION"]], indirect=True)
|
||||
def test_google_ai_text_embedding_init_with_use_vertexai_missing_region(google_ai_unit_test_env) -> None:
|
||||
"""Test initialization of GoogleAITextEmbedding with use_vertexai true but missing region"""
|
||||
with pytest.raises(ServiceInitializationError, match="Region must be provided when use_vertexai is True."):
|
||||
GoogleAITextEmbedding(use_vertexai=True, env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
def test_prompt_execution_settings_class(google_ai_unit_test_env) -> None:
|
||||
google_ai_text_embedding = GoogleAITextEmbedding()
|
||||
assert google_ai_text_embedding.get_prompt_execution_settings_class() == GoogleAIEmbeddingPromptExecutionSettings
|
||||
|
||||
|
||||
# endregion init
|
||||
|
||||
|
||||
@patch.object(AsyncModels, "embed_content", new_callable=AsyncMock)
|
||||
async def test_embedding(mock_google_model_embed_content, google_ai_unit_test_env, prompt):
|
||||
"""Test that the service initializes and generates embeddings correctly."""
|
||||
model_id = google_ai_unit_test_env["GOOGLE_AI_EMBEDDING_MODEL_ID"]
|
||||
|
||||
mock_google_model_embed_content.return_value = EmbedContentResponse(
|
||||
embeddings=[ContentEmbedding(values=[0.1, 0.2, 0.3])]
|
||||
)
|
||||
settings = GoogleAIEmbeddingPromptExecutionSettings()
|
||||
|
||||
google_ai_text_embedding = GoogleAITextEmbedding()
|
||||
response: ndarray = await google_ai_text_embedding.generate_embeddings(
|
||||
[prompt],
|
||||
settings=settings,
|
||||
)
|
||||
|
||||
assert len(response) == 1
|
||||
assert response.all() == array([0.1, 0.2, 0.3]).all()
|
||||
mock_google_model_embed_content.assert_called_once_with(
|
||||
model=model_id,
|
||||
contents=[prompt],
|
||||
config=EmbedContentConfigDict(output_dimensionality=settings.output_dimensionality),
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncModels, "embed_content", new_callable=AsyncMock)
|
||||
async def test_embedding_with_settings(mock_google_model_embed_content, google_ai_unit_test_env, prompt):
|
||||
"""Test that the service initializes and generates embeddings correctly."""
|
||||
model_id = google_ai_unit_test_env["GOOGLE_AI_EMBEDDING_MODEL_ID"]
|
||||
|
||||
mock_google_model_embed_content.return_value = EmbedContentResponse(
|
||||
embeddings=[ContentEmbedding(values=[0.1, 0.2, 0.3])]
|
||||
)
|
||||
settings = GoogleAIEmbeddingPromptExecutionSettings()
|
||||
settings.output_dimensionality = 3
|
||||
|
||||
google_ai_text_embedding = GoogleAITextEmbedding()
|
||||
response: ndarray = await google_ai_text_embedding.generate_embeddings(
|
||||
[prompt],
|
||||
settings=settings,
|
||||
)
|
||||
|
||||
assert len(response) == 1
|
||||
assert response.all() == array([0.1, 0.2, 0.3]).all()
|
||||
mock_google_model_embed_content.assert_called_once_with(
|
||||
model=model_id,
|
||||
contents=[prompt],
|
||||
config=EmbedContentConfigDict(**settings.prepare_settings_dict()),
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncModels, "embed_content", new_callable=AsyncMock)
|
||||
async def test_embedding_without_settings(mock_google_model_embed_content, google_ai_unit_test_env, prompt):
|
||||
"""Test that the service initializes and generates embeddings correctly without settings."""
|
||||
model_id = google_ai_unit_test_env["GOOGLE_AI_EMBEDDING_MODEL_ID"]
|
||||
|
||||
mock_google_model_embed_content.return_value = EmbedContentResponse(
|
||||
embeddings=[ContentEmbedding(values=[0.1, 0.2, 0.3])]
|
||||
)
|
||||
google_ai_text_embedding = GoogleAITextEmbedding()
|
||||
response: ndarray = await google_ai_text_embedding.generate_embeddings([prompt])
|
||||
|
||||
assert len(response) == 1
|
||||
assert response.all() == array([0.1, 0.2, 0.3]).all()
|
||||
mock_google_model_embed_content.assert_called_once_with(
|
||||
model=model_id,
|
||||
contents=[prompt],
|
||||
config=EmbedContentConfigDict(output_dimensionality=None),
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncModels, "embed_content", new_callable=AsyncMock)
|
||||
async def test_embedding_list_input(mock_google_model_embed_content, google_ai_unit_test_env, prompt):
|
||||
"""Test that the service initializes and generates embeddings correctly with a list of prompts."""
|
||||
model_id = google_ai_unit_test_env["GOOGLE_AI_EMBEDDING_MODEL_ID"]
|
||||
|
||||
mock_google_model_embed_content.return_value = EmbedContentResponse(
|
||||
embeddings=[ContentEmbedding(values=[0.1, 0.2, 0.3]), ContentEmbedding(values=[0.1, 0.2, 0.3])]
|
||||
)
|
||||
settings = GoogleAIEmbeddingPromptExecutionSettings()
|
||||
|
||||
google_ai_text_embedding = GoogleAITextEmbedding()
|
||||
response: ndarray = await google_ai_text_embedding.generate_embeddings(
|
||||
[prompt, prompt],
|
||||
settings=settings,
|
||||
)
|
||||
|
||||
assert len(response) == 2
|
||||
assert response.all() == array([[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]).all()
|
||||
mock_google_model_embed_content.assert_called_once_with(
|
||||
model=model_id,
|
||||
contents=[prompt, prompt],
|
||||
config=EmbedContentConfigDict(output_dimensionality=settings.output_dimensionality),
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncModels, "embed_content", new_callable=AsyncMock)
|
||||
async def test_raw_embedding(mock_google_model_embed_content, google_ai_unit_test_env, prompt):
|
||||
"""Test that the service initializes and generates embeddings correctly."""
|
||||
model_id = google_ai_unit_test_env["GOOGLE_AI_EMBEDDING_MODEL_ID"]
|
||||
|
||||
mock_google_model_embed_content.return_value = EmbedContentResponse(
|
||||
embeddings=[ContentEmbedding(values=[0.1, 0.2, 0.3])]
|
||||
)
|
||||
settings = GoogleAIEmbeddingPromptExecutionSettings()
|
||||
|
||||
google_ai_text_embedding = GoogleAITextEmbedding()
|
||||
response: list[list[float]] = await google_ai_text_embedding.generate_raw_embeddings(
|
||||
[prompt],
|
||||
settings=settings,
|
||||
)
|
||||
|
||||
assert len(response) == 1
|
||||
assert response[0] == [0.1, 0.2, 0.3]
|
||||
mock_google_model_embed_content.assert_called_once_with(
|
||||
model=model_id,
|
||||
contents=[prompt],
|
||||
config=EmbedContentConfigDict(output_dimensionality=settings.output_dimensionality),
|
||||
)
|
||||
|
||||
|
||||
async def test_embedding_with_custom_client(google_ai_unit_test_env, prompt) -> None:
|
||||
"""Test embedding with GoogleAITextEmbedding using a custom client"""
|
||||
# Create a custom client with a fake API key for testing
|
||||
custom_client = Client(api_key="fake-api-key-for-testing")
|
||||
|
||||
# Mock the custom client's embed_content method
|
||||
mock_embed_content = AsyncMock(
|
||||
return_value=EmbedContentResponse(embeddings=[ContentEmbedding(values=[0.1, 0.2, 0.3])])
|
||||
)
|
||||
custom_client.aio.models.embed_content = mock_embed_content
|
||||
|
||||
google_ai_text_embedding = GoogleAITextEmbedding(client=custom_client)
|
||||
response: list[list[float]] = await google_ai_text_embedding.generate_embeddings(
|
||||
[prompt],
|
||||
GoogleAIEmbeddingPromptExecutionSettings(),
|
||||
)
|
||||
|
||||
custom_client.close()
|
||||
|
||||
# Verify that the custom client was used and returned the expected response
|
||||
assert len(response) == 1
|
||||
assert response.all() == array([0.1, 0.2, 0.3]).all()
|
||||
|
||||
# Verify that the custom client's method was called
|
||||
mock_embed_content.assert_called_once()
|
||||
@@ -0,0 +1,201 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import pytest
|
||||
from google.genai.types import FinishReason, Part
|
||||
|
||||
from semantic_kernel.connectors.ai.google.google_ai.services.utils import (
|
||||
finish_reason_from_google_ai_to_semantic_kernel,
|
||||
format_assistant_message,
|
||||
format_user_message,
|
||||
kernel_function_metadata_to_google_ai_function_call_format,
|
||||
)
|
||||
from semantic_kernel.contents.chat_message_content import ChatMessageContent
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
from semantic_kernel.contents.function_result_content import FunctionResultContent
|
||||
from semantic_kernel.contents.image_content import ImageContent
|
||||
from semantic_kernel.contents.text_content import TextContent
|
||||
from semantic_kernel.contents.utils.author_role import AuthorRole
|
||||
from semantic_kernel.contents.utils.finish_reason import FinishReason as SemanticKernelFinishReason
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInvalidRequestError
|
||||
from semantic_kernel.functions.kernel_function_metadata import KernelFunctionMetadata
|
||||
from semantic_kernel.functions.kernel_parameter_metadata import KernelParameterMetadata
|
||||
|
||||
|
||||
def test_finish_reason_from_google_ai_to_semantic_kernel():
|
||||
"""Test finish_reason_from_google_ai_to_semantic_kernel."""
|
||||
assert finish_reason_from_google_ai_to_semantic_kernel(FinishReason.STOP) == SemanticKernelFinishReason.STOP
|
||||
assert finish_reason_from_google_ai_to_semantic_kernel(FinishReason.MAX_TOKENS) == SemanticKernelFinishReason.LENGTH
|
||||
assert (
|
||||
finish_reason_from_google_ai_to_semantic_kernel(FinishReason.SAFETY)
|
||||
== SemanticKernelFinishReason.CONTENT_FILTER
|
||||
)
|
||||
assert finish_reason_from_google_ai_to_semantic_kernel(FinishReason.OTHER) is None
|
||||
assert finish_reason_from_google_ai_to_semantic_kernel(None) is None
|
||||
|
||||
|
||||
def test_format_user_message():
|
||||
"""Test format_user_message."""
|
||||
user_message = ChatMessageContent(role=AuthorRole.USER, content="User message")
|
||||
formatted_user_message = format_user_message(user_message)
|
||||
|
||||
assert len(formatted_user_message) == 1
|
||||
assert isinstance(formatted_user_message[0], Part)
|
||||
assert formatted_user_message[0].text == "User message"
|
||||
|
||||
# Test with an image content
|
||||
image_content = ImageContent(data="image data", mime_type="image/png")
|
||||
user_message = ChatMessageContent(
|
||||
role=AuthorRole.USER,
|
||||
items=[
|
||||
TextContent(text="Text content"),
|
||||
image_content,
|
||||
],
|
||||
)
|
||||
formatted_user_message = format_user_message(user_message)
|
||||
|
||||
assert len(formatted_user_message) == 2
|
||||
assert isinstance(formatted_user_message[0], Part)
|
||||
assert formatted_user_message[0].text == "Text content"
|
||||
assert isinstance(formatted_user_message[1], Part)
|
||||
assert formatted_user_message[1].inline_data.mime_type == "image/png"
|
||||
assert formatted_user_message[1].inline_data.data == image_content.data
|
||||
|
||||
|
||||
def test_format_user_message_throws_with_unsupported_items() -> None:
|
||||
"""Test format_user_message with unsupported items."""
|
||||
# Test with unsupported items, any item other than TextContent and ImageContent should raise an error
|
||||
user_message = ChatMessageContent(
|
||||
role=AuthorRole.USER,
|
||||
items=[
|
||||
FunctionCallContent(),
|
||||
],
|
||||
)
|
||||
with pytest.raises(ServiceInvalidRequestError):
|
||||
format_user_message(user_message)
|
||||
|
||||
# Test with an ImageContent that has no data_uri
|
||||
user_message = ChatMessageContent(
|
||||
role=AuthorRole.USER,
|
||||
items=[
|
||||
ImageContent(data_uri=""),
|
||||
],
|
||||
)
|
||||
with pytest.raises(ServiceInvalidRequestError):
|
||||
format_user_message(user_message)
|
||||
|
||||
|
||||
def test_format_assistant_message() -> None:
|
||||
assistant_message = ChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT,
|
||||
items=[
|
||||
TextContent(text="test"),
|
||||
FunctionCallContent(name="test_function", arguments={}),
|
||||
ImageContent(data="image data", mime_type="image/png"),
|
||||
],
|
||||
)
|
||||
|
||||
formatted_assistant_message = format_assistant_message(assistant_message)
|
||||
assert isinstance(formatted_assistant_message, list)
|
||||
assert len(formatted_assistant_message) == 3
|
||||
assert isinstance(formatted_assistant_message[0], Part)
|
||||
assert formatted_assistant_message[0].text == "test"
|
||||
assert isinstance(formatted_assistant_message[1], Part)
|
||||
assert formatted_assistant_message[1].function_call.name == "test_function"
|
||||
assert formatted_assistant_message[1].function_call.args == {}
|
||||
assert isinstance(formatted_assistant_message[2], Part)
|
||||
assert formatted_assistant_message[2].inline_data
|
||||
|
||||
|
||||
def test_format_assistant_message_with_unsupported_items() -> None:
|
||||
assistant_message = ChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT,
|
||||
items=[
|
||||
FunctionResultContent(id="test_id", function_name="test_function"),
|
||||
],
|
||||
)
|
||||
|
||||
with pytest.raises(ServiceInvalidRequestError):
|
||||
format_assistant_message(assistant_message)
|
||||
|
||||
|
||||
def test_format_assistant_message_with_thought_signature() -> None:
|
||||
"""Test that thought_signature is preserved in function call parts."""
|
||||
import base64
|
||||
|
||||
thought_sig = base64.b64encode(b"test_thought_signature_data")
|
||||
assistant_message = ChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT,
|
||||
items=[
|
||||
FunctionCallContent(
|
||||
name="test_function",
|
||||
arguments={"arg1": "value1"},
|
||||
metadata={"thought_signature": thought_sig},
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
formatted = format_assistant_message(assistant_message)
|
||||
assert len(formatted) == 1
|
||||
assert isinstance(formatted[0], Part)
|
||||
assert formatted[0].function_call.name == "test_function"
|
||||
assert formatted[0].function_call.args == {"arg1": "value1"}
|
||||
assert formatted[0].thought_signature == thought_sig
|
||||
|
||||
|
||||
def test_format_assistant_message_without_thought_signature() -> None:
|
||||
"""Test that function calls without thought_signature still work."""
|
||||
assistant_message = ChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT,
|
||||
items=[
|
||||
FunctionCallContent(
|
||||
name="test_function",
|
||||
arguments={"arg1": "value1"},
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
formatted = format_assistant_message(assistant_message)
|
||||
assert len(formatted) == 1
|
||||
assert isinstance(formatted[0], Part)
|
||||
assert formatted[0].function_call.name == "test_function"
|
||||
assert formatted[0].function_call.args == {"arg1": "value1"}
|
||||
assert not getattr(formatted[0], "thought_signature", None)
|
||||
|
||||
|
||||
def test_google_ai_function_call_format_sanitizes_anyof_schema() -> None:
|
||||
"""Integration test: anyOf in param schema_data is sanitized in the output dict."""
|
||||
metadata = KernelFunctionMetadata(
|
||||
name="test_func",
|
||||
description="A test function",
|
||||
is_prompt=False,
|
||||
parameters=[
|
||||
KernelParameterMetadata(
|
||||
name="messages",
|
||||
description="The user messages",
|
||||
is_required=True,
|
||||
schema_data={
|
||||
"anyOf": [
|
||||
{"type": "string"},
|
||||
{"type": "array", "items": {"type": "string"}},
|
||||
],
|
||||
"description": "The user messages",
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
result = kernel_function_metadata_to_google_ai_function_call_format(metadata)
|
||||
param_schema = result["parameters"]["properties"]["messages"]
|
||||
assert "anyOf" not in param_schema
|
||||
assert param_schema["type"] == "string"
|
||||
|
||||
|
||||
def test_google_ai_function_call_format_empty_parameters() -> None:
|
||||
"""Integration test: metadata with no parameters produces parameters=None."""
|
||||
metadata = KernelFunctionMetadata(
|
||||
name="no_params_func",
|
||||
description="No parameters",
|
||||
is_prompt=False,
|
||||
parameters=[],
|
||||
)
|
||||
result = kernel_function_metadata_to_google_ai_function_call_format(metadata)
|
||||
assert result["parameters"] is None
|
||||
@@ -0,0 +1,133 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from semantic_kernel.connectors.ai.google.google_ai import (
|
||||
GoogleAIChatPromptExecutionSettings,
|
||||
GoogleAIEmbeddingPromptExecutionSettings,
|
||||
GoogleAIPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
|
||||
|
||||
def test_default_google_ai_prompt_execution_settings():
|
||||
settings = GoogleAIPromptExecutionSettings()
|
||||
|
||||
assert settings.stop_sequences is None
|
||||
assert settings.response_mime_type is None
|
||||
assert settings.response_schema is None
|
||||
assert settings.candidate_count is None
|
||||
assert settings.max_output_tokens is None
|
||||
assert settings.temperature is None
|
||||
assert settings.top_p is None
|
||||
assert settings.top_k is None
|
||||
|
||||
|
||||
def test_custom_google_ai_prompt_execution_settings():
|
||||
settings = GoogleAIPromptExecutionSettings(
|
||||
stop_sequences=["world"],
|
||||
response_mime_type="text/plain",
|
||||
candidate_count=1,
|
||||
max_output_tokens=128,
|
||||
temperature=0.5,
|
||||
top_p=0.5,
|
||||
top_k=10,
|
||||
)
|
||||
|
||||
assert settings.stop_sequences == ["world"]
|
||||
assert settings.response_mime_type == "text/plain"
|
||||
assert settings.candidate_count == 1
|
||||
assert settings.max_output_tokens == 128
|
||||
assert settings.temperature == 0.5
|
||||
assert settings.top_p == 0.5
|
||||
assert settings.top_k == 10
|
||||
|
||||
|
||||
def test_google_ai_prompt_execution_settings_from_default_completion_config():
|
||||
settings = PromptExecutionSettings(service_id="test_service")
|
||||
chat_settings = GoogleAIChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
|
||||
assert chat_settings.service_id == "test_service"
|
||||
assert chat_settings.stop_sequences is None
|
||||
assert chat_settings.response_mime_type is None
|
||||
assert chat_settings.response_schema is None
|
||||
assert chat_settings.candidate_count is None
|
||||
assert chat_settings.max_output_tokens is None
|
||||
assert chat_settings.temperature is None
|
||||
assert chat_settings.top_p is None
|
||||
assert chat_settings.top_k is None
|
||||
|
||||
|
||||
def test_google_ai_prompt_execution_settings_from_openai_prompt_execution_settings():
|
||||
chat_settings = GoogleAIChatPromptExecutionSettings(service_id="test_service", temperature=1.0)
|
||||
new_settings = GoogleAIPromptExecutionSettings(service_id="test_2", temperature=0.0)
|
||||
chat_settings.update_from_prompt_execution_settings(new_settings)
|
||||
|
||||
assert chat_settings.service_id == "test_2"
|
||||
assert chat_settings.temperature == 0.0
|
||||
|
||||
|
||||
def test_google_ai_prompt_execution_settings_from_custom_completion_config():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"stop_sequences": ["world"],
|
||||
"response_mime_type": "text/plain",
|
||||
"candidate_count": 1,
|
||||
"max_output_tokens": 128,
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"top_k": 10,
|
||||
},
|
||||
)
|
||||
chat_settings = GoogleAIChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
|
||||
assert chat_settings.stop_sequences == ["world"]
|
||||
assert chat_settings.response_mime_type == "text/plain"
|
||||
assert chat_settings.candidate_count == 1
|
||||
assert chat_settings.max_output_tokens == 128
|
||||
assert chat_settings.temperature == 0.5
|
||||
assert chat_settings.top_p == 0.5
|
||||
assert chat_settings.top_k == 10
|
||||
|
||||
|
||||
def test_google_ai_chat_prompt_execution_settings_from_custom_completion_config_with_functions():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"tools": [{"function": {}}],
|
||||
},
|
||||
)
|
||||
chat_settings = GoogleAIChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
|
||||
assert chat_settings.tools == [{"function": {}}]
|
||||
|
||||
|
||||
def test_create_options():
|
||||
settings = GoogleAIChatPromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"stop_sequences": ["world"],
|
||||
"response_mime_type": "text/plain",
|
||||
"candidate_count": 1,
|
||||
"max_output_tokens": 128,
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"top_k": 10,
|
||||
},
|
||||
)
|
||||
options = settings.prepare_settings_dict()
|
||||
|
||||
assert options["stop_sequences"] == ["world"]
|
||||
assert options["response_mime_type"] == "text/plain"
|
||||
assert options["candidate_count"] == 1
|
||||
assert options["max_output_tokens"] == 128
|
||||
assert options["temperature"] == 0.5
|
||||
assert options["top_p"] == 0.5
|
||||
assert options["top_k"] == 10
|
||||
assert "tools" not in options
|
||||
assert "tool_config" not in options
|
||||
|
||||
|
||||
def test_default_google_ai_embedding_prompt_execution_settings():
|
||||
settings = GoogleAIEmbeddingPromptExecutionSettings()
|
||||
|
||||
assert settings.output_dimensionality is None
|
||||
@@ -0,0 +1,268 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceType
|
||||
from semantic_kernel.connectors.ai.google.shared_utils import (
|
||||
FUNCTION_CHOICE_TYPE_TO_GOOGLE_FUNCTION_CALLING_MODE,
|
||||
GEMINI_FUNCTION_NAME_SEPARATOR,
|
||||
collapse_function_call_results_in_chat_history,
|
||||
filter_system_message,
|
||||
format_gemini_function_name_to_kernel_function_fully_qualified_name,
|
||||
sanitize_schema_for_google_ai,
|
||||
)
|
||||
from semantic_kernel.contents.chat_history import ChatHistory
|
||||
from semantic_kernel.contents.chat_message_content import ChatMessageContent
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
from semantic_kernel.contents.function_result_content import FunctionResultContent
|
||||
from semantic_kernel.contents.utils.author_role import AuthorRole
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInvalidRequestError
|
||||
|
||||
|
||||
def test_first_system_message():
|
||||
"""Test filter_system_message."""
|
||||
# Test with a single system message
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_system_message("System message")
|
||||
chat_history.add_user_message("User message")
|
||||
assert filter_system_message(chat_history) == "System message"
|
||||
|
||||
# Test with no system message
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_user_message("User message")
|
||||
assert filter_system_message(chat_history) is None
|
||||
|
||||
# Test with multiple system messages
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_system_message("System message 1")
|
||||
chat_history.add_system_message("System message 2")
|
||||
with pytest.raises(ServiceInvalidRequestError):
|
||||
filter_system_message(chat_history)
|
||||
|
||||
|
||||
def test_function_choice_type_to_google_function_calling_mode_contain_all_types() -> None:
|
||||
assert FunctionChoiceType.AUTO in FUNCTION_CHOICE_TYPE_TO_GOOGLE_FUNCTION_CALLING_MODE
|
||||
assert FunctionChoiceType.NONE in FUNCTION_CHOICE_TYPE_TO_GOOGLE_FUNCTION_CALLING_MODE
|
||||
assert FunctionChoiceType.REQUIRED in FUNCTION_CHOICE_TYPE_TO_GOOGLE_FUNCTION_CALLING_MODE
|
||||
|
||||
|
||||
def test_format_gemini_function_name_to_kernel_function_fully_qualified_name() -> None:
|
||||
# Contains the separator
|
||||
gemini_function_name = f"plugin{GEMINI_FUNCTION_NAME_SEPARATOR}function"
|
||||
assert (
|
||||
format_gemini_function_name_to_kernel_function_fully_qualified_name(gemini_function_name) == "plugin-function"
|
||||
)
|
||||
|
||||
# Doesn't contain the separator
|
||||
gemini_function_name = "function"
|
||||
assert format_gemini_function_name_to_kernel_function_fully_qualified_name(gemini_function_name) == "function"
|
||||
|
||||
|
||||
def test_collapse_function_call_results_in_chat_history() -> None:
|
||||
chat_history = ChatHistory()
|
||||
chat_history.extend([
|
||||
ChatMessageContent(
|
||||
AuthorRole.ASSISTANT,
|
||||
items=[
|
||||
FunctionCallContent(id="function1", name="function1"),
|
||||
FunctionCallContent(id="function2", name="function2"),
|
||||
],
|
||||
),
|
||||
# The following two messages should be collapsed into a single message
|
||||
ChatMessageContent(
|
||||
AuthorRole.TOOL,
|
||||
items=[FunctionResultContent(id="function1", name="function1", result="result1")],
|
||||
),
|
||||
ChatMessageContent(
|
||||
AuthorRole.TOOL,
|
||||
items=[FunctionResultContent(id="function2", name="function2", result="result2")],
|
||||
),
|
||||
ChatMessageContent(AuthorRole.ASSISTANT, content="Assistant message"),
|
||||
ChatMessageContent(AuthorRole.USER, content="User message"),
|
||||
ChatMessageContent(
|
||||
AuthorRole.ASSISTANT,
|
||||
items=[FunctionCallContent(id="function3", name="function3")],
|
||||
),
|
||||
ChatMessageContent(
|
||||
AuthorRole.TOOL,
|
||||
items=[FunctionResultContent(id="function3", name="function3", result="result3")],
|
||||
),
|
||||
ChatMessageContent(AuthorRole.ASSISTANT, content="Assistant message"),
|
||||
])
|
||||
|
||||
assert len(chat_history.messages) == 8
|
||||
collapse_function_call_results_in_chat_history(chat_history)
|
||||
assert len(chat_history.messages) == 7
|
||||
assert len(chat_history.messages[1].items) == 2
|
||||
|
||||
|
||||
# --- sanitize_schema_for_google_ai tests ---
|
||||
|
||||
|
||||
def test_sanitize_schema_none():
|
||||
"""Test that None input returns None."""
|
||||
assert sanitize_schema_for_google_ai(None) is None
|
||||
|
||||
|
||||
def test_sanitize_schema_simple_passthrough():
|
||||
"""Test that a simple schema passes through unchanged."""
|
||||
schema = {"type": "string", "description": "A name"}
|
||||
result = sanitize_schema_for_google_ai(schema)
|
||||
assert result == {"type": "string", "description": "A name"}
|
||||
|
||||
|
||||
def test_sanitize_schema_type_as_list_with_null():
|
||||
"""type: ["string", "null"] should become type: "string" + nullable: true."""
|
||||
schema = {"type": ["string", "null"], "description": "Optional field"}
|
||||
result = sanitize_schema_for_google_ai(schema)
|
||||
assert result == {"type": "string", "nullable": True, "description": "Optional field"}
|
||||
|
||||
|
||||
def test_sanitize_schema_type_as_list_without_null():
|
||||
"""type: ["string", "integer"] should pick the first type."""
|
||||
schema = {"type": ["string", "integer"]}
|
||||
result = sanitize_schema_for_google_ai(schema)
|
||||
assert result == {"type": "string"}
|
||||
|
||||
|
||||
def test_sanitize_schema_anyof_with_null():
|
||||
"""AnyOf with null variant should become the non-null type + nullable."""
|
||||
schema = {
|
||||
"anyOf": [{"type": "string"}, {"type": "null"}],
|
||||
"description": "Optional param",
|
||||
}
|
||||
result = sanitize_schema_for_google_ai(schema)
|
||||
assert result == {"type": "string", "nullable": True, "description": "Optional param"}
|
||||
|
||||
|
||||
def test_sanitize_schema_anyof_without_null():
|
||||
"""AnyOf without null should pick the first variant."""
|
||||
schema = {
|
||||
"anyOf": [
|
||||
{"type": "string"},
|
||||
{"type": "array", "items": {"type": "string"}},
|
||||
],
|
||||
}
|
||||
result = sanitize_schema_for_google_ai(schema)
|
||||
assert result == {"type": "string"}
|
||||
|
||||
|
||||
def test_sanitize_schema_oneof():
|
||||
"""OneOf should be handled the same as anyOf."""
|
||||
schema = {
|
||||
"oneOf": [{"type": "integer"}, {"type": "null"}],
|
||||
}
|
||||
result = sanitize_schema_for_google_ai(schema)
|
||||
assert result == {"type": "integer", "nullable": True}
|
||||
|
||||
|
||||
def test_sanitize_schema_nested_properties():
|
||||
"""AnyOf inside nested properties should be sanitized recursively."""
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"value": {"anyOf": [{"type": "number"}, {"type": "null"}]},
|
||||
},
|
||||
}
|
||||
result = sanitize_schema_for_google_ai(schema)
|
||||
assert result == {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"value": {"type": "number", "nullable": True},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_sanitize_schema_nested_items():
|
||||
"""AnyOf inside array items should be sanitized recursively."""
|
||||
schema = {
|
||||
"type": "array",
|
||||
"items": {"anyOf": [{"type": "string"}, {"type": "integer"}]},
|
||||
}
|
||||
result = sanitize_schema_for_google_ai(schema)
|
||||
assert result == {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
}
|
||||
|
||||
|
||||
def test_sanitize_schema_does_not_mutate_original():
|
||||
"""The original schema dict should not be modified."""
|
||||
schema = {
|
||||
"anyOf": [{"type": "string"}, {"type": "null"}],
|
||||
"description": "test",
|
||||
}
|
||||
original = {"anyOf": [{"type": "string"}, {"type": "null"}], "description": "test"}
|
||||
sanitize_schema_for_google_ai(schema)
|
||||
assert schema == original
|
||||
|
||||
|
||||
def test_sanitize_schema_agent_messages_param():
|
||||
"""Reproducer for issue #12442: str | list[str] parameter schema."""
|
||||
schema = {
|
||||
"anyOf": [
|
||||
{"type": "string"},
|
||||
{"type": "array", "items": {"type": "string"}},
|
||||
],
|
||||
"description": "The user messages for the agent.",
|
||||
}
|
||||
result = sanitize_schema_for_google_ai(schema)
|
||||
assert "anyOf" not in result
|
||||
assert result["type"] == "string"
|
||||
assert result["description"] == "The user messages for the agent."
|
||||
|
||||
|
||||
def test_sanitize_schema_allof():
|
||||
"""AllOf should be handled like anyOf/oneOf, picking the first variant."""
|
||||
schema = {
|
||||
"allOf": [
|
||||
{"type": "object", "properties": {"name": {"type": "string"}}},
|
||||
{"type": "object", "properties": {"age": {"type": "integer"}}},
|
||||
],
|
||||
}
|
||||
result = sanitize_schema_for_google_ai(schema)
|
||||
assert "allOf" not in result
|
||||
assert result["type"] == "object"
|
||||
assert "name" in result["properties"]
|
||||
|
||||
|
||||
def test_sanitize_schema_allof_with_null():
|
||||
"""AllOf with a null variant should produce nullable: true."""
|
||||
schema = {
|
||||
"allOf": [{"type": "string"}, {"type": "null"}],
|
||||
}
|
||||
result = sanitize_schema_for_google_ai(schema)
|
||||
assert "allOf" not in result
|
||||
assert result["type"] == "string"
|
||||
assert result["nullable"] is True
|
||||
|
||||
|
||||
def test_sanitize_schema_all_null_type_list():
|
||||
"""type: ["null"] should fall back to type: "string" + nullable: true."""
|
||||
schema = {"type": ["null"]}
|
||||
result = sanitize_schema_for_google_ai(schema)
|
||||
assert result == {"type": "string", "nullable": True}
|
||||
|
||||
|
||||
def test_sanitize_schema_all_null_anyof():
|
||||
"""AnyOf where all variants are null should fall back to type: "string"."""
|
||||
schema = {"anyOf": [{"type": "null"}]}
|
||||
result = sanitize_schema_for_google_ai(schema)
|
||||
assert result == {"type": "string", "nullable": True}
|
||||
|
||||
|
||||
def test_sanitize_schema_chosen_variant_keeps_own_description():
|
||||
"""When the chosen anyOf variant has its own description, do not overwrite it."""
|
||||
schema = {
|
||||
"anyOf": [
|
||||
{"type": "string", "description": "inner desc"},
|
||||
{"type": "null"},
|
||||
],
|
||||
"description": "outer desc",
|
||||
}
|
||||
result = sanitize_schema_for_google_ai(schema)
|
||||
assert result["description"] == "inner desc"
|
||||
assert result["nullable"] is True
|
||||
@@ -0,0 +1,189 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from collections.abc import AsyncGenerator, AsyncIterable
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from google.cloud.aiplatform_v1beta1.types.content import Candidate, Content, Part
|
||||
from google.cloud.aiplatform_v1beta1.types.prediction_service import GenerateContentResponse
|
||||
from google.cloud.aiplatform_v1beta1.types.tool import FunctionCall
|
||||
from vertexai.generative_models import GenerationResponse
|
||||
from vertexai.language_models import TextEmbedding
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def vertex_ai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for Vertex AI Unit Tests."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"VERTEX_AI_GEMINI_MODEL_ID": "test-gemini-model-id",
|
||||
"VERTEX_AI_EMBEDDING_MODEL_ID": "test-embedding-model-id",
|
||||
"VERTEX_AI_PROJECT_ID": "test-project-id",
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_vertex_ai_chat_completion_response() -> GenerationResponse:
|
||||
"""Mock Vertex AI Chat Completion response."""
|
||||
candidate = Candidate()
|
||||
candidate.index = 0
|
||||
candidate.content = Content(role="user", parts=[Part(text="Test content")])
|
||||
candidate.finish_reason = Candidate.FinishReason.STOP
|
||||
|
||||
response = GenerateContentResponse()
|
||||
response.candidates.append(candidate)
|
||||
response.usage_metadata = GenerateContentResponse.UsageMetadata(
|
||||
prompt_token_count=0,
|
||||
candidates_token_count=0,
|
||||
total_token_count=0,
|
||||
)
|
||||
|
||||
return GenerationResponse._from_gapic(response)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_vertex_ai_chat_completion_response_with_tool_call() -> GenerationResponse:
|
||||
"""Mock Vertex AI Chat Completion response."""
|
||||
candidate = Candidate()
|
||||
candidate.index = 0
|
||||
candidate.content = Content(
|
||||
role="user",
|
||||
parts=[
|
||||
Part(
|
||||
function_call=FunctionCall(
|
||||
name="test_function",
|
||||
args={"test_arg": "test_value"},
|
||||
)
|
||||
)
|
||||
],
|
||||
)
|
||||
candidate.finish_reason = Candidate.FinishReason.STOP
|
||||
|
||||
response = GenerateContentResponse()
|
||||
response.candidates.append(candidate)
|
||||
response.usage_metadata = GenerateContentResponse.UsageMetadata(
|
||||
prompt_token_count=0,
|
||||
candidates_token_count=0,
|
||||
total_token_count=0,
|
||||
)
|
||||
|
||||
return GenerationResponse._from_gapic(response)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_vertex_ai_streaming_chat_completion_response() -> AsyncIterable[GenerationResponse]:
|
||||
"""Mock Vertex AI streaming Chat Completion response."""
|
||||
candidate = Candidate()
|
||||
candidate.index = 0
|
||||
candidate.content = Content(role="user", parts=[Part(text="Test content")])
|
||||
candidate.finish_reason = Candidate.FinishReason.STOP
|
||||
|
||||
response = GenerateContentResponse()
|
||||
response.candidates.append(candidate)
|
||||
response.usage_metadata = GenerateContentResponse.UsageMetadata(
|
||||
prompt_token_count=0,
|
||||
candidates_token_count=0,
|
||||
total_token_count=0,
|
||||
)
|
||||
|
||||
iterable = MagicMock(spec=AsyncGenerator)
|
||||
iterable.__aiter__.return_value = [GenerationResponse._from_gapic(response)]
|
||||
|
||||
return iterable
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_vertex_ai_streaming_chat_completion_response_with_tool_call() -> AsyncIterable[GenerationResponse]:
|
||||
"""Mock Vertex AI streaming Chat Completion response."""
|
||||
candidate = Candidate()
|
||||
candidate.index = 0
|
||||
candidate.content = Content(
|
||||
role="user",
|
||||
parts=[
|
||||
Part(
|
||||
function_call=FunctionCall(
|
||||
name="getLightStatus",
|
||||
args={"arg1": "test_value"},
|
||||
)
|
||||
)
|
||||
],
|
||||
)
|
||||
candidate.finish_reason = Candidate.FinishReason.STOP
|
||||
|
||||
response = GenerateContentResponse()
|
||||
response.candidates.append(candidate)
|
||||
response.usage_metadata = GenerateContentResponse.UsageMetadata(
|
||||
prompt_token_count=0,
|
||||
candidates_token_count=0,
|
||||
total_token_count=0,
|
||||
)
|
||||
|
||||
iterable = MagicMock(spec=AsyncGenerator)
|
||||
iterable.__aiter__.return_value = [GenerationResponse._from_gapic(response)]
|
||||
|
||||
return iterable
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_vertex_ai_text_completion_response() -> GenerationResponse:
|
||||
"""Mock Vertex AI Text Completion response."""
|
||||
candidate = Candidate()
|
||||
candidate.index = 0
|
||||
candidate.content = Content(parts=[Part(text="Test content")])
|
||||
|
||||
response = GenerateContentResponse()
|
||||
response.candidates.append(candidate)
|
||||
response.usage_metadata = GenerateContentResponse.UsageMetadata(
|
||||
prompt_token_count=0,
|
||||
candidates_token_count=0,
|
||||
total_token_count=0,
|
||||
)
|
||||
|
||||
return GenerationResponse._from_gapic(response)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_vertex_ai_streaming_text_completion_response() -> AsyncIterable[GenerationResponse]:
|
||||
"""Mock Vertex AI streaming Text Completion response."""
|
||||
candidate = Candidate()
|
||||
candidate.index = 0
|
||||
candidate.content = Content(parts=[Part(text="Test content")])
|
||||
|
||||
response = GenerateContentResponse()
|
||||
response.candidates.append(candidate)
|
||||
response.usage_metadata = GenerateContentResponse.UsageMetadata(
|
||||
prompt_token_count=0,
|
||||
candidates_token_count=0,
|
||||
total_token_count=0,
|
||||
)
|
||||
|
||||
iterable = MagicMock(spec=AsyncGenerator)
|
||||
iterable.__aiter__.return_value = [GenerationResponse._from_gapic(response)]
|
||||
|
||||
return iterable
|
||||
|
||||
|
||||
class MockTextEmbeddingModel:
|
||||
async def get_embeddings_async(
|
||||
self,
|
||||
texts: list[str],
|
||||
*,
|
||||
auto_truncate: bool = True,
|
||||
output_dimensionality: int | None = None,
|
||||
) -> list[TextEmbedding]:
|
||||
pass
|
||||
+545
@@ -0,0 +1,545 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from vertexai.generative_models import Content, GenerativeModel
|
||||
|
||||
from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior
|
||||
from semantic_kernel.connectors.ai.google.vertex_ai.services.vertex_ai_chat_completion import VertexAIChatCompletion
|
||||
from semantic_kernel.connectors.ai.google.vertex_ai.vertex_ai_prompt_execution_settings import (
|
||||
VertexAIChatPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.google.vertex_ai.vertex_ai_settings import VertexAISettings
|
||||
from semantic_kernel.contents.chat_history import ChatHistory
|
||||
from semantic_kernel.contents.chat_message_content import ChatMessageContent
|
||||
from semantic_kernel.contents.utils.finish_reason import FinishReason
|
||||
from semantic_kernel.exceptions.service_exceptions import (
|
||||
ServiceInitializationError,
|
||||
ServiceInvalidExecutionSettingsError,
|
||||
)
|
||||
from semantic_kernel.kernel import Kernel
|
||||
|
||||
|
||||
# region init
|
||||
def test_vertex_ai_chat_completion_init(vertex_ai_unit_test_env) -> None:
|
||||
"""Test initialization of VertexAIChatCompletion"""
|
||||
model_id = vertex_ai_unit_test_env["VERTEX_AI_GEMINI_MODEL_ID"]
|
||||
project_id = vertex_ai_unit_test_env["VERTEX_AI_PROJECT_ID"]
|
||||
vertex_ai_chat_completion = VertexAIChatCompletion()
|
||||
|
||||
assert vertex_ai_chat_completion.ai_model_id == model_id
|
||||
assert vertex_ai_chat_completion.service_id == model_id
|
||||
|
||||
assert isinstance(vertex_ai_chat_completion.service_settings, VertexAISettings)
|
||||
assert vertex_ai_chat_completion.service_settings.gemini_model_id == model_id
|
||||
assert vertex_ai_chat_completion.service_settings.project_id == project_id
|
||||
|
||||
|
||||
def test_vertex_ai_chat_completion_init_with_service_id(vertex_ai_unit_test_env, service_id) -> None:
|
||||
"""Test initialization of VertexAIChatCompletion with a service id that is not the model id"""
|
||||
vertex_ai_chat_completion = VertexAIChatCompletion(service_id=service_id)
|
||||
|
||||
assert vertex_ai_chat_completion.service_id == service_id
|
||||
|
||||
|
||||
def test_vertex_ai_chat_completion_init_with_model_id_in_argument(vertex_ai_unit_test_env) -> None:
|
||||
"""Test initialization of VertexAIChatCompletion with model id in argument"""
|
||||
vertex_ai_chat_completion = VertexAIChatCompletion(gemini_model_id="custom_model_id")
|
||||
|
||||
assert vertex_ai_chat_completion.ai_model_id == "custom_model_id"
|
||||
assert vertex_ai_chat_completion.service_id == "custom_model_id"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["VERTEX_AI_GEMINI_MODEL_ID"]], indirect=True)
|
||||
def test_vertex_ai_chat_completion_init_with_empty_model_id(vertex_ai_unit_test_env) -> None:
|
||||
"""Test initialization of VertexAIChatCompletion with an empty model id"""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
VertexAIChatCompletion(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["VERTEX_AI_PROJECT_ID"]], indirect=True)
|
||||
def test_vertex_ai_chat_completion_init_with_empty_project_id(vertex_ai_unit_test_env) -> None:
|
||||
"""Test initialization of VertexAIChatCompletion with an empty project id"""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
VertexAIChatCompletion(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
def test_prompt_execution_settings_class(vertex_ai_unit_test_env) -> None:
|
||||
vertex_ai_chat_completion = VertexAIChatCompletion()
|
||||
assert vertex_ai_chat_completion.get_prompt_execution_settings_class() == VertexAIChatPromptExecutionSettings
|
||||
|
||||
|
||||
# endregion init
|
||||
|
||||
|
||||
# region chat completion
|
||||
|
||||
|
||||
@patch.object(GenerativeModel, "generate_content_async", new_callable=AsyncMock)
|
||||
async def test_vertex_ai_chat_completion(
|
||||
mock_vertex_ai_model_generate_content_async,
|
||||
vertex_ai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_vertex_ai_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test chat completion with VertexAIChatCompletion"""
|
||||
settings = VertexAIChatPromptExecutionSettings()
|
||||
|
||||
mock_vertex_ai_model_generate_content_async.return_value = mock_vertex_ai_chat_completion_response
|
||||
|
||||
vertex_ai_chat_completion = VertexAIChatCompletion()
|
||||
responses: list[ChatMessageContent] = await vertex_ai_chat_completion.get_chat_message_contents(
|
||||
chat_history, settings
|
||||
)
|
||||
|
||||
# Verify the call was made once
|
||||
mock_vertex_ai_model_generate_content_async.assert_called_once()
|
||||
|
||||
# Get the actual call arguments
|
||||
call_args = mock_vertex_ai_model_generate_content_async.call_args
|
||||
|
||||
# Verify the contents
|
||||
contents = call_args.kwargs["contents"]
|
||||
assert len(contents) == 1
|
||||
assert contents[0].role == "user"
|
||||
assert len(contents[0].parts) == 1
|
||||
assert contents[0].parts[0].text == "test_prompt"
|
||||
|
||||
# Verify other arguments
|
||||
assert call_args.kwargs["generation_config"] == settings.prepare_settings_dict()
|
||||
assert call_args.kwargs["tools"] is None
|
||||
assert call_args.kwargs["tool_config"] is None
|
||||
assert len(responses) == 1
|
||||
assert responses[0].role == "assistant"
|
||||
assert responses[0].content == mock_vertex_ai_chat_completion_response.candidates[0].content.parts[0].text
|
||||
assert responses[0].finish_reason == FinishReason.STOP
|
||||
assert "usage" in responses[0].metadata
|
||||
assert "prompt_feedback" in responses[0].metadata
|
||||
assert responses[0].inner_content == mock_vertex_ai_chat_completion_response
|
||||
|
||||
|
||||
async def test_vertex_ai_chat_completion_with_function_choice_behavior_fail_verification(
|
||||
chat_history: ChatHistory,
|
||||
vertex_ai_unit_test_env,
|
||||
) -> None:
|
||||
"""Test completion of VertexAIChatCompletion with function choice behavior expect verification failure"""
|
||||
|
||||
# Missing kernel
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
settings = VertexAIChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
|
||||
vertex_ai_chat_completion = VertexAIChatCompletion()
|
||||
|
||||
await vertex_ai_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=settings,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(GenerativeModel, "generate_content_async", new_callable=AsyncMock)
|
||||
async def test_vertex_ai_chat_completion_with_function_choice_behavior(
|
||||
mock_vertex_ai_model_generate_content_async,
|
||||
vertex_ai_unit_test_env,
|
||||
kernel,
|
||||
chat_history: ChatHistory,
|
||||
mock_vertex_ai_chat_completion_response_with_tool_call,
|
||||
) -> None:
|
||||
"""Test completion of VertexAIChatCompletion with function choice behavior"""
|
||||
mock_vertex_ai_model_generate_content_async.return_value = mock_vertex_ai_chat_completion_response_with_tool_call
|
||||
|
||||
settings = VertexAIChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
settings.function_choice_behavior.maximum_auto_invoke_attempts = 1
|
||||
|
||||
vertex_ai_chat_completion = VertexAIChatCompletion()
|
||||
|
||||
responses = await vertex_ai_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=settings,
|
||||
kernel=kernel,
|
||||
)
|
||||
|
||||
# The function should be called twice:
|
||||
# One for the tool call and one for the last completion
|
||||
# after the maximum_auto_invoke_attempts is reached
|
||||
assert mock_vertex_ai_model_generate_content_async.call_count == 2
|
||||
assert len(responses) == 1
|
||||
assert responses[0].role == "assistant"
|
||||
# Google doesn't return STOP as the finish reason for tool calls
|
||||
assert responses[0].finish_reason == FinishReason.STOP
|
||||
|
||||
|
||||
@patch.object(GenerativeModel, "generate_content_async", new_callable=AsyncMock)
|
||||
async def test_vertex_ai_chat_completion_with_function_choice_behavior_no_tool_call(
|
||||
mock_vertex_ai_model_generate_content_async,
|
||||
vertex_ai_unit_test_env,
|
||||
kernel,
|
||||
chat_history: ChatHistory,
|
||||
mock_vertex_ai_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test completion of VertexAIChatCompletion with function choice behavior but no tool call returned"""
|
||||
mock_vertex_ai_model_generate_content_async.return_value = mock_vertex_ai_chat_completion_response
|
||||
|
||||
settings = VertexAIChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
settings.function_choice_behavior.maximum_auto_invoke_attempts = 1
|
||||
|
||||
vertex_ai_chat_completion = VertexAIChatCompletion()
|
||||
|
||||
responses = await vertex_ai_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=settings,
|
||||
kernel=kernel,
|
||||
)
|
||||
|
||||
# Verify the call was made once
|
||||
mock_vertex_ai_model_generate_content_async.assert_awaited_once()
|
||||
|
||||
# Get the actual call arguments
|
||||
call_args = mock_vertex_ai_model_generate_content_async.await_args
|
||||
|
||||
# Verify the contents
|
||||
contents = call_args.kwargs["contents"]
|
||||
assert len(contents) == 1
|
||||
assert contents[0].role == "user"
|
||||
assert len(contents[0].parts) == 1
|
||||
assert contents[0].parts[0].text == "test_prompt"
|
||||
|
||||
# Verify other arguments
|
||||
assert call_args.kwargs["generation_config"] == settings.prepare_settings_dict()
|
||||
assert call_args.kwargs["tools"] is None
|
||||
assert call_args.kwargs["tool_config"] is None
|
||||
assert len(responses) == 1
|
||||
assert responses[0].role == "assistant"
|
||||
assert responses[0].content == mock_vertex_ai_chat_completion_response.candidates[0].content.parts[0].text
|
||||
|
||||
|
||||
# endregion chat completion
|
||||
|
||||
|
||||
# region streaming chat completion
|
||||
|
||||
|
||||
@patch.object(GenerativeModel, "generate_content_async", new_callable=AsyncMock)
|
||||
async def test_vertex_ai_streaming_chat_completion(
|
||||
mock_vertex_ai_model_generate_content_async,
|
||||
vertex_ai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_vertex_ai_streaming_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test chat completion with VertexAIChatCompletion"""
|
||||
settings = VertexAIChatPromptExecutionSettings()
|
||||
|
||||
mock_vertex_ai_model_generate_content_async.return_value = mock_vertex_ai_streaming_chat_completion_response
|
||||
|
||||
vertex_ai_chat_completion = VertexAIChatCompletion()
|
||||
async for messages in vertex_ai_chat_completion.get_streaming_chat_message_contents(chat_history, settings):
|
||||
assert len(messages) == 1
|
||||
assert messages[0].role == "assistant"
|
||||
assert messages[0].finish_reason == FinishReason.STOP
|
||||
assert "usage" in messages[0].metadata
|
||||
assert "prompt_feedback" in messages[0].metadata
|
||||
|
||||
# Verify the call was made once
|
||||
mock_vertex_ai_model_generate_content_async.assert_called_once()
|
||||
|
||||
# Get the actual call arguments
|
||||
call_args = mock_vertex_ai_model_generate_content_async.call_args
|
||||
|
||||
# Verify the contents
|
||||
contents = call_args.kwargs["contents"]
|
||||
assert len(contents) == 1
|
||||
assert contents[0].role == "user"
|
||||
assert len(contents[0].parts) == 1
|
||||
assert contents[0].parts[0].text == "test_prompt"
|
||||
|
||||
# Verify other arguments
|
||||
assert call_args.kwargs["generation_config"] == settings.prepare_settings_dict()
|
||||
assert call_args.kwargs["tools"] is None
|
||||
assert call_args.kwargs["tool_config"] is None
|
||||
assert call_args.kwargs["stream"] is True
|
||||
|
||||
|
||||
async def test_vertex_ai_streaming_chat_completion_with_function_choice_behavior_fail_verification(
|
||||
chat_history: ChatHistory,
|
||||
vertex_ai_unit_test_env,
|
||||
) -> None:
|
||||
"""Test streaming chat completion of VertexAIChatCompletion with function choice
|
||||
behavior expect verification failure"""
|
||||
|
||||
# Missing kernel
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
settings = VertexAIChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
|
||||
vertex_ai_chat_completion = VertexAIChatCompletion()
|
||||
|
||||
async for _ in vertex_ai_chat_completion.get_streaming_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=settings,
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
@patch.object(GenerativeModel, "generate_content_async", new_callable=AsyncMock)
|
||||
async def test_vertex_ai_streaming_chat_completion_with_function_choice_behavior(
|
||||
mock_vertex_ai_model_generate_content_async,
|
||||
vertex_ai_unit_test_env,
|
||||
kernel: Kernel,
|
||||
chat_history: ChatHistory,
|
||||
mock_vertex_ai_streaming_chat_completion_response_with_tool_call,
|
||||
decorated_native_function,
|
||||
) -> None:
|
||||
"""Test streaming chat completion of VertexAIChatCompletion with function choice behavior"""
|
||||
mock_vertex_ai_model_generate_content_async.return_value = (
|
||||
mock_vertex_ai_streaming_chat_completion_response_with_tool_call
|
||||
)
|
||||
|
||||
settings = VertexAIChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
settings.function_choice_behavior.maximum_auto_invoke_attempts = 1
|
||||
|
||||
vertex_ai_chat_completion = VertexAIChatCompletion()
|
||||
|
||||
kernel.add_function(plugin_name="TestPlugin", function=decorated_native_function)
|
||||
|
||||
all_messages = []
|
||||
async for messages in vertex_ai_chat_completion.get_streaming_chat_message_contents(
|
||||
chat_history,
|
||||
settings,
|
||||
kernel=kernel,
|
||||
):
|
||||
all_messages.extend(messages)
|
||||
|
||||
assert len(all_messages) == 2, f"Expected 2 messages, got {len(all_messages)}"
|
||||
|
||||
# Validate the first message
|
||||
assert all_messages[0].role == "assistant", f"Unexpected role for first message: {all_messages[0].role}"
|
||||
assert all_messages[0].content == "", f"Unexpected content for first message: {all_messages[0].content}"
|
||||
assert all_messages[0].finish_reason == FinishReason.STOP, (
|
||||
f"Unexpected finish reason for first message: {all_messages[0].finish_reason}"
|
||||
)
|
||||
|
||||
# Validate the second message
|
||||
assert all_messages[1].role == "tool", f"Unexpected role for second message: {all_messages[1].role}"
|
||||
assert all_messages[1].content == "", f"Unexpected content for second message: {all_messages[1].content}"
|
||||
assert all_messages[1].finish_reason is None
|
||||
|
||||
|
||||
@patch.object(GenerativeModel, "generate_content_async", new_callable=AsyncMock)
|
||||
async def test_vertex_ai_streaming_chat_completion_with_function_choice_behavior_no_tool_call(
|
||||
mock_vertex_ai_model_generate_content_async,
|
||||
vertex_ai_unit_test_env,
|
||||
kernel,
|
||||
chat_history: ChatHistory,
|
||||
mock_vertex_ai_streaming_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test completion of VertexAIChatCompletion with function choice behavior but no tool call returned"""
|
||||
mock_vertex_ai_model_generate_content_async.return_value = mock_vertex_ai_streaming_chat_completion_response
|
||||
|
||||
settings = VertexAIChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
)
|
||||
settings.function_choice_behavior.maximum_auto_invoke_attempts = 1
|
||||
|
||||
vertex_ai_chat_completion = VertexAIChatCompletion()
|
||||
|
||||
async for messages in vertex_ai_chat_completion.get_streaming_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=settings,
|
||||
kernel=kernel,
|
||||
):
|
||||
assert len(messages) == 1
|
||||
assert messages[0].role == "assistant"
|
||||
|
||||
# Verify the call was made once
|
||||
mock_vertex_ai_model_generate_content_async.assert_awaited_once()
|
||||
|
||||
# Get the actual call arguments
|
||||
call_args = mock_vertex_ai_model_generate_content_async.await_args
|
||||
|
||||
# Verify the contents
|
||||
contents = call_args.kwargs["contents"]
|
||||
assert len(contents) == 1
|
||||
assert contents[0].role == "user"
|
||||
assert len(contents[0].parts) == 1
|
||||
assert contents[0].parts[0].text == "test_prompt"
|
||||
|
||||
# Verify other arguments
|
||||
assert call_args.kwargs["generation_config"] == settings.prepare_settings_dict()
|
||||
assert call_args.kwargs["tools"] is None
|
||||
assert call_args.kwargs["tool_config"] is None
|
||||
assert call_args.kwargs["stream"] is True
|
||||
|
||||
|
||||
# endregion streaming chat completion
|
||||
|
||||
|
||||
def test_vertex_ai_chat_completion_parse_chat_history_correctly(vertex_ai_unit_test_env) -> None:
|
||||
"""Test _prepare_chat_history_for_request method"""
|
||||
vertex_ai_chat_completion = VertexAIChatCompletion()
|
||||
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_system_message("test_system_message")
|
||||
chat_history.add_user_message("test_user_message")
|
||||
chat_history.add_assistant_message("test_assistant_message")
|
||||
|
||||
parsed_chat_history = vertex_ai_chat_completion._prepare_chat_history_for_request(chat_history)
|
||||
|
||||
assert isinstance(parsed_chat_history, list)
|
||||
# System message should be ignored
|
||||
assert len(parsed_chat_history) == 2
|
||||
assert all(isinstance(message, Content) for message in parsed_chat_history)
|
||||
assert parsed_chat_history[0].role == "user"
|
||||
assert parsed_chat_history[0].parts[0].text == "test_user_message"
|
||||
assert parsed_chat_history[1].role == "model"
|
||||
assert parsed_chat_history[1].parts[0].text == "test_assistant_message"
|
||||
|
||||
|
||||
# region thought_signature deserialization tests
|
||||
|
||||
|
||||
def test_create_chat_message_content_with_thought_signature(vertex_ai_unit_test_env) -> None:
|
||||
"""Test that thought_signature from a Part dict is deserialized into FunctionCallContent.metadata."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from google.cloud.aiplatform_v1beta1.types.content import Candidate as GapicCandidate
|
||||
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
|
||||
thought_sig_value = "vertex-test-thought-sig"
|
||||
|
||||
# Create a mock Part whose to_dict() returns thought_signature
|
||||
mock_part = MagicMock()
|
||||
mock_part.text = None
|
||||
mock_part.to_dict.return_value = {
|
||||
"function_call": {"name": "test_function", "args": {"key": "value"}},
|
||||
"thought_signature": thought_sig_value,
|
||||
}
|
||||
mock_part.function_call.name = "test_function"
|
||||
mock_part.function_call.args = {"key": "value"}
|
||||
|
||||
# Build a mock candidate with the mock part
|
||||
mock_candidate = MagicMock()
|
||||
mock_candidate.index = 0
|
||||
mock_candidate.content.parts = [mock_part]
|
||||
mock_candidate.finish_reason = GapicCandidate.FinishReason.STOP
|
||||
|
||||
# Build a mock response
|
||||
mock_response = MagicMock()
|
||||
|
||||
completion = VertexAIChatCompletion()
|
||||
result = completion._create_chat_message_content(mock_response, mock_candidate)
|
||||
|
||||
fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
|
||||
assert len(fc_items) == 1
|
||||
assert fc_items[0].metadata is not None
|
||||
assert fc_items[0].metadata["thought_signature"] == thought_sig_value
|
||||
|
||||
|
||||
def test_create_chat_message_content_without_thought_signature(vertex_ai_unit_test_env) -> None:
|
||||
"""Test that FunctionCallContent works when Part dict has no thought_signature."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from google.cloud.aiplatform_v1beta1.types.content import Candidate as GapicCandidate
|
||||
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
|
||||
mock_part = MagicMock()
|
||||
mock_part.text = None
|
||||
mock_part.to_dict.return_value = {
|
||||
"function_call": {"name": "test_function", "args": {"key": "value"}},
|
||||
}
|
||||
mock_part.function_call.name = "test_function"
|
||||
mock_part.function_call.args = {"key": "value"}
|
||||
|
||||
mock_candidate = MagicMock()
|
||||
mock_candidate.index = 0
|
||||
mock_candidate.content.parts = [mock_part]
|
||||
mock_candidate.finish_reason = GapicCandidate.FinishReason.STOP
|
||||
|
||||
mock_response = MagicMock()
|
||||
|
||||
completion = VertexAIChatCompletion()
|
||||
result = completion._create_chat_message_content(mock_response, mock_candidate)
|
||||
|
||||
fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
|
||||
assert len(fc_items) == 1
|
||||
assert "thought_signature" not in fc_items[0].metadata
|
||||
|
||||
|
||||
def test_create_streaming_chat_message_content_with_thought_signature(vertex_ai_unit_test_env) -> None:
|
||||
"""Test that thought_signature from a Part dict is deserialized in streaming path."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from google.cloud.aiplatform_v1beta1.types.content import Candidate as GapicCandidate
|
||||
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
|
||||
thought_sig_value = "vertex-streaming-thought-sig"
|
||||
|
||||
mock_part = MagicMock()
|
||||
mock_part.text = None
|
||||
mock_part.to_dict.return_value = {
|
||||
"function_call": {"name": "stream_func", "args": {"a": "b"}},
|
||||
"thought_signature": thought_sig_value,
|
||||
}
|
||||
mock_part.function_call.name = "stream_func"
|
||||
mock_part.function_call.args = {"a": "b"}
|
||||
|
||||
mock_candidate = MagicMock()
|
||||
mock_candidate.index = 0
|
||||
mock_candidate.content.parts = [mock_part]
|
||||
mock_candidate.finish_reason = GapicCandidate.FinishReason.STOP
|
||||
|
||||
mock_chunk = MagicMock()
|
||||
|
||||
completion = VertexAIChatCompletion()
|
||||
result = completion._create_streaming_chat_message_content(mock_chunk, mock_candidate, function_invoke_attempt=0)
|
||||
|
||||
fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
|
||||
assert len(fc_items) == 1
|
||||
assert fc_items[0].metadata is not None
|
||||
assert fc_items[0].metadata["thought_signature"] == thought_sig_value
|
||||
|
||||
|
||||
def test_create_streaming_chat_message_content_without_thought_signature(vertex_ai_unit_test_env) -> None:
|
||||
"""Test that streaming FunctionCallContent works when Part dict lacks thought_signature."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from google.cloud.aiplatform_v1beta1.types.content import Candidate as GapicCandidate
|
||||
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
|
||||
mock_part = MagicMock()
|
||||
mock_part.text = None
|
||||
mock_part.to_dict.return_value = {
|
||||
"function_call": {"name": "stream_func", "args": {"a": "b"}},
|
||||
}
|
||||
mock_part.function_call.name = "stream_func"
|
||||
mock_part.function_call.args = {"a": "b"}
|
||||
|
||||
mock_candidate = MagicMock()
|
||||
mock_candidate.index = 0
|
||||
mock_candidate.content.parts = [mock_part]
|
||||
mock_candidate.finish_reason = GapicCandidate.FinishReason.STOP
|
||||
|
||||
mock_chunk = MagicMock()
|
||||
|
||||
completion = VertexAIChatCompletion()
|
||||
result = completion._create_streaming_chat_message_content(mock_chunk, mock_candidate, function_invoke_attempt=0)
|
||||
|
||||
fc_items = [item for item in result.items if isinstance(item, FunctionCallContent)]
|
||||
assert len(fc_items) == 1
|
||||
assert "thought_signature" not in fc_items[0].metadata
|
||||
|
||||
|
||||
# endregion thought_signature deserialization tests
|
||||
+130
@@ -0,0 +1,130 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from vertexai.generative_models import GenerativeModel
|
||||
|
||||
from semantic_kernel.connectors.ai.google.vertex_ai.services.vertex_ai_text_completion import VertexAITextCompletion
|
||||
from semantic_kernel.connectors.ai.google.vertex_ai.vertex_ai_prompt_execution_settings import (
|
||||
VertexAITextPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.google.vertex_ai.vertex_ai_settings import VertexAISettings
|
||||
from semantic_kernel.contents.text_content import TextContent
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
|
||||
|
||||
|
||||
# region init
|
||||
def test_vertex_ai_text_completion_init(vertex_ai_unit_test_env) -> None:
|
||||
"""Test initialization of VertexAITextCompletion"""
|
||||
model_id = vertex_ai_unit_test_env["VERTEX_AI_GEMINI_MODEL_ID"]
|
||||
project_id = vertex_ai_unit_test_env["VERTEX_AI_PROJECT_ID"]
|
||||
vertex_ai_text_completion = VertexAITextCompletion()
|
||||
|
||||
assert vertex_ai_text_completion.ai_model_id == model_id
|
||||
assert vertex_ai_text_completion.service_id == model_id
|
||||
|
||||
assert isinstance(vertex_ai_text_completion.service_settings, VertexAISettings)
|
||||
assert vertex_ai_text_completion.service_settings.gemini_model_id == model_id
|
||||
assert vertex_ai_text_completion.service_settings.project_id == project_id
|
||||
|
||||
|
||||
def test_vertex_ai_text_completion_init_with_service_id(vertex_ai_unit_test_env, service_id) -> None:
|
||||
"""Test initialization of VertexAITextCompletion with a service id that is not the model id"""
|
||||
vertex_ai_text_completion = VertexAITextCompletion(service_id=service_id)
|
||||
|
||||
assert vertex_ai_text_completion.service_id == service_id
|
||||
|
||||
|
||||
def test_vertex_ai_text_completion_init_with_model_id_in_argument(vertex_ai_unit_test_env) -> None:
|
||||
"""Test initialization of VertexAIChatCompletion with model id in argument"""
|
||||
vertex_ai_text_completion = VertexAITextCompletion(gemini_model_id="custom_model_id")
|
||||
|
||||
assert vertex_ai_text_completion.ai_model_id == "custom_model_id"
|
||||
assert vertex_ai_text_completion.service_id == "custom_model_id"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["VERTEX_AI_GEMINI_MODEL_ID"]], indirect=True)
|
||||
def test_vertex_ai_text_completion_init_with_empty_model_id(vertex_ai_unit_test_env) -> None:
|
||||
"""Test initialization of VertexAITextCompletion with an empty model id"""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
VertexAITextCompletion(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["VERTEX_AI_PROJECT_ID"]], indirect=True)
|
||||
def test_vertex_ai_text_completion_init_with_empty_project_id(vertex_ai_unit_test_env) -> None:
|
||||
"""Test initialization of VertexAITextCompletion with an empty project id"""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
VertexAITextCompletion(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
def test_prompt_execution_settings_class(vertex_ai_unit_test_env) -> None:
|
||||
vertex_ai_text_completion = VertexAITextCompletion()
|
||||
assert vertex_ai_text_completion.get_prompt_execution_settings_class() == VertexAITextPromptExecutionSettings
|
||||
|
||||
|
||||
# endregion init
|
||||
|
||||
|
||||
# region text completion
|
||||
|
||||
|
||||
@patch.object(GenerativeModel, "generate_content_async", new_callable=AsyncMock)
|
||||
async def test_vertex_ai_text_completion(
|
||||
mock_vertex_ai_model_generate_content_async,
|
||||
vertex_ai_unit_test_env,
|
||||
prompt: str,
|
||||
mock_vertex_ai_text_completion_response,
|
||||
) -> None:
|
||||
"""Test text completion with VertexAITextCompletion"""
|
||||
settings = VertexAITextPromptExecutionSettings()
|
||||
|
||||
mock_vertex_ai_model_generate_content_async.return_value = mock_vertex_ai_text_completion_response
|
||||
|
||||
vertex_ai_text_completion = VertexAITextCompletion()
|
||||
responses: list[TextContent] = await vertex_ai_text_completion.get_text_contents(prompt, settings)
|
||||
|
||||
mock_vertex_ai_model_generate_content_async.assert_called_once_with(
|
||||
contents=prompt,
|
||||
generation_config=settings.prepare_settings_dict(),
|
||||
)
|
||||
assert len(responses) == 1
|
||||
assert responses[0].text == mock_vertex_ai_text_completion_response.candidates[0].content.parts[0].text
|
||||
assert "usage" in responses[0].metadata
|
||||
assert "prompt_feedback" in responses[0].metadata
|
||||
assert responses[0].inner_content == mock_vertex_ai_text_completion_response
|
||||
|
||||
|
||||
# endregion text completion
|
||||
|
||||
|
||||
# region streaming text completion
|
||||
|
||||
|
||||
@patch.object(GenerativeModel, "generate_content_async", new_callable=AsyncMock)
|
||||
async def test_vertex_ai_streaming_text_completion(
|
||||
mock_vertex_ai_model_generate_content_async,
|
||||
vertex_ai_unit_test_env,
|
||||
prompt: str,
|
||||
mock_vertex_ai_streaming_text_completion_response,
|
||||
) -> None:
|
||||
"""Test streaming text completion with VertexAITextCompletion"""
|
||||
settings = VertexAITextPromptExecutionSettings()
|
||||
|
||||
mock_vertex_ai_model_generate_content_async.return_value = mock_vertex_ai_streaming_text_completion_response
|
||||
|
||||
vertex_ai_text_completion = VertexAITextCompletion()
|
||||
async for chunks in vertex_ai_text_completion.get_streaming_text_contents(prompt, settings):
|
||||
assert len(chunks) == 1
|
||||
assert "usage" in chunks[0].metadata
|
||||
assert "prompt_feedback" in chunks[0].metadata
|
||||
|
||||
mock_vertex_ai_model_generate_content_async.assert_called_once_with(
|
||||
contents=prompt,
|
||||
generation_config=settings.prepare_settings_dict(),
|
||||
stream=True,
|
||||
)
|
||||
|
||||
|
||||
# endregion streaming text completion
|
||||
+160
@@ -0,0 +1,160 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from numpy import array, ndarray
|
||||
from vertexai.language_models import TextEmbedding, TextEmbeddingModel
|
||||
|
||||
from semantic_kernel.connectors.ai.google.vertex_ai.services.vertex_ai_text_embedding import VertexAITextEmbedding
|
||||
from semantic_kernel.connectors.ai.google.vertex_ai.vertex_ai_prompt_execution_settings import (
|
||||
VertexAIEmbeddingPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.google.vertex_ai.vertex_ai_settings import VertexAISettings
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
|
||||
from tests.unit.connectors.ai.google.vertex_ai.conftest import MockTextEmbeddingModel
|
||||
|
||||
|
||||
# region init
|
||||
def test_vertex_ai_text_embedding_init(vertex_ai_unit_test_env) -> None:
|
||||
"""Test initialization of VertexAITextEmbedding"""
|
||||
model_id = vertex_ai_unit_test_env["VERTEX_AI_EMBEDDING_MODEL_ID"]
|
||||
project_id = vertex_ai_unit_test_env["VERTEX_AI_PROJECT_ID"]
|
||||
vertex_ai_text_embedding = VertexAITextEmbedding()
|
||||
|
||||
assert vertex_ai_text_embedding.ai_model_id == model_id
|
||||
assert vertex_ai_text_embedding.service_id == model_id
|
||||
|
||||
assert isinstance(vertex_ai_text_embedding.service_settings, VertexAISettings)
|
||||
assert vertex_ai_text_embedding.service_settings.embedding_model_id == model_id
|
||||
assert vertex_ai_text_embedding.service_settings.project_id == project_id
|
||||
|
||||
|
||||
def test_vertex_ai_text_embedding_init_with_service_id(vertex_ai_unit_test_env, service_id) -> None:
|
||||
"""Test initialization of VertexAITextEmbedding with a service id that is not the model id"""
|
||||
vertex_ai_text_embedding = VertexAITextEmbedding(service_id=service_id)
|
||||
|
||||
assert vertex_ai_text_embedding.service_id == service_id
|
||||
|
||||
|
||||
def test_vertex_ai_text_embedding_init_with_model_id_in_argument(vertex_ai_unit_test_env) -> None:
|
||||
"""Test initialization of VertexAITextEmbedding with model id in argument"""
|
||||
vertex_ai_chat_completion = VertexAITextEmbedding(embedding_model_id="custom_model_id")
|
||||
|
||||
assert vertex_ai_chat_completion.ai_model_id == "custom_model_id"
|
||||
assert vertex_ai_chat_completion.service_id == "custom_model_id"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["VERTEX_AI_EMBEDDING_MODEL_ID"]], indirect=True)
|
||||
def test_vertex_ai_text_embedding_init_with_empty_model_id(vertex_ai_unit_test_env) -> None:
|
||||
"""Test initialization of VertexAITextEmbedding with an empty model id"""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
VertexAITextEmbedding(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["VERTEX_AI_PROJECT_ID"]], indirect=True)
|
||||
def test_vertex_ai_text_embedding_init_with_empty_project_id(vertex_ai_unit_test_env) -> None:
|
||||
"""Test initialization of VertexAITextEmbedding with an empty project id"""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
VertexAITextEmbedding(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
def test_prompt_execution_settings_class(vertex_ai_unit_test_env) -> None:
|
||||
vertex_ai_text_embedding = VertexAITextEmbedding()
|
||||
assert vertex_ai_text_embedding.get_prompt_execution_settings_class() == VertexAIEmbeddingPromptExecutionSettings
|
||||
|
||||
|
||||
# endregion init
|
||||
|
||||
|
||||
@patch.object(TextEmbeddingModel, "from_pretrained")
|
||||
@patch.object(MockTextEmbeddingModel, "get_embeddings_async", new_callable=AsyncMock)
|
||||
async def test_embedding(mock_embedding_client, mock_from_pretrained, vertex_ai_unit_test_env, prompt):
|
||||
"""Test that the service initializes and generates embeddings correctly."""
|
||||
mock_from_pretrained.return_value = MockTextEmbeddingModel()
|
||||
mock_embedding_client.return_value = [TextEmbedding(values=[0.1, 0.2, 0.3])]
|
||||
|
||||
settings = VertexAIEmbeddingPromptExecutionSettings()
|
||||
|
||||
vertex_ai_text_embedding = VertexAITextEmbedding()
|
||||
response: ndarray = await vertex_ai_text_embedding.generate_embeddings(
|
||||
[prompt],
|
||||
settings=settings,
|
||||
)
|
||||
|
||||
assert len(response) == 1
|
||||
assert response.all() == array([0.1, 0.2, 0.3]).all()
|
||||
mock_embedding_client.assert_called_once_with([prompt])
|
||||
|
||||
|
||||
@patch.object(TextEmbeddingModel, "from_pretrained")
|
||||
@patch.object(MockTextEmbeddingModel, "get_embeddings_async", new_callable=AsyncMock)
|
||||
async def test_embedding_with_settings(mock_embedding_client, mock_from_pretrained, vertex_ai_unit_test_env, prompt):
|
||||
"""Test that the service initializes and generates embeddings correctly."""
|
||||
mock_from_pretrained.return_value = MockTextEmbeddingModel()
|
||||
mock_embedding_client.return_value = [TextEmbedding(values=[0.1, 0.2, 0.3])]
|
||||
|
||||
settings = VertexAIEmbeddingPromptExecutionSettings()
|
||||
settings.output_dimensionality = 3
|
||||
settings.auto_truncate = True
|
||||
|
||||
vertex_ai_text_embedding = VertexAITextEmbedding()
|
||||
response: ndarray = await vertex_ai_text_embedding.generate_embeddings(
|
||||
[prompt],
|
||||
settings=settings,
|
||||
)
|
||||
|
||||
assert len(response) == 1
|
||||
assert response.all() == array([0.1, 0.2, 0.3]).all()
|
||||
mock_embedding_client.assert_called_once_with(
|
||||
[prompt],
|
||||
**settings.prepare_settings_dict(),
|
||||
)
|
||||
|
||||
|
||||
@patch.object(TextEmbeddingModel, "from_pretrained")
|
||||
@patch.object(MockTextEmbeddingModel, "get_embeddings_async", new_callable=AsyncMock)
|
||||
async def test_embedding_without_settings(mock_embedding_client, mock_from_pretrained, vertex_ai_unit_test_env, prompt):
|
||||
"""Test that the service initializes and generates embeddings correctly without settings."""
|
||||
mock_from_pretrained.return_value = MockTextEmbeddingModel()
|
||||
mock_embedding_client.return_value = [TextEmbedding(values=[0.1, 0.2, 0.3])]
|
||||
|
||||
vertex_ai_text_embedding = VertexAITextEmbedding()
|
||||
response: ndarray = await vertex_ai_text_embedding.generate_embeddings([prompt])
|
||||
|
||||
assert len(response) == 1
|
||||
assert response.all() == array([0.1, 0.2, 0.3]).all()
|
||||
mock_embedding_client.assert_called_once_with([prompt])
|
||||
|
||||
|
||||
@patch.object(TextEmbeddingModel, "from_pretrained")
|
||||
@patch.object(MockTextEmbeddingModel, "get_embeddings_async", new_callable=AsyncMock)
|
||||
async def test_embedding_list_input(mock_embedding_client, mock_from_pretrained, vertex_ai_unit_test_env, prompt):
|
||||
"""Test that the service initializes and generates embeddings correctly with a list of prompts."""
|
||||
mock_from_pretrained.return_value = MockTextEmbeddingModel()
|
||||
mock_embedding_client.return_value = [TextEmbedding(values=[0.1, 0.2, 0.3]), TextEmbedding(values=[0.1, 0.2, 0.3])]
|
||||
|
||||
vertex_ai_text_embedding = VertexAITextEmbedding()
|
||||
response: ndarray = await vertex_ai_text_embedding.generate_embeddings([prompt, prompt])
|
||||
|
||||
assert len(response) == 2
|
||||
assert response.all() == array([[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]).all()
|
||||
mock_embedding_client.assert_called_once_with([prompt, prompt])
|
||||
|
||||
|
||||
@patch.object(TextEmbeddingModel, "from_pretrained")
|
||||
@patch.object(MockTextEmbeddingModel, "get_embeddings_async", new_callable=AsyncMock)
|
||||
async def test_raw_embedding(mock_embedding_client, mock_from_pretrained, vertex_ai_unit_test_env, prompt):
|
||||
"""Test that the service initializes and generates embeddings correctly."""
|
||||
mock_from_pretrained.return_value = MockTextEmbeddingModel()
|
||||
mock_embedding_client.return_value = [TextEmbedding(values=[0.1, 0.2, 0.3])]
|
||||
|
||||
settings = VertexAIEmbeddingPromptExecutionSettings()
|
||||
|
||||
vertex_ai_text_embedding = VertexAITextEmbedding()
|
||||
response: ndarray = await vertex_ai_text_embedding.generate_raw_embeddings([prompt], settings)
|
||||
|
||||
assert len(response) == 1
|
||||
assert response[0] == [0.1, 0.2, 0.3]
|
||||
mock_embedding_client.assert_called_once_with([prompt])
|
||||
@@ -0,0 +1,199 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import pytest
|
||||
from google.cloud.aiplatform_v1beta1.types.content import Candidate
|
||||
from vertexai.generative_models import FunctionDeclaration, Part
|
||||
|
||||
from semantic_kernel.connectors.ai.google.vertex_ai.services.utils import (
|
||||
finish_reason_from_vertex_ai_to_semantic_kernel,
|
||||
format_assistant_message,
|
||||
format_user_message,
|
||||
kernel_function_metadata_to_vertex_ai_function_call_format,
|
||||
)
|
||||
from semantic_kernel.contents.chat_message_content import ChatMessageContent
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
from semantic_kernel.contents.function_result_content import FunctionResultContent
|
||||
from semantic_kernel.contents.image_content import ImageContent
|
||||
from semantic_kernel.contents.text_content import TextContent
|
||||
from semantic_kernel.contents.utils.author_role import AuthorRole
|
||||
from semantic_kernel.contents.utils.finish_reason import FinishReason
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInvalidRequestError
|
||||
from semantic_kernel.functions.kernel_function_metadata import KernelFunctionMetadata
|
||||
from semantic_kernel.functions.kernel_parameter_metadata import KernelParameterMetadata
|
||||
|
||||
|
||||
def test_finish_reason_from_vertex_ai_to_semantic_kernel():
|
||||
"""Test finish_reason_from_vertex_ai_to_semantic_kernel."""
|
||||
assert finish_reason_from_vertex_ai_to_semantic_kernel(Candidate.FinishReason.STOP) == FinishReason.STOP
|
||||
assert finish_reason_from_vertex_ai_to_semantic_kernel(Candidate.FinishReason.MAX_TOKENS) == FinishReason.LENGTH
|
||||
assert finish_reason_from_vertex_ai_to_semantic_kernel(Candidate.FinishReason.SAFETY) == FinishReason.CONTENT_FILTER
|
||||
assert finish_reason_from_vertex_ai_to_semantic_kernel(Candidate.FinishReason.OTHER) is None
|
||||
|
||||
|
||||
def test_format_user_message():
|
||||
"""Test format_user_message."""
|
||||
user_message = ChatMessageContent(role=AuthorRole.USER, content="User message")
|
||||
formatted_user_message = format_user_message(user_message)
|
||||
|
||||
assert len(formatted_user_message) == 1
|
||||
assert isinstance(formatted_user_message[0], Part)
|
||||
assert formatted_user_message[0].text == "User message"
|
||||
|
||||
# Test with an image content
|
||||
image_content = ImageContent(data="image data", mime_type="image/png")
|
||||
user_message = ChatMessageContent(
|
||||
role=AuthorRole.USER,
|
||||
items=[
|
||||
TextContent(text="Text content"),
|
||||
image_content,
|
||||
],
|
||||
)
|
||||
formatted_user_message = format_user_message(user_message)
|
||||
|
||||
assert len(formatted_user_message) == 2
|
||||
assert isinstance(formatted_user_message[0], Part)
|
||||
assert formatted_user_message[0].text == "Text content"
|
||||
assert isinstance(formatted_user_message[1], Part)
|
||||
assert formatted_user_message[1].inline_data.mime_type == "image/png"
|
||||
assert formatted_user_message[1].inline_data.data == image_content.data
|
||||
|
||||
|
||||
def test_format_user_message_throws_with_unsupported_items() -> None:
|
||||
"""Test format_user_message with unsupported items."""
|
||||
# Test with unsupported items, any item other than TextContent and ImageContent should raise an error
|
||||
user_message = ChatMessageContent(
|
||||
role=AuthorRole.USER,
|
||||
items=[
|
||||
FunctionCallContent(),
|
||||
],
|
||||
)
|
||||
with pytest.raises(ServiceInvalidRequestError):
|
||||
format_user_message(user_message)
|
||||
|
||||
# Test with an ImageContent that has no data_uri
|
||||
user_message = ChatMessageContent(
|
||||
role=AuthorRole.USER,
|
||||
items=[
|
||||
ImageContent(data_uri=""),
|
||||
],
|
||||
)
|
||||
with pytest.raises(ServiceInvalidRequestError):
|
||||
format_user_message(user_message)
|
||||
|
||||
|
||||
def test_format_assistant_message() -> None:
|
||||
assistant_message = ChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT,
|
||||
items=[
|
||||
TextContent(text="test"),
|
||||
FunctionCallContent(name="test_function", arguments={}),
|
||||
ImageContent(data="image data", mime_type="image/png"),
|
||||
],
|
||||
)
|
||||
|
||||
formatted_assistant_message = format_assistant_message(assistant_message)
|
||||
assert isinstance(formatted_assistant_message, list)
|
||||
assert len(formatted_assistant_message) == 3
|
||||
assert isinstance(formatted_assistant_message[0], Part)
|
||||
assert formatted_assistant_message[0].text == "test"
|
||||
assert isinstance(formatted_assistant_message[1], Part)
|
||||
assert formatted_assistant_message[1].function_call.name == "test_function"
|
||||
assert formatted_assistant_message[1].function_call.args == {}
|
||||
assert isinstance(formatted_assistant_message[2], Part)
|
||||
assert formatted_assistant_message[2].inline_data
|
||||
|
||||
|
||||
def test_format_assistant_message_with_unsupported_items() -> None:
|
||||
assistant_message = ChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT,
|
||||
items=[
|
||||
FunctionResultContent(id="test_id", function_name="test_function"),
|
||||
],
|
||||
)
|
||||
|
||||
with pytest.raises(ServiceInvalidRequestError):
|
||||
format_assistant_message(assistant_message)
|
||||
|
||||
|
||||
def test_format_assistant_message_with_thought_signature() -> None:
|
||||
"""Test that thought_signature is preserved in function call parts for Vertex AI."""
|
||||
import base64
|
||||
|
||||
thought_sig = base64.b64encode(b"test_thought_signature_data").decode("utf-8")
|
||||
assistant_message = ChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT,
|
||||
items=[
|
||||
FunctionCallContent(
|
||||
name="test_function",
|
||||
arguments={"arg1": "value1"},
|
||||
metadata={"thought_signature": thought_sig},
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
formatted = format_assistant_message(assistant_message)
|
||||
assert len(formatted) == 1
|
||||
assert isinstance(formatted[0], Part)
|
||||
assert formatted[0].function_call.name == "test_function"
|
||||
assert formatted[0].function_call.args == {"arg1": "value1"}
|
||||
part_dict = formatted[0].to_dict()
|
||||
assert "thought_signature" in part_dict
|
||||
assert part_dict["thought_signature"] == thought_sig
|
||||
|
||||
|
||||
def test_format_assistant_message_without_thought_signature() -> None:
|
||||
"""Test that function calls without thought_signature still work for Vertex AI."""
|
||||
assistant_message = ChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT,
|
||||
items=[
|
||||
FunctionCallContent(
|
||||
name="test_function",
|
||||
arguments={"arg1": "value1"},
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
formatted = format_assistant_message(assistant_message)
|
||||
assert len(formatted) == 1
|
||||
assert isinstance(formatted[0], Part)
|
||||
assert formatted[0].function_call.name == "test_function"
|
||||
assert formatted[0].function_call.args == {"arg1": "value1"}
|
||||
part_dict = formatted[0].to_dict()
|
||||
assert "thought_signature" not in part_dict
|
||||
|
||||
|
||||
def test_vertex_ai_function_call_format_sanitizes_anyof_schema() -> None:
|
||||
"""Integration test: anyOf in param schema_data is sanitized in the FunctionDeclaration."""
|
||||
metadata = KernelFunctionMetadata(
|
||||
name="test_func",
|
||||
description="A test function",
|
||||
is_prompt=False,
|
||||
parameters=[
|
||||
KernelParameterMetadata(
|
||||
name="messages",
|
||||
description="The user messages",
|
||||
is_required=True,
|
||||
schema_data={
|
||||
"anyOf": [
|
||||
{"type": "string"},
|
||||
{"type": "array", "items": {"type": "string"}},
|
||||
],
|
||||
"description": "The user messages",
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
result = kernel_function_metadata_to_vertex_ai_function_call_format(metadata)
|
||||
assert isinstance(result, FunctionDeclaration)
|
||||
|
||||
|
||||
def test_vertex_ai_function_call_format_empty_parameters() -> None:
|
||||
"""Integration test: metadata with no parameters produces empty properties, no crash."""
|
||||
metadata = KernelFunctionMetadata(
|
||||
name="no_params_func",
|
||||
description="No parameters",
|
||||
is_prompt=False,
|
||||
parameters=[],
|
||||
)
|
||||
result = kernel_function_metadata_to_vertex_ai_function_call_format(metadata)
|
||||
assert isinstance(result, FunctionDeclaration)
|
||||
@@ -0,0 +1,136 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from semantic_kernel.connectors.ai.google.vertex_ai import (
|
||||
VertexAIChatPromptExecutionSettings,
|
||||
VertexAIPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.google.vertex_ai.vertex_ai_prompt_execution_settings import (
|
||||
VertexAIEmbeddingPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
|
||||
|
||||
def test_default_vertex_ai_prompt_execution_settings():
|
||||
settings = VertexAIPromptExecutionSettings()
|
||||
|
||||
assert settings.stop_sequences is None
|
||||
assert settings.response_mime_type is None
|
||||
assert settings.response_schema is None
|
||||
assert settings.candidate_count is None
|
||||
assert settings.max_output_tokens is None
|
||||
assert settings.temperature is None
|
||||
assert settings.top_p is None
|
||||
assert settings.top_k is None
|
||||
|
||||
|
||||
def test_custom_vertex_ai_prompt_execution_settings():
|
||||
settings = VertexAIPromptExecutionSettings(
|
||||
stop_sequences=["world"],
|
||||
response_mime_type="text/plain",
|
||||
candidate_count=1,
|
||||
max_output_tokens=128,
|
||||
temperature=0.5,
|
||||
top_p=0.5,
|
||||
top_k=10,
|
||||
)
|
||||
|
||||
assert settings.stop_sequences == ["world"]
|
||||
assert settings.response_mime_type == "text/plain"
|
||||
assert settings.candidate_count == 1
|
||||
assert settings.max_output_tokens == 128
|
||||
assert settings.temperature == 0.5
|
||||
assert settings.top_p == 0.5
|
||||
assert settings.top_k == 10
|
||||
|
||||
|
||||
def test_vertex_ai_prompt_execution_settings_from_default_completion_config():
|
||||
settings = PromptExecutionSettings(service_id="test_service")
|
||||
chat_settings = VertexAIChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
|
||||
assert chat_settings.service_id == "test_service"
|
||||
assert chat_settings.stop_sequences is None
|
||||
assert chat_settings.response_mime_type is None
|
||||
assert chat_settings.response_schema is None
|
||||
assert chat_settings.candidate_count is None
|
||||
assert chat_settings.max_output_tokens is None
|
||||
assert chat_settings.temperature is None
|
||||
assert chat_settings.top_p is None
|
||||
assert chat_settings.top_k is None
|
||||
|
||||
|
||||
def test_vertex_ai_prompt_execution_settings_from_openai_prompt_execution_settings():
|
||||
chat_settings = VertexAIChatPromptExecutionSettings(service_id="test_service", temperature=1.0)
|
||||
new_settings = VertexAIPromptExecutionSettings(service_id="test_2", temperature=0.0)
|
||||
chat_settings.update_from_prompt_execution_settings(new_settings)
|
||||
|
||||
assert chat_settings.service_id == "test_2"
|
||||
assert chat_settings.temperature == 0.0
|
||||
|
||||
|
||||
def test_vertex_ai_prompt_execution_settings_from_custom_completion_config():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"stop_sequences": ["world"],
|
||||
"response_mime_type": "text/plain",
|
||||
"candidate_count": 1,
|
||||
"max_output_tokens": 128,
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"top_k": 10,
|
||||
},
|
||||
)
|
||||
chat_settings = VertexAIChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
|
||||
assert chat_settings.stop_sequences == ["world"]
|
||||
assert chat_settings.response_mime_type == "text/plain"
|
||||
assert chat_settings.candidate_count == 1
|
||||
assert chat_settings.max_output_tokens == 128
|
||||
assert chat_settings.temperature == 0.5
|
||||
assert chat_settings.top_p == 0.5
|
||||
assert chat_settings.top_k == 10
|
||||
|
||||
|
||||
def test_vertex_ai_chat_prompt_execution_settings_from_custom_completion_config_with_functions():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"tools": [],
|
||||
},
|
||||
)
|
||||
chat_settings = VertexAIChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
|
||||
assert chat_settings.tools == []
|
||||
|
||||
|
||||
def test_create_options():
|
||||
settings = VertexAIChatPromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"stop_sequences": ["world"],
|
||||
"response_mime_type": "text/plain",
|
||||
"candidate_count": 1,
|
||||
"max_output_tokens": 128,
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"top_k": 10,
|
||||
},
|
||||
)
|
||||
options = settings.prepare_settings_dict()
|
||||
|
||||
assert options["stop_sequences"] == ["world"]
|
||||
assert options["response_mime_type"] == "text/plain"
|
||||
assert options["candidate_count"] == 1
|
||||
assert options["max_output_tokens"] == 128
|
||||
assert options["temperature"] == 0.5
|
||||
assert options["top_p"] == 0.5
|
||||
assert options["top_k"] == 10
|
||||
assert "tools" not in options
|
||||
assert "tool_config" not in options
|
||||
|
||||
|
||||
def test_default_vertex_ai_embedding_prompt_execution_settings():
|
||||
settings = VertexAIEmbeddingPromptExecutionSettings()
|
||||
|
||||
assert settings.output_dimensionality is None
|
||||
assert settings.auto_truncate is None
|
||||
@@ -0,0 +1,219 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from threading import Thread
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
from transformers import AutoTokenizer, TextIteratorStreamer
|
||||
|
||||
from semantic_kernel.connectors.ai.hugging_face.services.hf_text_completion import HuggingFaceTextCompletion
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
from semantic_kernel.exceptions import KernelInvokeException, ServiceResponseException
|
||||
from semantic_kernel.functions.kernel_arguments import KernelArguments
|
||||
from semantic_kernel.kernel import Kernel
|
||||
from semantic_kernel.prompt_template.prompt_template_config import PromptTemplateConfig
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model_name", "task", "input_str"),
|
||||
[
|
||||
(
|
||||
"patrickvonplaten/t5-tiny-random",
|
||||
"text2text-generation",
|
||||
"translate English to Dutch: Hello, how are you?",
|
||||
),
|
||||
(
|
||||
"Falconsai/text_summarization",
|
||||
"summarization",
|
||||
"""
|
||||
Summarize: Whales are fully aquatic, open-ocean animals:
|
||||
they can feed, mate, give birth, suckle and raise their young at sea.
|
||||
Whales range in size from the 2.6 metres (8.5 ft) and 135 kilograms (298 lb)
|
||||
dwarf sperm whale to the 29.9 metres (98 ft) and 190 tonnes (210 short tons) blue whale,
|
||||
which is the largest known animal that has ever lived. The sperm whale is the largest
|
||||
toothed predator on Earth. Several whale species exhibit sexual dimorphism,
|
||||
in that the females are larger than males.
|
||||
""",
|
||||
),
|
||||
("HuggingFaceM4/tiny-random-LlamaForCausalLM", "text-generation", "Hello, I like sleeping and "),
|
||||
],
|
||||
ids=["text2text-generation", "summarization", "text-generation"],
|
||||
)
|
||||
async def test_text_completion(model_name, task, input_str):
|
||||
kernel = Kernel()
|
||||
|
||||
ret = {"summary_text": "test"} if task == "summarization" else {"generated_text": "test"}
|
||||
mock_pipeline = Mock(return_value=ret)
|
||||
|
||||
# Configure LLM service
|
||||
with patch("semantic_kernel.connectors.ai.hugging_face.services.hf_text_completion.pipeline") as patched_pipeline:
|
||||
patched_pipeline.return_value = mock_pipeline
|
||||
service = HuggingFaceTextCompletion(service_id=model_name, ai_model_id=model_name, task=task)
|
||||
kernel.add_service(
|
||||
service=service,
|
||||
)
|
||||
|
||||
exec_settings = PromptExecutionSettings(service_id=model_name, extension_data={"max_new_tokens": 25})
|
||||
|
||||
# Define semantic function using SK prompt template language
|
||||
prompt = "{{$input}}"
|
||||
|
||||
prompt_template_config = PromptTemplateConfig(template=prompt, execution_settings=exec_settings)
|
||||
|
||||
kernel.add_function(
|
||||
prompt_template_config=prompt_template_config,
|
||||
function_name="TestFunction",
|
||||
plugin_name="TestPlugin",
|
||||
prompt_execution_settings=exec_settings,
|
||||
)
|
||||
|
||||
arguments = KernelArguments(input=input_str)
|
||||
|
||||
await kernel.invoke(function_name="TestFunction", plugin_name="TestPlugin", arguments=arguments)
|
||||
assert mock_pipeline.call_args.args[0] == input_str
|
||||
|
||||
|
||||
async def test_text_completion_throws():
|
||||
kernel = Kernel()
|
||||
|
||||
model_name = "patrickvonplaten/t5-tiny-random"
|
||||
task = "text2text-generation"
|
||||
input_str = "translate English to Dutch: Hello, how are you?"
|
||||
|
||||
with patch("semantic_kernel.connectors.ai.hugging_face.services.hf_text_completion.pipeline") as patched_pipeline:
|
||||
mock_generator = Mock()
|
||||
mock_generator.side_effect = Exception("Test exception")
|
||||
patched_pipeline.return_value = mock_generator
|
||||
service = HuggingFaceTextCompletion(service_id=model_name, ai_model_id=model_name, task=task)
|
||||
kernel.add_service(service=service)
|
||||
|
||||
exec_settings = PromptExecutionSettings(service_id=model_name, extension_data={"max_new_tokens": 25})
|
||||
|
||||
prompt = "{{$input}}"
|
||||
prompt_template_config = PromptTemplateConfig(template=prompt, execution_settings=exec_settings)
|
||||
|
||||
kernel.add_function(
|
||||
prompt_template_config=prompt_template_config,
|
||||
function_name="TestFunction",
|
||||
plugin_name="TestPlugin",
|
||||
prompt_execution_settings=exec_settings,
|
||||
)
|
||||
|
||||
arguments = KernelArguments(input=input_str)
|
||||
|
||||
with pytest.raises(
|
||||
KernelInvokeException, match="Error occurred while invoking function: 'TestPlugin-TestFunction'"
|
||||
):
|
||||
await kernel.invoke(function_name="TestFunction", plugin_name="TestPlugin", arguments=arguments)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model_name", "task", "input_str"),
|
||||
[
|
||||
(
|
||||
"patrickvonplaten/t5-tiny-random",
|
||||
"text2text-generation",
|
||||
"translate English to Dutch: Hello, how are you?",
|
||||
),
|
||||
("HuggingFaceM4/tiny-random-LlamaForCausalLM", "text-generation", "Hello, I like sleeping and "),
|
||||
],
|
||||
ids=["text2text-generation", "text-generation"],
|
||||
)
|
||||
async def test_text_completion_streaming(model_name, task, input_str):
|
||||
ret = {"summary_text": "test"} if task == "summarization" else {"generated_text": "test"}
|
||||
mock_pipeline = Mock(return_value=ret)
|
||||
|
||||
mock_streamer = MagicMock(spec=TextIteratorStreamer)
|
||||
mock_streamer.__iter__.return_value = iter(["mocked_text"])
|
||||
|
||||
with (
|
||||
patch(
|
||||
"semantic_kernel.connectors.ai.hugging_face.services.hf_text_completion.pipeline",
|
||||
return_value=mock_pipeline,
|
||||
),
|
||||
patch(
|
||||
"semantic_kernel.connectors.ai.hugging_face.services.hf_text_completion.Thread",
|
||||
side_effect=Mock(spec=Thread),
|
||||
),
|
||||
patch(
|
||||
"semantic_kernel.connectors.ai.hugging_face.services.hf_text_completion.AutoTokenizer",
|
||||
side_effect=Mock(spec=AutoTokenizer),
|
||||
),
|
||||
patch(
|
||||
"semantic_kernel.connectors.ai.hugging_face.services.hf_text_completion.TextIteratorStreamer",
|
||||
return_value=mock_streamer,
|
||||
) as mock_stream,
|
||||
):
|
||||
mock_stream.return_value = mock_streamer
|
||||
service = HuggingFaceTextCompletion(service_id=model_name, ai_model_id=model_name, task=task)
|
||||
prompt = "test prompt"
|
||||
exec_settings = PromptExecutionSettings(service_id=model_name, extension_data={"max_new_tokens": 25})
|
||||
|
||||
result = []
|
||||
async for content in service.get_streaming_text_contents(prompt, exec_settings):
|
||||
result.append(content)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0][0].inner_content == "mocked_text"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model_name", "task", "input_str"),
|
||||
[
|
||||
(
|
||||
"patrickvonplaten/t5-tiny-random",
|
||||
"text2text-generation",
|
||||
"translate English to Dutch: Hello, how are you?",
|
||||
),
|
||||
("HuggingFaceM4/tiny-random-LlamaForCausalLM", "text-generation", "Hello, I like sleeping and "),
|
||||
],
|
||||
ids=["text2text-generation", "text-generation"],
|
||||
)
|
||||
async def test_text_completion_streaming_throws(model_name, task, input_str):
|
||||
ret = {"summary_text": "test"} if task == "summarization" else {"generated_text": "test"}
|
||||
mock_pipeline = Mock(return_value=ret)
|
||||
|
||||
mock_streamer = MagicMock(spec=TextIteratorStreamer)
|
||||
mock_streamer.__iter__.return_value = Exception()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"semantic_kernel.connectors.ai.hugging_face.services.hf_text_completion.pipeline",
|
||||
return_value=mock_pipeline,
|
||||
),
|
||||
patch(
|
||||
"semantic_kernel.connectors.ai.hugging_face.services.hf_text_completion.Thread",
|
||||
side_effect=Exception(),
|
||||
),
|
||||
patch(
|
||||
"semantic_kernel.connectors.ai.hugging_face.services.hf_text_completion.TextIteratorStreamer",
|
||||
return_value=mock_streamer,
|
||||
) as mock_stream,
|
||||
):
|
||||
mock_stream.return_value = mock_streamer
|
||||
service = HuggingFaceTextCompletion(service_id=model_name, ai_model_id=model_name, task=task)
|
||||
prompt = "test prompt"
|
||||
exec_settings = PromptExecutionSettings(service_id=model_name, extension_data={"max_new_tokens": 25})
|
||||
|
||||
with pytest.raises(ServiceResponseException, match=("Hugging Face completion failed")):
|
||||
async for _ in service.get_streaming_text_contents(prompt, exec_settings):
|
||||
pass
|
||||
|
||||
|
||||
def test_hugging_face_text_completion_init():
|
||||
with (
|
||||
patch("semantic_kernel.connectors.ai.hugging_face.services.hf_text_completion.pipeline") as patched_pipeline,
|
||||
patch(
|
||||
"semantic_kernel.connectors.ai.hugging_face.services.hf_text_completion.torch.cuda.is_available"
|
||||
) as mock_torch_cuda_is_available,
|
||||
):
|
||||
patched_pipeline.return_value = patched_pipeline
|
||||
mock_torch_cuda_is_available.return_value = False
|
||||
|
||||
ai_model_id = "test-model"
|
||||
task = "summarization"
|
||||
device = -1
|
||||
|
||||
service = HuggingFaceTextCompletion(service_id="test", ai_model_id=ai_model_id, task=task, device=device)
|
||||
|
||||
assert service is not None
|
||||
@@ -0,0 +1,56 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from numpy import array, ndarray
|
||||
|
||||
from semantic_kernel.connectors.ai.hugging_face.services.hf_text_embedding import HuggingFaceTextEmbedding
|
||||
from semantic_kernel.exceptions import ServiceResponseException
|
||||
|
||||
|
||||
def test_huggingface_text_embedding_initialization():
|
||||
model_name = "sentence-transformers/all-MiniLM-L6-v2"
|
||||
device = -1
|
||||
|
||||
with patch("sentence_transformers.SentenceTransformer") as mock_transformer:
|
||||
mock_instance = mock_transformer.return_value
|
||||
service = HuggingFaceTextEmbedding(service_id="test", ai_model_id=model_name, device=device)
|
||||
|
||||
assert service.ai_model_id == model_name
|
||||
assert service.device == "cpu"
|
||||
assert service.generator == mock_instance
|
||||
mock_transformer.assert_called_once_with(model_name_or_path=model_name, device="cpu")
|
||||
|
||||
|
||||
async def test_generate_embeddings_success():
|
||||
model_name = "sentence-transformers/all-MiniLM-L6-v2"
|
||||
device = -1
|
||||
texts = ["Hello world!", "How are you?"]
|
||||
mock_embeddings = array([[0.1, 0.2], [0.3, 0.4]])
|
||||
|
||||
with patch("sentence_transformers.SentenceTransformer") as mock_transformer:
|
||||
mock_instance = mock_transformer.return_value
|
||||
mock_instance.encode.return_value = mock_embeddings
|
||||
|
||||
service = HuggingFaceTextEmbedding(service_id="test", ai_model_id=model_name, device=device)
|
||||
embeddings = await service.generate_embeddings(texts)
|
||||
|
||||
assert isinstance(embeddings, ndarray)
|
||||
assert embeddings.shape == (2, 2)
|
||||
assert (embeddings == mock_embeddings).all()
|
||||
|
||||
|
||||
async def test_generate_embeddings_throws():
|
||||
model_name = "sentence-transformers/all-MiniLM-L6-v2"
|
||||
device = -1
|
||||
texts = ["Hello world!", "How are you?"]
|
||||
|
||||
with patch("sentence_transformers.SentenceTransformer") as mock_transformer:
|
||||
mock_instance = mock_transformer.return_value
|
||||
mock_instance.encode.side_effect = Exception("Test exception")
|
||||
|
||||
service = HuggingFaceTextEmbedding(service_id="test", ai_model_id=model_name, device=device)
|
||||
|
||||
with pytest.raises(ServiceResponseException, match="Hugging Face embeddings failed"):
|
||||
await service.generate_embeddings(texts)
|
||||
+616
@@ -0,0 +1,616 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from mistralai import CompletionEvent, Mistral
|
||||
from mistralai.models import (
|
||||
AssistantMessage,
|
||||
ChatCompletionChoice,
|
||||
ChatCompletionResponse,
|
||||
CompletionChunk,
|
||||
CompletionResponseStreamChoice,
|
||||
DeltaMessage,
|
||||
UsageInfo,
|
||||
)
|
||||
|
||||
from semantic_kernel.connectors.ai.chat_completion_client_base import ChatCompletionClientBase
|
||||
from semantic_kernel.connectors.ai.function_call_choice_configuration import FunctionCallChoiceConfiguration
|
||||
from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior, FunctionChoiceType
|
||||
from semantic_kernel.connectors.ai.mistral_ai.prompt_execution_settings.mistral_ai_prompt_execution_settings import (
|
||||
MistralAIChatPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.mistral_ai.services.mistral_ai_chat_completion import MistralAIChatCompletion
|
||||
from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.open_ai_prompt_execution_settings import (
|
||||
OpenAIChatPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.contents import FunctionCallContent, TextContent
|
||||
from semantic_kernel.contents.chat_history import ChatHistory
|
||||
from semantic_kernel.contents.chat_message_content import ChatMessageContent
|
||||
from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent
|
||||
from semantic_kernel.contents.utils.author_role import AuthorRole
|
||||
from semantic_kernel.exceptions import (
|
||||
ServiceInitializationError,
|
||||
ServiceInvalidExecutionSettingsError,
|
||||
ServiceResponseException,
|
||||
)
|
||||
from semantic_kernel.functions.kernel_arguments import KernelArguments
|
||||
from semantic_kernel.functions.kernel_function import KernelFunction
|
||||
from semantic_kernel.functions.kernel_function_decorator import kernel_function
|
||||
from semantic_kernel.kernel import Kernel
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_settings() -> MistralAIChatPromptExecutionSettings:
|
||||
return MistralAIChatPromptExecutionSettings()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_mistral_ai_client_completion() -> Mistral:
|
||||
client = MagicMock(spec=Mistral)
|
||||
client.chat = MagicMock()
|
||||
|
||||
# Use proper ChatCompletionResponse type so isinstance checks pass
|
||||
chat_completion_response = ChatCompletionResponse(
|
||||
id="test_id",
|
||||
object="object",
|
||||
created=12345,
|
||||
usage=UsageInfo(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
||||
model="test_model_id",
|
||||
choices=[
|
||||
ChatCompletionChoice(
|
||||
index=0,
|
||||
message=AssistantMessage(role="assistant", content="Test"),
|
||||
finish_reason="stop",
|
||||
)
|
||||
],
|
||||
)
|
||||
client.chat.complete_async = AsyncMock(return_value=chat_completion_response)
|
||||
return client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_mistral_ai_client_completion_stream() -> Mistral:
|
||||
client = MagicMock(spec=Mistral)
|
||||
client.chat = MagicMock()
|
||||
|
||||
# Use proper CompletionEvent/CompletionChunk types so isinstance checks pass
|
||||
mock_chunk = CompletionEvent(
|
||||
data=CompletionChunk(
|
||||
id="test_chunk",
|
||||
created=12345,
|
||||
model="test_model_id",
|
||||
choices=[
|
||||
CompletionResponseStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(role="assistant", content="Test"),
|
||||
finish_reason="stop",
|
||||
)
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
async def mock_stream_async(**kwargs):
|
||||
yield mock_chunk
|
||||
|
||||
client.chat.stream_async = AsyncMock(return_value=mock_stream_async())
|
||||
return client
|
||||
|
||||
|
||||
async def test_complete_chat_contents(
|
||||
kernel: Kernel,
|
||||
mock_settings: MistralAIChatPromptExecutionSettings,
|
||||
mock_mistral_ai_client_completion: Mistral,
|
||||
):
|
||||
chat_history = MagicMock()
|
||||
arguments = KernelArguments()
|
||||
chat_completion_base = MistralAIChatCompletion(
|
||||
ai_model_id="test_model_id", service_id="test", api_key="", async_client=mock_mistral_ai_client_completion
|
||||
)
|
||||
|
||||
content: list[ChatMessageContent] = await chat_completion_base.get_chat_message_contents(
|
||||
chat_history=chat_history, settings=mock_settings, kernel=kernel, arguments=arguments
|
||||
)
|
||||
assert content is not None
|
||||
|
||||
|
||||
mock_message_text_content = ChatMessageContent(role=AuthorRole.ASSISTANT, items=[TextContent(text="test")])
|
||||
|
||||
mock_message_function_call = ChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT,
|
||||
items=[
|
||||
FunctionCallContent(
|
||||
name="test",
|
||||
arguments={"key": "test"},
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"function_choice_behavior,model_responses,expected_result",
|
||||
[
|
||||
pytest.param(
|
||||
FunctionChoiceBehavior.Auto(),
|
||||
[[mock_message_function_call], [mock_message_text_content]],
|
||||
TextContent,
|
||||
id="auto",
|
||||
),
|
||||
pytest.param(
|
||||
FunctionChoiceBehavior.Auto(auto_invoke=False),
|
||||
[[mock_message_function_call]],
|
||||
FunctionCallContent,
|
||||
id="auto_none_invoke",
|
||||
),
|
||||
pytest.param(
|
||||
FunctionChoiceBehavior.Required(auto_invoke=False),
|
||||
[[mock_message_function_call]],
|
||||
FunctionCallContent,
|
||||
id="required_none_invoke",
|
||||
),
|
||||
pytest.param(FunctionChoiceBehavior.NoneInvoke(), [[mock_message_text_content]], TextContent, id="none"),
|
||||
],
|
||||
)
|
||||
async def test_complete_chat_contents_function_call_behavior_tool_call(
|
||||
kernel: Kernel,
|
||||
mock_settings: MistralAIChatPromptExecutionSettings,
|
||||
function_choice_behavior: FunctionChoiceBehavior,
|
||||
model_responses,
|
||||
expected_result,
|
||||
):
|
||||
kernel.add_function("test", kernel_function(lambda key: "test", name="test"))
|
||||
mock_settings.function_choice_behavior = function_choice_behavior
|
||||
|
||||
arguments = KernelArguments()
|
||||
chat_completion_base = MistralAIChatCompletion(ai_model_id="test_model_id", service_id="test", api_key="")
|
||||
|
||||
with (
|
||||
patch.object(chat_completion_base, "_inner_get_chat_message_contents", side_effect=model_responses),
|
||||
):
|
||||
response: list[ChatMessageContent] = await chat_completion_base.get_chat_message_contents(
|
||||
chat_history=ChatHistory(system_message="Test"), settings=mock_settings, kernel=kernel, arguments=arguments
|
||||
)
|
||||
|
||||
assert all(isinstance(content, expected_result) for content in response[0].items)
|
||||
|
||||
|
||||
async def test_complete_chat_contents_function_call_behavior_without_kernel(
|
||||
mock_settings: MistralAIChatPromptExecutionSettings,
|
||||
mock_mistral_ai_client_completion: Mistral,
|
||||
):
|
||||
chat_history = MagicMock()
|
||||
chat_completion_base = MistralAIChatCompletion(
|
||||
ai_model_id="test_model_id", service_id="test", api_key="", async_client=mock_mistral_ai_client_completion
|
||||
)
|
||||
|
||||
mock_settings.function_choice_behavior = FunctionChoiceBehavior.Auto()
|
||||
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
await chat_completion_base.get_chat_message_contents(chat_history=chat_history, settings=mock_settings)
|
||||
|
||||
|
||||
async def test_complete_chat_stream_contents(
|
||||
kernel: Kernel,
|
||||
mock_settings: MistralAIChatPromptExecutionSettings,
|
||||
mock_mistral_ai_client_completion_stream: Mistral,
|
||||
):
|
||||
chat_history = MagicMock()
|
||||
arguments = KernelArguments()
|
||||
|
||||
chat_completion_base = MistralAIChatCompletion(
|
||||
ai_model_id="test_model_id",
|
||||
service_id="test",
|
||||
api_key="",
|
||||
async_client=mock_mistral_ai_client_completion_stream,
|
||||
)
|
||||
|
||||
async for content in chat_completion_base.get_streaming_chat_message_contents(
|
||||
chat_history, mock_settings, kernel=kernel, arguments=arguments
|
||||
):
|
||||
assert content is not None
|
||||
|
||||
|
||||
mock_message_function_call = StreamingChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT, items=[FunctionCallContent(name="test")], choice_index="0"
|
||||
)
|
||||
|
||||
mock_message_text_content = StreamingChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT, items=[TextContent(text="test")], choice_index="0"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"function_choice_behavior,model_responses,expected_result",
|
||||
[
|
||||
pytest.param(
|
||||
FunctionChoiceBehavior.Auto(),
|
||||
[[mock_message_function_call], [mock_message_text_content]],
|
||||
TextContent,
|
||||
id="auto",
|
||||
),
|
||||
pytest.param(
|
||||
FunctionChoiceBehavior.Auto(auto_invoke=False),
|
||||
[[mock_message_function_call]],
|
||||
FunctionCallContent,
|
||||
id="auto_none_invoke",
|
||||
),
|
||||
pytest.param(
|
||||
FunctionChoiceBehavior.Required(auto_invoke=False),
|
||||
[[mock_message_function_call]],
|
||||
FunctionCallContent,
|
||||
id="required_none_invoke",
|
||||
),
|
||||
pytest.param(FunctionChoiceBehavior.NoneInvoke(), [[mock_message_text_content]], TextContent, id="none"),
|
||||
],
|
||||
)
|
||||
async def test_complete_chat_contents_streaming_function_call_behavior_tool_call(
|
||||
kernel: Kernel,
|
||||
mock_settings: MistralAIChatPromptExecutionSettings,
|
||||
function_choice_behavior: FunctionChoiceBehavior,
|
||||
model_responses,
|
||||
expected_result,
|
||||
):
|
||||
mock_settings.function_choice_behavior = function_choice_behavior
|
||||
|
||||
# Mock sequence of model responses
|
||||
generator_mocks = []
|
||||
for mock_message in model_responses:
|
||||
generator_mock = MagicMock()
|
||||
generator_mock.__aiter__.return_value = [mock_message]
|
||||
generator_mocks.append(generator_mock)
|
||||
|
||||
arguments = KernelArguments()
|
||||
chat_completion_base = MistralAIChatCompletion(ai_model_id="test_model_id", service_id="test", api_key="")
|
||||
|
||||
with patch.object(chat_completion_base, "_inner_get_streaming_chat_message_contents", side_effect=generator_mocks):
|
||||
messages = []
|
||||
async for chunk in chat_completion_base.get_streaming_chat_message_contents(
|
||||
chat_history=ChatHistory(system_message="Test"), settings=mock_settings, kernel=kernel, arguments=arguments
|
||||
):
|
||||
messages.append(chunk)
|
||||
|
||||
response = messages[-1]
|
||||
assert all(isinstance(content, expected_result) for content in response[0].items)
|
||||
|
||||
|
||||
async def test_mistral_ai_sdk_exception(kernel: Kernel, mock_settings: MistralAIChatPromptExecutionSettings):
|
||||
chat_history = MagicMock()
|
||||
arguments = KernelArguments()
|
||||
client = MagicMock(spec=Mistral)
|
||||
client.chat = MagicMock()
|
||||
client.chat.complete_async.side_effect = Exception("Test Exception")
|
||||
|
||||
chat_completion_base = MistralAIChatCompletion(
|
||||
ai_model_id="test_model_id", service_id="test", api_key="", async_client=client
|
||||
)
|
||||
|
||||
with pytest.raises(ServiceResponseException):
|
||||
await chat_completion_base.get_chat_message_contents(
|
||||
chat_history=chat_history, settings=mock_settings, kernel=kernel, arguments=arguments
|
||||
)
|
||||
|
||||
|
||||
async def test_mistral_ai_sdk_exception_streaming(kernel: Kernel, mock_settings: MistralAIChatPromptExecutionSettings):
|
||||
chat_history = MagicMock()
|
||||
arguments = KernelArguments()
|
||||
client = MagicMock(spec=Mistral)
|
||||
client.chat = MagicMock()
|
||||
client.chat.chat_stream.side_effect = Exception("Test Exception")
|
||||
|
||||
chat_completion_base = MistralAIChatCompletion(
|
||||
ai_model_id="test_model_id", service_id="test", api_key="", async_client=client
|
||||
)
|
||||
|
||||
with pytest.raises(ServiceResponseException):
|
||||
async for content in chat_completion_base.get_streaming_chat_message_contents(
|
||||
chat_history, mock_settings, kernel=kernel, arguments=arguments
|
||||
):
|
||||
assert content is not None
|
||||
|
||||
|
||||
def test_mistral_ai_chat_completion_init(mistralai_unit_test_env) -> None:
|
||||
# Test successful initialization
|
||||
mistral_ai_chat_completion = MistralAIChatCompletion()
|
||||
|
||||
assert mistral_ai_chat_completion.ai_model_id == mistralai_unit_test_env["MISTRALAI_CHAT_MODEL_ID"]
|
||||
api_key = mistralai_unit_test_env["MISTRALAI_API_KEY"]
|
||||
assert mistral_ai_chat_completion.async_client.sdk_configuration.security.api_key == api_key
|
||||
assert isinstance(mistral_ai_chat_completion, ChatCompletionClientBase)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["MISTRALAI_API_KEY", "MISTRALAI_CHAT_MODEL_ID"]], indirect=True)
|
||||
def test_mistral_ai_chat_completion_init_constructor(mistralai_unit_test_env) -> None:
|
||||
# Test successful initialization
|
||||
mistral_ai_chat_completion = MistralAIChatCompletion(
|
||||
api_key="overwrite_api_key",
|
||||
ai_model_id="overwrite_model_id",
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
assert mistral_ai_chat_completion.ai_model_id == "overwrite_model_id"
|
||||
assert mistral_ai_chat_completion.async_client.sdk_configuration.security.api_key == "overwrite_api_key"
|
||||
assert isinstance(mistral_ai_chat_completion, ChatCompletionClientBase)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["MISTRALAI_API_KEY", "MISTRALAI_CHAT_MODEL_ID"]], indirect=True)
|
||||
def test_mistral_ai_chat_completion_init_constructor_missing_model(mistralai_unit_test_env) -> None:
|
||||
# Test successful initialization
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
MistralAIChatCompletion(api_key="overwrite_api_key", env_file_path="test.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["MISTRALAI_API_KEY", "MISTRALAI_CHAT_MODEL_ID"]], indirect=True)
|
||||
def test_mistral_ai_chat_completion_init_constructor_missing_api_key(mistralai_unit_test_env) -> None:
|
||||
# Test successful initialization
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
MistralAIChatCompletion(ai_model_id="overwrite_model_id", env_file_path="test.env")
|
||||
|
||||
|
||||
def test_mistral_ai_chat_completion_init_hybrid(mistralai_unit_test_env) -> None:
|
||||
mistral_ai_chat_completion = MistralAIChatCompletion(
|
||||
ai_model_id="overwrite_model_id",
|
||||
env_file_path="test.env",
|
||||
)
|
||||
assert mistral_ai_chat_completion.ai_model_id == "overwrite_model_id"
|
||||
assert mistral_ai_chat_completion.async_client.sdk_configuration.security.api_key == "test_api_key"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["MISTRALAI_CHAT_MODEL_ID"]], indirect=True)
|
||||
def test_mistral_ai_chat_completion_init_with_empty_model_id(mistralai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
MistralAIChatCompletion(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
def test_prompt_execution_settings_class(mistralai_unit_test_env):
|
||||
mistral_ai_chat_completion = MistralAIChatCompletion()
|
||||
prompt_execution_settings = mistral_ai_chat_completion.get_prompt_execution_settings_class()
|
||||
assert prompt_execution_settings == MistralAIChatPromptExecutionSettings
|
||||
|
||||
|
||||
async def test_with_different_execution_settings(kernel: Kernel, mock_mistral_ai_client_completion: MagicMock):
|
||||
chat_history = MagicMock()
|
||||
settings = OpenAIChatPromptExecutionSettings(temperature=0.2, seed=2)
|
||||
arguments = KernelArguments()
|
||||
chat_completion_base = MistralAIChatCompletion(
|
||||
ai_model_id="test_model_id", service_id="test", api_key="", async_client=mock_mistral_ai_client_completion
|
||||
)
|
||||
|
||||
await chat_completion_base.get_chat_message_contents(
|
||||
chat_history=chat_history, settings=settings, kernel=kernel, arguments=arguments
|
||||
)
|
||||
assert mock_mistral_ai_client_completion.chat.complete_async.call_args.kwargs["temperature"] == 0.2
|
||||
assert mock_mistral_ai_client_completion.chat.complete_async.call_args.kwargs["seed"] == 2
|
||||
|
||||
|
||||
async def test_with_different_execution_settings_stream(
|
||||
kernel: Kernel, mock_mistral_ai_client_completion_stream: MagicMock
|
||||
):
|
||||
chat_history = MagicMock()
|
||||
settings = OpenAIChatPromptExecutionSettings(temperature=0.2, seed=2)
|
||||
arguments = KernelArguments()
|
||||
chat_completion_base = MistralAIChatCompletion(
|
||||
ai_model_id="test_model_id",
|
||||
service_id="test",
|
||||
api_key="",
|
||||
async_client=mock_mistral_ai_client_completion_stream,
|
||||
)
|
||||
|
||||
async for chunk in chat_completion_base.get_streaming_chat_message_contents(
|
||||
chat_history, settings, kernel=kernel, arguments=arguments
|
||||
):
|
||||
continue
|
||||
assert mock_mistral_ai_client_completion_stream.chat.stream_async.call_args.kwargs["temperature"] == 0.2
|
||||
assert mock_mistral_ai_client_completion_stream.chat.stream_async.call_args.kwargs["seed"] == 2
|
||||
|
||||
|
||||
async def test_mistral_ai_chat_completion_get_chat_message_contents_success():
|
||||
"""Test get_chat_message_contents with a successful ChatCompletionResponse."""
|
||||
|
||||
# Mock the response from the Mistral chat complete_async.
|
||||
mock_response = ChatCompletionResponse(
|
||||
id="some_id",
|
||||
object="object",
|
||||
created=12345,
|
||||
usage=UsageInfo(prompt_tokens=10, completion_tokens=20, total_tokens=30),
|
||||
model="test-model",
|
||||
choices=[
|
||||
ChatCompletionChoice(
|
||||
index=0,
|
||||
message=AssistantMessage(role="assistant", content="Hello!"),
|
||||
finish_reason="stop",
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
async_mock_client = MagicMock(spec=Mistral)
|
||||
async_mock_client.chat = MagicMock()
|
||||
async_mock_client.chat.complete_async = AsyncMock(return_value=mock_response)
|
||||
|
||||
chat_completion = MistralAIChatCompletion(
|
||||
ai_model_id="test-model",
|
||||
api_key="test_key",
|
||||
async_client=async_mock_client,
|
||||
)
|
||||
|
||||
# We create a ChatHistory.
|
||||
chat_history = ChatHistory()
|
||||
settings = MistralAIChatPromptExecutionSettings()
|
||||
|
||||
results = await chat_completion.get_chat_message_contents(chat_history, settings)
|
||||
|
||||
# We should have exactly one ChatMessageContent.
|
||||
assert len(results) == 1
|
||||
assert results[0].role.value == "assistant"
|
||||
assert results[0].finish_reason is not None
|
||||
assert results[0].finish_reason.value == "stop"
|
||||
assert "Hello!" in results[0].content
|
||||
async_mock_client.chat.complete_async.assert_awaited_once()
|
||||
|
||||
|
||||
async def test_mistral_ai_chat_completion_get_chat_message_contents_failure():
|
||||
"""Test get_chat_message_contents should raise ServiceResponseException if Mistral call fails."""
|
||||
async_mock_client = MagicMock(spec=Mistral)
|
||||
async_mock_client.chat = MagicMock()
|
||||
async_mock_client.chat.complete_async = AsyncMock(side_effect=Exception("API error"))
|
||||
|
||||
chat_completion = MistralAIChatCompletion(
|
||||
ai_model_id="test-model",
|
||||
api_key="test_key",
|
||||
async_client=async_mock_client,
|
||||
)
|
||||
|
||||
chat_history = ChatHistory()
|
||||
settings = MistralAIChatPromptExecutionSettings()
|
||||
|
||||
with pytest.raises(ServiceResponseException) as exc:
|
||||
await chat_completion.get_chat_message_contents(chat_history, settings)
|
||||
assert "service failed to complete the prompt" in str(exc.value)
|
||||
|
||||
|
||||
async def test_mistral_ai_chat_completion_get_streaming_chat_message_contents_success():
|
||||
"""Test get_streaming_chat_message_contents when streaming successfully."""
|
||||
|
||||
# We'll yield multiple chunks to simulate streaming.
|
||||
mock_chunk1 = CompletionEvent(
|
||||
data=CompletionChunk(
|
||||
id="chunk1",
|
||||
created=1,
|
||||
model="test-model",
|
||||
choices=[
|
||||
CompletionResponseStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(role="assistant", content="Hello "),
|
||||
finish_reason=None,
|
||||
)
|
||||
],
|
||||
)
|
||||
)
|
||||
mock_chunk2 = CompletionEvent(
|
||||
data=CompletionChunk(
|
||||
id="chunk1",
|
||||
created=1,
|
||||
model="test-model",
|
||||
choices=[
|
||||
CompletionResponseStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(content="World!"),
|
||||
finish_reason="stop",
|
||||
)
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
async def mock_stream_async(**kwargs):
|
||||
yield mock_chunk1
|
||||
yield mock_chunk2
|
||||
|
||||
async_mock_client = MagicMock(spec=Mistral)
|
||||
async_mock_client.chat = MagicMock()
|
||||
async_mock_client.chat.stream_async = AsyncMock(return_value=mock_stream_async())
|
||||
|
||||
chat_completion = MistralAIChatCompletion(
|
||||
ai_model_id="test-model",
|
||||
api_key="test_key",
|
||||
async_client=async_mock_client,
|
||||
)
|
||||
|
||||
chat_history = ChatHistory()
|
||||
settings = MistralAIChatPromptExecutionSettings()
|
||||
|
||||
collected_chunks = []
|
||||
async for chunk_list in chat_completion.get_streaming_chat_message_contents(chat_history, settings):
|
||||
collected_chunks.append(chunk_list)
|
||||
|
||||
# We expect two sets of chunk_list yields.
|
||||
assert len(collected_chunks) == 2
|
||||
assert len(collected_chunks[0]) == 1
|
||||
assert len(collected_chunks[1]) == 1
|
||||
|
||||
# First chunk contains "Hello ", second chunk "World!".
|
||||
assert collected_chunks[0][0].items[0].text == "Hello "
|
||||
assert collected_chunks[1][0].items[0].text == "World!"
|
||||
|
||||
|
||||
async def test_mistral_ai_chat_completion_get_streaming_chat_message_contents_failure():
|
||||
"""Test get_streaming_chat_message_contents raising a ServiceResponseException on failure."""
|
||||
async_mock_client = MagicMock(spec=Mistral)
|
||||
async_mock_client.chat = MagicMock()
|
||||
async_mock_client.chat.stream_async = AsyncMock(side_effect=Exception("Streaming error"))
|
||||
|
||||
chat_completion = MistralAIChatCompletion(
|
||||
ai_model_id="test-model",
|
||||
api_key="test_key",
|
||||
async_client=async_mock_client,
|
||||
)
|
||||
|
||||
chat_history = ChatHistory()
|
||||
settings = MistralAIChatPromptExecutionSettings()
|
||||
|
||||
with pytest.raises(ServiceResponseException) as exc:
|
||||
async for _ in chat_completion.get_streaming_chat_message_contents(chat_history, settings):
|
||||
pass
|
||||
assert "service failed to complete the prompt" in str(exc.value)
|
||||
|
||||
|
||||
def test_mistral_ai_chat_completion_update_settings_from_function_call_configuration_mistral():
|
||||
"""Test update_settings_from_function_call_configuration_mistral sets tools etc."""
|
||||
|
||||
chat_completion = MistralAIChatCompletion(
|
||||
ai_model_id="test-model",
|
||||
api_key="test_key",
|
||||
)
|
||||
|
||||
# Create a mock settings object.
|
||||
settings = MistralAIChatPromptExecutionSettings()
|
||||
# Create a function choice config with some available functions.
|
||||
config = FunctionCallChoiceConfiguration()
|
||||
mock_func = MagicMock(
|
||||
spec=KernelFunction,
|
||||
)
|
||||
mock_func.name = "my_func"
|
||||
mock_func.description = "some desc"
|
||||
mock_func.fully_qualified_name = "mod.my_func"
|
||||
mock_func.parameters = []
|
||||
config.available_functions = [mock_func]
|
||||
|
||||
# Call the update_settings_from_function_call_configuration_mistral with type=ANY.
|
||||
chat_completion.update_settings_from_function_call_configuration_mistral(
|
||||
function_choice_configuration=config,
|
||||
settings=settings,
|
||||
type=FunctionChoiceType.AUTO,
|
||||
)
|
||||
|
||||
assert settings.tool_choice == FunctionChoiceType.AUTO.value
|
||||
assert settings.tools is not None
|
||||
assert len(settings.tools) == 1
|
||||
assert settings.tools[0]["function"]["name"] == "mod.my_func"
|
||||
|
||||
|
||||
def test_mistral_ai_chat_completion_reset_function_choice_settings():
|
||||
"""Test that _reset_function_choice_settings resets specific attributes."""
|
||||
chat_completion = MistralAIChatCompletion(
|
||||
ai_model_id="test-model",
|
||||
api_key="test_key",
|
||||
)
|
||||
settings = MistralAIChatPromptExecutionSettings(tool_choice="any", tools=[{"name": "func1"}])
|
||||
|
||||
chat_completion._reset_function_choice_settings(settings)
|
||||
assert settings.tool_choice is None
|
||||
assert settings.tools is None
|
||||
|
||||
|
||||
def test_mistral_ai_chat_completion_service_url():
|
||||
"""Test that service_url attempts to use _endpoint from the async_client."""
|
||||
async_mock_client = MagicMock(spec=Mistral)
|
||||
async_mock_client._endpoint = "mistral"
|
||||
|
||||
chat_completion = MistralAIChatCompletion(
|
||||
ai_model_id="test-model",
|
||||
api_key="test_key",
|
||||
async_client=async_mock_client,
|
||||
)
|
||||
|
||||
url = chat_completion.service_url()
|
||||
assert url == "mistral"
|
||||
+114
@@ -0,0 +1,114 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from mistralai import Mistral
|
||||
from mistralai.models import EmbeddingResponse
|
||||
|
||||
from semantic_kernel.connectors.ai.mistral_ai.services.mistral_ai_text_embedding import MistralAITextEmbedding
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceResponseException
|
||||
|
||||
|
||||
def test_embedding_with_env_variables(mistralai_unit_test_env):
|
||||
text_embedding = MistralAITextEmbedding()
|
||||
assert text_embedding.ai_model_id == "test_embedding_model_id"
|
||||
assert text_embedding.async_client.sdk_configuration.security.api_key == "test_api_key"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["MISTRALAI_API_KEY", "MISTRALAI_EMBEDDING_MODEL_ID"]], indirect=True)
|
||||
def test_embedding_with_constructor(mistralai_unit_test_env):
|
||||
text_embedding = MistralAITextEmbedding(
|
||||
api_key="overwrite-api-key",
|
||||
ai_model_id="overwrite-model",
|
||||
)
|
||||
assert text_embedding.ai_model_id == "overwrite-model"
|
||||
assert text_embedding.async_client.sdk_configuration.security.api_key == "overwrite-api-key"
|
||||
|
||||
|
||||
def test_embedding_with_client(mistralai_unit_test_env):
|
||||
client = MagicMock(spec=Mistral)
|
||||
text_embedding = MistralAITextEmbedding(async_client=client)
|
||||
assert text_embedding.async_client == client
|
||||
assert text_embedding.ai_model_id == "test_embedding_model_id"
|
||||
|
||||
|
||||
def test_embedding_with_api_key(mistralai_unit_test_env):
|
||||
text_embedding = MistralAITextEmbedding(api_key="overwrite-api-key")
|
||||
assert text_embedding.async_client.sdk_configuration.security.api_key == "overwrite-api-key"
|
||||
assert text_embedding.ai_model_id == "test_embedding_model_id"
|
||||
|
||||
|
||||
def test_embedding_with_model(mistralai_unit_test_env):
|
||||
text_embedding = MistralAITextEmbedding(ai_model_id="overwrite-model")
|
||||
assert text_embedding.ai_model_id == "overwrite-model"
|
||||
assert text_embedding.async_client.sdk_configuration.security.api_key == "test_api_key"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["MISTRALAI_EMBEDDING_MODEL_ID"]], indirect=True)
|
||||
def test_embedding_with_model_without_env(mistralai_unit_test_env):
|
||||
text_embedding = MistralAITextEmbedding(ai_model_id="overwrite-model")
|
||||
assert text_embedding.ai_model_id == "overwrite-model"
|
||||
assert text_embedding.async_client.sdk_configuration.security.api_key == "test_api_key"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["MISTRALAI_EMBEDDING_MODEL_ID"]], indirect=True)
|
||||
def test_embedding_missing_model(mistralai_unit_test_env):
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
MistralAITextEmbedding(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["MISTRALAI_API_KEY"]], indirect=True)
|
||||
def test_embedding_missing_api_key(mistralai_unit_test_env):
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
MistralAITextEmbedding(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["MISTRALAI_API_KEY", "MISTRALAI_EMBEDDING_MODEL_ID"]], indirect=True)
|
||||
def test_embedding_missing_api_key_constructor(mistralai_unit_test_env):
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
MistralAITextEmbedding(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["MISTRALAI_API_KEY", "MISTRALAI_EMBEDDING_MODEL_ID"]], indirect=True)
|
||||
def test_embedding_missing_model_constructor(mistralai_unit_test_env):
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
MistralAITextEmbedding(
|
||||
api_key="test_api_key",
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
async def test_embedding_generate_raw_embedding(mistralai_unit_test_env):
|
||||
mock_client = AsyncMock(spec=Mistral)
|
||||
mock_client.embeddings = AsyncMock()
|
||||
mock_embedding_response = MagicMock(spec=EmbeddingResponse, data=[MagicMock(embedding=[1, 2, 3, 4, 5])])
|
||||
mock_client.embeddings.create_async.return_value = mock_embedding_response
|
||||
text_embedding = MistralAITextEmbedding(async_client=mock_client)
|
||||
embedding = await text_embedding.generate_raw_embeddings(["test"])
|
||||
assert embedding == [[1, 2, 3, 4, 5]]
|
||||
|
||||
|
||||
async def test_embedding_generate_embedding(mistralai_unit_test_env):
|
||||
mock_client = AsyncMock(spec=Mistral)
|
||||
mock_client.embeddings = AsyncMock()
|
||||
mock_embedding_response = MagicMock(spec=EmbeddingResponse, data=[MagicMock(embedding=[1, 2, 3, 4, 5])])
|
||||
mock_client.embeddings.create_async.return_value = mock_embedding_response
|
||||
text_embedding = MistralAITextEmbedding(async_client=mock_client)
|
||||
embedding = await text_embedding.generate_embeddings(["test"])
|
||||
assert embedding.tolist() == [[1, 2, 3, 4, 5]]
|
||||
|
||||
|
||||
async def test_embedding_generate_embedding_exception(mistralai_unit_test_env):
|
||||
mock_client = AsyncMock(spec=Mistral)
|
||||
mock_client.embeddings = AsyncMock()
|
||||
mock_client.embeddings.create_async.side_effect = Exception("Test Exception")
|
||||
text_embedding = MistralAITextEmbedding(async_client=mock_client)
|
||||
with pytest.raises(ServiceResponseException):
|
||||
await text_embedding.generate_embeddings(["test"])
|
||||
@@ -0,0 +1,125 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from semantic_kernel.connectors.ai.mistral_ai.prompt_execution_settings.mistral_ai_prompt_execution_settings import (
|
||||
MistralAIChatPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
|
||||
|
||||
def test_default_mistralai_chat_prompt_execution_settings():
|
||||
settings = MistralAIChatPromptExecutionSettings()
|
||||
assert settings.temperature is None
|
||||
assert settings.top_p is None
|
||||
assert settings.max_tokens is None
|
||||
assert settings.messages is None
|
||||
|
||||
|
||||
def test_custom_mistralai_chat_prompt_execution_settings():
|
||||
settings = MistralAIChatPromptExecutionSettings(
|
||||
temperature=0.5,
|
||||
top_p=0.5,
|
||||
max_tokens=128,
|
||||
messages=[{"role": "system", "content": "Hello"}],
|
||||
)
|
||||
assert settings.temperature == 0.5
|
||||
assert settings.top_p == 0.5
|
||||
assert settings.max_tokens == 128
|
||||
assert settings.messages == [{"role": "system", "content": "Hello"}]
|
||||
|
||||
|
||||
def test_mistralai_chat_prompt_execution_settings_from_default_completion_config():
|
||||
settings = PromptExecutionSettings(service_id="test_service")
|
||||
chat_settings = MistralAIChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
assert chat_settings.service_id == "test_service"
|
||||
assert chat_settings.temperature is None
|
||||
assert chat_settings.top_p is None
|
||||
assert chat_settings.max_tokens is None
|
||||
|
||||
|
||||
def test_mistral_chat_prompt_execution_settings_from_openai_prompt_execution_settings():
|
||||
chat_settings = MistralAIChatPromptExecutionSettings(service_id="test_service", temperature=1.0)
|
||||
new_settings = MistralAIChatPromptExecutionSettings(service_id="test_2", temperature=0.0)
|
||||
chat_settings.update_from_prompt_execution_settings(new_settings)
|
||||
assert chat_settings.service_id == "test_2"
|
||||
assert chat_settings.temperature == 0.0
|
||||
|
||||
|
||||
def test_mistral_chat_prompt_execution_settings_from_custom_completion_config():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"max_tokens": 128,
|
||||
"messages": [{"role": "system", "content": "Hello"}],
|
||||
},
|
||||
)
|
||||
chat_settings = MistralAIChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
assert chat_settings.temperature == 0.5
|
||||
assert chat_settings.top_p == 0.5
|
||||
assert chat_settings.max_tokens == 128
|
||||
|
||||
|
||||
def test_openai_chat_prompt_execution_settings_from_custom_completion_config_with_none():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"max_tokens": 128,
|
||||
"messages": [{"role": "system", "content": "Hello"}],
|
||||
},
|
||||
)
|
||||
chat_settings = MistralAIChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
assert chat_settings.temperature == 0.5
|
||||
assert chat_settings.top_p == 0.5
|
||||
assert chat_settings.max_tokens == 128
|
||||
|
||||
|
||||
def test_openai_chat_prompt_execution_settings_from_custom_completion_config_with_functions():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"max_tokens": 128,
|
||||
"tools": [{}],
|
||||
"messages": [{"role": "system", "content": "Hello"}],
|
||||
},
|
||||
)
|
||||
chat_settings = MistralAIChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
assert chat_settings.temperature == 0.5
|
||||
assert chat_settings.top_p == 0.5
|
||||
assert chat_settings.max_tokens == 128
|
||||
|
||||
|
||||
def test_create_options():
|
||||
settings = MistralAIChatPromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"max_tokens": 128,
|
||||
"tools": [{}],
|
||||
"messages": [{"role": "system", "content": "Hello"}],
|
||||
},
|
||||
)
|
||||
options = settings.prepare_settings_dict()
|
||||
assert options["temperature"] == 0.5
|
||||
assert options["top_p"] == 0.5
|
||||
assert options["max_tokens"] == 128
|
||||
|
||||
|
||||
def test_create_options_with_function_choice_behavior():
|
||||
settings = MistralAIChatPromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
function_choice_behavior="auto",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"max_tokens": 128,
|
||||
"tools": [{}],
|
||||
"messages": [{"role": "system", "content": "Hello"}],
|
||||
},
|
||||
)
|
||||
assert settings.function_choice_behavior
|
||||
+101
@@ -0,0 +1,101 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from semantic_kernel.connectors.ai.nvidia.prompt_execution_settings.nvidia_prompt_execution_settings import (
|
||||
NvidiaChatPromptExecutionSettings,
|
||||
NvidiaEmbeddingPromptExecutionSettings,
|
||||
NvidiaPromptExecutionSettings,
|
||||
)
|
||||
|
||||
|
||||
class TestNvidiaPromptExecutionSettings:
|
||||
"""Test cases for NvidiaPromptExecutionSettings."""
|
||||
|
||||
def test_init_with_defaults(self):
|
||||
"""Test initialization with default values."""
|
||||
settings = NvidiaPromptExecutionSettings()
|
||||
assert settings.format is None
|
||||
assert settings.options is None
|
||||
|
||||
def test_init_with_values(self):
|
||||
"""Test initialization with specific values."""
|
||||
settings = NvidiaPromptExecutionSettings(
|
||||
format="json",
|
||||
options={"key": "value"},
|
||||
)
|
||||
assert settings.format == "json"
|
||||
assert settings.options == {"key": "value"}
|
||||
|
||||
def test_validation_format_values(self):
|
||||
"""Test format validation values."""
|
||||
# Valid values
|
||||
settings = NvidiaPromptExecutionSettings(format="json")
|
||||
assert settings.format == "json"
|
||||
|
||||
|
||||
class TestNvidiaChatPromptExecutionSettings:
|
||||
"""Test cases for NvidiaChatPromptExecutionSettings."""
|
||||
|
||||
def test_init_with_defaults(self):
|
||||
"""Test initialization with default values."""
|
||||
settings = NvidiaChatPromptExecutionSettings()
|
||||
assert settings.messages is None
|
||||
assert settings.response_format is None
|
||||
|
||||
def test_response_format_with_pydantic_model(self):
|
||||
"""Test response_format with Pydantic model."""
|
||||
|
||||
class TestModel(BaseModel):
|
||||
name: str
|
||||
value: int
|
||||
|
||||
settings = NvidiaChatPromptExecutionSettings(response_format=TestModel)
|
||||
|
||||
assert settings.response_format == TestModel
|
||||
|
||||
def test_response_format_with_dict(self):
|
||||
"""Test response_format with dictionary."""
|
||||
settings = NvidiaChatPromptExecutionSettings(response_format={"type": "json_object"})
|
||||
|
||||
assert settings.response_format == {"type": "json_object"}
|
||||
|
||||
|
||||
class TestNvidiaEmbeddingPromptExecutionSettings:
|
||||
"""Test cases for NvidiaEmbeddingPromptExecutionSettings."""
|
||||
|
||||
def test_init_with_defaults(self):
|
||||
"""Test initialization with default values."""
|
||||
settings = NvidiaEmbeddingPromptExecutionSettings()
|
||||
assert settings.input is None
|
||||
assert settings.encoding_format == "float"
|
||||
assert settings.input_type == "query"
|
||||
assert settings.truncate == "NONE"
|
||||
|
||||
def test_init_with_values(self):
|
||||
"""Test initialization with specific values."""
|
||||
settings = NvidiaEmbeddingPromptExecutionSettings(
|
||||
input=["hello", "world"],
|
||||
encoding_format="base64",
|
||||
input_type="passage",
|
||||
truncate="START",
|
||||
)
|
||||
|
||||
assert settings.input == ["hello", "world"]
|
||||
assert settings.encoding_format == "base64"
|
||||
assert settings.input_type == "passage"
|
||||
assert settings.truncate == "START"
|
||||
|
||||
def test_validation_encoding_format(self):
|
||||
"""Test encoding_format validation."""
|
||||
# Valid values
|
||||
settings = NvidiaEmbeddingPromptExecutionSettings(encoding_format="float")
|
||||
assert settings.encoding_format == "float"
|
||||
|
||||
settings = NvidiaEmbeddingPromptExecutionSettings(encoding_format="base64")
|
||||
assert settings.encoding_format == "base64"
|
||||
|
||||
# Invalid values
|
||||
with pytest.raises(ValidationError):
|
||||
NvidiaEmbeddingPromptExecutionSettings(encoding_format="invalid")
|
||||
@@ -0,0 +1,151 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from openai.resources.chat.completions import AsyncCompletions
|
||||
from openai.types.chat import ChatCompletion, ChatCompletionMessage
|
||||
from openai.types.chat.chat_completion import Choice
|
||||
from openai.types.completion_usage import CompletionUsage
|
||||
from pydantic import BaseModel
|
||||
|
||||
from semantic_kernel.connectors.ai.nvidia import NvidiaChatCompletion
|
||||
from semantic_kernel.connectors.ai.nvidia.prompt_execution_settings.nvidia_prompt_execution_settings import (
|
||||
NvidiaChatPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.nvidia.services.nvidia_chat_completion import DEFAULT_NVIDIA_CHAT_MODEL
|
||||
from semantic_kernel.contents import ChatHistory
|
||||
from semantic_kernel.exceptions import ServiceInitializationError, ServiceResponseException
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def nvidia_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for NvidiaChatCompletion."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {"NVIDIA_API_KEY": "test_api_key", "NVIDIA_CHAT_MODEL_ID": "meta/llama-3.1-8b-instruct"}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
def _create_mock_chat_completion(content: str = "Hello!") -> ChatCompletion:
|
||||
"""Helper function to create a mock ChatCompletion response."""
|
||||
message = ChatCompletionMessage(role="assistant", content=content)
|
||||
choice = Choice(
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
message=message,
|
||||
)
|
||||
usage = CompletionUsage(completion_tokens=20, prompt_tokens=10, total_tokens=30)
|
||||
return ChatCompletion(
|
||||
id="test-id",
|
||||
choices=[choice],
|
||||
created=1234567890,
|
||||
model="meta/llama-3.1-8b-instruct",
|
||||
object="chat.completion",
|
||||
usage=usage,
|
||||
)
|
||||
|
||||
|
||||
class TestNvidiaChatCompletion:
|
||||
"""Test cases for NvidiaChatCompletion."""
|
||||
|
||||
def test_init_with_defaults(self, nvidia_unit_test_env):
|
||||
"""Test initialization with default values."""
|
||||
service = NvidiaChatCompletion()
|
||||
assert service.ai_model_id == nvidia_unit_test_env["NVIDIA_CHAT_MODEL_ID"]
|
||||
|
||||
def test_get_prompt_execution_settings_class(self, nvidia_unit_test_env):
|
||||
"""Test getting the prompt execution settings class."""
|
||||
service = NvidiaChatCompletion()
|
||||
from semantic_kernel.connectors.ai.nvidia.prompt_execution_settings.nvidia_prompt_execution_settings import (
|
||||
NvidiaChatPromptExecutionSettings,
|
||||
)
|
||||
|
||||
assert service.get_prompt_execution_settings_class() == NvidiaChatPromptExecutionSettings
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["NVIDIA_API_KEY"]], indirect=True)
|
||||
def test_init_with_empty_api_key(self, nvidia_unit_test_env):
|
||||
"""Test initialization fails with empty API key."""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
NvidiaChatCompletion()
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["NVIDIA_CHAT_MODEL_ID"]], indirect=True)
|
||||
def test_init_with_empty_model_id(self, nvidia_unit_test_env):
|
||||
"""Test initialization with empty model ID uses default."""
|
||||
service = NvidiaChatCompletion()
|
||||
assert service.ai_model_id == DEFAULT_NVIDIA_CHAT_MODEL
|
||||
|
||||
def test_init_with_custom_model_id(self, nvidia_unit_test_env):
|
||||
"""Test initialization with custom model ID."""
|
||||
custom_model = "custom/nvidia-model"
|
||||
service = NvidiaChatCompletion(ai_model_id=custom_model)
|
||||
assert service.ai_model_id == custom_model
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch.object(AsyncCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_get_chat_message_contents(self, mock_create, nvidia_unit_test_env):
|
||||
"""Test basic chat completion."""
|
||||
mock_create.return_value = _create_mock_chat_completion("Hello!")
|
||||
|
||||
service = NvidiaChatCompletion()
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_user_message("Hello")
|
||||
settings = NvidiaChatPromptExecutionSettings()
|
||||
|
||||
result = await service.get_chat_message_contents(chat_history, settings)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].content == "Hello!"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch.object(AsyncCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_structured_output_with_pydantic_model(self, mock_create, nvidia_unit_test_env):
|
||||
"""Test structured output with Pydantic model."""
|
||||
|
||||
# Define test model
|
||||
class TestModel(BaseModel):
|
||||
name: str
|
||||
value: int
|
||||
|
||||
mock_create.return_value = _create_mock_chat_completion('{"name": "test", "value": 42}')
|
||||
|
||||
service = NvidiaChatCompletion()
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_user_message("Give me structured data")
|
||||
settings = NvidiaChatPromptExecutionSettings()
|
||||
settings.response_format = TestModel
|
||||
|
||||
await service.get_chat_message_contents(chat_history, settings)
|
||||
|
||||
# Verify nvext was passed
|
||||
call_args = mock_create.call_args[1]
|
||||
assert "extra_body" in call_args
|
||||
assert "nvext" in call_args["extra_body"]
|
||||
assert "guided_json" in call_args["extra_body"]["nvext"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch.object(AsyncCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_error_handling(self, mock_create, nvidia_unit_test_env):
|
||||
"""Test error handling."""
|
||||
mock_create.side_effect = Exception("API Error")
|
||||
|
||||
service = NvidiaChatCompletion()
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_user_message("Hello")
|
||||
settings = NvidiaChatPromptExecutionSettings()
|
||||
|
||||
with pytest.raises(ServiceResponseException):
|
||||
await service.get_chat_message_contents(chat_history, settings)
|
||||
@@ -0,0 +1,152 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from semantic_kernel.connectors.ai.nvidia.prompt_execution_settings.nvidia_prompt_execution_settings import (
|
||||
NvidiaChatPromptExecutionSettings,
|
||||
NvidiaEmbeddingPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.nvidia.services.nvidia_handler import NvidiaHandler
|
||||
from semantic_kernel.connectors.ai.nvidia.services.nvidia_model_types import NvidiaModelTypes
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_openai_client():
|
||||
"""Create a mock OpenAI client."""
|
||||
return AsyncMock(spec=AsyncOpenAI)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def nvidia_handler(mock_openai_client):
|
||||
"""Create a NvidiaHandler instance with mocked client."""
|
||||
return NvidiaHandler(
|
||||
client=mock_openai_client,
|
||||
ai_model_type=NvidiaModelTypes.CHAT,
|
||||
ai_model_id="test-model",
|
||||
api_key="test-key",
|
||||
)
|
||||
|
||||
|
||||
class TestNvidiaHandler:
|
||||
"""Test cases for NvidiaHandler."""
|
||||
|
||||
def test_init(self, mock_openai_client):
|
||||
"""Test initialization."""
|
||||
handler = NvidiaHandler(
|
||||
client=mock_openai_client,
|
||||
ai_model_type=NvidiaModelTypes.CHAT,
|
||||
)
|
||||
|
||||
assert handler.client == mock_openai_client
|
||||
assert handler.ai_model_type == NvidiaModelTypes.CHAT
|
||||
assert handler.MODEL_PROVIDER_NAME == "nvidia"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_chat_completion_request(self, nvidia_handler, mock_openai_client):
|
||||
"""Test sending chat completion request."""
|
||||
# Mock the response
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [
|
||||
MagicMock(
|
||||
message=MagicMock(role="assistant", content="Hello!"),
|
||||
finish_reason="stop",
|
||||
)
|
||||
]
|
||||
mock_response.usage = MagicMock(prompt_tokens=10, completion_tokens=20, total_tokens=30)
|
||||
mock_openai_client.chat.completions.create = AsyncMock(return_value=mock_response)
|
||||
|
||||
# Create settings
|
||||
settings = NvidiaChatPromptExecutionSettings(
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
model="test-model",
|
||||
)
|
||||
|
||||
# Test the method
|
||||
result = await nvidia_handler._send_chat_completion_request(settings)
|
||||
assert result == mock_response
|
||||
|
||||
# Verify usage was stored
|
||||
assert nvidia_handler.prompt_tokens == 10
|
||||
assert nvidia_handler.completion_tokens == 20
|
||||
assert nvidia_handler.total_tokens == 30
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_chat_completion_request_with_nvext(self, nvidia_handler, mock_openai_client):
|
||||
"""Test sending chat completion request with nvext parameter."""
|
||||
# Mock the response
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [
|
||||
MagicMock(
|
||||
message=MagicMock(role="assistant", content='{"result": "success"}'),
|
||||
finish_reason="stop",
|
||||
)
|
||||
]
|
||||
mock_response.usage = MagicMock(prompt_tokens=10, completion_tokens=20, total_tokens=30)
|
||||
mock_openai_client.chat.completions.create = AsyncMock(return_value=mock_response)
|
||||
|
||||
# Create settings with nvext
|
||||
settings = NvidiaChatPromptExecutionSettings(
|
||||
messages=[{"role": "user", "content": "Give me JSON"}],
|
||||
model="test-model",
|
||||
extra_body={"nvext": {"guided_json": {"type": "object"}}},
|
||||
)
|
||||
|
||||
# Test the method
|
||||
result = await nvidia_handler._send_chat_completion_request(settings)
|
||||
assert result == mock_response
|
||||
|
||||
# Verify the client was called with nvext in extra_body
|
||||
call_args = mock_openai_client.chat.completions.create.call_args[1]
|
||||
assert "extra_body" in call_args
|
||||
assert "nvext" in call_args["extra_body"]
|
||||
assert call_args["extra_body"]["nvext"] == {"guided_json": {"type": "object"}}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_embedding_request(self, mock_openai_client):
|
||||
"""Test sending embedding request."""
|
||||
handler = NvidiaHandler(
|
||||
client=mock_openai_client,
|
||||
ai_model_type=NvidiaModelTypes.EMBEDDING,
|
||||
ai_model_id="test-model",
|
||||
)
|
||||
|
||||
# Mock the response
|
||||
mock_response = MagicMock()
|
||||
mock_response.data = [
|
||||
MagicMock(embedding=[0.1, 0.2, 0.3]),
|
||||
MagicMock(embedding=[0.4, 0.5, 0.6]),
|
||||
]
|
||||
mock_response.usage = MagicMock(prompt_tokens=10, total_tokens=10)
|
||||
mock_openai_client.embeddings.create = AsyncMock(return_value=mock_response)
|
||||
|
||||
# Create settings
|
||||
settings = NvidiaEmbeddingPromptExecutionSettings(
|
||||
input=["hello", "world"],
|
||||
model="test-model",
|
||||
)
|
||||
|
||||
# Test the method
|
||||
result = await handler._send_embedding_request(settings)
|
||||
assert result == [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_request_unsupported_model_type(self, mock_openai_client):
|
||||
"""Test send_request with unsupported model type."""
|
||||
# Create a handler with invalid model type by bypassing validation
|
||||
handler = NvidiaHandler(
|
||||
client=mock_openai_client,
|
||||
ai_model_type=NvidiaModelTypes.CHAT,
|
||||
)
|
||||
# Manually set the attribute to bypass Pydantic validation
|
||||
object.__setattr__(handler, "ai_model_type", "UNSUPPORTED")
|
||||
|
||||
settings = NvidiaChatPromptExecutionSettings(
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
model="test-model",
|
||||
)
|
||||
|
||||
with pytest.raises(NotImplementedError, match="Model type UNSUPPORTED is not supported"):
|
||||
await handler._send_request(settings)
|
||||
@@ -0,0 +1,134 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from openai import AsyncClient
|
||||
from openai.resources.embeddings import AsyncEmbeddings
|
||||
|
||||
from semantic_kernel.connectors.ai.nvidia.prompt_execution_settings.nvidia_prompt_execution_settings import (
|
||||
NvidiaEmbeddingPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.nvidia.services.nvidia_text_embedding import NvidiaTextEmbedding
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceResponseException
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def nvidia_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for NvidiaTextEmbedding."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {"NVIDIA_API_KEY": "test_api_key", "NVIDIA_EMBEDDING_MODEL_ID": "test_embedding_model_id"}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
def test_init(nvidia_unit_test_env):
|
||||
nvidia_text_embedding = NvidiaTextEmbedding()
|
||||
|
||||
assert nvidia_text_embedding.client is not None
|
||||
assert isinstance(nvidia_text_embedding.client, AsyncClient)
|
||||
assert nvidia_text_embedding.ai_model_id == nvidia_unit_test_env["NVIDIA_EMBEDDING_MODEL_ID"]
|
||||
|
||||
assert nvidia_text_embedding.get_prompt_execution_settings_class() == NvidiaEmbeddingPromptExecutionSettings
|
||||
|
||||
|
||||
def test_init_validation_fail() -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
NvidiaTextEmbedding(api_key="34523", ai_model_id={"test": "dict"})
|
||||
|
||||
|
||||
def test_init_to_from_dict(nvidia_unit_test_env):
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
settings = {
|
||||
"ai_model_id": nvidia_unit_test_env["NVIDIA_EMBEDDING_MODEL_ID"],
|
||||
"api_key": nvidia_unit_test_env["NVIDIA_API_KEY"],
|
||||
"default_headers": default_headers,
|
||||
}
|
||||
text_embedding = NvidiaTextEmbedding.from_dict(settings)
|
||||
dumped_settings = text_embedding.to_dict()
|
||||
assert dumped_settings["ai_model_id"] == settings["ai_model_id"]
|
||||
assert dumped_settings["api_key"] == settings["api_key"]
|
||||
|
||||
|
||||
@patch.object(AsyncEmbeddings, "create", new_callable=AsyncMock)
|
||||
async def test_embedding_calls_with_parameters(mock_create, nvidia_unit_test_env) -> None:
|
||||
ai_model_id = "NV-Embed-QA"
|
||||
texts = ["hello world", "goodbye world"]
|
||||
embedding_dimensions = 1536
|
||||
|
||||
nvidia_text_embedding = NvidiaTextEmbedding(
|
||||
ai_model_id=ai_model_id,
|
||||
)
|
||||
|
||||
await nvidia_text_embedding.generate_embeddings(texts, dimensions=embedding_dimensions)
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
input=texts,
|
||||
model=ai_model_id,
|
||||
dimensions=embedding_dimensions,
|
||||
encoding_format="float",
|
||||
extra_body={"input_type": "query", "truncate": "NONE"},
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncEmbeddings, "create", new_callable=AsyncMock)
|
||||
async def test_embedding_calls_with_settings(mock_create, nvidia_unit_test_env) -> None:
|
||||
ai_model_id = "test_model_id"
|
||||
texts = ["hello world", "goodbye world"]
|
||||
settings = NvidiaEmbeddingPromptExecutionSettings(service_id="default")
|
||||
nvidia_text_embedding = NvidiaTextEmbedding(service_id="default", ai_model_id=ai_model_id)
|
||||
|
||||
await nvidia_text_embedding.generate_embeddings(texts, settings=settings)
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
input=texts,
|
||||
model=ai_model_id,
|
||||
encoding_format="float",
|
||||
extra_body={"input_type": "query", "truncate": "NONE"},
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncEmbeddings, "create", new_callable=AsyncMock, side_effect=Exception)
|
||||
async def test_embedding_fail(mock_create, nvidia_unit_test_env) -> None:
|
||||
ai_model_id = "test_model_id"
|
||||
texts = ["hello world", "goodbye world"]
|
||||
|
||||
nvidia_text_embedding = NvidiaTextEmbedding(
|
||||
ai_model_id=ai_model_id,
|
||||
)
|
||||
with pytest.raises(ServiceResponseException):
|
||||
await nvidia_text_embedding.generate_embeddings(texts)
|
||||
|
||||
|
||||
@patch.object(AsyncEmbeddings, "create", new_callable=AsyncMock)
|
||||
async def test_embedding_pes(mock_create, nvidia_unit_test_env) -> None:
|
||||
ai_model_id = "test_model_id"
|
||||
texts = ["hello world", "goodbye world"]
|
||||
|
||||
pes = PromptExecutionSettings(service_id="x", ai_model_id=ai_model_id)
|
||||
|
||||
nvidia_text_embedding = NvidiaTextEmbedding(ai_model_id=ai_model_id)
|
||||
|
||||
await nvidia_text_embedding.generate_raw_embeddings(texts, pes)
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
input=texts,
|
||||
model=ai_model_id,
|
||||
encoding_format="float",
|
||||
extra_body={"input_type": "query", "truncate": "NONE"},
|
||||
)
|
||||
@@ -0,0 +1,56 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from semantic_kernel.connectors.ai.nvidia.settings.nvidia_settings import NvidiaSettings
|
||||
|
||||
|
||||
class TestNvidiaSettings:
|
||||
"""Test cases for NvidiaSettings."""
|
||||
|
||||
def test_init_with_defaults(self):
|
||||
"""Test initialization with default values."""
|
||||
settings = NvidiaSettings()
|
||||
assert settings.api_key is None
|
||||
assert settings.base_url == "https://integrate.api.nvidia.com/v1"
|
||||
assert settings.embedding_model_id is None
|
||||
assert settings.chat_model_id is None
|
||||
|
||||
def test_init_with_values(self):
|
||||
"""Test initialization with specific values."""
|
||||
settings = NvidiaSettings(
|
||||
api_key="test-api-key",
|
||||
base_url="https://custom.nvidia.com/v1",
|
||||
embedding_model_id="test-embedding-model",
|
||||
chat_model_id="test-chat-model",
|
||||
)
|
||||
|
||||
assert settings.api_key.get_secret_value() == "test-api-key"
|
||||
assert settings.base_url == "https://custom.nvidia.com/v1"
|
||||
assert settings.embedding_model_id == "test-embedding-model"
|
||||
assert settings.chat_model_id == "test-chat-model"
|
||||
|
||||
def test_env_prefix(self):
|
||||
"""Test environment variable prefix."""
|
||||
assert NvidiaSettings.env_prefix == "NVIDIA_"
|
||||
|
||||
def test_api_key_secret_str(self):
|
||||
"""Test that api_key is properly handled as SecretStr."""
|
||||
settings = NvidiaSettings(api_key="secret-key")
|
||||
|
||||
# Should be SecretStr type
|
||||
assert hasattr(settings.api_key, "get_secret_value")
|
||||
assert settings.api_key.get_secret_value() == "secret-key"
|
||||
|
||||
# Should not expose the secret in string representation
|
||||
str_repr = str(settings)
|
||||
assert "secret-key" not in str_repr
|
||||
|
||||
def test_environment_variables(self, monkeypatch):
|
||||
"""Test that environment variables override defaults."""
|
||||
monkeypatch.setenv("NVIDIA_API_KEY", "env-key")
|
||||
monkeypatch.setenv("NVIDIA_CHAT_MODEL_ID", "env-chat")
|
||||
|
||||
settings = NvidiaSettings()
|
||||
|
||||
assert settings.api_key.get_secret_value() == "env-key"
|
||||
assert settings.chat_model_id == "env-chat"
|
||||
@@ -0,0 +1,84 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from collections.abc import AsyncGenerator, AsyncIterator
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from ollama import AsyncClient
|
||||
|
||||
from semantic_kernel.contents.chat_history import ChatHistory
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def model_id() -> str:
|
||||
return "test_model_id"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def service_id() -> str:
|
||||
return "test_service_id"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def host() -> str:
|
||||
return "http://localhost:5000"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def custom_client() -> AsyncClient:
|
||||
return AsyncClient("http://localhost:5001")
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def chat_history() -> ChatHistory:
|
||||
chat_history = ChatHistory()
|
||||
chat_history.add_user_message("test_prompt")
|
||||
return chat_history
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def prompt() -> str:
|
||||
return "test_prompt"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def default_options() -> dict:
|
||||
return {"test": "test"}
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def ollama_unit_test_env(monkeypatch, host, exclude_list):
|
||||
"""Fixture to set environment variables for OllamaSettings."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
env_vars = {
|
||||
"OLLAMA_CHAT_MODEL_ID": "test_chat_model_id",
|
||||
"OLLAMA_TEXT_MODEL_ID": "test_text_model_id",
|
||||
"OLLAMA_EMBEDDING_MODEL_ID": "test_embedding_model_id",
|
||||
"OLLAMA_HOST": host,
|
||||
}
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_streaming_text_response() -> AsyncIterator:
|
||||
streaming_text_response = MagicMock(spec=AsyncGenerator)
|
||||
streaming_text_response.__aiter__.return_value = [{"response": "test_response"}]
|
||||
|
||||
return streaming_text_response
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_streaming_chat_response() -> AsyncIterator:
|
||||
streaming_chat_response = MagicMock(spec=AsyncGenerator)
|
||||
streaming_chat_response.__aiter__.return_value = [{"message": {"content": "test_response"}}]
|
||||
|
||||
return streaming_chat_response
|
||||
@@ -0,0 +1,453 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from ollama import AsyncClient
|
||||
|
||||
import semantic_kernel.connectors.ai.ollama.services.ollama_chat_completion as occ_module
|
||||
from semantic_kernel.connectors.ai.completion_usage import CompletionUsage
|
||||
from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior
|
||||
from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import OllamaChatPromptExecutionSettings
|
||||
from semantic_kernel.connectors.ai.ollama.services.ollama_chat_completion import OllamaChatCompletion
|
||||
from semantic_kernel.contents.chat_history import ChatHistory
|
||||
from semantic_kernel.contents.chat_message_content import ChatMessageContent
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent
|
||||
from semantic_kernel.contents.streaming_text_content import StreamingTextContent
|
||||
from semantic_kernel.exceptions.service_exceptions import (
|
||||
ServiceInitializationError,
|
||||
ServiceInvalidExecutionSettingsError,
|
||||
ServiceInvalidResponseError,
|
||||
)
|
||||
|
||||
|
||||
def test_settings(model_id):
|
||||
"""Test that the settings class is correct."""
|
||||
ollama = OllamaChatCompletion(ai_model_id=model_id)
|
||||
settings = ollama.get_prompt_execution_settings_class()
|
||||
assert settings == OllamaChatPromptExecutionSettings
|
||||
|
||||
|
||||
def test_init_empty_service_id(model_id):
|
||||
"""Test that the service initializes correctly with an empty service id."""
|
||||
ollama = OllamaChatCompletion(ai_model_id=model_id)
|
||||
assert ollama.service_id == model_id
|
||||
|
||||
|
||||
def test_init_empty_string_ai_model_id():
|
||||
"""Test that the service initializes with a error if there is no ai_model_id."""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
_ = OllamaChatCompletion(ai_model_id="")
|
||||
|
||||
|
||||
def test_custom_client(model_id, custom_client):
|
||||
"""Test that the service initializes correctly with a custom client."""
|
||||
ollama = OllamaChatCompletion(ai_model_id=model_id, client=custom_client)
|
||||
assert ollama.client == custom_client
|
||||
|
||||
|
||||
def test_invalid_ollama_settings():
|
||||
"""Test that the service initializes incorrectly with invalid settings."""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
_ = OllamaChatCompletion(ai_model_id=123)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OLLAMA_CHAT_MODEL_ID"]], indirect=True)
|
||||
def test_init_empty_model_id_in_env(ollama_unit_test_env):
|
||||
"""Test that the service initializes incorrectly with an empty model id."""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
_ = OllamaChatCompletion(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
def test_function_choice_settings(ollama_unit_test_env):
|
||||
"""Test that REQUIRED and NONE function choice settings are unsupported."""
|
||||
ollama = OllamaChatCompletion()
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
ollama._verify_function_choice_settings(
|
||||
OllamaChatPromptExecutionSettings(function_choice_behavior=FunctionChoiceBehavior.Required())
|
||||
)
|
||||
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
ollama._verify_function_choice_settings(
|
||||
OllamaChatPromptExecutionSettings(function_choice_behavior=FunctionChoiceBehavior.NoneInvoke())
|
||||
)
|
||||
|
||||
|
||||
def test_service_url(ollama_unit_test_env):
|
||||
"""Test that the service URL is correct."""
|
||||
ollama = OllamaChatCompletion()
|
||||
assert ollama.service_url() == ollama_unit_test_env["OLLAMA_HOST"]
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.__init__", return_value=None) # mock_client
|
||||
@patch("ollama.AsyncClient.chat") # mock_chat_client
|
||||
async def test_custom_host(
|
||||
mock_chat_client,
|
||||
mock_client,
|
||||
model_id,
|
||||
service_id,
|
||||
host,
|
||||
chat_history,
|
||||
prompt,
|
||||
default_options,
|
||||
):
|
||||
"""Test that the service initializes and generates content correctly with a custom host."""
|
||||
mock_chat_client.return_value = {"message": {"content": "test_response"}}
|
||||
|
||||
ollama = OllamaChatCompletion(ai_model_id=model_id, host=host)
|
||||
|
||||
chat_responses = await ollama.get_chat_message_contents(
|
||||
chat_history,
|
||||
OllamaChatPromptExecutionSettings(service_id=service_id, options=default_options),
|
||||
)
|
||||
|
||||
# Check that the client was initialized once with the correct host
|
||||
assert mock_client.call_count == 1
|
||||
mock_client.assert_called_with(host=host)
|
||||
# Check that the chat client was called once and the responses are correct
|
||||
assert mock_chat_client.call_count == 1
|
||||
assert len(chat_responses) == 1
|
||||
assert chat_responses[0].content == "test_response"
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.__init__", return_value=None) # mock_client
|
||||
@patch("ollama.AsyncClient.chat") # mock_chat_client
|
||||
async def test_custom_host_streaming(
|
||||
mock_chat_client,
|
||||
mock_client,
|
||||
mock_streaming_chat_response,
|
||||
model_id,
|
||||
service_id,
|
||||
host,
|
||||
chat_history,
|
||||
prompt,
|
||||
default_options,
|
||||
):
|
||||
"""Test that the service initializes and generates streaming content correctly with a custom host."""
|
||||
mock_chat_client.return_value = mock_streaming_chat_response
|
||||
|
||||
ollama = OllamaChatCompletion(ai_model_id=model_id, host=host)
|
||||
|
||||
async for messages in ollama.get_streaming_chat_message_contents(
|
||||
chat_history,
|
||||
OllamaChatPromptExecutionSettings(service_id=service_id, options=default_options),
|
||||
):
|
||||
assert len(messages) == 1
|
||||
assert messages[0].role == "assistant"
|
||||
assert messages[0].content == "test_response"
|
||||
|
||||
# Check that the client was initialized once with the correct host
|
||||
assert mock_client.call_count == 1
|
||||
mock_client.assert_called_with(host=host)
|
||||
# Check that the chat client was called once
|
||||
assert mock_chat_client.call_count == 1
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.chat")
|
||||
async def test_chat_completion(mock_chat_client, model_id, service_id, chat_history, default_options):
|
||||
"""Test that the chat completion service completes correctly."""
|
||||
mock_chat_client.return_value = {"message": {"content": "test_response"}}
|
||||
|
||||
ollama = OllamaChatCompletion(ai_model_id=model_id)
|
||||
response = await ollama.get_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=OllamaChatPromptExecutionSettings(service_id=service_id, options=default_options),
|
||||
)
|
||||
|
||||
assert len(response) == 1
|
||||
assert response[0].content == "test_response"
|
||||
mock_chat_client.assert_called_once_with(
|
||||
model=model_id,
|
||||
messages=ollama._prepare_chat_history_for_request(chat_history),
|
||||
options=default_options,
|
||||
stream=False,
|
||||
)
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.chat")
|
||||
async def test_chat_completion_wrong_return_type(
|
||||
mock_chat_client,
|
||||
mock_streaming_chat_response,
|
||||
model_id,
|
||||
service_id,
|
||||
chat_history,
|
||||
default_options,
|
||||
):
|
||||
"""Test that the chat completion service fails when the return type is incorrect."""
|
||||
mock_chat_client.return_value = mock_streaming_chat_response # should not be a streaming response
|
||||
|
||||
ollama = OllamaChatCompletion(ai_model_id=model_id)
|
||||
with pytest.raises(ServiceInvalidResponseError):
|
||||
await ollama.get_chat_message_contents(
|
||||
chat_history,
|
||||
OllamaChatPromptExecutionSettings(service_id=service_id, options=default_options),
|
||||
)
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.chat")
|
||||
async def test_streaming_chat_completion(
|
||||
mock_chat_client,
|
||||
mock_streaming_chat_response,
|
||||
model_id,
|
||||
service_id,
|
||||
chat_history,
|
||||
default_options,
|
||||
):
|
||||
"""Test that the streaming chat completion service completes correctly."""
|
||||
mock_chat_client.return_value = mock_streaming_chat_response
|
||||
|
||||
ollama = OllamaChatCompletion(ai_model_id=model_id)
|
||||
response = ollama.get_streaming_chat_message_contents(
|
||||
chat_history,
|
||||
OllamaChatPromptExecutionSettings(service_id=service_id, options=default_options),
|
||||
)
|
||||
|
||||
responses = []
|
||||
async for line in response:
|
||||
if line:
|
||||
assert line[0].content == "test_response"
|
||||
responses.append(line[0].content)
|
||||
assert len(responses) == 1
|
||||
|
||||
mock_chat_client.assert_called_once_with(
|
||||
model=model_id,
|
||||
messages=ollama._prepare_chat_history_for_request(chat_history),
|
||||
options=default_options,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.chat")
|
||||
async def test_streaming_chat_completion_wrong_return_type(
|
||||
mock_chat_client,
|
||||
model_id,
|
||||
service_id,
|
||||
chat_history,
|
||||
default_options,
|
||||
):
|
||||
"""Test that the chat completion streaming service fails when the return type is incorrect."""
|
||||
mock_chat_client.return_value = {"message": {"content": "test_response"}} # should not be a non-streaming response
|
||||
|
||||
ollama = OllamaChatCompletion(ai_model_id=model_id)
|
||||
with pytest.raises(ServiceInvalidResponseError):
|
||||
async for _ in ollama.get_streaming_chat_message_contents(
|
||||
chat_history,
|
||||
OllamaChatPromptExecutionSettings(service_id=service_id, options=default_options),
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def setup_ollama_chat_completion():
|
||||
async_client_mock = AsyncMock(spec=AsyncClient)
|
||||
async_client_mock.chat = AsyncMock()
|
||||
ollama_chat_completion = OllamaChatCompletion(
|
||||
service_id="test_service", ai_model_id="test_model_id", client=async_client_mock
|
||||
)
|
||||
return ollama_chat_completion, async_client_mock
|
||||
|
||||
|
||||
async def test_service_url_new(setup_ollama_chat_completion):
|
||||
ollama_chat_completion, async_client_mock = setup_ollama_chat_completion
|
||||
# Mock the client's internal structure
|
||||
async_client_mock._client = AsyncMock(spec=httpx.AsyncClient)
|
||||
async_client_mock._client.base_url = "http://mocked_base_url"
|
||||
|
||||
service_url = ollama_chat_completion.service_url()
|
||||
assert service_url == "http://mocked_base_url"
|
||||
|
||||
|
||||
async def test_prepare_chat_history_for_request(setup_ollama_chat_completion):
|
||||
ollama_chat_completion, _ = setup_ollama_chat_completion
|
||||
chat_history = MagicMock(spec=ChatHistory)
|
||||
chat_history.messages = []
|
||||
|
||||
prepared_history = ollama_chat_completion._prepare_chat_history_for_request(chat_history)
|
||||
assert prepared_history == []
|
||||
|
||||
|
||||
async def test_service_url_with_httpx_client(model_id: str) -> None:
|
||||
"""
|
||||
Test that service_url returns the base_url of the underlying httpx.AsyncClient.
|
||||
"""
|
||||
# Initialize an AsyncClient and manually set its _client attribute to an httpx.AsyncClient
|
||||
client = AsyncClient(host="unused")
|
||||
base = httpx.AsyncClient(base_url="http://example.com:8000")
|
||||
client._client = base # simulate underlying httpx client
|
||||
|
||||
ollama = OllamaChatCompletion(ai_model_id=model_id, client=client)
|
||||
# service_url should reflect the base_url of the httpx client
|
||||
assert ollama.service_url() == "http://example.com:8000"
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.chat", new_callable=AsyncMock)
|
||||
async def test_chat_response_branch(
|
||||
mock_chat: AsyncMock,
|
||||
model_id: str,
|
||||
service_id: str,
|
||||
default_options: dict,
|
||||
chat_history,
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
"""
|
||||
Test get_chat_message_contents when AsyncClient.chat returns a ChatResponse instance.
|
||||
"""
|
||||
|
||||
class DummyFunction:
|
||||
def __init__(self, name, arguments):
|
||||
self.name = name
|
||||
self.arguments = arguments
|
||||
|
||||
class DummyToolCall:
|
||||
def __init__(self, function):
|
||||
self.function = function
|
||||
|
||||
class DummyMessage:
|
||||
def __init__(self, content: str, tool_calls=None) -> None:
|
||||
self.content = content
|
||||
self.tool_calls = tool_calls or []
|
||||
|
||||
class DummyChatResponse:
|
||||
def __init__(
|
||||
self,
|
||||
content: str,
|
||||
model: str,
|
||||
prompt_eval_count: int,
|
||||
eval_count: int,
|
||||
tool_calls=None,
|
||||
) -> None:
|
||||
function_calls = [
|
||||
DummyToolCall(DummyFunction(tc["function"]["name"], tc["function"]["arguments"])) for tc in tool_calls
|
||||
]
|
||||
self.message = DummyMessage(content, function_calls)
|
||||
self.model = model
|
||||
self.prompt_eval_count = prompt_eval_count
|
||||
self.eval_count = eval_count
|
||||
|
||||
# Monkeypatch the ChatResponse type in the module so isinstance works
|
||||
monkeypatch.setattr(occ_module, "ChatResponse", DummyChatResponse)
|
||||
|
||||
# Prepare a dummy ChatResponse return value
|
||||
dummy_resp = DummyChatResponse(
|
||||
content="resp_text",
|
||||
model="mdl",
|
||||
prompt_eval_count=2,
|
||||
eval_count=3,
|
||||
tool_calls=[{"function": {"name": "fn", "arguments": {"x": 1}}}],
|
||||
)
|
||||
mock_chat.return_value = dummy_resp
|
||||
|
||||
ollama = OllamaChatCompletion(ai_model_id=model_id)
|
||||
settings = OllamaChatPromptExecutionSettings(service_id=service_id, options=default_options)
|
||||
|
||||
results = await ollama.get_chat_message_contents(chat_history, settings)
|
||||
# Only one response expected
|
||||
assert len(results) == 1
|
||||
msg = results[0]
|
||||
# Assert it's a ChatMessageContent
|
||||
assert isinstance(msg, ChatMessageContent)
|
||||
# The content property should return the response text
|
||||
assert msg.content == "resp_text"
|
||||
|
||||
# The second item should be a FunctionCallContent
|
||||
func_item = msg.items[1]
|
||||
assert isinstance(func_item, FunctionCallContent)
|
||||
# Validate function call details
|
||||
assert func_item.name == "fn"
|
||||
assert func_item.arguments == {"x": 1}
|
||||
|
||||
# Check metadata
|
||||
assert "model" in msg.metadata and msg.metadata["model"] == "mdl"
|
||||
# Access usage directly, key should exist
|
||||
usage = msg.metadata["usage"]
|
||||
assert isinstance(usage, CompletionUsage)
|
||||
assert usage.prompt_tokens == 2 and usage.completion_tokens == 3
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.chat", new_callable=AsyncMock)
|
||||
async def test_streaming_chat_response_branch(
|
||||
mock_chat: AsyncMock,
|
||||
model_id: str,
|
||||
service_id: str,
|
||||
default_options: dict,
|
||||
chat_history,
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
"""
|
||||
Test get_streaming_chat_message_contents when AsyncClient.chat yields ChatResponse instances.
|
||||
"""
|
||||
|
||||
class DummyFunction:
|
||||
def __init__(self, name, arguments):
|
||||
self.name = name
|
||||
self.arguments = arguments
|
||||
|
||||
class DummyToolCall:
|
||||
def __init__(self, function):
|
||||
self.function = function
|
||||
|
||||
class DummyMessage:
|
||||
def __init__(self, content: str, tool_calls=None) -> None:
|
||||
self.content = content
|
||||
self.tool_calls = tool_calls or []
|
||||
|
||||
class DummyChatResponse:
|
||||
def __init__(
|
||||
self,
|
||||
content: str,
|
||||
model: str,
|
||||
prompt_eval_count: int,
|
||||
eval_count: int,
|
||||
tool_calls=None,
|
||||
) -> None:
|
||||
function_calls = [
|
||||
DummyToolCall(DummyFunction(tc["function"]["name"], tc["function"]["arguments"])) for tc in tool_calls
|
||||
]
|
||||
self.message = DummyMessage(content, function_calls)
|
||||
self.model = model
|
||||
self.prompt_eval_count = prompt_eval_count
|
||||
self.eval_count = eval_count
|
||||
|
||||
# Monkeypatch ChatResponse type
|
||||
monkeypatch.setattr(occ_module, "ChatResponse", DummyChatResponse)
|
||||
|
||||
# Prepare an async generator yielding DummyChatResponse
|
||||
async def fake_stream() -> AsyncGenerator[DummyChatResponse, None]:
|
||||
yield DummyChatResponse(
|
||||
content="stream_text",
|
||||
model="m2",
|
||||
prompt_eval_count=1,
|
||||
eval_count=1,
|
||||
tool_calls=[{"function": {"name": "f2", "arguments": {}}}],
|
||||
)
|
||||
|
||||
mock_chat.return_value = fake_stream()
|
||||
|
||||
ollama = OllamaChatCompletion(ai_model_id=model_id)
|
||||
settings = OllamaChatPromptExecutionSettings(service_id=service_id, options=default_options)
|
||||
|
||||
collected = []
|
||||
# Iterate over streamed batches
|
||||
async for batch in ollama.get_streaming_chat_message_contents(chat_history, settings):
|
||||
# We expect a list with a single StreamingChatMessageContent
|
||||
assert len(batch) == 1
|
||||
sc = batch[0]
|
||||
assert isinstance(sc, StreamingChatMessageContent)
|
||||
|
||||
# First item should be text content
|
||||
text_item = sc.items[0]
|
||||
assert isinstance(text_item, StreamingTextContent)
|
||||
assert text_item.text == "stream_text"
|
||||
|
||||
# Next item should be a FunctionCallContent
|
||||
func_item = sc.items[1]
|
||||
assert isinstance(func_item, FunctionCallContent)
|
||||
assert func_item.name == "f2"
|
||||
|
||||
collected.append(sc)
|
||||
|
||||
# Only one batch should be collected
|
||||
assert len(collected) == 1
|
||||
@@ -0,0 +1,178 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import OllamaTextPromptExecutionSettings
|
||||
from semantic_kernel.connectors.ai.ollama.services.ollama_text_completion import OllamaTextCompletion
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidResponseError
|
||||
|
||||
|
||||
def test_settings(model_id):
|
||||
"""Test that the settings class is correct"""
|
||||
ollama = OllamaTextCompletion(ai_model_id=model_id)
|
||||
settings = ollama.get_prompt_execution_settings_class()
|
||||
assert settings == OllamaTextPromptExecutionSettings
|
||||
|
||||
|
||||
def test_init_empty_service_id(model_id):
|
||||
"""Test that the service initializes correctly with an empty service_id"""
|
||||
ollama = OllamaTextCompletion(ai_model_id=model_id)
|
||||
assert ollama.service_id == model_id
|
||||
|
||||
|
||||
def test_custom_client(model_id, custom_client):
|
||||
"""Test that the service initializes correctly with a custom client."""
|
||||
ollama = OllamaTextCompletion(ai_model_id=model_id, client=custom_client)
|
||||
assert ollama.client == custom_client
|
||||
|
||||
|
||||
def test_invalid_ollama_settings():
|
||||
"""Test that the service initializes incorrectly with invalid settings."""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
_ = OllamaTextCompletion(ai_model_id=123)
|
||||
|
||||
|
||||
def test_service_url(ollama_unit_test_env):
|
||||
"""Test that the service URL is correct."""
|
||||
ollama = OllamaTextCompletion()
|
||||
assert ollama.service_url() == ollama_unit_test_env["OLLAMA_HOST"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OLLAMA_TEXT_MODEL_ID"]], indirect=True)
|
||||
def test_init_empty_model_id(ollama_unit_test_env):
|
||||
"""Test that the service initializes incorrectly with an empty model_id"""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
_ = OllamaTextCompletion(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.__init__", return_value=None) # mock_client
|
||||
@patch("ollama.AsyncClient.generate") # mock_completion_client
|
||||
async def test_custom_host(
|
||||
mock_completion_client, mock_client, model_id, service_id, host, chat_history, default_options
|
||||
):
|
||||
"""Test that the service initializes and generates content correctly with a custom host."""
|
||||
mock_completion_client.return_value = {"response": "test_response"}
|
||||
|
||||
ollama = OllamaTextCompletion(ai_model_id=model_id, host=host)
|
||||
_ = await ollama.get_text_contents(
|
||||
chat_history,
|
||||
OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options),
|
||||
)
|
||||
|
||||
mock_client.assert_called_once_with(host=host)
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.__init__", return_value=None) # mock_client
|
||||
@patch("ollama.AsyncClient.generate") # mock_completion_client
|
||||
async def test_custom_host_streaming(
|
||||
mock_completion_client, mock_client, model_id, service_id, host, chat_history, default_options
|
||||
):
|
||||
"""Test that the service initializes and generates streaming content correctly with a custom host."""
|
||||
# Create a proper async iterator mock for streaming
|
||||
streaming_response = MagicMock(spec=AsyncGenerator)
|
||||
streaming_response.__aiter__.return_value = [{"response": "test_response"}]
|
||||
mock_completion_client.return_value = streaming_response
|
||||
|
||||
ollama = OllamaTextCompletion(ai_model_id=model_id, host=host)
|
||||
async for _ in ollama.get_streaming_text_contents(
|
||||
chat_history,
|
||||
OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options),
|
||||
):
|
||||
pass
|
||||
|
||||
mock_client.assert_called_once_with(host=host)
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.generate")
|
||||
async def test_completion(mock_completion_client, model_id, service_id, prompt, default_options):
|
||||
"""Test that the service generates content correctly."""
|
||||
mock_completion_client.return_value = {"response": "test_response"}
|
||||
|
||||
ollama = OllamaTextCompletion(ai_model_id=model_id)
|
||||
response = await ollama.get_text_contents(
|
||||
prompt=prompt,
|
||||
settings=OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options),
|
||||
)
|
||||
|
||||
assert response[0].text == "test_response"
|
||||
mock_completion_client.assert_called_once_with(
|
||||
model=model_id,
|
||||
prompt=prompt,
|
||||
options=default_options,
|
||||
stream=False,
|
||||
)
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.generate")
|
||||
async def test_completion_wrong_return_type(
|
||||
mock_completion_client,
|
||||
mock_streaming_text_response,
|
||||
model_id,
|
||||
service_id,
|
||||
chat_history,
|
||||
default_options,
|
||||
):
|
||||
"""Test that the completion service fails when the return type is incorrect."""
|
||||
mock_completion_client.return_value = mock_streaming_text_response # should not be a streaming response
|
||||
|
||||
ollama = OllamaTextCompletion(ai_model_id=model_id)
|
||||
with pytest.raises(ServiceInvalidResponseError):
|
||||
await ollama.get_text_contents(
|
||||
chat_history,
|
||||
OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options),
|
||||
)
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.generate")
|
||||
async def test_streaming_completion(
|
||||
mock_completion_client,
|
||||
mock_streaming_text_response,
|
||||
model_id,
|
||||
service_id,
|
||||
prompt,
|
||||
default_options,
|
||||
):
|
||||
"""Test that the service generates streaming content correctly."""
|
||||
mock_completion_client.return_value = mock_streaming_text_response
|
||||
|
||||
ollama = OllamaTextCompletion(ai_model_id=model_id)
|
||||
response = ollama.get_streaming_text_contents(
|
||||
prompt,
|
||||
OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options),
|
||||
)
|
||||
|
||||
responses = []
|
||||
async for line in response:
|
||||
assert line[0].text == "test_response"
|
||||
responses.append(line)
|
||||
assert len(responses) == 1
|
||||
|
||||
mock_completion_client.assert_called_once_with(
|
||||
model=model_id,
|
||||
prompt=prompt,
|
||||
options=default_options,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.generate")
|
||||
async def test_streaming_completion_wrong_return_type(
|
||||
mock_completion_client,
|
||||
model_id,
|
||||
service_id,
|
||||
chat_history,
|
||||
default_options,
|
||||
):
|
||||
"""Test that the streaming completion service fails when the return type is incorrect."""
|
||||
mock_completion_client.return_value = {"response": "test_response"} # should not be a non-streaming response
|
||||
|
||||
ollama = OllamaTextCompletion(ai_model_id=model_id)
|
||||
with pytest.raises(ServiceInvalidResponseError):
|
||||
async for _ in ollama.get_streaming_text_contents(
|
||||
chat_history,
|
||||
OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options),
|
||||
):
|
||||
pass
|
||||
@@ -0,0 +1,122 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy
|
||||
import pytest
|
||||
from numpy import array
|
||||
|
||||
from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import OllamaEmbeddingPromptExecutionSettings
|
||||
from semantic_kernel.connectors.ai.ollama.services.ollama_text_embedding import OllamaTextEmbedding
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
|
||||
|
||||
|
||||
def test_init_empty_service_id(model_id):
|
||||
"""Test that the service initializes correctly with an empty service id."""
|
||||
ollama = OllamaTextEmbedding(ai_model_id=model_id)
|
||||
assert ollama.service_id == model_id
|
||||
|
||||
|
||||
def test_custom_client(model_id, custom_client):
|
||||
"""Test that the service initializes correctly with a custom client."""
|
||||
ollama = OllamaTextEmbedding(ai_model_id=model_id, client=custom_client)
|
||||
assert ollama.client == custom_client
|
||||
|
||||
|
||||
def test_invalid_ollama_settings():
|
||||
"""Test that the service initializes incorrectly with invalid settings."""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
_ = OllamaTextEmbedding(ai_model_id=123)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OLLAMA_EMBEDDING_MODEL_ID"]], indirect=True)
|
||||
def test_init_empty_model_id(ollama_unit_test_env):
|
||||
"""Test that the service initializes incorrectly with an empty model id."""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
_ = OllamaTextEmbedding(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.__init__", return_value=None) # mock_client
|
||||
@patch("ollama.AsyncClient.embeddings") # mock_embedding_client
|
||||
async def test_custom_host(mock_embedding_client, mock_client, model_id, host, prompt):
|
||||
"""Test that the service initializes and generates embeddings correctly with a custom host."""
|
||||
mock_embedding_client.return_value = {"embedding": [0.1, 0.2, 0.3]}
|
||||
|
||||
ollama = OllamaTextEmbedding(ai_model_id=model_id, host=host)
|
||||
_ = await ollama.generate_embeddings(
|
||||
[prompt],
|
||||
)
|
||||
|
||||
mock_client.assert_called_once_with(host=host)
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.embeddings")
|
||||
async def test_embedding(mock_embedding_client, model_id, prompt):
|
||||
"""Test that the service initializes and generates embeddings correctly."""
|
||||
mock_embedding_client.return_value = {"embedding": [0.1, 0.2, 0.3]}
|
||||
settings = OllamaEmbeddingPromptExecutionSettings()
|
||||
settings.options = {"test_key": "test_value"}
|
||||
|
||||
ollama = OllamaTextEmbedding(ai_model_id=model_id)
|
||||
response = await ollama.generate_embeddings(
|
||||
[prompt],
|
||||
settings=settings,
|
||||
)
|
||||
|
||||
assert response.all() == array([0.1, 0.2, 0.3]).all()
|
||||
mock_embedding_client.assert_called_once_with(model=model_id, prompt=prompt, options=settings.options)
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.embeddings")
|
||||
async def test_embedding_list_input(mock_embedding_client, model_id, prompt):
|
||||
"""Test that the service initializes and generates embeddings correctly with a list of prompts."""
|
||||
mock_embedding_client.return_value = {"embedding": [0.1, 0.2, 0.3]}
|
||||
settings = OllamaEmbeddingPromptExecutionSettings()
|
||||
settings.options = {"test_key": "test_value"}
|
||||
|
||||
ollama = OllamaTextEmbedding(ai_model_id=model_id)
|
||||
responses = await ollama.generate_embeddings(
|
||||
[prompt, prompt],
|
||||
settings=settings,
|
||||
)
|
||||
|
||||
assert len(responses) == 2
|
||||
assert type(responses) is numpy.ndarray
|
||||
assert all(type(response) is numpy.ndarray for response in responses)
|
||||
assert mock_embedding_client.call_count == 2
|
||||
mock_embedding_client.assert_called_with(model=model_id, prompt=prompt, options=settings.options)
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.embeddings")
|
||||
async def test_raw_embedding(mock_embedding_client, model_id, prompt):
|
||||
"""Test that the service initializes and generates embeddings correctly."""
|
||||
mock_embedding_client.return_value = {"embedding": [0.1, 0.2, 0.3]}
|
||||
settings = OllamaEmbeddingPromptExecutionSettings()
|
||||
settings.options = {"test_key": "test_value"}
|
||||
|
||||
ollama = OllamaTextEmbedding(ai_model_id=model_id)
|
||||
response = await ollama.generate_raw_embeddings(
|
||||
[prompt],
|
||||
settings=settings,
|
||||
)
|
||||
|
||||
assert response == [[0.1, 0.2, 0.3]]
|
||||
mock_embedding_client.assert_called_once_with(model=model_id, prompt=prompt, options=settings.options)
|
||||
|
||||
|
||||
@patch("ollama.AsyncClient.embeddings")
|
||||
async def test_raw_embedding_list_input(mock_embedding_client, model_id, prompt):
|
||||
"""Test that the service initializes and generates embeddings correctly with a list of prompts."""
|
||||
mock_embedding_client.return_value = {"embedding": [0.1, 0.2, 0.3]}
|
||||
settings = OllamaEmbeddingPromptExecutionSettings()
|
||||
settings.options = {"test_key": "test_value"}
|
||||
|
||||
ollama = OllamaTextEmbedding(ai_model_id=model_id)
|
||||
responses = await ollama.generate_raw_embeddings(
|
||||
[prompt, prompt],
|
||||
settings=settings,
|
||||
)
|
||||
|
||||
assert responses == [[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]
|
||||
assert mock_embedding_client.call_count == 2
|
||||
mock_embedding_client.assert_called_with(model=model_id, prompt=prompt, options=settings.options)
|
||||
@@ -0,0 +1,240 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.ai.function_call_choice_configuration import FunctionCallChoiceConfiguration
|
||||
from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceType
|
||||
|
||||
# The code under test
|
||||
from semantic_kernel.connectors.ai.ollama.services.utils import (
|
||||
MESSAGE_CONVERTERS,
|
||||
update_settings_from_function_choice_configuration,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
from semantic_kernel.contents.chat_message_content import ChatMessageContent
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
from semantic_kernel.contents.function_result_content import FunctionResultContent
|
||||
from semantic_kernel.contents.image_content import ImageContent
|
||||
from semantic_kernel.contents.utils.author_role import AuthorRole
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_chat_message_content() -> ChatMessageContent:
|
||||
"""Fixture to create a basic ChatMessageContent object with role=USER and simple text content."""
|
||||
return ChatMessageContent(
|
||||
role=AuthorRole.USER,
|
||||
content="Hello, I am a user message.", # The text content
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_system_message_content() -> ChatMessageContent:
|
||||
"""Fixture to create a ChatMessageContent object with role=SYSTEM."""
|
||||
return ChatMessageContent(role=AuthorRole.SYSTEM, content="This is a system message.")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_assistant_message_content() -> ChatMessageContent:
|
||||
"""Fixture to create a ChatMessageContent object with role=ASSISTANT."""
|
||||
return ChatMessageContent(role=AuthorRole.ASSISTANT, content="This is an assistant message.")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tool_message_content() -> ChatMessageContent:
|
||||
"""Fixture to create a ChatMessageContent object with role=TOOL."""
|
||||
return ChatMessageContent(role=AuthorRole.TOOL, content="This is a tool message.")
|
||||
|
||||
|
||||
def test_message_converters_system(mock_system_message_content: ChatMessageContent) -> None:
|
||||
"""Test that passing a system message returns the correct dictionary structure for 'system' role."""
|
||||
# Act
|
||||
converter = MESSAGE_CONVERTERS[AuthorRole.SYSTEM]
|
||||
result = converter(mock_system_message_content)
|
||||
|
||||
# Assert
|
||||
assert result["role"] == "system", "Expected role to be 'system' on the returned message."
|
||||
assert result["content"] == mock_system_message_content.content, (
|
||||
"Expected content to match the system message content."
|
||||
)
|
||||
|
||||
|
||||
def test_message_converters_user_no_images(mock_chat_message_content: ChatMessageContent) -> None:
|
||||
"""Test that passing a user message without images returns correct dictionary structure for 'user' role."""
|
||||
# Act
|
||||
converter = MESSAGE_CONVERTERS[AuthorRole.USER]
|
||||
result = converter(mock_chat_message_content)
|
||||
|
||||
# Assert
|
||||
assert result["role"] == "user", "Expected role to be 'user' on the returned message."
|
||||
assert result["content"] == mock_chat_message_content.content, "Expected content to match the user message content."
|
||||
# Ensure that no 'images' field is added
|
||||
assert "images" not in result, "No images should be present if no ImageContent is added."
|
||||
|
||||
|
||||
def test_message_converters_user_with_images() -> None:
|
||||
"""Test user message with multiple images, verifying the 'images' field is populated."""
|
||||
# Arrange
|
||||
img1 = ImageContent(data="some_base64_data")
|
||||
img2 = ImageContent(data="other_base64_data")
|
||||
content = ChatMessageContent(role=AuthorRole.USER, items=[img1, img2], content="User with images")
|
||||
|
||||
# Act
|
||||
converter = MESSAGE_CONVERTERS[AuthorRole.USER]
|
||||
result = converter(content)
|
||||
|
||||
# Assert
|
||||
assert result["role"] == "user"
|
||||
assert result["content"] == content.content
|
||||
assert "images" in result, "Images field expected when ImageContent is present."
|
||||
assert len(result["images"]) == 2, "Two images should be in the 'images' field."
|
||||
assert result["images"] == [b"some_base64_data", b"other_base64_data"], (
|
||||
"Image data should match the content from ImageContent."
|
||||
)
|
||||
|
||||
|
||||
def test_message_converters_user_with_image_missing_data() -> None:
|
||||
"""Test user message with image content that has missing data, expecting ValueError."""
|
||||
# Arrange
|
||||
bad_image = ImageContent(data="") # empty data for image
|
||||
content = ChatMessageContent(role=AuthorRole.USER, items=[bad_image])
|
||||
|
||||
# Act & Assert
|
||||
converter = MESSAGE_CONVERTERS[AuthorRole.USER]
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
converter(content)
|
||||
|
||||
assert "Image item must contain data encoded as base64." in str(exc_info.value), (
|
||||
"Should raise ValueError for missing base64 data in image."
|
||||
)
|
||||
|
||||
|
||||
def test_message_converters_assistant_basic(mock_assistant_message_content: ChatMessageContent) -> None:
|
||||
"""Test assistant message without images or tool calls."""
|
||||
# Act
|
||||
converter = MESSAGE_CONVERTERS[AuthorRole.ASSISTANT]
|
||||
result = converter(mock_assistant_message_content)
|
||||
|
||||
# Assert
|
||||
assert result["role"] == "assistant", "Assistant role expected."
|
||||
assert result["content"] == mock_assistant_message_content.content
|
||||
assert "images" not in result, "No images included, so should not have an 'images' field."
|
||||
assert "tool_calls" not in result, "No FunctionCallContent, so 'tool_calls' field shouldn't be present."
|
||||
|
||||
|
||||
def test_message_converters_assistant_with_image() -> None:
|
||||
"""Test assistant message containing images. Verify 'images' field is added."""
|
||||
# Arrange
|
||||
img = ImageContent(data="assistant_base64_data")
|
||||
content = ChatMessageContent(role=AuthorRole.ASSISTANT, items=[img], content="Assistant image message")
|
||||
|
||||
# Act
|
||||
converter = MESSAGE_CONVERTERS[AuthorRole.ASSISTANT]
|
||||
result = converter(content)
|
||||
|
||||
# Assert
|
||||
assert result["role"] == "assistant"
|
||||
assert result["content"] == content.content
|
||||
assert "images" in result, "Images should be included for assistant messages with ImageContent."
|
||||
assert result["images"] == [b"assistant_base64_data"], "Expected matching base64 data in images."
|
||||
|
||||
|
||||
def test_message_converters_assistant_with_tool_calls() -> None:
|
||||
"""Test assistant message with FunctionCallContent should populate 'tool_calls'."""
|
||||
# Arrange
|
||||
tool_call_1 = FunctionCallContent(function_name="foo", arguments='{"key": "value"}')
|
||||
tool_call_2 = FunctionCallContent(function_name="bar", arguments='{"another": "123"}')
|
||||
|
||||
content = ChatMessageContent(
|
||||
role=AuthorRole.ASSISTANT, items=[tool_call_1, tool_call_2], content="Assistant with tools"
|
||||
)
|
||||
|
||||
# Act
|
||||
converter = MESSAGE_CONVERTERS[AuthorRole.ASSISTANT]
|
||||
result = converter(content)
|
||||
|
||||
# Assert
|
||||
assert result["role"] == "assistant"
|
||||
assert result["content"] == content.content
|
||||
assert "tool_calls" in result, "tool_calls field should be present for assistant messages with FunctionCallContent."
|
||||
assert len(result["tool_calls"]) == 2, "Expected two tool calls in the result."
|
||||
assert result["tool_calls"][0]["function"]["name"] == "foo", "First tool call function name mismatched."
|
||||
assert result["tool_calls"][0]["function"]["arguments"] == {"key": "value"}, "Expected arguments to be JSON loaded."
|
||||
assert result["tool_calls"][1]["function"]["name"] == "bar", "Second tool call function name mismatched."
|
||||
assert result["tool_calls"][1]["function"]["arguments"] == {"another": "123"}, (
|
||||
"Expected arguments to be JSON loaded."
|
||||
)
|
||||
|
||||
|
||||
def test_message_converters_tool_with_result() -> None:
|
||||
"""Test tool message with a FunctionResultContent, verifying the message content is set."""
|
||||
# Arrange
|
||||
fr_content = FunctionResultContent(id="some_id", result="some result", function_name="test_func")
|
||||
tool_message = ChatMessageContent(role=AuthorRole.TOOL, items=[fr_content])
|
||||
|
||||
# Act
|
||||
converter = MESSAGE_CONVERTERS[AuthorRole.TOOL]
|
||||
result = converter(tool_message)
|
||||
|
||||
# Assert
|
||||
assert result["role"] == "tool", "Expected role to be 'tool' for a tool message."
|
||||
# The code takes the first FunctionResultContent's result as the content
|
||||
assert result["content"] == fr_content.result, "Expected content to match the function result."
|
||||
|
||||
|
||||
def test_message_converters_tool_missing_function_result_content(mock_tool_message_content: ChatMessageContent) -> None:
|
||||
"""Test that if no FunctionResultContent is present, ValueError is raised."""
|
||||
# Arrange
|
||||
mock_tool_message_content.items = [] # no FunctionResultContent in items
|
||||
converter = MESSAGE_CONVERTERS[AuthorRole.TOOL]
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
converter(mock_tool_message_content)
|
||||
assert "Tool message must have a function result content item." in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("choice_type", [FunctionChoiceType.AUTO, FunctionChoiceType.NONE, FunctionChoiceType.REQUIRED])
|
||||
def test_update_settings_from_function_choice_configuration(choice_type: FunctionChoiceType) -> None:
|
||||
"""Test that update_settings_from_function_choice_configuration updates the settings with the correct tools."""
|
||||
# Arrange
|
||||
# We'll create a mock configuration with some available functions.
|
||||
mock_config = FunctionCallChoiceConfiguration()
|
||||
mock_config.available_functions = [MagicMock() for _ in range(2)]
|
||||
|
||||
# We also patch the kernel_function_metadata_to_function_call_format function.
|
||||
# The function returns a dict object describing each function.
|
||||
mock_tool_description = {"type": "function", "function": {"name": "mocked_function"}}
|
||||
|
||||
with patch(
|
||||
"semantic_kernel.connectors.ai.ollama.services.utils.kernel_function_metadata_to_function_call_format",
|
||||
return_value=mock_tool_description,
|
||||
):
|
||||
settings = PromptExecutionSettings()
|
||||
|
||||
# Act
|
||||
update_settings_from_function_choice_configuration(
|
||||
function_choice_configuration=mock_config,
|
||||
settings=settings,
|
||||
type=choice_type,
|
||||
)
|
||||
|
||||
# Assert
|
||||
# After the call, either settings.tools or settings.extension_data["tools"] should be set.
|
||||
# The code tries settings.tools first and if it fails, it sets extension_data["tools"].
|
||||
# We'll check both possibilities.
|
||||
possible_tools = getattr(settings, "tools", None)
|
||||
|
||||
if possible_tools is not None:
|
||||
# If settings.tools exists, ensure it got updated
|
||||
assert len(possible_tools) == 2, "Should have exactly two tools set in the settings.tools attribute."
|
||||
assert possible_tools[0]["function"]["name"] == "mocked_function", (
|
||||
"Expected mocked function name in settings.tools."
|
||||
)
|
||||
else:
|
||||
# Otherwise check for extension_data
|
||||
assert "tools" in settings.extension_data, "Expected 'tools' in extension_data if settings.tools not present."
|
||||
assert len(settings.extension_data["tools"]) == 2, "Should have exactly two tools in extension_data."
|
||||
assert settings.extension_data["tools"][0]["function"]["name"] == "mocked_function", (
|
||||
"Expected mocked function name in extension_data."
|
||||
)
|
||||
@@ -0,0 +1,95 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from semantic_kernel.connectors.ai.ollama import OllamaPromptExecutionSettings
|
||||
from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import (
|
||||
OllamaChatPromptExecutionSettings,
|
||||
OllamaTextPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
|
||||
|
||||
def test_default_ollama_prompt_execution_settings():
|
||||
settings = OllamaPromptExecutionSettings()
|
||||
|
||||
assert settings.format is None
|
||||
assert settings.options is None
|
||||
|
||||
|
||||
def test_custom_ollama_prompt_execution_settings():
|
||||
settings = OllamaPromptExecutionSettings(
|
||||
format="json",
|
||||
options={
|
||||
"key": "value",
|
||||
},
|
||||
)
|
||||
|
||||
assert settings.format == "json"
|
||||
assert settings.options == {"key": "value"}
|
||||
|
||||
|
||||
def test_ollama_prompt_execution_settings_from_default_completion_config():
|
||||
settings = PromptExecutionSettings(service_id="test_service")
|
||||
chat_settings = OllamaChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
|
||||
assert chat_settings.service_id == "test_service"
|
||||
assert chat_settings.format is None
|
||||
assert chat_settings.options is None
|
||||
|
||||
|
||||
def test_ollama_prompt_execution_settings_from_openai_prompt_execution_settings():
|
||||
chat_settings = OllamaChatPromptExecutionSettings(service_id="test_service", options={"temperature": 0.5})
|
||||
new_settings = OllamaPromptExecutionSettings(service_id="test_2", options={"temperature": 0.0})
|
||||
chat_settings.update_from_prompt_execution_settings(new_settings)
|
||||
|
||||
assert chat_settings.service_id == "test_2"
|
||||
assert chat_settings.options["temperature"] == 0.0
|
||||
|
||||
|
||||
def test_ollama_prompt_execution_settings_from_custom_completion_config():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"format": "json",
|
||||
"options": {
|
||||
"key": "value",
|
||||
},
|
||||
},
|
||||
)
|
||||
chat_settings = OllamaChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
|
||||
assert chat_settings.service_id == "test_service"
|
||||
assert chat_settings.format == "json"
|
||||
assert chat_settings.options == {"key": "value"}
|
||||
|
||||
|
||||
def test_ollama_chat_prompt_execution_settings_from_custom_completion_config_with_functions():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"tools": [{"function": {}}],
|
||||
},
|
||||
)
|
||||
chat_settings = OllamaChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
|
||||
assert chat_settings.tools == [{"function": {}}]
|
||||
|
||||
|
||||
def test_create_options():
|
||||
settings = OllamaChatPromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"format": "json",
|
||||
},
|
||||
)
|
||||
options = settings.prepare_settings_dict()
|
||||
|
||||
assert options["format"] == "json"
|
||||
|
||||
|
||||
def test_default_ollama_text_prompt_execution_settings():
|
||||
settings = OllamaTextPromptExecutionSettings()
|
||||
|
||||
assert settings.system is None
|
||||
assert settings.template is None
|
||||
assert settings.context is None
|
||||
assert settings.raw is None
|
||||
@@ -0,0 +1,29 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import json
|
||||
|
||||
|
||||
class MockResponse:
|
||||
def __init__(self, response, status=200):
|
||||
self._response = response
|
||||
self.status = status
|
||||
|
||||
async def text(self):
|
||||
return self._response
|
||||
|
||||
async def json(self):
|
||||
return self._response
|
||||
|
||||
def raise_for_status(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
async def content(self):
|
||||
yield json.dumps(self._response).encode("utf-8")
|
||||
yield json.dumps({"done": True}).encode("utf-8")
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
@@ -0,0 +1,32 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
from pytest import fixture
|
||||
|
||||
|
||||
@fixture()
|
||||
def onnx_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for OnnxGenAISettings."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"ONNX_GEN_AI_CHAT_MODEL_FOLDER": "test",
|
||||
"ONNX_GEN_AI_TEXT_MODEL_FOLDER": "test",
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
gen_ai_config = {"model": {"test": "test"}}
|
||||
|
||||
gen_ai_config_vision = {"model": {"vision": "test"}}
|
||||
@@ -0,0 +1,187 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
import json
|
||||
import os
|
||||
from unittest.mock import MagicMock, mock_open, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.ai.onnx import OnnxGenAIChatCompletion, OnnxGenAIPromptExecutionSettings, ONNXTemplate
|
||||
from semantic_kernel.contents import AuthorRole, ChatHistory, ChatMessageContent, ImageContent
|
||||
from semantic_kernel.exceptions import ServiceInitializationError, ServiceInvalidExecutionSettingsError
|
||||
from semantic_kernel.kernel import Kernel
|
||||
from tests.unit.connectors.ai.onnx.conftest import gen_ai_config, gen_ai_config_vision
|
||||
|
||||
|
||||
@patch("builtins.open", new_callable=mock_open, read_data=json.dumps(gen_ai_config))
|
||||
@patch("onnxruntime_genai.Model")
|
||||
@patch("onnxruntime_genai.Tokenizer")
|
||||
def test_onnx_chat_completion_with_valid_env_variable(gen_ai_config, model, tokenizer, onnx_unit_test_env):
|
||||
service = OnnxGenAIChatCompletion(template=ONNXTemplate.PHI3, env_file_path="test.env")
|
||||
assert not service.enable_multi_modality
|
||||
|
||||
|
||||
@patch("builtins.open", new_callable=mock_open, read_data=json.dumps(gen_ai_config_vision))
|
||||
@patch("onnxruntime_genai.Model")
|
||||
@patch("onnxruntime_genai.Tokenizer")
|
||||
def test_onnx_chat_completion_with_vision_valid_env_variable(
|
||||
gen_ai_vision_config, model, tokenizer, onnx_unit_test_env
|
||||
):
|
||||
service = OnnxGenAIChatCompletion(template=ONNXTemplate.PHI3, env_file_path="test.env")
|
||||
assert service.enable_multi_modality
|
||||
|
||||
|
||||
@patch("builtins.open", new_callable=mock_open, read_data=json.dumps(gen_ai_config))
|
||||
@patch("onnxruntime_genai.Model")
|
||||
@patch("onnxruntime_genai.Tokenizer")
|
||||
def test_onnx_chat_completion_with_valid_parameter(gen_ai_config, model, tokenizer):
|
||||
assert OnnxGenAIChatCompletion(ai_model_path="/valid_path", template=ONNXTemplate.PHI3)
|
||||
|
||||
|
||||
@patch("builtins.open", new_callable=mock_open, read_data=json.dumps(gen_ai_config))
|
||||
@patch("onnxruntime_genai.Model")
|
||||
@patch("onnxruntime_genai.Tokenizer")
|
||||
def test_onnx_chat_completion_with_str_template(gen_ai_config, model, tokenizer):
|
||||
assert OnnxGenAIChatCompletion(ai_model_path="/valid_path", template="phi3")
|
||||
|
||||
|
||||
def test_onnx_chat_completion_with_invalid_model():
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OnnxGenAIChatCompletion(
|
||||
ai_model_path="/invalid_path",
|
||||
template=ONNXTemplate.PHI3,
|
||||
)
|
||||
|
||||
|
||||
@patch("builtins.open", new_callable=mock_open, read_data=json.dumps(gen_ai_config_vision))
|
||||
def test_onnx_chat_completion_with_multimodality_without_prompt_template(gen_ai_config_vision):
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OnnxGenAIChatCompletion()
|
||||
|
||||
|
||||
def test_onnx_chat_completion_with_invalid_env_variable(onnx_unit_test_env):
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OnnxGenAIChatCompletion(
|
||||
template=ONNXTemplate.PHI3,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["ONNX_GEN_AI_CHAT_MODEL_FOLDER"]], indirect=True)
|
||||
def test_onnx_chat_completion_with_missing_ai_path(onnx_unit_test_env):
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OnnxGenAIChatCompletion(template=ONNXTemplate.PHI3, env_file_path="test.env")
|
||||
|
||||
|
||||
@patch("builtins.open", new_callable=mock_open, read_data=json.dumps(gen_ai_config))
|
||||
@patch("onnxruntime_genai.Model")
|
||||
@patch("onnxruntime_genai.Tokenizer")
|
||||
async def test_onnx_chat_completion(gen_ai_config, model, tokenizer):
|
||||
generator_mock = MagicMock()
|
||||
generator_mock.__aiter__.return_value = [["H"], ["e"], ["l"], ["l"], ["o"]]
|
||||
|
||||
chat_completion = OnnxGenAIChatCompletion(template=ONNXTemplate.PHI3, ai_model_path="test")
|
||||
|
||||
history = ChatHistory()
|
||||
history.add_system_message("test")
|
||||
history.add_user_message("test")
|
||||
|
||||
with patch.object(chat_completion, "_generate_next_token_async", return_value=generator_mock):
|
||||
completed_text: ChatMessageContent = await chat_completion.get_chat_message_content(
|
||||
prompt="test", chat_history=history, settings=OnnxGenAIPromptExecutionSettings(), kernel=Kernel()
|
||||
)
|
||||
|
||||
assert str(completed_text) == "Hello"
|
||||
|
||||
|
||||
@patch("builtins.open", new_callable=mock_open, read_data=json.dumps(gen_ai_config))
|
||||
@patch("onnxruntime_genai.Model")
|
||||
@patch("onnxruntime_genai.Tokenizer")
|
||||
async def test_onnx_chat_completion_streaming(gen_ai_config, model, tokenizer):
|
||||
generator_mock = MagicMock()
|
||||
generator_mock.__aiter__.return_value = [["H"], ["e"], ["l"], ["l"], ["o"]]
|
||||
|
||||
chat_completion = OnnxGenAIChatCompletion(template=ONNXTemplate.PHI3, ai_model_path="test")
|
||||
|
||||
history = ChatHistory()
|
||||
history.add_system_message("test")
|
||||
history.add_user_message("test")
|
||||
|
||||
completed_text: str = ""
|
||||
|
||||
with patch.object(chat_completion, "_generate_next_token_async", return_value=generator_mock):
|
||||
async for chunk in chat_completion.get_streaming_chat_message_content(
|
||||
prompt="test", chat_history=history, settings=OnnxGenAIPromptExecutionSettings(), kernel=Kernel()
|
||||
):
|
||||
completed_text += str(chunk)
|
||||
|
||||
assert completed_text == "Hello"
|
||||
|
||||
|
||||
@patch("onnxruntime_genai.Model")
|
||||
def test_onnx_chat_get_image_history(model):
|
||||
builtin_open = open # save the unpatched version
|
||||
|
||||
def patch_open(*args, **kwargs):
|
||||
if "genai_config.json" in str(args[0]):
|
||||
# mocked open for path "genai_config.json"
|
||||
return mock_open(read_data=json.dumps(gen_ai_config_vision))(*args, **kwargs)
|
||||
# unpatched version for every other path
|
||||
return builtin_open(*args, **kwargs)
|
||||
|
||||
with patch("builtins.open", patch_open):
|
||||
chat_completion = OnnxGenAIChatCompletion(
|
||||
template=ONNXTemplate.PHI3,
|
||||
ai_model_path="test",
|
||||
)
|
||||
|
||||
image_content = ImageContent.from_image_path(
|
||||
image_path=os.path.join(os.path.dirname(__file__), "../../../../../", "assets/sample_image.jpg")
|
||||
)
|
||||
|
||||
history = ChatHistory()
|
||||
history.add_system_message("test")
|
||||
history.add_user_message("test")
|
||||
history.add_message(
|
||||
ChatMessageContent(
|
||||
role=AuthorRole.USER,
|
||||
items=[image_content],
|
||||
),
|
||||
)
|
||||
|
||||
last_image = chat_completion._get_images_from_history(history)
|
||||
assert last_image == [image_content]
|
||||
|
||||
|
||||
@patch("onnxruntime_genai.Model")
|
||||
@patch("onnxruntime_genai.Tokenizer")
|
||||
async def test_onnx_chat_get_image_history_with_not_multimodal(model, tokenizer):
|
||||
builtin_open = open # save the unpatched version
|
||||
|
||||
def patch_open(*args, **kwargs):
|
||||
if "genai_config.json" in str(args[0]):
|
||||
# mocked open for path "genai_config.json"
|
||||
return mock_open(read_data=json.dumps(gen_ai_config))(*args, **kwargs)
|
||||
# unpatched version for every other path
|
||||
return builtin_open(*args, **kwargs)
|
||||
|
||||
with patch("builtins.open", patch_open):
|
||||
chat_completion = OnnxGenAIChatCompletion(
|
||||
template=ONNXTemplate.PHI3,
|
||||
ai_model_path="test",
|
||||
)
|
||||
|
||||
image_content = ImageContent.from_image_path(
|
||||
image_path=os.path.join(os.path.dirname(__file__), "../../../../../", "assets/sample_image.jpg")
|
||||
)
|
||||
|
||||
history = ChatHistory()
|
||||
history.add_system_message("test")
|
||||
history.add_user_message("test")
|
||||
history.add_message(
|
||||
ChatMessageContent(
|
||||
role=AuthorRole.USER,
|
||||
items=[image_content],
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
_ = await chat_completion._get_images_from_history(history)
|
||||
@@ -0,0 +1,74 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
import json
|
||||
from unittest.mock import MagicMock, mock_open, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.ai.onnx import OnnxGenAIPromptExecutionSettings, OnnxGenAITextCompletion # noqa: E402
|
||||
from semantic_kernel.contents import TextContent
|
||||
from semantic_kernel.exceptions import ServiceInitializationError
|
||||
from tests.unit.connectors.ai.onnx.conftest import gen_ai_config
|
||||
|
||||
|
||||
@patch("builtins.open", new_callable=mock_open, read_data=json.dumps(gen_ai_config))
|
||||
@patch("onnxruntime_genai.Model")
|
||||
@patch("onnxruntime_genai.Tokenizer")
|
||||
def test_onnx_chat_completion_with_valid_env_variable(gen_ai_config, model, tokenizer, onnx_unit_test_env):
|
||||
assert OnnxGenAITextCompletion(env_file_path="test.env")
|
||||
|
||||
|
||||
@patch("builtins.open", new_callable=mock_open, read_data=json.dumps(gen_ai_config))
|
||||
@patch("onnxruntime_genai.Model")
|
||||
@patch("onnxruntime_genai.Tokenizer")
|
||||
def test_onnx_chat_completion_with_valid_parameter(gen_ai_config, model, tokenizer):
|
||||
assert OnnxGenAITextCompletion(ai_model_path="/valid_path")
|
||||
|
||||
|
||||
def test_onnx_chat_completion_with_invalid_model():
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OnnxGenAITextCompletion(ai_model_path="/invalid_path")
|
||||
|
||||
|
||||
def test_onnx_chat_completion_with_invalid_env_variable(onnx_unit_test_env):
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OnnxGenAITextCompletion()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["ONNX_GEN_AI_TEXT_MODEL_FOLDER"]], indirect=True)
|
||||
def test_onnx_chat_completion_with_missing_ai_path(onnx_unit_test_env):
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OnnxGenAITextCompletion(env_file_path="test.env")
|
||||
|
||||
|
||||
@patch("builtins.open", new_callable=mock_open, read_data=json.dumps(gen_ai_config))
|
||||
@patch("onnxruntime_genai.Model")
|
||||
@patch("onnxruntime_genai.Tokenizer")
|
||||
async def test_onnx_text_completion(gen_ai_config, model, tokenizer):
|
||||
generator_mock = MagicMock()
|
||||
generator_mock.__aiter__.return_value = [["H"], ["e"], ["l"], ["l"], ["o"]]
|
||||
|
||||
text_completion = OnnxGenAITextCompletion(ai_model_path="test")
|
||||
with patch.object(text_completion, "_generate_next_token_async", return_value=generator_mock):
|
||||
completed_text: TextContent = await text_completion.get_text_content(
|
||||
prompt="test", settings=OnnxGenAIPromptExecutionSettings()
|
||||
)
|
||||
|
||||
assert completed_text.text == "Hello"
|
||||
|
||||
|
||||
@patch("builtins.open", new_callable=mock_open, read_data=json.dumps(gen_ai_config))
|
||||
@patch("onnxruntime_genai.Model")
|
||||
@patch("onnxruntime_genai.Tokenizer")
|
||||
async def test_onnx_text_completion_streaming(gen_ai_config, model, tokenizer):
|
||||
generator_mock = MagicMock()
|
||||
generator_mock.__aiter__.return_value = [["H"], ["e"], ["l"], ["l"], ["o"]]
|
||||
|
||||
text_completion = OnnxGenAITextCompletion(ai_model_path="test")
|
||||
completed_text: str = ""
|
||||
with patch.object(text_completion, "_generate_next_token_async", return_value=generator_mock):
|
||||
async for chunk in text_completion.get_streaming_text_content(
|
||||
prompt="test", settings=OnnxGenAIPromptExecutionSettings()
|
||||
):
|
||||
completed_text += chunk.text
|
||||
|
||||
assert completed_text == "Hello"
|
||||
@@ -0,0 +1,132 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from semantic_kernel.connectors.ai.onnx.utils import (
|
||||
gemma_template,
|
||||
llama_template,
|
||||
phi3_template,
|
||||
phi3v_template,
|
||||
phi4_template,
|
||||
phi4mm_template,
|
||||
)
|
||||
from semantic_kernel.contents import AudioContent, AuthorRole, ChatHistory, ImageContent, TextContent
|
||||
|
||||
|
||||
def test_phi3v_template_with_text_and_image():
|
||||
history = ChatHistory(
|
||||
messages=[
|
||||
{"role": AuthorRole.SYSTEM, "content": "System message"},
|
||||
{
|
||||
"role": AuthorRole.USER,
|
||||
"items": [TextContent(text="User text message"), ImageContent(url="http://example.com/image.png")],
|
||||
},
|
||||
{"role": AuthorRole.ASSISTANT, "content": "Assistant message"},
|
||||
]
|
||||
)
|
||||
|
||||
expected_output = (
|
||||
"<|system|>\nSystem message<|end|>\n"
|
||||
"<|user|>\nUser text message<|end|>\n"
|
||||
"<|image_1|>\n"
|
||||
"<|assistant|>\nAssistant message<|end|>\n"
|
||||
"<|assistant|>\n"
|
||||
)
|
||||
|
||||
assert phi3v_template(history) == expected_output
|
||||
|
||||
|
||||
def test_phi4mm_template_with_text_and_image():
|
||||
history = ChatHistory(
|
||||
messages=[
|
||||
{"role": AuthorRole.SYSTEM, "content": "System message"},
|
||||
{
|
||||
"role": AuthorRole.USER,
|
||||
"items": [
|
||||
TextContent(text="User text message"),
|
||||
ImageContent(url="http://example.com/image.png"),
|
||||
AudioContent(url="http://example.com/audio.mp3"),
|
||||
],
|
||||
},
|
||||
{"role": AuthorRole.ASSISTANT, "content": "Assistant message"},
|
||||
]
|
||||
)
|
||||
|
||||
expected_output = (
|
||||
"<|system|>\nSystem message<|end|>\n"
|
||||
"<|user|>\nUser text message<|end|>\n"
|
||||
"<|image_1|>\n"
|
||||
"<|audio_1|>\n"
|
||||
"<|assistant|>\nAssistant message<|end|>\n"
|
||||
"<|assistant|>\n"
|
||||
)
|
||||
|
||||
assert phi4mm_template(history) == expected_output
|
||||
|
||||
|
||||
def test_phi3_template_with_only_text():
|
||||
history = ChatHistory(messages=[{"role": AuthorRole.USER, "items": [TextContent(text="User text message")]}])
|
||||
|
||||
expected_output = "<|user|>\nUser text message<|end|>\n<|assistant|>\n"
|
||||
|
||||
assert phi3_template(history) == expected_output
|
||||
|
||||
|
||||
def test_phi4_template_with_only_text():
|
||||
history = ChatHistory(messages=[{"role": AuthorRole.USER, "items": [TextContent(text="User text message")]}])
|
||||
|
||||
expected_output = "<|user|>\nUser text message<|end|>\n<|assistant|>\n"
|
||||
|
||||
assert phi4_template(history) == expected_output
|
||||
|
||||
|
||||
def test_gemma_template_with_user_and_assistant_messages():
|
||||
history = ChatHistory(
|
||||
messages=[
|
||||
{"role": AuthorRole.USER, "content": "User text message"},
|
||||
{"role": AuthorRole.ASSISTANT, "content": "Assistant message"},
|
||||
]
|
||||
)
|
||||
|
||||
expected_output = (
|
||||
"<bos>"
|
||||
"<start_of_turn>user\nUser text message<end_of_turn>\n"
|
||||
"<start_of_turn>model\nAssistant message<end_of_turn>\n"
|
||||
"<start_of_turn>model\n"
|
||||
)
|
||||
|
||||
assert gemma_template(history) == expected_output
|
||||
|
||||
|
||||
def test_gemma_template_with_only_user_message():
|
||||
history = ChatHistory(messages=[{"role": AuthorRole.USER, "content": "User text message"}])
|
||||
|
||||
expected_output = "<bos><start_of_turn>user\nUser text message<end_of_turn>\n<start_of_turn>model\n"
|
||||
|
||||
assert gemma_template(history) == expected_output
|
||||
|
||||
|
||||
def test_llama_template_with_user_and_assistant_messages():
|
||||
history = ChatHistory(
|
||||
messages=[
|
||||
{"role": AuthorRole.USER, "content": "User text message"},
|
||||
{"role": AuthorRole.ASSISTANT, "content": "Assistant message"},
|
||||
]
|
||||
)
|
||||
|
||||
expected_output = (
|
||||
"<|start_header_id|>user<|end_header_id|>\n\nUser text message<|eot_id|>"
|
||||
"<|start_header_id|>assistant<|end_header_id|>\n\nAssistant message<|eot_id|>"
|
||||
"<|start_header_id|>assistant<|end_header_id|>"
|
||||
)
|
||||
|
||||
assert llama_template(history) == expected_output
|
||||
|
||||
|
||||
def test_llama_template_with_only_user_message():
|
||||
history = ChatHistory(messages=[{"role": AuthorRole.USER, "content": "User text message"}])
|
||||
|
||||
expected_output = (
|
||||
"<|start_header_id|>user<|end_header_id|>\n\nUser text message<|eot_id|>"
|
||||
"<|start_header_id|>assistant<|end_header_id|>"
|
||||
)
|
||||
|
||||
assert llama_template(history) == expected_output
|
||||
@@ -0,0 +1,84 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from semantic_kernel.connectors.ai.onnx.onnx_gen_ai_prompt_execution_settings import (
|
||||
OnnxGenAIPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
|
||||
|
||||
def test_default_onnx_chat_prompt_execution_settings():
|
||||
settings = OnnxGenAIPromptExecutionSettings()
|
||||
assert settings.temperature is None
|
||||
assert settings.top_p is None
|
||||
|
||||
|
||||
def test_custom_onnx_chat_prompt_execution_settings():
|
||||
settings = OnnxGenAIPromptExecutionSettings(
|
||||
temperature=0.5,
|
||||
top_p=0.5,
|
||||
max_length=128,
|
||||
)
|
||||
assert settings.temperature == 0.5
|
||||
assert settings.top_p == 0.5
|
||||
assert settings.max_length == 128
|
||||
|
||||
|
||||
def test_onnx_chat_prompt_execution_settings_from_default_completion_config():
|
||||
settings = PromptExecutionSettings(service_id="test_service")
|
||||
chat_settings = OnnxGenAIPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
assert chat_settings.service_id == "test_service"
|
||||
assert chat_settings.temperature is None
|
||||
assert chat_settings.top_p is None
|
||||
|
||||
|
||||
def test_onnx_chat_prompt_execution_settings_from_onnx_prompt_execution_settings():
|
||||
chat_settings = OnnxGenAIPromptExecutionSettings(service_id="test_service", temperature=1.0)
|
||||
new_settings = OnnxGenAIPromptExecutionSettings(service_id="test_2", temperature=0.0)
|
||||
chat_settings.update_from_prompt_execution_settings(new_settings)
|
||||
assert chat_settings.service_id == "test_2"
|
||||
assert chat_settings.temperature == 0.0
|
||||
|
||||
|
||||
def test_onnx_chat_prompt_execution_settings_from_custom_completion_config():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"max_length": 128,
|
||||
},
|
||||
)
|
||||
chat_settings = OnnxGenAIPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
assert chat_settings.temperature == 0.5
|
||||
assert chat_settings.top_p == 0.5
|
||||
assert chat_settings.max_length == 128
|
||||
|
||||
|
||||
def test_create_options():
|
||||
settings = OnnxGenAIPromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"max_length": 128,
|
||||
},
|
||||
)
|
||||
options = settings.prepare_settings_dict()
|
||||
assert options["temperature"] == 0.5
|
||||
assert options["top_p"] == 0.5
|
||||
assert options["max_length"] == 128
|
||||
|
||||
|
||||
def test_create_options_with_wrong_parameter():
|
||||
with pytest.raises(ValidationError):
|
||||
OnnxGenAIPromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
function_choice_behavior="auto",
|
||||
extension_data={
|
||||
"temperature": 10.0,
|
||||
"top_p": 0.5,
|
||||
"max_length": 128,
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1,99 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from openai import AsyncAzureOpenAI
|
||||
from openai.resources.audio.transcriptions import AsyncTranscriptions
|
||||
from openai.types.audio import Transcription
|
||||
|
||||
from semantic_kernel.connectors.ai.open_ai import AzureAudioToText
|
||||
from semantic_kernel.contents import AudioContent
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidRequestError
|
||||
|
||||
|
||||
def test_azure_audio_to_text_init(azure_openai_unit_test_env) -> None:
|
||||
azure_audio_to_text = AzureAudioToText()
|
||||
|
||||
assert azure_audio_to_text.client is not None
|
||||
assert isinstance(azure_audio_to_text.client, AsyncAzureOpenAI)
|
||||
assert azure_audio_to_text.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_AUDIO_TO_TEXT_DEPLOYMENT_NAME"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_AUDIO_TO_TEXT_DEPLOYMENT_NAME"]], indirect=True)
|
||||
def test_azure_audio_to_text_init_with_empty_deployment_name(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError, match="The Azure OpenAI audio to text deployment name is required."):
|
||||
AzureAudioToText(env_file_path="test.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_API_KEY"]], indirect=True)
|
||||
def test_azure_audio_to_text_init_with_empty_api_key(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureAudioToText(env_file_path="test.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_BASE_URL"]], indirect=True)
|
||||
def test_azure_audio_to_text_init_with_empty_endpoint_and_base_url(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError, match="Please provide an endpoint or a base_url"):
|
||||
AzureAudioToText(env_file_path="test.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("override_env_param_dict", [{"AZURE_OPENAI_ENDPOINT": "http://test.com"}], indirect=True)
|
||||
def test_azure_audio_to_text_init_with_invalid_http_endpoint(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError, match="Invalid settings: "):
|
||||
AzureAudioToText()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"override_env_param_dict",
|
||||
[{"AZURE_OPENAI_BASE_URL": "https://test_audio_to_text_deployment.test-base-url.com"}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_azure_audio_to_text_init_with_from_dict(azure_openai_unit_test_env) -> None:
|
||||
default_headers = {"test_header": "test_value"}
|
||||
|
||||
settings = {
|
||||
"deployment_name": azure_openai_unit_test_env["AZURE_OPENAI_AUDIO_TO_TEXT_DEPLOYMENT_NAME"],
|
||||
"endpoint": azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
|
||||
"api_key": azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
|
||||
"api_version": azure_openai_unit_test_env["AZURE_OPENAI_API_VERSION"],
|
||||
"default_headers": default_headers,
|
||||
}
|
||||
|
||||
azure_audio_to_text = AzureAudioToText.from_dict(settings=settings)
|
||||
|
||||
assert azure_audio_to_text.client is not None
|
||||
assert isinstance(azure_audio_to_text.client, AsyncAzureOpenAI)
|
||||
assert azure_audio_to_text.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_AUDIO_TO_TEXT_DEPLOYMENT_NAME"]
|
||||
assert settings["deployment_name"] in str(azure_audio_to_text.client.base_url)
|
||||
assert azure_audio_to_text.client.api_key == azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"]
|
||||
|
||||
# Assert that the default header we added is present in the client's default headers
|
||||
for key, value in default_headers.items():
|
||||
assert key in azure_audio_to_text.client.default_headers
|
||||
assert azure_audio_to_text.client.default_headers[key] == value
|
||||
|
||||
|
||||
async def test_azure_audio_to_text_get_text_contents(azure_openai_unit_test_env) -> None:
|
||||
audio_content = AudioContent.from_audio_file(
|
||||
os.path.join(os.path.dirname(__file__), "../../../../../", "assets/sample_audio.mp3")
|
||||
)
|
||||
|
||||
with patch.object(AsyncTranscriptions, "create", new_callable=AsyncMock) as mock_transcription_create:
|
||||
mock_transcription_create.return_value = Transcription(text="This is a test audio file.")
|
||||
|
||||
openai_audio_to_text = AzureAudioToText()
|
||||
|
||||
text_contents = await openai_audio_to_text.get_text_contents(audio_content)
|
||||
assert len(text_contents) == 1
|
||||
assert text_contents[0].text == "This is a test audio file."
|
||||
assert text_contents[0].ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_AUDIO_TO_TEXT_DEPLOYMENT_NAME"]
|
||||
|
||||
|
||||
async def test_azure_audio_to_text_get_text_contents_invalid_audio_content(azure_openai_unit_test_env):
|
||||
audio_content = AudioContent()
|
||||
|
||||
openai_audio_to_text = AzureAudioToText()
|
||||
with pytest.raises(ServiceInvalidRequestError, match="Audio content uri must be a string to a local file."):
|
||||
await openai_audio_to_text.get_text_contents(audio_content)
|
||||
@@ -0,0 +1,960 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import json
|
||||
import os
|
||||
from copy import deepcopy
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
from httpx import Request, Response
|
||||
from openai import AsyncAzureOpenAI, AsyncStream
|
||||
from openai.resources.chat.completions import AsyncCompletions as AsyncChatCompletions
|
||||
from openai.types.chat import ChatCompletion, ChatCompletionChunk
|
||||
from openai.types.chat.chat_completion import Choice
|
||||
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
|
||||
from openai.types.chat.chat_completion_chunk import ChoiceDelta as ChunkChoiceDelta
|
||||
from openai.types.chat.chat_completion_message import ChatCompletionMessage
|
||||
|
||||
from semantic_kernel.connectors.ai.chat_completion_client_base import ChatCompletionClientBase
|
||||
from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior
|
||||
from semantic_kernel.connectors.ai.open_ai import AzureChatCompletion
|
||||
from semantic_kernel.connectors.ai.open_ai.exceptions.content_filter_ai_exception import (
|
||||
ContentFilterAIException,
|
||||
ContentFilterResultSeverity,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.azure_chat_prompt_execution_settings import (
|
||||
AzureChatPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.const import USER_AGENT
|
||||
from semantic_kernel.contents.chat_history import ChatHistory
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
from semantic_kernel.contents.function_result_content import FunctionResultContent
|
||||
from semantic_kernel.contents.text_content import TextContent
|
||||
from semantic_kernel.exceptions import ServiceInitializationError, ServiceInvalidExecutionSettingsError
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceResponseException
|
||||
from semantic_kernel.functions.kernel_arguments import KernelArguments
|
||||
from semantic_kernel.functions.kernel_function_decorator import kernel_function
|
||||
from semantic_kernel.kernel import Kernel
|
||||
|
||||
# region Service Setup
|
||||
|
||||
|
||||
def test_init(azure_openai_unit_test_env) -> None:
|
||||
# Test successful initialization
|
||||
azure_chat_completion = AzureChatCompletion(service_id="test_service_id")
|
||||
|
||||
assert azure_chat_completion.client is not None
|
||||
assert isinstance(azure_chat_completion.client, AsyncAzureOpenAI)
|
||||
assert azure_chat_completion.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
|
||||
assert isinstance(azure_chat_completion, ChatCompletionClientBase)
|
||||
assert azure_chat_completion.get_prompt_execution_settings_class() == AzureChatPromptExecutionSettings
|
||||
|
||||
|
||||
def test_init_client(azure_openai_unit_test_env) -> None:
|
||||
# Test successful initialization with client
|
||||
client = MagicMock(spec=AsyncAzureOpenAI)
|
||||
azure_chat_completion = AzureChatCompletion(async_client=client)
|
||||
|
||||
assert azure_chat_completion.client is not None
|
||||
assert isinstance(azure_chat_completion.client, AsyncAzureOpenAI)
|
||||
|
||||
|
||||
def test_init_base_url(azure_openai_unit_test_env) -> None:
|
||||
# Custom header for testing
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
azure_chat_completion = AzureChatCompletion(
|
||||
default_headers=default_headers,
|
||||
)
|
||||
|
||||
assert azure_chat_completion.client is not None
|
||||
assert isinstance(azure_chat_completion.client, AsyncAzureOpenAI)
|
||||
assert azure_chat_completion.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
|
||||
assert isinstance(azure_chat_completion, ChatCompletionClientBase)
|
||||
for key, value in default_headers.items():
|
||||
assert key in azure_chat_completion.client.default_headers
|
||||
assert azure_chat_completion.client.default_headers[key] == value
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_BASE_URL"]], indirect=True)
|
||||
def test_init_endpoint(azure_openai_unit_test_env) -> None:
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
|
||||
assert azure_chat_completion.client is not None
|
||||
assert isinstance(azure_chat_completion.client, AsyncAzureOpenAI)
|
||||
assert azure_chat_completion.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
|
||||
assert isinstance(azure_chat_completion, ChatCompletionClientBase)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]], indirect=True)
|
||||
def test_init_with_empty_deployment_name(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureChatCompletion(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_BASE_URL"]], indirect=True)
|
||||
def test_init_with_empty_endpoint_and_base_url(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureChatCompletion(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("override_env_param_dict", [{"AZURE_OPENAI_ENDPOINT": "http://test.com"}], indirect=True)
|
||||
def test_init_with_invalid_endpoint(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureChatCompletion()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_BASE_URL"]], indirect=True)
|
||||
def test_serialize(azure_openai_unit_test_env) -> None:
|
||||
default_headers = {"X-Test": "test"}
|
||||
|
||||
settings = {
|
||||
"deployment_name": azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
"endpoint": azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
|
||||
"api_key": azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
|
||||
"api_version": azure_openai_unit_test_env["AZURE_OPENAI_API_VERSION"],
|
||||
"default_headers": default_headers,
|
||||
}
|
||||
|
||||
azure_chat_completion = AzureChatCompletion.from_dict(settings)
|
||||
dumped_settings = azure_chat_completion.to_dict()
|
||||
assert dumped_settings["ai_model_id"] == settings["deployment_name"]
|
||||
assert settings["endpoint"] in str(dumped_settings["base_url"])
|
||||
assert settings["deployment_name"] in str(dumped_settings["base_url"])
|
||||
assert settings["api_key"] == dumped_settings["api_key"]
|
||||
assert settings["api_version"] == dumped_settings["api_version"]
|
||||
|
||||
# Assert that the default header we added is present in the dumped_settings default headers
|
||||
for key, value in default_headers.items():
|
||||
assert key in dumped_settings["default_headers"]
|
||||
assert dumped_settings["default_headers"][key] == value
|
||||
|
||||
# Assert that the 'User-agent' header is not present in the dumped_settings default headers
|
||||
assert USER_AGENT not in dumped_settings["default_headers"]
|
||||
|
||||
|
||||
# endregion
|
||||
# region CMC
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_chat_completion_response() -> ChatCompletion:
|
||||
return ChatCompletion(
|
||||
id="test_id",
|
||||
choices=[
|
||||
Choice(index=0, message=ChatCompletionMessage(content="test", role="assistant"), finish_reason="stop")
|
||||
],
|
||||
created=0,
|
||||
model="test",
|
||||
object="chat.completion",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_streaming_chat_completion_response() -> AsyncStream[ChatCompletionChunk]:
|
||||
content = ChatCompletionChunk(
|
||||
id="test_id",
|
||||
choices=[ChunkChoice(index=0, delta=ChunkChoiceDelta(content="test", role="assistant"), finish_reason="stop")],
|
||||
created=0,
|
||||
model="test",
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
stream = MagicMock(spec=AsyncStream)
|
||||
stream.__aiter__.return_value = [content]
|
||||
return stream
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc(
|
||||
mock_create,
|
||||
kernel: Kernel,
|
||||
azure_openai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
chat_history.add_user_message("hello world")
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings(service_id="test_service_id")
|
||||
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history, settings=complete_prompt_execution_settings, kernel=kernel
|
||||
)
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
stream=False,
|
||||
messages=azure_chat_completion._prepare_chat_history_for_request(chat_history),
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc_with_developer_instruction_role_propagates(
|
||||
mock_create,
|
||||
kernel: Kernel,
|
||||
azure_openai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
chat_history.add_user_message("hello world")
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings(service_id="test_service_id")
|
||||
|
||||
azure_chat_completion = AzureChatCompletion(instruction_role="developer")
|
||||
await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history, settings=complete_prompt_execution_settings, kernel=kernel
|
||||
)
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
stream=False,
|
||||
messages=azure_chat_completion._prepare_chat_history_for_request(chat_history),
|
||||
)
|
||||
assert azure_chat_completion.instruction_role == "developer"
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc_with_logit_bias(
|
||||
mock_create,
|
||||
kernel: Kernel,
|
||||
azure_openai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
prompt = "hello world"
|
||||
chat_history.add_user_message(prompt)
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings()
|
||||
|
||||
token_bias = {"1": -100}
|
||||
complete_prompt_execution_settings.logit_bias = token_bias
|
||||
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
|
||||
await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history, settings=complete_prompt_execution_settings, kernel=kernel
|
||||
)
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
messages=azure_chat_completion._prepare_chat_history_for_request(chat_history),
|
||||
stream=False,
|
||||
logit_bias=token_bias,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc_with_stop(
|
||||
mock_create,
|
||||
azure_openai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings()
|
||||
|
||||
stop = ["!"]
|
||||
complete_prompt_execution_settings.stop = stop
|
||||
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
|
||||
await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history, settings=complete_prompt_execution_settings
|
||||
)
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
messages=azure_chat_completion._prepare_chat_history_for_request(chat_history),
|
||||
stream=False,
|
||||
stop=stop,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_azure_on_your_data(
|
||||
mock_create,
|
||||
kernel: Kernel,
|
||||
azure_openai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_chat_completion_response.choices = [
|
||||
Choice(
|
||||
index=0,
|
||||
message=ChatCompletionMessage(
|
||||
content="test",
|
||||
role="assistant",
|
||||
context={
|
||||
"citations": {
|
||||
"content": "test content",
|
||||
"title": "test title",
|
||||
"url": "test url",
|
||||
"filepath": "test filepath",
|
||||
"chunk_id": "test chunk_id",
|
||||
},
|
||||
"intent": "query used",
|
||||
},
|
||||
),
|
||||
finish_reason="stop",
|
||||
)
|
||||
]
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
prompt = "hello world"
|
||||
messages_in = chat_history
|
||||
messages_in.add_user_message(prompt)
|
||||
messages_out = ChatHistory()
|
||||
messages_out.add_user_message(prompt)
|
||||
|
||||
expected_data_settings = {
|
||||
"data_sources": [
|
||||
{
|
||||
"type": "AzureCognitiveSearch",
|
||||
"parameters": {
|
||||
"indexName": "test_index",
|
||||
"endpoint": "https://test-endpoint-search.com",
|
||||
"key": "test_key",
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings(extra_body=expected_data_settings)
|
||||
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
|
||||
content = await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=messages_in, settings=complete_prompt_execution_settings, kernel=kernel
|
||||
)
|
||||
assert isinstance(content[0].items[0], FunctionCallContent)
|
||||
assert isinstance(content[0].items[1], FunctionResultContent)
|
||||
assert isinstance(content[0].items[2], TextContent)
|
||||
assert content[0].items[2].text == "test"
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
messages=azure_chat_completion._prepare_chat_history_for_request(messages_out),
|
||||
stream=False,
|
||||
extra_body=expected_data_settings,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_azure_on_your_data_string(
|
||||
mock_create,
|
||||
kernel: Kernel,
|
||||
azure_openai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_chat_completion_response.choices = [
|
||||
Choice(
|
||||
index=0,
|
||||
message=ChatCompletionMessage(
|
||||
content="test",
|
||||
role="assistant",
|
||||
context=json.dumps({
|
||||
"citations": {
|
||||
"content": "test content",
|
||||
"title": "test title",
|
||||
"url": "test url",
|
||||
"filepath": "test filepath",
|
||||
"chunk_id": "test chunk_id",
|
||||
},
|
||||
"intent": "query used",
|
||||
}),
|
||||
),
|
||||
finish_reason="stop",
|
||||
)
|
||||
]
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
prompt = "hello world"
|
||||
messages_in = chat_history
|
||||
messages_in.add_user_message(prompt)
|
||||
messages_out = ChatHistory()
|
||||
messages_out.add_user_message(prompt)
|
||||
|
||||
expected_data_settings = {
|
||||
"data_sources": [
|
||||
{
|
||||
"type": "AzureCognitiveSearch",
|
||||
"parameters": {
|
||||
"indexName": "test_index",
|
||||
"endpoint": "https://test-endpoint-search.com",
|
||||
"key": "test_key",
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings(extra_body=expected_data_settings)
|
||||
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
|
||||
content = await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=messages_in, settings=complete_prompt_execution_settings, kernel=kernel
|
||||
)
|
||||
assert isinstance(content[0].items[0], FunctionCallContent)
|
||||
assert isinstance(content[0].items[1], FunctionResultContent)
|
||||
assert isinstance(content[0].items[2], TextContent)
|
||||
assert content[0].items[2].text == "test"
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
messages=azure_chat_completion._prepare_chat_history_for_request(messages_out),
|
||||
stream=False,
|
||||
extra_body=expected_data_settings,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_azure_on_your_data_fail(
|
||||
mock_create,
|
||||
kernel: Kernel,
|
||||
azure_openai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_chat_completion_response.choices = [
|
||||
Choice(
|
||||
index=0,
|
||||
message=ChatCompletionMessage(
|
||||
content="test",
|
||||
role="assistant",
|
||||
context="not a dictionary",
|
||||
),
|
||||
finish_reason="stop",
|
||||
)
|
||||
]
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
prompt = "hello world"
|
||||
messages_in = chat_history
|
||||
messages_in.add_user_message(prompt)
|
||||
messages_out = ChatHistory()
|
||||
messages_out.add_user_message(prompt)
|
||||
|
||||
expected_data_settings = {
|
||||
"data_sources": [
|
||||
{
|
||||
"type": "AzureCognitiveSearch",
|
||||
"parameters": {
|
||||
"indexName": "test_index",
|
||||
"endpoint": "https://test-endpoint-search.com",
|
||||
"key": "test_key",
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings(extra_body=expected_data_settings)
|
||||
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
|
||||
content = await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=messages_in, settings=complete_prompt_execution_settings, kernel=kernel
|
||||
)
|
||||
assert isinstance(content[0].items[0], TextContent)
|
||||
assert content[0].items[0].text == "test"
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
messages=azure_chat_completion._prepare_chat_history_for_request(messages_out),
|
||||
stream=False,
|
||||
extra_body=expected_data_settings,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_azure_on_your_data_split_messages(
|
||||
mock_create,
|
||||
kernel: Kernel,
|
||||
azure_openai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_chat_completion_response.choices = [
|
||||
Choice(
|
||||
index=0,
|
||||
message=ChatCompletionMessage(
|
||||
content="test",
|
||||
role="assistant",
|
||||
context={
|
||||
"citations": {
|
||||
"content": "test content",
|
||||
"title": "test title",
|
||||
"url": "test url",
|
||||
"filepath": "test filepath",
|
||||
"chunk_id": "test chunk_id",
|
||||
},
|
||||
"intent": "query used",
|
||||
},
|
||||
),
|
||||
finish_reason="stop",
|
||||
)
|
||||
]
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
prompt = "hello world"
|
||||
messages_in = chat_history
|
||||
messages_in.add_user_message(prompt)
|
||||
messages_out = ChatHistory()
|
||||
messages_out.add_user_message(prompt)
|
||||
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings()
|
||||
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
|
||||
content = await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=messages_in, settings=complete_prompt_execution_settings, kernel=kernel
|
||||
)
|
||||
messages = azure_chat_completion.split_message(content[0])
|
||||
assert len(messages) == 3
|
||||
assert isinstance(messages[0].items[0], FunctionCallContent)
|
||||
assert isinstance(messages[1].items[0], FunctionResultContent)
|
||||
assert isinstance(messages[2].items[0], TextContent)
|
||||
assert messages[2].items[0].text == "test"
|
||||
message = azure_chat_completion.split_message(messages[0])
|
||||
assert message == [messages[0]]
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc_function_calling(
|
||||
mock_create,
|
||||
kernel: Kernel,
|
||||
azure_openai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_chat_completion_response.choices = [
|
||||
Choice(
|
||||
index=0,
|
||||
message=ChatCompletionMessage(
|
||||
content=None,
|
||||
role="assistant",
|
||||
function_call={"name": "test-function", "arguments": '{"key": "value"}'},
|
||||
),
|
||||
finish_reason="stop",
|
||||
)
|
||||
]
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
prompt = "hello world"
|
||||
chat_history.add_user_message(prompt)
|
||||
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
|
||||
functions = [{"name": "test-function", "description": "test-description"}]
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings(
|
||||
function_call="test-function",
|
||||
functions=functions,
|
||||
)
|
||||
|
||||
content = await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=complete_prompt_execution_settings,
|
||||
kernel=kernel,
|
||||
)
|
||||
assert isinstance(content[0].items[0], FunctionCallContent)
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
messages=azure_chat_completion._prepare_chat_history_for_request(chat_history),
|
||||
stream=False,
|
||||
functions=functions,
|
||||
function_call=complete_prompt_execution_settings.function_call,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc_tool_calling(
|
||||
mock_create,
|
||||
kernel: Kernel,
|
||||
azure_openai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_chat_completion_response.choices = [
|
||||
Choice(
|
||||
index=0,
|
||||
message=ChatCompletionMessage(
|
||||
content=None,
|
||||
role="assistant",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "test id",
|
||||
"function": {"name": "test-tool", "arguments": '{"key": "value"}'},
|
||||
"type": "function",
|
||||
}
|
||||
],
|
||||
),
|
||||
finish_reason="stop",
|
||||
)
|
||||
]
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
prompt = "hello world"
|
||||
chat_history.add_user_message(prompt)
|
||||
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings()
|
||||
|
||||
content = await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=complete_prompt_execution_settings,
|
||||
kernel=kernel,
|
||||
)
|
||||
assert isinstance(content[0].items[0], FunctionCallContent)
|
||||
assert content[0].items[0].id == "test id"
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
messages=azure_chat_completion._prepare_chat_history_for_request(chat_history),
|
||||
stream=False,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc_tool_calling_parallel_tool_calls(
|
||||
mock_create,
|
||||
kernel: Kernel,
|
||||
azure_openai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_chat_completion_response.choices = [
|
||||
Choice(
|
||||
index=0,
|
||||
message=ChatCompletionMessage(
|
||||
content=None,
|
||||
role="assistant",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "test id",
|
||||
"function": {"name": "test-tool", "arguments": '{"key": "value"}'},
|
||||
"type": "function",
|
||||
}
|
||||
],
|
||||
),
|
||||
finish_reason="stop",
|
||||
)
|
||||
]
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
prompt = "hello world"
|
||||
chat_history.add_user_message(prompt)
|
||||
|
||||
class MockPlugin:
|
||||
@kernel_function(name="test_tool")
|
||||
def test_tool(self, key: str):
|
||||
return "test"
|
||||
|
||||
kernel.add_plugin(MockPlugin(), plugin_name="test_tool")
|
||||
|
||||
orig_chat_history = deepcopy(chat_history)
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings(
|
||||
service_id="test_service_id", function_choice_behavior=FunctionChoiceBehavior.Auto()
|
||||
)
|
||||
|
||||
with patch(
|
||||
"semantic_kernel.kernel.Kernel.invoke_function_call",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_process_function_call:
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=complete_prompt_execution_settings,
|
||||
kernel=kernel,
|
||||
arguments=KernelArguments(),
|
||||
)
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
stream=False,
|
||||
messages=azure_chat_completion._prepare_chat_history_for_request(orig_chat_history),
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "test_tool-test_tool",
|
||||
"description": "",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"key": {"type": "string"}},
|
||||
"required": ["key"],
|
||||
},
|
||||
},
|
||||
}
|
||||
],
|
||||
tool_choice="auto",
|
||||
)
|
||||
mock_process_function_call.assert_awaited()
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc_tool_calling_parallel_tool_calls_disabled(
|
||||
mock_create,
|
||||
kernel: Kernel,
|
||||
azure_openai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_chat_completion_response: ChatCompletion,
|
||||
) -> None:
|
||||
mock_chat_completion_response.choices = [
|
||||
Choice(
|
||||
index=0,
|
||||
message=ChatCompletionMessage(
|
||||
content=None,
|
||||
role="assistant",
|
||||
tool_calls=[
|
||||
{
|
||||
"id": "test id",
|
||||
"function": {"name": "test-tool", "arguments": '{"key": "value"}'},
|
||||
"type": "function",
|
||||
}
|
||||
],
|
||||
),
|
||||
finish_reason="stop",
|
||||
)
|
||||
]
|
||||
mock_create.return_value = mock_chat_completion_response
|
||||
prompt = "hello world"
|
||||
chat_history.add_user_message(prompt)
|
||||
|
||||
class MockPlugin:
|
||||
@kernel_function(name="test_tool")
|
||||
def test_tool(self, key: str):
|
||||
return "test"
|
||||
|
||||
kernel.add_plugin(MockPlugin(), plugin_name="test_tool")
|
||||
|
||||
orig_chat_history = deepcopy(chat_history)
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings(
|
||||
service_id="test_service_id",
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(),
|
||||
parallel_tool_calls=False,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"semantic_kernel.kernel.Kernel.invoke_function_call",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_process_function_call:
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history,
|
||||
settings=complete_prompt_execution_settings,
|
||||
kernel=kernel,
|
||||
arguments=KernelArguments(),
|
||||
)
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
stream=False,
|
||||
messages=azure_chat_completion._prepare_chat_history_for_request(orig_chat_history),
|
||||
parallel_tool_calls=False,
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "test_tool-test_tool",
|
||||
"description": "",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"key": {"type": "string"}},
|
||||
"required": ["key"],
|
||||
},
|
||||
},
|
||||
}
|
||||
],
|
||||
tool_choice="auto",
|
||||
)
|
||||
mock_process_function_call.assert_awaited()
|
||||
|
||||
|
||||
CONTENT_FILTERED_ERROR_MESSAGE = (
|
||||
"The response was filtered due to the prompt triggering Azure OpenAI's content management policy. Please "
|
||||
"modify your prompt and retry. To learn more about our content filtering policies please read our "
|
||||
"documentation: https://go.microsoft.com/fwlink/?linkid=2198766"
|
||||
)
|
||||
CONTENT_FILTERED_ERROR_FULL_MESSAGE = (
|
||||
"Error code: 400 - {'error': {'message': \"%s\", 'type': null, 'param': 'prompt', 'code': 'content_filter', "
|
||||
"'status': 400, 'innererror': {'code': 'ResponsibleAIPolicyViolation', 'content_filter_result': {'hate': "
|
||||
"{'filtered': True, 'severity': 'high'}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': "
|
||||
"{'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}}}"
|
||||
) % CONTENT_FILTERED_ERROR_MESSAGE
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create")
|
||||
async def test_content_filtering_raises_correct_exception(
|
||||
mock_create, kernel: Kernel, azure_openai_unit_test_env, chat_history: ChatHistory
|
||||
) -> None:
|
||||
prompt = "some prompt that would trigger the content filtering"
|
||||
chat_history.add_user_message(prompt)
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings()
|
||||
|
||||
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
|
||||
mock_create.side_effect = openai.BadRequestError(
|
||||
CONTENT_FILTERED_ERROR_FULL_MESSAGE,
|
||||
response=Response(400, request=Request("POST", test_endpoint)),
|
||||
body={
|
||||
"message": CONTENT_FILTERED_ERROR_MESSAGE,
|
||||
"type": None,
|
||||
"param": "prompt",
|
||||
"code": "content_filter",
|
||||
"status": 400,
|
||||
"innererror": {
|
||||
"code": "ResponsibleAIPolicyViolation",
|
||||
"content_filter_result": {
|
||||
"hate": {"filtered": True, "severity": "high"},
|
||||
"self_harm": {"filtered": False, "severity": "safe"},
|
||||
"sexual": {"filtered": False, "severity": "safe"},
|
||||
"violence": {"filtered": False, "severity": "safe"},
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
|
||||
with pytest.raises(ContentFilterAIException, match="service encountered a content error") as exc_info:
|
||||
await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history, settings=complete_prompt_execution_settings, kernel=kernel
|
||||
)
|
||||
|
||||
content_filter_exc = exc_info.value
|
||||
assert content_filter_exc.param == "prompt"
|
||||
assert content_filter_exc.content_filter_result["hate"].filtered
|
||||
assert content_filter_exc.content_filter_result["hate"].severity == ContentFilterResultSeverity.HIGH
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create")
|
||||
async def test_content_filtering_without_response_code_raises_with_default_code(
|
||||
mock_create, kernel: Kernel, azure_openai_unit_test_env, chat_history: ChatHistory
|
||||
) -> None:
|
||||
prompt = "some prompt that would trigger the content filtering"
|
||||
chat_history.add_user_message(prompt)
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings()
|
||||
|
||||
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
|
||||
mock_create.side_effect = openai.BadRequestError(
|
||||
CONTENT_FILTERED_ERROR_FULL_MESSAGE,
|
||||
response=Response(400, request=Request("POST", test_endpoint)),
|
||||
body={
|
||||
"message": CONTENT_FILTERED_ERROR_MESSAGE,
|
||||
"type": None,
|
||||
"param": "prompt",
|
||||
"code": "content_filter",
|
||||
"status": 400,
|
||||
"innererror": {
|
||||
"content_filter_result": {
|
||||
"hate": {"filtered": True, "severity": "high"},
|
||||
"self_harm": {"filtered": False, "severity": "safe"},
|
||||
"sexual": {"filtered": False, "severity": "safe"},
|
||||
"violence": {"filtered": False, "severity": "safe"},
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
|
||||
with pytest.raises(ContentFilterAIException, match="service encountered a content error"):
|
||||
await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history, settings=complete_prompt_execution_settings, kernel=kernel
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create")
|
||||
async def test_bad_request_non_content_filter(
|
||||
mock_create, kernel: Kernel, azure_openai_unit_test_env, chat_history: ChatHistory
|
||||
) -> None:
|
||||
prompt = "some prompt that would trigger the content filtering"
|
||||
chat_history.add_user_message(prompt)
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings()
|
||||
|
||||
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
|
||||
mock_create.side_effect = openai.BadRequestError(
|
||||
"The request was bad.", response=Response(400, request=Request("POST", test_endpoint)), body={}
|
||||
)
|
||||
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
|
||||
with pytest.raises(ServiceResponseException, match="service failed to complete the prompt"):
|
||||
await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history, settings=complete_prompt_execution_settings, kernel=kernel
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create")
|
||||
async def test_no_kernel_provided_throws_error(
|
||||
mock_create, azure_openai_unit_test_env, chat_history: ChatHistory
|
||||
) -> None:
|
||||
prompt = "some prompt that would trigger the content filtering"
|
||||
chat_history.add_user_message(prompt)
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto()
|
||||
)
|
||||
|
||||
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
|
||||
mock_create.side_effect = openai.BadRequestError(
|
||||
"The request was bad.", response=Response(400, request=Request("POST", test_endpoint)), body={}
|
||||
)
|
||||
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
|
||||
with pytest.raises(
|
||||
ServiceInvalidExecutionSettingsError,
|
||||
match="The kernel is required for function calls.",
|
||||
):
|
||||
await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history, settings=complete_prompt_execution_settings
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create")
|
||||
async def test_auto_invoke_false_no_kernel_provided_throws_error(
|
||||
mock_create, azure_openai_unit_test_env, chat_history: ChatHistory
|
||||
) -> None:
|
||||
prompt = "some prompt that would trigger the content filtering"
|
||||
chat_history.add_user_message(prompt)
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings(
|
||||
function_choice_behavior=FunctionChoiceBehavior.Auto(auto_invoke=False)
|
||||
)
|
||||
|
||||
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
|
||||
mock_create.side_effect = openai.BadRequestError(
|
||||
"The request was bad.", response=Response(400, request=Request("POST", test_endpoint)), body={}
|
||||
)
|
||||
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
|
||||
with pytest.raises(
|
||||
ServiceInvalidExecutionSettingsError,
|
||||
match="The kernel is required for function calls.",
|
||||
):
|
||||
await azure_chat_completion.get_chat_message_contents(
|
||||
chat_history=chat_history, settings=complete_prompt_execution_settings
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_cmc_streaming(
|
||||
mock_create,
|
||||
kernel: Kernel,
|
||||
azure_openai_unit_test_env,
|
||||
chat_history: ChatHistory,
|
||||
mock_streaming_chat_completion_response: AsyncStream[ChatCompletionChunk],
|
||||
) -> None:
|
||||
mock_create.return_value = mock_streaming_chat_completion_response
|
||||
chat_history.add_user_message("hello world")
|
||||
complete_prompt_execution_settings = AzureChatPromptExecutionSettings(service_id="test_service_id")
|
||||
|
||||
azure_chat_completion = AzureChatCompletion()
|
||||
async for msg in azure_chat_completion.get_streaming_chat_message_contents(
|
||||
chat_history=chat_history, settings=complete_prompt_execution_settings, kernel=kernel
|
||||
):
|
||||
assert msg is not None
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
|
||||
stream=True,
|
||||
messages=azure_chat_completion._prepare_chat_history_for_request(chat_history),
|
||||
# NOTE: The `stream_options={"include_usage": True}` is explicitly enforced in
|
||||
# `OpenAIChatCompletionBase._inner_get_streaming_chat_message_contents`.
|
||||
# To ensure consistency, we align the arguments here accordingly.
|
||||
stream_options={"include_usage": True},
|
||||
)
|
||||
@@ -0,0 +1,111 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
from openai import AsyncAzureOpenAI
|
||||
from openai.types import Completion
|
||||
|
||||
from semantic_kernel.connectors.ai.open_ai.services.azure_text_completion import AzureTextCompletion
|
||||
from semantic_kernel.connectors.ai.text_completion_client_base import TextCompletionClientBase
|
||||
from semantic_kernel.exceptions import ServiceInitializationError
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_text_completion_response() -> Mock:
|
||||
mock_response = Mock(spec=Completion)
|
||||
mock_response.id = "test_id"
|
||||
mock_response.created = "time"
|
||||
mock_response.usage = None
|
||||
mock_response.choices = []
|
||||
return mock_response
|
||||
|
||||
|
||||
def test_init(azure_openai_unit_test_env) -> None:
|
||||
# Test successful initialization
|
||||
azure_text_completion = AzureTextCompletion()
|
||||
|
||||
assert azure_text_completion.client is not None
|
||||
assert isinstance(azure_text_completion.client, AsyncAzureOpenAI)
|
||||
assert azure_text_completion.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_TEXT_DEPLOYMENT_NAME"]
|
||||
assert isinstance(azure_text_completion, TextCompletionClientBase)
|
||||
|
||||
|
||||
def test_init_with_custom_header(azure_openai_unit_test_env) -> None:
|
||||
# Custom header for testing
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
# Test successful initialization
|
||||
azure_text_completion = AzureTextCompletion(
|
||||
default_headers=default_headers,
|
||||
)
|
||||
|
||||
assert azure_text_completion.client is not None
|
||||
assert isinstance(azure_text_completion.client, AsyncAzureOpenAI)
|
||||
assert azure_text_completion.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_TEXT_DEPLOYMENT_NAME"]
|
||||
assert isinstance(azure_text_completion, TextCompletionClientBase)
|
||||
for key, value in default_headers.items():
|
||||
assert key in azure_text_completion.client.default_headers
|
||||
assert azure_text_completion.client.default_headers[key] == value
|
||||
|
||||
|
||||
def test_azure_text_embedding_generates_no_token_with_api_key_in_env(azure_openai_unit_test_env) -> None:
|
||||
with (
|
||||
patch(
|
||||
"semantic_kernel.utils.authentication.entra_id_authentication.get_entra_auth_token",
|
||||
) as mock_get_token,
|
||||
):
|
||||
azure_text_completion = AzureTextCompletion()
|
||||
|
||||
assert azure_text_completion.client is not None
|
||||
# API key is provided in env var, so the ad_token should be None
|
||||
assert mock_get_token.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_TEXT_DEPLOYMENT_NAME"]], indirect=True)
|
||||
def test_init_with_empty_deployment_name(monkeypatch, azure_openai_unit_test_env) -> None:
|
||||
monkeypatch.delenv("AZURE_OPENAI_TEXT_DEPLOYMENT_NAME", raising=False)
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureTextCompletion(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_BASE_URL"]], indirect=True)
|
||||
def test_init_with_empty_endpoint_and_base_url(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureTextCompletion(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("override_env_param_dict", [{"AZURE_OPENAI_ENDPOINT": "http://test.com"}], indirect=True)
|
||||
def test_init_with_invalid_endpoint(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureTextCompletion()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_BASE_URL"]], indirect=True)
|
||||
def test_serialize(azure_openai_unit_test_env) -> None:
|
||||
default_headers = {"X-Test": "test"}
|
||||
|
||||
settings = {
|
||||
"deployment_name": azure_openai_unit_test_env["AZURE_OPENAI_TEXT_DEPLOYMENT_NAME"],
|
||||
"endpoint": azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
|
||||
"api_key": azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
|
||||
"api_version": azure_openai_unit_test_env["AZURE_OPENAI_API_VERSION"],
|
||||
"default_headers": default_headers,
|
||||
}
|
||||
|
||||
azure_text_completion = AzureTextCompletion.from_dict(settings)
|
||||
dumped_settings = azure_text_completion.to_dict()
|
||||
assert dumped_settings["ai_model_id"] == settings["deployment_name"]
|
||||
assert settings["endpoint"] in str(dumped_settings["base_url"])
|
||||
assert settings["deployment_name"] in str(dumped_settings["base_url"])
|
||||
assert settings["api_key"] == dumped_settings["api_key"]
|
||||
assert settings["api_version"] == dumped_settings["api_version"]
|
||||
|
||||
# Assert that the default header we added is present in the dumped_settings default headers
|
||||
for key, value in default_headers.items():
|
||||
assert key in dumped_settings["default_headers"]
|
||||
assert dumped_settings["default_headers"][key] == value
|
||||
@@ -0,0 +1,126 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, call, patch
|
||||
|
||||
import pytest
|
||||
from openai import AsyncAzureOpenAI
|
||||
from openai.resources.embeddings import AsyncEmbeddings
|
||||
|
||||
from semantic_kernel.connectors.ai.embedding_generator_base import EmbeddingGeneratorBase
|
||||
from semantic_kernel.connectors.ai.open_ai.services.azure_text_embedding import AzureTextEmbedding
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
|
||||
|
||||
|
||||
def test_azure_text_embedding_init(azure_openai_unit_test_env) -> None:
|
||||
# Test successful initialization
|
||||
azure_text_embedding = AzureTextEmbedding()
|
||||
|
||||
assert azure_text_embedding.client is not None
|
||||
assert isinstance(azure_text_embedding.client, AsyncAzureOpenAI)
|
||||
assert azure_text_embedding.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"]
|
||||
assert isinstance(azure_text_embedding, EmbeddingGeneratorBase)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"]], indirect=True)
|
||||
def test_azure_text_embedding_init_with_empty_deployment_name(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureTextEmbedding(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_BASE_URL"]], indirect=True)
|
||||
def test_azure_text_embedding_init_with_empty_endpoint_and_base_url(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureTextEmbedding(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("override_env_param_dict", [{"AZURE_OPENAI_ENDPOINT": "http://test.com"}], indirect=True)
|
||||
def test_azure_text_embedding_init_with_invalid_endpoint(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureTextEmbedding()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"override_env_param_dict",
|
||||
[{"AZURE_OPENAI_BASE_URL": "https://test_embedding_deployment.test-base-url.com"}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_azure_text_embedding_init_with_from_dict(azure_openai_unit_test_env) -> None:
|
||||
default_headers = {"test_header": "test_value"}
|
||||
|
||||
settings = {
|
||||
"deployment_name": azure_openai_unit_test_env["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"],
|
||||
"endpoint": azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
|
||||
"api_key": azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
|
||||
"api_version": azure_openai_unit_test_env["AZURE_OPENAI_API_VERSION"],
|
||||
"default_headers": default_headers,
|
||||
}
|
||||
|
||||
azure_text_embedding = AzureTextEmbedding.from_dict(settings=settings)
|
||||
|
||||
assert azure_text_embedding.client is not None
|
||||
assert isinstance(azure_text_embedding.client, AsyncAzureOpenAI)
|
||||
assert azure_text_embedding.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"]
|
||||
assert isinstance(azure_text_embedding, EmbeddingGeneratorBase)
|
||||
assert settings["deployment_name"] in str(azure_text_embedding.client.base_url)
|
||||
assert azure_text_embedding.client.api_key == azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"]
|
||||
|
||||
# Assert that the default header we added is present in the client's default headers
|
||||
for key, value in default_headers.items():
|
||||
assert key in azure_text_embedding.client.default_headers
|
||||
assert azure_text_embedding.client.default_headers[key] == value
|
||||
|
||||
|
||||
def test_azure_text_embedding_generates_no_token_with_api_key_in_env(azure_openai_unit_test_env) -> None:
|
||||
with (
|
||||
patch(
|
||||
"semantic_kernel.utils.authentication.entra_id_authentication.get_entra_auth_token",
|
||||
) as mock_get_token,
|
||||
):
|
||||
azure_text_embedding = AzureTextEmbedding()
|
||||
|
||||
assert azure_text_embedding.client is not None
|
||||
# API key is provided in env var, so the ad_token should be None
|
||||
assert mock_get_token.call_count == 0
|
||||
|
||||
|
||||
@patch.object(AsyncEmbeddings, "create", new_callable=AsyncMock)
|
||||
async def test_azure_text_embedding_calls_with_parameters(mock_create, azure_openai_unit_test_env) -> None:
|
||||
texts = ["hello world", "goodbye world"]
|
||||
embedding_dimensions = 1536
|
||||
|
||||
azure_text_embedding = AzureTextEmbedding()
|
||||
|
||||
await azure_text_embedding.generate_embeddings(texts, dimensions=embedding_dimensions)
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
input=texts,
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"],
|
||||
dimensions=embedding_dimensions,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncEmbeddings, "create", new_callable=AsyncMock)
|
||||
async def test_azure_text_embedding_calls_with_batches(mock_create, azure_openai_unit_test_env) -> None:
|
||||
texts = [i for i in range(0, 5)]
|
||||
|
||||
azure_text_embedding = AzureTextEmbedding()
|
||||
|
||||
await azure_text_embedding.generate_embeddings(texts, batch_size=3)
|
||||
|
||||
mock_create.assert_has_awaits(
|
||||
[
|
||||
call(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"],
|
||||
input=texts[0:3],
|
||||
),
|
||||
call(
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"],
|
||||
input=texts[3:5],
|
||||
),
|
||||
],
|
||||
any_order=False,
|
||||
)
|
||||
@@ -0,0 +1,82 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from openai import AsyncAzureOpenAI, _legacy_response
|
||||
from openai.resources.audio.speech import AsyncSpeech
|
||||
|
||||
from semantic_kernel.connectors.ai.open_ai import AzureTextToAudio
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
|
||||
|
||||
|
||||
def test_azure_text_to_audio_init(azure_openai_unit_test_env) -> None:
|
||||
azure_text_to_audio = AzureTextToAudio()
|
||||
|
||||
assert azure_text_to_audio.client is not None
|
||||
assert isinstance(azure_text_to_audio.client, AsyncAzureOpenAI)
|
||||
assert azure_text_to_audio.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_TEXT_TO_AUDIO_DEPLOYMENT_NAME"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_TEXT_TO_AUDIO_DEPLOYMENT_NAME"]], indirect=True)
|
||||
def test_azure_text_to_audio_init_with_empty_deployment_name(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError, match="The Azure OpenAI text to audio deployment name is required."):
|
||||
AzureTextToAudio(env_file_path="test.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_API_KEY"]], indirect=True)
|
||||
def test_azure_text_to_audio_init_with_empty_api_key(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureTextToAudio(env_file_path="test.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_BASE_URL"]], indirect=True)
|
||||
def test_azure_text_to_audio_init_with_empty_endpoint_and_base_url(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError, match="Please provide an endpoint or a base_url"):
|
||||
AzureTextToAudio(env_file_path="test.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("override_env_param_dict", [{"AZURE_OPENAI_ENDPOINT": "http://test.com"}], indirect=True)
|
||||
def test_azure_text_to_audio_init_with_invalid_http_endpoint(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError, match="Invalid settings: "):
|
||||
AzureTextToAudio()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"override_env_param_dict",
|
||||
[{"AZURE_OPENAI_BASE_URL": "https://test_text_to_audio_deployment.test-base-url.com"}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_azure_text_to_audio_init_with_from_dict(azure_openai_unit_test_env) -> None:
|
||||
default_headers = {"test_header": "test_value"}
|
||||
|
||||
settings = {
|
||||
"deployment_name": azure_openai_unit_test_env["AZURE_OPENAI_TEXT_TO_AUDIO_DEPLOYMENT_NAME"],
|
||||
"endpoint": azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
|
||||
"api_key": azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
|
||||
"api_version": azure_openai_unit_test_env["AZURE_OPENAI_API_VERSION"],
|
||||
"default_headers": default_headers,
|
||||
}
|
||||
|
||||
azure_text_to_audio = AzureTextToAudio.from_dict(settings=settings)
|
||||
|
||||
assert azure_text_to_audio.client is not None
|
||||
assert isinstance(azure_text_to_audio.client, AsyncAzureOpenAI)
|
||||
assert azure_text_to_audio.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_TEXT_TO_AUDIO_DEPLOYMENT_NAME"]
|
||||
assert settings["deployment_name"] in str(azure_text_to_audio.client.base_url)
|
||||
assert azure_text_to_audio.client.api_key == azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"]
|
||||
|
||||
# Assert that the default header we added is present in the client's default headers
|
||||
for key, value in default_headers.items():
|
||||
assert key in azure_text_to_audio.client.default_headers
|
||||
assert azure_text_to_audio.client.default_headers[key] == value
|
||||
|
||||
|
||||
@patch.object(AsyncSpeech, "create", return_value=_legacy_response.HttpxBinaryResponseContent(httpx.Response(200)))
|
||||
async def test_azure_text_to_audio_get_audio_contents(mock_speech_create, azure_openai_unit_test_env) -> None:
|
||||
openai_audio_to_text = AzureTextToAudio()
|
||||
|
||||
audio_contents = await openai_audio_to_text.get_audio_contents("Hello World!")
|
||||
assert len(audio_contents) == 1
|
||||
assert audio_contents[0].ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_TEXT_TO_AUDIO_DEPLOYMENT_NAME"]
|
||||
@@ -0,0 +1,99 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from openai import AsyncAzureOpenAI
|
||||
from openai.resources.images import AsyncImages
|
||||
from openai.types.image import Image
|
||||
from openai.types.images_response import ImagesResponse
|
||||
|
||||
from semantic_kernel.connectors.ai.open_ai.services.azure_text_to_image import AzureTextToImage
|
||||
from semantic_kernel.connectors.ai.text_to_image_client_base import TextToImageClientBase
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
|
||||
|
||||
|
||||
def test_azure_text_to_image_init(azure_openai_unit_test_env) -> None:
|
||||
# Test successful initialization
|
||||
azure_text_to_image = AzureTextToImage()
|
||||
|
||||
assert azure_text_to_image.client is not None
|
||||
assert isinstance(azure_text_to_image.client, AsyncAzureOpenAI)
|
||||
assert azure_text_to_image.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_TEXT_TO_IMAGE_DEPLOYMENT_NAME"]
|
||||
assert isinstance(azure_text_to_image, TextToImageClientBase)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_TEXT_TO_IMAGE_DEPLOYMENT_NAME"]], indirect=True)
|
||||
def test_azure_text_to_image_init_with_empty_deployment_name(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureTextToImage(env_file_path="test.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_API_KEY"]], indirect=True)
|
||||
def test_azure_text_to_image_init_with_empty_api_key(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureTextToImage(env_file_path="test.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_BASE_URL"]], indirect=True)
|
||||
def test_azure_text_to_image_init_with_empty_endpoint_and_base_url(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureTextToImage(env_file_path="test.env")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("override_env_param_dict", [{"AZURE_OPENAI_ENDPOINT": "http://test.com"}], indirect=True)
|
||||
def test_azure_text_to_image_init_with_invalid_endpoint(azure_openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
AzureTextToImage()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"override_env_param_dict",
|
||||
[{"AZURE_OPENAI_BASE_URL": "https://test_text_to_image_deployment.test-base-url.com"}],
|
||||
indirect=True,
|
||||
)
|
||||
def test_azure_text_to_image_init_with_from_dict(azure_openai_unit_test_env) -> None:
|
||||
default_headers = {"test_header": "test_value"}
|
||||
|
||||
settings = {
|
||||
"deployment_name": azure_openai_unit_test_env["AZURE_OPENAI_TEXT_TO_IMAGE_DEPLOYMENT_NAME"],
|
||||
"endpoint": azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
|
||||
"api_key": azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
|
||||
"api_version": azure_openai_unit_test_env["AZURE_OPENAI_API_VERSION"],
|
||||
"default_headers": default_headers,
|
||||
}
|
||||
|
||||
azure_text_to_image = AzureTextToImage.from_dict(settings=settings)
|
||||
|
||||
assert azure_text_to_image.client is not None
|
||||
assert isinstance(azure_text_to_image.client, AsyncAzureOpenAI)
|
||||
assert azure_text_to_image.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_TEXT_TO_IMAGE_DEPLOYMENT_NAME"]
|
||||
assert isinstance(azure_text_to_image, TextToImageClientBase)
|
||||
assert settings["deployment_name"] in str(azure_text_to_image.client.base_url)
|
||||
assert azure_text_to_image.client.api_key == azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"]
|
||||
|
||||
# Assert that the default header we added is present in the client's default headers
|
||||
for key, value in default_headers.items():
|
||||
assert key in azure_text_to_image.client.default_headers
|
||||
assert azure_text_to_image.client.default_headers[key] == value
|
||||
|
||||
|
||||
@patch.object(AsyncImages, "generate", new_callable=AsyncMock)
|
||||
async def test_azure_text_to_image_calls_with_parameters(mock_generate, azure_openai_unit_test_env) -> None:
|
||||
mock_response = ImagesResponse(created=1, data=[Image(url="abc")], usage=None)
|
||||
mock_generate.return_value = mock_response
|
||||
|
||||
prompt = "A painting of a vase with flowers"
|
||||
width = 512
|
||||
|
||||
azure_text_to_image = AzureTextToImage(
|
||||
deployment_name=azure_openai_unit_test_env["AZURE_OPENAI_TEXT_TO_IMAGE_DEPLOYMENT_NAME"]
|
||||
)
|
||||
await azure_text_to_image.generate_image(prompt, width=width, height=width)
|
||||
|
||||
mock_generate.assert_awaited_once_with(
|
||||
prompt=prompt,
|
||||
model=azure_openai_unit_test_env["AZURE_OPENAI_TEXT_TO_IMAGE_DEPLOYMENT_NAME"],
|
||||
size=f"{width}x{width}",
|
||||
n=1,
|
||||
)
|
||||
@@ -0,0 +1,87 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from openai import AsyncClient
|
||||
from openai.resources.audio.transcriptions import AsyncTranscriptions
|
||||
from openai.types.audio import Transcription
|
||||
|
||||
from semantic_kernel.connectors.ai.open_ai import OpenAIAudioToTextExecutionSettings
|
||||
from semantic_kernel.connectors.ai.open_ai.services.open_ai_audio_to_text import OpenAIAudioToText
|
||||
from semantic_kernel.contents import AudioContent
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidRequestError
|
||||
|
||||
|
||||
def test_init(openai_unit_test_env):
|
||||
openai_audio_to_text = OpenAIAudioToText()
|
||||
|
||||
assert openai_audio_to_text.client is not None
|
||||
assert isinstance(openai_audio_to_text.client, AsyncClient)
|
||||
assert openai_audio_to_text.ai_model_id == openai_unit_test_env["OPENAI_AUDIO_TO_TEXT_MODEL_ID"]
|
||||
|
||||
|
||||
def test_init_validation_fail() -> None:
|
||||
with pytest.raises(ServiceInitializationError, match="Failed to create OpenAI settings."):
|
||||
OpenAIAudioToText(api_key="34523", ai_model_id={"test": "dict"})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_AUDIO_TO_TEXT_MODEL_ID"]], indirect=True)
|
||||
def test_init_audio_to_text_model_not_provided(openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError, match="The OpenAI audio to text model ID is required."):
|
||||
OpenAIAudioToText(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True)
|
||||
def test_init_with_empty_api_key(openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAIAudioToText(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
def test_init_to_from_dict(openai_unit_test_env):
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
settings = {
|
||||
"ai_model_id": openai_unit_test_env["OPENAI_AUDIO_TO_TEXT_MODEL_ID"],
|
||||
"api_key": openai_unit_test_env["OPENAI_API_KEY"],
|
||||
"default_headers": default_headers,
|
||||
}
|
||||
audio_to_text = OpenAIAudioToText.from_dict(settings)
|
||||
dumped_settings = audio_to_text.to_dict()
|
||||
assert dumped_settings["ai_model_id"] == settings["ai_model_id"]
|
||||
assert dumped_settings["api_key"] == settings["api_key"]
|
||||
|
||||
|
||||
def test_prompt_execution_settings_class(openai_unit_test_env) -> None:
|
||||
openai_audio_to_text = OpenAIAudioToText()
|
||||
assert openai_audio_to_text.get_prompt_execution_settings_class() == OpenAIAudioToTextExecutionSettings
|
||||
|
||||
|
||||
async def test_get_text_contents(openai_unit_test_env):
|
||||
audio_content = AudioContent.from_audio_file(
|
||||
os.path.join(os.path.dirname(__file__), "../../../../../", "assets/sample_audio.mp3")
|
||||
)
|
||||
|
||||
with patch.object(AsyncTranscriptions, "create", new_callable=AsyncMock) as mock_transcription_create:
|
||||
mock_transcription_create.return_value = Transcription(text="This is a test audio file.")
|
||||
|
||||
openai_audio_to_text = OpenAIAudioToText()
|
||||
|
||||
text_contents = await openai_audio_to_text.get_text_contents(audio_content)
|
||||
assert len(text_contents) == 1
|
||||
assert text_contents[0].text == "This is a test audio file."
|
||||
assert text_contents[0].ai_model_id == openai_unit_test_env["OPENAI_AUDIO_TO_TEXT_MODEL_ID"]
|
||||
|
||||
|
||||
async def test_get_text_contents_invalid_audio_content(openai_unit_test_env):
|
||||
audio_content = AudioContent()
|
||||
|
||||
openai_audio_to_text = OpenAIAudioToText()
|
||||
with pytest.raises(ServiceInvalidRequestError, match="Audio content uri must be a string to a local file."):
|
||||
await openai_audio_to_text.get_text_contents(audio_content)
|
||||
@@ -0,0 +1,105 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.ai.chat_completion_client_base import ChatCompletionClientBase
|
||||
from semantic_kernel.connectors.ai.open_ai.services.open_ai_chat_completion import OpenAIChatCompletion
|
||||
from semantic_kernel.const import USER_AGENT
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
|
||||
|
||||
|
||||
def test_init(openai_unit_test_env) -> None:
|
||||
# Test successful initialization
|
||||
open_ai_chat_completion = OpenAIChatCompletion()
|
||||
|
||||
assert open_ai_chat_completion.ai_model_id == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"]
|
||||
assert isinstance(open_ai_chat_completion, ChatCompletionClientBase)
|
||||
|
||||
|
||||
def test_init_validation_fail() -> None:
|
||||
# Test successful initialization
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAIChatCompletion(api_key="34523", ai_model_id={"test": "dict"})
|
||||
|
||||
|
||||
def test_init_ai_model_id_constructor(openai_unit_test_env) -> None:
|
||||
# Test successful initialization
|
||||
ai_model_id = "test_model_id"
|
||||
open_ai_chat_completion = OpenAIChatCompletion(ai_model_id=ai_model_id)
|
||||
|
||||
assert open_ai_chat_completion.ai_model_id == ai_model_id
|
||||
assert isinstance(open_ai_chat_completion, ChatCompletionClientBase)
|
||||
|
||||
|
||||
def test_init_with_default_header(openai_unit_test_env) -> None:
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
# Test successful initialization
|
||||
open_ai_chat_completion = OpenAIChatCompletion(
|
||||
default_headers=default_headers,
|
||||
)
|
||||
|
||||
assert open_ai_chat_completion.ai_model_id == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"]
|
||||
assert isinstance(open_ai_chat_completion, ChatCompletionClientBase)
|
||||
|
||||
# Assert that the default header we added is present in the client's default headers
|
||||
for key, value in default_headers.items():
|
||||
assert key in open_ai_chat_completion.client.default_headers
|
||||
assert open_ai_chat_completion.client.default_headers[key] == value
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_CHAT_MODEL_ID"]], indirect=True)
|
||||
def test_init_with_empty_model_id(openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAIChatCompletion(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True)
|
||||
def test_init_with_empty_api_key(openai_unit_test_env) -> None:
|
||||
ai_model_id = "test_model_id"
|
||||
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAIChatCompletion(
|
||||
ai_model_id=ai_model_id,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
def test_serialize(openai_unit_test_env) -> None:
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
settings = {
|
||||
"ai_model_id": openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
|
||||
"api_key": openai_unit_test_env["OPENAI_API_KEY"],
|
||||
"default_headers": default_headers,
|
||||
}
|
||||
|
||||
open_ai_chat_completion = OpenAIChatCompletion.from_dict(settings)
|
||||
dumped_settings = open_ai_chat_completion.to_dict()
|
||||
assert dumped_settings["ai_model_id"] == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"]
|
||||
assert dumped_settings["api_key"] == openai_unit_test_env["OPENAI_API_KEY"]
|
||||
# Assert that the default header we added is present in the dumped_settings default headers
|
||||
for key, value in default_headers.items():
|
||||
assert key in dumped_settings["default_headers"]
|
||||
assert dumped_settings["default_headers"][key] == value
|
||||
# Assert that the 'User-agent' header is not present in the dumped_settings default headers
|
||||
assert USER_AGENT not in dumped_settings["default_headers"]
|
||||
|
||||
|
||||
def test_serialize_with_org_id(openai_unit_test_env) -> None:
|
||||
settings = {
|
||||
"ai_model_id": openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
|
||||
"api_key": openai_unit_test_env["OPENAI_API_KEY"],
|
||||
"org_id": openai_unit_test_env["OPENAI_ORG_ID"],
|
||||
}
|
||||
|
||||
open_ai_chat_completion = OpenAIChatCompletion.from_dict(settings)
|
||||
dumped_settings = open_ai_chat_completion.to_dict()
|
||||
assert dumped_settings["ai_model_id"] == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"]
|
||||
assert dumped_settings["api_key"] == openai_unit_test_env["OPENAI_API_KEY"]
|
||||
assert dumped_settings["org_id"] == openai_unit_test_env["OPENAI_ORG_ID"]
|
||||
# Assert that the 'User-agent' header is not present in the dumped_settings default headers
|
||||
assert USER_AGENT not in dumped_settings["default_headers"]
|
||||
+1111
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,324 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from openai import AsyncStream
|
||||
from openai.resources import AsyncCompletions
|
||||
from openai.types import Completion as TextCompletion
|
||||
from openai.types import CompletionChoice as TextCompletionChoice
|
||||
|
||||
from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.open_ai_prompt_execution_settings import (
|
||||
OpenAITextPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.open_ai.services.open_ai_text_completion import OpenAITextCompletion
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
from semantic_kernel.connectors.ai.text_completion_client_base import TextCompletionClientBase
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
|
||||
|
||||
|
||||
def test_init(openai_unit_test_env) -> None:
|
||||
# Test successful initialization
|
||||
open_ai_text_completion = OpenAITextCompletion()
|
||||
|
||||
assert open_ai_text_completion.ai_model_id == openai_unit_test_env["OPENAI_TEXT_MODEL_ID"]
|
||||
assert isinstance(open_ai_text_completion, TextCompletionClientBase)
|
||||
|
||||
|
||||
def test_init_with_ai_model_id(openai_unit_test_env) -> None:
|
||||
# Test successful initialization
|
||||
ai_model_id = "test_model_id"
|
||||
open_ai_text_completion = OpenAITextCompletion(ai_model_id=ai_model_id)
|
||||
|
||||
assert open_ai_text_completion.ai_model_id == ai_model_id
|
||||
assert isinstance(open_ai_text_completion, TextCompletionClientBase)
|
||||
|
||||
|
||||
def test_init_with_default_header(openai_unit_test_env) -> None:
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
# Test successful initialization
|
||||
open_ai_text_completion = OpenAITextCompletion(
|
||||
default_headers=default_headers,
|
||||
)
|
||||
|
||||
assert open_ai_text_completion.ai_model_id == openai_unit_test_env["OPENAI_TEXT_MODEL_ID"]
|
||||
assert isinstance(open_ai_text_completion, TextCompletionClientBase)
|
||||
for key, value in default_headers.items():
|
||||
assert key in open_ai_text_completion.client.default_headers
|
||||
assert open_ai_text_completion.client.default_headers[key] == value
|
||||
|
||||
|
||||
def test_init_validation_fail() -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAITextCompletion(api_key="34523", ai_model_id={"test": "dict"})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True)
|
||||
def test_init_with_empty_api_key(openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAITextCompletion(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_TEXT_MODEL_ID"]], indirect=True)
|
||||
def test_init_with_empty_model(openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAITextCompletion(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
def test_serialize(openai_unit_test_env) -> None:
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
settings = {
|
||||
"ai_model_id": openai_unit_test_env["OPENAI_TEXT_MODEL_ID"],
|
||||
"api_key": openai_unit_test_env["OPENAI_API_KEY"],
|
||||
"default_headers": default_headers,
|
||||
}
|
||||
|
||||
open_ai_text_completion = OpenAITextCompletion.from_dict(settings)
|
||||
dumped_settings = open_ai_text_completion.to_dict()
|
||||
assert dumped_settings["ai_model_id"] == openai_unit_test_env["OPENAI_TEXT_MODEL_ID"]
|
||||
assert dumped_settings["api_key"] == openai_unit_test_env["OPENAI_API_KEY"]
|
||||
# Assert that the default header we added is present in the dumped_settings default headers
|
||||
for key, value in default_headers.items():
|
||||
assert key in dumped_settings["default_headers"]
|
||||
assert dumped_settings["default_headers"][key] == value
|
||||
|
||||
|
||||
def test_serialize_def_headers_string(openai_unit_test_env) -> None:
|
||||
default_headers = '{"X-Unit-Test": "test-guid"}'
|
||||
|
||||
settings = {
|
||||
"ai_model_id": openai_unit_test_env["OPENAI_TEXT_MODEL_ID"],
|
||||
"api_key": openai_unit_test_env["OPENAI_API_KEY"],
|
||||
"default_headers": default_headers,
|
||||
}
|
||||
|
||||
open_ai_text_completion = OpenAITextCompletion.from_dict(settings)
|
||||
dumped_settings = open_ai_text_completion.to_dict()
|
||||
assert dumped_settings["ai_model_id"] == openai_unit_test_env["OPENAI_TEXT_MODEL_ID"]
|
||||
assert dumped_settings["api_key"] == openai_unit_test_env["OPENAI_API_KEY"]
|
||||
# Assert that the default header we added is present in the dumped_settings default headers
|
||||
for key, value in json.loads(default_headers).items():
|
||||
assert key in dumped_settings["default_headers"]
|
||||
assert dumped_settings["default_headers"][key] == value
|
||||
|
||||
|
||||
def test_serialize_with_org_id(openai_unit_test_env) -> None:
|
||||
settings = {
|
||||
"ai_model_id": openai_unit_test_env["OPENAI_TEXT_MODEL_ID"],
|
||||
"api_key": openai_unit_test_env["OPENAI_API_KEY"],
|
||||
"org_id": openai_unit_test_env["OPENAI_ORG_ID"],
|
||||
}
|
||||
|
||||
open_ai_text_completion = OpenAITextCompletion.from_dict(settings)
|
||||
dumped_settings = open_ai_text_completion.to_dict()
|
||||
assert dumped_settings["ai_model_id"] == openai_unit_test_env["OPENAI_TEXT_MODEL_ID"]
|
||||
assert dumped_settings["api_key"] == openai_unit_test_env["OPENAI_API_KEY"]
|
||||
assert dumped_settings["org_id"] == openai_unit_test_env["OPENAI_ORG_ID"]
|
||||
|
||||
|
||||
# region Get Text Contents
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def completion_response() -> TextCompletion:
|
||||
return TextCompletion(
|
||||
id="test",
|
||||
choices=[TextCompletionChoice(text="test", index=0, finish_reason="stop")],
|
||||
created=0,
|
||||
model="test",
|
||||
object="text_completion",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def streaming_completion_response() -> AsyncStream[TextCompletion]:
|
||||
content = TextCompletion(
|
||||
id="test",
|
||||
choices=[TextCompletionChoice(text="test", index=0, finish_reason="stop")],
|
||||
created=0,
|
||||
model="test",
|
||||
object="text_completion",
|
||||
)
|
||||
stream = MagicMock(spec=AsyncStream)
|
||||
stream.__aiter__.return_value = [content]
|
||||
return stream
|
||||
|
||||
|
||||
@patch.object(AsyncCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_tc(
|
||||
mock_create,
|
||||
openai_unit_test_env,
|
||||
completion_response,
|
||||
) -> None:
|
||||
mock_create.return_value = completion_response
|
||||
complete_prompt_execution_settings = OpenAITextPromptExecutionSettings(service_id="test_service_id")
|
||||
|
||||
openai_text_completion = OpenAITextCompletion()
|
||||
await openai_text_completion.get_text_contents(prompt="test", settings=complete_prompt_execution_settings)
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=openai_unit_test_env["OPENAI_TEXT_MODEL_ID"],
|
||||
stream=False,
|
||||
prompt="test",
|
||||
echo=False,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_tc_singular(
|
||||
mock_create,
|
||||
openai_unit_test_env,
|
||||
completion_response,
|
||||
) -> None:
|
||||
mock_create.return_value = completion_response
|
||||
complete_prompt_execution_settings = OpenAITextPromptExecutionSettings(service_id="test_service_id")
|
||||
|
||||
openai_text_completion = OpenAITextCompletion()
|
||||
await openai_text_completion.get_text_content(prompt="test", settings=complete_prompt_execution_settings)
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=openai_unit_test_env["OPENAI_TEXT_MODEL_ID"],
|
||||
stream=False,
|
||||
prompt="test",
|
||||
echo=False,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_tc_prompt_execution_settings(
|
||||
mock_create,
|
||||
openai_unit_test_env,
|
||||
completion_response,
|
||||
) -> None:
|
||||
mock_create.return_value = completion_response
|
||||
complete_prompt_execution_settings = PromptExecutionSettings(service_id="test_service_id")
|
||||
|
||||
openai_text_completion = OpenAITextCompletion()
|
||||
await openai_text_completion.get_text_contents(prompt="test", settings=complete_prompt_execution_settings)
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=openai_unit_test_env["OPENAI_TEXT_MODEL_ID"],
|
||||
stream=False,
|
||||
prompt="test",
|
||||
echo=False,
|
||||
)
|
||||
|
||||
|
||||
# region Streaming
|
||||
|
||||
|
||||
@patch.object(AsyncCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_stc(
|
||||
mock_create,
|
||||
openai_unit_test_env,
|
||||
streaming_completion_response,
|
||||
) -> None:
|
||||
mock_create.return_value = streaming_completion_response
|
||||
complete_prompt_execution_settings = OpenAITextPromptExecutionSettings(service_id="test_service_id")
|
||||
|
||||
openai_text_completion = OpenAITextCompletion()
|
||||
[
|
||||
text
|
||||
async for text in openai_text_completion.get_streaming_text_contents(
|
||||
prompt="test", settings=complete_prompt_execution_settings
|
||||
)
|
||||
]
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=openai_unit_test_env["OPENAI_TEXT_MODEL_ID"],
|
||||
stream=True,
|
||||
prompt="test",
|
||||
echo=False,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_stc_singular(
|
||||
mock_create,
|
||||
openai_unit_test_env,
|
||||
streaming_completion_response,
|
||||
) -> None:
|
||||
mock_create.return_value = streaming_completion_response
|
||||
complete_prompt_execution_settings = OpenAITextPromptExecutionSettings(service_id="test_service_id")
|
||||
|
||||
openai_text_completion = OpenAITextCompletion()
|
||||
[
|
||||
text
|
||||
async for text in openai_text_completion.get_streaming_text_content(
|
||||
prompt="test", settings=complete_prompt_execution_settings
|
||||
)
|
||||
]
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=openai_unit_test_env["OPENAI_TEXT_MODEL_ID"],
|
||||
stream=True,
|
||||
prompt="test",
|
||||
echo=False,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_stc_prompt_execution_settings(
|
||||
mock_create,
|
||||
openai_unit_test_env,
|
||||
streaming_completion_response,
|
||||
) -> None:
|
||||
mock_create.return_value = streaming_completion_response
|
||||
complete_prompt_execution_settings = PromptExecutionSettings(service_id="test_service_id")
|
||||
|
||||
openai_text_completion = OpenAITextCompletion()
|
||||
[
|
||||
text
|
||||
async for text in openai_text_completion.get_streaming_text_contents(
|
||||
prompt="test", settings=complete_prompt_execution_settings
|
||||
)
|
||||
]
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=openai_unit_test_env["OPENAI_TEXT_MODEL_ID"],
|
||||
stream=True,
|
||||
prompt="test",
|
||||
echo=False,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncCompletions, "create", new_callable=AsyncMock)
|
||||
async def test_stc_empty_choices(
|
||||
mock_create,
|
||||
openai_unit_test_env,
|
||||
) -> None:
|
||||
content1 = TextCompletion(
|
||||
id="test",
|
||||
choices=[],
|
||||
created=0,
|
||||
model="test",
|
||||
object="text_completion",
|
||||
)
|
||||
content2 = TextCompletion(
|
||||
id="test",
|
||||
choices=[TextCompletionChoice(text="test", index=0, finish_reason="stop")],
|
||||
created=0,
|
||||
model="test",
|
||||
object="text_completion",
|
||||
)
|
||||
stream = MagicMock(spec=AsyncStream)
|
||||
stream.__aiter__.return_value = [content1, content2]
|
||||
mock_create.return_value = stream
|
||||
complete_prompt_execution_settings = OpenAITextPromptExecutionSettings(service_id="test_service_id")
|
||||
|
||||
openai_text_completion = OpenAITextCompletion()
|
||||
results = [
|
||||
text
|
||||
async for text in openai_text_completion.get_streaming_text_contents(
|
||||
prompt="test", settings=complete_prompt_execution_settings
|
||||
)
|
||||
]
|
||||
assert len(results) == 1
|
||||
mock_create.assert_awaited_once_with(
|
||||
model=openai_unit_test_env["OPENAI_TEXT_MODEL_ID"],
|
||||
stream=True,
|
||||
prompt="test",
|
||||
echo=False,
|
||||
)
|
||||
@@ -0,0 +1,126 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from openai import AsyncClient
|
||||
from openai.resources.embeddings import AsyncEmbeddings
|
||||
|
||||
from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.open_ai_prompt_execution_settings import (
|
||||
OpenAIEmbeddingPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.open_ai.services.open_ai_text_embedding import OpenAITextEmbedding
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceResponseException
|
||||
|
||||
|
||||
def test_init(openai_unit_test_env):
|
||||
openai_text_embedding = OpenAITextEmbedding()
|
||||
|
||||
assert openai_text_embedding.client is not None
|
||||
assert isinstance(openai_text_embedding.client, AsyncClient)
|
||||
assert openai_text_embedding.ai_model_id == openai_unit_test_env["OPENAI_EMBEDDING_MODEL_ID"]
|
||||
|
||||
assert openai_text_embedding.get_prompt_execution_settings_class() == OpenAIEmbeddingPromptExecutionSettings
|
||||
|
||||
|
||||
def test_init_validation_fail() -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAITextEmbedding(api_key="34523", ai_model_id={"test": "dict"})
|
||||
|
||||
|
||||
def test_init_to_from_dict(openai_unit_test_env):
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
settings = {
|
||||
"ai_model_id": openai_unit_test_env["OPENAI_EMBEDDING_MODEL_ID"],
|
||||
"api_key": openai_unit_test_env["OPENAI_API_KEY"],
|
||||
"default_headers": default_headers,
|
||||
}
|
||||
text_embedding = OpenAITextEmbedding.from_dict(settings)
|
||||
dumped_settings = text_embedding.to_dict()
|
||||
assert dumped_settings["ai_model_id"] == settings["ai_model_id"]
|
||||
assert dumped_settings["api_key"] == settings["api_key"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True)
|
||||
def test_init_with_empty_api_key(openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAITextEmbedding(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_EMBEDDING_MODEL_ID"]], indirect=True)
|
||||
def test_init_with_no_model_id(openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAITextEmbedding(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncEmbeddings, "create", new_callable=AsyncMock)
|
||||
async def test_embedding_calls_with_parameters(mock_create, openai_unit_test_env) -> None:
|
||||
ai_model_id = "test_model_id"
|
||||
texts = ["hello world", "goodbye world"]
|
||||
embedding_dimensions = 1536
|
||||
|
||||
openai_text_embedding = OpenAITextEmbedding(
|
||||
ai_model_id=ai_model_id,
|
||||
)
|
||||
|
||||
await openai_text_embedding.generate_embeddings(texts, dimensions=embedding_dimensions)
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
input=texts,
|
||||
model=ai_model_id,
|
||||
dimensions=embedding_dimensions,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncEmbeddings, "create", new_callable=AsyncMock)
|
||||
async def test_embedding_calls_with_settings(mock_create, openai_unit_test_env) -> None:
|
||||
ai_model_id = "test_model_id"
|
||||
texts = ["hello world", "goodbye world"]
|
||||
settings = OpenAIEmbeddingPromptExecutionSettings(service_id="default", dimensions=1536)
|
||||
openai_text_embedding = OpenAITextEmbedding(service_id="default", ai_model_id=ai_model_id)
|
||||
|
||||
await openai_text_embedding.generate_embeddings(texts, settings=settings, timeout=10)
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
input=texts,
|
||||
model=ai_model_id,
|
||||
dimensions=1536,
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(AsyncEmbeddings, "create", new_callable=AsyncMock, side_effect=Exception)
|
||||
async def test_embedding_fail(mock_create, openai_unit_test_env) -> None:
|
||||
ai_model_id = "test_model_id"
|
||||
texts = ["hello world", "goodbye world"]
|
||||
embedding_dimensions = 1536
|
||||
|
||||
openai_text_embedding = OpenAITextEmbedding(
|
||||
ai_model_id=ai_model_id,
|
||||
)
|
||||
with pytest.raises(ServiceResponseException):
|
||||
await openai_text_embedding.generate_embeddings(texts, dimensions=embedding_dimensions)
|
||||
|
||||
|
||||
@patch.object(AsyncEmbeddings, "create", new_callable=AsyncMock)
|
||||
async def test_embedding_pes(mock_create, openai_unit_test_env) -> None:
|
||||
ai_model_id = "test_model_id"
|
||||
texts = ["hello world", "goodbye world"]
|
||||
embedding_dimensions = 1536
|
||||
pes = PromptExecutionSettings(ai_model_id=ai_model_id, dimensions=embedding_dimensions)
|
||||
|
||||
openai_text_embedding = OpenAITextEmbedding(ai_model_id=ai_model_id)
|
||||
|
||||
await openai_text_embedding.generate_raw_embeddings(texts, pes)
|
||||
|
||||
mock_create.assert_awaited_once_with(
|
||||
input=texts,
|
||||
model=ai_model_id,
|
||||
dimensions=embedding_dimensions,
|
||||
)
|
||||
@@ -0,0 +1,69 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from openai import AsyncClient, _legacy_response
|
||||
from openai.resources.audio.speech import AsyncSpeech
|
||||
|
||||
from semantic_kernel.connectors.ai.open_ai import OpenAITextToAudio, OpenAITextToAudioExecutionSettings
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
|
||||
|
||||
|
||||
def test_init(openai_unit_test_env):
|
||||
openai_text_to_audio = OpenAITextToAudio()
|
||||
|
||||
assert openai_text_to_audio.client is not None
|
||||
assert isinstance(openai_text_to_audio.client, AsyncClient)
|
||||
assert openai_text_to_audio.ai_model_id == openai_unit_test_env["OPENAI_TEXT_TO_AUDIO_MODEL_ID"]
|
||||
|
||||
|
||||
def test_init_validation_fail() -> None:
|
||||
with pytest.raises(ServiceInitializationError, match="Failed to create OpenAI settings."):
|
||||
OpenAITextToAudio(api_key="34523", ai_model_id={"test": "dict"})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_TEXT_TO_AUDIO_MODEL_ID"]], indirect=True)
|
||||
def test_init_text_to_audio_model_not_provided(openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError, match="The OpenAI text to audio model ID is required."):
|
||||
OpenAITextToAudio(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True)
|
||||
def test_init_with_empty_api_key(openai_unit_test_env) -> None:
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAITextToAudio(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
def test_init_to_from_dict(openai_unit_test_env):
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
settings = {
|
||||
"ai_model_id": openai_unit_test_env["OPENAI_TEXT_TO_AUDIO_MODEL_ID"],
|
||||
"api_key": openai_unit_test_env["OPENAI_API_KEY"],
|
||||
"default_headers": default_headers,
|
||||
}
|
||||
audio_to_text = OpenAITextToAudio.from_dict(settings)
|
||||
dumped_settings = audio_to_text.to_dict()
|
||||
assert dumped_settings["ai_model_id"] == settings["ai_model_id"]
|
||||
assert dumped_settings["api_key"] == settings["api_key"]
|
||||
|
||||
|
||||
def test_prompt_execution_settings_class(openai_unit_test_env) -> None:
|
||||
openai_text_to_audio = OpenAITextToAudio()
|
||||
assert openai_text_to_audio.get_prompt_execution_settings_class() == OpenAITextToAudioExecutionSettings
|
||||
|
||||
|
||||
@patch.object(AsyncSpeech, "create", return_value=_legacy_response.HttpxBinaryResponseContent(httpx.Response(200)))
|
||||
async def test_get_text_contents(mock_speech_create, openai_unit_test_env):
|
||||
openai_text_to_audio = OpenAITextToAudio()
|
||||
|
||||
audio_contents = await openai_text_to_audio.get_audio_contents("Hello World!")
|
||||
assert len(audio_contents) == 1
|
||||
assert audio_contents[0].ai_model_id == openai_unit_test_env["OPENAI_TEXT_TO_AUDIO_MODEL_ID"]
|
||||
@@ -0,0 +1,375 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import os
|
||||
import warnings
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pydantic
|
||||
import pytest
|
||||
from openai import AsyncClient
|
||||
from openai.resources.images import AsyncImages
|
||||
from openai.types.image import Image
|
||||
from openai.types.images_response import ImagesResponse
|
||||
|
||||
from semantic_kernel.connectors.ai.open_ai import OpenAITextToImage, OpenAITextToImageExecutionSettings
|
||||
from semantic_kernel.connectors.ai.open_ai.services.open_ai_text_to_image_base import OpenAITextToImageBase
|
||||
from semantic_kernel.exceptions.service_exceptions import (
|
||||
ServiceInitializationError,
|
||||
ServiceInvalidExecutionSettingsError,
|
||||
ServiceInvalidRequestError,
|
||||
ServiceResponseException,
|
||||
)
|
||||
|
||||
sample_img = os.path.join(os.path.dirname(__file__), "../../../../../assets/sample_image.jpg")
|
||||
|
||||
|
||||
def test_init(openai_unit_test_env):
|
||||
"""Test that OpenAITextToImage initializes with the correct model id and client."""
|
||||
openai_text_to_image = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
|
||||
assert openai_text_to_image.client is not None
|
||||
assert isinstance(openai_text_to_image.client, AsyncClient)
|
||||
assert openai_text_to_image.ai_model_id == openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_TEXT_TO_IMAGE_MODEL_ID"]], indirect=True)
|
||||
def test_init_validation_fail(openai_unit_test_env) -> None:
|
||||
"""Test that initialization fails when required parameters are missing."""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAITextToImage(api_key="34523", ai_model_id=None, env_file_path="test.env")
|
||||
|
||||
|
||||
def test_init_to_from_dict(openai_unit_test_env):
|
||||
"""Test to_dict and from_dict methods for correct serialization and deserialization."""
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
settings = {
|
||||
"ai_model_id": openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"],
|
||||
"api_key": openai_unit_test_env["OPENAI_API_KEY"],
|
||||
"default_headers": default_headers,
|
||||
}
|
||||
text_to_image = OpenAITextToImage.from_dict(settings)
|
||||
dumped_settings = text_to_image.to_dict()
|
||||
assert dumped_settings["ai_model_id"] == settings["ai_model_id"]
|
||||
assert dumped_settings["api_key"] == settings["api_key"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True)
|
||||
def test_init_with_empty_api_key(openai_unit_test_env) -> None:
|
||||
"""Test that initialization fails when API key is missing."""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAITextToImage(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_TEXT_TO_IMAGE_MODEL_ID"]], indirect=True)
|
||||
def test_init_with_no_model_id(openai_unit_test_env) -> None:
|
||||
"""Test that initialization fails when model id is missing."""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAITextToImage(
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
def test_prompt_execution_settings_class(openai_unit_test_env) -> None:
|
||||
"""Test that the correct prompt execution settings class is returned."""
|
||||
openai_text_to_image = OpenAITextToImage()
|
||||
assert openai_text_to_image.get_prompt_execution_settings_class() == OpenAITextToImageExecutionSettings
|
||||
|
||||
|
||||
@patch.object(AsyncImages, "generate", new_callable=AsyncMock)
|
||||
async def test_generate_calls_with_parameters(mock_generate, openai_unit_test_env) -> None:
|
||||
"""Test that generate_image calls the OpenAI API with correct parameters."""
|
||||
mock_response = ImagesResponse(created=1, data=[Image(url="abc")], usage=None)
|
||||
mock_generate.return_value = mock_response
|
||||
|
||||
ai_model_id = "test_model_id"
|
||||
prompt = "painting of flowers in vase"
|
||||
width = 512
|
||||
|
||||
openai_text_to_image = OpenAITextToImage(ai_model_id=ai_model_id)
|
||||
|
||||
with warnings.catch_warnings(record=True) as w:
|
||||
await openai_text_to_image.generate_image(description=prompt, width=width, height=width)
|
||||
|
||||
mock_generate.assert_awaited_once_with(
|
||||
prompt=prompt,
|
||||
model=ai_model_id,
|
||||
size=f"{width}x{width}",
|
||||
n=1,
|
||||
)
|
||||
assert len(w) == 3
|
||||
|
||||
|
||||
@patch.object(AsyncImages, "generate", new_callable=AsyncMock, side_effect=Exception)
|
||||
async def test_generate_fail(mock_generate, openai_unit_test_env) -> None:
|
||||
"""Test that generate_image raises ServiceResponseException on API failure."""
|
||||
ai_model_id = "test_model_id"
|
||||
width = 512
|
||||
|
||||
openai_text_to_image = OpenAITextToImage(ai_model_id=ai_model_id)
|
||||
with pytest.raises(ServiceResponseException):
|
||||
await openai_text_to_image.generate_image(description="painting of flowers in vase", width=width, height=width)
|
||||
|
||||
|
||||
async def test_generate_invalid_image_size(openai_unit_test_env) -> None:
|
||||
"""Test that invalid image size raises ServiceInvalidExecutionSettingsError."""
|
||||
ai_model_id = "test_model_id"
|
||||
width = 100
|
||||
|
||||
openai_text_to_image = OpenAITextToImage(ai_model_id=ai_model_id)
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
await openai_text_to_image.generate_image(description="painting of flowers in vase", width=width, height=width)
|
||||
|
||||
|
||||
async def test_generate_empty_description(openai_unit_test_env) -> None:
|
||||
"""Test that empty description raises ServiceInvalidExecutionSettingsError."""
|
||||
ai_model_id = "test_model_id"
|
||||
width = 100
|
||||
|
||||
openai_text_to_image = OpenAITextToImage(ai_model_id=ai_model_id)
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
await openai_text_to_image.generate_image(description="", width=width, height=width)
|
||||
|
||||
|
||||
@patch.object(AsyncImages, "generate", new_callable=AsyncMock)
|
||||
async def test_generate_no_result(mock_generate, openai_unit_test_env) -> None:
|
||||
"""Test that no result from API raises ServiceResponseException."""
|
||||
mock_generate.return_value = ImagesResponse(created=0, data=[], usage=None)
|
||||
ai_model_id = "test_model_id"
|
||||
width = 512
|
||||
|
||||
openai_text_to_image = OpenAITextToImage(ai_model_id=ai_model_id)
|
||||
with pytest.raises(ServiceResponseException):
|
||||
await openai_text_to_image.generate_image(description="painting of flowers in vase", width=width, height=width)
|
||||
|
||||
|
||||
@patch.object(OpenAITextToImageBase, "_send_image_edit_request", new_callable=AsyncMock)
|
||||
async def test_edit_image_with_path_success(mock_edit, openai_unit_test_env):
|
||||
"""Test editing an image using a file path returns the expected URL."""
|
||||
mock_edit.return_value = ImagesResponse(created=1, data=[Image(url="edited_url")], usage=None)
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
result = await service.edit_image(
|
||||
prompt="edit this image",
|
||||
image_paths=[sample_img],
|
||||
)
|
||||
assert result == ["edited_url"]
|
||||
mock_edit.assert_awaited()
|
||||
|
||||
|
||||
@patch.object(OpenAITextToImageBase, "_send_image_edit_request", new_callable=AsyncMock)
|
||||
async def test_edit_image_with_file_success(mock_edit, openai_unit_test_env):
|
||||
"""Test editing an image using a file object returns the expected URL."""
|
||||
mock_edit.return_value = ImagesResponse(created=1, data=[Image(url="edited_url")], usage=None)
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
with open(sample_img, "rb") as f:
|
||||
result = await service.edit_image(
|
||||
prompt="edit this image",
|
||||
image_files=[f],
|
||||
)
|
||||
assert result == ["edited_url"]
|
||||
mock_edit.assert_awaited()
|
||||
|
||||
|
||||
@patch.object(OpenAITextToImageBase, "_send_image_edit_request", new_callable=AsyncMock)
|
||||
async def test_edit_image_with_mask_path_and_file(mock_edit, openai_unit_test_env):
|
||||
"""Test editing an image with both mask path and mask file returns the expected URL."""
|
||||
mock_edit.return_value = ImagesResponse(created=1, data=[Image(url="edited_url")], usage=None)
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
# mask_path
|
||||
result = await service.edit_image(
|
||||
prompt="edit with mask",
|
||||
image_paths=[sample_img],
|
||||
mask_path=sample_img,
|
||||
)
|
||||
assert result == ["edited_url"]
|
||||
# mask_file
|
||||
with open(sample_img, "rb") as mf:
|
||||
result2 = await service.edit_image(
|
||||
prompt="edit with mask",
|
||||
image_paths=[sample_img],
|
||||
mask_file=mf,
|
||||
)
|
||||
assert result2 == ["edited_url"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_image_prompt_required(openai_unit_test_env):
|
||||
"""Test that an empty prompt raises ServiceInvalidRequestError."""
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
with pytest.raises(ServiceInvalidRequestError):
|
||||
await service.edit_image(prompt="", image_paths=[sample_img])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_image_both_path_and_file_error(openai_unit_test_env):
|
||||
"""Test that providing both image_paths and image_files raises ServiceInvalidRequestError."""
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
with (
|
||||
open(sample_img, "rb") as f,
|
||||
pytest.raises(ServiceInvalidRequestError),
|
||||
):
|
||||
await service.edit_image(
|
||||
prompt="edit",
|
||||
image_paths=[sample_img],
|
||||
image_files=[f],
|
||||
)
|
||||
|
||||
|
||||
@patch.object(OpenAITextToImageBase, "_send_image_edit_request", new_callable=AsyncMock)
|
||||
async def test_edit_image_no_valid_data_in_response(mock_edit, openai_unit_test_env):
|
||||
"""Test that no valid data in edit response raises ServiceResponseException."""
|
||||
mock_edit.return_value = ImagesResponse(created=1, data=[], usage=None)
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
with pytest.raises(ServiceResponseException):
|
||||
await service.edit_image(
|
||||
prompt="edit",
|
||||
image_paths=[sample_img],
|
||||
)
|
||||
|
||||
|
||||
@patch.object(OpenAITextToImageBase, "_send_request", new_callable=AsyncMock)
|
||||
async def test_generate_images_with_n_parameter(mock_generate, openai_unit_test_env):
|
||||
"""Test that generate_images returns correct URLs when n parameter is set."""
|
||||
mock_generate.return_value = ImagesResponse(created=3, data=[Image(url=f"url_{i}") for i in range(3)], usage=None)
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
settings = OpenAITextToImageExecutionSettings(n=3)
|
||||
result = await service.generate_images("prompt", settings=settings)
|
||||
assert result == [f"url_{i}" for i in range(3)]
|
||||
|
||||
|
||||
@patch.object(OpenAITextToImageBase, "_send_request", new_callable=AsyncMock)
|
||||
async def test_generate_images_with_output_compression_and_background(mock_generate, openai_unit_test_env):
|
||||
"""Test that output_compression and background parameters are handled correctly."""
|
||||
mock_generate.return_value = ImagesResponse(created=1, data=[Image(url="url")], usage=None)
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
settings = OpenAITextToImageExecutionSettings(output_compression=5, background="transparent")
|
||||
await service.generate_images("prompt", settings=settings)
|
||||
called_settings = mock_generate.call_args[0][0]
|
||||
assert called_settings.output_compression == 5
|
||||
assert called_settings.background == "transparent"
|
||||
|
||||
|
||||
@patch.object(OpenAITextToImageBase, "store_usage")
|
||||
def test_store_usage_for_images_response(mock_store_usage, openai_unit_test_env):
|
||||
"""Test that store_usage is called for ImagesResponse."""
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
response = ImagesResponse(created=1, data=[Image(url="url")], usage=None)
|
||||
service.store_usage(response)
|
||||
mock_store_usage.assert_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_image_invalid_n_parameter():
|
||||
"""Test that invalid n parameter raises pydantic.ValidationError."""
|
||||
with pytest.raises(pydantic.ValidationError):
|
||||
OpenAITextToImageExecutionSettings(n=0)
|
||||
with pytest.raises(pydantic.ValidationError):
|
||||
OpenAITextToImageExecutionSettings(n=11)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_images_empty_prompt(openai_unit_test_env):
|
||||
"""Test that empty prompt raises ServiceInvalidRequestError."""
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
with pytest.raises(ServiceInvalidRequestError):
|
||||
await service.generate_images("")
|
||||
|
||||
|
||||
@patch.object(OpenAITextToImageBase, "_send_request", new_callable=AsyncMock)
|
||||
async def test_generate_images_no_result(mock_generate, openai_unit_test_env):
|
||||
"""Test that empty response data raises ServiceResponseException."""
|
||||
mock_generate.return_value = ImagesResponse(created=0, data=[], usage=None)
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
with pytest.raises(ServiceResponseException):
|
||||
await service.generate_images("prompt")
|
||||
|
||||
|
||||
@patch.object(OpenAITextToImageBase, "_send_request", new_callable=AsyncMock)
|
||||
async def test_generate_images_b64_json_response(mock_generate, openai_unit_test_env):
|
||||
"""Test that generate_images returns b64_json when url is not present."""
|
||||
mock_generate.return_value = ImagesResponse(created=1, data=[Image(b64_json="base64encodeddata")], usage=None)
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
result = await service.generate_images("prompt")
|
||||
assert result == ["base64encodeddata"]
|
||||
|
||||
|
||||
@patch.object(OpenAITextToImageBase, "_send_request", new_callable=AsyncMock)
|
||||
async def test_generate_images_mixed_url_and_b64_response(mock_generate, openai_unit_test_env):
|
||||
"""Test that generate_images handles mixed url and b64_json responses."""
|
||||
mock_generate.return_value = ImagesResponse(
|
||||
created=2,
|
||||
data=[Image(url="http://example.com/img1.png"), Image(b64_json="base64data")],
|
||||
usage=None,
|
||||
)
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
result = await service.generate_images("prompt")
|
||||
assert result == ["http://example.com/img1.png", "base64data"]
|
||||
|
||||
|
||||
@patch.object(OpenAITextToImageBase, "_send_request", new_callable=AsyncMock)
|
||||
async def test_generate_images_with_default_settings(mock_generate, openai_unit_test_env):
|
||||
"""Test that generate_images works when no settings are provided."""
|
||||
mock_generate.return_value = ImagesResponse(created=1, data=[Image(url="url")], usage=None)
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
result = await service.generate_images("a beautiful sunset")
|
||||
assert result == ["url"]
|
||||
mock_generate.assert_awaited_once()
|
||||
|
||||
|
||||
@patch.object(OpenAITextToImageBase, "_send_request", new_callable=AsyncMock)
|
||||
async def test_generate_images_no_valid_image_data(mock_generate, openai_unit_test_env):
|
||||
"""Test that generate_images raises error when images have neither url nor b64_json."""
|
||||
mock_generate.return_value = ImagesResponse(created=1, data=[Image()], usage=None)
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
with pytest.raises(ServiceResponseException, match="No valid image data found"):
|
||||
await service.generate_images("prompt")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edit_image_neither_path_nor_file(openai_unit_test_env):
|
||||
"""Test that providing neither image_paths nor image_files raises ServiceInvalidRequestError."""
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
with pytest.raises(ServiceInvalidRequestError):
|
||||
await service.edit_image(prompt="edit this")
|
||||
|
||||
|
||||
@patch.object(OpenAITextToImageBase, "_send_image_edit_request", new_callable=AsyncMock)
|
||||
async def test_edit_image_b64_json_response(mock_edit, openai_unit_test_env):
|
||||
"""Test editing an image returns b64_json when url is not present."""
|
||||
mock_edit.return_value = ImagesResponse(created=1, data=[Image(b64_json="edited_b64")], usage=None)
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
result = await service.edit_image(
|
||||
prompt="edit this image",
|
||||
image_paths=[sample_img],
|
||||
)
|
||||
assert result == ["edited_b64"]
|
||||
|
||||
|
||||
@patch.object(OpenAITextToImageBase, "_send_image_edit_request", new_callable=AsyncMock)
|
||||
async def test_edit_image_mixed_response(mock_edit, openai_unit_test_env):
|
||||
"""Test editing images handles mixed b64_json and url responses."""
|
||||
mock_edit.return_value = ImagesResponse(
|
||||
created=2,
|
||||
data=[Image(b64_json="b64data"), Image(url="http://example.com/edited.png")],
|
||||
usage=None,
|
||||
)
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
result = await service.edit_image(
|
||||
prompt="edit these images",
|
||||
image_paths=[sample_img],
|
||||
)
|
||||
assert result == ["b64data", "http://example.com/edited.png"]
|
||||
|
||||
|
||||
@patch.object(OpenAITextToImageBase, "_send_image_edit_request", new_callable=AsyncMock)
|
||||
async def test_edit_image_response_no_data_attribute(mock_edit, openai_unit_test_env):
|
||||
"""Test that edit_image raises error when response has no valid data."""
|
||||
mock_edit.return_value = ImagesResponse(created=1, data=None, usage=None)
|
||||
service = OpenAITextToImage(ai_model_id=openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"])
|
||||
with pytest.raises(ServiceResponseException):
|
||||
await service.edit_image(
|
||||
prompt="edit",
|
||||
image_paths=[sample_img],
|
||||
)
|
||||
@@ -0,0 +1,384 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.azure_chat_prompt_execution_settings import (
|
||||
AzureAISearchDataSource,
|
||||
AzureAISearchDataSourceParameters,
|
||||
AzureChatPromptExecutionSettings,
|
||||
ExtraBody,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.open_ai_prompt_execution_settings import (
|
||||
OpenAIChatPromptExecutionSettings,
|
||||
OpenAITextPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
|
||||
from semantic_kernel.connectors.azure_ai_search import AzureAISearchSettings
|
||||
from semantic_kernel.exceptions import ServiceInvalidExecutionSettingsError
|
||||
from semantic_kernel.kernel_pydantic import KernelBaseModel
|
||||
|
||||
|
||||
############################################
|
||||
# Test classes for structured output
|
||||
class ClassTest:
|
||||
attribute: str
|
||||
|
||||
|
||||
class ClassTestPydantic(KernelBaseModel):
|
||||
attribute: str
|
||||
|
||||
|
||||
############################################
|
||||
|
||||
|
||||
def test_default_openai_chat_prompt_execution_settings():
|
||||
settings = OpenAIChatPromptExecutionSettings()
|
||||
assert settings.temperature is None
|
||||
assert settings.top_p is None
|
||||
assert settings.presence_penalty is None
|
||||
assert settings.frequency_penalty is None
|
||||
assert settings.max_tokens is None
|
||||
assert settings.stop is None
|
||||
assert settings.number_of_responses is None
|
||||
assert settings.logit_bias is None
|
||||
assert settings.messages is None
|
||||
|
||||
|
||||
def test_custom_openai_chat_prompt_execution_settings():
|
||||
settings = OpenAIChatPromptExecutionSettings(
|
||||
temperature=0.5,
|
||||
top_p=0.5,
|
||||
presence_penalty=0.5,
|
||||
frequency_penalty=0.5,
|
||||
max_tokens=128,
|
||||
stop="\n",
|
||||
number_of_responses=2,
|
||||
logit_bias={"1": 1},
|
||||
messages=[{"role": "system", "content": "Hello"}],
|
||||
)
|
||||
assert settings.temperature == 0.5
|
||||
assert settings.top_p == 0.5
|
||||
assert settings.presence_penalty == 0.5
|
||||
assert settings.frequency_penalty == 0.5
|
||||
assert settings.max_tokens == 128
|
||||
assert settings.stop == "\n"
|
||||
assert settings.number_of_responses == 2
|
||||
assert settings.logit_bias == {"1": 1}
|
||||
assert settings.messages == [{"role": "system", "content": "Hello"}]
|
||||
|
||||
|
||||
def test_openai_chat_prompt_execution_settings_from_default_completion_config():
|
||||
settings = PromptExecutionSettings(service_id="test_service")
|
||||
chat_settings = OpenAIChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
assert chat_settings.service_id == "test_service"
|
||||
assert chat_settings.temperature is None
|
||||
assert chat_settings.top_p is None
|
||||
assert chat_settings.presence_penalty is None
|
||||
assert chat_settings.frequency_penalty is None
|
||||
assert chat_settings.max_tokens is None
|
||||
assert chat_settings.stop is None
|
||||
assert chat_settings.number_of_responses is None
|
||||
assert chat_settings.logit_bias is None
|
||||
|
||||
|
||||
def test_openai_chat_prompt_execution_settings_from_openai_prompt_execution_settings():
|
||||
chat_settings = OpenAIChatPromptExecutionSettings(service_id="test_service", temperature=1.0)
|
||||
new_settings = OpenAIChatPromptExecutionSettings(service_id="test_2", temperature=0.0)
|
||||
chat_settings.update_from_prompt_execution_settings(new_settings)
|
||||
assert chat_settings.service_id == "test_2"
|
||||
assert chat_settings.temperature == 0.0
|
||||
|
||||
|
||||
def test_openai_text_prompt_execution_settings_validation():
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError, match="best_of must be greater than number_of_responses"):
|
||||
OpenAITextPromptExecutionSettings(best_of=1, number_of_responses=2)
|
||||
|
||||
|
||||
def test_openai_text_prompt_execution_settings_validation_manual():
|
||||
text_oai = OpenAITextPromptExecutionSettings(best_of=1, number_of_responses=1)
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError, match="best_of must be greater than number_of_responses"):
|
||||
text_oai.number_of_responses = 2
|
||||
|
||||
|
||||
def test_openai_chat_prompt_execution_settings_from_custom_completion_config():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"presence_penalty": 0.5,
|
||||
"frequency_penalty": 0.5,
|
||||
"max_tokens": 128,
|
||||
"stop": ["\n"],
|
||||
"number_of_responses": 2,
|
||||
"logprobs": 1,
|
||||
"logit_bias": {"1": 1},
|
||||
"messages": [{"role": "system", "content": "Hello"}],
|
||||
},
|
||||
)
|
||||
chat_settings = OpenAIChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
assert chat_settings.temperature == 0.5
|
||||
assert chat_settings.top_p == 0.5
|
||||
assert chat_settings.presence_penalty == 0.5
|
||||
assert chat_settings.frequency_penalty == 0.5
|
||||
assert chat_settings.max_tokens == 128
|
||||
assert chat_settings.stop == ["\n"]
|
||||
assert chat_settings.number_of_responses == 2
|
||||
assert chat_settings.logit_bias == {"1": 1}
|
||||
|
||||
|
||||
def test_openai_chat_prompt_execution_settings_from_custom_completion_config_with_none():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"presence_penalty": 0.5,
|
||||
"frequency_penalty": 0.5,
|
||||
"max_tokens": 128,
|
||||
"stop": ["\n"],
|
||||
"number_of_responses": 2,
|
||||
"functions": None,
|
||||
"logit_bias": {"1": 1},
|
||||
"messages": [{"role": "system", "content": "Hello"}],
|
||||
},
|
||||
)
|
||||
chat_settings = OpenAIChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
assert chat_settings.temperature == 0.5
|
||||
assert chat_settings.top_p == 0.5
|
||||
assert chat_settings.presence_penalty == 0.5
|
||||
assert chat_settings.frequency_penalty == 0.5
|
||||
assert chat_settings.max_tokens == 128
|
||||
assert chat_settings.stop == ["\n"]
|
||||
assert chat_settings.number_of_responses == 2
|
||||
assert chat_settings.logit_bias == {"1": 1}
|
||||
assert chat_settings.functions is None
|
||||
|
||||
|
||||
def test_openai_chat_prompt_execution_settings_from_custom_completion_config_with_functions():
|
||||
settings = PromptExecutionSettings(
|
||||
service_id="test_service",
|
||||
extension_data={
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.5,
|
||||
"presence_penalty": 0.5,
|
||||
"frequency_penalty": 0.5,
|
||||
"max_tokens": 128,
|
||||
"stop": ["\n"],
|
||||
"number_of_responses": 2,
|
||||
"functions": [{}],
|
||||
"function_call": "auto",
|
||||
"logit_bias": {"1": 1},
|
||||
"messages": [{"role": "system", "content": "Hello"}],
|
||||
},
|
||||
)
|
||||
chat_settings = OpenAIChatPromptExecutionSettings.from_prompt_execution_settings(settings)
|
||||
assert chat_settings.temperature == 0.5
|
||||
assert chat_settings.top_p == 0.5
|
||||
assert chat_settings.presence_penalty == 0.5
|
||||
assert chat_settings.frequency_penalty == 0.5
|
||||
assert chat_settings.max_tokens == 128
|
||||
assert chat_settings.stop == ["\n"]
|
||||
assert chat_settings.number_of_responses == 2
|
||||
assert chat_settings.logit_bias == {"1": 1}
|
||||
assert chat_settings.functions == [{}]
|
||||
|
||||
|
||||
def test_create_options():
|
||||
settings = OpenAIChatPromptExecutionSettings(
|
||||
temperature=0.5,
|
||||
top_p=0.5,
|
||||
presence_penalty=0.5,
|
||||
frequency_penalty=0.5,
|
||||
max_tokens=128,
|
||||
stop=["\n"],
|
||||
number_of_responses=2,
|
||||
logit_bias={"1": 1},
|
||||
messages=[{"role": "system", "content": "Hello"}],
|
||||
function_call="auto",
|
||||
)
|
||||
options = settings.prepare_settings_dict()
|
||||
assert options["temperature"] == 0.5
|
||||
assert options["top_p"] == 0.5
|
||||
assert options["presence_penalty"] == 0.5
|
||||
assert options["frequency_penalty"] == 0.5
|
||||
assert options["max_tokens"] == 128
|
||||
assert options["stop"] == ["\n"]
|
||||
assert options["n"] == 2
|
||||
assert options["logit_bias"] == {"1": 1}
|
||||
assert not options["stream"]
|
||||
|
||||
|
||||
def test_create_options_azure_data():
|
||||
az_source = AzureAISearchDataSource(
|
||||
parameters={
|
||||
"indexName": "test-index",
|
||||
"endpoint": "test-endpoint",
|
||||
"authentication": {"type": "api_key", "key": "test-key"},
|
||||
}
|
||||
)
|
||||
extra = ExtraBody(data_sources=[az_source])
|
||||
assert extra["data_sources"] is not None
|
||||
assert extra.data_sources is not None
|
||||
settings = AzureChatPromptExecutionSettings(extra_body=extra)
|
||||
options = settings.prepare_settings_dict()
|
||||
assert options["extra_body"] == extra.model_dump(exclude_none=True, by_alias=True)
|
||||
assert options["extra_body"]["data_sources"][0]["type"] == "azure_search"
|
||||
|
||||
|
||||
def test_create_options_azure_data_from_azure_ai_settings(azure_ai_search_unit_test_env):
|
||||
az_source = AzureAISearchDataSource.from_azure_ai_search_settings(AzureAISearchSettings())
|
||||
extra = ExtraBody(data_sources=[az_source])
|
||||
assert extra["data_sources"] is not None
|
||||
settings = AzureChatPromptExecutionSettings(extra_body=extra)
|
||||
options = settings.prepare_settings_dict()
|
||||
assert options["extra_body"] == extra.model_dump(exclude_none=True, by_alias=True)
|
||||
assert options["extra_body"]["data_sources"][0]["type"] == "azure_search"
|
||||
|
||||
|
||||
def test_azure_open_ai_chat_prompt_execution_settings_with_cosmosdb_data_sources():
|
||||
input_dict = {
|
||||
"messages": [{"role": "system", "content": "Hello"}],
|
||||
"extra_body": {
|
||||
"dataSources": [
|
||||
{
|
||||
"type": "AzureCosmosDB",
|
||||
"parameters": {
|
||||
"authentication": {
|
||||
"type": "ConnectionString",
|
||||
"connectionString": "mongodb+srv://onyourdatatest:{password}$@{cluster-name}.mongocluster.cosmos.azure.com/?tls=true&authMechanism=SCRAM-SHA-256&retrywrites=false&maxIdleTimeMS=120000",
|
||||
},
|
||||
"databaseName": "vectordb",
|
||||
"containerName": "azuredocs",
|
||||
"indexName": "azuredocindex",
|
||||
"embeddingDependency": {
|
||||
"type": "DeploymentName",
|
||||
"deploymentName": "{embedding deployment name}",
|
||||
},
|
||||
"fieldsMapping": {"vectorFields": ["contentvector"]},
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
settings = AzureChatPromptExecutionSettings.model_validate(input_dict, strict=True, from_attributes=True)
|
||||
assert settings.extra_body["dataSources"][0]["type"] == "AzureCosmosDB"
|
||||
|
||||
|
||||
def test_azure_open_ai_chat_prompt_execution_settings_with_aisearch_data_sources():
|
||||
input_dict = {
|
||||
"messages": [{"role": "system", "content": "Hello"}],
|
||||
"extra_body": {
|
||||
"dataSources": [
|
||||
{
|
||||
"type": "AzureCognitiveSearch",
|
||||
"parameters": {
|
||||
"authentication": {
|
||||
"type": "APIKey",
|
||||
"key": "****",
|
||||
},
|
||||
"endpoint": "https://****.search.windows.net/",
|
||||
"indexName": "azuredocindex",
|
||||
"queryType": "vector",
|
||||
"embeddingDependency": {
|
||||
"type": "DeploymentName",
|
||||
"deploymentName": "{embedding deployment name}",
|
||||
},
|
||||
"fieldsMapping": {"vectorFields": ["contentvector"]},
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
settings = AzureChatPromptExecutionSettings.model_validate(input_dict, strict=True, from_attributes=True)
|
||||
assert settings.extra_body["dataSources"][0]["type"] == "AzureCognitiveSearch"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"authentication",
|
||||
[
|
||||
{"type": "APIKey", "key": "test_key"},
|
||||
{"type": "api_key", "key": "test_key"},
|
||||
pytest.param({"type": "api_key"}, marks=pytest.mark.xfail),
|
||||
{"type": "SystemAssignedManagedIdentity"},
|
||||
{"type": "system_assigned_managed_identity"},
|
||||
{"type": "UserAssignedManagedIdentity", "managed_identity_resource_id": "test_id"},
|
||||
{"type": "user_assigned_managed_identity", "managed_identity_resource_id": "test_id"},
|
||||
pytest.param({"type": "user_assigned_managed_identity"}, marks=pytest.mark.xfail),
|
||||
{"type": "AccessToken", "access_token": "test_token"},
|
||||
{"type": "access_token", "access_token": "test_token"},
|
||||
pytest.param({"type": "access_token"}, marks=pytest.mark.xfail),
|
||||
pytest.param({"type": "invalid"}, marks=pytest.mark.xfail),
|
||||
],
|
||||
ids=[
|
||||
"APIKey",
|
||||
"api_key",
|
||||
"api_key_no_key",
|
||||
"SystemAssignedManagedIdentity",
|
||||
"system_assigned_managed_identity",
|
||||
"UserAssignedManagedIdentity",
|
||||
"user_assigned_managed_identity",
|
||||
"user_assigned_managed_identity_no_id",
|
||||
"AccessToken",
|
||||
"access_token",
|
||||
"access_token_no_token",
|
||||
"invalid",
|
||||
],
|
||||
)
|
||||
def test_aisearch_data_source_parameters(authentication) -> None:
|
||||
AzureAISearchDataSourceParameters(index_name="test_index", authentication=authentication)
|
||||
|
||||
|
||||
def test_azure_open_ai_chat_prompt_execution_settings_with_response_format_json():
|
||||
response_format = {"type": "json_object"}
|
||||
settings = AzureChatPromptExecutionSettings(response_format=response_format)
|
||||
options = settings.prepare_settings_dict()
|
||||
assert options["response_format"] == response_format
|
||||
|
||||
|
||||
def test_openai_chat_prompt_execution_settings_with_json_structured_output():
|
||||
settings = OpenAIChatPromptExecutionSettings()
|
||||
settings.response_format = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "math_response",
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"steps": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {"explanation": {"type": "string"}, "output": {"type": "string"}},
|
||||
"required": ["explanation", "output"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
"final_answer": {"type": "string"},
|
||||
},
|
||||
"required": ["steps", "final_answer"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"strict": True,
|
||||
},
|
||||
}
|
||||
assert isinstance(settings.response_format, dict)
|
||||
|
||||
|
||||
def test_openai_chat_prompt_execution_settings_with_nonpydantic_type_structured_output():
|
||||
settings = OpenAIChatPromptExecutionSettings()
|
||||
settings.response_format = ClassTest
|
||||
assert isinstance(settings.response_format, type)
|
||||
|
||||
|
||||
def test_openai_chat_prompt_execution_settings_with_pydantic_type_structured_output():
|
||||
settings = OpenAIChatPromptExecutionSettings()
|
||||
settings.response_format = ClassTestPydantic
|
||||
assert issubclass(settings.response_format, BaseModel)
|
||||
|
||||
|
||||
def test_openai_chat_prompt_execution_settings_with_invalid_structured_output():
|
||||
settings = OpenAIChatPromptExecutionSettings()
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
settings.response_format = "invalid"
|
||||
@@ -0,0 +1,43 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from semantic_kernel.connectors.ai.completion_usage import CompletionUsage
|
||||
|
||||
|
||||
def test_completion_usage() -> None:
|
||||
"""Test the CompletionUsage class."""
|
||||
# Create a CompletionUsage object
|
||||
usage = CompletionUsage(prompt_tokens=10, completion_tokens=20)
|
||||
|
||||
# Check that the prompt tokens and completion tokens are set correctly
|
||||
assert usage.prompt_tokens == 10
|
||||
assert usage.completion_tokens == 20
|
||||
|
||||
# Create another CompletionUsage object
|
||||
other_usage = CompletionUsage(prompt_tokens=5, completion_tokens=15)
|
||||
|
||||
# Add the two CompletionUsage objects together
|
||||
total_usage = usage + other_usage
|
||||
|
||||
# Check that the total prompt tokens and completion tokens are correct
|
||||
assert total_usage.prompt_tokens == 15
|
||||
assert total_usage.completion_tokens == 35
|
||||
|
||||
|
||||
def test_completion_usage_empty() -> None:
|
||||
"""Test the CompletionUsage class with empty values."""
|
||||
# Create a CompletionUsage object with empty values
|
||||
usage = CompletionUsage()
|
||||
|
||||
# Check that the prompt tokens and completion tokens are None
|
||||
assert usage.prompt_tokens is None
|
||||
assert usage.completion_tokens is None
|
||||
|
||||
# Create another CompletionUsage object with empty values
|
||||
other_usage = CompletionUsage(prompt_tokens=5, completion_tokens=None)
|
||||
|
||||
# Add the two CompletionUsage objects together
|
||||
total_usage = usage + other_usage
|
||||
|
||||
# Check that the total prompt tokens and completion tokens are None
|
||||
assert total_usage.prompt_tokens == 5
|
||||
assert total_usage.completion_tokens == 0
|
||||
@@ -0,0 +1,232 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from semantic_kernel.kernel import Kernel
|
||||
|
||||
from semantic_kernel.connectors.ai.function_calling_utils import _combine_filter_dicts
|
||||
from semantic_kernel.connectors.ai.function_choice_behavior import (
|
||||
DEFAULT_MAX_AUTO_INVOKE_ATTEMPTS,
|
||||
FunctionChoiceBehavior,
|
||||
FunctionChoiceType,
|
||||
)
|
||||
from semantic_kernel.exceptions import ServiceInitializationError
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def function_choice_behavior():
|
||||
return FunctionChoiceBehavior()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def update_settings_callback():
|
||||
mock = Mock()
|
||||
mock.return_value = None
|
||||
return mock
|
||||
|
||||
|
||||
def test_function_choice_behavior_auto():
|
||||
behavior = FunctionChoiceBehavior.Auto(auto_invoke=True)
|
||||
assert behavior.type_ == FunctionChoiceType.AUTO
|
||||
assert behavior.maximum_auto_invoke_attempts == DEFAULT_MAX_AUTO_INVOKE_ATTEMPTS
|
||||
|
||||
|
||||
def test_function_choice_behavior_none_invoke():
|
||||
behavior = FunctionChoiceBehavior.NoneInvoke()
|
||||
assert behavior.type_ == FunctionChoiceType.NONE
|
||||
assert behavior.maximum_auto_invoke_attempts == 0
|
||||
|
||||
|
||||
def test_function_choice_behavior_required():
|
||||
expected_filters = {"included_functions": ["plugin1-func1"]}
|
||||
behavior = FunctionChoiceBehavior.Required(auto_invoke=True, filters=expected_filters)
|
||||
assert behavior.type_ == FunctionChoiceType.REQUIRED
|
||||
assert behavior.maximum_auto_invoke_attempts == 1
|
||||
assert behavior.filters == expected_filters
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("type", "max_auto_invoke_attempts"), [("auto", 5), ("none", 0), ("required", 1)])
|
||||
def test_auto_function_choice_behavior_from_dict(type: str, max_auto_invoke_attempts: int):
|
||||
data = {
|
||||
"type": type,
|
||||
"filters": {"included_functions": ["plugin1-func1", "plugin2-func2"]},
|
||||
"maximum_auto_invoke_attempts": max_auto_invoke_attempts,
|
||||
}
|
||||
behavior = FunctionChoiceBehavior.from_dict(data)
|
||||
assert behavior.type_ == FunctionChoiceType(type)
|
||||
assert behavior.filters == {"included_functions": ["plugin1-func1", "plugin2-func2"]}
|
||||
assert behavior.maximum_auto_invoke_attempts == max_auto_invoke_attempts
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("type", "max_auto_invoke_attempts"), [("auto", 5), ("none", 0), ("required", 1)])
|
||||
def test_auto_function_choice_behavior_from_dict_with_same_filters_and_functions(
|
||||
type: str, max_auto_invoke_attempts: int
|
||||
):
|
||||
data = {
|
||||
"type": type,
|
||||
"filters": {"included_functions": ["plugin1-func1", "plugin2-func2"]},
|
||||
"functions": ["plugin1-func1", "plugin2-func2"],
|
||||
"maximum_auto_invoke_attempts": max_auto_invoke_attempts,
|
||||
}
|
||||
behavior = FunctionChoiceBehavior.from_dict(data)
|
||||
assert behavior.type_ == FunctionChoiceType(type)
|
||||
assert behavior.filters == {"included_functions": ["plugin1-func1", "plugin2-func2"]}
|
||||
assert behavior.maximum_auto_invoke_attempts == max_auto_invoke_attempts
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("type", "max_auto_invoke_attempts"), [("auto", 5), ("none", 0), ("required", 1)])
|
||||
def test_auto_function_choice_behavior_from_dict_with_different_filters_and_functions(
|
||||
type: str, max_auto_invoke_attempts: int
|
||||
):
|
||||
data = {
|
||||
"type": type,
|
||||
"filters": {"included_functions": ["plugin1-func1", "plugin2-func2"]},
|
||||
"functions": ["plugin3-func3"],
|
||||
"maximum_auto_invoke_attempts": max_auto_invoke_attempts,
|
||||
}
|
||||
behavior = FunctionChoiceBehavior.from_dict(data)
|
||||
assert behavior.type_ == FunctionChoiceType(type)
|
||||
assert behavior.filters == {"included_functions": ["plugin1-func1", "plugin2-func2", "plugin3-func3"]}
|
||||
assert behavior.maximum_auto_invoke_attempts == max_auto_invoke_attempts
|
||||
|
||||
|
||||
def test_function_choice_behavior_get_set(function_choice_behavior: FunctionChoiceBehavior):
|
||||
function_choice_behavior.enable_kernel_functions = False
|
||||
assert function_choice_behavior.enable_kernel_functions is False
|
||||
function_choice_behavior.maximum_auto_invoke_attempts = 10
|
||||
assert function_choice_behavior.maximum_auto_invoke_attempts == 10
|
||||
assert function_choice_behavior.auto_invoke_kernel_functions is True
|
||||
function_choice_behavior.auto_invoke_kernel_functions = False
|
||||
assert function_choice_behavior.auto_invoke_kernel_functions is False
|
||||
assert function_choice_behavior.maximum_auto_invoke_attempts == 0
|
||||
function_choice_behavior.auto_invoke_kernel_functions = True
|
||||
assert function_choice_behavior.auto_invoke_kernel_functions is True
|
||||
assert function_choice_behavior.maximum_auto_invoke_attempts == 5
|
||||
|
||||
|
||||
def test_auto_invoke_kernel_functions():
|
||||
fcb = FunctionChoiceBehavior.Auto(auto_invoke=True)
|
||||
assert fcb is not None
|
||||
assert fcb.enable_kernel_functions is True
|
||||
assert fcb.maximum_auto_invoke_attempts == 5
|
||||
assert fcb.auto_invoke_kernel_functions is True
|
||||
|
||||
|
||||
def test_none_invoke_kernel_functions():
|
||||
fcb = FunctionChoiceBehavior.NoneInvoke()
|
||||
assert fcb is not None
|
||||
assert fcb.enable_kernel_functions is True
|
||||
assert fcb.maximum_auto_invoke_attempts == 0
|
||||
assert fcb.auto_invoke_kernel_functions is False
|
||||
|
||||
|
||||
def test_enable_functions():
|
||||
fcb = FunctionChoiceBehavior.Auto(auto_invoke=True, filters={"excluded_plugins": ["test"]})
|
||||
assert fcb is not None
|
||||
assert fcb.enable_kernel_functions is True
|
||||
assert fcb.maximum_auto_invoke_attempts == 5
|
||||
assert fcb.auto_invoke_kernel_functions is True
|
||||
assert fcb.filters == {"excluded_plugins": ["test"]}
|
||||
|
||||
|
||||
def test_required_function():
|
||||
fcb = FunctionChoiceBehavior.Required(auto_invoke=True, filters={"included_functions": ["test"]})
|
||||
assert fcb is not None
|
||||
assert fcb.enable_kernel_functions is True
|
||||
assert fcb.maximum_auto_invoke_attempts == 1
|
||||
assert fcb.auto_invoke_kernel_functions is True
|
||||
|
||||
|
||||
def test_configure_auto_invoke_kernel_functions(update_settings_callback, kernel: "Kernel"):
|
||||
fcb = FunctionChoiceBehavior.Auto(auto_invoke=True)
|
||||
fcb.configure(kernel, update_settings_callback, None)
|
||||
assert update_settings_callback.called
|
||||
|
||||
|
||||
def test_configure_auto_invoke_kernel_functions_skip(update_settings_callback, kernel: "Kernel"):
|
||||
fcb = FunctionChoiceBehavior.Auto(auto_invoke=True)
|
||||
fcb.enable_kernel_functions = False
|
||||
fcb.configure(kernel, update_settings_callback, None)
|
||||
assert not update_settings_callback.called
|
||||
|
||||
|
||||
def test_configure_none_invoke_kernel_functions(update_settings_callback, kernel: "Kernel"):
|
||||
fcb = FunctionChoiceBehavior.NoneInvoke()
|
||||
fcb.configure(kernel, update_settings_callback, None)
|
||||
assert update_settings_callback.called
|
||||
|
||||
|
||||
def test_configure_none_invoke_kernel_functions_skip(update_settings_callback, kernel: "Kernel"):
|
||||
fcb = FunctionChoiceBehavior.NoneInvoke()
|
||||
fcb.enable_kernel_functions = False
|
||||
fcb.configure(kernel, update_settings_callback, None)
|
||||
assert not update_settings_callback.called
|
||||
|
||||
|
||||
def test_configure_enable_functions(update_settings_callback, kernel: "Kernel"):
|
||||
fcb = FunctionChoiceBehavior.Auto(auto_invoke=True, filters={"excluded_plugins": ["test"]})
|
||||
fcb.configure(kernel, update_settings_callback, None)
|
||||
assert update_settings_callback.called
|
||||
|
||||
|
||||
def test_configure_enable_functions_skip(update_settings_callback, kernel: "Kernel"):
|
||||
fcb = FunctionChoiceBehavior.Auto(auto_invoke=True, filters={"excluded_plugins": ["test"]})
|
||||
fcb.enable_kernel_functions = False
|
||||
fcb.configure(kernel, update_settings_callback, None)
|
||||
assert not update_settings_callback.called
|
||||
|
||||
|
||||
def test_configure_required_function(update_settings_callback, kernel: "Kernel"):
|
||||
fcb = FunctionChoiceBehavior.Required(auto_invoke=True, filters={"included_functions": ["plugin1-func1"]})
|
||||
fcb.configure(kernel, update_settings_callback, None)
|
||||
assert update_settings_callback.called
|
||||
|
||||
|
||||
def test_configure_required_function_max_invoke_updated(update_settings_callback, kernel: "Kernel"):
|
||||
fcb = FunctionChoiceBehavior.Required(auto_invoke=True, filters={"included_functions": ["plugin1-func1"]})
|
||||
fcb.maximum_auto_invoke_attempts = 10
|
||||
fcb.configure(kernel, update_settings_callback, None)
|
||||
assert update_settings_callback.called
|
||||
assert fcb.maximum_auto_invoke_attempts == 10
|
||||
|
||||
|
||||
def test_configure_required_function_skip(update_settings_callback, kernel: "Kernel"):
|
||||
fcb = FunctionChoiceBehavior.Required(auto_invoke=True, filters={"included_functions": ["test"]})
|
||||
fcb.enable_kernel_functions = False
|
||||
fcb.configure(kernel, update_settings_callback, None)
|
||||
assert not update_settings_callback.called
|
||||
|
||||
|
||||
def test_service_initialization_error():
|
||||
dict1 = {"filter1": ["a", "b", "c"]}
|
||||
dict2 = {"filter1": "not_a_list"} # This should trigger the error
|
||||
|
||||
with pytest.raises(ServiceInitializationError, match="Values for filter key 'filter1' are not lists."):
|
||||
_combine_filter_dicts(dict1, dict2)
|
||||
|
||||
|
||||
def test_from_string_auto():
|
||||
auto = FunctionChoiceBehavior.from_string("auto")
|
||||
assert auto == FunctionChoiceBehavior.Auto()
|
||||
|
||||
|
||||
def test_from_string_none():
|
||||
none = FunctionChoiceBehavior.from_string("none")
|
||||
assert none == FunctionChoiceBehavior.NoneInvoke()
|
||||
|
||||
|
||||
def test_from_string_required():
|
||||
required = FunctionChoiceBehavior.from_string("required")
|
||||
assert required == FunctionChoiceBehavior.Required()
|
||||
|
||||
|
||||
def test_from_string_invalid():
|
||||
with pytest.raises(
|
||||
ServiceInitializationError,
|
||||
match="The specified type `invalid` is not supported. Allowed types are: `auto`, `none`, `required`.",
|
||||
):
|
||||
FunctionChoiceBehavior.from_string("invalid")
|
||||
@@ -0,0 +1,16 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from semantic_kernel.connectors.ai import PromptExecutionSettings
|
||||
|
||||
|
||||
def test_init():
|
||||
settings = PromptExecutionSettings()
|
||||
assert settings.service_id is None
|
||||
assert settings.extension_data == {}
|
||||
|
||||
|
||||
def test_init_with_data():
|
||||
ext_data = {"test": "test"}
|
||||
settings = PromptExecutionSettings(service_id="test", extension_data=ext_data)
|
||||
assert settings.service_id == "test"
|
||||
assert settings.extension_data["test"] == "test"
|
||||
@@ -0,0 +1,150 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from pytest import fixture
|
||||
|
||||
|
||||
@fixture()
|
||||
def mistralai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for MistralAISettings."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"MISTRALAI_CHAT_MODEL_ID": "test_chat_model_id",
|
||||
"MISTRALAI_API_KEY": "test_api_key",
|
||||
"MISTRALAI_EMBEDDING_MODEL_ID": "test_embedding_model_id",
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@fixture()
|
||||
def anthropic_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for AnthropicSettings."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {"ANTHROPIC_CHAT_MODEL_ID": "test_chat_model_id", "ANTHROPIC_API_KEY": "test_api_key"}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@fixture
|
||||
def azure_ai_search_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for ACA Python Unit Tests."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"AZURE_AI_SEARCH_API_KEY": "test-api-key",
|
||||
"AZURE_AI_SEARCH_ENDPOINT": "https://test-endpoint.com",
|
||||
"AZURE_AI_SEARCH_INDEX_NAME": "test-index-name",
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@fixture()
|
||||
def bing_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for BingConnector."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"BING_API_KEY": "test_api_key",
|
||||
"BING_CUSTOM_CONFIG": "test_org_id",
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@fixture()
|
||||
def brave_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for BraveConnector."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {"BRAVE_API_KEY": "test_api_key"}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@fixture()
|
||||
def google_search_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for the Google Search Connector."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"GOOGLE_SEARCH_API_KEY": "test_api_key",
|
||||
"GOOGLE_SEARCH_ENGINE_ID": "test_id",
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
@@ -0,0 +1,480 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
import re
|
||||
from typing import TYPE_CHECKING
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from mcp import ClientSession, ListToolsResult, StdioServerParameters, Tool, types
|
||||
|
||||
from semantic_kernel.connectors.mcp import MCPSsePlugin, MCPStdioPlugin, MCPStreamableHttpPlugin, MCPWebsocketPlugin
|
||||
from semantic_kernel.exceptions import KernelPluginInvalidConfigurationError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from semantic_kernel import Kernel
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def list_tool_calls_with_slash() -> ListToolsResult:
|
||||
return ListToolsResult(
|
||||
tools=[
|
||||
Tool(
|
||||
name="nasa/get-astronomy-picture",
|
||||
description="func with slash",
|
||||
inputSchema={"properties": {}, "required": []},
|
||||
),
|
||||
Tool(
|
||||
name="weird\\name with spaces",
|
||||
description="func with backslash and spaces",
|
||||
inputSchema={"properties": {}, "required": []},
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def list_tool_calls() -> ListToolsResult:
|
||||
return ListToolsResult(
|
||||
tools=[
|
||||
Tool(
|
||||
name="func1",
|
||||
description="func1",
|
||||
inputSchema={
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
},
|
||||
"required": ["name"],
|
||||
},
|
||||
),
|
||||
Tool(
|
||||
name="func2",
|
||||
description="func2",
|
||||
inputSchema={},
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"plugin_class,plugin_args",
|
||||
[
|
||||
(MCPSsePlugin, {"url": "http://localhost:8080/sse"}),
|
||||
(MCPStreamableHttpPlugin, {"url": "http://localhost:8080/mcp"}),
|
||||
],
|
||||
)
|
||||
async def test_mcp_plugin_session_not_initialize(plugin_class, plugin_args):
|
||||
# Test if Client can insert it's own Session
|
||||
mock_session = AsyncMock(spec=ClientSession)
|
||||
mock_session._request_id = 0
|
||||
mock_session.initialize = AsyncMock()
|
||||
async with plugin_class(name="test", session=mock_session, **plugin_args) as plugin:
|
||||
assert plugin.session is mock_session
|
||||
assert mock_session.initialize.called
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"plugin_class,plugin_args",
|
||||
[
|
||||
(MCPSsePlugin, {"url": "http://localhost:8080/sse"}),
|
||||
(MCPStreamableHttpPlugin, {"url": "http://localhost:8080/mcp"}),
|
||||
],
|
||||
)
|
||||
async def test_mcp_plugin_session_initialized(plugin_class, plugin_args):
|
||||
# Test if Client can insert it's own initialized Session
|
||||
mock_session = AsyncMock(spec=ClientSession)
|
||||
mock_session._request_id = 1
|
||||
mock_session.initialize = AsyncMock()
|
||||
async with plugin_class(name="test", session=mock_session, **plugin_args) as plugin:
|
||||
assert plugin.session is mock_session
|
||||
assert not mock_session.initialize.called
|
||||
|
||||
|
||||
async def test_mcp_sampling_denied_by_consent_callback():
|
||||
sampling_consent_callback = AsyncMock(return_value=False)
|
||||
plugin = MCPSsePlugin(
|
||||
name="TestMCPPlugin",
|
||||
url="http://localhost:8080/sse",
|
||||
sampling_consent_callback=sampling_consent_callback,
|
||||
)
|
||||
params = types.CreateMessageRequestParams(
|
||||
messages=[types.SamplingMessage(role="user", content=types.TextContent(type="text", text="hello"))],
|
||||
systemPrompt="server instructions",
|
||||
maxTokens=100,
|
||||
)
|
||||
|
||||
result = await plugin.sampling_callback(MagicMock(), params)
|
||||
|
||||
sampling_consent_callback.assert_awaited_once_with("TestMCPPlugin", params)
|
||||
assert isinstance(result, types.ErrorData)
|
||||
assert result.message == "Sampling denied by policy."
|
||||
|
||||
|
||||
async def test_mcp_sampling_consent_callback_error_denies_request(caplog):
|
||||
sampling_consent_callback = AsyncMock(side_effect=RuntimeError("policy failure"))
|
||||
plugin = MCPSsePlugin(
|
||||
name="TestMCPPlugin",
|
||||
url="http://localhost:8080/sse",
|
||||
sampling_consent_callback=sampling_consent_callback,
|
||||
)
|
||||
params = types.CreateMessageRequestParams(
|
||||
messages=[types.SamplingMessage(role="user", content=types.TextContent(type="text", text="hello"))],
|
||||
systemPrompt="server instructions",
|
||||
maxTokens=100,
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.ERROR, logger="semantic_kernel.connectors.mcp"):
|
||||
result = await plugin.sampling_callback(MagicMock(), params)
|
||||
|
||||
sampling_consent_callback.assert_awaited_once_with("TestMCPPlugin", params)
|
||||
assert isinstance(result, types.ErrorData)
|
||||
assert result.message == "Sampling denied by policy."
|
||||
assert "MCP sampling consent callback failed" in caplog.text
|
||||
|
||||
|
||||
async def test_mcp_sampling_without_consent_callback_denies_by_default(caplog):
|
||||
plugin = MCPSsePlugin(name="TestMCPPlugin", url="http://localhost:8080/sse")
|
||||
params = types.CreateMessageRequestParams(
|
||||
messages=[types.SamplingMessage(role="user", content=types.TextContent(type="text", text="hello"))],
|
||||
systemPrompt="server instructions",
|
||||
maxTokens=100,
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="semantic_kernel.connectors.mcp"):
|
||||
result = await plugin.sampling_callback(MagicMock(), params)
|
||||
|
||||
assert isinstance(result, types.ErrorData)
|
||||
assert result.message == "Sampling denied: no consent callback configured."
|
||||
assert "denied because no sampling consent callback was configured" in caplog.text
|
||||
|
||||
|
||||
async def test_mcp_sampling_auto_approve_logs_warning(caplog):
|
||||
plugin = MCPSsePlugin(
|
||||
name="TestMCPPlugin",
|
||||
url="http://localhost:8080/sse",
|
||||
sampling_auto_approve=True,
|
||||
)
|
||||
params = types.CreateMessageRequestParams(
|
||||
messages=[types.SamplingMessage(role="user", content=types.TextContent(type="text", text="hello"))],
|
||||
systemPrompt="server instructions",
|
||||
maxTokens=100,
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="semantic_kernel.connectors.mcp"):
|
||||
result = await plugin.sampling_callback(MagicMock(), params)
|
||||
|
||||
# No kernel configured, so the request is approved but then fails for lack of a chat service.
|
||||
assert isinstance(result, types.ErrorData)
|
||||
assert "auto-approved because sampling_auto_approve is enabled" in caplog.text
|
||||
|
||||
|
||||
async def test_mcp_tool_and_prompt_names_do_not_shadow_plugin_attributes():
|
||||
kernel = MagicMock()
|
||||
plugin = MCPSsePlugin(name="TestMCPPlugin", url="http://localhost:8080/sse", kernel=kernel)
|
||||
session = AsyncMock(spec=ClientSession)
|
||||
session.list_tools.return_value = ListToolsResult(
|
||||
tools=[
|
||||
Tool(name="kernel", description="reserved", inputSchema={}),
|
||||
Tool(name="safe_tool", description="safe", inputSchema={}),
|
||||
]
|
||||
)
|
||||
session.list_prompts.return_value = types.ListPromptsResult(
|
||||
prompts=[
|
||||
types.Prompt(name="session", description="reserved", arguments=[]),
|
||||
types.Prompt(name="safe_prompt", description="safe", arguments=[]),
|
||||
]
|
||||
)
|
||||
plugin.session = session
|
||||
|
||||
await plugin.load_tools()
|
||||
|
||||
assert plugin.kernel is kernel
|
||||
assert hasattr(plugin, "safe_tool")
|
||||
|
||||
await plugin.load_prompts()
|
||||
|
||||
assert plugin.session is session
|
||||
assert hasattr(plugin, "safe_prompt")
|
||||
|
||||
|
||||
async def test_mcp_tool_and_prompt_names_can_reload_existing_mcp_functions():
|
||||
plugin = MCPSsePlugin(name="TestMCPPlugin", url="http://localhost:8080/sse")
|
||||
session = AsyncMock(spec=ClientSession)
|
||||
session.list_tools.side_effect = [
|
||||
ListToolsResult(tools=[Tool(name="safe_tool", description="first tool", inputSchema={})]),
|
||||
ListToolsResult(tools=[Tool(name="safe_tool", description="second tool", inputSchema={})]),
|
||||
]
|
||||
session.list_prompts.side_effect = [
|
||||
types.ListPromptsResult(prompts=[types.Prompt(name="safe_prompt", description="first prompt", arguments=[])]),
|
||||
types.ListPromptsResult(prompts=[types.Prompt(name="safe_prompt", description="second prompt", arguments=[])]),
|
||||
]
|
||||
plugin.session = session
|
||||
|
||||
await plugin.load_tools()
|
||||
first_tool = plugin.safe_tool
|
||||
await plugin.load_tools()
|
||||
|
||||
assert plugin.safe_tool is not first_tool
|
||||
assert plugin.safe_tool.__kernel_function_description__ == "second tool"
|
||||
|
||||
await plugin.load_prompts()
|
||||
first_prompt = plugin.safe_prompt
|
||||
await plugin.load_prompts()
|
||||
|
||||
assert plugin.safe_prompt is not first_prompt
|
||||
assert plugin.safe_prompt.__kernel_function_description__ == "second prompt"
|
||||
|
||||
|
||||
async def test_mcp_plugin_failed_get_session():
|
||||
with (
|
||||
patch("semantic_kernel.connectors.mcp.stdio_client") as mock_stdio_client,
|
||||
):
|
||||
mock_read = MagicMock()
|
||||
mock_write = MagicMock()
|
||||
|
||||
mock_generator = MagicMock()
|
||||
# Make the mock_stdio_client return an AsyncMock for the context manager
|
||||
mock_generator.__aenter__.side_effect = Exception("Connection failed")
|
||||
mock_generator.__aexit__.return_value = (mock_read, mock_write)
|
||||
|
||||
# Make the mock_stdio_client return an AsyncMock for the context manager
|
||||
mock_stdio_client.return_value = mock_generator
|
||||
|
||||
with pytest.raises(KernelPluginInvalidConfigurationError):
|
||||
async with MCPStdioPlugin(
|
||||
name="test",
|
||||
command="echo",
|
||||
args=["Hello"],
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
@patch("semantic_kernel.connectors.mcp.stdio_client")
|
||||
@patch("semantic_kernel.connectors.mcp.ClientSession")
|
||||
async def test_with_kwargs_stdio(mock_session, mock_client, list_tool_calls, kernel: "Kernel"):
|
||||
mock_read = MagicMock()
|
||||
mock_write = MagicMock()
|
||||
|
||||
mock_generator = MagicMock()
|
||||
# Make the mock_stdio_client return an AsyncMock for the context manager
|
||||
mock_generator.__aenter__.return_value = (mock_read, mock_write)
|
||||
mock_generator.__aexit__.return_value = (mock_read, mock_write)
|
||||
|
||||
# Make the mock_stdio_client return an AsyncMock for the context manager
|
||||
mock_client.return_value = mock_generator
|
||||
mock_session.return_value.__aenter__.return_value.list_tools.return_value = list_tool_calls
|
||||
async with MCPStdioPlugin(
|
||||
name="TestMCPPlugin",
|
||||
description="Test MCP Plugin",
|
||||
command="uv",
|
||||
args=["--directory", "path", "run", "file.py"],
|
||||
) as plugin:
|
||||
mock_client.assert_called_once_with(
|
||||
server=StdioServerParameters(command="uv", args=["--directory", "path", "run", "file.py"])
|
||||
)
|
||||
loaded_plugin = kernel.add_plugin(plugin)
|
||||
assert loaded_plugin is not None
|
||||
assert loaded_plugin.name == "TestMCPPlugin"
|
||||
assert loaded_plugin.description == "Test MCP Plugin"
|
||||
assert loaded_plugin.functions.get("func1") is not None
|
||||
assert loaded_plugin.functions["func1"].parameters[0].name == "name"
|
||||
assert loaded_plugin.functions["func1"].parameters[0].is_required
|
||||
assert loaded_plugin.functions.get("func2") is not None
|
||||
assert len(loaded_plugin.functions["func2"].parameters) == 0
|
||||
|
||||
|
||||
@patch("semantic_kernel.connectors.mcp.websocket_client")
|
||||
@patch("semantic_kernel.connectors.mcp.ClientSession")
|
||||
async def test_with_kwargs_websocket(mock_session, mock_client, list_tool_calls, kernel: "Kernel"):
|
||||
mock_read = MagicMock()
|
||||
mock_write = MagicMock()
|
||||
|
||||
mock_generator = MagicMock()
|
||||
# Make the mock_stdio_client return an AsyncMock for the context manager
|
||||
mock_generator.__aenter__.return_value = (mock_read, mock_write)
|
||||
mock_generator.__aexit__.return_value = (mock_read, mock_write)
|
||||
|
||||
# Make the mock_stdio_client return an AsyncMock for the context manager
|
||||
mock_client.return_value = mock_generator
|
||||
mock_session.return_value.__aenter__.return_value.list_tools.return_value = list_tool_calls
|
||||
async with MCPWebsocketPlugin(
|
||||
name="TestMCPPlugin",
|
||||
description="Test MCP Plugin",
|
||||
url="http://localhost:8080/websocket",
|
||||
) as plugin:
|
||||
mock_client.assert_called_once_with(url="http://localhost:8080/websocket")
|
||||
loaded_plugin = kernel.add_plugin(plugin)
|
||||
assert loaded_plugin is not None
|
||||
assert loaded_plugin.name == "TestMCPPlugin"
|
||||
assert loaded_plugin.description == "Test MCP Plugin"
|
||||
assert loaded_plugin.functions.get("func1") is not None
|
||||
assert loaded_plugin.functions["func1"].parameters[0].name == "name"
|
||||
assert loaded_plugin.functions["func1"].parameters[0].is_required
|
||||
assert loaded_plugin.functions.get("func2") is not None
|
||||
assert len(loaded_plugin.functions["func2"].parameters) == 0
|
||||
|
||||
|
||||
@patch("semantic_kernel.connectors.mcp.sse_client")
|
||||
@patch("semantic_kernel.connectors.mcp.ClientSession")
|
||||
async def test_with_kwargs_sse(mock_session, mock_client, list_tool_calls, kernel: "Kernel"):
|
||||
mock_read = MagicMock()
|
||||
mock_write = MagicMock()
|
||||
|
||||
mock_generator = MagicMock()
|
||||
# Make the mock_stdio_client return an AsyncMock for the context manager
|
||||
mock_generator.__aenter__.return_value = (mock_read, mock_write)
|
||||
mock_generator.__aexit__.return_value = (mock_read, mock_write)
|
||||
|
||||
# Make the mock_stdio_client return an AsyncMock for the context manager
|
||||
mock_client.return_value = mock_generator
|
||||
mock_session.return_value.__aenter__.return_value.list_tools.return_value = list_tool_calls
|
||||
async with MCPSsePlugin(
|
||||
name="TestMCPPlugin",
|
||||
description="Test MCP Plugin",
|
||||
url="http://localhost:8080/sse",
|
||||
) as plugin:
|
||||
mock_client.assert_called_once_with(url="http://localhost:8080/sse")
|
||||
loaded_plugin = kernel.add_plugin(plugin)
|
||||
assert loaded_plugin is not None
|
||||
assert loaded_plugin.name == "TestMCPPlugin"
|
||||
assert loaded_plugin.description == "Test MCP Plugin"
|
||||
assert loaded_plugin.functions.get("func1") is not None
|
||||
assert loaded_plugin.functions["func1"].parameters[0].name == "name"
|
||||
assert loaded_plugin.functions["func1"].parameters[0].is_required
|
||||
assert loaded_plugin.functions.get("func2") is not None
|
||||
assert len(loaded_plugin.functions["func2"].parameters) == 0
|
||||
|
||||
|
||||
@patch("semantic_kernel.connectors.mcp.streamablehttp_client")
|
||||
@patch("semantic_kernel.connectors.mcp.ClientSession")
|
||||
async def test_with_kwargs_streamablehttp(mock_session, mock_client, list_tool_calls, kernel: "Kernel"):
|
||||
mock_read = MagicMock()
|
||||
mock_write = MagicMock()
|
||||
mock_callback = MagicMock()
|
||||
|
||||
mock_generator = MagicMock()
|
||||
# Make the mock_streamablehttp_client return an AsyncMock for the context manager
|
||||
mock_generator.__aenter__.return_value = (mock_read, mock_write, mock_callback)
|
||||
mock_generator.__aexit__.return_value = (mock_read, mock_write, mock_callback)
|
||||
|
||||
# Make the mock_streamablehttp_client return an AsyncMock for the context manager
|
||||
mock_client.return_value = mock_generator
|
||||
mock_session.return_value.__aenter__.return_value.list_tools.return_value = list_tool_calls
|
||||
async with MCPStreamableHttpPlugin(
|
||||
name="TestMCPPlugin",
|
||||
description="Test MCP Plugin",
|
||||
url="http://localhost:8080/mcp",
|
||||
) as plugin:
|
||||
mock_client.assert_called_once_with(url="http://localhost:8080/mcp")
|
||||
loaded_plugin = kernel.add_plugin(plugin)
|
||||
assert loaded_plugin is not None
|
||||
assert loaded_plugin.name == "TestMCPPlugin"
|
||||
assert loaded_plugin.description == "Test MCP Plugin"
|
||||
assert loaded_plugin.functions.get("func1") is not None
|
||||
assert loaded_plugin.functions["func1"].parameters[0].name == "name"
|
||||
assert loaded_plugin.functions["func1"].parameters[0].is_required
|
||||
assert loaded_plugin.functions.get("func2") is not None
|
||||
assert len(loaded_plugin.functions["func2"].parameters) == 0
|
||||
|
||||
|
||||
async def test_kernel_as_mcp_server(kernel: "Kernel", decorated_native_function, custom_plugin_class):
|
||||
kernel.add_plugin(custom_plugin_class, "test")
|
||||
kernel.add_functions("test", [decorated_native_function])
|
||||
server = kernel.as_mcp_server()
|
||||
assert server is not None
|
||||
assert types.PingRequest in server.request_handlers
|
||||
assert types.ListToolsRequest in server.request_handlers
|
||||
assert types.CallToolRequest in server.request_handlers
|
||||
assert server.name == "Semantic Kernel MCP Server"
|
||||
|
||||
|
||||
@patch("semantic_kernel.connectors.mcp.sse_client")
|
||||
@patch("semantic_kernel.connectors.mcp.ClientSession")
|
||||
async def test_mcp_tool_name_normalization(mock_session, mock_client, list_tool_calls_with_slash, kernel: "Kernel"):
|
||||
"""Test that MCP tool names with illegal characters are normalized."""
|
||||
mock_read = MagicMock()
|
||||
mock_write = MagicMock()
|
||||
mock_generator = MagicMock()
|
||||
mock_generator.__aenter__.return_value = (mock_read, mock_write)
|
||||
mock_generator.__aexit__.return_value = (mock_read, mock_write)
|
||||
mock_client.return_value = mock_generator
|
||||
mock_session.return_value.__aenter__.return_value.list_tools.return_value = list_tool_calls_with_slash
|
||||
|
||||
async with MCPSsePlugin(
|
||||
name="TestMCPPlugin",
|
||||
description="Test MCP Plugin",
|
||||
url="http://localhost:8080/sse",
|
||||
) as plugin:
|
||||
loaded_plugin = kernel.add_plugin(plugin)
|
||||
# The normalized names:
|
||||
assert "nasa-get-astronomy-picture" in loaded_plugin.functions
|
||||
assert "weird-name-with-spaces" in loaded_plugin.functions
|
||||
# They should not exist with their original (invalid) names:
|
||||
assert "nasa/get-astronomy-picture" not in loaded_plugin.functions
|
||||
assert "weird\\name with spaces" not in loaded_plugin.functions
|
||||
|
||||
normalized_names = list(loaded_plugin.functions.keys())
|
||||
for name in normalized_names:
|
||||
assert re.match(r"^[A-Za-z0-9_.-]+$", name)
|
||||
|
||||
|
||||
@patch("semantic_kernel.connectors.mcp.ClientSession")
|
||||
async def test_mcp_normalization_function(mock_session, list_tool_calls_with_slash):
|
||||
"""Unit test for the normalize_mcp_name function (should exist in codebase)."""
|
||||
from semantic_kernel.connectors.mcp import _normalize_mcp_name
|
||||
|
||||
assert _normalize_mcp_name("nasa/get-astronomy-picture") == "nasa-get-astronomy-picture"
|
||||
assert _normalize_mcp_name("weird\\name with spaces") == "weird-name-with-spaces"
|
||||
assert _normalize_mcp_name("simple_name") == "simple_name"
|
||||
assert _normalize_mcp_name("Name-With.Dots_And-Hyphens") == "Name-With.Dots_And-Hyphens"
|
||||
|
||||
|
||||
async def test_excluded_function_cannot_be_called(kernel: "Kernel"):
|
||||
"""Test that excluded functions are rejected at call time, not just hidden from listing."""
|
||||
from semantic_kernel.connectors.mcp import create_mcp_server_from_kernel
|
||||
from semantic_kernel.functions.kernel_function_decorator import kernel_function
|
||||
|
||||
side_effect_called = False
|
||||
|
||||
@kernel_function(name="public_echo")
|
||||
def public_echo(message: str) -> str:
|
||||
return f"echo: {message}"
|
||||
|
||||
@kernel_function(name="secret_admin")
|
||||
def secret_admin(target: str) -> str:
|
||||
nonlocal side_effect_called
|
||||
side_effect_called = True
|
||||
return f"privileged action on {target}"
|
||||
|
||||
kernel.add_function(plugin_name="tools", function=public_echo)
|
||||
kernel.add_function(plugin_name="tools", function=secret_admin)
|
||||
|
||||
server = create_mcp_server_from_kernel(kernel, excluded_functions=["secret_admin"])
|
||||
|
||||
# Verify the server was created with handlers
|
||||
assert types.ListToolsRequest in server.request_handlers
|
||||
assert types.CallToolRequest in server.request_handlers
|
||||
|
||||
# Mock _get_cached_tool_definition to bypass SDK request context requirements
|
||||
# (normally set by a real MCP session transport)
|
||||
async def _fake_get_cached_tool_definition(tool_name):
|
||||
return None
|
||||
|
||||
server._get_cached_tool_definition = _fake_get_cached_tool_definition
|
||||
|
||||
# Build a proper CallToolRequest as the MCP SDK would send
|
||||
call_tool_request = types.CallToolRequest(
|
||||
method="tools/call",
|
||||
params=types.CallToolRequestParams(name="secret_admin", arguments={}),
|
||||
)
|
||||
|
||||
# The internal handler wraps our _call_tool; invoke via the registered handler
|
||||
handler = server.request_handlers[types.CallToolRequest]
|
||||
result = await handler(call_tool_request)
|
||||
|
||||
# The call must fail (isError=True) with the correct error message
|
||||
assert result.root.isError is True, "Calling an excluded function should return an error"
|
||||
assert any("Unknown tool" in c.text for c in result.root.content if hasattr(c, "text")), (
|
||||
f"Expected 'Unknown tool' error, got: {result.root.content}"
|
||||
)
|
||||
assert not side_effect_called, "Excluded function's side effect should not have fired"
|
||||
@@ -0,0 +1,86 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def database_name():
|
||||
"""Fixture for the database name."""
|
||||
return "test_database"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def collection_name():
|
||||
"""Fixture for the collection name."""
|
||||
return "test_collection"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def url():
|
||||
"""Fixture for the url."""
|
||||
return "https://test.cosmos.azure.com/"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def key():
|
||||
"""Fixture for the key."""
|
||||
return "test_key"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def azure_cosmos_db_mongo_db_unit_test_env(monkeypatch, url, key, database_name, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for Azure Cosmos DB NoSQL unit tests."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"AZURE_COSMOS_DB_MONGODB_CONNECTION_STRING": url,
|
||||
"AZURE_COSMOS_DB_MONGODB_DATABASE_NAME": database_name,
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def azure_cosmos_db_no_sql_unit_test_env(monkeypatch, url, key, database_name, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for Azure Cosmos DB NoSQL unit tests."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"AZURE_COSMOS_DB_NO_SQL_URL": url,
|
||||
"AZURE_COSMOS_DB_NO_SQL_KEY": key,
|
||||
"AZURE_COSMOS_DB_NO_SQL_DATABASE_NAME": database_name,
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def clear_azure_cosmos_db_no_sql_env(monkeypatch):
|
||||
"""Fixture to clear the environment variables for Weaviate unit tests."""
|
||||
monkeypatch.delenv("AZURE_COSMOS_DB_NO_SQL_URL", raising=False)
|
||||
monkeypatch.delenv("AZURE_COSMOS_DB_NO_SQL_KEY", raising=False)
|
||||
monkeypatch.delenv("AZURE_COSMOS_DB_NO_SQL_DATABASE_NAME", raising=False)
|
||||
+156
@@ -0,0 +1,156 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from pymongo import AsyncMongoClient
|
||||
|
||||
from semantic_kernel.connectors.azure_cosmos_db import CosmosMongoCollection
|
||||
from semantic_kernel.data.vector import VectorStoreCollectionDefinition, VectorStoreField
|
||||
from semantic_kernel.exceptions import VectorStoreInitializationException
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_model() -> VectorStoreCollectionDefinition:
|
||||
return VectorStoreCollectionDefinition(
|
||||
fields=[
|
||||
VectorStoreField("key", name="id"),
|
||||
VectorStoreField("data", name="content"),
|
||||
VectorStoreField("vector", name="vector", dimensions=5),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
async def test_constructor_with_mongo_client_provided(mock_model) -> None:
|
||||
"""
|
||||
Test the constructor of AzureCosmosDBforMongoDBCollection when a mongo_client
|
||||
is directly provided. Expect that the class is successfully initialized
|
||||
and doesn't attempt to manage the client.
|
||||
"""
|
||||
mock_client = AsyncMock(spec=AsyncMongoClient)
|
||||
collection_name = "test_collection"
|
||||
|
||||
collection = CosmosMongoCollection(
|
||||
collection_name=collection_name,
|
||||
record_type=dict,
|
||||
mongo_client=mock_client,
|
||||
definition=mock_model,
|
||||
)
|
||||
|
||||
assert collection.mongo_client == mock_client
|
||||
assert collection.collection_name == collection_name
|
||||
assert not collection.managed_client, "Should not be managing client when provided"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_COSMOS_DB_MONGODB_CONNECTION_STRING"]], indirect=True)
|
||||
async def test_constructor_raises_exception_on_validation_error(
|
||||
azure_cosmos_db_mongo_db_unit_test_env, definition
|
||||
) -> None:
|
||||
"""
|
||||
Test that the constructor raises VectorStoreInitializationException when
|
||||
AzureCosmosDBforMongoDBSettings fails with ValidationError.
|
||||
"""
|
||||
with pytest.raises(VectorStoreInitializationException) as exc_info:
|
||||
CosmosMongoCollection(
|
||||
collection_name="test_collection",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
database_name="",
|
||||
env_file_path=".no.env",
|
||||
)
|
||||
assert "The Azure CosmosDB for MongoDB connection string is required." in str(exc_info.value)
|
||||
|
||||
|
||||
async def test_ensure_collection_exists_calls_database_methods(definition) -> None:
|
||||
"""
|
||||
Test ensure_collection_exists to verify that it first creates a collection, then
|
||||
calls the appropriate command to create a vector index.
|
||||
"""
|
||||
# Setup
|
||||
mock_database = AsyncMock()
|
||||
mock_database.create_collection = AsyncMock()
|
||||
mock_database.command = AsyncMock()
|
||||
|
||||
mock_client = AsyncMock(spec=AsyncMongoClient)
|
||||
mock_client.get_database = MagicMock(return_value=mock_database)
|
||||
|
||||
# Instantiate
|
||||
collection = CosmosMongoCollection(
|
||||
collection_name="test_collection",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
mongo_client=mock_client,
|
||||
database_name="test_db",
|
||||
)
|
||||
|
||||
# Act
|
||||
await collection.ensure_collection_exists(customArg="customValue")
|
||||
|
||||
# Assert
|
||||
mock_database.create_collection.assert_awaited_once_with("test_collection", customArg="customValue")
|
||||
mock_database.command.assert_awaited()
|
||||
command_args = mock_database.command.call_args.kwargs["command"]
|
||||
|
||||
assert command_args["createIndexes"] == "test_collection"
|
||||
assert len(command_args["indexes"]) == 2, "One for the data field, one for the vector field"
|
||||
# Check the data field index
|
||||
assert command_args["indexes"][0]["name"] == "content_"
|
||||
# Check the vector field index creation
|
||||
assert command_args["indexes"][1]["name"] == "vector_"
|
||||
assert command_args["indexes"][1]["key"] == {"vector": "cosmosSearch"}
|
||||
assert command_args["indexes"][1]["cosmosSearchOptions"]["kind"] == "COS"
|
||||
assert command_args["indexes"][1]["cosmosSearchOptions"]["similarity"] is not None
|
||||
assert command_args["indexes"][1]["cosmosSearchOptions"]["dimensions"] == 5
|
||||
|
||||
|
||||
async def test_context_manager_calls_aconnect_and_close_when_managed(mock_model) -> None:
|
||||
"""
|
||||
Test that the context manager in AzureCosmosDBforMongoDBCollection calls 'aconnect' and
|
||||
'close' when the client is managed (i.e., created internally).
|
||||
"""
|
||||
mock_client = AsyncMock(spec=AsyncMongoClient)
|
||||
|
||||
with patch(
|
||||
"semantic_kernel.connectors.azure_cosmos_db.AsyncMongoClient",
|
||||
return_value=mock_client,
|
||||
):
|
||||
collection = CosmosMongoCollection(
|
||||
collection_name="test_collection",
|
||||
record_type=dict,
|
||||
connection_string="mongodb://fake",
|
||||
definition=mock_model,
|
||||
)
|
||||
|
||||
# "__aenter__" should call 'aconnect'
|
||||
async with collection as c:
|
||||
mock_client.aconnect.assert_awaited_once()
|
||||
assert c is collection
|
||||
|
||||
# "__aexit__" should call 'close' if managed
|
||||
mock_client.close.assert_awaited_once()
|
||||
|
||||
|
||||
async def test_context_manager_does_not_close_when_not_managed(mock_model) -> None:
|
||||
"""
|
||||
Test that the context manager in AzureCosmosDBforMongoDBCollection does not call 'close'
|
||||
when the client is not managed (i.e., provided externally).
|
||||
"""
|
||||
|
||||
external_client = AsyncMock(spec=AsyncMongoClient, name="external_client", value=None)
|
||||
external_client.aconnect = AsyncMock(name="aconnect")
|
||||
external_client.close = AsyncMock(name="close")
|
||||
|
||||
collection = CosmosMongoCollection(
|
||||
collection_name="test_collection",
|
||||
record_type=dict,
|
||||
mongo_client=external_client,
|
||||
definition=mock_model,
|
||||
)
|
||||
|
||||
# "__aenter__" scenario
|
||||
async with collection as c:
|
||||
external_client.aconnect.assert_awaited()
|
||||
assert c is collection
|
||||
|
||||
# "__aexit__" should NOT call "close" when not managed
|
||||
external_client.close.assert_not_awaited()
|
||||
+554
@@ -0,0 +1,554 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from azure.core.credentials_async import AsyncTokenCredential
|
||||
from azure.cosmos.aio import CosmosClient
|
||||
from azure.cosmos.exceptions import CosmosHttpResponseError, CosmosResourceNotFoundError
|
||||
|
||||
from semantic_kernel.connectors.azure_cosmos_db import (
|
||||
COSMOS_ITEM_ID_PROPERTY_NAME,
|
||||
CosmosNoSqlCollection,
|
||||
_create_default_indexing_policy_nosql,
|
||||
_create_default_vector_embedding_policy,
|
||||
)
|
||||
from semantic_kernel.data._shared import default_dynamic_filter_function
|
||||
from semantic_kernel.exceptions import (
|
||||
VectorStoreInitializationException,
|
||||
VectorStoreModelException,
|
||||
VectorStoreOperationException,
|
||||
)
|
||||
from semantic_kernel.functions import KernelParameterMetadata
|
||||
|
||||
|
||||
def test_azure_cosmos_db_no_sql_collection_init(
|
||||
clear_azure_cosmos_db_no_sql_env,
|
||||
record_type,
|
||||
database_name: str,
|
||||
collection_name: str,
|
||||
url: str,
|
||||
key: str,
|
||||
) -> None:
|
||||
"""Test the initialization of an AzureCosmosDBNoSQLCollection object."""
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
database_name=database_name,
|
||||
url=url,
|
||||
key=key,
|
||||
)
|
||||
|
||||
assert vector_collection is not None
|
||||
assert vector_collection.database_name == database_name
|
||||
assert vector_collection.collection_name == collection_name
|
||||
assert vector_collection.cosmos_client is not None
|
||||
assert vector_collection.partition_key.path == f"/{vector_collection.definition.key_name}"
|
||||
assert vector_collection.create_database is False
|
||||
|
||||
|
||||
def test_azure_cosmos_db_no_sql_collection_init_env(
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test the initialization of an AzureCosmosDBNoSQLCollection object with environment variables."""
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
|
||||
assert vector_collection is not None
|
||||
assert (
|
||||
vector_collection.database_name == azure_cosmos_db_no_sql_unit_test_env["AZURE_COSMOS_DB_NO_SQL_DATABASE_NAME"]
|
||||
)
|
||||
assert vector_collection.collection_name == collection_name
|
||||
assert vector_collection.partition_key.path == f"/{vector_collection.definition.key_name}"
|
||||
assert vector_collection.create_database is False
|
||||
|
||||
|
||||
def test_azure_cosmos_db_no_sql_build_filter_escapes_apostrophes(
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test Cosmos DB filter building escapes apostrophes in string literals."""
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
|
||||
filter_string = vector_collection._build_filter('lambda x: x.content == "O\'Reilly"')
|
||||
|
||||
assert filter_string == "c.content = 'O''Reilly'"
|
||||
|
||||
|
||||
def test_azure_cosmos_db_no_sql_build_filter_escapes_injection_payload(
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test Cosmos DB filter building keeps injection-shaped strings inside the literal."""
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
|
||||
filter_string = vector_collection._build_filter("lambda x: x.content == \"test' OR '1'='1\"")
|
||||
|
||||
assert filter_string == "c.content = 'test'' OR ''1''=''1'"
|
||||
|
||||
|
||||
def test_azure_cosmos_db_no_sql_dynamic_filter_injection_payload_stays_literal(
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test default_dynamic_filter_function does not let user values alter Cosmos filter syntax."""
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
generated_filter = default_dynamic_filter_function(
|
||||
filter=None,
|
||||
parameters=[
|
||||
KernelParameterMetadata(
|
||||
name="content",
|
||||
description="Content filter",
|
||||
type="str",
|
||||
is_required=False,
|
||||
type_object=str,
|
||||
)
|
||||
],
|
||||
content="test' OR '1'='1",
|
||||
)
|
||||
|
||||
assert isinstance(generated_filter, str)
|
||||
assert vector_collection._build_filter(generated_filter) == "c.content = 'test'' OR ''1''=''1'"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_COSMOS_DB_NO_SQL_URL"]], indirect=True)
|
||||
def test_azure_cosmos_db_no_sql_collection_init_no_url(
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test the initialization of an AzureCosmosDBNoSQLCollection object with missing URL."""
|
||||
with pytest.raises(VectorStoreInitializationException):
|
||||
CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
env_file_path="fake_path",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_COSMOS_DB_NO_SQL_DATABASE_NAME"]], indirect=True)
|
||||
def test_azure_cosmos_db_no_sql_collection_init_no_database_name(
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test the initialization of an AzureCosmosDBNoSQLCollection object with missing database name."""
|
||||
with pytest.raises(
|
||||
VectorStoreInitializationException, match="The name of the Azure Cosmos DB NoSQL database is missing."
|
||||
):
|
||||
CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
env_file_path="fake_path",
|
||||
)
|
||||
|
||||
|
||||
def test_azure_cosmos_db_no_sql_collection_invalid_settings(
|
||||
clear_azure_cosmos_db_no_sql_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test the initialization of an AzureCosmosDBNoSQLCollection object with invalid settings."""
|
||||
with pytest.raises(VectorStoreInitializationException):
|
||||
CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
url="invalid_url",
|
||||
)
|
||||
|
||||
|
||||
@patch.object(CosmosClient, "__init__", return_value=None)
|
||||
def test_azure_cosmos_db_no_sql_get_cosmos_client(
|
||||
mock_cosmos_client_init,
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test the creation of a cosmos client."""
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
|
||||
assert vector_collection.cosmos_client is not None
|
||||
mock_cosmos_client_init.assert_called_once_with(
|
||||
str(azure_cosmos_db_no_sql_unit_test_env["AZURE_COSMOS_DB_NO_SQL_URL"]),
|
||||
credential=azure_cosmos_db_no_sql_unit_test_env["AZURE_COSMOS_DB_NO_SQL_KEY"],
|
||||
)
|
||||
|
||||
|
||||
@patch.object(CosmosClient, "__init__", return_value=None)
|
||||
def test_azure_cosmos_db_no_sql_get_cosmos_client_without_key(
|
||||
mock_cosmos_client_init,
|
||||
clear_azure_cosmos_db_no_sql_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
database_name: str,
|
||||
url: str,
|
||||
) -> None:
|
||||
"""Test the creation of a cosmos client."""
|
||||
credential = AsyncMock(spec=AsyncTokenCredential)
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
database_name=database_name,
|
||||
url=url,
|
||||
credential=credential,
|
||||
)
|
||||
|
||||
assert vector_collection.cosmos_client is not None
|
||||
mock_cosmos_client_init.assert_called_once_with(url, credential=ANY)
|
||||
|
||||
|
||||
@patch("azure.cosmos.aio.CosmosClient", spec=True)
|
||||
async def test_azure_cosmos_db_no_sql_collection_create_database_if_not_exists(
|
||||
mock_cosmos_client,
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test the creation of a cosmos DB NoSQL database if it does not exist when create_database=True."""
|
||||
mock_cosmos_client.get_database_client.side_effect = CosmosResourceNotFoundError
|
||||
mock_cosmos_client.create_database = AsyncMock()
|
||||
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
cosmos_client=mock_cosmos_client,
|
||||
create_database=True,
|
||||
)
|
||||
|
||||
assert vector_collection.create_database is True
|
||||
|
||||
await vector_collection._get_database_proxy()
|
||||
|
||||
mock_cosmos_client.get_database_client.assert_called_once_with(
|
||||
azure_cosmos_db_no_sql_unit_test_env["AZURE_COSMOS_DB_NO_SQL_DATABASE_NAME"]
|
||||
)
|
||||
mock_cosmos_client.create_database.assert_called_once_with(
|
||||
azure_cosmos_db_no_sql_unit_test_env["AZURE_COSMOS_DB_NO_SQL_DATABASE_NAME"]
|
||||
)
|
||||
|
||||
|
||||
@patch("azure.cosmos.aio.CosmosClient", spec=True)
|
||||
async def test_azure_cosmos_db_no_sql_collection_create_database_raise_if_database_not_exists(
|
||||
mock_cosmos_client,
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test _get_database_proxy raises an error if the database does not exist when create_database=False."""
|
||||
mock_cosmos_client.get_database_client.side_effect = CosmosResourceNotFoundError
|
||||
mock_cosmos_client.create_database = AsyncMock()
|
||||
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
cosmos_client=mock_cosmos_client,
|
||||
create_database=False,
|
||||
)
|
||||
|
||||
assert vector_collection.create_database is False
|
||||
|
||||
with pytest.raises(VectorStoreOperationException):
|
||||
await vector_collection._get_database_proxy()
|
||||
|
||||
|
||||
@patch("azure.cosmos.aio.CosmosClient")
|
||||
@patch("azure.cosmos.aio.DatabaseProxy")
|
||||
@pytest.mark.parametrize("index_kind, distance_function", [("flat", "cosine_similarity")])
|
||||
async def test_azure_cosmos_db_no_sql_collection_ensure_collection_exists(
|
||||
mock_database_proxy,
|
||||
mock_cosmos_client,
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
):
|
||||
"""Test the creation of a cosmos DB NoSQL collection."""
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
|
||||
vector_collection._get_database_proxy = AsyncMock(return_value=mock_database_proxy)
|
||||
|
||||
mock_database_proxy.create_container_if_not_exists = AsyncMock(return_value=None)
|
||||
|
||||
await vector_collection.ensure_collection_exists()
|
||||
|
||||
mock_database_proxy.create_container_if_not_exists.assert_called_once_with(
|
||||
id=collection_name,
|
||||
partition_key=vector_collection.partition_key,
|
||||
indexing_policy=_create_default_indexing_policy_nosql(vector_collection.definition),
|
||||
vector_embedding_policy=_create_default_vector_embedding_policy(vector_collection.definition),
|
||||
)
|
||||
|
||||
|
||||
@patch("azure.cosmos.aio.CosmosClient")
|
||||
@patch("azure.cosmos.aio.DatabaseProxy")
|
||||
@pytest.mark.parametrize("index_kind, distance_function", [("flat", "cosine_similarity")])
|
||||
async def test_azure_cosmos_db_no_sql_collection_ensure_collection_exists_allow_custom_indexing_policy(
|
||||
mock_database_proxy,
|
||||
mock_cosmos_client,
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
):
|
||||
"""Test the creation of a cosmos DB NoSQL collection with a custom indexing policy."""
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
|
||||
vector_collection._get_database_proxy = AsyncMock(return_value=mock_database_proxy)
|
||||
|
||||
mock_database_proxy.create_container_if_not_exists = AsyncMock(return_value=None)
|
||||
|
||||
await vector_collection.ensure_collection_exists(indexing_policy={"automatic": False})
|
||||
|
||||
mock_database_proxy.create_container_if_not_exists.assert_called_once_with(
|
||||
id=collection_name,
|
||||
partition_key=vector_collection.partition_key,
|
||||
indexing_policy={"automatic": False},
|
||||
vector_embedding_policy=_create_default_vector_embedding_policy(vector_collection.definition),
|
||||
)
|
||||
|
||||
|
||||
@patch("azure.cosmos.aio.CosmosClient")
|
||||
@patch("azure.cosmos.aio.DatabaseProxy")
|
||||
@pytest.mark.parametrize("index_kind, distance_function", [("flat", "cosine_similarity")])
|
||||
async def test_azure_cosmos_db_no_sql_collection_ensure_collection_exists_allow_custom_vector_embedding_policy(
|
||||
mock_database_proxy,
|
||||
mock_cosmos_client,
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
):
|
||||
"""Test the creation of a cosmos DB NoSQL collection with a custom vector embedding policy."""
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
|
||||
vector_collection._get_database_proxy = AsyncMock(return_value=mock_database_proxy)
|
||||
|
||||
mock_database_proxy.create_container_if_not_exists = AsyncMock(return_value=None)
|
||||
|
||||
await vector_collection.ensure_collection_exists(vector_embedding_policy={"vectorEmbeddings": []})
|
||||
|
||||
mock_database_proxy.create_container_if_not_exists.assert_called_once_with(
|
||||
id=collection_name,
|
||||
partition_key=vector_collection.partition_key,
|
||||
indexing_policy=_create_default_indexing_policy_nosql(vector_collection.definition),
|
||||
vector_embedding_policy={"vectorEmbeddings": []},
|
||||
)
|
||||
|
||||
|
||||
@patch("azure.cosmos.aio.CosmosClient")
|
||||
@patch("azure.cosmos.aio.DatabaseProxy")
|
||||
@pytest.mark.parametrize(
|
||||
"index_kind, distance_function, vector_property_type",
|
||||
[
|
||||
("hnsw", "cosine_similarity", "float"), # unsupported index kind
|
||||
("flat", "hamming", "float"), # unsupported distance function
|
||||
],
|
||||
)
|
||||
async def test_azure_cosmos_db_no_sql_collection_ensure_collection_exists_unsupported_vector_field_property(
|
||||
mock_database_proxy,
|
||||
mock_cosmos_client,
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
):
|
||||
"""Test the creation of a cosmos DB NoSQL collection with an unsupported index kind."""
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
|
||||
vector_collection._get_database_proxy = AsyncMock(return_value=mock_database_proxy)
|
||||
|
||||
mock_database_proxy.create_container_if_not_exists = AsyncMock(return_value=None)
|
||||
|
||||
with pytest.raises(VectorStoreModelException):
|
||||
await vector_collection.ensure_collection_exists()
|
||||
|
||||
|
||||
@patch("azure.cosmos.aio.DatabaseProxy")
|
||||
async def test_azure_cosmos_db_no_sql_collection_ensure_collection_deleted(
|
||||
mock_database_proxy,
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test the deletion of a cosmos DB NoSQL collection."""
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
|
||||
vector_collection._get_database_proxy = AsyncMock(return_value=mock_database_proxy)
|
||||
|
||||
mock_database_proxy.delete_container = AsyncMock()
|
||||
|
||||
await vector_collection.ensure_collection_deleted()
|
||||
|
||||
mock_database_proxy.delete_container.assert_called_once_with(collection_name)
|
||||
|
||||
|
||||
@patch("azure.cosmos.aio.DatabaseProxy")
|
||||
async def test_azure_cosmos_db_no_sql_collection_ensure_collection_deleted_fail(
|
||||
mock_database_proxy,
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test the deletion of a cosmos DB NoSQL collection that does not exist."""
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
|
||||
vector_collection._get_database_proxy = AsyncMock(return_value=mock_database_proxy)
|
||||
mock_database_proxy.delete_container = AsyncMock(side_effect=CosmosHttpResponseError)
|
||||
|
||||
with pytest.raises(VectorStoreOperationException, match="Container could not be deleted."):
|
||||
await vector_collection.ensure_collection_deleted()
|
||||
|
||||
|
||||
@patch("azure.cosmos.aio.ContainerProxy")
|
||||
async def test_azure_cosmos_db_no_sql_upsert(
|
||||
mock_container_proxy,
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test the upsert of a document in a cosmos DB NoSQL collection."""
|
||||
item = {"content": "test_content", "vector": [1.0, 2.0, 3.0], "id": "test_id"}
|
||||
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
|
||||
vector_collection._get_container_proxy = AsyncMock(return_value=mock_container_proxy)
|
||||
|
||||
mock_container_proxy.upsert_item = AsyncMock(return_value={COSMOS_ITEM_ID_PROPERTY_NAME: item["id"]})
|
||||
|
||||
result = await vector_collection.upsert(item)
|
||||
|
||||
mock_container_proxy.upsert_item.assert_called_once_with(item)
|
||||
assert result == item["id"]
|
||||
|
||||
|
||||
@patch("azure.cosmos.aio.ContainerProxy")
|
||||
async def test_azure_cosmos_db_no_sql_upsert_without_id(
|
||||
mock_container_proxy,
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type_with_key_as_key_field,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test the upsert of a document in a cosmos DB NoSQL collection where the name of the key field is 'key'."""
|
||||
item = {"content": "test_content", "vector": [1.0, 2.0, 3.0], "key": "test_key"}
|
||||
item_with_id = {"content": "test_content", "vector": [1.0, 2.0, 3.0], COSMOS_ITEM_ID_PROPERTY_NAME: "test_key"}
|
||||
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type_with_key_as_key_field,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
|
||||
vector_collection._get_container_proxy = AsyncMock(return_value=mock_container_proxy)
|
||||
|
||||
mock_container_proxy.upsert_item = AsyncMock(return_value={COSMOS_ITEM_ID_PROPERTY_NAME: item["key"]})
|
||||
|
||||
result = await vector_collection.upsert(item)
|
||||
|
||||
mock_container_proxy.upsert_item.assert_called_once_with(item_with_id)
|
||||
assert result == item["key"]
|
||||
|
||||
|
||||
@patch("azure.cosmos.aio.ContainerProxy")
|
||||
async def test_azure_cosmos_db_no_sql_get(
|
||||
mock_container_proxy,
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test the retrieval of a document from a cosmos DB NoSQL collection."""
|
||||
vector_collection: CosmosNoSqlCollection[str, record_type] = CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
|
||||
vector_collection._get_container_proxy = AsyncMock(return_value=mock_container_proxy)
|
||||
|
||||
get_results = MagicMock(spec=AsyncGenerator)
|
||||
get_results.__aiter__.return_value = [{"content": "test_content", "id": "test_id"}]
|
||||
mock_container_proxy.query_items.return_value = get_results
|
||||
|
||||
record = await vector_collection.get("test_id")
|
||||
assert isinstance(record, record_type)
|
||||
assert record.content == "test_content"
|
||||
assert record.id == "test_id"
|
||||
|
||||
|
||||
@patch("azure.cosmos.aio.ContainerProxy")
|
||||
async def test_azure_cosmos_db_no_sql_get_without_id(
|
||||
mock_container_proxy,
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type_with_key_as_key_field,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test the retrieval of a document from a cosmos DB NoSQL collection where the name of the key field is 'key'."""
|
||||
vector_collection = CosmosNoSqlCollection(
|
||||
record_type=record_type_with_key_as_key_field,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
|
||||
vector_collection._get_container_proxy = AsyncMock(return_value=mock_container_proxy)
|
||||
|
||||
get_results = MagicMock(spec=AsyncGenerator)
|
||||
get_results.__aiter__.return_value = [
|
||||
{"content": "test_content", "vector": [1.0, 2.0, 3.0], COSMOS_ITEM_ID_PROPERTY_NAME: "test_key"}
|
||||
]
|
||||
mock_container_proxy.query_items.return_value = get_results
|
||||
|
||||
record = await vector_collection.get("test_key")
|
||||
assert isinstance(record, record_type_with_key_as_key_field)
|
||||
assert record.content == "test_content"
|
||||
assert record.vector == [1.0, 2.0, 3.0]
|
||||
assert record.key == "test_key"
|
||||
|
||||
|
||||
@patch.object(CosmosClient, "close", return_value=None)
|
||||
async def test_client_is_closed(
|
||||
mock_cosmos_client_close,
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
record_type,
|
||||
collection_name: str,
|
||||
) -> None:
|
||||
"""Test the close method of an AzureCosmosDBNoSQLCollection object."""
|
||||
async with CosmosNoSqlCollection(
|
||||
record_type=record_type,
|
||||
collection_name=collection_name,
|
||||
) as collection:
|
||||
assert collection.cosmos_client is not None
|
||||
|
||||
mock_cosmos_client_close.assert_called()
|
||||
+101
@@ -0,0 +1,101 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from azure.cosmos.aio import CosmosClient
|
||||
|
||||
from semantic_kernel.connectors.azure_cosmos_db import CosmosNoSqlCollection, CosmosNoSqlStore
|
||||
from semantic_kernel.exceptions import VectorStoreInitializationException
|
||||
|
||||
|
||||
def test_azure_cosmos_db_no_sql_store_init(
|
||||
clear_azure_cosmos_db_no_sql_env,
|
||||
database_name: str,
|
||||
url: str,
|
||||
key: str,
|
||||
) -> None:
|
||||
"""Test the initialization of an AzureCosmosDBNoSQLStore object."""
|
||||
vector_store = CosmosNoSqlStore(url=url, key=key, database_name=database_name)
|
||||
|
||||
assert vector_store is not None
|
||||
assert vector_store.database_name == database_name
|
||||
assert vector_store.cosmos_client is not None
|
||||
assert vector_store.create_database is False
|
||||
|
||||
|
||||
def test_azure_cosmos_db_no_sql_store_init_env(azure_cosmos_db_no_sql_unit_test_env) -> None:
|
||||
"""Test the initialization of an AzureCosmosDBNoSQLStore object with environment variables."""
|
||||
vector_store = CosmosNoSqlStore()
|
||||
|
||||
assert vector_store is not None
|
||||
assert vector_store.database_name == azure_cosmos_db_no_sql_unit_test_env["AZURE_COSMOS_DB_NO_SQL_DATABASE_NAME"]
|
||||
assert vector_store.cosmos_client is not None
|
||||
assert vector_store.create_database is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_COSMOS_DB_NO_SQL_URL"]], indirect=True)
|
||||
def test_azure_cosmos_db_no_sql_store_init_no_url(
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
) -> None:
|
||||
"""Test the initialization of an AzureCosmosDBNoSQLStore object with missing URL."""
|
||||
with pytest.raises(VectorStoreInitializationException):
|
||||
CosmosNoSqlStore(env_file_path="fake_path")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["AZURE_COSMOS_DB_NO_SQL_DATABASE_NAME"]], indirect=True)
|
||||
def test_azure_cosmos_db_no_sql_store_init_no_database_name(
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
) -> None:
|
||||
"""Test the initialization of an AzureCosmosDBNoSQLStore object with missing database name."""
|
||||
with pytest.raises(
|
||||
VectorStoreInitializationException, match="The name of the Azure Cosmos DB NoSQL database is missing."
|
||||
):
|
||||
CosmosNoSqlStore(env_file_path="fake_path")
|
||||
|
||||
|
||||
def test_azure_cosmos_db_no_sql_store_invalid_settings(
|
||||
clear_azure_cosmos_db_no_sql_env,
|
||||
) -> None:
|
||||
"""Test the initialization of an AzureCosmosDBNoSQLStore object with invalid settings."""
|
||||
with pytest.raises(VectorStoreInitializationException, match="Failed to validate Azure Cosmos DB NoSQL settings."):
|
||||
CosmosNoSqlStore(url="invalid_url")
|
||||
|
||||
|
||||
@patch.object(CosmosNoSqlCollection, "__init__", return_value=None)
|
||||
def test_azure_cosmos_db_no_sql_store_get_collection(
|
||||
mock_azure_cosmos_db_no_sql_collection_init,
|
||||
azure_cosmos_db_no_sql_unit_test_env,
|
||||
collection_name: str,
|
||||
record_type,
|
||||
) -> None:
|
||||
"""Test the get_collection method of an AzureCosmosDBNoSQLStore object."""
|
||||
vector_store = CosmosNoSqlStore()
|
||||
|
||||
collection = vector_store.get_collection(collection_name=collection_name, record_type=record_type)
|
||||
|
||||
assert collection is not None
|
||||
mock_azure_cosmos_db_no_sql_collection_init.assert_called_once_with(
|
||||
record_type=record_type,
|
||||
definition=None,
|
||||
collection_name=collection_name,
|
||||
database_name=azure_cosmos_db_no_sql_unit_test_env["AZURE_COSMOS_DB_NO_SQL_DATABASE_NAME"],
|
||||
embedding_generator=None,
|
||||
url=azure_cosmos_db_no_sql_unit_test_env["AZURE_COSMOS_DB_NO_SQL_URL"],
|
||||
key=azure_cosmos_db_no_sql_unit_test_env["AZURE_COSMOS_DB_NO_SQL_KEY"],
|
||||
cosmos_client=vector_store.cosmos_client,
|
||||
partition_key=None,
|
||||
create_database=vector_store.create_database,
|
||||
env_file_path=None,
|
||||
env_file_encoding=None,
|
||||
)
|
||||
|
||||
|
||||
@patch.object(CosmosClient, "close", return_value=None)
|
||||
async def test_client_is_closed(mock_cosmos_client_close, azure_cosmos_db_no_sql_unit_test_env) -> None:
|
||||
"""Test the close method of an AzureCosmosDBNoSQLStore object."""
|
||||
async with CosmosNoSqlStore() as vector_store:
|
||||
assert vector_store.cosmos_client is not None
|
||||
|
||||
mock_cosmos_client_close.assert_called()
|
||||
@@ -0,0 +1,308 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from _pytest.mark.structures import ParameterSet
|
||||
from pytest import fixture, param
|
||||
|
||||
from semantic_kernel.exceptions.vector_store_exceptions import VectorStoreOperationException
|
||||
|
||||
|
||||
@fixture()
|
||||
def mongodb_atlas_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for MongoDB Atlas Unit Tests."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {"MONGODB_ATLAS_CONNECTION_STRING": "mongodb://test", "MONGODB_ATLAS_DATABASE_NAME": "test-database"}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@fixture
|
||||
def postgres_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for Postgres connector."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {"POSTGRES_CONNECTION_STRING": "host=localhost port=5432 dbname=postgres user=testuser password=example"}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@fixture
|
||||
def qdrant_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for QdrantConnector."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {"QDRANT_LOCATION": "http://localhost:6333"}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@fixture
|
||||
def redis_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for Redis."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {"REDIS_CONNECTION_STRING": "redis://localhost:6379"}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@fixture
|
||||
def pinecone_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for Pinecone."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {"PINECONE_API_KEY": "test_key"}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@fixture
|
||||
def sql_server_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for SQL Server."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"SQL_SERVER_CONNECTION_STRING": "Driver={ODBC Driver 18 for SQL Server};Server=localhost;Database=testdb;User Id=testuser;Password=example;" # noqa: E501
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
def filter_lambda_list(store: str) -> list[ParameterSet]:
|
||||
"""Fixture to provide a list of filter lambdas for testing."""
|
||||
sets = [
|
||||
(
|
||||
lambda x: x.content == "value",
|
||||
{
|
||||
"ai_search": "content eq 'value'",
|
||||
},
|
||||
"equal with string",
|
||||
),
|
||||
(
|
||||
lambda x: x.id == 0,
|
||||
{
|
||||
"ai_search": "id eq 0",
|
||||
},
|
||||
"equal with int",
|
||||
),
|
||||
(
|
||||
lambda x: x.content != "value",
|
||||
{
|
||||
"ai_search": "content ne 'value'",
|
||||
},
|
||||
"not equal",
|
||||
),
|
||||
(
|
||||
lambda x: x.id > 0,
|
||||
{
|
||||
"ai_search": "id gt 0",
|
||||
},
|
||||
"greater than",
|
||||
),
|
||||
(
|
||||
lambda x: x.id >= 0,
|
||||
{
|
||||
"ai_search": "id ge 0",
|
||||
},
|
||||
"greater than or equal",
|
||||
),
|
||||
(
|
||||
lambda x: x.id == +0,
|
||||
{
|
||||
"ai_search": "id eq +0",
|
||||
},
|
||||
"equal with explicit positive",
|
||||
),
|
||||
(
|
||||
lambda x: x.id < 0,
|
||||
{
|
||||
"ai_search": "id lt 0",
|
||||
},
|
||||
"less than",
|
||||
),
|
||||
(
|
||||
lambda x: x.id <= 0,
|
||||
{
|
||||
"ai_search": "id le 0",
|
||||
},
|
||||
"less than or equal",
|
||||
),
|
||||
(
|
||||
lambda x: -10 <= x.id <= 0,
|
||||
{
|
||||
"ai_search": "(-10 le id and id le 0)",
|
||||
},
|
||||
"between inclusive",
|
||||
),
|
||||
(
|
||||
lambda x: -10 < x.id < 0,
|
||||
{
|
||||
"ai_search": "(-10 lt id and id lt 0)",
|
||||
},
|
||||
"between exclusive",
|
||||
),
|
||||
(
|
||||
lambda x: x.content == "value" and x.id == 0,
|
||||
{
|
||||
"ai_search": "(content eq 'value' and id eq 0)",
|
||||
},
|
||||
"and",
|
||||
),
|
||||
(
|
||||
lambda x: x.content == "value" or x.id == 0,
|
||||
{
|
||||
"ai_search": "(content eq 'value' or id eq 0)",
|
||||
},
|
||||
"or",
|
||||
),
|
||||
(
|
||||
lambda x: not x.content,
|
||||
{
|
||||
"ai_search": "not content",
|
||||
},
|
||||
"not with truthy",
|
||||
),
|
||||
(
|
||||
lambda x: not (x.content == "value"), # noqa: SIM201
|
||||
{
|
||||
"ai_search": "not content eq 'value'",
|
||||
},
|
||||
"not with equal",
|
||||
),
|
||||
(
|
||||
lambda x: not (x.content != "value"), # noqa: SIM202
|
||||
{
|
||||
"ai_search": "not content ne 'value'",
|
||||
},
|
||||
"not with not equal",
|
||||
),
|
||||
(
|
||||
lambda x: "value" in x.content,
|
||||
{
|
||||
"ai_search": "search.ismatch('value', 'content')",
|
||||
},
|
||||
"contains",
|
||||
),
|
||||
(
|
||||
lambda x: "value" not in x.content,
|
||||
{
|
||||
"ai_search": "not search.ismatch('value', 'content')",
|
||||
},
|
||||
"not contains",
|
||||
),
|
||||
(
|
||||
lambda x: (x.id > 0 and x.id < 3) or (x.id > 7 and x.id < 10),
|
||||
{
|
||||
"ai_search": "((id gt 0 and id lt 3) or (id gt 7 and id lt 10))",
|
||||
},
|
||||
"complex",
|
||||
),
|
||||
(
|
||||
lambda x: x.unknown_field == "value",
|
||||
{
|
||||
"ai_search": VectorStoreOperationException,
|
||||
},
|
||||
"fail unknown field",
|
||||
),
|
||||
(
|
||||
lambda x: any(x == "a" for x in x.content),
|
||||
{
|
||||
"ai_search": NotImplementedError,
|
||||
},
|
||||
"comprehension",
|
||||
),
|
||||
(
|
||||
lambda x: ~x.id,
|
||||
{
|
||||
"ai_search": NotImplementedError,
|
||||
},
|
||||
"invert",
|
||||
),
|
||||
(
|
||||
lambda x: constant, # noqa: F821
|
||||
{
|
||||
"ai_search": NotImplementedError,
|
||||
},
|
||||
"constant",
|
||||
),
|
||||
(
|
||||
lambda x: x.content.city == "Seattle",
|
||||
{
|
||||
"ai_search": "content/city eq 'Seattle'",
|
||||
},
|
||||
"nested property",
|
||||
),
|
||||
]
|
||||
return [param(s[0], s[1][store], id=s[2]) for s in sets if store in s[1]]
|
||||
@@ -0,0 +1,37 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from pymongo import AsyncMongoClient
|
||||
from pymongo.asynchronous.collection import AsyncCollection
|
||||
from pymongo.asynchronous.database import AsyncDatabase
|
||||
|
||||
BASE_PATH = "pymongo.asynchronous.mongo_client.AsyncMongoClient"
|
||||
DATABASE_PATH = "pymongo.asynchronous.database.AsyncDatabase"
|
||||
COLLECTION_PATH = "pymongo.asynchronous.collection.AsyncCollection"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_mongo_client():
|
||||
with patch(BASE_PATH, spec=AsyncMongoClient) as mock:
|
||||
yield mock
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_get_database(mock_mongo_client):
|
||||
with (
|
||||
patch(DATABASE_PATH, spec=AsyncDatabase) as mock_db,
|
||||
patch.object(mock_mongo_client, "get_database", new_callable=lambda: mock_db) as mock,
|
||||
):
|
||||
yield mock
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_get_collection(mock_get_database):
|
||||
with (
|
||||
patch(COLLECTION_PATH, spec=AsyncCollection) as mock_collection,
|
||||
patch.object(mock_get_database, "get_collection", new_callable=lambda: mock_collection) as mock,
|
||||
):
|
||||
yield mock
|
||||
@@ -0,0 +1,93 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from pymongo import AsyncMongoClient
|
||||
from pymongo.asynchronous.cursor import AsyncCursor
|
||||
from pymongo.results import UpdateResult
|
||||
from pytest import mark, raises
|
||||
|
||||
from semantic_kernel.connectors.mongodb import DEFAULT_DB_NAME, DEFAULT_SEARCH_INDEX_NAME, MongoDBAtlasCollection
|
||||
from semantic_kernel.exceptions.vector_store_exceptions import VectorStoreInitializationException
|
||||
|
||||
|
||||
def test_mongodb_atlas_collection_initialization(mongodb_atlas_unit_test_env, definition, mock_mongo_client):
|
||||
collection = MongoDBAtlasCollection(
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
collection_name="test_collection",
|
||||
mongo_client=mock_mongo_client,
|
||||
)
|
||||
assert collection.mongo_client is not None
|
||||
assert isinstance(collection.mongo_client, AsyncMongoClient)
|
||||
|
||||
|
||||
@mark.parametrize("exclude_list", [["MONGODB_ATLAS_CONNECTION_STRING"]], indirect=True)
|
||||
def test_mongodb_atlas_collection_initialization_fail(mongodb_atlas_unit_test_env, definition):
|
||||
with raises(VectorStoreInitializationException):
|
||||
MongoDBAtlasCollection(
|
||||
collection_name="test_collection",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
|
||||
|
||||
@mark.parametrize("exclude_list", [["MONGODB_ATLAS_DATABASE_NAME", "MONGODB_ATLAS_INDEX_NAME"]], indirect=True)
|
||||
def test_mongodb_atlas_collection_initialization_defaults(mongodb_atlas_unit_test_env, definition):
|
||||
collection = MongoDBAtlasCollection(
|
||||
collection_name="test_collection",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
assert collection.database_name == DEFAULT_DB_NAME
|
||||
assert collection.index_name == DEFAULT_SEARCH_INDEX_NAME
|
||||
|
||||
|
||||
async def test_mongodb_atlas_collection_upsert(mongodb_atlas_unit_test_env, definition, mock_get_collection):
|
||||
collection = MongoDBAtlasCollection(
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
collection_name="test_collection",
|
||||
)
|
||||
with patch.object(collection, "_get_collection", new=mock_get_collection) as mock_get:
|
||||
result_mock = AsyncMock(spec=UpdateResult)
|
||||
result_mock.upserted_ids = {0: "test_id"}
|
||||
mock_get.return_value.bulk_write.return_value = result_mock
|
||||
result = await collection._inner_upsert([{"_id": "test_id", "data": "test_data"}])
|
||||
assert result == ["test_id"]
|
||||
|
||||
|
||||
async def test_mongodb_atlas_collection_get(mongodb_atlas_unit_test_env, definition, mock_get_collection):
|
||||
collection = MongoDBAtlasCollection(
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
collection_name="test_collection",
|
||||
)
|
||||
with patch.object(collection, "_get_collection", new=mock_get_collection) as mock_get:
|
||||
result_mock = AsyncMock(spec=AsyncCursor)
|
||||
result_mock.to_list.return_value = [{"_id": "test_id", "data": "test_data"}]
|
||||
mock_get.return_value.find.return_value = result_mock
|
||||
result = await collection._inner_get(["test_id"])
|
||||
assert result == [{"_id": "test_id", "data": "test_data"}]
|
||||
|
||||
|
||||
async def test_mongodb_atlas_collection_delete(mongodb_atlas_unit_test_env, definition, mock_get_collection):
|
||||
collection = MongoDBAtlasCollection(
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
collection_name="test_collection",
|
||||
)
|
||||
with patch.object(collection, "_get_collection", new=mock_get_collection) as mock_get:
|
||||
await collection._inner_delete(["test_id"])
|
||||
mock_get.return_value.delete_many.assert_called_with({"_id": {"$in": ["test_id"]}})
|
||||
|
||||
|
||||
async def test_mongodb_atlas_collection_collection_exists(mongodb_atlas_unit_test_env, definition, mock_get_database):
|
||||
collection = MongoDBAtlasCollection(
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
collection_name="test_collection",
|
||||
)
|
||||
with patch.object(collection, "_get_database", new=mock_get_database) as mock_get:
|
||||
mock_get.return_value.list_collection_names.return_value = ["test_collection"]
|
||||
assert await collection.collection_exists()
|
||||
@@ -0,0 +1,30 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from pymongo import AsyncMongoClient
|
||||
|
||||
from semantic_kernel.connectors.mongodb import MongoDBAtlasCollection, MongoDBAtlasStore
|
||||
|
||||
|
||||
def test_mongodb_atlas_store_initialization(mongodb_atlas_unit_test_env):
|
||||
store = MongoDBAtlasStore()
|
||||
assert store.mongo_client is not None
|
||||
assert isinstance(store.mongo_client, AsyncMongoClient)
|
||||
|
||||
|
||||
def test_mongodb_atlas_store_get_collection(mongodb_atlas_unit_test_env, definition):
|
||||
store = MongoDBAtlasStore()
|
||||
collection = store.get_collection(
|
||||
collection_name="test_collection",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
assert collection is not None
|
||||
assert isinstance(collection, MongoDBAtlasCollection)
|
||||
|
||||
|
||||
async def test_mongodb_atlas_store_list_collection_names(mongodb_atlas_unit_test_env, mock_mongo_client):
|
||||
store = MongoDBAtlasStore(mongo_client=mock_mongo_client, database_name="test_db")
|
||||
store.mongo_client.get_database().list_collection_names.return_value = ["test_collection"]
|
||||
result = await store.list_collection_names()
|
||||
assert result == ["test_collection"]
|
||||
@@ -0,0 +1,450 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
import numpy as np
|
||||
from azure.search.documents.aio import SearchClient
|
||||
from azure.search.documents.indexes.aio import SearchIndexClient
|
||||
from pytest import fixture, mark, param, raises
|
||||
|
||||
from semantic_kernel.connectors.ai.embedding_generator_base import EmbeddingGeneratorBase
|
||||
from semantic_kernel.connectors.azure_ai_search import (
|
||||
AzureAISearchCollection,
|
||||
AzureAISearchSettings,
|
||||
AzureAISearchStore,
|
||||
_definition_to_azure_ai_search_index,
|
||||
_get_search_index_client,
|
||||
_resolve_credential,
|
||||
)
|
||||
from semantic_kernel.exceptions import (
|
||||
ServiceInitializationError,
|
||||
VectorStoreInitializationException,
|
||||
VectorStoreOperationException,
|
||||
)
|
||||
from semantic_kernel.utils.list_handler import desync_list
|
||||
from tests.unit.connectors.memory.conftest import filter_lambda_list
|
||||
|
||||
BASE_PATH_SEARCH_CLIENT = "azure.search.documents.aio.SearchClient"
|
||||
BASE_PATH_INDEX_CLIENT = "azure.search.documents.indexes.aio.SearchIndexClient"
|
||||
|
||||
|
||||
@fixture
|
||||
def vector_store(azure_ai_search_unit_test_env):
|
||||
"""Fixture to instantiate AzureCognitiveSearchMemoryStore with basic configuration."""
|
||||
return AzureAISearchStore()
|
||||
|
||||
|
||||
@fixture
|
||||
def mock_ensure_collection_exists():
|
||||
"""Fixture to patch 'SearchIndexClient' and its 'create_index' method."""
|
||||
with patch(f"{BASE_PATH_INDEX_CLIENT}.create_index") as mock_create_index:
|
||||
yield mock_create_index
|
||||
|
||||
|
||||
@fixture
|
||||
def mock_ensure_collection_deleted():
|
||||
"""Fixture to patch 'SearchIndexClient' and its 'create_index' method."""
|
||||
with patch(f"{BASE_PATH_INDEX_CLIENT}.delete_index") as mock_delete_index:
|
||||
yield mock_delete_index
|
||||
|
||||
|
||||
@fixture
|
||||
def mock_list_collection_names():
|
||||
"""Fixture to patch 'SearchIndexClient' and its 'create_index' method."""
|
||||
with patch(f"{BASE_PATH_INDEX_CLIENT}.list_index_names") as mock_list_index_names:
|
||||
# Setup the mock to return a specific SearchIndex instance when called
|
||||
mock_list_index_names.return_value = desync_list(["test"])
|
||||
yield mock_list_index_names
|
||||
|
||||
|
||||
@fixture
|
||||
def mock_upsert():
|
||||
with patch(f"{BASE_PATH_SEARCH_CLIENT}.merge_or_upload_documents") as mock_merge_or_upload_documents:
|
||||
from azure.search.documents.models import IndexingResult
|
||||
|
||||
result = MagicMock(spec=IndexingResult)
|
||||
result.key = "id1"
|
||||
mock_merge_or_upload_documents.return_value = [result]
|
||||
yield mock_merge_or_upload_documents
|
||||
|
||||
|
||||
@fixture
|
||||
def mock_get():
|
||||
with patch(f"{BASE_PATH_SEARCH_CLIENT}.get_document") as mock_get_document:
|
||||
mock_get_document.return_value = {"id": "id1", "content": "content", "vector": [1.0, 2.0, 3.0]}
|
||||
yield mock_get_document
|
||||
|
||||
|
||||
@fixture
|
||||
def mock_search():
|
||||
async def iter_search_results(*args, **kwargs):
|
||||
yield {"id": "id1", "content": "content", "vector": [1.0, 2.0, 3.0]}
|
||||
await asyncio.sleep(0.0)
|
||||
|
||||
with patch(f"{BASE_PATH_SEARCH_CLIENT}.search") as mock_search:
|
||||
mock_search.side_effect = iter_search_results
|
||||
yield mock_search
|
||||
|
||||
|
||||
@fixture
|
||||
def mock_delete():
|
||||
with patch(f"{BASE_PATH_SEARCH_CLIENT}.delete_documents") as mock_delete_documents:
|
||||
yield mock_delete_documents
|
||||
|
||||
|
||||
@fixture
|
||||
def collection(azure_ai_search_unit_test_env, definition):
|
||||
return AzureAISearchCollection(record_type=dict, definition=definition)
|
||||
|
||||
|
||||
async def test_init(azure_ai_search_unit_test_env, definition):
|
||||
async with AzureAISearchCollection(record_type=dict, definition=definition) as collection:
|
||||
assert collection is not None
|
||||
assert collection.record_type is dict
|
||||
assert collection.definition == definition
|
||||
assert collection.collection_name == "test-index-name"
|
||||
assert collection.search_index_client is not None
|
||||
assert collection.search_client is not None
|
||||
|
||||
|
||||
def test_init_with_type(azure_ai_search_unit_test_env, record_type):
|
||||
collection = AzureAISearchCollection(record_type=record_type)
|
||||
assert collection is not None
|
||||
assert collection.record_type is record_type
|
||||
assert collection.collection_name == "test-index-name"
|
||||
assert collection.search_index_client is not None
|
||||
assert collection.search_client is not None
|
||||
|
||||
|
||||
@mark.parametrize("exclude_list", [["AZURE_AI_SEARCH_ENDPOINT"]], indirect=True)
|
||||
def test_init_endpoint_fail(azure_ai_search_unit_test_env, definition):
|
||||
with raises(VectorStoreInitializationException):
|
||||
AzureAISearchCollection(record_type=dict, definition=definition, env_file_path="test.env")
|
||||
|
||||
|
||||
@mark.parametrize("exclude_list", [["AZURE_AI_SEARCH_INDEX_NAME"]], indirect=True)
|
||||
def test_init_index_fail(azure_ai_search_unit_test_env, definition):
|
||||
with raises(VectorStoreInitializationException):
|
||||
AzureAISearchCollection(record_type=dict, definition=definition, env_file_path="test.env")
|
||||
|
||||
|
||||
def test_init_with_clients(azure_ai_search_unit_test_env, definition):
|
||||
search_index_client = MagicMock(spec=SearchIndexClient)
|
||||
search_client = MagicMock(spec=SearchClient)
|
||||
search_client._index_name = "test-index-name"
|
||||
|
||||
collection = AzureAISearchCollection(
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
search_index_client=search_index_client,
|
||||
search_client=search_client,
|
||||
)
|
||||
assert collection is not None
|
||||
assert collection.record_type is dict
|
||||
assert collection.definition == definition
|
||||
assert collection.collection_name == "test-index-name"
|
||||
assert collection.search_index_client == search_index_client
|
||||
assert collection.search_client == search_client
|
||||
|
||||
|
||||
def test_init_with_search_index_client(azure_ai_search_unit_test_env, definition):
|
||||
search_index_client = MagicMock(spec=SearchIndexClient)
|
||||
with patch("semantic_kernel.connectors.azure_ai_search._get_search_client") as get_search_client:
|
||||
search_client = MagicMock(spec=SearchClient)
|
||||
get_search_client.return_value = search_client
|
||||
|
||||
collection = AzureAISearchCollection(
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
collection_name="test",
|
||||
search_index_client=search_index_client,
|
||||
)
|
||||
assert collection is not None
|
||||
assert collection.record_type is dict
|
||||
assert collection.definition == definition
|
||||
assert collection.collection_name == "test"
|
||||
assert collection.search_index_client == search_index_client
|
||||
assert collection.search_client == search_client
|
||||
|
||||
|
||||
@mark.parametrize("exclude_list", [["AZURE_AI_SEARCH_INDEX_NAME"]], indirect=True)
|
||||
def test_init_with_search_index_client_fail(azure_ai_search_unit_test_env, definition):
|
||||
search_index_client = MagicMock(spec=SearchIndexClient)
|
||||
with raises(VectorStoreInitializationException):
|
||||
AzureAISearchCollection(
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
search_index_client=search_index_client,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
async def test_upsert(collection, mock_upsert):
|
||||
ids = await collection._inner_upsert({"id": "id1", "name": "test"})
|
||||
assert ids[0] == "id1"
|
||||
|
||||
ids = await collection.upsert(records={"id": "id1", "content": "content", "vector": [1.0, 2.0, 3.0]})
|
||||
assert ids == "id1"
|
||||
|
||||
|
||||
async def test_get(collection, mock_get):
|
||||
records = await collection._inner_get(["id1"])
|
||||
assert records is not None
|
||||
|
||||
records = await collection.get("id1")
|
||||
assert records is not None
|
||||
|
||||
|
||||
@mark.parametrize(
|
||||
"order_by, ordering",
|
||||
[
|
||||
param("id", ["id"], id="single id"),
|
||||
param({"id": True}, ["id"], id="ascending id"),
|
||||
param({"id": False}, ["id desc"], id="descending id"),
|
||||
param(["id"], ["id"], id="ascending id list"),
|
||||
param(["id", "content"], ["id", "content"], id="multiple"),
|
||||
param([{"id": True}, {"content": False}], ["id", "content desc"], id="multiple desc"),
|
||||
param(["id", {"content": False}], ["id", "content desc"], id="multiple mix"),
|
||||
],
|
||||
)
|
||||
async def test_get_without_key(collection, mock_get, mock_search, order_by, ordering):
|
||||
records = await collection.get(top=10, order_by=order_by)
|
||||
assert records is not None
|
||||
mock_search.assert_called_once_with(
|
||||
search_text="*",
|
||||
top=10,
|
||||
skip=0,
|
||||
select=["id", "content"],
|
||||
order_by=ordering,
|
||||
)
|
||||
|
||||
|
||||
async def test_delete(collection, mock_delete):
|
||||
await collection._inner_delete(["id1"])
|
||||
|
||||
|
||||
async def test_collection_exists(collection, mock_list_collection_names):
|
||||
await collection.collection_exists()
|
||||
|
||||
|
||||
async def test_ensure_collection_deleted(collection, mock_ensure_collection_deleted):
|
||||
await collection.ensure_collection_deleted()
|
||||
|
||||
|
||||
@mark.parametrize("distance_function", [("cosine_distance")])
|
||||
async def test_create_index_from_index(collection, mock_ensure_collection_exists):
|
||||
from azure.search.documents.indexes.models import SearchIndex
|
||||
|
||||
index = MagicMock(spec=SearchIndex)
|
||||
await collection.ensure_collection_exists(index=index)
|
||||
|
||||
|
||||
@mark.parametrize("distance_function", [("cosine_distance")])
|
||||
async def test_create_index_from_definition(collection, mock_ensure_collection_exists):
|
||||
from azure.search.documents.indexes.models import SearchIndex
|
||||
|
||||
with patch(
|
||||
"semantic_kernel.connectors.azure_ai_search._definition_to_azure_ai_search_index",
|
||||
return_value=MagicMock(spec=SearchIndex),
|
||||
):
|
||||
await collection.ensure_collection_exists()
|
||||
|
||||
|
||||
async def test_create_index_from_index_fail(collection, mock_ensure_collection_exists):
|
||||
index = Mock()
|
||||
with raises(VectorStoreOperationException):
|
||||
await collection.ensure_collection_exists(index=index)
|
||||
|
||||
|
||||
@mark.parametrize("distance_function", [("cosine_distance")])
|
||||
def test_definition_to_azure_ai_search_index(definition):
|
||||
index = _definition_to_azure_ai_search_index("test", definition)
|
||||
assert index is not None
|
||||
assert index.name == "test"
|
||||
assert len(index.fields) == 3
|
||||
|
||||
|
||||
@mark.parametrize("exclude_list", [["AZURE_AI_SEARCH_ENDPOINT"]], indirect=True)
|
||||
async def test_vector_store_fail(azure_ai_search_unit_test_env):
|
||||
with raises(VectorStoreInitializationException):
|
||||
AzureAISearchStore(env_file_path="test.env")
|
||||
|
||||
|
||||
async def test_vector_store_list_collection_names(vector_store, mock_list_collection_names):
|
||||
assert vector_store.search_index_client is not None
|
||||
collection_names = await vector_store.list_collection_names()
|
||||
assert collection_names == ["test"]
|
||||
mock_list_collection_names.assert_called_once()
|
||||
|
||||
|
||||
async def test_vector_store_collection_existss(vector_store, mock_list_collection_names):
|
||||
assert vector_store.search_index_client is not None
|
||||
exists = await vector_store.collection_exists("test")
|
||||
assert exists
|
||||
mock_list_collection_names.assert_called_once()
|
||||
|
||||
|
||||
async def test_vector_store_ensure_collection_deleted(vector_store, mock_ensure_collection_deleted):
|
||||
assert vector_store.search_index_client is not None
|
||||
await vector_store.ensure_collection_deleted("test")
|
||||
mock_ensure_collection_deleted.assert_called_once()
|
||||
|
||||
|
||||
def test_get_collection(vector_store, definition):
|
||||
collection = vector_store.get_collection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
assert collection is not None
|
||||
assert collection.collection_name == "test"
|
||||
assert collection.search_index_client == vector_store.search_index_client
|
||||
assert collection.search_client is not None
|
||||
assert collection.search_endpoint == vector_store.search_endpoint
|
||||
assert collection.search_credential == vector_store.search_credential
|
||||
|
||||
|
||||
def test_get_collection_with_provided_search_index_client(azure_ai_search_unit_test_env, definition):
|
||||
"""Test that get_collection works when AzureAISearchStore is created with a pre-built search_index_client.
|
||||
|
||||
When search_index_client is provided directly, search_endpoint and search_credential
|
||||
are not resolved at store creation time. get_collection() should still succeed
|
||||
by falling back to environment variables for endpoint/credential resolution.
|
||||
"""
|
||||
search_index_client = MagicMock(spec=SearchIndexClient)
|
||||
store = AzureAISearchStore(search_index_client=search_index_client)
|
||||
assert store.search_endpoint is None
|
||||
assert store.search_credential is None
|
||||
|
||||
collection = store.get_collection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
assert collection is not None
|
||||
assert collection.collection_name == "test"
|
||||
assert collection.search_index_client == search_index_client
|
||||
assert collection.search_client is not None
|
||||
|
||||
|
||||
@mark.parametrize("exclude_list", [["AZURE_AI_SEARCH_API_KEY"]], indirect=True)
|
||||
def test_get_search_index_client(azure_ai_search_unit_test_env):
|
||||
from azure.core.credentials import AzureKeyCredential
|
||||
from azure.core.credentials_async import AsyncTokenCredential
|
||||
|
||||
settings = AzureAISearchSettings(**azure_ai_search_unit_test_env, env_file_path="test.env")
|
||||
|
||||
azure_credential = MagicMock(spec=AzureKeyCredential)
|
||||
client = _get_search_index_client(settings, azure_credential=azure_credential)
|
||||
assert client is not None
|
||||
|
||||
token_credential = MagicMock(spec=AsyncTokenCredential)
|
||||
client2 = _get_search_index_client(
|
||||
settings,
|
||||
token_credential=token_credential,
|
||||
)
|
||||
assert client2 is not None
|
||||
|
||||
with raises(ServiceInitializationError):
|
||||
_get_search_index_client(settings)
|
||||
|
||||
|
||||
@mark.parametrize("exclude_list", [["AZURE_AI_SEARCH_API_KEY"]], indirect=True)
|
||||
def test_resolve_credential(azure_ai_search_unit_test_env):
|
||||
from azure.core.credentials import AzureKeyCredential
|
||||
from azure.core.credentials_async import AsyncTokenCredential
|
||||
|
||||
settings = AzureAISearchSettings(**azure_ai_search_unit_test_env, env_file_path="test.env")
|
||||
|
||||
azure_credential = MagicMock(spec=AzureKeyCredential)
|
||||
resolved = _resolve_credential(settings, azure_credential=azure_credential)
|
||||
assert resolved == azure_credential
|
||||
|
||||
token_credential = MagicMock(spec=AsyncTokenCredential)
|
||||
resolved = _resolve_credential(settings, token_credential=token_credential)
|
||||
assert resolved == token_credential
|
||||
|
||||
with raises(ServiceInitializationError):
|
||||
_resolve_credential(settings)
|
||||
|
||||
|
||||
@mark.parametrize("include_vectors", [True, False])
|
||||
async def test_search_vectorized_search(collection, mock_search, include_vectors):
|
||||
results = await collection.search(vector=[0.1, 0.2, 0.3], include_vectors=include_vectors)
|
||||
assert results is not None
|
||||
async for result in results.results:
|
||||
assert result is not None
|
||||
assert result.record is not None
|
||||
assert result.record["id"] == "id1"
|
||||
assert result.record["content"] == "content"
|
||||
if include_vectors:
|
||||
assert result.record["vector"] == [1.0, 2.0, 3.0]
|
||||
for call in mock_search.call_args_list:
|
||||
assert call[1]["top"] == 3
|
||||
assert call[1]["skip"] == 0
|
||||
assert call[1]["include_total_count"] is False
|
||||
assert call[1]["select"] == ["*"] if include_vectors else ["id", "content"]
|
||||
assert call[1]["vector_queries"][0].vector == [0.1, 0.2, 0.3]
|
||||
assert call[1]["vector_queries"][0].fields == "vector"
|
||||
|
||||
|
||||
@mark.parametrize("include_vectors", [True, False])
|
||||
async def test_search_vectorizable_search(collection, mock_search, include_vectors):
|
||||
collection.embedding_generator = AsyncMock(spec=EmbeddingGeneratorBase)
|
||||
collection.embedding_generator.generate_embeddings.return_value = np.array([[0.1, 0.2, 0.3]])
|
||||
results = await collection.search("test", include_vectors=include_vectors)
|
||||
assert results is not None
|
||||
async for result in results.results:
|
||||
assert result is not None
|
||||
assert result.record is not None
|
||||
assert result.record["id"] == "id1"
|
||||
assert result.record["content"] == "content"
|
||||
if include_vectors:
|
||||
assert result.record["vector"] == [1.0, 2.0, 3.0]
|
||||
for call in mock_search.call_args_list:
|
||||
assert call[1]["top"] == 3
|
||||
assert call[1]["skip"] == 0
|
||||
assert call[1]["include_total_count"] is False
|
||||
assert call[1]["select"] == ["*"] if include_vectors else ["id", "content"]
|
||||
assert call[1]["vector_queries"][0].vector == [0.1, 0.2, 0.3]
|
||||
assert call[1]["vector_queries"][0].fields == "vector"
|
||||
|
||||
|
||||
@mark.parametrize("include_vectors", [True, False])
|
||||
@mark.parametrize("keywords", ["test", ["test1", "test2"]], ids=["single", "multiple"])
|
||||
async def test_search_keyword_hybrid_search(collection, mock_search, include_vectors, keywords):
|
||||
results = await collection.hybrid_search(
|
||||
values=keywords,
|
||||
vector=[0.1, 0.2, 0.3],
|
||||
include_vectors=include_vectors,
|
||||
additional_property_name="content",
|
||||
)
|
||||
assert results is not None
|
||||
async for result in results.results:
|
||||
assert result is not None
|
||||
assert result.record is not None
|
||||
assert result.record["id"] == "id1"
|
||||
assert result.record["content"] == "content"
|
||||
if include_vectors:
|
||||
assert result.record["vector"] == [1.0, 2.0, 3.0]
|
||||
for call in mock_search.call_args_list:
|
||||
assert call[1]["top"] == 3
|
||||
assert call[1]["skip"] == 0
|
||||
assert call[1]["include_total_count"] is False
|
||||
assert call[1]["select"] == ["*"] if include_vectors else ["id", "content"]
|
||||
assert call[1]["search_fields"] == ["content"]
|
||||
assert call[1]["search_text"] == "test" if keywords == "test" else "test1, test2"
|
||||
assert call[1]["vector_queries"][0].vector == [0.1, 0.2, 0.3]
|
||||
assert call[1]["vector_queries"][0].fields == "vector"
|
||||
|
||||
|
||||
@mark.parametrize("filter, result", filter_lambda_list("ai_search"))
|
||||
def test_lambda_filter(collection, filter, result):
|
||||
if isinstance(result, type) and issubclass(result, Exception):
|
||||
with raises(result):
|
||||
collection._build_filter(filter)
|
||||
else:
|
||||
filter_string = collection._build_filter(filter)
|
||||
assert filter_string == result
|
||||
@@ -0,0 +1,115 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from chromadb.api import ClientAPI
|
||||
from chromadb.api.models.Collection import Collection
|
||||
|
||||
from semantic_kernel.connectors.chroma import ChromaCollection, ChromaStore
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_client():
|
||||
return MagicMock(spec=ClientAPI)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def chroma_collection(mock_client, definition):
|
||||
return ChromaCollection(
|
||||
collection_name="test_collection",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
client=mock_client,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def chroma_store(mock_client):
|
||||
return ChromaStore(client=mock_client)
|
||||
|
||||
|
||||
def test_chroma_collection_initialization(chroma_collection):
|
||||
assert chroma_collection.collection_name == "test_collection"
|
||||
assert chroma_collection.record_type is dict
|
||||
|
||||
|
||||
def test_chroma_store_initialization(chroma_store):
|
||||
assert chroma_store.client is not None
|
||||
|
||||
|
||||
def test_chroma_collection_get_collection(chroma_collection, mock_client):
|
||||
mock_client.get_collection.return_value = "mock_collection"
|
||||
collection = chroma_collection._get_collection()
|
||||
assert collection == "mock_collection"
|
||||
|
||||
|
||||
def test_chroma_store_get_collection(chroma_store, mock_client, definition):
|
||||
collection = chroma_store.get_collection(collection_name="test_collection", record_type=dict, definition=definition)
|
||||
assert collection is not None
|
||||
assert isinstance(collection, ChromaCollection)
|
||||
|
||||
|
||||
async def test_chroma_collection_collection_exists(chroma_collection, mock_client):
|
||||
mock_client.get_collection.return_value = "mock_collection"
|
||||
exists = await chroma_collection.collection_exists()
|
||||
assert exists
|
||||
|
||||
|
||||
async def test_chroma_store_list_collection_names(chroma_store, mock_client):
|
||||
mock_collection = MagicMock(spec=Collection)
|
||||
mock_collection.name = "test_collection"
|
||||
mock_client.list_collections.return_value = [mock_collection]
|
||||
collections = await chroma_store.list_collection_names()
|
||||
assert collections == ["test_collection"]
|
||||
|
||||
|
||||
async def test_chroma_collection_ensure_collection_exists(chroma_collection, mock_client):
|
||||
await chroma_collection.ensure_collection_exists()
|
||||
mock_client.create_collection.assert_called_once_with(
|
||||
name="test_collection", embedding_function=None, configuration={"hnsw": {"space": "cosine"}}, get_or_create=True
|
||||
)
|
||||
|
||||
|
||||
async def test_chroma_collection_ensure_collection_deleted(chroma_collection, mock_client):
|
||||
await chroma_collection.ensure_collection_deleted()
|
||||
mock_client.delete_collection.assert_called_once_with(name="test_collection")
|
||||
|
||||
|
||||
async def test_chroma_collection_upsert(chroma_collection, mock_client):
|
||||
records = [{"id": "1", "vector": [0.1, 0.2, 0.3, 0.4, 0.5], "content": "test document"}]
|
||||
ids = await chroma_collection.upsert(records)
|
||||
assert ids == ["1"]
|
||||
mock_client.get_collection().add.assert_called_once()
|
||||
|
||||
|
||||
async def test_chroma_collection_get(chroma_collection, mock_client):
|
||||
mock_client.get_collection().get.return_value = {
|
||||
"ids": [["1"]],
|
||||
"documents": [["test document"]],
|
||||
"embeddings": [[[0.1, 0.2, 0.3, 0.4, 0.5]]],
|
||||
"metadatas": [[{}]],
|
||||
}
|
||||
records = await chroma_collection._inner_get(["1"])
|
||||
assert len(records) == 1
|
||||
assert records[0]["id"] == "1"
|
||||
|
||||
|
||||
async def test_chroma_collection_delete(chroma_collection, mock_client):
|
||||
await chroma_collection._inner_delete(["1"])
|
||||
mock_client.get_collection().delete.assert_called_once_with(ids=["1"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("include_vectors", [True, False])
|
||||
async def test_chroma_collection_search(chroma_collection, mock_client, include_vectors):
|
||||
mock_client.get_collection().query.return_value = {
|
||||
"ids": [["1"]],
|
||||
"documents": [["test document"]],
|
||||
"embeddings": [[[0.1, 0.2, 0.3, 0.4, 0.5]]],
|
||||
"metadatas": [[{}]],
|
||||
"distances": [[0.1]],
|
||||
}
|
||||
results = await chroma_collection.search(vector=[0.1, 0.2, 0.3, 0.4, 0.5], top=1, include_vectors=include_vectors)
|
||||
async for res in results.results:
|
||||
assert res.record["id"] == "1"
|
||||
assert res.score == 0.1
|
||||
@@ -0,0 +1,176 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import faiss
|
||||
from pytest import fixture, mark, raises
|
||||
|
||||
from semantic_kernel.connectors.faiss import FaissCollection, FaissStore
|
||||
from semantic_kernel.data.vector import DistanceFunction, VectorStoreCollectionDefinition, VectorStoreField
|
||||
from semantic_kernel.exceptions import VectorStoreInitializationException
|
||||
|
||||
|
||||
@fixture(scope="function")
|
||||
def data_model_def() -> VectorStoreCollectionDefinition:
|
||||
return VectorStoreCollectionDefinition(
|
||||
fields=[
|
||||
VectorStoreField("key", name="id"),
|
||||
VectorStoreField("data", name="content"),
|
||||
VectorStoreField(
|
||||
"vector",
|
||||
name="vector",
|
||||
dimensions=5,
|
||||
index_kind="flat",
|
||||
distance_function="dot_prod",
|
||||
type="float",
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@fixture(scope="function")
|
||||
def store() -> FaissStore:
|
||||
return FaissStore()
|
||||
|
||||
|
||||
@fixture(scope="function")
|
||||
def faiss_collection(data_model_def):
|
||||
return FaissCollection(record_type=dict, definition=data_model_def, collection_name="test")
|
||||
|
||||
|
||||
async def test_store_get_collection(store, data_model_def):
|
||||
collection = store.get_collection(dict, definition=data_model_def, collection_name="test")
|
||||
assert collection.collection_name == "test"
|
||||
assert collection.record_type is dict
|
||||
assert collection.definition == data_model_def
|
||||
assert collection.inner_storage == {}
|
||||
|
||||
|
||||
@mark.parametrize(
|
||||
"dist",
|
||||
[
|
||||
DistanceFunction.EUCLIDEAN_SQUARED_DISTANCE,
|
||||
DistanceFunction.DOT_PROD,
|
||||
],
|
||||
)
|
||||
async def test_ensure_collection_exists(store, data_model_def, dist):
|
||||
for field in data_model_def.fields:
|
||||
if field.name == "vector":
|
||||
field.distance_function = dist
|
||||
collection = store.get_collection(collection_name="test", record_type=dict, definition=data_model_def)
|
||||
await collection.ensure_collection_exists()
|
||||
assert collection.inner_storage == {}
|
||||
assert collection.indexes
|
||||
assert collection.indexes["vector"] is not None
|
||||
|
||||
|
||||
async def test_ensure_collection_exists_incompatible_dist(store, data_model_def):
|
||||
for field in data_model_def.fields:
|
||||
if field.name == "vector":
|
||||
field.distance_function = "cosine_distance"
|
||||
collection = store.get_collection(collection_name="test", record_type=dict, definition=data_model_def)
|
||||
with raises(VectorStoreInitializationException):
|
||||
await collection.ensure_collection_exists()
|
||||
|
||||
|
||||
async def test_ensure_collection_exists_custom(store, data_model_def):
|
||||
index = faiss.IndexFlat(5)
|
||||
collection = store.get_collection(collection_name="test", record_type=dict, definition=data_model_def)
|
||||
await collection.ensure_collection_exists(index=index)
|
||||
assert collection.inner_storage == {}
|
||||
assert collection.indexes
|
||||
assert collection.indexes["vector"] is not None
|
||||
assert collection.indexes["vector"] == index
|
||||
assert collection.indexes["vector"].is_trained is True
|
||||
await collection.ensure_collection_deleted()
|
||||
|
||||
|
||||
async def test_ensure_collection_exists_custom_untrained(store, data_model_def):
|
||||
index = faiss.IndexIVFFlat(faiss.IndexFlat(5), 5, 10)
|
||||
collection = store.get_collection(collection_name="test", record_type=dict, definition=data_model_def)
|
||||
with raises(VectorStoreInitializationException):
|
||||
await collection.ensure_collection_exists(index=index)
|
||||
del index
|
||||
|
||||
|
||||
async def test_ensure_collection_exists_custom_dict(store, data_model_def):
|
||||
index = faiss.IndexFlat(5)
|
||||
collection = store.get_collection(collection_name="test", record_type=dict, definition=data_model_def)
|
||||
await collection.ensure_collection_exists(indexes={"vector": index})
|
||||
assert collection.inner_storage == {}
|
||||
assert collection.indexes
|
||||
assert collection.indexes["vector"] is not None
|
||||
assert collection.indexes["vector"] == index
|
||||
await collection.ensure_collection_deleted()
|
||||
|
||||
|
||||
async def test_upsert(faiss_collection):
|
||||
await faiss_collection.ensure_collection_exists()
|
||||
record = {"id": "testid", "content": "test content", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]}
|
||||
key = await faiss_collection.upsert(record)
|
||||
assert key == "testid"
|
||||
assert faiss_collection.inner_storage == {"testid": record}
|
||||
await faiss_collection.ensure_collection_deleted()
|
||||
|
||||
|
||||
async def test_get(faiss_collection):
|
||||
await faiss_collection.ensure_collection_exists()
|
||||
record = {"id": "testid", "content": "test content", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]}
|
||||
await faiss_collection.upsert(record)
|
||||
result = await faiss_collection.get("testid")
|
||||
assert result["id"] == record["id"]
|
||||
assert result["content"] == record["content"]
|
||||
await faiss_collection.ensure_collection_deleted()
|
||||
|
||||
|
||||
async def test_get_missing(faiss_collection):
|
||||
await faiss_collection.ensure_collection_exists()
|
||||
result = await faiss_collection.get("testid")
|
||||
assert result is None
|
||||
await faiss_collection.ensure_collection_deleted()
|
||||
|
||||
|
||||
async def test_delete(faiss_collection):
|
||||
await faiss_collection.ensure_collection_exists()
|
||||
record = {"id": "testid", "content": "test content", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]}
|
||||
await faiss_collection.upsert(record)
|
||||
await faiss_collection.delete("testid")
|
||||
assert faiss_collection.inner_storage == {}
|
||||
await faiss_collection.ensure_collection_deleted()
|
||||
|
||||
|
||||
async def test_collection_exists(faiss_collection):
|
||||
assert await faiss_collection.collection_exists() is False
|
||||
await faiss_collection.ensure_collection_exists()
|
||||
assert await faiss_collection.collection_exists() is True
|
||||
await faiss_collection.ensure_collection_deleted()
|
||||
|
||||
|
||||
async def test_ensure_collection_deleted(faiss_collection):
|
||||
await faiss_collection.ensure_collection_exists()
|
||||
record = {"id": "testid", "content": "test content", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]}
|
||||
await faiss_collection.upsert(record)
|
||||
assert faiss_collection.inner_storage == {"testid": record}
|
||||
await faiss_collection.ensure_collection_deleted()
|
||||
assert faiss_collection.inner_storage == {}
|
||||
|
||||
|
||||
@mark.parametrize("dist", [DistanceFunction.EUCLIDEAN_SQUARED_DISTANCE, DistanceFunction.DOT_PROD])
|
||||
async def test_ensure_collection_exists_and_search(faiss_collection, dist):
|
||||
for field in faiss_collection.definition.fields:
|
||||
if field.name == "vector":
|
||||
field.distance_function = dist
|
||||
await faiss_collection.ensure_collection_exists()
|
||||
record1 = {"id": "testid1", "content": "test content", "vector": [1.0, 1.0, 1.0, 1.0, 1.0]}
|
||||
record2 = {"id": "testid2", "content": "test content", "vector": [-1.0, -1.0, -1.0, -1.0, -1.0]}
|
||||
await faiss_collection.upsert([record1, record2])
|
||||
results = await faiss_collection.search(
|
||||
vector=[0.9, 0.9, 0.9, 0.9, 0.9],
|
||||
vector_property_name="vector",
|
||||
include_total_count=True,
|
||||
include_vectors=True,
|
||||
)
|
||||
assert results.total_count == 2
|
||||
idx = 0
|
||||
async for res in results.results:
|
||||
assert res.record == record1 if idx == 0 else record2
|
||||
idx += 1
|
||||
await faiss_collection.ensure_collection_deleted()
|
||||
@@ -0,0 +1,294 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import ast
|
||||
|
||||
from pytest import fixture, mark, raises
|
||||
|
||||
from semantic_kernel.connectors.in_memory import InMemoryCollection, InMemoryStore
|
||||
from semantic_kernel.data._shared import default_dynamic_filter_function
|
||||
from semantic_kernel.data.vector import DistanceFunction
|
||||
from semantic_kernel.exceptions.vector_store_exceptions import VectorStoreOperationException
|
||||
|
||||
|
||||
@fixture
|
||||
def collection(definition):
|
||||
return InMemoryCollection(collection_name="test", record_type=dict, definition=definition)
|
||||
|
||||
|
||||
def test_store_init():
|
||||
store = InMemoryStore()
|
||||
assert store is not None
|
||||
|
||||
|
||||
def test_store_get_collection(definition):
|
||||
store = InMemoryStore()
|
||||
collection = store.get_collection(collection_name="test", record_type=dict, definition=definition)
|
||||
assert collection.collection_name == "test"
|
||||
assert collection.record_type is dict
|
||||
assert collection.definition == definition
|
||||
|
||||
|
||||
async def test_upsert(collection):
|
||||
record = {"id": "testid", "content": "test content", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]}
|
||||
key = await collection.upsert(record)
|
||||
assert key == "testid"
|
||||
assert collection.inner_storage == {"testid": record}
|
||||
|
||||
|
||||
async def test_get(collection):
|
||||
record = {"id": "testid", "content": "test content", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]}
|
||||
await collection.upsert(record)
|
||||
result = await collection.get("testid")
|
||||
assert result["id"] == record["id"]
|
||||
assert result["content"] == record["content"]
|
||||
|
||||
|
||||
async def test_get_missing(collection):
|
||||
result = await collection.get("testid")
|
||||
assert result is None
|
||||
|
||||
|
||||
async def test_delete(collection):
|
||||
record = {"id": "testid", "content": "test content", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]}
|
||||
await collection.upsert(record)
|
||||
await collection.delete("testid")
|
||||
assert collection.inner_storage == {}
|
||||
|
||||
|
||||
async def test_collection_exists(collection):
|
||||
assert await collection.collection_exists() is True
|
||||
|
||||
|
||||
async def test_ensure_collection_deleted(collection):
|
||||
record = {"id": "testid", "content": "test content", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]}
|
||||
await collection.upsert(record)
|
||||
assert collection.inner_storage == {"testid": record}
|
||||
await collection.ensure_collection_deleted()
|
||||
assert collection.inner_storage == {}
|
||||
|
||||
|
||||
async def test_ensure_collection_exists(collection):
|
||||
await collection.ensure_collection_exists()
|
||||
|
||||
|
||||
@mark.parametrize(
|
||||
"distance_function",
|
||||
[
|
||||
DistanceFunction.COSINE_DISTANCE,
|
||||
DistanceFunction.COSINE_SIMILARITY,
|
||||
DistanceFunction.EUCLIDEAN_DISTANCE,
|
||||
DistanceFunction.MANHATTAN,
|
||||
DistanceFunction.EUCLIDEAN_SQUARED_DISTANCE,
|
||||
DistanceFunction.DOT_PROD,
|
||||
DistanceFunction.HAMMING,
|
||||
],
|
||||
)
|
||||
async def test_vectorized_search_similar(collection, distance_function):
|
||||
for field in collection.definition.fields:
|
||||
if field.name == "vector":
|
||||
field.distance_function = distance_function
|
||||
record1 = {"id": "testid1", "content": "test content", "vector": [1.0, 1.0, 1.0, 1.0, 1.0]}
|
||||
record2 = {"id": "testid2", "content": "test content", "vector": [-1.0, -1.0, -1.0, -1.0, -1.0]}
|
||||
await collection.upsert([record1, record2])
|
||||
results = await collection.search(
|
||||
vector=[0.9, 0.9, 0.9, 0.9, 0.9],
|
||||
vector_property_name="vector",
|
||||
include_total_count=True,
|
||||
include_vectors=True,
|
||||
)
|
||||
assert results.total_count == 2
|
||||
idx = 0
|
||||
async for res in results.results:
|
||||
assert res.record == record1 if idx == 0 else record2
|
||||
idx += 1
|
||||
|
||||
|
||||
async def test_valid_lambda_filter(collection):
|
||||
record1 = {"id": "1", "vector": [1, 2, 3, 4, 5]}
|
||||
record2 = {"id": "2", "vector": [5, 4, 3, 2, 1]}
|
||||
await collection.upsert([record1, record2])
|
||||
# Filter to select only record with id == '1'
|
||||
results = collection._get_filtered_records(type("opt", (), {"filter": "lambda x: x.id == '1'"})())
|
||||
assert len(results) == 1
|
||||
assert "1" in results
|
||||
|
||||
|
||||
async def test_valid_lambda_filter_attribute_access(collection):
|
||||
record1 = {"id": "1", "vector": [1, 2, 3, 4, 5]}
|
||||
record2 = {"id": "2", "vector": [5, 4, 3, 2, 1]}
|
||||
await collection.upsert([record1, record2])
|
||||
# Filter to select only record with id == '2' using attribute access
|
||||
results = collection._get_filtered_records(type("opt", (), {"filter": "lambda x: x['id'] == '2'"})())
|
||||
assert len(results) == 1
|
||||
assert "2" in results
|
||||
|
||||
|
||||
async def test_invalid_filter_not_lambda(collection):
|
||||
with raises(VectorStoreOperationException, match="must be a lambda expression"):
|
||||
collection._get_filtered_records(type("opt", (), {"filter": "x.id == '1'"})())
|
||||
|
||||
|
||||
async def test_invalid_filter_syntax(collection):
|
||||
with raises(VectorStoreOperationException, match="not valid Python"):
|
||||
collection._get_filtered_records(type("opt", (), {"filter": "lambda x: x.id == '1' and"})())
|
||||
|
||||
|
||||
async def test_malicious_filter_import(collection):
|
||||
# Should not allow import statement
|
||||
with raises(VectorStoreOperationException):
|
||||
collection._get_filtered_records(
|
||||
type("opt", (), {"filter": "lambda x: __import__('os').system('echo malicious')"})()
|
||||
)
|
||||
|
||||
|
||||
async def test_malicious_filter_exec(collection):
|
||||
# Should not allow exec or similar
|
||||
with raises(VectorStoreOperationException):
|
||||
collection._get_filtered_records(type("opt", (), {"filter": "lambda x: exec('print(1)')"})())
|
||||
|
||||
|
||||
async def test_malicious_filter_builtins(collection):
|
||||
# Should not allow access to builtins
|
||||
with raises(VectorStoreOperationException):
|
||||
collection._get_filtered_records(
|
||||
type("opt", (), {"filter": "lambda x: __builtins__.__import__('os').system('echo malicious')"})()
|
||||
)
|
||||
|
||||
|
||||
async def test_malicious_filter_open(collection):
|
||||
# Should not allow open()
|
||||
with raises(VectorStoreOperationException):
|
||||
collection._get_filtered_records(type("opt", (), {"filter": "lambda x: open('somefile.txt', 'w')"})())
|
||||
|
||||
|
||||
async def test_malicious_filter_eval(collection):
|
||||
# Should not allow eval()
|
||||
with raises(VectorStoreOperationException):
|
||||
collection._get_filtered_records(type("opt", (), {"filter": "lambda x: eval('2+2')"})())
|
||||
|
||||
|
||||
async def test_multiple_filters(collection):
|
||||
record1 = {"id": "1", "vector": [1, 2, 3, 4, 5]}
|
||||
record2 = {"id": "2", "vector": [5, 4, 3, 2, 1]}
|
||||
await collection.upsert([record1, record2])
|
||||
filters = ["lambda x: x.id == '1'", "lambda x: x.vector[0] == 1"]
|
||||
results = collection._get_filtered_records(type("opt", (), {"filter": filters})())
|
||||
assert len(results) == 1
|
||||
assert "1" in results
|
||||
|
||||
|
||||
@mark.parametrize(
|
||||
"filter_str",
|
||||
[
|
||||
"lambda x: [x.clear][0]() or True",
|
||||
"lambda x: [x.update][0]({'role': 'admin'}) or True",
|
||||
"lambda x: [x.pop][0]('secret', '') or True",
|
||||
"lambda x: [x.__setitem__][0]('leaked', ['{0.__class__.__mro__}'.format][0](x)) or True",
|
||||
],
|
||||
)
|
||||
def test_malicious_subscript_call_patterns_blocked(collection, filter_str):
|
||||
with raises(VectorStoreOperationException, match="Call target node type 'Subscript' is not allowed"):
|
||||
collection._parse_and_validate_filter(filter_str)
|
||||
|
||||
|
||||
def test_direct_mutating_method_call_remains_blocked(collection):
|
||||
with raises(VectorStoreOperationException, match="Function 'clear' is not allowed"):
|
||||
collection._parse_and_validate_filter("lambda x: x.clear() or True")
|
||||
|
||||
|
||||
@mark.parametrize(
|
||||
"attr",
|
||||
[
|
||||
"__base__",
|
||||
"__bases__",
|
||||
"__class__",
|
||||
"__mro__",
|
||||
"__subclasses__",
|
||||
"__globals__",
|
||||
],
|
||||
)
|
||||
def test_blocked_dunder_attributes_rejected(collection, attr):
|
||||
with raises(VectorStoreOperationException, match=f"Access to attribute '{attr}' is not allowed"):
|
||||
collection._parse_and_validate_filter(f"lambda x: x.{attr}")
|
||||
|
||||
|
||||
async def test_valid_lambda_filter_with_get_method(collection):
|
||||
record1 = {"id": "1", "vector": [1, 2, 3, 4, 5]}
|
||||
record2 = {"id": "2", "vector": [5, 4, 3, 2, 1]}
|
||||
await collection.upsert([record1, record2])
|
||||
results = collection._get_filtered_records(type("opt", (), {"filter": "lambda x: x.get('id') == '1'"})())
|
||||
assert len(results) == 1
|
||||
assert "1" in results
|
||||
|
||||
|
||||
async def test_valid_lambda_filter_with_bounded_sequence_repeat(collection):
|
||||
record = {"id": "1", "vector": [1, 2, 3, 4, 5]}
|
||||
await collection.upsert(record)
|
||||
|
||||
results = collection._get_filtered_records(type("opt", (), {"filter": "lambda x: ([0] * 2)[1] == 0"})())
|
||||
|
||||
assert len(results) == 1
|
||||
assert "1" in results
|
||||
|
||||
|
||||
async def test_sequence_repeat_limit_can_be_overridden(collection):
|
||||
record = {"id": "1", "vector": [1, 2, 3, 4, 5]}
|
||||
await collection.upsert(record)
|
||||
filter_options = type("opt", (), {"filter": "lambda x: ([0] * 2)[1] == 0"})()
|
||||
|
||||
collection.max_filter_sequence_repeat_size = 1
|
||||
with raises(VectorStoreOperationException, match="Sequence repetition in filter expressions exceeds the maximum"):
|
||||
collection._get_filtered_records(filter_options)
|
||||
|
||||
collection.max_filter_sequence_repeat_size = 2
|
||||
results = collection._get_filtered_records(filter_options)
|
||||
|
||||
assert len(results) == 1
|
||||
assert "1" in results
|
||||
|
||||
|
||||
async def test_callable_filter_cannot_mutate_stored_record(collection):
|
||||
record = {"id": "1", "content": "value", "vector": [1, 2, 3, 4, 5]}
|
||||
await collection.upsert(record)
|
||||
|
||||
def mutating_filter(x):
|
||||
x["role"] = "admin"
|
||||
return True
|
||||
|
||||
with raises(VectorStoreOperationException, match="Error running filter"):
|
||||
collection._get_filtered_records(type("opt", (), {"filter": mutating_filter})())
|
||||
|
||||
assert "role" not in collection.inner_storage["1"]
|
||||
assert collection.inner_storage["1"]["content"] == "value"
|
||||
|
||||
|
||||
def test_default_dynamic_filter_injection_payload_remains_string_literal(collection):
|
||||
class Param:
|
||||
def __init__(self, name, default_value=None):
|
||||
self.name = name
|
||||
self.default_value = default_value
|
||||
|
||||
injected_value = "' or [x.update][0]({'role':'admin'}) or x.name=='"
|
||||
generated_filter = default_dynamic_filter_function(
|
||||
filter=None,
|
||||
parameters=[Param("category")],
|
||||
category=injected_value,
|
||||
)
|
||||
|
||||
assert isinstance(generated_filter, str)
|
||||
tree = ast.parse(generated_filter, mode="eval")
|
||||
assert isinstance(tree.body, ast.Lambda)
|
||||
assert isinstance(tree.body.body, ast.Compare)
|
||||
assert isinstance(tree.body.body.comparators[0], ast.Constant)
|
||||
assert tree.body.body.comparators[0].value == injected_value
|
||||
|
||||
filter_func = collection._parse_and_validate_filter(generated_filter)
|
||||
assert filter_func({"category": "finance", "name": "alice", "vector": [0.1] * 5}) is False
|
||||
|
||||
|
||||
async def test_large_sequence_repeat_filter_is_blocked(collection):
|
||||
record = {"id": "1", "content": "value", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]}
|
||||
await collection.upsert(record)
|
||||
|
||||
with raises(VectorStoreOperationException, match="Sequence repetition in filter expressions exceeds the maximum"):
|
||||
collection._get_filtered_records(type("opt", (), {"filter": "lambda x: [0] * 2000000000"})())
|
||||
@@ -0,0 +1,421 @@
|
||||
# Copyright (c) 2025, Oracle Corporation. All rights reserved. # noqa: CPY001
|
||||
|
||||
from array import array
|
||||
from dataclasses import dataclass
|
||||
from types import SimpleNamespace
|
||||
from typing import Annotated
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import oracledb
|
||||
import pandas as pd
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from semantic_kernel.connectors.oracle import OracleCollection, OracleStore
|
||||
from semantic_kernel.data.vector import (
|
||||
DistanceFunction,
|
||||
IndexKind,
|
||||
VectorStoreCollectionDefinition,
|
||||
VectorStoreField,
|
||||
vectorstoremodel,
|
||||
)
|
||||
|
||||
|
||||
@vectorstoremodel
|
||||
@dataclass
|
||||
class SimpleModel:
|
||||
id: Annotated[int, VectorStoreField("key")]
|
||||
vector: Annotated[
|
||||
list[float] | None,
|
||||
VectorStoreField(
|
||||
"vector",
|
||||
type="float",
|
||||
dimensions=3,
|
||||
index_kind=IndexKind.HNSW,
|
||||
distance_function=DistanceFunction.EUCLIDEAN_SQUARED_DISTANCE,
|
||||
),
|
||||
] = None
|
||||
|
||||
|
||||
def PandasDataframeModel(record) -> tuple:
|
||||
definition = VectorStoreCollectionDefinition(
|
||||
fields=[
|
||||
VectorStoreField("key", name="id", type="int"),
|
||||
VectorStoreField(
|
||||
"vector",
|
||||
name="embedding",
|
||||
type="float32",
|
||||
dimensions=5,
|
||||
distance_function=DistanceFunction.COSINE_DISTANCE,
|
||||
index_kind=IndexKind.IVF_FLAT,
|
||||
),
|
||||
],
|
||||
to_dict=lambda record, **_: record.to_dict(orient="records"),
|
||||
from_dict=lambda records, **_: pd.DataFrame(records),
|
||||
container_mode=True,
|
||||
)
|
||||
df = pd.DataFrame([record]) if isinstance(record, dict) else pd.DataFrame(record)
|
||||
return definition, df
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def mock_connection_pool():
|
||||
mock_conn = AsyncMock()
|
||||
mock_conn.fetchall = AsyncMock(return_value=[("COLL1",), ("COLL2",)])
|
||||
|
||||
mock_context_manager = AsyncMock()
|
||||
mock_context_manager.__aenter__.return_value = mock_conn
|
||||
mock_context_manager.__aexit__.return_value = None
|
||||
|
||||
pool = MagicMock(spec=oracledb.AsyncConnectionPool)
|
||||
pool.acquire.return_value = mock_context_manager
|
||||
|
||||
return pool
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def oracle_store(mock_connection_pool):
|
||||
return OracleStore(
|
||||
connection_pool=mock_connection_pool,
|
||||
db_schema="MY_SCHEMA",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_collection_names_with_schema(oracle_store):
|
||||
names = await oracle_store.list_collection_names()
|
||||
assert names == ["COLL1", "COLL2"]
|
||||
|
||||
|
||||
def test_get_collection_returns_oracle_collection(oracle_store):
|
||||
collection = oracle_store.get_collection(
|
||||
SimpleModel,
|
||||
collection_name="TEST",
|
||||
)
|
||||
assert isinstance(collection, OracleCollection)
|
||||
assert collection.collection_name == "TEST"
|
||||
assert collection.db_schema == "MY_SCHEMA"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_collection_exists_true(oracle_store, mock_connection_pool):
|
||||
conn = AsyncMock()
|
||||
conn.fetchone = AsyncMock(return_value=(1,))
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = conn
|
||||
|
||||
collection = oracle_store.get_collection(
|
||||
SimpleModel,
|
||||
collection_name="EXISTING",
|
||||
)
|
||||
result = await collection.collection_exists()
|
||||
assert result is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_collection_exists_false(oracle_store, mock_connection_pool):
|
||||
conn = AsyncMock()
|
||||
conn.fetchone = AsyncMock(return_value=None)
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = conn
|
||||
|
||||
collection = oracle_store.get_collection(
|
||||
SimpleModel,
|
||||
collection_name="MISSING",
|
||||
)
|
||||
result = await collection.collection_exists()
|
||||
assert result is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_collection_exists_creates_when_missing(oracle_store, mock_connection_pool):
|
||||
conn = AsyncMock()
|
||||
conn.fetchone = AsyncMock(return_value=False)
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = conn
|
||||
collection = oracle_store.get_collection(SimpleModel, collection_name="NEW")
|
||||
await collection.ensure_collection_exists()
|
||||
conn.execute.assert_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_table_with_get_collection(oracle_store, mock_connection_pool):
|
||||
mock_conn = AsyncMock()
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = mock_conn
|
||||
|
||||
collection = oracle_store.get_collection(SimpleModel, collection_name="MY_COLLECTION")
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
|
||||
mock_conn.execute.assert_awaited()
|
||||
sql_statements = [args[0].lower() for name, args, _ in mock_conn.mock_calls if name == "execute"]
|
||||
|
||||
assert any("create table" in sql for sql in sql_statements)
|
||||
assert any("my_collection" in sql for sql in sql_statements)
|
||||
assert any("vector(3 , float64)" in sql for sql in sql_statements)
|
||||
|
||||
mock_conn.commit.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pandasDataframe_with_get_collection(oracle_store, mock_connection_pool):
|
||||
mock_conn = AsyncMock()
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = mock_conn
|
||||
|
||||
records = {
|
||||
"id": 1,
|
||||
"embedding": [1.1, 2.2, 3.3],
|
||||
}
|
||||
|
||||
definition, _ = PandasDataframeModel(records)
|
||||
|
||||
collection = oracle_store.get_collection(
|
||||
collection_name="MY_COLLECTION",
|
||||
record_type=pd.DataFrame,
|
||||
definition=definition,
|
||||
)
|
||||
await collection.ensure_collection_exists()
|
||||
|
||||
mock_conn.execute.assert_awaited()
|
||||
sql_statements = [args[0].lower() for name, args, _ in mock_conn.mock_calls if name == "execute"]
|
||||
|
||||
assert any("create table" in sql for sql in sql_statements)
|
||||
assert any("my_collection" in sql for sql in sql_statements)
|
||||
assert any("vector(5 , float32)" in sql for sql in sql_statements)
|
||||
|
||||
mock_conn.commit.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_index_creation_distance_with_get_collection(oracle_store, mock_connection_pool):
|
||||
mock_conn = AsyncMock()
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = mock_conn
|
||||
|
||||
collection = oracle_store.get_collection(SimpleModel, collection_name="COLLECTION_WITH_INDEX")
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
|
||||
mock_conn.execute.assert_awaited()
|
||||
called_sql = mock_conn.execute.call_args[0][0].lower()
|
||||
|
||||
assert "create vector index" in called_sql
|
||||
assert "collection_with_index_vector_idx" in called_sql
|
||||
assert "inmemory neighbor graph" in called_sql
|
||||
assert "distance euclidean_squared" in called_sql
|
||||
|
||||
mock_conn.commit.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_collection_deleted(oracle_store, mock_connection_pool):
|
||||
conn = AsyncMock()
|
||||
conn.fetchone = AsyncMock(return_value=(1,))
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = conn
|
||||
|
||||
collection = oracle_store.get_collection(SimpleModel, collection_name="TO_DELETE")
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
await collection.ensure_collection_deleted()
|
||||
|
||||
assert any("DROP TABLE" in str(call.args[0]) for call in conn.execute.call_args_list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert(oracle_store, mock_connection_pool):
|
||||
mock_conn = AsyncMock()
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = mock_conn
|
||||
|
||||
collection = oracle_store.get_collection(SimpleModel, collection_name="MY_COLLECTION")
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
await collection.upsert(SimpleModel(id=1, vector=[0.1, 0.2, 0.3]))
|
||||
|
||||
mock_conn.executemany.assert_called_once()
|
||||
|
||||
merge_sql, params = mock_conn.executemany.call_args[0]
|
||||
|
||||
assert merge_sql.startswith('MERGE INTO "MY_SCHEMA"."MY_COLLECTION"')
|
||||
assert 'UPDATE SET t."vector"' in merge_sql
|
||||
assert "WHEN NOT MATCHED THEN" in merge_sql
|
||||
assert 'INSERT ("id", "vector")' in merge_sql
|
||||
|
||||
expected_param = (1, array("d", [0.1, 0.2, 0.3]))
|
||||
assert params[0] == expected_param
|
||||
assert mock_conn.commit.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_with_include_vectors(oracle_store, mock_connection_pool):
|
||||
mock_conn = AsyncMock()
|
||||
mock_conn.description = [("id",), ("vector",)]
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = mock_conn
|
||||
|
||||
collection = oracle_store.get_collection(SimpleModel, collection_name="MY_COLLECTION")
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
await collection.upsert(SimpleModel(id=1, vector=[0.1, 0.2, 0.3]))
|
||||
|
||||
mock_conn.fetchall.return_value = [(1, [0.1, 0.2, 0.3])]
|
||||
results = await collection.get([1], include_vectors=True)
|
||||
assert results.id == 1
|
||||
assert results.vector == [0.1, 0.2, 0.3]
|
||||
|
||||
executed_sql = [args[0] for args, _ in mock_conn.fetchall.call_args_list]
|
||||
assert any('SELECT "id" AS "id", "vector" AS "vector"' in sql for sql in executed_sql), (
|
||||
"Expected vector column to be selected when include_vectors=True"
|
||||
)
|
||||
|
||||
mock_conn.fetchall.reset_mock()
|
||||
mock_conn.fetchall.return_value = [(1, None)]
|
||||
results = await collection.get([1], include_vectors=False)
|
||||
assert results.id == 1
|
||||
assert results.vector is None
|
||||
|
||||
executed_sql = [args[0] for args, _ in mock_conn.fetchall.call_args_list]
|
||||
assert any('SELECT "id" AS "id", "vector" AS "vector"' not in sql for sql in executed_sql), (
|
||||
"Vector column should not be selected when include_vectors=False"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_and_get(oracle_store, mock_connection_pool):
|
||||
conn = AsyncMock()
|
||||
conn.description = [("id",), ("vector",)]
|
||||
conn.fetchall = AsyncMock(return_value=[(1, [0.1, 0.2, 0.3])])
|
||||
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = conn
|
||||
|
||||
collection = oracle_store.get_collection(SimpleModel, collection_name="TEST")
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
await collection.upsert(SimpleModel(id=1, vector=[0.1, 0.2, 0.3]))
|
||||
assert conn.executemany.await_count >= 1 or conn.execute.await_count >= 1, (
|
||||
"Expected upsert to call executemany() or execute()"
|
||||
)
|
||||
|
||||
results = await collection.get([1], include_vectors=True)
|
||||
|
||||
assert results.id == 1
|
||||
assert results.vector == [0.1, 0.2, 0.3]
|
||||
conn.fetchall.assert_awaited_once()
|
||||
conn.fetchall.reset_mock()
|
||||
conn.fetchall.return_value = [(1, None)]
|
||||
conn.description = [("id",)]
|
||||
|
||||
results = await collection.get([1], include_vectors=False)
|
||||
assert results.id == 1
|
||||
assert results.vector is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_record(oracle_store, mock_connection_pool):
|
||||
conn = AsyncMock()
|
||||
mock_connection_pool.acquire.return_value.__aenter__.return_value = conn
|
||||
|
||||
collection = oracle_store.get_collection(
|
||||
SimpleModel,
|
||||
collection_name="MY_COLLECTION",
|
||||
)
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
await collection.delete([1])
|
||||
conn.executemany.assert_awaited()
|
||||
called_sql = conn.executemany.call_args[0][0].lower()
|
||||
assert "delete" in called_sql
|
||||
assert "my_collection" in called_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search(oracle_store, mock_connection_pool):
|
||||
class MockCursor:
|
||||
def __init__(self):
|
||||
self.execute_called_with = []
|
||||
self._rows = [(1, [0.1, 0.2, 0.3])]
|
||||
self.description = [SimpleNamespace(name="id"), SimpleNamespace(name="vector")]
|
||||
self._i = 0
|
||||
|
||||
async def execute(self, sql, binds=None):
|
||||
self.execute_called_with.append((sql, binds))
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
return None
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
return None
|
||||
|
||||
def __aiter__(self):
|
||||
self._i = 0
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._i >= len(self._rows):
|
||||
raise StopAsyncIteration
|
||||
|
||||
r = self._rows[self._i]
|
||||
self._i += 1
|
||||
return r
|
||||
|
||||
mock_cursor = MockCursor()
|
||||
|
||||
class MockConnection:
|
||||
def __init__(self, cur):
|
||||
self._cur = cur
|
||||
self.inputtypehandler = None
|
||||
self.outputtypehandler = None
|
||||
self.execute = AsyncMock()
|
||||
self.commit = AsyncMock()
|
||||
|
||||
def cursor(self):
|
||||
return self._cur
|
||||
|
||||
mock_conn = MockConnection(mock_cursor)
|
||||
|
||||
class MockAcquire:
|
||||
def __init__(self, conn):
|
||||
self._conn = conn
|
||||
|
||||
def __await__(self):
|
||||
async def _():
|
||||
return self
|
||||
|
||||
return _().__await__()
|
||||
|
||||
async def __aenter__(self):
|
||||
return self._conn
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
return None
|
||||
|
||||
mock_connection_pool.acquire = lambda **kwargs: MockAcquire(mock_conn)
|
||||
|
||||
collection = oracle_store.get_collection(
|
||||
model=SimpleModel, record_type=SimpleModel, collection_name="MY_COLLECTION"
|
||||
)
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
assert mock_conn.execute.await_count >= 1
|
||||
|
||||
ks_results = await collection.search(
|
||||
vector_property_name="vector",
|
||||
vector=[0.1, 0.2, 0.3],
|
||||
top=1,
|
||||
filter=lambda x: x.id in [1, 7, 9],
|
||||
include_vectors=True,
|
||||
)
|
||||
results = [r async for r in ks_results.results]
|
||||
|
||||
assert mock_cursor.execute_called_with, "cursor.execute was not called"
|
||||
sql, binds = mock_cursor.execute_called_with[0]
|
||||
assert "SELECT" in sql.upper()
|
||||
assert '"id", "vector", VECTOR_DISTANCE' in sql
|
||||
expected_where = 'WHERE "id" IN (:bind_val1, :bind_val2, :bind_val3)'
|
||||
assert expected_where in sql
|
||||
|
||||
assert binds is None or isinstance(binds[0], array)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].record.id == 1
|
||||
assert results[0].record.vector == [0.1, 0.2, 0.3]
|
||||
@@ -0,0 +1,337 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
from pinecone import FetchResponse, IndexModel, Metric, QueryResponse, ServerlessSpec, Vector
|
||||
from pinecone.core.openapi.db_data.models import (
|
||||
Hit,
|
||||
ScoredVector,
|
||||
SearchRecordsResponse,
|
||||
SearchRecordsResponseResult,
|
||||
SearchUsage,
|
||||
)
|
||||
from pinecone.db_data.index_asyncio import _IndexAsyncio
|
||||
from pytest import fixture, mark, raises
|
||||
|
||||
from semantic_kernel.connectors.pinecone import PineconeCollection, PineconeStore
|
||||
from semantic_kernel.exceptions.vector_store_exceptions import VectorStoreInitializationException
|
||||
|
||||
BASE_PATH_ASYNCIO = "pinecone.PineconeAsyncio"
|
||||
BASE_PATH_INDEX_CLIENT_ASYNCIO = "pinecone.db_data.index_asyncio._IndexAsyncio"
|
||||
|
||||
|
||||
@fixture
|
||||
def embed(request) -> dict[str, Any] | None:
|
||||
if hasattr(request, "param"):
|
||||
return request.param
|
||||
return None
|
||||
|
||||
|
||||
@fixture
|
||||
def mock_index_model(embed: dict[str, Any] | None):
|
||||
"""Mock IndexModel for testing."""
|
||||
mock_index_model = Mock(spec=IndexModel)
|
||||
mock_index_model.name = "test"
|
||||
mock_index_model.embed = embed
|
||||
mock_index_model.host = "test_host"
|
||||
return mock_index_model
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_list_collection_names(mock_index_model):
|
||||
with patch(f"{BASE_PATH_ASYNCIO}.list_indexes") as mock_list_indexes:
|
||||
mock_list_indexes.return_value = [mock_index_model]
|
||||
yield mock_list_indexes
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_create_index(mock_index_model):
|
||||
with patch(f"{BASE_PATH_ASYNCIO}.create_index") as mock_create_index:
|
||||
mock_create_index.return_value = mock_index_model
|
||||
yield mock_create_index
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_create_index_for_model(mock_index_model):
|
||||
with patch(f"{BASE_PATH_ASYNCIO}.create_index_for_model") as mock_create_index_for_model:
|
||||
mock_create_index_for_model.return_value = mock_index_model
|
||||
yield mock_create_index_for_model
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_describe_index(mock_index_model):
|
||||
with patch(f"{BASE_PATH_ASYNCIO}.describe_index") as mock_describe_index:
|
||||
mock_describe_index.return_value = mock_index_model
|
||||
yield mock_describe_index
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_has_index():
|
||||
with patch(f"{BASE_PATH_ASYNCIO}.has_index") as mock_has_index:
|
||||
mock_create_index.return_value = True
|
||||
yield mock_has_index
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_index_asyncio():
|
||||
mock_index_asyncio = AsyncMock(spec=_IndexAsyncio)
|
||||
mock_index_asyncio.close.return_value = None
|
||||
with patch(f"{BASE_PATH_ASYNCIO}.IndexAsyncio") as mock_index:
|
||||
mock_index.return_value = mock_index_asyncio
|
||||
yield mock_index
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_delete_index():
|
||||
with patch(f"{BASE_PATH_ASYNCIO}.delete_index") as mock_delete:
|
||||
yield mock_delete
|
||||
|
||||
|
||||
@fixture
|
||||
async def store(pinecone_unit_test_env) -> PineconeStore:
|
||||
"""Fixture to create a Pinecone store."""
|
||||
async with PineconeStore() as store:
|
||||
yield store
|
||||
|
||||
|
||||
@fixture
|
||||
async def collection(pinecone_unit_test_env, definition) -> PineconeCollection:
|
||||
"""Fixture to create a Pinecone store."""
|
||||
async with PineconeCollection(
|
||||
collection_name="test_collection",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
) as collection:
|
||||
yield collection
|
||||
|
||||
|
||||
async def test_create_store(pinecone_unit_test_env):
|
||||
"""Test the creation of a Pinecone store."""
|
||||
# Create a Pinecone store
|
||||
store = PineconeStore()
|
||||
assert store is not None
|
||||
assert store.client is not None
|
||||
|
||||
|
||||
@mark.parametrize("exclude_list", [["PINECONE_API_KEY"]], indirect=True)
|
||||
async def test_create_store_fail(pinecone_unit_test_env):
|
||||
"""Test the creation of a Pinecone store."""
|
||||
with raises(VectorStoreInitializationException):
|
||||
PineconeStore(env_file_path="test.env")
|
||||
|
||||
|
||||
def test_create_store_grpc(pinecone_unit_test_env):
|
||||
"""Test the creation of a Pinecone store."""
|
||||
|
||||
# Create a Pinecone store
|
||||
store = PineconeStore(use_grpc=True)
|
||||
assert store is not None
|
||||
assert store.client is not None
|
||||
|
||||
|
||||
@mark.parametrize("exclude_list", [["PINECONE_API_KEY"]], indirect=True)
|
||||
async def test_ensure_collection_exists_fail(pinecone_unit_test_env, definition):
|
||||
with raises(VectorStoreInitializationException):
|
||||
PineconeCollection(
|
||||
collection_name="test_collection",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
async def test_get_collection(store: PineconeStore, definition):
|
||||
"""Test the creation of a Pinecone collection."""
|
||||
# Create a collection
|
||||
collection = store.get_collection(collection_name="test_collection", record_type=dict, definition=definition)
|
||||
assert collection is not None
|
||||
assert collection.collection_name == "test_collection"
|
||||
|
||||
|
||||
async def test_list_collection_names(store: PineconeStore):
|
||||
"""Test the listing of Pinecone collections."""
|
||||
# List collections
|
||||
collections = await store.list_collection_names()
|
||||
assert collections is not None
|
||||
assert len(collections) == 1
|
||||
assert collections[0] == "test"
|
||||
|
||||
|
||||
@mark.parametrize("embed", [None, {"model": "test-model"}])
|
||||
async def test_load_index_client(collection, mock_index_asyncio):
|
||||
# Test loading the index client
|
||||
await collection._load_index_client()
|
||||
assert collection.index is not None
|
||||
assert collection.index_client is not None
|
||||
assert isinstance(collection.index_client, _IndexAsyncio)
|
||||
assert collection.embed_settings == collection.index.embed
|
||||
|
||||
|
||||
async def test_ensure_collection_exists(collection, mock_create_index):
|
||||
await collection.ensure_collection_exists()
|
||||
assert collection.index is not None
|
||||
assert collection.index_client is not None
|
||||
mock_create_index.assert_awaited_once_with(
|
||||
name=collection.collection_name,
|
||||
spec=ServerlessSpec(cloud="aws", region="us-east-1"),
|
||||
dimension=5,
|
||||
metric=Metric.COSINE,
|
||||
vector_type="dense",
|
||||
)
|
||||
|
||||
|
||||
@mark.parametrize("embed", [{"model": "test-model"}])
|
||||
async def test_ensure_collection_exists_integrated(collection, mock_create_index_for_model):
|
||||
await collection.ensure_collection_exists(embed={"model": "test-model"})
|
||||
assert collection.index is not None
|
||||
assert collection.index_client is not None
|
||||
mock_create_index_for_model.assert_awaited_once_with(
|
||||
name=collection.collection_name,
|
||||
cloud="aws",
|
||||
region="us-east-1",
|
||||
embed={"model": "test-model", "metric": Metric.COSINE, "field_map": {"text": "vector"}},
|
||||
)
|
||||
|
||||
|
||||
async def test_ensure_collection_deleted(collection):
|
||||
# Test deleting the collection
|
||||
await collection.ensure_collection_deleted()
|
||||
assert collection.index is None
|
||||
assert collection.index_client is None
|
||||
|
||||
|
||||
async def test_upsert(collection):
|
||||
record = {
|
||||
"id": "test_id",
|
||||
"vector": [0.1, 0.2, 0.3, 0.4, 0.5],
|
||||
"content": "test_content",
|
||||
}
|
||||
pinecone_vector = Vector(values=record["vector"], id=record["id"], metadata={"content": record["content"]})
|
||||
await collection._load_index_client()
|
||||
with patch.object(collection.index_client, "upsert", new_callable=AsyncMock) as mock_upsert:
|
||||
await collection.upsert(record)
|
||||
mock_upsert.assert_awaited_once_with(
|
||||
[pinecone_vector],
|
||||
namespace=collection.namespace,
|
||||
)
|
||||
|
||||
|
||||
@mark.parametrize("embed", [{"model": "test-model"}])
|
||||
async def test_upsert_embed(collection):
|
||||
record = {
|
||||
"id": "test_id",
|
||||
"content": "test_content",
|
||||
"vector": [0.1, 0.2, 0.3, 0.4, 0.5],
|
||||
}
|
||||
await collection._load_index_client()
|
||||
with patch.object(collection.index_client, "upsert_records", new_callable=AsyncMock) as mock_upsert:
|
||||
await collection.upsert(record)
|
||||
mock_upsert.assert_awaited_once_with(
|
||||
records=[{"_id": record["id"], "content": record["content"]}],
|
||||
namespace=collection.namespace,
|
||||
)
|
||||
|
||||
|
||||
async def test_get(collection):
|
||||
record = {
|
||||
"id": "test_id",
|
||||
"vector": [0.1, 0.2, 0.3, 0.4, 0.5],
|
||||
"content": "test_content",
|
||||
}
|
||||
fetch_response = FetchResponse(
|
||||
namespace="",
|
||||
vectors={
|
||||
record["id"]: Vector(values=record["vector"], id=record["id"], metadata={"content": record["content"]})
|
||||
},
|
||||
usage={},
|
||||
)
|
||||
await collection._load_index_client()
|
||||
with patch.object(collection.index_client, "fetch", new_callable=AsyncMock) as mock_fetch:
|
||||
mock_fetch.return_value = fetch_response
|
||||
get_record = await collection.get(record["id"])
|
||||
mock_fetch.assert_awaited_once_with(
|
||||
ids=[record["id"]],
|
||||
namespace=collection.namespace,
|
||||
)
|
||||
assert record["id"] == get_record["id"]
|
||||
assert record["content"] == get_record["content"]
|
||||
|
||||
|
||||
async def test_delete(collection):
|
||||
await collection._load_index_client()
|
||||
with patch.object(collection.index_client, "delete", new_callable=AsyncMock) as mock_delete:
|
||||
await collection.delete("test_id")
|
||||
mock_delete.assert_awaited_once_with(
|
||||
ids=["test_id"],
|
||||
namespace=collection.namespace,
|
||||
)
|
||||
|
||||
|
||||
async def test_search(collection):
|
||||
record = {
|
||||
"id": "test_id",
|
||||
"vector": [0.1, 0.2, 0.3, 0.4, 0.5],
|
||||
"content": "test_content",
|
||||
}
|
||||
query_response = QueryResponse._from_openapi_data(
|
||||
namespace="",
|
||||
matches=[
|
||||
ScoredVector(**{
|
||||
"values": record["vector"],
|
||||
"id": record["id"],
|
||||
"metadata": {"content": record["content"]},
|
||||
"score": 0.1,
|
||||
})
|
||||
],
|
||||
)
|
||||
await collection._load_index_client()
|
||||
with patch.object(collection.index_client, "query", new_callable=AsyncMock) as mock_query:
|
||||
mock_query.return_value = query_response
|
||||
query_response = await collection.search(
|
||||
vector=[0.1, 0.2, 0.3, 0.4, 0.5],
|
||||
top=1,
|
||||
include_vectors=True,
|
||||
filter=lambda x: x.content == "test_content",
|
||||
)
|
||||
mock_query.assert_awaited_once_with(
|
||||
vector=[0.1, 0.2, 0.3, 0.4, 0.5],
|
||||
top_k=1,
|
||||
include_metadata=True,
|
||||
include_values=True,
|
||||
namespace=collection.namespace,
|
||||
filter={"content": "test_content"},
|
||||
)
|
||||
assert query_response.total_count == 1
|
||||
async for result in query_response.results:
|
||||
assert result.record == record
|
||||
assert result.score == 0.1
|
||||
|
||||
|
||||
@mark.parametrize("embed", [{"model": "test-model"}])
|
||||
async def test_search_embed(collection):
|
||||
record = {"id": "test_id", "content": "test_content", "vector": None}
|
||||
query_response = SearchRecordsResponse._from_openapi_data(
|
||||
result=SearchRecordsResponseResult._from_openapi_data(**{
|
||||
"hits": [
|
||||
Hit(**{
|
||||
"_id": record["id"],
|
||||
"fields": {"id": record["id"], "content": record["content"]},
|
||||
"_score": 0.1,
|
||||
})
|
||||
]
|
||||
}),
|
||||
usage=SearchUsage(read_units=0),
|
||||
)
|
||||
await collection._load_index_client()
|
||||
with patch.object(collection.index_client, "search_records", new_callable=AsyncMock) as mock_query:
|
||||
mock_query.return_value = query_response
|
||||
query_response = await collection.search(values="test", top=1, include_vectors=True)
|
||||
mock_query.assert_awaited_once_with(
|
||||
query={"inputs": {"text": "test"}, "top_k": 1},
|
||||
namespace=collection.namespace,
|
||||
)
|
||||
assert query_response.total_count == 1
|
||||
async for result in query_response.results:
|
||||
assert result.record == record
|
||||
assert result.score == 0.1
|
||||
@@ -0,0 +1,415 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from dataclasses import dataclass
|
||||
from typing import Annotated, Any
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from psycopg import AsyncConnection, AsyncCursor
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
from pytest import fixture
|
||||
|
||||
from semantic_kernel.connectors.postgres import (
|
||||
DISTANCE_COLUMN_NAME,
|
||||
PostgresCollection,
|
||||
PostgresSettings,
|
||||
PostgresStore,
|
||||
)
|
||||
from semantic_kernel.data.vector import DistanceFunction, IndexKind, VectorStoreField, vectorstoremodel
|
||||
|
||||
|
||||
@fixture(scope="function")
|
||||
def mock_cursor():
|
||||
return AsyncMock(spec=AsyncCursor)
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_connection_pool(mock_cursor: Mock):
|
||||
with (
|
||||
patch(
|
||||
f"{AsyncConnectionPool.__module__}.{AsyncConnectionPool.__qualname__}.connection",
|
||||
) as mock_pool_connection,
|
||||
patch(
|
||||
f"{AsyncConnectionPool.__module__}.{AsyncConnectionPool.__qualname__}.open",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_pool_open,
|
||||
):
|
||||
mock_conn = AsyncMock(spec=AsyncConnection)
|
||||
|
||||
mock_pool_connection.return_value.__aenter__.return_value = mock_conn
|
||||
mock_conn.cursor.return_value.__aenter__.return_value = mock_cursor
|
||||
|
||||
mock_pool_open.return_value = None
|
||||
|
||||
yield mock_pool_connection, mock_pool_open
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def vector_store(postgres_unit_test_env) -> AsyncGenerator[PostgresStore, None]:
|
||||
async with await PostgresSettings(env_file_path="test.env").create_connection_pool() as pool:
|
||||
yield PostgresStore(connection_pool=pool)
|
||||
|
||||
|
||||
@vectorstoremodel
|
||||
@dataclass
|
||||
class SimpleDataModel:
|
||||
id: Annotated[int, VectorStoreField("key")]
|
||||
data: Annotated[
|
||||
list[float] | str | None,
|
||||
VectorStoreField(
|
||||
"vector",
|
||||
type="float",
|
||||
dimensions=1536,
|
||||
index_kind=IndexKind.HNSW,
|
||||
distance_function=DistanceFunction.COSINE_SIMILARITY,
|
||||
),
|
||||
] = None
|
||||
|
||||
|
||||
# region VectorStore Tests
|
||||
|
||||
|
||||
async def test_vector_store_defaults(vector_store: PostgresStore) -> None:
|
||||
assert vector_store.connection_pool is not None
|
||||
async with vector_store.connection_pool.connection() as conn:
|
||||
assert isinstance(conn, Mock)
|
||||
|
||||
|
||||
def test_vector_store_with_connection_pool(vector_store: PostgresStore) -> None:
|
||||
connection_pool = MagicMock(spec=AsyncConnectionPool)
|
||||
vector_store = PostgresStore(connection_pool=connection_pool)
|
||||
assert vector_store.connection_pool == connection_pool
|
||||
|
||||
|
||||
async def test_list_collection_names(vector_store: PostgresStore, mock_cursor: Mock) -> None:
|
||||
mock_cursor.fetchall.return_value = [
|
||||
("test_collection",),
|
||||
("test_collection_2",),
|
||||
]
|
||||
names = await vector_store.list_collection_names()
|
||||
assert names == ["test_collection", "test_collection_2"]
|
||||
|
||||
|
||||
def test_get_collection(vector_store: PostgresStore) -> None:
|
||||
collection = vector_store.get_collection(collection_name="test_collection", record_type=SimpleDataModel)
|
||||
assert collection.collection_name == "test_collection"
|
||||
|
||||
|
||||
async def test_collection_exists(vector_store: PostgresStore, mock_cursor: Mock) -> None:
|
||||
mock_cursor.fetchall.return_value = [("test_collection",)]
|
||||
collection = vector_store.get_collection(collection_name="test_collection", record_type=SimpleDataModel)
|
||||
result = await collection.collection_exists()
|
||||
assert result is True
|
||||
|
||||
|
||||
async def test_ensure_collection_deleted(vector_store: PostgresStore, mock_cursor: Mock) -> None:
|
||||
collection = vector_store.get_collection(collection_name="test_collection", record_type=SimpleDataModel)
|
||||
await collection.ensure_collection_deleted()
|
||||
|
||||
assert mock_cursor.execute.call_count == 1
|
||||
execute_args, _ = mock_cursor.execute.call_args
|
||||
statement = execute_args[0]
|
||||
statement_str = statement.as_string()
|
||||
|
||||
assert statement_str == 'DROP TABLE "public"."test_collection" CASCADE'
|
||||
|
||||
|
||||
async def test_delete_records(vector_store: PostgresStore, mock_cursor: Mock) -> None:
|
||||
collection = vector_store.get_collection(collection_name="test_collection", record_type=SimpleDataModel)
|
||||
await collection.delete([1, 2])
|
||||
|
||||
assert mock_cursor.execute.call_count == 1
|
||||
execute_args, _ = mock_cursor.execute.call_args
|
||||
statement = execute_args[0]
|
||||
statement_str = statement.as_string()
|
||||
|
||||
assert statement_str == """DELETE FROM "public"."test_collection" WHERE "id" IN (1, 2)"""
|
||||
|
||||
|
||||
async def test_ensure_collection_exists_simple_model(vector_store: PostgresStore, mock_cursor: Mock) -> None:
|
||||
collection = vector_store.get_collection(collection_name="test_collection", record_type=SimpleDataModel)
|
||||
await collection.ensure_collection_exists()
|
||||
|
||||
# 2 calls, once for the table creation and once for the index creation
|
||||
assert mock_cursor.execute.call_count == 2
|
||||
|
||||
# Check the table creation statement
|
||||
execute_args, _ = mock_cursor.execute.call_args_list[0]
|
||||
statement = execute_args[0]
|
||||
statement_str = statement.as_string()
|
||||
assert statement_str == ('CREATE TABLE "public"."test_collection" ("id" INTEGER PRIMARY KEY, "data" VECTOR(1536))')
|
||||
|
||||
# Check the index creation statement
|
||||
execute_args, _ = mock_cursor.execute.call_args_list[1]
|
||||
statement = execute_args[0]
|
||||
statement_str = statement.as_string()
|
||||
assert statement_str == (
|
||||
'CREATE INDEX "test_collection_data_idx" ON "public"."test_collection" USING hnsw ("data" vector_cosine_ops)'
|
||||
)
|
||||
|
||||
|
||||
async def test_ensure_collection_exists_model_with_python_types(vector_store: PostgresStore, mock_cursor: Mock) -> None:
|
||||
@vectorstoremodel
|
||||
@dataclass
|
||||
class ModelWithImplicitTypes:
|
||||
name: Annotated[str, VectorStoreField("key")]
|
||||
age: Annotated[int, VectorStoreField("data")]
|
||||
data: Annotated[dict[str, Any], VectorStoreField("data")]
|
||||
embedding: Annotated[list[float], VectorStoreField("vector", dimensions=20)]
|
||||
scores: Annotated[list[float], VectorStoreField("data")]
|
||||
tags: Annotated[list[str], VectorStoreField("data")]
|
||||
|
||||
collection = vector_store.get_collection(collection_name="test_collection", record_type=ModelWithImplicitTypes)
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
|
||||
assert mock_cursor.execute.call_count == 2
|
||||
|
||||
# Check the table creation statement
|
||||
execute_args, _ = mock_cursor.execute.call_args_list[0]
|
||||
statement = execute_args[0]
|
||||
statement_str = statement.as_string()
|
||||
assert statement_str == (
|
||||
'CREATE TABLE "public"."test_collection" '
|
||||
'("name" TEXT PRIMARY KEY, "age" INTEGER, "data" JSONB, '
|
||||
'"embedding" VECTOR(20), "scores" DOUBLE PRECISION[], "tags" TEXT[])'
|
||||
)
|
||||
|
||||
# Check the index creation statement
|
||||
execute_args, _ = mock_cursor.execute.call_args_list[1]
|
||||
statement = execute_args[0]
|
||||
statement_str = statement.as_string()
|
||||
assert statement_str == (
|
||||
'CREATE INDEX "test_collection_embedding_idx" ON "public"."test_collection" '
|
||||
'USING hnsw ("embedding" vector_cosine_ops)'
|
||||
)
|
||||
|
||||
|
||||
async def test_upsert_records(vector_store: PostgresStore, mock_cursor: Mock) -> None:
|
||||
collection = vector_store.get_collection(collection_name="test_collection", record_type=SimpleDataModel)
|
||||
await collection.upsert([
|
||||
SimpleDataModel(id=1, data=[1.0, 2.0, 3.0]),
|
||||
SimpleDataModel(id=2, data=[4.0, 5.0, 6.0]),
|
||||
SimpleDataModel(id=3, data=[5.0, 6.0, 1.0]),
|
||||
])
|
||||
|
||||
assert mock_cursor.executemany.call_count == 1
|
||||
execute_args, _ = mock_cursor.executemany.call_args
|
||||
statement_str = execute_args[0].as_string()
|
||||
values = execute_args[1]
|
||||
assert len(values) == 3
|
||||
|
||||
assert statement_str == (
|
||||
'INSERT INTO "public"."test_collection" ("id", "data") '
|
||||
"VALUES (%s, %s) "
|
||||
'ON CONFLICT ("id") DO UPDATE SET "data" = EXCLUDED."data"'
|
||||
)
|
||||
|
||||
assert values[0] == (1, [1.0, 2.0, 3.0])
|
||||
assert values[1] == (2, [4.0, 5.0, 6.0])
|
||||
assert values[2] == (3, [5.0, 6.0, 1.0])
|
||||
|
||||
|
||||
async def test_get_records(vector_store: PostgresStore, mock_cursor: Mock) -> None:
|
||||
mock_cursor.fetchall.return_value = [
|
||||
(1, "[1.0, 2.0, 3.0]", {"key": "value1"}),
|
||||
(2, "[4.0, 5.0, 6.0]", {"key": "value2"}),
|
||||
(3, "[5.0, 6.0, 1.0]", {"key": "value3"}),
|
||||
]
|
||||
|
||||
collection = vector_store.get_collection(collection_name="test_collection", record_type=SimpleDataModel)
|
||||
records = await collection.get([1, 2, 3])
|
||||
|
||||
assert len(records) == 3
|
||||
assert records[0].id == 1
|
||||
assert records[1].id == 2
|
||||
assert records[2].id == 3
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
# region Vector Search tests
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"distance_function, operator, subquery_distance, include_vectors, include_total_count",
|
||||
[
|
||||
(DistanceFunction.COSINE_SIMILARITY, "<=>", f'1 - subquery."{DISTANCE_COLUMN_NAME}"', False, False),
|
||||
(DistanceFunction.COSINE_DISTANCE, "<=>", None, False, False),
|
||||
(DistanceFunction.DOT_PROD, "<#>", f'-1 * subquery."{DISTANCE_COLUMN_NAME}"', True, False),
|
||||
(DistanceFunction.EUCLIDEAN_DISTANCE, "<->", None, False, True),
|
||||
(DistanceFunction.MANHATTAN, "<+>", None, True, True),
|
||||
],
|
||||
)
|
||||
async def test_vector_search(
|
||||
vector_store: PostgresStore,
|
||||
mock_cursor: Mock,
|
||||
distance_function: DistanceFunction,
|
||||
operator: str,
|
||||
subquery_distance: str | None,
|
||||
include_vectors: bool,
|
||||
include_total_count: bool,
|
||||
) -> None:
|
||||
@vectorstoremodel
|
||||
@dataclass
|
||||
class SimpleDataModel:
|
||||
id: Annotated[int, VectorStoreField("key")]
|
||||
embedding: Annotated[
|
||||
list[float] | str | None,
|
||||
VectorStoreField(
|
||||
"vector",
|
||||
index_kind=IndexKind.HNSW,
|
||||
dimensions=1536,
|
||||
distance_function=distance_function,
|
||||
type="float",
|
||||
),
|
||||
]
|
||||
data: Annotated[
|
||||
dict[str, Any],
|
||||
VectorStoreField("data", type="JSONB"),
|
||||
]
|
||||
|
||||
def model_post_init(self, context: Any) -> None:
|
||||
if self.embedding is None:
|
||||
self.embedding = self.data
|
||||
|
||||
collection = vector_store.get_collection(collection_name="test_collection", record_type=SimpleDataModel)
|
||||
assert isinstance(collection, PostgresCollection)
|
||||
|
||||
search_results = await collection.search(
|
||||
vector=[1.0, 2.0, 3.0],
|
||||
top=10,
|
||||
skip=5,
|
||||
include_vectors=include_vectors,
|
||||
include_total_count=include_total_count,
|
||||
)
|
||||
if include_total_count:
|
||||
# Including total count issues query directly
|
||||
assert mock_cursor.execute.call_count == 1
|
||||
else:
|
||||
# Total count is not included, query is issued when iterating over results
|
||||
assert mock_cursor.execute.call_count == 0
|
||||
async for _ in search_results.results:
|
||||
pass
|
||||
assert mock_cursor.execute.call_count == 1
|
||||
|
||||
execute_args, _ = mock_cursor.execute.call_args
|
||||
|
||||
assert (search_results.total_count is not None) == include_total_count
|
||||
|
||||
statement = execute_args[0]
|
||||
statement_str = statement.as_string()
|
||||
|
||||
expected_columns = '"id", "data"'
|
||||
if include_vectors:
|
||||
expected_columns = '"id", "embedding", "data"'
|
||||
|
||||
expected_statement = (
|
||||
f'SELECT {expected_columns}, "embedding" {operator} %s as "{DISTANCE_COLUMN_NAME}" '
|
||||
'FROM "public"."test_collection" '
|
||||
f'ORDER BY "{DISTANCE_COLUMN_NAME}" LIMIT 10 OFFSET 5'
|
||||
)
|
||||
|
||||
if subquery_distance:
|
||||
expected_statement = (
|
||||
f'SELECT subquery.*, {subquery_distance} AS "{DISTANCE_COLUMN_NAME}" FROM ('
|
||||
+ expected_statement
|
||||
+ ") AS subquery"
|
||||
)
|
||||
|
||||
assert statement_str == expected_statement
|
||||
|
||||
|
||||
async def test_model_post_init_conflicting_distance_column_name(vector_store: PostgresStore) -> None:
|
||||
@vectorstoremodel
|
||||
@dataclass
|
||||
class ConflictingDataModel:
|
||||
id: Annotated[int, VectorStoreField("key")]
|
||||
sk_pg_distance: Annotated[
|
||||
float, VectorStoreField("data")
|
||||
] # Note: test depends on value of DISTANCE_COLUMN_NAME constant
|
||||
|
||||
embedding: Annotated[
|
||||
list[float],
|
||||
VectorStoreField(
|
||||
"vector",
|
||||
index_kind=IndexKind.HNSW,
|
||||
dimensions=1536,
|
||||
distance_function=DistanceFunction.COSINE_SIMILARITY,
|
||||
type="float",
|
||||
),
|
||||
]
|
||||
data: Annotated[
|
||||
dict[str, Any],
|
||||
VectorStoreField("data", type="JSONB"),
|
||||
]
|
||||
|
||||
collection = vector_store.get_collection(collection_name="test_collection", record_type=ConflictingDataModel)
|
||||
assert isinstance(collection, PostgresCollection)
|
||||
|
||||
# Ensure that the distance column name has been changed to avoid conflict
|
||||
assert collection._distance_column_name != DISTANCE_COLUMN_NAME
|
||||
assert collection._distance_column_name.startswith(f"{DISTANCE_COLUMN_NAME}_")
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
# region Settings tests
|
||||
|
||||
|
||||
def test_settings_connection_string(monkeypatch) -> None:
|
||||
monkeypatch.delenv("PGHOST", raising=False)
|
||||
monkeypatch.delenv("PGPORT", raising=False)
|
||||
monkeypatch.delenv("PGDATABASE", raising=False)
|
||||
monkeypatch.delenv("PGUSER", raising=False)
|
||||
monkeypatch.delenv("PGPASSWORD", raising=False)
|
||||
|
||||
settings = PostgresSettings(connection_string="host=localhost port=5432 dbname=dbname user=user password=password")
|
||||
conn_info = settings.get_connection_args()
|
||||
|
||||
assert conn_info["host"] == "localhost"
|
||||
assert conn_info["port"] == 5432
|
||||
assert conn_info["dbname"] == "dbname"
|
||||
assert conn_info["user"] == "user"
|
||||
assert conn_info["password"] == "password"
|
||||
|
||||
|
||||
def test_settings_env_connection_string(monkeypatch) -> None:
|
||||
monkeypatch.delenv("PGHOST", raising=False)
|
||||
monkeypatch.delenv("PGPORT", raising=False)
|
||||
monkeypatch.delenv("PGDATABASE", raising=False)
|
||||
monkeypatch.delenv("PGUSER", raising=False)
|
||||
monkeypatch.delenv("PGPASSWORD", raising=False)
|
||||
|
||||
monkeypatch.setenv(
|
||||
"POSTGRES_CONNECTION_STRING", "host=localhost port=5432 dbname=dbname user=user password=password"
|
||||
)
|
||||
|
||||
settings = PostgresSettings()
|
||||
conn_info = settings.get_connection_args()
|
||||
assert conn_info["host"] == "localhost"
|
||||
assert conn_info["port"] == 5432
|
||||
assert conn_info["dbname"] == "dbname"
|
||||
assert conn_info["user"] == "user"
|
||||
assert conn_info["password"] == "password"
|
||||
|
||||
|
||||
def test_settings_env_vars(monkeypatch) -> None:
|
||||
monkeypatch.setenv("PGHOST", "localhost")
|
||||
monkeypatch.setenv("PGPORT", "5432")
|
||||
monkeypatch.setenv("PGDATABASE", "dbname")
|
||||
monkeypatch.setenv("PGUSER", "user")
|
||||
monkeypatch.setenv("PGPASSWORD", "password")
|
||||
|
||||
settings = PostgresSettings()
|
||||
conn_info = settings.get_connection_args()
|
||||
assert conn_info["host"] == "localhost"
|
||||
assert conn_info["port"] == 5432
|
||||
assert conn_info["dbname"] == "dbname"
|
||||
assert conn_info["user"] == "user"
|
||||
assert conn_info["password"] == "password"
|
||||
|
||||
|
||||
# endregion
|
||||
@@ -0,0 +1,341 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from pytest import fixture, mark, raises
|
||||
from qdrant_client.async_qdrant_client import AsyncQdrantClient
|
||||
from qdrant_client.models import Datatype, Distance, FieldCondition, MatchValue, VectorParams
|
||||
|
||||
from semantic_kernel.connectors.qdrant import QdrantCollection, QdrantStore
|
||||
from semantic_kernel.data.vector import DistanceFunction, VectorStoreField
|
||||
from semantic_kernel.exceptions import (
|
||||
VectorSearchExecutionException,
|
||||
VectorStoreInitializationException,
|
||||
VectorStoreModelValidationError,
|
||||
VectorStoreOperationException,
|
||||
)
|
||||
|
||||
BASE_PATH = "qdrant_client.async_qdrant_client.AsyncQdrantClient"
|
||||
|
||||
|
||||
@fixture
|
||||
def vector_store(qdrant_unit_test_env):
|
||||
return QdrantStore(env_file_path="test.env")
|
||||
|
||||
|
||||
@fixture
|
||||
def collection(qdrant_unit_test_env, definition):
|
||||
return QdrantCollection(
|
||||
record_type=dict,
|
||||
collection_name="test",
|
||||
definition=definition,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
def collection_without_named_vectors(qdrant_unit_test_env, definition):
|
||||
return QdrantCollection(
|
||||
record_type=dict,
|
||||
collection_name="test",
|
||||
definition=definition,
|
||||
named_vectors=False,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_list_collection_names():
|
||||
with patch(f"{BASE_PATH}.get_collections") as mock_get_collections:
|
||||
from qdrant_client.conversions.common_types import CollectionsResponse
|
||||
from qdrant_client.http.models import CollectionDescription
|
||||
|
||||
response = MagicMock(spec=CollectionsResponse)
|
||||
response.collections = [CollectionDescription(name="test")]
|
||||
mock_get_collections.return_value = response
|
||||
yield mock_get_collections
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_collection_exists():
|
||||
with patch(f"{BASE_PATH}.collection_exists") as mock_collection_exists:
|
||||
mock_collection_exists.return_value = True
|
||||
yield mock_collection_exists
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_ensure_collection_exists():
|
||||
with patch(f"{BASE_PATH}.create_collection") as mock_ensure_collection_exists:
|
||||
yield mock_ensure_collection_exists
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_ensure_collection_deleted():
|
||||
with patch(f"{BASE_PATH}.delete_collection") as mock_ensure_collection_deleted:
|
||||
mock_ensure_collection_deleted.return_value = True
|
||||
yield mock_ensure_collection_deleted
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_upsert():
|
||||
with patch(f"{BASE_PATH}.upsert") as mock_upsert:
|
||||
from qdrant_client.conversions.common_types import UpdateResult
|
||||
|
||||
result = MagicMock(spec=UpdateResult)
|
||||
result.status = "completed"
|
||||
mock_upsert.return_value = result
|
||||
yield mock_upsert
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_get(collection):
|
||||
with patch(f"{BASE_PATH}.retrieve") as mock_retrieve:
|
||||
from qdrant_client.http.models import Record
|
||||
|
||||
if collection.named_vectors:
|
||||
mock_retrieve.return_value = [
|
||||
Record(id="id1", payload={"content": "content"}, vector={"vector": [1.0, 2.0, 3.0]})
|
||||
]
|
||||
else:
|
||||
mock_retrieve.return_value = [Record(id="id1", payload={"content": "content"}, vector=[1.0, 2.0, 3.0])]
|
||||
yield mock_retrieve
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_delete():
|
||||
with patch(f"{BASE_PATH}.delete") as mock_delete:
|
||||
yield mock_delete
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_search():
|
||||
with patch(f"{BASE_PATH}.search") as mock_search:
|
||||
from qdrant_client.models import ScoredPoint
|
||||
|
||||
response1 = ScoredPoint(id="id1", version=1, score=0.0, payload={"content": "content"})
|
||||
response2 = ScoredPoint(id="id2", version=1, score=0.0, payload={"content": "content"})
|
||||
mock_search.return_value = [response1, response2]
|
||||
yield mock_search
|
||||
|
||||
|
||||
async def test_vector_store_defaults(vector_store):
|
||||
async with vector_store:
|
||||
assert vector_store.qdrant_client is not None
|
||||
assert vector_store.qdrant_client._client.rest_uri == "http://localhost:6333"
|
||||
|
||||
|
||||
def test_vector_store_with_client():
|
||||
qdrant_store = QdrantStore(client=AsyncQdrantClient())
|
||||
assert qdrant_store.qdrant_client is not None
|
||||
assert qdrant_store.qdrant_client._client.rest_uri == "http://localhost:6333"
|
||||
|
||||
|
||||
@mark.parametrize("exclude_list", [["QDRANT_LOCATION"]], indirect=True)
|
||||
def test_vector_store_in_memory(qdrant_unit_test_env):
|
||||
from qdrant_client.local.async_qdrant_local import AsyncQdrantLocal
|
||||
|
||||
qdrant_store = QdrantStore(api_key="supersecretkey", env_file_path="test.env")
|
||||
assert qdrant_store.qdrant_client is not None
|
||||
assert isinstance(qdrant_store.qdrant_client._client, AsyncQdrantLocal)
|
||||
assert qdrant_store.qdrant_client._client.location == ":memory:"
|
||||
|
||||
|
||||
def test_vector_store_fail():
|
||||
with raises(VectorStoreInitializationException, match="Failed to create Qdrant settings."):
|
||||
QdrantStore(location="localhost", url="localhost", env_file_path="test.env")
|
||||
|
||||
with raises(VectorStoreInitializationException, match="Failed to create Qdrant client."):
|
||||
QdrantStore(location="localhost", url="http://localhost", env_file_path="test.env")
|
||||
|
||||
|
||||
async def test_store_list_collection_names(vector_store):
|
||||
collections = await vector_store.list_collection_names()
|
||||
assert collections == ["test"]
|
||||
|
||||
|
||||
def test_get_collection(vector_store: QdrantStore, definition, qdrant_unit_test_env):
|
||||
collection = vector_store.get_collection(collection_name="test", record_type=dict, definition=definition)
|
||||
assert collection.collection_name == "test"
|
||||
assert collection.qdrant_client == vector_store.qdrant_client
|
||||
assert collection.record_type is dict
|
||||
assert collection.definition == definition
|
||||
|
||||
|
||||
async def test_collection_init(definition, qdrant_unit_test_env):
|
||||
async with QdrantCollection(
|
||||
record_type=dict,
|
||||
collection_name="test",
|
||||
definition=definition,
|
||||
env_file_path="test.env",
|
||||
) as collection:
|
||||
assert collection.collection_name == "test"
|
||||
assert collection.qdrant_client is not None
|
||||
assert collection.record_type is dict
|
||||
assert collection.definition == definition
|
||||
assert collection.named_vectors
|
||||
|
||||
|
||||
def test_collection_init_fail(definition):
|
||||
with raises(VectorStoreInitializationException, match="Failed to create Qdrant settings."):
|
||||
QdrantCollection(
|
||||
record_type=dict,
|
||||
collection_name="test",
|
||||
definition=definition,
|
||||
url="localhost",
|
||||
env_file_path="test.env",
|
||||
)
|
||||
with raises(VectorStoreInitializationException, match="Failed to create Qdrant client."):
|
||||
QdrantCollection(
|
||||
record_type=dict,
|
||||
collection_name="test",
|
||||
definition=definition,
|
||||
location="localhost",
|
||||
url="http://localhost",
|
||||
env_file_path="test.env",
|
||||
)
|
||||
with raises(
|
||||
VectorStoreModelValidationError, match="Only one vector field is allowed when not using named vectors."
|
||||
):
|
||||
definition.fields.append(VectorStoreField("vector", name="vector2", dimensions=3))
|
||||
QdrantCollection(
|
||||
record_type=dict,
|
||||
collection_name="test",
|
||||
definition=definition,
|
||||
named_vectors=False,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@mark.parametrize("collection_to_use", ["collection", "collection_without_named_vectors"])
|
||||
async def test_upsert(collection_to_use, request):
|
||||
from qdrant_client.models import PointStruct
|
||||
|
||||
collection = request.getfixturevalue(collection_to_use)
|
||||
if collection.named_vectors:
|
||||
record = PointStruct(id="id1", payload={"content": "content"}, vector={"vector": [1.0, 2.0, 3.0]})
|
||||
else:
|
||||
record = PointStruct(id="id1", payload={"content": "content"}, vector=[1.0, 2.0, 3.0])
|
||||
ids = await collection._inner_upsert([record])
|
||||
assert ids[0] == "id1"
|
||||
|
||||
ids = await collection.upsert(records={"id": "id1", "content": "content", "vector": [1.0, 2.0, 3.0]})
|
||||
assert ids == "id1"
|
||||
|
||||
|
||||
async def test_get(collection):
|
||||
records = await collection._inner_get(["id1"])
|
||||
assert records is not None
|
||||
|
||||
records = await collection.get("id1")
|
||||
assert records is not None
|
||||
|
||||
|
||||
async def test_delete(collection):
|
||||
await collection._inner_delete(["id1"])
|
||||
|
||||
|
||||
async def test_collection_exists(collection):
|
||||
await collection.collection_exists()
|
||||
|
||||
|
||||
async def test_ensure_collection_deleted(collection):
|
||||
await collection.ensure_collection_deleted()
|
||||
|
||||
|
||||
@mark.parametrize(
|
||||
"collection_to_use, results",
|
||||
[
|
||||
(
|
||||
"collection",
|
||||
{
|
||||
"collection_name": "test",
|
||||
"vectors_config": {"vector": VectorParams(size=5, distance=Distance.COSINE, datatype=Datatype.FLOAT32)},
|
||||
},
|
||||
),
|
||||
(
|
||||
"collection_without_named_vectors",
|
||||
{
|
||||
"collection_name": "test",
|
||||
"vectors_config": VectorParams(size=5, distance=Distance.COSINE, datatype=Datatype.FLOAT32),
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_create_index_with_named_vectors(collection_to_use, results, mock_ensure_collection_exists, request):
|
||||
await request.getfixturevalue(collection_to_use).ensure_collection_exists()
|
||||
mock_ensure_collection_exists.assert_called_once_with(**results)
|
||||
|
||||
|
||||
@mark.parametrize("collection_to_use", ["collection", "collection_without_named_vectors"])
|
||||
async def test_create_index_fail(collection_to_use, request):
|
||||
collection = request.getfixturevalue(collection_to_use)
|
||||
for field in collection.definition.vector_fields:
|
||||
field.distance_function = DistanceFunction.HAMMING
|
||||
with raises(VectorStoreOperationException):
|
||||
await collection.ensure_collection_exists()
|
||||
|
||||
|
||||
async def test_search(collection, mock_search):
|
||||
collection.named_vectors = False
|
||||
results = await collection.search(vector=[1.0, 2.0, 3.0], include_vectors=False)
|
||||
async for result in results.results:
|
||||
assert result.record["id"] == "id1"
|
||||
break
|
||||
|
||||
assert mock_search.call_count == 1
|
||||
mock_search.assert_called_with(
|
||||
collection_name="test",
|
||||
query_vector=[1.0, 2.0, 3.0],
|
||||
query_filter=None,
|
||||
with_vectors=False,
|
||||
limit=3,
|
||||
offset=0,
|
||||
)
|
||||
|
||||
|
||||
async def test_search_named_vectors(collection, mock_search):
|
||||
collection.named_vectors = True
|
||||
results = await collection.search(
|
||||
vector=[1.0, 2.0, 3.0],
|
||||
vector_property_name="vector",
|
||||
include_vectors=False,
|
||||
)
|
||||
async for result in results.results:
|
||||
assert result.record["id"] == "id1"
|
||||
break
|
||||
|
||||
assert mock_search.call_count == 1
|
||||
mock_search.assert_called_with(
|
||||
collection_name="test",
|
||||
query_vector=("vector", [1.0, 2.0, 3.0]),
|
||||
query_filter=None,
|
||||
with_vectors=False,
|
||||
limit=3,
|
||||
offset=0,
|
||||
)
|
||||
|
||||
|
||||
async def test_search_filter(collection, mock_search):
|
||||
results = await collection.search(
|
||||
vector=[1.0, 2.0, 3.0],
|
||||
include_vectors=False,
|
||||
filter=lambda x: x.id == "id1",
|
||||
)
|
||||
async for result in results.results:
|
||||
assert result.record["id"] == "id1"
|
||||
break
|
||||
|
||||
assert mock_search.call_count == 1
|
||||
mock_search.assert_called_with(
|
||||
collection_name="test",
|
||||
query_vector=("vector", [1.0, 2.0, 3.0]),
|
||||
query_filter=FieldCondition(key="id", match=MatchValue(value="id1")),
|
||||
with_vectors=False,
|
||||
limit=3,
|
||||
offset=0,
|
||||
)
|
||||
|
||||
|
||||
async def test_search_fail(collection):
|
||||
with raises(VectorSearchExecutionException, match="Search requires a vector."):
|
||||
await collection.search(include_vectors=False)
|
||||
@@ -0,0 +1,308 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import numpy as np
|
||||
from pytest import fixture, mark, raises
|
||||
from redis.asyncio.client import Redis
|
||||
|
||||
from semantic_kernel.connectors.redis import (
|
||||
RedisCollectionTypes,
|
||||
RedisHashsetCollection,
|
||||
RedisJsonCollection,
|
||||
RedisStore,
|
||||
)
|
||||
from semantic_kernel.exceptions import VectorStoreInitializationException, VectorStoreOperationException
|
||||
|
||||
BASE_PATH = "redis.asyncio.client.Redis"
|
||||
BASE_PATH_FT = "redis.commands.search.AsyncSearch"
|
||||
BASE_PATH_JSON = "redis.commands.json.commands.JSONCommands"
|
||||
|
||||
|
||||
@fixture
|
||||
def vector_store(redis_unit_test_env):
|
||||
return RedisStore(env_file_path="test.env")
|
||||
|
||||
|
||||
@fixture
|
||||
def collection_hash(redis_unit_test_env, definition):
|
||||
return RedisHashsetCollection(
|
||||
record_type=dict,
|
||||
collection_name="test",
|
||||
definition=definition,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
def collection_json(redis_unit_test_env, definition):
|
||||
return RedisJsonCollection(
|
||||
record_type=dict,
|
||||
collection_name="test",
|
||||
definition=definition,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
def collection_with_prefix_hash(redis_unit_test_env, definition):
|
||||
return RedisHashsetCollection(
|
||||
record_type=dict,
|
||||
collection_name="test",
|
||||
definition=definition,
|
||||
prefix_collection_name_to_key_names=True,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
def collection_with_prefix_json(redis_unit_test_env, definition):
|
||||
return RedisJsonCollection(
|
||||
record_type=dict,
|
||||
collection_name="test",
|
||||
definition=definition,
|
||||
prefix_collection_name_to_key_names=True,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def moc_list_collection_names():
|
||||
with patch(f"{BASE_PATH}.execute_command") as mock_get_collections:
|
||||
mock_get_collections.return_value = [b"test"]
|
||||
yield mock_get_collections
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_collection_exists():
|
||||
with patch(f"{BASE_PATH_FT}.info", new=AsyncMock()) as mock_collection_exists:
|
||||
mock_collection_exists.return_value = True
|
||||
yield mock_collection_exists
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_ensure_collection_exists():
|
||||
with patch(f"{BASE_PATH_FT}.create_index", new=AsyncMock()) as mock_reensure_collection_exists:
|
||||
yield mock_reensure_collection_exists
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_ensure_collection_deleted():
|
||||
with patch(f"{BASE_PATH_FT}.dropindex", new=AsyncMock()) as mock_ensure_collection_deleted:
|
||||
yield mock_ensure_collection_deleted
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_upsert_hash():
|
||||
with patch(f"{BASE_PATH}.hset", new=AsyncMock()) as mock_upsert:
|
||||
yield mock_upsert
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_upsert_json():
|
||||
with patch(f"{BASE_PATH_JSON}.set", new=AsyncMock()) as mock_upsert:
|
||||
yield mock_upsert
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_get_hash():
|
||||
with patch(f"{BASE_PATH}.hgetall", new=AsyncMock()) as mock_get:
|
||||
mock_get.return_value = {
|
||||
b"content": b"content",
|
||||
b"vector": np.array([1.0, 2.0, 3.0]).tobytes(),
|
||||
}
|
||||
yield mock_get
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_get_json():
|
||||
with patch(f"{BASE_PATH_JSON}.mget", new=AsyncMock()) as mock_get:
|
||||
mock_get.return_value = [
|
||||
[
|
||||
{
|
||||
"content": "content",
|
||||
"vector": [1.0, 2.0, 3.0],
|
||||
}
|
||||
]
|
||||
]
|
||||
yield mock_get
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_delete_hash():
|
||||
with patch(f"{BASE_PATH}.delete", new=AsyncMock()) as mock_delete:
|
||||
yield mock_delete
|
||||
|
||||
|
||||
@fixture(autouse=True)
|
||||
def mock_delete_json():
|
||||
with patch(f"{BASE_PATH_JSON}.delete", new=AsyncMock()) as mock_delete:
|
||||
yield mock_delete
|
||||
|
||||
|
||||
def test_vector_store_defaults(vector_store):
|
||||
assert vector_store.redis_database is not None
|
||||
assert vector_store.redis_database.connection_pool.connection_kwargs["host"] == "localhost"
|
||||
|
||||
|
||||
def test_vector_store_with_client(redis_unit_test_env):
|
||||
vector_store = RedisStore(redis_database=Redis.from_url(redis_unit_test_env["REDIS_CONNECTION_STRING"]))
|
||||
assert vector_store.redis_database is not None
|
||||
assert vector_store.redis_database.connection_pool.connection_kwargs["host"] == "localhost"
|
||||
|
||||
|
||||
@mark.parametrize("exclude_list", [["REDIS_CONNECTION_STRING"]], indirect=True)
|
||||
def test_vector_store_fail(redis_unit_test_env):
|
||||
with raises(VectorStoreInitializationException, match="Failed to create Redis settings."):
|
||||
RedisStore(env_file_path="test.env")
|
||||
|
||||
|
||||
async def test_store_list_collection_names(vector_store, moc_list_collection_names):
|
||||
collections = await vector_store.list_collection_names()
|
||||
assert collections == ["test"]
|
||||
|
||||
|
||||
@mark.parametrize("type_", ["hashset", "json"])
|
||||
def test_get_collection(vector_store, definition, type_):
|
||||
if type_ == "hashset":
|
||||
collection = vector_store.get_collection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
collection_type=RedisCollectionTypes.HASHSET,
|
||||
)
|
||||
assert isinstance(collection, RedisHashsetCollection)
|
||||
else:
|
||||
collection = vector_store.get_collection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
collection_type=RedisCollectionTypes.JSON,
|
||||
)
|
||||
assert isinstance(collection, RedisJsonCollection)
|
||||
assert collection.collection_name == "test"
|
||||
assert collection.redis_database == vector_store.redis_database
|
||||
assert collection.record_type is dict
|
||||
assert collection.definition == definition
|
||||
|
||||
|
||||
@mark.parametrize("type_", ["hashset", "json"])
|
||||
def test_collection_init(redis_unit_test_env, definition, type_):
|
||||
if type_ == "hashset":
|
||||
collection = RedisHashsetCollection(
|
||||
record_type=dict,
|
||||
collection_name="test",
|
||||
definition=definition,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
else:
|
||||
collection = RedisJsonCollection(
|
||||
record_type=dict,
|
||||
collection_name="test",
|
||||
definition=definition,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
assert collection.collection_name == "test"
|
||||
assert collection.redis_database is not None
|
||||
assert collection.record_type is dict
|
||||
assert collection.definition == definition
|
||||
assert collection.prefix_collection_name_to_key_names is False
|
||||
|
||||
|
||||
@mark.parametrize("type_", ["hashset", "json"])
|
||||
def test_init_with_type(redis_unit_test_env, record_type, type_):
|
||||
if type_ == "hashset":
|
||||
collection = RedisHashsetCollection(record_type=record_type, collection_name="test")
|
||||
else:
|
||||
collection = RedisJsonCollection(record_type=record_type, collection_name="test")
|
||||
assert collection is not None
|
||||
assert collection.record_type is record_type
|
||||
assert collection.collection_name == "test"
|
||||
|
||||
|
||||
@mark.parametrize("exclude_list", [["REDIS_CONNECTION_STRING"]], indirect=True)
|
||||
def test_collection_fail(redis_unit_test_env, definition):
|
||||
with raises(VectorStoreInitializationException, match="Failed to create Redis settings."):
|
||||
RedisHashsetCollection(
|
||||
record_type=dict,
|
||||
collection_name="test",
|
||||
definition=definition,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
with raises(VectorStoreInitializationException, match="Failed to create Redis settings."):
|
||||
RedisJsonCollection(
|
||||
record_type=dict,
|
||||
collection_name="test",
|
||||
definition=definition,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
|
||||
@mark.parametrize("type_", ["hashset", "json"])
|
||||
async def test_upsert(collection_hash, collection_json, type_):
|
||||
collection = collection_hash if type_ == "hashset" else collection_json
|
||||
ids = await collection.upsert(records={"id": "id1", "content": "content", "vector": [1.0, 2.0, 3.0]})
|
||||
assert ids == "id1"
|
||||
|
||||
|
||||
async def test_upsert_with_prefix(collection_with_prefix_hash, collection_with_prefix_json):
|
||||
ids = await collection_with_prefix_hash.upsert(
|
||||
records={"id": "id1", "content": "content", "vector": [1.0, 2.0, 3.0]}
|
||||
)
|
||||
assert ids == "id1"
|
||||
ids = await collection_with_prefix_json.upsert(
|
||||
records={"id": "id1", "content": "content", "vector": [1.0, 2.0, 3.0]}
|
||||
)
|
||||
assert ids == "id1"
|
||||
|
||||
|
||||
@mark.parametrize("prefix", [True, False])
|
||||
@mark.parametrize("type_", ["hashset", "json"])
|
||||
async def test_get(
|
||||
collection_hash, collection_json, collection_with_prefix_hash, collection_with_prefix_json, type_, prefix
|
||||
):
|
||||
if prefix:
|
||||
collection = collection_with_prefix_hash if type_ == "hashset" else collection_with_prefix_json
|
||||
else:
|
||||
collection = collection_hash if type_ == "hashset" else collection_json
|
||||
|
||||
records = await collection.get("id1")
|
||||
assert records is not None
|
||||
|
||||
|
||||
@mark.parametrize("type_", ["hashset", "json"])
|
||||
async def test_delete(collection_hash, collection_json, type_):
|
||||
collection = collection_hash if type_ == "hashset" else collection_json
|
||||
await collection._inner_delete(["id1"])
|
||||
|
||||
|
||||
async def test_collection_exists(collection_hash, mock_collection_exists):
|
||||
await collection_hash.collection_exists()
|
||||
|
||||
|
||||
async def test_collection_exists_false(collection_hash, mock_collection_exists):
|
||||
mock_collection_exists.side_effect = Exception
|
||||
exists = await collection_hash.collection_exists()
|
||||
assert not exists
|
||||
|
||||
|
||||
async def test_ensure_collection_deleted(collection_hash, mock_ensure_collection_deleted):
|
||||
await collection_hash.ensure_collection_deleted()
|
||||
await collection_hash.ensure_collection_deleted()
|
||||
|
||||
|
||||
async def test_create_index(collection_hash, mock_ensure_collection_exists):
|
||||
await collection_hash.ensure_collection_exists()
|
||||
|
||||
|
||||
async def test_create_index_manual(collection_hash, mock_ensure_collection_exists):
|
||||
from redis.commands.search.index_definition import IndexDefinition, IndexType
|
||||
|
||||
fields = ["fields"]
|
||||
index_definition = IndexDefinition(prefix="test:", index_type=IndexType.HASH)
|
||||
await collection_hash.ensure_collection_exists(index_definition=index_definition, fields=fields)
|
||||
|
||||
|
||||
async def test_create_index_fail(collection_hash, mock_ensure_collection_exists):
|
||||
with raises(VectorStoreOperationException, match="Invalid index type supplied."):
|
||||
await collection_hash.ensure_collection_exists(index_definition="index_definition", fields="fields")
|
||||
@@ -0,0 +1,601 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import json
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from typing import NamedTuple
|
||||
from unittest.mock import AsyncMock, MagicMock, NonCallableMagicMock, patch
|
||||
|
||||
from pytest import fixture, mark, param, raises
|
||||
|
||||
from semantic_kernel.connectors.sql_server import (
|
||||
QueryBuilder,
|
||||
SqlCommand,
|
||||
SqlServerCollection,
|
||||
SqlServerStore,
|
||||
_build_create_table_query,
|
||||
_build_delete_query,
|
||||
_build_delete_table_query,
|
||||
_build_merge_query,
|
||||
_build_search_query,
|
||||
_build_select_query,
|
||||
_build_select_table_names_query,
|
||||
)
|
||||
from semantic_kernel.data.vector import DistanceFunction, IndexKind, VectorSearchOptions, VectorStoreField
|
||||
from semantic_kernel.exceptions.vector_store_exceptions import (
|
||||
VectorStoreInitializationException,
|
||||
VectorStoreOperationException,
|
||||
)
|
||||
|
||||
|
||||
class TestQueryBuilder:
|
||||
def test_query_builder_append(self):
|
||||
qb = QueryBuilder()
|
||||
qb.append("SELECT * FROM")
|
||||
qb.append(" table", suffix=";")
|
||||
result = str(qb).strip()
|
||||
assert result == "SELECT * FROM table;"
|
||||
|
||||
def test_query_builder_append_list(self):
|
||||
qb = QueryBuilder()
|
||||
qb.append_list(["id", "name", "age"], sep=", ", suffix=";")
|
||||
result = str(qb).strip()
|
||||
assert result == "id, name, age;"
|
||||
|
||||
def test_query_builder_escape_identifier(self):
|
||||
assert QueryBuilder.escape_identifier("simple") == "[simple]"
|
||||
assert QueryBuilder.escape_identifier("has]bracket") == "[has]]bracket]"
|
||||
assert QueryBuilder.escape_identifier("two]]brackets") == "[two]]]]brackets]"
|
||||
assert QueryBuilder.escape_identifier("") == "[]"
|
||||
|
||||
def test_query_builder_append_table_name(self):
|
||||
qb = QueryBuilder()
|
||||
qb.append_table_name("dbo", "Users", prefix="SELECT * FROM", suffix=";", newline=False)
|
||||
result = str(qb).strip()
|
||||
assert result == "SELECT * FROM [dbo].[Users] ;"
|
||||
|
||||
def test_query_builder_append_table_name_escapes_closing_bracket(self):
|
||||
qb = QueryBuilder()
|
||||
qb.append_table_name("my]schema", "my]table", prefix="SELECT * FROM", suffix=";")
|
||||
result = str(qb).strip()
|
||||
assert result == "SELECT * FROM [my]]schema].[my]]table] ;"
|
||||
|
||||
def test_query_builder_append_table_name_prevents_sql_injection(self):
|
||||
qb = QueryBuilder()
|
||||
qb.append("DROP TABLE IF EXISTS")
|
||||
qb.append_table_name("dbo", "]; EXEC xp_cmdshell('whoami'); --", suffix=";")
|
||||
result = str(qb)
|
||||
assert "EXEC xp_cmdshell" not in result.split("].[")[0], "SQL injection should not escape bracket quoting"
|
||||
assert "[dbo].[]]; EXEC xp_cmdshell('whoami'); --]" in result
|
||||
|
||||
def test_query_builder_remove_last(self):
|
||||
qb = QueryBuilder("SELECT * FROM table;")
|
||||
qb.remove_last(1) # remove trailing semicolon
|
||||
result = str(qb).strip()
|
||||
assert result == "SELECT * FROM table"
|
||||
|
||||
def test_query_builder_in_parenthesis(self):
|
||||
qb = QueryBuilder("INSERT INTO table")
|
||||
with qb.in_parenthesis():
|
||||
qb.append("id, name, age")
|
||||
result = str(qb).strip()
|
||||
assert result == "INSERT INTO table (id, name, age)"
|
||||
|
||||
def test_query_builder_in_parenthesis_with_prefix_suffix(self):
|
||||
qb = QueryBuilder()
|
||||
with qb.in_parenthesis(prefix="VALUES", suffix=";"):
|
||||
qb.append_list(["1", "'John'", "30"])
|
||||
result = str(qb).strip()
|
||||
assert result == "VALUES (1, 'John', 30) ;"
|
||||
|
||||
def test_query_builder_in_logical_group(self):
|
||||
qb = QueryBuilder()
|
||||
with qb.in_logical_group():
|
||||
qb.append("UPDATE Users SET name = 'John'")
|
||||
result = str(qb).strip()
|
||||
lines = result.splitlines()
|
||||
assert lines[0] == "BEGIN"
|
||||
assert lines[1] == "UPDATE Users SET name = 'John'"
|
||||
assert lines[2] == "END"
|
||||
|
||||
|
||||
class TestSqlCommand:
|
||||
def test_sql_command_initial_query(self):
|
||||
cmd = SqlCommand("SELECT 1")
|
||||
assert str(cmd.query) == "SELECT 1"
|
||||
|
||||
def test_sql_command_add_parameter(self):
|
||||
cmd = SqlCommand("SELECT * FROM Test WHERE id = ?")
|
||||
cmd.add_parameter("42")
|
||||
assert cmd.parameters[0] == "42"
|
||||
|
||||
def test_sql_command_add_parameters(self):
|
||||
cmd = SqlCommand("SELECT * FROM Test WHERE id = ?")
|
||||
cmd.add_parameters(["42", "43"])
|
||||
assert cmd.parameters[0] == "42"
|
||||
assert cmd.parameters[1] == "43"
|
||||
|
||||
def test_parameter_limit(self):
|
||||
cmd = SqlCommand()
|
||||
cmd.add_parameters(["42"] * 2100)
|
||||
with raises(VectorStoreOperationException):
|
||||
cmd.add_parameter("43")
|
||||
with raises(VectorStoreOperationException):
|
||||
cmd.add_parameters(["43", "44"])
|
||||
|
||||
|
||||
class TestQueryBuildFunctions:
|
||||
def test_build_create_table_query(self):
|
||||
schema = "dbo"
|
||||
table = "Test"
|
||||
key_field = VectorStoreField("key", name="id", type="str")
|
||||
data_fields = [
|
||||
VectorStoreField("data", name="name", type="str"),
|
||||
VectorStoreField("data", name="age", type="int"),
|
||||
]
|
||||
vector_fields = [
|
||||
VectorStoreField("vector", name="embedding", type="float", dimensions=1536),
|
||||
]
|
||||
cmd = _build_create_table_query(schema, table, key_field, data_fields, vector_fields)
|
||||
assert not cmd.parameters
|
||||
cmd_str = str(cmd.query)
|
||||
assert (
|
||||
cmd_str
|
||||
== "BEGIN\nCREATE TABLE [dbo].[Test] \n ([id] nvarchar(255) NOT NULL,\n[name] nvarchar(max) NULL,\n[age] "
|
||||
"int NULL,\n[embedding] VECTOR(1536) NULL,\nPRIMARY KEY ([id]) \n) ;\nEND\n"
|
||||
)
|
||||
|
||||
def test_build_create_table_query_escapes_single_quote_in_object_id(self):
|
||||
key_field = VectorStoreField("key", name="id", type="str")
|
||||
cmd = _build_create_table_query("dbo", "test'table", key_field, [], [], if_not_exists=True)
|
||||
cmd_str = str(cmd.query)
|
||||
# Single quote must be escaped inside the OBJECT_ID N'...' string literal
|
||||
assert "OBJECT_ID(N'[dbo].[test''table]'" in cmd_str
|
||||
|
||||
def test_build_create_table_query_uses_storage_name_for_primary_key(self):
|
||||
key_field = VectorStoreField("key", name="id", type="str", storage_name="pk_id")
|
||||
cmd = _build_create_table_query("dbo", "Test", key_field, [], [])
|
||||
cmd_str = str(cmd.query)
|
||||
assert "[pk_id] nvarchar" in cmd_str
|
||||
assert "PRIMARY KEY ([pk_id])" in cmd_str
|
||||
|
||||
def test_build_merge_query_output_uses_storage_name(self):
|
||||
key_field = VectorStoreField("key", name="id", type="str", storage_name="pk_id")
|
||||
records = [{"pk_id": "1"}]
|
||||
cmd = _build_merge_query("dbo", "Test", key_field, [], [], records)
|
||||
cmd_str = str(cmd.query)
|
||||
assert "OUTPUT inserted.[pk_id] INTO @UpsertedKeys" in cmd_str
|
||||
|
||||
def test_delete_table_query(self):
|
||||
schema = "dbo"
|
||||
table = "Test"
|
||||
cmd = _build_delete_table_query(schema, table)
|
||||
assert str(cmd.query) == f"DROP TABLE IF EXISTS [{schema}].[{table}] ;"
|
||||
|
||||
@mark.parametrize("schema", ["dbo", None])
|
||||
def test_build_select_table_names_query(self, schema):
|
||||
cmd = _build_select_table_names_query(schema)
|
||||
if schema:
|
||||
assert cmd.parameters == [schema]
|
||||
assert str(cmd) == (
|
||||
"SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES "
|
||||
"WHERE TABLE_TYPE = 'BASE TABLE' "
|
||||
"AND (@schema is NULL or TABLE_SCHEMA = ?);"
|
||||
)
|
||||
else:
|
||||
assert str(cmd) == "SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_TYPE = 'BASE TABLE';"
|
||||
|
||||
def test_build_merge_query(self):
|
||||
schema = "dbo"
|
||||
table = "Test"
|
||||
key_field = VectorStoreField("key", name="id", type="str")
|
||||
data_fields = [
|
||||
VectorStoreField("data", name="name", type="str"),
|
||||
VectorStoreField("data", name="age", type="int"),
|
||||
]
|
||||
vector_fields = [
|
||||
VectorStoreField("vector", name="embedding", type="float", dimensions=5),
|
||||
]
|
||||
records = [
|
||||
{
|
||||
"id": "test",
|
||||
"name": "name",
|
||||
"age": 50,
|
||||
"embedding": [0.1, 0.2, 0.3, 0.4, 0.5],
|
||||
}
|
||||
]
|
||||
cmd = _build_merge_query(schema, table, key_field, data_fields, vector_fields, records)
|
||||
assert cmd.parameters[0] == records[0]["id"]
|
||||
assert cmd.parameters[1] == records[0]["name"]
|
||||
assert cmd.parameters[2] == str(records[0]["age"])
|
||||
assert cmd.parameters[3] == json.dumps(records[0]["embedding"])
|
||||
str_cmd = str(cmd)
|
||||
assert str_cmd == (
|
||||
"DECLARE @UpsertedKeys TABLE (KeyColumn nvarchar(255));\n"
|
||||
"MERGE INTO [dbo].[Test] AS t\n"
|
||||
"USING ( VALUES (?, ?, ?, ?) ) AS s ([id], [name], [age], [embedding]) "
|
||||
"ON (t.[id] = s.[id]) \n"
|
||||
"WHEN MATCHED THEN\n"
|
||||
"UPDATE SET t.[name] = s.[name], t.[age] = s.[age], t.[embedding] = s.[embedding]\n"
|
||||
"WHEN NOT MATCHED THEN\n"
|
||||
"INSERT ([id], [name], [age], [embedding]) "
|
||||
"VALUES (s.[id], s.[name], s.[age], s.[embedding]) \n"
|
||||
"OUTPUT inserted.[id] INTO @UpsertedKeys (KeyColumn);\n"
|
||||
"SELECT KeyColumn FROM @UpsertedKeys;\n"
|
||||
)
|
||||
|
||||
def test_build_select_query(self):
|
||||
schema = "dbo"
|
||||
table = "Test"
|
||||
key_field = VectorStoreField("key", name="id", type="str")
|
||||
data_fields = [
|
||||
VectorStoreField("data", name="name", type="str"),
|
||||
VectorStoreField("data", name="age", type="int"),
|
||||
]
|
||||
vector_fields = [
|
||||
VectorStoreField("vector", name="embedding", type="float", dimensions=5),
|
||||
]
|
||||
keys = ["test"]
|
||||
cmd = _build_select_query(schema, table, key_field, data_fields, vector_fields, keys)
|
||||
assert cmd.parameters == ["test"]
|
||||
str_cmd = str(cmd)
|
||||
assert str_cmd == "SELECT\n[id], [name], [age], [embedding] FROM [dbo].[Test] \nWHERE [id] IN\n (?) ;"
|
||||
|
||||
def test_build_delete_query(self):
|
||||
schema = "dbo"
|
||||
table = "Test"
|
||||
key_field = VectorStoreField("key", name="id", type="str")
|
||||
keys = ["test"]
|
||||
cmd = _build_delete_query(schema, table, key_field, keys)
|
||||
str_cmd = str(cmd)
|
||||
assert cmd.parameters[0] == "test"
|
||||
assert str_cmd == "DELETE FROM [dbo].[Test] WHERE [id] IN (?) ;"
|
||||
|
||||
def test_build_search_query(self):
|
||||
schema = "dbo"
|
||||
table = "Test"
|
||||
key_field = VectorStoreField("key", name="id", type="str")
|
||||
data_fields = [
|
||||
VectorStoreField("data", name="name", type="str"),
|
||||
VectorStoreField("data", name="age", type="int"),
|
||||
]
|
||||
vector_fields = [
|
||||
VectorStoreField(
|
||||
"vector",
|
||||
name="embedding",
|
||||
type="float",
|
||||
dimensions=5,
|
||||
distance_function=DistanceFunction.COSINE_DISTANCE,
|
||||
),
|
||||
]
|
||||
vector = [0.1, 0.2, 0.3, 0.4, 0.5]
|
||||
options = VectorSearchOptions(
|
||||
vector_property_name="embedding",
|
||||
)
|
||||
|
||||
cmd = _build_search_query(schema, table, key_field, data_fields, vector_fields, vector, options)
|
||||
assert cmd.parameters[0] == json.dumps(vector)
|
||||
str_cmd = str(cmd)
|
||||
assert (
|
||||
str_cmd == "SELECT [id], [name], [age], VECTOR_DISTANCE('cosine', [embedding], CAST(? AS VECTOR(5))) as "
|
||||
"_vector_distance_value\n FROM [dbo].[Test] \nORDER BY "
|
||||
"_vector_distance_value ASC\nOFFSET 0 ROWS FETCH NEXT 3 ROWS ONLY;"
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
async def mock_connection(*args, **kwargs):
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@mark.parametrize(
|
||||
"connection_string",
|
||||
[
|
||||
param(
|
||||
"Driver={ODBC Driver 18 for SQL Server};Server=localhost;Database=testdb;uid=testuserLongAsMax=yes;",
|
||||
id="with uid",
|
||||
),
|
||||
param(
|
||||
"Driver={ODBC Driver 18 for SQL Server};Server=localhost;Database=testdb;LongAsMax=yes;", id="credential"
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_get_mssql_connection(connection_string):
|
||||
mock_pyodbc = NonCallableMagicMock()
|
||||
sys.modules["pyodbc"] = mock_pyodbc
|
||||
|
||||
with patch("pyodbc.connect") as patched_connection:
|
||||
from azure.core.credentials_async import AsyncTokenCredential
|
||||
|
||||
from semantic_kernel.connectors.sql_server import SqlSettings, _get_mssql_connection
|
||||
|
||||
token = MagicMock()
|
||||
token.token.return_value = "test_token"
|
||||
token.token.encode.return_value = b"test_token"
|
||||
credential = AsyncMock(spec=AsyncTokenCredential)
|
||||
credential.__aenter__.return_value = credential
|
||||
credential.get_token.return_value = token
|
||||
|
||||
settings = SqlSettings(connection_string=connection_string)
|
||||
with patch("semantic_kernel.connectors.sql_server.AsyncTokenCredential", return_value=credential):
|
||||
connection = await _get_mssql_connection(settings, credential=credential)
|
||||
assert connection is not None
|
||||
assert isinstance(connection, MagicMock)
|
||||
if "uid" in connection_string:
|
||||
assert patched_connection.call_args.kwargs["attrs_before"] is None
|
||||
else:
|
||||
assert patched_connection.call_args.kwargs["attrs_before"] == {
|
||||
1256: b"\n\x00\x00\x00test_token",
|
||||
}
|
||||
|
||||
|
||||
class TestSqlServerStore:
|
||||
async def test_create_store(self, sql_server_unit_test_env):
|
||||
store = SqlServerStore()
|
||||
assert store is not None
|
||||
assert store.settings is not None
|
||||
assert store.settings.connection_string is not None
|
||||
assert "LongAsMax=yes;" in store.settings.connection_string.get_secret_value()
|
||||
|
||||
with patch("semantic_kernel.connectors.sql_server._get_mssql_connection") as mock_get_connection:
|
||||
mock_get_connection.return_value = AsyncMock()
|
||||
await store.__aenter__()
|
||||
assert store.connection is not None
|
||||
|
||||
@mark.parametrize(
|
||||
"override_env_param_dict",
|
||||
[
|
||||
{
|
||||
"SQL_SERVER_CONNECTION_STRING": "Driver={ODBC Driver 18 for SQL Server};Server=localhost;Database=testdb;User Id=testuser;Password=example;LongAsMax=yes;" # noqa: E501
|
||||
}
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
def test_create_store_with_long_as_max(self, sql_server_unit_test_env):
|
||||
store = SqlServerStore()
|
||||
assert store is not None
|
||||
assert store.settings is not None
|
||||
assert store.settings.connection_string is not None
|
||||
|
||||
@mark.parametrize("exclude_list", ["SQL_SERVER_CONNECTION_STRING"], indirect=True)
|
||||
def test_create_without_connection_string(self, sql_server_unit_test_env):
|
||||
with raises(VectorStoreInitializationException):
|
||||
SqlServerStore(env_file_path="test.env")
|
||||
|
||||
def test_get_collection(self, sql_server_unit_test_env, definition):
|
||||
store = SqlServerStore()
|
||||
collection = store.get_collection(collection_name="test", record_type=dict, definition=definition)
|
||||
assert collection is not None
|
||||
|
||||
async def test_list_collection_names(self, sql_server_unit_test_env, mock_connection):
|
||||
async with SqlServerStore(connection=mock_connection) as store:
|
||||
mock_connection.cursor.return_value.__enter__.return_value.fetchall.return_value = [
|
||||
["Test1"],
|
||||
["Test2"],
|
||||
]
|
||||
collection_names = await store.list_collection_names()
|
||||
assert collection_names == ["Test1", "Test2"]
|
||||
|
||||
async def test_no_connection(self, sql_server_unit_test_env):
|
||||
store = SqlServerStore()
|
||||
with raises(VectorStoreOperationException):
|
||||
await store.list_collection_names()
|
||||
|
||||
|
||||
class TestSqlServerCollection:
|
||||
@mark.parametrize("exclude_list", ["SQL_SERVER_CONNECTION_STRING"], indirect=True)
|
||||
def test_create_without_connection_string(self, sql_server_unit_test_env, definition):
|
||||
with raises(VectorStoreInitializationException):
|
||||
SqlServerCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
env_file_path="test.env",
|
||||
)
|
||||
|
||||
async def test_create(self, sql_server_unit_test_env, definition):
|
||||
collection = SqlServerCollection(collection_name="test", record_type=dict, definition=definition)
|
||||
assert collection is not None
|
||||
assert collection.collection_name == "test"
|
||||
assert collection.settings is not None
|
||||
assert collection.settings.connection_string is not None
|
||||
|
||||
with patch("semantic_kernel.connectors.sql_server._get_mssql_connection") as mock_get_connection:
|
||||
mock_get_connection.return_value = AsyncMock()
|
||||
await collection.__aenter__()
|
||||
assert collection.connection is not None
|
||||
|
||||
async def test_upsert(
|
||||
self,
|
||||
sql_server_unit_test_env,
|
||||
mock_connection,
|
||||
definition,
|
||||
):
|
||||
collection = SqlServerCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
connection=mock_connection,
|
||||
)
|
||||
record = {"id": "1", "content": "test", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]}
|
||||
mock_connection.cursor.return_value.__enter__.return_value.nextset.side_effect = [True, False]
|
||||
mock_connection.cursor.return_value.__enter__.return_value.fetchall.return_value = [
|
||||
["1"],
|
||||
]
|
||||
await collection.upsert(record)
|
||||
mock_connection.cursor.return_value.__enter__.return_value.execute.assert_called_with(
|
||||
(
|
||||
"DECLARE @UpsertedKeys TABLE (KeyColumn nvarchar(255));\n"
|
||||
"MERGE INTO [dbo].[test] AS t\n"
|
||||
"USING ( VALUES (?, ?, ?) ) AS s ([id], [content], [vector]) "
|
||||
"ON (t.[id] = s.[id]) \n"
|
||||
"WHEN MATCHED THEN\n"
|
||||
"UPDATE SET t.[content] = s.[content], t.[vector] = s.[vector]\n"
|
||||
"WHEN NOT MATCHED THEN\n"
|
||||
"INSERT ([id], [content], [vector]) "
|
||||
"VALUES (s.[id], s.[content], s.[vector]) \n"
|
||||
"OUTPUT inserted.[id] INTO @UpsertedKeys (KeyColumn);\n"
|
||||
"SELECT KeyColumn FROM @UpsertedKeys;\n"
|
||||
),
|
||||
("1", "test", json.dumps([0.1, 0.2, 0.3, 0.4, 0.5])),
|
||||
)
|
||||
|
||||
async def test_get(
|
||||
self,
|
||||
sql_server_unit_test_env,
|
||||
mock_connection,
|
||||
definition,
|
||||
):
|
||||
class MockRow(NamedTuple):
|
||||
id: str
|
||||
content: str
|
||||
vector: str
|
||||
|
||||
mock_cursor = MagicMock()
|
||||
mock_connection.cursor.return_value.__enter__.return_value = mock_cursor
|
||||
|
||||
collection = SqlServerCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
connection=mock_connection,
|
||||
)
|
||||
key = "1"
|
||||
|
||||
row = MockRow("1", "test", "[0.1, 0.2, 0.3, 0.4, 0.5]")
|
||||
mock_cursor.description = [["id"], ["content"], ["vector"]]
|
||||
|
||||
mock_cursor.__iter__.return_value = [row]
|
||||
record = await collection.get(key, include_vectors=True)
|
||||
mock_cursor.execute.assert_called_with(
|
||||
"SELECT\n[id], [content], [vector] FROM [dbo].[test] \nWHERE [id] IN\n (?) ;", ("1",)
|
||||
)
|
||||
assert record["id"] == "1"
|
||||
assert record["content"] == "test"
|
||||
assert record["vector"] == [0.1, 0.2, 0.3, 0.4, 0.5]
|
||||
|
||||
async def test_delete(
|
||||
self,
|
||||
sql_server_unit_test_env,
|
||||
mock_connection,
|
||||
definition,
|
||||
):
|
||||
collection = SqlServerCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
connection=mock_connection,
|
||||
)
|
||||
key = "1"
|
||||
await collection.delete(key)
|
||||
mock_connection.cursor.return_value.__enter__.return_value.execute.assert_called_with(
|
||||
"DELETE FROM [dbo].[test] WHERE [id] IN (?) ;", ("1",)
|
||||
)
|
||||
|
||||
async def test_search(
|
||||
self,
|
||||
sql_server_unit_test_env,
|
||||
mock_connection,
|
||||
definition,
|
||||
):
|
||||
mock_cursor = MagicMock()
|
||||
mock_connection.cursor.return_value.__enter__.return_value = mock_cursor
|
||||
for field in definition.vector_fields:
|
||||
field.distance_function = DistanceFunction.COSINE_DISTANCE
|
||||
collection = SqlServerCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
connection=mock_connection,
|
||||
)
|
||||
vector = [0.1, 0.2, 0.3, 0.4, 0.5]
|
||||
|
||||
@dataclass
|
||||
class MockRow:
|
||||
id: str
|
||||
content: str
|
||||
_vector_distance_value: float
|
||||
|
||||
row = MockRow("1", "test", 0.1)
|
||||
mock_cursor.description = [["id"], ["content"], ["_vector_distance_value"]]
|
||||
|
||||
mock_cursor.__iter__.return_value = [row]
|
||||
search_result = await collection.search(
|
||||
vector=vector,
|
||||
vector_property_name="vector",
|
||||
filter=lambda x: x.content == "test",
|
||||
)
|
||||
async for record in search_result.results:
|
||||
assert record.record["id"] == "1"
|
||||
assert record.record["content"] == "test"
|
||||
assert record.score == 0.1
|
||||
mock_cursor.execute.assert_called_with(
|
||||
(
|
||||
"SELECT [id], [content], VECTOR_DISTANCE('cosine', [vector], CAST(? AS VECTOR(5))) as "
|
||||
"_vector_distance_value\n FROM [dbo].[test] \n WHERE [content] = ? \nORDER BY _vector_distance_value "
|
||||
"ASC\nOFFSET 0 ROWS FETCH NEXT 3 ROWS ONLY;"
|
||||
),
|
||||
(json.dumps(vector), "test"),
|
||||
)
|
||||
|
||||
async def test_ensure_collection_exists(
|
||||
self,
|
||||
sql_server_unit_test_env,
|
||||
mock_connection,
|
||||
definition,
|
||||
):
|
||||
for field in definition.vector_fields:
|
||||
field.index_kind = IndexKind.FLAT
|
||||
collection = SqlServerCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
connection=mock_connection,
|
||||
)
|
||||
await collection.ensure_collection_exists()
|
||||
mock_connection.cursor.return_value.__enter__.return_value.execute.assert_called_with(
|
||||
(
|
||||
"IF OBJECT_ID(N'[dbo].[test]', N'U') IS NULL\n"
|
||||
"BEGIN\nCREATE TABLE [dbo].[test] \n ([id] nvarchar"
|
||||
"(255) NOT NULL,\n[content] nvarchar(max) NULL,\n[vector] VECTOR(5) NULL,\n"
|
||||
"PRIMARY KEY ([id]) \n) ;"
|
||||
"\nEND\n"
|
||||
),
|
||||
(),
|
||||
)
|
||||
|
||||
async def test_ensure_collection_deleted(
|
||||
self,
|
||||
sql_server_unit_test_env,
|
||||
mock_connection,
|
||||
definition,
|
||||
):
|
||||
collection = SqlServerCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
connection=mock_connection,
|
||||
)
|
||||
await collection.ensure_collection_deleted()
|
||||
mock_connection.cursor.return_value.__enter__.return_value.execute.assert_called_with(
|
||||
"DROP TABLE IF EXISTS [dbo].[test] ;", ()
|
||||
)
|
||||
|
||||
async def test_no_connection(self, sql_server_unit_test_env, definition):
|
||||
collection = SqlServerCollection(
|
||||
collection_name="test",
|
||||
record_type=dict,
|
||||
definition=definition,
|
||||
)
|
||||
with raises(VectorStoreOperationException):
|
||||
await collection.ensure_collection_exists()
|
||||
with raises(VectorStoreOperationException):
|
||||
await collection.ensure_collection_deleted()
|
||||
with raises(VectorStoreOperationException):
|
||||
await collection.collection_exists()
|
||||
with raises(VectorStoreOperationException):
|
||||
await collection.upsert({"id": "1", "content": "test", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]})
|
||||
with raises(VectorStoreOperationException):
|
||||
await collection.get("1")
|
||||
with raises(VectorStoreOperationException):
|
||||
await collection.delete("1")
|
||||
@@ -0,0 +1,70 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from weaviate import WeaviateAsyncClient
|
||||
from weaviate.collections.collections.async_ import _CollectionsAsync
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def collections_side_effects(request):
|
||||
"""Fixture that returns a dictionary of side effects for the mock async client methods."""
|
||||
return request.param if hasattr(request, "param") else {}
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_async_client(collections_side_effects) -> AsyncMock:
|
||||
"""Fixture to create a mock async client."""
|
||||
async_mock = AsyncMock(spec=WeaviateAsyncClient)
|
||||
async_mock.collections = AsyncMock(spec=_CollectionsAsync)
|
||||
async_mock.collections.create = AsyncMock()
|
||||
async_mock.collections.delete = AsyncMock()
|
||||
async_mock.collections.exists = AsyncMock()
|
||||
|
||||
if collections_side_effects:
|
||||
for method_name, exception in collections_side_effects.items():
|
||||
getattr(async_mock.collections, method_name).side_effect = exception
|
||||
|
||||
return async_mock
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def weaviate_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
|
||||
"""Fixture to set environment variables for Weaviate unit tests."""
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"WEAVIATE_URL": "test-api-key",
|
||||
"WEAVIATE_API_KEY": "https://test-endpoint.com",
|
||||
"WEAVIATE_LOCAL_HOST": "localhost",
|
||||
"WEAVIATE_LOCAL_PORT": 8080,
|
||||
"WEAVIATE_LOCAL_GRPC_PORT": 8081,
|
||||
"WEAVIATE_USE_EMBED": True,
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict)
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key not in exclude_list:
|
||||
monkeypatch.setenv(key, value)
|
||||
else:
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def clear_weaviate_env(monkeypatch):
|
||||
"""Fixture to clear the environment variables for Weaviate unit tests."""
|
||||
monkeypatch.delenv("WEAVIATE_URL", raising=False)
|
||||
monkeypatch.delenv("WEAVIATE_API_KEY", raising=False)
|
||||
monkeypatch.delenv("WEAVIATE_LOCAL_HOST", raising=False)
|
||||
monkeypatch.delenv("WEAVIATE_LOCAL_PORT", raising=False)
|
||||
monkeypatch.delenv("WEAVIATE_LOCAL_GRPC_PORT", raising=False)
|
||||
monkeypatch.delenv("WEAVIATE_USE_EMBED", raising=False)
|
||||
@@ -0,0 +1,429 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import ANY, AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from weaviate import WeaviateAsyncClient
|
||||
from weaviate.classes.config import Configure, DataType, Property
|
||||
from weaviate.collections.classes.config_vectorizers import VectorDistances
|
||||
from weaviate.collections.classes.data import DataObject
|
||||
|
||||
from semantic_kernel.connectors.weaviate import WeaviateCollection
|
||||
from semantic_kernel.exceptions import (
|
||||
ServiceInvalidExecutionSettingsError,
|
||||
VectorStoreInitializationException,
|
||||
VectorStoreOperationException,
|
||||
)
|
||||
|
||||
|
||||
@patch(
|
||||
"semantic_kernel.connectors.weaviate.use_async_with_weaviate_cloud",
|
||||
return_value=AsyncMock(spec=WeaviateAsyncClient),
|
||||
)
|
||||
def test_weaviate_collection_init_with_weaviate_cloud(
|
||||
mock_use_weaviate_cloud,
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
) -> None:
|
||||
"""Test the initialization of a WeaviateCollection object with Weaviate Cloud."""
|
||||
collection_name = "TestCollection"
|
||||
|
||||
collection = WeaviateCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
url="https://test.cloud.weaviate.com/",
|
||||
api_key="test_api_key",
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
assert collection.collection_name == collection_name
|
||||
assert collection.async_client is not None
|
||||
mock_use_weaviate_cloud.assert_called_once_with(
|
||||
cluster_url="https://test.cloud.weaviate.com/",
|
||||
auth_credentials=ANY,
|
||||
)
|
||||
|
||||
|
||||
@patch(
|
||||
"semantic_kernel.connectors.weaviate.use_async_with_local",
|
||||
return_value=AsyncMock(spec=WeaviateAsyncClient),
|
||||
)
|
||||
def test_weaviate_collection_init_with_local(
|
||||
mock_use_weaviate_local,
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
) -> None:
|
||||
"""Test the initialization of a WeaviateCollection object with Weaviate local deployment."""
|
||||
collection_name = "TestCollection"
|
||||
|
||||
collection = WeaviateCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
local_host="localhost",
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
assert collection.collection_name == collection_name
|
||||
assert collection.async_client is not None
|
||||
mock_use_weaviate_local.assert_called_once_with(host="localhost")
|
||||
|
||||
|
||||
@patch(
|
||||
"semantic_kernel.connectors.weaviate.use_async_with_embedded",
|
||||
return_value=AsyncMock(spec=WeaviateAsyncClient),
|
||||
)
|
||||
def test_weaviate_collection_init_with_embedded(
|
||||
mock_use_weaviate_embedded,
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
) -> None:
|
||||
"""Test the initialization of a WeaviateCollection object with Weaviate embedded deployment."""
|
||||
collection_name = "TestCollection"
|
||||
|
||||
collection = WeaviateCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
use_embed=True,
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
assert collection.collection_name == collection_name
|
||||
assert collection.async_client is not None
|
||||
mock_use_weaviate_embedded.assert_called_once()
|
||||
|
||||
|
||||
def test_weaviate_collection_init_with_invalid_settings_more_than_one_backends(
|
||||
weaviate_unit_test_env,
|
||||
record_type,
|
||||
definition,
|
||||
) -> None:
|
||||
"""Test the initialization of a WeaviateCollection object with multiple backend options enabled."""
|
||||
collection_name = "TestCollection"
|
||||
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
WeaviateCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
|
||||
def test_weaviate_collection_init_with_invalid_settings_no_backends(
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
) -> None:
|
||||
"""Test the initialization of a WeaviateCollection object with no backend options enabled."""
|
||||
collection_name = "TestCollection"
|
||||
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
WeaviateCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
|
||||
def test_weaviate_collection_init_with_custom_client(
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
) -> None:
|
||||
"""Test the initialization of a WeaviateCollection object with a custom client."""
|
||||
collection_name = "TestCollection"
|
||||
|
||||
collection = WeaviateCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
async_client=AsyncMock(spec=WeaviateAsyncClient),
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
assert collection.collection_name == collection_name
|
||||
assert collection.async_client is not None
|
||||
|
||||
|
||||
@patch(
|
||||
"semantic_kernel.connectors.weaviate.use_async_with_local",
|
||||
side_effect=Exception,
|
||||
)
|
||||
def test_weaviate_collection_init_fail_to_create_client(
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
) -> None:
|
||||
"""Test the initialization of a WeaviateCollection object raises an error when failing to create a client."""
|
||||
collection_name = "TestCollection"
|
||||
|
||||
with pytest.raises(VectorStoreInitializationException):
|
||||
WeaviateCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
local_host="localhost",
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
|
||||
@patch(
|
||||
"semantic_kernel.connectors.weaviate.use_async_with_weaviate_cloud",
|
||||
return_value=AsyncMock(spec=WeaviateAsyncClient),
|
||||
)
|
||||
def test_weaviate_collection_init_with_lower_case_collection_name(
|
||||
mock_use_weaviate_cloud,
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
) -> None:
|
||||
"""Test a collection name with lower case letters."""
|
||||
collection_name = "testCollection"
|
||||
|
||||
collection = WeaviateCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
url="https://test.cloud.weaviate.com",
|
||||
api_key="test_api_key",
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
assert collection.collection_name[0].isupper()
|
||||
assert collection.async_client is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("index_kind, distance_function", [("hnsw", "cosine_distance")])
|
||||
async def test_weaviate_collection_ensure_collection_exists(
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
mock_async_client,
|
||||
) -> None:
|
||||
"""Test the creation of a collection in Weaviate."""
|
||||
collection_name = "TestCollection"
|
||||
|
||||
collection = WeaviateCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
async_client=mock_async_client,
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
await collection.ensure_collection_exists()
|
||||
|
||||
mock_async_client.collections.create.assert_called_once_with(
|
||||
name=collection_name,
|
||||
properties=[
|
||||
Property(
|
||||
name="content",
|
||||
data_type=DataType.TEXT,
|
||||
)
|
||||
],
|
||||
vector_index_config=None,
|
||||
vectorizer_config=[
|
||||
Configure.NamedVectors.none(
|
||||
name="vector",
|
||||
vector_index_config=Configure.VectorIndex.hnsw(distance_metric=VectorDistances.COSINE),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"collections_side_effects",
|
||||
[
|
||||
{"create": Exception},
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_weaviate_collection_ensure_collection_exists_fail(
|
||||
mock_async_client,
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
) -> None:
|
||||
"""Test failing to create a collection in Weaviate."""
|
||||
collection_name = "TestCollection"
|
||||
|
||||
collection = WeaviateCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
async_client=mock_async_client,
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
with pytest.raises(VectorStoreOperationException):
|
||||
await collection.ensure_collection_exists()
|
||||
|
||||
|
||||
async def test_weaviate_collection_ensure_collection_deleted(
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
mock_async_client,
|
||||
) -> None:
|
||||
"""Test the deletion of a collection in Weaviate."""
|
||||
collection_name = "TestCollection"
|
||||
|
||||
collection = WeaviateCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
async_client=mock_async_client,
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
await collection.ensure_collection_deleted()
|
||||
|
||||
mock_async_client.collections.delete.assert_called_once_with(collection_name)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"collections_side_effects",
|
||||
[
|
||||
{"delete": Exception},
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_weaviate_collection_ensure_collection_deleted_fail(
|
||||
mock_async_client,
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
) -> None:
|
||||
"""Test failing to delete a collection in Weaviate."""
|
||||
collection_name = "TestCollection"
|
||||
|
||||
collection = WeaviateCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
async_client=mock_async_client,
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
with pytest.raises(VectorStoreOperationException):
|
||||
await collection.ensure_collection_deleted()
|
||||
|
||||
|
||||
async def test_weaviate_collection_collection_exist(
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
mock_async_client,
|
||||
) -> None:
|
||||
"""Test checking if a collection exists in Weaviate."""
|
||||
collection_name = "TestCollection"
|
||||
|
||||
collection = WeaviateCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
async_client=mock_async_client,
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
await collection.collection_exists()
|
||||
|
||||
mock_async_client.collections.exists.assert_called_once_with(collection_name)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"collections_side_effects",
|
||||
[
|
||||
{"exists": Exception},
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_weaviate_collection_collection_exist_fail(
|
||||
mock_async_client,
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
) -> None:
|
||||
"""Test failing to check if a collection exists in Weaviate."""
|
||||
collection_name = "TestCollection"
|
||||
|
||||
collection = WeaviateCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
async_client=mock_async_client,
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
with pytest.raises(VectorStoreOperationException):
|
||||
await collection.collection_exists()
|
||||
|
||||
|
||||
async def test_weaviate_collection_serialize_data(
|
||||
mock_async_client,
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
dataclass_vector_data_model,
|
||||
) -> None:
|
||||
"""Test upserting data into a collection in Weaviate."""
|
||||
collection_name = "TestCollection"
|
||||
|
||||
collection = WeaviateCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
async_client=mock_async_client,
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
with patch.object(collection, "_inner_upsert") as mock_inner_upsert:
|
||||
data = dataclass_vector_data_model()
|
||||
await collection.upsert(data)
|
||||
|
||||
mock_inner_upsert.assert_called_once_with([
|
||||
DataObject(
|
||||
properties={"content": "content1"},
|
||||
uuid=data.id,
|
||||
vector={"vector": data.vector},
|
||||
references=None,
|
||||
)
|
||||
])
|
||||
|
||||
|
||||
async def test_weaviate_collection_deserialize_data(
|
||||
mock_async_client,
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
dataclass_vector_data_model,
|
||||
) -> None:
|
||||
"""Test getting data from a collection in Weaviate."""
|
||||
collection_name = "TestCollection"
|
||||
|
||||
collection = WeaviateCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
async_client=mock_async_client,
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
data = dataclass_vector_data_model()
|
||||
weaviate_data_object = DataObject(
|
||||
properties={"content": "content1"},
|
||||
uuid=data.id,
|
||||
vector={"vector": data.vector or [1, 2, 3]},
|
||||
)
|
||||
|
||||
with patch.object(collection, "_inner_get", return_value=[weaviate_data_object]) as mock_inner_get:
|
||||
await collection.get(key=data.id)
|
||||
|
||||
mock_inner_get.assert_called_once_with([data.id], include_vectors=False, options=None)
|
||||
@@ -0,0 +1,120 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import ANY, AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from weaviate import WeaviateAsyncClient
|
||||
|
||||
from semantic_kernel.connectors.weaviate import WeaviateStore
|
||||
from semantic_kernel.exceptions import ServiceInvalidExecutionSettingsError, VectorStoreInitializationException
|
||||
|
||||
|
||||
@patch(
|
||||
"semantic_kernel.connectors.weaviate.use_async_with_weaviate_cloud",
|
||||
return_value=AsyncMock(spec=WeaviateAsyncClient),
|
||||
)
|
||||
def test_weaviate_store_init_with_weaviate_cloud(
|
||||
mock_use_weaviate_cloud,
|
||||
clear_weaviate_env,
|
||||
) -> None:
|
||||
"""Test the initialization of a WeaviateStore object with Weaviate Cloud."""
|
||||
store = WeaviateStore(
|
||||
url="https://test.cloud.weaviate.com/",
|
||||
api_key="test_api_key",
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
assert store.async_client is not None
|
||||
mock_use_weaviate_cloud.assert_called_once_with(
|
||||
cluster_url="https://test.cloud.weaviate.com/",
|
||||
auth_credentials=ANY,
|
||||
)
|
||||
|
||||
|
||||
@patch(
|
||||
"semantic_kernel.connectors.weaviate.use_async_with_local",
|
||||
return_value=AsyncMock(spec=WeaviateAsyncClient),
|
||||
)
|
||||
def test_weaviate_store_init_with_local(
|
||||
mock_use_weaviate_local,
|
||||
clear_weaviate_env,
|
||||
) -> None:
|
||||
"""Test the initialization of a WeaviateStore object with Weaviate local deployment."""
|
||||
store = WeaviateStore(
|
||||
local_host="localhost",
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
assert store.async_client is not None
|
||||
mock_use_weaviate_local.assert_called_once_with(host="localhost")
|
||||
|
||||
|
||||
@patch(
|
||||
"semantic_kernel.connectors.weaviate.use_async_with_embedded",
|
||||
return_value=AsyncMock(spec=WeaviateAsyncClient),
|
||||
)
|
||||
def test_weaviate_store_init_with_embedded(
|
||||
mock_use_weaviate_embedded,
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
) -> None:
|
||||
"""Test the initialization of a WeaviateStore object with Weaviate embedded deployment."""
|
||||
store = WeaviateStore(
|
||||
use_embed=True,
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
assert store.async_client is not None
|
||||
mock_use_weaviate_embedded.assert_called_once()
|
||||
|
||||
|
||||
def test_weaviate_store_init_with_invalid_settings_more_than_one_backends(
|
||||
weaviate_unit_test_env,
|
||||
) -> None:
|
||||
"""Test the initialization of a WeaviateStore object with multiple backend options enabled."""
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
WeaviateStore(
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
|
||||
def test_weaviate_store_init_with_invalid_settings_no_backends(
|
||||
clear_weaviate_env,
|
||||
) -> None:
|
||||
"""Test the initialization of a WeaviateStore object with no backend options enabled."""
|
||||
with pytest.raises(ServiceInvalidExecutionSettingsError):
|
||||
WeaviateStore(
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
|
||||
def test_weaviate_store_init_with_custom_client(
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
) -> None:
|
||||
"""Test the initialization of a WeaviateStore object with a custom client."""
|
||||
store = WeaviateStore(
|
||||
async_client=AsyncMock(spec=WeaviateAsyncClient),
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
|
||||
assert store.async_client is not None
|
||||
|
||||
|
||||
@patch(
|
||||
"semantic_kernel.connectors.weaviate.use_async_with_local",
|
||||
side_effect=Exception,
|
||||
)
|
||||
def test_weaviate_store_init_fail_to_create_client(
|
||||
clear_weaviate_env,
|
||||
record_type,
|
||||
definition,
|
||||
) -> None:
|
||||
"""Test the initialization of a WeaviateStore object raises an error when failing to create a client."""
|
||||
with pytest.raises(VectorStoreInitializationException):
|
||||
WeaviateStore(
|
||||
local_host="localhost",
|
||||
env_file_path="fake_env_file_path.env",
|
||||
)
|
||||
@@ -0,0 +1,79 @@
|
||||
{
|
||||
"openapi": "3.0.1",
|
||||
"info": {
|
||||
"title": "Semantic Kernel Open API Sample",
|
||||
"description": "A sample Open API schema with endpoints which have security requirements defined.",
|
||||
"version": "1.0"
|
||||
},
|
||||
"servers": [
|
||||
{
|
||||
"url": "https://example.org"
|
||||
}
|
||||
],
|
||||
"paths": {
|
||||
"/use_global_security": {
|
||||
"get": {
|
||||
"summary": "No security defined on operation",
|
||||
"description": "",
|
||||
"operationId": "NoSecurity",
|
||||
"parameters": [],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "default"
|
||||
}
|
||||
}
|
||||
},
|
||||
"post": {
|
||||
"summary": "Security defined on operation",
|
||||
"description": "",
|
||||
"operationId": "Security",
|
||||
"parameters": [],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "default"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"ApiKeyAuthQuery": []
|
||||
}
|
||||
]
|
||||
},
|
||||
"put": {
|
||||
"summary": "Security defined on operation with new scope",
|
||||
"description": "",
|
||||
"operationId": "SecurityAndScope",
|
||||
"parameters": [],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "default"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"ApiKeyAuthQuery": []
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
"components": {
|
||||
"securitySchemes": {
|
||||
"ApiKeyAuthHeader": {
|
||||
"type": "apiKey",
|
||||
"in": "header",
|
||||
"name": "X-API-KEY"
|
||||
},
|
||||
"ApiKeyAuthQuery": {
|
||||
"type": "apiKey",
|
||||
"in": "query",
|
||||
"name": "apiKey"
|
||||
}
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"ApiKeyAuthHeader": []
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
openapi: 3.0.3
|
||||
info:
|
||||
title: Simple HTTPBin API
|
||||
version: 1.0.0
|
||||
|
||||
servers:
|
||||
- url: https://httpbin.org
|
||||
|
||||
paths:
|
||||
/get:
|
||||
get:
|
||||
operationId: duplicateId
|
||||
summary: Simple GET request to httpbin.org
|
||||
responses:
|
||||
'200':
|
||||
description: Successful response from httpbin
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
type: object
|
||||
|
||||
/ip:
|
||||
get:
|
||||
operationId: duplicateId
|
||||
summary: Get client IP address from httpbin
|
||||
responses:
|
||||
'200':
|
||||
description: Successful response with IP
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
type: object
|
||||
@@ -0,0 +1,21 @@
|
||||
openapi: 3.0.0
|
||||
info:
|
||||
title: Test API
|
||||
version: 1.0.0
|
||||
servers:
|
||||
- url: http://example.com
|
||||
paths:
|
||||
/todos:
|
||||
get:
|
||||
summary: Get all todos
|
||||
operationId: getTodos
|
||||
responses:
|
||||
'200':
|
||||
description: OK
|
||||
parameters:
|
||||
- name: Authorization
|
||||
in: header
|
||||
required: true
|
||||
schema:
|
||||
type: string
|
||||
description: The authorization token
|
||||
@@ -0,0 +1,17 @@
|
||||
openapi: 3.0.3
|
||||
info:
|
||||
title: Simple HTTPBin API
|
||||
version: 1.0.0
|
||||
servers:
|
||||
- url: https://httpbin.org
|
||||
paths:
|
||||
/get:
|
||||
get:
|
||||
summary: Simple GET request to httpbin.org
|
||||
responses:
|
||||
'200':
|
||||
description: Successful response from httpbin
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
type: object
|
||||
@@ -0,0 +1,74 @@
|
||||
{
|
||||
"openapi": "3.0.1",
|
||||
"info": {
|
||||
"title": "Semantic Kernel Open API Sample",
|
||||
"description": "A sample Open API schema with endpoints which have security requirements defined.",
|
||||
"version": "1.0"
|
||||
},
|
||||
"servers": [
|
||||
{
|
||||
"url": "https://example.org"
|
||||
}
|
||||
],
|
||||
"paths": {
|
||||
"/use_global_security": {
|
||||
"get": {
|
||||
"summary": "No security defined on operation",
|
||||
"description": "",
|
||||
"operationId": "NoSecurity",
|
||||
"parameters": [],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "default"
|
||||
}
|
||||
}
|
||||
},
|
||||
"post": {
|
||||
"summary": "Security defined on operation",
|
||||
"description": "",
|
||||
"operationId": "Security",
|
||||
"parameters": [],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "default"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"ApiKeyAuthQuery": []
|
||||
}
|
||||
]
|
||||
},
|
||||
"put": {
|
||||
"summary": "Security defined on operation with new scope",
|
||||
"description": "",
|
||||
"operationId": "SecurityAndScope",
|
||||
"parameters": [],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "default"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"ApiKeyAuthQuery": ["new_scope"]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
"components": {
|
||||
"securitySchemes": {
|
||||
"ApiKeyAuthHeader": {
|
||||
"type": "apiKey",
|
||||
"in": "header",
|
||||
"name": "X-API-KEY"
|
||||
},
|
||||
"ApiKeyAuthQuery": {
|
||||
"type": "apiKey",
|
||||
"in": "query",
|
||||
"name": "apiKey"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
{
|
||||
"openapi": "3.0.1",
|
||||
"info": {
|
||||
"title": "Semantic Kernel Open API Sample",
|
||||
"description": "A sample Open API schema with endpoints which have security requirements defined.",
|
||||
"version": "1.0"
|
||||
},
|
||||
"servers": [
|
||||
{
|
||||
"url": "https://example.org"
|
||||
}
|
||||
],
|
||||
"paths": {
|
||||
"/use_global_security": {
|
||||
"get": {
|
||||
"summary": "No security defined on operation",
|
||||
"description": "",
|
||||
"operationId": "NoSecurity",
|
||||
"parameters": [],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "default"
|
||||
}
|
||||
}
|
||||
},
|
||||
"post": {
|
||||
"summary": "Security defined on operation",
|
||||
"description": "",
|
||||
"operationId": "Security",
|
||||
"parameters": [],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "default"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"OAuth2Auth": []
|
||||
}
|
||||
]
|
||||
},
|
||||
"put": {
|
||||
"summary": "Security defined on operation with new scope",
|
||||
"description": "",
|
||||
"operationId": "SecurityAndScope",
|
||||
"parameters": [],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "default"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"OAuth2Auth": [ "new_scope" ]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
"components": {
|
||||
"securitySchemes": {
|
||||
"OAuth2Auth": {
|
||||
"type": "oauth2",
|
||||
"flows": {
|
||||
"authorizationCode": {
|
||||
"authorizationUrl": "https://login.windows.net/common/oauth2/authorize",
|
||||
"tokenUrl": "https://login.windows.net/common/oauth2/authorize",
|
||||
"scopes": {}
|
||||
}
|
||||
}
|
||||
},
|
||||
"ApiKeyAuthHeader": {
|
||||
"type": "apiKey",
|
||||
"in": "header",
|
||||
"name": "X-API-KEY"
|
||||
},
|
||||
"ApiKeyAuthQuery": {
|
||||
"type": "apiKey",
|
||||
"in": "query",
|
||||
"name": "apiKey"
|
||||
}
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"OAuth2Auth": []
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,129 @@
|
||||
openapi: 3.0.0
|
||||
info:
|
||||
title: Test API
|
||||
version: 1.0.0
|
||||
servers:
|
||||
- url: http://example.com
|
||||
- url: https://example-two.com
|
||||
paths:
|
||||
/todos:
|
||||
get:
|
||||
summary: Get all todos
|
||||
operationId: getTodos
|
||||
responses:
|
||||
'200':
|
||||
description: OK
|
||||
parameters:
|
||||
- name: Authorization
|
||||
in: header
|
||||
required: true
|
||||
schema:
|
||||
type: string
|
||||
description: The authorization token
|
||||
post:
|
||||
summary: Add a new todo
|
||||
operationId: addTodo
|
||||
requestBody:
|
||||
required: true
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
type: object
|
||||
properties:
|
||||
title:
|
||||
type: string
|
||||
description: The title of the todo
|
||||
example: Buy milk
|
||||
completed:
|
||||
type: boolean
|
||||
description: Whether the todo is completed or not
|
||||
example: false
|
||||
responses:
|
||||
'201':
|
||||
description: Created
|
||||
parameters:
|
||||
- name: Authorization
|
||||
in: header
|
||||
required: true
|
||||
schema:
|
||||
type: string
|
||||
description: The authorization token
|
||||
/todos/{id}:
|
||||
get:
|
||||
summary: Get a todo by ID
|
||||
operationId: getTodoById
|
||||
parameters:
|
||||
- name: id
|
||||
in: path
|
||||
required: true
|
||||
schema:
|
||||
type: integer
|
||||
minimum: 1
|
||||
- name: Authorization
|
||||
in: header
|
||||
required: true
|
||||
schema:
|
||||
type: string
|
||||
description: The authorization token
|
||||
responses:
|
||||
'200':
|
||||
description: OK
|
||||
put:
|
||||
summary: Update a todo by ID
|
||||
operationId: updateTodoById
|
||||
parameters:
|
||||
- name: id
|
||||
in: path
|
||||
required: true
|
||||
schema:
|
||||
type: integer
|
||||
minimum: 1
|
||||
- name: Authorization
|
||||
in: header
|
||||
required: true
|
||||
schema:
|
||||
type: string
|
||||
description: The authorization token
|
||||
- name: completed
|
||||
in: query
|
||||
required: false
|
||||
schema:
|
||||
type: boolean
|
||||
description: Whether the todo is completed or not
|
||||
requestBody:
|
||||
required: true
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
type: object
|
||||
properties:
|
||||
title:
|
||||
type: string
|
||||
description: The title of the todo
|
||||
example: Buy milk
|
||||
completed:
|
||||
type: boolean
|
||||
description: Whether the todo is completed or not
|
||||
example: false
|
||||
responses:
|
||||
'200':
|
||||
description: OK
|
||||
delete:
|
||||
summary: Delete a todo by ID
|
||||
operationId: deleteTodoById
|
||||
parameters:
|
||||
- name: id
|
||||
in: path
|
||||
required: true
|
||||
schema:
|
||||
type: integer
|
||||
minimum: 1
|
||||
- name: Authorization
|
||||
in: header
|
||||
required: true
|
||||
schema:
|
||||
type: string
|
||||
description: The authorization token
|
||||
responses:
|
||||
'204':
|
||||
description: No Content
|
||||
@@ -0,0 +1,57 @@
|
||||
openapi: 3.0.0
|
||||
info:
|
||||
title: Todo List API
|
||||
version: 1.0.0
|
||||
description: API for managing todo lists
|
||||
paths:
|
||||
/list:
|
||||
get:
|
||||
summary: Get todo list
|
||||
operationId: get_todo_list
|
||||
description: get todo list from specific group
|
||||
parameters:
|
||||
- name: listName
|
||||
in: query
|
||||
required: true
|
||||
description: todo list group name description
|
||||
schema:
|
||||
type: string
|
||||
description: todo list group name
|
||||
responses:
|
||||
"200":
|
||||
description: Successful response
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
type: array
|
||||
items:
|
||||
type: object
|
||||
properties:
|
||||
task:
|
||||
type: string
|
||||
listName:
|
||||
type: string
|
||||
|
||||
/add:
|
||||
post:
|
||||
summary: Add a task to a list
|
||||
operationId: add_todo_list
|
||||
description: add todo to specific group
|
||||
requestBody:
|
||||
required: true
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
type: object
|
||||
required:
|
||||
- task
|
||||
properties:
|
||||
task:
|
||||
type: string
|
||||
description: task name
|
||||
listName:
|
||||
type: string
|
||||
description: task group name
|
||||
responses:
|
||||
"201":
|
||||
description: Task added successfully
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user