mistral[patch]: add usage_metadata to (a)invoke and (a)stream (#22781)

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

@@ -1,11 +1,12 @@
 """Test ChatMistral chat model."""
 
 import json
-from typing import Any
+from typing import Any, Optional
 
 from langchain_core.messages import (
     AIMessage,
     AIMessageChunk,
+    BaseMessageChunk,
     HumanMessage,
 )
 from langchain_core.pydantic_v1 import BaseModel
@@ -25,8 +26,28 @@ async def test_astream() -> None:
     """Test streaming tokens from ChatMistralAI."""
     llm = ChatMistralAI()
 
+    full: Optional[BaseMessageChunk] = None
+    chunks_with_token_counts = 0
     async for token in llm.astream("I'm Pickle Rick"):
+        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"]
+    )
 
 
 async def test_abatch() -> None:

libs/partners/mistralai/tests/integration_tests/test_standard.py

@@ -20,14 +20,3 @@ class TestMistralStandard(ChatModelIntegrationTests):
             "model": "mistral-large-latest",
             "temperature": 0,
         }
-
-    @pytest.mark.xfail(reason="Not implemented.")
-    def test_usage_metadata(
-        self,
-        chat_model_class: Type[BaseChatModel],
-        chat_model_params: dict,
-    ) -> None:
-        super().test_usage_metadata(
-            chat_model_class,
-            chat_model_params,
-        )