fix: OpenAI prompt details and completion tokens details missing from total usage

[ivanbelenky]

Describe your changes

The change solve create_with_completion failing to output the correct token usage: total_usage: CompletionUsage now is initialized with both PromptTokenDetails and CompletionTokenDetails fields. I am aware that this implies overriding to some degree the optionality of this attributes, nevertheless I think 0 is a good alias for the information that None in this context conveys.

def initialize_usage(mode: Mode) -> CompletionUsage | Any:
    ...
    total_usage = CompletionUsage(completion_tokens=0, prompt_tokens=0, total_tokens=0,
        completion_tokens_details = CompletionTokensDetails(audio_tokens=0, reasoning_tokens=0),
        prompt_token_details = PromptTokensDetails(audio_tokens=0, cached_tokens=0)
    )
    ...

and this values are respectfully filled out whenever possible

def update_total_usage(...) -> ...:
    ...
    if isinstance(response_usage, OpenAIUsage) and isinstance(total_usage, OpenAIUsage):
        total_usage.completion_tokens += response_usage.completion_tokens or 0
        total_usage.prompt_tokens += response_usage.prompt_tokens or 0
        total_usage.total_tokens += response_usage.total_tokens or 0

        if (rtd := response_usage.completion_tokens_details) and (ttd := total_usage.completion_tokens_details):
            ttd.audio_tokens = (ttd.audio_tokens or 0) + (rtd.audio_tokens or 0)
            ttd.reasoning_tokens = (ttd.reasoning_tokens or 0) + (rtd.reasoning_tokens or 0)

        if (rpd := response_usage.prompt_tokens_details) and (tpd := total_usage.prompt_tokens_details):
            tpd.audio_tokens = (tpd.audio_tokens or 0) + (rpd.audio_tokens or 0)
            tpd.cached_tokens = (tpd.cached_tokens or 0) + (rpd.cached_tokens or 0)

        response.usage = total_usage  # Replace each response usage with the total usage

Related Issue

---

Important

Fixes token usage reporting by initializing CompletionUsage with token details and updating aggregation in retry.py and utils.py.

• Behavior:
  • Fixes incorrect token usage reporting by initializing CompletionUsage with CompletionTokensDetails and PromptTokensDetails in initialize_usage() in retry.py.
  • Updates update_total_usage() in utils.py to aggregate audio_tokens, reasoning_tokens, and cached_tokens.
• Imports:
  • Adds CompletionTokensDetails and PromptTokensDetails to imports in retry.py.

^{This description was created by for 60d549a. It will automatically update as commits are pushed.}

[ellipsis-dev]

👍 Looks good to me! Reviewed everything up to 60d549a in 18 seconds

More details

• Looked at 43 lines of code in 2 files
• Skipped 0 files when reviewing.
• Skipped posting 3 drafted comments based on config settings.

1. instructor/utils.py:143

• Draft comment:
  Ensure the use of the walrus operator := is compatible with the minimum Python version supported by the project. This operator is only available in Python 3.8 and later.
• Reason this comment was not posted:
  Confidence changes required: 50%
  The code in update_total_usage uses the walrus operator := which is only available in Python 3.8 and later. This should be noted if the codebase supports older Python versions.

2. instructor/retry.py:17

• Draft comment:
  The import of CompletionTokensDetails and PromptTokensDetails is necessary for the changes in initialize_usage. Ensure these classes are defined and available in the openai.types.completion_usage module.
• Reason this comment was not posted:
  Confidence changes required: 20%
  The import statement in instructor/retry.py for CompletionTokensDetails and PromptTokensDetails is correct and necessary for the changes made in the initialize_usage function.

3. instructor/utils.py:140

• Draft comment:
  Assertions should always have an error message that is formatted well. Please add error messages to the assertions in update_total_usage. This applies to other assertions in the function as well.
• Reason this comment was not posted:
  Confidence changes required: 80%
  The function update_total_usage in instructor/utils.py has multiple instances of assertions without error messages. This violates the rule that assertions should always have an error message that is formatted well.

Workflow ID: wflow_gfxY8ZwYYs4u5kQM

---

You can customize Ellipsis with 👍 / 👎 feedback, review rules, user-specific overrides, quiet mode, and more.

[ivanleomk]

Looks good to me! Testing it using O1 and the reasoning tokens are now updated well.

from openai import OpenAI
from pydantic import BaseModel, field_validator
import instructor
from rich import print
from instructor.mode import Mode

client = instructor.from_openai(OpenAI(), mode=Mode.JSON_O1)


class Person(BaseModel):
    name: str
    age: int


resp, usage = client.chat.completions.create_with_completion(
    model="o1-mini",
    response_model=Person,
    messages=[
        {
            "role": "user",
            "content": "What's a plausible name for a 20 year old person living in Mongolia?",
        },
    ],
)

print(resp)
# Person(name='Bat-Erdene', age=20
print(usage.usage)
# CompletionUsage(
#     completion_tokens=546,
#     prompt_tokens=153,
#     total_tokens=699,
#     completion_tokens_details=CompletionTokensDetails(
#         accepted_prediction_tokens=None,
#         audio_tokens=0,
#         reasoning_tokens=512,
#         rejected_prediction_tokens=None,
#     ),
#     prompt_tokens_details=None,
#     prompt_token_details=PromptTokensDetails(audio_tokens=0, cached_tokens=0),
# )