dirty demo hacks

.env

@@ -1 +1 @@
-INVARIANT_API_URL=https://explorer.invariantlabs.ai
+INVARIANT_API_URL=http://explorer-app-api-1:8000

docker-compose.local.yml

@@ -4,6 +4,7 @@ services:
     build:
       context: ./proxy
       dockerfile: ../proxy/Dockerfile.proxy
+    platform: linux/amd64
     working_dir: /srv/proxy
     env_file:
       - .env

proxy/requirements.txt

@@ -2,4 +2,12 @@ fastapi==0.115.7
 httpx==0.28.1
 invariant-sdk>=0.0.10
 starlette-compress==1.4.0
-uvicorn==0.34.0
+uvicorn==0.34.0
+invariant-ai[dev]==0.2
+presidio-analyzer
+spacy
+transformers
+torch
+python-dotenv
+numpy
+invariant_sdk

proxy/routes/open_ai.py

@@ -7,7 +7,7 @@ import httpx
 from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response
 from starlette.responses import StreamingResponse
 from utils.constants import CLIENT_TIMEOUT, IGNORED_HEADERS
-from utils.explorer import push_trace
+from utils.explorer import error_label, push_trace, validate_guardrails
 
 ALLOWED_OPEN_AI_ENDPOINTS = {"chat/completions"}
 
@@ -308,21 +308,28 @@ async def push_to_explorer(
 ) -> None:
     """Pushes the full trace to the Invariant Explorer"""
     # Only push the trace to explorer if the message is an end turn message
+    only_push_if_blocked = False
+
     if (
         merged_response.get("choices")
         and merged_response["choices"][0].get("finish_reason")
         not in FINISH_REASON_TO_PUSH_TRACE
     ):
-        return
+        only_push_if_blocked = True
+
     # Combine the messages from the request body and the choices from the OpenAI response
     messages = request_body.get("messages", [])
     messages += [choice["message"] for choice in merged_response.get("choices", [])]
