vertext-ai-proxy 0.2.0__py3-none-any.whl

vertext_ai_proxy/__init__.py

@@ -0,0 +1,3 @@
+from vertext_ai_proxy.cli import main
+
+__all__ = ["main"]

vertext_ai_proxy/app.py

@@ -0,0 +1,143 @@
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+import logging
+from fastapi import FastAPI, Request
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse
+import httpx
+import uvicorn
+
+from vertext_ai_proxy.core.auth import GCPTokenManager
+from vertext_ai_proxy.core.config import AppConfig, load_config, validate_config
+from vertext_ai_proxy.core.exceptions import (
+    OpenAIAPIException,
+    openai_exception_handler,
+)
+from vertext_ai_proxy.core.security import APIKeyValidator
+from vertext_ai_proxy.proxy import VertexAIProxy
+from vertext_ai_proxy.routers.health import router as health_router
+from vertext_ai_proxy.routers.v1.chat import router as chat_router
+from vertext_ai_proxy.routers.v1.embeddings import router as embeddings_router
+from vertext_ai_proxy.routers.v1.models import router as models_router
+
+logger = logging.getLogger(__name__)
+
+
+def create_app(
+    config: AppConfig | None = None,
+    token_manager: GCPTokenManager | None = None,
+    http_client: httpx.AsyncClient | None = None,
+) -> FastAPI:
+    """FastAPI 应用工厂函数，支持依赖注入以便利单元测试。"""
+    active_config = config or load_config()
+
+    # 启动前强制执行配置安全校验（阻断示例占位 Key 和空 Key）
+    validate_config(active_config)
+
+    # 初始化 Token 管理器
+    active_token_manager = token_manager or GCPTokenManager(active_config.gcp)
+
+    created_client = False
+    effective_http_client = http_client
+    if effective_http_client is None:
+        proxy_url = active_config.proxy.http_proxy
+        effective_http_client = httpx.AsyncClient(
+            proxy=proxy_url if proxy_url else None,
+            follow_redirects=True,
+        )
+        created_client = True
+
+    proxy_instance = VertexAIProxy(
+        config=active_config,
+        token_manager=active_token_manager,
+        http_client=effective_http_client,
+    )
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
+        logger.info(
+            f"Vertex AI OpenAI Proxy started on http://{active_config.server.host}:{active_config.server.port} "
+            f"(GCP Project: {active_token_manager.project_id}, Location: {active_token_manager.location})"
+        )
+        try:
+            yield
+        finally:
+            if created_client and effective_http_client:
+                await effective_http_client.aclose()
+                logger.info("HTTP client closed.")
+
+    app = FastAPI(
+        title="Vertex AI OpenAI Proxy",
+        description="Lightweight reverse proxy translating Google Cloud Vertex AI OpenAPI to OpenAI format",
+        version="0.2.0",
+        lifespan=lifespan,
+    )
+
+    # 注册依赖服务与组件至 app.state
+    app.state.config = active_config
+    app.state.proxy = proxy_instance
+    app.state.client = proxy_instance.client
+    app.state.chat_service = proxy_instance.chat_service
+    app.state.embedding_service = proxy_instance.embedding_service
+    app.state.model_service = proxy_instance.model_service
+    app.state.api_key_validator = APIKeyValidator(active_config)
+
+    # 注册符合 OpenAI 标准的异常拦截器
+    app.add_exception_handler(OpenAIAPIException, openai_exception_handler)
+
+    @app.exception_handler(RequestValidationError)
+    async def validation_exception_handler(
+        request: Request, exc: RequestValidationError
+    ) -> JSONResponse:
+        errors = exc.errors()
+        err_msg = "; ".join([f"{err.get('loc', [])}: {err.get('msg')}" for err in errors])
+        return JSONResponse(
+            status_code=400,
+            content={
+                "error": {
+                    "message": err_msg,
+                    "type": "invalid_request_error",
+                    "param": str(errors[0].get("loc", [None])[-1]) if errors else None,
+                    "code": "invalid_parameter",
+                }
+            },
+        )
+
+    @app.exception_handler(Exception)
+    async def generic_exception_handler(request: Request, exc: Exception) -> JSONResponse:
+        logger.exception(f"Unhandled error processing {request.method} {request.url}: {exc}")
+        return JSONResponse(
+            status_code=500,
+            content={
+                "error": {
+                    "message": f"Internal proxy error: {str(exc)}",
+                    "type": "api_error",
+                    "param": None,
+                    "code": "internal_error",
+                }
+            },
+        )
+
+    # 挂载模块化路由
+    app.include_router(health_router)
+    app.include_router(models_router)
+    app.include_router(chat_router)
+    app.include_router(embeddings_router)
+
+    return app
+
+
+def main() -> None:
+    """服务主入口命令行。"""
+    config = load_config()
+    app = create_app(config=config)
+    uvicorn.run(
+        app,
+        host=config.server.host,
+        port=config.server.port,
+        log_level=config.server.log_level.lower(),
+    )
+
+
+if __name__ == "__main__":
+    main()

