From e587009b247823c33ef76a7255d7d965e3da8bbb Mon Sep 17 00:00:00 2001 From: Thota shashank Date: Wed, 7 Oct 2026 19:07:33 +0530 Subject: [PATCH 1/2] fix(ext): normalize Semantic Kernel token usage --- .../_sk_chat_completion_adapter.py | 28 ++++-- .../models/test_sk_chat_completion_adapter.py | 89 ++++++++++++++++++- 2 files changed, 109 insertions(+), 8 deletions(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/semantic_kernel/_sk_chat_completion_adapter.py b/python/packages/autogen-ext/src/autogen_ext/models/semantic_kernel/_sk_chat_completion_adapter.py index b9267057cd49..963b9e724c45 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/semantic_kernel/_sk_chat_completion_adapter.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/semantic_kernel/_sk_chat_completion_adapter.py @@ -1,7 +1,7 @@ import json import logging import warnings -from typing import Any, Literal, Mapping, Optional, Sequence, Union +from typing import Any, Literal, Mapping, Optional, Sequence, Union, cast from autogen_core import EVENT_LOGGER_NAME, FunctionCall from autogen_core._cancellation_token import CancellationToken @@ -38,6 +38,20 @@ logger = logging.getLogger(EVENT_LOGGER_NAME) +def _get_token_usage(usage: Any) -> RequestUsage: + """Normalize OpenAI and Anthropic usage objects or dictionaries from Semantic Kernel.""" + if isinstance(usage, Mapping): + usage = cast(Mapping[str, Any], usage) + return RequestUsage( + prompt_tokens=usage.get("prompt_tokens", usage.get("input_tokens", 0)), + completion_tokens=usage.get("completion_tokens", usage.get("output_tokens", 0)), + ) + return RequestUsage( + prompt_tokens=getattr(usage, "prompt_tokens", getattr(usage, "input_tokens", 0)), + completion_tokens=getattr(usage, "completion_tokens", getattr(usage, "output_tokens", 0)), + ) + + def ensure_serializable(data: BaseModel) -> BaseModel: """ Workaround for https://github.com/pydantic/pydantic/issues/7713, see https://github.com/pydantic/pydantic/issues/7713#issuecomment-2604574418 @@ -496,9 +510,9 @@ async def create( completion_tokens = 0 if result[0].metadata and "usage" in result[0].metadata: - usage = result[0].metadata["usage"] - prompt_tokens = getattr(usage, "prompt_tokens", 0) - completion_tokens = getattr(usage, "completion_tokens", 0) + usage = _get_token_usage(result[0].metadata["usage"]) + prompt_tokens = usage.prompt_tokens + completion_tokens = usage.completion_tokens logger.info( LLMCallEvent( @@ -635,9 +649,9 @@ async def create_stream( for msg in streaming_messages: # Track token usage if msg.metadata and "usage" in msg.metadata: - usage = msg.metadata["usage"] - prompt_tokens = getattr(usage, "prompt_tokens", 0) - completion_tokens = getattr(usage, "completion_tokens", 0) + usage = _get_token_usage(msg.metadata["usage"]) + prompt_tokens = usage.prompt_tokens + completion_tokens = usage.completion_tokens # Process function call deltas for item in msg.items: diff --git a/python/packages/autogen-ext/tests/models/test_sk_chat_completion_adapter.py b/python/packages/autogen-ext/tests/models/test_sk_chat_completion_adapter.py index 300ae0982904..36e5c6f3b37e 100644 --- a/python/packages/autogen-ext/tests/models/test_sk_chat_completion_adapter.py +++ b/python/packages/autogen-ext/tests/models/test_sk_chat_completion_adapter.py @@ -1,10 +1,13 @@ import logging import os +from types import SimpleNamespace from typing import Any, AsyncGenerator from unittest.mock import AsyncMock import pytest -from autogen_core import CancellationToken, FunctionCall +from anthropic.types import Usage +from autogen_core import EVENT_LOGGER_NAME, CancellationToken, FunctionCall +from autogen_core.logging import LLMCallEvent, LLMStreamEndEvent from autogen_core.models import ( AssistantMessage, CreateResult, @@ -13,6 +16,7 @@ LLMMessage, ModelFamily, ModelInfo, + RequestUsage, SystemMessage, UserMessage, ) @@ -41,6 +45,89 @@ from semantic_kernel.memory.null_memory import NullMemory +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True], ids=["create", "create_stream"]) +@pytest.mark.parametrize( + "metadata, expected_usage", + [ + pytest.param( + {"usage": Usage(input_tokens=147, output_tokens=1024)}, + RequestUsage(prompt_tokens=147, completion_tokens=1024), + id="anthropic", + ), + pytest.param( + {"usage": CompletionUsage(prompt_tokens=20, completion_tokens=9, total_tokens=29)}, + RequestUsage(prompt_tokens=20, completion_tokens=9), + id="openai", + ), + pytest.param( + {"usage": {"input_tokens": 147, "output_tokens": 1024}}, + RequestUsage(prompt_tokens=147, completion_tokens=1024), + id="anthropic-dict", + ), + pytest.param( + {"usage": {"prompt_tokens": 20, "completion_tokens": 9}}, + RequestUsage(prompt_tokens=20, completion_tokens=9), + id="openai-dict", + ), + pytest.param( + {"usage": SimpleNamespace(prompt_tokens=0, completion_tokens=0, input_tokens=147, output_tokens=1024)}, + RequestUsage(prompt_tokens=0, completion_tokens=0), + id="openai-fields-take-precedence", + ), + pytest.param( + {"usage": {"prompt_tokens": 0, "completion_tokens": 0, "input_tokens": 147, "output_tokens": 1024}}, + RequestUsage(prompt_tokens=0, completion_tokens=0), + id="openai-dict-fields-take-precedence", + ), + pytest.param({}, RequestUsage(prompt_tokens=0, completion_tokens=0), id="missing-usage"), + pytest.param({"usage": None}, RequestUsage(prompt_tokens=0, completion_tokens=0), id="null-usage"), + ], +) +async def test_sk_chat_completion_token_usage( + stream: bool, + metadata: dict[str, Any], + expected_usage: RequestUsage, + caplog: pytest.LogCaptureFixture, +) -> None: + mock_client = AsyncMock(spec=AzureChatCompletion) + mock_client.get_chat_message_contents.return_value = [ + ChatMessageContent(role=AuthorRole.ASSISTANT, content="Hello!", metadata=metadata) + ] + + async def mock_stream(*args: Any, **kwargs: Any) -> AsyncGenerator[list[StreamingChatMessageContent], None]: + yield [StreamingChatMessageContent(choice_index=0, role=AuthorRole.ASSISTANT, content="Hello!")] + yield [StreamingChatMessageContent(choice_index=0, role=AuthorRole.ASSISTANT, items=[], metadata=metadata)] + + mock_client.get_streaming_chat_message_contents = mock_stream + adapter = SKChatCompletionAdapter(mock_client, kernel=Kernel(memory=NullMemory())) + messages = [UserMessage(content="Say hello!", source="user")] + + with caplog.at_level(logging.INFO, logger=EVENT_LOGGER_NAME): + for call_count in (1, 2): + if stream: + chunks = [chunk async for chunk in adapter.create_stream(messages)] + assert chunks[:-1] == ["Hello!"] + result = chunks[-1] + assert isinstance(result, CreateResult) + else: + result = await adapter.create(messages) + + assert result.content == "Hello!" + assert result.usage == expected_usage + total_usage = RequestUsage( + prompt_tokens=expected_usage.prompt_tokens * call_count, + completion_tokens=expected_usage.completion_tokens * call_count, + ) + assert adapter.actual_usage() == total_usage + assert adapter.total_usage() == total_usage + + events = [record.msg for record in caplog.records if isinstance(record.msg, (LLMCallEvent, LLMStreamEndEvent))] + assert len(events) == 2 + assert all(event.prompt_tokens == expected_usage.prompt_tokens for event in events) + assert all(event.completion_tokens == expected_usage.completion_tokens for event in events) + + class CalculatorArgs(BaseModel): a: float b: float From 321a096c09c76d1e0c19d7312275469d4076a139 Mon Sep 17 00:00:00 2001 From: Thota shashank Date: Sat, 10 Oct 2026 17:58:54 +0530 Subject: [PATCH 2/2] test(ext): cover shared streaming usage metadata --- .../models/test_sk_chat_completion_adapter.py | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/python/packages/autogen-ext/tests/models/test_sk_chat_completion_adapter.py b/python/packages/autogen-ext/tests/models/test_sk_chat_completion_adapter.py index 36e5c6f3b37e..ededc4b025a4 100644 --- a/python/packages/autogen-ext/tests/models/test_sk_chat_completion_adapter.py +++ b/python/packages/autogen-ext/tests/models/test_sk_chat_completion_adapter.py @@ -128,6 +128,47 @@ async def mock_stream(*args: Any, **kwargs: Any) -> AsyncGenerator[list[Streamin assert all(event.completion_tokens == expected_usage.completion_tokens for event in events) +@pytest.mark.asyncio +async def test_sk_chat_completion_stream_shared_usage_metadata(caplog: pytest.LogCaptureFixture) -> None: + mock_client = AsyncMock(spec=AzureChatCompletion) + + async def mock_stream(*args: Any, **kwargs: Any) -> AsyncGenerator[list[StreamingChatMessageContent], None]: + # Anthropic's SK connector reuses this dictionary and adds output usage at the end. + metadata = {"usage": {"input_tokens": 147}} + for text in ("Hello", "!"): + chunk = StreamingChatMessageContent(choice_index=0, role=AuthorRole.ASSISTANT, content=text) + chunk.metadata = metadata + yield [chunk] + + metadata["usage"]["output_tokens"] = 1024 + final_chunk = StreamingChatMessageContent( + choice_index=0, role=AuthorRole.ASSISTANT, items=[], finish_reason=FinishReason.STOP + ) + final_chunk.metadata = metadata + yield [final_chunk] + + mock_client.get_streaming_chat_message_contents = mock_stream + adapter = SKChatCompletionAdapter(mock_client, kernel=Kernel(memory=NullMemory())) + + with caplog.at_level(logging.INFO, logger=EVENT_LOGGER_NAME): + chunks = [chunk async for chunk in adapter.create_stream([UserMessage(content="Say hello!", source="user")])] + + assert chunks[:-1] == ["Hello", "!"] + result = chunks[-1] + assert isinstance(result, CreateResult) + assert result.content == "Hello!" + assert result.finish_reason == "stop" + expected_usage = RequestUsage(prompt_tokens=147, completion_tokens=1024) + assert result.usage == expected_usage + assert adapter.actual_usage() == expected_usage + assert adapter.total_usage() == expected_usage + + events = [record.msg for record in caplog.records if isinstance(record.msg, LLMStreamEndEvent)] + assert len(events) == 1 + assert events[0].prompt_tokens == 147 + assert events[0].completion_tokens == 1024 + + class CalculatorArgs(BaseModel): a: float b: float