From 5ecab37feb46035077c89aecb7ab017e5ff630ff Mon Sep 17 00:00:00 2001 From: Hemang Date: Wed, 26 Feb 2025 15:00:35 +0100 Subject: [PATCH] Support passing a policy file to the proxy on startup. --- docker-compose.local.yml | 5 ++- proxy/__init__.py | 0 proxy/common/__init__.py | 0 proxy/common/config_manager.py | 63 ++++++++++++++++++++++++++++++++++ proxy/requirements.txt | 1 + proxy/routes/anthropic.py | 2 ++ proxy/routes/open_ai.py | 3 +- proxy/run.sh | 12 +++++-- proxy/serve.py | 4 +-- proxy/validate_config.py | 13 +++++++ run.sh | 29 ++++++++++++++-- 11 files changed, 123 insertions(+), 9 deletions(-) create mode 100644 proxy/__init__.py create mode 100644 proxy/common/__init__.py create mode 100644 proxy/common/config_manager.py create mode 100644 proxy/validate_config.py diff --git a/docker-compose.local.yml b/docker-compose.local.yml index 3f0eb27..09cf7d4 100644 --- a/docker-compose.local.yml +++ b/docker-compose.local.yml @@ -9,10 +9,14 @@ services: - .env environment: - DEV_MODE=true + - POLICIES_FILE_PATH=${POLICIES_FILE_PATH:+/srv/resources/policies.py} volumes: - type: bind source: ./proxy target: /srv/proxy + - type: bind + source: ${POLICIES_FILE_PATH:-/dev/null} + target: /srv/resources/policies.py networks: - invariant-explorer-web ports: @@ -24,7 +28,6 @@ services: - "traefik.http.routers.invariant-proxy-api.entrypoints=invariant-explorer-web" - "traefik.http.services.invariant-proxy-api.loadbalancer.server.port=8000" - "traefik.docker.network=invariant-explorer-web" - networks: invariant-explorer-web: external: true \ No newline at end of file diff --git a/proxy/__init__.py b/proxy/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/proxy/common/__init__.py b/proxy/common/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/proxy/common/config_manager.py b/proxy/common/config_manager.py new file mode 100644 index 0000000..7190746 --- /dev/null +++ b/proxy/common/config_manager.py @@ -0,0 +1,63 @@ +"""Common Configurations for the Proxy Server.""" + +import os +import threading + +from invariant.analyzer import Policy + + +class ProxyConfig: + """Common configurations for the Proxy Server.""" + + def __init__(self): + print("hello in init of ProxyConfig") + self.policies = self._load_policies() + + def _load_policies(self) -> str: + """ + Loads and validates policies from the file specified in POLICIES_FILE_PATH. + Returns the policy file content as a string if valid; otherwise, raises an error. + """ + print("hello in load_policies") + policies_file = os.getenv("POLICIES_FILE_PATH", "") + + if not policies_file: + print("Warning: POLICIES_FILE_PATH is not set. Using empty policies.") + return "" + + try: + with open(policies_file, "r", encoding="utf-8") as f: + policy_file_content = f.read() + _ = Policy.from_string(policy_file_content) + return policy_file_content + + except (FileNotFoundError, PermissionError, OSError) as e: + raise ValueError( + f"Error: Unable to read policies file ({policies_file}): {e}" + ) from e + + except Exception as e: + raise ValueError(f"Invalid policy content in {policies_file}: {e}") from e + + def __repr__(self) -> str: + return f"ProxyConfig(policies={repr(self.policies)})" + + +class ProxyConfigManager: + """Manager for Proxy Configuration.""" + + _config_instance = None + _lock = threading.Lock() + + @classmethod + def get_config(cls): + """Initializes and returns the proxy configuration using double-checked locking.""" + local_config = cls._config_instance + + if local_config is None: + with cls._lock: + local_config = cls._config_instance + if local_config is None: + local_config = ProxyConfig() + cls._config_instance = local_config + return local_config diff --git a/proxy/requirements.txt b/proxy/requirements.txt index a564b13..2aabbbf 100644 --- a/proxy/requirements.txt +++ b/proxy/requirements.txt @@ -1,5 +1,6 @@ fastapi==0.115.7 httpx==0.28.1 +invariant-ai>=0.2.1 invariant-sdk>=0.0.10 starlette-compress==1.4.0 uvicorn==0.34.0 \ No newline at end of file diff --git a/proxy/routes/anthropic.py b/proxy/routes/anthropic.py index 182485c..c30461d 100644 --- a/proxy/routes/anthropic.py +++ b/proxy/routes/anthropic.py @@ -4,6 +4,7 @@ import json from typing import Any, Optional import httpx +from common.config_manager import ProxyConfig, ProxyConfigManager from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response from starlette.responses import StreamingResponse from utils.constants import CLIENT_TIMEOUT, IGNORED_HEADERS @@ -43,6 +44,7 @@ def validate_headers(x_api_key: str = Header(None)): async def anthropic_v1_messages_proxy( request: Request, dataset_name: str = None, + config: ProxyConfig = Depends(ProxyConfigManager.get_config), # pylint: disable=unused-argument ): """Proxy calls to the Anthropic APIs""" headers = { diff --git a/proxy/routes/open_ai.py b/proxy/routes/open_ai.py index 11569a1..5357045 100644 --- a/proxy/routes/open_ai.py +++ b/proxy/routes/open_ai.py @@ -4,6 +4,7 @@ import json from typing import Any, Optional import httpx +from common.config_manager import ProxyConfig, ProxyConfigManager from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response from fastapi.responses import StreamingResponse from utils.constants import CLIENT_TIMEOUT, IGNORED_HEADERS @@ -33,9 +34,9 @@ def validate_headers(authorization: str = Header(None)): async def openai_chat_completions_proxy( request: Request, dataset_name: str = None, + config: ProxyConfig = Depends(ProxyConfigManager.get_config), # pylint: disable=unused-argument ) -> Response: """Proxy calls to the OpenAI APIs""" - headers = { k: v for k, v in request.headers.items() if k.lower() not in IGNORED_HEADERS } diff --git a/proxy/run.sh b/proxy/run.sh index 583ef78..0663eea 100755 --- a/proxy/run.sh +++ b/proxy/run.sh @@ -1,8 +1,16 @@ #!/bin/bash -# if DEV_MODE is true, then run the app with auto-reload +# Validate configuration +python validate_config.py + +# Check exit code of validation script +if [ $? -ne 0 ]; then + echo "Configuration validation failed. Exiting." + exit 1 +fi + if [ "$DEV_MODE" = "true" ]; then uvicorn serve:app --host 0.0.0.0 --port 8000 --reload else uvicorn serve:app --host 0.0.0.0 --port 8000 -fi +fi \ No newline at end of file diff --git a/proxy/serve.py b/proxy/serve.py index b2d923e..29ae04c 100644 --- a/proxy/serve.py +++ b/proxy/serve.py @@ -19,7 +19,7 @@ router = fastapi.APIRouter(prefix="/api/v1") @router.get("/proxy/health", tags=["health_check"], include_in_schema=False) async def get_proxy_info(): """Health check""" - return {"message": "Hello v1/proxy"} + return {"message": "Hello from Invariant v1/proxy"} router.include_router(open_ai_proxy, prefix="/proxy", tags=["open_ai_proxy"]) @@ -28,6 +28,6 @@ router.include_router(anthropic_proxy, prefix="/proxy", tags=["anthropic_proxy"] app.include_router(router) -# Serve the API + if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=8000) diff --git a/proxy/validate_config.py b/proxy/validate_config.py new file mode 100644 index 0000000..2b11a54 --- /dev/null +++ b/proxy/validate_config.py @@ -0,0 +1,13 @@ +"""Validates the ProxyConfigManager configuration.""" + +import sys + +from common.config_manager import ProxyConfigManager + +try: + _ = ProxyConfigManager.get_config() + print("ProxyConfig validated successfully.") + sys.exit(0) +except Exception as e: + print(f"Error loading ProxyConfig error: {e}") + sys.exit(1) diff --git a/run.sh b/run.sh index 5313499..6fb27a7 100755 --- a/run.sh +++ b/run.sh @@ -3,11 +3,33 @@ up() { docker network inspect invariant-explorer-web >/dev/null 2>&1 || \ docker network create invariant-explorer-web - # Start your local docker-compose services - docker compose -f docker-compose.local.yml up -d + # Default values + POLICIES_FILE_PATH="" + + # Parse command-line arguments + while [[ "$#" -gt 0 ]]; do + case "$1" in + --policies-file=*) + POLICIES_FILE_PATH="${1#*=}" + ;; + *) + echo "Unknown parameter: $1" + exit 1 + ;; + esac + shift + done + + if [[ -n "$POLICIES_FILE_PATH" ]]; then + POLICIES_FILE_PATH=$(realpath "$POLICIES_FILE_PATH") + fi + + # Start Docker Compose with the correct environment variable + POLICIES_FILE_PATH="$POLICIES_FILE_PATH" docker compose -f docker-compose.local.yml up -d echo "Proxy started at http://localhost:8005/api/v1/proxy/" echo "See http://localhost:8005/api/v1/proxy/docs for API documentation" + echo "Using Policies File: ${POLICIES_FILE_PATH:-None}" } build() { @@ -78,7 +100,8 @@ tests() { # ----------------------------- case "$1" in "up") - up + shift + up $@ ;; "build") build