mirror of
https://github.com/hwchase17/langchain.git
synced 2026-06-09 18:50:33 +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:
@@ -89,7 +89,16 @@ def test__format_output() -> None:
|
||||
)
|
||||
expected = ChatResult(
|
||||
generations=[
|
||||
ChatGeneration(message=AIMessage("bar")),
|
||||
ChatGeneration(
|
||||
message=AIMessage(
|
||||
"bar",
|
||||
usage_metadata={
|
||||
"input_tokens": 2,
|
||||
"output_tokens": 1,
|
||||
"total_tokens": 3,
|
||||
},
|
||||
)
|
||||
),
|
||||
],
|
||||
llm_output={
|
||||
"id": "foo",
|
||||
|
||||
Reference in New Issue
Block a user