fix(llm-gateway): improve type safety and datetime consistency

- Add type annotations for mypy compliance
- Use UTC-aware datetimes consistently (datetime.now(UTC))
- Add type: ignore comments for LiteLLM incomplete stubs
- Fix import ordering and formatting
- Update pyproject.toml mypy configuration

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
2026-01-03 20:56:05 +01:00
parent 6e8b0b022a
commit f482559e15
15 changed files with 111 additions and 105 deletions

View File

@@ -108,19 +108,19 @@ class TestCostTracker:
)
# Verify by getting usage report
report = asyncio.run(
tracker.get_project_usage("proj-123", period="day")
)
report = asyncio.run(tracker.get_project_usage("proj-123", period="day"))
assert report.total_requests == 1
assert report.total_cost_usd == pytest.approx(0.01, rel=0.01)
def test_record_usage_disabled(self, tracker_settings: Settings) -> None:
"""Test recording is skipped when disabled."""
settings = Settings(**{
**tracker_settings.model_dump(),
"cost_tracking_enabled": False,
})
settings = Settings(
**{
**tracker_settings.model_dump(),
"cost_tracking_enabled": False,
}
)
fake_redis = fakeredis.aioredis.FakeRedis(decode_responses=True)
disabled_tracker = CostTracker(redis_client=fake_redis, settings=settings)
@@ -157,9 +157,7 @@ class TestCostTracker:
)
# Verify session usage
session_usage = asyncio.run(
tracker.get_session_usage("session-789")
)
session_usage = asyncio.run(tracker.get_session_usage("session-789"))
assert session_usage["session_id"] == "session-789"
assert session_usage["total_cost_usd"] == pytest.approx(0.01, rel=0.01)
@@ -199,9 +197,7 @@ class TestCostTracker:
)
)
report = asyncio.run(
tracker.get_project_usage("proj-123", period="day")
)
report = asyncio.run(tracker.get_project_usage("proj-123", period="day"))
assert report.total_requests == 2
assert len(report.by_model) == 2
@@ -221,9 +217,7 @@ class TestCostTracker:
)
)
report = asyncio.run(
tracker.get_agent_usage("agent-456", period="day")
)
report = asyncio.run(tracker.get_agent_usage("agent-456", period="day"))
assert report.entity_id == "agent-456"
assert report.entity_type == "agent"
@@ -243,12 +237,8 @@ class TestCostTracker:
)
# Check different periods
hour_report = asyncio.run(
tracker.get_project_usage("proj-123", period="hour")
)
day_report = asyncio.run(
tracker.get_project_usage("proj-123", period="day")
)
hour_report = asyncio.run(tracker.get_project_usage("proj-123", period="hour"))
day_report = asyncio.run(tracker.get_project_usage("proj-123", period="day"))
month_report = asyncio.run(
tracker.get_project_usage("proj-123", period="month")
)
@@ -306,9 +296,7 @@ class TestCostTracker:
def test_check_budget_default_limit(self, tracker: CostTracker) -> None:
"""Test budget check with default limit."""
within, current, limit = asyncio.run(
tracker.check_budget("proj-123")
)
within, current, limit = asyncio.run(tracker.check_budget("proj-123"))
assert limit == 1000.0 # Default from settings