From f321c20118d88779463ce7649ae56077a3c92aca Mon Sep 17 00:00:00 2001 From: Fabian Franz Date: Thu, 7 Aug 2025 00:42:35 +0100 Subject: [PATCH 1/3] feat: Add bedrock_proxy --- src/api/app.py | 3 +- src/api/routers/bedrock_proxy.py | 111 +++++++++++++++++++++++++++++++ src/requirements.txt | 1 + 3 files changed, 114 insertions(+), 1 deletion(-) create mode 100644 src/api/routers/bedrock_proxy.py diff --git a/src/api/app.py b/src/api/app.py index 5ea7ae78..d3f9c312 100644 --- a/src/api/app.py +++ b/src/api/app.py @@ -7,7 +7,7 @@ from fastapi.responses import PlainTextResponse from mangum import Mangum -from api.routers import chat, embeddings, model +from api.routers import chat, embeddings, model, bedrock_proxy from api.setting import API_ROUTE_PREFIX, DESCRIPTION, SUMMARY, TITLE, VERSION config = { @@ -35,6 +35,7 @@ app.include_router(model.router, prefix=API_ROUTE_PREFIX) app.include_router(chat.router, prefix=API_ROUTE_PREFIX) app.include_router(embeddings.router, prefix=API_ROUTE_PREFIX) +app.include_router(bedrock_proxy.router, prefix=API_ROUTE_PREFIX) @app.get("/health") diff --git a/src/api/routers/bedrock_proxy.py b/src/api/routers/bedrock_proxy.py new file mode 100644 index 00000000..c5104cda --- /dev/null +++ b/src/api/routers/bedrock_proxy.py @@ -0,0 +1,111 @@ +import os +import logging +from typing import Dict, Any +from urllib.parse import quote + +import httpx +from fastapi import APIRouter, Depends, HTTPException, Request, BackgroundTasks +from fastapi.responses import StreamingResponse, Response + +from api.auth import api_key_auth +from api.setting import AWS_REGION, DEBUG + +logger = logging.getLogger(__name__) + +router = APIRouter(prefix="/bedrock") + +# Get AWS bearer token from environment +AWS_BEARER_TOKEN = os.environ.get("AWS_BEARER_TOKEN_BEDROCK") + +if not AWS_BEARER_TOKEN: + logger.warning("AWS_BEARER_TOKEN_BEDROCK not set - bedrock proxy endpoints will not work") + + +def get_aws_url(model_id: str, endpoint_path: str) -> str: + """Convert proxy path to AWS Bedrock URL""" + encoded_model_id = quote(model_id, safe='') + base_url = f"https://bedrock-runtime.{AWS_REGION}.amazonaws.com" + return f"{base_url}/model/{encoded_model_id}/{endpoint_path}" + + +def get_proxy_headers(request: Request) -> Dict[str, str]: + """Get headers to forward to AWS, replacing Authorization""" + headers = dict(request.headers) + + # Remove proxy authorization and add AWS bearer token + headers.pop("authorization", None) + headers.pop("host", None) # Let httpx set the correct host + + if AWS_BEARER_TOKEN: + headers["Authorization"] = f"Bearer {AWS_BEARER_TOKEN}" + + return headers + + +@router.api_route("/model/{model_id}/{endpoint_path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"]) +async def transparent_proxy( + request: Request, + background_tasks: BackgroundTasks, + model_id: str, + endpoint_path: str, + _: None = Depends(api_key_auth) +): + """ + Transparent HTTP proxy to AWS Bedrock. + Forwards all requests as-is, only changing auth and URL. + """ + if not AWS_BEARER_TOKEN: + raise HTTPException( + status_code=503, + detail="AWS_BEARER_TOKEN_BEDROCK not configured" + ) + + # Build AWS URL + aws_url = get_aws_url(model_id, endpoint_path) + + # Get headers to forward + proxy_headers = get_proxy_headers(request) + + # Get request body + body = await request.body() + + if DEBUG: + logger.info(f"Proxying {request.method} to: {aws_url}") + logger.info(f"Headers: {dict(proxy_headers)}") + if body: + logger.info(f"Body length: {len(body)} bytes") + + try: + # Always use streaming for transparent pass-through + client = httpx.AsyncClient() + + # Add cleanup task + async def cleanup_client(): + await client.aclose() + + background_tasks.add_task(cleanup_client) + + async def stream_generator(): + async with client.stream( + method=request.method, + url=aws_url, + headers=proxy_headers, + content=body, + params=request.query_params, + timeout=120.0 + ) as response: + async for chunk in response.aiter_bytes(): + if chunk: # Only yield non-empty chunks + yield chunk + + return StreamingResponse(content=stream_generator()) + + except httpx.RequestError as e: + logger.error(f"Proxy request failed: {e}") + raise HTTPException(status_code=502, detail=f"Upstream request failed: {str(e)}") + except httpx.HTTPStatusError as e: + logger.error(f"AWS returned error: {e.response.status_code}") + raise HTTPException(status_code=e.response.status_code, detail=e.response.text) + except Exception as e: + logger.error(f"Proxy error: {e}") + raise HTTPException(status_code=500, detail="Proxy error") \ No newline at end of file diff --git a/src/requirements.txt b/src/requirements.txt index c8139beb..7c022b53 100644 --- a/src/requirements.txt +++ b/src/requirements.txt @@ -4,6 +4,7 @@ uvicorn==0.29.0 mangum==0.17.0 tiktoken==0.6.0 requests==2.32.4 +httpx==0.27.0 numpy==1.26.4 boto3==1.37.0 botocore==1.37.0 \ No newline at end of file From 6375126773576a0c494c4d7bdaf054d8c5e76cfc Mon Sep 17 00:00:00 2001 From: Fabian Franz Date: Thu, 7 Aug 2025 20:15:44 +0100 Subject: [PATCH 2/3] feat: Automatically generate bedrock token --- src/api/routers/bedrock_proxy.py | 30 ++++++++++++++++++++---------- src/requirements.txt | 3 ++- 2 files changed, 22 insertions(+), 11 deletions(-) diff --git a/src/api/routers/bedrock_proxy.py b/src/api/routers/bedrock_proxy.py index c5104cda..decd14c5 100644 --- a/src/api/routers/bedrock_proxy.py +++ b/src/api/routers/bedrock_proxy.py @@ -6,6 +6,7 @@ import httpx from fastapi import APIRouter, Depends, HTTPException, Request, BackgroundTasks from fastapi.responses import StreamingResponse, Response +from aws_bedrock_token_generator import provide_token from api.auth import api_key_auth from api.setting import AWS_REGION, DEBUG @@ -14,11 +15,23 @@ router = APIRouter(prefix="/bedrock") -# Get AWS bearer token from environment +# Get static token if provided (convenience feature) AWS_BEARER_TOKEN = os.environ.get("AWS_BEARER_TOKEN_BEDROCK") -if not AWS_BEARER_TOKEN: - logger.warning("AWS_BEARER_TOKEN_BEDROCK not set - bedrock proxy endpoints will not work") +def get_aws_bearer_token() -> str: + """Get AWS bearer token - static if provided, otherwise auto-generate""" + if AWS_BEARER_TOKEN: + logger.debug("Using static AWS bearer token") + return AWS_BEARER_TOKEN + + # Default: auto-generate token using AWS SDK credentials + try: + token = provide_token(region=AWS_REGION) + logger.debug("Generated fresh AWS Bedrock token") + return token + except Exception as e: + logger.error(f"Failed to generate AWS token: {e}") + raise HTTPException(status_code=503, detail="Failed to generate AWS authentication token. Ensure AWS credentials are configured or set AWS_BEARER_TOKEN_BEDROCK") def get_aws_url(model_id: str, endpoint_path: str) -> str: @@ -36,8 +49,9 @@ def get_proxy_headers(request: Request) -> Dict[str, str]: headers.pop("authorization", None) headers.pop("host", None) # Let httpx set the correct host - if AWS_BEARER_TOKEN: - headers["Authorization"] = f"Bearer {AWS_BEARER_TOKEN}" + # Get fresh AWS token (static or auto-generated) + aws_token = get_aws_bearer_token() + headers["Authorization"] = f"Bearer {aws_token}" return headers @@ -53,12 +67,8 @@ async def transparent_proxy( """ Transparent HTTP proxy to AWS Bedrock. Forwards all requests as-is, only changing auth and URL. + Supports both static tokens and auto-refresh tokens. """ - if not AWS_BEARER_TOKEN: - raise HTTPException( - status_code=503, - detail="AWS_BEARER_TOKEN_BEDROCK not configured" - ) # Build AWS URL aws_url = get_aws_url(model_id, endpoint_path) diff --git a/src/requirements.txt b/src/requirements.txt index 7c022b53..b71a9fb8 100644 --- a/src/requirements.txt +++ b/src/requirements.txt @@ -7,4 +7,5 @@ requests==2.32.4 httpx==0.27.0 numpy==1.26.4 boto3==1.37.0 -botocore==1.37.0 \ No newline at end of file +botocore==1.37.0 +aws-bedrock-token-generator From 58a26b85ea7013bf1c646cb28e8cbdeb9435a46c Mon Sep 17 00:00:00 2001 From: Fabian Franz Date: Tue, 12 Aug 2025 16:05:58 +0100 Subject: [PATCH 3/3] fix: Pass through AWS Bedrock response headers in transparent proxy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously the transparent proxy was using FastAPI's default headers. Now properly forwards all AWS response headers including content-type and x-amzn-requestid for better API compatibility. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude --- src/api/routers/bedrock_proxy.py | 41 ++++++++++++++++++++++---------- 1 file changed, 28 insertions(+), 13 deletions(-) diff --git a/src/api/routers/bedrock_proxy.py b/src/api/routers/bedrock_proxy.py index decd14c5..7db63385 100644 --- a/src/api/routers/bedrock_proxy.py +++ b/src/api/routers/bedrock_proxy.py @@ -95,20 +95,35 @@ async def cleanup_client(): background_tasks.add_task(cleanup_client) + # Use a custom response class that captures headers from the stream + stream_request = client.stream( + method=request.method, + url=aws_url, + headers=proxy_headers, + content=body, + params=request.query_params, + timeout=120.0 + ) + + # Start the stream to get response object + response = await stream_request.__aenter__() + + # Schedule cleanup + async def cleanup_stream(): + await stream_request.__aexit__(None, None, None) + background_tasks.add_task(cleanup_stream) + async def stream_generator(): - async with client.stream( - method=request.method, - url=aws_url, - headers=proxy_headers, - content=body, - params=request.query_params, - timeout=120.0 - ) as response: - async for chunk in response.aiter_bytes(): - if chunk: # Only yield non-empty chunks - yield chunk - - return StreamingResponse(content=stream_generator()) + async for chunk in response.aiter_bytes(): + if chunk: # Only yield non-empty chunks + yield chunk + + # Create StreamingResponse with AWS response headers and status + return StreamingResponse( + content=stream_generator(), + status_code=response.status_code, + headers=dict(response.headers) + ) except httpx.RequestError as e: logger.error(f"Proxy request failed: {e}")