initial draft: pipelined guardrails

.env

@@ -3,4 +3,4 @@
 # If you want to push to a local instance of explorer, then specify the app-api docker container name like: 
 # http://<app-api-docker-container-name>:8000 to push to the local explorer instance.
 INVARIANT_API_URL=https://explorer.invariantlabs.ai
-GUADRAILS_API_URL=https://guardrail.invariantnet.com
+GUADRAILS_API_URL=https://explorer.invariantlabs.ai

example_policy.gr

@@ -0,0 +1,6 @@
+from invariant.detectors import prompt_injection
+
+raise "Don't say 'Hello'" if:
+    (msg: Message)
+    msg.role == "user"
+    prompt_injection(msg.content)

gateway/integrations/guardails.py

@@ -8,7 +8,7 @@ from functools import wraps
 
 import httpx
 
-DEFAULT_API_URL = "https://guardrail.invariantnet.com"
+DEFAULT_API_URL = "https://explorer.invariantlabs.ai"
 
 
 # Timestamps of last API calls per guardrails string
@@ -99,6 +99,214 @@ async def preload_guardrails(context: "RequestContextData") -> None:
         print(f"Error scheduling preload_guardrails task: {e}")
 
 
+class YieldException(Exception):
+    """
+    Raise this exception in stream instrumentor listeners to
+    end the stream early, or to emit additional items in a stream.
+    """
+
+    def __init__(self, value, end_of_stream=False):
+        super().__init__(value)
+        self.value = value
+        self.end_of_stream = end_of_stream
+
+    def __str__(self):
+        return f"YieldException: {self.value}"
+
+
+class StreamInstrumentor:
+    """
+    A class to instrument async iterables with hooks for processing
+    chunks, before processing, and on completion.
+    """
+
+    def __init__(self):
+        # called on every chunk (async)
+        self.on_chunk_listeners = []
+        # called once before the first chunk is processed, or even earlier (async)
+        self.before_listeners = []
+        # called once on stream completion (async)
+        self.on_complete_listeners = []
+
+        self.stat_token_times = []
+        self.stat_before_time = None
+        self.stat_after_time = None
+
+        self.stat_first_item_time = None
+
+    # decorator
+    def on(self, event: str):
+        """
+        Decorator to register listeners for different events.
+
+        Args:
+            event (str): The event to listen for. Can be 'on_chunk',
+                         'before', or 'on_complete'.
+
+        Returns:
+            Callable: A decorator to register the listener.
+        """
+
+        def decorator(func):
+            if event == "chunk":
+                if self.on_chunk_listeners is None:
+                    self.on_chunk_listeners = []
+                self.on_chunk_listeners.append(func)
+            elif event == "start":
+                if self.before_listeners is None:
+                    self.before_listeners = []
+                self.before_listeners.append(func)
+            elif event == "end":
+                if self.on_complete_listeners is None:
+                    self.on_complete_listeners = []
+                self.on_complete_listeners.append(func)
+            else:
+                raise ValueError("Invalid event type. Use 'chunk', 'before', or 'end'.")
+
+            return func
+
+        return decorator
+
+    async def stream(self, async_iterable):
+        """
+        Streams the async iterable and invokes all instrumented hooks.
+
+        Args:
+            async_iterable: An async iterable to stream.
+
+        Yields:
+            The streamed data.
+        """
+        try:
+            start = time.time()
+
+            # schedule all before listeners which can be run concurrently
+            before_tasks = [
+                asyncio.create_task(listener()) for listener in self.before_listeners
+            ]
+
+            # create async iterator from async_iterable
+            aiterable = aiter(async_iterable)
+
+            # [STAT] capture start time of first item
+            start_first_item_request = time.time()
+
+            # waits for first item of the iterable
+            async def wait_for_first_item():
+                nonlocal start_first_item_request, aiterable
+
+                r = await aiterable.__anext__()
+                self.stat_first_item_time = time.time() - start_first_item_request
+                return r
+
+            next_item_task = asyncio.create_task(wait_for_first_item())
+
+            # wait for all before listeners to finish
+            for before_task in before_tasks:
+                try:
+                    await before_task
+                except YieldException as e:
+                    # yield extra value before any real items
+                    yield e.value
+                    # stop the stream if end_of_stream is True
+                    if e.end_of_stream:
+                        # if first item is already available
+                        if not next_item_task.done():
+                            # cancel the task
+                            next_item_task.cancel()
+                            # [STAT] capture time to first item to be now +0.01
+                            if self.stat_first_item_time is None:
+                                self.stat_first_item_time = (
+                                    time.time() - start_first_item_request + 0.01
+                                )
+                        else:
+                            print(
+                                "before yields, but next item already ready", flush=True
+                            )
+
+            # [STAT] capture before time stamp
+            self.stat_before_time = time.time() - start
+
+            while True:
+                # wait for first item
+                try:
+                    item = await next_item_task
+                except StopAsyncIteration:
+                    break
+
+                # schedule next item
+                next_item_task = asyncio.create_task(aiterable.__anext__())
+
+                # [STAT] capture token time stamp
+                if len(self.stat_token_times) == 0:
+                    self.stat_token_times.append(time.time() - start)
+                else:
+                    self.stat_token_times.append(
+                        time.time() - start - sum(self.stat_token_times)
+                    )
+
+                # invoke on_chunk listeners
+                for listener in self.on_chunk_listeners:
+                    any_end_of_stream = False
+                    try:
+                        await listener(item)
+                    except YieldException as e:
+                        yield e.value
+                        # if end_of_stream is True, stop the stream
+                        if e.end_of_stream:
+                            any_end_of_stream = True
+
+                # if end_of_stream is True, stop the stream
+                if any_end_of_stream:
+                    break
+
+                # yield item
+                yield item
+            # execute on complete listeners
+            on_complete_tasks = [
+                asyncio.create_task(listener())
+                for listener in self.on_complete_listeners
+            ]
+            for result in asyncio.as_completed(on_complete_tasks):
+                try:
+                    await result
+                except YieldException as e:
+                    # yield extra value before any real items
+                    yield e.value
+                    # we ignore end_of_stream here, because we are already at the end
+
+            # [STAT] capture after time stamp
+            self.stat_after_time = time.time() - start
+
+        finally:
+            # [STAT] end all open intervals if not already closed
+            if self.stat_after_time is None:
+                self.stat_before_time = time.time() - start
+            if self.stat_after_time is None:
+                self.stat_after_time = 0
+            if self.stat_first_item_time is None:
+                self.stat_first_item_time = 0
+
+            token_times_5_decimale = str([f"{x:.5f}" for x in self.stat_token_times])
+            print(
+                f"[STATS]\n [token times: {token_times_5_decimale} ({len(self.stat_token_times)})]"
+            )
+            print(f" [before:             {self.stat_before_time:.2f}s] ")
+            print(f" [time-to-first-item: {self.stat_first_item_time:.2f}s]")
+            print(
+                f" [zero-latency:       {' TRUE' if self.stat_before_time < self.stat_first_item_time else 'FALSE'}]"
+            )
+            print(
+                f" [extra-latency:      {self.stat_before_time - self.stat_first_item_time:.2f}s]"
+            )
+            print(f" [after:              {self.stat_after_time:.2f}s]")
+            if len(self.stat_token_times) > 0:
+                print(
+                    f" [average token time: {sum(self.stat_token_times) / len(self.stat_token_times):.2f}s]"
+                )
+            print(f" [total: {time.time() - start:.2f}s]")
+
+
 async def check_guardrails(
     messages: List[Dict[str, Any]], guardrails: str, invariant_authorization: str
 ) -> Dict[str, Any]:

