diff --git a/.env b/.env index 410fc9a..af1265e 100644 --- a/.env +++ b/.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://: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 diff --git a/example_policy.gr b/example_policy.gr new file mode 100644 index 0000000..29f9c03 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/gateway/integrations/guardails.py b/gateway/integrations/guardails.py index 4c571fc..6067dbc 100644 --- a/gateway/integrations/guardails.py +++ b/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]: diff --git a/gateway/routes/open_ai.py b/gateway/routes/open_ai.py index 897dbb8..1a8d13e 100644 --- a/gateway/routes/open_ai.py +++ b/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 diff --git a/gateway/run.sh b/gateway/run.sh index ebc4b1b..4cee98c 100755 --- a/gateway/run.sh +++ b/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 \ No newline at end of file