Skip to content

Commit 5476ac3

Browse files
authored
Python: Update model diagnostics (microsoft#8346)
### 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. --> We have a module that performs tracing for model invocations. However, that module only works with OpenAI. We need it to work for all AI connectors. This PR does the refactoring necessary for the module to become more generic. ### Description <!-- Describe your changes, the overall approach, the underlying design. These notes will help understanding how your code works. Thanks! --> 1. Rename the tracing module to `model_diagnostics`. 2. Restructure and optimize code. 3. Add tracing to all AI connectors except the HuggingFace connector, because the HuggingFace connector will need some refactoring in the near future. ### Contribution Checklist <!-- 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 - [x] I didn't break anyone 😄
1 parent b58180a commit 5476ac3

31 files changed

Lines changed: 854 additions & 613 deletions

‎python/.cspell.json‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,8 @@
5555
"huggingface",
5656
"pytestmark",
5757
"contoso",
58-
"opentelemetry"
58+
"opentelemetry",
59+
"SEMANTICKERNEL",
60+
"OTEL"
5961
]
6062
}
Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1,3 @@
1-
TELEMETRY_SAMPLE_CONNECTION_STRING="..."
1+
TELEMETRY_SAMPLE_CONNECTION_STRING="..."
2+
SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS=true
3+
SEMANTICKERNEL_EXPERIMENTAL_GENAI_ENABLE_OTEL_DIAGNOSTICS_SENSITIVE=true

‎python/semantic_kernel/connectors/ai/anthropic/services/anthropic_chat_completion.py‎

Lines changed: 44 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,14 @@
11
# Copyright (c) Microsoft. All rights reserved.
22

33
import logging
4+
import sys
45
from collections.abc import AsyncGenerator
5-
from typing import Any
6+
from typing import Any, ClassVar
7+
8+
if sys.version_info >= (3, 12):
9+
from typing import override # pragma: no cover
10+
else:
11+
from typing_extensions import override # pragma: no cover
612