vertext_ai_proxy/auth.py

@@ -0,0 +1,10 @@
+# Backward-compatibility re-exports
+from vertext_ai_proxy.core.auth import (  # noqa: F401
+    CLOUD_PLATFORM_SCOPE,
+    GCPTokenManager,
+)
+
+__all__ = [
+    "CLOUD_PLATFORM_SCOPE",
+    "GCPTokenManager",
+]

vertext_ai_proxy/cli.py

@@ -0,0 +1,158 @@
+import argparse
+import logging
+from pathlib import Path
+import sys
+
+from vertext_ai_proxy.app import create_app
+from vertext_ai_proxy.core.config import load_config, validate_config
+from vertext_ai_proxy.core.templates import DEFAULT_CONFIG_TEMPLATE
+
+__version__ = "0.2.0"
+logger = logging.getLogger("vertext_ai_proxy.cli")
+
+
+def handle_init(args: argparse.Namespace) -> int:
+    """处理 'vap init' 命令：在当前目录创建 config.local.yaml。"""
+    target = Path(args.output)
+    if target.exists() and not args.force:
+        print(
+            f"[WARN] '{target.name}' already exists at: {target.resolve()}\n"
+            f"       To overwrite the existing file, run: vap init --force",
+            file=sys.stderr,
+        )
+        return 1
+
+    try:
+        target.write_text(DEFAULT_CONFIG_TEMPLATE, encoding="utf-8")
+        print(
+            f"[SUCCESS] Initialized configuration file at:\n"
+            f"          {target.resolve()}\n\n"
+            f"Next steps:\n"
+            f"  1. Edit '{target.name}':\n"
+            f"     - proxy.api_keys: Replace placeholder with your secret key (e.g. sk-my-secret-key)\n"
+            f"     - gcp.service_account_file: Path to your Google Cloud Service Account JSON file\n"
+            f"  2. Start the proxy server:\n"
+            f"     vap run\n"
+        )
+        return 0
+    except Exception as e:
+        print(f"[ERROR] Failed to write '{target}': {e}", file=sys.stderr)
+        return 1
+
+
+def handle_run(args: argparse.Namespace) -> int:
+    """处理 'vap run' 命令：加载配置并启动 Uvicorn 反代服务。"""
+    import uvicorn
+
+    try:
+        config = load_config(args.config)
+        # 命令行参数覆盖
+        if args.host:
+            config.server.host = args.host
+        if args.port:
+            config.server.port = args.port
+        if args.log_level:
+            config.server.log_level = args.log_level
+
+        # 启动安全校验
+        validate_config(config)
+    except Exception as e:
+        print(f"\n{e}\n", file=sys.stderr)
+        return 1
+
+    app = create_app(config=config)
+    uvicorn.run(
+        app,
+        host=config.server.host,
+        port=config.server.port,
+        log_level=config.server.log_level.lower(),
+    )
+    return 0
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """构建 CLI 参数解析器。"""
+    parser = argparse.ArgumentParser(
+        prog="vap",
+        description="Vertex AI OpenAI Proxy (vap) - Lightweight OpenAI-compatible gateway for Google Vertex AI.",
+    )
+    parser.add_argument(
+        "-v",
+        "--version",
+        action="version",
+        version=f"%(prog)s {__version__}",
+        help="Show program version and exit.",
+    )
+
+    subparsers = parser.add_subparsers(dest="subcommand", title="Commands")
+
+    # vap init
+    init_parser = subparsers.add_parser(
+        "init",
+        help="Initialize a new config.local.yaml in the current working directory.",
+        description="Scaffold config.local.yaml from default template in current directory.",
+    )
+    init_parser.add_argument(
+        "-f",
+        "--force",
+        action="store_true",
+        help="Overwrite existing config.local.yaml if it already exists.",
+    )
+    init_parser.add_argument(
+        "-o",
+        "--output",
+        default="config.local.yaml",
+        help="Target output filename (default: config.local.yaml).",
+    )
+    init_parser.set_defaults(func=handle_init)
+
+    # vap run
+    run_parser = subparsers.add_parser(
+        "run",
+        help="Start the reverse proxy server.",
+        description="Run Vertex AI OpenAI Proxy server with Uvicorn.",
+    )
+    run_parser.add_argument(
+        "-c",
+        "--config",
+        default=None,
+        help="Path to YAML configuration file (default: searches config.local.yaml then config.yaml).",
+    )
+    run_parser.add_argument(
+        "-H",
+        "--host",
+        default=None,
+        help="Bind host (default: from config, e.g. 127.0.0.1).",
+    )
+    run_parser.add_argument(
+        "-p",
+        "--port",
+        type=int,
+        default=None,
+        help="Bind port (default: from config, e.g. 8766).",
+    )
+    run_parser.add_argument(
+        "--log-level",
+        default=None,
+        choices=["debug", "info", "warning", "error"],
+        help="Uvicorn log level (default: from config, e.g. info).",
+    )
+    run_parser.set_defaults(func=handle_run)
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    """CLI 主入口。"""
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    if not args.subcommand:
+        parser.print_help()
+        return 0
+
+    return args.func(args)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

