mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
feat(pass_through_endpoints.py): only show available pass through endpoints based on what teams has access to
makes it easier to control access
This commit is contained in:
Krrish Dholakia
parent
5aa5a3b425
commit
f2d4f03339
5 changed files with 521 additions and 141 deletions
|
|
@ -1818,23 +1818,91 @@ async def _get_pass_through_endpoints_from_db(
|
|||
return returned_endpoints
|
||||
|
||||
|
||||
async def _filter_endpoints_by_team_allowed_routes(
|
||||
team_id: str,
|
||||
pass_through_endpoints: List[PassThroughGenericEndpoint],
|
||||
prisma_client,
|
||||
) -> List[PassThroughGenericEndpoint]:
|
||||
"""
|
||||
Filter pass-through endpoints based on team's allowed_passthrough_routes metadata.
|
||||
|
||||
Args:
|
||||
team_id: The team ID to check permissions for
|
||||
pass_through_endpoints: List of endpoints to filter
|
||||
prisma_client: Database client
|
||||
|
||||
Returns:
|
||||
Filtered list of endpoints based on team permissions
|
||||
|
||||
Raises:
|
||||
HTTPException: If team is not found
|
||||
"""
|
||||
# retrieve team from db
|
||||
team = await prisma_client.db.litellm_teamtable.find_unique(
|
||||
where={"team_id": team_id},
|
||||
)
|
||||
if team is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": "Team not found"},
|
||||
)
|
||||
|
||||
# retrieve team metadata
|
||||
team_metadata = team.metadata
|
||||
if (
|
||||
team_metadata is not None
|
||||
and team_metadata.get("allowed_passthrough_routes") is not None
|
||||
):
|
||||
## FILTER pass_through_endpoints by allowed_passthrough_routes
|
||||
pass_through_endpoints = [
|
||||
endpoint
|
||||
for endpoint in pass_through_endpoints
|
||||
if endpoint.path in team_metadata.get("allowed_passthrough_routes")
|
||||
]
|
||||
|
||||
return pass_through_endpoints
|
||||
|
||||
|
||||
@router.get(
|
||||
"/config/pass_through_endpoint",
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=PassThroughEndpointResponse,
|
||||
)
|
||||
@router.get(
|
||||
"/config/pass_through_endpoint/team/{team_id}",
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=PassThroughEndpointResponse,
|
||||
)
|
||||
async def get_pass_through_endpoints(
|
||||
endpoint_id: Optional[str] = None,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
team_id: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
GET configured pass through endpoint.
|
||||
|
||||
If no endpoint_id given, return all configured endpoints.
|
||||
""" ## Get existing pass-through endpoint field value
|
||||
from litellm.proxy._types import CommonProxyErrors
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={"error": CommonProxyErrors.db_not_connected_error.value},
|
||||
)
|
||||
|
||||
pass_through_endpoints = await _get_pass_through_endpoints_from_db(
|
||||
endpoint_id=endpoint_id, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
|
||||
if team_id is not None:
|
||||
pass_through_endpoints = await _filter_endpoints_by_team_allowed_routes(
|
||||
team_id=team_id,
|
||||
pass_through_endpoints=pass_through_endpoints,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
return PassThroughEndpointResponse(endpoints=pass_through_endpoints)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -670,7 +670,7 @@ async def test_create_pass_through_route_with_cost_per_request():
|
|||
# Create a proper UserAPIKeyAuth mock
|
||||
mock_user_api_key_dict = MagicMock()
|
||||
mock_user_api_key_dict.api_key = "test-key"
|
||||
|
||||
|
||||
await endpoint_func(
|
||||
request=mock_request,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
|
|
@ -747,12 +747,12 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs():
|
|||
"""
|
||||
Test that pass_through_request (parent method) correctly includes proxy_server_request
|
||||
in kwargs passed to the success handler.
|
||||
|
||||
Critical Test: Ensures that when pass_through_request is called, the kwargs passed to
|
||||
|
||||
Critical Test: Ensures that when pass_through_request is called, the kwargs passed to
|
||||
downstream methods contain the proxy server request details (url, method, body).
|
||||
"""
|
||||
print("running test_pass_through_request_contains_proxy_server_request_in_kwargs")
|
||||
|
||||
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.HttpPassThroughEndpointHelpers.non_streaming_http_request_handler"
|
||||
|
|
@ -766,38 +766,44 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs():
|
|||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_response_body"
|
||||
) as mock_get_response_body:
|
||||
# Setup mock for pre_call_hook and post_call_failure_hook
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={"test": "data"})
|
||||
# Setup mock for pre_call_hook and post_call_failure_hook
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(
|
||||
return_value={"test": "data"}
|
||||
)
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
|
||||
|
||||
# Setup mock for http response
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {}
|
||||
mock_response.aread = AsyncMock(return_value=b'{"success": true}')
|
||||
mock_response.aread = AsyncMock(
|
||||
return_value=b'{"success": true}'
|
||||
)
|
||||
mock_response.text = '{"success": true}'
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
|
||||
|
||||
# Mock the HTTP request handler directly
|
||||
mock_http_handler.return_value = mock_response
|
||||
|
||||
|
||||
# Mock response body parser
|
||||
mock_get_response_body.return_value = {"success": True}
|
||||
|
||||
|
||||
# Mock headers for custom headers
|
||||
mock_processing.get_custom_headers.return_value = {}
|
||||
|
||||
|
||||
# Mock success handler to capture kwargs
|
||||
mock_success_handler.return_value = None
|
||||
|
||||
|
||||
# Create mock request
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.method = "POST"
|
||||
mock_request.url = "http://test-proxy.com/api/endpoint"
|
||||
mock_request.body = AsyncMock(return_value=b'{"message": "test request"}')
|
||||
mock_request.body = AsyncMock(
|
||||
return_value=b'{"message": "test request"}'
|
||||
)
|
||||
mock_request.headers = Headers({})
|
||||
mock_request.query_params = QueryParams({})
|
||||
|
||||
|
||||
# Create mock user API key dict
|
||||
mock_user_api_key_dict = MagicMock()
|
||||
mock_user_api_key_dict.api_key = "test-api-key"
|
||||
|
|
@ -809,7 +815,7 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs():
|
|||
mock_user_api_key_dict.team_alias = "test-team-alias"
|
||||
mock_user_api_key_dict.end_user_id = "test-end-user-id"
|
||||
mock_user_api_key_dict.request_route = "/api/endpoint"
|
||||
|
||||
|
||||
# Call pass_through_request (the parent method)
|
||||
result = await pass_through_request(
|
||||
request=mock_request,
|
||||
|
|
@ -817,38 +823,38 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs():
|
|||
custom_headers={"X-Custom": "header"},
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
|
||||
# Verify the success handler was called
|
||||
mock_success_handler.assert_called_once()
|
||||
|
||||
|
||||
# Extract the kwargs passed to the success handler
|
||||
call_kwargs = mock_success_handler.call_args[1]
|
||||
|
||||
|
||||
# Verify that litellm_params exists in kwargs
|
||||
assert "litellm_params" in call_kwargs
|
||||
litellm_params = call_kwargs["litellm_params"]
|
||||
|
||||
|
||||
# Verify that proxy_server_request exists in litellm_params
|
||||
assert "proxy_server_request" in litellm_params
|
||||
proxy_server_request = litellm_params["proxy_server_request"]
|
||||
|
||||
|
||||
# Verify the proxy_server_request contains expected fields
|
||||
assert "url" in proxy_server_request
|
||||
assert "method" in proxy_server_request
|
||||
assert "body" in proxy_server_request
|
||||
|
||||
|
||||
# Verify the values match the original request
|
||||
assert proxy_server_request["url"] == str(mock_request.url)
|
||||
assert proxy_server_request["method"] == mock_request.method
|
||||
# The body should be the value returned by pre_call_hook, not the original request body
|
||||
assert proxy_server_request["body"] == {"test": "data"}
|
||||
|
||||
|
||||
# Verify other required kwargs are present
|
||||
assert "call_type" in call_kwargs
|
||||
assert call_kwargs["call_type"] == "pass_through_endpoint"
|
||||
assert "litellm_call_id" in call_kwargs
|
||||
assert "passthrough_logging_payload" in call_kwargs
|
||||
|
||||
|
||||
# Verify metadata contains user information
|
||||
assert "metadata" in litellm_params
|
||||
metadata = litellm_params["metadata"]
|
||||
|
|
@ -862,7 +868,7 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs():
|
|||
async def test_create_pass_through_endpoint():
|
||||
"""
|
||||
Test creating a new pass-through endpoint
|
||||
|
||||
|
||||
This test verifies that the create_pass_through_endpoints function:
|
||||
1. Accepts a PassThroughGenericEndpoint object
|
||||
2. Auto-generates an ID if not provided
|
||||
|
|
@ -881,36 +887,38 @@ async def test_create_pass_through_endpoint():
|
|||
)
|
||||
|
||||
# Mock the database functions
|
||||
with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
|
||||
with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config:
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.get_config_general_settings"
|
||||
) as mock_get_config:
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.update_config_general_settings"
|
||||
) as mock_update_config:
|
||||
# Mock existing config (empty list)
|
||||
mock_get_config.return_value = ConfigFieldInfo(
|
||||
field_name="pass_through_endpoints",
|
||||
field_value=[]
|
||||
field_name="pass_through_endpoints", field_value=[]
|
||||
)
|
||||
|
||||
|
||||
# Create test endpoint data
|
||||
test_endpoint = PassThroughGenericEndpoint(
|
||||
path="/test/endpoint",
|
||||
target="http://example.com/api",
|
||||
headers={"Authorization": "Bearer test-token"},
|
||||
include_subpath=True,
|
||||
cost_per_request=0.50
|
||||
cost_per_request=0.50,
|
||||
)
|
||||
|
||||
|
||||
# Mock user API key dict
|
||||
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
|
||||
|
||||
|
||||
# Call the create function
|
||||
result = await create_pass_through_endpoints(
|
||||
data=test_endpoint,
|
||||
user_api_key_dict=mock_user_api_key_dict
|
||||
data=test_endpoint, user_api_key_dict=mock_user_api_key_dict
|
||||
)
|
||||
|
||||
|
||||
# Verify the result
|
||||
assert isinstance(result, PassThroughEndpointResponse)
|
||||
assert len(result.endpoints) == 1
|
||||
|
||||
|
||||
created_endpoint = result.endpoints[0]
|
||||
assert created_endpoint.path == "/test/endpoint"
|
||||
assert created_endpoint.target == "http://example.com/api"
|
||||
|
|
@ -918,13 +926,13 @@ async def test_create_pass_through_endpoint():
|
|||
assert created_endpoint.include_subpath is True
|
||||
assert created_endpoint.cost_per_request == 0.50
|
||||
assert created_endpoint.id is not None # Should be auto-generated
|
||||
|
||||
|
||||
# Verify database calls
|
||||
mock_get_config.assert_called_once_with(
|
||||
field_name="pass_through_endpoints",
|
||||
user_api_key_dict=mock_user_api_key_dict
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
|
||||
mock_update_config.assert_called_once()
|
||||
update_call_args = mock_update_config.call_args[1]
|
||||
assert update_call_args["data"].field_name == "pass_through_endpoints"
|
||||
|
|
@ -937,7 +945,7 @@ async def test_create_pass_through_endpoint():
|
|||
async def test_update_pass_through_endpoint():
|
||||
"""
|
||||
Test updating an existing pass-through endpoint
|
||||
|
||||
|
||||
This test verifies that the update_pass_through_endpoints function:
|
||||
1. Finds the existing endpoint by ID
|
||||
2. Updates only the provided fields (partial updates)
|
||||
|
|
@ -957,8 +965,12 @@ async def test_update_pass_through_endpoint():
|
|||
)
|
||||
|
||||
# Mock the database functions
|
||||
with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
|
||||
with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config:
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.get_config_general_settings"
|
||||
) as mock_get_config:
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.update_config_general_settings"
|
||||
) as mock_update_config:
|
||||
# Create existing endpoint data
|
||||
existing_endpoint_id = "test-endpoint-123"
|
||||
existing_endpoints = [
|
||||
|
|
@ -968,53 +980,58 @@ async def test_update_pass_through_endpoint():
|
|||
"target": "http://example.com/api",
|
||||
"headers": {"Authorization": "Bearer old-token"},
|
||||
"include_subpath": False,
|
||||
"cost_per_request": 0.25
|
||||
"cost_per_request": 0.25,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
# Mock existing config
|
||||
mock_get_config.return_value = ConfigFieldInfo(
|
||||
field_name="pass_through_endpoints",
|
||||
field_value=existing_endpoints
|
||||
field_name="pass_through_endpoints", field_value=existing_endpoints
|
||||
)
|
||||
|
||||
|
||||
# Create update data (partial update)
|
||||
update_data = PassThroughGenericEndpoint(
|
||||
path="/test/endpoint", # Keep same path
|
||||
target="http://newapi.com/v2", # Update target
|
||||
headers={"Authorization": "Bearer new-token", "X-Custom": "header"}, # Update headers
|
||||
cost_per_request=0.75 # Update cost
|
||||
headers={
|
||||
"Authorization": "Bearer new-token",
|
||||
"X-Custom": "header",
|
||||
}, # Update headers
|
||||
cost_per_request=0.75, # Update cost
|
||||
# include_subpath not provided - should preserve existing value
|
||||
)
|
||||
|
||||
|
||||
# Mock user API key dict
|
||||
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
|
||||
|
||||
|
||||
# Call the update function
|
||||
result = await update_pass_through_endpoints(
|
||||
endpoint_id=existing_endpoint_id,
|
||||
data=update_data,
|
||||
user_api_key_dict=mock_user_api_key_dict
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
|
||||
# Verify the result
|
||||
assert isinstance(result, PassThroughEndpointResponse)
|
||||
assert len(result.endpoints) == 1
|
||||
|
||||
|
||||
updated_endpoint = result.endpoints[0]
|
||||
assert updated_endpoint.id == existing_endpoint_id # ID preserved
|
||||
assert updated_endpoint.path == "/test/endpoint"
|
||||
assert updated_endpoint.target == "http://newapi.com/v2" # Updated
|
||||
assert updated_endpoint.headers == {"Authorization": "Bearer new-token", "X-Custom": "header"} # Updated
|
||||
assert updated_endpoint.headers == {
|
||||
"Authorization": "Bearer new-token",
|
||||
"X-Custom": "header",
|
||||
} # Updated
|
||||
assert updated_endpoint.include_subpath is False # Preserved existing value
|
||||
assert updated_endpoint.cost_per_request == 0.75 # Updated
|
||||
|
||||
|
||||
# Verify database calls
|
||||
mock_get_config.assert_called_once_with(
|
||||
field_name="pass_through_endpoints",
|
||||
user_api_key_dict=mock_user_api_key_dict
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
|
||||
mock_update_config.assert_called_once()
|
||||
update_call_args = mock_update_config.call_args[1]
|
||||
assert update_call_args["data"].field_name == "pass_through_endpoints"
|
||||
|
|
@ -1042,7 +1059,9 @@ async def test_update_pass_through_endpoint_not_found():
|
|||
)
|
||||
|
||||
# Mock the database functions
|
||||
with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.get_config_general_settings"
|
||||
) as mock_get_config:
|
||||
# Mock existing config with different endpoint
|
||||
existing_endpoints = [
|
||||
{
|
||||
|
|
@ -1051,32 +1070,30 @@ async def test_update_pass_through_endpoint_not_found():
|
|||
"target": "http://different.com/api",
|
||||
"headers": {},
|
||||
"include_subpath": False,
|
||||
"cost_per_request": 0.0
|
||||
"cost_per_request": 0.0,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
mock_get_config.return_value = ConfigFieldInfo(
|
||||
field_name="pass_through_endpoints",
|
||||
field_value=existing_endpoints
|
||||
field_name="pass_through_endpoints", field_value=existing_endpoints
|
||||
)
|
||||
|
||||
|
||||
# Create update data
|
||||
update_data = PassThroughGenericEndpoint(
|
||||
path="/test/endpoint",
|
||||
target="http://newapi.com/v2"
|
||||
path="/test/endpoint", target="http://newapi.com/v2"
|
||||
)
|
||||
|
||||
|
||||
# Mock user API key dict
|
||||
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
|
||||
|
||||
|
||||
# Call the update function with non-existent ID
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await update_pass_through_endpoints(
|
||||
endpoint_id="non-existent-endpoint-123",
|
||||
data=update_data,
|
||||
user_api_key_dict=mock_user_api_key_dict
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
|
||||
# Verify the exception
|
||||
assert exc_info.value.status_code == 404
|
||||
assert "not found" in str(exc_info.value.detail).lower()
|
||||
|
|
@ -1086,7 +1103,7 @@ async def test_update_pass_through_endpoint_not_found():
|
|||
async def test_delete_pass_through_endpoint():
|
||||
"""
|
||||
Test deleting an existing pass-through endpoint
|
||||
|
||||
|
||||
This test verifies that the delete_pass_through_endpoints function:
|
||||
1. Finds the existing endpoint by ID
|
||||
2. Removes it from the database
|
||||
|
|
@ -1103,8 +1120,12 @@ async def test_delete_pass_through_endpoint():
|
|||
)
|
||||
|
||||
# Mock the database functions
|
||||
with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
|
||||
with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config:
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.get_config_general_settings"
|
||||
) as mock_get_config:
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.update_config_general_settings"
|
||||
) as mock_update_config:
|
||||
# Create existing endpoint data
|
||||
endpoint_to_delete_id = "test-endpoint-123"
|
||||
other_endpoint_id = "other-endpoint-456"
|
||||
|
|
@ -1115,7 +1136,7 @@ async def test_delete_pass_through_endpoint():
|
|||
"target": "http://example.com/api",
|
||||
"headers": {"Authorization": "Bearer test-token"},
|
||||
"include_subpath": True,
|
||||
"cost_per_request": 0.50
|
||||
"cost_per_request": 0.50,
|
||||
},
|
||||
{
|
||||
"id": other_endpoint_id,
|
||||
|
|
@ -1123,29 +1144,28 @@ async def test_delete_pass_through_endpoint():
|
|||
"target": "http://other.com/api",
|
||||
"headers": {},
|
||||
"include_subpath": False,
|
||||
"cost_per_request": 0.25
|
||||
}
|
||||
"cost_per_request": 0.25,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
# Mock existing config
|
||||
mock_get_config.return_value = ConfigFieldInfo(
|
||||
field_name="pass_through_endpoints",
|
||||
field_value=existing_endpoints
|
||||
field_name="pass_through_endpoints", field_value=existing_endpoints
|
||||
)
|
||||
|
||||
|
||||
# Mock user API key dict
|
||||
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
|
||||
|
||||
|
||||
# Call the delete function
|
||||
result = await delete_pass_through_endpoints(
|
||||
endpoint_id=endpoint_to_delete_id,
|
||||
user_api_key_dict=mock_user_api_key_dict
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
|
||||
# Verify the result
|
||||
assert isinstance(result, PassThroughEndpointResponse)
|
||||
assert len(result.endpoints) == 1
|
||||
|
||||
|
||||
deleted_endpoint = result.endpoints[0]
|
||||
assert deleted_endpoint.id == endpoint_to_delete_id
|
||||
assert deleted_endpoint.path == "/test/endpoint"
|
||||
|
|
@ -1153,13 +1173,13 @@ async def test_delete_pass_through_endpoint():
|
|||
assert deleted_endpoint.headers == {"Authorization": "Bearer test-token"}
|
||||
assert deleted_endpoint.include_subpath is True
|
||||
assert deleted_endpoint.cost_per_request == 0.50
|
||||
|
||||
|
||||
# Verify database calls
|
||||
mock_get_config.assert_called_once_with(
|
||||
field_name="pass_through_endpoints",
|
||||
user_api_key_dict=mock_user_api_key_dict
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
|
||||
mock_update_config.assert_called_once()
|
||||
update_call_args = mock_update_config.call_args[1]
|
||||
assert update_call_args["data"].field_name == "pass_through_endpoints"
|
||||
|
|
@ -1183,7 +1203,9 @@ async def test_delete_pass_through_endpoint_not_found():
|
|||
)
|
||||
|
||||
# Mock the database functions
|
||||
with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.get_config_general_settings"
|
||||
) as mock_get_config:
|
||||
# Mock existing config with different endpoint
|
||||
existing_endpoints = [
|
||||
{
|
||||
|
|
@ -1192,25 +1214,24 @@ async def test_delete_pass_through_endpoint_not_found():
|
|||
"target": "http://different.com/api",
|
||||
"headers": {},
|
||||
"include_subpath": False,
|
||||
"cost_per_request": 0.0
|
||||
"cost_per_request": 0.0,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
mock_get_config.return_value = ConfigFieldInfo(
|
||||
field_name="pass_through_endpoints",
|
||||
field_value=existing_endpoints
|
||||
field_name="pass_through_endpoints", field_value=existing_endpoints
|
||||
)
|
||||
|
||||
|
||||
# Mock user API key dict
|
||||
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
|
||||
|
||||
|
||||
# Call the delete function with non-existent ID
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await delete_pass_through_endpoints(
|
||||
endpoint_id="non-existent-endpoint-123",
|
||||
user_api_key_dict=mock_user_api_key_dict
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
|
||||
# Verify the exception
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "not found" in str(exc_info.value.detail).lower()
|
||||
|
|
@ -1229,34 +1250,33 @@ async def test_delete_pass_through_endpoint_empty_list():
|
|||
)
|
||||
|
||||
# Mock the database functions
|
||||
with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.get_config_general_settings"
|
||||
) as mock_get_config:
|
||||
# Mock empty config
|
||||
mock_get_config.return_value = ConfigFieldInfo(
|
||||
field_name="pass_through_endpoints",
|
||||
field_value=None
|
||||
field_name="pass_through_endpoints", field_value=None
|
||||
)
|
||||
|
||||
|
||||
# Mock user API key dict
|
||||
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
|
||||
|
||||
|
||||
# Call the delete function
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await delete_pass_through_endpoints(
|
||||
endpoint_id="any-endpoint-123",
|
||||
user_api_key_dict=mock_user_api_key_dict
|
||||
endpoint_id="any-endpoint-123", user_api_key_dict=mock_user_api_key_dict
|
||||
)
|
||||
|
||||
|
||||
# Verify the exception
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "no pass-through endpoints setup" in str(exc_info.value.detail).lower()
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_request_query_params_forwarding():
|
||||
"""
|
||||
Test that query parameters from the original request are properly forwarded to the target URL.
|
||||
|
||||
|
||||
This test verifies the fix for the bug where query parameters like api-version were being lost
|
||||
when forwarding requests to Azure OpenAI and other pass-through endpoints.
|
||||
"""
|
||||
|
|
@ -1275,42 +1295,57 @@ async def test_pass_through_request_query_params_forwarding():
|
|||
) as mock_get_response_body:
|
||||
# Setup mock for pre_call_hook
|
||||
test_body = {"name": "Azure Assistant", "model": "gpt-4o"}
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(return_value=test_body)
|
||||
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(
|
||||
return_value=test_body
|
||||
)
|
||||
|
||||
# Setup mock for http response
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {"content-type": "application/json"}
|
||||
mock_response.aread = AsyncMock(return_value=b'{"id": "asst_123", "object": "assistant"}')
|
||||
mock_response.aread = AsyncMock(
|
||||
return_value=b'{"id": "asst_123", "object": "assistant"}'
|
||||
)
|
||||
mock_response.text = '{"id": "asst_123", "object": "assistant"}'
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
|
||||
|
||||
# Mock the HTTP request handler to capture the call
|
||||
mock_http_handler.return_value = mock_response
|
||||
|
||||
|
||||
# Mock response body parser
|
||||
mock_get_response_body.return_value = {"id": "asst_123", "object": "assistant"}
|
||||
|
||||
mock_get_response_body.return_value = {
|
||||
"id": "asst_123",
|
||||
"object": "assistant",
|
||||
}
|
||||
|
||||
# Mock headers for custom headers
|
||||
mock_processing.get_custom_headers.return_value = {}
|
||||
|
||||
|
||||
# Mock success handler
|
||||
mock_success_handler.return_value = None
|
||||
|
||||
|
||||
# Create mock request with query parameters (Azure API version)
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.method = "POST"
|
||||
mock_request.url = "http://localhost:4000/azure-assistant/openai/assistants"
|
||||
mock_request.body = AsyncMock(return_value=json.dumps(test_body).encode())
|
||||
mock_request.headers = Headers({"Content-Type": "application/json"})
|
||||
|
||||
mock_request.url = (
|
||||
"http://localhost:4000/azure-assistant/openai/assistants"
|
||||
)
|
||||
mock_request.body = AsyncMock(
|
||||
return_value=json.dumps(test_body).encode()
|
||||
)
|
||||
mock_request.headers = Headers(
|
||||
{"Content-Type": "application/json"}
|
||||
)
|
||||
|
||||
# Create QueryParams with api-version parameter
|
||||
mock_request.query_params = QueryParams([("api-version", "2025-01-01-preview")])
|
||||
|
||||
mock_request.query_params = QueryParams(
|
||||
[("api-version", "2025-01-01-preview")]
|
||||
)
|
||||
|
||||
# Create mock user API key dict
|
||||
mock_user_api_key_dict = MagicMock()
|
||||
mock_user_api_key_dict.api_key = "sk-1234"
|
||||
|
||||
|
||||
# Call pass_through_request
|
||||
result = await pass_through_request(
|
||||
request=mock_request,
|
||||
|
|
@ -1318,20 +1353,25 @@ async def test_pass_through_request_query_params_forwarding():
|
|||
custom_headers={"Authorization": "Bearer azure_token"},
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
|
||||
# Verify the HTTP handler was called
|
||||
mock_http_handler.assert_called_once()
|
||||
|
||||
|
||||
# Extract the call arguments to verify query parameters were passed
|
||||
call_kwargs = mock_http_handler.call_args[1]
|
||||
|
||||
|
||||
# The key assertion: query parameters should be preserved and passed to the HTTP handler
|
||||
assert "requested_query_params" in call_kwargs
|
||||
assert call_kwargs["requested_query_params"] == {"api-version": "2025-01-01-preview"}
|
||||
|
||||
assert call_kwargs["requested_query_params"] == {
|
||||
"api-version": "2025-01-01-preview"
|
||||
}
|
||||
|
||||
# Verify the target URL is correct
|
||||
assert str(call_kwargs["url"]) == "https://krris-m2f9a9i7-eastus2.openai.azure.com/openai/assistants"
|
||||
|
||||
assert (
|
||||
str(call_kwargs["url"])
|
||||
== "https://krris-m2f9a9i7-eastus2.openai.azure.com/openai/assistants"
|
||||
)
|
||||
|
||||
# Verify the request body is preserved
|
||||
assert call_kwargs["_parsed_body"] == test_body
|
||||
|
||||
|
|
@ -1356,15 +1396,16 @@ async def test_pass_through_with_httpbin_redirect():
|
|||
mock_request.method = "GET"
|
||||
mock_request.headers = Headers({})
|
||||
mock_request.query_params = QueryParams("")
|
||||
|
||||
|
||||
# Mock the body method to return empty bytes for GET request
|
||||
async def mock_body():
|
||||
return b""
|
||||
|
||||
mock_request.body = mock_body
|
||||
|
||||
|
||||
# Mock user API key dict
|
||||
mock_user_api_key_dict = MagicMock()
|
||||
|
||||
|
||||
try:
|
||||
# Test with httpbin.org redirect endpoint
|
||||
# This will redirect to httpbin.org/get
|
||||
|
|
@ -1372,19 +1413,283 @@ async def test_pass_through_with_httpbin_redirect():
|
|||
request=mock_request,
|
||||
target="https://httpbin.org/redirect/1",
|
||||
custom_headers={},
|
||||
user_api_key_dict=mock_user_api_key_dict
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
|
||||
# Should get the final response (200) from /get endpoint, not the redirect (302)
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
# The response should be from the /get endpoint
|
||||
response_content = response.body.decode('utf-8')
|
||||
|
||||
response_content = response.body.decode("utf-8")
|
||||
|
||||
# httpbin.org/get returns JSON with info about the request
|
||||
assert '"url": "https://httpbin.org/get"' in response_content
|
||||
print("GOT A Response from HTTPBIN=", response_content)
|
||||
except Exception as e:
|
||||
# If httpbin.org is not accessible, skip the test
|
||||
import pytest
|
||||
|
||||
pytest.skip(f"Could not reach httpbin.org for integration test: {e}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_endpoints_by_team_allowed_routes_with_filter():
|
||||
"""
|
||||
Test that _filter_endpoints_by_team_allowed_routes correctly filters endpoints
|
||||
when team has allowed_passthrough_routes in metadata
|
||||
"""
|
||||
from litellm.proxy._types import PassThroughGenericEndpoint
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
_filter_endpoints_by_team_allowed_routes,
|
||||
)
|
||||
|
||||
# Create test endpoints
|
||||
endpoints = [
|
||||
PassThroughGenericEndpoint(
|
||||
id="endpoint-1", path="/api/allowed1", target="http://example.com/api1"
|
||||
),
|
||||
PassThroughGenericEndpoint(
|
||||
id="endpoint-2", path="/api/allowed2", target="http://example.com/api2"
|
||||
),
|
||||
PassThroughGenericEndpoint(
|
||||
id="endpoint-3", path="/api/notallowed", target="http://example.com/api3"
|
||||
),
|
||||
]
|
||||
|
||||
# Mock prisma client
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_team = MagicMock()
|
||||
mock_team.metadata = {
|
||||
"allowed_passthrough_routes": ["/api/allowed1", "/api/allowed2"]
|
||||
}
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
|
||||
return_value=mock_team
|
||||
)
|
||||
|
||||
# Call the function
|
||||
result = await _filter_endpoints_by_team_allowed_routes(
|
||||
team_id="test-team-123",
|
||||
pass_through_endpoints=endpoints,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
# Should only return allowed endpoints
|
||||
assert len(result) == 2
|
||||
assert result[0].path == "/api/allowed1"
|
||||
assert result[1].path == "/api/allowed2"
|
||||
|
||||
# Verify database call
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique.assert_called_once_with(
|
||||
where={"team_id": "test-team-123"}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_endpoints_by_team_allowed_routes_team_not_found():
|
||||
"""
|
||||
Test that _filter_endpoints_by_team_allowed_routes raises HTTPException
|
||||
when team is not found
|
||||
"""
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import PassThroughGenericEndpoint
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
_filter_endpoints_by_team_allowed_routes,
|
||||
)
|
||||
|
||||
# Create test endpoints
|
||||
endpoints = [
|
||||
PassThroughGenericEndpoint(
|
||||
id="endpoint-1", path="/api/test", target="http://example.com/api"
|
||||
),
|
||||
]
|
||||
|
||||
# Mock prisma client to return None (team not found)
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
# Call the function and expect HTTPException
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _filter_endpoints_by_team_allowed_routes(
|
||||
team_id="non-existent-team",
|
||||
pass_through_endpoints=endpoints,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
# Verify the exception
|
||||
assert exc_info.value.status_code == 404
|
||||
assert "Team not found" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_endpoints_by_team_allowed_routes_no_metadata():
|
||||
"""
|
||||
Test that _filter_endpoints_by_team_allowed_routes returns all endpoints
|
||||
when team has no metadata
|
||||
"""
|
||||
from litellm.proxy._types import PassThroughGenericEndpoint
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
_filter_endpoints_by_team_allowed_routes,
|
||||
)
|
||||
|
||||
# Create test endpoints
|
||||
endpoints = [
|
||||
PassThroughGenericEndpoint(
|
||||
id="endpoint-1", path="/api/test1", target="http://example.com/api1"
|
||||
),
|
||||
PassThroughGenericEndpoint(
|
||||
id="endpoint-2", path="/api/test2", target="http://example.com/api2"
|
||||
),
|
||||
]
|
||||
|
||||
# Mock prisma client with team that has None metadata
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_team = MagicMock()
|
||||
mock_team.metadata = None
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
|
||||
return_value=mock_team
|
||||
)
|
||||
|
||||
# Call the function
|
||||
result = await _filter_endpoints_by_team_allowed_routes(
|
||||
team_id="test-team-123",
|
||||
pass_through_endpoints=endpoints,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
# Should return all endpoints when no metadata
|
||||
assert len(result) == 2
|
||||
assert result[0].path == "/api/test1"
|
||||
assert result[1].path == "/api/test2"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_endpoints_by_team_allowed_routes_no_allowed_routes_key():
|
||||
"""
|
||||
Test that _filter_endpoints_by_team_allowed_routes returns all endpoints
|
||||
when team metadata doesn't have allowed_passthrough_routes key
|
||||
"""
|
||||
from litellm.proxy._types import PassThroughGenericEndpoint
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
_filter_endpoints_by_team_allowed_routes,
|
||||
)
|
||||
|
||||
# Create test endpoints
|
||||
endpoints = [
|
||||
PassThroughGenericEndpoint(
|
||||
id="endpoint-1", path="/api/test1", target="http://example.com/api1"
|
||||
),
|
||||
PassThroughGenericEndpoint(
|
||||
id="endpoint-2", path="/api/test2", target="http://example.com/api2"
|
||||
),
|
||||
]
|
||||
|
||||
# Mock prisma client with team that has metadata but no allowed_passthrough_routes
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_team = MagicMock()
|
||||
mock_team.metadata = {"some_other_key": "some_value"}
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
|
||||
return_value=mock_team
|
||||
)
|
||||
|
||||
# Call the function
|
||||
result = await _filter_endpoints_by_team_allowed_routes(
|
||||
team_id="test-team-123",
|
||||
pass_through_endpoints=endpoints,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
# Should return all endpoints when allowed_passthrough_routes key doesn't exist
|
||||
assert len(result) == 2
|
||||
assert result[0].path == "/api/test1"
|
||||
assert result[1].path == "/api/test2"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_endpoints_by_team_allowed_routes_empty_allowed_list():
|
||||
"""
|
||||
Test that _filter_endpoints_by_team_allowed_routes returns empty list
|
||||
when team has empty allowed_passthrough_routes list
|
||||
"""
|
||||
from litellm.proxy._types import PassThroughGenericEndpoint
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
_filter_endpoints_by_team_allowed_routes,
|
||||
)
|
||||
|
||||
# Create test endpoints
|
||||
endpoints = [
|
||||
PassThroughGenericEndpoint(
|
||||
id="endpoint-1", path="/api/test1", target="http://example.com/api1"
|
||||
),
|
||||
PassThroughGenericEndpoint(
|
||||
id="endpoint-2", path="/api/test2", target="http://example.com/api2"
|
||||
),
|
||||
]
|
||||
|
||||
# Mock prisma client with team that has empty allowed_passthrough_routes
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_team = MagicMock()
|
||||
mock_team.metadata = {"allowed_passthrough_routes": []}
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
|
||||
return_value=mock_team
|
||||
)
|
||||
|
||||
# Call the function
|
||||
result = await _filter_endpoints_by_team_allowed_routes(
|
||||
team_id="test-team-123",
|
||||
pass_through_endpoints=endpoints,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
# Should return empty list when allowed_passthrough_routes is empty
|
||||
assert len(result) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_endpoints_by_team_allowed_routes_partial_match():
|
||||
"""
|
||||
Test that _filter_endpoints_by_team_allowed_routes correctly filters
|
||||
when only some endpoints match allowed routes
|
||||
"""
|
||||
from litellm.proxy._types import PassThroughGenericEndpoint
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
_filter_endpoints_by_team_allowed_routes,
|
||||
)
|
||||
|
||||
# Create test endpoints
|
||||
endpoints = [
|
||||
PassThroughGenericEndpoint(
|
||||
id="endpoint-1", path="/api/openai", target="http://example.com/openai"
|
||||
),
|
||||
PassThroughGenericEndpoint(
|
||||
id="endpoint-2",
|
||||
path="/api/anthropic",
|
||||
target="http://example.com/anthropic",
|
||||
),
|
||||
PassThroughGenericEndpoint(
|
||||
id="endpoint-3", path="/api/azure", target="http://example.com/azure"
|
||||
),
|
||||
PassThroughGenericEndpoint(
|
||||
id="endpoint-4", path="/api/cohere", target="http://example.com/cohere"
|
||||
),
|
||||
]
|
||||
|
||||
# Mock prisma client with team that allows only 2 routes
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_team = MagicMock()
|
||||
mock_team.metadata = {"allowed_passthrough_routes": ["/api/openai", "/api/azure"]}
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
|
||||
return_value=mock_team
|
||||
)
|
||||
|
||||
# Call the function
|
||||
result = await _filter_endpoints_by_team_allowed_routes(
|
||||
team_id="test-team-123",
|
||||
pass_through_endpoints=endpoints,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
# Should return only the 2 allowed endpoints
|
||||
assert len(result) == 2
|
||||
assert result[0].path == "/api/openai"
|
||||
assert result[1].path == "/api/azure"
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ interface PassThroughRoutesSelectorProps {
|
|||
accessToken: string;
|
||||
placeholder?: string;
|
||||
disabled?: boolean;
|
||||
teamId?: string | null;
|
||||
}
|
||||
|
||||
const PassThroughRoutesSelector: React.FC<PassThroughRoutesSelectorProps> = ({
|
||||
|
|
@ -18,6 +19,7 @@ const PassThroughRoutesSelector: React.FC<PassThroughRoutesSelectorProps> = ({
|
|||
accessToken,
|
||||
placeholder = "Select pass through routes",
|
||||
disabled = false,
|
||||
teamId,
|
||||
}) => {
|
||||
const [passThroughRoutes, setPassThroughRoutes] = useState<string[]>([]);
|
||||
const [loading, setLoading] = useState(false);
|
||||
|
|
@ -28,7 +30,7 @@ const PassThroughRoutesSelector: React.FC<PassThroughRoutesSelectorProps> = ({
|
|||
|
||||
setLoading(true);
|
||||
try {
|
||||
const response = await getPassThroughEndpointsCall(accessToken);
|
||||
const response = await getPassThroughEndpointsCall(accessToken, teamId);
|
||||
if (response.endpoints) {
|
||||
const routes = response.endpoints.map((route: { path: string }) => route.path);
|
||||
setPassThroughRoutes(routes);
|
||||
|
|
@ -41,7 +43,7 @@ const PassThroughRoutesSelector: React.FC<PassThroughRoutesSelectorProps> = ({
|
|||
};
|
||||
|
||||
fetchPassThroughRoutes();
|
||||
}, [accessToken]);
|
||||
}, [accessToken, teamId]);
|
||||
|
||||
return (
|
||||
<Select
|
||||
|
|
|
|||
|
|
@ -4453,10 +4453,14 @@ export const getGeneralSettingsCall = async (accessToken: String) => {
|
|||
}
|
||||
};
|
||||
|
||||
export const getPassThroughEndpointsCall = async (accessToken: String) => {
|
||||
export const getPassThroughEndpointsCall = async (accessToken: String, teamId?: string | null) => {
|
||||
try {
|
||||
let url = proxyBaseUrl ? `${proxyBaseUrl}/config/pass_through_endpoint` : `/config/pass_through_endpoint`;
|
||||
|
||||
if (teamId) {
|
||||
url += `/team/${teamId}`;
|
||||
}
|
||||
|
||||
//NotificationsManager.info("Requesting model data");
|
||||
const response = await fetch(url, {
|
||||
method: "GET",
|
||||
|
|
|
|||
|
|
@ -947,6 +947,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({
|
|||
!premiumUser ? "Premium feature - Upgrade to set pass through routes by key" : "Select or enter pass through routes"
|
||||
}
|
||||
disabled={!premiumUser}
|
||||
teamId={selectedCreateKeyTeam ? selectedCreateKeyTeam.team_id : null}
|
||||
/>
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue