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:
ccurme
2024-05-23 14:21:58 -04:00
committed by GitHub
parent 3d26807b92
commit fbfed65fb1
13 changed files with 364 additions and 3 deletions

View File

@@ -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",