mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
test: migrate legacy provider tests to tests/unit (wave 2, phase 13)
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yuneng
parent
b9a2d441cf
commit
e1556ce32b
40 changed files with 94 additions and 621 deletions
|
|
@ -1,214 +0,0 @@
|
|||
"""
|
||||
Test Vertex AI files integration with main files API
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import litellm
|
||||
from litellm.types.llms.openai import HttpxBinaryResponseContent
|
||||
|
||||
|
||||
class TestVertexAIFilesIntegration:
|
||||
"""Test integration of Vertex AI files with main litellm API"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_afile_content_vertex_ai_provider(self):
|
||||
"""Test litellm.afile_content with vertex_ai provider"""
|
||||
file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt"
|
||||
expected_content = b"test file content"
|
||||
|
||||
# Create a mock HttpxBinaryResponseContent response
|
||||
import httpx
|
||||
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=expected_content,
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"),
|
||||
)
|
||||
mock_result = HttpxBinaryResponseContent(response=mock_response)
|
||||
|
||||
# Mock the base_llm_http_handler.retrieve_file_content since the code
|
||||
# now routes through ProviderConfigManager -> base_llm_http_handler
|
||||
with patch(
|
||||
"litellm.files.main.base_llm_http_handler.retrieve_file_content",
|
||||
new_callable=MagicMock,
|
||||
) as mock_retrieve:
|
||||
# Make it return a coroutine for async path
|
||||
mock_retrieve.return_value = mock_result
|
||||
|
||||
result = await litellm.afile_content(
|
||||
file_id=file_id,
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
vertex_credentials=None,
|
||||
)
|
||||
|
||||
# Verify the result
|
||||
assert isinstance(result, HttpxBinaryResponseContent)
|
||||
assert result.response.content == expected_content
|
||||
assert result.response.status_code == 200
|
||||
|
||||
# Verify the mock was called
|
||||
mock_retrieve.assert_called_once()
|
||||
|
||||
def test_litellm_file_content_vertex_ai_provider(self):
|
||||
"""Test litellm.file_content with vertex_ai provider (sync)"""
|
||||
file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt"
|
||||
expected_content = b"test file content"
|
||||
|
||||
# Create a mock HttpxBinaryResponseContent response
|
||||
import httpx
|
||||
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=expected_content,
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"),
|
||||
)
|
||||
mock_result = HttpxBinaryResponseContent(response=mock_response)
|
||||
|
||||
# Mock the base_llm_http_handler.retrieve_file_content
|
||||
with patch(
|
||||
"litellm.files.main.base_llm_http_handler.retrieve_file_content",
|
||||
return_value=mock_result,
|
||||
) as mock_retrieve:
|
||||
result = litellm.file_content(
|
||||
file_id=file_id,
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
vertex_credentials=None,
|
||||
)
|
||||
|
||||
# Verify the result
|
||||
assert isinstance(result, HttpxBinaryResponseContent)
|
||||
assert result.response.content == expected_content
|
||||
assert result.response.status_code == 200
|
||||
|
||||
# Verify the mock was called
|
||||
mock_retrieve.assert_called_once()
|
||||
|
||||
def test_litellm_file_content_vertex_ai_with_model_provider_detection(self):
|
||||
"""Test litellm.file_content with model parameter for provider detection"""
|
||||
file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt"
|
||||
expected_content = b"test file content"
|
||||
|
||||
# Create a mock HttpxBinaryResponseContent response
|
||||
import httpx
|
||||
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=expected_content,
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"),
|
||||
)
|
||||
mock_result = HttpxBinaryResponseContent(response=mock_response)
|
||||
|
||||
# Mock the base_llm_http_handler.retrieve_file_content
|
||||
with patch(
|
||||
"litellm.files.main.base_llm_http_handler.retrieve_file_content",
|
||||
return_value=mock_result,
|
||||
):
|
||||
# Mock get_llm_provider to return vertex_ai
|
||||
with patch("litellm.files.main.get_llm_provider") as mock_get_provider:
|
||||
mock_get_provider.return_value = (
|
||||
"vertex_ai/gemini-pro",
|
||||
"vertex_ai",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
|
||||
# Call litellm.file_content with model to trigger provider detection
|
||||
result = litellm.file_content(
|
||||
file_id=file_id,
|
||||
model="vertex_ai/gemini-pro",
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
)
|
||||
|
||||
# Verify the result
|
||||
assert isinstance(result, HttpxBinaryResponseContent)
|
||||
assert result.response.content == expected_content
|
||||
|
||||
# Verify provider detection was called
|
||||
mock_get_provider.assert_called_once()
|
||||
|
||||
def test_litellm_file_content_vertex_ai_error_cases(self):
|
||||
"""Test error handling in vertex_ai file_content"""
|
||||
# Test missing file_id - the VertexAI provider config's
|
||||
# transform_file_content_request should handle empty file_id.
|
||||
# Since the code now goes through base_llm_http_handler, we mock
|
||||
# ProviderConfigManager to return None so it falls through to the
|
||||
# old vertex_ai code path that validates file_id.
|
||||
with patch(
|
||||
"litellm.files.main.ProviderConfigManager.get_provider_files_config",
|
||||
return_value=None,
|
||||
):
|
||||
with pytest.raises(ValueError, match="file_id is required"):
|
||||
litellm.file_content(
|
||||
file_id="", # Empty file_id should cause error
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex_project="test-project",
|
||||
)
|
||||
|
||||
def test_vertex_ai_provider_in_supported_providers_list(self):
|
||||
"""Test that vertex_ai is included in supported providers for file_content"""
|
||||
# This test ensures the type annotations and error messages include vertex_ai
|
||||
|
||||
# Test that calling with unsupported provider raises appropriate error
|
||||
with pytest.raises(Exception, match="unsupported_provider' is not a valid LlmProviders") as exc_info:
|
||||
litellm.file_content(
|
||||
file_id="test-file-id",
|
||||
custom_llm_provider="unsupported_provider", # This should fail
|
||||
)
|
||||
|
||||
# The error message should mention supported providers including vertex_ai
|
||||
error_message = str(exc_info.value)
|
||||
assert "vertex_ai" in error_message or "supported" in error_message.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_ai_file_content_with_timeout_and_retries(self):
|
||||
"""Test vertex_ai file_content with timeout and retry configuration"""
|
||||
file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt"
|
||||
expected_content = b"test file content"
|
||||
|
||||
# Create a mock HttpxBinaryResponseContent response
|
||||
import httpx
|
||||
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=expected_content,
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"),
|
||||
)
|
||||
mock_result = HttpxBinaryResponseContent(response=mock_response)
|
||||
|
||||
# Mock the base_llm_http_handler.retrieve_file_content
|
||||
with patch(
|
||||
"litellm.files.main.base_llm_http_handler.retrieve_file_content",
|
||||
new_callable=MagicMock,
|
||||
) as mock_retrieve:
|
||||
mock_retrieve.return_value = mock_result
|
||||
|
||||
# Call with custom timeout and max_retries
|
||||
result = await litellm.afile_content(
|
||||
file_id=file_id,
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
timeout=120,
|
||||
max_retries=5,
|
||||
)
|
||||
|
||||
# Verify the result
|
||||
assert isinstance(result, HttpxBinaryResponseContent)
|
||||
assert result.response.content == expected_content
|
||||
|
||||
# Verify the mock was called
|
||||
mock_retrieve.assert_called_once()
|
||||
# Verify the timeout was passed through
|
||||
call_kwargs = mock_retrieve.call_args.kwargs
|
||||
assert call_kwargs["timeout"] == 120
|
||||
0
tests/unit/llms/tencent/__init__.py
Normal file
0
tests/unit/llms/tencent/__init__.py
Normal file
0
tests/unit/llms/tencent/messages/__init__.py
Normal file
0
tests/unit/llms/tencent/messages/__init__.py
Normal file
0
tests/unit/llms/together_ai/__init__.py
Normal file
0
tests/unit/llms/together_ai/__init__.py
Normal file
0
tests/unit/llms/together_ai/chat/__init__.py
Normal file
0
tests/unit/llms/together_ai/chat/__init__.py
Normal file
0
tests/unit/llms/valkey/__init__.py
Normal file
0
tests/unit/llms/valkey/__init__.py
Normal file
0
tests/unit/llms/valkey/vector_stores/__init__.py
Normal file
0
tests/unit/llms/valkey/vector_stores/__init__.py
Normal file
|
|
@ -1,8 +1,7 @@
|
|||
import struct
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import MagicMock
|
||||
from urllib.parse import unquote, urlsplit
|
||||
|
||||
import httpx
|
||||
|
|
@ -355,13 +354,6 @@ def test_search_treats_an_explicit_null_max_num_results_as_the_default():
|
|||
assert client.index.searched_query.query_string() == "*=>[KNN 10 @embedding $vec AS vector_distance]"
|
||||
|
||||
|
||||
def test_missing_redis_dependency_raises_actionable_error():
|
||||
config = ValkeyVectorStoreConfig(sync_client=FakeRedis(), embedding_fn=FakeEmbeddingFn([1.0]))
|
||||
blocked = {name: None for name in list(sys.modules) if name == "redis" or name.startswith("redis.")}
|
||||
|
||||
with patch.dict(sys.modules, blocked):
|
||||
with pytest.raises(ValueError, match="pip install redis"):
|
||||
_search(config)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
0
tests/unit/llms/vercel_ai_gateway/__init__.py
Normal file
0
tests/unit/llms/vercel_ai_gateway/__init__.py
Normal file
0
tests/unit/llms/vercel_ai_gateway/chat/__init__.py
Normal file
0
tests/unit/llms/vercel_ai_gateway/chat/__init__.py
Normal file
0
tests/unit/llms/vercel_ai_gateway/embedding/__init__.py
Normal file
0
tests/unit/llms/vercel_ai_gateway/embedding/__init__.py
Normal file
0
tests/unit/llms/vertex_ai/__init__.py
Normal file
0
tests/unit/llms/vertex_ai/__init__.py
Normal file
0
tests/unit/llms/vertex_ai/agent_engine/__init__.py
Normal file
0
tests/unit/llms/vertex_ai/agent_engine/__init__.py
Normal file
0
tests/unit/llms/vertex_ai/context_caching/__init__.py
Normal file
0
tests/unit/llms/vertex_ai/context_caching/__init__.py
Normal file
|
|
@ -9,7 +9,6 @@ from litellm.litellm_core_utils.litellm_logging import Logging
|
|||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError
|
||||
from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import (
|
||||
MAX_PAGINATION_PAGES,
|
||||
ContextCachingEndpoints,
|
||||
)
|
||||
|
||||
|
|
@ -1892,100 +1891,7 @@ class TestCheckCachePagination:
|
|||
assert result is None
|
||||
assert self.mock_async_client.get.call_count == 1
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
|
||||
)
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
def test_check_cache_pagination_max_pages_limit(
|
||||
self, mock_get_token_url, custom_llm_provider
|
||||
):
|
||||
"""Test that pagination stops after MAX_PAGINATION_PAGES iterations"""
|
||||
# Setup
|
||||
mock_get_token_url.return_value = ("token", "https://test-url.com")
|
||||
cache_key_to_find = "nonexistent_cache_key"
|
||||
|
||||
# Create mock response that always has nextPageToken (infinite pagination scenario)
|
||||
def create_page_response(page_num):
|
||||
response = MagicMock()
|
||||
response.json.return_value = {
|
||||
"cachedContents": [
|
||||
{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}
|
||||
],
|
||||
"nextPageToken": f"token_page_{page_num + 1}",
|
||||
}
|
||||
return response
|
||||
|
||||
# Create MAX_PAGINATION_PAGES responses, each with a nextPageToken
|
||||
self.mock_client.get.side_effect = [
|
||||
create_page_response(i) for i in range(MAX_PAGINATION_PAGES)
|
||||
]
|
||||
|
||||
# Execute
|
||||
result = self.context_caching.check_cache(
|
||||
cache_key=cache_key_to_find,
|
||||
client=self.mock_client,
|
||||
headers={"Authorization": "Bearer token"},
|
||||
api_key="test_key",
|
||||
api_base=None,
|
||||
logging_obj=self.mock_logging,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project="test_project",
|
||||
vertex_location="us-central1",
|
||||
vertex_auth_header="Bearer test-token",
|
||||
)
|
||||
|
||||
# Assert - should return None after exhausting all pages without finding match
|
||||
assert result is None
|
||||
# Verify exactly MAX_PAGINATION_PAGES API calls were made (not more)
|
||||
assert self.mock_client.get.call_count == MAX_PAGINATION_PAGES
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
|
||||
)
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
async def test_async_check_cache_pagination_max_pages_limit(
|
||||
self, mock_get_token_url, custom_llm_provider
|
||||
):
|
||||
"""Test that async pagination stops after MAX_PAGINATION_PAGES iterations"""
|
||||
# Setup
|
||||
mock_get_token_url.return_value = ("token", "https://test-url.com")
|
||||
cache_key_to_find = "nonexistent_cache_key"
|
||||
|
||||
# Create mock response that always has nextPageToken (infinite pagination scenario)
|
||||
def create_page_response(page_num):
|
||||
response = MagicMock()
|
||||
response.json.return_value = {
|
||||
"cachedContents": [
|
||||
{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}
|
||||
],
|
||||
"nextPageToken": f"token_page_{page_num + 1}",
|
||||
}
|
||||
return response
|
||||
|
||||
# Create MAX_PAGINATION_PAGES responses, each with a nextPageToken
|
||||
self.mock_async_client.get = AsyncMock(
|
||||
side_effect=[create_page_response(i) for i in range(MAX_PAGINATION_PAGES)]
|
||||
)
|
||||
|
||||
# Execute
|
||||
result = await self.context_caching.async_check_cache(
|
||||
cache_key=cache_key_to_find,
|
||||
client=self.mock_async_client,
|
||||
headers={"Authorization": "Bearer token"},
|
||||
api_key="test_key",
|
||||
api_base=None,
|
||||
logging_obj=self.mock_logging,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project="test_project",
|
||||
vertex_location="us-central1",
|
||||
vertex_auth_header="Bearer test-token",
|
||||
)
|
||||
|
||||
# Assert - should return None after exhausting all pages without finding match
|
||||
assert result is None
|
||||
# Verify exactly MAX_PAGINATION_PAGES async API calls were made (not more)
|
||||
assert self.mock_async_client.get.call_count == MAX_PAGINATION_PAGES
|
||||
|
||||
|
||||
class TestVertexAIGlobalLocation:
|
||||
0
tests/unit/llms/vertex_ai/files/__init__.py
Normal file
0
tests/unit/llms/vertex_ai/files/__init__.py
Normal file
|
|
@ -11,8 +11,6 @@ import io
|
|||
import json
|
||||
import pytest
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.custom_httpx.llm_http_handler import AsyncHTTPHandler
|
||||
from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig
|
||||
from litellm.types.llms.openai import CreateFileRequest
|
||||
|
|
@ -96,39 +94,6 @@ class TestVertexAIBinaryFileUpload:
|
|||
assert isinstance(transformed_request, bytes)
|
||||
assert transformed_request == mock_png_content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_handler_accepts_bytes_without_decoding(self):
|
||||
"""
|
||||
Test that httpx correctly accepts binary data without decoding.
|
||||
|
||||
This test verifies that bytes can be passed to httpx's post/put methods
|
||||
without needing UTF-8 decoding, which is the core of our fix.
|
||||
"""
|
||||
# Create mock binary data with non-UTF-8 bytes
|
||||
mock_binary_data = b"\x00\x01\x02\x03\xff\xfe\xfd\xc4\xe5\xf2"
|
||||
|
||||
# Test that httpx accepts bytes in the data parameter
|
||||
# We're testing the behavior, not making an actual request
|
||||
|
||||
# Verify that attempting to decode would fail (proving it's binary)
|
||||
with pytest.raises(UnicodeDecodeError):
|
||||
mock_binary_data.decode("utf-8")
|
||||
|
||||
# Verify that httpx Request accepts bytes
|
||||
try:
|
||||
request = httpx.Request(
|
||||
method="POST",
|
||||
url="https://example.com/upload",
|
||||
data=mock_binary_data,
|
||||
headers={"Content-Type": "application/octet-stream"},
|
||||
)
|
||||
# If we get here, httpx accepts bytes - which is what we need
|
||||
assert request.content == mock_binary_data
|
||||
except Exception as e:
|
||||
pytest.fail(f"httpx should accept bytes in data parameter: {e}")
|
||||
|
||||
# Document the expected behavior
|
||||
assert isinstance(mock_binary_data, bytes), "Binary file data should remain as bytes"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_jsonl_file_upload_returns_streaming_body(self):
|
||||
|
|
@ -224,36 +189,3 @@ class TestVertexAIBinaryFileUpload:
|
|||
litellm_params={},
|
||||
)
|
||||
assert isinstance(result3, bytes)
|
||||
|
||||
def test_bytes_type_preservation_documentation(self):
|
||||
"""
|
||||
Documentation test: Verify that bytes are the correct type for binary uploads.
|
||||
|
||||
This test documents the expected behavior:
|
||||
- Binary files (PDF, images, etc.) should remain as bytes
|
||||
- Text files (JSONL) should be strings
|
||||
- httpx accepts both bytes and strings in the 'data' parameter
|
||||
- bytes should NEVER be decoded to UTF-8 for binary files
|
||||
"""
|
||||
# This is a documentation test - it always passes
|
||||
# but serves as a reference for the expected behavior
|
||||
|
||||
expected_behavior = {
|
||||
"binary_files": {
|
||||
"input_type": "bytes",
|
||||
"output_type": "bytes",
|
||||
"examples": ["PDF", "PNG", "JPEG", "binary data"],
|
||||
"http_method": "POST or PUT",
|
||||
"encoding": "none - preserve raw bytes",
|
||||
},
|
||||
"text_files": {
|
||||
"input_type": "str or bytes",
|
||||
"output_type": "bytes",
|
||||
"examples": ["JSONL", "CSV", "TXT"],
|
||||
"http_method": "POST",
|
||||
"encoding": "UTF-8",
|
||||
},
|
||||
}
|
||||
|
||||
assert expected_behavior["binary_files"]["encoding"] == "none - preserve raw bytes"
|
||||
assert expected_behavior["text_files"]["encoding"] == "UTF-8"
|
||||
|
|
@ -2,14 +2,11 @@
|
|||
Test Vertex AI files handler functionality
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
from types import MappingProxyType
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.vertex_ai.files.handler import VertexAIFilesHandler
|
||||
from litellm.types.llms.openai import FileContentRequest, HttpxBinaryResponseContent
|
||||
|
||||
|
|
@ -312,104 +309,3 @@ class TestVertexAIFilesHandler:
|
|||
assert isinstance(result, HttpxBinaryResponseContent)
|
||||
dynamic_params = mock_download.call_args.kwargs["standard_callback_dynamic_params"]
|
||||
assert dynamic_params["gcs_bucket_name"] == "my-model-bucket"
|
||||
|
||||
def test_file_content_sync_success(self):
|
||||
"""Test successful sync file content retrieval"""
|
||||
file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt"
|
||||
expected_content = b"test file content"
|
||||
|
||||
file_content_request = FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None)
|
||||
|
||||
# Create expected response
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=expected_content,
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"),
|
||||
)
|
||||
expected_result = HttpxBinaryResponseContent(response=mock_response)
|
||||
|
||||
# Mock asyncio.run to return our expected result
|
||||
with patch("asyncio.run") as mock_run:
|
||||
mock_run.return_value = expected_result
|
||||
|
||||
result = self.handler.file_content(
|
||||
_is_async=False,
|
||||
file_content_request=file_content_request,
|
||||
api_base="",
|
||||
vertex_credentials=None,
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
timeout=60.0,
|
||||
max_retries=3,
|
||||
)
|
||||
|
||||
# Verify the result
|
||||
assert result == expected_result
|
||||
|
||||
# Verify asyncio.run was called (indicating sync execution)
|
||||
mock_run.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_content_async_mode(self):
|
||||
"""Test async file content retrieval when _is_async=True"""
|
||||
file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt"
|
||||
expected_content = b"test file content"
|
||||
|
||||
file_content_request = FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None)
|
||||
|
||||
# Mock the afile_content method
|
||||
with patch.object(self.handler, "afile_content", new_callable=AsyncMock) as mock_afile_content:
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=expected_content,
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"),
|
||||
)
|
||||
mock_afile_content.return_value = HttpxBinaryResponseContent(response=mock_response)
|
||||
|
||||
# Call the method with _is_async=True
|
||||
result = self.handler.file_content(
|
||||
_is_async=True,
|
||||
file_content_request=file_content_request,
|
||||
api_base="",
|
||||
vertex_credentials=None,
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
timeout=60.0,
|
||||
max_retries=3,
|
||||
)
|
||||
|
||||
# Should return a coroutine since _is_async=True
|
||||
assert asyncio.iscoroutine(result)
|
||||
|
||||
# Await the result
|
||||
final_result = await result
|
||||
assert isinstance(final_result, HttpxBinaryResponseContent)
|
||||
assert final_result.response.content == expected_content
|
||||
|
||||
def test_httpx_response_compatibility(self):
|
||||
"""Test that the created HttpxBinaryResponseContent is compatible with expected interface"""
|
||||
# Test the mock response creation logic
|
||||
expected_content = b"test file content"
|
||||
decoded_path = "gs://test-bucket/test-file.txt"
|
||||
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=expected_content,
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request(method="GET", url=decoded_path),
|
||||
)
|
||||
|
||||
result = HttpxBinaryResponseContent(response=mock_response)
|
||||
|
||||
# Verify the response properties
|
||||
assert result.response.status_code == 200
|
||||
assert result.response.content == expected_content
|
||||
assert result.response.headers["content-type"] == "application/octet-stream"
|
||||
|
||||
# Verify it has the expected interface (matching OpenAI file content response)
|
||||
assert hasattr(result, "response")
|
||||
assert hasattr(result.response, "content")
|
||||
assert hasattr(result.response, "status_code")
|
||||
assert hasattr(result.response, "headers")
|
||||
|
|
@ -0,0 +1,93 @@
|
|||
"""
|
||||
Test Vertex AI files integration with main files API
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import litellm
|
||||
from litellm.types.llms.openai import HttpxBinaryResponseContent
|
||||
|
||||
|
||||
class TestVertexAIFilesIntegration:
|
||||
"""Test integration of Vertex AI files with main litellm API"""
|
||||
|
||||
|
||||
|
||||
|
||||
def test_litellm_file_content_vertex_ai_error_cases(self):
|
||||
"""Test error handling in vertex_ai file_content"""
|
||||
# Test missing file_id - the VertexAI provider config's
|
||||
# transform_file_content_request should handle empty file_id.
|
||||
# Since the code now goes through base_llm_http_handler, we mock
|
||||
# ProviderConfigManager to return None so it falls through to the
|
||||
# old vertex_ai code path that validates file_id.
|
||||
with patch(
|
||||
"litellm.files.main.ProviderConfigManager.get_provider_files_config",
|
||||
return_value=None,
|
||||
):
|
||||
with pytest.raises(ValueError, match="file_id is required"):
|
||||
litellm.file_content(
|
||||
file_id="", # Empty file_id should cause error
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex_project="test-project",
|
||||
)
|
||||
|
||||
def test_vertex_ai_provider_in_supported_providers_list(self):
|
||||
"""Test that vertex_ai is included in supported providers for file_content"""
|
||||
# This test ensures the type annotations and error messages include vertex_ai
|
||||
|
||||
# Test that calling with unsupported provider raises appropriate error
|
||||
with pytest.raises(Exception, match="unsupported_provider' is not a valid LlmProviders") as exc_info:
|
||||
litellm.file_content(
|
||||
file_id="test-file-id",
|
||||
custom_llm_provider="unsupported_provider", # This should fail
|
||||
)
|
||||
|
||||
# The error message should mention supported providers including vertex_ai
|
||||
error_message = str(exc_info.value)
|
||||
assert "vertex_ai" in error_message or "supported" in error_message.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_ai_file_content_with_timeout_and_retries(self):
|
||||
"""Test vertex_ai file_content with timeout and retry configuration"""
|
||||
file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt"
|
||||
expected_content = b"test file content"
|
||||
|
||||
# Create a mock HttpxBinaryResponseContent response
|
||||
import httpx
|
||||
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=expected_content,
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"),
|
||||
)
|
||||
mock_result = HttpxBinaryResponseContent(response=mock_response)
|
||||
|
||||
# Mock the base_llm_http_handler.retrieve_file_content
|
||||
with patch(
|
||||
"litellm.files.main.base_llm_http_handler.retrieve_file_content",
|
||||
new_callable=MagicMock,
|
||||
) as mock_retrieve:
|
||||
mock_retrieve.return_value = mock_result
|
||||
|
||||
# Call with custom timeout and max_retries
|
||||
result = await litellm.afile_content(
|
||||
file_id=file_id,
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
timeout=120,
|
||||
max_retries=5,
|
||||
)
|
||||
|
||||
# Verify the result
|
||||
assert isinstance(result, HttpxBinaryResponseContent)
|
||||
assert result.response.content == expected_content
|
||||
|
||||
# Verify the mock was called
|
||||
mock_retrieve.assert_called_once()
|
||||
# Verify the timeout was passed through
|
||||
call_kwargs = mock_retrieve.call_args.kwargs
|
||||
assert call_kwargs["timeout"] == 120
|
||||
|
|
@ -860,94 +860,6 @@ class TestVertexBatchOutputTransformation:
|
|||
binary = b"%PDF-1.4\n%\xc4\xe5\xf2\xe5\xeb\xa7\n" + b"\x00\x01\x02\xff\xfe" * 64
|
||||
assert config._try_transform_vertex_batch_output_to_openai(binary) == binary
|
||||
|
||||
def test_streaming_transform_peaks_below_list_pipeline(self, config):
|
||||
"""The output transform must stream row-by-row, not build a list of every
|
||||
parsed row and a second list of transformed rows. This guards against a
|
||||
regression to the list pipeline, which peaks at several full copies and
|
||||
OOMs on large result files. The relative comparison cancels shared noise
|
||||
(per-row transform cost, GC timing) and only the list overhead differs.
|
||||
"""
|
||||
import gc
|
||||
import tracemalloc
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexGeminiConfig,
|
||||
)
|
||||
|
||||
def vertex_row(index: int) -> dict:
|
||||
return {
|
||||
"status": "",
|
||||
"processed_time": "2024-11-01T18:13:16.826+00:00",
|
||||
"request": {
|
||||
"contents": [{"role": "user", "parts": [{"text": "hi"}]}],
|
||||
"labels": {"litellm_custom_id": f"r-{index}"},
|
||||
},
|
||||
"response": {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {
|
||||
"parts": [{"text": "hello " * 20}],
|
||||
"role": "model",
|
||||
},
|
||||
"finishReason": "STOP",
|
||||
}
|
||||
],
|
||||
"modelVersion": "gemini-2.0-flash-001",
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 10,
|
||||
"candidatesTokenCount": 20,
|
||||
"totalTokenCount": 30,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
content = ("\n".join(json.dumps(vertex_row(i)) for i in range(4000))).encode("utf-8")
|
||||
|
||||
def list_pipeline() -> bytes:
|
||||
gemini_config = VertexGeminiConfig()
|
||||
logging_obj = Logging(
|
||||
model="",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="batch_transform",
|
||||
start_time=0.1,
|
||||
litellm_call_id="",
|
||||
function_id="",
|
||||
)
|
||||
logging_obj.optional_params = {}
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
headers={"content-type": "application/json"},
|
||||
request=httpx.Request("POST", "https://example.com"),
|
||||
)
|
||||
rows = content.decode("utf-8").strip().split("\n")
|
||||
transformed = [
|
||||
json.dumps(
|
||||
config._transform_single_vertex_batch_output_to_openai(
|
||||
json.loads(row), gemini_config, logging_obj, mock_response
|
||||
)
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
return "\n".join(transformed).encode("utf-8")
|
||||
|
||||
def peak_of(fn) -> int:
|
||||
gc.collect()
|
||||
tracemalloc.start()
|
||||
try:
|
||||
fn()
|
||||
return tracemalloc.get_traced_memory()[1]
|
||||
finally:
|
||||
tracemalloc.stop()
|
||||
|
||||
streaming_peak = peak_of(lambda: config._try_transform_vertex_batch_output_to_openai(content))
|
||||
list_peak = peak_of(list_pipeline)
|
||||
|
||||
assert streaming_peak < list_peak * 0.75, (
|
||||
f"streaming peak {streaming_peak} is not a clear win over the list "
|
||||
f"pipeline {list_peak} (ratio {streaming_peak / list_peak:.2f})"
|
||||
)
|
||||
|
||||
|
||||
class TestTryTransformDoesNotMutateCallerLoggingObj:
|
||||
0
tests/unit/llms/vertex_ai/gemini_embeddings/__init__.py
Normal file
0
tests/unit/llms/vertex_ai/gemini_embeddings/__init__.py
Normal file
0
tests/unit/llms/vertex_ai/image_edit/__init__.py
Normal file
0
tests/unit/llms/vertex_ai/image_edit/__init__.py
Normal file
0
tests/unit/llms/vertex_ai/interactions/__init__.py
Normal file
0
tests/unit/llms/vertex_ai/interactions/__init__.py
Normal file
0
tests/unit/llms/vertex_ai/realtime/__init__.py
Normal file
0
tests/unit/llms/vertex_ai/realtime/__init__.py
Normal file
|
|
@ -387,50 +387,6 @@ def test_vertex_does_not_warn_when_dropping_non_guardrail_session_update(caplog)
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_realtime_does_not_forward_client_query_params_to_vertex_backend(
|
||||
monkeypatch,
|
||||
):
|
||||
"""Regression: forwarding client ?model=/?intent= to the Vertex Live WSS URL causes 1007 errors.
|
||||
|
||||
Exercises ``async_realtime`` end-to-end so that re-adding ``_append_query_params``
|
||||
(the reverted bug) would push ``model=``/``intent=`` onto the backend URL and fail here.
|
||||
"""
|
||||
import websockets
|
||||
|
||||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
|
||||
cfg = VertexAIRealtimeConfig(
|
||||
access_token="tok", project="my-proj", location="us-central1"
|
||||
)
|
||||
|
||||
captured: dict = {}
|
||||
|
||||
def fake_connect(url, *args, **kwargs):
|
||||
captured["url"] = url
|
||||
raise RuntimeError("stop before establishing the backend connection")
|
||||
|
||||
monkeypatch.setattr(websockets, "connect", fake_connect)
|
||||
|
||||
try:
|
||||
await BaseLLMHTTPHandler().async_realtime(
|
||||
model="gemini-live-2.5-flash-preview-native-audio-09-2025",
|
||||
websocket=AsyncMock(),
|
||||
logging_obj=MagicMock(),
|
||||
provider_config=cfg,
|
||||
headers={},
|
||||
query_params={
|
||||
"model": "gemini-live-2.5-flash-preview-native-audio-09-2025",
|
||||
"intent": "chat",
|
||||
},
|
||||
)
|
||||
except (RuntimeError, Exception):
|
||||
pass
|
||||
|
||||
assert "url" in captured, "websockets.connect was never called"
|
||||
assert "?" not in captured["url"]
|
||||
assert "model=" not in captured["url"]
|
||||
assert "intent=" not in captured["url"]
|
||||
|
||||
|
||||
def test_vertex_function_call_output_omits_id():
|
||||
0
tests/unit/llms/vertex_ai/text_to_speech/__init__.py
Normal file
0
tests/unit/llms/vertex_ai/text_to_speech/__init__.py
Normal file
Loading…
Add table
Reference in a new issue