mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
e2971e0af4
* refactor(litellm): migrate private usage in llms Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(litellm): preserve Bedrock batch signature marker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(litellm): migrate llms private usage symbols Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(litellm): preserve UUID Watsonx project IDs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(litellm): retarget llms mocks to public names Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(litellm): limit llms changes to renames and forwarders Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(litellm): keep GCS mock client patching private Vertex auth methods Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(litellm): assert forwarder arguments and type forwarder signatures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
102 lines
3.5 KiB
Python
102 lines
3.5 KiB
Python
"""
|
|
Tests for Lambda AI provider integration
|
|
"""
|
|
|
|
import os
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.llms.lambda_ai.chat.transformation import LambdaAIChatConfig
|
|
|
|
|
|
def test_lambda_ai_config_initialization():
|
|
"""Test LambdaAIChatConfig initializes correctly"""
|
|
config = LambdaAIChatConfig()
|
|
assert config.custom_llm_provider == "lambda_ai"
|
|
|
|
|
|
def test_lambda_ai_get_openai_compatible_provider_info():
|
|
"""Test Lambda AI provider info retrieval"""
|
|
config = LambdaAIChatConfig()
|
|
|
|
# Test with default values (no env vars set)
|
|
with mock.patch.dict(os.environ, {}, clear=True):
|
|
api_base, api_key = config.get_openai_compatible_provider_info(None, None)
|
|
assert api_base == "https://api.lambda.ai/v1"
|
|
assert api_key is None
|
|
|
|
# Test with environment variables
|
|
with mock.patch.dict(
|
|
os.environ,
|
|
{
|
|
"LAMBDA_API_KEY": "test-key",
|
|
"LAMBDA_API_BASE": "https://custom.lambda.ai/v1",
|
|
},
|
|
):
|
|
api_base, api_key = config.get_openai_compatible_provider_info(None, None)
|
|
assert api_base == "https://custom.lambda.ai/v1"
|
|
assert api_key == "test-key"
|
|
|
|
# Test with explicit parameters (should override env vars)
|
|
with mock.patch.dict(
|
|
os.environ,
|
|
{"LAMBDA_API_KEY": "env-key", "LAMBDA_API_BASE": "https://env.lambda.ai/v1"},
|
|
):
|
|
api_base, api_key = config.get_openai_compatible_provider_info("https://param.lambda.ai/v1", "param-key")
|
|
assert api_base == "https://param.lambda.ai/v1"
|
|
assert api_key == "param-key"
|
|
|
|
|
|
def test_get_llm_provider_lambda_ai():
|
|
"""Test that get_llm_provider correctly identifies Lambda AI"""
|
|
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
|
|
|
# Test with lambda_ai/model-name format
|
|
model, provider, api_key, api_base = get_llm_provider(
|
|
"lambda_ai/llama3.1-8b-instruct"
|
|
)
|
|
assert model == "llama3.1-8b-instruct"
|
|
assert provider == "lambda_ai"
|
|
|
|
# Test with api_base containing Lambda AI endpoint
|
|
model, provider, api_key, api_base = get_llm_provider(
|
|
"llama3.1-8b-instruct", api_base="https://api.lambda.ai/v1"
|
|
)
|
|
assert model == "llama3.1-8b-instruct"
|
|
assert provider == "lambda_ai"
|
|
assert api_base == "https://api.lambda.ai/v1"
|
|
|
|
|
|
def test_lambda_ai_in_provider_lists():
|
|
"""Test that Lambda AI is registered in all necessary provider lists"""
|
|
assert "lambda_ai" in litellm.openai_compatible_providers
|
|
assert "lambda_ai" in litellm.provider_list
|
|
assert "https://api.lambda.ai/v1" in litellm.openai_compatible_endpoints
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lambda_ai_completion_call():
|
|
"""Test completion call with Lambda AI provider (requires LAMBDA_API_KEY)"""
|
|
# Skip if no API key is available
|
|
if not os.getenv("LAMBDA_API_KEY"):
|
|
pytest.skip("LAMBDA_API_KEY not set")
|
|
|
|
try:
|
|
response = await litellm.acompletion(
|
|
model="lambda_ai/llama3.1-8b-instruct",
|
|
messages=[{"role": "user", "content": "Hello, this is a test"}],
|
|
max_tokens=10,
|
|
)
|
|
assert response.choices[0].message.content
|
|
assert response.model
|
|
assert response.usage
|
|
except Exception as e:
|
|
# If the API key is invalid or there's a network issue, that's okay
|
|
# The important thing is that the provider was recognized
|
|
if "lambda_ai" not in str(e) and "provider" not in str(e).lower():
|
|
# Re-raise if it's not a provider-related error
|
|
raise
|
|
|
|
|