magentic icon indicating copy to clipboard operation
magentic copied to clipboard

Return usage stats on AssistantMessage

Open jackmpcollins opened this issue 1 month ago • 1 comments

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)

jackmpcollins avatar May 16 '24 07:05 jackmpcollins