magentic
magentic copied to clipboard
Return usage stats on AssistantMessage
Resolves #74
- Create new
Usage
class - Add this as an optional attribute of
AssistantMessage
- Wrap openai response stream to set usage when encountered
- Add usage to Anthropic messages directly from sync response
Example of non-streamed response with usage immediately available
from magentic import OpenaiChatModel, UserMessage
chat_model = OpenaiChatModel("gpt-3.5-turbo", seed=42)
message = chat_model.complete(messages=[UserMessage("Say hello!")])
print(message.usage)
# > Usage(input_tokens=10, output_tokens=9)
Example of streamed response where usage only becomes available after the stream has been processed
from magentic import OpenaiChatModel, UserMessage
from magentic.streaming import StreamedStr
chat_model = OpenaiChatModel("gpt-3.5-turbo", seed=42)
message = chat_model.complete(messages=[UserMessage("Say hello!")], output_types=[StreamedStr])
print(message.usage)
# > `None` because stream has not be processed yet
# Process the stream (convert StreamedStr to str)
str(message.content)
print(message.usage)
# > Usage(input_tokens=10, output_tokens=9)