vertext_ai_proxy/config.py

@@ -0,0 +1,22 @@
+# Backward-compatibility re-exports
+from vertext_ai_proxy.core.config import (  # noqa: F401
+    AppConfig,
+    DEFAULT_MODELS,
+    FORBIDDEN_PLACEHOLDER_KEYS,
+    GCPConfig,
+    ProxyConfig,
+    ServerConfig,
+    load_config,
+    validate_config,
+)
+
+__all__ = [
+    "AppConfig",
+    "DEFAULT_MODELS",
+    "FORBIDDEN_PLACEHOLDER_KEYS",
+    "GCPConfig",
+    "ProxyConfig",
+    "ServerConfig",
+    "load_config",
+    "validate_config",
+]

vertext_ai_proxy/core/__init__.py

@@ -0,0 +1 @@
+"""Core infrastructure components: config, auth, security, and exceptions."""

vertext_ai_proxy/core/auth.py

@@ -0,0 +1,122 @@
+import asyncio
+from datetime import datetime, timezone
+import json
+import logging
+from pathlib import Path
+
+from google.auth.transport.requests import Request
+from google.oauth2 import service_account
+from vertext_ai_proxy.core.config import GCPConfig
+
+logger = logging.getLogger(__name__)
+
+CLOUD_PLATFORM_SCOPE = "https://www.googleapis.com/auth/cloud-platform"
+
+
+class GCPTokenManager:
+    """管理 Google Cloud Service Account 凭据与 OAuth 2.0 Access Token 的生命周期。
+
+    - 支持原生 GCP 服务账号 JSON，以及带有外层封装（如 service_account 嵌套结构）的 JSON
+    - 自动提取 project_id（若未手动指定）
+    - 内存中缓存 Access Token，并在过期前通过线程池异步安全刷新
+    - 使用 asyncio.Lock 避免高并发下的重复刷新击穿
+    """
+
+    def __init__(
+        self,
+        config: GCPConfig,
+        credentials: service_account.Credentials | None = None,
+    ) -> None:
+        self.config = config
+        self._lock = asyncio.Lock()
+        self._request = Request()
+
+        if credentials is not None:
+            self.credentials = credentials
+        else:
+            sa_path = Path(config.service_account_file)
+            if not sa_path.exists():
+                raise FileNotFoundError(
+                    f"Google Cloud Service Account file not found at: {sa_path.resolve()}"
+                )
+            with open(sa_path, "r", encoding="utf-8") as f:
+                raw_data = json.load(f)
+
+            # 兼容两种常见格式：
+            # 1. 纯 Google 官方下载的 Service Account JSON
+            # 2. 外部工具封装格式（如 cli-proxy-api 格式，嵌套在 'service_account' 键下）
+            if (
+                isinstance(raw_data, dict)
+                and "service_account" in raw_data
+                and isinstance(raw_data["service_account"], dict)
+            ):
+                logger.info("Detected wrapped service_account format, extracting inner credentials.")
+                sa_info = raw_data["service_account"]
+                # 若外层显式声明了 project_id 或 location，智能回退继承
+                if not self.config.project_id and raw_data.get("project_id"):
+                    self.config.project_id = str(raw_data["project_id"])
+                if (not self.config.location or self.config.location == "global") and raw_data.get("location"):
+                    self.config.location = str(raw_data["location"])
+            else:
+                sa_info = raw_data
+
+            self.credentials = service_account.Credentials.from_service_account_info(
+                sa_info,
+                scopes=[CLOUD_PLATFORM_SCOPE],
+            )
+
+        # 优先使用配置中显式指定的 project_id，否则从 Service Account JSON 中提取
+        self.project_id: str = (
+            self.config.project_id.strip()
+            if self.config.project_id and self.config.project_id.strip()
+            else (getattr(self.credentials, "project_id", "") or "")
+        )
+
+        if not self.project_id:
+            logger.warning(
+                "No project_id could be determined from GCPConfig or credentials."
+            )
+
+    @property
+    def location(self) -> str:
+        return self.config.location
+
+    @property
+    def base_upstream_url(self) -> str:
+        """根据 location 和 project_id 构造 Vertex AI OpenAPI 的目标 Base URL。"""
+        # Google 统一端点路由
+        return (
+            f"https://aiplatform.googleapis.com/v1/projects/{self.project_id}"
+            f"/locations/{self.location}/endpoints/openapi"
+        )
+
+    def _is_token_expiring(self, buffer_seconds: int = 300) -> bool:
+        """检查 Token 是否已失效或即将失效（默认预留 5 分钟缓冲）。"""
+        if not self.credentials.token or not self.credentials.expiry:
+            return True
+
+        now = datetime.now(timezone.utc)
+        expiry = self.credentials.expiry
+        if expiry.tzinfo is None:
+            expiry = expiry.replace(tzinfo=timezone.utc)
+
+        return (expiry - now).total_seconds() < buffer_seconds
+
+    def _refresh_sync(self) -> None:
+        """同步刷新 Google OAuth Token（在线程池中调用）。"""
+        self.credentials.refresh(self._request)
+
+    async def get_access_token(self) -> str:
+        """异步获取当前有效的 Access Token，必要时执行刷新。"""
+        if not self._is_token_expiring():
+            return self.credentials.token
+
+        async with self._lock:
+            # 双重检查锁（DCL）
+            if not self._is_token_expiring():
+                return self.credentials.token
+
+            logger.info("Google Cloud OAuth token expiring or absent; refreshing...")
+            await asyncio.to_thread(self._refresh_sync)
+            logger.info("Google Cloud OAuth token successfully refreshed.")
+            return self.credentials.token

