scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-03 02:22:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 82
litellm/tests/unit/llms/sagemaker/test_sagemaker_embedding_voyage.py
yuneng-jiang 5e6dc89ba1
test: move tests/test_litellm/llms into tests/unit/llms (#43191) 
* ci: run the unit_selection.sh shard files on every event instead of only fork pull requests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* ci: rename fork-flag to unit-flag now that it applies on every event

* test: move tests/test_litellm root and small trees into tests/unit

Pure renames, no content changes. Follow-up commits in this PR fix
references, merge the three files that already existed in tests/unit,
keep live-provider tests in tests/test_litellm and wire CI.

* test: carry tests/test_litellm conftest isolation into tests/unit

Callback lists, routing fallbacks, cached HTTP clients, logger state, AWS,
proxy-URL and keychain env, and session-end client cleanup now reset for
unit tests too. The environment isolation owns its MonkeyPatch so a test's
own monkeypatch is undone before the model-cost teardown runs.

* test: merge, split and prune the moved root and small-tree tests

Merge batches/test_batch_utils.py and the chat_completions and messages
dispatch tests into the files that already existed in tests/unit. Keep
the live Gemini interactions tests, the async image-fetch format test and
the OpenAI embedding scorer test in tests/test_litellm since they need
real network or keys. Put test_router.py under tests/unit/test_router so
the existing package no longer shadows it. Delete eight tests the audit
found superseded by stronger ones kept in this move.

* ci: run the moved root and small-tree tests under their legacy flags

Add the misc and responses-caching-types flags to unit_selection.sh and
CircleCI, extend enterprise-routing and mcp-integration, and point the
legacy GHA shards, Makefile, redis-compat workflow, merge smoke manifest
and change classifier at the new paths.

* test: make the new tests/unit directories packages

tests/unit/test_package_layout.py requires every directory to carry an
__init__.py, and without one the moved and retained
test_litellm_responses_bridge.py modules collide on import.

* test: scope the unit socket block to tests/unit in shared sessions

The GHA shards collect the legacy test-path and the unit selection in one
pytest session. The unit conftest's loopback-only block leaked into legacy
modules that reach the network at import. The legacy conftest now lifts the
restriction at collect and setup time, and the unit conftest re-applies it
when collecting its own modules.

* test: move tests/test_litellm/llms into tests/unit/llms

Rename-only. Moves the provider tests and the fine-tuning fixtures they
load, mirroring the old paths. Follow-up commits merge, split and wire them.

* test: merge, split and prune the moved llms tests

Merges the Databricks chat transformation tests into the existing unit
file, keeps the tests that need real keys or the network in
tests/test_litellm, deletes the audited tests a stronger unit test
already covers, and points imports at tests.unit.llms.

* ci: run the moved llms tests under their legacy flags

The Vertex AI and All Other Providers shards keep their legacy test-path
for the retained files and add the llm-vertex-ai and llm-other-providers
unit selections. CircleCI gets matching unit jobs.

* test: make the tests/unit/llms directories packages

Adds __init__.py to the moved dirs and drops the legacy ones whose
directories no longer hold tests.

* test: drop script runners and path hacks the llms split left dangling

The __main__ runners in the split openai_like files and the Databricks e2e
runner called tests that now live in the other half of the split or were
deleted. The retained legacy halves also no longer need sys.path edits.

* test: give the shard-script tests their own GITHUB_OUTPUT

They only passed where the runner set it. The CircleCI unit job's env
allowlist drops it, so the script's redirect failed there.

* test: point the router and module-deletion checks at tests/unit

router_code_coverage and code_qa_check_tests only searched tests/test_litellm,
so the moved router tests no longer counted. The two silent-experiment tests
the audit deleted were the only direct callers of those methods; they are
replaced with tests that assert the forwarded shadow request and the
recursion guard.

* test: keep the Databricks manual e2e runner and fix the SageMaker Nova run path

The Databricks e2e file is a manual script whose main() calls the tests
that were pruned, so pruning them broke the documented run. It is back to
its main version. The SageMaker Nova docstring now points at the file's
real location in tests/local_testing.

* test: keep the job's UNIT_FLAG out of the shard-script tests

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-09-25 12:43:23 -07:00

628 lines
24 KiB
Python
Raw Blame History

"""
Test cases for SageMaker Voyage embedding model integration
This module tests the factory pattern implementation for Voyage embedding models
deployed on AWS SageMaker endpoints, including parameter handling, request/response
transformation, and model type detection.
"""
import json
from unittest.mock import MagicMock, patch
import httpx
import pytest
from litellm import embedding
from litellm.llms.sagemaker.embedding.cohere_transformation import (
SagemakerCohereEmbeddingConfig,
)
from litellm.llms.sagemaker.embedding.transformation import SagemakerEmbeddingConfig
from litellm.llms.voyage.embedding.transformation import VoyageEmbeddingConfig
from litellm.types.utils import EmbeddingResponse, Usage
class TestSagemakerEmbeddingFactory:
"""Test the factory pattern for SageMaker embedding configurations"""
def test_get_model_config_voyage_model(self):
"""Test that Voyage models return VoyageEmbeddingConfig"""
config = SagemakerEmbeddingConfig.get_model_config("voyage-3-5-embedding")
assert isinstance(config, VoyageEmbeddingConfig)
assert config.get_supported_openai_params("voyage-3-5-embedding") == [
"encoding_format",
"dimensions",
]
def test_get_model_config_hf_model(self):
"""Test that non-Voyage models return base SagemakerEmbeddingConfig"""
config = SagemakerEmbeddingConfig.get_model_config(
"sentence-transformers-model"
)
assert isinstance(config, SagemakerEmbeddingConfig)
assert config.get_supported_openai_params("sentence-transformers-model") == []
def test_get_model_config_case_insensitive(self):
"""Test that model detection is case insensitive"""
config1 = SagemakerEmbeddingConfig.get_model_config("VOYAGE-3-5-embedding")
config2 = SagemakerEmbeddingConfig.get_model_config("Voyage-3-5-Embedding")
config3 = SagemakerEmbeddingConfig.get_model_config("voyage-3-5-embedding")
assert isinstance(config1, VoyageEmbeddingConfig)
assert isinstance(config2, VoyageEmbeddingConfig)
assert isinstance(config3, VoyageEmbeddingConfig)
def test_get_model_config_cohere_model(self):
"""Cohere SageMaker endpoints route to SagemakerCohereEmbeddingConfig"""
for endpoint_name in (
"cohere.embed-multilingual-v3",
"cohere-embed-english-v3-prod",
"my-cohere-marketplace-endpoint",
"COHERE-EMBED-V4",
):
config = SagemakerEmbeddingConfig.get_model_config(endpoint_name)
assert isinstance(config, SagemakerCohereEmbeddingConfig), endpoint_name
class TestSagemakerCohereEmbeddingConfig:
"""Cohere-specific SageMaker embedding request/response transforms"""
def setup_method(self):
self.config = SagemakerCohereEmbeddingConfig()
MODEL = "cohere.embed-multilingual-v3"
def test_transform_request_uses_cohere_payload(self):
"""Bug repro: request must use `texts` + `input_type`, not HF `inputs`"""
result = self.config.transform_embedding_request(
model=self.MODEL,
input=["hello"],
optional_params={"input_type": "search_query"},
headers={},
)
assert "inputs" not in result
assert result["texts"] == ["hello"]
assert result["input_type"] == "search_query"
def test_transform_request_default_input_type(self):
result = self.config.transform_embedding_request(
model=self.MODEL,
input=["hello"],
optional_params={},
headers={},
)
assert result["texts"] == ["hello"]
assert result["input_type"] == "search_document"
def test_transform_request_normalizes_string_input(self):
result = self.config.transform_embedding_request(
model=self.MODEL,
input="hello",
optional_params={},
headers={},
)
assert result["texts"] == ["hello"]
def test_map_openai_params_dimensions_to_output_dimension(self):
params = self.config.map_openai_params(
non_default_params={"dimensions": 512, "encoding_format": "float"},
optional_params={},
model=self.MODEL,
drop_params=False,
)
assert params["output_dimension"] == 512
assert params["embedding_types"] == ["float"]
def test_map_openai_params_input_type_from_non_default_params(self):
params = self.config.map_openai_params(
non_default_params={"input_type": "search_query"},
optional_params={},
model=self.MODEL,
drop_params=False,
)
assert params["input_type"] == "search_query"
def test_get_optional_params_embeddings_preserves_input_type(self):
"""Exercises get_optional_params_embeddings, not transform in isolation."""
from litellm.utils import get_optional_params_embeddings
optional_params = get_optional_params_embeddings(
model=self.MODEL,
custom_llm_provider="sagemaker",
input_type="search_query",
)
assert optional_params.get("input_type") == "search_query"
body = self.config.transform_embedding_request(
model=self.MODEL,
input=["hello"],
optional_params=optional_params,
headers={},
)
assert body["texts"] == ["hello"]
assert body["input_type"] == "search_query"
def test_get_optional_params_embeddings_maps_dimensions_without_duplicate(self):
"""dimensions must map to output_dimension only, not also stay as dimensions."""
from litellm.utils import get_optional_params_embeddings
optional_params = get_optional_params_embeddings(
model=self.MODEL,
custom_llm_provider="sagemaker",
dimensions=512,
input_type="search_query",
)
assert optional_params.get("output_dimension") == 512
assert "dimensions" not in optional_params
assert optional_params.get("input_type") == "search_query"
def test_transform_response_parses_cohere_payload(self):
cohere_response = {
"embeddings": [[0.1, 0.2, 0.3]],
"meta": {"billed_units": {"input_tokens": 2}},
}
mock_response = httpx.Response(
status_code=200,
content=json.dumps(cohere_response).encode("utf-8"),
headers={"content-type": "application/json"},
)
logging_obj = MagicMock()
logging_obj.model_call_details = {"input": ["hello"]}
result = self.config.transform_embedding_response(
model=self.MODEL,
raw_response=mock_response,
model_response=EmbeddingResponse(),
logging_obj=logging_obj,
api_key=None,
request_data={"texts": ["hello"], "input_type": "search_query"},
optional_params={},
litellm_params={},
)
assert result.object == "list"
assert len(result.data) == 1
assert result.data[0]["embedding"] == [0.1, 0.2, 0.3]
assert result.usage.prompt_tokens == 2
def test_transform_response_does_not_double_call_post_call(self):
"""
Greptile review fix: SageMaker handler already calls
`logging_obj.post_call` once before invoking
`transform_embedding_response`. The transform must NOT call it again,
otherwise callbacks, cost calculators, and log handlers double-fire
for every Cohere SageMaker embedding call.
"""
cohere_response = {
"embeddings": [[0.1, 0.2, 0.3]],
"meta": {"billed_units": {"input_tokens": 2}},
}
mock_response = httpx.Response(
status_code=200,
content=json.dumps(cohere_response).encode("utf-8"),
headers={"content-type": "application/json"},
)
logging_obj = MagicMock()
logging_obj.model_call_details = {"input": ["hello"]}
self.config.transform_embedding_response(
model=self.MODEL,
raw_response=mock_response,
model_response=EmbeddingResponse(),
logging_obj=logging_obj,
api_key=None,
request_data={"texts": ["hello"], "input_type": "search_query"},
optional_params={},
litellm_params={},
)
logging_obj.post_call.assert_not_called()
class TestVoyageEmbeddingConfig:
"""Test Voyage-specific embedding configuration"""
def setup_method(self):
self.config = VoyageEmbeddingConfig()
def test_get_supported_openai_params(self):
"""Test supported parameters for Voyage models"""
params = self.config.get_supported_openai_params("voyage-3-5-embedding")
assert params == ["encoding_format", "dimensions"]
def test_map_openai_params_encoding_format(self):
"""Test mapping of encoding_format parameter"""
result = self.config.map_openai_params(
non_default_params={"encoding_format": "float"},
optional_params={},
model="voyage-3-5-embedding",
drop_params=False,
)
assert result == {"encoding_format": "float"}
def test_map_openai_params_dimensions(self):
"""Test mapping of dimensions parameter to output_dimension"""
result = self.config.map_openai_params(
non_default_params={"dimensions": 1024},
optional_params={},
model="voyage-3-5-embedding",
drop_params=False,
)
assert result == {"output_dimension": 1024}
def test_map_openai_params_unsupported_encoding(self):
"""Test handling of unsupported encoding_format values - VoyageEmbeddingConfig passes through without validation"""
result = self.config.map_openai_params(
non_default_params={"encoding_format": "invalid"},
optional_params={},
model="voyage-3-5-embedding",
drop_params=False,
)
assert result == {"encoding_format": "invalid"}
def test_map_openai_params_drop_unsupported(self):
"""Test that VoyageEmbeddingConfig doesn't drop parameters based on drop_params flag"""
result = self.config.map_openai_params(
non_default_params={"encoding_format": "invalid", "dimensions": 512},
optional_params={},
model="voyage-3-5-embedding",
drop_params=True,
)
assert result == {"encoding_format": "invalid", "output_dimension": 512}
def test_transform_embedding_request(self):
"""Test Voyage request transformation"""
result = self.config.transform_embedding_request(
model="voyage-3-5-embedding",
input=["Hello", "World"],
optional_params={"encoding_format": "float"},
headers={},
)
expected = {
"input": ["Hello", "World"],
"model": "voyage-3-5-embedding",
"encoding_format": "float",
}
assert result == expected
def test_transform_embedding_response(self):
"""Test Voyage response transformation to OpenAI format"""
# Mock Voyage response
voyage_response = {
"data": [
{"object": "embedding", "embedding": [0.1, 0.2, 0.3], "index": 0},
{"object": "embedding", "embedding": [0.4, 0.5, 0.6], "index": 1},
],
"object": "list",
"model": "voyage-3-5-embedding",
"usage": {"total_tokens": 10},
}
# Create mock httpx Response
mock_response = httpx.Response(
status_code=200,
content=json.dumps(voyage_response).encode("utf-8"),
headers={"content-type": "application/json"},
)
model_response = EmbeddingResponse()
result = self.config.transform_embedding_response(
model="voyage-3-5-embedding",
raw_response=mock_response,
model_response=model_response,
logging_obj=None,
request_data={"input": ["Hello", "World"]},
)
# Verify response structure
assert result.object == "list"
assert result.model == "voyage-3-5-embedding"
assert len(result.data) == 2
assert result.data[0]["object"] == "embedding"
assert result.data[0]["index"] == 0
assert result.data[0]["embedding"] == [0.1, 0.2, 0.3]
assert result.data[1]["object"] == "embedding"
assert result.data[1]["index"] == 1
assert result.data[1]["embedding"] == [0.4, 0.5, 0.6]
assert isinstance(result.usage, Usage)
assert result.usage.total_tokens == 10
class TestHFSagemakerEmbeddingConfig:
"""Test Hugging Face embedding configuration"""
def setup_method(self):
self.config = SagemakerEmbeddingConfig()
def test_get_supported_openai_params_hf(self):
"""Test that HF models don't support embedding parameters"""
params = self.config.get_supported_openai_params("sentence-transformers-model")
assert params == []
def test_transform_embedding_request_hf(self):
"""Test HF request transformation"""
result = self.config.transform_embedding_request(
model="sentence-transformers-model",
input=["Hello", "World"],
optional_params={},
headers={},
)
expected = {"inputs": ["Hello", "World"]}
assert result == expected
def test_transform_embedding_response_hf(self):
"""Test HF response transformation to OpenAI format"""
# Mock HF response
hf_response = {"embedding": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]}
# Create mock httpx Response
mock_response = httpx.Response(
status_code=200,
content=json.dumps(hf_response).encode("utf-8"),
headers={"content-type": "application/json"},
)
model_response = EmbeddingResponse()
result = self.config.transform_embedding_response(
model="sentence-transformers-model",
raw_response=mock_response,
model_response=model_response,
logging_obj=None,
request_data={"inputs": ["Hello", "World"]},
)
# Verify response structure
assert result.object == "list"
assert result.model == "sentence-transformers-model"
assert len(result.data) == 2
assert result.data[0]["object"] == "embedding"
assert result.data[0]["index"] == 0
assert result.data[0]["embedding"] == [0.1, 0.2, 0.3]
assert result.data[1]["object"] == "embedding"
assert result.data[1]["index"] == 1
assert result.data[1]["embedding"] == [0.4, 0.5, 0.6]
assert isinstance(result.usage, Usage)
class TestSagemakerEmbeddingIntegration:
"""Integration tests for SageMaker embedding with factory pattern"""
def test_voyage_embedding_request_format(self):
"""Test that Voyage models use correct request format"""
with patch(
"litellm.llms.sagemaker.completion.handler.SagemakerLLM.embedding"
) as mock_embedding:
# Mock the actual SageMaker call to avoid AWS credentials
mock_embedding.return_value = EmbeddingResponse(
object="list",
data=[
{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]},
{"object": "embedding", "index": 1, "embedding": [0.4, 0.5, 0.6]},
],
model="voyage-3-5-embedding",
usage=Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10),
)
# Test Voyage model
response = embedding(
model="sagemaker/voyage-3-5-embedding-endpoint",
input=["Hello", "World"],
encoding_format="float",
dimensions=1024,
)
# Verify the request was made with correct format
mock_embedding.assert_called_once()
call_args = mock_embedding.call_args
assert call_args[1]["model"] == "voyage-3-5-embedding-endpoint"
assert call_args[1]["input"] == ["Hello", "World"]
# Check that the parameters are in the optional_params
optional_params = call_args[1].get("optional_params", {})
assert optional_params.get("encoding_format") == "float"
assert (
optional_params.get("output_dimension") == 1024
) # dimensions is mapped to output_dimension
def test_hf_embedding_request_format(self):
"""Test that HF models use correct request format"""
with patch(
"litellm.llms.sagemaker.completion.handler.SagemakerLLM.embedding"
) as mock_embedding:
# Mock the actual SageMaker call to avoid AWS credentials
mock_embedding.return_value = EmbeddingResponse(
object="list",
data=[
{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]},
{"object": "embedding", "index": 1, "embedding": [0.4, 0.5, 0.6]},
],
model="sentence-transformers-model",
usage=Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10),
)
# Test HF model with drop_params=True to ignore unsupported parameters
response = embedding(
model="sagemaker/sentence-transformers-endpoint",
input=["Hello", "World"],
encoding_format="float", # Should be ignored
dimensions=1024, # Should be ignored
drop_params=True,
)
# Verify the request was made
mock_embedding.assert_called_once()
call_args = mock_embedding.call_args
assert call_args[1]["model"] == "sentence-transformers-endpoint"
assert call_args[1]["input"] == ["Hello", "World"]
# HF models should ignore these parameters in optional_params
optional_params = call_args[1].get("optional_params", {})
assert (
"encoding_format" not in optional_params
or optional_params["encoding_format"] is None
)
assert (
"dimensions" not in optional_params
or optional_params["dimensions"] is None
)
def test_parameter_validation_voyage(self):
"""Test parameter validation for Voyage models"""
# Test valid parameters
config = VoyageEmbeddingConfig()
result = config.map_openai_params(
non_default_params={"encoding_format": "float", "dimensions": 512},
optional_params={},
model="voyage-3-5-embedding",
drop_params=False,
)
assert result == {"encoding_format": "float", "output_dimension": 512}
def test_parameter_validation_hf(self):
"""Test parameter validation for HF models"""
# Test that HF models ignore embedding parameters
config = SagemakerEmbeddingConfig()
result = config.map_openai_params(
non_default_params={"encoding_format": "float", "dimensions": 512},
optional_params={},
model="sentence-transformers-model",
drop_params=False,
)
assert result == {} # HF models should ignore these parameters
class TestErrorHandling:
"""Test error handling in the embedding integration"""
def test_voyage_response_missing_data(self):
"""Test handling of Voyage response missing data field"""
config = VoyageEmbeddingConfig()
# Mock response without data field
mock_response = httpx.Response(
status_code=200,
content=json.dumps({"object": "list"}).encode("utf-8"),
headers={"content-type": "application/json"},
)
model_response = EmbeddingResponse()
# VoyageEmbeddingConfig doesn't validate for missing data field, it just sets it to None
result = config.transform_embedding_response(
model="voyage-3-5-embedding",
raw_response=mock_response,
model_response=model_response,
logging_obj=None,
request_data={"input": ["Hello"]},
)
assert result.data is None
def test_hf_response_missing_embedding(self):
"""Test handling of HF response missing embedding field"""
config = SagemakerEmbeddingConfig()
# Mock response without embedding field
mock_response = httpx.Response(
status_code=200,
content=json.dumps({"object": "list"}).encode("utf-8"),
headers={"content-type": "application/json"},
)
model_response = EmbeddingResponse()
with pytest.raises(Exception, match="Unexpected response format"):
config.transform_embedding_response(
model="sentence-transformers-model",
raw_response=mock_response,
model_response=model_response,
logging_obj=None,
request_data={"inputs": ["Hello"]},
)
class TestTEIEmbeddingResponse:
"""Test HuggingFace Text Embeddings Inference (TEI) response format support"""
def setup_method(self):
self.config = SagemakerEmbeddingConfig()
def test_transform_embedding_response_tei_raw_array(self):
"""Test TEI response transformation - raw array format [[...]]"""
# TEI returns raw embedding arrays without wrapper
tei_response = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
mock_response = httpx.Response(
status_code=200,
content=json.dumps(tei_response).encode("utf-8"),
headers={"content-type": "application/json"},
)
model_response = EmbeddingResponse()
result = self.config.transform_embedding_response(
model="tei-qwen-embedding",
raw_response=mock_response,
model_response=model_response,
logging_obj=None,
request_data={"inputs": ["Hello", "World"]},
)
# Verify response structure
assert result.object == "list"
assert result.model == "tei-qwen-embedding"
assert len(result.data) == 2
assert result.data[0]["object"] == "embedding"
assert result.data[0]["index"] == 0
assert result.data[0]["embedding"] == [0.1, 0.2, 0.3]
assert result.data[1]["object"] == "embedding"
assert result.data[1]["index"] == 1
assert result.data[1]["embedding"] == [0.4, 0.5, 0.6]
assert isinstance(result.usage, Usage)
def test_transform_embedding_response_tei_single_input(self):
"""Test TEI response with single input"""
tei_response = [[0.1, 0.2, 0.3, 0.4, 0.5]]
mock_response = httpx.Response(
status_code=200,
content=json.dumps(tei_response).encode("utf-8"),
headers={"content-type": "application/json"},
)
model_response = EmbeddingResponse()
result = self.config.transform_embedding_response(
model="tei-model",
raw_response=mock_response,
model_response=model_response,
logging_obj=None,
request_data={"inputs": ["Hello"]},
)
assert len(result.data) == 1
assert result.data[0]["embedding"] == [0.1, 0.2, 0.3, 0.4, 0.5]
def test_transform_embedding_response_wrapped_format_still_works(self):
"""Test that wrapped format {"embedding": [...]} still works"""
hf_response = {"embedding": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]}
mock_response = httpx.Response(
status_code=200,
content=json.dumps(hf_response).encode("utf-8"),
headers={"content-type": "application/json"},
)
model_response = EmbeddingResponse()
result = self.config.transform_embedding_response(
model="hf-model",
raw_response=mock_response,
model_response=model_response,
logging_obj=None,
request_data={"inputs": ["Hello", "World"]},
)
assert len(result.data) == 2
assert result.data[0]["embedding"] == [0.1, 0.2, 0.3]
assert result.data[1]["embedding"] == [0.4, 0.5, 0.6]
if __name__ == "__main__":
pytest.main([__file__])
Reference in a new issue View git blame Copy permalink