diff --git a/router/payment/cost_caculation.py b/router/payment/cost_caculation.py index 2ee7195f..6eedaf5c 100644 --- a/router/payment/cost_caculation.py +++ b/router/payment/cost_caculation.py @@ -98,35 +98,20 @@ def calculate_cost( }, ) - model = None - - if response_model in [model.id for model in MODELS]: - model = next(model for model in MODELS if model.id == response_model) - else: - canonical_models = [] - for m in MODELS: - if "/" in m.canonical_slug: - canonical_name = m.canonical_slug.split("/", 1)[1] - canonical_models.append((canonical_name, m)) - - canonical_dict = {name: m for name, m in canonical_models} - - if response_model in canonical_dict: - model = canonical_dict[response_model] - else: - logger.error( - "Invalid model in response", - extra={ - "response_model": response_model, - "available_models": [model.id for model in MODELS], - "available_canonical_models": list(canonical_dict.keys()), - }, - ) - return CostDataError( - message=f"Invalid model in response: {response_model}", - code="model_not_found", - ) + if response_model not in [model.id for model in MODELS]: + logger.error( + "Invalid model in response", + extra={ + "response_model": response_model, + "available_models": [model.id for model in MODELS], + }, + ) + return CostDataError( + message=f"Invalid model in response: {response_model}", + code="model_not_found", + ) + model = next(model for model in MODELS if model.id == response_model) if model.sats_pricing is None: logger.error( "Model pricing not defined", diff --git a/router/payment/models.py b/router/payment/models.py index 6ab3b7d2..7b56b9ff 100644 --- a/router/payment/models.py +++ b/router/payment/models.py @@ -37,7 +37,6 @@ class TopProvider(BaseModel): class Model(BaseModel): id: str - canonical_slug: str name: str created: int description: str diff --git a/tests/test_models.py b/tests/test_models.py index fe704eb8..cb101693 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -19,7 +19,6 @@ def sample_model() -> Model: """Create a sample model for testing.""" return Model( id="test-model", - canonical_slug="test/test-model", name="Test Model", created=1700000000, description="A test model", @@ -122,7 +121,6 @@ async def test_update_sats_pricing_without_top_provider() -> None: """Test sats pricing calculation for models without top_provider.""" model_without_top = Model( id="test-model-no-top", - canonical_slug="test/test-model-no-top", name="Test Model No Top", created=1700000000, description="A test model without top provider",