vertext_ai_proxy/core/config.py

@@ -0,0 +1,104 @@
+import logging
+from pathlib import Path
+import yaml
+from pydantic import BaseModel, Field
+
+logger = logging.getLogger(__name__)
+
+# 已知的示例或占位符 Key 集合
+FORBIDDEN_PLACEHOLDER_KEYS: set[str] = {
+    "sk-change-this-to-your-own-secret-key",
+    "sk-vertex-local-test",
+    "sk-change-me",
+    "sk-example-key",
+    "sk-my-vertex-proxy-key",
+    "sk-placeholder",
+}
+
+DEFAULT_MODELS: list[str] = [
+    "gemini-2.5-flash",
+    "gemini-2.5-pro",
+    "gemini-2.0-flash",
+    "gemini-2.0-flash-lite",
+    "gemini-1.5-pro",
+    "gemini-1.5-flash",
+]
+
+
+class ServerConfig(BaseModel):
+    host: str = "127.0.0.1"
+    port: int = 8766
+    log_level: str = "info"
+
+
+class ProxyConfig(BaseModel):
+    api_keys: list[str] = Field(default_factory=list)
+    auto_prefix_model: bool = True
+    http_proxy: str | None = None
+    # 可在 yaml 配置文件中自由增减模型，供 /v1/models 接口返回给客户端
+    models: list[str] = Field(default_factory=lambda: list(DEFAULT_MODELS))
+
+
+class GCPConfig(BaseModel):
+    service_account_file: str = "service_account.json"
+    project_id: str = ""
+    location: str = "global"
+    timeout_seconds: int = 120
+
+
+class AppConfig(BaseModel):
+    server: ServerConfig = Field(default_factory=ServerConfig)
+    proxy: ProxyConfig = Field(default_factory=ProxyConfig)
+    gcp: GCPConfig = Field(default_factory=GCPConfig)
+
+
+def validate_config(config: AppConfig) -> None:
+    """启动前强制进行安全性校验：
+
+    1. 校验 api_keys 是否非空
+    2. 校验是否仍在使用示例/占位 Key
+    """
+    if not config.proxy.api_keys:
+        raise ValueError(
+            "[FATAL SECURITY ERROR] No API key configured in 'proxy.api_keys'. "
+            "Please configure at least one custom secret key in config.local.yaml."
+        )
+
+    for raw_key in config.proxy.api_keys:
+        key = raw_key.strip()
+        if not key:
+            continue
+        if key in FORBIDDEN_PLACEHOLDER_KEYS or key.startswith("sk-change-this"):
+            raise ValueError(
+                f"[FATAL SECURITY ERROR] Detected placeholder API Key: '{key}'. "
+                "The server refuses to start with default/example keys. "
+                "Please copy config.example.yaml to config.local.yaml and set your own secure API key."
+            )
+
+
+def load_config(config_path: str | Path | None = None) -> AppConfig:
+    """加载配置文件：
+
+    优先级：指定路径 -> config.local.yaml -> config.yaml -> 默认配置
+    """
+    target_path: Path | None = None
+
+    if config_path:
+        target_path = Path(config_path)
+    else:
+        local_config = Path("config.local.yaml")
+        standard_config = Path("config.yaml")
+
+        if local_config.exists():
+            target_path = local_config
+        elif standard_config.exists():
+            target_path = standard_config
+
+    if target_path and target_path.exists():
+        logger.info(f"Loading configuration from {target_path}")
+        with open(target_path, "r", encoding="utf-8") as f:
+            data = yaml.safe_load(f) or {}
+        return AppConfig(**data)
+
+    logger.info("No configuration file found; initializing with default configuration.")
+    return AppConfig()

