diff --git a/litellm/llms/voyage/embedding/transformation.py b/litellm/llms/voyage/embedding/transformation.py index 1ba66e5ce65..ccdc62519e2 100644 --- a/litellm/llms/voyage/embedding/transformation.py +++ b/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)}", } diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index 62c03139fbe..093eb6961b8 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/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)}", } diff --git a/litellm/llms/voyage/embedding/transformation_multimodal.py b/litellm/llms/voyage/embedding/transformation_multimodal.py index c04a2e1f26b..23864d89dca 100644 --- a/litellm/llms/voyage/embedding/transformation_multimodal.py +++ b/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") diff --git a/litellm/llms/voyage/rerank/transformation.py b/litellm/llms/voyage/rerank/transformation.py index 21c6e6a82b9..cc6a23643d8 100644 --- a/litellm/llms/voyage/rerank/transformation.py +++ b/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", } diff --git a/tests/unit/llms/voyage/embedding/test_transformation_contextual.py b/tests/unit/llms/voyage/embedding/test_transformation_contextual.py index c64f899a9f4..22f69d2fa8b 100644 --- a/tests/unit/llms/voyage/embedding/test_transformation_contextual.py +++ b/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"} diff --git a/tests/unit/llms/voyage/embedding/test_transformation_multimodal.py b/tests/unit/llms/voyage/embedding/test_transformation_multimodal.py index 4b8d27bba7d..fe05254fe2e 100644 --- a/tests/unit/llms/voyage/embedding/test_transformation_multimodal.py +++ b/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"} diff --git a/tests/unit/llms/voyage/embedding/test_voyage_embedding_transformation.py b/tests/unit/llms/voyage/embedding/test_voyage_embedding_transformation.py index cfbc8dc24fb..a61397e62d8 100644 --- a/tests/unit/llms/voyage/embedding/test_voyage_embedding_transformation.py +++ b/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"} diff --git a/tests/unit/llms/voyage/rerank/test_voyage_rerank_transformation.py b/tests/unit/llms/voyage/rerank/test_voyage_rerank_transformation.py index 15689f43bd0..772efced8ec 100644 --- a/tests/unit/llms/voyage/rerank/test_voyage_rerank_transformation.py +++ b/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"}