"""Usage Tracker 测试""" from datetime import datetime, timedelta, timezone import pytest from agentkit.llm.protocol import TokenUsage from agentkit.llm.providers.tracker import UsageSummary, UsageTracker from agentkit.llm.providers.usage_store import UsageRecord 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, ) # Verify via get_usage() instead of internal _records summary = tracker.get_usage() assert len(summary.records) == 1 rec = summary.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) summary = tracker.get_usage() assert len(summary.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) tracker.record("agent_a", "gpt-4o", usage2, 0.010, 200.0) # Manually set timestamp of second record to 2 hours ago via store store = tracker._store store._records[-1].timestamp = (now - timedelta(hours=2)).isoformat() # 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 == []