Add token usage to streamed chat

I have created a MR to start implementing this functionality.

MR currently contains logic for the first action point:

Add so that you can collect tokens in the StreamedChatMessage, but do not add it to the interface (no breaking changes).

This allows this to be used in the ai_provider_openai module for example:

class OpenAiChatMessageIterator extends StreamedChatMessageIterator {

  /**
   * {@inheritdoc}
   */
  public function getIterator(): \Generator {
    foreach ($this->iterator->getIterator() as $data) {
      $metadata = $data->usage ? $data->usage->toArray() : [];
      $streamedChatMessage = new StreamedChatMessage(
        $data->choices[0]->delta->role ?? '',
        $data->choices[0]->delta->content ?? '',
        $metadata
      );
      if (!empty($metadata)) {
        $streamedChatMessage->setInputTokenUsage($metadata['prompt_tokens']);
        $streamedChatMessage->setOutputTokenUsage($metadata['completion_tokens']);
        $streamedChatMessage->setTotalTokenUsage($metadata['total_tokens']);
        $streamedChatMessage->setCachedTokenUsage($metadata['prompt_tokens_details']['cached_tokens']);
        $streamedChatMessage->setReasoningTokenUsage($metadata['completion_tokens_details']['reasoning_tokens']);
      }
      yield $streamedChatMessage;
    }
  }

}

I am a bit stuck on the second action point:

Add so we can collect it in the PostStreamingResponseEvent

I've been looking through the code in the AI module, but can't find where the PostStreamingResponseEvent gets triggered.
Well, it gets triggered in \Drupal\ai\OperationType\Chat\StreamedChatMessageIterator::triggerEvent(), but I cannot find where that method is called. Is that supposed to be called by the OpenAI provider for example (because it doesn't currently)?
I'm a bit confused as to what the exact code flow should look like...

As per the proposed resolution, i have:
- Added token collection methods to PostStreamingResponseEvent
- Added token injection into PostStreamingResponseEvent using setter methods (no constructor changes)

Kindly review and let me know if anything else needs to be addressed or adjusted here.

Thanks!

Thanks for the added logic sarvjeetsingh.
I am unable to make this work on my local installation though, as mentioned above.

Test scenario:
- Enable AI CKEditor integration module
- Add AI button to CKEditor
- Enable "Fix spelling" for example
- Select a text in CKEditor that needs spelling fixes.
- The action gets performed, but my debugger never enters the method `\Drupal\ai\OperationType\Chat\StreamedChatMessageIterator::triggerEvent`

Can also be tested by installing dev version of https://www.drupal.org/project/ai_usage_limits → module, and configuring it with token usage limits. The usage stats will not go up in the test scenario mentioned above, even though they should...

Seems like the failure of PostStreamingResponseEvent not being triggered, is a separate issue.
See 🐛 PostStreamingResponseEvent never gets triggered Active .

The code in this MR can probably be already reviewed, but will not be functional until this other issue is fixed first.

I have not tested the functionality, but I have resolved all PHPCS issues.