diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index f84cdd0c222..fc6f4eedae9 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1014,7 +1014,9 @@ async def bedrock_proxy_route( raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") aws_region_name: Final = litellm.utils.get_secret(secret_name="AWS_REGION_NAME") - if _is_bedrock_agent_runtime_route(endpoint=endpoint): # handle bedrock agents + if _is_bedrock_agent_control_plane_route(endpoint=endpoint): + base_target_url: Final = f"https://bedrock-agent.{aws_region_name}.amazonaws.com" + elif _is_bedrock_agent_runtime_route(endpoint=endpoint): # handle bedrock agents base_target_url: Final = f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com" else: return await bedrock_llm_proxy_route( @@ -1048,7 +1050,7 @@ async def bedrock_proxy_route( data: Final = await request.json() except Exception as e: raise HTTPException(status_code=400, detail={"error": e}) - _request: Final = AWSRequest(method="POST", url=str(updated_url), data=json.dumps(data), headers=headers) + _request: Final = AWSRequest(method=request.method, url=str(updated_url), data=json.dumps(data), headers=headers) sigv4.add_auth(_request) prepped: Final = _request.prepare() @@ -1166,6 +1168,11 @@ def _is_bedrock_agent_runtime_route(endpoint: str) -> bool: return False +def _is_bedrock_agent_control_plane_route(endpoint: str) -> bool: + normalized_endpoint: Final = endpoint.strip("/") + return re.fullmatch(r"knowledgebases/[^/]+/datasources/[^/]+/documents", normalized_endpoint) is not None + + @router.api_route( "/assemblyai/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], diff --git a/tests/pass_through_unit_tests/test_pass_through_unit_tests.py b/tests/pass_through_unit_tests/test_pass_through_unit_tests.py index c263b8ce381..4c0dd4ad4ef 100644 --- a/tests/pass_through_unit_tests/test_pass_through_unit_tests.py +++ b/tests/pass_through_unit_tests/test_pass_through_unit_tests.py @@ -3,7 +3,7 @@ import os import sys from datetime import datetime from unittest.mock import AsyncMock, Mock, patch, MagicMock -from typing import Optional +from typing import Final, Optional sys.path.insert( 0, os.path.abspath("../..") @@ -468,6 +468,120 @@ def test_is_bedrock_agent_runtime_route(): assert _is_bedrock_agent_runtime_route("/some/random/endpoint") is False +def test_is_bedrock_agent_control_plane_route(): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _is_bedrock_agent_control_plane_route, + _is_bedrock_agent_runtime_route, + ) + + direct_ingestion_endpoint: Final = ( + "/knowledgebases/kb-123/datasources/ds-456/documents" + ) + + assert _is_bedrock_agent_control_plane_route(direct_ingestion_endpoint) is True + assert _is_bedrock_agent_runtime_route(direct_ingestion_endpoint) is True + assert ( + _is_bedrock_agent_control_plane_route( + "/knowledgebases/kb-123/retrieve" + ) + is False + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("endpoint", "method", "expected_url"), + ( + ( + "knowledgebases/kb-123/retrieve", + "POST", + "https://bedrock-agent-runtime.ap-northeast-2.amazonaws.com/knowledgebases/kb-123/retrieve", + ), + ( + "knowledgebases/kb-123/datasources/ds-456/documents", + "PUT", + "https://bedrock-agent.ap-northeast-2.amazonaws.com/knowledgebases/kb-123/datasources/ds-456/documents", + ), + ), +) +async def test_bedrock_agent_route_target_and_sigv4_method( + endpoint: str, method: str, expected_url: str +): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + bedrock_proxy_route, + ) + + body: Final = json.dumps({"documents": []}).encode() + + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": body, "more_body": False} + + request: Final = Request( + { + "type": "http", + "method": method, + "path": f"/bedrock/{endpoint}", + "headers": ((b"content-type", b"application/json"),), + "query_string": b"", + "server": ("testserver", 80), + "client": ("testclient", 50000), + "scheme": "http", + }, + receive, + ) + prepared_request: Final = MagicMock( + url=expected_url, + headers={"Authorization": "AWS4-HMAC-SHA256 signed"}, + body=body, + ) + upstream_response: Final = MagicMock() + endpoint_func: Final = AsyncMock(return_value=upstream_response) + credentials: Final = MagicMock() + + with ( + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.litellm.utils.get_secret", + return_value="ap-northeast-2", + ), + patch( + "litellm.llms.bedrock.chat.BedrockConverseLLM.get_credentials", + return_value=credentials, + ), + patch("botocore.auth.SigV4Auth") as sigv4_auth, + patch("botocore.awsrequest.AWSRequest") as aws_request, + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route", + return_value=endpoint_func, + ) as create_route, + ): + aws_request.return_value.prepare.return_value = prepared_request + + response: Final = await bedrock_proxy_route( + endpoint=endpoint, + request=request, + fastapi_response=fastapi.Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + ) + + assert response is upstream_response + aws_request.assert_called_once_with( + method=method, + url=expected_url, + data=json.dumps({"documents": []}), + headers={"Content-Type": "application/json"}, + ) + sigv4_auth.assert_called_once_with(credentials, "bedrock", "ap-northeast-2") + sigv4_auth.return_value.add_auth.assert_called_once_with(aws_request.return_value) + create_route.assert_called_once_with( + endpoint=endpoint, + target=expected_url, + custom_headers=prepared_request.headers, + is_streaming_request=False, + _forward_headers=True, + ) + endpoint_func.assert_awaited_once() + + def test_init_kwargs_filters_pricing_params(mock_request, mock_user_api_key_dict): """ Test that pricing parameters are properly filtered out from the request body