From 8ed3898c7442611157f43392f80f8c586278cc5d Mon Sep 17 00:00:00 2001
From: Jack Collins <6640905+jackmpcollins@users.noreply.github.com>
Date: Fri, 29 Nov 2024 14:44:55 -0800
Subject: [PATCH] Remove assumption of one usage per stream

---
 src/magentic/chat_model/openai_chat_model.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/magentic/chat_model/openai_chat_model.py b/src/magentic/chat_model/openai_chat_model.py
index 9461c6f..2ece0b3 100644
--- a/src/magentic/chat_model/openai_chat_model.py
+++ b/src/magentic/chat_model/openai_chat_model.py
@@ -293,7 +293,10 @@ def update(self, item: ChatCompletionChunk) -> None:
                 tool_call_chunk.index = self._current_tool_call_index
         self._chat_completion_stream_state.handle_chunk(item)
         if item.usage:
-            assert not self.usage_ref  # noqa: S101
+            # Only keep the last usage
+            # Gemini openai-compatible API includes usage in all streamed chunks
+            # but OpenAI only includes this in the last chunk
+            self.usage_ref.clear()
             self.usage_ref.append(
                 Usage(
                     input_tokens=item.usage.prompt_tokens,