diff --git a/.github/workflows/publish-images.yml b/.github/workflows/publish-images.yml index 04e8fb9..b4c15dd 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: ./gateway/Dockerfile.gateway + file: Dockerfile.gateway platforms: linux/amd64 push: true tags: ${{ env.tags }} diff --git a/.python-version b/.python-version new file mode 100644 index 0000000..e4fba21 --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.12 diff --git a/Dockerfile.gateway b/Dockerfile.gateway new file mode 100644 index 0000000..24e6cbd --- /dev/null +++ b/Dockerfile.gateway @@ -0,0 +1,16 @@ +FROM python:3.12 + +WORKDIR /srv/gateway + +COPY pyproject.toml ./ + +COPY gateway/ ./ + +COPY README.md /srv/gateway/README.md + +RUN pip install --no-cache-dir . + +RUN chmod +x /srv/gateway/run.sh + +# Run the application +CMD ["./run.sh"] \ No newline at end of file diff --git a/docker-compose.local.yml b/docker-compose.local.yml index 60ea5ad..f570640 100644 --- a/docker-compose.local.yml +++ b/docker-compose.local.yml @@ -2,8 +2,8 @@ services: invariant-gateway: container_name: invariant-gateway build: - context: ./gateway - dockerfile: ../gateway/Dockerfile.gateway + context: . + dockerfile: Dockerfile.gateway working_dir: /srv/gateway env_file: - .env diff --git a/gateway/Dockerfile.gateway b/gateway/Dockerfile.gateway deleted file mode 100644 index 431bfba..0000000 --- a/gateway/Dockerfile.gateway +++ /dev/null @@ -1,14 +0,0 @@ -FROM python:3.10 - -COPY ./requirements.txt /srv/gateway/requirements.txt -WORKDIR /srv/gateway - -RUN pip install --no-cache-dir -r requirements.txt -COPY . /srv/gateway - -# Ensure run.sh is executable -RUN chmod +x /srv/gateway/run.sh -WORKDIR /srv/gateway - -# Run the application -CMD ["./run.sh"] \ No newline at end of file diff --git a/gateway/common/config_manager.py b/gateway/common/config_manager.py index 951e41b..9499ac8 100644 --- a/gateway/common/config_manager.py +++ b/gateway/common/config_manager.py @@ -8,8 +8,7 @@ from typing import Optional import fastapi from httpx import HTTPStatusError -from common.guardrails import Guardrail, GuardrailAction, GuardrailRuleSet -from common.authorization import extract_authorization_from_headers +from gateway.common.guardrails import Guardrail, GuardrailAction, GuardrailRuleSet def extract_policy_from_headers(request: Optional[fastapi.Request]) -> Optional[str]: @@ -40,7 +39,7 @@ class GatewayConfig: Loads the guardrails from the file specified in GUARDRAILS_FILE_PATH. Returns the guardrails file content as a string. """ - from integrations.guardrails import _preload + from gateway.integrations.guardrails import _preload guardrails_file = os.getenv("GUARDRAILS_FILE_PATH", "") if not guardrails_file: diff --git a/gateway/common/request_context.py b/gateway/common/request_context.py index 68cb1d7..945c8b1 100644 --- a/gateway/common/request_context.py +++ b/gateway/common/request_context.py @@ -5,9 +5,9 @@ from typing import Any, Dict, Optional import fastapi -from common.config_manager import GatewayConfig -from common.guardrails import GuardrailRuleSet, Guardrail, GuardrailAction -from common.authorization import ( +from gateway.common.config_manager import GatewayConfig +from gateway.common.guardrails import GuardrailRuleSet, Guardrail, GuardrailAction +from gateway.common.authorization import ( extract_guardrail_service_authorization_from_headers, ) diff --git a/gateway/integrations/explorer.py b/gateway/integrations/explorer.py index fed63c2..85f820d 100644 --- a/gateway/integrations/explorer.py +++ b/gateway/integrations/explorer.py @@ -5,7 +5,7 @@ from typing import Any, Dict, List from fastapi import HTTPException -from common.guardrails import GuardrailRuleSet, Guardrail, GuardrailAction +from gateway.common.guardrails import GuardrailRuleSet, Guardrail, GuardrailAction from invariant_sdk.async_client import AsyncClient from invariant_sdk.types.push_traces import PushTracesRequest, PushTracesResponse from invariant_sdk.types.annotations import AnnotationCreate diff --git a/gateway/integrations/guardrails.py b/gateway/integrations/guardrails.py index d0f47ac..c84e44d 100644 --- a/gateway/integrations/guardrails.py +++ b/gateway/integrations/guardrails.py @@ -8,9 +8,10 @@ from functools import wraps from fastapi import HTTPException import httpx -from common.guardrails import Guardrail -from common.request_context import RequestContext -from common.authorization import ( + +from gateway.common.guardrails import Guardrail +from gateway.common.request_context import RequestContext +from gateway.common.authorization import ( INVARIANT_GUARDRAIL_SERVICE_AUTHORIZATION_HEADER, ) diff --git a/gateway/requirements.txt b/gateway/requirements.txt deleted file mode 100644 index 2aabbbf..0000000 --- a/gateway/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ -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/gateway/routes/anthropic.py b/gateway/routes/anthropic.py index 13b5f20..950f15b 100644 --- a/gateway/routes/anthropic.py +++ b/gateway/routes/anthropic.py @@ -8,27 +8,27 @@ import httpx from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response from starlette.responses import StreamingResponse -from common.authorization import extract_authorization_from_headers -from common.config_manager import ( +from gateway.common.authorization import extract_authorization_from_headers +from gateway.common.config_manager import ( GatewayConfig, GatewayConfigManager, GuardrailsInHeader, ) -from common.constants import ( +from gateway.common.constants import ( CLIENT_TIMEOUT, IGNORED_HEADERS, ) -from common.guardrails import GuardrailAction, GuardrailRuleSet -from common.request_context import RequestContext -from converters.anthropic_to_invariant import ( +from gateway.common.guardrails import GuardrailAction, GuardrailRuleSet +from gateway.common.request_context import RequestContext +from gateway.converters.anthropic_to_invariant import ( convert_anthropic_to_invariant_message_format, ) -from integrations.explorer import ( +from gateway.integrations.explorer import ( create_annotations_from_guardrails_errors, fetch_guardrails_from_explorer, push_trace, ) -from integrations.guardrails import ( +from gateway.integrations.guardrails import ( ExtraItem, InstrumentedResponse, InstrumentedStreamingResponse, diff --git a/gateway/routes/gemini.py b/gateway/routes/gemini.py index 00101f9..06814c7 100644 --- a/gateway/routes/gemini.py +++ b/gateway/routes/gemini.py @@ -8,25 +8,25 @@ import httpx from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response from fastapi.responses import StreamingResponse -from common.authorization import extract_authorization_from_headers -from common.config_manager import ( +from gateway.common.authorization import extract_authorization_from_headers +from gateway.common.config_manager import ( GatewayConfig, GatewayConfigManager, GuardrailsInHeader, ) -from common.constants import ( +from gateway.common.constants import ( CLIENT_TIMEOUT, IGNORED_HEADERS, ) -from common.guardrails import GuardrailAction, GuardrailRuleSet -from common.request_context import RequestContext -from converters.gemini_to_invariant import convert_request, convert_response -from integrations.explorer import ( +from gateway.common.guardrails import GuardrailAction, GuardrailRuleSet +from gateway.common.request_context import RequestContext +from gateway.converters.gemini_to_invariant import convert_request, convert_response +from gateway.integrations.explorer import ( create_annotations_from_guardrails_errors, fetch_guardrails_from_explorer, push_trace, ) -from integrations.guardrails import ( +from gateway.integrations.guardrails import ( ExtraItem, InstrumentedResponse, InstrumentedStreamingResponse, diff --git a/gateway/routes/open_ai.py b/gateway/routes/open_ai.py index 5695c03..b849dab 100644 --- a/gateway/routes/open_ai.py +++ b/gateway/routes/open_ai.py @@ -8,24 +8,24 @@ import httpx from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response from fastapi.responses import StreamingResponse -from common.authorization import extract_authorization_from_headers -from common.config_manager import ( +from gateway.common.authorization import extract_authorization_from_headers +from gateway.common.config_manager import ( GatewayConfig, GatewayConfigManager, GuardrailsInHeader, ) -from common.constants import ( +from gateway.common.constants import ( CLIENT_TIMEOUT, IGNORED_HEADERS, ) -from common.guardrails import GuardrailAction, GuardrailRuleSet -from common.request_context import RequestContext -from integrations.explorer import ( +from gateway.common.guardrails import GuardrailAction, GuardrailRuleSet +from gateway.common.request_context import RequestContext +from gateway.integrations.explorer import ( create_annotations_from_guardrails_errors, fetch_guardrails_from_explorer, push_trace, ) -from integrations.guardrails import ( +from gateway.integrations.guardrails import ( ExtraItem, InstrumentedResponse, InstrumentedStreamingResponse, diff --git a/gateway/run.sh b/gateway/run.sh index 4cee98c..38fc835 100755 --- a/gateway/run.sh +++ b/gateway/run.sh @@ -1,5 +1,7 @@ #!/bin/bash +export PYTHONPATH=/srv + # Validate configuration python validate_config.py diff --git a/gateway/serve.py b/gateway/serve.py index e3683f9..2568d6c 100644 --- a/gateway/serve.py +++ b/gateway/serve.py @@ -2,11 +2,12 @@ import fastapi import uvicorn -from routes.anthropic import gateway as anthropic_gateway -from routes.gemini import gateway as gemini_gateway -from routes.open_ai import gateway as open_ai_gateway from starlette_compress import CompressMiddleware +from gateway.routes.anthropic import gateway as anthropic_gateway +from gateway.routes.gemini import gateway as gemini_gateway +from gateway.routes.open_ai import gateway as open_ai_gateway + app = fastapi.app = fastapi.FastAPI( docs_url="/api/v1/gateway/docs", redoc_url="/api/v1/gateway/redoc", diff --git a/gateway/validate_config.py b/gateway/validate_config.py index 6ad9392..67d77c9 100644 --- a/gateway/validate_config.py +++ b/gateway/validate_config.py @@ -2,7 +2,7 @@ import sys -from common.config_manager import GatewayConfigManager +from gateway.common.config_manager import GatewayConfigManager try: _ = GatewayConfigManager.get_config() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..befa096 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "invariant-gateway" +version = "0.1.0" +description = "LLM proxy to observe and debug what your AI agents are doing" +readme = "README.md" +requires-python = ">=3.12" +dependencies = [ + "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" +] + +[tool.setuptools.packages.find] +where = ["."] + +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" \ No newline at end of file diff --git a/run.sh b/run.sh index 251f68d..7a7ad47 100755 --- a/run.sh +++ b/run.sh @@ -64,13 +64,12 @@ build() { down() { # Bring down local services docker compose -f docker-compose.local.yml down - GATEWAY_PATH=$(pwd)/gateway docker compose -f tests/integration/docker-compose.test.yml down + GATEWAY_ROOT_PATH=$(pwd) docker compose -f tests/integration/docker-compose.test.yml down } unit_tests() { echo "Running unit tests..." - - pytest tests/unit_tests $@ + PYTHONPATH=. pytest tests/unit_tests $@ } integration_tests() { @@ -108,7 +107,7 @@ integration_tests() { fi fi - export GATEWAY_PATH=$(pwd)/gateway + export GATEWAY_ROOT_PATH=$(pwd) export GUARDRAILS_FILE_PATH="$TEST_GUARDRAILS_FILE_PATH" # Start containers @@ -151,7 +150,7 @@ integration_tests() { --env-file ./tests/integration/.env.test \ invariant-gateway-tests $@ - unset GATEWAY_PATH + unset GATEWAY_ROOT_PATH unset GUARDRAILS_FILE_PATH } diff --git a/tests/integration/docker-compose.test.yml b/tests/integration/docker-compose.test.yml index 8c1afe3..42cb340 100644 --- a/tests/integration/docker-compose.test.yml +++ b/tests/integration/docker-compose.test.yml @@ -24,8 +24,8 @@ services: invariant-gateway: container_name: invariant-gateway-test build: - context: ${GATEWAY_PATH} - dockerfile: ${GATEWAY_PATH}/Dockerfile.gateway + context: ${GATEWAY_ROOT_PATH} + dockerfile: ${GATEWAY_ROOT_PATH}/Dockerfile.gateway depends_on: app-api: condition: service_healthy @@ -38,7 +38,7 @@ services: - ${INVARIANT_API_KEY:+INVARIANT_API_KEY=${INVARIANT_API_KEY}} volumes: - type: bind - source: ${GATEWAY_PATH} + source: ${GATEWAY_ROOT_PATH}/gateway target: /srv/gateway - type: bind source: ${GUARDRAILS_FILE_PATH:-/dev/null} diff --git a/tests/integration/guardrails/test_guardrails_anthropic.py b/tests/integration/guardrails/test_guardrails_anthropic.py index 035a845..14f06e5 100644 --- a/tests/integration/guardrails/test_guardrails_anthropic.py +++ b/tests/integration/guardrails/test_guardrails_anthropic.py @@ -457,12 +457,12 @@ async def test_with_guardrails_from_explorer(explorer_api_url, gateway_url, do_s assert ( annotations[0]["content"] == "ogre detected in response" and annotations[0]["extra_metadata"]["source"] == "guardrails-error" - and annotations[0]["extra_metadata"]["guardrail-action"] == "block" + and annotations[0]["extra_metadata"]["guardrail"]["action"] == "block" ) assert ( annotations[1]["content"] == "Fiona detected in response" and annotations[1]["extra_metadata"]["source"] == "guardrails-error" - and annotations[1]["extra_metadata"]["guardrail-action"] == "log" + and annotations[1]["extra_metadata"]["guardrail"]["action"] == "log" ) @@ -584,7 +584,7 @@ async def test_preguardrailing_with_guardrails_from_explorer( assert ( annotations[0]["content"] == "pun detected in user message" and annotations[0]["extra_metadata"]["source"] == "guardrails-error" - and annotations[0]["extra_metadata"]["guardrail-action"] == "block" + and annotations[0]["extra_metadata"]["guardrail"]["action"] == "block" if is_block_action else "log" ) diff --git a/tests/integration/guardrails/test_guardrails_gemini.py b/tests/integration/guardrails/test_guardrails_gemini.py index b3ac35e..292d473 100644 --- a/tests/integration/guardrails/test_guardrails_gemini.py +++ b/tests/integration/guardrails/test_guardrails_gemini.py @@ -435,12 +435,12 @@ async def test_with_guardrails_from_explorer(explorer_api_url, gateway_url, do_s assert ( annotations[0]["content"] == "ogre detected in response" and annotations[0]["extra_metadata"]["source"] == "guardrails-error" - and annotations[0]["extra_metadata"]["guardrail-action"] == "block" + and annotations[0]["extra_metadata"]["guardrail"]["action"] == "block" ) assert ( annotations[1]["content"] == "Fiona detected in response" and annotations[1]["extra_metadata"]["source"] == "guardrails-error" - and annotations[1]["extra_metadata"]["guardrail-action"] == "log" + and annotations[1]["extra_metadata"]["guardrail"]["action"] == "log" ) @@ -550,7 +550,7 @@ async def test_preguardrailing_with_guardrails_from_explorer( assert ( annotations[0]["content"] == "pun detected in user message" and annotations[0]["extra_metadata"]["source"] == "guardrails-error" - and annotations[0]["extra_metadata"]["guardrail-action"] == "block" + and annotations[0]["extra_metadata"]["guardrail"]["action"] == "block" if is_block_action else "log" ) diff --git a/tests/integration/guardrails/test_guardrails_open_ai.py b/tests/integration/guardrails/test_guardrails_open_ai.py index 6031778..a418f25 100644 --- a/tests/integration/guardrails/test_guardrails_open_ai.py +++ b/tests/integration/guardrails/test_guardrails_open_ai.py @@ -457,12 +457,12 @@ async def test_with_guardrails_from_explorer(explorer_api_url, gateway_url, do_s assert ( annotations[0]["content"] == "ogre detected in response" and annotations[0]["extra_metadata"]["source"] == "guardrails-error" - and annotations[0]["extra_metadata"]["guardrail-action"] == "block" + and annotations[0]["extra_metadata"]["guardrail"]["action"] == "block" ) assert ( annotations[1]["content"] == "Fiona detected in response" and annotations[1]["extra_metadata"]["source"] == "guardrails-error" - and annotations[1]["extra_metadata"]["guardrail-action"] == "log" + and annotations[1]["extra_metadata"]["guardrail"]["action"] == "log" ) @@ -581,7 +581,7 @@ async def test_preguardrailing_with_guardrails_from_explorer( assert ( annotations[0]["content"] == "pun detected in user message" and annotations[0]["extra_metadata"]["source"] == "guardrails-error" - and annotations[0]["extra_metadata"]["guardrail-action"] == "block" + and annotations[0]["extra_metadata"]["guardrail"]["action"] == "block" if is_block_action else "log" ) diff --git a/tests/unit_tests/common/test_authorization.py b/tests/unit_tests/common/test_authorization.py index 60512aa..2014bc1 100644 --- a/tests/unit_tests/common/test_authorization.py +++ b/tests/unit_tests/common/test_authorization.py @@ -7,10 +7,6 @@ import random import string import pytest -from gateway.common.config_manager import GatewayConfig -from gateway.common.request_context import RequestContext - - # Add root folder (parent) to sys.path sys.path.append( os.path.dirname( @@ -18,6 +14,8 @@ sys.path.append( ) ) +from gateway.common.config_manager import GatewayConfig +from gateway.common.request_context import RequestContext from gateway.common.authorization import ( INVARIANT_GUARDRAIL_SERVICE_AUTHORIZATION_HEADER, extract_authorization_from_headers,