chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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"},
|
||||
)
|
||||
Reference in New Issue
Block a user