76 lines
3.0 KiB
Python
76 lines
3.0 KiB
Python
|
"""Tests for the token-usage HTTP response hook and attach_usage_hook wiring."""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import json
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
|
||
|
|
from rag_eval.metrics.factory import _usage_response_hook, attach_usage_hook
|
||
|
|
from rag_eval.metrics.token_tracker import track_token_usage
|
||
|
|
|
||
|
|
|
||
|
|
def _fake_response(payload: dict | None) -> httpx.Response:
|
||
|
|
"""Build a real httpx.Response with a JSON (or broken) body for hook testing."""
|
||
|
|
content = b"not json" if payload is None else json.dumps(payload).encode("utf-8")
|
||
|
|
return httpx.Response(200, content=content, request=httpx.Request("POST", "http://test/x"))
|
||
|
|
|
||
|
|
|
||
|
|
class TestUsageResponseHook:
|
||
|
|
def test_records_usage_when_tracker_active(self):
|
||
|
|
with track_token_usage() as tracker:
|
||
|
|
response = _fake_response({
|
||
|
|
"model": "gpt-5",
|
||
|
|
"usage": {"prompt_tokens": 120, "completion_tokens": 45},
|
||
|
|
})
|
||
|
|
asyncio.run(_usage_response_hook(response))
|
||
|
|
assert tracker.as_dict() == {
|
||
|
|
"gpt-5": {"input_tokens": 120, "output_tokens": 45, "calls": 1}
|
||
|
|
}
|
||
|
|
|
||
|
|
def test_noop_when_no_tracker_active(self):
|
||
|
|
response = _fake_response({"model": "gpt-5", "usage": {"prompt_tokens": 1, "completion_tokens": 1}})
|
||
|
|
# Must not raise even though no tracker is active.
|
||
|
|
asyncio.run(_usage_response_hook(response))
|
||
|
|
|
||
|
|
def test_noop_when_response_has_no_usage_field(self):
|
||
|
|
with track_token_usage() as tracker:
|
||
|
|
response = _fake_response({"model": "gpt-5"})
|
||
|
|
asyncio.run(_usage_response_hook(response))
|
||
|
|
assert tracker.as_dict() == {}
|
||
|
|
|
||
|
|
def test_noop_on_non_json_response(self):
|
||
|
|
with track_token_usage() as tracker:
|
||
|
|
response = _fake_response(None)
|
||
|
|
asyncio.run(_usage_response_hook(response))
|
||
|
|
assert tracker.as_dict() == {}
|
||
|
|
|
||
|
|
def test_embedding_response_without_completion_tokens_defaults_output_to_zero(self):
|
||
|
|
"""Embeddings responses omit completion_tokens; output should default to 0."""
|
||
|
|
with track_token_usage() as tracker:
|
||
|
|
response = _fake_response({
|
||
|
|
"model": "Qwen/Qwen3-Embedding-4B",
|
||
|
|
"usage": {"prompt_tokens": 30, "total_tokens": 30},
|
||
|
|
})
|
||
|
|
asyncio.run(_usage_response_hook(response))
|
||
|
|
assert tracker.as_dict() == {
|
||
|
|
"Qwen/Qwen3-Embedding-4B": {"input_tokens": 30, "output_tokens": 0, "calls": 1}
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
class TestAttachUsageHook:
|
||
|
|
def test_attaches_hook_to_client_event_hooks(self):
|
||
|
|
from openai import AsyncOpenAI
|
||
|
|
|
||
|
|
client = AsyncOpenAI(api_key="sk-test", base_url="http://localhost:1")
|
||
|
|
attach_usage_hook(client)
|
||
|
|
assert _usage_response_hook in client._client.event_hooks["response"]
|
||
|
|
|
||
|
|
def test_idempotent_when_called_twice_on_same_client(self):
|
||
|
|
from openai import AsyncOpenAI
|
||
|
|
|
||
|
|
client = AsyncOpenAI(api_key="sk-test", base_url="http://localhost:1")
|
||
|
|
attach_usage_hook(client)
|
||
|
|
attach_usage_hook(client)
|
||
|
|
assert client._client.event_hooks["response"].count(_usage_response_hook) == 1
|