forked from cardosofelipe/fast-next-template
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:
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user