import json import re import traceback from typing import AsyncGenerator import httpx from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Request from fastapi.responses import Response, StreamingResponse from router.payment.helpers import ( UPSTREAM_BASE_URL, check_token_balance, create_error_response, prepare_upstream_headers, ) from router.payment.x_cashu import x_cashu_handler from .auth import adjust_payment_for_tokens, pay_for_request, validate_bearer_key from .cashu import x_cashu_refund from .db import ApiKey, AsyncSession, create_session, get_session proxy_router = APIRouter() async def handle_streaming_chat_completion( response: httpx.Response, key: ApiKey, session: AsyncSession ) -> StreamingResponse: """Handle streaming chat completion responses with token-based pricing.""" async def stream_with_cost() -> AsyncGenerator[bytes, None]: # Store all chunks to analyze stored_chunks = [] async for chunk in response.aiter_bytes(): # Store chunk for later analysis stored_chunks.append(chunk) # Pass through each chunk to client yield chunk # Process stored chunks to find usage data # Start from the end and work backwards for i in range(len(stored_chunks) - 1, -1, -1): chunk = stored_chunks[i] if not chunk or chunk == b"": continue try: # Split by "data: " to get individual SSE events events = re.split(b"data: ", chunk) for event_data in events: if ( not event_data or event_data.strip() == b"[DONE]" or event_data.strip() == b"" ): continue try: data = json.loads(event_data) if ( "usage" in data and data["usage"] is not None and isinstance(data["usage"], dict) ): # Found usage data, calculate cost # Create a new session for this operation async with create_session() as new_session: # Re-fetch the key in the new session fresh_key = await new_session.get( key.__class__, key.hashed_key ) if fresh_key: cost_data = await adjust_payment_for_tokens( fresh_key, data, new_session ) # Format as SSE and yield cost_json = json.dumps({"cost": cost_data}) yield f"data: {cost_json}\n\n".encode() break except json.JSONDecodeError: continue except Exception as e: print(f"Error processing streaming response for cost: {e}") return StreamingResponse( stream_with_cost(), status_code=response.status_code, headers=dict(response.headers), ) async def handle_non_streaming_chat_completion( response: httpx.Response, key: ApiKey, session: AsyncSession ) -> Response: """Handle non-streaming chat completion responses with token-based pricing.""" try: content = await response.aread() response_json = json.loads(content) cost_data = await adjust_payment_for_tokens(key, response_json, session) response_json["cost"] = cost_data # Keep only standard headers that are safe to pass through allowed_headers = { "content-type", "cache-control", "date", "vary", "access-control-allow-origin", "access-control-allow-methods", "access-control-allow-headers", "access-control-allow-credentials", "access-control-expose-headers", "access-control-max-age", } response_headers = { k: v for k, v in response.headers.items() if k.lower() in allowed_headers } return Response( content=json.dumps(response_json).encode(), status_code=response.status_code, headers=response_headers, media_type="application/json", ) except json.JSONDecodeError as e: print(f"Failed to parse JSON from upstream response: {e}") raise except Exception as e: print(f"Error adjusting payment for tokens: {e}") raise async def forward_to_upstream( request: Request, path: str, headers: dict, request_body: bytes | None, key: ApiKey, session: AsyncSession, ) -> Response | StreamingResponse: """Forward request to upstream and handle the response.""" if path.startswith("v1/"): path = path.replace("v1/", "") url = f"{UPSTREAM_BASE_URL}/{path}" client = httpx.AsyncClient( transport=httpx.AsyncHTTPTransport(retries=1), timeout=None, # No timeout - requests can take as long as needed ) try: # Use the pre-read body if available, otherwise stream if request_body is not None: response = await client.send( client.build_request( request.method, url, headers=headers, content=request_body, params=request.query_params, ), stream=True, ) else: response = await client.send( client.build_request( request.method, url, headers=headers, content=request.stream(), params=request.query_params, ), stream=True, ) # For chat completions, we need to handle token-based pricing if path.endswith("chat/completions"): # Check if client requested streaming client_wants_streaming = False if request_body: try: request_data = json.loads(request_body) client_wants_streaming = request_data.get("stream", False) except json.JSONDecodeError: pass # Handle both streaming and non-streaming responses content_type = response.headers.get("content-type", "") upstream_is_streaming = "text/event-stream" in content_type is_streaming = client_wants_streaming and upstream_is_streaming if is_streaming and response.status_code == 200: # Process streaming response and extract cost from the last chunk result = await handle_streaming_chat_completion(response, key, session) background_tasks = BackgroundTasks() background_tasks.add_task(response.aclose) background_tasks.add_task(client.aclose) result.background = background_tasks return result elif response.status_code == 200: # Handle non-streaming response try: return await handle_non_streaming_chat_completion( response, key, session ) finally: await response.aclose() await client.aclose() # For all other responses, stream the response background_tasks = BackgroundTasks() background_tasks.add_task(response.aclose) background_tasks.add_task(client.aclose) return StreamingResponse( response.aiter_bytes(), status_code=response.status_code, headers=dict(response.headers), background=background_tasks, ) except httpx.RequestError as exc: await client.aclose() error_type = type(exc).__name__ error_details = str(exc) print( f"Error forwarding request to upstream: {error_type}: {error_details}\n" f"Request details: method={request.method}, url={url}, headers={headers}, " f"path={path}, query_params={dict(request.query_params)}" ) # Provide more specific error messages based on the error type if isinstance(exc, httpx.ConnectError): error_message = "Unable to connect to upstream service" elif isinstance(exc, httpx.TimeoutException): error_message = "Upstream service request timed out" elif isinstance(exc, httpx.NetworkError): error_message = "Network error while connecting to upstream service" else: error_message = f"Error connecting to upstream service: {error_type}" return create_error_response("upstream_error", error_message, 502) except Exception as exc: await client.aclose() import traceback tb = traceback.format_exc() print( f"Unexpected error: {exc}\n" f"Request details: method={request.method}, url={url}, headers={headers}, " f"path={path}, query_params={dict(request.query_params)}\n" f"Traceback:\n{tb}" ) return create_error_response( "internal_error", "An unexpected server error occurred", 500 ) @proxy_router.api_route("/{path:path}", methods=["GET", "POST"], response_model=None) async def proxy( request: Request, path: str, session: AsyncSession = Depends(get_session) ) -> Response | StreamingResponse: request_body = await request.body() headers = dict(request.headers) # Parse JSON body if present, handle empty/invalid JSON request_body_dict = {} if request_body: try: request_body_dict = json.loads(request_body) except json.JSONDecodeError as e: print(f"Error: failed to parse request body '{e}'") return Response( content=json.dumps( {"error": {"type": "invalid_request_error", "code": "invalid_json"}} ), status_code=400, media_type="application/json", ) # Check token balance for all requests to get currency unit unit = check_token_balance(headers, request_body_dict) # Handle authentication if x_cashu := headers.get("x-cashu", None): return await x_cashu_handler(request, x_cashu, path) elif auth := headers.get("authorization", None): key = await get_bearer_token_key(headers, path, session, auth) else: if request.method not in ["GET"]: return Response( content=json.dumps({"detail": "Unauthorized"}), status_code=401, media_type="application/json", ) # Prepare headers for upstream headers = prepare_upstream_headers(dict(request.headers)) return await forward_get_to_upstream(request, path, headers) # Only pay for request if we have request body data (for completions endpoints) if request_body_dict: await pay_for_request(key, session, request_body_dict) # Prepare headers for upstream headers = prepare_upstream_headers(dict(request.headers)) # Forward to upstream and handle response response = await forward_to_upstream( request, path, headers, request_body, key, session ) if response.status_code != 200 and key.refund_address == "X-CASHU": refund_token = await x_cashu_refund(key, session, unit) response = Response( content=json.dumps( { "error": { "message": "Error forwarding request to upstream", "type": "upstream_error", "code": response.status_code, "refund_token": refund_token, } } ), status_code=response.status_code, media_type="application/json", ) response.headers["X-Cashu"] = refund_token return response if key.refund_address == "X-CASHU": refund_token = await x_cashu_refund(key, session, unit) response.headers["X-Cashu"] = refund_token return response async def get_bearer_token_key( headers: dict, path: str, session: AsyncSession, auth: str ) -> ApiKey: """Handle bearer token authentication proxy requests.""" bearer_key = auth.replace("Bearer ", "") if auth.startswith("Bearer ") else "" refund_address = headers.get("Refund-LNURL", None) key_expiry_time = headers.get("Key-Expiry-Time", None) # Validate key_expiry_time header if key_expiry_time: try: key_expiry_time = int(key_expiry_time) # type: ignore except ValueError: print("Invalid Key-Expiry-Time: must be a valid Unix timestamp") raise HTTPException( status_code=400, detail="Invalid Key-Expiry-Time: must be a valid Unix timestamp", ) if not refund_address: print("Error: Refund-LNURL header required when using Key-Expiry-Time") raise HTTPException( status_code=400, detail="Error: Refund-LNURL header required when using Key-Expiry-Time", ) else: key_expiry_time = None return await validate_bearer_key( bearer_key, session, refund_address, key_expiry_time, # type: ignore ) async def forward_get_to_upstream( request: Request, path: str, headers: dict, ) -> Response | StreamingResponse: """Forward request to upstream and handle the response.""" if path.startswith("v1/"): path = path.replace("v1/", "") url = f"{UPSTREAM_BASE_URL}/{path}" async with httpx.AsyncClient( transport=httpx.AsyncHTTPTransport(retries=1), timeout=None, ) as client: try: response = await client.send( client.build_request( request.method, url, headers=headers, content=request.stream(), params=request.query_params, ), ) return StreamingResponse( response.aiter_bytes(), status_code=response.status_code, headers=dict(response.headers), ) except Exception as exc: tb = traceback.format_exc() print( f"Unexpected error: {exc}\n" f"Request details: method={request.method}, url={url}, headers={headers}, " f"path={path}, query_params={dict(request.query_params)}\n" f"Traceback:\n{tb}" ) return create_error_response( "internal_error", "An unexpected server error occurred", 500 )