chore: import upstream snapshot with attribution
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user