From a5ea86a64ecdd365ee0905d7c5ae54037a8c7149 Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Fri, 14 Mar 2025 14:37:23 +0100
Subject: [PATCH] Show logs if the docker container doesn't start on doing
 up(). Tested when the guardrails file contained unparseable policy.

---
 gateway/common/request_context_data.py |  2 +-
 gateway/integrations/explorer.py       | 21 ++++++++++-----------
 gateway/routes/open_ai.py              |  7 ++++++-
 run.sh                                 | 13 ++++++++++++-
 4 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/gateway/common/request_context_data.py b/gateway/common/request_context_data.py
index da967bf..8ae98ba 100644
--- a/gateway/common/request_context_data.py
+++ b/gateway/common/request_context_data.py
@@ -3,7 +3,7 @@
 from dataclasses import dataclass
 from typing import Any, Dict, Optional
 
-from config_manager import GatewayConfig
+from common.config_manager import GatewayConfig
 
 
 @dataclass(frozen=True)
diff --git a/gateway/integrations/explorer.py b/gateway/integrations/explorer.py
index b557008..b61e610 100644
--- a/gateway/integrations/explorer.py
+++ b/gateway/integrations/explorer.py
@@ -17,18 +17,14 @@ def create_annotations_from_guardrails_errors(
     annotations = []
     for error in guardrails_errors:
         content = error.get("args")[0]
-        address = None
         for r in error.get("ranges", []):
-            # Choose the longest path as the address
-            if address is None or len(r) > len(address):
-                address = r
-        annotations.append(
-            AnnotationCreate(
-                content=content,
-                address=address,
-                extra_metadata={"source": "guardrails-error"},
+            annotations.append(
+                AnnotationCreate(
+                    content=content,
+                    address=r,
+                    extra_metadata={"source": "guardrails-error"},
+                )
             )
-        )
     return annotations
 
 
@@ -36,6 +32,7 @@ async def push_trace(
     messages: List[List[Dict[str, Any]]],
     dataset_name: str,
     invariant_authorization: str,
+    annotations: List[List[AnnotationCreate]] = None,
 ) -> PushTracesResponse:
     """Pushes traces to the dataset on the Invariant Explorer.
 
@@ -55,7 +52,9 @@ async def push_trace(
         [{k: v for k, v in msg.items() if v is not None} for msg in msg_list]
         for msg_list in messages
     ]
-    request = PushTracesRequest(messages=update_messages, dataset=dataset_name)
+    request = PushTracesRequest(
+        messages=update_messages, annotations=annotations, dataset=dataset_name
+    )
     client = AsyncClient(
         api_url=os.getenv("INVARIANT_API_URL", DEFAULT_API_URL).rstrip("/"),
         api_key=invariant_authorization.split("Bearer ")[1],
diff --git a/gateway/routes/open_ai.py b/gateway/routes/open_ai.py
index c21474a..0e7338d 100644
--- a/gateway/routes/open_ai.py
+++ b/gateway/routes/open_ai.py
@@ -332,9 +332,14 @@ async def handle_non_streaming_response(
     response_code = response.status_code
 
     if context.config and context.config.guardrails:
-        # Block on the guardrails check
         messages = list(context.request_json.get("messages", []))
         messages += [choice["message"] for choice in json_response.get("choices", [])]
+        # TODO: Remove this once the guardrails API is fixed
+        for message in messages:
+            if "tool_calls" in message and message["tool_calls"] is None:
+                message["tool_calls"] = []
+
+        # Block on the guardrails check
         guardrails_execution_result = await check_guardrails(
             messages=messages,
             guardrails=context.config.guardrails,
diff --git a/run.sh b/run.sh
index f7bea4a..db3f273 100755
--- a/run.sh
+++ b/run.sh
@@ -32,9 +32,20 @@ up() {
   # Start Docker Compose with the correct environment variable
   GUARDRAILS_FILE_PATH="$GUARDRAILS_FILE_PATH" docker compose -f docker-compose.local.yml up -d
 
+  # Get the status of the container
+  sleep 2
+
+  if [ -z "$(docker ps -qf 'name=invariant-gateway')" ]; then
+    echo "The invariant-gateway container failed to start."
+    docker logs invariant-gateway | tail -20  # Show last 20 lines of logs
+    exit 1
+  fi
+
   echo "Gateway started at http://localhost:8005/api/v1/gateway/"
   echo "See http://localhost:8005/api/v1/gateway/docs for API documentation"
-  echo "Using Guardrails File: ${GUARDRAILS_FILE_PATH:-None}"
+  if [ -n "$GUARDRAILS_FILE_PATH" ]; then
+    echo "Using Guardrails File: $GUARDRAILS_FILE_PATH"
+  fi
 }
 
 build() {