fix(ext): normalize Semantic Kernel token usage

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:

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