713
from anthropic import AsyncAnthropic
814
from anthropic.types import (
@@ -29,11 +35,9 @@
2935
from semantic_kernel.contents.text_content import TextContent
3036
from semantic_kernel.contents.utils.author_role import AuthorRole
3137
from semantic_kernel.contents.utils.finish_reason import FinishReason as SemanticKernelFinishReason
32-
from semantic_kernel.exceptions.service_exceptions import (
33-
ServiceInitializationError,
34-
ServiceResponseException,
35-
)
38+
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceResponseException
3639
from semantic_kernel.utils.experimental_decorator import experimental_class
40+
from semantic_kernel.utils.telemetry.model_diagnostics.decorators import trace_chat_completion
3741

3842
# map finish reasons from Anthropic to Semantic Kernel
3943
ANTHROPIC_TO_SEMANTIC_KERNEL_FINISH_REASON_MAP = {
@@ -49,8 +53,10 @@
4953
class AnthropicChatCompletion(ChatCompletionClientBase):
5054
"""Antropic ChatCompletion class."""
5155

56+
MODEL_PROVIDER_NAME: ClassVar[str] = "anthropic"
57+
5258
async_client: AsyncAnthropic
53-
59+
5460
def __init__(
5561
self,
5662
ai_model_id: str | None = None,
@@ -68,10 +74,10 @@ def __init__(
6874
service_id: Service ID tied to the execution settings.
6975
api_key: The optional API key to use. If provided will override,
7076
the env vars or .env file value.
71-
async_client: An existing client to use.
77+
async_client: An existing client to use.
7278
env_file_path: Use the environment settings file as a fallback
73-
to environment variables.
74-
env_file_encoding: The encoding of the environment settings file.
79+
to environment variables.
80+
env_file_encoding: The encoding of the environment settings file.
7581
"""
7682
try:
7783
anthropic_settings = AnthropicSettings.create(
@@ -82,7 +88,7 @@ def __init__(
8288
)
8389
except ValidationError as ex:
8490
raise ServiceInitializationError("Failed to create Anthropic settings.", ex) from ex
85-
91+
8692
if not anthropic_settings.chat_model_id:
8793
raise ServiceInitializationError("The Anthropic chat model ID is required.")
8894

@@ -97,12 +103,14 @@ def __init__(
97103
ai_model_id=anthropic_settings.chat_model_id,
98104
)
99105

106+
@override
107+
@trace_chat_completion(MODEL_PROVIDER_NAME)
100108
async def get_chat_message_contents(
101109
self,
102110
chat_history: "ChatHistory",
103111
settings: "PromptExecutionSettings",
104112
**kwargs: Any,
105-
) -> list["ChatMessageContent"]:
113+
) -> list["ChatMessageContent"]:
106114
"""Executes a chat completion request and returns the result.
107115
108116
Args:
@@ -127,22 +135,23 @@ async def get_chat_message_contents(
127135
raise ServiceResponseException(
128136
f"{type(self)} service failed to complete the prompt",
129137
ex,
130-
) from ex
131-
138+
) from ex
139+
132140
metadata: dict[str, Any] = {"id": response.id}
133141
# Check if usage exists and has a value, then add it to the metadata
134142
if hasattr(response, "usage") and response.usage is not None:
135143
metadata["usage"] = response.usage
136144

137-
return [self._create_chat_message_content(response, content_block, metadata)
138-
for content_block in response.content]
139-
145+
return [
146+
self._create_chat_message_content(response, content_block, metadata) for content_block in response.content
147+
]
148+
140149
async def get_streaming_chat_message_contents(
141150
self,
142151
chat_history: ChatHistory,
143-
settings: PromptExecutionSettings,
152+
settings: PromptExecutionSettings,
144153
**kwargs: Any,
145-
) -> AsyncGenerator[list[StreamingChatMessageContent], Any]:
154+
) -> AsyncGenerator[list[StreamingChatMessageContent], Any]:
146155
"""Executes a streaming chat completion request and returns the result.
147156
148157
Args:
@@ -166,17 +175,18 @@ async def get_streaming_chat_message_contents(
166175
author_role = None
167176
metadata: dict[str, Any] = {"usage": {}, "id": None}
168177
content_block_idx = 0
169-
178+
170179
async for stream_event in stream:
171180
if isinstance(stream_event, RawMessageStartEvent):
172181
author_role = stream_event.message.role
173182
metadata["usage"]["input_tokens"] = stream_event.message.usage.input_tokens
174183
metadata["id"] = stream_event.message.id
175184
elif isinstance(stream_event, (RawContentBlockDeltaEvent, RawMessageDeltaEvent)):
176-
yield [self._create_streaming_chat_message_content(stream_event,
177-
content_block_idx,
178-
author_role,
179-
metadata)]
185+
yield [
186+
self._create_streaming_chat_message_content(
187+
stream_event, content_block_idx, author_role, metadata
188+
)
189+
]
180190
elif isinstance(stream_event, ContentBlockStopEvent):
181191
content_block_idx += 1
182192

@@ -187,21 +197,18 @@ async def get_streaming_chat_message_contents(
187197
) from ex
188198

189199
def _create_chat_message_content(
190-
self,
191-
response: Message,
192-
content: TextBlock,
193-
response_metadata: dict[str, Any]
200+
self, response: Message, content: TextBlock, response_metadata: dict[str, Any]
194201
) -> "ChatMessageContent":
195202
"""Create a chat message content object."""
196203
items: list[ITEM_TYPES] = []
197-
204+
198205
if content.text:
199206
items.append(TextContent(text=content.text))
200207

201208
finish_reason = None
202209
if response.stop_reason:
203210
finish_reason = ANTHROPIC_TO_SEMANTIC_KERNEL_FINISH_REASON_MAP[response.stop_reason]
204-
211+
205212
return ChatMessageContent(
206213
inner_content=response,
207214
ai_model_id=self.ai_model_id,
@@ -212,20 +219,20 @@ def _create_chat_message_content(
212219
)
213220

214221
def _create_streaming_chat_message_content(
215-
self,
216-
stream_event: RawContentBlockDeltaEvent | RawMessageDeltaEvent,
217-
content_block_idx: int,
218-
role: str | None = None,
219-
metadata: dict[str, Any] = {}
222+
self,
223+
stream_event: RawContentBlockDeltaEvent | RawMessageDeltaEvent,
224+
content_block_idx: int,
225+
role: str | None = None,
226+
metadata: dict[str, Any] = {},
220227
) -> StreamingChatMessageContent:
221228
"""Create a streaming chat message content object from a choice."""
222229
text_content = ""
223-
230+
224231
if stream_event.delta and hasattr(stream_event.delta, "text"):
225232
text_content = stream_event.delta.text
226-
233+
227234
items: list[STREAMING_ITEM_TYPES] = [StreamingTextContent(choice_index=content_block_idx, text=text_content)]
228-
235+
229236
finish_reason = None
230237
if isinstance(stream_event, RawMessageDeltaEvent):
231238
if stream_event.delta.stop_reason:
@@ -246,4 +253,3 @@ def _create_streaming_chat_message_content(
246253
def get_prompt_execution_settings_class(self) -> "type[AnthropicChatPromptExecutionSettings]":
247254
"""Create a request settings object."""
248255
return AnthropicChatPromptExecutionSettings
249-

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import asyncio
44
import contextlib
55
from abc import ABC
6+
from typing import ClassVar
67

78
from azure.ai.inference.aio import ChatCompletionsClient, EmbeddingsClient
89

@@ -14,6 +15,8 @@
1415
class AzureAIInferenceBase(KernelBaseModel, ABC):
1516
"""Azure AI Inference Chat Completion Service."""
1617

18+
MODEL_PROVIDER_NAME: ClassVar[str] = "azureai"
19+
1720
client: ChatCompletionsClient | EmbeddingsClient
1821

1922
def __del__(self) -> None:

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from functools import reduce
88
from typing import TYPE_CHECKING, Any
99

10+
from semantic_kernel.utils.telemetry.model_diagnostics.decorators import trace_chat_completion
1011
from semantic_kernel.utils.telemetry.user_agent import SEMANTIC_KERNEL_USER_AGENT
1112

1213
if sys.version_info >= (3, 12):
@@ -119,6 +120,8 @@ def __init__(
119120
)
120121

121122
# region Non-streaming
123+
@override
124+
@trace_chat_completion(AzureAIInferenceBase.MODEL_PROVIDER_NAME)
122125
async def get_chat_message_contents(
123126
self,
124127
chat_history: ChatHistory,

‎python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_base.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# Copyright (c) Microsoft. All rights reserved.
22

33
from abc import ABC
4+
from typing import ClassVar
45

56
from semantic_kernel.connectors.ai.google.google_ai.google_ai_settings import GoogleAISettings
67
from semantic_kernel.kernel_pydantic import KernelBaseModel
@@ -9,4 +10,6 @@
910
class GoogleAIBase(KernelBaseModel, ABC):
1011
"""Google AI Service."""
1112

13+
MODEL_PROVIDER_NAME: ClassVar[str] = "googleai"
14+
1215
service_settings: GoogleAISettings

‎python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_chat_completion.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
from semantic_kernel.contents.utils.finish_reason import FinishReason
4141
from semantic_kernel.functions.kernel_arguments import KernelArguments
4242
from semantic_kernel.kernel import Kernel
43+
from semantic_kernel.utils.telemetry.model_diagnostics.decorators import trace_chat_completion
4344

4445
if sys.version_info >= (3, 12):
4546
from typing import override # pragma: no cover
@@ -109,6 +110,7 @@ def __init__(
109110

110111
# region Non-streaming
111112
@override
113+
@trace_chat_completion(GoogleAIBase.MODEL_PROVIDER_NAME)
112114
async def get_chat_message_contents(
113115
self,
114116
chat_history: ChatHistory,

‎python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_text_completion.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
)
1616
from semantic_kernel.connectors.ai.google.google_ai.services.google_ai_base import GoogleAIBase
1717
from semantic_kernel.connectors.ai.text_completion_client_base import TextCompletionClientBase
18+
from semantic_kernel.utils.telemetry.model_diagnostics.decorators import trace_text_completion
1819

1920
if sys.version_info >= (3, 12):
2021
from typing import override # pragma: no cover
@@ -78,6 +79,7 @@ def __init__(
7879

7980
# region Non-streaming
8081
@override
82+
@trace_text_completion(GoogleAIBase.MODEL_PROVIDER_NAME)
8183
async def get_text_contents(
8284
self,
8385
prompt: str,

‎python/semantic_kernel/connectors/ai/google/vertex_ai/services/vertex_ai_base.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# Copyright (c) Microsoft. All rights reserved.
22

33
from abc import ABC
4+
from typing import ClassVar
45

56
from semantic_kernel.connectors.ai.google.vertex_ai.vertex_ai_settings import VertexAISettings
67
from semantic_kernel.kernel_pydantic import KernelBaseModel
@@ -9,4 +10,6 @@
910
class VertexAIBase(KernelBaseModel, ABC):
1011
"""Vertex AI Service."""
1112

13+
MODEL_PROVIDER_NAME: ClassVar[str] = "vertexai"
14+
1215
service_settings: VertexAISettings

‎python/semantic_kernel/connectors/ai/google/vertex_ai/services/vertex_ai_chat_completion.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
)
4646
from semantic_kernel.functions.kernel_arguments import KernelArguments
4747
from semantic_kernel.kernel import Kernel
48+
from semantic_kernel.utils.telemetry.model_diagnostics.decorators import trace_chat_completion
4849

4950
if sys.version_info >= (3, 12):
5051
from typing import override # pragma: no cover
@@ -103,6 +104,7 @@ def __init__(
103104

104105
# region Non-streaming
105106
@override
107+
@trace_chat_completion(VertexAIBase.MODEL_PROVIDER_NAME)
106108
async def get_chat_message_contents(
107109
self,
108110
chat_history: ChatHistory,

0 commit comments

Comments
 (0)