Skip to content

Commit fb5aa6f

Browse files
TaoChenOSUTaoChenOSU
authoredNov 19, 2024
Python: Azure AI Inference tracing SDK (microsoft#9693)
### Motivation and Context <!-- Thank you for your contribution to the semantic-kernel repo! Please help reviewers and future users, providing the following information: 1. Why is this change required? 2. What problem does it solve? 3. What scenario does it contribute to? 4. If it fixes an open issue, please link to the issue here. --> Addresses: microsoft#9413 ### Description <!-- Describe your changes, the overall approach, the underlying design. These notes will help understanding how your code works. Thanks! --> The latest Azure AI Inference SDK has been released with the tracing package. We have decided to upgrade to the latest so that we will no longer need to instrument the Azure AI Inference connector with our own model diagnostics module. ### Contribution Checklist 1. Upgrade to the latest Azure AI Inference SDK with the tracing package. 2. Refactor the AI Inference connector to reduce duplicated code. 3. Some other minor fixes. <!-- Before submitting this PR, please make sure: --> - [x] The code builds clean without any errors or warnings - [x] The PR follows the [SK Contribution Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md) and the [pre-submission formatting script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts) raises no violations - [x] All unit tests pass, and I have added new tests where possible - [ ] I didn't break anyone 😄
1 parent 1ce4769 commit fb5aa6f

21 files changed

+1406
-1020
lines changed
 

‎python/.cspell.json

+2
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
"hnsw",
4040
"httpx",
4141
"huggingface",
42+
"Instrumentor",
4243
"kernelfunction",
4344
"logit",
4445
"logprobs",
@@ -61,6 +62,7 @@
6162
"serde",
6263
"skprompt",
6364
"templating",
65+
"uninstrument",
6466
"vectordb",
6567
"vectorizer",
6668
"vectorstoremodel",

‎python/pyproject.toml

+2-1
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,8 @@ dependencies = [
5050
### Optional dependencies
5151
[project.optional-dependencies]
5252
azure = [
53-
"azure-ai-inference >= 1.0.0b4",
53+
"azure-ai-inference >= 1.0.0b6",
54+
"azure-core-tracing-opentelemetry >= 1.0.0b11",
5455
"azure-search-documents >= 11.6.0b4",
5556
"azure-identity ~= 1.13",
5657
"azure-cosmos ~= 4.7"

‎python/samples/demos/telemetry/main.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ async def main(scenario: Literal["ai_service", "kernel_function", "auto_function
140140
with tracer.start_as_current_span("main") as current_span:
141141
print(f"Trace ID: {format_trace_id(current_span.get_span_context().trace_id)}")
142142

143-
stream = True
143+
stream = False
144144

145145
# Scenarios where telemetry is collected in the SDK, from the most basic to the most complex.
146146
if scenario == "ai_service" or scenario == "all":

‎python/samples/demos/telemetry/scenarios.py

+6
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,12 @@ def set_up_kernel() -> Kernel:
2222
# All built-in AI services are instrumented with telemetry.
2323
# Select any AI service to see the telemetry in action.
2424
kernel.add_service(OpenAIChatCompletion(service_id="open_ai"))
25+
# kernel.add_service(
26+
# AzureAIInferenceChatCompletion(
27+
# ai_model_id="serverless-deployment",
28+
# service_id="azure-ai-inference",
29+
# )
30+
# )
2531
# kernel.add_service(GoogleAIChatCompletion(service_id="google_ai"))
2632

2733
if (sample_plugin_path := get_sample_plugin_path()) is None:

‎python/semantic_kernel/connectors/ai/azure_ai_inference/azure_ai_inference_settings.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -34,4 +34,4 @@ class AzureAIInferenceSettings(KernelBaseSettings):
3434
env_prefix: ClassVar[str] = "AZURE_AI_INFERENCE_"
3535

3636
endpoint: HttpsUrl
37-
api_key: SecretStr
37+
api_key: SecretStr | None = None

‎python/semantic_kernel/connectors/ai/azure_ai_inference/services/azure_ai_inference_base.py

+93-5
Original file line numberDiff line numberDiff line change
@@ -3,23 +3,111 @@
33
import asyncio
44
import contextlib
55
from abc import ABC
6-
from typing import ClassVar
6+
from enum import Enum
7+
from typing import Any
78

89
from azure.ai.inference.aio import ChatCompletionsClient, EmbeddingsClient
10+
from azure.core.credentials import AzureKeyCredential
11+
from pydantic import ValidationError
912

13+
from semantic_kernel.connectors.ai.azure_ai_inference.azure_ai_inference_settings import AzureAIInferenceSettings
14+
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
1015
from semantic_kernel.kernel_pydantic import KernelBaseModel
16+
from semantic_kernel.utils.authentication.async_default_azure_credential_wrapper import (
17+
AsyncDefaultAzureCredentialWrapper,
18+
)
1119
from semantic_kernel.utils.experimental_decorator import experimental_class
20+
from semantic_kernel.utils.telemetry.user_agent import SEMANTIC_KERNEL_USER_AGENT
21+
22+
23+
class AzureAIInferenceClientType(Enum):
24+
"""Client type for Azure AI Inference."""
25+
26+
ChatCompletions = "ChatCompletions"
27+
Embeddings = "Embeddings"
28+
29+
@classmethod
30+
def get_client_class(cls, client_type: "AzureAIInferenceClientType") -> Any:
31+
"""Get the client class based on the client type."""
32+
class_mapping = {
33+
cls.ChatCompletions: ChatCompletionsClient,
34+
cls.Embeddings: EmbeddingsClient,
35+
}
36+
37+
return class_mapping[client_type]
1238

1339

1440
@experimental_class
1541
class AzureAIInferenceBase(KernelBaseModel, ABC):
1642
"""Azure AI Inference Chat Completion Service."""
1743

18-
MODEL_PROVIDER_NAME: ClassVar[str] = "azureai"
19-
2044
client: ChatCompletionsClient | EmbeddingsClient
45+
managed_client: bool = False
46+
47+
def __init__(
48+
self,
49+
client_type: AzureAIInferenceClientType,
50+
api_key: str | None = None,
51+
endpoint: str | None = None,
52+
env_file_path: str | None = None,
53+
env_file_encoding: str | None = None,
54+
client: ChatCompletionsClient | EmbeddingsClient | None = None,
55+
**kwargs: Any,
56+
) -> None:
57+
"""Initialize the Azure AI Inference Chat Completion service.
58+
59+
If no arguments are provided, the service will attempt to load the settings from the environment.
60+
The following environment variables are used:
61+
- AZURE_AI_INFERENCE_API_KEY
62+
- AZURE_AI_INFERENCE_ENDPOINT
63+
64+
Args:
65+
client_type (AzureAIInferenceClientType): The client type to use.
66+
api_key (str | None): The API key for the Azure AI Inference service deployment. (Optional)
67+
endpoint (str | None): The endpoint of the Azure AI Inference service deployment. (Optional)
68+
env_file_path (str | None): The path to the environment file. (Optional)
69+
env_file_encoding (str | None): The encoding of the environment file. (Optional)
70+
client (ChatCompletionsClient | None): The Azure AI Inference client to use. (Optional)
71+
**kwargs: Additional keyword arguments.
72+
73+
Raises:
74+
ServiceInitializationError: If an error occurs during initialization.
75+
"""
76+
managed_client = client is None
77+
if not client:
78+
try:
79+
azure_ai_inference_settings = AzureAIInferenceSettings.create(
80+
api_key=api_key,
81+
endpoint=endpoint,
82+
env_file_path=env_file_path,
83+
env_file_encoding=env_file_encoding,
84+
)
85+
except ValidationError as e:
86+
raise ServiceInitializationError(f"Failed to validate Azure AI Inference settings: {e}") from e
87+
88+
endpoint = str(azure_ai_inference_settings.endpoint)
89+
if azure_ai_inference_settings.api_key is not None:
90+
client = AzureAIInferenceClientType.get_client_class(client_type)(
91+
endpoint=endpoint,
92+
credential=AzureKeyCredential(azure_ai_inference_settings.api_key.get_secret_value()),
93+
user_agent=SEMANTIC_KERNEL_USER_AGENT,
94+
)
95+
else:
96+
# Try to create the client with a DefaultAzureCredential
97+
client = AzureAIInferenceClientType.get_client_class(client_type)(
98+
endpoint=endpoint,
99+
credential=AsyncDefaultAzureCredentialWrapper(),
100+
user_agent=SEMANTIC_KERNEL_USER_AGENT,
101+
)
102+
103+
super().__init__(
104+
client=client,
105+
managed_client=managed_client,
106+
**kwargs,
107+
)
21108

22109
def __del__(self) -> None:
23110
"""Close the client when the object is deleted."""
24-
with contextlib.suppress(Exception):
25-
asyncio.get_running_loop().create_task(self.client.close())
111+
if self.managed_client:
112+
with contextlib.suppress(Exception):
113+
asyncio.get_running_loop().create_task(self.client.close())

‎python/semantic_kernel/connectors/ai/azure_ai_inference/services/azure_ai_inference_chat_completion.py

+24-60
Original file line numberDiff line numberDiff line change
@@ -20,22 +20,19 @@
2020
StreamingChatChoiceUpdate,
2121
StreamingChatCompletionsUpdate,
2222
)
23-
from azure.core.credentials import AzureKeyCredential
24-
from azure.identity import DefaultAzureCredential
25-
from pydantic import ValidationError
2623

27-
from semantic_kernel.connectors.ai.azure_ai_inference import (
28-
AzureAIInferenceChatPromptExecutionSettings,
29-
AzureAIInferenceSettings,
24+
from semantic_kernel.connectors.ai.azure_ai_inference import AzureAIInferenceChatPromptExecutionSettings
25+
from semantic_kernel.connectors.ai.azure_ai_inference.services.azure_ai_inference_base import (
26+
AzureAIInferenceBase,
27+
AzureAIInferenceClientType,
3028
)
31-
from semantic_kernel.connectors.ai.azure_ai_inference.services.azure_ai_inference_base import AzureAIInferenceBase
29+
from semantic_kernel.connectors.ai.azure_ai_inference.services.azure_ai_inference_tracing import AzureAIInferenceTracing
3230
from semantic_kernel.connectors.ai.azure_ai_inference.services.utils import MESSAGE_CONVERTERS
3331
from semantic_kernel.connectors.ai.chat_completion_client_base import ChatCompletionClientBase
3432
from semantic_kernel.connectors.ai.completion_usage import CompletionUsage
3533
from semantic_kernel.connectors.ai.function_call_choice_configuration import FunctionCallChoiceConfiguration
3634
from semantic_kernel.connectors.ai.function_calling_utils import update_settings_from_function_call_configuration
3735
from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceType
38-
from semantic_kernel.connectors.ai.open_ai.const import DEFAULT_AZURE_API_VERSION
3936
from semantic_kernel.contents.chat_history import ChatHistory
4037
from semantic_kernel.contents.chat_message_content import ITEM_TYPES, ChatMessageContent
4138
from semantic_kernel.contents.function_call_content import FunctionCallContent
@@ -45,16 +42,8 @@
4542
from semantic_kernel.contents.text_content import TextContent
4643
from semantic_kernel.contents.utils.author_role import AuthorRole
4744
from semantic_kernel.contents.utils.finish_reason import FinishReason
48-
from semantic_kernel.exceptions.service_exceptions import (
49-
ServiceInitializationError,
50-
ServiceInvalidExecutionSettingsError,
51-
)
45+
from semantic_kernel.exceptions.service_exceptions import ServiceInvalidExecutionSettingsError
5246
from semantic_kernel.utils.experimental_decorator import experimental_class
53-
from semantic_kernel.utils.telemetry.model_diagnostics.decorators import (
54-
trace_chat_completion,
55-
trace_streaming_chat_completion,
56-
)
57-
from semantic_kernel.utils.telemetry.user_agent import SEMANTIC_KERNEL_USER_AGENT
5847

5948
if TYPE_CHECKING:
6049
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
@@ -97,39 +86,14 @@ def __init__(
9786
Raises:
9887
ServiceInitializationError: If an error occurs during initialization.
9988
"""
100-
if not client:
101-
try:
102-
azure_ai_inference_settings = AzureAIInferenceSettings.create(
103-
api_key=api_key,
104-
endpoint=endpoint,
105-
env_file_path=env_file_path,
106-
env_file_encoding=env_file_encoding,
107-
)
108-
except ValidationError as e:
109-
raise ServiceInitializationError(f"Failed to validate Azure AI Inference settings: {e}") from e
110-
111-
endpoint_to_use: str = str(azure_ai_inference_settings.endpoint)
112-
if azure_ai_inference_settings.api_key is not None:
113-
client = ChatCompletionsClient(
114-
endpoint=endpoint_to_use,
115-
credential=AzureKeyCredential(azure_ai_inference_settings.api_key.get_secret_value()),
116-
user_agent=SEMANTIC_KERNEL_USER_AGENT,
117-
)
118-
else:
119-
# Try to create the client with a DefaultAzureCredential
120-
client = (
121-
ChatCompletionsClient(
122-
endpoint=endpoint_to_use,
123-
credential=DefaultAzureCredential(),
124-
credential_scopes=["https://cognitiveservices.azure.com/.default"],
125-
api_version=DEFAULT_AZURE_API_VERSION,
126-
user_agent=SEMANTIC_KERNEL_USER_AGENT,
127-
),
128-
)
129-
13089
super().__init__(
13190
ai_model_id=ai_model_id,
13291
service_id=service_id or ai_model_id,
92+
client_type=AzureAIInferenceClientType.ChatCompletions,
93+
api_key=api_key,
94+
endpoint=endpoint,
95+
env_file_path=env_file_path,
96+
env_file_encoding=env_file_encoding,
13397
client=client,
13498
)
13599

@@ -149,7 +113,6 @@ def service_url(self) -> str | None:
149113
return None
150114

151115
@override
152-
@trace_chat_completion(AzureAIInferenceBase.MODEL_PROVIDER_NAME)
153116
async def _inner_get_chat_message_contents(
154117
self,
155118
chat_history: "ChatHistory",
@@ -160,17 +123,17 @@ async def _inner_get_chat_message_contents(
160123
assert isinstance(settings, AzureAIInferenceChatPromptExecutionSettings) # nosec
161124

162125
assert isinstance(self.client, ChatCompletionsClient) # nosec
163-
response: ChatCompletions = await self.client.complete(
164-
messages=self._prepare_chat_history_for_request(chat_history),
165-
model_extras=settings.extra_parameters,
166-
**settings.prepare_settings_dict(),
167-
)
126+
with AzureAIInferenceTracing():
127+
response: ChatCompletions = await self.client.complete(
128+
messages=self._prepare_chat_history_for_request(chat_history),
129+
model_extras=settings.extra_parameters,
130+
**settings.prepare_settings_dict(),
131+
)
168132
response_metadata = self._get_metadata_from_response(response)
169133

170134
return [self._create_chat_message_content(response, choice, response_metadata) for choice in response.choices]
171135

172136
@override
173-
@trace_streaming_chat_completion(AzureAIInferenceBase.MODEL_PROVIDER_NAME)
174137
async def _inner_get_streaming_chat_message_contents(
175138
self,
176139
chat_history: "ChatHistory",
@@ -181,12 +144,13 @@ async def _inner_get_streaming_chat_message_contents(
181144
assert isinstance(settings, AzureAIInferenceChatPromptExecutionSettings) # nosec
182145

183146
assert isinstance(self.client, ChatCompletionsClient) # nosec
184-
response: AsyncStreamingChatCompletions = await self.client.complete(
185-
stream=True,
186-
messages=self._prepare_chat_history_for_request(chat_history),
187-
model_extras=settings.extra_parameters,
188-
**settings.prepare_settings_dict(),
189-
)
147+
with AzureAIInferenceTracing():
148+
response: AsyncStreamingChatCompletions = await self.client.complete(
149+
stream=True,
150+
messages=self._prepare_chat_history_for_request(chat_history),
151+
model_extras=settings.extra_parameters,
152+
**settings.prepare_settings_dict(),
153+
)
190154

191155
async for chunk in response:
192156
if len(chunk.choices) == 0:

‎python/semantic_kernel/connectors/ai/azure_ai_inference/services/azure_ai_inference_text_embedding.py

+9-36
Original file line numberDiff line numberDiff line change
@@ -10,21 +10,17 @@
1010

1111
from azure.ai.inference.aio import EmbeddingsClient
1212
from azure.ai.inference.models import EmbeddingsResult
13-
from azure.core.credentials import AzureKeyCredential
14-
from azure.identity import DefaultAzureCredential
1513
from numpy import array, ndarray
16-
from pydantic import ValidationError
1714

1815
from semantic_kernel.connectors.ai.azure_ai_inference.azure_ai_inference_prompt_execution_settings import (
1916
AzureAIInferenceEmbeddingPromptExecutionSettings,
2017
)
21-
from semantic_kernel.connectors.ai.azure_ai_inference.azure_ai_inference_settings import AzureAIInferenceSettings
22-
from semantic_kernel.connectors.ai.azure_ai_inference.services.azure_ai_inference_base import AzureAIInferenceBase
18+
from semantic_kernel.connectors.ai.azure_ai_inference.services.azure_ai_inference_base import (
19+
AzureAIInferenceBase,
20+
AzureAIInferenceClientType,
21+
)
2322
from semantic_kernel.connectors.ai.embeddings.embedding_generator_base import EmbeddingGeneratorBase
24-
from semantic_kernel.connectors.ai.open_ai.const import DEFAULT_AZURE_API_VERSION
25-
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
2623
from semantic_kernel.utils.experimental_decorator import experimental_class
27-
from semantic_kernel.utils.telemetry.user_agent import SEMANTIC_KERNEL_USER_AGENT
2824

2925
if TYPE_CHECKING:
3026
from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
@@ -63,37 +59,14 @@ def __init__(
6359
Raises:
6460
ServiceInitializationError: If an error occurs during initialization.
6561
"""
66-
if not client:
67-
try:
68-
azure_ai_inference_settings = AzureAIInferenceSettings.create(
69-
api_key=api_key,
70-
endpoint=endpoint,
71-
env_file_path=env_file_path,
72-
env_file_encoding=env_file_encoding,
73-
)
74-
except ValidationError as e:
75-
raise ServiceInitializationError(f"Failed to validate Azure AI Inference settings: {e}") from e
76-
77-
endpoint = str(azure_ai_inference_settings.endpoint)
78-
if azure_ai_inference_settings.api_key is not None:
79-
client = EmbeddingsClient(
80-
endpoint=endpoint,
81-
credential=AzureKeyCredential(azure_ai_inference_settings.api_key.get_secret_value()),
82-
user_agent=SEMANTIC_KERNEL_USER_AGENT,
83-
)
84-
else:
85-
# Try to create the client with a DefaultAzureCredential
86-
client = EmbeddingsClient(
87-
endpoint=endpoint,
88-
credential=DefaultAzureCredential(),
89-
credential_scopes=["https://cognitiveservices.azure.com/.default"],
90-
api_version=DEFAULT_AZURE_API_VERSION,
91-
user_agent=SEMANTIC_KERNEL_USER_AGENT,
92-
)
93-
9462
super().__init__(
9563
ai_model_id=ai_model_id,
9664
service_id=service_id or ai_model_id,
65+
client_type=AzureAIInferenceClientType.Embeddings,
66+
api_key=api_key,
67+
endpoint=endpoint,
68+
env_file_path=env_file_path,
69+
env_file_encoding=env_file_encoding,
9770
client=client,
9871
)
9972

‎python/semantic_kernel/connectors/ai/azure_ai_inference/services/azure_ai_inference_tracing.py

+48
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
# Copyright (c) Microsoft. All rights reserved.
2+
3+
from azure.ai.inference.tracing import AIInferenceInstrumentor
4+
from azure.core.settings import settings
5+
6+
from semantic_kernel.kernel_pydantic import KernelBaseModel
7+
from semantic_kernel.utils.telemetry.model_diagnostics.model_diagnostics_settings import ModelDiagnosticSettings
8+
9+
10+
class AzureAIInferenceTracing(KernelBaseModel):
11+
"""Enable tracing for Azure AI Inference.
12+
13+
This class is intended to be used as a context manager.
14+
The instrument() call effect should be scoped to the context manager.
15+
"""
16+
17+
diagnostics_settings: ModelDiagnosticSettings
18+
19+
def __init__(self, diagnostics_settings: ModelDiagnosticSettings | None = None) -> None:
20+
"""Initialize the Azure AI Inference Tracing.
21+
22+
Args:
23+
diagnostics_settings (ModelDiagnosticSettings, optional): Model diagnostics settings. Defaults to None.
24+
"""
25+
settings.tracing_implementation = "opentelemetry"
26+
super().__init__(diagnostics_settings=diagnostics_settings or ModelDiagnosticSettings.create())
27+
28+
def __enter__(self) -> None:
29+
"""Enable tracing.
30+
31+
Both enable_otel_diagnostics and enable_otel_diagnostics_sensitive will enable tracing.
32+
enable_otel_diagnostics_sensitive will also enable content recording.
33+
"""
34+
if (
35+
self.diagnostics_settings.enable_otel_diagnostics
36+
or self.diagnostics_settings.enable_otel_diagnostics_sensitive
37+
):
38+
AIInferenceInstrumentor().instrument(
39+
enable_content_recording=self.diagnostics_settings.enable_otel_diagnostics_sensitive
40+
)
41+
42+
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
43+
"""Disable tracing."""
44+
if (
45+
self.diagnostics_settings.enable_otel_diagnostics
46+
or self.diagnostics_settings.enable_otel_diagnostics_sensitive
47+
):
48+
AIInferenceInstrumentor().uninstrument()

‎python/semantic_kernel/connectors/memory/azure_cosmos_db/azure_cosmos_db_no_sql_base.py

+5-2
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,15 @@
55
from pydantic import ValidationError
66

77
from semantic_kernel.connectors.memory.azure_cosmos_db.azure_cosmos_db_no_sql_settings import AzureCosmosDBNoSQLSettings
8-
from semantic_kernel.connectors.memory.azure_cosmos_db.utils import CosmosClientWrapper, DefaultAzureCredentialWrapper
8+
from semantic_kernel.connectors.memory.azure_cosmos_db.utils import CosmosClientWrapper
99
from semantic_kernel.exceptions.memory_connector_exceptions import (
1010
MemoryConnectorInitializationError,
1111
MemoryConnectorResourceNotFound,
1212
)
1313
from semantic_kernel.kernel_pydantic import KernelBaseModel
14+
from semantic_kernel.utils.authentication.async_default_azure_credential_wrapper import (
15+
AsyncDefaultAzureCredentialWrapper,
16+
)
1417
from semantic_kernel.utils.experimental_decorator import experimental_class
1518

1619

@@ -72,7 +75,7 @@ def __init__(
7275
)
7376
else:
7477
cosmos_client = CosmosClientWrapper(
75-
str(cosmos_db_nosql_settings.url), credential=DefaultAzureCredentialWrapper()
78+
str(cosmos_db_nosql_settings.url), credential=AsyncDefaultAzureCredentialWrapper()
7679
)
7780

7881
super().__init__(

‎python/semantic_kernel/connectors/memory/azure_cosmos_db/utils.py

-10
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55
from typing import Any
66

77
from azure.cosmos.aio import CosmosClient
8-
from azure.identity.aio import DefaultAzureCredential
98

109
from semantic_kernel.connectors.memory.azure_cosmos_db.azure_cosmos_db_no_sql_composite_key import (
1110
AzureCosmosDBNoSQLCompositeKey,
@@ -179,12 +178,3 @@ def __del__(self) -> None:
179178
"""Close the CosmosClient."""
180179
with contextlib.suppress(Exception):
181180
asyncio.get_running_loop().create_task(self.close())
182-
183-
184-
class DefaultAzureCredentialWrapper(DefaultAzureCredential):
185-
"""Wrapper to make sure the DefaultAzureCredential is closed properly."""
186-
187-
def __del__(self) -> None:
188-
"""Close the DefaultAzureCredential."""
189-
with contextlib.suppress(Exception):
190-
asyncio.get_running_loop().create_task(self.close())

‎python/semantic_kernel/utils/authentication/async_default_azure_credential_wrapper.py

+15
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# Copyright (c) Microsoft. All rights reserved.
2+
3+
import asyncio
4+
import contextlib
5+
6+
from azure.identity.aio import DefaultAzureCredential
7+
8+
9+
class AsyncDefaultAzureCredentialWrapper(DefaultAzureCredential):
10+
"""Wrapper to make sure the async version of the DefaultAzureCredential is closed properly."""
11+
12+
def __del__(self) -> None:
13+
"""Close the DefaultAzureCredential."""
14+
with contextlib.suppress(Exception):
15+
asyncio.get_running_loop().create_task(self.close())

‎python/semantic_kernel/utils/telemetry/model_diagnostics/model_diagnostics_settings.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ class ModelDiagnosticSettings(KernelBaseSettings):
1111
"""Settings for model diagnostics.
1212
1313
The settings are first loaded from environment variables with
14-
the prefix 'AZURE_AI_INFERENCE_'.
14+
the prefix 'SEMANTICKERNEL_EXPERIMENTAL_GENAI_'.
1515
If the environment variables are not found, the settings can
1616
be loaded from a .env file with the encoding 'utf-8'.
1717
If the settings are not found in the .env file, the settings

‎python/tests/integration/completions/chat_completion_test_base.py

-2
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,6 @@
3939
from semantic_kernel.connectors.ai.mistral_ai.services.mistral_ai_chat_completion import MistralAIChatCompletion
4040
from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import OllamaChatPromptExecutionSettings
4141
from semantic_kernel.connectors.ai.ollama.services.ollama_chat_completion import OllamaChatCompletion
42-
from semantic_kernel.connectors.ai.open_ai.const import DEFAULT_AZURE_API_VERSION
4342
from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.azure_chat_prompt_execution_settings import (
4443
AzureChatPromptExecutionSettings,
4544
)
@@ -134,7 +133,6 @@ def services(self) -> dict[str, tuple[ServiceType, type[PromptExecutionSettings]
134133
endpoint=f"{str(endpoint).strip('/')}/openai/deployments/{deployment_name}",
135134
credential=DefaultAzureCredential(),
136135
credential_scopes=["https://cognitiveservices.azure.com/.default"],
137-
api_version=DEFAULT_AZURE_API_VERSION,
138136
),
139137
)
140138

‎python/tests/integration/embeddings/test_embedding_service_base.py

-2
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828
from semantic_kernel.connectors.ai.mistral_ai.services.mistral_ai_text_embedding import MistralAITextEmbedding
2929
from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import OllamaEmbeddingPromptExecutionSettings
3030
from semantic_kernel.connectors.ai.ollama.services.ollama_text_embedding import OllamaTextEmbedding
31-
from semantic_kernel.connectors.ai.open_ai.const import DEFAULT_AZURE_API_VERSION
3231
from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.open_ai_prompt_execution_settings import (
3332
OpenAIEmbeddingPromptExecutionSettings,
3433
)
@@ -76,7 +75,6 @@ def services(self) -> dict[str, tuple[EmbeddingGeneratorBase, type[PromptExecuti
7675
endpoint=f"{str(endpoint).strip('/')}/openai/deployments/{deployment_name}",
7776
credential=DefaultAzureCredential(),
7877
credential_scopes=["https://cognitiveservices.azure.com/.default"],
79-
api_version=DEFAULT_AZURE_API_VERSION,
8078
),
8179
)
8280

‎python/tests/unit/connectors/ai/azure_ai_inference/conftest.py

+25
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,31 @@ def azure_ai_inference_unit_test_env(monkeypatch, exclude_list, override_env_par
5959
return env_vars
6060

6161

62+
@pytest.fixture()
63+
def model_diagnostics_test_env(monkeypatch, exclude_list, override_env_param_dict):
64+
"""Fixture to set environment variables for Azure AI Inference Unit Tests."""
65+
if exclude_list is None:
66+
exclude_list = []
67+
68+
if override_env_param_dict is None:
69+
override_env_param_dict = {}
70+
71+
env_vars = {
72+
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS": "true",
73+
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE": "true",
74+
}
75+
76+
env_vars.update(override_env_param_dict)
77+
78+
for key, value in env_vars.items():
79+
if key not in exclude_list:
80+
monkeypatch.setenv(key, value)
81+
else:
82+
monkeypatch.delenv(key, raising=False)
83+
84+
return env_vars
85+
86+
6287
@pytest.fixture(scope="function")
6388
def azure_ai_inference_client(azure_ai_inference_unit_test_env, request) -> ChatCompletionsClient | EmbeddingsClient:
6489
"""Fixture to create Azure AI Inference client for unit tests."""

‎python/tests/unit/connectors/ai/azure_ai_inference/services/test_azure_ai_inference_chat_completion.py

+1-8
Original file line numberDiff line numberDiff line change
@@ -78,20 +78,13 @@ def test_azure_ai_inference_chat_completion_init_with_custom_client(azure_ai_inf
7878
assert azure_ai_inference.client == client
7979

8080

81-
@pytest.mark.parametrize("exclude_list", [["AZURE_AI_INFERENCE_API_KEY"]], indirect=True)
82-
def test_azure_ai_inference_chat_completion_init_with_empty_api_key(azure_ai_inference_unit_test_env, model_id) -> None:
83-
"""Test initialization of AzureAIInferenceChatCompletion with empty API key"""
84-
with pytest.raises(ServiceInitializationError):
85-
AzureAIInferenceChatCompletion(model_id)
86-
87-
8881
@pytest.mark.parametrize("exclude_list", [["AZURE_AI_INFERENCE_ENDPOINT"]], indirect=True)
8982
def test_azure_ai_inference_chat_completion_init_with_empty_endpoint(
9083
azure_ai_inference_unit_test_env, model_id
9184
) -> None:
9285
"""Test initialization of AzureAIInferenceChatCompletion with empty endpoint"""
9386
with pytest.raises(ServiceInitializationError):
94-
AzureAIInferenceChatCompletion(model_id)
87+
AzureAIInferenceChatCompletion(model_id, env_file_path="fake_path")
9588

9689

9790
def test_prompt_execution_settings_class(azure_ai_inference_unit_test_env, model_id) -> None:

‎python/tests/unit/connectors/ai/azure_ai_inference/services/test_azure_ai_inference_text_embedding.py

+1-8
Original file line numberDiff line numberDiff line change
@@ -67,18 +67,11 @@ def test_azure_ai_inference_chat_completion_init_with_custom_client(azure_ai_inf
6767
assert azure_ai_inference.client == client
6868

6969

70-
@pytest.mark.parametrize("exclude_list", [["AZURE_AI_INFERENCE_API_KEY"]], indirect=True)
71-
def test_azure_ai_inference_text_embedding_init_with_empty_api_key(azure_ai_inference_unit_test_env, model_id) -> None:
72-
"""Test initialization of AzureAIInferenceTextEmbedding with empty API key"""
73-
with pytest.raises(ServiceInitializationError):
74-
AzureAIInferenceTextEmbedding(model_id)
75-
76-
7770
@pytest.mark.parametrize("exclude_list", [["AZURE_AI_INFERENCE_ENDPOINT"]], indirect=True)
7871
def test_azure_ai_inference_text_embedding_init_with_empty_endpoint(azure_ai_inference_unit_test_env, model_id) -> None:
7972
"""Test initialization of AzureAIInferenceTextEmbedding with empty endpoint"""
8073
with pytest.raises(ServiceInitializationError):
81-
AzureAIInferenceTextEmbedding(model_id)
74+
AzureAIInferenceTextEmbedding(model_id, env_file_path="fake_path")
8275

8376

8477
@pytest.mark.asyncio

‎python/tests/unit/connectors/ai/azure_ai_inference/services/test_azure_ai_inference_tracing.py

+245
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,245 @@
1+
# Copyright (c) Microsoft. All rights reserved.
2+
3+
from unittest.mock import AsyncMock, patch
4+
5+
import pytest
6+
from azure.ai.inference.aio import ChatCompletionsClient
7+
8+
from semantic_kernel.connectors.ai.azure_ai_inference.azure_ai_inference_prompt_execution_settings import (
9+
AzureAIInferenceChatPromptExecutionSettings,
10+
)
11+
from semantic_kernel.connectors.ai.azure_ai_inference.services.azure_ai_inference_chat_completion import (
12+
AzureAIInferenceChatCompletion,
13+
)
14+
from semantic_kernel.contents.chat_history import ChatHistory
15+
16+
17+
@pytest.mark.asyncio
18+
@pytest.mark.parametrize(
19+
"azure_ai_inference_service",
20+
[AzureAIInferenceChatCompletion.__name__],
21+
indirect=True,
22+
)
23+
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
24+
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.uninstrument")
25+
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.instrument")
26+
async def test_azure_ai_inference_chat_completion_instrumentation(
27+
mock_instrument,
28+
mock_uninstrument,
29+
mock_complete,
30+
azure_ai_inference_service,
31+
chat_history: ChatHistory,
32+
mock_azure_ai_inference_chat_completion_response,
33+
model_diagnostics_test_env,
34+
) -> None:
35+
"""Test completion of AzureAIInferenceChatCompletion"""
36+
settings = AzureAIInferenceChatPromptExecutionSettings()
37+
38+
mock_complete.return_value = mock_azure_ai_inference_chat_completion_response
39+
40+
await azure_ai_inference_service.get_chat_message_contents(chat_history=chat_history, settings=settings)
41+
42+
mock_instrument.assert_called_once_with(enable_content_recording=True)
43+
mock_uninstrument.assert_called_once()
44+
45+
46+
@pytest.mark.asyncio
47+
@pytest.mark.parametrize(
48+
"azure_ai_inference_service",
49+
[
50+
AzureAIInferenceChatCompletion.__name__,
51+
],
52+
indirect=True,
53+
)
54+
@pytest.mark.parametrize(
55+
"override_env_param_dict",
56+
[
57+
{
58+
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS": "False",
59+
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE": "False",
60+
},
61+
],
62+
indirect=True,
63+
)
64+
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
65+
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.uninstrument")
66+
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.instrument")
67+
async def test_azure_ai_inference_chat_completion_not_instrumentation(
68+
mock_instrument,
69+
mock_uninstrument,
70+
mock_complete,
71+
azure_ai_inference_service,
72+
chat_history: ChatHistory,
73+
mock_azure_ai_inference_chat_completion_response,
74+
model_diagnostics_test_env,
75+
) -> None:
76+
"""Test completion of AzureAIInferenceChatCompletion"""
77+
settings = AzureAIInferenceChatPromptExecutionSettings()
78+
79+
mock_complete.return_value = mock_azure_ai_inference_chat_completion_response
80+
81+
await azure_ai_inference_service.get_chat_message_contents(chat_history=chat_history, settings=settings)
82+
83+
mock_instrument.assert_not_called()
84+
mock_uninstrument.assert_not_called()
85+
86+
87+
@pytest.mark.asyncio
88+
@pytest.mark.parametrize(
89+
"azure_ai_inference_service",
90+
[
91+
AzureAIInferenceChatCompletion.__name__,
92+
],
93+
indirect=True,
94+
)
95+
@pytest.mark.parametrize(
96+
"override_env_param_dict",
97+
[
98+
{
99+
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS": "True",
100+
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE": "False",
101+
},
102+
],
103+
indirect=True,
104+
)
105+
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
106+
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.uninstrument")
107+
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.instrument")
108+
async def test_azure_ai_inference_chat_completion_instrumentation_without_sensitive(
109+
mock_instrument,
110+
mock_uninstrument,
111+
mock_complete,
112+
azure_ai_inference_service,
113+
chat_history: ChatHistory,
114+
mock_azure_ai_inference_chat_completion_response,
115+
model_diagnostics_test_env,
116+
) -> None:
117+
"""Test completion of AzureAIInferenceChatCompletion"""
118+
settings = AzureAIInferenceChatPromptExecutionSettings()
119+
120+
mock_complete.return_value = mock_azure_ai_inference_chat_completion_response
121+
122+
await azure_ai_inference_service.get_chat_message_contents(chat_history=chat_history, settings=settings)
123+
124+
mock_instrument.assert_called_once_with(enable_content_recording=False)
125+
mock_uninstrument.assert_called_once()
126+
127+
128+
@pytest.mark.asyncio
129+
@pytest.mark.parametrize(
130+
"azure_ai_inference_service",
131+
[AzureAIInferenceChatCompletion.__name__],
132+
indirect=True,
133+
)
134+
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
135+
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.uninstrument")
136+
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.instrument")
137+
async def test_azure_ai_inference_streaming_chat_completion_instrumentation(
138+
mock_instrument,
139+
mock_uninstrument,
140+
mock_complete,
141+
azure_ai_inference_service,
142+
chat_history: ChatHistory,
143+
mock_azure_ai_inference_streaming_chat_completion_response,
144+
model_diagnostics_test_env,
145+
) -> None:
146+
"""Test completion of AzureAIInferenceChatCompletion"""
147+
settings = AzureAIInferenceChatPromptExecutionSettings()
148+
149+
mock_complete.return_value = mock_azure_ai_inference_streaming_chat_completion_response
150+
151+
async for _ in azure_ai_inference_service.get_streaming_chat_message_contents(
152+
chat_history=chat_history, settings=settings
153+
):
154+
pass
155+
156+
mock_instrument.assert_called_once_with(enable_content_recording=True)
157+
mock_uninstrument.assert_called_once()
158+
159+
160+
@pytest.mark.asyncio
161+
@pytest.mark.parametrize(
162+
"azure_ai_inference_service",
163+
[
164+
AzureAIInferenceChatCompletion.__name__,
165+
],
166+
indirect=True,
167+
)
168+
@pytest.mark.parametrize(
169+
"override_env_param_dict",
170+
[
171+
{
172+
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS": "False",
173+
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE": "False",
174+
},
175+
],
176+
indirect=True,
177+
)
178+
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
179+
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.uninstrument")
180+
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.instrument")
181+
async def test_azure_ai_inference_streaming_chat_completion_not_instrumentation(
182+
mock_instrument,
183+
mock_uninstrument,
184+
mock_complete,
185+
azure_ai_inference_service,
186+
chat_history: ChatHistory,
187+
mock_azure_ai_inference_streaming_chat_completion_response,
188+
model_diagnostics_test_env,
189+
) -> None:
190+
"""Test completion of AzureAIInferenceChatCompletion"""
191+
settings = AzureAIInferenceChatPromptExecutionSettings()
192+
193+
mock_complete.return_value = mock_azure_ai_inference_streaming_chat_completion_response
194+
195+
async for _ in azure_ai_inference_service.get_streaming_chat_message_contents(
196+
chat_history=chat_history, settings=settings
197+
):
198+
pass
199+
200+
mock_instrument.assert_not_called()
201+
mock_uninstrument.assert_not_called()
202+
203+
204+
@pytest.mark.asyncio
205+
@pytest.mark.parametrize(
206+
"azure_ai_inference_service",
207+
[
208+
AzureAIInferenceChatCompletion.__name__,
209+
],
210+
indirect=True,
211+
)
212+
@pytest.mark.parametrize(
213+
"override_env_param_dict",
214+
[
215+
{
216+
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS": "True",
217+
"SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE": "False",
218+
},
219+
],
220+
indirect=True,
221+
)
222+
@patch.object(ChatCompletionsClient, "complete", new_callable=AsyncMock)
223+
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.uninstrument")
224+
@patch("azure.ai.inference.tracing.AIInferenceInstrumentor.instrument")
225+
async def test_azure_ai_inference_streaming_chat_completion_instrumentation_without_sensitive(
226+
mock_instrument,
227+
mock_uninstrument,
228+
mock_complete,
229+
azure_ai_inference_service,
230+
chat_history: ChatHistory,
231+
mock_azure_ai_inference_streaming_chat_completion_response,
232+
model_diagnostics_test_env,
233+
) -> None:
234+
"""Test completion of AzureAIInferenceChatCompletion"""
235+
settings = AzureAIInferenceChatPromptExecutionSettings()
236+
237+
mock_complete.return_value = mock_azure_ai_inference_streaming_chat_completion_response
238+
239+
async for _ in azure_ai_inference_service.get_streaming_chat_message_contents(
240+
chat_history=chat_history, settings=settings
241+
):
242+
pass
243+
244+
mock_instrument.assert_called_once_with(enable_content_recording=False)
245+
mock_uninstrument.assert_called_once()

‎python/tests/unit/utils/model_diagnostics/test_decorated.py

-14
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,6 @@
33
import pytest
44

55
from semantic_kernel.connectors.ai.anthropic.services.anthropic_chat_completion import AnthropicChatCompletion
6-
from semantic_kernel.connectors.ai.azure_ai_inference.services.azure_ai_inference_chat_completion import (
7-
AzureAIInferenceChatCompletion,
8-
)
96
from semantic_kernel.connectors.ai.google.google_ai.services.google_ai_chat_completion import GoogleAIChatCompletion
107
from semantic_kernel.connectors.ai.google.google_ai.services.google_ai_text_completion import GoogleAITextCompletion
118
from semantic_kernel.connectors.ai.google.vertex_ai.services.vertex_ai_chat_completion import VertexAIChatCompletion
@@ -118,17 +115,6 @@
118115
"__model_diagnostics_streaming_text_completion__",
119116
id="GoogleAITextCompletion._inner_get_streaming_text_contents",
120117
),
121-
# AzureAIInferenceChatCompletion
122-
pytest.param(
123-
AzureAIInferenceChatCompletion._inner_get_chat_message_contents,
124-
"__model_diagnostics_chat_completion__",
125-
id="AzureAIInferenceChatCompletion._inner_get_chat_message_contents",
126-
),
127-
pytest.param(
128-
AzureAIInferenceChatCompletion._inner_get_streaming_chat_message_contents,
129-
"__model_diagnostics_streaming_chat_completion__",
130-
id="AzureAIInferenceChatCompletion._inner_get_streaming_chat_message_contents",
131-
),
132118
# AnthropicChatCompletion
133119
pytest.param(
134120
AnthropicChatCompletion._inner_get_chat_message_contents,

‎python/uv.lock

+927-869
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)
Please sign in to comment.