routstr-core/router/proxy.py

import asyncio
import os
import json
from fastapi import APIRouter, Request, BackgroundTasks, Depends
from fastapi.responses import Response, StreamingResponse
import httpx
import re

from router.cashu import pay_out_with_new_session

from .auth import validate_bearer_key, pay_for_request, adjust_payment_for_tokens
from .db import AsyncSession, get_session

UPSTREAM_BASE_URL = os.environ["UPSTREAM_BASE_URL"]
UPSTREAM_API_KEY = os.environ.get("UPSTREAM_API_KEY", "")

proxy_router = APIRouter()


@proxy_router.api_route(
    "/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "OPTIONS", "HEAD", "PATCH"]
)
async def proxy(
    request: Request, path: str, session: AsyncSession = Depends(get_session)
):
    auth = request.headers.get("Authorization", "")
    bearer_key = auth.replace("Bearer ", "") if auth.startswith("Bearer ") else ""
    refund_address = request.headers.get("Refund-LNURL", None)
    key_expiry_time = request.headers.get("Key-Expiry-Time", None)

    # Validate key_expiry_time header
    if key_expiry_time:
        try:
            key_expiry_time = int(key_expiry_time)
        except ValueError:
            return Response(
                content="Invalid Key-Expiry-Time: must be a valid Unix timestamp",
                status_code=400,
            )
    else:
        key_expiry_time = None

    if(key_expiry_time and not refund_address):
        return Response(
                content=f"Error: Refund-LNURL header required when using Key-Expiry-Time",
                status_code=400,
            )
    

    key = await validate_bearer_key(bearer_key, session, refund_address, key_expiry_time)
    
    # Pre-validate JSON for requests that require it
    request_body = None
    if request.method in ["POST", "PUT", "PATCH"] and path.endswith("chat/completions"):
        try:
            request_body = await request.body()
            # Try to parse JSON to validate it
            if request_body:
                json.loads(request_body)
        except json.JSONDecodeError as e:
            return Response(
                content=json.dumps({
                    "error": {
                        "message": f"Invalid JSON in request body: {str(e)}",
                        "type": "invalid_request_error",
                        "code": "invalid_json"
                    }
                }),
                status_code=400,
                media_type="application/json"
            )
        except Exception as e:
            return Response(
                content=json.dumps({
                    "error": {
                        "message": "Error reading request body",
                        "type": "invalid_request_error",
                        "code": "request_error"
                    }
                }),
                status_code=400,
                media_type="application/json"
            )
    
    await pay_for_request(key, session, request, request_body)

    # Prepare headers, removing sensitive/problematic ones
    headers = dict(request.headers)
    headers.pop("host", None)
    headers.pop("content-length", None)
    headers.pop("refund-lnurl", None)
    headers.pop("key-expiry-time", None)

    if UPSTREAM_API_KEY:
        headers["Authorization"] = f"Bearer {UPSTREAM_API_KEY}"
        headers.pop("authorization", None)
    else:
        headers.pop("Authorization", None)
        headers.pop("authorization", None)

    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"):
            # Handle both streaming and non-streaming responses
            content_type = response.headers.get("content-type", "")
            is_streaming = "text/event-stream" in content_type

            if is_streaming and response.status_code == 200:
                # Process streaming response and extract cost from the last chunk
                async def stream_with_cost():
                    # Store all chunks to analyze
                    stored_chunks = []
                    usage_data_found = False

                    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
                                        cost_data = await adjust_payment_for_tokens(
                                            key, data, session
                                        )
                                        # Format as SSE and yield
                                        cost_json = json.dumps({"cost": cost_data})
                                        yield f"data: {cost_json}\n\n".encode()
                                        usage_data_found = True
                                        break
                                except json.JSONDecodeError:
                                    continue
                            
                        except Exception as e:
                            print(f"Error processing streaming response for cost: {e}")

                background_tasks = BackgroundTasks()
                background_tasks.add_task(response.aclose)
                background_tasks.add_task(client.aclose)

                return StreamingResponse(
                    stream_with_cost(),
                    status_code=response.status_code,
                    headers=dict(response.headers),
                    background=background_tasks,
                )

            elif response.status_code == 200 and "application/json" in content_type:
                # Handle non-streaming response
                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
                    return Response(
                        content=json.dumps(response_json).encode(),
                        status_code=response.status_code,
                        headers=dict(response.headers),
                        media_type="application/json",
                    )
                except json.JSONDecodeError as e:
                    print(f"Failed to parse JSON from upstream response: {e}")
                except Exception as e:
                    print(f"Error adjusting payment for tokens: {e}")
                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)
        background_tasks.add_task(pay_out_with_new_session)

        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 Response(
            content=json.dumps({
                "error": {
                    "message": error_message,
                    "type": "upstream_error",
                    "code": 502
                }
            }),
            status_code=502,
            media_type="application/json"
        )
    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 Response(
            content=json.dumps({
                "error": {
                    "message": "An unexpected server error occurred",
                    "type": "internal_error",
                    "code": 500
                }
            }),
            status_code=500,
            media_type="application/json"
        )