119 lines
4.2 KiB
Python
119 lines
4.2 KiB
Python
"""Usage Tracker 测试"""
|
|
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
import pytest
|
|
|
|
from agentkit.llm.protocol import TokenUsage
|
|
from agentkit.llm.providers.tracker import UsageRecord, UsageSummary, UsageTracker
|
|
|
|
|
|
class TestUsageTrackerRecord:
|
|
"""record() 方法测试"""
|
|
|
|
def test_record_stores_usage(self):
|
|
tracker = UsageTracker()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
|
|
tracker.record(
|
|
agent_name="test_agent",
|
|
model="gpt-4o",
|
|
usage=usage,
|
|
cost=0.005,
|
|
latency_ms=200.0,
|
|
)
|
|
|
|
assert len(tracker._records) == 1
|
|
rec = tracker._records[0]
|
|
assert rec.agent_name == "test_agent"
|
|
assert rec.model == "gpt-4o"
|
|
assert rec.prompt_tokens == 100
|
|
assert rec.completion_tokens == 50
|
|
assert rec.total_tokens == 150
|
|
assert rec.cost == 0.005
|
|
assert rec.latency_ms == 200.0
|
|
|
|
def test_record_multiple_entries(self):
|
|
tracker = UsageTracker()
|
|
usage1 = TokenUsage(prompt_tokens=10, completion_tokens=5)
|
|
usage2 = TokenUsage(prompt_tokens=20, completion_tokens=10)
|
|
|
|
tracker.record("agent_a", "gpt-4o", usage1, 0.001, 100.0)
|
|
tracker.record("agent_b", "deepseek-chat", usage2, 0.002, 150.0)
|
|
|
|
assert len(tracker._records) == 2
|
|
|
|
|
|
class TestUsageTrackerGetUsage:
|
|
"""get_usage() 方法测试"""
|
|
|
|
def test_get_usage_aggregates_totals(self):
|
|
tracker = UsageTracker()
|
|
usage1 = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
usage2 = TokenUsage(prompt_tokens=200, completion_tokens=100)
|
|
|
|
tracker.record("agent_a", "gpt-4o", usage1, 0.005, 100.0)
|
|
tracker.record("agent_a", "gpt-4o", usage2, 0.010, 200.0)
|
|
|
|
summary = tracker.get_usage()
|
|
assert summary.total_tokens == 450
|
|
assert summary.total_cost == pytest.approx(0.015)
|
|
assert len(summary.records) == 2
|
|
|
|
def test_get_usage_filters_by_agent_name(self):
|
|
tracker = UsageTracker()
|
|
usage1 = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
usage2 = TokenUsage(prompt_tokens=200, completion_tokens=100)
|
|
|
|
tracker.record("agent_a", "gpt-4o", usage1, 0.005, 100.0)
|
|
tracker.record("agent_b", "gpt-4o", usage2, 0.010, 200.0)
|
|
|
|
summary = tracker.get_usage(agent_name="agent_a")
|
|
assert summary.total_tokens == 150
|
|
assert len(summary.records) == 1
|
|
assert summary.records[0].agent_name == "agent_a"
|
|
|
|
def test_get_usage_filters_by_time_range(self):
|
|
tracker = UsageTracker()
|
|
now = datetime.now(timezone.utc)
|
|
usage1 = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
usage2 = TokenUsage(prompt_tokens=200, completion_tokens=100)
|
|
|
|
tracker.record("agent_a", "gpt-4o", usage1, 0.005, 100.0)
|
|
|
|
# Manually set timestamp of second record to 2 hours ago
|
|
tracker.record("agent_a", "gpt-4o", usage2, 0.010, 200.0)
|
|
tracker._records[-1].timestamp = now - timedelta(hours=2)
|
|
|
|
# Query last hour only
|
|
summary = tracker.get_usage(start_time=now - timedelta(hours=1), end_time=now + timedelta(hours=1))
|
|
assert len(summary.records) == 1
|
|
assert summary.total_tokens == 150
|
|
|
|
def test_get_usage_by_model(self):
|
|
tracker = UsageTracker()
|
|
usage1 = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
usage2 = TokenUsage(prompt_tokens=200, completion_tokens=100)
|
|
|
|
tracker.record("agent_a", "gpt-4o", usage1, 0.005, 100.0)
|
|
tracker.record("agent_a", "deepseek-chat", usage2, 0.002, 200.0)
|
|
|
|
summary = tracker.get_usage()
|
|
assert "gpt-4o" in summary.by_model
|
|
assert "deepseek-chat" in summary.by_model
|
|
assert summary.by_model["gpt-4o"]["total_tokens"] == 150
|
|
assert summary.by_model["deepseek-chat"]["total_tokens"] == 300
|
|
|
|
|
|
class TestUsageSummaryEmpty:
|
|
"""空记录 UsageSummary 测试"""
|
|
|
|
def test_empty_records_return_zero_summary(self):
|
|
tracker = UsageTracker()
|
|
summary = tracker.get_usage()
|
|
assert isinstance(summary, UsageSummary)
|
|
assert summary.total_tokens == 0
|
|
assert summary.total_cost == 0.0
|
|
assert summary.by_model == {}
|
|
assert summary.records == []
|