diff --git a/routstr/proxy.py b/routstr/proxy.py index 394c92c0..6f8d976e 100644 --- a/routstr/proxy.py +++ b/routstr/proxy.py @@ -13,16 +13,13 @@ from .payment.helpers import ( create_error_response, get_max_cost_for_model, ) -from .upstream import UpstreamProvider +from .upstream import init_upstreams logger = get_logger(__name__) proxy_router = APIRouter() -upstream = UpstreamProvider( - base_url=settings.upstream_base_url, - api_key=settings.upstream_api_key, - chat_completions_api_version=settings.chat_completions_api_version, -) +upstreams = init_upstreams(settings.upstream_base_url, settings.upstream_api_key) +upstream = upstreams[0] @proxy_router.api_route("/{path:path}", methods=["GET", "POST"], response_model=None) diff --git a/routstr/upstream.py b/routstr/upstream.py index 1f9dce6b..d9728017 100644 --- a/routstr/upstream.py +++ b/routstr/upstream.py @@ -18,19 +18,67 @@ from .auth import adjust_payment_for_tokens from .core import get_logger from .core.db import ApiKey, AsyncSession, create_session from .payment.helpers import create_error_response +from .payment.models import Model logger = get_logger(__name__) +def init_upstreams( + base_url: str, api_key: str, api_version: str | None = None +) -> list[UpstreamProvider]: + """Initialize upstream providers based on settings. + + Args: + base_url: Base URL of the upstream API endpoint + api_key: API key for authenticating with the upstream service + api_version: API version for Azure OpenAI + """ + from .core.settings import settings + + upstreams: list[UpstreamProvider] = [] + if settings.chat_completions_api_version: + upstreams.append( + AzureUpstreamProvider( + settings.upstream_base_url, + settings.upstream_api_key, + settings.chat_completions_api_version, + ) + ) + + if "api.openai.com" in settings.upstream_base_url.lower(): + upstreams.append(OpenAIUpstreamProvider(settings.upstream_api_key)) + elif "openrouter.ai/api/v1" in settings.upstream_base_url.lower(): + upstreams.append(OpenRouterUpstreamProvider(settings.upstream_api_key)) + else: + upstreams.append( + UpstreamProvider(settings.upstream_base_url, settings.upstream_api_key) + ) + + return upstreams + + class UpstreamProvider: - def __init__( - self, base_url: str, api_key: str, chat_completions_api_version: str = "" - ): + """Provider for forwarding requests to an upstream AI service API.""" + + def __init__(self, base_url: str, api_key: str): + """Initialize the upstream provider. + + Args: + base_url: Base URL of the upstream API endpoint + api_key: API key for authenticating with the upstream service + """ self.base_url = base_url self.api_key = api_key - self.chat_completions_api_version = chat_completions_api_version def prepare_headers(self, request_headers: dict) -> dict: + """Prepare headers for upstream request by removing proxy-specific headers and adding authentication. + + Args: + request_headers: Original request headers from the client + + Returns: + Headers dict ready for upstream forwarding with authentication added + """ logger.debug( "Preparing upstream headers", extra={ @@ -74,15 +122,31 @@ class UpstreamProvider: def prepare_params( self, path: str, query_params: Mapping[str, str] | None - ) -> dict[str, str]: - params: dict[str, str] = dict(query_params or {}) - if path.endswith("chat/completions") and self.chat_completions_api_version: - params["api-version"] = self.chat_completions_api_version - return params + ) -> Mapping[str, str]: + """Prepare query parameters for upstream request. + + Base implementation passes through query params unchanged. Override in subclasses for provider-specific params. + + Args: + path: Request path + query_params: Original query parameters from the client + + Returns: + Query parameters dict ready for upstream forwarding + """ + return query_params or {} def _extract_upstream_error_message( self, body_bytes: bytes ) -> tuple[str, str | None]: + """Extract error message and code from upstream error response body. + + Args: + body_bytes: Raw response body bytes from upstream + + Returns: + Tuple of (error_message, error_code), where error_code may be None + """ message: str = "Upstream request failed" upstream_code: str | None = None if not body_bytes: @@ -115,6 +179,16 @@ class UpstreamProvider: async def map_upstream_error_response( self, request: Request, path: str, upstream_response: httpx.Response ) -> Response: + """Map upstream error responses to appropriate proxy error responses. + + Args: + request: Original FastAPI request + path: Request path + upstream_response: Response from upstream service + + Returns: + Mapped error response with appropriate status code and error type + """ status_code = upstream_response.status_code headers = dict(upstream_response.headers) content_type = headers.get("content-type", "") @@ -176,6 +250,16 @@ class UpstreamProvider: async def handle_streaming_chat_completion( self, response: httpx.Response, key: ApiKey, max_cost_for_model: int ) -> StreamingResponse: + """Handle streaming chat completion responses with token usage tracking and cost adjustment. + + Args: + response: Streaming response from upstream + key: API key for the authenticated user + max_cost_for_model: Maximum cost deducted upfront for the model + + Returns: + StreamingResponse with cost data injected at the end + """ logger.info( "Processing streaming chat completion", extra={ @@ -351,6 +435,17 @@ class UpstreamProvider: session: AsyncSession, deducted_max_cost: int, ) -> Response: + """Handle non-streaming chat completion responses with token usage tracking and cost adjustment. + + Args: + response: Response from upstream + key: API key for the authenticated user + session: Database session for updating balance + deducted_max_cost: Maximum cost deducted upfront + + Returns: + Response with cost data added to JSON body + """ logger.info( "Processing non-streaming chat completion", extra={ @@ -446,6 +541,20 @@ class UpstreamProvider: max_cost_for_model: int, session: AsyncSession, ) -> Response | StreamingResponse: + """Forward authenticated request to upstream service with cost tracking. + + Args: + request: Original FastAPI request + path: Request path + headers: Prepared headers for upstream + request_body: Request body bytes, if any + key: API key for authenticated user + max_cost_for_model: Maximum cost deducted upfront + session: Database session for balance updates + + Returns: + Response or StreamingResponse from upstream with cost tracking + """ if path.startswith("v1/"): path = path.replace("v1/", "") @@ -647,6 +756,16 @@ class UpstreamProvider: path: str, headers: dict, ) -> Response | StreamingResponse: + """Forward unauthenticated GET request to upstream service. + + Args: + request: Original FastAPI request + path: Request path + headers: Prepared headers for upstream + + Returns: + StreamingResponse from upstream + """ if path.startswith("v1/"): path = path.replace("v1/", "") @@ -714,6 +833,15 @@ class UpstreamProvider: async def get_x_cashu_cost( self, response_data: dict, max_cost_for_model: int ) -> MaxCostData | CostData | None: + """Calculate cost for X-Cashu payment based on response data. + + Args: + response_data: Response data containing model and usage information + max_cost_for_model: Maximum cost for the model + + Returns: + Cost data object (MaxCostData or CostData) or None if calculation fails + """ from .payment.cost_caculation import ( CostData, CostDataError, @@ -768,6 +896,16 @@ class UpstreamProvider: return None async def send_refund(self, amount: int, unit: str, mint: str | None = None) -> str: + """Create and send a refund token to the user. + + Args: + amount: Refund amount + unit: Unit of the refund (sat or msat) + mint: Optional mint URL for the refund token + + Returns: + Refund token string + """ from .wallet import send_token logger.debug( @@ -844,6 +982,18 @@ class UpstreamProvider: unit: str, max_cost_for_model: int, ) -> StreamingResponse: + """Handle streaming response for X-Cashu payment, calculating refund if needed. + + Args: + content_str: Response content as string + response: Original httpx response + amount: Payment amount received + unit: Payment unit (sat or msat) + max_cost_for_model: Maximum cost for the model + + Returns: + StreamingResponse with refund token in header if applicable + """ logger.debug( "Processing streaming response", extra={ @@ -964,6 +1114,18 @@ class UpstreamProvider: unit: str, max_cost_for_model: int, ) -> Response: + """Handle non-streaming response for X-Cashu payment, calculating refund if needed. + + Args: + content_str: Response content as string + response: Original httpx response + amount: Payment amount received + unit: Payment unit (sat or msat) + max_cost_for_model: Maximum cost for the model + + Returns: + Response with refund token in header if applicable + """ logger.debug( "Processing non-streaming response", extra={"amount": amount, "unit": unit, "content_length": len(content_str)}, @@ -1079,6 +1241,17 @@ class UpstreamProvider: async def handle_x_cashu_chat_completion( self, response: httpx.Response, amount: int, unit: str, max_cost_for_model: int ) -> StreamingResponse | Response: + """Handle chat completion response for X-Cashu payment, detecting streaming vs non-streaming. + + Args: + response: Response from upstream + amount: Payment amount received + unit: Payment unit (sat or msat) + max_cost_for_model: Maximum cost for the model + + Returns: + StreamingResponse or Response depending on response type + """ logger.debug( "Handling chat completion response", extra={"amount": amount, "unit": unit, "status_code": response.status_code}, @@ -1135,6 +1308,19 @@ class UpstreamProvider: unit: str, max_cost_for_model: int, ) -> Response | StreamingResponse: + """Forward request paid with X-Cashu token to upstream service. + + Args: + request: Original FastAPI request + path: Request path + headers: Prepared headers for upstream + amount: Payment amount from X-Cashu token + unit: Payment unit (sat or msat) + max_cost_for_model: Maximum cost for the model + + Returns: + Response or StreamingResponse with refund if applicable + """ if path.startswith("v1/"): path = path.replace("v1/", "") @@ -1271,6 +1457,17 @@ class UpstreamProvider: async def handle_x_cashu( self, request: Request, x_cashu_token: str, path: str, max_cost_for_model: int ) -> Response | StreamingResponse: + """Handle request with X-Cashu token payment, redeeming token and forwarding request. + + Args: + request: Original FastAPI request + x_cashu_token: X-Cashu token from request header + path: Request path + max_cost_for_model: Maximum cost for the model + + Returns: + Response or StreamingResponse from upstream with refund if applicable + """ from .wallet import recieve_token logger.info( @@ -1348,3 +1545,82 @@ class UpstreamProvider: request=request, token=x_cashu_token, ) + + +class OpenAIUpstreamProvider(UpstreamProvider): + """Upstream provider specifically configured for OpenAI API.""" + + def __init__(self, api_key: str): + super().__init__(base_url="https://api.openai.com", api_key=api_key) + + +class AzureUpstreamProvider(UpstreamProvider): + """Upstream provider specifically configured for Azure OpenAI Service.""" + + def __init__(self, base_url: str, api_key: str, api_version: str): + """Initialize Azure provider with API key and version. + + Args: + base_url: Azure OpenAI endpoint base URL + api_key: Azure OpenAI API key for authentication + api_version: Azure OpenAI API version (e.g., "2024-02-15-preview") + """ + super().__init__(base_url=base_url, api_key=api_key) + self.api_version = api_version + + def prepare_params( + self, path: str, query_params: Mapping[str, str] | None + ) -> Mapping[str, str]: + """Prepare query parameters for Azure OpenAI, adding API version. + + Args: + path: Request path + query_params: Original query parameters from the client + + Returns: + Query parameters dict with Azure API version added for chat completions + """ + params = dict(query_params or {}) + if path.endswith("chat/completions"): + params["api-version"] = self.api_version + return params + + +class OpenRouterUpstreamProvider(UpstreamProvider): + """Upstream provider specifically configured for OpenRouter API.""" + + def __init__(self, api_key: str): + """Initialize OpenRouter provider with API key. + + Args: + api_key: OpenRouter API key for authentication + """ + super().__init__(base_url="https://openrouter.ai/api/v1", api_key=api_key) + + async def fetch_models(self) -> dict: + """Fetch available models from OpenRouter API. + + Returns: + Raw JSON response containing model data + """ + async with httpx.AsyncClient() as client: + response = await client.get( + "https://openrouter.ai/api/v1/models", + headers={"Authorization": f"Bearer {self.api_key}"}, + ) + return response.json() + + async def models(self) -> list[Model]: + """Get list of available models from OpenRouter. + + Returns: + List of Model objects representing available models + """ + response_data = await self.fetch_models() + models_list: list[Model] = [] + for model_data in response_data.get("data", []): + try: + models_list.append(Model(**model_data)) # type: ignore + except Exception: + continue + return models_list