403Webshell
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 :
current_dir [ Writeable ] document_root [ Writeable ]

 

Command :


[ Back ]     

Current File : /tsai/repo/api/tests/unit/test_usage_agent.py
"""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

Youez - 2016 - github.com/yon3zu
LinuXploit