diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index 3771eacccd0..886c60b1102 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -1,4 +1,14 @@ -from typing import Any, Iterator, AsyncIterator, Coroutine, Dict, List, Optional, Union, cast +from typing import ( + Any, + Iterator, + AsyncIterator, + Coroutine, + Dict, + List, + Optional, + Union, + cast, +) import litellm from litellm.types.router import GenericLiteLLMParams diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 1818515ea87..a6bcf230e4f 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -321,7 +321,7 @@ class GoogleGenAIAdapter: ) # Return the SSE-wrapped version for proper event formatting return google_genai_wrapper.async_google_genai_sse_wrapper() - + def sync_translate_completion_output_params_streaming( self, completion_stream: Any, @@ -330,7 +330,7 @@ class GoogleGenAIAdapter: google_genai_wrapper = GoogleGenAIStreamWrapper( completion_stream=completion_stream ) - return google_genai_wrapper.google_genai_sse_wrapper() + return google_genai_wrapper.google_genai_sse_wrapper() def _transform_google_genai_tools_to_openai( self, diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 92c038e112e..d0f0c51889d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -229,7 +229,7 @@ class AnthropicAdapter: ) # Return the SSE-wrapped version for proper event formatting return anthropic_wrapper.async_anthropic_sse_wrapper() - + def sync_translate_completion_output_params_streaming( self, completion_stream: Any, @@ -568,9 +568,9 @@ class LiteLLMAnthropicMessagesAdapter: ## ASSISTANT MESSAGE ## assistant_message_str: Optional[str] = None - assistant_content_list: List[ - Dict[str, Any] - ] = [] # For content blocks with cache_control + assistant_content_list: List[Dict[str, Any]] = ( + [] + ) # For content blocks with cache_control has_cache_control_in_text = False tool_calls: List[ChatCompletionAssistantToolCall] = [] thinking_blocks: List[ @@ -681,7 +681,7 @@ class LiteLLMAnthropicMessagesAdapter: @staticmethod def translate_anthropic_thinking_to_reasoning_effort( - thinking: Dict[str, Any] + thinking: Dict[str, Any], ) -> Optional[str]: """ Translate Anthropic's thinking parameter to OpenAI's reasoning_effort. diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 74cf6844197..9e60aa065f3 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -8,7 +8,17 @@ import asyncio import contextvars from functools import partial -from typing import Any, Iterator, AsyncIterator, Coroutine, Dict, List, Optional, Union, cast +from typing import ( + Any, + Iterator, + AsyncIterator, + Coroutine, + Dict, + List, + Optional, + Union, + cast, +) import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj diff --git a/tests/test_litellm/google_genai/test_google_genai_handler.py b/tests/test_litellm/google_genai/test_google_genai_handler.py index 17a6cba2d63..64a354b6672 100644 --- a/tests/test_litellm/google_genai/test_google_genai_handler.py +++ b/tests/test_litellm/google_genai/test_google_genai_handler.py @@ -2,6 +2,7 @@ """ Test to verify the Google GenAI generate_content handler functionality """ + import json import os import sys @@ -32,24 +33,21 @@ def test_non_stream_response_when_stream_requested_sync(): choices=[ Choices( index=0, - message={ - "role": "assistant", - "content": "Hello, world!" - }, - finish_reason="stop" + message={"role": "assistant", "content": "Hello, world!"}, + finish_reason="stop", ) ], created=1234567890, model="gpt-3.5-turbo", - object="chat.completion" + object="chat.completion", ) - + # Create an instance of the adapter adapter = GoogleGenAIAdapter() - + # Test the adapter's translate_completion_to_generate_content method directly result = adapter.translate_completion_to_generate_content(mock_response) - + # Verify the result is a valid Google GenAI format response assert "candidates" in result assert isinstance(result["candidates"], list) @@ -77,24 +75,21 @@ async def test_non_stream_response_when_stream_requested_async(): choices=[ Choices( index=0, - message={ - "role": "assistant", - "content": "Hello, world!" - }, - finish_reason="stop" + message={"role": "assistant", "content": "Hello, world!"}, + finish_reason="stop", ) ], created=1234567890, model="gpt-3.5-turbo", - object="chat.completion" + object="chat.completion", ) - + # Create an instance of the adapter adapter = GoogleGenAIAdapter() - + # Test the adapter's translate_completion_to_generate_content method directly result = adapter.translate_completion_to_generate_content(mock_response) - + # Verify the result is a valid Google GenAI format response assert "candidates" in result assert isinstance(result["candidates"], list) @@ -116,12 +111,12 @@ def test_stream_response_when_stream_requested_sync(): # Mock a stream response mock_stream = MagicMock() mock_stream.__iter__ = MagicMock(return_value=iter([])) - - # Mock the GoogleGenAIAdapter's translate_completion_output_params_streaming method + + # Mock the GoogleGenAIAdapter's sync_translate_completion_output_params_streaming method with patch.object( - GoogleGenAIAdapter, - "translate_completion_output_params_streaming", - return_value=mock_stream + GoogleGenAIAdapter, + "sync_translate_completion_output_params_streaming", + return_value=mock_stream, ) as mock_translate: with patch("litellm.completion", return_value=mock_stream): # Call the handler with stream=True @@ -129,10 +124,10 @@ def test_stream_response_when_stream_requested_sync(): model="gemini-pro", contents=[{"role": "user", "parts": [{"text": "Hello"}]}], litellm_params={}, # Empty dict for params - stream=True + stream=True, ) - - # Verify that translate_completion_output_params_streaming was called + + # Verify that sync_translate_completion_output_params_streaming was called mock_translate.assert_called_once_with(mock_stream) # Verify the result is the transformed stream assert result == mock_stream @@ -146,23 +141,27 @@ async def test_stream_response_when_stream_requested_async(): """ # Mock a stream response mock_stream = MagicMock() - mock_stream.__aiter__ = AsyncMock(return_value=iter([])) # Return an empty async iterator - + mock_stream.__aiter__ = AsyncMock( + return_value=iter([]) + ) # Return an empty async iterator + # Mock the GoogleGenAIAdapter's translate_completion_output_params_streaming method with patch.object( - GoogleGenAIAdapter, - "translate_completion_output_params_streaming", - return_value=mock_stream + GoogleGenAIAdapter, + "translate_completion_output_params_streaming", + return_value=mock_stream, ) as mock_translate: with patch("litellm.acompletion", return_value=mock_stream): # Call the handler with stream=True - result = await GenerateContentToCompletionHandler.async_generate_content_handler( - model="gemini-pro", - contents=[{"role": "user", "parts": [{"text": "Hello"}]}], - litellm_params={}, # Empty dict for params - stream=True + result = ( + await GenerateContentToCompletionHandler.async_generate_content_handler( + model="gemini-pro", + contents=[{"role": "user", "parts": [{"text": "Hello"}]}], + litellm_params={}, # Empty dict for params + stream=True, + ) ) - + # Verify that translate_completion_output_params_streaming was called mock_translate.assert_called_once_with(mock_stream) # Verify the result is the transformed stream @@ -176,22 +175,24 @@ def test_stream_transformation_error_sync(): # Mock a stream response mock_stream = MagicMock() mock_stream.__iter__ = MagicMock(return_value=iter([])) - - # Mock the GoogleGenAIAdapter's translate_completion_output_params_streaming method to return None + + # Mock the GoogleGenAIAdapter's sync_translate_completion_output_params_streaming method to return None with patch.object( - GoogleGenAIAdapter, - "translate_completion_output_params_streaming", - return_value=None + GoogleGenAIAdapter, + "sync_translate_completion_output_params_streaming", + return_value=None, ): # Patch litellm.completion directly to prevent real API calls with patch("litellm.completion", return_value=mock_stream): # Call the handler with stream=True and expect a ValueError - with pytest.raises(ValueError, match="Failed to transform streaming response"): + with pytest.raises( + ValueError, match="Failed to transform streaming response" + ): GenerateContentToCompletionHandler.generate_content_handler( model="gemini-pro", contents=[{"role": "user", "parts": [{"text": "Hello"}]}], litellm_params={}, # Empty dict for params - stream=True + stream=True, ) @@ -203,12 +204,12 @@ async def test_stream_transformation_error_async(): # Mock a stream response mock_stream = MagicMock() mock_stream.__aiter__ = AsyncMock(return_value=mock_stream) - + # Mock the GoogleGenAIAdapter's translate_completion_output_params_streaming method to return None with patch.object( - GoogleGenAIAdapter, - "translate_completion_output_params_streaming", - return_value=None + GoogleGenAIAdapter, + "translate_completion_output_params_streaming", + return_value=None, ): # Mock litellm.acompletion at the module level where it's imported # We need to patch it in the handler module, not in litellm itself @@ -216,12 +217,14 @@ async def test_stream_transformation_error_async(): # Use AsyncMock for async function mock_litellm.acompletion = AsyncMock(return_value=mock_stream) # Call the handler with stream=True and expect a ValueError - with pytest.raises(ValueError, match="Failed to transform streaming response"): + with pytest.raises( + ValueError, match="Failed to transform streaming response" + ): await GenerateContentToCompletionHandler.async_generate_content_handler( model="gemini-pro", contents=[{"role": "user", "parts": [{"text": "Hello"}]}], litellm_params={}, # Empty dict for params - stream=True + stream=True, ) @@ -247,7 +250,7 @@ def test_citation_metadata_transformation(): "text": "This is a video analysis response with citation metadata." } ], - "role": "model" + "role": "model", }, "finishReason": "STOP", "index": 0, @@ -260,7 +263,7 @@ def test_citation_metadata_transformation(): "uri": "https://example.com/video-source", "license": "MIT", "title": "Video Analysis Source", - "publicationDate": "2024-01-15" + "publicationDate": "2024-01-15", }, { "startIndex": 6200, @@ -268,26 +271,26 @@ def test_citation_metadata_transformation(): "uri": "https://another-source.com/reference", "license": "CC-BY", "title": "Another Reference", - "publicationDate": "2024-02-01" - } + "publicationDate": "2024-02-01", + }, ] - } + }, } ], "usageMetadata": { "promptTokenCount": 150, "candidatesTokenCount": 200, - "totalTokenCount": 350 + "totalTokenCount": 350, }, - "responseId": "test-response-123" + "responseId": "test-response-123", } - + # Create mock httpx response mock_httpx_response = MagicMock(spec=httpx.Response) mock_httpx_response.json.return_value = mock_response_data mock_httpx_response.status_code = 200 mock_httpx_response.headers = {} - + # Create logging object logging_obj = LiteLLMLoggingObj( model="gemini-2.5-flash", @@ -296,40 +299,53 @@ def test_citation_metadata_transformation(): call_type="generate_content", start_time=1234567890, litellm_call_id="test-call-123", - function_id="test-function-123" + function_id="test-function-123", ) - + # Create GoogleGenAI config config = GoogleGenAIConfig() - + # Test the transformation try: result = config.transform_generate_content_response( model="gemini-2.5-flash", raw_response=mock_httpx_response, - logging_obj=logging_obj + logging_obj=logging_obj, ) - + # Verify the transformation worked assert result is not None - + # Check that citationSources was transformed to citations - if hasattr(result, 'candidates') and result.candidates: + if hasattr(result, "candidates") and result.candidates: candidate = result.candidates[0] - if hasattr(candidate, 'citationMetadata') and candidate.citationMetadata: + if hasattr(candidate, "citationMetadata") and candidate.citationMetadata: # The citationMetadata should now have 'citations' instead of 'citationSources' citation_metadata = candidate.citationMetadata - + # Check that citations field exists - assert hasattr(citation_metadata, 'citations'), "citations field should exist after transformation" - + assert hasattr( + citation_metadata, "citations" + ), "citations field should exist after transformation" + # Verify the citations data is preserved - if hasattr(citation_metadata, 'citations') and citation_metadata.citations: - assert len(citation_metadata.citations) == 2, "Should have 2 citations" - assert citation_metadata.citations[0]['uri'] == "https://example.com/video-source" - assert citation_metadata.citations[1]['uri'] == "https://another-source.com/reference" - + if ( + hasattr(citation_metadata, "citations") + and citation_metadata.citations + ): + assert ( + len(citation_metadata.citations) == 2 + ), "Should have 2 citations" + assert ( + citation_metadata.citations[0]["uri"] + == "https://example.com/video-source" + ) + assert ( + citation_metadata.citations[1]["uri"] + == "https://another-source.com/reference" + ) + print("✅ Citation metadata transformation test passed!") - + except Exception as e: - pytest.fail(f"Citation metadata transformation failed: {e}") \ No newline at end of file + pytest.fail(f"Citation metadata transformation failed: {e}") diff --git a/tests/test_litellm/test_sync_streaming_adapters.py b/tests/test_litellm/test_sync_streaming_adapters.py index 8209fd3b63b..3eb93c63ae6 100644 --- a/tests/test_litellm/test_sync_streaming_adapters.py +++ b/tests/test_litellm/test_sync_streaming_adapters.py @@ -230,7 +230,7 @@ def test_google_genai_sync_translate_sse_contains_candidates(): for chunk in result: chunk_str = chunk.decode("utf-8") if chunk_str.startswith("data: "): - json_str = chunk_str[len("data: "):].strip() + json_str = chunk_str[len("data: ") :].strip() if json_str: data = json.loads(json_str) if "candidates" in data: