mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-02 18:21:56 +00:00
Add support for voyage-context-3 embedding model
This commit is contained in:
Sameer Kankute
@@ -1151,6 +1151,7 @@ from .llms.topaz.image_variations.transformation import TopazImageVariationConfi
|
||||
from litellm.llms.openai.completion.transformation import OpenAITextCompletionConfig
|
||||
from .llms.groq.chat.transformation import GroqChatConfig
|
||||
from .llms.voyage.embedding.transformation import VoyageEmbeddingConfig
|
||||
from .llms.voyage.embedding.transformation_contextual import VoyageContextualEmbeddingConfig
|
||||
from .llms.infinity.embedding.transformation import InfinityEmbeddingConfig
|
||||
from .llms.azure_ai.chat.transformation import AzureAIStudioConfig
|
||||
from .llms.mistral.chat.transformation import MistralConfig
|
||||
|
||||
@@ -0,0 +1,153 @@
|
||||
"""
|
||||
This module is used to transform the request and response for the Voyage contextualized embeddings API.
|
||||
This would be used for all the contextualized embeddings models in Voyage.
|
||||
"""
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues
|
||||
from litellm.types.utils import EmbeddingResponse, Usage
|
||||
|
||||
|
||||
class VoyageError(BaseLLMException):
|
||||
def __init__(
|
||||
self,
|
||||
status_code: int,
|
||||
message: str,
|
||||
headers: Union[dict, httpx.Headers] = {},
|
||||
):
|
||||
self.status_code = status_code
|
||||
self.message = message
|
||||
self.request = httpx.Request(
|
||||
method="POST", url="https://api.voyageai.com/v1/contextualizedembeddings"
|
||||
)
|
||||
self.response = httpx.Response(status_code=status_code, request=self.request)
|
||||
super().__init__(
|
||||
status_code=status_code,
|
||||
message=message,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
|
||||
class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
|
||||
"""
|
||||
Reference: https://docs.voyageai.com/reference/embeddings-api
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
api_key: Optional[str],
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
stream: Optional[bool] = None,
|
||||
) -> str:
|
||||
if api_base:
|
||||
if not api_base.endswith("/contextualizedembeddings"):
|
||||
api_base = f"{api_base}/contextualizedembeddings"
|
||||
return api_base
|
||||
return "https://api.voyageai.com/v1/contextualizedembeddings"
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
return ["encoding_format", "dimensions"]
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
optional_params: dict,
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
"""
|
||||
Map OpenAI params to Voyage params
|
||||
|
||||
Reference: https://docs.voyageai.com/reference/contextualized-embeddings-api
|
||||
"""
|
||||
if "encoding_format" in non_default_params:
|
||||
optional_params["encoding_format"] = non_default_params["encoding_format"]
|
||||
if "dimensions" in non_default_params:
|
||||
optional_params["output_dimension"] = non_default_params["dimensions"]
|
||||
return optional_params
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
if api_key is None:
|
||||
api_key = (
|
||||
get_secret_str("VOYAGE_API_KEY")
|
||||
or get_secret_str("VOYAGE_AI_API_KEY")
|
||||
or get_secret_str("VOYAGE_AI_TOKEN")
|
||||
)
|
||||
return {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
|
||||
def transform_embedding_request(
|
||||
self,
|
||||
model: str,
|
||||
input: Union[AllEmbeddingInputValues, List[List[str]]],
|
||||
optional_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
return {
|
||||
"inputs": input,
|
||||
"model": model,
|
||||
**optional_params,
|
||||
}
|
||||
|
||||
def transform_embedding_response(
|
||||
self,
|
||||
model: str,
|
||||
raw_response: httpx.Response,
|
||||
model_response: EmbeddingResponse,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
api_key: Optional[str] = None,
|
||||
request_data: dict = {},
|
||||
optional_params: dict = {},
|
||||
litellm_params: dict = {},
|
||||
) -> EmbeddingResponse:
|
||||
try:
|
||||
raw_response_json = raw_response.json()
|
||||
except Exception:
|
||||
raise VoyageError(
|
||||
message=raw_response.text, status_code=raw_response.status_code
|
||||
)
|
||||
|
||||
# model_response.usage
|
||||
model_response.model = raw_response_json.get("model")
|
||||
model_response.data = raw_response_json.get("data")
|
||||
model_response.object = raw_response_json.get("object")
|
||||
|
||||
usage = Usage(
|
||||
prompt_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0),
|
||||
total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0),
|
||||
)
|
||||
model_response.usage = usage
|
||||
return model_response
|
||||
|
||||
def get_error_class(
|
||||
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
|
||||
) -> BaseLLMException:
|
||||
return VoyageError(
|
||||
message=error_message, status_code=status_code, headers=headers
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def is_contextualized_embeddings(model: str) -> bool:
|
||||
return "context" in model.lower()
|
||||
@@ -16704,6 +16704,14 @@
|
||||
"litellm_provider": "voyage",
|
||||
"mode": "embedding"
|
||||
},
|
||||
"voyage/voyage-context-3": {
|
||||
"max_tokens": 120000,
|
||||
"max_input_tokens": 120000,
|
||||
"input_cost_per_token": 1.8e-07,
|
||||
"output_cost_per_token": 0.0,
|
||||
"litellm_provider": "voyage",
|
||||
"mode": "embedding"
|
||||
},
|
||||
"voyage/rerank-2": {
|
||||
"max_tokens": 16000,
|
||||
"max_input_tokens": 16000,
|
||||
|
||||
+24
-7
@@ -2802,12 +2802,22 @@ def get_optional_params_embeddings( # noqa: PLR0915
|
||||
request_type="embeddings",
|
||||
)
|
||||
_check_valid_arg(supported_params=supported_params)
|
||||
optional_params = litellm.VoyageEmbeddingConfig().map_openai_params(
|
||||
non_default_params=non_default_params,
|
||||
optional_params={},
|
||||
model=model,
|
||||
drop_params=drop_params if drop_params is not None else False,
|
||||
)
|
||||
if litellm.VoyageContextualEmbeddingConfig.is_contextualized_embeddings(model):
|
||||
optional_params = (
|
||||
litellm.VoyageContextualEmbeddingConfig().map_openai_params(
|
||||
non_default_params=non_default_params,
|
||||
optional_params={},
|
||||
model=model,
|
||||
drop_params=drop_params if drop_params is not None else False,
|
||||
)
|
||||
)
|
||||
else:
|
||||
optional_params = litellm.VoyageEmbeddingConfig().map_openai_params(
|
||||
non_default_params=non_default_params,
|
||||
optional_params={},
|
||||
model=model,
|
||||
drop_params=drop_params if drop_params is not None else False,
|
||||
)
|
||||
elif custom_llm_provider == "infinity":
|
||||
supported_params = get_supported_openai_params(
|
||||
model=model,
|
||||
@@ -7013,7 +7023,14 @@ class ProviderConfigManager:
|
||||
model: str,
|
||||
provider: LlmProviders,
|
||||
) -> Optional[BaseEmbeddingConfig]:
|
||||
if litellm.LlmProviders.VOYAGE == provider:
|
||||
if (
|
||||
litellm.LlmProviders.VOYAGE == provider
|
||||
and litellm.VoyageContextualEmbeddingConfig.is_contextualized_embeddings(
|
||||
model
|
||||
)
|
||||
):
|
||||
return litellm.VoyageContextualEmbeddingConfig()
|
||||
elif litellm.LlmProviders.VOYAGE == provider:
|
||||
return litellm.VoyageEmbeddingConfig()
|
||||
elif litellm.LlmProviders.TRITON == provider:
|
||||
return litellm.TritonEmbeddingConfig()
|
||||
|
||||
@@ -16674,6 +16674,14 @@
|
||||
"litellm_provider": "voyage",
|
||||
"mode": "embedding"
|
||||
},
|
||||
"voyage/voyage-context-3": {
|
||||
"max_tokens": 120000,
|
||||
"max_input_tokens": 120000,
|
||||
"input_cost_per_token": 1.8e-07,
|
||||
"output_cost_per_token": 0.0,
|
||||
"litellm_provider": "voyage",
|
||||
"mode": "embedding"
|
||||
},
|
||||
"voyage/rerank-2": {
|
||||
"max_tokens": 16000,
|
||||
"max_input_tokens": 16000,
|
||||
|
||||
@@ -1,8 +1,7 @@
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
@@ -10,10 +9,11 @@ sys.path.insert(
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from base_embedding_unit_tests import BaseLLMEmbeddingTest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
|
||||
class TestVoyageAI(BaseLLMEmbeddingTest):
|
||||
@@ -25,56 +25,409 @@ class TestVoyageAI(BaseLLMEmbeddingTest):
|
||||
"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
|
||||
|
||||
client = HTTPHandler()
|
||||
litellm.set_verbose = True
|
||||
|
||||
with patch.object(client, "post") as mock_client:
|
||||
response = litellm.embedding(
|
||||
litellm.embedding(
|
||||
model="voyage/voyage-3-lite",
|
||||
input=["a"],
|
||||
dimensions=512,
|
||||
input_type="document",
|
||||
client=client,
|
||||
)
|
||||
|
||||
mock_client.assert_called_once()
|
||||
json_data = json.loads(mock_client.call_args.kwargs["data"])
|
||||
|
||||
print("request data to voyage ai", json.dumps(json_data, indent=4))
|
||||
|
||||
# Assert the request parameters
|
||||
assert json_data["input"] == ["a"]
|
||||
assert json_data["model"] == "voyage-3-lite"
|
||||
assert json_data["output_dimension"] == 512
|
||||
assert json_data["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
|
||||
|
||||
client = HTTPHandler()
|
||||
litellm.set_verbose = True
|
||||
|
||||
with patch.object(client, "post", return_value=MagicMock(status_code=200, json=lambda: {"usage": {"total_tokens": 120}})) as mock_client:
|
||||
response = litellm.embedding(
|
||||
model="voyage/voyage-3-lite",
|
||||
input=["a"],
|
||||
dimensions=512,
|
||||
input_type="document",
|
||||
client=client,
|
||||
)
|
||||
|
||||
mock_client.assert_called_once()
|
||||
# Assert the response
|
||||
# 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}")
|
||||
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}")
|
||||
|
||||
Reference in New Issue
Block a user