mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-11 21:46:22 +00:00
Merge pull request #23474 from BerriAI/litellm_litellm-ci-stability-c0eb
Litellm ci stability
This commit is contained in:
yuneng-jiang
@@ -202,11 +202,6 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
|
||||
if "reasoning_effort" in optional_params:
|
||||
optional_params["reasoning_effort"] = normalized
|
||||
|
||||
reasoning_effort = (
|
||||
non_default_params.get("reasoning_effort")
|
||||
or optional_params.get("reasoning_effort")
|
||||
or raw_reasoning_effort
|
||||
)
|
||||
if effective_effort is not None and effective_effort == "xhigh":
|
||||
if not self._supports_reasoning_effort_level(model, "xhigh"):
|
||||
if litellm.drop_params or drop_params:
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import importlib
|
||||
from datetime import datetime, timezone
|
||||
from datetime import datetime
|
||||
from typing import Any, Awaitable, Callable, Dict, List, Optional, Set, Union
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Request
|
||||
@@ -360,6 +360,103 @@ if MCP_AVAILABLE:
|
||||
allowed_mcp_servers.append(server)
|
||||
return allowed_mcp_servers
|
||||
|
||||
async def _list_tools_for_single_server(
|
||||
server_id: str,
|
||||
allowed_server_ids: List[str],
|
||||
rest_client_ip: Optional[str],
|
||||
mcp_server_auth_headers: dict,
|
||||
mcp_auth_header: Optional[str],
|
||||
raw_headers_from_request: dict,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
) -> dict:
|
||||
"""
|
||||
Resolve and fetch tools for a single specified MCP server.
|
||||
|
||||
Returns the full REST response dict (tools / error / message).
|
||||
Raises HTTPException on access / IP-filter errors.
|
||||
"""
|
||||
# Resolve a server name to its UUID if needed
|
||||
_name_resolved = None
|
||||
if server_id not in allowed_server_ids:
|
||||
_name_resolved = global_mcp_server_manager.get_mcp_server_by_name(
|
||||
server_id
|
||||
)
|
||||
if _name_resolved is not None and _name_resolved.server_id in set(
|
||||
allowed_server_ids
|
||||
):
|
||||
server_id = _name_resolved.server_id
|
||||
|
||||
if server_id not in allowed_server_ids:
|
||||
_server = (
|
||||
global_mcp_server_manager.get_mcp_server_by_id(server_id)
|
||||
or _name_resolved
|
||||
)
|
||||
if (
|
||||
_server is not None
|
||||
and rest_client_ip is not None
|
||||
and not global_mcp_server_manager._is_server_accessible_from_ip(
|
||||
_server, rest_client_ip
|
||||
)
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "ip_filtering",
|
||||
"message": (
|
||||
f"MCP server '{server_id}' is not accessible from your IP address "
|
||||
f"({rest_client_ip}). This server is restricted to internal "
|
||||
"networks only. To make it externally accessible, set "
|
||||
"'available_on_public_internet: true' in the server configuration."
|
||||
),
|
||||
},
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "access_denied",
|
||||
"message": f"The key is not allowed to access server {server_id}",
|
||||
},
|
||||
)
|
||||
|
||||
server = global_mcp_server_manager.get_mcp_server_by_id(server_id)
|
||||
if server is None:
|
||||
return {
|
||||
"tools": [],
|
||||
"error": "server_not_found",
|
||||
"message": f"Server with id {server_id} not found",
|
||||
}
|
||||
|
||||
server_auth_header = _get_server_auth_header(
|
||||
server, mcp_server_auth_headers, mcp_auth_header
|
||||
)
|
||||
user_oauth_extra_headers = await _get_user_oauth_extra_headers(
|
||||
server, user_api_key_dict
|
||||
)
|
||||
|
||||
try:
|
||||
tools = await _get_tools_for_single_server(
|
||||
server,
|
||||
server_auth_header,
|
||||
raw_headers_from_request,
|
||||
user_api_key_dict,
|
||||
extra_headers=user_oauth_extra_headers,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
f"Error getting tools from {server.name}: {e}"
|
||||
)
|
||||
return {
|
||||
"tools": [],
|
||||
"error": "server_error",
|
||||
"message": f"Failed to get tools from server {server.name}: {str(e)}",
|
||||
}
|
||||
|
||||
return {
|
||||
"tools": tools,
|
||||
"error": None,
|
||||
"message": "Successfully retrieved tools",
|
||||
}
|
||||
|
||||
########################################################
|
||||
@router.get("/tools/list", dependencies=[Depends(user_api_key_auth)])
|
||||
async def list_tool_rest_api(
|
||||
@@ -427,86 +524,15 @@ if MCP_AVAILABLE:
|
||||
|
||||
# If server_id is specified, only query that specific server
|
||||
if server_id:
|
||||
# Resolve a server name to its UUID if needed (MCPConnectPicker passes
|
||||
# server_name strings, but allowed_server_ids_set contains UUIDs).
|
||||
# _name_resolved is kept so the second check can reuse it for accurate
|
||||
# IP-filter error reporting if the resolved UUID is not in allowed_server_ids.
|
||||
_name_resolved = None
|
||||
if server_id not in allowed_server_ids:
|
||||
_name_resolved = global_mcp_server_manager.get_mcp_server_by_name(
|
||||
server_id
|
||||
)
|
||||
if _name_resolved is not None and _name_resolved.server_id in set(
|
||||
allowed_server_ids
|
||||
):
|
||||
server_id = _name_resolved.server_id
|
||||
|
||||
if server_id not in allowed_server_ids:
|
||||
# Try UUID lookup first; fall back to the name-resolved server so that
|
||||
# IP-filter reporting works correctly even when server_id is a name string.
|
||||
_server = (
|
||||
global_mcp_server_manager.get_mcp_server_by_id(server_id)
|
||||
or _name_resolved
|
||||
)
|
||||
if (
|
||||
_server is not None
|
||||
and _rest_client_ip is not None
|
||||
and not global_mcp_server_manager._is_server_accessible_from_ip(
|
||||
_server, _rest_client_ip
|
||||
)
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "ip_filtering",
|
||||
"message": (
|
||||
f"MCP server '{server_id}' is not accessible from your IP address "
|
||||
f"({_rest_client_ip}). This server is restricted to internal "
|
||||
"networks only. To make it externally accessible, set "
|
||||
"'available_on_public_internet: true' in the server configuration."
|
||||
),
|
||||
},
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "access_denied",
|
||||
"message": f"The key is not allowed to access server {server_id}",
|
||||
},
|
||||
)
|
||||
server = global_mcp_server_manager.get_mcp_server_by_id(server_id)
|
||||
if server is None:
|
||||
return {
|
||||
"tools": [],
|
||||
"error": "server_not_found",
|
||||
"message": f"Server with id {server_id} not found",
|
||||
}
|
||||
|
||||
server_auth_header = _get_server_auth_header(
|
||||
server, mcp_server_auth_headers, mcp_auth_header
|
||||
return await _list_tools_for_single_server(
|
||||
server_id=server_id,
|
||||
allowed_server_ids=allowed_server_ids,
|
||||
rest_client_ip=_rest_client_ip,
|
||||
mcp_server_auth_headers=mcp_server_auth_headers,
|
||||
mcp_auth_header=mcp_auth_header,
|
||||
raw_headers_from_request=raw_headers_from_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
# Single-server request: targeted lookup is more efficient than a bulk fetch.
|
||||
user_oauth_extra_headers = await _get_user_oauth_extra_headers(
|
||||
server, user_api_key_dict
|
||||
)
|
||||
|
||||
try:
|
||||
list_tools_result = await _get_tools_for_single_server(
|
||||
server,
|
||||
server_auth_header,
|
||||
raw_headers_from_request,
|
||||
user_api_key_dict,
|
||||
extra_headers=user_oauth_extra_headers,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
f"Error getting tools from {server.name}: {e}"
|
||||
)
|
||||
return {
|
||||
"tools": [],
|
||||
"error": "server_error",
|
||||
"message": f"Failed to get tools from server {server.name}: {str(e)}",
|
||||
}
|
||||
else:
|
||||
if not allowed_server_ids:
|
||||
if _ip_blocked_count > 0:
|
||||
|
||||
@@ -8,7 +8,7 @@ import contextlib
|
||||
import time
|
||||
import traceback
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from datetime import datetime
|
||||
from typing import (
|
||||
Any,
|
||||
AsyncIterator,
|
||||
|
||||
@@ -11,7 +11,6 @@ from fastapi import status as http_status
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
|
||||
@@ -260,6 +260,9 @@ if TYPE_CHECKING:
|
||||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.llms.base_llm.files.transformation import BaseFilesConfig
|
||||
from litellm.llms.base_llm.realtime.http_transformation import (
|
||||
BaseRealtimeHTTPConfig,
|
||||
)
|
||||
from litellm.proxy._types import AllowedModelRegion
|
||||
|
||||
# Type stubs for lazy-loaded functions to help mypy understand their types
|
||||
|
||||
@@ -1832,8 +1832,12 @@ def test_get_max_tokens_for_model_claude_35():
|
||||
config = AnthropicConfig()
|
||||
|
||||
# Claude 3.5 Sonnet should return 8192
|
||||
max_tokens = config.get_max_tokens_for_model("claude-3-5-sonnet-20241022")
|
||||
assert max_tokens == 8192
|
||||
with patch(
|
||||
"litellm.llms.anthropic.chat.transformation.get_max_tokens",
|
||||
return_value=8192,
|
||||
):
|
||||
max_tokens = config.get_max_tokens_for_model("claude-3-5-sonnet-20241022")
|
||||
assert max_tokens == 8192
|
||||
|
||||
|
||||
def test_get_max_tokens_for_model_claude_37():
|
||||
@@ -1879,17 +1883,34 @@ def test_get_config_with_model_uses_dynamic_max_tokens():
|
||||
|
||||
Fixes: https://github.com/BerriAI/litellm/issues/8835
|
||||
"""
|
||||
# Claude 3 model should get 4096
|
||||
config_claude3 = AnthropicConfig.get_config(model="claude-3-sonnet-20240229")
|
||||
assert config_claude3["max_tokens"] == 4096
|
||||
|
||||
# Claude 3.5 model should get 8192
|
||||
config_claude35 = AnthropicConfig.get_config(model="claude-3-5-sonnet-20241022")
|
||||
assert config_claude35["max_tokens"] == 8192
|
||||
def _mock_get_max_tokens(model):
|
||||
"""Return expected max_output_tokens for each model."""
|
||||
model_map = {
|
||||
"claude-3-sonnet-20240229": 4096,
|
||||
"claude-3-5-sonnet-20241022": 8192,
|
||||
"claude-3-7-sonnet-20250219": 64000,
|
||||
}
|
||||
result = model_map.get(model)
|
||||
if result is None:
|
||||
raise Exception(f"Model {model} not found")
|
||||
return result
|
||||
|
||||
# Claude 3.7 model should get 64000 (64K default, 128K requires beta header)
|
||||
config_claude37 = AnthropicConfig.get_config(model="claude-3-7-sonnet-20250219")
|
||||
assert config_claude37["max_tokens"] == 64000
|
||||
with patch(
|
||||
"litellm.llms.anthropic.chat.transformation.get_max_tokens",
|
||||
side_effect=_mock_get_max_tokens,
|
||||
):
|
||||
# Claude 3 model should get 4096
|
||||
config_claude3 = AnthropicConfig.get_config(model="claude-3-sonnet-20240229")
|
||||
assert config_claude3["max_tokens"] == 4096
|
||||
|
||||
# Claude 3.5 model should get 8192
|
||||
config_claude35 = AnthropicConfig.get_config(model="claude-3-5-sonnet-20241022")
|
||||
assert config_claude35["max_tokens"] == 8192
|
||||
|
||||
# Claude 3.7 model should get 64000 (64K default, 128K requires beta header)
|
||||
config_claude37 = AnthropicConfig.get_config(model="claude-3-7-sonnet-20250219")
|
||||
assert config_claude37["max_tokens"] == 64000
|
||||
|
||||
|
||||
def test_get_config_without_model_uses_fallback():
|
||||
|
||||
@@ -3,7 +3,7 @@ Test Bedrock files integration with main files API
|
||||
"""
|
||||
|
||||
import base64
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -21,25 +21,26 @@ class TestBedrockFilesIntegration:
|
||||
file_id = "s3://test-bucket/test-file.jsonl"
|
||||
expected_content = b'{"recordId": "request-1", "modelInput": {}, "modelOutput": {}}'
|
||||
|
||||
# Mock the bedrock_files_instance.file_content method
|
||||
with patch(
|
||||
"litellm.files.main.bedrock_files_instance.file_content",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_file_content:
|
||||
# Create a mock HttpxBinaryResponseContent response
|
||||
import httpx
|
||||
# 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="s3://test-bucket/test-file.jsonl"
|
||||
),
|
||||
)
|
||||
mock_file_content.return_value = HttpxBinaryResponseContent(
|
||||
response=mock_response
|
||||
)
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=expected_content,
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request(
|
||||
method="GET", url="s3://test-bucket/test-file.jsonl"
|
||||
),
|
||||
)
|
||||
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:
|
||||
mock_retrieve.return_value = mock_result
|
||||
|
||||
# Call litellm.afile_content
|
||||
result = await litellm.afile_content(
|
||||
@@ -54,8 +55,8 @@ class TestBedrockFilesIntegration:
|
||||
assert result.response.status_code == 200
|
||||
|
||||
# Verify the mock was called with correct parameters
|
||||
mock_file_content.assert_called_once()
|
||||
call_kwargs = mock_file_content.call_args.kwargs
|
||||
mock_retrieve.assert_called_once()
|
||||
call_kwargs = mock_retrieve.call_args.kwargs
|
||||
assert call_kwargs["_is_async"] is True
|
||||
assert call_kwargs["file_content_request"]["file_id"] == file_id
|
||||
|
||||
@@ -66,29 +67,29 @@ class TestBedrockFilesIntegration:
|
||||
s3_uri = "s3://test-bucket/batch-outputs/output.jsonl"
|
||||
unified_id = "test-unified-id-123"
|
||||
model_id = "test-model-id-456"
|
||||
|
||||
|
||||
unified_file_id_str = f"litellm_proxy:application/json;unified_id,{unified_id};target_model_names,;llm_output_file_id,{s3_uri};llm_output_file_model_id,{model_id}"
|
||||
encoded_file_id = base64.urlsafe_b64encode(unified_file_id_str.encode()).decode().rstrip("=")
|
||||
|
||||
|
||||
expected_content = b'{"recordId": "request-1", "modelInput": {}, "modelOutput": {}}'
|
||||
|
||||
# Mock the bedrock_files_instance.file_content method
|
||||
with patch(
|
||||
"litellm.files.main.bedrock_files_instance.file_content",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_file_content:
|
||||
# Create a mock HttpxBinaryResponseContent response
|
||||
import httpx
|
||||
# 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=s3_uri),
|
||||
)
|
||||
mock_file_content.return_value = HttpxBinaryResponseContent(
|
||||
response=mock_response
|
||||
)
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=expected_content,
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request(method="GET", url=s3_uri),
|
||||
)
|
||||
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 litellm.afile_content with unified file ID
|
||||
result = await litellm.afile_content(
|
||||
@@ -102,9 +103,9 @@ class TestBedrockFilesIntegration:
|
||||
assert result.response.content == expected_content
|
||||
assert result.response.status_code == 200
|
||||
|
||||
# Verify the mock was called - the handler should extract S3 URI from unified file ID
|
||||
mock_file_content.assert_called_once()
|
||||
call_kwargs = mock_file_content.call_args.kwargs
|
||||
# Verify the mock was called
|
||||
mock_retrieve.assert_called_once()
|
||||
call_kwargs = mock_retrieve.call_args.kwargs
|
||||
assert call_kwargs["_is_async"] is True
|
||||
# The handler extracts S3 URI from the unified file ID
|
||||
# The handler passes the encoded file_id as-is
|
||||
assert call_kwargs["file_content_request"]["file_id"] == encoded_file_id
|
||||
|
||||
@@ -600,7 +600,8 @@ class TestGeminiVideoCostTracking:
|
||||
cost_veo2 = video_generation_cost(
|
||||
model="gemini/veo-2.0-generate-001",
|
||||
duration_seconds=5.0,
|
||||
custom_llm_provider="gemini"
|
||||
custom_llm_provider="gemini",
|
||||
model_info={"output_cost_per_second": 0.35},
|
||||
)
|
||||
expected_veo2 = 0.35 * 5.0 # $1.75
|
||||
assert abs(cost_veo2 - expected_veo2) < 0.001, f"Expected ${expected_veo2}, got ${cost_veo2}"
|
||||
@@ -609,7 +610,8 @@ class TestGeminiVideoCostTracking:
|
||||
cost_veo3 = video_generation_cost(
|
||||
model="gemini/veo-3.0-generate-preview",
|
||||
duration_seconds=8.0,
|
||||
custom_llm_provider="gemini"
|
||||
custom_llm_provider="gemini",
|
||||
model_info={"output_cost_per_second": 0.75},
|
||||
)
|
||||
expected_veo3 = 0.75 * 8.0 # $6.00
|
||||
assert abs(cost_veo3 - expected_veo3) < 0.001, f"Expected ${expected_veo3}, got ${cost_veo3}"
|
||||
@@ -618,7 +620,8 @@ class TestGeminiVideoCostTracking:
|
||||
cost_veo31 = video_generation_cost(
|
||||
model="gemini/veo-3.1-generate-preview",
|
||||
duration_seconds=10.0,
|
||||
custom_llm_provider="gemini"
|
||||
custom_llm_provider="gemini",
|
||||
model_info={"output_cost_per_second": 0.40},
|
||||
)
|
||||
expected_veo31 = 0.40 * 10.0 # $4.00
|
||||
assert abs(cost_veo31 - expected_veo31) < 0.001, f"Expected ${expected_veo31}, got ${cost_veo31}"
|
||||
@@ -627,7 +630,8 @@ class TestGeminiVideoCostTracking:
|
||||
cost_veo31_fast = video_generation_cost(
|
||||
model="gemini/veo-3.1-fast-generate-preview",
|
||||
duration_seconds=6.0,
|
||||
custom_llm_provider="gemini"
|
||||
custom_llm_provider="gemini",
|
||||
model_info={"output_cost_per_second": 0.15},
|
||||
)
|
||||
expected_veo31_fast = 0.15 * 6.0 # $0.90
|
||||
assert abs(cost_veo31_fast - expected_veo31_fast) < 0.001, f"Expected ${expected_veo31_fast}, got ${cost_veo31_fast}"
|
||||
@@ -667,7 +671,8 @@ class TestGeminiVideoCostTracking:
|
||||
cost = video_generation_cost(
|
||||
model="gemini/veo-3.0-generate-preview",
|
||||
duration_seconds=duration,
|
||||
custom_llm_provider="gemini"
|
||||
custom_llm_provider="gemini",
|
||||
model_info={"output_cost_per_second": 0.75},
|
||||
)
|
||||
|
||||
# Verify cost calculation (VEO 3.0 is $0.75/second)
|
||||
|
||||
@@ -3,7 +3,7 @@ Test Vertex AI files integration with main files API
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import litellm
|
||||
from litellm.types.llms.openai import HttpxBinaryResponseContent
|
||||
@@ -18,27 +18,28 @@ class TestVertexAIFilesIntegration:
|
||||
file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt"
|
||||
expected_content = b"test file content"
|
||||
|
||||
# Mock the vertex_ai_files_instance.file_content method
|
||||
# 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.vertex_ai_files_instance.file_content",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_file_content:
|
||||
# Create a mock HttpxBinaryResponseContent response
|
||||
import httpx
|
||||
"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
|
||||
|
||||
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_file_content.return_value = HttpxBinaryResponseContent(
|
||||
response=mock_response
|
||||
)
|
||||
|
||||
# Call litellm.afile_content
|
||||
result = await litellm.afile_content(
|
||||
file_id=file_id,
|
||||
custom_llm_provider="vertex_ai",
|
||||
@@ -52,39 +53,32 @@ class TestVertexAIFilesIntegration:
|
||||
assert result.response.content == expected_content
|
||||
assert result.response.status_code == 200
|
||||
|
||||
# Verify the mock was called with correct parameters
|
||||
mock_file_content.assert_called_once()
|
||||
call_kwargs = mock_file_content.call_args.kwargs
|
||||
assert call_kwargs["_is_async"] is True
|
||||
assert call_kwargs["file_content_request"]["file_id"] == file_id
|
||||
assert call_kwargs["vertex_project"] == "test-project"
|
||||
assert call_kwargs["vertex_location"] == "us-central1"
|
||||
# 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"
|
||||
|
||||
# Mock the vertex_ai_files_instance.file_content method
|
||||
# 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.vertex_ai_files_instance.file_content"
|
||||
) as mock_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_file_content.return_value = HttpxBinaryResponseContent(
|
||||
response=mock_response
|
||||
)
|
||||
|
||||
# Call litellm.file_content
|
||||
"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",
|
||||
@@ -98,23 +92,32 @@ class TestVertexAIFilesIntegration:
|
||||
assert result.response.content == expected_content
|
||||
assert result.response.status_code == 200
|
||||
|
||||
# Verify the mock was called with correct parameters
|
||||
mock_file_content.assert_called_once()
|
||||
call_kwargs = mock_file_content.call_args.kwargs
|
||||
assert call_kwargs["_is_async"] is False
|
||||
assert call_kwargs["file_content_request"]["file_id"] == file_id
|
||||
assert call_kwargs["vertex_project"] == "test-project"
|
||||
assert call_kwargs["vertex_location"] == "us-central1"
|
||||
# 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"
|
||||
|
||||
# Mock the vertex_ai_files_instance.file_content method
|
||||
# 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.vertex_ai_files_instance.file_content"
|
||||
) as mock_file_content:
|
||||
"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 = (
|
||||
@@ -124,25 +127,10 @@ class TestVertexAIFilesIntegration:
|
||||
None,
|
||||
)
|
||||
|
||||
# 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_file_content.return_value = HttpxBinaryResponseContent(
|
||||
response=mock_response
|
||||
)
|
||||
|
||||
# Call litellm.file_content with model to trigger provider detection
|
||||
result = litellm.file_content(
|
||||
file_id=file_id,
|
||||
model="vertex_ai/gemini-pro", # This should trigger provider detection
|
||||
model="vertex_ai/gemini-pro",
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
)
|
||||
@@ -156,13 +144,21 @@ class TestVertexAIFilesIntegration:
|
||||
|
||||
def test_litellm_file_content_vertex_ai_error_cases(self):
|
||||
"""Test error handling in vertex_ai file_content"""
|
||||
# Test missing file_id
|
||||
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",
|
||||
)
|
||||
# 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"""
|
||||
@@ -185,25 +181,25 @@ class TestVertexAIFilesIntegration:
|
||||
file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt"
|
||||
expected_content = b"test file content"
|
||||
|
||||
# Mock the vertex_ai_files_instance.file_content method
|
||||
with patch(
|
||||
"litellm.files.main.vertex_ai_files_instance.file_content",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_file_content:
|
||||
# Create a mock HttpxBinaryResponseContent response
|
||||
import httpx
|
||||
# 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_file_content.return_value = HttpxBinaryResponseContent(
|
||||
response=mock_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"
|
||||
),
|
||||
)
|
||||
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(
|
||||
@@ -219,7 +215,8 @@ class TestVertexAIFilesIntegration:
|
||||
assert isinstance(result, HttpxBinaryResponseContent)
|
||||
assert result.response.content == expected_content
|
||||
|
||||
# Verify the timeout and max_retries were passed through
|
||||
call_kwargs = mock_file_content.call_args.kwargs
|
||||
# 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
|
||||
assert call_kwargs["max_retries"] == 5
|
||||
|
||||
@@ -55,33 +55,30 @@ def test_completion_pydantic_obj_2():
|
||||
],
|
||||
"generationConfig": {
|
||||
"response_mime_type": "application/json",
|
||||
"response_schema": {
|
||||
"response_json_schema": {
|
||||
"$defs": {
|
||||
"CalendarEvent": {
|
||||
"properties": {
|
||||
"name": {"title": "Name", "type": "string"},
|
||||
"date": {"title": "Date", "type": "string"},
|
||||
"participants": {
|
||||
"items": {"type": "string"},
|
||||
"title": "Participants",
|
||||
"type": "array",
|
||||
},
|
||||
},
|
||||
"required": ["name", "date", "participants"],
|
||||
"title": "CalendarEvent",
|
||||
"type": "object",
|
||||
}
|
||||
},
|
||||
"properties": {
|
||||
"events": {
|
||||
"items": {
|
||||
"properties": {
|
||||
"name": {"title": "Name", "type": "string"},
|
||||
"date": {"title": "Date", "type": "string"},
|
||||
"participants": {
|
||||
"items": {"type": "string"},
|
||||
"title": "Participants",
|
||||
"type": "array",
|
||||
},
|
||||
},
|
||||
"propertyOrdering": [
|
||||
"name",
|
||||
"date",
|
||||
"participants",
|
||||
],
|
||||
"required": ["name", "date", "participants"],
|
||||
"title": "CalendarEvent",
|
||||
"type": "object",
|
||||
},
|
||||
"items": {"$ref": "#/$defs/CalendarEvent"},
|
||||
"title": "Events",
|
||||
"type": "array",
|
||||
}
|
||||
},
|
||||
"propertyOrdering": ["events"],
|
||||
"required": ["events"],
|
||||
"title": "EventsList",
|
||||
"type": "object",
|
||||
@@ -93,7 +90,7 @@ def test_completion_pydantic_obj_2():
|
||||
mock_post.return_value = expected_request_body
|
||||
try:
|
||||
response = litellm.completion(
|
||||
model="gemini/gemini-1.5-pro",
|
||||
model="gemini/gemini-2.5-flash",
|
||||
messages=messages,
|
||||
response_format=EventsList,
|
||||
client=client,
|
||||
|
||||
Reference in New Issue
Block a user