mirror of
https://github.com/invariantlabs-ai/invariant-gateway.git
synced 2026-08-11 11:30:20 +02:00
364 lines
12 KiB
Python
364 lines
12 KiB
Python
"""Utility functions for Guardrails execution."""
|
|
|
|
import asyncio
|
|
import os
|
|
import time
|
|
from typing import Any, Dict, List
|
|
from functools import wraps
|
|
|
|
import httpx
|
|
|
|
DEFAULT_API_URL = "https://explorer.invariantlabs.ai"
|
|
|
|
|
|
# Timestamps of last API calls per guardrails string
|
|
_guardrails_cache = {}
|
|
# Locks per guardrails string
|
|
_guardrails_locks = {}
|
|
|
|
|
|
def rate_limit(expiration_time: int = 3600):
|
|
"""
|
|
Decorator to limit API calls to once per expiration_time seconds
|
|
per unique guardrails string.
|
|
|
|
Args:
|
|
expiration_time (int): Time in seconds to cache the guardrails.
|
|
"""
|
|
|
|
def decorator(func):
|
|
@wraps(func)
|
|
async def wrapper(guardrails: str, *args, **kwargs):
|
|
now = time.time()
|
|
|
|
# Get or create a per-guardrail lock
|
|
if guardrails not in _guardrails_locks:
|
|
_guardrails_locks[guardrails] = asyncio.Lock()
|
|
guardrail_lock = _guardrails_locks[guardrails]
|
|
|
|
async with guardrail_lock:
|
|
last_called = _guardrails_cache.get(guardrails)
|
|
|
|
if last_called and (now - last_called < expiration_time):
|
|
# Skipping API call: Guardrails '{guardrails}' already
|
|
# preloaded within expiration_time
|
|
return
|
|
|
|
# Update cache timestamp
|
|
_guardrails_cache[guardrails] = now
|
|
|
|
try:
|
|
await func(guardrails, *args, **kwargs)
|
|
finally:
|
|
_guardrails_locks.pop(guardrails, None)
|
|
|
|
return wrapper
|
|
|
|
return decorator
|
|
|
|
|
|
@rate_limit(3600) # Don't preload the same guardrails string more than once per hour
|
|
async def _preload(guardrails: str, invariant_authorization: str) -> None:
|
|
"""
|
|
Calls the Guardrails API to preload the provided policy for faster checking later.
|
|
|
|
Args:
|
|
guardrails (str): The guardrails to preload.
|
|
invariant_authorization (str): Value of the
|
|
invariant-authorization header.
|
|
"""
|
|
async with httpx.AsyncClient() as client:
|
|
url = os.getenv("GUADRAILS_API_URL", DEFAULT_API_URL).rstrip("/")
|
|
result = await client.post(
|
|
f"{url}/api/v1/policy/load",
|
|
json={"policy": guardrails},
|
|
headers={
|
|
"Authorization": invariant_authorization,
|
|
"Accept": "application/json",
|
|
},
|
|
)
|
|
result.raise_for_status()
|
|
|
|
|
|
async def preload_guardrails(context: "RequestContextData") -> None:
|
|
"""
|
|
Preloads the guardrails for faster checking later.
|
|
|
|
Args:
|
|
context: RequestContextData object.
|
|
"""
|
|
if not context.config or not context.config.guardrails:
|
|
return
|
|
|
|
try:
|
|
task = asyncio.create_task(
|
|
_preload(context.config.guardrails, context.invariant_authorization)
|
|
)
|
|
asyncio.shield(task)
|
|
except Exception as e:
|
|
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.
|
|
|
|
Use `@on('chunk')`, `@on('start')`, and `@on('end')` decorators
|
|
to register listeners for different events.
|
|
|
|
Listeners can simply process data, or alternatively raise a designated
|
|
YieldException to yield additional values or stop the stream.
|
|
|
|
Example usage:
|
|
|
|
```
|
|
instrumentor = StreamInstrumentor()
|
|
|
|
@instrumentor.on('chunk')
|
|
async def process_chunk(chunk):
|
|
# Process the chunk
|
|
print(f"Processing chunk: {chunk}")
|
|
|
|
if some_condition:
|
|
# Yield an additional value that will be interleaved in the stream
|
|
# Pass `end_of_stream=True` to stop the stream after yielding
|
|
# Pass `end_of_stream=False` to continue the stream after the interleaved value
|
|
raise YieldException("Extra value", end_of_stream=True)
|
|
```
|
|
"""
|
|
|
|
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]:
|
|
"""
|
|
Checks guardrails on the list of messages.
|
|
|
|
Args:
|
|
messages (List[Dict[str, Any]]): List of messages to verify the guardrails against.
|
|
guardrails (str): The guardrails to check against.
|
|
invariant_authorization (str): Value of the
|
|
invariant-authorization header.
|
|
|
|
Returns:
|
|
Dict: Response containing guardrail check results.
|
|
"""
|
|
async with httpx.AsyncClient() as client:
|
|
url = os.getenv("GUADRAILS_API_URL", DEFAULT_API_URL).rstrip("/")
|
|
try:
|
|
result = await client.post(
|
|
f"{url}/api/v1/policy/check",
|
|
json={"messages": messages, "policy": guardrails},
|
|
headers={
|
|
"Authorization": invariant_authorization,
|
|
"Accept": "application/json",
|
|
},
|
|
)
|
|
print(f"Guardrail check response: {result.json()}")
|
|
return result.json()
|
|
except Exception as e:
|
|
print(f"Failed to verify guardrails: {e}")
|
|
return {"error": str(e)}
|