Files
567-labs--instructor/tests/test_update_total_usage.py
wehub-resource-sync 97e91a83f3
Ruff / Ruff (push) Waiting to run
Test / Core Tests (push) Waiting to run
Test / Offline Coverage Tests (Python 3.10) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.11) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.12) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.13) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.9) (push) Waiting to run
Test / Full Coverage (Python 3.11) (push) Waiting to run
Test / Core Provider Tests (OpenAI) (push) Blocked by required conditions
Test / Core Provider Tests (Anthropic) (push) Blocked by required conditions
Test / Core Provider Tests (Google) (push) Blocked by required conditions
Test / Core Provider Tests (Other) (push) Blocked by required conditions
Test / Anthropic Tests (push) Blocked by required conditions
Test / Gemini Tests (push) Blocked by required conditions
Test / Google GenAI Tests (push) Blocked by required conditions
Test / Vertex AI Tests (push) Blocked by required conditions
Test / OpenAI Tests (push) Blocked by required conditions
Test / Writer Tests (push) Blocked by required conditions
Test / Auto Client Tests (push) Blocked by required conditions
ty / type-check (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:36:38 +08:00

44 lines
1.4 KiB
Python

from openai.types.completion_usage import CompletionUsage
from openai.types.chat.chat_completion import ChatCompletion, Choice
from openai.types.chat.chat_completion_message import ChatCompletionMessage
from instructor.utils.core import update_total_usage
class ProviderSpecificUsage(CompletionUsage):
def get(self, key: str, default=None):
return getattr(self, key, default)
def make_response(usage: CompletionUsage) -> ChatCompletion:
return ChatCompletion(
id="test-id",
choices=[
Choice(
index=0,
finish_reason="stop",
logprobs=None,
message=ChatCompletionMessage(role="assistant", content="ok"),
)
],
created=123,
model="gpt-4o-mini",
object="chat.completion",
usage=usage,
)
def test_update_total_usage_preserves_openai_usage_subclass() -> None:
response = make_response(
ProviderSpecificUsage(prompt_tokens=11, completion_tokens=7, total_tokens=18)
)
total_usage = CompletionUsage(prompt_tokens=5, completion_tokens=3, total_tokens=8)
updated = update_total_usage(response, total_usage)
assert updated is not None
assert isinstance(updated.usage, ProviderSpecificUsage)
assert updated.usage.get("prompt_tokens") == 16
assert updated.usage.get("completion_tokens") == 10
assert updated.usage.total_tokens == 26