test(cost): type the off-peak fixture helpers with OffPeakPricing

This commit is contained in:
mateo-berri 2026-09-03 15:03:10 -07:00
parent e65e3d0e2b
commit f9e41470d6
2 changed files with 4 additions and 3 deletions

View file

@ -7,7 +7,7 @@ import pytest
import litellm
from litellm.llms.fireworks_ai.cost_calculator import cost_per_token
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
from litellm.types.utils import OffPeakPricing, PromptTokensDetailsWrapper, Usage
MODEL = "accounts/fireworks/models/glm-5p2"
INPUT_COST = 1.4e-06
@ -78,7 +78,7 @@ STANDARD_OUTPUT_COST = 6e-07
STANDARD_CACHE_READ_COST = 1.5e-08
def _register_off_peak_model(off_peak_pricing: dict, cache_read_cost: float | None = STANDARD_CACHE_READ_COST) -> None:
def _register_off_peak_model(off_peak_pricing: OffPeakPricing, cache_read_cost: float | None = STANDARD_CACHE_READ_COST) -> None:
litellm.model_cost[f"fireworks_ai/{OFF_PEAK_MODEL}"] = {
"litellm_provider": "fireworks_ai",
"mode": "chat",

View file

@ -22,6 +22,7 @@ from litellm.llms.perplexity.cost_calculator import (
)
from litellm.types.utils import (
CompletionTokensDetailsWrapper,
OffPeakPricing,
Usage,
PromptTokensDetailsWrapper,
)
@ -530,7 +531,7 @@ class TestPerplexityCostCalculator:
INSIDE_WINDOW = datetime(2026, 9, 3, 17, 25, tzinfo=timezone.utc)
OUTSIDE_WINDOW = datetime(2026, 9, 3, 9, 0, tzinfo=timezone.utc)
def _register_off_peak_model(self, off_peak_pricing: dict) -> None:
def _register_off_peak_model(self, off_peak_pricing: OffPeakPricing) -> None:
litellm.model_cost[f"perplexity/{self.OFF_PEAK_MODEL}"] = {
"litellm_provider": "perplexity",
"mode": "chat",