groq[patch]: add usage_metadata to (a)invoke and (a)stream (#22834)

libs/partners/groq/tests/integration_tests/test_chat_models.py

@@ -1,7 +1,7 @@
 """Test ChatGroq chat model."""
 
 import json
-from typing import Any
+from typing import Any, Optional
 
 import pytest
 from langchain_core.messages import (
@@ -93,9 +93,28 @@ async def test_astream() -> None:
     """Test streaming tokens from Groq."""
     chat = ChatGroq(max_tokens=10)
 
+    full: Optional[BaseMessageChunk] = None
+    chunks_with_token_counts = 0
     async for token in chat.astream("Welcome to the Groqetship!"):
-        assert isinstance(token, BaseMessageChunk)
+        assert isinstance(token, AIMessageChunk)
         assert isinstance(token.content, str)
+        full = token if full is None else full + token
+        if token.usage_metadata is not None:
+            chunks_with_token_counts += 1
+    if chunks_with_token_counts != 1:
+        raise AssertionError(
+            "Expected exactly one chunk with token counts. "
+            "AIMessageChunk aggregation adds counts. Check that "
+            "this is behaving properly."
+        )
+    assert isinstance(full, AIMessageChunk)
+    assert full.usage_metadata is not None
+    assert full.usage_metadata["input_tokens"] > 0
+    assert full.usage_metadata["output_tokens"] > 0
+    assert (
+        full.usage_metadata["input_tokens"] + full.usage_metadata["output_tokens"]
+        == full.usage_metadata["total_tokens"]
+    )
 
 
 #