Repository: microsoft/autogen
Pull request: Default usage statistics for streaming responses
## Why are these changes needed?

Enables usage statistics for streaming responses by default.

There is a similar bug in the AzureAI client. Theoretically adding the parameter
```
model_extras={"stream_options": {"include_usage": True}}
```
should fix the problem, but I'm currently unable to test that workflow

## Related issue number

closes https://github.com/microsoft/autogen/issues/6548

## Checks

- [ ] I've included any doc changes needed for <https://microsoft.github.io/autogen/>. See <https://github.com/microsoft/autogen/blob/main/CONTRIBUTING.md> to build and test documentation locally.
- [ ] I've added tests (if relevant) corresponding to the changes introduced in this PR.
- [ ] I've made sure all auto checks have passed.


Problem statement:
models_usage shows prompt_tokens and completion_tokens as 0
### What happened?

**Describe the bug**
I was previously using v0.5.1 and sometime over the last few weeks, I had upgraded to v0.5.6. I was checking the history data and noticed that the `prompt_tokens` and `completion_tokens` under `models_usage` are coming as 0. 
I remember they used to report token count but I can only see zero. 

**To Reproduce**
```
response = await agent.on_messages(
	messages=[request], 
	cancellation_token=CancellationToken()
)
print('>>>>', response.chat_message.json())

```
**Expected behavior**
Expecting token usage counts

**Screenshots**
I see output like: `{"source":"assistant","models_usage":{"prompt_tokens":0,"completion_tokens":0},"metadata":{},"content":"Hello! How can I assist you today?","type":"TextMessage"}`




### Which packages was the bug in?

Python Core (autogen-core)

### AutoGen library version.

Python 0.5.6

### Other library version.

_No response_

### Model used

gpt-4o-mini, gemini-2.0-flash

### Model provider

OpenAI

### Other model provider

Gemini

### Python version

3.12

### .NET version

None

### Operating system

None

Hints from issue discussion:
v0.5.7, cannot reproduce

```python
from autogen_agentchat.agents import AssistantAgent
from autogen_ext.models.openai import OpenAIChatCompletionClient

async def main():
    client = OpenAIChatCompletionClient(model="gemini-2.0-flash")
    assistant = AssistantAgent(
        "assistant",
        model_client=client,
        system_message="You are a helpful assistant.",
    )
    result = await assistant.run(task="What is the capital of France?")
    print(result)

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())
```

```
messages=[TextMessage(source='user', models_usage=None, metadata={}, content='What is the capital of France?', type='TextMessage'), TextMessage(source='assistant', models_usage=RequestUsage(prompt_tokens=13, completion_tokens=8), metadata={}, content='The capital of France is Paris.\n', type='TextMessage')] stop_reason=None
```
Hi, below is a sample code using Azure OpenAI but I noticed the same behaviour with Gemini as well. Please let me know what am I missing:

```
from autogen_agentchat.agents import AssistantAgent
from autogen_core.models import ChatCompletionClient, UserMessage
from autogen_agentchat.messages import TextMessage
# Print loaded versions
import importlib.metadata

def print_versions():
    packages = ["autogen-agentchat", "autogen-core"]
    for pkg in packages:
        try:
            version = importlib.metadata.version(pkg)
            print(f"{pkg} version: {version}")
        except importlib.metadata.PackageNotFoundError:
            print(f"{pkg} is not installed or version not found")

print_versions()
```
autogen-agentchat version: 0.5.7
autogen-core version: 0.5.7
```
model_config = {
    'provider': 'autogen_ext.models.openai.AzureOpenAIChatCompletionClient', 
    'config': 
        {
            'model': 'gpt-4o-mini', 
            'azure_endpoint': os.getenv('AZURE_ENDPOINT_URL'),
            'api_version': '2025-02-01-preview', 
            'max_tokens': 500,
            'temperature': 0.1
        }
}
model_client = ChatCompletionClient.load_component(model_config)

async def main():
    agent = AssistantAgent(
            name="assistant",
            model_client=model_client,
            model_client_stream=True,
            system_message = "You are a helpful assistant",
            reflect_on_tool_use=True
        )
    request = TextMessage(content="What is the capital of France?", source="user")
    stream = agent.run_stream(task=request)
    async for message in stream:
        if isinstance(message, TextMessage):
            print(message.model_dump())
            
await main()
```
{'source': 'user', 'models_usage': None, 'metadata': {}, 'content': 'What is the capital of France?', 'type': 'TextMessage'}
{'source': 'assistant', 'models_usage': {'prompt_tokens': 0, 'completion_tokens': 0}, 'metadata': {}, 'content': 'The capital of France is Paris.', 'type': 'TextMessage'}

The problem appears to be limited to streaming responses. Investigating
The problem here is a missing parameter. As a temporary workaround while I fix the issue, you can add the following parameter on your side:

OpenAI/Azure OpenAI:
```
model_client = OpenAIChatCompletionClient(
    ...,
    stream_options={"include_usage": True},
)
```

