mirror of
https://github.com/tiennm99/litellm.git
synced 2026-06-18 00:48:01 +00:00
yeahyung
a92bf8173e
* (#13284) add avector_store_create to route_type which doesn't require model * (#13284) exclude hidden params in metadata when create vector store * (#13284) fix lint error * (#13284) keep metadata None if metadata is None(not empty dict) * (#13284) add test code * (#13284) change test code name * (#13284) add avector_store_search to route_type which doesn't require model
127 lines
4.2 KiB
Python
127 lines
4.2 KiB
Python
import json
|
|
import os
|
|
import sys
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../../..")
|
|
) # Adds the parent directory to the system path
|
|
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.proxy.route_llm_request import route_request
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"route_type",
|
|
[
|
|
"atext_completion",
|
|
"acompletion",
|
|
"aembedding",
|
|
"aimage_generation",
|
|
"aspeech",
|
|
"atranscription",
|
|
"amoderation",
|
|
"arerank",
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_route_request_dynamic_credentials(route_type):
|
|
data = {
|
|
"model": "openai/gpt-4o-mini-2024-07-18",
|
|
"api_key": "my-bad-key",
|
|
"api_base": "https://api.openai.com/v1 ",
|
|
}
|
|
llm_router = MagicMock()
|
|
# Ensure that the dynamic method exists on the llm_router mock.
|
|
getattr(llm_router, route_type).return_value = "fake_response"
|
|
|
|
response = await route_request(data, llm_router, None, route_type)
|
|
# Optionally verify the response if needed:
|
|
assert response == "fake_response"
|
|
# Now assert that the dynamic method was called once with the expected kwargs.
|
|
getattr(llm_router, route_type).assert_called_once_with(**data)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_route_request_no_model_required():
|
|
"""Test route types that don't require model parameter"""
|
|
test_cases = ["amoderation", "aget_responses", "adelete_responses", "avector_store_create", "avector_store_search"]
|
|
|
|
for route_type in test_cases:
|
|
# Test data without model parameter
|
|
data = {"input": "test input", "api_key": "test-key"}
|
|
|
|
llm_router = MagicMock()
|
|
getattr(llm_router, route_type).return_value = "fake_response"
|
|
|
|
response = await route_request(data, llm_router, None, route_type)
|
|
|
|
# Verify response
|
|
assert response == "fake_response"
|
|
# Verify the method was called with correct parameters
|
|
getattr(llm_router, route_type).assert_called_once_with(**data)
|
|
|
|
# Reset mock for next iteration
|
|
llm_router.reset_mock()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_route_request_no_model_required_with_router_settings():
|
|
"""Test route types that don't require model parameter with router settings"""
|
|
test_cases = ["amoderation", "aget_responses", "adelete_responses", "avector_store_create", "avector_store_search"]
|
|
|
|
for route_type in test_cases:
|
|
# Test data with model parameter (it will be ignored for these route types)
|
|
data = {
|
|
"input": "test input",
|
|
"model": "test-model", # Include dummy model to avoid KeyError
|
|
}
|
|
|
|
llm_router = MagicMock()
|
|
# Set up router settings
|
|
llm_router.router_general_settings.pass_through_all_models = False
|
|
llm_router.default_deployment = None
|
|
llm_router.pattern_router.patterns = []
|
|
llm_router.model_names = [] # Empty model names list
|
|
llm_router.get_model_ids.return_value = [] # Empty model IDs
|
|
llm_router.model_group_alias = None # No model group alias
|
|
|
|
# Mock the async route call
|
|
getattr(llm_router, route_type).return_value = "fake_response"
|
|
|
|
# Run the request
|
|
response = await route_request(data, llm_router, None, route_type)
|
|
|
|
# Assert the mocked method was called with expected input
|
|
assert response == "fake_response"
|
|
getattr(llm_router, route_type).assert_called_once_with(**data)
|
|
|
|
# Reset the mock for the next route
|
|
llm_router.reset_mock()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_route_request_no_model_required_with_router_settings_and_no_router():
|
|
"""Test route types that don't require model parameter with router settings and no router"""
|
|
from unittest.mock import patch
|
|
|
|
import litellm
|
|
from litellm.proxy.route_llm_request import route_request
|
|
|
|
data = {
|
|
"model": "my-model-id",
|
|
"api_key": "my-api-key",
|
|
"messages": [{"role": "user", "content": "what llm are you"}],
|
|
}
|
|
|
|
with patch.object(
|
|
litellm, "acompletion", return_value="fake_response"
|
|
) as mock_completion:
|
|
response = await route_request(data, None, "gpt-3.5-turbo", "acompletion")
|
|
|
|
mock_completion.assert_called_once_with(**data)
|