#!/usr/bin/env python3 # -*- coding: utf-8 -*- """Proxy request controller for OpenAI and Claude compatible routes.""" from __future__ import annotations from typing import Any, Dict, Iterable, Optional from flask import Response, jsonify, request from ..application.app_context import AppContext from ..config import ConfigManager, ProviderManager from ..hooks import HookAbortError from ..services import LogService, ProxyService, UserService from ..services.proxy_service import ProxyErrorInfo from ..utils import normalize_ip from ..utils.local_time import now_local_datetime class ProxyController: """Expose downstream OpenAI-compatible proxy routes.""" def __init__( self, ctx: AppContext, proxy_service: ProxyService, user_service: UserService, log_service: LogService, provider_manager: ProviderManager, ): self._ctx = ctx self._app = ctx.flask_app self._logger = ctx.logger self._config_manager: ConfigManager = ctx.config_manager self._proxy_service = proxy_service self._user_service = user_service self._log_service = log_service self._provider_manager = provider_manager self._register_routes() def _register_routes(self) -> None: self._app.route("/v1/chat/completions", methods=["POST"])(self.chat_completions) self._app.route("/v1/responses", methods=["POST"])(self.responses) self._app.route("/v1/messages", methods=["POST"])(self.messages) self._app.route("/v1/models", methods=["GET"])(self.list_models) def _get_user_by_ip(self, ip_address: str) -> Optional[Dict[str, Any]]: return self._user_service.get_user_by_ip(ip_address, require_whitelist_access=True) def _is_whitelist_required(self) -> bool: return self._config_manager.is_chat_whitelist_enabled() def _get_authorized_user_for_request( self, client_ip: str, *, error_format: str, ) -> tuple[Optional[Dict[str, Any]], Optional[tuple[Response, int]]]: """在开启白名单时解析当前请求对应用户。""" if not self._is_whitelist_required(): return None, None user = self._get_user_by_ip(client_ip) if user: return user, None self._logger.warning("Proxy denied: ip=%s is not in whitelist", client_ip) return None, self._error_response( f"IP address {client_ip} is not in whitelist", 403, error_type="permission_error", code="ip_not_whitelisted", error_format=error_format, ) @staticmethod def _build_error_payload( message: str, *, error_type: str, status_code: int = 400, code: Optional[str] = None, error_format: str = "openai_chat", ) -> Dict[str, Any]: normalized_format = str(error_format or "").strip().lower() if normalized_format == "claude_chat": return { "type": "error", "error": { "type": error_type, "message": message, }, } return { "error": { "message": message, "type": error_type, "param": None, "code": code, } } def _error_response( self, message: str, status_code: int, *, error_type: str, code: Optional[str] = None, error_format: str = "openai_chat", ) -> tuple[Response, int]: return ( jsonify( self._build_error_payload( message, error_type=error_type, status_code=status_code, code=code, error_format=error_format, ) ), status_code, ) def chat_completions(self) -> Response: return self._proxy_completion_request( route_name="chat_completions", expected_target_formats=("openai_chat",), inspect_stream_usage=True, ) def responses(self) -> Response: return self._proxy_completion_request( route_name="responses", expected_target_formats=("openai_responses", "codex"), inspect_stream_usage=False, ) def messages(self) -> Response: return self._proxy_completion_request( route_name="messages", expected_target_formats=("claude_chat",), inspect_stream_usage=False, ) def _proxy_completion_request( self, *, route_name: str, expected_target_formats: Iterable[str], inspect_stream_usage: bool, error_format: Optional[str] = None, ) -> Response: normalized_expected_target_formats = tuple( str(item or "").strip().lower() for item in expected_target_formats if str(item or "").strip() ) if not normalized_expected_target_formats: raise ValueError("expected_target_formats must not be empty") resolved_error_format = error_format or normalized_expected_target_formats[0] try: client_ip = normalize_ip(request.remote_addr) self._logger.info("Proxy request received: route=%s ip=%s", route_name, client_ip) user, denial_response = self._get_authorized_user_for_request( client_ip, error_format=resolved_error_format, ) if denial_response is not None: return denial_response request_data = request.get_json(silent=True) if request_data is None: request_data = {} if not isinstance(request_data, dict): self._logger.warning("Proxy rejected: request body is not a JSON object route=%s", route_name) return self._error_response( "Request body must be a JSON object", 400, error_type="invalid_request_error", code="invalid_request_body", error_format=resolved_error_format, ) request_data = dict(request_data) if "model" not in request_data: self._logger.warning("Proxy rejected: missing model in request body route=%s", route_name) return self._error_response( "Missing 'model' in request body", 400, error_type="invalid_request_error", code="missing_model", error_format=resolved_error_format, ) model_name = request_data["model"] provider = self._provider_manager.get_provider_for_model(model_name) if not provider: self._logger.warning("Proxy rejected: unknown model=%r route=%s", model_name, route_name) return self._error_response( f"Unknown model: {model_name}", 400, error_type="invalid_request_error", code="unknown_model", error_format=resolved_error_format, ) if self._is_whitelist_required() and not self._user_service.can_user_access_model( user, model_name, available_models=self._provider_manager.list_model_names(), ): self._logger.warning( "Proxy denied: ip=%s is not allowed to access model=%s route=%s", client_ip, model_name, route_name, ) return self._error_response( f"IP address {client_ip} is not allowed to access model {model_name}", 403, error_type="permission_error", code="model_not_allowed", error_format=resolved_error_format, ) provider_target_format = str(getattr(provider, "target_format", "") or "").strip().lower() if provider_target_format not in normalized_expected_target_formats: self._logger.warning( "Proxy rejected: model=%s configured for target_format=%s route=%s", model_name, provider_target_format or "", route_name, ) if len(normalized_expected_target_formats) == 1: expected_hint = normalized_expected_target_formats[0] mismatch_message = ( f"Model {model_name} is configured for downstream format " f"{provider_target_format or ''}, not {expected_hint}" ) else: expected_hint = ", ".join(normalized_expected_target_formats) mismatch_message = ( f"Model {model_name} is configured for downstream format " f"{provider_target_format or ''}, not one of {expected_hint}" ) return self._error_response( mismatch_message, 400, error_type="invalid_request_error", code="target_format_mismatch", error_format=resolved_error_format, ) client_requested_usage_chunk = False if inspect_stream_usage: stream_options = request_data.get("stream_options") client_requested_usage_chunk = isinstance(stream_options, dict) and ( stream_options.get("include_usage") is True ) headers = self._filter_request_headers(request.headers) start_time = now_local_datetime() def on_proxy_complete(response_meta: Dict[str, Any]) -> None: self._logger.info( "Proxy completed: route=%s model=%s response_model=%s total_tokens=%s ip=%s", route_name, model_name, response_meta.get("response_model"), response_meta.get("total_tokens", 0), client_ip, ) self._log_service.log_request( request_model=model_name, response_model=response_meta.get("response_model"), total_tokens=response_meta.get("total_tokens", 0), prompt_tokens=response_meta.get("prompt_tokens", 0), completion_tokens=response_meta.get("completion_tokens", 0), start_time=start_time, end_time=now_local_datetime(), ip_address=client_ip, ) result, status_code, failure_info = self._proxy_service.proxy_request( provider, request_data, headers, on_complete=on_proxy_complete, forward_stream_usage=client_requested_usage_chunk, ) if result is None: failure_info = failure_info or ProxyErrorInfo( message="Upstream request failed after retries", status_code=status_code, error_type="upstream_error", error_code="upstream_request_failed", ) self._logger.error( "Proxy failed after retries: route=%s model=%s ip=%s status=%s upstream_error=%s", route_name, model_name, client_ip, status_code, failure_info.message, ) return self._error_response( failure_info.message, status_code, error_type=failure_info.error_type, code=failure_info.error_code, error_format=resolved_error_format, ) return result except HookAbortError as exc: self._logger.warning( "Proxy blocked by hook: route=%s ip=%s status=%s type=%s message=%s", route_name, normalize_ip(request.remote_addr), exc.status_code, exc.error_type, exc.message, ) return self._error_response( exc.message, exc.status_code, error_type=exc.error_type, code=exc.error_type, error_format=resolved_error_format, ) except Exception as exc: self._logger.error("Error in %s: %s", route_name, exc) return self._error_response( str(exc), 500, error_type="server_error", code="internal_error", error_format=resolved_error_format, ) def list_models(self) -> Response: try: client_ip = normalize_ip(request.remote_addr) user, denial_response = self._get_authorized_user_for_request( client_ip, error_format="openai_chat", ) if denial_response is not None: return denial_response model_names = list(self._provider_manager.list_model_names()) if self._is_whitelist_required(): allowed_models = set( self._user_service.get_accessible_models_for_user( user, available_models=model_names, ) ) model_names = [model_name for model_name in model_names if model_name in allowed_models] data = [] for model_key in model_names: provider_name, _, _ = str(model_key).partition("/") provider_view = self._provider_manager.get_provider_view(provider_name) source_format = getattr(provider_view, "source_format", None) target_format = getattr(provider_view, "target_format", None) transport = getattr(provider_view, "transport", None) data.append( { "id": model_key, "object": "model", "owned_by": provider_name or "proxy", "provider_name": provider_name or "proxy", "source_format": source_format, "target_format": target_format, "transport": transport, } ) return jsonify({"object": "list", "data": data}) except Exception as exc: self._logger.error("Error listing models: %s", exc) return self._error_response( str(exc), 500, error_type="server_error", code="internal_error", ) @staticmethod def _filter_request_headers(headers: Any) -> Dict[str, str]: excluded = { "host", "content-length", "connection", "keep-alive", "proxy-authenticate", "proxy-authorization", "te", "trailers", "transfer-encoding", "upgrade", } return {k: v for k, v in headers.items() if k.lower() not in excluded}