From 1c7cbf64efbfb231c9ff1ac905cdb6b8e95a4ae3 Mon Sep 17 00:00:00 2001 From: Shroominic Date: Mon, 3 Nov 2025 22:24:27 +0800 Subject: [PATCH] ruff fmt --- routstr/algorithm.py | 136 ++++++++++++++++++----------------- tests/unit/test_algorithm.py | 73 +++++++++++-------- 2 files changed, 115 insertions(+), 94 deletions(-) diff --git a/routstr/algorithm.py b/routstr/algorithm.py index b4899228..66f25d2b 100644 --- a/routstr/algorithm.py +++ b/routstr/algorithm.py @@ -13,72 +13,74 @@ logger = get_logger(__name__) def calculate_model_cost_score(model: "Model") -> float: """Calculate a representative cost score for a model. - + This score is used to compare models when multiple providers offer the same model. Lower scores indicate cheaper models. - + The score is calculated as a weighted average of: - Input token cost (weighted by typical input usage) - Output token cost (weighted by typical output usage) - Fixed request cost - + Args: model: Model instance with pricing information - + Returns: Float representing the cost score. Lower is better. """ pricing = model.pricing - + # Weight costs by typical usage patterns # Assume average request: 1000 input tokens, 500 output tokens TYPICAL_INPUT_TOKENS = 1000.0 TYPICAL_OUTPUT_TOKENS = 500.0 - + # Calculate weighted cost in USD input_cost = pricing.prompt * (TYPICAL_INPUT_TOKENS / 1000.0) output_cost = pricing.completion * (TYPICAL_OUTPUT_TOKENS / 1000.0) request_cost = pricing.request - + # Include additional costs if present - image_cost = getattr(pricing, "image", 0.0) * 0.1 # Weight lower as not every request uses images + image_cost = ( + getattr(pricing, "image", 0.0) * 0.1 + ) # Weight lower as not every request uses images web_search_cost = getattr(pricing, "web_search", 0.0) * 0.1 reasoning_cost = getattr(pricing, "internal_reasoning", 0.0) * 0.2 - + total_cost = ( - input_cost + - output_cost + - request_cost + - image_cost + - web_search_cost + - reasoning_cost + input_cost + + output_cost + + request_cost + + image_cost + + web_search_cost + + reasoning_cost ) - + return total_cost def get_provider_penalty(provider: "UpstreamProvider") -> float: """Calculate a penalty multiplier for certain providers. - + This allows applying policy-based adjustments beyond pure cost. For example, preferring certain providers for reliability or features. - + Args: provider: UpstreamProvider instance - + Returns: Float multiplier to apply to cost (1.0 = no penalty, >1.0 = penalize) """ # Default: no penalty penalty = 1.0 - + # Check if this is OpenRouter (can be identified by base URL) base_url = getattr(provider, "base_url", "") if "openrouter.ai" in base_url.lower(): # Small penalty for OpenRouter to prefer other providers when costs are very close # This maintains the original behavior of preferring non-OpenRouter providers penalty = 1.001 # 0.1% penalty - + return penalty @@ -90,30 +92,30 @@ def should_prefer_model( alias: str, ) -> bool: """Determine if candidate model should replace current model for an alias. - + This is the core decision function for model prioritization. It considers: 1. Alias matching quality (exact match vs. canonical slug match) 2. Model cost (lower is better) 3. Provider penalties (e.g., slight preference against OpenRouter) - + Args: candidate_model: The new model being considered candidate_provider: Provider offering the candidate model current_model: The currently selected model for this alias current_provider: Provider offering the current model alias: The model alias being mapped - + Returns: True if candidate should replace current, False otherwise """ - + def get_base_model_id(model_id: str) -> str: """Get base model ID by removing provider prefix.""" return model_id.split("/", 1)[1] if "/" in model_id else model_id - + def alias_priority(model: "Model") -> int: """Rank how strong the mapping of alias->model is. - + Highest priority when alias exactly equals the model ID without provider prefix. Next when alias equals canonical slug without prefix. Otherwise lowest. """ @@ -125,39 +127,41 @@ def should_prefer_model( if canonical_base == alias: return 2 return 1 - + candidate_alias_priority = alias_priority(candidate_model) current_alias_priority = alias_priority(current_model) - + # If candidate has better alias match, prefer it regardless of cost if candidate_alias_priority > current_alias_priority: return True - + # If current has better alias match, keep it regardless of cost if current_alias_priority > candidate_alias_priority: return False - + # Same alias priority - compare costs candidate_cost = calculate_model_cost_score(candidate_model) current_cost = calculate_model_cost_score(current_model) - + # Apply provider penalties candidate_adjusted = candidate_cost * get_provider_penalty(candidate_provider) current_adjusted = current_cost * get_provider_penalty(current_provider) - + # Prefer lower adjusted cost should_replace = candidate_adjusted < current_adjusted - + # Log provider changes when candidate wins if should_replace: - candidate_provider_name = getattr(candidate_provider, "upstream_name", "unknown") + candidate_provider_name = getattr( + candidate_provider, "upstream_name", "unknown" + ) current_provider_name = getattr(current_provider, "upstream_name", "unknown") logger.debug( f"Model selection for alias '{alias}': choosing {candidate_provider_name} " f"(cost: ${candidate_adjusted:.6f}) over {current_provider_name} " f"(cost: ${current_adjusted:.6f})" ) - + return should_replace @@ -167,50 +171,52 @@ def create_model_mappings( disabled_model_ids: set[str], ) -> tuple[dict[str, "Model"], dict[str, "UpstreamProvider"], dict[str, "Model"]]: """Create optimal model mappings based on cost and provider preferences. - + This is the main entry point for the algorithm. It processes all upstream providers and creates three mappings based on cost optimization: - + 1. model_instances: alias -> Model (all model aliases mapped to their Model objects) 2. provider_map: alias -> UpstreamProvider (which provider to use for each alias) 3. unique_models: base_id -> Model (unique models without provider prefixes) - + The algorithm: - Processes non-OpenRouter providers first (they're typically cheaper) - Then processes OpenRouter models (they can still win if cheaper) - For each model alias, uses should_prefer_model() to select the best provider - + Args: upstreams: List of all upstream provider instances overrides_by_id: Dict of model overrides from database {model_id: (ModelRow, fee)} disabled_model_ids: Set of model IDs that should be excluded - + Returns: Tuple of (model_instances, provider_map, unique_models) """ from .payment.models import _row_to_model from .upstream import resolve_model_alias - + model_instances: dict[str, "Model"] = {} provider_map: dict[str, "UpstreamProvider"] = {} unique_models: dict[str, "Model"] = {} - + # Separate OpenRouter from other providers openrouter: "UpstreamProvider" | None = None other_upstreams: list["UpstreamProvider"] = [] - + for upstream in upstreams: base_url = getattr(upstream, "base_url", "") if base_url == "https://openrouter.ai/api/v1": openrouter = upstream else: other_upstreams.append(upstream) - + def get_base_model_id(model_id: str) -> str: """Get base model ID by removing provider prefix.""" return model_id.split("/", 1)[1] if "/" in model_id else model_id - - def _maybe_set_alias(alias: str, model: "Model", provider: "UpstreamProvider") -> None: + + def _maybe_set_alias( + alias: str, model: "Model", provider: "UpstreamProvider" + ) -> None: """Set alias to model/provider if not set or if new model is preferred.""" existing_model = model_instances.get(alias) if not existing_model: @@ -220,21 +226,22 @@ def create_model_mappings( else: # Check if candidate should replace existing existing_provider = provider_map[alias] - if should_prefer_model(model, provider, existing_model, existing_provider, alias): + if should_prefer_model( + model, provider, existing_model, existing_provider, alias + ): model_instances[alias] = model provider_map[alias] = provider - + def process_provider_models( - upstream: "UpstreamProvider", - is_openrouter: bool = False + upstream: "UpstreamProvider", is_openrouter: bool = False ) -> None: """Process all models from a given provider.""" upstream_prefix = getattr(upstream, "upstream_name", None) - + for model in upstream.get_cached_models(): if not model.enabled or model.id in disabled_model_ids: continue - + # Apply overrides if present if model.id in overrides_by_id: override_row, provider_fee = overrides_by_id[model.id] @@ -243,42 +250,40 @@ def create_model_mappings( ) else: model_to_use = model - + # Add to unique models base_id = get_base_model_id(model_to_use.id) if not is_openrouter or base_id not in unique_models: unique_model = model_to_use.copy(update={"id": base_id}) unique_models[base_id] = unique_model - + # Get all aliases for this model - aliases = resolve_model_alias( - model_to_use.id, model_to_use.canonical_slug - ) - + aliases = resolve_model_alias(model_to_use.id, model_to_use.canonical_slug) + # Add prefixed alias if applicable if upstream_prefix and "/" not in model_to_use.id: prefixed_id = f"{upstream_prefix}/{model_to_use.id}" if prefixed_id not in aliases: aliases.append(prefixed_id) - + # Try to set each alias for alias in aliases: _maybe_set_alias(alias, model_to_use, upstream) - + # Process non-OpenRouter providers first (they're typically cheaper) for upstream in other_upstreams: process_provider_models(upstream, is_openrouter=False) - + # Process OpenRouter last - models only win if they're cheaper or better matched if openrouter: process_provider_models(openrouter, is_openrouter=True) - + # Log provider distribution provider_counts: dict[str, int] = {} for provider in provider_map.values(): provider_name = getattr(provider, "upstream_name", "unknown") provider_counts[provider_name] = provider_counts.get(provider_name, 0) + 1 - + logger.debug( "Created model mappings", extra={ @@ -287,6 +292,5 @@ def create_model_mappings( "provider_distribution": provider_counts, }, ) - - return model_instances, provider_map, unique_models + return model_instances, provider_map, unique_models diff --git a/tests/unit/test_algorithm.py b/tests/unit/test_algorithm.py index ecd89294..22155e79 100644 --- a/tests/unit/test_algorithm.py +++ b/tests/unit/test_algorithm.py @@ -58,7 +58,7 @@ def test_calculate_model_cost_score_basic() -> None: """Test basic cost calculation.""" model = create_test_model("test-model", prompt_price=0.001, completion_price=0.002) cost = calculate_model_cost_score(model) - + # Expected: (1000 tokens * 0.001) + (500 tokens * 0.002) = 0.001 + 0.001 = 0.002 assert cost == 0.002 @@ -72,16 +72,18 @@ def test_calculate_model_cost_score_with_request_fee() -> None: request_price=0.0005, ) cost = calculate_model_cost_score(model) - + # Expected: 0.001 + 0.001 + 0.0005 = 0.0025 assert cost == 0.0025 def test_calculate_model_cost_score_expensive_model() -> None: """Test cost calculation for expensive model.""" - model = create_test_model("expensive-model", prompt_price=0.03, completion_price=0.06) + model = create_test_model( + "expensive-model", prompt_price=0.03, completion_price=0.06 + ) cost = calculate_model_cost_score(model) - + # Expected: (1000 * 0.03) + (500 * 0.06) = 0.03 + 0.03 = 0.06 assert cost == 0.06 @@ -103,16 +105,18 @@ def test_get_provider_penalty_openrouter() -> None: def test_should_prefer_model_cheaper_wins() -> None: """Test that cheaper model is preferred.""" cheap_model = create_test_model("cheap", prompt_price=0.001, completion_price=0.002) - expensive_model = create_test_model("expensive", prompt_price=0.03, completion_price=0.06) - + expensive_model = create_test_model( + "expensive", prompt_price=0.03, completion_price=0.06 + ) + provider1 = create_test_provider("provider1") provider2 = create_test_provider("provider2") - + # Cheaper model should win assert should_prefer_model( cheap_model, provider1, expensive_model, provider2, "test-alias" ) - + # More expensive model should not win assert not should_prefer_model( expensive_model, provider2, cheap_model, provider1, "test-alias" @@ -122,12 +126,16 @@ def test_should_prefer_model_cheaper_wins() -> None: def test_should_prefer_model_exact_match_wins() -> None: """Test that exact alias match beats cheaper price.""" # Make model IDs match the alias differently - exact_match = create_test_model("test-model", prompt_price=0.03, completion_price=0.06) - no_match = create_test_model("other-model", prompt_price=0.001, completion_price=0.002) - + exact_match = create_test_model( + "test-model", prompt_price=0.03, completion_price=0.06 + ) + no_match = create_test_model( + "other-model", prompt_price=0.001, completion_price=0.002 + ) + provider1 = create_test_provider("provider1") provider2 = create_test_provider("provider2") - + # Exact match should win even though it's more expensive assert should_prefer_model( exact_match, provider1, no_match, provider2, "test-model" @@ -138,15 +146,17 @@ def test_should_prefer_model_openrouter_slight_penalty() -> None: """Test that OpenRouter has slight penalty compared to other providers.""" model1 = create_test_model("model1", prompt_price=0.001, completion_price=0.002) model2 = create_test_model("model2", prompt_price=0.001, completion_price=0.002) - + regular_provider = create_test_provider("regular", "http://provider.com") - openrouter_provider = create_test_provider("openrouter", "https://openrouter.ai/api/v1") - + openrouter_provider = create_test_provider( + "openrouter", "https://openrouter.ai/api/v1" + ) + # Regular provider should be preferred over OpenRouter at same cost assert should_prefer_model( model1, regular_provider, model2, openrouter_provider, "test-alias" ) - + # OpenRouter should not replace regular provider at same cost assert not should_prefer_model( model2, openrouter_provider, model1, regular_provider, "test-alias" @@ -155,15 +165,25 @@ def test_should_prefer_model_openrouter_slight_penalty() -> None: def test_should_prefer_model_openrouter_can_win_if_cheaper() -> None: """Test that OpenRouter can still win if significantly cheaper.""" - cheap_model = create_test_model("cheap", prompt_price=0.0001, completion_price=0.0002) - expensive_model = create_test_model("expensive", prompt_price=0.03, completion_price=0.06) - + cheap_model = create_test_model( + "cheap", prompt_price=0.0001, completion_price=0.0002 + ) + expensive_model = create_test_model( + "expensive", prompt_price=0.03, completion_price=0.06 + ) + regular_provider = create_test_provider("regular", "http://provider.com") - openrouter_provider = create_test_provider("openrouter", "https://openrouter.ai/api/v1") - + openrouter_provider = create_test_provider( + "openrouter", "https://openrouter.ai/api/v1" + ) + # OpenRouter should win if it's much cheaper (even with penalty) assert should_prefer_model( - cheap_model, openrouter_provider, expensive_model, regular_provider, "test-alias" + cheap_model, + openrouter_provider, + expensive_model, + regular_provider, + "test-alias", ) @@ -171,12 +191,9 @@ def test_should_prefer_model_same_cost_first_wins() -> None: """Test that when costs are identical, current model is kept.""" model1 = create_test_model("model1", prompt_price=0.001, completion_price=0.002) model2 = create_test_model("model2", prompt_price=0.001, completion_price=0.002) - + provider1 = create_test_provider("provider1") provider2 = create_test_provider("provider2") - - # When costs are equal, should not replace - assert not should_prefer_model( - model2, provider2, model1, provider1, "test-alias" - ) + # When costs are equal, should not replace + assert not should_prefer_model(model2, provider2, model1, provider1, "test-alias")