Merge 70db783d13 into aeede53112

litellm/llms/voyage/embedding/transformation.py

@@ -88,6 +88,7 @@ class VoyageEmbeddingConfig(BaseEmbeddingConfig):
         api_base: str | None = None,
     ) -> dict:
         return {
+            **headers,
             "Authorization": f"Bearer {get_voyage_api_key(api_key)}",
         }
 

litellm/llms/voyage/embedding/transformation_contextual.py

@@ -93,6 +93,7 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
         api_base: str | None = None,
     ) -> dict:
         return {
+            **headers,
             "Authorization": f"Bearer {get_voyage_api_key(api_key)}",
         }
 

litellm/llms/voyage/embedding/transformation_multimodal.py

@@ -93,7 +93,7 @@ class VoyageMultimodalEmbeddingConfig(BaseEmbeddingConfig):
                 "Set VOYAGE_API_KEY / VOYAGE_AI_API_KEY / VOYAGE_AI_TOKEN "
                 "or pass `api_key` explicitly."
             )
-        return {"Authorization": f"Bearer {resolved_api_key}"}
+        return {**headers, "Authorization": f"Bearer {resolved_api_key}"}
 
     def _normalize_content_item(self, item: dict[str, object]) -> dict[str, object]:
         item_type: Final = item.get("type")

litellm/llms/voyage/rerank/transformation.py

@@ -165,6 +165,7 @@ class VoyageRerankConfig(BaseRerankConfig):
             raise ValueError("Voyage AI API key is required. Set via `api_key` parameter or `VOYAGE_API_KEY` env var.")
         self._api_key = resolved_api_key
         return {
+            **headers,
             "Authorization": f"Bearer {resolved_api_key}",
             "content-type": "application/json",
         }

tests/unit/llms/voyage/embedding/test_transformation_contextual.py

@@ -1,5 +1,6 @@
 import asyncio
 import datetime
+from typing import Final
 import importlib
 import json
 import os
@@ -579,3 +580,14 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str):
         )
         == "https://gateway.example.com/voyage/v1/contextualizedembeddings"
     )
+
+
+def test_validate_environment_preserves_custom_gateway_headers() -> None:
+    original: Final = {"X-Gateway-Token": "gateway-token", "X-Trace-ID": "trace-id"}
+    headers: Final = VoyageContextualEmbeddingConfig().validate_environment(
+        headers=original, model="test-model", messages=[], optional_params={}, litellm_params={}, api_key="provider-key"
+    )
+    assert headers["X-Gateway-Token"] == "gateway-token"
+    assert headers["X-Trace-ID"] == "trace-id"
+    assert headers["Authorization"] == "Bearer provider-key"
+    assert original == {"X-Gateway-Token": "gateway-token", "X-Trace-ID": "trace-id"}

tests/unit/llms/voyage/embedding/test_transformation_multimodal.py

@@ -1,4 +1,5 @@
 import datetime
+from typing import Final
 
 import httpx
 import pytest
@@ -106,3 +107,14 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str):
         )
         == "https://gateway.example.com/voyage/v1/multimodalembeddings"
     )
+
+
+def test_validate_environment_preserves_custom_gateway_headers() -> None:
+    original: Final = {"X-Gateway-Token": "gateway-token", "X-Trace-ID": "trace-id"}
+    headers: Final = VoyageMultimodalEmbeddingConfig().validate_environment(
+        headers=original, model="test-model", messages=[], optional_params={}, litellm_params={}, api_key="provider-key"
+    )
+    assert headers["X-Gateway-Token"] == "gateway-token"
+    assert headers["X-Trace-ID"] == "trace-id"
+    assert headers["Authorization"] == "Bearer provider-key"
+    assert original == {"X-Gateway-Token": "gateway-token", "X-Trace-ID": "trace-id"}

tests/unit/llms/voyage/embedding/test_voyage_embedding_transformation.py

@@ -1,4 +1,5 @@
 import datetime
+from typing import Final
 
 import httpx
 import pytest
@@ -106,3 +107,14 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str):
         )
         == "https://gateway.example.com/voyage/v1/embeddings"
     )
+
+
+def test_validate_environment_preserves_custom_gateway_headers() -> None:
+    original: Final = {"X-Gateway-Token": "gateway-token", "X-Trace-ID": "trace-id"}
+    headers: Final = VoyageEmbeddingConfig().validate_environment(
+        headers=original, model="test-model", messages=[], optional_params={}, litellm_params={}, api_key="provider-key"
+    )
+    assert headers["X-Gateway-Token"] == "gateway-token"
+    assert headers["X-Trace-ID"] == "trace-id"
+    assert headers["Authorization"] == "Bearer provider-key"
+    assert original == {"X-Gateway-Token": "gateway-token", "X-Trace-ID": "trace-id"}

tests/unit/llms/voyage/rerank/test_voyage_rerank_transformation.py

@@ -4,6 +4,7 @@ Tests for Voyage AI rerank transformation functionality.
 
 import json
 import uuid
+from typing import Final
 from unittest.mock import MagicMock, patch
 
 import httpx
@@ -189,9 +190,7 @@ class TestVoyageRerankTransform:
         )
 
         assert len(result.results) == 2
-        assert (
-            result.results[0]["document"]["text"] == "Paris is the capital of France."
-        )
+        assert result.results[0]["document"]["text"] == "Paris is the capital of France."
         assert result.results[1]["document"]["text"] == "France is a country in Europe."
 
     def test_transform_rerank_response_missing_data(self):
@@ -229,7 +228,7 @@ class TestVoyageRerankTransform:
         mock_logging = MagicMock()
         model_response = RerankResponse()
 
-        with pytest.raises(Exception, match='Unauthorized') as exc_info:
+        with pytest.raises(Exception, match="Unauthorized") as exc_info:
             self.config.transform_rerank_response(
                 model=self.model,
                 raw_response=mock_response,
@@ -250,7 +249,7 @@ class TestVoyageRerankTransform:
         mock_logging = MagicMock()
         model_response = RerankResponse()
 
-        with pytest.raises(Exception, match='Failed to parse response: Invalid JSON response') as exc_info:
+        with pytest.raises(Exception, match="Failed to parse response: Invalid JSON response") as exc_info:
             self.config.transform_rerank_response(
                 model=self.model,
                 raw_response=mock_response,
@@ -421,3 +420,14 @@ def test_transform_rerank_response_shape_errors_do_not_echo_the_payload():
         _transform({"data": [["leaked document text"]]})
 
     assert "leaked document text" not in str(exc_info.value)
+
+
+def test_validate_environment_preserves_custom_gateway_headers() -> None:
+    original: Final = {"X-Gateway-Token": "gateway-token", "X-Trace-ID": "trace-id"}
+    headers: Final = VoyageRerankConfig().validate_environment(
+        headers=original, model="test-model", api_key="provider-key"
+    )
+    assert headers["X-Gateway-Token"] == "gateway-token"
+    assert headers["X-Trace-ID"] == "trace-id"
+    assert headers["Authorization"] == "Bearer provider-key"
+    assert original == {"X-Gateway-Token": "gateway-token", "X-Trace-ID": "trace-id"}