mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
test(custom_openai): cover model discovery branches
This commit is contained in:
parent
cf1445448b
commit
db61d672b1
2 changed files with 57 additions and 37 deletions
|
|
@ -1656,42 +1656,6 @@ def test_get_valid_models_openai_proxy(monkeypatch):
|
|||
assert "litellm_proxy/gpt-5.5" in valid_models
|
||||
|
||||
|
||||
def test_get_valid_models_custom_openai(monkeypatch):
|
||||
from litellm.utils import get_valid_models
|
||||
import litellm
|
||||
|
||||
mock_response_data = {
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": "my-custom-model",
|
||||
"object": "model",
|
||||
"created": 1686935002,
|
||||
"owned_by": "organization-owner",
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = mock_response_data
|
||||
|
||||
with patch.object(litellm.module_level_client, "get", return_value=mock_response):
|
||||
valid_models = get_valid_models(
|
||||
check_provider_endpoint=True,
|
||||
custom_llm_provider="custom_openai",
|
||||
api_key="sk-1234",
|
||||
api_base="https://my-openai-compatible-endpoint/v1",
|
||||
)
|
||||
assert "custom_openai/my-custom-model" in valid_models
|
||||
|
||||
|
||||
def test_custom_openai_get_models_requires_api_base():
|
||||
from litellm.llms.openai_like.chat.transformation import CustomOpenAIChatConfig
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
CustomOpenAIChatConfig().get_models(api_base=None)
|
||||
|
||||
|
||||
def test_get_valid_models_fireworks_ai(monkeypatch):
|
||||
from litellm.utils import get_valid_models
|
||||
|
|
|
|||
|
|
@ -1,5 +1,14 @@
|
|||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import litellm
|
||||
import pytest
|
||||
from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig
|
||||
from litellm.llms.openai_like.chat.transformation import (
|
||||
CustomOpenAIChatConfig,
|
||||
OpenAILikeChatConfig,
|
||||
)
|
||||
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.utils import ProviderConfigManager, get_valid_models
|
||||
|
||||
|
||||
def test_sanitize_usage_obj_handles_null_tokens():
|
||||
|
|
@ -47,3 +56,50 @@ def test_sanitize_usage_obj_valid_usage():
|
|||
|
||||
# Assert
|
||||
assert sanitized_json == original_json # The object should be unchanged
|
||||
|
||||
|
||||
def test_get_valid_models_custom_openai():
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": "my-custom-model",
|
||||
"object": "model",
|
||||
"created": 1686935002,
|
||||
"owned_by": "organization-owner",
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
with patch.object(litellm.module_level_client, "get", return_value=mock_response) as mock_get:
|
||||
valid_models = get_valid_models(
|
||||
check_provider_endpoint=True,
|
||||
custom_llm_provider="custom_openai",
|
||||
api_key="sk-1234",
|
||||
api_base="https://my-openai-compatible-endpoint/v1",
|
||||
)
|
||||
|
||||
assert valid_models == ["custom_openai/my-custom-model"]
|
||||
mock_get.assert_called_once_with(
|
||||
url="https://my-openai-compatible-endpoint/v1/models",
|
||||
headers={"Authorization": "Bearer sk-1234"},
|
||||
)
|
||||
|
||||
|
||||
def test_custom_openai_get_models_requires_api_base():
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="api_base must be set to discover models for the custom_openai provider",
|
||||
):
|
||||
CustomOpenAIChatConfig().get_models(api_base=None)
|
||||
|
||||
|
||||
def test_openai_model_info_uses_openai_config():
|
||||
config = ProviderConfigManager.get_provider_model_info(
|
||||
model=None,
|
||||
provider=LlmProviders.OPENAI,
|
||||
)
|
||||
|
||||
assert type(config) is OpenAIGPTConfig
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue