diff --git a/.env b/.env deleted file mode 100644 index ae81d93..0000000 --- a/.env +++ /dev/null @@ -1,6 +0,0 @@ -# This specifies the Invariant Explorer instance where the gateway will push the traces -# Set this to https://preview-explorer.invariantlabs.ai if you want to push to Preview. -# If you want to push to a local instance of explorer, then specify the app-api docker container name like: -# http://:8000 to push to the local explorer instance. -INVARIANT_API_URL=https://explorer.invariantlabs.ai -GUARDRAILS_API_URL=https://explorer.invariantlabs.ai \ No newline at end of file diff --git a/.github/workflows/publish-images.yml b/.github/workflows/publish-images.yml index b4c15dd..04e8fb9 100644 --- a/.github/workflows/publish-images.yml +++ b/.github/workflows/publish-images.yml @@ -48,7 +48,7 @@ jobs: uses: docker/build-push-action@v5 with: context: ./gateway - file: Dockerfile.gateway + file: ./gateway/Dockerfile.gateway platforms: linux/amd64 push: true tags: ${{ env.tags }} diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..2e13d6c --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,2 @@ +include gateway/docker-compose.local.yml +include gateway/Dockerfile.gateway diff --git a/gateway/.env b/gateway/.env index c84df8b..093ad91 100644 --- a/gateway/.env +++ b/gateway/.env @@ -1,4 +1,11 @@ POSTGRES_USER=postgres POSTGRES_PASSWORD=postgres POSTGRES_DB=invariantmonitor -POSTGRES_HOST=database \ No newline at end of file +POSTGRES_HOST=database + +# This specifies the Invariant Explorer instance where the gateway will push the traces +# Set this to https://preview-explorer.invariantlabs.ai if you want to push to Preview. +# If you want to push to a local instance of explorer, then specify the app-api docker container name like: +# http://:8000 to push to the local explorer instance. +INVARIANT_API_URL=https://explorer.invariantlabs.ai +GUARDRAILS_API_URL=https://explorer.invariantlabs.ai \ No newline at end of file diff --git a/Dockerfile.gateway b/gateway/Dockerfile.gateway similarity index 65% rename from Dockerfile.gateway rename to gateway/Dockerfile.gateway index 24e6cbd..0642c60 100644 --- a/Dockerfile.gateway +++ b/gateway/Dockerfile.gateway @@ -2,11 +2,11 @@ FROM python:3.12 WORKDIR /srv/gateway -COPY pyproject.toml ./ +COPY ../pyproject.toml ./pyproject.toml -COPY gateway/ ./ +COPY . . -COPY README.md /srv/gateway/README.md +COPY ../README.md ./README.md RUN pip install --no-cache-dir . diff --git a/gateway/__main__.py b/gateway/__main__.py index 942e3df..9bc9935 100644 --- a/gateway/__main__.py +++ b/gateway/__main__.py @@ -1,12 +1,20 @@ """Script is used to run actions using the Invariant Gateway.""" import asyncio +import os import signal +import subprocess import sys +import time + +from typing import Optional from gateway.mcp import mcp +LOCAL_COMPOSE_FILE = "gateway/docker-compose.local.yml" + + # Handle signals to ensure clean shutdown def signal_handler(sig, frame): """Handle signals for graceful shutdown.""" @@ -16,8 +24,14 @@ def signal_handler(sig, frame): def print_help(): """Prints the help message.""" actions = { - "mcp": "Runs the Invariant Gateway against MCP (Model Context Protocol) servers with guardrailing and push to Explorer features.", - "llm": "Runs the Invariant Gateway against LLM providers with guardrailing and push to Explorer features. Not implemented yet.", + "mcp": """ + Runs the Invariant Gateway against MCP (Model Context Protocol) stdio servers with guardrailing and push to Explorer features. + """, + "server": """ + Runs the Invariant Gateway server locally providing guardrailing and push to Explorer features. + Should be called with one of the following subcommands: build, up, down, logs. + A guardrails file can be passed with the flag: --guardrails-file=/path/to/guardrails/file. + """, "help": "Shows this help message.", } @@ -25,6 +39,181 @@ def print_help(): print(f"{verb}: {description}") +def ensure_network_exists(network_name: str = "invariant-explorer-web") -> bool: + """Ensure the Docker network exists.""" + try: + # Check if network exists + result = subprocess.run( + ["docker", "network", "inspect", network_name], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=False, + ) + + if result.returncode != 0: + print(f"Creating Docker network: {network_name}") + subprocess.run( + ["docker", "network", "create", network_name], + check=True, + ) + + return True + except subprocess.CalledProcessError as e: + print(f"Error creating Docker network: {e}") + return False + + +def setup_guardrails(guardrails_file_path: Optional[str] = None) -> bool: + """Configure guardrails if specified.""" + if not guardrails_file_path: + return True + + if not os.path.isfile(guardrails_file_path): + print( + f"Error: Specified guardrails file does not exist: {guardrails_file_path}" + ) + return False + + # Convert to absolute path + guardrails_file_path = os.path.realpath(guardrails_file_path) + os.environ["GUARDRAILS_FILE_PATH"] = guardrails_file_path + + # Check if INVARIANT_API_KEY is set + if not os.environ.get("INVARIANT_API_KEY"): + print( + "Error: A guardrails file is specified, but INVARIANT_API_KEY env var is not set. " + "This is required to validate guardrails." + ) + return False + + return True + + +def build(): + """Build Docker containers using docker-compose.""" + try: + print(f"Building using docker-compose file: {LOCAL_COMPOSE_FILE}") + subprocess.run( + ["docker", "compose", "-f", str(LOCAL_COMPOSE_FILE), "build"], + check=True, + ) + print("Build completed successfully") + return True + except subprocess.CalledProcessError as e: + print(f"Error building containers: {e}") + return False + + +def up(guardrails_file_path: Optional[str] = None): + """Set up the local server for the Invariant Gateway.""" + # Ensure network exists + if not ensure_network_exists(): + return 1 + + # Set up guardrails + if not setup_guardrails(guardrails_file_path=guardrails_file_path): + return 1 + + # Run the server + try: + # Start containers + print(f"Starting containers using docker-compose file: {LOCAL_COMPOSE_FILE}") + subprocess.run( + ["docker", "compose", "-f", str(LOCAL_COMPOSE_FILE), "up", "-d"], + check=True, + ) + + # Wait for containers to start + time.sleep(2) + + # Check if gateway container is running + result = subprocess.run( + ["docker", "ps", "-qf", "name=invariant-gateway"], + capture_output=True, + text=True, + check=True, + ) + + if not result.stdout.strip(): + print("The invariant-gateway container failed to start.") + logs = subprocess.run( + ["docker", "logs", "invariant-gateway"], + capture_output=True, + text=True, + check=False, + ) + print("Last 20 lines of logs:") + print("\n".join(logs.stdout.strip().split("\n")[-20:])) + return False + + print("Gateway started at http://localhost:8005/api/v1/gateway/") + print("See http://localhost:8005/api/v1/gateway/docs for API documentation") + + if guardrails_file_path: + print(f"Using Guardrails File: {guardrails_file_path}") + + return True + except subprocess.CalledProcessError as e: + print(f"Error starting containers: {e}") + return False + + +def down(): + """Stop the Docker containers.""" + try: + print(f"Stopping containers using docker-compose file: {LOCAL_COMPOSE_FILE}") + subprocess.run( + ["docker", "compose", "-f", str(LOCAL_COMPOSE_FILE), "down"], + check=True, + ) + print("Containers stopped successfully") + return True + except subprocess.CalledProcessError as e: + print(f"Error stopping containers: {e}") + return False + + +def logs(): + """Show container logs.""" + try: + subprocess.run( + ["docker", "compose", "-f", str(LOCAL_COMPOSE_FILE), "logs", "-f"], + check=True, + ) + return True + except subprocess.CalledProcessError as e: + print(f"Error showing logs: {e}") + return False + except KeyboardInterrupt: + print("\nExiting logs view") + return True + + +def run_server_command(command, args=None): + """Run a server command.""" + if args is None: + args = [] + + if command == "build": + return build() + elif command == "up": + # Parse guardrails file from args + guardrails_file = None + for arg in args: + if arg.startswith("--guardrails-file="): + guardrails_file = arg.split("=", 1)[1] + + return up(guardrails_file) + elif command == "down": + return down() + elif command == "logs": + return logs() + else: + print(f"Unknown server command: {command}") + print("Available commands: build, up, down, logs") + return False + + def main(): """Entry point for the Invariant Gateway.""" signal.signal(signal.SIGINT, signal_handler) @@ -37,9 +226,19 @@ def main(): verb = sys.argv[1] if verb == "mcp": return asyncio.run(mcp.execute(sys.argv[2:])) - if verb == "llm": - print("[gateway/__main__.py] 'llm' action is not implemented yet.") - return 1 + if verb == "server": + if len(sys.argv) < 3: + print( + "Error: Missing command for server. Should be one of: build, up, down, logs" + ) + print_help() + return 1 + + command = sys.argv[2] + args = sys.argv[3:] + if not run_server_command(command, args): + return 1 + return 0 if verb == "help": print_help() return 0 diff --git a/docker-compose.local.yml b/gateway/docker-compose.local.yml similarity index 92% rename from docker-compose.local.yml rename to gateway/docker-compose.local.yml index f570640..0f92ca9 100644 --- a/docker-compose.local.yml +++ b/gateway/docker-compose.local.yml @@ -2,8 +2,8 @@ services: invariant-gateway: container_name: invariant-gateway build: - context: . - dockerfile: Dockerfile.gateway + context: .. + dockerfile: gateway/Dockerfile.gateway working_dir: /srv/gateway env_file: - .env @@ -13,7 +13,7 @@ services: - ${INVARIANT_API_KEY:+INVARIANT_API_KEY=${INVARIANT_API_KEY}} volumes: - type: bind - source: ./gateway + source: ../gateway target: /srv/gateway - type: bind source: ${GUARDRAILS_FILE_PATH:-/dev/null} diff --git a/gateway/integrations/explorer.py b/gateway/integrations/explorer.py index dfd8bee..01bbb15 100644 --- a/gateway/integrations/explorer.py +++ b/gateway/integrations/explorer.py @@ -164,7 +164,7 @@ async def fetch_guardrails_from_explorer( ) # Get the user details. - user_info_response = await client.get("/api/v1/user/identity") + user_info_response = await client.get("/api/v1/user/identity", timeout=5) if user_info_response.status_code == 401: raise HTTPException( status_code=401, @@ -179,7 +179,7 @@ async def fetch_guardrails_from_explorer( # Get the dataset policies. policies_response = await client.get( - f"/api/v1/dataset/byuser/{username}/{dataset_name}/policy" + f"/api/v1/dataset/byuser/{username}/{dataset_name}/policy", timeout=5 ) if policies_response.status_code != 200: if policies_response.status_code == 404: diff --git a/gateway/integrations/guardrails.py b/gateway/integrations/guardrails.py index ca25016..31d4319 100644 --- a/gateway/integrations/guardrails.py +++ b/gateway/integrations/guardrails.py @@ -370,6 +370,7 @@ async def check_guardrails( "Authorization": context.get_guardrailing_authorization(), "Accept": "application/json", }, + timeout=5, ) if not result.is_success: if result.status_code == 401: diff --git a/run.sh b/run.sh index 636d634..6fff8c0 100755 --- a/run.sh +++ b/run.sh @@ -37,7 +37,7 @@ up() { fi # Start Docker Compose with the correct environment variable - docker compose -f docker-compose.local.yml up -d + docker compose -f gateway/docker-compose.local.yml up -d # Get the status of the container sleep 2 @@ -58,12 +58,12 @@ up() { build() { # Build local services - docker compose -f docker-compose.local.yml build + docker compose -f gateway/docker-compose.local.yml build } down() { # Bring down local services - docker compose -f docker-compose.local.yml down + docker compose -f gateway/docker-compose.local.yml down GATEWAY_ROOT_PATH=$(pwd) docker compose -f tests/integration/docker-compose.test.yml down } @@ -143,6 +143,7 @@ integration_tests() { # Generate latest whl file for the invariant-gateway package. # This is required to run the integration tests. pip install build + rm -rf dist python -m build WHEEL_FILE=$(ls dist/*.whl | head -n 1) echo "WHEEL_FILE: $WHEEL_FILE" @@ -151,9 +152,13 @@ integration_tests() { docker build -t 'invariant-gateway-tests' -f ./tests/integration/Dockerfile.test ./tests - docker run \ + docker rm -f invariant-gateway-integration-tests-container >/dev/null 2>&1 || true + mkdir -p /tmp/docker-logs/.invariant + + docker run --name invariant-gateway-integration-tests-container \ --mount type=bind,source=./tests/integration,target=/tests \ --mount type=bind,source=$(realpath $WHEEL_FILE),target=/package/$(basename $WHEEL_FILE) \ + --mount type=bind,source=/tmp/docker-logs/.invariant,target=/root/.invariant \ --network invariant-gateway-web-test \ -e OPENAI_API_KEY="$OPENAI_API_KEY" \ -e ANTHROPIC_API_KEY="$ANTHROPIC_API_KEY"\ @@ -185,7 +190,7 @@ case "$1" in down ;; "logs") - docker compose -f docker-compose.local.yml logs -f + docker compose -f gateway/docker-compose.local.yml logs -f ;; "unit-tests") shift diff --git a/tests/integration/docker-compose.test.yml b/tests/integration/docker-compose.test.yml index 42cb340..3695cc0 100644 --- a/tests/integration/docker-compose.test.yml +++ b/tests/integration/docker-compose.test.yml @@ -25,7 +25,7 @@ services: container_name: invariant-gateway-test build: context: ${GATEWAY_ROOT_PATH} - dockerfile: ${GATEWAY_ROOT_PATH}/Dockerfile.gateway + dockerfile: ${GATEWAY_ROOT_PATH}/gateway/Dockerfile.gateway depends_on: app-api: condition: service_healthy diff --git a/tests/integration/resources/mcp/stdio/client/main.py b/tests/integration/resources/mcp/stdio/client/main.py index 4238949..9feb2bb 100644 --- a/tests/integration/resources/mcp/stdio/client/main.py +++ b/tests/integration/resources/mcp/stdio/client/main.py @@ -2,7 +2,7 @@ import os import time - +from datetime import timedelta from contextlib import AsyncExitStack from typing import Any, Optional @@ -68,7 +68,9 @@ class MCPClient: ) self.stdio, self.write = stdio_transport self.session = await self.exit_stack.enter_async_context( - ClientSession(self.stdio, self.write) + ClientSession( + self.stdio, self.write, read_timeout_seconds=timedelta(minutes=0.5) + ) ) await self.session.initialize() @@ -130,5 +132,6 @@ async def run( finally: # Sleep for a while to allow the server to process the background tasks # like pushing traces to the explorer - time.sleep(2) + if push_to_explorer: + time.sleep(2) await client.cleanup()