diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index 2d661610eaa..d555a1a2266 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -2,7 +2,7 @@ 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 +from typing import Any, Dict, List, Optional, Tuple, Union import httpx @@ -98,6 +98,11 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): "Authorization": f"Bearer {api_key}", } + # Chunk size (in tokens) used when the API auto-chunks a flat ``list[str]``. + # Matches the voyage-context-4 context window so each string stays a single + # chunk instead of being split. + AUTO_CHUNK_SIZE = 32000 + def transform_embedding_request( self, model: str, @@ -105,42 +110,74 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): optional_params: dict, headers: dict, ) -> dict: + inputs, extra_params = self._transform_contextual_inputs(input, optional_params) return { - "inputs": self._transform_contextual_inputs(input, optional_params), + "inputs": inputs, "model": model, **optional_params, + **extra_params, } - @staticmethod + @classmethod def _transform_contextual_inputs( + cls, input: Union[AllEmbeddingInputValues, List[List[str]]], optional_params: dict, - ) -> List[List[str]]: + ) -> Tuple[Union[List[str], List[List[str]]], dict]: """ - Voyage's contextualized embeddings API expects ``inputs`` to be a - ``list[list[str]]`` (each inner list is a document made of chunks that - share context). + Normalize ``input`` for Voyage's contextualized embeddings API and + return ``(inputs, extra_params)`` where ``extra_params`` carries any + request fields (e.g. auto-chunking) needed for the chosen shape. - It also accepts a flat ``list[str]`` - but only when - ``input_type == "query"``. In every other case a flat list must be - wrapped so each string becomes its own single-chunk document, otherwise - the API rejects the request with a 400. + The API contract (verified against the live endpoint) is: + + - A flat ``list[str]`` is only accepted with ``input_type="query"`` or + with ``enable_auto_chunking=True`` (which itself requires + ``input_type="document"``). + - A ``list[list[str]]`` (each inner list = one document's chunks) is + always accepted. + + So we prefer to send a flat ``list[str]`` and let the API auto-chunk, + instead of pre-wrapping into ``list[list[str]]``: + + - ``str`` -> ``[str]`` + ``enable_auto_chunking`` (input_type=document) + - flat ``list[str]`` + ``input_type="query"`` -> kept flat, as-is + - flat ``list[str]`` otherwise -> kept flat + ``enable_auto_chunking`` + (input_type=document) + - ``list[list[str]]`` -> passed through unchanged Reference: https://docs.voyageai.com/reference/contextualized-embeddings-api """ - # Single string -> one document with a single chunk. + # Single string -> a one-element flat list, auto-chunked. if isinstance(input, str): - return [[input]] + return [input], cls._auto_chunk_params(optional_params) - # Flat list[str]: keep as list[str] when the API allows it - # (input_type="query"), otherwise wrap each string as its own document. + # Flat list[str]. if isinstance(input, list) and all(isinstance(i, str) for i in input): if optional_params.get("input_type") == "query": - return input # type: ignore[return-value] - return [[i] for i in input] + # The API accepts a flat query list as-is. + return input, {} # type: ignore[return-value] + # Otherwise let the API auto-chunk the flat list. + return input, cls._auto_chunk_params(optional_params) # type: ignore[return-value] # Already list[list[str]] (or another shape) -> pass through unchanged. - return input # type: ignore[return-value] + return input, {} # type: ignore[return-value] + + @classmethod + def _auto_chunk_params(cls, optional_params: dict) -> dict: + """ + Params required to send a flat ``list[str]`` to the contextualized API. + + ``enable_auto_chunking=True`` requires ``input_type="document"``, so set + it unless the caller already provided an ``input_type``. + """ + params: Dict[str, Any] = { + "enable_auto_chunking": True, + "chunk_size": cls.AUTO_CHUNK_SIZE, + } + if not optional_params.get("input_type"): + params["input_type"] = "document" + return 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 577ee808035..b7a2c41a979 100644 --- a/tests/llm_translation/test_voyage_ai.py +++ b/tests/llm_translation/test_voyage_ai.py @@ -219,8 +219,13 @@ class TestVoyageContextualEmbeddings: assert transformed["model"] == "voyage-context-3" assert transformed["encoding_format"] == "float" - def test_contextual_flat_list_str_is_wrapped(self): - """Flat list[str] without input_type must be wrapped to list[list[str]].""" + def test_contextual_flat_list_str_is_auto_chunked(self): + """Flat list[str] without input_type stays flat and is auto-chunked. + + The API accepts a flat list[str] only with input_type="query" or with + enable_auto_chunking=True (which requires input_type="document"), so we + keep the list flat and let the API auto-chunk it. + """ from litellm.llms.voyage.embedding.transformation_contextual import ( VoyageContextualEmbeddingConfig, ) @@ -231,10 +236,11 @@ class TestVoyageContextualEmbeddings: "voyage-context-4", ["Hello", "world"], {}, {} ) - # Voyage rejects a flat list[str] unless input_type="query", so each - # string becomes its own single-chunk document. - assert transformed["inputs"] == [["Hello"], ["world"]] + assert transformed["inputs"] == ["Hello", "world"] assert transformed["model"] == "voyage-context-4" + assert transformed["enable_auto_chunking"] is True + assert transformed["chunk_size"] == 32000 + assert transformed["input_type"] == "document" def test_contextual_flat_list_str_query_stays_flat(self): """Flat list[str] with input_type='query' is sent as-is (list[str]).""" @@ -253,9 +259,31 @@ class TestVoyageContextualEmbeddings: assert transformed["inputs"] == ["Hello", "world"] assert transformed["input_type"] == "query" + # A query list is accepted as-is, no auto-chunking needed. + assert "enable_auto_chunking" not in transformed - def test_contextual_single_string_is_wrapped(self): - """A single string is wrapped to a single document with a single chunk.""" + def test_contextual_flat_list_str_document_input_type_preserved(self): + """Explicit input_type='document' is preserved while auto-chunking.""" + from litellm.llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, + ) + + config = VoyageContextualEmbeddingConfig() + + transformed = config.transform_embedding_request( + "voyage-context-4", + ["Hello", "world"], + {"input_type": "document"}, + {}, + ) + + assert transformed["inputs"] == ["Hello", "world"] + assert transformed["input_type"] == "document" + assert transformed["enable_auto_chunking"] is True + assert transformed["chunk_size"] == 32000 + + def test_contextual_single_string_is_auto_chunked(self): + """A single string becomes a one-element flat list and is auto-chunked.""" from litellm.llms.voyage.embedding.transformation_contextual import ( VoyageContextualEmbeddingConfig, ) @@ -266,7 +294,10 @@ class TestVoyageContextualEmbeddings: "voyage-context-4", "Hello", {}, {} ) - assert transformed["inputs"] == [["Hello"]] + assert transformed["inputs"] == ["Hello"] + assert transformed["enable_auto_chunking"] is True + assert transformed["chunk_size"] == 32000 + assert transformed["input_type"] == "document" def test_contextual_nested_input_passthrough(self): """Already-nested list[list[str]] input is passed through unchanged."""