vertext_ai_proxy/core/exceptions.py

@@ -0,0 +1,40 @@
+from fastapi import HTTPException, Request, status
+from fastapi.responses import JSONResponse
+
+
+class OpenAIAPIException(HTTPException):
+    """标准的 OpenAI 格式 API 异常。"""
+
+    def __init__(
+        self,
+        status_code: int,
+        message: str,
+        code: str = "invalid_api_key",
+        param: str | None = None,
+        err_type: str = "invalid_request_error",
+    ) -> None:
+        super().__init__(status_code=status_code, detail=message)
+        self.status_code = status_code
+        self.error_dict = {
+            "error": {
+                "message": message,
+                "type": err_type,
+                "param": param,
+                "code": code,
+            }
+        }
+
+
+async def openai_exception_handler(
+    request: Request, exc: OpenAIAPIException
+) -> JSONResponse:
+    """FastAPI 异常处理器：确保直接返回根级 {"error": {...}} 结构而非嵌套在 {"detail": ...} 中。"""
+    return JSONResponse(status_code=exc.status_code, content=exc.error_dict)
+
+
+def raise_unauthorized(message: str = "Incorrect API key provided.") -> None:
+    raise OpenAIAPIException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        message=message,
+        code="invalid_api_key",
+    )

vertext_ai_proxy/core/security.py

@@ -0,0 +1,49 @@
+import secrets
+from fastapi import Header
+from vertext_ai_proxy.core.config import AppConfig
+from vertext_ai_proxy.core.exceptions import OpenAIAPIException, raise_unauthorized
+
+
+class APIKeyValidator:
+    """验证客户端请求中的 API Key 是否在允许列表中。"""
+
+    def __init__(self, config: AppConfig) -> None:
+        self.config = config
+
+    def __call__(
+        self,
+        authorization: str | None = Header(None),
+        api_key_header: str | None = Header(None, alias="api-key"),
+    ) -> str:
+        token: str | None = None
+
+        if authorization:
+            parts = authorization.strip().split(" ", 1)
+            if len(parts) == 2 and parts[0].lower() == "bearer":
+                token = parts[1].strip()
+            else:
+                token = authorization.strip()
+
+        if not token and api_key_header:
+            token = api_key_header.strip()
+
+        if not token:
+            raise_unauthorized(
+                "Missing API key. Provide authorization header: 'Authorization: Bearer <key>' or 'api-key: <key>'."
+            )
+
+        valid_keys = [k.strip() for k in self.config.proxy.api_keys if k.strip()]
+        if not valid_keys:
+            raise_unauthorized("No authorized API keys configured on proxy server.")
+
+        # 使用 secrets.compare_digest 防止时序攻击
+        matched = False
+        for valid_key in valid_keys:
+            if secrets.compare_digest(token, valid_key):
+                matched = True
+                break
+
+        if not matched:
+            raise_unauthorized("Invalid API key provided.")
+
+        return token