#!/usr/bin/env python3 # -*- coding: utf-8 -*- """代理控制器。""" from typing import Any, Dict, Optional from datetime import datetime from flask import Response, jsonify, request from ..application.app_context import AppContext from ..config import ConfigManager, ProviderManager from ..services import LogService, ProxyService, UserService from ..utils import normalize_ip class ProxyController: """处理模型代理相关路由。""" 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/models", methods=["GET"])(self.list_models) def _get_user_by_ip(self, ip_address: str) -> Optional[Dict[str, Any]]: """按 IP 查询白名单用户。""" 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 chat_completions(self) -> Response: """代理上游 chat/completions 请求。""" try: client_ip = normalize_ip(request.remote_addr) self._logger.info(f"Proxy request received: ip={client_ip}") whitelist_required = self._is_whitelist_required() user = self._get_user_by_ip(client_ip) if whitelist_required else None if whitelist_required and not user: self._logger.warning(f"Proxy denied: ip={client_ip} is not in whitelist") return ( jsonify( { "error": { "message": f"IP address {client_ip} is not in whitelist", "type": "access_denied", } } ), 403, ) request_data = request.get_json(silent=True) if not request_data or "model" not in request_data: self._logger.warning("Proxy rejected: missing model in request body") return ( jsonify( { "error": { "message": "Missing 'model' in request body", "type": "invalid_request_error", } } ), 400, ) stream_options = request_data.get("stream_options") if not isinstance(stream_options, dict): stream_options = {} if stream_options.get("include_usage") is not True: stream_options["include_usage"] = True request_data["stream_options"] = stream_options model_name = request_data["model"] provider = self._provider_manager.find_provider_by_model(model_name) if not provider: self._logger.warning(f"Proxy rejected: unknown model={model_name!r}") return ( jsonify( { "error": { "message": f"Unknown model: {model_name}", "type": "invalid_request_error", } } ), 400, ) headers = self._filter_request_headers(request.headers) start_time = datetime.now() def on_proxy_complete(response_meta: Dict[str, Any]) -> None: self._logger.info( "Proxy completed: model=%s response_model=%s total_tokens=%s ip=%s", 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=datetime.now(), ip_address=client_ip, ) result, status_code = self._proxy_service.proxy_request( provider, request_data, headers, on_complete=on_proxy_complete ) if result is None: self._logger.error(f"Proxy failed after retries: model={model_name!r}, ip={client_ip}") return ( jsonify( { "error": { "message": "Request failed after retries", "type": "proxy_error", } } ), status_code, ) return result except Exception as exc: self._logger.error(f"Error in chat_completions: {exc}") return ( jsonify({"error": {"message": str(exc), "type": "internal_error"}}), 500, ) def list_models(self) -> Response: """返回当前可用模型列表。""" try: data = [{"id": model_key} for model_key in self._provider_manager.get_all_models()] return jsonify({"object": "list", "data": data}) except Exception as exc: self._logger.error(f"Error listing models: {exc}") return jsonify({"error": str(exc)}), 500 @staticmethod def _filter_request_headers(headers: Any) -> Dict[str, str]: """过滤不应转发到上游的 hop-by-hop 请求头。""" 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}