feat: implement model routing engine#98
Conversation
Add strategy-based LLM model routing with resolver, 4 strategies (manual, role_based, cost_aware, smart), and 83 unit tests at 94%+ coverage. Closes #6.
📝 WalkthroughSummary by CodeRabbitRelease Notes
WalkthroughThis PR introduces a complete model routing engine subsystem for the AI company providers module. It adds observability event constants, domain models, error hierarchy, a model resolver and router, four routing strategies (Manual, RoleBasedStrategy, CostAwareStrategy, SmartStrategy), comprehensive test coverage, and public API exports. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant ModelRouter
participant RoutingStrategy
participant ModelResolver
participant ProviderConfig
Client->>ModelRouter: route(RoutingRequest)
ModelRouter->>ModelRouter: initialize (from_config)
ModelRouter->>ModelResolver: build index
ModelResolver->>ProviderConfig: iterate models
ModelResolver-->>ModelRouter: ready
ModelRouter->>RoutingStrategy: select(request, config, resolver)
RoutingStrategy->>ModelResolver: resolve(model_ref)
RoutingStrategy->>ModelResolver: all_models_sorted_by_cost()
ModelResolver-->>RoutingStrategy: ResolvedModel(s)
RoutingStrategy->>RoutingStrategy: apply strategy logic
RoutingStrategy-->>ModelRouter: RoutingDecision
ModelRouter->>ModelRouter: log ROUTING_DECISION_MADE
ModelRouter-->>Client: RoutingDecision
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes The PR introduces substantial new functionality across multiple files with heterogeneous logic. The strategies module contains complex fallback handling, budget awareness, and multi-signal decision logic. The resolver implements model indexing and cost-based sorting. Test coverage is comprehensive but adds review overhead. Domain models and error hierarchy are straightforward, reducing overall complexity somewhat. Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Dependency Review✅ No vulnerabilities or license issues or OpenSSF Scorecard issues found.Scanned FilesNone |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a robust and flexible model routing engine designed to intelligently select the most appropriate Large Language Model (LLM) based on various criteria. It centralizes model resolution and decision-making, allowing for dynamic adaptation to different use cases and operational constraints. This enhancement significantly improves the system's ability to manage and utilize diverse LLM resources efficiently. Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive model routing engine, a significant and well-structured feature. A security audit found no vulnerabilities, noting the use of Pydantic for input validation, structured logging, and automatic redaction of sensitive keys. My review suggests simplifying and making the logic more robust, particularly around fallback strategies to avoid redundancy and improve maintainability. Minor suggestions for code conciseness and a type hint correction are also provided.
| def select( | ||
| self, | ||
| request: RoutingRequest, | ||
| config: RoutingConfig, | ||
| resolver: ModelResolver, | ||
| ) -> RoutingDecision: | ||
| """Select model based on role level. | ||
|
|
||
| Raises: | ||
| ModelResolutionError: If no agent_level is set. | ||
| NoAvailableModelError: If all candidates are exhausted. | ||
| """ | ||
| if request.agent_level is None: | ||
| msg = "RoleBasedStrategy requires agent_level to be set" | ||
| raise ModelResolutionError(msg) | ||
|
|
||
| # Try matching rules | ||
| for rule in config.rules: | ||
| if rule.role_level == request.agent_level: | ||
| model, tried = _try_resolve_with_fallback( | ||
| rule.preferred_model, | ||
| rule, | ||
| config, | ||
| resolver, | ||
| ) | ||
| return RoutingDecision( | ||
| resolved_model=model, | ||
| strategy_used=self.name, | ||
| reason=( | ||
| f"Role rule match: " | ||
| f"level={request.agent_level.value}" | ||
| f", model={model.model_id}" | ||
| ), | ||
| fallbacks_tried=tried, | ||
| ) | ||
|
|
||
| # No rule matched — use seniority default | ||
| logger.debug( | ||
| ROUTING_NO_RULE_MATCHED, | ||
| level=request.agent_level.value, | ||
| strategy=self.name, | ||
| ) | ||
| tier = get_seniority_info( | ||
| request.agent_level, | ||
| ).typical_model_tier | ||
| result = _try_resolve_with_fallback_safe( | ||
| tier, | ||
| None, | ||
| config, | ||
| resolver, | ||
| ) | ||
| if result is not None: | ||
| model, tried = result | ||
| return RoutingDecision( | ||
| resolved_model=model, | ||
| strategy_used=self.name, | ||
| reason=( | ||
| f"Seniority default: level={request.agent_level.value}, tier={tier}" | ||
| ), | ||
| fallbacks_tried=tried, | ||
| ) | ||
|
|
||
| # Last resort: global fallback chain | ||
| chain_result = _walk_fallback_chain(config, resolver) | ||
| if chain_result is not None: | ||
| model, tried = chain_result | ||
| return RoutingDecision( | ||
| resolved_model=model, | ||
| strategy_used=self.name, | ||
| reason=(f"Global fallback for level={request.agent_level.value}"), | ||
| fallbacks_tried=tried, | ||
| ) | ||
|
|
||
| msg = ( | ||
| f"No model available for " | ||
| f"level={request.agent_level.value} " | ||
| f"(tier={tier}, no rules, no fallback chain)" | ||
| ) | ||
| raise NoAvailableModelError(msg) |
There was a problem hiding this comment.
This strategy has some redundant fallback logic. The _try_resolve_with_fallback and _try_resolve_with_fallback_safe helpers already attempt to use the global fallback chain. This makes the final explicit call to _walk_fallback_chain redundant. If the earlier calls fail, it means the global chain has already been tried and was unsuccessful.
This is related to my more detailed comment on SmartStrategy.select. Simplifying the helper functions to separate rule-specific fallbacks from the global fallback chain would make this strategy's logic clearer and more efficient.
| def select( | ||
| self, | ||
| request: RoutingRequest, | ||
| config: RoutingConfig, | ||
| resolver: ModelResolver, | ||
| ) -> RoutingDecision: | ||
| """Select a model using all available signals. | ||
|
|
||
| Raises: | ||
| NoAvailableModelError: If all candidates are exhausted. | ||
| """ | ||
| return ( | ||
| self._try_override(request, resolver) | ||
| or _try_task_type_rules( | ||
| request, | ||
| config, | ||
| resolver, | ||
| self.name, | ||
| ) | ||
| or _try_role_rules( | ||
| request, | ||
| config, | ||
| resolver, | ||
| self.name, | ||
| ) | ||
| or _try_seniority_default( | ||
| request, | ||
| resolver, | ||
| self.name, | ||
| ) | ||
| or self._try_cheapest(request, resolver) | ||
| or self._try_global_chain(config, resolver) | ||
| or self._raise_exhausted() | ||
| ) |
There was a problem hiding this comment.
There's some redundant logic in this strategy due to how the helper functions are structured. Both _try_task_type_rules and _try_role_rules call _try_resolve_with_fallback_safe, which already attempts to use the global fallback chain. If those helpers fail and return None, it implies the global chain was already exhausted. Therefore, the subsequent call to self._try_global_chain(config, resolver) at the end of the or sequence is redundant.
I recommend refactoring the fallback logic. The _try_resolve_with_fallback helper should likely not handle the global fallback chain. That responsibility should lie with the strategy itself as the final step. This would simplify the control flow, remove the redundancy, and make the priority order of fallbacks more explicit and easier to maintain.
| def _raise_exhausted(self) -> RoutingDecision: | ||
| msg = "SmartStrategy: no model available from any signal" | ||
| raise NoAvailableModelError(msg) |
There was a problem hiding this comment.
This method always raises an exception, so its return type should be typing.NoReturn to correctly reflect that it never returns. You'll need to add from typing import NoReturn to the file's imports.
| def _raise_exhausted(self) -> RoutingDecision: | |
| msg = "SmartStrategy: no model available from any signal" | |
| raise NoAvailableModelError(msg) | |
| def _raise_exhausted(self) -> NoReturn: | |
| msg = "SmartStrategy: no model available from any signal" | |
| raise NoAvailableModelError(msg) |
| seen_ids: set[str] = set() | ||
| unique: list[ResolvedModel] = [] | ||
| for model in self._index.values(): | ||
| if model.model_id not in seen_ids: | ||
| seen_ids.add(model.model_id) | ||
| unique.append(model) | ||
| return tuple(unique) |
There was a problem hiding this comment.
The current implementation for deduplicating models is a bit verbose. You can achieve the same result more concisely and potentially more efficiently using a dictionary comprehension to leverage the fact that dictionary keys are unique.
| seen_ids: set[str] = set() | |
| unique: list[ResolvedModel] = [] | |
| for model in self._index.values(): | |
| if model.model_id not in seen_ids: | |
| seen_ids.add(model.model_id) | |
| unique.append(model) | |
| return tuple(unique) | |
| unique_models_by_id = {model.model_id: model for model in self._index.values()} | |
| return tuple(unique_models_by_id.values()) |
There was a problem hiding this comment.
Pull request overview
Adds a new providers.routing subpackage that selects an LLM model per request using configurable routing strategies, with resolver/router infrastructure and observability events, plus a comprehensive unit-test suite to validate behavior.
Changes:
- Introduces
ModelResolver(immutable alias/ID → model index) andModelRouter(strategy dispatcher). - Implements routing strategies (
ManualStrategy,RoleBasedStrategy,CostAwareStrategy,SmartStrategy) and routing domain models/errors. - Adds routing-related observability event constants and re-exports routing APIs from
ai_company.providers.
Reviewed changes
Copilot reviewed 14 out of 15 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/providers/routing/test_strategies.py | Unit tests for all routing strategies and edge cases. |
| tests/unit/providers/routing/test_router.py | Unit tests for router construction, routing, and logging events. |
| tests/unit/providers/routing/test_resolver.py | Unit tests for resolver indexing, lookup, sorting, and immutability. |
| tests/unit/providers/routing/test_models.py | Unit tests for Pydantic routing models validation/frozen behavior. |
| tests/unit/providers/routing/test_errors.py | Unit tests for routing error hierarchy and non-retryable semantics. |
| tests/unit/providers/routing/conftest.py | Fixtures + Polyfactory factories for routing tests and sample provider configs. |
| tests/unit/providers/routing/init.py | Declares routing test package. |
| src/ai_company/providers/routing/strategies.py | Strategy implementations and shared fallback/budget helpers. |
| src/ai_company/providers/routing/router.py | ModelRouter entry point and routing decision logging. |
| src/ai_company/providers/routing/resolver.py | Immutable resolver building + alias/ID resolution and logging. |
| src/ai_company/providers/routing/models.py | ResolvedModel, RoutingRequest, RoutingDecision domain models. |
| src/ai_company/providers/routing/errors.py | Routing-specific error hierarchy under ProviderError. |
| src/ai_company/providers/routing/init.py | Public routing API exports. |
| src/ai_company/providers/init.py | Re-exports routing types/errors/strategies from ai_company.providers. |
| src/ai_company/observability/events.py | Adds routing lifecycle event constants for structured logging. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
|
||
| logger.info( | ||
| ROUTING_ROUTER_BUILT, | ||
| strategy=strategy_name, |
There was a problem hiding this comment.
ModelRouter.__init__ logs ROUTING_ROUTER_BUILT with the configured strategy_name, which can differ from the effective strategy when aliases are used (e.g. config strategy="cheapest" but router.strategy_name resolves to "cost_aware"). This can make observability confusing. Consider logging both the configured strategy and the resolved/effective strategy (or log the effective one).
| strategy=strategy_name, | |
| strategy_configured=strategy_name, | |
| strategy=self._strategy.name, |
| strategy = STRATEGY_MAP.get(strategy_name) | ||
| if strategy is None: | ||
| msg = ( | ||
| f"Unknown routing strategy {strategy_name!r}. " | ||
| f"Available: {sorted(STRATEGY_MAP)}" | ||
| ) | ||
| raise UnknownStrategyError( | ||
| msg, | ||
| context={"strategy": strategy_name}, | ||
| ) |
There was a problem hiding this comment.
When raising UnknownStrategyError, the constructor doesn't emit an ERROR/WARNING log before raising. Other provider-layer config errors (e.g. ProviderRegistry.get) log first. Consider logging a dedicated routing event (or at least logger.error(...) with the unknown strategy + available strategies) before raising to aid diagnosis in production.
| logger.info( | ||
| ROUTING_RESOLVER_BUILT, | ||
| model_count=len(index), |
There was a problem hiding this comment.
ModelResolver.from_config logs model_count=len(index), but index includes both model IDs and aliases, so this value is actually a ref/entry count (and will be ~2x the number of models when aliases are present). Consider either renaming the field (e.g. ref_count) or logging both model_count=len(resolver.all_models()) and ref_count=len(index) to avoid misleading metrics.
| logger.info( | |
| ROUTING_RESOLVER_BUILT, | |
| model_count=len(index), | |
| model_count = sum(len(provider_config.models) for provider_config in providers.values()) | |
| logger.info( | |
| ROUTING_RESOLVER_BUILT, | |
| model_count=model_count, | |
| ref_count=len(index), |
| # Last resort: global fallback chain | ||
| chain_result = _walk_fallback_chain(config, resolver) | ||
| if chain_result is not None: | ||
| model, tried = chain_result | ||
| return RoutingDecision( | ||
| resolved_model=model, | ||
| strategy_used=self.name, | ||
| reason=(f"Global fallback for level={request.agent_level.value}"), | ||
| fallbacks_tried=tried, | ||
| ) | ||
|
|
||
| msg = ( | ||
| f"No model available for " | ||
| f"level={request.agent_level.value} " | ||
| f"(tier={tier}, no rules, no fallback chain)" |
There was a problem hiding this comment.
In RoleBasedStrategy.select, the call to _try_resolve_with_fallback_safe(tier, ...) already walks config.fallback_chain (via _try_resolve_with_fallback). That means the subsequent _walk_fallback_chain block is effectively unreachable and can never return a model. Consider either (a) removing the extra _walk_fallback_chain section, or (b) changing the seniority-tier attempt to only try resolver.resolve_safe(tier) so the explicit global-chain step remains meaningful (and fallbacks_tried is accurate).
| # Last resort: global fallback chain | |
| chain_result = _walk_fallback_chain(config, resolver) | |
| if chain_result is not None: | |
| model, tried = chain_result | |
| return RoutingDecision( | |
| resolved_model=model, | |
| strategy_used=self.name, | |
| reason=(f"Global fallback for level={request.agent_level.value}"), | |
| fallbacks_tried=tried, | |
| ) | |
| msg = ( | |
| f"No model available for " | |
| f"level={request.agent_level.value} " | |
| f"(tier={tier}, no rules, no fallback chain)" | |
| msg = ( | |
| f"No model available for " | |
| f"level={request.agent_level.value} " | |
| f"(tier={tier}, no rules, fallback chain exhausted)" |
| or self._try_cheapest(request, resolver) | ||
| or self._try_global_chain(config, resolver) |
There was a problem hiding this comment.
SmartStrategy currently tries the cheapest model before consulting config.fallback_chain. As implemented, if any models are registered, _try_cheapest will always return a decision and the global fallback chain will never be used (even when the config intends a specific fallback order). Consider moving _try_global_chain before _try_cheapest, or updating the strategy docstring/tests to reflect that fallback_chain is only used when no models are available.
| or self._try_cheapest(request, resolver) | |
| or self._try_global_chain(config, resolver) | |
| or self._try_global_chain(config, resolver) | |
| or self._try_cheapest(request, resolver) |
| the model cannot be resolved. | ||
| """ | ||
| if request.model_override is None: | ||
| msg = "ManualStrategy requires model_override to be set" |
There was a problem hiding this comment.
ManualStrategy.select raises ModelResolutionError when model_override is missing, but doesn't emit a WARNING/ERROR log first. Project logging conventions call for logging error paths with context before raising (see CLAUDE.md logging section). Consider logging (at least at WARNING) with the strategy name and missing field before raising.
| msg = "ManualStrategy requires model_override to be set" | |
| msg = "ManualStrategy requires model_override to be set" | |
| logger.warning( | |
| "Routing strategy '%s' missing required field '%s'; raising ModelResolutionError", | |
| self.name, | |
| "model_override", | |
| ) |
| NoAvailableModelError: If all candidates are exhausted. | ||
| """ | ||
| if request.agent_level is None: | ||
| msg = "RoleBasedStrategy requires agent_level to be set" |
There was a problem hiding this comment.
RoleBasedStrategy.select raises ModelResolutionError when agent_level is missing, but doesn't emit a WARNING/ERROR log before raising. To align with the repo's logging rules (CLAUDE.md), consider logging the invalid request context (e.g., missing agent_level, strategy name) before raising.
| msg = "RoleBasedStrategy requires agent_level to be set" | |
| msg = "RoleBasedStrategy requires agent_level to be set" | |
| logger.warning( | |
| msg, | |
| strategy=self.name, | |
| request=request, | |
| ) |
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/ai_company/providers/routing/resolver.py`:
- Around line 84-87: The log currently uses model_count=len(index) which counts
IDs and aliases, inflating the number; compute the unique model count by
deriving the canonical model identifier for each entry (resolve aliases to their
target or use the entry's canonical id field) and use len(set(...)) for that
collection, then pass that value as model_count to LOGGER.info (the call around
ROUTING_RESOLVER_BUILT that also references providers and index).
- Around line 70-83: When building the index in the resolver (looping providers
→ provider_config.models and creating ResolvedModel), detect collisions before
assigning index[model_config.id] or index[model_config.alias]: if either key
already exists in index (or the same key maps to a different
ResolvedModel/provider_name), raise an explicit error or log and abort so you
don't silently overwrite; update the logic around the ResolvedModel
creation/assignment in resolver.py to check index.get(model_config.id) and
index.get(model_config.alias) and handle duplicates (include provider_name and
model_config.id/alias in the error message).
In `@src/ai_company/providers/routing/router.py`:
- Around line 61-69: The UnknownStrategyError is raised without logging; update
the routing logic that checks strategy (the block referencing strategy_name,
STRATEGY_MAP and raising UnknownStrategyError) to emit a warning or error log
with context (e.g., include strategy_name and available STRATEGY_MAP) via the
module/class logger before raising; ensure the log level is WARNING or ERROR and
include the same context fields passed to UnknownStrategyError so production
telemetry captures the misconfiguration prior to the exception being raised.
In `@src/ai_company/providers/routing/strategies.py`:
- Around line 322-400: The select method in RoleBasedStrategy is too long and
contains a redundant global fallback traversal: remove the final
_walk_fallback_chain(...) branch because _try_resolve_with_fallback_safe(tier,
None, config, resolver) already walks the global fallback chain; instead extract
the seniority-default resolution into a small helper (e.g.,
_resolve_seniority_default(request, config, resolver)) so select stays under ~50
lines, and return a RoutingDecision when that helper finds a model; keep error
handling (raise NoAvailableModelError) unchanged and preserve references to
RoutingDecision, _try_resolve_with_fallback_safe, and ModelResolutionError for
locating the code.
- Around line 162-165: Before raising routing exceptions (e.g.,
NoAvailableModelError) in this module, log contextual information at WARNING or
ERROR level first: capture and log variables like all_models, requested model
id/name, and relevant request context; use the module logger (logger =
logging.getLogger(__name__) or existing logger) and call logger.warning(...) for
expected error paths and logger.error(...) for unexpected/fatal ones, then raise
the exception (apply the same pattern around the other raise sites in this file
that throw routing exceptions).
In `@tests/unit/providers/routing/conftest.py`:
- Around line 59-81: The fixtures HAIKU_MODEL, SONNET_MODEL, and OPUS_MODEL are
using real Claude model IDs; update their ProviderModelConfig instances to use
fake/test model IDs (e.g., replace "claude-haiku-4-5", "claude-sonnet-4-6",
"claude-opus-4-6" with test names like "test-haiku:latest",
"test-sonnet:latest", "test-opus:latest") while keeping the rest of the fields
(alias, cost_per_1k_input, cost_per_1k_output, max_context) unchanged so tests
are decoupled from external models.
In `@tests/unit/providers/routing/test_resolver.py`:
- Around line 13-109: Tests in tests/unit/providers/routing/test_resolver.py use
real vendor model IDs (e.g., "claude-sonnet-4-6", "gpt-4o") which makes fixtures
brittle; update the ProviderConfig/ProviderModelConfig instances used by
ModelResolver.from_config in tests (functions test_indexes_model_ids,
test_indexes_aliases, test_multiple_providers, and the resolver fixture used by
TestResolverResolve/TestResolverAllModels) to use synthetic fake IDs and aliases
(e.g., "test-model:8b", "fake-writer:latest", "sonnet-test") everywhere instead
of vendor names so assertions still validate resolution, aliasing, deduplication
and cost-sorting but without real model identifiers.
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (15)
src/ai_company/observability/events.pysrc/ai_company/providers/__init__.pysrc/ai_company/providers/routing/__init__.pysrc/ai_company/providers/routing/errors.pysrc/ai_company/providers/routing/models.pysrc/ai_company/providers/routing/resolver.pysrc/ai_company/providers/routing/router.pysrc/ai_company/providers/routing/strategies.pytests/unit/providers/routing/__init__.pytests/unit/providers/routing/conftest.pytests/unit/providers/routing/test_errors.pytests/unit/providers/routing/test_models.pytests/unit/providers/routing/test_resolver.pytests/unit/providers/routing/test_router.pytests/unit/providers/routing/test_strategies.py
📜 Review details
🧰 Additional context used
📓 Path-based instructions (3)
**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.py: Do not usefrom __future__ import annotations— Python 3.14 has native lazy annotations via PEP 649
Use PEP 758 except syntax:except A, B:(no parentheses) — ruff enforces this on Python 3.14
All public functions must include type hints; enforce via mypy strict mode
Use Google-style docstrings on all public classes and functions; enforced by ruff D rules
Use Pydantic v2 models withBaseModel,model_validator, andConfigDict
Enforce 88-character line length (enforced by ruff)
Keep functions under 50 lines; keep files under 800 lines
Files:
tests/unit/providers/routing/test_strategies.pytests/unit/providers/routing/test_router.pysrc/ai_company/observability/events.pysrc/ai_company/providers/routing/models.pytests/unit/providers/routing/conftest.pysrc/ai_company/providers/routing/resolver.pytests/unit/providers/routing/test_resolver.pysrc/ai_company/providers/routing/router.pysrc/ai_company/providers/routing/errors.pytests/unit/providers/routing/test_models.pysrc/ai_company/providers/__init__.pysrc/ai_company/providers/routing/strategies.pytests/unit/providers/routing/test_errors.pysrc/ai_company/providers/routing/__init__.py
tests/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
tests/**/*.py: Mark all test cases with pytest markers:@pytest.mark.unit,@pytest.mark.integration,@pytest.mark.e2e, or@pytest.mark.slow
Configure pytest async mode asasyncio_mode = "auto"— do not manually add@pytest.mark.asyncioto test functions
Files:
tests/unit/providers/routing/test_strategies.pytests/unit/providers/routing/test_router.pytests/unit/providers/routing/conftest.pytests/unit/providers/routing/test_resolver.pytests/unit/providers/routing/test_models.pytests/unit/providers/routing/test_errors.py
src/ai_company/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
src/ai_company/**/*.py: Every module with business logic must import:from ai_company.observability import get_loggerthenlogger = get_logger(__name__)
Never useimport loggingorlogging.getLogger()orprint()in application code — use the ai_company logger
Always useloggeras the variable name (not_loggerorlog)
Always use event name constants fromai_company.observability.eventsin logger calls
Use structured logging format:logger.info(EVENT, key=value)— never use format strings likelogger.info("msg %s", val)
All error paths must log at WARNING or ERROR level with context before raising an exception
All state transitions must log at INFO level
Use DEBUG level logging for object creation, internal flow, and entry/exit of key functions
Files:
src/ai_company/observability/events.pysrc/ai_company/providers/routing/models.pysrc/ai_company/providers/routing/resolver.pysrc/ai_company/providers/routing/router.pysrc/ai_company/providers/routing/errors.pysrc/ai_company/providers/__init__.pysrc/ai_company/providers/routing/strategies.pysrc/ai_company/providers/routing/__init__.py
🧠 Learnings (3)
📚 Learning: 2026-03-01T19:59:46.936Z
Learnt from: CR
Repo: Aureliolo/ai-company PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-01T19:59:46.936Z
Learning: Applies to src/ai_company/**/*.py : Always use event name constants from `ai_company.observability.events` in logger calls
Applied to files:
src/ai_company/observability/events.py
📚 Learning: 2026-02-26T17:43:50.902Z
Learnt from: CR
Repo: Aureliolo/story-factory PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-02-26T17:43:50.902Z
Learning: Applies to tests/**/*.py : Tests must use fake model names (e.g., `test-model:8b`, `fake-writer:latest`)—never use real model IDs from `RECOMMENDED_MODELS`.
Applied to files:
tests/unit/providers/routing/test_resolver.py
📚 Learning: 2026-03-01T19:59:46.936Z
Learnt from: CR
Repo: Aureliolo/ai-company PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-01T19:59:46.936Z
Learning: Applies to src/ai_company/**/*.py : All error paths must log at WARNING or ERROR level with context before raising an exception
Applied to files:
src/ai_company/providers/routing/errors.pysrc/ai_company/providers/routing/strategies.py
🧬 Code graph analysis (8)
tests/unit/providers/routing/test_strategies.py (7)
src/ai_company/core/enums.py (1)
SeniorityLevel(6-21)src/ai_company/providers/routing/errors.py (2)
ModelResolutionError(16-19)NoAvailableModelError(22-25)src/ai_company/providers/routing/models.py (1)
RoutingRequest(43-71)src/ai_company/providers/routing/router.py (1)
resolver(80-82)tests/unit/providers/routing/conftest.py (3)
resolver(97-101)standard_routing_config(105-130)three_model_provider(85-93)src/ai_company/providers/routing/resolver.py (2)
ModelResolver(28-151)from_config(56-89)src/ai_company/providers/routing/strategies.py (9)
CostAwareStrategy(406-450)ManualStrategy(271-303)RoleBasedStrategy(309-400)SmartStrategy(456-558)select(42-62)select(282-303)select(322-400)select(418-450)select(469-502)
src/ai_company/providers/routing/models.py (1)
src/ai_company/core/enums.py (1)
SeniorityLevel(6-21)
tests/unit/providers/routing/conftest.py (5)
src/ai_company/config/schema.py (4)
ProviderConfig(52-95)ProviderModelConfig(17-49)RoutingConfig(127-149)RoutingRuleConfig(98-124)src/ai_company/core/enums.py (1)
SeniorityLevel(6-21)src/ai_company/providers/routing/models.py (3)
ResolvedModel(9-40)RoutingDecision(74-92)RoutingRequest(43-71)src/ai_company/providers/routing/router.py (1)
resolver(80-82)src/ai_company/providers/routing/resolver.py (2)
ModelResolver(28-151)from_config(56-89)
src/ai_company/providers/routing/resolver.py (3)
src/ai_company/providers/routing/errors.py (1)
ModelResolutionError(16-19)src/ai_company/providers/routing/models.py (1)
ResolvedModel(9-40)src/ai_company/config/schema.py (1)
ProviderConfig(52-95)
src/ai_company/providers/routing/errors.py (1)
src/ai_company/providers/errors.py (1)
ProviderError(22-69)
tests/unit/providers/routing/test_models.py (2)
src/ai_company/core/enums.py (1)
SeniorityLevel(6-21)src/ai_company/providers/routing/models.py (3)
ResolvedModel(9-40)RoutingDecision(74-92)RoutingRequest(43-71)
src/ai_company/providers/__init__.py (5)
src/ai_company/providers/routing/strategies.py (5)
CostAwareStrategy(406-450)ManualStrategy(271-303)RoleBasedStrategy(309-400)RoutingStrategy(34-62)SmartStrategy(456-558)src/ai_company/providers/routing/errors.py (4)
ModelResolutionError(16-19)NoAvailableModelError(22-25)RoutingError(10-13)UnknownStrategyError(28-31)src/ai_company/providers/routing/resolver.py (1)
ModelResolver(28-151)src/ai_company/providers/routing/router.py (1)
ModelRouter(27-115)src/ai_company/providers/routing/models.py (3)
ResolvedModel(9-40)RoutingDecision(74-92)RoutingRequest(43-71)
tests/unit/providers/routing/test_errors.py (2)
src/ai_company/providers/errors.py (1)
ProviderError(22-69)src/ai_company/providers/routing/errors.py (4)
ModelResolutionError(16-19)NoAvailableModelError(22-25)RoutingError(10-13)UnknownStrategyError(28-31)
🔇 Additional comments (9)
src/ai_company/providers/routing/errors.py (1)
10-31: Error hierarchy is clean and consistent.Inheritance and non-retryable semantics are implemented clearly and match provider-layer expectations.
src/ai_company/observability/events.py (1)
79-89: Routing event additions are well-scoped and consistent.The new constants follow the existing observability taxonomy and keep event naming uniform.
src/ai_company/providers/routing/models.py (1)
9-92: Model schemas are strongly typed and immutable.The three routing domain models are cleanly defined with useful constraints and frozen configs.
src/ai_company/providers/__init__.py (1)
37-52: Routing re-exports look complete and coherent.The package surface now exposes the full routing API in a consistent way.
Also applies to: 62-97
tests/unit/providers/routing/test_errors.py (1)
16-61: Test coverage for routing error behavior is solid.Hierarchy, retryability, ProviderError catchability, and context serialization are all validated.
tests/unit/providers/routing/test_strategies.py (1)
33-39: The review comment warns about aTypeErrorthat would occur fromisinstance(cls(), RoutingStrategy)without proper protocol decoration. However, verification confirms thatRoutingStrategyis already correctly decorated with@runtime_checkable(line 33), andruntime_checkableis properly imported (line 8). The test code is safe as written.Likely an incorrect or invalid review comment.
tests/unit/providers/routing/test_router.py (1)
21-146: Strong router test coverage and observability assertions.This suite validates strategy selection, alias behavior, error paths, and structured routing events end-to-end for
ModelRouter.tests/unit/providers/routing/test_models.py (1)
21-125: Model validation and immutability tests are well targeted.The cases cover the important Pydantic constraints and frozen-model behavior for
ResolvedModel,RoutingRequest, andRoutingDecision.src/ai_company/providers/routing/__init__.py (1)
7-39: Public routing API exports are clean and consistent.The re-export surface (
errors,models,resolver,router,strategies) is coherent and explicit via__all__.
| for provider_name, provider_config in providers.items(): | ||
| for model_config in provider_config.models: | ||
| resolved = ResolvedModel( | ||
| provider_name=provider_name, | ||
| model_id=model_config.id, | ||
| alias=model_config.alias, | ||
| cost_per_1k_input=model_config.cost_per_1k_input, | ||
| cost_per_1k_output=model_config.cost_per_1k_output, | ||
| max_context=model_config.max_context, | ||
| ) | ||
| index[model_config.id] = resolved | ||
| if model_config.alias is not None: | ||
| index[model_config.alias] = resolved | ||
|
|
There was a problem hiding this comment.
Prevent silent alias/ID collisions during index construction.
The current assignments can overwrite an existing ref when two providers share the same model ID or alias. That can silently route requests to the wrong provider/model.
🔧 Suggested fix
for provider_name, provider_config in providers.items():
for model_config in provider_config.models:
resolved = ResolvedModel(
provider_name=provider_name,
model_id=model_config.id,
alias=model_config.alias,
cost_per_1k_input=model_config.cost_per_1k_input,
cost_per_1k_output=model_config.cost_per_1k_output,
max_context=model_config.max_context,
)
- index[model_config.id] = resolved
- if model_config.alias is not None:
- index[model_config.alias] = resolved
+ for ref in (model_config.id, model_config.alias):
+ if ref is None:
+ continue
+ existing = index.get(ref)
+ if existing is not None and existing.model_id != resolved.model_id:
+ logger.error(
+ ROUTING_MODEL_RESOLUTION_FAILED,
+ ref=ref,
+ existing_provider=existing.provider_name,
+ existing_model_id=existing.model_id,
+ new_provider=provider_name,
+ new_model_id=resolved.model_id,
+ )
+ msg = f"Duplicate model reference {ref!r} across providers"
+ raise ModelResolutionError(
+ msg,
+ context={
+ "ref": ref,
+ "existing_provider": existing.provider_name,
+ "existing_model_id": existing.model_id,
+ "new_provider": provider_name,
+ "new_model_id": resolved.model_id,
+ },
+ )
+ index[ref] = resolved🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/ai_company/providers/routing/resolver.py` around lines 70 - 83, When
building the index in the resolver (looping providers → provider_config.models
and creating ResolvedModel), detect collisions before assigning
index[model_config.id] or index[model_config.alias]: if either key already
exists in index (or the same key maps to a different
ResolvedModel/provider_name), raise an explicit error or log and abort so you
don't silently overwrite; update the logic around the ResolvedModel
creation/assignment in resolver.py to check index.get(model_config.id) and
index.get(model_config.alias) and handle duplicates (include provider_name and
model_config.id/alias in the error message).
| logger.info( | ||
| ROUTING_RESOLVER_BUILT, | ||
| model_count=len(index), | ||
| providers=sorted(providers), |
There was a problem hiding this comment.
model_count currently logs reference count, not unique model count.
len(index) includes both IDs and aliases, so this field is inflated when aliases exist.
🔧 Suggested fix
logger.info(
ROUTING_RESOLVER_BUILT,
- model_count=len(index),
+ model_count=len({m.model_id for m in index.values()}),
+ reference_count=len(index),
providers=sorted(providers),
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| logger.info( | |
| ROUTING_RESOLVER_BUILT, | |
| model_count=len(index), | |
| providers=sorted(providers), | |
| logger.info( | |
| ROUTING_RESOLVER_BUILT, | |
| model_count=len({m.model_id for m in index.values()}), | |
| reference_count=len(index), | |
| providers=sorted(providers), | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/ai_company/providers/routing/resolver.py` around lines 84 - 87, The log
currently uses model_count=len(index) which counts IDs and aliases, inflating
the number; compute the unique model count by deriving the canonical model
identifier for each entry (resolve aliases to their target or use the entry's
canonical id field) and use len(set(...)) for that collection, then pass that
value as model_count to LOGGER.info (the call around ROUTING_RESOLVER_BUILT that
also references providers and index).
| if strategy is None: | ||
| msg = ( | ||
| f"Unknown routing strategy {strategy_name!r}. " | ||
| f"Available: {sorted(STRATEGY_MAP)}" | ||
| ) | ||
| raise UnknownStrategyError( | ||
| msg, | ||
| context={"strategy": strategy_name}, | ||
| ) |
There was a problem hiding this comment.
Log unknown strategy failures before raising.
Line 66 raises UnknownStrategyError without a preceding WARNING/ERROR log, so this configuration failure is under-observed in production telemetry.
As per coding guidelines: "All error paths must log at WARNING or ERROR level with context before raising an exception."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/ai_company/providers/routing/router.py` around lines 61 - 69, The
UnknownStrategyError is raised without logging; update the routing logic that
checks strategy (the block referencing strategy_name, STRATEGY_MAP and raising
UnknownStrategyError) to emit a warning or error log with context (e.g., include
strategy_name and available STRATEGY_MAP) via the module/class logger before
raising; ensure the log level is WARNING or ERROR and include the same context
fields passed to UnknownStrategyError so production telemetry captures the
misconfiguration prior to the exception being raised.
| if not all_models: | ||
| msg = "No models registered in resolver" | ||
| raise NoAvailableModelError(msg) | ||
|
|
There was a problem hiding this comment.
Several exception paths raise without required WARNING/ERROR logs.
Lines 164, 296, 336, 400, and 558 raise routing exceptions without logging contextual warning/error first.
Based on learnings: "Applies to src/ai_company/**/*.py : All error paths must log at WARNING or ERROR level with context before raising an exception."
Also applies to: 294-297, 334-337, 395-400, 556-558
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/ai_company/providers/routing/strategies.py` around lines 162 - 165,
Before raising routing exceptions (e.g., NoAvailableModelError) in this module,
log contextual information at WARNING or ERROR level first: capture and log
variables like all_models, requested model id/name, and relevant request
context; use the module logger (logger = logging.getLogger(__name__) or existing
logger) and call logger.warning(...) for expected error paths and
logger.error(...) for unexpected/fatal ones, then raise the exception (apply the
same pattern around the other raise sites in this file that throw routing
exceptions).
| def select( | ||
| self, | ||
| request: RoutingRequest, | ||
| config: RoutingConfig, | ||
| resolver: ModelResolver, | ||
| ) -> RoutingDecision: | ||
| """Select model based on role level. | ||
|
|
||
| Raises: | ||
| ModelResolutionError: If no agent_level is set. | ||
| NoAvailableModelError: If all candidates are exhausted. | ||
| """ | ||
| if request.agent_level is None: | ||
| msg = "RoleBasedStrategy requires agent_level to be set" | ||
| raise ModelResolutionError(msg) | ||
|
|
||
| # Try matching rules | ||
| for rule in config.rules: | ||
| if rule.role_level == request.agent_level: | ||
| model, tried = _try_resolve_with_fallback( | ||
| rule.preferred_model, | ||
| rule, | ||
| config, | ||
| resolver, | ||
| ) | ||
| return RoutingDecision( | ||
| resolved_model=model, | ||
| strategy_used=self.name, | ||
| reason=( | ||
| f"Role rule match: " | ||
| f"level={request.agent_level.value}" | ||
| f", model={model.model_id}" | ||
| ), | ||
| fallbacks_tried=tried, | ||
| ) | ||
|
|
||
| # No rule matched — use seniority default | ||
| logger.debug( | ||
| ROUTING_NO_RULE_MATCHED, | ||
| level=request.agent_level.value, | ||
| strategy=self.name, | ||
| ) | ||
| tier = get_seniority_info( | ||
| request.agent_level, | ||
| ).typical_model_tier | ||
| result = _try_resolve_with_fallback_safe( | ||
| tier, | ||
| None, | ||
| config, | ||
| resolver, | ||
| ) | ||
| if result is not None: | ||
| model, tried = result | ||
| return RoutingDecision( | ||
| resolved_model=model, | ||
| strategy_used=self.name, | ||
| reason=( | ||
| f"Seniority default: level={request.agent_level.value}, tier={tier}" | ||
| ), | ||
| fallbacks_tried=tried, | ||
| ) | ||
|
|
||
| # Last resort: global fallback chain | ||
| chain_result = _walk_fallback_chain(config, resolver) | ||
| if chain_result is not None: | ||
| model, tried = chain_result | ||
| return RoutingDecision( | ||
| resolved_model=model, | ||
| strategy_used=self.name, | ||
| reason=(f"Global fallback for level={request.agent_level.value}"), | ||
| fallbacks_tried=tried, | ||
| ) | ||
|
|
||
| msg = ( | ||
| f"No model available for " | ||
| f"level={request.agent_level.value} " | ||
| f"(tier={tier}, no rules, no fallback chain)" | ||
| ) | ||
| raise NoAvailableModelError(msg) |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Refactor RoleBasedStrategy.select (oversized + redundant fallback traversal).
This method exceeds the 50-line limit, and the _walk_fallback_chain() branch after _try_resolve_with_fallback_safe(..., rule=None, ...) is redundant because the helper already walks the global chain.
Refactor direction (remove redundant branch as part of decomposition)
- # Last resort: global fallback chain
- chain_result = _walk_fallback_chain(config, resolver)
- if chain_result is not None:
- model, tried = chain_result
- return RoutingDecision(
- resolved_model=model,
- strategy_used=self.name,
- reason=(f"Global fallback for level={request.agent_level.value}"),
- fallbacks_tried=tried,
- )As per coding guidelines: "Keep functions under 50 lines; keep files under 800 lines."
Also applies to: 367-386
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/ai_company/providers/routing/strategies.py` around lines 322 - 400, The
select method in RoleBasedStrategy is too long and contains a redundant global
fallback traversal: remove the final _walk_fallback_chain(...) branch because
_try_resolve_with_fallback_safe(tier, None, config, resolver) already walks the
global fallback chain; instead extract the seniority-default resolution into a
small helper (e.g., _resolve_seniority_default(request, config, resolver)) so
select stays under ~50 lines, and return a RoutingDecision when that helper
finds a model; keep error handling (raise NoAvailableModelError) unchanged and
preserve references to RoutingDecision, _try_resolve_with_fallback_safe, and
ModelResolutionError for locating the code.
| HAIKU_MODEL = ProviderModelConfig( | ||
| id="claude-haiku-4-5", | ||
| alias="haiku", | ||
| cost_per_1k_input=0.001, | ||
| cost_per_1k_output=0.005, | ||
| max_context=200_000, | ||
| ) | ||
|
|
||
| SONNET_MODEL = ProviderModelConfig( | ||
| id="claude-sonnet-4-6", | ||
| alias="sonnet", | ||
| cost_per_1k_input=0.003, | ||
| cost_per_1k_output=0.015, | ||
| max_context=200_000, | ||
| ) | ||
|
|
||
| OPUS_MODEL = ProviderModelConfig( | ||
| id="claude-opus-4-6", | ||
| alias="opus", | ||
| cost_per_1k_input=0.015, | ||
| cost_per_1k_output=0.075, | ||
| max_context=200_000, | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect ProviderModelConfig declaration and key fields
echo "=== ProviderModelConfig class definition ==="
rg -n -A 20 'class\s+ProviderModelConfig' src/ai_company/config/schema.py | head -40
echo ""
echo "=== Field definitions (id, model_id) ==="
rg -n -B 2 -A 1 '(id|model_id)\s*:' src/ai_company/config/schema.py | head -30
echo ""
echo "=== ConfigDict and alias/population settings ==="
rg -n -B 2 -A 2 '(ConfigDict|validation_alias|populate_by_name)' src/ai_company/config/schema.py | head -40Repository: Aureliolo/ai-company
Length of output: 2169
Use fake model names in test fixtures; replace real Claude IDs with test-model variants.
Lines 60, 68, and 76 use real Claude model IDs (claude-haiku-4-5, claude-sonnet-4-6, claude-opus-4-6). Test fixtures must use fake model names like test-haiku:latest, test-sonnet:latest, test-opus:latest to avoid coupling tests to external model availability and pricing.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/unit/providers/routing/conftest.py` around lines 59 - 81, The fixtures
HAIKU_MODEL, SONNET_MODEL, and OPUS_MODEL are using real Claude model IDs;
update their ProviderModelConfig instances to use fake/test model IDs (e.g.,
replace "claude-haiku-4-5", "claude-sonnet-4-6", "claude-opus-4-6" with test
names like "test-haiku:latest", "test-sonnet:latest", "test-opus:latest") while
keeping the rest of the fields (alias, cost_per_1k_input, cost_per_1k_output,
max_context) unchanged so tests are decoupled from external models.
| def test_indexes_model_ids( | ||
| self, | ||
| three_model_provider: dict[str, ProviderConfig], | ||
| ) -> None: | ||
| resolver = ModelResolver.from_config(three_model_provider) | ||
| model = resolver.resolve("claude-sonnet-4-6") | ||
| assert model.model_id == "claude-sonnet-4-6" | ||
| assert model.provider_name == "anthropic" | ||
|
|
||
| def test_indexes_aliases( | ||
| self, | ||
| three_model_provider: dict[str, ProviderConfig], | ||
| ) -> None: | ||
| resolver = ModelResolver.from_config(three_model_provider) | ||
| model = resolver.resolve("sonnet") | ||
| assert model.model_id == "claude-sonnet-4-6" | ||
|
|
||
| def test_empty_providers(self) -> None: | ||
| resolver = ModelResolver.from_config({}) | ||
| assert resolver.all_models() == () | ||
|
|
||
| def test_multiple_providers(self) -> None: | ||
| providers = { | ||
| "anthropic": ProviderConfig( | ||
| models=( | ||
| ProviderModelConfig( | ||
| id="claude-sonnet-4-6", | ||
| alias="sonnet", | ||
| cost_per_1k_input=0.003, | ||
| cost_per_1k_output=0.015, | ||
| ), | ||
| ), | ||
| ), | ||
| "openai": ProviderConfig( | ||
| models=( | ||
| ProviderModelConfig( | ||
| id="gpt-4o", | ||
| alias="gpt4", | ||
| cost_per_1k_input=0.005, | ||
| cost_per_1k_output=0.015, | ||
| ), | ||
| ), | ||
| ), | ||
| } | ||
| resolver = ModelResolver.from_config(providers) | ||
| assert len(resolver.all_models()) == 2 | ||
|
|
||
|
|
||
| class TestResolverResolve: | ||
| def test_resolve_by_id(self, resolver: ModelResolver) -> None: | ||
| model = resolver.resolve("claude-haiku-4-5") | ||
| assert model.model_id == "claude-haiku-4-5" | ||
|
|
||
| def test_resolve_by_alias(self, resolver: ModelResolver) -> None: | ||
| model = resolver.resolve("opus") | ||
| assert model.model_id == "claude-opus-4-6" | ||
|
|
||
| def test_resolve_unknown_raises(self, resolver: ModelResolver) -> None: | ||
| with pytest.raises(ModelResolutionError, match="not found"): | ||
| resolver.resolve("nonexistent") | ||
|
|
||
| def test_resolve_error_contains_context(self, resolver: ModelResolver) -> None: | ||
| with pytest.raises(ModelResolutionError) as exc_info: | ||
| resolver.resolve("nonexistent") | ||
| assert exc_info.value.context["ref"] == "nonexistent" | ||
|
|
||
|
|
||
| class TestResolverResolveSafe: | ||
| def test_resolve_safe_found(self, resolver: ModelResolver) -> None: | ||
| model = resolver.resolve_safe("sonnet") | ||
| assert model is not None | ||
| assert model.model_id == "claude-sonnet-4-6" | ||
|
|
||
| def test_resolve_safe_not_found(self, resolver: ModelResolver) -> None: | ||
| assert resolver.resolve_safe("nonexistent") is None | ||
|
|
||
|
|
||
| class TestResolverAllModels: | ||
| def test_all_models_deduplicates(self, resolver: ModelResolver) -> None: | ||
| models = resolver.all_models() | ||
| ids = [m.model_id for m in models] | ||
| assert len(ids) == len(set(ids)) | ||
| assert len(models) == 3 | ||
|
|
||
| def test_all_models_sorted_by_cost(self, resolver: ModelResolver) -> None: | ||
| models = resolver.all_models_sorted_by_cost() | ||
| costs = [m.cost_per_1k_input + m.cost_per_1k_output for m in models] | ||
| assert costs == sorted(costs) | ||
|
|
||
| def test_cheapest_is_haiku(self, resolver: ModelResolver) -> None: | ||
| models = resolver.all_models_sorted_by_cost() | ||
| assert models[0].alias == "haiku" | ||
|
|
||
| def test_most_expensive_is_opus(self, resolver: ModelResolver) -> None: | ||
| models = resolver.all_models_sorted_by_cost() | ||
| assert models[-1].alias == "opus" | ||
|
|
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Prefer synthetic model identifiers in unit tests.
These tests use real vendor model IDs (claude-*, gpt-4o), which can make test fixtures brittle to catalog churn. Use fake IDs/aliases instead.
Based on learnings: Tests must use fake model names (e.g., test-model:8b, fake-writer:latest)—never use real model IDs from RECOMMENDED_MODELS.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/unit/providers/routing/test_resolver.py` around lines 13 - 109, Tests
in tests/unit/providers/routing/test_resolver.py use real vendor model IDs
(e.g., "claude-sonnet-4-6", "gpt-4o") which makes fixtures brittle; update the
ProviderConfig/ProviderModelConfig instances used by ModelResolver.from_config
in tests (functions test_indexes_model_ids, test_indexes_aliases,
test_multiple_providers, and the resolver fixture used by
TestResolverResolve/TestResolverAllModels) to use synthetic fake IDs and aliases
(e.g., "test-model:8b", "fake-writer:latest", "sonnet-test") everywhere instead
of vendor names so assertions still validate resolution, aliasing, deduplication
and cost-sorting but without real model identifiers.
Summary
providers/routing/subpackage with strategy-based LLM model routingModelResolverfor O(1) alias/ID → model lookup with immutable indexManualStrategy,RoleBasedStrategy,CostAwareStrategy,SmartStrategyModelRouterentry point that delegates to configured strategyproviders/__init__.pyCloses #6
Test plan
🤖 Generated with Claude Code