mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
430 lines
16 KiB
Python
430 lines
16 KiB
Python
import json
|
|
import os
|
|
|
|
import pytest
|
|
|
|
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from base_embedding_unit_tests import BaseLLMEmbeddingTest
|
|
|
|
import litellm
|
|
|
|
|
|
class TestVoyageAI(BaseLLMEmbeddingTest):
|
|
def get_custom_llm_provider(self) -> litellm.LlmProviders:
|
|
return litellm.LlmProviders.VOYAGE
|
|
|
|
def get_base_embedding_call_args(self) -> dict:
|
|
return {
|
|
"model": "voyage/voyage-3-lite",
|
|
}
|
|
|
|
@pytest.mark.asyncio()
|
|
@pytest.mark.parametrize("sync_mode", [True, False])
|
|
async def test_basic_embedding(self, sync_mode):
|
|
"""Override base test to handle Voyage embeddings properly"""
|
|
litellm.set_verbose = True
|
|
embedding_call_args = self.get_base_embedding_call_args()
|
|
|
|
# Mock the embedding function to avoid API calls
|
|
with (
|
|
patch("litellm.embedding") as mock_embedding,
|
|
patch("litellm.aembedding") as mock_aembedding,
|
|
):
|
|
# Create a mock response that matches Voyage format
|
|
mock_response = MagicMock()
|
|
mock_response.model = "voyage-3-lite"
|
|
mock_response.object = "list"
|
|
mock_response.data = [
|
|
{"object": "embedding", "embedding": [0.1, 0.2, 0.3], "index": 0}
|
|
]
|
|
mock_response.usage.prompt_tokens = 24
|
|
mock_response.usage.total_tokens = 24
|
|
|
|
mock_embedding.return_value = mock_response
|
|
mock_aembedding.return_value = mock_response
|
|
|
|
if sync_mode is True:
|
|
response = litellm.embedding(
|
|
**embedding_call_args,
|
|
input=["hello", "world"],
|
|
)
|
|
# Verify the response structure
|
|
assert response.model == "voyage-3-lite"
|
|
assert response.object == "list"
|
|
assert len(response.data) > 0
|
|
assert response.usage.total_tokens > 0
|
|
else:
|
|
response = await litellm.aembedding(
|
|
**embedding_call_args,
|
|
input=["hello", "world"],
|
|
)
|
|
# Verify the response structure
|
|
assert response.model == "voyage-3-lite"
|
|
assert response.object == "list"
|
|
assert len(response.data) > 0
|
|
assert response.usage.total_tokens > 0
|
|
|
|
|
|
def test_voyage_ai_embedding_extra_params():
|
|
"""Test Voyage AI embedding with extra parameters"""
|
|
try:
|
|
# Mock the entire embedding function to avoid API calls
|
|
with patch("litellm.embedding") as mock_embedding:
|
|
# Create a mock response
|
|
mock_response = MagicMock()
|
|
mock_response.usage.prompt_tokens = 24
|
|
mock_response.usage.total_tokens = 24
|
|
mock_response.model = "voyage-3-lite"
|
|
mock_embedding.return_value = mock_response
|
|
|
|
litellm.embedding(
|
|
model="voyage/voyage-3-lite",
|
|
input=["a"],
|
|
dimensions=512,
|
|
input_type="document",
|
|
)
|
|
|
|
# Verify the function was called with correct parameters
|
|
mock_embedding.assert_called_once()
|
|
call_args = mock_embedding.call_args
|
|
assert call_args[1]["model"] == "voyage/voyage-3-lite"
|
|
assert call_args[1]["input"] == ["a"]
|
|
assert call_args[1]["dimensions"] == 512
|
|
assert call_args[1]["input_type"] == "document"
|
|
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
def test_voyage_ai_embedding_prompt_token_mapping():
|
|
"""Test Voyage AI embedding token mapping"""
|
|
try:
|
|
# Mock the entire embedding function
|
|
with patch("litellm.embedding") as mock_embedding:
|
|
# Create a mock response with usage
|
|
mock_response = MagicMock()
|
|
mock_response.usage.prompt_tokens = 120
|
|
mock_response.usage.total_tokens = 120
|
|
mock_embedding.return_value = mock_response
|
|
|
|
response = litellm.embedding(
|
|
model="voyage/voyage-3-lite",
|
|
input=["a"],
|
|
dimensions=512,
|
|
input_type="document",
|
|
)
|
|
|
|
# Verify the response
|
|
assert response.usage.prompt_tokens == 120
|
|
assert response.usage.total_tokens == 120
|
|
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
# Tests for Voyage Contextual Embeddings
|
|
class TestVoyageContextualEmbeddings:
|
|
"""Test suite for Voyage contextual embeddings functionality"""
|
|
|
|
def test_contextual_embedding_model_detection(self):
|
|
"""Test that contextual models are correctly identified"""
|
|
from litellm.llms.voyage.embedding.transformation_contextual import (
|
|
VoyageContextualEmbeddingConfig,
|
|
)
|
|
|
|
config = VoyageContextualEmbeddingConfig()
|
|
|
|
# Test contextual model detection
|
|
assert config.is_contextualized_embeddings("voyage-context-3") is True
|
|
assert config.is_contextualized_embeddings("voyage-context-2") is True
|
|
assert config.is_contextualized_embeddings("context-model") is True
|
|
|
|
# Test regular model detection
|
|
assert config.is_contextualized_embeddings("voyage-3-lite") is False
|
|
assert config.is_contextualized_embeddings("voyage-2") is False
|
|
assert config.is_contextualized_embeddings("regular-model") is False
|
|
|
|
def test_contextual_embedding_url_generation(self):
|
|
"""Test URL generation for contextual embeddings"""
|
|
from litellm.llms.voyage.embedding.transformation_contextual import (
|
|
VoyageContextualEmbeddingConfig,
|
|
)
|
|
|
|
config = VoyageContextualEmbeddingConfig()
|
|
|
|
# Test default URL
|
|
url = config.get_complete_url(None, None, "voyage-context-3", {}, {})
|
|
assert url == "https://api.voyageai.com/v1/contextualizedembeddings"
|
|
|
|
# Test custom API base
|
|
url = config.get_complete_url(
|
|
"https://custom.api.com", None, "voyage-context-3", {}, {}
|
|
)
|
|
assert url == "https://custom.api.com/contextualizedembeddings"
|
|
|
|
# Test API base that already ends with endpoint
|
|
url = config.get_complete_url(
|
|
"https://custom.api.com/contextualizedembeddings",
|
|
None,
|
|
"voyage-context-3",
|
|
{},
|
|
{},
|
|
)
|
|
assert url == "https://custom.api.com/contextualizedembeddings"
|
|
|
|
def test_contextual_embedding_request_transformation(self):
|
|
"""Test request transformation for contextual embeddings"""
|
|
from litellm.llms.voyage.embedding.transformation_contextual import (
|
|
VoyageContextualEmbeddingConfig,
|
|
)
|
|
|
|
config = VoyageContextualEmbeddingConfig()
|
|
|
|
# Test with nested input structure
|
|
input_data = [["Hello", "world"], ["Test", "sentence"]]
|
|
optional_params = {"encoding_format": "float"}
|
|
|
|
transformed = config.transform_embedding_request(
|
|
"voyage-context-3", input_data, optional_params, {}
|
|
)
|
|
|
|
assert transformed["inputs"] == input_data
|
|
assert transformed["model"] == "voyage-context-3"
|
|
assert transformed["encoding_format"] == "float"
|
|
|
|
def test_contextual_embedding_response_transformation(self):
|
|
"""Test response transformation for contextual embeddings"""
|
|
from litellm.llms.voyage.embedding.transformation_contextual import (
|
|
VoyageContextualEmbeddingConfig,
|
|
)
|
|
from litellm.types.utils import EmbeddingResponse
|
|
|
|
config = VoyageContextualEmbeddingConfig()
|
|
|
|
# Mock the nested response structure from Voyage contextual embeddings
|
|
mock_response_data = {
|
|
"object": "list",
|
|
"data": [
|
|
{
|
|
"object": "list",
|
|
"data": [
|
|
{
|
|
"object": "embedding",
|
|
"embedding": [0.1, 0.2, 0.3],
|
|
"index": 0,
|
|
}
|
|
],
|
|
"index": 0,
|
|
}
|
|
],
|
|
"model": "voyage-context-3",
|
|
"usage": {"total_tokens": 24},
|
|
}
|
|
|
|
# Create mock response
|
|
mock_response = MagicMock()
|
|
mock_response.json.return_value = mock_response_data
|
|
mock_response.status_code = 200
|
|
mock_response.text = json.dumps(mock_response_data)
|
|
|
|
# Create model response
|
|
model_response = EmbeddingResponse()
|
|
|
|
# Transform response
|
|
transformed = config.transform_embedding_response(
|
|
"voyage-context-3", mock_response, model_response, MagicMock()
|
|
)
|
|
|
|
# Assert the transformation preserves the nested structure
|
|
assert transformed.model == "voyage-context-3"
|
|
assert transformed.object == "list"
|
|
assert transformed.data == mock_response_data["data"]
|
|
assert transformed.usage.prompt_tokens == 24
|
|
assert transformed.usage.total_tokens == 24
|
|
|
|
def test_contextual_embedding_parameter_mapping(self):
|
|
"""Test parameter mapping for contextual embeddings"""
|
|
from litellm.llms.voyage.embedding.transformation_contextual import (
|
|
VoyageContextualEmbeddingConfig,
|
|
)
|
|
|
|
config = VoyageContextualEmbeddingConfig()
|
|
|
|
non_default_params = {"encoding_format": "float", "dimensions": 512}
|
|
optional_params = {}
|
|
|
|
mapped = config.map_openai_params(
|
|
non_default_params, optional_params, "voyage-context-3", False
|
|
)
|
|
|
|
assert mapped["encoding_format"] == "float"
|
|
assert mapped["output_dimension"] == 512
|
|
|
|
def test_contextual_embedding_environment_validation(self):
|
|
"""Test environment validation for contextual embeddings"""
|
|
from litellm.llms.voyage.embedding.transformation_contextual import (
|
|
VoyageContextualEmbeddingConfig,
|
|
)
|
|
|
|
config = VoyageContextualEmbeddingConfig()
|
|
|
|
# Test with API key in environment
|
|
os.environ["VOYAGE_API_KEY"] = "test-key"
|
|
|
|
headers = config.validate_environment({}, "voyage-context-3", [], {}, {})
|
|
assert headers["Authorization"] == "Bearer test-key"
|
|
|
|
# Test with custom API key
|
|
headers = config.validate_environment(
|
|
{}, "voyage-context-3", [], {}, {}, api_key="custom-key"
|
|
)
|
|
assert headers["Authorization"] == "Bearer custom-key"
|
|
|
|
def test_contextual_embedding_error_handling(self):
|
|
"""Test error handling for contextual embeddings"""
|
|
from litellm.llms.voyage.embedding.transformation_contextual import (
|
|
VoyageContextualEmbeddingConfig,
|
|
VoyageError,
|
|
)
|
|
|
|
config = VoyageContextualEmbeddingConfig()
|
|
|
|
# Test error class creation
|
|
error = config.get_error_class("Test error", 400, {})
|
|
assert isinstance(error, VoyageError)
|
|
assert error.status_code == 400
|
|
assert error.message == "Test error"
|
|
|
|
def test_contextual_vs_regular_embedding_differences(self):
|
|
"""Test that contextual and regular embeddings are handled differently"""
|
|
from litellm.llms.voyage.embedding.transformation import VoyageEmbeddingConfig
|
|
from litellm.llms.voyage.embedding.transformation_contextual import (
|
|
VoyageContextualEmbeddingConfig,
|
|
)
|
|
|
|
regular_config = VoyageEmbeddingConfig()
|
|
contextual_config = VoyageContextualEmbeddingConfig()
|
|
|
|
# Test URL differences
|
|
regular_url = regular_config.get_complete_url(
|
|
None, None, "voyage-3-lite", {}, {}
|
|
)
|
|
contextual_url = contextual_config.get_complete_url(
|
|
None, None, "voyage-context-3", {}, {}
|
|
)
|
|
|
|
assert regular_url == "https://api.voyageai.com/v1/embeddings"
|
|
assert contextual_url == "https://api.voyageai.com/v1/contextualizedembeddings"
|
|
|
|
# Test request transformation differences
|
|
regular_transformed = regular_config.transform_embedding_request(
|
|
"voyage-3-lite", ["Hello"], {}, {}
|
|
)
|
|
contextual_transformed = contextual_config.transform_embedding_request(
|
|
"voyage-context-3", [["Hello"]], {}, {}
|
|
)
|
|
|
|
assert regular_transformed["input"] == ["Hello"]
|
|
assert contextual_transformed["inputs"] == [["Hello"]]
|
|
|
|
def test_contextual_embedding_integration(self):
|
|
"""Test full integration of contextual embeddings"""
|
|
try:
|
|
# Mock the entire embedding function to avoid API calls
|
|
with patch("litellm.embedding") as mock_embedding:
|
|
# Create a mock response that matches the expected structure
|
|
mock_response = MagicMock()
|
|
mock_response.model = "voyage-context-3"
|
|
mock_response.usage.total_tokens = 24
|
|
mock_response.data = [
|
|
{
|
|
"object": "list",
|
|
"data": [
|
|
{
|
|
"object": "embedding",
|
|
"embedding": [0.1, 0.2, 0.3],
|
|
"index": 0,
|
|
}
|
|
],
|
|
"index": 0,
|
|
}
|
|
]
|
|
mock_embedding.return_value = mock_response
|
|
|
|
response = litellm.embedding(
|
|
model="voyage/voyage-context-3",
|
|
input=[["Hello", "world"]],
|
|
input_type="document",
|
|
)
|
|
|
|
# Verify the function was called with correct parameters
|
|
mock_embedding.assert_called_once()
|
|
call_args = mock_embedding.call_args
|
|
assert call_args[1]["model"] == "voyage/voyage-context-3"
|
|
assert call_args[1]["input"] == [["Hello", "world"]]
|
|
assert call_args[1]["input_type"] == "document"
|
|
|
|
# Assert the response structure
|
|
assert response.model == "voyage-context-3"
|
|
assert response.usage.total_tokens == 24
|
|
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
def test_contextual_embedding_multiple_inputs(self):
|
|
"""Test contextual embeddings with multiple input groups"""
|
|
try:
|
|
# Mock the entire embedding function
|
|
with patch("litellm.embedding") as mock_embedding:
|
|
# Create a mock response for multiple input groups
|
|
mock_response = MagicMock()
|
|
mock_response.model = "voyage-context-3"
|
|
mock_response.usage.total_tokens = 48
|
|
mock_response.data = [
|
|
{
|
|
"object": "list",
|
|
"data": [
|
|
{
|
|
"object": "embedding",
|
|
"embedding": [0.1, 0.2],
|
|
"index": 0,
|
|
},
|
|
{
|
|
"object": "embedding",
|
|
"embedding": [0.3, 0.4],
|
|
"index": 1,
|
|
},
|
|
],
|
|
"index": 0,
|
|
},
|
|
{
|
|
"object": "list",
|
|
"data": [
|
|
{"object": "embedding", "embedding": [0.5, 0.6], "index": 0}
|
|
],
|
|
"index": 1,
|
|
},
|
|
]
|
|
mock_embedding.return_value = mock_response
|
|
|
|
response = litellm.embedding(
|
|
model="voyage/voyage-context-3",
|
|
input=[["Hello", "world"], ["Test"]],
|
|
)
|
|
|
|
# Verify the function was called with correct parameters
|
|
mock_embedding.assert_called_once()
|
|
call_args = mock_embedding.call_args
|
|
assert call_args[1]["model"] == "voyage/voyage-context-3"
|
|
assert call_args[1]["input"] == [["Hello", "world"], ["Test"]]
|
|
|
|
# Assert response structure
|
|
assert len(response.data) == 2
|
|
assert response.data[0]["index"] == 0
|
|
assert response.data[1]["index"] == 1
|
|
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|