gateway/routes/open_ai.py

@@ -13,7 +13,12 @@ from common.constants import (
     IGNORED_HEADERS,
 )
 from integrations.explorer import create_annotations_from_guardrails_errors, push_trace
-from integrations.guardails import check_guardrails, preload_guardrails
+from integrations.guardails import (
+    StreamInstrumentor,
+    YieldException,
+    check_guardrails,
+    preload_guardrails,
+)
 from common.authorization import extract_authorization_from_headers
 from common.request_context_data import RequestContextData
 
@@ -98,15 +103,23 @@ async def stream_response(
     It is sent to the Invariant Explorer at the end of the stream
     """
 
-    response = await client.send(open_ai_request, stream=True)
-    if response.status_code != 200:
-        error_content = await response.aread()
-        try:
-            error_json = json.loads(error_content.decode("utf-8"))
-            error_detail = error_json.get("error", "Unknown error from OpenAI API")
-        except json.JSONDecodeError:
-            error_detail = {"error": "Failed to parse OpenAI error response"}
-        raise HTTPException(status_code=response.status_code, detail=error_detail)
+    async def request_and_stream():
+        """
+        Sets of the request and then streams the result.
+        """
+        response = await client.send(open_ai_request, stream=True)
+        if response.status_code != 200:
+            error_content = await response.aread()
+            try:
+                error_json = json.loads(error_content.decode("utf-8"))
+                error_detail = error_json.get("error", "Unknown error from OpenAI API")
+            except json.JSONDecodeError:
+                error_detail = {"error": "Failed to parse OpenAI error response"}
+            raise HTTPException(status_code=response.status_code, detail=error_detail)
+
+        # stream out chunks
+        async for chunk in response.aiter_bytes():
+            yield chunk
 
     async def event_generator() -> Any:
         # merged_response will be updated with the data from the chunks in the stream
@@ -119,6 +132,7 @@ async def stream_response(
             "choices": [],
             "usage": None,
         }
+
         # Each chunk in the stream contains a list called "choices" each entry in the list
         # has an index.
         # A choice has a field called "delta" which may contain a list called "tool_calls".
@@ -127,10 +141,37 @@ async def stream_response(
         # Combines the choice index and tool call index to uniquely identify a tool call
         tool_call_mapping_by_index = {}
 
-        async for chunk in response.aiter_bytes():
+        # prepare stream instrumentor
+        instrumentor = StreamInstrumentor()
+
+        @instrumentor.on("start")
+        async def precheck_guardrails() -> None:
+            # Check guardrails on the first chunk
+            if context.config and context.config.guardrails:
+                # Block on the guardrails check
+                guardrails_execution_result = await get_guardrails_check_result(
+                    context, merged_response
+                )
+                if guardrails_execution_result.get("errors", []):
+                    error_chunk = json.dumps(
+                        {
+                            "error": {
+                                "message": "[Invariant] The response did not pass the guardrails",
+                                "details": guardrails_execution_result,
+                            }
+                        }
+                    )
+                    # if we find something, we end the stream prematurely (end_of_stream=True)
+                    # and yield an error chunk instead of actually beginning the stream
+                    raise YieldException(
+                        f"data: {error_chunk}\n\n".encode(), end_of_stream=True
+                    )
+
+        @instrumentor.on("chunk")
+        async def process_chunk(chunk: bytes) -> None:
             chunk_text = chunk.decode().strip()
             if not chunk_text:
-                continue
+                return
 
             # Process the chunk
             # This will update merged_response with the data from the chunk
@@ -141,7 +182,7 @@ async def stream_response(
                 tool_call_mapping_by_index,
             )
 
-            # Check guardrails on the last chunk.
+            # Check guardrails on the 'DONE' SSE chunk.
             if (
                 "data: [DONE]" in chunk_text
                 and context.config
@@ -169,9 +210,11 @@ async def stream_response(
                                 guardrails_execution_result,
                             )
                         )
-                    yield f"data: {error_chunk}\n\n".encode()
-                    return
 
+                    # yield an extra error chunk (without preventing the original chunk to go through after)
+                    raise YieldException(f"data: {error_chunk}\n\n".encode())
+
+        async for chunk in instrumentor.stream(request_and_stream()):
             # Yield chunk to the client
             yield chunk
 

gateway/run.sh

@@ -15,7 +15,7 @@ UVICORN_PORT=${PORT:-8000}
 # using 'exec' belows ensures that signals like SIGTERM are passed to the child process
 # and not the shell script itself (important when running in a container)
 if [ "$DEV_MODE" = "true" ]; then
-    exec uvicorn serve:app --host 0.0.0.0 --port $UVICORN_PORT --reload
+    exec uvicorn serve:app --host 0.0.0.0 --port $UVICORN_PORT --reload --reload-dir /srv/resources --reload-dir /srv/gateway
 else
     exec uvicorn serve:app --host 0.0.0.0 --port $UVICORN_PORT
 fi