-    _ = await push_trace(
+
+    blocked, response = await push_trace(
         dataset_name=dataset_name,
         messages=[messages],
         invariant_authorization=invariant_authorization,
+        dry=only_push_if_blocked,
     )
 
+    return blocked[0]
+
 
 async def handle_non_streaming_response(
     response: httpx.Response,
@@ -332,10 +339,19 @@ async def handle_non_streaming_response(
 ):
     """Handles non-streaming OpenAI responses"""
     json_response = response.json()
-    await push_to_explorer(
+    blocked = await push_to_explorer(
         dataset_name, json_response, request_body_json, invariant_authorization
     )
 
+    if blocked:
+        # json_response["c
+        json_response["choices"][-1]["finish_reason"] = "blocked"
+        json_response["choices"][-1]["message"]["content"] = (
+            f"[Agent execution blocked by guardrail: {error_label(blocked)}]"
+        )
+        # remove tool calls
+        json_response["choices"][-1]["message"]["tool_calls"] = None
+
     return Response(
         content=json.dumps(json_response),
         status_code=response.status_code,

proxy/utils/explorer.py

@@ -2,9 +2,17 @@
 
 import os
 from typing import Any, Dict, List
+import re
 
 from invariant_sdk.async_client import AsyncClient
-from invariant_sdk.types.push_traces import PushTracesRequest, PushTracesResponse
+from invariant_sdk.types.push_traces import (
+    PushTracesRequest,
+    PushTracesResponse,
+    AnnotationCreate,
+)
+import json
+
+from invariant.analyzer import Policy
 
 DEFAULT_API_URL = "https://explorer.invariantlabs.ai"
 
@@ -13,13 +21,14 @@ async def push_trace(
     messages: List[List[Dict[str, Any]]],
     dataset_name: str,
     invariant_authorization: str,
+    dry: bool = False,
 ) -> PushTracesResponse:
     """Pushes traces to the dataset on the Invariant Explorer.
 
     If a dataset with the given name does not exist, it will be created.
 
     Args:
-        messages (List[List[Dict[str, Any]]]): List of messages to push.
+        messages (List[List[Dict[str, Any]]]): List of messages to push
         dataset_name (str): Name of the dataset.
         invariant_authorization (str): Value of the
                                        invariant-authorization header.
@@ -32,13 +41,157 @@ 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)
+
     client = AsyncClient(
         api_url=os.getenv("INVARIANT_API_URL", DEFAULT_API_URL).rstrip("/"),
         api_key=invariant_authorization.split("Bearer ")[1],
     )
+
+    # validate guardrails (and get annotations)
+    blocked_and_annotations = [
+        await validate_guardrails(client, messages, dataset_name)
+        for messages in update_messages
+    ]
+
+    annotations = [annotation for _, annotation in blocked_and_annotations]
+    blocked = [block for block, _ in blocked_and_annotations]
+
+    # for blocked messages histories, apply "system" message with [blocked] content
+    for i, block in enumerate(blocked):
+        if block:
+            update_messages[i].append(
+                {
+                    "content": f"[Agent execution blocked by guardrail: {error_label(block)}]",
+                    "role": "system",
+                }
+            )
+
+    request = PushTracesRequest(
+        messages=update_messages, dataset=dataset_name, annotations=annotations
+    )
+
     try:
-        return await client.push_trace(request)
+        # if dry run, don't push the trace (but still validate guardrails)
+        if dry and not any(blocked):
+            result = {"dry_run": True}
+        else:
+            result = await client.push_trace(request)
     except Exception as e:
         print(f"Failed to push trace: {e}")
-        return {"error": str(e)}
+        result = {"error": str(e)}
+
+    return blocked, result
+
+
+async def validate_guardrails(
+    client: AsyncClient,
+    messages: List[List[Dict[str, Any]]],
+    dataset_name: str,
+) -> PushTracesResponse:
+    """Fetches and validates the guardrails for the given dataset.
+
+    Args:
+        messages (List[List[Dict[str, Any]]]): List of messages to push.
+        dataset_name (str): Name of the dataset.
+        invariant_authorization (str): Value of the
+                                       invariant-authorization header.
+
+    Returns:
+        PushTracesResponse: Response containing the trace ID details.
+    """
+    try:
+        metadata = await client.get_dataset_metadata(dataset_name=dataset_name)
+    except Exception as e:
+        print(f"Failed to get dataset metadata: {e}")
+        return False, []
+    guardrails = json.loads(metadata.get("guardrails", "[]"))
+
+    blocked = False
+
+    # preprocess messages
+    trace = [{**msg} for msg in messages]
+
+    # if content is missing in a msg, set it to empty string
+    for msg in trace:
+        if "content" not in msg:
+            msg["content"] = ""
+
+    annotations = []
+
+    for guardrail in guardrails:
+        if not guardrail.get("enabled", False):
+            print(f"Skipping guardrail {guardrail['name']} as it is disabled")
+            continue
+
+        policy = Policy.from_string(guardrail["policy"])
+        results = policy.analyze(trace)
+
+        for error in results.errors:
+            label = error_label(error)
+            # check for action=[block|warn|...]
+            if "action=" in label:
+                action = label.split("action=")[1].split()[0]
+                if action == "block":
+                    blocked = error
+
+            ranges = [range for range in error.ranges]
+
+            if len(ranges) > 1:
+                prefixfree_ranges = []
+                # remove prefixes
+                for range in ranges:
+                    if not any(
+                        range.json_path.startswith(prefix.json_path)
+                        for prefix in ranges
+                        if range != prefix
+                    ):
+                        prefixfree_ranges.append(range)
+                ranges = prefixfree_ranges
+
+            for range in ranges:
+                address = "messages." + range.json_path
+                old_address = address
+                try:
+                    # if there is .start and .end, also include them as :start-end
+                    if (
+                        range.start is not None
+                        and range.end is not None
+                        and "-" not in address
+                    ):
+                        address += f":{range.start}-{range.end}"
+                    # handle special cases
+                    # 1. if the address is messages.[0-9]+, change it to messages.[0-9]+.content:0-<end of content field>
+                    elif re.match(r"messages\.\d+$", address) and "-" not in address:
+                        end = len(trace[int(address.split(".")[1])]["content"])
+                        address += ".content:0-" + str(end)
+                    # 2. if address is messages.4.tool_calls.0, change it to messages.4.tool_calls.0.function.name:0-<end of function name>
+                    elif (
+                        re.match(r"messages\.\d+\.tool_calls\.\d+$", address)
+                        and "-" not in address
+                    ):
+                        end = len(
+                            trace[int(address.split(".")[1])]["tool_calls"][
+                                int(address.split(".")[3])
+                            ]["function"]["name"]
+                        )
+                        address += ".function.name:0-" + str(end)
+                except Exception as e:
+                    print(f"Error while handling special cases: {e}")
+                    pass
+
+                print("fixed", old_address, "->", address)
+
+                annotations.append(
+                    AnnotationCreate(
+                        content=label,
+                        address=address,
+                        extra_metadata={"source": "guardrail"},
+                    )
+                )
+    return blocked, annotations
+
+
+def error_label(error):
+    if str(error).startswith("ErrorInformation("):
+        return str(error).split("(", 1)[1].rsplit(")", 1)[0]
+    return str(error)