style: Apply Black formatting to pass CI linting check

Co-authored-by: l.tingting@pku.edu.cn
This commit is contained in:
GeGeeWhy 2026-03-27 00:07:16 +08:00
parent be4d86e60d
commit 1e55d35c48
6 changed files with 123 additions and 87 deletions

View file

@ -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

View file

@ -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,

View file

@ -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.

View file

@ -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

View file

@ -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}")
pytest.fail(f"Citation metadata transformation failed: {e}")

View file

@ -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: