mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(cost-map): keep register_model url fetch to a single attempt
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
aa0a9ab3ea
commit
536a85b429
2 changed files with 24 additions and 1 deletions
|
|
@ -3079,7 +3079,7 @@ def register_model(
|
|||
# Convert stringified numbers to appropriate numeric types
|
||||
loaded_model_cost = model_cost
|
||||
elif isinstance(model_cost, str):
|
||||
loaded_model_cost = litellm.get_model_cost_map(url=model_cost)
|
||||
loaded_model_cost = litellm.get_model_cost_map(url=model_cost, max_attempts=1)
|
||||
|
||||
if persist_across_reloads:
|
||||
_registrations: Final[Mapping[str, Mapping[str, object]]] = loaded_model_cost
|
||||
|
|
|
|||
|
|
@ -2,10 +2,12 @@ import asyncio
|
|||
import json
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from jsonschema import validate
|
||||
|
|
@ -2382,6 +2384,27 @@ def test_register_model_with_scientific_notation():
|
|||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_register_model_url_fetch_uses_single_attempt(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost))
|
||||
before = dict(litellm.model_cost)
|
||||
threads_before = {thread.name for thread in threading.enumerate()}
|
||||
route = respx.get("https://example.invalid/custom_pricing.json").mock(
|
||||
return_value=httpx.Response(503)
|
||||
)
|
||||
|
||||
litellm.register_model(model_cost="https://example.invalid/custom_pricing.json")
|
||||
|
||||
threads_after = {thread.name for thread in threading.enumerate()}
|
||||
assert route.call_count == 1
|
||||
assert not (threads_after - threads_before) & {"litellm-model-cost-map-retry"}
|
||||
assert not any(
|
||||
thread.name == "litellm-model-cost-map-retry" and thread.is_alive()
|
||||
for thread in threading.enumerate()
|
||||
)
|
||||
assert litellm.model_cost.keys() >= before.keys()
|
||||
|
||||
|
||||
def test_register_model_openrouter_without_slash():
|
||||
"""
|
||||
Test that register_model handles openrouter models without '/' in the name.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue