Skip to content

feat: add in-memory route cache to Router #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
itsarbit wants to merge 1 commit into
base: master
Choose a base branch
from
Open

feat: add in-memory route cache to Router #1

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 19 additions & 3 deletions src/tokenwise/router.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,9 @@
"complex": (2000, 1000),
}

# Simple in-memory route cache: key = (query, strategy, budget) -> model_id

@itsarbit itsarbit Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[MEDIUM] Global mutable state for cache

The use of a global dictionary _route_cache for caching introduces mutable global state, which can lead to issues in concurrent environments and makes the function's behavior dependent on external state.

Suggestion: Consider encapsulating the cache within a class or using a thread-safe structure to manage the cache, such as a functools.lru_cache or an instance variable if route is part of a class.

Endorsed by: security, performance, readability, testing

Found by architecture | Consensus score: 10

_route_cache: dict[tuple, str] = {}


def _detect_capabilities(query: str) -> list[str]:
"""Detect likely required capabilities from query text."""
Expand Down Expand Up @@ -127,6 +130,15 @@ def route(
if isinstance(strategy, str):

@itsarbit itsarbit Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[MEDIUM] Concurrency issue with in-memory cache

The in-memory cache _route_cache is not thread-safe. Concurrent access to this cache could lead to race conditions, especially in a multi-threaded environment.

Suggestion: Consider using a thread-safe data structure like collections.defaultdict with a threading.Lock or use a library like cachetools that provides thread-safe cache implementations.

Endorsed by: security, readability, architecture, testing

Found by performance | Consensus score: 10

@itsarbit itsarbit Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[MEDIUM] Missing test coverage for cache mechanism

The newly introduced in-memory cache for routing results is not covered by tests. There should be tests verifying that the cache is correctly storing and retrieving values, and that it handles cache hits and misses appropriately.

Suggestion: Add unit tests that specifically test the caching behavior, including scenarios where the cache is hit and where it is missed.

Endorsed by: security, performance, readability, architecture

Found by testing | Consensus score: 10

@itsarbit itsarbit Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[MEDIUM] Potential cache invalidation issue

The cache does not appear to handle invalidation. If the underlying model registry changes (e.g., models are added or removed), the cache could return stale or invalid results.

Suggestion: Implement a cache invalidation strategy that clears or updates the cache when the model registry changes.

Endorsed by: security, performance, readability, architecture

Found by testing | Consensus score: 10

@itsarbit itsarbit Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[MEDIUM] Cache key construction

The cache key is constructed using a tuple of (query, str(strategy), budget). If query or strategy can have high variability or large size, this might lead to inefficient cache usage or memory issues.

Suggestion: Ensure that the components of the cache key are appropriately hashed or reduced in size, and consider the potential for high cardinality in the cache keys.

Found by architecture | Consensus score: 2

@itsarbit itsarbit Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[LOW] Potential unnecessary model search

When a cache hit occurs, the code iterates over all models to find the one with the matching ID. This could be inefficient if the number of models is large.

Suggestion: Store the entire model object in the cache instead of just the model ID to avoid the need to search through the models again.

Challenged by security: This is a performance optimization suggestion rather than a critical issue. The current implementation may be acceptable depending on the typical size of the model list.

Challenged by readability: Storing the entire model object in the cache might increase memory usage significantly, especially if models are large. The trade-off between memory usage and search efficiency needs careful consideration.

Found by performance | Consensus score: -1

@itsarbit itsarbit Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[MEDIUM] Cache Key Construction

The cache key is constructed using the 'query', 'strategy', and 'budget'. If these parameters are mutable or if their string representations are not unique, it could lead to cache misses or incorrect cache hits.

Suggestion: Ensure that the parameters used in the cache key are immutable and have unique string representations. Consider using a more robust method to generate cache keys if necessary.

Endorsed by: performance

Challenged by security: While the construction of cache keys should be considered, the current implementation may be sufficient if the parameters are not highly mutable or variable.

Challenged by architecture: The concern about cache key uniqueness was already addressed in my original finding regarding high variability and size of cache keys.

Challenged by testing: The concern about cache key construction is already covered by my original finding regarding the cache key not considering all relevant factors. The focus should be on ensuring the key includes all necessary factors and is immutable.

Found by readability | Consensus score: -2

@itsarbit itsarbit Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[LOW] Cache key does not consider all relevant factors

The cache key is based on query, strategy, and budget, but does not consider other potential factors that might affect routing, such as user-specific settings or environmental conditions.

Suggestion: Review the factors that should influence routing and ensure they are all included in the cache key if necessary.

Challenged by security: The current cache key factors may be sufficient for the intended use case, and adding more factors could unnecessarily complicate the cache logic.

Challenged by performance: Without specific context on what other factors might affect routing, this finding seems speculative. The current cache key construction seems reasonable based on the provided parameters.

Challenged by readability: The current cache key construction might be sufficient if the routing logic is primarily influenced by query, strategy, and budget. Adding more factors could unnecessarily complicate the cache key without clear benefits.

Challenged by architecture: The current cache key construction seems to be based on the primary factors affecting routing. Additional factors should be considered only if they are proven to have a significant impact on routing decisions.

Found by testing | Consensus score: -3

@itsarbit itsarbit Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[MEDIUM] required_capability missing from cache key

The cache key does not include required_capability, which is a parameter of the route function. If required_capability changes between calls with the same query, strategy, and budget, the cache may return an incorrect model. For example, if a user requests a different capability but the same query, strategy, and budget, they might receive a cached result that doesn't meet the new capability requirement.

Suggestion: Include required_capability in the cache key by modifying the cache_key tuple to (query, str(strategy), budget, required_capability).

Endorsed by: correctness

Challenged by architecture: The claim is incorrect because the cache is designed to differentiate based on budget, which is a valid requirement. Normalizing budget values could lead to incorrect cache hits if different budgets should indeed lead to different routing outcomes.

Found by architecture, testing | Consensus score: 2

@itsarbit itsarbit Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[MEDIUM] Cache key does not include required_capability

The cache key is constructed using (query, strategy, budget), but it does not include required_capability. If required_capability changes, the cache may return an incorrect model that does not satisfy the new capability requirements. This can lead to incorrect routing decisions.

Suggestion: Include required_capability in the cache key by modifying the line to: cache_key = (query, str(strategy), budget, required_capability).

Endorsed by: performance, readability, architecture, testing, correctness

Found by security, performance, architecture | Consensus score: 12

strategy = RoutingStrategy(strategy)

# Check cache first
cache_key = (query, str(strategy), budget)

@itsarbit itsarbit Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[LOW] Inefficient model lookup after cache hit

After a cache hit, the code iterates over all models to find the one with the cached ID. This can be inefficient if the number of models is large. The performance impact increases with the size of the model registry.

Suggestion: Consider maintaining a dictionary mapping model IDs to model instances in self.registry to allow O(1) lookups by ID.

Endorsed by: readability, architecture, testing

Found by performance | Consensus score: 4

@itsarbit itsarbit Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[MEDIUM] Potential stale cache entry due to model registry changes

The cache stores a model_id based on the current state of the model registry. If the registry changes (e.g., models are added or removed), the cached model_id may become invalid, leading to incorrect routing results. This issue can occur if the model registry is dynamic and changes over time.

Suggestion: Invalidate or update the cache whenever the model registry changes. Alternatively, include a version or timestamp of the registry as part of the cache key to ensure cache entries are only valid for the current state of the registry.

Endorsed by: security, readability, architecture, testing, correctness

Found by performance | Consensus score: 12

if cache_key in _route_cache:
cached_id = _route_cache[cache_key]
models = self.registry.find_models()
for m in models:
if m.id == cached_id:
return m

# ── Stage 1: Scenario detection ──────────────────────────────
primary_cap, complexity = self._detect_scenario(query, required_capability)

Expand Down Expand Up @@ -156,11 +168,15 @@ def route(

@itsarbit itsarbit Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[LOW] Cache Storage

The cache stores the 'result.id' without checking if 'result' is None. If no result is found, this could lead to a KeyError when accessing 'result.id'.

Suggestion: Add a check to ensure 'result' is not None before attempting to store 'result.id' in the cache.

Endorsed by: testing

Found by readability | Consensus score: 2

@itsarbit itsarbit Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[LOW] Inconsistent return pattern after caching logic

The function initially returns directly from the strategy-specific routing functions, but after caching logic is added, it assigns the result to a variable and returns it later. This inconsistency can make the control flow harder to follow, especially if further modifications are needed.

Suggestion: Standardize the return pattern by always assigning the result to a variable and returning it at the end, even before the caching logic is added. This will make the function easier to modify and understand.

Endorsed by: security, performance, architecture, testing

Challenged by correctness: The quoted claim is not a correctness issue; it is a style preference. The current implementation does not affect the correctness of the function's behavior.

Challenged by readability: The quoted claim is incorrect because the cache logic is designed to store and retrieve IDs, not full model objects, which is a valid approach if the lookup logic correctly maps IDs back to objects.

Found by readability, correctness | Consensus score: 3

@itsarbit itsarbit Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[LOW] Cache storage should verify model existence

The cache stores the result's id without verifying if the model is still valid or exists in the registry. If models are removed or changed, the cache might return stale or invalid results. For instance, if a model is deleted after being cached, subsequent cache hits will return an ID that no longer corresponds to a valid model.

Suggestion: Before storing the result in the cache, verify that the model is still valid by checking its existence in the registry. Alternatively, implement a cache invalidation strategy when models are updated or removed.

Found by readability | Consensus score: 1

# Apply strategy preference
if strategy == RoutingStrategy.CHEAPEST:
return self._route_cheapest(candidates)
result = self._route_cheapest(candidates)
elif strategy == RoutingStrategy.BEST_QUALITY:
return self._route_best_quality(candidates)
result = self._route_best_quality(candidates)
else: # balanced
return self._route_balanced(candidates, complexity)
result = self._route_balanced(candidates, complexity)

# Store in cache
_route_cache[cache_key] = result.id
return result

def route_with_trace(
self,
Expand Down
Loading