Commit message: Merge branch 'main' into default_stream_usage

Patch to review:
diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py
index 936b3a6c7a23..e6a800d6891a 100644
--- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py
+++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py
@@ -770,6 +770,12 @@ async def create_stream(
             extra_create_args,
         )
 
+        # Force usage statistics if the user hasn't specified any stream options.
+        if "stream_options" not in create_params.create_args:
+            create_params.create_args["stream_options"] = {
+                "include_usage": True,
+            }
+
         if max_consecutive_empty_chunk_tolerance != 0:
             warnings.warn(
                 "The 'max_consecutive_empty_chunk_tolerance' parameter is deprecated and will be removed in the future releases. All of empty chunks will be skipped with a warning.",
diff --git a/python/packages/autogen-ext/tests/models/test_openai_model_client.py b/python/packages/autogen-ext/tests/models/test_openai_model_client.py
index 1f824274e436..a23ddeeba17c 100644
--- a/python/packages/autogen-ext/tests/models/test_openai_model_client.py
+++ b/python/packages/autogen-ext/tests/models/test_openai_model_client.py
@@ -315,7 +315,7 @@ async def test_openai_chat_completion_client_create_stream_no_usage_default(monk
     assert chunks[2] == " Yet Another Hello"
     assert isinstance(chunks[-1], CreateResult)
     assert chunks[-1].content == "Hello Another Hello Yet Another Hello"
-    assert chunks[-1].usage == RequestUsage(prompt_tokens=0, completion_tokens=0)
+    assert chunks[-1].usage == RequestUsage(prompt_tokens=3, completion_tokens=3)
 
 
 @pytest.mark.asyncio


Code review comments:
diff_hunk: @@ -770,6 +770,12 @@ async def create_stream(
             extra_create_args,
         )
 
+        # Force usage statistics if the user hasn't specified any stream options.
+        if "stream_options" not in create_params.create_args:
+            create_params.create_args["stream_options"] = {
+                "include_usage": True,
path: python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py
text: @user2:
This may or may not break scenarios using this with "OpenAI-compatible" endpoints. @user1 what's your thought on this? 

@user1:
@user2
I think so... I am not sure about whole "OpenAI-compatible" endpoint exist "stream_options"

@author Could you diff with older version of AutoGen? #6548 says, older version could response of token-count, How could it do that?

@author:
> @author Could you diff with older version of AutoGen? #6548 says, older version could response of token-count, How could it do that?

I'm wondering if older version's streaming interface returned token counts? AFAIK this is a relatively new feature that was added to openai -- [this article](https://community.openai.com/t/usage-stats-now-available-when-using-streaming-with-the-chat-completions-api-or-completions-api/738156) puts it at May last year

I'll take a look at some of our early versions, but unless we accidentally deleted the `include_usage` flag, I don't think it should have worked.

This was the case as far back as 7 months ago, which can be seen here
https://github.com/microsoft/autogen/blob/2cc8c73d3b83911892903001d64425298cc94da4/python/packages/autogen-ext/tests/models/test_openai_model_client.py#L238

The top test explicitly includes the `include_usage` flag, which gives the token count. The test under it does not, which gives a 0 token count

@user1:
@author
I see. You are right.

Just #6548 said, 

> Describe the bug
> I was previously using v0.5.1 and sometime over the last few weeks, I had upgraded to v0.5.6. I was checking the history data and noticed that the prompt_tokens and completion_tokens under models_usage are coming as 0.
I remember they used to report token count but I can only see zero.

So, I think we need to check how to v0.5.1 could response prompt_tokens without this option.
Or when stream, response zero at v0.5.1 too? (Maybe that issue is wrong?)

And, I have worry about "OpenAI-compatible" endpoints could not compatible with this options. So they could response error instead of right response. Cause this patch is unconditionally adding that option -`"include_usage": True,`-

@author:
@user1 I see your point. Would you feel more comfortable if I flip the logic so the default is off, and a user can pass in a flag to enable usage? I can limit the change specifically to the stream interface

@user1:
@author
Right that is my opinion. 

And I could check v0.5.1 has that this error too.
> just `model_client_stream=True` was return token count 0, v0.5.1, v0.5.6 both.

diff_hunk: @@ -315,7 +315,7 @@ async def test_openai_chat_completion_client_create_stream_no_usage_default(monk
     assert chunks[2] == " Yet Another Hello"
     assert isinstance(chunks[-1], CreateResult)
     assert chunks[-1].content == "Hello Another Hello Yet Another Hello"
-    assert chunks[-1].usage == RequestUsage(prompt_tokens=0, completion_tokens=0)
+    assert chunks[-1].usage == RequestUsage(prompt_tokens=3, completion_tokens=3)
path: python/packages/autogen-ext/tests/models/test_openai_model_client.py
text: @user1:
This may change depending on the server's implementation of tokenization. 

@author:
the previous test uses these same values. I can check to see if they're not 0 if you're more comfortable with that