diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index 260806fc164..bcfb55b278c 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -104,27 +104,41 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): optional_params: dict, headers: dict, ) -> dict: + inputs, contextual_params = self._prepare_contextual_inputs(input, optional_params) return { - "inputs": self._normalize_contextual_inputs(input), + "inputs": inputs, "model": model, **optional_params, + **contextual_params, } @staticmethod - def _normalize_contextual_inputs( + def _prepare_contextual_inputs( input: AllEmbeddingInputValues | list[list[str]], - ) -> AllEmbeddingInputValues | list[list[str]]: + optional_params: dict, + ) -> tuple[AllEmbeddingInputValues | list[list[str]], dict[str, str | bool]]: """ - Voyage's contextualized embeddings API accepts ``inputs`` as either a - flat ``list[str]`` (one document's chunks) or a nested ``list[list[str]]`` - (multiple documents). Both are sent through unchanged; a bare ``str`` is - wrapped into a single-element list so the payload always matches the spec. + Shape ``inputs`` and the auto-chunking params to match Voyage's + contextualized embeddings contract. + + - ``list[list[str]]`` (pre-chunked documents) is always valid and passes through. + - A flat ``list[str]`` or bare ``str`` is only valid as documents when + ``enable_auto_chunking=True`` with ``input_type="document"``, or as + queries with ``input_type="query"``. So a non-query flat input is sent + with those two params defaulted (caller-set values win). Reference: https://docs.voyageai.com/docs/contextualized-chunk-embeddings """ - if isinstance(input, str): - return [input] - return input + if isinstance(input, list) and len(input) > 0 and isinstance(input[0], list): + return input, {} + flat: Final = [input] if isinstance(input, str) else input + if optional_params.get("input_type") == "query": + return flat, {} + contextual_params: Final = { + **({"input_type": "document"} if "input_type" not in optional_params else {}), + **({"enable_auto_chunking": True} if "enable_auto_chunking" not in optional_params else {}), + } + return flat, contextual_params def transform_embedding_response( self, diff --git a/tests/llm_translation/test_voyage_ai.py b/tests/llm_translation/test_voyage_ai.py index 33705bf7b34..96fb5974ee4 100644 --- a/tests/llm_translation/test_voyage_ai.py +++ b/tests/llm_translation/test_voyage_ai.py @@ -195,8 +195,8 @@ class TestVoyageContextualEmbeddings: assert transformed["model"] == "voyage-context-3" assert transformed["encoding_format"] == "float" - def test_contextual_embedding_flat_list_input(self): - """A flat list[str] is passed through unchanged as inputs (spec allows list[str])""" + def test_contextual_embedding_flat_list_defaults_to_document_auto_chunk(self): + """A flat list[str] with no input_type is documents, so it needs auto-chunking + input_type=document""" from litellm.llms.voyage.embedding.transformation_contextual import ( VoyageContextualEmbeddingConfig, ) @@ -210,9 +210,28 @@ class TestVoyageContextualEmbeddings: assert transformed["inputs"] == flat_input assert transformed["model"] == "voyage-context-4" + assert transformed["input_type"] == "document" + assert transformed["enable_auto_chunking"] is True - def test_contextual_embedding_nested_list_input(self): - """A nested list[list[str]] is passed through unchanged as inputs""" + def test_contextual_embedding_flat_list_query_stays_flat_without_auto_chunk(self): + """A flat list[str] of queries is valid as-is, so no auto-chunking must be forced on""" + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + flat_input = ["what is voyage", "who owns voyage"] + + transformed = config.transform_embedding_request( + "voyage-context-4", flat_input, {"input_type": "query"}, {} + ) + + assert transformed["inputs"] == flat_input + assert transformed["input_type"] == "query" + assert "enable_auto_chunking" not in transformed + + def test_contextual_embedding_nested_list_input_passes_through(self): + """A nested list[list[str]] is pre-chunked documents, valid unchanged with no extra params""" from litellm.llms.voyage.embedding.transformation_contextual import ( VoyageContextualEmbeddingConfig, ) @@ -225,9 +244,11 @@ class TestVoyageContextualEmbeddings: ) assert transformed["inputs"] == nested_input + assert "enable_auto_chunking" not in transformed + assert "input_type" not in transformed - def test_contextual_embedding_str_input_wrapped(self): - """A bare str is wrapped into a single-element list so inputs is always a list""" + def test_contextual_embedding_str_input_wrapped_with_auto_chunk(self): + """A bare str is wrapped to a one-element list and, as documents, gets auto-chunking + input_type=document""" from litellm.llms.voyage.embedding.transformation_contextual import ( VoyageContextualEmbeddingConfig, ) @@ -239,6 +260,28 @@ class TestVoyageContextualEmbeddings: ) assert transformed["inputs"] == ["just one chunk"] + assert transformed["input_type"] == "document" + assert transformed["enable_auto_chunking"] is True + + def test_contextual_embedding_caller_params_win(self): + """Caller-set input_type=document with explicit auto-chunking off must be respected, not overridden""" + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + flat_input = ["chunk a", "chunk b"] + + transformed = config.transform_embedding_request( + "voyage-context-4", + flat_input, + {"input_type": "document", "enable_auto_chunking": False}, + {}, + ) + + assert transformed["inputs"] == flat_input + assert transformed["input_type"] == "document" + assert transformed["enable_auto_chunking"] is False def test_contextual_embedding_response_transformation(self): """Test response transformation for contextual embeddings"""