mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +00:00
d2ac51893b
Regenerated every touched file from origin/main applying only the B1 test deletions and the unused import and helper cleanup they leave behind, without running the formatter across untouched code. CI only checks ruff format under litellm/, so the earlier reflows of test files were pure diff noise for reviewers Also drops the tests/local_testing/test_prompt_caching.py entry from the caching-local shard in test-unit.yml since that file is deleted Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
104 lines
3.5 KiB
Python
104 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
|
|
|
|
|