mirror of
https://github.com/hwchase17/langchain.git
synced 2026-06-09 10:17:00 +00:00
core, partners: add token usage attribute to AIMessage (#21944)
```python
class UsageMetadata(TypedDict):
"""Usage metadata for a message, such as token counts.
Attributes:
input_tokens: (int) count of input (or prompt) tokens
output_tokens: (int) count of output (or completion) tokens
total_tokens: (int) total token count
"""
input_tokens: int
output_tokens: int
total_tokens: int
```
```python
class AIMessage(BaseMessage):
...
usage_metadata: Optional[UsageMetadata] = None
"""If provided, token usage information associated with the message."""
...
```
This commit is contained in:
@@ -20,3 +20,14 @@ 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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user