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

@@ -79,14 +79,23 @@ class TestModelRouter:
def test_get_preferred_group_for_agent(self, router: ModelRouter) -> None:
"""Test getting preferred group for agent types."""
assert router.get_preferred_group_for_agent("product_owner") == ModelGroup.REASONING
assert router.get_preferred_group_for_agent("software_engineer") == ModelGroup.CODE
assert router.get_preferred_group_for_agent("devops_engineer") == ModelGroup.FAST
assert (
router.get_preferred_group_for_agent("product_owner")
== ModelGroup.REASONING
)
assert (
router.get_preferred_group_for_agent("software_engineer") == ModelGroup.CODE
)
assert (
router.get_preferred_group_for_agent("devops_engineer") == ModelGroup.FAST
)
def test_get_preferred_group_unknown_agent(self, router: ModelRouter) -> None:
"""Test getting preferred group for unknown agent."""
# Should default to REASONING
assert router.get_preferred_group_for_agent("unknown_type") == ModelGroup.REASONING
assert (
router.get_preferred_group_for_agent("unknown_type") == ModelGroup.REASONING
)
def test_select_model_by_group(self, router: ModelRouter) -> None:
"""Test selecting model by group."""
@@ -99,9 +108,7 @@ class TestModelRouter:
def test_select_model_by_group_string(self, router: ModelRouter) -> None:
"""Test selecting model by group string."""
model_name, config = asyncio.run(
router.select_model(model_group="code")
)
model_name, config = asyncio.run(router.select_model(model_group="code"))
assert model_name == "claude-sonnet-4"
@@ -174,9 +181,7 @@ class TestModelRouter:
limited_router = ModelRouter(settings=settings, circuit_registry=registry)
with pytest.raises(AllProvidersFailedError) as exc_info:
asyncio.run(
limited_router.select_model(model_group=ModelGroup.REASONING)
)
asyncio.run(limited_router.select_model(model_group=ModelGroup.REASONING))
assert exc_info.value.model_group == "reasoning"
assert len(exc_info.value.attempted_models) > 0
@@ -195,18 +200,14 @@ class TestModelRouter:
def test_get_available_models_for_group_string(self, router: ModelRouter) -> None:
"""Test getting available models with string group."""
models = asyncio.run(
router.get_available_models_for_group("code")
)
models = asyncio.run(router.get_available_models_for_group("code"))
assert len(models) > 0
def test_get_available_models_invalid_group(self, router: ModelRouter) -> None:
"""Test getting models for invalid group."""
with pytest.raises(InvalidModelGroupError):
asyncio.run(
router.get_available_models_for_group("invalid")
)
asyncio.run(router.get_available_models_for_group("invalid"))
def test_get_all_model_groups(self, router: ModelRouter) -> None:
"""Test getting all model groups info."""