| Server IP : 3.147.158.171 / Your IP : 216.73.216.216 Web Server : Apache/2.4.67 (Amazon Linux) OpenSSL/3.5.5 System : Linux ip-172-31-2-178.us-east-2.compute.internal 6.1.172-216.329.amzn2023.x86_64 #1 SMP PREEMPT_DYNAMIC Wed May 20 06:31:34 UTC 2026 x86_64 User : ec2-user ( 1000) PHP Version : 8.4.21 Disable Function : NONE MySQL : OFF | cURL : ON | WGET : ON | Perl : ON | Python : OFF | Sudo : ON | Pkexec : OFF Directory : /tsai/repo/api/tests/unit/ |
Upload File : |
"""Unit tests for the interactive (agent / studio / chat) usage recording added to
close the LLM + Tavily tracking gaps.
These cover the read-side-agnostic recording layer only — they patch
``record_usage`` and assert what would be appended, so nothing touches the JSONL
file or the network.
Guarantees under test:
1. ``record_agent_usage`` turns a LangChain UsageMetadataCallbackHandler's
per-model aggregate into one ``record_usage`` call per model, with the right
provider / operation / slug.
2. It is a safe no-op when the handler collected nothing.
3. The AMA web-search wrapper's crediting is correct: an advanced Tavily search is
two credits, recorded under the ``web_search`` operation.
"""
from unittest.mock import patch
from app.services import usage_history
class _FakeCallback:
"""Stands in for langchain_core's UsageMetadataCallbackHandler."""
def __init__(self, usage_metadata):
self.usage_metadata = usage_metadata
def test_record_agent_usage_records_one_event_per_model():
cb = _FakeCallback({
"claude-sonnet-4-6": {"input_tokens": 1200, "output_tokens": 340},
"gpt-4.1-mini": {"input_tokens": 50, "output_tokens": 10},
})
with patch.object(usage_history, "record_usage") as rec:
usage_history.record_agent_usage(
cb, provider="Claude", operation="article_studio", site_slug="foo"
)
assert rec.call_count == 2
by_model = {c.kwargs["model"]: c.kwargs for c in rec.call_args_list}
assert by_model["claude-sonnet-4-6"] == {
"provider": "Claude", "model": "claude-sonnet-4-6",
"operation": "article_studio", "input_tokens": 1200,
"output_tokens": 340, "site_slug": "foo",
}
assert by_model["gpt-4.1-mini"]["input_tokens"] == 50
assert by_model["gpt-4.1-mini"]["output_tokens"] == 10
def test_record_agent_usage_noop_when_empty():
with patch.object(usage_history, "record_usage") as rec:
usage_history.record_agent_usage(
_FakeCallback({}), provider="OpenAI", operation="chat", site_slug="foo"
)
usage_history.record_agent_usage(
_FakeCallback(None), provider="OpenAI", operation="chat", site_slug="foo"
)
rec.assert_not_called()
def test_record_agent_usage_tolerates_missing_token_keys():
cb = _FakeCallback({"gpt-5.5": {}}) # model ran but handler had no counts
with patch.object(usage_history, "record_usage") as rec:
usage_history.record_agent_usage(
cb, provider="OpenAI", operation="chat", site_slug="foo"
)
assert rec.call_count == 1
assert rec.call_args.kwargs["input_tokens"] == 0
assert rec.call_args.kwargs["output_tokens"] == 0
def test_web_search_records_two_credits():
"""The AMA agent wraps its Tavily tool with record_tavily_search(advanced) under
the web_search operation — an advanced search is two credits."""
with patch.object(usage_history, "record_usage") as rec:
usage_history.record_tavily_search(search_depth="advanced", operation="web_search")
rec.assert_called_once()
assert rec.call_args.kwargs["provider"] == usage_history.PROVIDER_TAVILY
assert rec.call_args.kwargs["operation"] == "web_search"
assert rec.call_args.kwargs["tavily_credits"] == 2