Retrieving token usage in Azure OpenAI response when streaming is enabled

Thanks for the question, Here is the sample for token count for stream enabled. Jupyter notebooks to calculate tokens usage with Tiktoken for scenarios with and without Token Streaming. https://github.com/LazaUK/AOAI-Streaming-TokenUsage/tree/main

Posting the answer here in case it helps others.

The stream_options: {"include_usage": True} option must be set using the model_extras keyword argument in the Azure client:

from azure.ai.inference import ChatCompletionsClient
from azure.core.credentials import AzureKeyCredential
from azure.ai.inference.models import SystemMessage, UserMessage


client = ChatCompletionsClient(
    endpoint=os.getenv("<YOUR ENDPOINT ENV VAR>"),
    credential=AzureKeyCredential("<YOUR AZURE KEY ENV VAR>")
)

for chunk in client.complete(
    messages=[
        SystemMessage(content="You are a helpful assistant."),
        UserMessage(content="I am going to Paris, what should I see?")
    ],
    stream=True,
    model_extras={"stream_options": {"include_usage": True}}
):
    if hasattr(chunk, "usage") and chunk.usage is not None:
        print(chunk.usage)

For the last chunk received, prints:

{'completion_tokens': 561, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens': 28, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}, 'total_tokens': 589}