chore: import upstream snapshot with attribution
CodeQL / Analyze (csharp) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled

This commit is contained in:
wehub-resource-sync
2026-07-13 13:21:23 +08:00
commit b957a53def
5423 changed files with 863745 additions and 0 deletions
@@ -0,0 +1,96 @@
# Copyright (c) Microsoft. All rights reserved.
import sys
from collections.abc import AsyncGenerator
from typing import Any, ClassVar
import pytest
if sys.version_info >= (3, 12):
from typing import override # pragma: no cover
else:
from typing_extensions import override # pragma: no cover
import semantic_kernel
from semantic_kernel.connectors.ai.chat_completion_client_base import ChatCompletionClientBase
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
from semantic_kernel.connectors.ai.text_completion_client_base import TextCompletionClientBase
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.streaming_text_content import StreamingTextContent
from semantic_kernel.contents.text_content import TextContent
from semantic_kernel.utils.telemetry.model_diagnostics.model_diagnostics_settings import ModelDiagnosticSettings
@pytest.fixture()
def model_diagnostics_unit_test_env(monkeypatch):
"""Fixture to set environment variables for Model Diagnostics Unit Tests."""
env_vars = {
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS": "true",
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE": "true",
}
for key, value in env_vars.items():
monkeypatch.setenv(key, value)
# Need to reload the settings to pick up the new environment variables since the
# settings are loaded at import time and this fixture is called after the import
semantic_kernel.utils.telemetry.model_diagnostics.decorators.MODEL_DIAGNOSTICS_SETTINGS = ModelDiagnosticSettings()
@pytest.fixture()
def service_env_vars(monkeypatch, request):
"""Fixture to set environment variables for AI Service Unit Tests."""
for key, value in request.param.items():
monkeypatch.setenv(key, value)
class MockChatCompletion(ChatCompletionClientBase):
MODEL_PROVIDER_NAME: ClassVar[str] = "mock"
@override
async def _inner_get_chat_message_contents(
self,
chat_history: "ChatHistory",
settings: "PromptExecutionSettings",
) -> list["ChatMessageContent"]:
return []
@override
async def _inner_get_streaming_chat_message_contents(
self,
chat_history: "ChatHistory",
settings: "PromptExecutionSettings",
) -> AsyncGenerator[list["StreamingChatMessageContent"], Any]:
yield []
@override
def service_url(self) -> str | None:
return "http://mock-service-url"
class MockTextCompletion(TextCompletionClientBase):
MODEL_PROVIDER_NAME: ClassVar[str] = "mock"
@override
async def _inner_get_text_contents(
self,
prompt: str,
settings: "PromptExecutionSettings",
) -> list["TextContent"]:
return []
@override
async def _inner_get_streaming_text_contents(
self,
prompt: str,
settings: "PromptExecutionSettings",
) -> AsyncGenerator[list["StreamingTextContent"], Any]:
yield []
@override
def service_url(self) -> str | None:
# Returning None to test the case where the service URL is not available
return None
@@ -0,0 +1,137 @@
# Copyright (c) Microsoft. All rights reserved.
import pytest
from semantic_kernel.connectors.ai.anthropic.services.anthropic_chat_completion import AnthropicChatCompletion
from semantic_kernel.connectors.ai.google.google_ai.services.google_ai_chat_completion import GoogleAIChatCompletion
from semantic_kernel.connectors.ai.google.google_ai.services.google_ai_text_completion import GoogleAITextCompletion
from semantic_kernel.connectors.ai.google.vertex_ai.services.vertex_ai_chat_completion import VertexAIChatCompletion
from semantic_kernel.connectors.ai.google.vertex_ai.services.vertex_ai_text_completion import VertexAITextCompletion
from semantic_kernel.connectors.ai.mistral_ai.services.mistral_ai_chat_completion import MistralAIChatCompletion
from semantic_kernel.connectors.ai.ollama.services.ollama_chat_completion import OllamaChatCompletion
from semantic_kernel.connectors.ai.ollama.services.ollama_text_completion import OllamaTextCompletion
from semantic_kernel.connectors.ai.open_ai.services.open_ai_chat_completion import OpenAIChatCompletion
from semantic_kernel.connectors.ai.open_ai.services.open_ai_text_completion import OpenAITextCompletion
pytestmark = pytest.mark.parametrize(
"decorated_method, expected_attribute",
[
# OpenAIChatCompletion
pytest.param(
OpenAIChatCompletion._inner_get_chat_message_contents,
"__model_diagnostics_chat_completion__",
id="OpenAIChatCompletion._inner_get_chat_message_contents",
),
pytest.param(
OpenAIChatCompletion._inner_get_streaming_chat_message_contents,
"__model_diagnostics_streaming_chat_completion__",
id="OpenAIChatCompletion._inner_get_streaming_chat_message_contents",
),
# OpenAITextCompletion
pytest.param(
OpenAITextCompletion._inner_get_text_contents,
"__model_diagnostics_text_completion__",
id="OpenAITextCompletion._inner_get_text_contents",
),
pytest.param(
OpenAITextCompletion._inner_get_streaming_text_contents,
"__model_diagnostics_streaming_text_completion__",
id="OpenAITextCompletion._inner_get_streaming_text_contents",
),
# OllamaChatCompletion
pytest.param(
OllamaChatCompletion._inner_get_chat_message_contents,
"__model_diagnostics_chat_completion__",
id="OllamaChatCompletion._inner_get_chat_message_contents",
),
pytest.param(
OllamaChatCompletion._inner_get_streaming_chat_message_contents,
"__model_diagnostics_streaming_chat_completion__",
id="OllamaChatCompletion._inner_get_streaming_chat_message_contents",
),
# OllamaTextCompletion
pytest.param(
OllamaTextCompletion._inner_get_text_contents,
"__model_diagnostics_text_completion__",
id="OllamaTextCompletion._inner_get_text_contents",
),
pytest.param(
OllamaTextCompletion._inner_get_streaming_text_contents,
"__model_diagnostics_streaming_text_completion__",
id="OllamaTextCompletion._inner_get_streaming_text_contents",
),
# MistralAIChatCompletion
pytest.param(
MistralAIChatCompletion._inner_get_chat_message_contents,
"__model_diagnostics_chat_completion__",
id="MistralAIChatCompletion._inner_get_chat_message_contents",
),
pytest.param(
MistralAIChatCompletion._inner_get_streaming_chat_message_contents,
"__model_diagnostics_streaming_chat_completion__",
id="MistralAIChatCompletion._inner_get_streaming_chat_message_contents",
),
# VertexAIChatCompletion
pytest.param(
VertexAIChatCompletion._inner_get_chat_message_contents,
"__model_diagnostics_chat_completion__",
id="VertexAIChatCompletion._inner_get_chat_message_contents",
),
pytest.param(
VertexAIChatCompletion._inner_get_streaming_chat_message_contents,
"__model_diagnostics_streaming_chat_completion__",
id="VertexAIChatCompletion._inner_get_streaming_chat_message_contents",
),
# VertexAITextCompletion
pytest.param(
VertexAITextCompletion._inner_get_text_contents,
"__model_diagnostics_text_completion__",
id="VertexAITextCompletion._inner_get_text_contents",
),
pytest.param(
VertexAITextCompletion._inner_get_streaming_text_contents,
"__model_diagnostics_streaming_text_completion__",
id="VertexAITextCompletion._inner_get_streaming_text_contents",
),
# GoogleAIChatCompletion
pytest.param(
GoogleAIChatCompletion._inner_get_chat_message_contents,
"__model_diagnostics_chat_completion__",
id="GoogleAIChatCompletion._inner_get_chat_message_contents",
),
pytest.param(
GoogleAIChatCompletion._inner_get_streaming_chat_message_contents,
"__model_diagnostics_streaming_chat_completion__",
id="GoogleAIChatCompletion._inner_get_streaming_chat_message_contents",
),
# GoogleAITextCompletion
pytest.param(
GoogleAITextCompletion._inner_get_text_contents,
"__model_diagnostics_text_completion__",
id="GoogleAITextCompletion._inner_get_text_contents",
),
pytest.param(
GoogleAITextCompletion._inner_get_streaming_text_contents,
"__model_diagnostics_streaming_text_completion__",
id="GoogleAITextCompletion._inner_get_streaming_text_contents",
),
# AnthropicChatCompletion
pytest.param(
AnthropicChatCompletion._inner_get_chat_message_contents,
"__model_diagnostics_chat_completion__",
id="AnthropicChatCompletion._inner_get_chat_message_contents",
),
pytest.param(
AnthropicChatCompletion._inner_get_streaming_chat_message_contents,
"__model_diagnostics_streaming_chat_completion__",
id="AnthropicChatCompletion._inner_get_streaming_chat_message_contents",
),
],
)
def test_decorated(decorated_method, expected_attribute):
"""Test that the connectors are being decorated properly with the model diagnostics decorators."""
assert hasattr(decorated_method, expected_attribute) and getattr(decorated_method, expected_attribute), (
f"{decorated_method} should be decorated with the appropriate model diagnostics decorator."
)
@@ -0,0 +1,205 @@
# Copyright (c) Microsoft. All rights reserved.
import json
from unittest.mock import call, patch
import pytest
from opentelemetry.trace import StatusCode
from semantic_kernel.connectors.ai.chat_completion_client_base import ChatCompletionClientBase
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
from semantic_kernel.contents.chat_message_content import ChatMessageContent
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 ServiceResponseException
from semantic_kernel.utils.telemetry.model_diagnostics import gen_ai_attributes
from semantic_kernel.utils.telemetry.model_diagnostics.decorators import (
CHAT_COMPLETION_OPERATION,
ChatHistoryMessageTimestampFilter,
trace_chat_completion,
)
from tests.unit.utils.model_diagnostics.conftest import MockChatCompletion
pytestmark = pytest.mark.parametrize(
"execution_settings, mock_response",
[
pytest.param(
PromptExecutionSettings(
extension_data={
"max_tokens": 1000,
"temperature": 0.5,
"top_p": 0.9,
}
),
[
ChatMessageContent(
role=AuthorRole.ASSISTANT,
ai_model_id="ai_model_id",
content="Test content",
metadata={"id": "test_id"},
finish_reason=FinishReason.STOP,
)
],
id="test_execution_settings_with_extension_data",
),
pytest.param(
PromptExecutionSettings(),
[
ChatMessageContent(
role=AuthorRole.ASSISTANT,
ai_model_id="ai_model_id",
metadata={"id": "test_id"},
finish_reason=FinishReason.STOP,
)
],
id="test_execution_settings_no_extension_data",
),
pytest.param(
PromptExecutionSettings(),
[
ChatMessageContent(
role=AuthorRole.ASSISTANT,
ai_model_id="ai_model_id",
metadata={},
finish_reason=FinishReason.STOP,
)
],
id="test_chat_message_content_no_metadata",
),
pytest.param(
PromptExecutionSettings(),
[
ChatMessageContent(
role=AuthorRole.ASSISTANT,
ai_model_id="ai_model_id",
metadata={"id": "test_id"},
)
],
id="test_chat_message_content_no_finish_reason",
),
],
)
@patch("semantic_kernel.utils.telemetry.model_diagnostics.decorators.logger")
@patch("semantic_kernel.utils.telemetry.model_diagnostics.decorators.tracer")
async def test_trace_chat_completion(
mock_tracer,
mock_logger,
execution_settings,
mock_response,
chat_history,
model_diagnostics_unit_test_env,
):
mock_span = mock_tracer.start_span.return_value
# Setup
chat_completion: ChatCompletionClientBase = MockChatCompletion(ai_model_id="ai_model_id")
with patch.object(MockChatCompletion, "_inner_get_chat_message_contents", return_value=mock_response):
# We need to reapply the decorator to the method since the mock will not have the decorator applied
MockChatCompletion._inner_get_chat_message_contents = trace_chat_completion(
MockChatCompletion.MODEL_PROVIDER_NAME
)(chat_completion._inner_get_chat_message_contents)
results: list[ChatMessageContent] = await chat_completion.get_chat_message_contents(
chat_history, execution_settings
)
assert results == mock_response
# Before the call to the model
mock_span.set_attributes.assert_called_with({
gen_ai_attributes.OPERATION: CHAT_COMPLETION_OPERATION,
gen_ai_attributes.SYSTEM: MockChatCompletion.MODEL_PROVIDER_NAME,
gen_ai_attributes.MODEL: chat_completion.ai_model_id,
})
mock_span.set_attribute.assert_any_call(gen_ai_attributes.ADDRESS, chat_completion.service_url())
# No all connectors take the same parameters
if execution_settings.extension_data.get("max_tokens") is not None:
mock_span.set_attribute.assert_any_call(
gen_ai_attributes.MAX_TOKENS, execution_settings.extension_data["max_tokens"]
)
if execution_settings.extension_data.get("temperature") is not None:
mock_span.set_attribute.assert_any_call(
gen_ai_attributes.TEMPERATURE, execution_settings.extension_data["temperature"]
)
if execution_settings.extension_data.get("top_p") is not None:
mock_span.set_attribute.assert_any_call(gen_ai_attributes.TOP_P, execution_settings.extension_data["top_p"])
mock_logger.info.assert_has_calls(
[
call(
json.dumps(message.to_dict()),
extra={
gen_ai_attributes.EVENT_NAME: gen_ai_attributes.ROLE_EVENT_MAP.get(message.role),
gen_ai_attributes.SYSTEM: MockChatCompletion.MODEL_PROVIDER_NAME,
ChatHistoryMessageTimestampFilter.INDEX_KEY: idx,
},
)
]
for idx, message in enumerate(chat_history)
)
# After the call to the model
# Not all connectors return the same metadata
if mock_response[0].metadata.get("id") is not None:
mock_span.set_attribute.assert_any_call(gen_ai_attributes.RESPONSE_ID, mock_response[0].metadata["id"])
if any(completion.finish_reason is not None for completion in mock_response):
mock_span.set_attribute.assert_any_call(
gen_ai_attributes.FINISH_REASON,
",".join([str(completion.finish_reason) for completion in mock_response]),
)
mock_logger.info.assert_any_call(
json.dumps({"message": results[0].to_dict(), "finish_reason": results[0].finish_reason}),
extra={
gen_ai_attributes.EVENT_NAME: gen_ai_attributes.CHOICE,
gen_ai_attributes.SYSTEM: MockChatCompletion.MODEL_PROVIDER_NAME,
},
)
@patch("semantic_kernel.utils.telemetry.model_diagnostics.decorators.logger")
@patch("semantic_kernel.utils.telemetry.model_diagnostics.decorators.tracer")
async def test_trace_chat_completion_exception(
mock_tracer,
mock_logger,
execution_settings,
mock_response,
chat_history,
model_diagnostics_unit_test_env,
):
mock_span = mock_tracer.start_span.return_value
# Setup
chat_completion: ChatCompletionClientBase = MockChatCompletion(ai_model_id="ai_model_id")
with patch.object(MockChatCompletion, "_inner_get_chat_message_contents", side_effect=ServiceResponseException()):
# We need to reapply the decorator to the method since the mock will not have the decorator applied
MockChatCompletion._inner_get_chat_message_contents = trace_chat_completion(
MockChatCompletion.MODEL_PROVIDER_NAME
)(chat_completion._inner_get_chat_message_contents)
with pytest.raises(ServiceResponseException):
await chat_completion.get_chat_message_contents(chat_history, execution_settings)
exception = ServiceResponseException()
mock_span.set_attribute.assert_any_call(gen_ai_attributes.ERROR_TYPE, str(type(exception)))
mock_span.set_status.assert_any_call(StatusCode.ERROR, repr(exception))
mock_span.end.assert_any_call()
mock_logger.info.assert_has_calls(
[
call(
json.dumps(message.to_dict()),
extra={
gen_ai_attributes.EVENT_NAME: gen_ai_attributes.ROLE_EVENT_MAP.get(message.role),
gen_ai_attributes.SYSTEM: MockChatCompletion.MODEL_PROVIDER_NAME,
ChatHistoryMessageTimestampFilter.INDEX_KEY: idx,
},
)
]
for idx, message in enumerate(chat_history)
)
@@ -0,0 +1,225 @@
# Copyright (c) Microsoft. All rights reserved.
import json
from collections.abc import AsyncGenerator
from functools import reduce
from unittest.mock import MagicMock, call, patch
import pytest
from opentelemetry.trace import StatusCode
from semantic_kernel.connectors.ai.chat_completion_client_base import ChatCompletionClientBase
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent
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 ServiceResponseException
from semantic_kernel.utils.telemetry.model_diagnostics import gen_ai_attributes
from semantic_kernel.utils.telemetry.model_diagnostics.decorators import (
CHAT_COMPLETION_OPERATION,
ChatHistoryMessageTimestampFilter,
trace_streaming_chat_completion,
)
from tests.unit.utils.model_diagnostics.conftest import MockChatCompletion
pytestmark = pytest.mark.parametrize(
"execution_settings, mock_response",
[
pytest.param(
PromptExecutionSettings(
extension_data={
"max_tokens": 1000,
"temperature": 0.5,
"top_p": 0.9,
}
),
[
StreamingChatMessageContent(
role=AuthorRole.ASSISTANT,
choice_index=0,
ai_model_id="ai_model_id",
content="Test content",
metadata={"id": "test_id"},
finish_reason=FinishReason.STOP,
)
],
id="test_execution_settings_with_extension_data",
),
pytest.param(
PromptExecutionSettings(),
[
StreamingChatMessageContent(
role=AuthorRole.ASSISTANT,
choice_index=0,
ai_model_id="ai_model_id",
metadata={"id": "test_id"},
finish_reason=FinishReason.STOP,
)
],
id="test_execution_settings_no_extension_data",
),
pytest.param(
PromptExecutionSettings(),
[
StreamingChatMessageContent(
role=AuthorRole.ASSISTANT,
choice_index=0,
ai_model_id="ai_model_id",
metadata={},
finish_reason=FinishReason.STOP,
)
],
id="test_chat_message_content_no_metadata",
),
pytest.param(
PromptExecutionSettings(),
[
StreamingChatMessageContent(
role=AuthorRole.ASSISTANT,
choice_index=0,
ai_model_id="ai_model_id",
metadata={"id": "test_id"},
)
],
id="test_chat_message_content_no_finish_reason",
),
],
)
@patch("semantic_kernel.utils.telemetry.model_diagnostics.decorators.logger")
@patch("semantic_kernel.utils.telemetry.model_diagnostics.decorators.tracer")
async def test_trace_streaming_chat_completion(
mock_tracer,
mock_logger,
execution_settings,
mock_response,
chat_history,
model_diagnostics_unit_test_env,
):
# Setup
mock_span = mock_tracer.start_span.return_value
chat_completion: ChatCompletionClientBase = MockChatCompletion(ai_model_id="ai_model_id")
iterable = MagicMock(spec=AsyncGenerator)
iterable.__aiter__.return_value = [mock_response]
with patch.object(MockChatCompletion, "_inner_get_streaming_chat_message_contents", return_value=iterable):
# We need to reapply the decorator to the method since the mock will not have the decorator applied
MockChatCompletion._inner_get_streaming_chat_message_contents = trace_streaming_chat_completion(
MockChatCompletion.MODEL_PROVIDER_NAME
)(chat_completion._inner_get_streaming_chat_message_contents)
updates = []
async for update in chat_completion._inner_get_streaming_chat_message_contents(
chat_history, execution_settings
):
updates.append(update)
updates_flatten = [reduce(lambda x, y: x + y, messages) for messages in updates]
assert updates_flatten == mock_response
# Before the call to the model
mock_span.set_attributes.assert_called_with({
gen_ai_attributes.OPERATION: CHAT_COMPLETION_OPERATION,
gen_ai_attributes.SYSTEM: MockChatCompletion.MODEL_PROVIDER_NAME,
gen_ai_attributes.MODEL: chat_completion.ai_model_id,
})
mock_span.set_attribute.assert_any_call(gen_ai_attributes.ADDRESS, chat_completion.service_url())
# No all connectors take the same parameters
if execution_settings.extension_data.get("max_tokens") is not None:
mock_span.set_attribute.assert_any_call(
gen_ai_attributes.MAX_TOKENS, execution_settings.extension_data["max_tokens"]
)
if execution_settings.extension_data.get("temperature") is not None:
mock_span.set_attribute.assert_any_call(
gen_ai_attributes.TEMPERATURE, execution_settings.extension_data["temperature"]
)
if execution_settings.extension_data.get("top_p") is not None:
mock_span.set_attribute.assert_any_call(gen_ai_attributes.TOP_P, execution_settings.extension_data["top_p"])
mock_logger.info.assert_has_calls(
[
call(
json.dumps(message.to_dict()),
extra={
gen_ai_attributes.EVENT_NAME: gen_ai_attributes.ROLE_EVENT_MAP.get(message.role),
gen_ai_attributes.SYSTEM: MockChatCompletion.MODEL_PROVIDER_NAME,
ChatHistoryMessageTimestampFilter.INDEX_KEY: idx,
},
)
]
for idx, message in enumerate(chat_history)
)
# After the call to the model
# Not all connectors return the same metadata
if mock_response[0].metadata.get("id") is not None:
mock_span.set_attribute.assert_any_call(gen_ai_attributes.RESPONSE_ID, mock_response[0].metadata["id"])
if any(completion.finish_reason is not None for completion in mock_response):
mock_span.set_attribute.assert_any_call(
gen_ai_attributes.FINISH_REASON,
",".join([str(completion.finish_reason) for completion in mock_response]),
)
mock_logger.info.assert_any_call(
json.dumps({
"message": updates_flatten[0].to_dict(),
"finish_reason": updates_flatten[0].finish_reason,
"index": 0,
}),
extra={
gen_ai_attributes.EVENT_NAME: gen_ai_attributes.CHOICE,
gen_ai_attributes.SYSTEM: MockChatCompletion.MODEL_PROVIDER_NAME,
},
)
@patch("semantic_kernel.utils.telemetry.model_diagnostics.decorators.logger")
@patch("semantic_kernel.utils.telemetry.model_diagnostics.decorators.tracer")
async def test_trace_streaming_chat_completion_exception(
mock_tracer,
mock_logger,
execution_settings,
mock_response,
chat_history,
model_diagnostics_unit_test_env,
):
# Setup
mock_span = mock_tracer.start_span.return_value
chat_completion: ChatCompletionClientBase = MockChatCompletion(ai_model_id="ai_model_id")
with patch.object(
MockChatCompletion, "_inner_get_streaming_chat_message_contents", side_effect=ServiceResponseException()
):
# We need to reapply the decorator to the method since the mock will not have the decorator applied
MockChatCompletion._inner_get_streaming_chat_message_contents = trace_streaming_chat_completion(
MockChatCompletion.MODEL_PROVIDER_NAME
)(chat_completion._inner_get_streaming_chat_message_contents)
with pytest.raises(ServiceResponseException):
async for update in chat_completion._inner_get_streaming_chat_message_contents(
chat_history, execution_settings
):
pass
exception = ServiceResponseException()
mock_span.set_attribute.assert_any_call(gen_ai_attributes.ERROR_TYPE, str(type(exception)))
mock_span.set_status.assert_any_call(StatusCode.ERROR, repr(exception))
mock_span.end.assert_any_call()
mock_logger.info.assert_has_calls(
[
call(
json.dumps(message.to_dict()),
extra={
gen_ai_attributes.EVENT_NAME: gen_ai_attributes.ROLE_EVENT_MAP.get(message.role),
gen_ai_attributes.SYSTEM: MockChatCompletion.MODEL_PROVIDER_NAME,
ChatHistoryMessageTimestampFilter.INDEX_KEY: idx,
},
)
]
for idx, message in enumerate(chat_history)
)
@@ -0,0 +1,183 @@
# Copyright (c) Microsoft. All rights reserved.
import json
from collections.abc import AsyncGenerator
from functools import reduce
from unittest.mock import ANY, MagicMock, patch
import pytest
from opentelemetry.trace import StatusCode
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
from semantic_kernel.connectors.ai.text_completion_client_base import TextCompletionClientBase
from semantic_kernel.contents.streaming_text_content import StreamingTextContent
from semantic_kernel.exceptions.service_exceptions import ServiceResponseException
from semantic_kernel.utils.telemetry.model_diagnostics import gen_ai_attributes
from semantic_kernel.utils.telemetry.model_diagnostics.decorators import (
TEXT_COMPLETION_OPERATION,
trace_streaming_text_completion,
)
from tests.unit.utils.model_diagnostics.conftest import MockTextCompletion
pytestmark = pytest.mark.parametrize(
"execution_settings, mock_response",
[
pytest.param(
PromptExecutionSettings(
extension_data={
"max_tokens": 1000,
"temperature": 0.5,
"top_p": 0.9,
}
),
[
StreamingTextContent(
choice_index=0,
ai_model_id="ai_model_id",
text="Test content",
metadata={"id": "test_id"},
)
],
id="test_execution_settings_with_extension_data",
),
pytest.param(
PromptExecutionSettings(),
[
StreamingTextContent(
choice_index=0,
ai_model_id="ai_model_id",
text="Test content",
metadata={"id": "test_id"},
)
],
id="test_execution_settings_no_extension_data",
),
pytest.param(
PromptExecutionSettings(),
[
StreamingTextContent(
choice_index=0,
ai_model_id="ai_model_id",
text="Test content",
metadata={},
)
],
id="test_text_content_no_metadata",
),
],
)
@patch("semantic_kernel.utils.telemetry.model_diagnostics.decorators.logger")
@patch("semantic_kernel.utils.telemetry.model_diagnostics.decorators.tracer")
async def test_trace_streaming_text_completion(
mock_tracer,
mock_logger,
execution_settings,
mock_response,
prompt,
model_diagnostics_unit_test_env,
):
mock_span = mock_tracer.start_span.return_value
# Setup
mock_span = mock_tracer.start_span.return_value
text_completion: TextCompletionClientBase = MockTextCompletion(ai_model_id="ai_model_id")
iterable = MagicMock(spec=AsyncGenerator)
iterable.__aiter__.return_value = [mock_response]
with patch.object(MockTextCompletion, "_inner_get_streaming_text_contents", return_value=iterable):
# We need to reapply the decorator to the method since the mock will not have the decorator applied
MockTextCompletion._inner_get_streaming_text_contents = trace_streaming_text_completion(
MockTextCompletion.MODEL_PROVIDER_NAME
)(text_completion._inner_get_streaming_text_contents)
updates = []
async for update in text_completion.get_streaming_text_contents(prompt=prompt, settings=execution_settings):
updates.append(update)
updates_flatten = [reduce(lambda x, y: x + y, update) for update in updates]
assert updates_flatten == mock_response
# Before the call to the model
mock_span.set_attributes.assert_called_with({
gen_ai_attributes.OPERATION: TEXT_COMPLETION_OPERATION,
gen_ai_attributes.SYSTEM: MockTextCompletion.MODEL_PROVIDER_NAME,
gen_ai_attributes.MODEL: text_completion.ai_model_id,
})
with pytest.raises(AssertionError):
# The service_url attribute is not set for text completion
mock_span.set_attribute.assert_any_call(gen_ai_attributes.ADDRESS, ANY)
# No all connectors take the same parameters
if execution_settings.extension_data.get("max_tokens") is not None:
mock_span.set_attribute.assert_any_call(
gen_ai_attributes.MAX_TOKENS, execution_settings.extension_data["max_tokens"]
)
if execution_settings.extension_data.get("temperature") is not None:
mock_span.set_attribute.assert_any_call(
gen_ai_attributes.TEMPERATURE, execution_settings.extension_data["temperature"]
)
if execution_settings.extension_data.get("top_p") is not None:
mock_span.set_attribute.assert_any_call(gen_ai_attributes.TOP_P, execution_settings.extension_data["top_p"])
mock_logger.info.assert_any_call(
prompt,
extra={
gen_ai_attributes.EVENT_NAME: gen_ai_attributes.PROMPT,
gen_ai_attributes.SYSTEM: MockTextCompletion.MODEL_PROVIDER_NAME,
},
)
# After the call to the model
# Not all connectors return the same metadata
if mock_response[0].metadata.get("id") is not None:
mock_span.set_attribute.assert_any_call(gen_ai_attributes.RESPONSE_ID, mock_response[0].metadata["id"])
mock_logger.info.assert_any_call(
json.dumps({"message": updates_flatten[0].to_dict(), "index": 0}),
extra={
gen_ai_attributes.EVENT_NAME: gen_ai_attributes.CHOICE,
gen_ai_attributes.SYSTEM: MockTextCompletion.MODEL_PROVIDER_NAME,
},
)
@patch("semantic_kernel.utils.telemetry.model_diagnostics.decorators.logger")
@patch("semantic_kernel.utils.telemetry.model_diagnostics.decorators.tracer")
async def test_trace_streaming_text_completion_exception(
mock_tracer,
mock_logger,
execution_settings,
mock_response,
prompt,
model_diagnostics_unit_test_env,
):
mock_span = mock_tracer.start_span.return_value
# Setup
mock_span = mock_tracer.start_span.return_value
text_completion: TextCompletionClientBase = MockTextCompletion(ai_model_id="ai_model_id")
with patch.object(MockTextCompletion, "_inner_get_streaming_text_contents", side_effect=ServiceResponseException()):
# We need to reapply the decorator to the method since the mock will not have the decorator applied
MockTextCompletion._inner_get_streaming_text_contents = trace_streaming_text_completion(
MockTextCompletion.MODEL_PROVIDER_NAME
)(text_completion._inner_get_streaming_text_contents)
with pytest.raises(ServiceResponseException):
async for update in text_completion.get_streaming_text_contents(prompt=prompt, settings=execution_settings):
pass
exception = ServiceResponseException()
mock_span.set_attribute.assert_any_call(gen_ai_attributes.ERROR_TYPE, str(type(exception)))
mock_span.set_status.assert_any_call(StatusCode.ERROR, repr(exception))
mock_span.end.assert_any_call()
mock_logger.info.assert_any_call(
prompt,
extra={
gen_ai_attributes.EVENT_NAME: gen_ai_attributes.PROMPT,
gen_ai_attributes.SYSTEM: MockTextCompletion.MODEL_PROVIDER_NAME,
},
)
@@ -0,0 +1,175 @@
# Copyright (c) Microsoft. All rights reserved.
import json
from unittest.mock import ANY, patch
import pytest
from opentelemetry.trace import StatusCode
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
from semantic_kernel.connectors.ai.text_completion_client_base import TextCompletionClientBase
from semantic_kernel.contents.chat_message_content import ChatMessageContent
from semantic_kernel.contents.text_content import TextContent
from semantic_kernel.exceptions.service_exceptions import ServiceResponseException
from semantic_kernel.utils.telemetry.model_diagnostics import gen_ai_attributes
from semantic_kernel.utils.telemetry.model_diagnostics.decorators import (
TEXT_COMPLETION_OPERATION,
trace_text_completion,
)
from tests.unit.utils.model_diagnostics.conftest import MockTextCompletion
pytestmark = pytest.mark.parametrize(
"execution_settings, mock_response",
[
pytest.param(
PromptExecutionSettings(
extension_data={
"max_tokens": 1000,
"temperature": 0.5,
"top_p": 0.9,
}
),
[
TextContent(
ai_model_id="ai_model_id",
text="Test content",
metadata={"id": "test_id"},
)
],
id="test_execution_settings_with_extension_data",
),
pytest.param(
PromptExecutionSettings(),
[
TextContent(
ai_model_id="ai_model_id",
text="Test content",
metadata={"id": "test_id"},
)
],
id="test_execution_settings_no_extension_data",
),
pytest.param(
PromptExecutionSettings(),
[
TextContent(
ai_model_id="ai_model_id",
text="Test content",
metadata={},
)
],
id="test_text_content_no_metadata",
),
],
)
@patch("semantic_kernel.utils.telemetry.model_diagnostics.decorators.logger")
@patch("semantic_kernel.utils.telemetry.model_diagnostics.decorators.tracer")
async def test_trace_text_completion(
mock_tracer,
mock_logger,
execution_settings,
mock_response,
prompt,
model_diagnostics_unit_test_env,
):
mock_span = mock_tracer.start_span.return_value
# Setup
mock_span = mock_tracer.start_span.return_value
text_completion: TextCompletionClientBase = MockTextCompletion(ai_model_id="ai_model_id")
with patch.object(MockTextCompletion, "_inner_get_text_contents", return_value=mock_response):
# We need to reapply the decorator to the method since the mock will not have the decorator applied
MockTextCompletion._inner_get_text_contents = trace_text_completion(MockTextCompletion.MODEL_PROVIDER_NAME)(
text_completion._inner_get_text_contents
)
results: list[ChatMessageContent] = await text_completion.get_text_contents(
prompt=prompt, settings=execution_settings
)
assert results == mock_response
# Before the call to the model
mock_span.set_attributes.assert_called_with({
gen_ai_attributes.OPERATION: TEXT_COMPLETION_OPERATION,
gen_ai_attributes.SYSTEM: MockTextCompletion.MODEL_PROVIDER_NAME,
gen_ai_attributes.MODEL: text_completion.ai_model_id,
})
with pytest.raises(AssertionError):
# The service_url attribute is not set for text completion
mock_span.set_attribute.assert_any_call(gen_ai_attributes.ADDRESS, ANY)
# No all connectors take the same parameters
if execution_settings.extension_data.get("max_tokens") is not None:
mock_span.set_attribute.assert_any_call(
gen_ai_attributes.MAX_TOKENS, execution_settings.extension_data["max_tokens"]
)
if execution_settings.extension_data.get("temperature") is not None:
mock_span.set_attribute.assert_any_call(
gen_ai_attributes.TEMPERATURE, execution_settings.extension_data["temperature"]
)
if execution_settings.extension_data.get("top_p") is not None:
mock_span.set_attribute.assert_any_call(gen_ai_attributes.TOP_P, execution_settings.extension_data["top_p"])
mock_logger.info.assert_any_call(
prompt,
extra={
gen_ai_attributes.EVENT_NAME: gen_ai_attributes.PROMPT,
gen_ai_attributes.SYSTEM: MockTextCompletion.MODEL_PROVIDER_NAME,
},
)
# After the call to the model
# Not all connectors return the same metadata
if mock_response[0].metadata.get("id") is not None:
mock_span.set_attribute.assert_any_call(gen_ai_attributes.RESPONSE_ID, mock_response[0].metadata["id"])
mock_logger.info.assert_any_call(
json.dumps({"message": results[0].to_dict()}),
extra={
gen_ai_attributes.EVENT_NAME: gen_ai_attributes.CHOICE,
gen_ai_attributes.SYSTEM: MockTextCompletion.MODEL_PROVIDER_NAME,
},
)
@patch("semantic_kernel.utils.telemetry.model_diagnostics.decorators.logger")
@patch("semantic_kernel.utils.telemetry.model_diagnostics.decorators.tracer")
async def test_trace_text_completion_exception(
mock_tracer,
mock_logger,
execution_settings,
mock_response,
prompt,
model_diagnostics_unit_test_env,
):
mock_span = mock_tracer.start_span.return_value
# Setup
mock_span = mock_tracer.start_span.return_value
text_completion: TextCompletionClientBase = MockTextCompletion(ai_model_id="ai_model_id")
with patch.object(MockTextCompletion, "_inner_get_text_contents", side_effect=ServiceResponseException()):
# We need to reapply the decorator to the method since the mock will not have the decorator applied
MockTextCompletion._inner_get_text_contents = trace_text_completion(MockTextCompletion.MODEL_PROVIDER_NAME)(
text_completion._inner_get_text_contents
)
with pytest.raises(ServiceResponseException):
await text_completion.get_text_contents(prompt=prompt, settings=execution_settings)
exception = ServiceResponseException()
mock_span.set_attribute.assert_any_call(gen_ai_attributes.ERROR_TYPE, str(type(exception)))
mock_span.set_status.assert_any_call(StatusCode.ERROR, repr(exception))
mock_span.end.assert_any_call()
mock_logger.info.assert_any_call(
prompt,
extra={
gen_ai_attributes.EVENT_NAME: gen_ai_attributes.PROMPT,
gen_ai_attributes.SYSTEM: MockTextCompletion.MODEL_PROVIDER_NAME,
},
)